[Qwen3.5-397B MLPerf] Consolidated top-ROI trainer & GDN/MoE optimizations on atwigg/mlperf

docs/guides/optimization/custom_mesh_and_rule.md

@@ -51,6 +51,12 @@ This rule utilizes the `data`, `stage`, `fsdp`, and `expert` axes. Its defining
 
 Similar in philosophy to `ep-as-cp.yml`, this configuration explicitly includes the `context` axis in the mesh layout alongside `data`, `stage`, `fsdp`, and `expert`. While context sharding is mapped to the `context` axis globally, within MoE components, this `context` axis dynamically shifts to perform expert parallelism instead of FSDP. This custom rule supports using CP and EP together.
 
+### `cp-as-ep-dense-no-cp.yml` and `cp-as-ep-dense-fsdp.yml`
+
+Variants of `cp-as-ep.yml` that take the dense weights (`embed`, `embed_attn`, `q_lora`, `kv_lora`: attention and shared-expert projections, the MoE router and the shared-expert gate) off the `context` axis. Under `cp-as-ep` they are sharded over `fsdp` x `context` x `expert`, so every all-gather of them spans all devices and also synchronizes the context ranks, which do unequal attention work under causal context parallelism. `cp-as-ep-dense-no-cp.yml` shards them over `fsdp` x `expert` and `cp-as-ep-dense-fsdp.yml` over `fsdp` only, so those all-gathers run in independent groups of `fsdp` x `expert` (or `fsdp`) devices, at the cost of replicating the weights and their gradients `context`-way (`context` x `expert`-way).
+
+Both rules also map `embed_grad` to the token axes. Left to GSPMD, the gradient of a weight replicated over `context` (a partial sum over every token axis) is all-reduced at full size over all devices and then sliced. With `embed_grad` mapped, `linears.route_dense_weight_grad` reduce-scatters it over all devices instead, as under `cp-as-ep`, and then all-gathers the shard over the replicated axes. The GatedDeltaNet projections (`embed_gdn`), activations, the vocabulary weights and the routed experts are sharded exactly as in `cp-as-ep.yml`, and `RoutedMoE` keeps `("context", "expert")` as its expert-parallel axes.
+
 ### `ep-as-dp.yml`
 
 Different with the rule in `base.yml`, this rule configures expert physical axis to function as data parallelism rather than FSDP. This removes the constraint where FSDPxEP is limited by specific model dimensions, particularly for small tensors such as attention projections. Ultimately, this change benefits large-scale training.

src/maxtext/checkpoint_conversion/standalone_scripts/convert_deepseek_family_ckpt.py

@@ -601,8 +601,6 @@ def _normalize(raw_key):
         moe["shared_experts"]["wo"]["kernel"][layer_idx, ...] = shared_wo
 
       # re-order
-      if q_lora_rank != 0:
-        moe["MoeBlock_0"]["gate"]["bias"] = np.transpose(moe["MoeBlock_0"]["gate"]["bias"], axes=(1, 0))
       moe["MoeBlock_0"]["gate"]["kernel"] = np.transpose(moe["MoeBlock_0"]["gate"]["kernel"], axes=(1, 0, 2))
       moe["shared_experts"]["wi_0"]["kernel"] = np.transpose(moe["shared_experts"]["wi_0"]["kernel"], axes=(1, 0, 2))
       moe["shared_experts"]["wi_1"]["kernel"] = np.transpose(moe["shared_experts"]["wi_1"]["kernel"], axes=(1, 0, 2))

src/maxtext/checkpoint_conversion/to_huggingface.py

@@ -238,6 +238,12 @@ def _validate_or_update_architecture(hf_config, max_config, override: bool):
       ("vocab_size", "vocab_size"),
       ("global_head_dim", "global_head_dim"),
       ("num_global_key_value_heads", "global_num_kv_heads"),
+      ("hc_mult", "mhc_expansion_rate"),
+      ("num_hash_layers", "first_num_hash_layers"),
+      ("index_n_heads", "indexer_n_heads"),
+      ("index_head_dim", "indexer_head_dim"),
+      ("o_lora_rank", "o_lora_rank"),
+      ("o_groups", "o_groups"),
   ]
 
   if max_config.attention_type == "mla":

src/maxtext/checkpoint_conversion/to_maxtext.py

@@ -59,20 +59,18 @@
 from typing import Any, Callable, List, Sequence
 import absl
 import ml_dtypes
-import flax.linen as nn
+from flax import nnx
+from flax.nnx import variablelib
 from huggingface_hub import hf_hub_download, list_repo_files
 import jax
 from maxtext.configs import pyconfig
 from maxtext.configs.types import DType
-from maxtext.common.common_types import MODEL_MODE_TRAIN
 from maxtext.checkpoint_conversion.utils.hf_model_configs import HF_MODEL_CONFIGS
 from maxtext.checkpoint_conversion.utils.param_mapping import HOOK_FNS, PARAM_MAPPING
 from maxtext.checkpoint_conversion.utils.tensor_handling import apply_hook_fns, nesting_depth, slice_shape, stacked_axes
 from maxtext.checkpoint_conversion.utils.utils import MemoryMonitorTqdm, load_hf_dict_from_transformers, load_hf_dict_from_safetensors, param_key_parts_from_path, print_peak_memory, print_ram_usage, save_weights_to_checkpoint, validate_and_filter_param_map_keys
 from maxtext.inference.inference_utils import str2bool
-from maxtext.layers import quantizations
-from maxtext.models import models
-from maxtext.utils import max_logging, max_utils, maxtext_utils
+from maxtext.utils import max_logging, max_utils, maxtext_utils, model_creation_utils
 from maxtext.utils.globals import HF_IDS
 import numpy as np
 from orbax.checkpoint import type_handlers
@@ -389,31 +387,30 @@ def get_maxtext_model_info(config):
   mesh = jax.sharding.Mesh(devices_array, config.mesh_axes)
 
   max_logging.log("Initializing MaxText abstract model...")
-  quant = quantizations.configure_quantization(config)
-  maxtext_model_flax = models.transformer_as_linen(config, mesh, quant=quant, model_mode=MODEL_MODE_TRAIN)
-
-  # Get abstract model structure (name, shape) without materializing the weights to save memory.
-  # Extract the 'params' collection from the abstract model state. This focuses checkpoint
-  # conversion on trainable model parameters; variables outside the 'params' collection
-  # (such as non-trainable state or optimizer buffers) are not included.
-  abstract_params_tree = maxtext_utils.get_abstract_param(maxtext_model_flax, config)["params"]
-
-  abstract_params_flat, abstract_params_treedef = jax.tree_util.tree_flatten_with_path(
-      abstract_params_tree,
-      is_leaf=lambda x: isinstance(x, nn.LogicallyPartitioned),
-  )
+  # Get abstract model structure (name, shape) without materializing the weights to save memory
+  _, abstract_model = model_creation_utils.create_nnx_abstract_model(config, mesh)
+
+  # Lay the variables out the way the checkpoint stores them: one collection per variable type
+  # (e.g. 'params', 'Tid2EidVar'), then the module path with Linen-style names ("layers_0").
+  abstract_params_tree = {}
+  for path, variable in nnx.to_flat_state(nnx.state(abstract_model)):
+    if isinstance(variable, nnx.RngState):
+      continue
+    node = abstract_params_tree.setdefault(variablelib.variable_name_from_type(type(variable), allow_register=True), {})
+    *parents, name = param_key_parts_from_path(path)
+    for part in parents:
+      node = node.setdefault(part, {})
+    node[name] = variable.get_value()
+
+  abstract_params_flat, abstract_params_treedef = jax.tree_util.tree_flatten_with_path(abstract_params_tree)
 
   max_logging.log("MaxText abstract model and state initialized.")
 
   # preprocess state
   maxtext_abstract_dict = {}
   for mt_target_idx, (path_tuple, abstract_leaf_value) in enumerate(abstract_params_flat):
-    mt_param_key = "params-" + "-".join(param_key_parts_from_path(path_tuple))
-    if isinstance(abstract_leaf_value, nn.LogicallyPartitioned):
-      mt_target_shape = abstract_leaf_value.value.shape
-    else:
-      mt_target_shape = abstract_leaf_value.shape
-    maxtext_abstract_dict[mt_param_key] = (mt_target_idx, mt_target_shape)
+    mt_param_key = "-".join(param_key_parts_from_path(path_tuple))
+    maxtext_abstract_dict[mt_param_key] = (mt_target_idx, abstract_leaf_value.shape)
 
   return maxtext_abstract_dict, abstract_params_treedef
 
@@ -481,6 +478,7 @@ def _build_single_axis_stacked_tensor(
     hook_fns: Any,
     target_shape: tuple,
     config,
+    mt_key: str = "",
 ) -> np.ndarray:
   """Builds a MaxText tensor by stacking HF weights along a single axis.
 
@@ -499,12 +497,16 @@ def _build_single_axis_stacked_tensor(
   """
   tensors_to_stack = []
 
+  axis_to_stack = 0
   if config.scan_layers:
-    # If the target tensor rank exceeds param_scan_axis (e.g., multidimensional weights or 2D block scales),
-    # stack along param_scan_axis (typically axis 1 in scanned layers). For 1D tensors (e.g., per-layer scalar
-    # quantization scales like Llama 3.1 FP8 `kernel_scale` with shape (num_layers,)), stack along axis 0
-    # because axis 1 does not exist on 1D tensors and stacking scalars along axis 0 produces the 1D vector.
-    axis_to_stack = config.param_scan_axis if len(target_shape) > config.param_scan_axis else 0
+    is_var_collection = "MoEBiasVar" in mt_key or "Tid2EidVar" in mt_key
+    is_unscanned_moe = "MoeBlock" in mt_key and "scanned_blocks" not in mt_key
+    if not (is_var_collection or is_unscanned_moe):
+      # If the target tensor rank exceeds param_scan_axis (e.g., multidimensional weights or 2D block scales),
+      # stack along param_scan_axis (typically axis 1 in scanned layers). For 1D tensors (e.g., per-layer scalar
+      # quantization scales like Llama 3.1 FP8 `kernel_scale` with shape (num_layers,)), stack along axis 0
+      # because axis 1 does not exist on 1D tensors and stacking scalars along axis 0 produces the 1D vector.
+      axis_to_stack = config.param_scan_axis if len(target_shape) > config.param_scan_axis else 0
   else:
     # Otherwise, if an unscanned MoE layer, and we stack along the expert axis (0).
     axis_to_stack = 0
@@ -545,6 +547,8 @@ def _get_hf_loading_function(hf_source_keys_or_key, tensor_getter, hook_fn, mt_t
   if not isinstance(hf_source_keys_or_key, list):
     # Case 1: Single hf key (str)
     def _loader(getter, key, shape, hook):
+      if key is None:
+        return apply_hook_fns(None, shape, hook)
       if isinstance(key, (list, tuple)):
         tensors = tuple(getter(k) for k in key)
         return apply_hook_fns(tensors, shape, hook)
@@ -567,6 +571,7 @@ def _loader(getter, key, shape, hook):
         hook_fn,
         mt_target_shape_or_shapes,
         config,
+        mt_key,
     )
   else:
     # isinstance(hf_source_keys_or_key[0], list)
@@ -1022,6 +1027,7 @@ def main(
     jax_weights = convert_lora_to_maxtext_adapter(config, lora_weights, save_dtype)
     adapter_name = os.path.basename(os.path.normpath(hf_lora_adapter_path))
     output_directory = os.path.join(output_directory, model_name_for_path, adapter_name)
+    state_params = jax_weights
   else:
 
     if lazy_load_tensors and config.use_multimodal:
@@ -1093,6 +1099,8 @@ def main(
           }
 
       def _eager_getter(key):
+        if key is None:
+          return None
         resolved_key = resolve_scale_key(key, hf_state_dict_numpy)
         if resolved_key not in hf_state_dict_numpy:
           raise ValueError(f"HuggingFace key {key} not found in state_dict.")
@@ -1180,11 +1188,19 @@ def _eager_getter(key):
     max_logging.log(f"Elapse for transform: {(time.time() - start) / 60:.2f} min")
     print_ram_usage("Before creating full JAX tree")
 
-    # Create final MaxText parameters tree
+    max_logging.log(f"Length of final_mt_weights: {len(final_mt_weights)}")
+    max_logging.log(f"Treedef num_leaves: {abstract_params_treedef.num_leaves}")
+    # Create final MaxText parameters tree containing all collections
     jax_weights = jax.tree_util.tree_unflatten(abstract_params_treedef, final_mt_weights)
     del final_mt_weights, abstract_params_treedef
 
+    state_params = jax_weights
+
   print_ram_usage("Before saving")
+  leaves = jax.tree_util.tree_leaves(state_params)
+  max_logging.log(f"Length of state_params leaves: {len(leaves)}")
+  if leaves:
+    max_logging.log(f"Type of first leaf before saving: {type(leaves[0])}")
   if lazy_load_tensors and not is_adapter_only:
     max_logging.log("Starting checkpoint save (loading weights just-in-time)...")
   else:
@@ -1195,7 +1211,7 @@ def _eager_getter(key):
   # and sharded across virtual devices.
   save_weights_to_checkpoint(
       output_directory,
-      jax_weights,
+      state_params,
       simulated_cpu_devices_count,
       config.checkpoint_storage_use_ocdbt,
       config.checkpoint_storage_use_zarr3,

src/maxtext/checkpoint_conversion/to_maxtext_diffusers.py

@@ -0,0 +1,15 @@
+# Copyright 2023–2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Universal Hugging Face Diffusers checkpoint converter and manager."""

src/maxtext/checkpoint_conversion/utils/hf_model_configs.py

@@ -1608,6 +1608,86 @@ def __init__(self, **kwargs):
   qwen3_5_35b_a3b_config = PTConfig(**qwen3_5_35b_a3b_dict)  # pytype: disable=wrong-arg-types
   qwen3_5_397b_a17b_config = PTConfig(**qwen3_5_397b_a17b_dict)  # pytype: disable=wrong-arg-types
 
+qwen3_5_9b_dict = {
+    "architectures": ["Qwen3_5ForConditionalGeneration"],
+    "image_token_id": 248056,
+    "model_type": "qwen3_5",
+    "text_config": {
+        "attention_bias": False,
+        "attention_dropout": 0.0,
+        "attn_output_gate": True,
+        "bos_token_id": None,
+        "dtype": "bfloat16",
+        "eos_token_id": 248044,
+        "full_attention_interval": 4,
+        "head_dim": 256,
+        "hidden_act": "silu",
+        "hidden_size": 4096,
+        "initializer_range": 0.02,
+        "intermediate_size": 12288,
+        "layer_types": [
+            "linear_attention",
+            "linear_attention",
+            "linear_attention",
+            "full_attention",
+        ]
+        * 8,
+        "linear_conv_kernel_dim": 4,
+        "linear_key_head_dim": 128,
+        "linear_num_key_heads": 16,
+        "linear_num_value_heads": 32,
+        "linear_value_head_dim": 128,
+        "mamba_ssm_dtype": "float32",
+        "max_position_embeddings": 262144,
+        "mlp_only_layers": [],
+        "model_type": "qwen3_5_text",
+        "mtp_num_hidden_layers": 1,
+        "mtp_use_dedicated_embeddings": False,
+        "num_attention_heads": 16,
+        "num_hidden_layers": 32,
+        "num_key_value_heads": 4,
+        "pad_token_id": None,
+        "partial_rotary_factor": 0.25,
+        "rms_norm_eps": 1e-06,
+        "rope_parameters": {
+            "mrope_interleaved": True,
+            "mrope_section": [11, 11, 10],
+            "partial_rotary_factor": 0.25,
+            "rope_theta": 10000000,
+            "rope_type": "default",
+        },
+        "tie_word_embeddings": False,
+        "use_cache": True,
+        "vocab_size": 248320,
+    },
+    "tie_word_embeddings": False,
+    "transformers_version": "5.16.0.dev0",
+    "video_token_id": 248057,
+    "vision_config": {
+        "deepstack_visual_indexes": [],
+        "depth": 27,
+        "hidden_act": "gelu_pytorch_tanh",
+        "hidden_size": 1152,
+        "in_channels": 3,
+        "initializer_range": 0.02,
+        "intermediate_size": 4304,
+        "model_type": "qwen3_5_vision",
+        "num_heads": 16,
+        "num_position_embeddings": 2304,
+        "out_hidden_size": 4096,
+        "patch_size": 16,
+        "spatial_merge_size": 2,
+        "temporal_patch_size": 2,
+    },
+    "vision_end_token_id": 248054,
+    "vision_start_token_id": 248053,
+}
+
+try:
+  qwen3_5_9b_config = transformers.Qwen3_5Config(**qwen3_5_9b_dict)  # pyrefly: ignore[missing-attribute]
+except AttributeError:
+  qwen3_5_9b_config = PTConfig(**qwen3_5_9b_dict)  # pytype: disable=wrong-arg-types
+
 
 # from https://huggingface.co/mistralai/Mixtral-8x7B-Instruct-v0.1/blob/main/config.json
 mixtral_8x7b_dict = {
@@ -1883,9 +1963,9 @@ def __init__(self, **kwargs):
 }
 qwen3_vl_30b_a3b_config = PTConfig(**qwen3_vl_30b_a3b_dict)
 
-cosmos3_nano_reasoner_dict = {
-    "architectures": ["Cosmos3OmniForConditionalGeneration"],
-    "model_type": "cosmos3_omni",
+weaver_mini_dict = {
+    "architectures": ["WeaverOmniForConditionalGeneration"],
+    "model_type": "weaver_omni",
     "text_config": {
         "attention_bias": False,
         "attention_dropout": 0.0,
@@ -1931,11 +2011,11 @@ def __init__(self, **kwargs):
         "temporal_patch_size": 2,
     },
 }
-cosmos3_nano_reasoner_config = PTConfig(**cosmos3_nano_reasoner_dict)
+weaver_mini_config = PTConfig(**weaver_mini_dict)
 
-cosmos3_super_reasoner_dict = {
-    "architectures": ["Cosmos3ForConditionalGeneration"],
-    "model_type": "cosmos3_omni",
+weaver_max_dict = {
+    "architectures": ["WeaverForConditionalGeneration"],
+    "model_type": "weaver_omni",
     "text_config": {
         "attention_bias": False,
         "attention_dropout": 0.0,
@@ -1981,7 +2061,7 @@ def __init__(self, **kwargs):
         "temporal_patch_size": 2,
     },
 }
-cosmos3_super_reasoner_config = PTConfig(**cosmos3_super_reasoner_dict)
+weaver_max_config = PTConfig(**weaver_max_dict)
 
 
 # {maxtext model name: hf model config}
@@ -2013,8 +2093,8 @@ def __init__(self, **kwargs):
     "qwen3-vl-2b": qwen3_vl_2b_config,
     "qwen3-vl-4b": qwen3_vl_4b_config,
     "qwen3-vl-30b-a3b": qwen3_vl_30b_a3b_config,
-    "cosmos3-nano-reasoner": cosmos3_nano_reasoner_config,
-    "cosmos3-super-reasoner": cosmos3_super_reasoner_config,
+    "weaver-mini": weaver_mini_config,
+    "weaver-max": weaver_max_config,
     "llama3.1-8b": llama31_8b_config,
     "llama3.1-8b-Instruct": llama31_8b_config,
     "llama3.1-70b": llama31_70b_config,
@@ -2036,6 +2116,7 @@ def __init__(self, **kwargs):
     "qwen3.5-35b-a3b": qwen3_5_35b_a3b_config,
     "qwen3.5-35b-a3b-fp8": qwen3_5_35b_a3b_config,
     "qwen3.5-35b-fp8": qwen3_5_35b_a3b_config,
+    "qwen3.5-9b": qwen3_5_9b_config,
     "mixtral-8x7b": mixtral_8x7b_config,
     "mixtral-8x22b": mixtral_8x22b_config,
     "olmo3-7b": olmo3_7b_config,