[GTP] NCCL Symmetric Memory Registration Support - #6155
Conversation
Register the GTP all-gather buffers and the DDP param buffer that feeds them into NCCL symmetric-memory pools on the GTP remat group, so both ends of the gather sit in the same window and NCCL can select NVLS. Gated by --gtp-nccl-ub / --egtp-nccl-ub, independent of --use-nccl-ub, which covers the DP group only. Co-Authored-By: Claude <noreply@anthropic.com> Signed-off-by: Prajwal Singhania <psinghania@nvidia.com>
|
/claude strict-review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| default=False, help='Register dense GTP param buffers with NCCL symmetric memory on the GTP ' | ||
| 'group, independent of --use-nccl-ub (which covers the DP group).') | ||
| group.add_argument('--egtp-nccl-ub', action='store_true', dest='egtp_nccl_ub', | ||
| default=False, help='Like --gtp-nccl-ub but for routed-expert (EGTP) groups.') |
There was a problem hiding this comment.
[IMPORTANT Compatibility] The new --gtp-nccl-ub / --egtp-nccl-ub flags create torch.cuda.MemPool objects (via create_nccl_mem_pool(symmetric=True) in gtp_symm.get_gtp_pool and the new distopt branch in param_and_grad_buffer.py), but they bypass the compatibility guard that --use-nccl-ub gets in DistributedDataParallelConfig.__post_init__:
if self.nccl_ub and not is_torch_min_version("2.11.0a0"):
if 'expandable_segments:True' in os.getenv('PYTORCH_CUDA_ALLOC_CONF', '').split(','):
raise ValueError("PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True is currently not supported with nccl_ub ...")Why it matters: A user running --gtp-nccl-ub on torch < 2.11 with PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True hits the same underlying torch.cuda.MemPool incompatibility, but instead of the friendly ValueError gets a confusing low-level crash. Since the gtp_symm module docstring already tells users to export NVLS env vars, this combination is realistic.
Suggestion: Add an equivalent validation for the GTP flags (in validate_args, since these live in the arg namespace rather than the DDP config) that raises the same clear error when gtp_nccl_ub or egtp_nccl_ub is set with expandable_segments:True on an unsupported torch version.
| # Expert vs dense by group identity check | ||
| _expert_group = ProcessGroupCollection.use_mpu_process_groups( | ||
| required_pgs=["expt_gtp_remat"] | ||
| ).expt_gtp_remat |
There was a problem hiding this comment.
[SUGGESTION Simplification] _gtp_attach_attrs runs once per GTP param, and each call constructs a fresh ProcessGroupCollection via use_mpu_process_groups(required_pgs=["expt_gtp_remat"]) just to read the constant expert group and do an identity compare. This rebuilds the collection and re-reads the MPU global for every parameter in the model.
Why it matters: It's construction-time (not per-forward), so throughput is unaffected — but it's redundant object churn and an extra parallel_state global read per param. The expert group is invariant across the whole model build.
Suggestion: Resolve the expert group once (e.g. cache it in GTP_CONFIG at configure_gtp_remat_from_recipe time, where a ProcessGroupCollection is already materialized) and have _gtp_attach_attrs compare against the cached group. This also aligns with the CLAUDE.md guidance to prefer passing groups through over repeated use_mpu_process_groups reads in megatron/core.
Strict Review Summary — [GTP] NCCL Symmetric Memory Registration SupportFindings by severity: CRITICAL: 0 · IMPORTANT: 1 · SUGGESTION: 1 OverviewThis PR adds NCCL symmetric-memory / user-buffer registration for GTP all-gather and reduce-scatter, gated behind two off-by-default flags ( I traced the correctness-critical paths and found them sound:
Findings
Risk assessment: LowFeature is entirely opt-in and off by default; no change to existing training paths beyond the correctly-guarded shutdown deregister. The one IMPORTANT item is a UX/error-message concern, not a numerical or data-integrity risk. Note the PR checklist marks unit tests as not added — given the intricate buffer-lifecycle logic, a small unit test around |
| default=False, help='Register dense GTP param buffers with NCCL symmetric memory on the GTP ' | ||
| 'group, independent of --use-nccl-ub (which covers the DP group).') | ||
| group.add_argument('--egtp-nccl-ub', action='store_true', dest='egtp_nccl_ub', | ||
| default=False, help='Like --gtp-nccl-ub but for routed-expert (EGTP) groups.') |
There was a problem hiding this comment.
[IMPORTANT Compatibility] The new --gtp-nccl-ub / --egtp-nccl-ub flags create torch.cuda.MemPool objects (via create_nccl_mem_pool(symmetric=True) in gtp_symm.get_gtp_pool and the new distopt branch in param_and_grad_buffer.py), but they bypass the compatibility guard that --use-nccl-ub gets in DistributedDataParallelConfig.__post_init__:
if self.nccl_ub and not is_torch_min_version("2.11.0a0"):
if 'expandable_segments:True' in os.getenv('PYTORCH_CUDA_ALLOC_CONF', '').split(','):
raise ValueError("PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True is currently not supported with nccl_ub ...")Why it matters: A user running --gtp-nccl-ub on torch < 2.11 with PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True will hit the same underlying torch.cuda.MemPool incompatibility, but instead of the friendly ValueError they get a confusing low-level crash. The --gtp-nccl-ub docs elsewhere in this MR also recommend NVLS env vars, so this combination is realistic.
Suggestion: Add an equivalent validation for the GTP flags (in validate_args / arguments.py, since these live in the arg namespace rather than the DDP config) that raises the same clear error when gtp_nccl_ub or egtp_nccl_ub is set with expandable_segments:True on an unsupported torch version.
| # Expert vs dense by group identity check | ||
| _expert_group = ProcessGroupCollection.use_mpu_process_groups( | ||
| required_pgs=["expt_gtp_remat"] | ||
| ).expt_gtp_remat |
There was a problem hiding this comment.
[SUGGESTION Simplification] _gtp_attach_attrs runs once per GTP param, and each call constructs a fresh ProcessGroupCollection via use_mpu_process_groups(required_pgs=["expt_gtp_remat"]) just to read the constant expert group and compare identity. This rebuilds the collection object and re-reads the MPU global for every parameter in the model.
Why it matters: It's construction-time (not per-forward), so throughput is unaffected, but it's redundant object churn and an extra parallel_state global read per param. The expert group is invariant across the whole model build.
Suggestion: Resolve the expert group once (e.g. cache it in GTP_CONFIG alongside the egtp_nccl_ub/gtp_nccl_ub flags at configure_gtp_remat_from_recipe time, where the collection is already materialized) and have _gtp_attach_attrs compare against that cached group. This also matches the CLAUDE.md guidance to prefer passing the group through over repeated use_mpu_process_groups reads in core code.
2 participants
What does this PR do?
Adds NCCL symmetric memory/user buffer registration support for GTP
Summary
Symmetric memory registration allows for the use of faster LSA/GIN based and CE based collectives in NCCL to improve all-gather and reduce-scatter communication performance for GTP. To enable these new kernels, NCCL requirements state:
Based on these requirements, this MR enables symmetric memory backed tensors gated by two flags:
--use-nccl-uband off by default.Design
One
ncclMemAlloc-backedtorch.cuda.MemPoolper GTP process group, registered once withregister_mem_pool(pool, symm=True). PyTorch's ProcessGroupNCCL segment hook then auto-registersevery later allocation made under
gtp_mem_pool_ctx(group), so the context manager itself iscollective-free and capture-safe.
Every buffer either collective touches, where it comes from, and what puts it in the pool:
--fp8-param-gatherparam_dataregister_ddp_buffers_on_gtp_groups_gtp_slice_one_paramGTPWeightCacheticket buffer_allocate_buffer(rs=False)under the pool contextRegisteredLifoPool, padded full shapealloc()under the pool contextGTPWeightCacheticket buffer_graphed_alloc, stays out of the windowContribution process
Pre-checks