pad_routing_map may return unaligned expert token counts
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 6h
- Merged PRs (30d)
- 271
Description
# `pad_routing_map` may return unaligned expert token counts
**Describe the bug**
`pad_routing_map` and `fused_pad_routing_map` are intended to pad each expert's token count to a multiple of `pad_multiple`.
However, both functions preserve the routing-map shape `[num_tokens, num_experts]` and pad by converting existing zero entries to one. If the next aligned count exceeds `num_tokens`, there are not enough zero entries to achieve the requested alignment.
In this case, the functions silently return an unaligned token count.
Affected code:
* https://github.com/NVIDIA/Megatron-LM/blob/4b18b260f012c8de51f729fb09771f99266bc675/megatron/core/transformer/moe/moe_utils.py#L675-L707
* https://github.com/NVIDIA/Megatron-LM/blob/4b18b260f012c8de51f729fb09771f99266bc675/megatron/core/fusions/fused_pad_routing_map.py#L31-L100
**Steps/Code to reproduce bug**
```python
import torch
from megatron.core.transformer.moe.moe_utils import pad_routing_map
routing_map = torch.zeros(
(100, 2),
dtype=torch.bool,
device="cuda",
)
routing_map[:90, 0] = True
padded_map = pad_routing_map(routing_map, pad_multiple=128)
tokens_per_expert = padded_map.sum(dim=0)
expert_0_tokens = int(tokens_per_expert[0].item())
print("tokens_per_expert:", tokens_per_expert)
print("expert 0 tokens:", expert_0_tokens)
print("expert 0 remainder:", expert_0_tokens % 128)
assert expert_0_tokens % 128 == 0, (
f"Expected expert token count to be divisible by 128, "
f"but got {expert_0_tokens}"
)
```
Actual output:
```text
tokens_per_expert: tensor([100, 0], device='cuda:0')
expert 0 tokens: 100
expert 0 remainder: 100
Traceback (most recent call last):
File "", line 21, in
AssertionError: Expected expert token count to be divisible by 128, but got 100
```
**Expected behavior**
The padding path should not silently return a routing map that does not satisfy the requested alignment.
When the target aligned count exceeds the routing-map capacity, this condition should be detected and reported consistently by the fused and non-fused paths.
The DeepEP path already performs this capacity check and emits a warning when routing-map padding is not possible. The standard all-to-all path currently has no equivalent check.
Contributor guide
Assessment
This issue has not been assessed yet.