Repository navigation
[NCCL EP] Support eager mode and overflow handling - #6229
Conversation
Signed-off-by: YangFei1990 <feiw@nvidia.com>
Signed-off-by: YangFei1990 <feiw@nvidia.com>
|
This PR has been automatically converted to draft because all PRs must start as drafts. When you are ready for review, click Ready for Review to begin the review process. This will:
See the contribution guide for more details. |
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
Signed-off-by: YangFei1990 <feiw@nvidia.com>
| if config.moe_pad_expert_input_to_capacity or config.moe_expert_capacity_factor is not None: | ||
| raise ValueError("drop and pad is not supported yet with the 'ncclep' backend.") | ||
|
|
||
| if self.eager: |
There was a problem hiding this comment.
So when we use zero copy now, do we still need to specify an appropriate recv-token-nums-per-rank? As I remember, BF16 used a symmetric memory pool, while FP8 required a fixed-size symmetric memory buffer. Why doesn’t BF16 zero-copy support eager mode either?
There was a problem hiding this comment.
We do not support zero-copy with eager mode at the moment due to the memory overhead introduced by the eager symm pool.
There was a problem hiding this comment.
Is this because the symmetric memory pool and the normal memory pool are independent of each other, so they cannot be reused simultaneously, leading to additional GPU memory usage?
There was a problem hiding this comment.
Yes symm buffer pool is independent with Pytorch normal pool. When using symm pool for eager, it will always be allocated to the outlier recv buffer size, which kind of defeats the purpose of eager.
There was a problem hiding this comment.
Double thinking, it is totally unacceptable for bf16 path, because the symm tensor will be saved per-layer per-mb as activations, so the pool will be super large. But for fp8 path, because only the quantized tensor will be save for backward, the symm tensor can be reused (in this case we allocate persistent tensor for graph mode), so the pool size probably can be acceptable, despite it is still allocated to the outlier recv size, but just for one tensor.
I think we should support eager mode zero copy with fp8 path, but nccl ep internal does not support it yet. I will work with the team for it later.
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| # Only after an actual rerun: rebootstrap NCCL EP | ||
| if num_tries > 1: | ||
| for mlp in self.moe_layers: | ||
| mlp.token_dispatcher.invalidate_ep_bootstrap() | ||
| nccl_ep_release_context() |
There was a problem hiding this comment.
[SUGGESTION Simplification] invalidate_ep_bootstrap() here is redundant with the nccl_ep_release_context() on the next line.
_ensure_bootstrap short-circuits on self._bootstrapped and is_nccl_ep_bootstrapped() (line 1617), so releasing the process-wide context is already sufficient to force every manager to rebuild — that is exactly what the new is_nccl_ep_bootstrapped() term in that condition buys you. The explicit per-layer invalidation adds a loop over every MoE layer for no behavioral change.
Either drop the loop, or keep it and note in the comment that it is belt-and-braces against a TE build where ep_finalize does not clear the bootstrap flag. As written the two mechanisms overlap without the reader being able to tell which one is load-bearing.
There was a problem hiding this comment.
No this is not redundant. nccl_ep_release_context is global, but invalidate_ep_bootstrap is per-layer. This is by design to have more fine-grained control.
| def check_required_capacity(self): | ||
| """Peak per-rank receive capacity the dispatcher needed, or None if it is not tracked.""" | ||
| if hasattr(self._comm_manager, 'required_recv'): | ||
| return self._comm_manager.required_recv | ||
| else: | ||
| return None | ||
|
|
||
| def grow_ep_recv_capacity(self, new_capacity: int) -> None: | ||
| """Grow the backend's receive budget to ``new_capacity`` so an overflowed step can be | ||
| replayed without dropping. No-op unless the backend supports it.""" | ||
| if hasattr(self._comm_manager, 'grow_recv_capacity'): | ||
| self._comm_manager.grow_recv_capacity(new_capacity) | ||
|
|
||
| def invalidate_ep_bootstrap(self) -> None: | ||
| """Force the backend to re-bootstrap on its next dispatch. Call after changing the mode: | ||
| every manager must re-derive it, not just the one that happens to dispatch first.""" | ||
| if hasattr(self._comm_manager, '_bootstrapped'): | ||
| self._comm_manager._bootstrapped = False |
There was a problem hiding this comment.
[SUGGESTION Simplification] These three new dispatcher methods each probe for a private attribute on _comm_manager with hasattr, which pushes backend capability detection into string literals.
hasattr(self._comm_manager, 'required_recv'), hasattr(..., 'grow_recv_capacity'), hasattr(..., '_bootstrapped') — the last is the worst of the three: it reaches into a private attribute of another object and assigns to it (self._comm_manager._bootstrapped = False). A rename of _bootstrapped inside _NCCLEPManager would make invalidate_ep_bootstrap() a silent no-op, and the overflow recovery would degrade to reusing a stale EP context with no error anywhere.
Prefer an explicit, overridable hook on the manager base class:
class _DispatchManager(ABC):
def invalidate_bootstrap(self) -> None:
"""No-op for backends without a bootstrapped context."""
def required_recv_capacity(self) -> Optional[torch.Tensor]:
return None
def grow_recv_capacity(self, new_capacity: int) -> None:
"""No-op for backends with no static receive budget."""Then the dispatcher just forwards, and adding a fourth backend gets the right default for free instead of relying on hasattr misses. This matches the existing check_over_budget pattern's problem rather than making it worse, so if you'd rather keep consistency with that, at minimum swap _bootstrapped for a public invalidate_bootstrap() method on _NCCLEPManager — mutating another class's private state across a module boundary is the part most likely to break quietly.
There was a problem hiding this comment.
This is to align with the existing code style of the flex dispatcher and comm manager
|
/ok to test b993cfe |
|
/claude strict-review |
svcnvidia-nemo-ci
added
Approved
|
/ok to test 2f67d31 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/31486939140 |
Add eager mode and overflow handling to the ncclEP flex dispatcher.
The ncclEP backend now selects between two modes based on
moe_expert_rank_capacity_factor: unset selects eager, where TE sizes the receive buffer per step from the actual received-token count; setting it selects the existing static path, which keeps a fixed budget and stays CUDA-graph capturable.In static mode the dispatcher enables NCCL EP's drop-on-overflow policy and accumulates over_budget and the peak required capacity device-side per dispatch, so the happy path never syncs — a MAX all-reduce runs only once overflow is reported. On overflow, PagedStashRunner grows the receive budget to that observed peak, rebuilds the EP context once, and replays the step, note this is different compared with hybrid where it use eager mode for rerun and fallback to same capacity afterwards.
Contribution process
Pre-checks
Code review
Feel free to message or comment @NVIDIA/mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!
All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.
Step 1: Mark PR as "Ready for Review"
.github/CODEOWNERS.Final Review might get declined if these requirements are not fulfilled.
Step 2: Final Review
For PRs that change
megatron/core, once all expert reviewers have approved, theFinal Reviewlabel is applied automatically and final reviewers are assigned.For PRs outside
megatron/core, this step is skipped.Step 3: Approved
Once all required reviewers have approved, the
Approvedlabel is applied automatically.Merge
Any member of mcore-engineers will be able to merge your PR.