Adds mesh rule for a3-megagpu-8g.

apple · Jan 21, 2025 · 13e8c44 · 13e8c44
1 parent 4678740
commit 13e8c44
Show file tree

Hide file tree

Showing 17 changed files with 18 additions and 18 deletions.
diff --git a/...riments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-7B-v1-flash-single-host.txt b/...riments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-7B-v1-flash-single-host.txt
@@ -195,7 +195,7 @@ mesh_rules[5][1][2]: 1
 mesh_rules[5][1][3]: 8
 mesh_rules[5][1][4]: 1
 mesh_rules[5][1][5]: 1
-mesh_rules[6][0]: 'gpu-(p5.48xlarge|p4de.24xlarge|a3-highgpu-8g)-(256|512|1024)'
+mesh_rules[6][0]: 'gpu-(p5.48xlarge|p4de.24xlarge|a3-highgpu-8g|a3-megagpu-8g)-(256|512|1024)'
 mesh_rules[6][1][0]: 1
 mesh_rules[6][1][1]: -1
 mesh_rules[6][1][2]: 1

diff --git a/axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-7B-v1-flash.txt b/axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-7B-v1-flash.txt
@@ -195,7 +195,7 @@ mesh_rules[5][1][2]: 1
 mesh_rules[5][1][3]: 8
 mesh_rules[5][1][4]: 1
 mesh_rules[5][1][5]: 1
-mesh_rules[6][0]: 'gpu-(p5.48xlarge|p4de.24xlarge|a3-highgpu-8g)-(256|512|1024)'
+mesh_rules[6][0]: 'gpu-(p5.48xlarge|p4de.24xlarge|a3-highgpu-8g|a3-megagpu-8g)-(256|512|1024)'
 mesh_rules[6][1][0]: 1
 mesh_rules[6][1][1]: -1
 mesh_rules[6][1][2]: 1

diff --git a/...n/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-7B-v1-single-host.txt b/...n/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-7B-v1-single-host.txt
@@ -195,7 +195,7 @@ mesh_rules[5][1][2]: 1
 mesh_rules[5][1][3]: 8
 mesh_rules[5][1][4]: 1
 mesh_rules[5][1][5]: 1
-mesh_rules[6][0]: 'gpu-(p5.48xlarge|p4de.24xlarge|a3-highgpu-8g)-(256|512|1024)'
+mesh_rules[6][0]: 'gpu-(p5.48xlarge|p4de.24xlarge|a3-highgpu-8g|a3-megagpu-8g)-(256|512|1024)'
 mesh_rules[6][1][0]: 1
 mesh_rules[6][1][1]: -1
 mesh_rules[6][1][2]: 1

diff --git a/axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-7B-v1.txt b/axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-7B-v1.txt
@@ -195,7 +195,7 @@ mesh_rules[5][1][2]: 1
 mesh_rules[5][1][3]: 8
 mesh_rules[5][1][4]: 1
 mesh_rules[5][1][5]: 1
-mesh_rules[6][0]: 'gpu-(p5.48xlarge|p4de.24xlarge|a3-highgpu-8g)-(256|512|1024)'
+mesh_rules[6][0]: 'gpu-(p5.48xlarge|p4de.24xlarge|a3-highgpu-8g|a3-megagpu-8g)-(256|512|1024)'
 mesh_rules[6][1][0]: 1
 mesh_rules[6][1][1]: -1
 mesh_rules[6][1][2]: 1

diff --git a/...riments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-7B-v2-flash-single-host.txt b/...riments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-7B-v2-flash-single-host.txt
@@ -195,7 +195,7 @@ mesh_rules[5][1][2]: 1
 mesh_rules[5][1][3]: 8
 mesh_rules[5][1][4]: 1
 mesh_rules[5][1][5]: 1
-mesh_rules[6][0]: 'gpu-(p5.48xlarge|p4de.24xlarge|a3-highgpu-8g)-(256|512|1024)'
+mesh_rules[6][0]: 'gpu-(p5.48xlarge|p4de.24xlarge|a3-highgpu-8g|a3-megagpu-8g)-(256|512|1024)'
 mesh_rules[6][1][0]: 1
 mesh_rules[6][1][1]: -1
 mesh_rules[6][1][2]: 1

diff --git a/axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-7B-v2-flash.txt b/axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-7B-v2-flash.txt
@@ -195,7 +195,7 @@ mesh_rules[5][1][2]: 1
 mesh_rules[5][1][3]: 8
 mesh_rules[5][1][4]: 1
 mesh_rules[5][1][5]: 1
-mesh_rules[6][0]: 'gpu-(p5.48xlarge|p4de.24xlarge|a3-highgpu-8g)-(256|512|1024)'
+mesh_rules[6][0]: 'gpu-(p5.48xlarge|p4de.24xlarge|a3-highgpu-8g|a3-megagpu-8g)-(256|512|1024)'
 mesh_rules[6][1][0]: 1
 mesh_rules[6][1][1]: -1
 mesh_rules[6][1][2]: 1

diff --git a/...n/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-7B-v2-single-host.txt b/...n/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-7B-v2-single-host.txt
@@ -195,7 +195,7 @@ mesh_rules[5][1][2]: 1
 mesh_rules[5][1][3]: 8
 mesh_rules[5][1][4]: 1
 mesh_rules[5][1][5]: 1
-mesh_rules[6][0]: 'gpu-(p5.48xlarge|p4de.24xlarge|a3-highgpu-8g)-(256|512|1024)'
+mesh_rules[6][0]: 'gpu-(p5.48xlarge|p4de.24xlarge|a3-highgpu-8g|a3-megagpu-8g)-(256|512|1024)'
 mesh_rules[6][1][0]: 1
 mesh_rules[6][1][1]: -1
 mesh_rules[6][1][2]: 1

diff --git a/axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-7B-v2.txt b/axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-7B-v2.txt
@@ -195,7 +195,7 @@ mesh_rules[5][1][2]: 1
 mesh_rules[5][1][3]: 8
 mesh_rules[5][1][4]: 1
 mesh_rules[5][1][5]: 1
-mesh_rules[6][0]: 'gpu-(p5.48xlarge|p4de.24xlarge|a3-highgpu-8g)-(256|512|1024)'
+mesh_rules[6][0]: 'gpu-(p5.48xlarge|p4de.24xlarge|a3-highgpu-8g|a3-megagpu-8g)-(256|512|1024)'
 mesh_rules[6][1][0]: 1
 mesh_rules[6][1][1]: -1
 mesh_rules[6][1][2]: 1

diff --git a/...riments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-7B-v3-flash-single-host.txt b/...riments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-7B-v3-flash-single-host.txt
@@ -195,7 +195,7 @@ mesh_rules[5][1][2]: 1
 mesh_rules[5][1][3]: 8
 mesh_rules[5][1][4]: 1
 mesh_rules[5][1][5]: 1
-mesh_rules[6][0]: 'gpu-(p5.48xlarge|p4de.24xlarge|a3-highgpu-8g)-(256|512|1024)'
+mesh_rules[6][0]: 'gpu-(p5.48xlarge|p4de.24xlarge|a3-highgpu-8g|a3-megagpu-8g)-(256|512|1024)'
 mesh_rules[6][1][0]: 1
 mesh_rules[6][1][1]: -1
 mesh_rules[6][1][2]: 1

diff --git a/axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-7B-v3-flash.txt b/axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-7B-v3-flash.txt
@@ -195,7 +195,7 @@ mesh_rules[5][1][2]: 1
 mesh_rules[5][1][3]: 8
 mesh_rules[5][1][4]: 1
 mesh_rules[5][1][5]: 1
-mesh_rules[6][0]: 'gpu-(p5.48xlarge|p4de.24xlarge|a3-highgpu-8g)-(256|512|1024)'
+mesh_rules[6][0]: 'gpu-(p5.48xlarge|p4de.24xlarge|a3-highgpu-8g|a3-megagpu-8g)-(256|512|1024)'
 mesh_rules[6][1][0]: 1
 mesh_rules[6][1][1]: -1
 mesh_rules[6][1][2]: 1

diff --git a/...n/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-7B-v3-single-host.txt b/...n/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-7B-v3-single-host.txt
@@ -195,7 +195,7 @@ mesh_rules[5][1][2]: 1
 mesh_rules[5][1][3]: 8
 mesh_rules[5][1][4]: 1
 mesh_rules[5][1][5]: 1
-mesh_rules[6][0]: 'gpu-(p5.48xlarge|p4de.24xlarge|a3-highgpu-8g)-(256|512|1024)'
+mesh_rules[6][0]: 'gpu-(p5.48xlarge|p4de.24xlarge|a3-highgpu-8g|a3-megagpu-8g)-(256|512|1024)'
 mesh_rules[6][1][0]: 1
 mesh_rules[6][1][1]: -1
 mesh_rules[6][1][2]: 1

diff --git a/axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-7B-v3.txt b/axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-7B-v3.txt
@@ -195,7 +195,7 @@ mesh_rules[5][1][2]: 1
 mesh_rules[5][1][3]: 8
 mesh_rules[5][1][4]: 1
 mesh_rules[5][1][5]: 1
-mesh_rules[6][0]: 'gpu-(p5.48xlarge|p4de.24xlarge|a3-highgpu-8g)-(256|512|1024)'
+mesh_rules[6][0]: 'gpu-(p5.48xlarge|p4de.24xlarge|a3-highgpu-8g|a3-megagpu-8g)-(256|512|1024)'
 mesh_rules[6][1][0]: 1
 mesh_rules[6][1][1]: -1
 mesh_rules[6][1][2]: 1

diff --git a/...estdata/axlearn.experiments.text.gpt.c4_trainer/fuji-8B-v3-tiktoken-flash-single-host.txt b/...estdata/axlearn.experiments.text.gpt.c4_trainer/fuji-8B-v3-tiktoken-flash-single-host.txt
@@ -179,7 +179,7 @@ mesh_rules[4][1][2]: 1
 mesh_rules[4][1][3]: 8
 mesh_rules[4][1][4]: 1
 mesh_rules[4][1][5]: 1
-mesh_rules[5][0]: 'gpu-(p5.48xlarge|p4de.24xlarge|a3-highgpu-8g)-(256|512|1024)'
+mesh_rules[5][0]: 'gpu-(p5.48xlarge|p4de.24xlarge|a3-highgpu-8g|a3-megagpu-8g)-(256|512|1024)'
 mesh_rules[5][1][0]: 1
 mesh_rules[5][1][1]: -1
 mesh_rules[5][1][2]: 1

diff --git a/...xperiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-8B-v3-tiktoken-flash.txt b/...xperiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-8B-v3-tiktoken-flash.txt
@@ -179,7 +179,7 @@ mesh_rules[4][1][2]: 1
 mesh_rules[4][1][3]: 8
 mesh_rules[4][1][4]: 1
 mesh_rules[4][1][5]: 1
-mesh_rules[5][0]: 'gpu-(p5.48xlarge|p4de.24xlarge|a3-highgpu-8g)-(256|512|1024)'
+mesh_rules[5][0]: 'gpu-(p5.48xlarge|p4de.24xlarge|a3-highgpu-8g|a3-megagpu-8g)-(256|512|1024)'
 mesh_rules[5][1][0]: 1
 mesh_rules[5][1][1]: -1
 mesh_rules[5][1][2]: 1

diff --git a/...ents/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-8B-v3-tiktoken-single-host.txt b/...ents/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-8B-v3-tiktoken-single-host.txt
@@ -179,7 +179,7 @@ mesh_rules[4][1][2]: 1
 mesh_rules[4][1][3]: 8
 mesh_rules[4][1][4]: 1
 mesh_rules[4][1][5]: 1
-mesh_rules[5][0]: 'gpu-(p5.48xlarge|p4de.24xlarge|a3-highgpu-8g)-(256|512|1024)'
+mesh_rules[5][0]: 'gpu-(p5.48xlarge|p4de.24xlarge|a3-highgpu-8g|a3-megagpu-8g)-(256|512|1024)'
 mesh_rules[5][1][0]: 1
 mesh_rules[5][1][1]: -1
 mesh_rules[5][1][2]: 1

diff --git a/axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-8B-v3-tiktoken.txt b/axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-8B-v3-tiktoken.txt
@@ -179,7 +179,7 @@ mesh_rules[4][1][2]: 1
 mesh_rules[4][1][3]: 8
 mesh_rules[4][1][4]: 1
 mesh_rules[4][1][5]: 1
-mesh_rules[5][0]: 'gpu-(p5.48xlarge|p4de.24xlarge|a3-highgpu-8g)-(256|512|1024)'
+mesh_rules[5][0]: 'gpu-(p5.48xlarge|p4de.24xlarge|a3-highgpu-8g|a3-megagpu-8g)-(256|512|1024)'
 mesh_rules[5][1][0]: 1
 mesh_rules[5][1][1]: -1
 mesh_rules[5][1][2]: 1

diff --git a/axlearn/experiments/text/gpt/fuji.py b/axlearn/experiments/text/gpt/fuji.py
@@ -332,7 +332,7 @@ def get_trainer_kwargs(
                 # v2 on gpu-p5.48xlarge-256, step time: 1.78s/step, MFU 39%.
                 # TODO(kelvin-zou): need to match 1.5s/step perf on TransformerEngine.
                 (
-                    "gpu-(p5.48xlarge|p4de.24xlarge|a3-highgpu-8g)-(256|512|1024)",
+                    "gpu-(p5.48xlarge|p4de.24xlarge|a3-highgpu-8g|a3-megagpu-8g)-(256|512|1024)",
                     mesh_shape_from_axes(data=-1, fsdp=8),
                 ),
             ),
@@ -412,7 +412,7 @@ def get_trainer_kwargs(
                 ),
                 ("tpu-v5p-.*", mesh_shape_from_axes(data=-1, fsdp=8)),
                 (
-                    "gpu-(p5.48xlarge|p4de.24xlarge|a3-highgpu-8g)-(256|512|1024)",
+                    "gpu-(p5.48xlarge|p4de.24xlarge|a3-highgpu-8g|a3-megagpu-8g)-(256|512|1024)",
                     mesh_shape_from_axes(data=-1, fsdp=8),
                 ),
             ),