dockyard_rl.models.jax.moe.load_balance

Aux-loss-free load-balance bias updater for MoE routing, JAX.

JAX mirror of models/dtensor/moe/load_balance.py. DeepSeek aux-loss-free load balancing (https://arxiv.org/abs/2408.15664): a per-expert bias expert_bias_E nudges the routing CHOICE (added to the router scores before top-k, not to the gating weights), stepped each optimizer update from the observed per-expert load — under-loaded experts get a positive bump, over-loaded a negative one, by a fixed magnitude load_balance_coeff.

compute_expert_bias_delta (the math) is pure and CPU-validatable. The cross-rank reduction that makes tokens_per_expert_E a GLOBAL load before the step is EP/data-parallel-aware (an all_reduce over the dp/cp axes) and hardware-deferred — injected as reduce_tokens_fn. With no reduction (single device / CPU) the step is exact for that device.

Module Contents

Functions

compute_expert_bias_delta	Sign-based, zero-centered expert-bias step.
iter_lb_moe_blocks	Yield every MoEBlock with aux-loss-free load balancing enabled.
update_expert_biases	Step each block’s expert_bias_E from its accumulated token load.

Data

Array
ReduceTokensFn

API

dockyard_rl.models.jax.moe.load_balance.Array

None

dockyard_rl.models.jax.moe.load_balance.ReduceTokensFn

None

dockyard_rl.models.jax.moe.load_balance.compute_expert_bias_delta(tokens_per_expert_E: dockyard_rl.models.jax.moe.load_balance.Array, load_balance_coeff: float) → dockyard_rl.models.jax.moe.load_balance.Array

Sign-based, zero-centered expert-bias step.

delta = coeff * sign(mean_load - load) moves under-loaded experts up and over-loaded down by coeff; subtracting the mean conserves the total bias (it only redistributes routing pressure).

dockyard_rl.models.jax.moe.load_balance.iter_lb_moe_blocks(model: flax.nnx.Module) → Iterator[dockyard_rl.models.jax.moe.block.MoEBlock]

Yield every MoEBlock with aux-loss-free load balancing enabled.

nnx.iter_modules recurses the whole graph, so nested blocks are found.

dockyard_rl.models.jax.moe.load_balance.update_expert_biases(blocks: Sequence[dockyard_rl.models.jax.moe.block.MoEBlock], reduce_tokens_fn: Optional[dockyard_rl.models.jax.moe.load_balance.ReduceTokensFn] = None) → None

Step each block’s expert_bias_E from its accumulated token load.

For every block: optionally reduce the local tokens_per_expert_E to a global load (reduce_tokens_fn, device-bound), apply the sign step, and zero the counter for the next window. When the reduction yields identical global counts on every device, all devices step expert_bias_E identically and stay in sync. The buffers are non-trained Buffer state (no grad).