Skip to content

Commit df49109

Browse files
Google-ML-AutomationGoogle-ML-Automation
committed
Merge pull request #3591 from CIeNET-International:charlesli/synthetic2
PiperOrigin-RevId: 896755640
2 parents 2b9cebb + 138bab0 commit df49109

2 files changed

Lines changed: 4 additions & 18 deletions

File tree

‎tests/integration/simple_decoder_layer_test.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@ def test_simple_decoder_layer(self):
3535
get_test_config_path(),
3636
"base_output_directory=gs://runner-maxtext-logs",
3737
"run_name=runner_simple_decoder_layer_test",
38-
"dataset_path=gs://maxtext-dataset",
38+
"dataset_type=synthetic",
3939
"decoder_block=simple",
4040
"enable_checkpointing=False",
4141
"enable_goodput_recording=False",
@@ -52,7 +52,7 @@ def test_mlp_decoder_layer(self):
5252
get_test_config_path(),
5353
"base_output_directory=gs://runner-maxtext-logs",
5454
"run_name=runner_simple_decoder_layer_test",
55-
"dataset_path=gs://maxtext-dataset",
55+
"dataset_type=synthetic",
5656
"decoder_block=simple_mlp",
5757
"enable_checkpointing=False",
5858
"enable_goodput_recording=False",

‎tests/integration/train_tests.py‎

Lines changed: 2 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -65,7 +65,6 @@ class TrainTests(unittest.TestCase):
6565
get_test_config_path(),
6666
f"base_output_directory={_base_output_directory}",
6767
"run_name=runner_test",
68-
f"dataset_path={dataset_path}",
6968
"steps=2",
7069
"enable_checkpointing=False",
7170
"enable_goodput_recording=False",
@@ -78,7 +77,6 @@ class TrainTests(unittest.TestCase):
7877
get_test_config_path(),
7978
f"base_output_directory={_base_output_directory}",
8079
"run_name=runner_test",
81-
f"dataset_path={dataset_path}",
8280
"dataset_type=synthetic", # use synthetic dataset_type to decrease training time
8381
"steps=2",
8482
"enable_checkpointing=False",
@@ -93,7 +91,6 @@ class TrainTests(unittest.TestCase):
9391
get_test_config_path(),
9492
f"base_output_directory={_base_output_directory}",
9593
"run_name=runner_test",
96-
f"dataset_path={dataset_path}",
9794
"dataset_type=synthetic", # use synthetic dataset_type to decrease training time
9895
"steps=2",
9996
"ici_tensor_transpose_parallelism=4",
@@ -106,7 +103,6 @@ class TrainTests(unittest.TestCase):
106103
get_test_config_path(),
107104
f"base_output_directory={_base_output_directory}",
108105
"run_name=runner_test",
109-
f"dataset_path={dataset_path}",
110106
"dataset_type=synthetic", # use synthetic dataset_type to decrease training time
111107
"quantization=int8",
112108
"steps=2",
@@ -120,7 +116,6 @@ class TrainTests(unittest.TestCase):
120116
get_test_config_path(),
121117
f"base_output_directory={_base_output_directory}",
122118
"run_name=runner_test",
123-
f"dataset_path={dataset_path}",
124119
"dataset_type=synthetic", # use synthetic dataset_type to decrease training time
125120
"quantization=fp8",
126121
"steps=2",
@@ -134,7 +129,6 @@ class TrainTests(unittest.TestCase):
134129
get_test_config_path(),
135130
f"base_output_directory={_base_output_directory}",
136131
"run_name=runner_test",
137-
f"dataset_path={dataset_path}",
138132
"dataset_type=synthetic", # use synthetic dataset_type to decrease training time
139133
"quantization=nanoo_fp8",
140134
"steps=2",
@@ -148,7 +142,6 @@ class TrainTests(unittest.TestCase):
148142
get_test_config_path(),
149143
f"base_output_directory={_base_output_directory}",
150144
"run_name=runner_test",
151-
f"dataset_path={dataset_path}",
152145
"dataset_type=synthetic", # use synthetic dataset_type to decrease training time
153146
"quantization=te_fp8_delayedscaling",
154147
"steps=2",
@@ -162,7 +155,6 @@ class TrainTests(unittest.TestCase):
162155
get_test_config_path(),
163156
f"base_output_directory={_base_output_directory}",
164157
"run_name=runner_test",
165-
f"dataset_path={dataset_path}",
166158
"dataset_type=synthetic", # use synthetic dataset_type to decrease training time
167159
"quantization=te_fp8_currentscaling",
168160
"steps=2",
@@ -176,7 +168,6 @@ class TrainTests(unittest.TestCase):
176168
get_test_config_path(),
177169
f"base_output_directory={_base_output_directory}",
178170
"run_name=runner_test",
179-
f"dataset_path={dataset_path}",
180171
"dataset_type=synthetic", # use synthetic dataset_type to decrease training time
181172
"quantization=te_mxfp8",
182173
"steps=2",
@@ -190,7 +181,6 @@ class TrainTests(unittest.TestCase):
190181
get_test_config_path(),
191182
f"base_output_directory={_base_output_directory}",
192183
"run_name=runner_test",
193-
f"dataset_path={dataset_path}",
194184
"dataset_type=synthetic", # use synthetic dataset_type to decrease training time
195185
"steps=2",
196186
"enable_checkpointing=False",
@@ -225,7 +215,7 @@ def test_tpu_base(self):
225215
@pytest.mark.integration_test
226216
@pytest.mark.tpu_only
227217
def test_tpu_tokamax(self):
228-
train_main(TrainTests.CONFIGS["base"] + ["use_tokamax_splash=true"])
218+
train_main(TrainTests.CONFIGS["synthetic"] + ["use_tokamax_splash=true"])
229219

230220
@pytest.mark.integration_test
231221
@pytest.mark.gpu_only
@@ -342,7 +332,6 @@ def test_gpu_cudnn_flash_te(self):
342332
get_test_config_path(),
343333
f"base_output_directory={self._base_output_directory}",
344334
"run_name=runner_test",
345-
f"dataset_path={self.dataset_path}",
346335
"dataset_type=synthetic", # use synthetic dataset_type to decrease training time
347336
"steps=2",
348337
"enable_checkpointing=False",
@@ -363,7 +352,6 @@ def test_gpu_context_parallelism(self):
363352
get_test_config_path(),
364353
f"base_output_directory={self._base_output_directory}",
365354
"run_name=runner_test",
366-
f"dataset_path={self.dataset_path}",
367355
"dataset_type=synthetic", # use synthetic dataset_type to decrease training time
368356
"steps=10",
369357
"enable_checkpointing=False",
@@ -490,7 +478,7 @@ def test_gpu_cudnn_flash_jax(self):
490478

491479
@pytest.mark.integration_test
492480
def test_base_model_shardy_false(self):
493-
train_main(TrainTests.CONFIGS["base"] + ["shardy=False"])
481+
train_main(TrainTests.CONFIGS["synthetic"] + ["shardy=False"])
494482

495483
@pytest.mark.integration_test
496484
@pytest.mark.tpu_only
@@ -570,7 +558,6 @@ def test_gpu_packed_attention(self):
570558
get_test_config_path(),
571559
f"base_output_directory={self._base_output_directory}",
572560
"run_name=runner_test",
573-
f"dataset_path={self.dataset_path}",
574561
"dataset_type=synthetic", # use synthetic dataset_type to decrease training time
575562
"steps=10",
576563
"enable_checkpointing=False",
@@ -595,7 +582,6 @@ def test_gpu_ring_attention(self):
595582
get_test_config_path(),
596583
f"base_output_directory={self._base_output_directory}",
597584
"run_name=runner_test",
598-
f"dataset_path={self.dataset_path}",
599585
"dataset_type=synthetic", # use synthetic dataset_type to decrease training time
600586
"steps=10",
601587
"enable_checkpointing=False",

0 commit comments

Comments
 (0)