Conversation
meta-cla
Bot
added
the
CLA Signed
pytorch-bot
Bot
added
ciflow/fake-pg
| hidden_dim=hidden_dim, | ||
| num_experts=8, | ||
| top_k=4, | ||
| load_balance_coeff=1e-3, |
There was a problem hiding this comment.
did you find out which load balancing scheme gpt-oss use from literature?
| comm_backend=moe_comm_backend, | ||
| non_blocking_capacity_factor=non_blocking_capacity_factor, | ||
| ), | ||
| aux_loss_coeff=1e-3, |
There was a problem hiding this comment.
I think we also need to make this None for qwen
https://github.com/pytorch/torchtitan/blob/main/torchtitan/models/common/config_utils.py#L316
Actually default should be None to avoid mistake in the future.
|
|
||
| def _post_optimizer_build_fn(optimizers, model_parts, parallel_dims): | ||
| """Register step pre-hooks for load balancing and aux-loss accumulators.""" | ||
| register_moe_load_balancing_hook(optimizers, model_parts, parallel_dims) |
| elif aux_loss_type == "global_batch_wise": | ||
| aux_loss = GlobalBatchWiseLoadBalanceLoss.Config(coeff=aux_loss_coeff) |
There was a problem hiding this comment.
sounds unambiguous to just
| elif aux_loss_type == "global_batch_wise": | |
| aux_loss = GlobalBatchWiseLoadBalanceLoss.Config(coeff=aux_loss_coeff) | |
| elif aux_loss_type == "batch_wise": | |
| aux_loss = BatchWiseLoadBalanceLoss.Config(coeff=aux_loss_coeff) |
| ) | ||
|
|
||
|
|
||
| def _reduce_token_partials( |
There was a problem hiding this comment.
does it make sense to put them under AuxLoss for now?
| if self._cumulative_expert_counts_E is None: | ||
| self._cumulative_expert_counts_E = torch.zeros_like( | ||
| global_counts_E, dtype=torch.float32 | ||
| ) |
There was a problem hiding this comment.
may not be cudagraph friendly if we reuse cudagraph for every spmd (outer) grad accum? But I think we probably will move to all spmd grad accum steps captured in one big cudagraph anyway. cc @fegin
That said, any reason we don't allocate memory for the final shape upfront?
| "GlobalBatchWiseLoadBalanceLoss with multiple microbatches is " | ||
| "incompatible with activation checkpointing because recomputation " | ||
| "would advance its rolling expert counts. Disable activation " | ||
| "checkpointing or use MicrobatchWiseLoadBalanceLoss." |
There was a problem hiding this comment.
Sounds too big of a limitation, as models using this loss cannot be trained with AC anymore.
On the other hand, I believe torch_remat can work around this? @anijain2305
2 participants
microbatching supported with algo from https://arxiv.org/abs/2501.11873
Stack from ghstack (oldest at bottom):