chore: add assert for dtensor v2 cpu offload

[yuki-97]

For world size 1, FSDP uses extra memory which is not necessary, so AutoModel disable it for now. This PR add an assert for it.

Summary by CodeRabbit

• Bug Fixes

  • Added validation to prevent unsupported CPU offload configuration on single-GPU setups.

• Improvements

  • Enhanced distributed setup configuration for improved generation performance when resources are not colocated.

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

[coderabbitai]

📝 Walkthrough

Walkthrough

Modified environment configuration for non-colocated generation by setting NCCL_CUMEM_ENABLE, and added validation to prevent CPU offload usage on single-GPU environments during distributed setup initialization.

Changes

Cohort / File(s)	Summary
Environment and Validation Setup nemo_rl/models/automodel/setup.py	Set NCCL_CUMEM_ENABLE environment variable when generation is not colocated with rationale documentation; added guard in setup_distributed to raise NotImplementedError if cpu_offload is enabled on single-GPU setups (world_size == 1)

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~8 minutes

🚥 Pre-merge checks | ✅ 3 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Title check	⚠️ Warning	The title mentions 'assert for dtensor v2 cpu offload' but the actual changes add a NotImplementedError guard (not an assert) and modify NCCL_CUMEM_ENABLE environment variable setup, which is not mentioned in the title.	Update the title to accurately reflect the main changes, such as 'Prevent single-GPU CPU offload in setup_distributed and configure NCCL_CUMEM_ENABLE' or similar.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Test Results For Major Changes	✅ Passed	PR contains only minor changes (7 lines added, 1 removed) classified as chore: setting environment variable and adding guard clause for unsupported configuration.

_{✏️ 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.}

[coderabbitai]

Actionable comments posted: 1

🤖 Fix all issues with AI agents

In `@nemo_rl/models/automodel/setup.py`:
- Around line 286-290: The NotImplementedError raised when cpu_offload is True
contains an incorrect GitHub URL; update the error message in the cpu_offload
guard (the NotImplementedError raised near cpu_offload and before
_setup_distributed()) to point to the correct repository URL (e.g.,
https://github.com/NVIDIA-NeMo/RL or the canonical
https://github.com/NVIDIA/NeMo) so the message directs users to the right issue
tracker.