Merge pull request #1770 from AI-Hypercomputer:checkpointing_integ_test

PiperOrigin-RevId: 762199934

Author: maxtext authors

MaxText/tests/integration_tests/checkpointing_test.py

@@ -14,46 +14,107 @@
 limitations under the License.
 """
 
-"""Integration tests for test_checkpointing.sh"""
+"""
+Integration tests for test_checkpointing.sh
+
+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 json
+from math import isclose
 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
 
 
-def run_checkpointing(attention_type):
-  """Tests grain checkpoint determinism."""
+def get_checkpointing_command(run_date, hardware, steps, metrics_file, attention_type):
+  model_params = [
+      "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",
+  ]
+  return [
+      None,
+      os.path.join(PKG_DIR, "configs", "base.yml"),
+      f"hardware={hardware}",
+      f"run_name=runner_{run_date}",
+      f"steps={steps}",
+      "max_target_length=128",
+      "per_device_batch_size=1",
+      f"metrics_file={metrics_file}",
+      "checkpoint_period=3",
+      "base_output_directory=gs://runner-maxtext-logs",
+      "dataset_path=/tmp/gcsfuse/",
+      "async_checkpointing=False",
+      f"attention={attention_type}",
+  ] + model_params
 
+
+def check_loss(metrics_file, target):
+  """Asserts over loss values from loaded checkpoint"""
+  metrics_file_saved = "saved_" + metrics_file
+  metrics_file_restored = "restored_" + metrics_file
+
+  with (
+      open(metrics_file_saved, "rt", encoding="utf8") as saved,
+      open(metrics_file_restored, "rt", encoding="utf8") as restored,
+  ):
+    saved_loss = json.loads(saved.readlines()[-1])[target]
+    restored_loss = json.loads(restored.readlines()[0])[target]
+    # Checks that checkpoint restore was successful by comparing loss of last
+    # step in saved checkpoint to loss of first step in restored checkpoint
+    print("saved loss: ", saved_loss)
+    print("restored loss: ", restored_loss)
+    assert isclose(saved_loss, restored_loss, rel_tol=0.1)
+
+
+def run_checkpointing(hardware, attention_type):
+  """Tests grain checkpoint determinism."""
   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_checkpointing.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
-      "False",  # collect_stack_trace
-      "grain",  # dataset_type
-      attention_type,
-      "False",  # async_checkpointing"
+  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(
+      get_checkpointing_command(
+          run_date,
+          hardware=hardware,
+          steps=5,
+          metrics_file="saved_metrics.txt",
+          attention_type=attention_type,
+      )
+      + grain_command
+  )
+
+  train_main(
+      get_checkpointing_command(
+          run_date,
+          hardware=hardware,
+          steps=10,
+          metrics_file="restored_metrics.txt",
+          attention_type=attention_type,
+      )
+      + grain_command
+  )
 
-  subprocess.run(command, check=True, cwd=os.path.dirname(PKG_DIR))
+  check_loss("metrics.txt", "learning/loss")
 
 
 @pytest.mark.integration_test
 @pytest.mark.tpu_only
-@pytest.mark.skipif(TEST_DISABLE_SUBPROCESS, reason=TEST_DISABLE_SUBPROCESS_STR)
 def test_autoselected_attention():
-  run_checkpointing("autoselected")
+  run_checkpointing("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_checkpointing("dot_product")
+  run_checkpointing("gpu", "dot_product")