checkpoint utility: shard checkpoint, monitor peak

Author: shuningjin

src/MaxText/utils/ckpt_conversion/README.md

@@ -50,7 +50,7 @@ python3 -m MaxText.utils.ckpt_conversion.to_maxtext src/MaxText/configs/base.yml
   * `hf_access_token`: Your Hugging Face token.
   * `base_output_directory`: The path where the converted Orbax checkpoint will be stored; it can be Googld Cloud Storage (GCS) or local. If not set, the default output directory is `Maxtext/tmp`.
   * `--lazy_load_tensors` (optional): If `true`, loads Hugging Face weights on-demand to minimize RAM usage.
-  * `--hf_model_path` (optional): Specifies a local directory containing the model weights. If unspecified, we use the [default Hugging Face repository ID](https://github.com/AI-Hypercomputer/maxtext/blob/2f77e7b5fcc4b580bc2d109525c362f3d9056ec9/src/MaxText/utils/ckpt_conversion/utils/utils.py#L54-L82) (e.g., openai/gpt-oss-20b). This is necessary for locally dequantized models like GPT-OSS or DeepSeek. 
+  * `--hf_model_path` (optional): Specifies a local or remote directory containing the model weights. If unspecified, we use the [default Hugging Face repository ID](https://github.com/AI-Hypercomputer/maxtext/blob/2f77e7b5fcc4b580bc2d109525c362f3d9056ec9/src/MaxText/utils/ckpt_conversion/utils/utils.py#L54-L82) (e.g., openai/gpt-oss-20b). This is necessary for locally dequantized models like GPT-OSS or DeepSeek. 
 
 
 ## MaxText to Hugging Face

src/MaxText/utils/ckpt_conversion/to_huggingface.py

@@ -55,7 +55,6 @@
 import os
 from typing import Sequence
 import time
-from tqdm import tqdm
 
 from transformers import AutoTokenizer, AutoProcessor
 
@@ -77,6 +76,8 @@
     load_orbax_checkpoint,
     detect_and_extract_checkpoint,
     HF_IDS,
+    MemoryMonitorTqdm,
+    print_peak_memory,
 )
 
 os.environ["JAX_PLATFORMS"] = "cpu"
@@ -179,7 +180,7 @@ def main(argv: Sequence[str]) -> None:
   start = time.time()
   processed_params_list = []
 
-  for key in tqdm(filtered_map_keys, total=len(filtered_map_keys)):
+  for key in MemoryMonitorTqdm(filtered_map_keys, total=len(filtered_map_keys), leave=True):
     if isinstance(key, tuple):
       # if key is tuple of param names, weight is list of param weights
       weight = [maxtext_state_dict[subkey] for subkey in key]
@@ -210,6 +211,7 @@ def main(argv: Sequence[str]) -> None:
   max_logging.log(f"✅ MaxText model successfully saved in HuggingFace format at {output_directory}")
   max_logging.log(f"Elapse for save: {(time.time() - start) / 60:.2f} min")
   max_logging.log(f"Overall Elapse: {(time.time() - overall_start) / 60:.2f} min")
+  print_peak_memory()
 
 
 if __name__ == "__main__":

src/MaxText/utils/ckpt_conversion/to_maxtext.py

@@ -28,8 +28,10 @@
   --lazy_load_tensors: (bool) If True, uses an on-demand loading strategy to minimize RAM
              usage during conversion. Recommended if, 2 * model_size (GB) >= system RAM
              Defaults to False.
-  --hf_model_path: (Optional) Specify a local HF path, rather than the default repo `HF_IDS[model_name]`. 
-              Useful for locally dequantized HF model like GPT-OSS or DeepSeek.
+  --hf_model_path: (Optional) Specifies a local or remote directory containing the model weights. 
+      If unspecified, we use the default Hugging Face repository ID 
+      (e.g., openai/gpt-oss-20b; see `HF_IDS[model_name]` in `utils/ckpt_conversion/utils`).
+      This is necessary for locally dequantized models like GPT-OSS or DeepSeek. 
 
 Environment Variables:
   HF_AUTH_TOKEN: (Required) HuggingFace authentication token, needed to
@@ -65,65 +67,41 @@
 from functools import partial
 from typing import Sequence, List, Any, Callable
 import numpy as np
-import jax
-import psutil
-from flax.training import train_state
-import flax.linen as nn
+import absl
+
 from transformers import AutoConfig
-from tqdm import tqdm
 from huggingface_hub import hf_hub_download, list_repo_files
 from safetensors import safe_open
-import absl
-
+import jax
+import flax.linen as nn
 from orbax.checkpoint import type_handlers
-from MaxText import checkpointing
 from MaxText import max_logging
 from MaxText import max_utils
 from MaxText import maxtext_utils
 from MaxText import pyconfig
 from MaxText.common_types import MODEL_MODE_TRAIN
 from MaxText.inference_utils import str2bool
 from MaxText.layers import models, quantizations
-from MaxText.checkpointing import save_checkpoint
+from MaxText.utils.ckpt_scripts.llama_or_mistral_ckpt import save_weights_to_checkpoint
 from MaxText.utils.ckpt_conversion.utils.param_mapping import HOOK_FNS, PARAM_MAPPING
-from MaxText.utils.ckpt_conversion.utils.utils import apply_hook_fns, HF_IDS, print_ram_usage, get_hf_model, validate_and_filter_param_map_keys
+from MaxText.utils.ckpt_conversion.utils.utils import (
+    apply_hook_fns,
+    HF_IDS,
+    get_hf_model,
+    validate_and_filter_param_map_keys,
+    MemoryMonitorTqdm,
+    print_ram_usage,
+    print_peak_memory,
+)
 
-jax.config.update("jax_platform_name", "cpu")
 
 absl.logging.set_verbosity(absl.logging.INFO)  # for max_logging.log
 
 
-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)
+# Used to convert numpy weights to sharded jax arrays across simulated cpu devices
+SIMULATED_CPU_DEVICES_COUNT = 16
+os.environ["JAX_PLATFORMS"] = "cpu"
+os.environ["XLA_FLAGS"] = f"--xla_force_host_platform_device_count={SIMULATED_CPU_DEVICES_COUNT}"
 
 
 class LazyHFLoader:
@@ -657,20 +635,12 @@ def _eager_getter(key):
   hook_fn_map_mt = HOOK_FNS[model_key](hf_config_obj.to_dict(), config, config.scan_layers, saving_to_hf=False)
   max_logging.log("Parameter mappings and hooks obtained.")
 
-  checkpoint_manager = checkpointing.create_orbax_checkpoint_manager(
-      output_directory,
-      enable_checkpointing=True,
-      use_async=False,  # Synchronous saving for simplicity in conversion script
-      save_interval_steps=1,  # Save at step 0
-      use_ocdbt=config.checkpoint_storage_use_ocdbt,
-      use_zarr3=config.checkpoint_storage_use_zarr3,
-  )
-
   maxtext_abstract_dict, abstract_params_treedef = get_maxtext_model_info(config)
 
   # Weight transformation
   max_logging.log("Starting weight transformation...")
   start = time.time()
+  # Stores MaxText weights: numpy.ndarray
   final_mt_weights = [None] * len(maxtext_abstract_dict)
 
   # Preprocess key
@@ -679,7 +649,7 @@ def _eager_getter(key):
   for mt_param_key_or_keys in MemoryMonitorTqdm(
       filtered_map_keys, desc="Transforming weights", unit="param", leave=True, dynamic_ncols=True
   ):
-    if not use_lazy_load and config.scan_layers:
+    if not use_lazy_load:
       max_logging.log(f"maxtext param: {mt_param_key_or_keys}")
 
     hf_source_keys_or_key = param_map_mt_to_hf.get(mt_param_key_or_keys)
@@ -718,31 +688,24 @@ def _eager_getter(key):
   jax_weights = jax.tree_util.tree_unflatten(abstract_params_treedef, final_mt_weights)
   del final_mt_weights, abstract_params_treedef
 
-  # Create TrainState for saving.
-  final_params_for_state = {"params": jax_weights}
-  final_save_state = train_state.TrainState(step=0, apply_fn=None, params=final_params_for_state, tx=None, opt_state={})
-  del final_params_for_state
-
   print_ram_usage("Before saving")
-  start = time.time()
-  if checkpoint_manager is not None:
-    if use_lazy_load:
-      max_logging.log("Starting checkpoint save (loading weights just-in-time)...")
-    else:
-      max_logging.log("Starting checkpoint save...")
-
-    if save_checkpoint(checkpoint_manager, 0, final_save_state):
-      max_logging.log("saved a checkpoint at step 0")
+  if use_lazy_load:
+    max_logging.log("Starting checkpoint save (loading weights just-in-time)...")
+  else:
+    max_logging.log("Starting checkpoint save...")
 
-    # Upon preemption, exit when and only when all ongoing saves are complete.
-    if checkpoint_manager.reached_preemption(0):
-      checkpoint_manager.wait_until_finished()
-      sys.exit()
+  save_weights_to_checkpoint(
+      output_directory,
+      jax_weights,
+      SIMULATED_CPU_DEVICES_COUNT,
+      config.checkpoint_storage_use_ocdbt,
+      config.checkpoint_storage_use_zarr3,
+  )
 
   print_ram_usage("Program Ends")
   max_logging.log(f"Conversion complete. Checkpoint saved to {output_directory}")
-  max_logging.log(f"Elapse for save: {(time.time() - start) / 60:.2f} min")
   max_logging.log(f"Overall Elapse: {(time.time() - overall_start) / 60:.2f} min")
+  print_peak_memory()
 
 
 if __name__ == "__main__":
@@ -757,7 +720,7 @@ def _eager_getter(key):
       default=False,
       help="Whether to use lazy loading of HF tensors.",
   )
-  # if not specified, default to MaxText.utils.ckpt_conversion.utils.utils.HF_IDS[model_name]
+  # If not specified, default to MaxText.utils.ckpt_conversion.utils.utils.HF_IDS[model_name]
   parser.add_argument(
       "--hf_model_path", type=str, required=False, default="", help="local path to hf model, or custom remote hf repo"
   )

src/MaxText/utils/ckpt_conversion/utils/utils.py

@@ -22,6 +22,8 @@
 import json
 from concurrent.futures import ThreadPoolExecutor
 from typing import Any
+from tqdm import tqdm
+import resource
 
 import jax
 from jax.experimental import multihost_utils
@@ -753,6 +755,45 @@ def print_ram_usage(stage=""):
   )
 
 
+def print_peak_memory():
+  # Returns peak usage in Kilobytes on Linux
+  peak_memory_kb = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
+  max_logging.log(f"Peak Memory: {peak_memory_kb / 1024**2:.2f} GB")
+
+
+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 load_orbax_checkpoint(config) -> dict:
   """Loads a full Orbax checkpoint from disk with unsharded arrays.
 
@@ -898,7 +939,9 @@ def get_hf_model(model_id: str, token: str):
   if model_id in ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]:
     from transformers import Qwen3OmniMoeForConditionalGeneration  # pylint: disable=import-outside-toplevel
 
-    hf_model = Qwen3OmniMoeForConditionalGeneration.from_pretrained(model_id, token=token)
+    model_class = Qwen3OmniMoeForConditionalGeneration.from_pretrained
   else:
-    hf_model = AutoModelForCausalLM.from_pretrained(model_id, token=token)
+    model_class = AutoModelForCausalLM
+
+  hf_model = model_class.from_pretrained(model_id, token=token)
   return hf_model

src/MaxText/utils/ckpt_scripts/convert_gpt_oss_ckpt.py

@@ -27,6 +27,7 @@
 import os
 import pathlib
 import absl
+import time
 
 os.environ["JAX_PLATFORMS"] = "cpu"
 
@@ -35,6 +36,7 @@
 import psutil
 from safetensors import safe_open
 from tqdm import tqdm
+from MaxText.utils.ckpt_conversion.utils.utils import MemoryMonitorTqdm, print_peak_memory
 
 from MaxText import max_logging
 from MaxText.inference_utils import str2bool
@@ -77,7 +79,7 @@ def _convert_huggingface_to_jax_weights(
   max_logging.log(f"Loading the base model from {base_model_path}")
   ckpt_paths = sorted(pathlib.Path(base_model_path).glob("[!.]*.safetensors"))
   chkpt_vars = {}
-  for i, ckpt_path in enumerate(ckpt_paths):
+  for i, ckpt_path in tqdm(enumerate(ckpt_paths), total=len(ckpt_paths)):
     max_logging.log(f"Loading checkpoint {i+1} of {len(ckpt_paths)} ...")
 
     with safe_open(ckpt_path, framework="pt", device="cpu") as f:
@@ -141,9 +143,9 @@ def _convert_huggingface_to_jax_weights(
 
   logging.debug("Memory usage: %f GB", mem_info.memory_info().rss / (1024**3))
 
-  # self attention ###############################################
+  # layer weight: self attention ###############################################
   max_logging.log("Processing self attention")
-  for layer_idx in tqdm(range(base_num_decoder_layers), desc="layers", leave=False):
+  for layer_idx in MemoryMonitorTqdm(range(base_num_decoder_layers), desc="layers", leave=True):
     block_layer_idx, block_idx = divmod(layer_idx, layer_cycle_interval)
     stack_shape = (base_num_decoder_layers // layer_cycle_interval,)
     self_attention = jax_weights["decoder"]["layers"][f"layers_{block_idx}"]["GptOssAttention"]
@@ -212,9 +214,9 @@ def _convert_huggingface_to_jax_weights(
 
   logging.debug("Memory usage: %f GB", mem_info.memory_info().rss / (1024**3))
 
-  # layer weight pre and post self attention norm ################
+  # layer weight: pre and post self attention norm ################
   max_logging.log("Processing pre and post self attention norms")
-  for layer_idx in tqdm(range(base_num_decoder_layers), desc="layers", leave=False):
+  for layer_idx in MemoryMonitorTqdm(range(base_num_decoder_layers), desc="layers", leave=True):
     block_layer_idx, block_idx = divmod(layer_idx, layer_cycle_interval)
     stack_shape = (base_num_decoder_layers // layer_cycle_interval,)
     layer_weight = jax_weights["decoder"]["layers"][f"layers_{block_idx}"]
@@ -246,10 +248,10 @@ def _convert_huggingface_to_jax_weights(
 
   logging.debug("Memory usage: %f GB", mem_info.memory_info().rss / (1024**3))
 
-  # layer weights ################################################
-  max_logging.log("Processing layer weights")
+  # layer weight: mlp ################################################
+  max_logging.log("Processing mlp weights")
 
-  for layer_idx in tqdm(range(base_num_decoder_layers), desc="layers", leave=False):
+  for layer_idx in MemoryMonitorTqdm(range(base_num_decoder_layers), desc="layers", leave=True):
     block_layer_idx, block_idx = divmod(layer_idx, layer_cycle_interval)
     stack_shape = (base_num_decoder_layers // layer_cycle_interval,)
     mlp_weight = jax_weights["decoder"]["layers"][f"layers_{block_idx}"]["GptOssMlp"]
@@ -337,17 +339,27 @@ def convert_to_jax_weights(base_model_path: str, model_size: str):
   parser.add_argument("--use-zarr3", type=str2bool, required=False, default=True)
   args = parser.parse_args()
 
+  overall_start = time.time()
+
   if args.model_size not in MODEL_PARAMS_DICT:
     raise NotImplementedError(f"Model '{args.model_size}' is not supported.")
 
   os.environ["XLA_FLAGS"] = f"--xla_force_host_platform_device_count={args.simulated_cpu_devices_count}"
   base_weights_path = args.maxtext_model_path
 
+  # transform
+  start = time.time()
+  weights = convert_to_jax_weights(args.base_model_path, args.model_size)
+  max_logging.log(f"Elapse for transform: {(time.time() - start) / 60:.2f} min")
+
+  # save
   save_weights_to_checkpoint(
       args.maxtext_model_path,
-      convert_to_jax_weights(args.base_model_path, args.model_size),
+      weights,
       args.simulated_cpu_devices_count,
       args.use_ocdbt,
       args.use_zarr3,
   )
   max_logging.log(f"Successfully saved base_weights to {base_weights_path}.")
+  max_logging.log(f"Overall Elapse: {(time.time() - overall_start) / 60:.2f} min")
+  print_peak_memory()