Support inserting and replaying router expert decisions in trainer

[khatwanimohit]

• In train.py loss_fn, extract forced_routed_experts from batch data and pass to model forward
• In RoutedMoE, override top_k_indices and route tokens with forced_routed_experts across sparse_matmul and dense_matmul
• In Decoder and NNXDecoder, plumb forced_routed_experts down across scanned/sequential decoder layers
• Plumbed forced_routed_experts across model implementations (Qwen3, Qwen3.5, DeepSeek, Gemma4, Mixtral)
• Added unit tests in tests/unit/forced_routing_test.py
• Added integration test in tests/test_trainer_router_replay.py verifying loss computation with forced routing

Description

Start with a short description of what the PR does and how this is a change from
the past.

The rest of the description includes relevant details and context, examples:

• why is this change being made,
• the problem being solved and any relevant context,
• why this is a good solution,
• some information about the specific implementation,
• shortcomings of the solution and possible future improvements.

If the change fixes a bug or a Github issue, please include a link, e.g.,:
FIXES: b/123456
FIXES: #123456

You can also provide a comma-separated list. If you don't want to close a bug but
simply to reference it, use BUGS, e.g.:
BUGS: b/123456

Notice 1: Once all tests pass, the "pull ready" label will automatically be assigned.
This label is used for administrative purposes. Please do not add it manually.

Notice 2: For external contributions, our settings currently require an approval from a MaxText maintainer to trigger CI tests.

Tests

Please describe how you tested this change, and include any instructions and/or
commands to reproduce.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[gemini-code-assist]

Code Review

This pull request introduces support for forced routing in Mixture of Experts (MoE) models across MaxText, allowing the replay of router decisions and expert selections by passing forced_routed_experts. This feature is integrated into both standard and NNX decoder paths, the trainer, and various model architectures, accompanied by unit and integration tests. The review feedback highlights critical improvements: adding missing MoE block types (DEEPSEEK4 and ENVY) to the decoder layers, robustly zeroing out weights for partially padded expert indices element-wise, and replacing .set() with .add() in JAX index updates to prevent undefined behavior when handling duplicate indices.

[gemini-code-assist]

The is_moe check is missing DecoderBlockType.DEEPSEEK4 and DecoderBlockType.ENVY. DeepSeek V4 is an MoE model, and Envy also interleaves MoE layers. Without including them, forced_routed_experts will not be passed to their MoE layers, breaking forced routing for these architectures.

            if cfg.decoder_block in (
                DecoderBlockType.MIXTRAL,
                DecoderBlockType.QWEN3_MOE,
                DecoderBlockType.QWEN3_NEXT,
                DecoderBlockType.QWEN3_5,
                DecoderBlockType.QWEN3_CUSTOM_MOE,
                DecoderBlockType.DEEPSEEK4,
            ):
              is_moe = True
            elif cfg.decoder_block == DecoderBlockType.LLAMA4:
              is_moe = llama4.determine_is_moe_layer(lyr, self.config.interleave_moe_layer_step)
            elif cfg.decoder_block == DecoderBlockType.ENVY:
              is_moe = (lyr + 1) % self.config.interleave_moe_layer_step == 0

[gemini-code-assist]

Similar to decoders.py, the is_moe check in nnx_decoders.py is missing DecoderBlockType.DEEPSEEK4 and DecoderBlockType.ENVY. We should include them to ensure forced routing works correctly for these models under the NNX decoder path.

Suggested change

	if cfg.decoder_block in (
	DecoderBlockType.MIXTRAL,
	DecoderBlockType.QWEN3_MOE,
	DecoderBlockType.QWEN3_NEXT,
	DecoderBlockType.QWEN3_5,
	DecoderBlockType.QWEN3_CUSTOM_MOE,
	):
	is_moe = True
	elif cfg.decoder_block == DecoderBlockType.DEEPSEEK:
	is_moe = lyr >= cfg.first_num_dense_layers
	elif cfg.decoder_block == DecoderBlockType.LLAMA4:
	is_moe = llama4.determine_is_moe_layer(lyr, self.config.interleave_moe_layer_step)
	if cfg.decoder_block in (
	DecoderBlockType.MIXTRAL,
	DecoderBlockType.QWEN3_MOE,
	DecoderBlockType.QWEN3_NEXT,
	DecoderBlockType.QWEN3_5,
	DecoderBlockType.QWEN3_CUSTOM_MOE,
	DecoderBlockType.DEEPSEEK4,
	):
	is_moe = True
	elif cfg.decoder_block == DecoderBlockType.DEEPSEEK:
	is_moe = lyr >= cfg.first_num_dense_layers
	elif cfg.decoder_block == DecoderBlockType.LLAMA4:
	is_moe = llama4.determine_is_moe_layer(lyr, self.config.interleave_moe_layer_step)
	elif cfg.decoder_block == DecoderBlockType.ENVY:
	is_moe = (lyr + 1) % self.config.interleave_moe_layer_step == 0

[gemini-code-assist]

Using top_k_indices[:, :, 0] != -1 only checks if the first expert index is valid. If a token is partially padded (where some expert indices are -1 and others are valid), the weights for the -1 indices will not be zeroed out. Checking top_k_indices != -1 element-wise is much more robust and correctly zeroes out weights for any individual padded expert index.

Suggested change

	if forced_routed_experts is not None:
	valid_token_mask = top_k_indices[:, :, 0] != -1
	top_k_weights = top_k_weights * valid_token_mask[:, :, None]
	if forced_routed_experts is not None:
	valid_expert_mask = top_k_indices != -1
	top_k_weights = top_k_weights * valid_expert_mask

[gemini-code-assist]

Using .set() with duplicate indices (which occurs when multiple -1 indices are mapped to 0 for the same token) has undefined behavior in JAX, and can silently overwrite valid routing weights with 0.0. Changing this to .add() avoids any overwriting issues and correctly preserves the valid weights since the dummy indices have a weight of 0.0.

Suggested change

	update_weights = update_weights.at[index_update].set(safe_weights, out_sharding=weight_sharding)
	update_weights = update_weights.at[index_update].add(safe_weights, out_sharding=weight_sharding)