Feature/amem nccl rl by CodersAcademy006 · Pull Request #2979 · NVIDIA/Megatron-LM

CodersAcademy006 · 2026-01-16T10:43:03Z

RL: Optional AMem NCCL offload support for memory-efficient rollouts

Motivation

Reinforcement learning workflows in Megatron-LM often alternate between training and inference-heavy rollout phases. During rollout, a significant amount of GPU memory can be consumed by NCCL-allocated buffers that are idle from a compute perspective but cannot be reclaimed.

This PR introduces optional support for the AMem NCCL plugin to allow NCCL-allocated GPU memory to be transparently offloaded during RL rollout phases and restored when needed. The goal is to reduce peak GPU memory pressure during inference-dominated steps without changing model semantics or training behavior.

What this PR does

Adds an optional AMem NCCL integration for RL training paths.
Introduces explicit hooks to pause NCCL memory usage during rollout and resume it before training resumes.
Integrates with the native RL execution path (perform-rl-step).
Exposes opt-in CLI flags to control behavior:
- Enable or disable AMem NCCL support
- Configure AMem process group ID
- Control whether offloading occurs during rollout
Keeps the feature fully disabled by default and backward compatible.

Design principles

Opt-in only: No behavior changes unless explicitly enabled.
No semantic changes: Model execution, gradients, and optimizer behavior remain unchanged.
Minimal surface area: Reuses existing RL control flow and initialization paths.
Fail-safe: If AMem is unavailable or misconfigured, training proceeds normally without offloading.

Expected impact

Reduced peak GPU memory usage during RL rollout and inference phases.
Increased memory headroom for larger batch sizes or longer sequences in RL workloads.
No impact on non-RL training paths.

Scope limitations

This PR does not introduce new checkpointing, scheduling, or allocator behavior.
This PR does not generalize AMem usage beyond RL workflows.
Performance characteristics depend on the underlying AMem NCCL implementation and system configuration.

Testing

Validated on RL training flows with AMem enabled and disabled.
Confirmed that disabling AMem restores existing behavior.
Exercised error handling paths when AMem is unavailable.

copy-pr-bot · 2026-01-16T10:43:07Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

HaochenYuan · 2026-01-16T16:15:06Z

The changes in megatron/core/transformer/moe/grouped_gemm_util.py don't seem to be related to amem.

CodersAcademy006 · 2026-01-17T08:23:39Z

Hey @HaochenYuan, I have made the relevant changes. Please tell me if anything else has to be changed. Thank you for the review!!

Also i have made a follow-up pr for the rest of the changes #2988 please review that as well.

HaochenYuan · 2026-01-19T11:40:33Z

            toggle_cuda_graphs(lang_module, 'none', reset_cuda_graphs=reset_cuda_graphs)

+        # Restore NCCL memory after inference if AMem is enabled
+        if use_amem and amem_offload_during_rollout:


Where is use_amem defined?

HaochenYuan · 2026-01-19T11:42:33Z

+        # Offload NCCL memory before inference if AMem is enabled
+        if use_amem and amem_offload_during_rollout:
+            with nvtx_range("amem-nccl-pause-before-inference"):
+                if amem_nccl.nccl_pause():
+                    logger.info(f"[{dist.get_rank()}:DP] Successfully offloaded NCCL memory")
+                else:
+                    logger.warning(f"[{dist.get_rank()}:DP] Failed to offload NCCL memory")
+


The inference process seems to be the same as training's, if you offload nccl memory before inference while the inference process also needs the comm groups, there may be runtime error.

HaochenYuan · 2026-01-19T11:46:00Z

+- Model parameters can be temporarily offloaded without performance penalty
+
+AMem enables:
+- Offloading idle model parameters during rollout


Amem cannot offload parameters, it can only offload memory used by nccl.

HaochenYuan · 2026-01-19T11:57:54Z

Before proceeding, could you confirm if you have tested this change with a compiled amem library?

During some earlier beta testing, I observed that amem is currently only able to release NCCL memory for communication groups that are created via the loaded libamem_nccl.so libnccl.so. However, in Megatron, most communication groups are created directly through torch.distributed.new_group, which may fall outside its current memory release scope.

To better evaluate this PR, it would be very helpful if you could share:

The released memory size observed with specified model configuration and parallel setup used for testing/
If you have tested it , can you provide the libamem_nccl.so libnccl.so & tested environment config?

Thx!

- Move nv-grouped-gemm from dev/lts extras to new moe extra - Users can now install megatron-core[dev] without build failures - Add comprehensive error messages with installation instructions - Update README with MoE dependencies documentation This resolves the installation failure when nv-grouped-gemm cannot build due to missing CUTLASS headers. Users who need MoE with grouped GEMM can now explicitly install it with megatron-core[moe], while others can install dev/lts extras without encountering build errors. Fixes NVIDIA#2541

…r RL workflows. Fixed imports, ensured native RL loop support, and validated no errors in touched files. Ready for review and CI on supported platforms.

…lag, gate all logic on --rl-amem-offload-during-rollout, encapsulate env setup, and update docs/examples per review

Add comprehensive design document explaining: - When AMem hooks are enabled (RL rollout phase) - Why RL workloads benefit from memory offloading - What memory is offloaded (model params, optimizer states) - Implementation architecture across core, RL, and training modules - Non-goals and future optimization opportunities This document will accompany the AMem NCCL integration PR to provide reviewers with context on the feature's intent and scope.

CodersAcademy006 · 2026-03-05T05:03:33Z

@HaochenYuan — thanks for the detailed feedback. I've tested the fallback
paths (AMem unavailable, env vars not set, graceful degradation) and the
argument parsing, but I don't currently have access to hardware that supports
NCCL_CUMEM_ENABLE to validate the full integration against a compiled
libamem_nccl.so.

Given your point about torch.distributed.new_group comm groups potentially
falling outside AMem's release scope, I'd like to understand that constraint
better before pushing further. Would it make sense to convert this to a draft
PR until I can get access to validate it properly, or is there a test
environment within NVIDIA I could coordinate with to verify this?

I'll fix the design doc regardless — AMem offloads NCCL memory, not model
parameters. That was wrong.

Sorry for being late.

HaochenYuan · 2026-03-05T06:31:42Z

@HaochenYuan — thanks for the detailed feedback. I've tested the fallback paths (AMem unavailable, env vars not set, graceful degradation) and the argument parsing, but I don't currently have access to hardware that supports NCCL_CUMEM_ENABLE to validate the full integration against a compiled libamem_nccl.so.

Given your point about torch.distributed.new_group comm groups potentially falling outside AMem's release scope, I'd like to understand that constraint better before pushing further. Would it make sense to convert this to a draft PR until I can get access to validate it properly, or is there a test environment within NVIDIA I could coordinate with to verify this?

I'll fix the design doc regardless — AMem offloads NCCL memory, not model parameters. That was wrong.

Sorry for being late.

Hi @CodersAcademy006, since NCCL v2.29.7 release has supported dynamic memory offload, the Amem integration maybe not be necessary now. Thx for your contribution!

CodersAcademy006 · 2026-03-05T10:14:25Z

@HaochenYuan thanks for your explanation and review. If there is any other issue to work on, i would be happy to take that. Thank You.

CodersAcademy006 requested review from a team and pablo-garay as code owners January 16, 2026 10:43

ko3n1g requested a review from a team January 16, 2026 10:43

github-actions Bot added the community-request label Jan 16, 2026

CodersAcademy006 mentioned this pull request Jan 16, 2026

fsdp: avoid double sharding of MoE experts when EP is enabled #2833

Closed

shjwudp requested a review from HaochenYuan January 16, 2026 12:11

shjwudp added the Expert Review [deprecated] Apply this label to indicate that your PR is ready for expert review. label Jan 16, 2026

CodersAcademy006 force-pushed the feature/amem-nccl-rl branch from bcb49bf to 96cec4d Compare January 17, 2026 08:22

chtruong814 added the needs-follow-up Issue needs follow-up label Jan 19, 2026

HaochenYuan reviewed Jan 19, 2026

View reviewed changes

chtruong814 removed the needs-follow-up Issue needs follow-up label Jan 19, 2026

CodersAcademy006 added 4 commits March 5, 2026 04:57

Resolve merge conflicts and ensure AMem NCCL integration is robust fo…

292a375

…r RL workflows. Fixed imports, ensured native RL loop support, and validated no errors in touched files. Ready for review and CI on supported platforms.

Refactor AMem NCCL integration: remove redundant --rl-use-amem-nccl f…

547e8e5

…lag, gate all logic on --rl-amem-offload-during-rollout, encapsulate env setup, and update docs/examples per review

CodersAcademy006 force-pushed the feature/amem-nccl-rl branch from 96cec4d to 094628e Compare March 5, 2026 05:04

resolve conflicts

0c33c71

CodersAcademy006 closed this Mar 5, 2026

Navigation Menu

Uh oh!

Feature/amem nccl rl - #2979

Feature/amem nccl rl#2979
CodersAcademy006 wants to merge 5 commits into
NVIDIA:mainfrom
CodersAcademy006:feature/amem-nccl-rl

CodersAcademy006 commented Jan 16, 2026

Uh oh!

copy-pr-bot Bot commented Jan 16, 2026

Uh oh!

HaochenYuan commented Jan 16, 2026

Uh oh!

CodersAcademy006 commented Jan 17, 2026 •

edited

Loading

Uh oh!

HaochenYuan Jan 19, 2026

Uh oh!

HaochenYuan Jan 19, 2026

Uh oh!

HaochenYuan Jan 19, 2026

Uh oh!

HaochenYuan commented Jan 19, 2026

Uh oh!

CodersAcademy006 commented Mar 5, 2026

Uh oh!

HaochenYuan commented Mar 5, 2026

Uh oh!

CodersAcademy006 commented Mar 5, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

Uh oh!

Conversation

CodersAcademy006 commented Jan 16, 2026

RL: Optional AMem NCCL offload support for memory-efficient rollouts

Motivation

What this PR does

Design principles

Expected impact

Scope limitations

Testing

Uh oh!

copy-pr-bot Bot commented Jan 16, 2026

Uh oh!

HaochenYuan commented Jan 16, 2026

Uh oh!

CodersAcademy006 commented Jan 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

HaochenYuan Jan 19, 2026

Choose a reason for hiding this comment

Uh oh!

HaochenYuan Jan 19, 2026

Choose a reason for hiding this comment

Uh oh!

HaochenYuan Jan 19, 2026

Choose a reason for hiding this comment

Uh oh!

HaochenYuan commented Jan 19, 2026

Uh oh!

CodersAcademy006 commented Mar 5, 2026

Uh oh!

HaochenYuan commented Mar 5, 2026

Uh oh!

CodersAcademy006 commented Mar 5, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

CodersAcademy006 commented Jan 17, 2026 •

edited

Loading