Skip to content

Support inserting and replaying router expert decisions in trainer #4826

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
khatwanimohit wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

Support inserting and replaying router expert decisions in trainer #4826

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions src/maxtext/integration/tunix/tunix_adapter.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -108,6 +108,7 @@ def __call__(
attention_mask: Optional[Array], # [B, L, L] or None
decoder_segment_ids: Optional[Array] = None,
output_hidden_states: bool = False, # ignored
forced_routed_experts: Optional[Array] = None,
) -> Tuple[Array, None]:
"""Forward compatible with Tunix Trainers default loss.
Returns logits, None.
Expand All @@ -118,6 +119,7 @@ def __call__(
decoder_input_tokens=input_tokens,
decoder_positions=positions,
decoder_segment_ids=decoder_segment_ids,
forced_routed_experts=forced_routed_experts,
)
return logits, None

Expand Down
49 changes: 44 additions & 5 deletions src/maxtext/layers/decoders.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -861,11 +861,15 @@ def __call__(
kv_caches: list[jax.Array] | None = None,
attention_metadata=None,
deepstack_visual_embeds: None | list[jnp.ndarray] = None,
forced_routed_experts: jnp.ndarray | None = None,
):
cfg = self.config
mesh = self.mesh
assert decoder_input_tokens.ndim == 2 # [batch, len]

if cfg.scan_layers and forced_routed_experts is not None:
raise NotImplementedError("Forced routing with scanned layers is not supported yet.")

# [batch, length] -> [batch, length, emb_dim]
y = self._apply_embedding(
shared_embedding,
Expand Down Expand Up @@ -1164,6 +1168,10 @@ def __call__(
global_layer_idx = global_layer_idx_offset + index
kv_cache = kv_caches[index] if kv_caches is not None else None
input_tokens = decoder_input_tokens if cfg.engram_layers else None
current_forced_routed_experts = None
if forced_routed_experts is not None and layer_prefix == "moe_layers":
current_forced_routed_experts = forced_routed_experts[:, :, index, :]

y, kv_cache = layer(
config=cfg,
mesh=mesh,
Expand All @@ -1182,6 +1190,7 @@ def __call__(
kv_cache=kv_cache,
attention_metadata=attention_metadata,
decoder_input_tokens=input_tokens,
forced_routed_experts=current_forced_routed_experts,
)
if kv_caches is not None and kv_cache is not None:
kv_caches[index] = kv_cache
Expand All @@ -1201,6 +1210,7 @@ def __call__(
slot=slot,
)
else:
moe_lyr_idx = 0
for lyr in range(cfg.num_decoder_layers):
RemattedBlockLayer = RemattedBlockLayers[0]
layer_kwargs = {}
Expand Down Expand Up @@ -1244,17 +1254,46 @@ def __call__(
layer = RemattedBlockLayer(
config=cfg, mesh=mesh, name=f"layers_{lyr}", quant=self.quant, model_mode=self.model_mode, **layer_kwargs
)
current_forced_routed_experts = None
is_moe = False
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.LLAMA4:
is_moe = llama4.determine_is_moe_layer(lyr, self.config.interleave_moe_layer_step)
Comment on lines +1259 to +1268

@gemini-code-assist gemini-code-assist Bot Aug 11, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

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


if is_moe and forced_routed_experts is not None:
if forced_routed_experts.ndim == 4:
current_forced_routed_experts = forced_routed_experts[:, :, moe_lyr_idx, :]
else:
current_forced_routed_experts = forced_routed_experts
moe_lyr_idx += 1
elif is_moe:
moe_lyr_idx += 1

call_kwargs = {
"previous_chunk": previous_chunk,
"slot": slot,
"kv_cache": kv_cache,
"attention_metadata": attention_metadata,
}
call_kwargs.update(layer_call_kwargs)

if is_moe and current_forced_routed_experts is not None:
call_kwargs["forced_routed_experts"] = current_forced_routed_experts

y, returned_cache = layer(
y,
decoder_segment_ids,
decoder_positions,
deterministic,
model_mode,
previous_chunk=previous_chunk,
slot=slot,
kv_cache=kv_cache,
attention_metadata=attention_metadata,
**layer_call_kwargs,
**call_kwargs,
)
if kv_caches is not None and returned_cache is not None:
kv_caches[lyr] = returned_cache
Expand Down
Loading
Loading