Merge pull request #1778 from AI-Hypercomputer:enable_ckpt_tests

maxtext authors · maxtext authors · commit e67180ea2cfd · 2025-05-29T10:54:42.000-07:00
PiperOrigin-RevId: 764797579
diff --git a/MaxText/tests/globals.py b/MaxText/tests/globals.py
diff --git a/MaxText/tests/integration_tests/checkpoint_compatibility_test.py b/MaxText/tests/integration_tests/checkpoint_compatibility_test.py
@@ -14,44 +14,94 @@
 limitations under the License.
 """
 
-"""Integration tests for test_checkpointing.sh"""
+"""
+Integration tests to check compatibility of checkpoints between different input pipelines.
+
+These tests verify that a checkpoint saved during a training run using one
+input pipeline (e.g., 'grain') can be successfully restored and continued
+by a subsequent training run using a different input pipeline (e.g., 'tfds').
+The tests confirm restoration by checking the starting step of the resumed runs.
+
+Note: Make sure to run
+  `bash setup_gcsfuse.sh DATASET_GCS_BUCKET=gs://maxtext-dataset MOUNT_PATH=/tmp/gcsfuse/`
+before running tests locally.
+"""
 
 from datetime import datetime
-import subprocess
-import os.path
+import json
 import pytest
-from MaxText.globals import PKG_DIR
-from MaxText.tests.globals import TEST_DISABLE_SUBPROCESS, TEST_DISABLE_SUBPROCESS_STR
+from MaxText.train import main as train_main
+from MaxText.tests.integration_tests.checkpointing_test import get_checkpointing_command
+
 
+def check_start_step(metrics_file, start_step_target):
+  with open(metrics_file, "rt", encoding="utf8") as metrics:
+    start_step = json.loads(metrics.readlines()[0])["step"]
+  print(f"Start step is {start_step}, start step target is {start_step_target}")
+  assert start_step == float(start_step_target)
 
-def run_checkpoint_compatibility(attention_type):
+
+def run_checkpoint_compatibility(hardware, attention_type):
   """Tests checkpoint compatibility."""
 
   run_date = datetime.now().strftime("%Y-%m-%d-%H-%M-%S")
-  script_path = os.path.join(os.path.dirname(PKG_DIR), "end_to_end", "test_checkpoint_compatibility.sh")
-  if not os.path.isfile(script_path):
-    raise FileNotFoundError(script_path)
-  command = [
-      "bash",
-      script_path,
-      f"runner_{run_date}",  # run_name
-      "gs://runner-maxtext-logs",  # output_path
-      "gs://maxtext-dataset",  # dataset_path
-      attention_type,
+  grain_command = [
+      "grain_worker_count=0",
+      "grain_train_files=/tmp/gcsfuse/array-record/c4/en/3.0.1/c4-train.array_record*",
   ]
 
-  subprocess.run(command, check=True, cwd=os.path.dirname(PKG_DIR))
+  # Run training using grain input pipeline
+  train_main(
+      get_checkpointing_command(
+          run_date,
+          hardware=hardware,
+          steps=3,
+          metrics_file="run_1_metrics.txt",
+          attention_type=attention_type,
+          dataset_type="grain",
+          dataset_path="/tmp/gcsfuse",
+      )
+      + grain_command
+  )
+
+  # Resume training using tfds input pipeline
+  train_main(
+      get_checkpointing_command(
+          run_date,
+          hardware=hardware,
+          steps=5,
+          metrics_file="run_2_metrics.txt",
+          attention_type=attention_type,
+          dataset_type="tfds",
+          dataset_path="/tmp/gcsfuse",
+      )
+  )
+
+  # Resume training again using grain input pipeline
+  train_main(
+      get_checkpointing_command(
+          run_date,
+          hardware=hardware,
+          steps=7,
+          metrics_file="run_3_metrics.txt",
+          attention_type=attention_type,
+          dataset_type="grain",
+          dataset_path="/tmp/gcsfuse",
+      )
+      + grain_command
+  )
+
+  check_start_step("run_2_metrics.txt", 3.0)
+  check_start_step("run_3_metrics.txt", 5.0)
 
 
 @pytest.mark.integration_test
 @pytest.mark.tpu_only
-@pytest.mark.skipif(TEST_DISABLE_SUBPROCESS, reason=TEST_DISABLE_SUBPROCESS_STR)
 def test_autoselected_attention():
-  run_checkpoint_compatibility("autoselected")
+  run_checkpoint_compatibility("tpu", "autoselected")
 
 
 @pytest.mark.integration_test
 @pytest.mark.gpu_only
-@pytest.mark.skipif(TEST_DISABLE_SUBPROCESS, reason=TEST_DISABLE_SUBPROCESS_STR)
 def test_with_dot_product():
-  run_checkpoint_compatibility("dot_product")
+  run_checkpoint_compatibility("gpu", "dot_product")
diff --git a/MaxText/tests/integration_tests/checkpointing_test.py b/MaxText/tests/integration_tests/checkpointing_test.py
@@ -15,7 +15,11 @@
 """
 
 """
-Integration tests for test_checkpointing.sh
+Integration tests for checkpointing functionality.
+
+These tests verify that a training run saves a checkpoint,
+and then a subsequent training run can correctly restore and
+continue from that saved checkpoint.
 
 Note: Make sure to run
   `bash setup_gcsfuse.sh DATASET_GCS_BUCKET=gs://maxtext-dataset MOUNT_PATH=/tmp/gcsfuse/`
@@ -31,7 +35,7 @@
 from MaxText.train import main as train_main
 
 
-def get_checkpointing_command(run_date, hardware, steps, metrics_file, attention_type):
+def get_checkpointing_command(run_date, hardware, steps, metrics_file, attention_type, dataset_type, dataset_path):
   model_params = [
       "base_emb_dim=384",
       "base_num_query_heads=8",
@@ -51,7 +55,8 @@ def get_checkpointing_command(run_date, hardware, steps, metrics_file, attention
       f"metrics_file={metrics_file}",
       "checkpoint_period=3",
       "base_output_directory=gs://runner-maxtext-logs",
-      "dataset_path=/tmp/gcsfuse/",
+      f"dataset_path={dataset_path}",
+      f"dataset_type={dataset_type}",
       "async_checkpointing=False",
       f"attention={attention_type}",
   ] + model_params
@@ -80,7 +85,6 @@ def run_checkpointing(hardware, attention_type):
   run_date = datetime.now().strftime("%Y-%m-%d-%H-%M-%S")
   grain_command = [
       "grain_worker_count=0",
-      "dataset_type=grain",
       "grain_train_files=/tmp/gcsfuse/array-record/c4/en/3.0.1/c4-train.array_record*",
   ]
   train_main(
@@ -90,6 +94,8 @@ def run_checkpointing(hardware, attention_type):
           steps=5,
           metrics_file="saved_metrics.txt",
           attention_type=attention_type,
+          dataset_type="grain",
+          dataset_path="/tmp/gcsfuse",
       )
       + grain_command
   )
@@ -101,6 +107,8 @@ def run_checkpointing(hardware, attention_type):
           steps=10,
           metrics_file="restored_metrics.txt",
           attention_type=attention_type,
+          dataset_type="grain",
+          dataset_path="/tmp/gcsfuse",
       )
       + grain_command
   )
diff --git a/MaxText/tests/integration_tests/generate_param_only_checkpoint_test.py b/MaxText/tests/integration_tests/generate_param_only_checkpoint_test.py
@@ -14,49 +14,96 @@
 limitations under the License.
 """
 
-"""Integration tests for test_generate_param_only_checkpoint.sh"""
+"""
+Integration tests for generating a decode-only checkpoint from a training checkpoint
+and then running decode with it.
+"""
 from datetime import datetime
-import subprocess
 import os.path
 import pytest
 
 from MaxText.globals import PKG_DIR
-from MaxText.tests.globals import TEST_DISABLE_SUBPROCESS, TEST_DISABLE_SUBPROCESS_STR
+from MaxText.train import main as train_main
+from MaxText.decode import main as decode_main
+from MaxText.generate_param_only_checkpoint import main as generate_param_only_ckpt_main
+from MaxText.tests.integration_tests.checkpointing_test import get_checkpointing_command
 
 
-def run_generate_param_only_checkpoint(attention_type, quantization):
+def run_generate_param_only_checkpoint(hardware, attention_type, quantization):
   """Tests generating a parameter-only checkpoint."""
 
   run_date = datetime.now().strftime("%Y-%m-%d-%H-%M-%S")
-  script_path = os.path.join(os.path.dirname(PKG_DIR), "end_to_end", "test_generate_param_only_checkpoint.sh")
-  if not os.path.isfile(script_path):
-    raise FileNotFoundError(script_path)
-  # fmt: off
-  command = [
-      "bash",
-      script_path,
-      "-r", f"runner_{run_date}",
-      "-o", "gs://runner-maxtext-logs",
-      "-d", "gs://maxtext-dataset",
-      "-i", "4",
-      "-a", attention_type,
-      "-q", quantization,
+  model_params = [
+      f"quantization={quantization}",
+      "base_emb_dim=384",
+      "base_num_query_heads=8",
+      "base_num_kv_heads=8",
+      "base_mlp_dim=192",
+      "base_num_decoder_layers=8",
+      "head_dim=128",
   ]
 
-  subprocess.run(command, check=True, cwd=os.path.dirname(PKG_DIR))
+  train_main(
+      get_checkpointing_command(
+          run_date,
+          hardware=hardware,
+          steps=5,
+          metrics_file="run_metrics.txt",
+          attention_type=attention_type,
+          dataset_type="tfds",
+          dataset_path="gs://maxtext-dataset",
+      )
+      + model_params
+  )
+
+  state_path = f"gs://runner-maxtext-logs/runner_{run_date}/checkpoints/4/items"
+  generate_param_only_ckpt_main(
+      [
+          None,
+          os.path.join(PKG_DIR, "configs", "base.yml"),
+          f"hardware={hardware}",
+          f"run_name=generate_param_{run_date}",
+          "base_output_directory=gs://runner-maxtext-logs",
+          "dataset_path=gs://maxtext-dataset",
+          "async_checkpointing=False",
+          f"attention={attention_type}",
+          f"load_full_state_path={state_path}",
+      ]
+      + model_params
+  )
+
+  decode_ckpt_path = f"gs://runner-maxtext-logs/generate_param_{run_date}/checkpoints/0/items"
+  decode_main(
+      [
+          None,
+          os.path.join(PKG_DIR, "configs", "base.yml"),
+          f"hardware={hardware}",
+          f"run_name=decode_{run_date}",
+          "base_output_directory=gs://runner-maxtext-logs",
+          "dataset_path=gs://maxtext-dataset",
+          f"load_parameters_path={decode_ckpt_path}",
+          f"attention={attention_type}",
+          "max_target_length=128",
+      ]
+      + model_params
+  )
 
 
 @pytest.mark.integration_test
 @pytest.mark.tpu_only
 @pytest.mark.parametrize("quantization", [(""), ("int8")])
-@pytest.mark.skipif(TEST_DISABLE_SUBPROCESS, reason=TEST_DISABLE_SUBPROCESS_STR)
-def test_autoselected_attention(quantization):
-  run_generate_param_only_checkpoint("autoselected", quantization)
+def test_autoselected_attention(quantization, capsys):
+  run_generate_param_only_checkpoint("tpu", "autoselected", quantization)
+  captured = capsys.readouterr()
+  expected_output = "Input `I love to`"
+  assert expected_output in captured.out
 
 
 @pytest.mark.integration_test
 @pytest.mark.gpu_only
 @pytest.mark.parametrize("quantization", [(""), ("int8")])
-@pytest.mark.skipif(TEST_DISABLE_SUBPROCESS, reason=TEST_DISABLE_SUBPROCESS_STR)
-def test_with_dot_product(quantization):
-  run_generate_param_only_checkpoint("dot_product", quantization)
+def test_with_dot_product(quantization, capsys):
+  run_generate_param_only_checkpoint("gpu", "dot_product", quantization)
+  captured = capsys.readouterr()
+  expected_output = "Input `I love to`"
+  assert expected_output in captured.out
