perf: selective activation checkpointing feature support by seonjinn · Pull Request #2280 · NVIDIA-NeMo/RL

seonjinn · 2026-04-17T06:57:33Z

Supports Megatron-Core recompute_granularity and recompute_modules through policy config and megatron setup so training can selectively recompute activations (core_attn, moe, moe_act, etc.) instead of full checkpointing.

What does this PR do ?

Add a one line overview of what this PR aims to accomplish.

Issues

List issues that this PR closes (syntax):

Usage

You can potentially add a usage example below

# Add a code snippet demonstrating how to use this

Before your PR is "Ready for review"

Pre checks:

Make sure you read and followed Contributor guidelines
Did you write any new necessary tests?
Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

...

Plumbs Megatron-Core recompute_granularity and recompute_modules through policy config and megatron setup so training can selectively recompute activations (core_attn, moe, moe_act, etc.) instead of full checkpointing. Documents the new fields in the grpo_math_1B base config. Signed-off-by: sna <sna@nvidia.com>

copy-pr-bot · 2026-04-17T06:57:36Z

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.

terrykong

Review Summary

Thanks for adding selective activation checkpointing support — this is a valuable feature that exposes a real MCore knob MoE users need. The implementation is well-scoped and the branching logic is clean.

Suggestions

Config-conventions guideline (inline): The .get("recompute_granularity", "full") pattern uses a hidden non-None default, which the repo guidelines forbid. The key should be accessed directly.

Incomplete module list (inline): The recompute_modules comments list 3 valid options but MCore allows 7.

Validation gap (inline): Invalid recompute_modules values (typos) are silently accepted because MCore's __post_init__ validation doesn’t re-run after attribute assignment. Filed #2291 to track a broader refactor.

Performance Evidence

Since this is a perf:-labeled PR and the TypedDict comment claims "∼10–18GB savings for MoE models," could you share a before/after comparison? For example, peak GPU memory and tokens/sec with selective vs full recompute on a representative MoE model would help future users understand the expected benefit and source the claim. Filling in the PR template sections would also be helpful.

Test Coverage (nit)

The new selective branch and ValueError path have no unit test coverage. Consider adding tests for: (a) selective with custom modules, (b) selective with None modules (MCore default), (c) invalid granularity raises ValueError. Happy to share three pre-verified tests that match the existing MagicMock-based style.

Linter: PASS

All hooks passed (ruff, ruff-format, taplo, pyrefly, end-of-files, trailing-whitespace, minimize-check).

Generated by Claude Code

Co-authored-by: Terry Kong <terryk@nvidia.com> Signed-off-by: Seonjin <sna@nvidia.com>

seonjinn · 2026-04-20T01:15:55Z

/ok to 150ab01

seonjinn · 2026-04-23T21:41:59Z

/ok to test c7db964

seonjinn · 2026-04-24T00:28:44Z

@terrykong I've made changes based on your feedback. Does this PR look good to you ?

seonjinn · 2026-05-14T21:05:56Z

/ok to test 139e004

seonjinn · 2026-05-15T05:39:18Z

/ok to test afa7450

Signed-off-by: Seonjin <sna@nvidia.com>

seonjinn · 2026-05-20T19:02:37Z

/ok to test 05ca762

Drops PR #2280 selective recompute overrides (structurally infeasible at TP=4 EP=8 PP=8 on Qwen3-235B per job 11941263 OOM finding: selective scope strictly < full uniform N=1 memory savings). Isolates PR #2514 (max_tokens clamp) effect vs HybridEP+to60 baseline 11891376 (median 413.95s/step). Target: cut exposed_gen long-tail from clamp logic in vllm_worker_async + rollouts.py active set. Signed-off-by: seonjinn <sna@nvidia.com>

seonjinn · 2026-05-20T21:59:21Z

/ok to test 638d22a12fa8fee0a1c83f1c47a4cec5fd2c3e6e0

seonjinn · 2026-05-20T22:25:47Z

/ok to test 638d22a

Signed-off-by: Seonjin <sna@nvidia.com>

seonjinn · 2026-05-22T01:12:19Z

/ok to test 9075f7d

Signed-off-by: seonjinn <sna@nvidia.com>

seonjinn · 2026-05-22T02:12:04Z

/ok to test c4dc71e

seonjinn · 2026-05-22T04:39:50Z

@terrykong can you review this PR ?

terrykong

lgtm. just that one comment about the other exemplar configs

Reviewer feedback on #2280: exemplar configs should contain defaults for typing.Required/NotRequired attributes. Add recompute_granularity (default 'full') and recompute_modules (default null) to the five exemplar configs with megatron_cfg sections that were missing them: grpo_math_1B_megatron, sft, dpo, distillation_math, distillation_math_megatron. Signed-off-by: seonjinn <sna@nvidia.com>

seonjinn · 2026-05-22T05:41:00Z

/ok to test dd16310

The previous commit added recompute_granularity/recompute_modules to dpo, sft, and distillation_math exemplars. test_reference_configs_up_to_date requires the matching reference_configs to carry the same keys. Signed-off-by: seonjinn <sna@nvidia.com>

seonjinn · 2026-05-22T06:19:59Z

/ok to test b88ffab

seonjinn requested review from a team as code owners April 17, 2026 06:57

seonjinn changed the title ~~Add selective activation recompute support and perf_test configs~~ perf: selective activation checkpointing feature support Apr 17, 2026

seonjinn self-assigned this Apr 17, 2026

seonjinn requested a review from terrykong April 17, 2026 21:11

terrykong mentioned this pull request Apr 19, 2026

refactor: leverage MCore TransformerConfig.__post_init__ validation for config fields set in _apply_performance_config #2291

Open

terrykong reviewed Apr 19, 2026

View reviewed changes

Comment thread nemo_rl/models/megatron/setup.py

Comment thread nemo_rl/models/megatron/setup.py

Comment thread nemo_rl/models/policy/__init__.py Outdated

Comment thread examples/configs/grpo_math_1B.yaml Outdated

seonjinn and others added 4 commits April 19, 2026 18:13

Update nemo_rl/models/policy/__init__.py

2c20d69

Co-authored-by: Terry Kong <terryk@nvidia.com> Signed-off-by: Seonjin <sna@nvidia.com>

Update nemo_rl/models/megatron/setup.py

391eb9b

Co-authored-by: Terry Kong <terryk@nvidia.com> Signed-off-by: Seonjin <sna@nvidia.com>

Update examples/configs/grpo_math_1B.yaml

38f92e5

Co-authored-by: Terry Kong <terryk@nvidia.com> Signed-off-by: Seonjin <sna@nvidia.com>

Update nemo_rl/models/megatron/setup.py

e5acdd8

Co-authored-by: Terry Kong <terryk@nvidia.com> Signed-off-by: Seonjin <sna@nvidia.com>

Merge branch 'main' into sj/selective-recompute

72d1d15

seonjinn added the CI:Lfast Runs a fast test suite and re-use nightly `main` container (but sync dependencies to PRs version) label Apr 23, 2026

seonjinn added 2 commits April 28, 2026 17:16

Merge branch 'main' into sj/selective-recompute

79486ec

Merge remote-tracking branch 'origin/main' into sj/selective-recompute

8102a0d

copy-pr-bot Bot temporarily deployed to public May 14, 2026 21:06 Inactive

copy-pr-bot Bot had a problem deploying to nemo-ci May 14, 2026 21:06 Failure

copy-pr-bot Bot temporarily deployed to nemo-ci May 14, 2026 21:06 Inactive

copy-pr-bot Bot temporarily deployed to public May 14, 2026 21:06 Inactive

copy-pr-bot Bot temporarily deployed to public May 14, 2026 21:07 Inactive

Merge remote-tracking branch 'origin/main' into sj/selective-recompute

7e54943

copy-pr-bot Bot temporarily deployed to public May 15, 2026 05:39 Inactive

copy-pr-bot Bot temporarily deployed to public May 20, 2026 16:48 Inactive

copy-pr-bot Bot temporarily deployed to public May 20, 2026 16:49 Inactive

copy-pr-bot Bot had a problem deploying to nemo-ci May 20, 2026 16:49 Failure

copy-pr-bot Bot temporarily deployed to nemo-ci May 20, 2026 16:49 Inactive

copy-pr-bot Bot temporarily deployed to public May 20, 2026 16:52 Inactive

Merge branch 'main' into sj/selective-recompute

05ca762

Signed-off-by: Seonjin <sna@nvidia.com>

seonjinn force-pushed the sj/selective-recompute branch from bd4b112 to 05ca762 Compare May 20, 2026 19:02

copy-pr-bot Bot temporarily deployed to public May 20, 2026 19:02 Inactive

copy-pr-bot Bot temporarily deployed to nemo-ci May 20, 2026 19:03 Inactive

copy-pr-bot Bot temporarily deployed to public May 20, 2026 19:03 Inactive

Merge branch 'main' into sj/selective-recompute

638d22a

Merge branch 'main' into sj/selective-recompute

9075f7d

Signed-off-by: Seonjin <sna@nvidia.com>

Retry CI on L1 grpo_megatron_lora_async reward flake

c4dc71e

Signed-off-by: seonjinn <sna@nvidia.com>

terrykong requested changes May 22, 2026

View reviewed changes

terrykong approved these changes May 22, 2026

View reviewed changes

kajalj22 mentioned this pull request May 29, 2026

Nightly test llm_dpo_qwen2_5_math7b_1n8g_megatron_chunked_linear_ce_loss failing metric check #2626

Closed

Uh oh!

Conversation

seonjinn commented Apr 17, 2026

What does this PR do ?

Issues

Usage

Before your PR is "Ready for review"

Additional Information

Uh oh!

copy-pr-bot Bot commented Apr 17, 2026

Uh oh!

terrykong left a comment

Choose a reason for hiding this comment

Review Summary

Suggestions

Performance Evidence

Test Coverage (nit)

Linter: PASS

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

seonjinn commented Apr 20, 2026

Uh oh!

seonjinn commented Apr 23, 2026

Uh oh!

seonjinn commented Apr 24, 2026

Uh oh!

seonjinn commented May 14, 2026

Uh oh!

seonjinn commented May 15, 2026

Uh oh!

seonjinn commented May 20, 2026

Uh oh!

seonjinn commented May 20, 2026

Uh oh!

seonjinn commented May 20, 2026

Uh oh!

seonjinn commented May 22, 2026

Uh oh!

seonjinn commented May 22, 2026

Uh oh!

seonjinn commented May 22, 2026

Uh oh!

terrykong left a comment

Choose a reason for hiding this comment

Uh oh!

seonjinn commented May 22, 2026

Uh oh!

seonjinn commented May 22, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants