update

Author: shuningjin

src/maxtext/checkpoint_conversion/inspect_checkpoint.py

@@ -23,7 +23,7 @@
 [Mode 1: HF/PyTorch]   
    python src/maxtext/checkpoint_conversion/inspect_checkpoint.py hf --path <local_hf_path> --format <safetensors | pth>
 [Mode 2: MaxText Arch] 
-  python src/maxtext/checkpoint_conversion/inspect_checkpoint.py maxtext --model_name <maxtext_model_name> --scan_layers <True | False>
+  python src/maxtext/checkpoint_conversion/inspect_checkpoint.py maxtext model_name <maxtext_model_name> scan_layers <True | False>
 [Mode 3: Orbax]        
   python src/maxtext/checkpoint_conversion/inspect_checkpoint.py orbax --path <local_orbax_path | gcs_orbax_path>
 """
@@ -43,8 +43,7 @@ def natural_sort_key(s: str):
 def print_structure(data_dict):
   """Utility to print sorted keys and shapes from a flat dictionary."""
   for key in sorted(data_dict.keys(), key=natural_sort_key):
-    shape = data_dict[key]
-    print(f"key: {key} | shape: {shape}")
+    print(f"key: {key} | {data_dict[key]}")
 
 
 # ==============================================================================
@@ -53,17 +52,11 @@ def print_structure(data_dict):
 def inspect_hf(args):
   print(f"\n--- Inspecting {args.format} files in {args.path} ---")
 
-  # Lazy imports
-  try:
-    import torch
-  except ImportError:
-    sys.exit("Error: 'torch' is required for this mode. `pip install torch`")
-
   ckpt_paths = sorted(pathlib.Path(args.path).glob(f"[!.]*.{args.format}"))
   if not ckpt_paths:
     sys.exit(f"No files with extension .{args.format} found in {args.path}")
 
-  chkpt_vars_raw = {}
+  param_dict = {}
 
   if args.format == "safetensors":
     try:
@@ -76,31 +69,34 @@ def inspect_hf(args):
       with safe_open(ckpt_path, framework="pt") as f:
         for k in f.keys():
           # Storing shape directly to save memory, rather than the full tensor
-          chkpt_vars_raw[k] = f.get_tensor(k).shape
+          shape = f.get_tensor(k).shape
+          param_dict[k] = f"shape: {shape}"
 
   elif args.format == "pth":
+    try:
+      import torch
+    except ImportError:
+      sys.exit("Error: 'torch' is required for this mode. `pip install torch`")
+
     for i, ckpt_path in enumerate(ckpt_paths):
       print(f"Loading {ckpt_path.name} ({i+1}/{len(ckpt_paths)})...")
       checkpoint = torch.load(ckpt_path, map_location="cpu")
       # Flatten logic might be needed depending on pth structure,
       # here we assume standard state_dict or handle the wrapper keys manually if needed.
       if isinstance(checkpoint, dict):
         for k, v in checkpoint.items():
-          if hasattr(v, "shape"):
-            chkpt_vars_raw[k] = v.shape
-          else:
-            # Handle nested state dicts or wrapper keys if common in your workflow
-            chkpt_vars_raw[k] = "Non-tensor found"
+          # Handle nested state dicts or wrapper keys if common in your workflow
+          shape = v.shape if hasattr(v, "shape") else "Non-tensor found"
+          param_dict[k] = f"shape: {shape}"
 
   print("\n=== Structure ===")
-  print_structure(chkpt_vars_raw)
+  print_structure(param_dict)
 
 
 # ==============================================================================
 # Mode 2: MaxText Architecture (On-the-fly)
 # ==============================================================================
-def inspect_maxtext(args):
-  print(f"\n--- Inspecting MaxText Architecture: {args.model_name} (Scan: {args.scan_layers}) ---")
+def inspect_maxtext(args, remaining_args):
 
   # Lazy imports
   import jax
@@ -113,17 +109,17 @@ def inspect_maxtext(args):
   Transformer = models.transformer_as_linen
 
   # Setup config
-  argv = [
-      "",  # First arg is usually script name in pyconfig
-      os.path.join(MAXTEXT_PKG_DIR, "configs", "base.yml"),
-      f"model_name={args.model_name}",
-      f"scan_layers={args.scan_layers}",
-      "attention=dot_product",
-      "skip_jax_distributed_system=true",
-  ]
+  argv = (
+      # First arg is usually script name in pyconfig
+      [None, os.path.join(MAXTEXT_PKG_DIR, "configs", "base.yml")]
+      + remaining_args
+      + ["attention=dot_product", "skip_jax_distributed_system=true"]
+  )
+  print(argv)
 
   # Initialize without heavyweight runtime
   config = pyconfig.initialize(argv)
+  print(f"\n--- Inspecting MaxText Architecture: {config.model_name} (Scan: {config.scan_layers}) ---")
   devices_array = maxtext_utils.create_device_mesh(config)
   mesh = jax.sharding.Mesh(devices_array, config.mesh_axes)
   quant = quantizations.configure_quantization(config)
@@ -133,19 +129,23 @@ def inspect_maxtext(args):
   abstract_param = maxtext_utils.get_abstract_param(model, config)
   num_params = max_utils.calculate_num_params_from_pytree(abstract_param)
 
-  print(f"\nTotal Parameters: {num_params} (~{num_params/1e9:.2f}B)")
+  print(f"\nTotal Parameters: {num_params} (~{num_params/1e9:.2f} B)")
   print("\n=== Structure ===")
 
   abstract_params_flat, _ = jax.tree_util.tree_flatten_with_path(abstract_param)
 
-  flat_shapes = {}
+  param_dict = {}
+  # abstract_leaf_value: ShapeDtypeStruct(shape=(128, 58), dtype=float32)
   for path_tuple, abstract_leaf_value in abstract_params_flat:
     key_parts = [k.key for k in path_tuple if hasattr(k, "key")]
     # Construct MaxText style parameter key
-    mt_param_key = "params-" + "-".join(key_parts)
-    flat_shapes[mt_param_key] = abstract_leaf_value.shape
+    param_key = "params-" + "-".join(key_parts)
+    shape = abstract_leaf_value.shape
+    param_dict[param_key] = f"shape: {shape}"
+    dtype = abstract_leaf_value.dtype
+    param_dict[param_key] += f" | dtype: {dtype}"
 
-  print_structure(flat_shapes)
+  print_structure(param_dict)
 
 
 # ==============================================================================
@@ -163,36 +163,38 @@ def inspect_orbax(args):
 
   path = epath.Path(args.path)
 
-  try:
-    # Depending on Orbax version, metadata access might vary slightly.
-    # This aligns with StandardCheckpointer usage.
-    metadata = ocp.StandardCheckpointer().metadata(path)
-    if hasattr(metadata, "item_metadata"):
-      metadata = metadata.item_metadata
-  except Exception as e:
-    sys.exit(f"Error reading Orbax metadata: {e}")
+  # Depending on Orbax version, metadata access might vary slightly.
+  # This aligns with StandardCheckpointer usage.
+  metadata = ocp.StandardCheckpointer().metadata(path)
+  if hasattr(metadata, "item_metadata"):
+    metadata = metadata.item_metadata
 
   # Convert to flat dict
   dictionary = ocp.tree.to_flat_dict(metadata)
 
   # Filter for params only and clean up keys
-  flat_shapes = {}
+  param_dict = {}
   for k, v in dictionary.items():
     # k is a tuple, join it. v is metadata object with .shape
-    key_str = ".".join(k)
-    if key_str.startswith("params"):
-      flat_shapes[key_str] = v.shape
+    param_key = ".".join(k)
+    if not param_key.startswith("params"):
+      continue
+    shape = v.shape
+    param_dict[param_key] = f"shape: {shape}"
+    dtype = v.dtype
+    param_dict[param_key] += f" | dtype: {dtype}"
+    print(v)
 
   print("\n=== Structure ===")
-  print_structure(flat_shapes)
+  print_structure(param_dict)
 
 
 # ==============================================================================
 # Main CLI Driver
 # ==============================================================================
 def main():
   parser = argparse.ArgumentParser(description="Consolidated Model Checkpoint Inspector")
-  subparsers = parser.add_subparsers(dest="mode", required=True, help="Inspection mode")
+  subparsers = parser.add_subparsers(dest="mode", required=True, help="Inspection mode: hf, maxtext, orbax")
 
   # Mode 1: HuggingFace / PyTorch
   parser_hf = subparsers.add_parser("hf", help="Inspect .safetensors or .pth files")
@@ -203,26 +205,17 @@ def main():
 
   # Mode 2: MaxText Architecture
   parser_mt = subparsers.add_parser("maxtext", help="Inspect MaxText theoretical architecture")
-  parser_mt.add_argument("--model_name", type=str, required=True, help="e.g. deepseek3-671b")
-  parser_mt.add_argument(
-      "--scan_layers",
-      type=str,
-      required=False,
-      default="true",
-      choices=["true", "false", "True", "False"],
-      help="Simulate scanned or unscanned structure",
-  )
 
   # Mode 3: Orbax
   parser_orbax = subparsers.add_parser("orbax", help="Inspect saved Orbax checkpoint metadata")
   parser_orbax.add_argument("--path", type=str, required=True, help="Path to checkpoint items (local or GCS)")
 
-  args = parser.parse_args()
+  args, remaining_args = parser.parse_known_args()
 
   if args.mode == "hf":
     inspect_hf(args)
   elif args.mode == "maxtext":
-    inspect_maxtext(args)
+    inspect_maxtext(args, remaining_args)
   elif args.mode == "orbax":
     inspect_orbax(args)