[sync] Fine-grained activation offloading

[ananthsub]

What does this PR do ?

Sync with changes from NVIDIA/Megatron-LM#1913
Enable CPU affinity when fine-grained activation offloading is enabled

Changelog

• Enable CPU affinity when fine-grained activation offloading is enabled

GitHub Actions CI

See the CI sectionin the Contributing doc for how to trigger the CI. A Nvidia developer will need to approve and trigger the CI for external contributors.

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 add or update any necessary documentation?
• Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  • Reviewer: Does the PR have correct import guards for all optional libraries?

If you haven't finished some of the above items you can still open "Draft" PR.

Additional Information

• Related to # (issue)

Summary by CodeRabbit

• New Features

  • Enhanced validation for fine-grained activation offloading to ensure compatibility with Transformer Engine and environment requirements.
  • Automatic CPU affinity configuration for GPU when fine-grained activation offloading is enabled.

• Tests

  • Added comprehensive unit tests for fine-grained activation offloading validation.

_{✏️ Tip: You can customize this high-level summary in your review settings.}

[copy-pr-bot]

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.

[ananthsub]

/ok to test fe3d2b2

[coderabbitai]

📝 Walkthrough

Walkthrough

This change introduces fine-grained activation offloading validation and runtime setup. A new validator ensures the feature is only enabled with transformer_engine and meets version-specific Transformer Engine requirements (NVTE_CPU_OFFLOAD_V1 environment variable for TE ≥ 2.10.0). During setup, CPU affinity is configured when the feature is active. Unit tests validate all constraint scenarios.

Changes

Cohort / File(s)	Summary
Fine-grained Activation Offloading Validation src/megatron/bridge/training/config.py	Added private validator _validate_fine_grained_activation_offloading() that enforces transformer_engine requirement and checks for NVTE_CPU_OFFLOAD_V1 environment variable when TE ≥ 2.10.0. Integrated into ConfigContainer.validate().
Runtime Setup for Fine-grained Offloading src/megatron/bridge/training/setup.py	Added conditional block to invoke set_ideal_affinity_for_current_gpu() when fine_grained_activation_offloading is enabled, with lazy import to avoid unnecessary dependencies.
Test Coverage tests/unit_tests/training/test_config.py	Added five unit tests covering validation scenarios: transformer_engine requirement, TE version checks, NVTE_CPU_OFFLOAD_V1 environment variable enforcement for TE ≥ 2.10.0, and disabled-state behavior.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

🚥 Pre-merge checks | ✅ 3 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Test Results For Major Changes	⚠️ Warning	PR introduces major changes for fine-grained activation offloading with new configuration validation and CPU affinity setup, but lacks documented test results in the PR description despite unit tests existing in code.	Update PR description to document unit tests added, confirm test pass status, and provide performance impact analysis for CPU affinity setup changes.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title '[sync] Fine-grained activation offloading' directly reflects the main change: syncing fine-grained activation offloading support from upstream Megatron-LM, including validation logic and CPU affinity setup.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing touches

• 📝 Generate docstrings

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}