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

[wwwjn]

Stack from ghstack (oldest at bottom):

• [WIP] Enable DP+EP for MoE inference in vLLM wrapper #3236
• [Not ready][rl] Enable TP2EP for unified MoE model in vLLM wrapper #3142
• -> Using "virtual padding" to calculate number_local_tokens per SP rank, and fix combine() shape mismatch #3595

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

[tianyu-l]

sounds right to me, one concrete issue before landing

[tianyu-l]

checking these in post_init is not safe, because CLI can override -- we have to do these check in update_from_config

[wwwjn]

I see, this is valid. Even TrainingConfig is not used in PolicyTrainer today, a user could also mistakenly override by CLI --trainer.training.seq_len to hack that.

Then the problem is passing BatchConfig from controller, into PolicyTrainer's TrainingConfig(), then calling self.model.update_from_config(config). So my updated plan is:

• remove the check in post_init
• Pass PolicyTrainer.TrianingConfig.seq_len to be BatchConfig.seq_len after parsing CLI override here
• Then we are good to only check in Decoder.update_config()

[tianyu-l]

Pass PolicyTrainer.TrianingConfig.seq_len to be BatchConfig.seq_len after parsing CLI override here

As discussed earlier, PolicyTrainer.TrainingConfig shouldn't have seq_len. In fact, for pretaining, we should have BatchConfig in Dataloader.Config, not in Trainer.TrainingConfig

[wwwjn]

I see, that's the right direction, let me add a TODO and also a github issue to track this

[pianpwk]

sorry, I saw this in the closed PR:

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.

what's the shape arg mentioned here?

[wwwjn]

sorry, I saw this in the closed PR:

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.

what's the shape arg mentioned here?

num_local_tokens_after_padding = (T + pad_tokens) // sp_size

This parameter here, num_local_tokens_after_padding we are passing this to combine() and doing a virtual padding ,

[pianpwk]

This parameter here, num_local_tokens_after_padding we are passing this to combine() and doing a virtual padding ,

oh my question was, redistribute() takes this?

[tianyu-l]

nice! some nit comments

[tianyu-l]

accidental?

[wwwjn]

CPU test failing might be cause of spmd_types version, and seems not related this PR
I rerun the test locally

• torch 2.13.0.dev20260609+cpu
• spmd_types==0.2.1

Command result:

CUDA_VISIBLE_DEVICES= titan-rl/bin/python -m pytest tests/unit_tests/
Result: 1 passed in 10.50s.