fix(post_train): Fix Tunix SFT signatures, Qwix LoRA mesh sharding, and scale Qwen/LLaMA parallelism

src/maxtext/checkpoint_conversion/utils/param_mapping.py

@@ -784,6 +784,16 @@ def reshape_kernel(input_tensor, target_shape):
     else:
       return input_tensor.T.reshape(target_shape)
 
+  def reshape_expert_kernel(input_tensor, target_shape=None):
+    """Transposes expert weights.
+
+    3D: (num_experts, in_dim, out_dim) -> (num_experts, out_dim, in_dim)
+    2D: (in_dim, out_dim) -> (out_dim, in_dim)
+    """
+    if input_tensor.ndim == 3:
+      return input_tensor.transpose(0, 2, 1)
+    return input_tensor.transpose(1, 0)
+
   def reshape_bias(input_tensor, target_shape=None):
     """Reshapes biases between MaxText 2D (heads, dim) and HF 1D (hidden)."""
     # saving_to_hf: MaxText [heads, head_dim] -> HF [hidden_dim] (flatten)
@@ -809,11 +819,13 @@ def reshape_bias(input_tensor, target_shape=None):
       "self_attention-key-bias",
       "self_attention-value-bias",
   ]
-  moe_kernel_hooks = [
+  moe_gate_hooks = [
       "moe_block-gate-kernel",
       "moe_block-wi_0-kernel",
       "moe_block-wi_1-kernel",
       "moe_block-wo-kernel",
+  ]
+  moe_expert_hooks = [
       "moe_block-wi_0",
       "moe_block-wi_1",
       "moe_block-wo",
@@ -825,17 +837,21 @@ def reshape_bias(input_tensor, target_shape=None):
     for key in bias_hooks:
       mapping[f"params-decoder-layers-{key}"] = reshape_bias
     if num_experts > 1:
-      for key in moe_kernel_hooks:
+      for key in moe_gate_hooks:
         mapping[f"params-decoder-layers-{key}"] = reshape_kernel
+      for key in moe_expert_hooks:
+        mapping[f"params-decoder-layers-{key}"] = reshape_expert_kernel
   else:
     for i in range(n_layers):
       for key in kernel_hooks:
         mapping[f"params-decoder-layers_{i}-{key}"] = reshape_kernel
       for key in bias_hooks:
         mapping[f"params-decoder-layers_{i}-{key}"] = reshape_bias
       if num_experts > 1:
-        for key in moe_kernel_hooks:
+        for key in moe_gate_hooks:
           mapping[f"params-decoder-layers_{i}-{key}"] = reshape_kernel
+        for key in moe_expert_hooks:
+          mapping[f"params-decoder-layers_{i}-{key}"] = reshape_expert_kernel
   return mapping
 
 

src/maxtext/layers/moe.py

@@ -747,6 +747,12 @@ def get_topk(self, gate_logits, pre_bias_logits, rngs=None, input_ids=None):
       if self.config.norm_topk_prob:
         top_k_weights /= top_k_weights.sum(axis=-1, keepdims=True)
 
+      if self.per_expert_scale is not None and not (
+          self.config.model_call_mode == "inference" and self.config.fuse_expert_scales
+      ):
+        per_expert_scale_topk = jnp.take_along_axis(self.per_expert_scale.value[None, None, :], top_k_indices, axis=-1)
+        top_k_weights = top_k_weights * per_expert_scale_topk.astype(top_k_weights.dtype)
+
     return top_k_weights, top_k_indices
 
   def deepseek_scale_weights(self, weights):

src/maxtext/trainers/post_train/sft/train_sft.py

@@ -41,6 +41,7 @@
 from absl import app
 import os
 import jax
+import jax.numpy as jnp
 import optax
 import pathwaysutils
 
@@ -111,8 +112,9 @@ def create_train_step_fn(self):
     def train_step(
         model: nnx.Module,
         optimizer: nnx.Optimizer,
+        grad_accumulator: Any,
         inputs: Any,
-        grad_accumulator: Any = None,
+        is_update_step: Any = True,
     ):
       inputs = gen_fn(inputs)
 
@@ -157,13 +159,35 @@ def loss_wrapper(diff_params, rest, **inputs_kw):
 
       nnx.update(model, new_rest)
 
-      # Apply optimizer update. grads has the same nnx.State(wrt) structure
-      # as diff_params, which is compatible with optimizer.update.
-      optimizer.update(model, grads)
+      # Handle gradient accumulation and conditional/direct optimizer update
+      if grad_accumulator is not None and hasattr(grad_accumulator, "add"):
+        grad_accumulator.add(grads)
+
+        def apply_updates(model, optimizer, grad_accumulator):
+          acc_grads = grad_accumulator.get()
+          norm = optax.global_norm(jax.tree_util.tree_map(lambda x: x.astype(jnp.float32), acc_grads))
+          optimizer.update(model, acc_grads)
+          grad_accumulator.reset()
+          return norm
+
+        def skip_updates(model, optimizer, grad_accumulator):
+          return jnp.array(0.0, dtype=jnp.float32)
+
+        grad_norm = nnx.cond(
+            is_update_step,
+            apply_updates,
+            skip_updates,
+            model,
+            optimizer,
+            grad_accumulator,
+        )
+      else:
+        optimizer.update(model, grads)
+        grad_norm = optax.global_norm(jax.tree_util.tree_map(lambda x: x.astype(jnp.float32), grads))
 
       aux_out = aux if has_aux else None
       if tunix_expects_grad_norm:
-        return out_val, aux_out, optax.global_norm(grads)
+        return out_val, aux_out, grad_norm
       return out_val, aux_out
 
     return train_step

src/maxtext/utils/lora_utils.py

@@ -470,10 +470,18 @@ def _build_lora_provider(mt_config: pyconfig.HyperParameters) -> qwix.LoraProvid
   return qwix.LoraProvider(**lora_kwargs)
 
 
-def _prepare_dummy_inputs(dummy_bs: int = 1) -> tuple[jnp.ndarray, jnp.ndarray]:
-  """Builds dummy decoder inputs used to materialize LoRA parameters."""
-  # Keep LoRA warmup as small as possible to minimize compile/memory overhead.
-  seq_len = 1
+def _prepare_dummy_inputs(
+    mesh: Optional[jax.sharding.Mesh] = None,
+    dummy_bs: int = 1,
+    seq_len: int = 1,
+) -> tuple[jnp.ndarray, jnp.ndarray]:
+  """Builds minimal dummy decoder inputs partitioned appropriately for the mesh."""
+  if mesh is not None:
+    for axis in ("data", "fsdp", "fsdp_transpose", "expert"):
+      dummy_bs *= mesh.shape.get(axis, 1)
+    for axis in ("tensor_sequence", "context"):
+      seq_len *= mesh.shape.get(axis, 1)
+
   decoder_input_tokens = jnp.zeros((dummy_bs, seq_len), dtype=jnp.int32)
   decoder_positions = jnp.zeros((dummy_bs, seq_len), dtype=jnp.int32)
   return decoder_input_tokens, decoder_positions
@@ -598,12 +606,8 @@ def apply_lora_to_model(
 
   lora_provider = _build_lora_provider(mt_config)
 
-  dp_size = 1
-  if mesh is not None and "data" in mesh.shape:
-    dp_size = mesh.shape["data"]
-
   model_rngs = getattr(model.decoder, "rngs", None)  # pyrefly: ignore[missing-attribute]
-  decoder_input_tokens, decoder_positions = _prepare_dummy_inputs(dummy_bs=dp_size)
+  decoder_input_tokens, decoder_positions = _prepare_dummy_inputs(mesh)
 
   lora_model = qwix.apply_lora_to_model(
       model,

tests/end_to_end/tpu/gemma3/4b/test_gemma3_to_mt.sh

@@ -89,7 +89,7 @@ if [ "${USE_MULTIMODAL}" = "false" ]; then
         use_multimodal=${USE_MULTIMODAL} \
         scan_layers=false \
         --hf_model_path=${HF_GOLDEN_MODEL} \
-        --max_kl_div=0.03 \
+        --max_kl_div=0.05 \
         --run_hf_model=true \
         attention=dot_product \
         hardware=cpu \

tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss.sh

@@ -62,7 +62,6 @@ python3 -m maxtext.trainers.pre_train.train \
     async_checkpointing=false \
     checkpoint_storage_use_zarr3=False \
     checkpoint_storage_use_ocdbt=False \
-    ici_fsdp_parallelism=64 \
     model_name=${MODEL_NAME} \
     scan_layers=false \
     use_multimodal=false

tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_rl.sh

@@ -59,6 +59,6 @@ python3 -m maxtext.trainers.post_train.rl.train_rl \
     enable_single_controller=${use_pathways} \
     checkpoint_storage_use_zarr3=False \
     checkpoint_storage_use_ocdbt=False \
-    rollout_data_parallelism=4 \
-    rollout_tensor_parallelism=8 \
+    rollout_data_parallelism=-1 \
+    rollout_tensor_parallelism=4 \
     hbm_utilization_vllm=0.8

tests/end_to_end/tpu/llama3.1/70b/test_llama3.1_70b_rl.sh

@@ -44,6 +44,7 @@ python3 -m maxtext.trainers.post_train.rl.train_rl \
     num_batches=5 batch_size=1 num_test_batches=5 \
     model_name=${MODEL_NAME} tokenizer_path='meta-llama/Llama-3.1-70B-Instruct' \
     enable_single_controller=${use_pathways} \
+    remat_policy=full \
     checkpoint_storage_use_zarr3=False checkpoint_storage_use_ocdbt=False \
     rollout_tensor_parallelism=4 \
     vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \

tests/end_to_end/tpu/llama3.1/70b/test_llama3.1_70b_sft.sh

@@ -44,6 +44,7 @@ python3 -m maxtext.trainers.post_train.sft.train_sft \
     steps=5 scan_layers=true \
     model_name=${MODEL_NAME} tokenizer_path='meta-llama/Llama-3.1-70B-Instruct' \
     enable_single_controller=${use_pathways} \
+    remat_policy=full \
     checkpoint_storage_use_zarr3=False checkpoint_storage_use_ocdbt=False
 
 # Step 3: Run inference on the checkpoint generated from the previous run

tests/end_to_end/tpu/qwen3/30b/test_qwen3.sh

@@ -59,8 +59,6 @@ python3 -m maxtext.trainers.pre_train.train \
     model_name=${MODEL_NAME} \
     scan_layers=true \
     remat_policy=full \
-    ici_tensor_parallelism=4 \
-    ici_fsdp_parallelism=16 \
     weight_dtype=bfloat16 \
     dtype=bfloat16 \
     opt_type=sgd

tests/end_to_end/tpu/qwen3/30b/test_qwen3_rl.sh

@@ -53,8 +53,8 @@ python3 -m maxtext.inference.vllm_decode \
     max_target_length=256 \
     max_num_batched_tokens=256 \
     ici_tensor_parallelism=4 \
-    ici_expert_parallelism=4 \
-    ici_data_parallelism=4 \
+    ici_expert_parallelism=2 \
+    ici_data_parallelism=2 \
     allow_split_physical_axes=True \
     prefuse_moe_weights=True \
     use_chat_template=True \
@@ -82,10 +82,7 @@ python3 -m maxtext.trainers.post_train.rl.train_rl \
     remat_policy=full \
     hbm_utilization_vllm=0.55 \
     use_pathways=True \
-    chips_per_vm=8 \
-    ici_tensor_parallelism=4 \
-    ici_fsdp_parallelism=4 \
-    ici_expert_parallelism=2 \
+    chips_per_vm=4 \
     max_target_length=512 \
     weight_dtype=bfloat16 \
     dtype=bfloat16 \
@@ -105,8 +102,8 @@ python3 -m maxtext.inference.vllm_decode \
     max_target_length=256 \
     max_num_batched_tokens=256 \
     ici_tensor_parallelism=4 \
-    ici_expert_parallelism=4 \
-    ici_data_parallelism=4 \
+    ici_expert_parallelism=2 \
+    ici_data_parallelism=2 \
     allow_split_physical_axes=True \
     prefuse_moe_weights=True \
     use_chat_template=True \

tests/end_to_end/tpu/qwen3/30b/test_qwen3_sft.sh

@@ -35,8 +35,8 @@ python3 -m maxtext.inference.vllm_decode \
     scan_layers=true \
     enable_single_controller=True \
     ici_tensor_parallelism=4 \
-    ici_expert_parallelism=4 \
-    ici_data_parallelism=4 \
+    ici_expert_parallelism=2 \
+    ici_data_parallelism=2 \
     prompt="Suggest some famous landmarks in London."
 
 # Step 2: Run SFT starting from the pre-converted checkpoint
@@ -51,9 +51,6 @@ python3 -m maxtext.trainers.post_train.sft.train_sft \
     checkpoint_storage_use_zarr3=False \
     checkpoint_storage_use_ocdbt=False \
     remat_policy=full \
-    ici_tensor_parallelism=4 \
-    ici_fsdp_parallelism=4 \
-    ici_expert_parallelism=4 \
     enable_single_controller=True \
     max_target_length=16 \
     weight_dtype=bfloat16 \
@@ -70,6 +67,6 @@ python3 -m maxtext.inference.vllm_decode \
     scan_layers=true \
     enable_single_controller=True \
     ici_tensor_parallelism=4 \
-    ici_expert_parallelism=4 \
-    ici_data_parallelism=4 \
+    ici_expert_parallelism=2 \
+    ici_data_parallelism=2 \
     prompt="Suggest some famous landmarks in London."

tests/end_to_end/tpu/qwen3/30b/test_qwen3_to_mt.sh

@@ -54,6 +54,8 @@ echo "Scanned checkpoint path: ${SCANNED_CKPT_PATH}"
 python3 -m tests.utils.forward_pass_logit_checker \
     load_parameters_path=${UNSCANNED_CKPT_PATH} \
     model_name=${MODEL_NAME} \
+    per_device_batch_size=1 \
+    dtype=float32 \
     scan_layers=false \
     --hf_model_path=${HF_GOLDEN_MODEL} \
     --max_kl_div=0.03 \

tests/end_to_end/tpu/qwen3/vl_2b/test_qwen3_to_mt.sh

@@ -23,7 +23,7 @@ HF_GOLDEN_MODEL=Qwen/Qwen3-VL-2B-Instruct
 BASE_OUTPUT_DIRECTORY=gs://runner-maxtext-logs/${MODEL_NAME}/to_maxtext
 
 # Step 1: Install torch
-python3 -m pip install torch --index-url https://download.pytorch.org/whl/cpu
+python3 -m pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu
 python3 -m pip install decord
 
 # Step 2: Convert to scanned multimodal checkpoint (for multimodal training)

tests/end_to_end/tpu/qwen3/vl_2b/test_qwen3_vl_2b_to_hf_e2e.sh

@@ -40,7 +40,7 @@ export LOCAL_PATH=<your_local_path>/hf/${MODEL_NAME}/${idx}
 
 
 # Installing torch for deps in forward_pass_logit_checker.py
-python3 -m pip install torch --index-url https://download.pytorch.org/whl/cpu
+python3 -m pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu
 python3 -m pip install decord
 
 # Check point conversion

tests/post_training/unit/lora_utils_test.py

@@ -141,6 +141,12 @@ def test_prepare_dummy_inputs(self):
     self.assertEqual(tokens.shape, (1, 1))
     self.assertEqual(positions.shape, (1, 1))
 
+    mock_mesh = mock.MagicMock()
+    mock_mesh.shape = {"data": 2, "fsdp": 16, "tensor_sequence": 4}
+    tokens, positions = lora_utils._prepare_dummy_inputs(mock_mesh)
+    self.assertEqual(tokens.shape, (32, 4))
+    self.assertEqual(positions.shape, (32, 4))
+
   def test_verify_lora_parameters_success(self):
     """Test verification of LoRA parameters with matches and enabled LoRA."""
     mock_model = mock.MagicMock()

tests/post_training/unit/train_sft_test.py

@@ -14,6 +14,7 @@
 
 """Unit tests for train_sft.py."""
 
+import inspect
 import unittest
 from unittest import mock
 from types import SimpleNamespace
@@ -27,6 +28,8 @@
 class TrainSFTTest(unittest.TestCase):
   """Tests for train_sft.py."""
 
+  # pylint: disable=protected-access
+
   def test_validate_config_valid(self):
     config = SimpleNamespace(
         optimizer_memory_host_offload=False,
@@ -81,6 +84,25 @@ def test_train_model_caching_dense(self):
         cache_nnx_graph=True,
     )
 
+  def test_maxtext_peft_trainer_train_step_signature(self):
+    """Test that MaxTextPeftTrainer train_step accepts Tunix args including is_update_step."""
+    mock_model = mock.MagicMock()
+
+    with mock.patch("flax.nnx.pop"), mock.patch("flax.nnx.split", return_value=(mock.MagicMock(), {}, {})):
+      trainer = mock.MagicMock()
+      trainer.loss_fn = mock.MagicMock()
+      trainer._has_aux = False
+      trainer.gen_model_input_fn = lambda x: x
+      trainer._lora_enabled = False
+      trainer.model = mock_model
+
+      train_step_fn = train_sft.MaxTextPeftTrainer.create_train_step_fn(trainer)
+
+      # Should accept positional and keyword args from Tunix PeftTrainer
+      sig = inspect.signature(train_step_fn)
+      params = list(sig.parameters.keys())
+      self.assertEqual(params, ["model", "optimizer", "grad_accumulator", "inputs", "is_update_step"])
+
 
 if __name__ == "__main__":
   unittest.main()