Merge pull request #2585 from AI-Hypercomputer:hengtaoguo-conv

PiperOrigin-RevId: 827661077

Author: Google-ML-Automation

src/MaxText/checkpointing.py

@@ -427,7 +427,9 @@ def _restore_grain_iterator(
   elif expansion_factor_real_data > 1 and process_count_stored == process_count_jax // expansion_factor_real_data:
     # Scaling up to a larger number of hosts.(e.g., 32 files -> 64 processes)
     # In this case, a subset of hosts restore the data iterator.
-    assert not isinstance(data_iterator, list), "when expansion_factor_real_data > 1, the data iterator should not be a list."
+    assert not isinstance(
+        data_iterator, list
+    ), "when expansion_factor_real_data > 1, the data iterator should not be a list."
     grain_restore_args = GrainCheckpointRestore(
         data_iterator.local_iterator, process_index=jax.process_index(), process_count=process_count_stored
     )

src/MaxText/utils/ckpt_conversion/to_maxtext.py

@@ -48,9 +48,11 @@
 
 import numpy as np
 import jax
+import psutil
 from absl import app
 from flax.training import train_state
 from transformers import AutoConfig, AutoModelForCausalLM
+from tqdm import tqdm
 
 from MaxText import checkpointing
 from MaxText import max_logging
@@ -67,6 +69,39 @@
 jax.config.update("jax_platform_name", "cpu")
 
 
+class MemoryMonitorTqdm(tqdm):
+  """Custom tqdm class that displays memory usage in the progress bar."""
+
+  def format_meter(
+      self,
+      n,
+      total,
+      elapsed,
+      postfix=None,
+      **extra_kwargs,
+  ):
+    """Override to add memory usage info to the postfix."""
+    # Get memory info
+    memory = psutil.virtual_memory()
+    used_gb = memory.used / (1024**3)
+    total_gb = memory.total / (1024**3)
+    memory_percent = memory.percent
+
+    # Create memory postfix
+    memory_info = f"RAM: {used_gb:.1f}/{total_gb:.1f}GB ({memory_percent:.1f}%)"
+
+    # Add memory info to postfix
+    if postfix:
+      if isinstance(postfix, dict):
+        postfix["memory"] = memory_info
+      else:
+        postfix = f"{postfix}, {memory_info}"
+    else:
+      postfix = memory_info
+
+    return super().format_meter(n=n, total=total, elapsed=elapsed, postfix=postfix, **extra_kwargs)
+
+
 def _build_multi_axis_stacked_tensor(
     hf_source_keys: List[List[str]], hf_state_dict: Dict[str, np.ndarray], hook_fns: Any
 ) -> np.ndarray:
@@ -229,7 +264,9 @@ def main(argv: Sequence[str]) -> None:
   max_logging.log("Starting weight transformation...")
   final_mt_weights = []
 
-  for path_tuple, abstract_leaf_value in abstract_params_flat:
+  for path_tuple, abstract_leaf_value in MemoryMonitorTqdm(
+      abstract_params_flat, desc="Transforming weights", unit="param", leave=True, dynamic_ncols=True
+  ):
     key_parts = [k.key for k in path_tuple]
     mt_param_key = "params-" + "-".join(key_parts)
     mt_target_shape_final = abstract_leaf_value.shape