Using "virtual padding" to calculate number_local_tokens per SP rank, and fix combine() shape mismatch

[wwwjn]

What's the problem

• [Currently] combine() wrongly assume the tokens are evenly sharded on each rank

  torchtitan/torchtitan/models/common/token_dispatcher.py

  Lines 439 to 456 in c0428bb

  	out_TD = torch.zeros(
  	x_TD.shape[0] * self.sp_size,
  	x_TD.shape[-1],
  	device=x_TD.device,
  	dtype=x_TD.dtype,
  	)
  	if not self.score_before_experts:
  	routed_output_RD = (
  	routed_output_RD.to(torch.float32)
  	* metadata.topk_scores_experts_sorted_N.reshape(-1, 1)
  	).to(routed_output_RD.dtype)
  	# With SP, token indices are 0-based within the local shard.
  	# Offset to global positions for the full-size scatter buffer.
  	if self.sp_size > 1:
  	token_indices_experts_sorted_N = (
  	metadata.token_indices_experts_sorted_N + x_TD.shape[0] * self.sp_rank

  (infer global SPMD in local SPMD region)
  • If uneven sharded, out_TD will have different shapes across SP ranks.
  • We should directly ban if input number of tokens in input batch can not be evenly sharded by SP ranks

• [Future] Router will use spmd_types soon, and router is per SP rank. Per SP rank should have even sharding

• [Future] we want to avoid dispatch/load_balacing the padded token, we should be able to do that by adding metadata field to record the actually local tokens for each sp rank

What does this PR do?

This PR is doing "virtual padding" , and passing metadata around

• Calculate num_local_tokens_after_padding = (T + pad_tokens) // sp_size in MoE module
• Pass num_local_tokens_after_padding to GroupedExperts module, then to combine()
• combine() returns a tensor with shape (num_local_tokens_after_padding * sp_rank, .... )
• slice the combined tensor to (T, ...) in MoE

[wwwjn]

Introducing this sp_rank, sp_size in LocalTokenDispatcher is not ideal, but it's used in local_num_valid_tokens() , I want to share local_num_valid_tokens() implementation across AllToAllTokenDispather, DeepEP/HybridEP

[tianyu-l]

My impression is that this "virtual padding" idea would work when GroupedExperts is in a local region inside a DTensor global region, where DTensor handles the shard / all-gather.

After we migrate to spmd_types, the all-gather would require this metadata as well (spmd.redistribute takes this shape arg), with which spmd_types could also achieve similar "pad / unpad only around collectives" effect. O/w we still have to do "real padding" in model code. cc @pianpwk

I'm OK with this temporary solution (after cleanup) to unblock your vLLM + MoE work.

[wwwjn]

Yes let me clean up.

For spmd_types in MoE region, we don't need "real padding" in our current setup, just passing this "vitural padded shape" as metadata around.

torchtitan/torchtitan/models/common/MOE_SHARDING.md

Line 4 in da230bb

[`moe_sharding.py`](moe_sharding.py).

[wwwjn]

Close this because of #3595 . Let's move discussion there

[tianyu-l]

sounds right to me, one concrete issue before landing