diff --git a/src/maxtext/checkpoint_conversion/utils/param_mapping.py b/src/maxtext/checkpoint_conversion/utils/param_mapping.py index 26359cdddc..f6afeff148 100644 --- a/src/maxtext/checkpoint_conversion/utils/param_mapping.py +++ b/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,8 +837,10 @@ 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: @@ -834,8 +848,10 @@ def reshape_bias(input_tensor, target_shape=None): 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 diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py index 16a2935d85..db0fda7c8f 100644 --- a/src/maxtext/layers/moe.py +++ b/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): diff --git a/src/maxtext/trainers/post_train/sft/train_sft.py b/src/maxtext/trainers/post_train/sft/train_sft.py index 4616cef1d4..f46d6f6141 100644 --- a/src/maxtext/trainers/post_train/sft/train_sft.py +++ b/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 diff --git a/src/maxtext/utils/lora_utils.py b/src/maxtext/utils/lora_utils.py index 3b29e89dab..4b313c178b 100644 --- a/src/maxtext/utils/lora_utils.py +++ b/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, diff --git a/tests/end_to_end/tpu/gemma3/4b/test_gemma3_to_mt.sh b/tests/end_to_end/tpu/gemma3/4b/test_gemma3_to_mt.sh index 18dfb80f16..f0a3321f36 100644 --- a/tests/end_to_end/tpu/gemma3/4b/test_gemma3_to_mt.sh +++ b/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 \ diff --git a/tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss.sh b/tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss.sh index 12edbc2592..e2f74d602b 100644 --- a/tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss.sh +++ b/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 diff --git a/tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_rl.sh b/tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_rl.sh index f00da367a3..f342afce9b 100644 --- a/tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_rl.sh +++ b/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 \ No newline at end of file diff --git a/tests/end_to_end/tpu/llama3.1/70b/test_llama3.1_70b_rl.sh b/tests/end_to_end/tpu/llama3.1/70b/test_llama3.1_70b_rl.sh index 49f84da45b..a76ddf0459 100644 --- a/tests/end_to_end/tpu/llama3.1/70b/test_llama3.1_70b_rl.sh +++ b/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"]}' \ diff --git a/tests/end_to_end/tpu/llama3.1/70b/test_llama3.1_70b_sft.sh b/tests/end_to_end/tpu/llama3.1/70b/test_llama3.1_70b_sft.sh index c12e299a1c..891a1a515d 100644 --- a/tests/end_to_end/tpu/llama3.1/70b/test_llama3.1_70b_sft.sh +++ b/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 diff --git a/tests/end_to_end/tpu/qwen3/30b/test_qwen3.sh b/tests/end_to_end/tpu/qwen3/30b/test_qwen3.sh index 30192fbeeb..e611fc9ee0 100755 --- a/tests/end_to_end/tpu/qwen3/30b/test_qwen3.sh +++ b/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 diff --git a/tests/end_to_end/tpu/qwen3/30b/test_qwen3_rl.sh b/tests/end_to_end/tpu/qwen3/30b/test_qwen3_rl.sh index 135f77bbc6..3b9a827f36 100755 --- a/tests/end_to_end/tpu/qwen3/30b/test_qwen3_rl.sh +++ b/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 \ diff --git a/tests/end_to_end/tpu/qwen3/30b/test_qwen3_sft.sh b/tests/end_to_end/tpu/qwen3/30b/test_qwen3_sft.sh index fb8331661f..637e73f9d9 100755 --- a/tests/end_to_end/tpu/qwen3/30b/test_qwen3_sft.sh +++ b/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." \ No newline at end of file diff --git a/tests/end_to_end/tpu/qwen3/30b/test_qwen3_to_mt.sh b/tests/end_to_end/tpu/qwen3/30b/test_qwen3_to_mt.sh index df32ba7218..98050e5fdd 100755 --- a/tests/end_to_end/tpu/qwen3/30b/test_qwen3_to_mt.sh +++ b/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 \ diff --git a/tests/end_to_end/tpu/qwen3/vl_2b/test_qwen3_to_mt.sh b/tests/end_to_end/tpu/qwen3/vl_2b/test_qwen3_to_mt.sh index e98d6f5ed7..051869c93d 100644 --- a/tests/end_to_end/tpu/qwen3/vl_2b/test_qwen3_to_mt.sh +++ b/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) diff --git a/tests/end_to_end/tpu/qwen3/vl_2b/test_qwen3_vl_2b_to_hf_e2e.sh b/tests/end_to_end/tpu/qwen3/vl_2b/test_qwen3_vl_2b_to_hf_e2e.sh index 56a8e12073..f65372f19c 100644 --- a/tests/end_to_end/tpu/qwen3/vl_2b/test_qwen3_vl_2b_to_hf_e2e.sh +++ b/tests/end_to_end/tpu/qwen3/vl_2b/test_qwen3_vl_2b_to_hf_e2e.sh @@ -40,7 +40,7 @@ export 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 diff --git a/tests/post_training/unit/lora_utils_test.py b/tests/post_training/unit/lora_utils_test.py index d0fb3d503a..f1c29c4f16 100644 --- a/tests/post_training/unit/lora_utils_test.py +++ b/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() diff --git a/tests/post_training/unit/train_sft_test.py b/tests/post_training/unit/train_sft_test.py index 3c927c6802..3e71ba6273 100644 --- a/tests/post_training/unit/train_sft_test.py +++ b/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()