Support v6e (apple#879)

ds-hwang · Dec 10, 2024 · 73625c9 · 73625c9
1 parent 0f03612
commit 73625c9
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/axlearn/cloud/gcp/tpu.py b/axlearn/cloud/gcp/tpu.py
@@ -718,7 +718,7 @@ def infer_tpu_workers(tpu_type: str) -> int:
             tpu_version, tpu_cores = match.groups()
             if tpu_version in {"v3", "v4", "v5p"}:
                 return int(tpu_cores) // 8
-            if tpu_version in {"v5litepod"}:
+            if tpu_version in {"v5litepod", "v6e"}:
                 return int(tpu_cores) // 4
     except Exception as e:  # pylint: disable=broad-except
         logging.error("Failed to parse tpu_type %s: %s", tpu_type, e)

diff --git a/axlearn/common/compiler_options.py b/axlearn/common/compiler_options.py
@@ -166,4 +166,4 @@ def infer_xsc_compiler_options(
     return options
 
 
-_TPU_VERSIONS = ("v3", "v4", "v5litepod", "v5p")
+_TPU_VERSIONS = ("v3", "v4", "v5litepod", "v5p", "v6e")
Original file line number	Diff line number	Diff line change
Expand Up		@@ -166,4 +166,4 @@ def infer_xsc_compiler_options(
		return options


		_TPU_VERSIONS = ("v3", "v4", "v5litepod", "v5p")
		_TPU_VERSIONS = ("v3", "v4", "v5litepod", "v5p", "v6e")