Layer stacking does not support MoE load balancing #22

Migrated from https://github.com/google-pytorch/torchtitan-archive/issues/391 (opened by jack8558 on 2026-09-12).

Stacking merges each parameter across layers into one [n_layers, ...] tensor. Buffers are not stacked, so every layer runs with layer 0's.

MoE keeps its load-balancing state in buffers. Load balancing here is the auxiliary-loss-free scheme (https://arxiv.org/abs/2408.15664): tokens_per_expert_E counts the tokens each expert received, and after every optimizer step a hook nudges expert_bias_E so over-used experts get a lower router score. Both are per-layer.

Shared across layers, every layer accumulates into and routes off layer 0's counts. No value comparison catches it -- the buffers start at zero on every layer and only diverge once training starts.

check_layers_are_stackable refuses stacking when a layer holds an MoE with load_balance_coeff set. With it unset there is no expert_bias_E and nothing reads tokens_per_expert_E, so sharing layer 0's is harmless.

Who this blocks

load_balance_coeff defaults to 1e-3, so every MoE recipe has it on except qwen3_30b_a3b_global_aux, which uses an aux loss instead. That includes stock qwen3-235b-a22b. Our scan A/B runs worked around it by overriding the coefficient to None, which is fine for a benchmark but not for a real run.

Two ways to fix this

A. Make the buffers per-layer.

Stack the MoE buffers the same way we stack weights, and pass row i into functional_call -- it already takes buffers, we just never send any. That alone fixes expert_bias_E, and fixes tokens_per_expert_E for the Python-loop path.

Scan needs one more step: tokens_per_expert_E.add_() writes into a scan input, which is not allowed. The count has to come out of the body as a carry or a per-layer output instead.

Keeps the aux-loss-free algorithm working. The scan part is unscoped -- nobody has tried it.

B. Use the aux-loss balancer instead.

It balances by adding a term to the loss, so it holds no per-layer state and there is nothing to stack. Set aux_loss and leave load_balance_coeff unset -- they are mutually exclusive, so expert_bias_E is never created.

Two small changes are needed first:

• Skip the tokens_per_expert_E accumulation when load_balance_coeff is None. Nothing reads it there, and the in-place write is what breaks scan.
• Set metrics.enable_moe_routing = False. moe.py:272 and :753 stash tensors on the module, which also fails inside a scan body.

Then add a 235B-A22B-global-aux flavor -- only 30B-A3B-global-aux exists today.

Verified on 8 chips, 2 MoE layers under scan: the bias scheme is refused, the aux scheme trains, and its loss differs from an unbalanced control, so the gradient really is applied.

MaxText's Qwen3-235B already uses B (load_balance_loss_weight: 0.001, routed_bias: false). So B is how this model is normally trained, not a workaround.

B does not replace A. load_balance_coeff defaults to 1e-3, so every other MoE recipe here still needs A, and MaxText itself uses the bias scheme for DeepSeek (routed_bias).

From https://github.com/google-pytorch/torchtitan-archive/pull/355 review, item 1.