NVIDIA / NVIDIA/Megatron-LM

pad_routing_map may return unaligned expert token counts

Open
#6,086 3 comments 0 reactions 1 assignee Claimed by @YangFei1990 View on GitHub
bug community-request waiting-on-maintainers
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.