Merge pull request #3555 from AI-Hypercomputer:ajkv/fault-tolerant-save-top-k

PiperOrigin-RevId: 902881912

Author: Google-ML-Automation

src/maxtext/trainers/post_train/distillation/save_top_k_teacher_logits.py

@@ -18,9 +18,9 @@
 
 Example command: 
 python3 src/maxtext/trainers/post_train/distillation/save_top_k_teacher_logits.py \
-src/maxtext/configs/post_train/distillation.yml \
---top_k=128 \
---gcs_upload_path=gs://my-bucket/teacher_logits/
+  src/maxtext/configs/post_train/distillation.yml \
+  --local_tmp_dir=/tmp/save_logits_dir \
+  --steps_per_file=10
 """
 
 import os
@@ -30,8 +30,10 @@
 import time
 import sys
 import tensorflow as tf
+import re
 
 import jax
+import jax.numpy as jnp
 import functools
 from itertools import islice
 
@@ -52,12 +54,41 @@ def get_top_k_logits(logits: jax.Array, k: int):
   return top_k_values, top_k_indices
 
 
+def get_start_step(config, local_args):
+  """Determines the starting step for the generation process."""
+  if jax.process_index() != 0:
+    return 0
+
+  output_dir = local_args.gcs_upload_path if local_args.gcs_upload_path else local_args.local_tmp_dir
+  if not tf.io.gfile.exists(output_dir):
+    tf.io.gfile.makedirs(output_dir)
+    return 0
+
+  existing_files = tf.io.gfile.glob(os.path.join(output_dir, "teacher_top_k_part_*.array_record"))
+  if not existing_files:
+    return 0
+
+  # Find the highest part number from the filenames
+  max_part_num = max(
+      (int(m.group(1)) for f in existing_files if (m := re.search(r"part_(\d+).array_record", os.path.basename(f)))),
+      default=-1,
+  )
+
+  if max_part_num == -1:
+    return 0
+
+  start_step = max_part_num * local_args.steps_per_file
+  max_logging.log(f"Found existing data, resuming from step {start_step}")
+  return start_step
+
+
 def generate_and_save_data(config, local_args):
-  """Generates top-k logits from the teacher model and saves them locally, optionally uploading to GCS."""
+  """Generates top-k logits from the teacher model and saves them locally, optionally uploading to GCS"""
   k_val = local_args.top_k
   optional_keys = local_args.optional_keys
   gcs_upload_path = local_args.gcs_upload_path
   local_tmp_dir = local_args.local_tmp_dir
+  steps_per_file = local_args.steps_per_file
 
   devices = jax.devices()
   devices_array = maxtext_utils.create_device_mesh(config, devices)
@@ -68,25 +99,43 @@ def generate_and_save_data(config, local_args):
   teacher_model, _ = model_creation_utils.create_nnx_model(config, mesh=mesh)
   train_iter, _ = input_pipeline_interface.create_data_iterator(config, mesh)
 
-  # Setup local tmp directory for Host 0
-  filename = "teacher_top_k_global.array_record"
-  local_output_path = os.path.join(local_tmp_dir, filename)
+  # Determine start_step for resuming
+  start_step = get_start_step(config, local_args)
+  start_step = int(multihost_utils.broadcast_one_to_all(jnp.array(start_step)))
 
   writer = None
+  local_output_path = None
   if jax.process_index() == 0:
     if not os.path.exists(local_tmp_dir):
       os.makedirs(local_tmp_dir)
-    max_logging.log(f"Process 0 writing globally gathered data to local path: {local_output_path}")
-    writer = array_record_module.ArrayRecordWriter(local_output_path, "group_size:1000")
 
   # Sync all hosts before starting the loop
   multihost_utils.sync_global_devices("start_generation_loop")
 
-  max_logging.log(f"Starting Top-K generation loop for {config.steps} steps...")
+  max_logging.log(f"Starting Top-K generation loop for {config.steps - start_step} steps...")
   loop_start = time.time()
 
-  for step, batch in enumerate(islice(train_iter, config.steps)):
+  for step, batch in enumerate(islice(train_iter, start_step, config.steps), start=start_step):
     step_start = time.time()
+
+    # Open a new writer for each file chunk on process 0
+    if jax.process_index() == 0 and step % steps_per_file == 0:
+      if writer:
+        writer.close()
+        if gcs_upload_path:
+          # Upload the previous file
+          gcs_file_path = os.path.join(gcs_upload_path, os.path.basename(local_output_path))
+          max_logging.log(f"Uploading {local_output_path} to {gcs_file_path}")
+          tf.io.gfile.copy(local_output_path, gcs_file_path, overwrite=True)
+          os.remove(local_output_path)
+          max_logging.log("Upload complete.")
+
+      file_index = step // steps_per_file
+      filename = f"teacher_top_k_part_{file_index:05d}.array_record"
+      local_output_path = os.path.join(local_tmp_dir, filename)
+      max_logging.log(f"Process 0 writing to new chunk: {local_output_path}")
+      writer = array_record_module.ArrayRecordWriter(local_output_path, "group_size:1000")
+
     tokens = batch["inputs"]
     logits = teacher_model(
         decoder_input_tokens=tokens,
@@ -122,19 +171,15 @@ def generate_and_save_data(config, local_args):
   multihost_utils.sync_global_devices("loop_finished")
 
   # Finalize writing and handle GCS upload on Host 0
-  if jax.process_index() == 0:
+  if jax.process_index() == 0 and writer:
     writer.close()
     max_logging.log(f"Finished writing to local disk: {local_output_path}")
 
     if gcs_upload_path:
-      gcs_file_path = os.path.join(gcs_upload_path, filename)
-      max_logging.log(f"Flag --gcs_upload_path detected. Uploading to: {gcs_file_path}")
-
-      if not tf.io.gfile.exists(gcs_upload_path):
-        tf.io.gfile.makedirs(gcs_upload_path)
-
-      # Perform the bulk copy to GCS
+      gcs_file_path = os.path.join(gcs_upload_path, os.path.basename(local_output_path))
+      max_logging.log(f"Uploading final chunk to: {gcs_file_path}")
       tf.io.gfile.copy(local_output_path, gcs_file_path, overwrite=True)
+      os.remove(local_output_path)
       max_logging.log("GCS Upload complete.")
 
   # Sync all hosts one last time so worker hosts don't terminate the job
@@ -182,6 +227,12 @@ def main(argv: Sequence[str], local_args):
       default="/tmp",
       help="Local temporary directory to write the ArrayRecord file before optional GCS upload.",
   )
+  parser.add_argument(
+      "--steps_per_file",
+      type=int,
+      default=1000,
+      help="Number of steps to save in each chunk.",
+  )
   local_arg, remaining_args = parser.parse_known_args()
 
   main_wrapper = functools.partial(main, local_args=local_arg)

src/maxtext/trainers/post_train/distillation/verify_saved_logits.py

@@ -16,10 +16,9 @@
 Verification script to check the correctness of saved top-k teacher logits.
 
 Example usage:
-python3 src/maxtext/trainers/post_train/distillation/verify_saved_logits.py \
-  --output_dir=/path/to/your/output \
-  --expected_steps=1000 \
-  --top_k=128
+python3 python3 src/maxtext/trainers/post_train/distillation/verify_saved_logits.py \
+  --output_dir=/tmp/save_logits_dir \
+  --expected_steps=140
 """
 
 import functools
@@ -40,44 +39,58 @@ def verify_array_records(output_dir, expected_steps, expected_k, expected_keys):
   files = tf.io.gfile.glob(file_pattern)
 
   if not files:
-    max_logging.log(f"Error: No ArrayRecord files found matching {file_pattern}")
-    return
+    assert False, f"Error: No ArrayRecord files found matching {file_pattern}"
 
   max_logging.log(f"Found {len(files)} ArrayRecord files. Starting verification...")
 
+  total_records_processed = 0
+  all_keys_verified = set()
+
   for file_path in files:
     max_logging.log(f"Verifying: {file_path}")
     reader = array_record_module.ArrayRecordReader(file_path)
-    num_records = reader.num_records()
+    num_records_in_file = reader.num_records()
+
+    if num_records_in_file == 0:
+      max_logging.log(f"Warning: {file_path} is empty.")
+      continue
 
-    step_count = 0
-    for _ in range(num_records):
+    for record_idx in range(num_records_in_file):
       record = reader.read()
       data = pickle.loads(record)
 
       # Verify all required keys are present
-      for key in ["tokens", "top_k_logits", "top_k_indices"]:
-        assert key in data, f"Missing required key '{key}' at step {step_count} in {file_path}"
+      required_keys = ["tokens", "top_k_logits", "top_k_indices"]
+      for key in required_keys:
+        assert key in data, f"Missing required key '{key}' in record {record_idx} in {file_path}"
 
       # Verify all optional keys are present
       for key in expected_keys:
-        assert key in data, f"Missing optional key '{key}' at step {step_count} in {file_path}"
+        assert key in data, f"Missing optional key '{key}' in record {record_idx} in {file_path}"
 
       # Verify shapes for Top-K outputs
       actual_k_logits = data["top_k_logits"].shape[-1]
       actual_k_indices = data["top_k_indices"].shape[-1]
       assert actual_k_logits == expected_k, f"Expected top_k={expected_k}, got {actual_k_logits} for logits"
       assert actual_k_indices == expected_k, f"Expected top_k={expected_k}, got {actual_k_indices} for indices"
 
-      step_count += 1
+      if not all_keys_verified:
+        all_keys_verified.update(data.keys())
+
+    total_records_processed += num_records_in_file
+    max_logging.log(f"Verified {num_records_in_file} records in {file_path}")
 
-    # Verify the total number of steps processed
-    assert step_count == expected_steps, f"Expected {expected_steps} steps, but found {step_count} in {file_path}."
+  # Verify the total number of steps processed across all files
+  assert (
+      total_records_processed == expected_steps
+  ), f"Expected a total of {expected_steps} steps across all files, but found {total_records_processed}."
 
-    max_logging.log(f"Successfully verified {file_path}")
-    max_logging.log(f"- Total steps: {step_count} (Matches expected)")
-    max_logging.log(f"- Top-K dimension: {expected_k}")
-    max_logging.log(f"- Keys verified: {list(data.keys())}")
+  max_logging.log("-----------------------------------------")
+  max_logging.log("Verification Successful!")
+  max_logging.log(f"- Total files verified: {len(files)}")
+  max_logging.log(f"- Total steps verified: {total_records_processed} (Matches expected)")
+  max_logging.log(f"- Top-K dimension: {expected_k}")
+  max_logging.log(f"- Keys verified in records: {sorted(list(all_keys_verified))}")
 
 
 def main(argv, local_args):