Fake balanced routing in MoE

[rakkit]

we can set DEBUG_FORCE_LOAD_BALANCED=1 to force each experts get same amount of tokens.

reprodue: DEBUG_FORCE_LOAD_BALANCED=1 CONFIG_FILE="./torchtitan/models/deepseek_v3/train_configs/debug_model.toml" NGPU=4 ./run_train.sh --compile.enable

here is test on 8layers, 8 activate and 64 total experts. Green one is vanilla one and purple one is with force load balance

[rakkit]

Thanks for the comments, I’ve updated the code as per your suggestions.

It seems that variables starting with _ broke CLI. so left it as training.debug_moe_force_load_balance.

[tianyu-l]

Thanks a lot, left some final nit comments.

[garrett361]

Out of curiosity, what's your use case for this @rakkit ? I've used forced MoE balancing before to ensure I can run my model setup without OOMing in (at the very least) the well-balanced routing regime and to find the best-case throughput, but wondered if you had other reasons.

Same question for @tianyu-l ; maybe the above and also for test cases?

PS: a little tangential, but I have (messy) code for aggregating MoE routing stats across PP ranks so that all layer stats get reported, in case you're using pipeline parallel and ever need that.

[tianyu-l]

@garrett361

Same question for @tianyu-l ; maybe the above and also for test cases?

I think also for getting an idea of throughput under some specific setup, without the impact of load imbalance, which is also along the lines of

to ensure I can run my model setup without OOMing in (at the very least) the well-balanced routing regime and to find the best-case throughput

[rakkit]

Thanks a lot for the comments @tianyu-l, these are solved now!

[rakkit]

Thanks @garrett361, sry for the late response.

what's your use case for this

As tianyu said, the main motivation is testing throughput more conveniently. We also noticed that enforcing load balance helps avoid OOM. We plan to run extensive throughput checks (similar to HF’s Ultra-Scale Playbook) once MoE+Compile becomes available in torchtitan.

PS: a little tangential, but I have (messy) code for aggregating MoE routing stats across PP ranks so that all layer stats get reported, in case you're using pipeline parallel and ever need that.

Thanks a lot! Can you pls share the link to the code? For our use case(here), we decide to log each PP stage separately, since torchtitan supports creating a wandb run per rank via save_for_all_ranks.

The reason are 1) we have lots of log entries on each PP-stage from the Muon/Scion optimizer (we have a preprint on arXiv maybe this week), which provide better insight into training dynamics than simply tracking the global gradient norm. and 2) Aggregating metrics across PP stages requires extra communication. Its ok for most PP strategies, but I feel its awkward for zb-v (zero-bubble PP), where optimizer.step runs sequentially along PP stages. In that case, communication would add sync overhead to async steps, or we’d need buffering for proper logging.

[tianyu-l]

Thanks a lot! Had one last nit comment.

[tianyu-l]

One last nit to let arg parser pick up the helper message.

Suggested change

	# if True, we force each experts get same amount of token via round-robin
	"""If True, we force each experts to get the same amount of tokens via round-robin. This option is for debugging usage only."""

[garrett361]

we decide to log each PP stage separately, since torchtitan supports creating a wandb run per rank via save_for_all_ranks

Oh, nice, that is probably the better option. Didn't know about save_for_all_ranks.

My implementation certainly has extra exposed comms. I thought they'd usually be pretty negligible for most models which are large enough to require PP, but maybe not the case for zb-v. In any case, the code is around here.

[garrett361]

Probably a nit, but I think the version like

selected_experts_indices = (torch.arange(num_tokens * self.top_k, device=device, dtype=torch.int64).reshape(
    num_tokens, self.top_k
) % self.num_experts)

is a little cleaner, as well as faster.

Also, we could cache this and avoid recomputing the selected indices every time, right?

[garrett361]

Just had one nit, think this is pretty useful!

[tianyu-l]

@rakkit would you like to address the final comment? Anyway it seems a rebase needs to happen before we merge.

[rakkit]

@tianyu-l which final comment? #1670 (comment) this one fixed. and i took the version #1670 (comment) garrett361 proposed.

also rebased.