Merge pull request #3551 from CIeNET-International:charlesli/synthetic_dataset

PiperOrigin-RevId: 895503401

Author: Google-ML-Automation

tests/integration/train_tests.py

@@ -79,6 +79,7 @@ class TrainTests(unittest.TestCase):
           f"base_output_directory={_base_output_directory}",
           "run_name=runner_test",
           f"dataset_path={dataset_path}",
+          "dataset_type=synthetic",  # use synthetic dataset_type to decrease training time
           "steps=2",
           "enable_checkpointing=False",
           "enable_goodput_recording=False",
@@ -93,6 +94,7 @@ class TrainTests(unittest.TestCase):
           f"base_output_directory={_base_output_directory}",
           "run_name=runner_test",
           f"dataset_path={dataset_path}",
+          "dataset_type=synthetic",  # use synthetic dataset_type to decrease training time
           "steps=2",
           "ici_tensor_transpose_parallelism=4",
           "enable_goodput_recording=False",
@@ -105,6 +107,7 @@ class TrainTests(unittest.TestCase):
           f"base_output_directory={_base_output_directory}",
           "run_name=runner_test",
           f"dataset_path={dataset_path}",
+          "dataset_type=synthetic",  # use synthetic dataset_type to decrease training time
           "quantization=int8",
           "steps=2",
           "enable_checkpointing=False",
@@ -118,6 +121,7 @@ class TrainTests(unittest.TestCase):
           f"base_output_directory={_base_output_directory}",
           "run_name=runner_test",
           f"dataset_path={dataset_path}",
+          "dataset_type=synthetic",  # use synthetic dataset_type to decrease training time
           "quantization=fp8",
           "steps=2",
           "enable_checkpointing=False",
@@ -131,6 +135,7 @@ class TrainTests(unittest.TestCase):
           f"base_output_directory={_base_output_directory}",
           "run_name=runner_test",
           f"dataset_path={dataset_path}",
+          "dataset_type=synthetic",  # use synthetic dataset_type to decrease training time
           "quantization=nanoo_fp8",
           "steps=2",
           "enable_checkpointing=False",
@@ -144,6 +149,7 @@ class TrainTests(unittest.TestCase):
           f"base_output_directory={_base_output_directory}",
           "run_name=runner_test",
           f"dataset_path={dataset_path}",
+          "dataset_type=synthetic",  # use synthetic dataset_type to decrease training time
           "quantization=te_fp8_delayedscaling",
           "steps=2",
           "enable_checkpointing=False",
@@ -157,6 +163,7 @@ class TrainTests(unittest.TestCase):
           f"base_output_directory={_base_output_directory}",
           "run_name=runner_test",
           f"dataset_path={dataset_path}",
+          "dataset_type=synthetic",  # use synthetic dataset_type to decrease training time
           "quantization=te_fp8_currentscaling",
           "steps=2",
           "enable_checkpointing=False",
@@ -170,6 +177,7 @@ class TrainTests(unittest.TestCase):
           f"base_output_directory={_base_output_directory}",
           "run_name=runner_test",
           f"dataset_path={dataset_path}",
+          "dataset_type=synthetic",  # use synthetic dataset_type to decrease training time
           "quantization=te_mxfp8",
           "steps=2",
           "enable_checkpointing=False",
@@ -183,6 +191,7 @@ class TrainTests(unittest.TestCase):
           f"base_output_directory={_base_output_directory}",
           "run_name=runner_test",
           f"dataset_path={dataset_path}",
+          "dataset_type=synthetic",  # use synthetic dataset_type to decrease training time
           "steps=2",
           "enable_checkpointing=False",
           "enable_goodput_recording=False",
@@ -334,6 +343,7 @@ def test_gpu_cudnn_flash_te(self):
         f"base_output_directory={self._base_output_directory}",
         "run_name=runner_test",
         f"dataset_path={self.dataset_path}",
+        "dataset_type=synthetic",  # use synthetic dataset_type to decrease training time
         "steps=2",
         "enable_checkpointing=False",
         "enable_goodput_recording=False",
@@ -354,6 +364,7 @@ def test_gpu_context_parallelism(self):
         f"base_output_directory={self._base_output_directory}",
         "run_name=runner_test",
         f"dataset_path={self.dataset_path}",
+        "dataset_type=synthetic",  # use synthetic dataset_type to decrease training time
         "steps=10",
         "enable_checkpointing=False",
         "enable_goodput_recording=False",
@@ -466,6 +477,7 @@ def test_gpu_cudnn_flash_jax(self):
         f"base_output_directory={self._base_output_directory}",
         "run_name=runner_test",
         f"dataset_path={self.dataset_path}",
+        "dataset_type=synthetic",  # use synthetic dataset_type to decrease training time
         "steps=2",
         "enable_checkpointing=False",
         "enable_goodput_recording=False",
@@ -559,6 +571,7 @@ def test_gpu_packed_attention(self):
         f"base_output_directory={self._base_output_directory}",
         "run_name=runner_test",
         f"dataset_path={self.dataset_path}",
+        "dataset_type=synthetic",  # use synthetic dataset_type to decrease training time
         "steps=10",
         "enable_checkpointing=False",
         "enable_goodput_recording=False",
@@ -583,6 +596,7 @@ def test_gpu_ring_attention(self):
         f"base_output_directory={self._base_output_directory}",
         "run_name=runner_test",
         f"dataset_path={self.dataset_path}",
+        "dataset_type=synthetic",  # use synthetic dataset_type to decrease training time
         "steps=10",
         "enable_checkpointing=False",
         "enable_goodput_recording=False",