[RFC] Add sequence-wise MoE load-balance auxiliary loss (DeepSeek-V3)#3864
[RFC] Add sequence-wise MoE load-balance auxiliary loss (DeepSeek-V3)#3864sdmyzlp wants to merge 4 commits into
Conversation
meta-cla
Bot
added
the
CLA Signed
|
The following ciflow label(s) have been added but CI has not been triggered yet because the workflows are awaiting approval:
Once a maintainer approves the workflows (scroll to the bottom of the PR page), the corresponding CI jobs will be triggered automatically. Please ping one of the reviewers if you do not have access to approve and run workflows. |
sdmyzlp
requested review from
IvanKobzarev,
SherlockNoMad,
aditvenk,
felipemello1,
sanketpurandare and
xmfan
as code owners
sdmyzlp
force-pushed
the
br_aux_loss
branch
3 times, most recently
from
96e4048 to
2def8b1
Compare
Thanks @tianyu-l for the pointer, and cc @pianpwk @shuhuayu. We have carefully studied #3000 and this PR is built on top of that existing work, with the main modifications summarized in the updated description — the key shift being a universal We understand that @pianpwk may currently be focusing on the spmd_types migration, and wonder if we could keep this PR open for further discussion. Would love to hear your thoughts. |
sdmyzlp
force-pushed
the
br_aux_loss
branch
2 times, most recently
from
42ce061 to
6307ed2
Compare
config.training.global_batch_size defaults to -1 and was independently resolved by each consumer. This caused update_from_config (which runs before the resolution) to read -1 and propagate it into sub-configs. Fix: resolve global_batch_size in-place on the config object right after batch_degree / dp_degree is known, before model_config.update_from_config or any downstream consumer reads it. All later -1 checks are now dead code because the value is always resolved.
…losses Auxiliary objectives computed inside the model (MoE load-balance, DSA indexer KL, ...) cannot reach the trainer's loss function under pipeline parallelism, where each stage output is a single tensor. The gradient injection decouples the gradient from the scalar output, and the logging framework accumulates per-step values with PP-safe collection.
Per-sequence f_i * P_i load-balance loss penalizes expert imbalance independently of batch composition, making gradients correct under PP microbatching and gradient accumulation.
When enabled, extra_metrics (e.g. aux losses) are printed to the console with per-key color cycling.
2 participants
Add a universal auxiliary loss framework. with MoE sequence-wise load balancing as the first example consumer. This work follows the approach from #3000 (gradient injection via
autograd.Functionfor PP compatibility).With several modification:
spmd_types-native reduction. This PR instead uses
ShardingConfigwithspmd.P → spmd.Iplacement on_SeqwiseCounts, letting the spmd_types framework handle the all-reduce at the module boundary instead of hand-rolled collectives.LoggedAuxLossas a generic base class. Extracts aux loss as a reusableLoggedAuxLossthat provides both gradient injection and per-microbatch metric accumulation with PP-safe collection viacollect_aux_loss_metrics. Future losses could subclass it.Built-in logging. Includes an logging attempt out of the box — aux losses appear in the console log line behind
metrics.console_log_extra_metricsand are collected automatically across all PP stages.@spmd.register_autograd_function. instead of a rawtorch.autograd.Function, for native spmd_types compatibility.Tested configurations
Side fixes
global_batch_sizeresolution beforeupdate_from_config()soLoggedAuxLoss.Config.global_batch_sizeis available during model config traversal.