.

Author: shuningjin

src/MaxText/utils/ckpt_conversion/to_huggingface.py

@@ -79,20 +79,14 @@
 os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=16"
 
 
-def _get_model_mappings(model_name: str, scan_layers: bool, hf_config_dict: dict, inhomogeneous_layer_cycle_interval: int):
+def _get_model_mappings(model_name: str, scan_layers: bool, hf_config_dict: dict, maxtext_config: Any):
   """Retrieves parameter, shape, and hook function mappings for the model.
 
   Args:
     model_name: The name of the model (e.g., "gemma2-2b").
     scan_layers: Boolean indicating if the model was trained with scanned layers.
     hf_config_dict: The Hugging Face model configuration dictionary.
-    inhomogeneous_layer_cycle_interval: For models with complex, non-uniform
-      layer structures (e.g., a repeating pattern of different layer types),
-      this specifies the number of unique layers in one cycle of the pattern.
-      For example, gpt-oss has 'sliding_attention' layer followed by a 
-      'full_attention' layer, this value would be 2. This allows
-      the conversion to correctly map parameters from a scanned MaxText model
-      where these inhomogeneous layers are packed into a single scanned block.
+    maxtext_config:  The maxtext model configuration.
 
   Returns:
     A dictionary containing the parameter mapping, shape mapping, and hook
@@ -105,11 +99,9 @@ def _get_model_mappings(model_name: str, scan_layers: bool, hf_config_dict: dict
     raise ValueError(f"Mappings not found for model: {model_name}. Available PARAM_MAPPING keys: {PARAM_MAPPING.keys()}")
 
   return {
-      "param_mapping": PARAM_MAPPING[model_name](hf_config_dict, scan_layers, inhomogeneous_layer_cycle_interval),
+      "param_mapping": PARAM_MAPPING[model_name](hf_config_dict, maxtext_config, scan_layers),
       "shape_mapping": HF_SHAPE[model_name](hf_config_dict),
-      "hook_fn_mapping": HOOK_FNS[model_name](
-          hf_config_dict, scan_layers, inhomogeneous_layer_cycle_interval, saving_to_hf=True
-      ),
+      "hook_fn_mapping": HOOK_FNS[model_name](hf_config_dict, maxtext_config, scan_layers, saving_to_hf=True),
   }
 
 
@@ -118,7 +110,7 @@ def _check_param_map_keys(param_map_keys, maxtext_state_keys):
 
   Ensures every MaxText checkpoint key (`maxtext_state_keys`) is covered by
   the flattened parameter map. Keys in the map that are not present in the
-  checkpoint (common for multi-variant maps like gemma3 or qwen3) are skipped.
+  checkpoint (common for multi-variant maps like gemma3, qwen3, deepseek) are skipped.
 
   Tuple keys represent N-to-1 mappings (multiple MaxText keys combining into one
   target key) and are only returned if all constituent keys exist in the checkpoint.

src/MaxText/utils/ckpt_conversion/to_maxtext.py

@@ -273,15 +273,11 @@ def main(argv: Sequence[str]) -> None:
   #   f"model.layers.{global_layer_idx}.input_layernorm.weight",
 
   model_key = config.model_name
-  param_map_mt_to_hf = PARAM_MAPPING[model_key](
-      hf_config_obj.to_dict(), config.scan_layers, config.inhomogeneous_layer_cycle_interval
-  )
+  param_map_mt_to_hf = PARAM_MAPPING[model_key](hf_config_obj.to_dict(), config, config.scan_layers)
 
   # Example of Hook FN mapping, to perform reshape:
   # f"params-decoder-layers_{maxtext_layer_idx}-self_attention_global-key-kernel": reshape_kernel,
-  hook_fn_map_mt = HOOK_FNS[model_key](
-      hf_config_obj.to_dict(), config.scan_layers, config.inhomogeneous_layer_cycle_interval, saving_to_hf=False
-  )
+  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.")
 
   # Transform weights

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

@@ -41,7 +41,7 @@
 import jax.numpy as jnp
 
 
-def GEMMA3_MAXTEXT_TO_HF_PARAM_MAPPING(config, scan_layers=False, layer_cycle_interval=1):
+def GEMMA3_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False):
   """Generates a parameter mapping from MaxText to Hugging Face for Gemma3.
 
   This function creates a dictionary that maps the parameter names from a
@@ -143,7 +143,7 @@ def GEMMA3_MAXTEXT_TO_HF_PARAM_MAPPING(config, scan_layers=False, layer_cycle_in
   return mapping
 
 
-def GEMMA3_MAXTEXT_TO_HF_PARAM_HOOK_FN(config, scan_layers=False, layer_cycle_interval=1, saving_to_hf=False):
+def GEMMA3_MAXTEXT_TO_HF_PARAM_HOOK_FN(config, maxtext_config, scan_layers=False, saving_to_hf=False):
   """Hook functions for Gemma3 parameter conversion.
 
   This function provides a dictionary of transformation functions (hooks) for
@@ -298,7 +298,7 @@ def pos_embed(x, target_shape):
   return hooks
 
 
-def GEMMA2_MAXTEXT_TO_HF_PARAM_MAPPING(config, scan_layers=False, layer_cycle_interval=1):
+def GEMMA2_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False):
   """Returns mapping between MaxText and HuggingFace Gemma2 weight paths.
 
   Args:
@@ -431,7 +431,7 @@ def GEMMA2_MAXTEXT_TO_HF_PARAM_MAPPING(config, scan_layers=False, layer_cycle_in
   return mapping
 
 
-def GEMMA2_MAXTEXT_TO_HF_PARAM_HOOK_FN(config, scan_layers=False, layer_cycle_interval=1, saving_to_hf=False):
+def GEMMA2_MAXTEXT_TO_HF_PARAM_HOOK_FN(config, maxtext_config, scan_layers=False, saving_to_hf=False):
   """Creates parameter transformation functions for Gemma2 conversion.
 
   This function generates a mapping of transformation functions that handle the
@@ -596,7 +596,7 @@ def from_hf():
   return mapping
 
 
-def QWEN3_MAXTEXT_TO_HF_PARAM_MAPPING(config, scan_layers=False, layer_cycle_interval=1):
+def QWEN3_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False):
   """Returns mapping from MaxText to HuggingFace Qwen3 weight paths.
 
   This function generates a dictionary that maps parameter names from a MaxText
@@ -729,7 +729,7 @@ def QWEN3_MAXTEXT_TO_HF_PARAM_MAPPING(config, scan_layers=False, layer_cycle_int
   return mapping
 
 
-def QWEN3_MAXTEXT_TO_HF_PARAM_HOOK_FN(config, scan_layers=False, layer_cycle_interval=1, saving_to_hf=False):
+def QWEN3_MAXTEXT_TO_HF_PARAM_HOOK_FN(config, maxtext_config, scan_layers=False, saving_to_hf=False):
   """Creates parameter transformation functions for Qwen3.
 
   This function provides a dictionary of transformation functions (hooks) for
@@ -814,7 +814,7 @@ def reshape_kernel(input_tensor, target_shape):
   return mapping
 
 
-def DEEPSEEK_MAXTEXT_TO_HF_PARAM_MAPPING(config, scan_layers=False, layer_cycle_interval=1):
+def DEEPSEEK_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False):
   """Returns mapping from MaxText to HuggingFace Deepseek weight paths using f-strings."""
   # TODO(shuningjin): add unscan support, b/457820735
   if not scan_layers:
@@ -885,7 +885,7 @@ def DEEPSEEK_MAXTEXT_TO_HF_PARAM_MAPPING(config, scan_layers=False, layer_cycle_
   return mapping
 
 
-def DEEPSEEK_MAXTEXT_TO_HF_PARAM_HOOK_FN(config, scan_layers=False, layer_cycle_interval=1, saving_to_hf=False):
+def DEEPSEEK_MAXTEXT_TO_HF_PARAM_HOOK_FN(config, maxtext_config, scan_layers=False, saving_to_hf=False):
   """Creates parameter transformation functions for Deepseek using f-strings."""
   # TODO(shuningjin): support hf->orbax(scan), b/457820372
   if not saving_to_hf:
@@ -937,14 +937,16 @@ def reshape_kernel(input_tensor, target_shape):
     mapping[key] = reshape_kernel
   return mapping
 
+
 def DEEPSEEK_NNX_TO_VLLM_PARAM_HOOK_FN():
   """Creates parameter transformation functions for Deepseek."""
   return {}
 
-def GPT_OSS_MAXTEXT_TO_HF_PARAM_MAPPING(hf_config, scan_layers=True, layer_cycle_interval=1):
+
+def GPT_OSS_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False):
   """Returns mapping from MaxText gpt-oss to Hugging Face weight paths.
 
-  Handles the inhomogeneous scan block structure (layer_cycle_interval)
+  Handles the inhomogeneous scan block structure (inhomogeneous_layer_cycle_interval)
 
   Handles N-to-1 mapping from maxtext to huggingface
   - (GptOssMlp-wi_0, GptOssMlp-wi_1): mlp.experts.gate_up_proj
@@ -954,7 +956,8 @@ def GPT_OSS_MAXTEXT_TO_HF_PARAM_MAPPING(hf_config, scan_layers=True, layer_cycle
   if not scan_layers:
     raise NotImplementedError("Current gpt-oss mapping only supports scan_layers=True")
 
-  n_layers = hf_config["num_hidden_layers"]
+  n_layers = config["num_hidden_layers"]  # hf config
+  layer_cycle_interval = maxtext_config.inhomogeneous_layer_cycle_interval
 
   # Base mapping for non-layer parameters (targeting standard HF keys)
   mapping = {
@@ -965,7 +968,7 @@ def GPT_OSS_MAXTEXT_TO_HF_PARAM_MAPPING(hf_config, scan_layers=True, layer_cycle
 
   for block_idx in range(layer_cycle_interval):
     # Identify all original HF layer indices that collapse into this block
-    hf_indices = list(range(block_idx, n_layers, layer_cycle_interval))
+    hf_indices = list(range(block_idx, n_layers, maxtext_config.layer_cycle_interval))
     prefix = f"params-decoder-layers-layers_{block_idx}"
 
     # Layer Norms
@@ -1024,10 +1027,10 @@ def GPT_OSS_MAXTEXT_TO_HF_PARAM_MAPPING(hf_config, scan_layers=True, layer_cycle
   return mapping
 
 
-def GPT_OSS_TO_HF_PARAM_HOOK_FN(hf_config, scan_layers=False, layer_cycle_interval=1, saving_to_hf=False):
+def GPT_OSS_TO_HF_PARAM_HOOK_FN(config, maxtext_config, scan_layers=False, saving_to_hf=False):
   """Transformation hooks for gpt-oss parameters.
 
-  Handles the inhomogeneous scan block structure (layer_cycle_interval)
+  Handles the inhomogeneous scan block structure (inhomogeneous_layer_cycle_interval)
 
   Handles N-to-1 mapping from maxtext to huggingface
   - (GptOssMlp-wi_0, GptOssMlp-wi_1): mlp.experts.gate_up_proj
@@ -1085,6 +1088,7 @@ def interleave(input_tensor, target_shape=None):
   }
 
   # Scan over blocks
+  layer_cycle_interval = maxtext_config.inhomogeneous_layer_cycle_interval
   for block_idx in range(layer_cycle_interval):
     prefix = f"params-decoder-layers-layers_{block_idx}"
     # Attention Kernels & Biases
@@ -1103,7 +1107,7 @@ def interleave(input_tensor, target_shape=None):
   return hooks
 
 
-def QWEN3_OMNI_MOE_MAXTEXT_TO_HF_PARAM_MAPPING(config, scan_layers=False, layer_cycle_interval=1):
+def QWEN3_OMNI_MOE_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False):
   """Returns mapping from MaxText to HuggingFace Qwen3-Omni weight paths.
 
   This function combines mappings from different modalities (text, vision, audio, etc.)
@@ -1137,7 +1141,7 @@ def QWEN3_OMNI_MOE_MAXTEXT_TO_HF_PARAM_MAPPING(config, scan_layers=False, layer_
   return mapping
 
 
-def QWEN3_OMNI_MOE_MAXTEXT_TO_HF_PARAM_HOOK_FN(config, scan_layers=False, layer_cycle_interval=1, saving_to_hf=False):
+def QWEN3_OMNI_MOE_MAXTEXT_TO_HF_PARAM_HOOK_FN(config, maxtext_config, scan_layers=False, saving_to_hf=False):
   """Creates parameter transformation functions for Qwen3-Omni.
 
   This function provides a dictionary of transformation functions (hooks) for
@@ -1188,7 +1192,7 @@ def QWEN3_NNX_TO_VLLM_PARAM_HOOK_FN(target_shape=None):
   return {}
 
 
-def LLAMA31_MAXTEXT_TO_HF_PARAM_MAPPING(config, scan_layers=False, layer_cycle_interval=1):
+def LLAMA31_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False):
   """
   Returns a dictionary mapping from MaxText parameter names to
   HuggingFace LLaMA3.1 parameter names.
@@ -1266,7 +1270,7 @@ def LLAMA31_MAXTEXT_TO_HF_PARAM_MAPPING(config, scan_layers=False, layer_cycle_i
   return mapping
 
 
-def LLAMA31_MAXTEXT_TO_HF_PARAM_HOOK_FN(config, scan_layers=False, layer_cycle_interval=1, saving_to_hf=False):
+def LLAMA31_MAXTEXT_TO_HF_PARAM_HOOK_FN(config, maxtext_config, scan_layers=False, saving_to_hf=False):
   """Creates parameter transformation functions for converting between MaxText and
   HuggingFace formats.