Merge pull request #1677 from AI-Hypercomputer:disable_goodput_logging

PiperOrigin-RevId: 754102205

Author: maxtext authors

MaxText/tests/integration_tests/gradient_accumulation_test.py

@@ -49,6 +49,7 @@ def test_grad_accumulate_same_loss(self):
         r"dataset_path=gs://maxtext-dataset",
         "gradient_clipping_threshold=0",  # Ensures we are testing raw scales of gradients (clipping off)
         "enable_checkpointing=False",
+        "enable_goodput_recording=False",
         "base_emb_dim=256",
         "base_num_decoder_layers=4",
         rf"tokenizer_path={os.path.join(os.path.dirname(PKG_DIR), 'assets', 'tokenizer.llama2')}",

MaxText/tests/integration_tests/standalone_dl_ckpt_test.py

@@ -49,6 +49,7 @@ def test_standalone_dataloader(self):
             "dataset_path=gs://maxtext-dataset",
             "steps=100",
             "enable_checkpointing=false",
+            "enable_goodput_recording=False",
             rf"tokenizer_path={os.path.join(os.path.dirname(PKG_DIR), 'assets', 'tokenizer.llama2')}",
         )
     )  # need to pass relative path to tokenizer
@@ -74,6 +75,7 @@ def test_standalone_checkpointer(self):
             "enable_checkpointing=True",
             "checkpoint_period=50",
             "async_checkpointing=False",
+            "enable_goodput_recording=False",
         )
     )
     # restore at 50 and checkpoint at 100
@@ -93,6 +95,7 @@ def test_standalone_checkpointer(self):
             "enable_checkpointing=True",
             "checkpoint_period=50",
             "async_checkpointing=False",
+            "enable_goodput_recording=False",
         )
     )
 

MaxText/tests/integration_tests/train_tests.py

@@ -45,6 +45,7 @@ class TrainTests(unittest.TestCase):
           r"dataset_path=gs://maxtext-dataset",
           "steps=2",
           "enable_checkpointing=False",
+          "enable_goodput_recording=False",
           "dataset_type=synthetic",
           rf"tokenizer_path={os.path.join(os.path.dirname(PKG_DIR), 'assets', 'tokenizer.llama2')}",
       ],
@@ -56,6 +57,7 @@ class TrainTests(unittest.TestCase):
           r"dataset_path=gs://maxtext-dataset",
           "steps=2",
           "enable_checkpointing=False",
+          "enable_goodput_recording=False",
           "per_device_batch_size=0.25",
           "ici_tensor_parallelism=4",
           rf"tokenizer_path={os.path.join(os.path.dirname(PKG_DIR), 'assets', 'tokenizer.llama2')}",
@@ -68,6 +70,7 @@ class TrainTests(unittest.TestCase):
           r"dataset_path=gs://maxtext-dataset",
           "steps=2",
           "ici_tensor_transpose_parallelism=4",
+          "enable_goodput_recording=False",
           rf"tokenizer_path={os.path.join(os.path.dirname(PKG_DIR), 'assets', 'tokenizer.llama2')}",
       ],
       "int8": [  # tests base config with int8
@@ -79,6 +82,7 @@ class TrainTests(unittest.TestCase):
           "quantization=int8",
           "steps=2",
           "enable_checkpointing=False",
+          "enable_goodput_recording=False",
           rf"tokenizer_path={os.path.join(os.path.dirname(PKG_DIR), 'assets', 'tokenizer.llama2')}",
       ],
       "fp8": [  # tests base config with fp8
@@ -90,6 +94,7 @@ class TrainTests(unittest.TestCase):
           "quantization=fp8",
           "steps=2",
           "enable_checkpointing=False",
+          "enable_goodput_recording=False",
           rf"tokenizer_path={os.path.join(os.path.dirname(PKG_DIR), 'assets', 'tokenizer.llama2')}",
       ],
       "nanoo_fp8": [  # tests base config with nanoo_fp8
@@ -101,6 +106,7 @@ class TrainTests(unittest.TestCase):
           "quantization=nanoo_fp8",
           "steps=2",
           "enable_checkpointing=False",
+          "enable_goodput_recording=False",
           rf"tokenizer_path={os.path.join(os.path.dirname(PKG_DIR), 'assets', 'tokenizer.llama2')}",
       ],
       "dropout": [  # tests base config with dropout
@@ -111,6 +117,7 @@ class TrainTests(unittest.TestCase):
           r"dataset_path=gs://maxtext-dataset",
           "steps=2",
           "enable_checkpointing=False",
+          "enable_goodput_recording=False",
           "max_target_length=128",
           "per_device_batch_size=1",
           "dropout_rate=0.02",
@@ -123,6 +130,7 @@ class TrainTests(unittest.TestCase):
           "run_name=runner_test",
           "steps=2",
           "enable_checkpointing=False",
+          "enable_goodput_recording=False",
           "dataset_type=hf",
           "hf_path=parquet",
           r"hf_train_files=gs://maxtext-dataset/hf/c4/c4-train-00000-of-01637.parquet",
@@ -235,6 +243,7 @@ def test_gpu_context_parallelism(self):
         r"dataset_path=gs://maxtext-dataset",
         "steps=10",
         "enable_checkpointing=False",
+        "enable_goodput_recording=False",
         "attention=cudnn_flash_te",
         "ici_fsdp_parallelism=2",
         "ici_context_parallelism=2",

MaxText/tests/pipeline_parallelism_test.py

@@ -157,6 +157,7 @@ def test_circular_minimum_microbatches_same_output_and_grad(self):
     config = pyconfig.initialize(
         [sys.argv[0], os.path.join(PKG_DIR, "configs", "base.yml")],
         enable_checkpointing=False,
+        enable_goodput_recording=False,
         run_name="circular_minimum_microbatches",
         max_target_length=128,
         base_emb_dim=28,
@@ -173,6 +174,7 @@ def test_circular_extra_microbatches_same_output_and_grad(self):
     config = pyconfig.initialize(
         [sys.argv[0], os.path.join(PKG_DIR, "configs", "base.yml")],
         enable_checkpointing=False,
+        enable_goodput_recording=False,
         run_name="circular_extra_microbatches",
         max_target_length=128,
         base_emb_dim=28,
@@ -189,6 +191,7 @@ def test_circular_ag_once(self):
     config = pyconfig.initialize(
         [sys.argv[0], os.path.join(PKG_DIR, "configs", "base.yml")],
         enable_checkpointing=False,
+        enable_goodput_recording=False,
         run_name="circular_ag_once",
         max_target_length=128,
         base_emb_dim=28,
@@ -239,6 +242,7 @@ def test_full_train_circular(self):
             "dataset_type=synthetic",
             "steps=3",
             "enable_checkpointing=False",
+            "enable_goodput_recording=False",
             "ici_pipeline_parallelism=4",
             "num_layers_per_pipeline_stage=2",
             "num_pipeline_microbatches=8",
@@ -253,6 +257,7 @@ def test_delay_activation_forwarding_same_output_and_grad(self):
     config = pyconfig.initialize(
         [sys.argv[0], os.path.join(PKG_DIR, "configs", "base.yml")],
         enable_checkpointing=False,
+        enable_goodput_recording=False,
         run_name="activation_forwarding",
         max_target_length=128,
         base_emb_dim=28,
@@ -287,6 +292,7 @@ def test_full_train_non_circular(self):
             "dataset_type=synthetic",
             "steps=3",
             "enable_checkpointing=False",
+            "enable_goodput_recording=False",
             "ici_pipeline_parallelism=4",
             "num_layers_per_pipeline_stage=8",
             "num_pipeline_microbatches=8",
@@ -318,6 +324,7 @@ def test_subset_layers(self):
             "dataset_type=synthetic",
             "steps=3",
             "enable_checkpointing=False",
+            "enable_goodput_recording=False",
             "ici_pipeline_parallelism=4",
             "num_layers_per_pipeline_stage=1",
             "num_pipeline_repeats=2",
@@ -351,6 +358,7 @@ def test_full_train_fp8(self):
             "dataset_type=synthetic",
             "steps=3",
             "enable_checkpointing=False",
+            "enable_goodput_recording=False",
             "ici_pipeline_parallelism=4",
             rf"tokenizer_path={os.path.join(os.path.dirname(PKG_DIR), 'assets', 'tokenizer.llama2')}",
             "quantization=fp8",
@@ -382,6 +390,7 @@ def test_full_train_nanoo_fp8(self):
             "dataset_type=synthetic",
             "steps=3",
             "enable_checkpointing=False",
+            "enable_goodput_recording=False",
             "ici_pipeline_parallelism=4",
             rf"tokenizer_path={os.path.join(os.path.dirname(PKG_DIR), 'assets', 'tokenizer.llama2')}",
             "quantization=nanoo_fp8",

MaxText/tests/simple_decoder_layer_test.py

@@ -32,6 +32,7 @@ def test_simple_decoder_layer(self):
             r"dataset_path=gs://maxtext-dataset",
             "decoder_block=simple",
             "enable_checkpointing=False",
+            "enable_goodput_recording=False",
             rf"tokenizer_path={os.path.join(os.path.dirname(PKG_DIR), 'assets', 'tokenizer.llama2')}",
             "steps=3",
         ]
@@ -48,6 +49,7 @@ def test_mlp_decoder_layer(self):
             r"dataset_path=gs://maxtext-dataset",
             "decoder_block=simple_mlp",
             "enable_checkpointing=False",
+            "enable_goodput_recording=False",
             rf"tokenizer_path={os.path.join(os.path.dirname(PKG_DIR), 'assets', 'tokenizer.llama2')}",
             "steps=3",
         ]