fix: Fix non-colocated refit when vLLM model parallel size is larger than 8

[youngeunkwon0405]

What does this PR do ?

See #1352 for the details.

Issues

List issues that this PR closes (syntax):

Closes #1352

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

• ...

Summary by CodeRabbit

• Bug Fixes
  • Improves reliability of multi-node inference when running non-colocated with cross-node model parallelism by automatically adjusting communication settings.
  • Adds an informational log to clarify when this configuration is applied.
  • Applies the adjustment only in relevant scenarios to avoid affecting colocated runs or other setups.
  • No changes to public interfaces; behavior is transparent and aims to prevent communication issues during distributed generation.

[coderabbitai]

📝 Walkthrough

Walkthrough

Introduces a conditional in vllm_generation.py to adjust NCCL environment variables during non-colocated, cross-node model-parallel inference: when NCCL_CUMEM_ENABLE is active and cross-node parallelism is required with colocated disabled, it sets NCCL_NVLS_ENABLE="0" and logs an informational message.

Changes

Cohort / File(s)	Summary of Changes
vLLM generation: cross-node non-colocated handling nemo_rl/models/generation/vllm/vllm_generation.py	Added conditional branch inside NCCL_CUMEM_ENABLE block to detect needs_cross_node_parallelism with colocated disabled and set NCCL_NVLS_ENABLE="0"; logs info. No API/signature changes.

Sequence Diagram(s)

sequenceDiagram
    autonumber
    participant Runner as Inference Runner
    participant VLLMGen as vllm_generation.py
    participant Env as OS Env Vars
    participant Log as Logger

    Runner->>VLLMGen: initialize_inference(...)
    activate VLLMGen
    VLLMGen->>Env: read NCCL_CUMEM_ENABLE
    alt NCCL_CUMEM_ENABLE == "1"
        VLLMGen->>VLLMGen: check colocated == False
        VLLMGen->>VLLMGen: check needs_cross_node_parallelism == True
        alt non-colocated and cross-node
            VLLMGen->>Env: set NCCL_NVLS_ENABLE="0"
            VLLMGen->>Log: info("Disabling NCCL NVLS for cross-node non-colocated")
        else otherwise
            Note over VLLMGen: No changes to NVLS
        end
    else
        Note over VLLMGen: No CUMEM block executed
    end
    VLLMGen-->>Runner: proceed with inference setup
    deactivate VLLMGen

Loading

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

Suggested reviewers

• yuki-97
• guyueh1

Pre-merge checks and finishing touches

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Test Results For Major Changes	⚠️ Warning	The change introduces a new code path that disables NCCL’s NVLS optimization for non-colocated cross-node inference, which can significantly impact distributed performance, yet the PR description contains no unit or integration test results nor any before-and-after performance measurements to demonstrate correctness or lack of regression.	Please add unit and integration tests that cover the non-colocated cross-node scenario to verify NCCL_NVLS_ENABLE behavior, and include before-and-after performance or convergence benchmarks (with configuration details) to demonstrate that disabling NVLS does not introduce regressions.

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title Check	✅ Passed	The pull request title "fix: Fix non-colocated refit when vLLM model parallel size is larger than 8" directly and specifically describes the main change being made. The title aligns with the code modification, which adds handling for cross-node model parallelism in non-colocated scenarios, and corresponds to the primary objective of resolving issue #1352 where unified placement with model parallel > 8 was causing crashes. The title is concise, clear, and provides enough context for someone scanning the history to understand the purpose of the change.
Linked Issues Check	✅ Passed	The pull request directly addresses the primary objective from linked issue #1352, which reports that unified placement with vLLM model parallel > 8 causes an error in the non-colocated code path during pynccl group initialization. The code changes add conditional handling for cross-node model parallelism by setting NCCL_NVLS_ENABLE to "0" when needed, which aligns with preventing the crash scenario described in the issue. The modification is targeted and only applies to non-colocated runs with cross-node parallelism requirements, which matches the scope of the reported bug.
Out of Scope Changes Check	✅ Passed	The pull request modifies only a single file (nemo_rl/models/generation/vllm/vllm_generation.py) with changes strictly limited to adding conditional handling for cross-node model parallelism within an existing NCCL_CUMEM_ENABLE block. The changes do not alter function signatures, introduce unrelated refactoring, or modify exported entities. All modifications are directly scoped to resolving the non-colocated refit issue with model parallel > 8, with no apparent extraneous changes.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

---

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: 2

📜 Review details

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 0a769cc and ff70518.

📒 Files selected for processing (1)

• nemo_rl/models/generation/vllm/vllm_generation.py (1 hunks)

🧰 Additional context used

📓 Path-based instructions (2)

**/*.py

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

**/*.py: Follow the Google Python Style Guide for all Python code
Target Python 3.12+ for all Python code in NeMo-RL
Indent Python code with 4 spaces; do not use tabs
Python filenames should be snake_case (e.g., some_file.py)
Class names should be PascalCase
Function and method names should be snake_case
Local variable names should be snake_case; if starting with a number, prefix with k (e.g., k_99th_percentile)
Global variables should be UPPER_SNAKE_CASE and prefixed with G_ (e.g., G_MY_GLOBAL)
Constants should be UPPER_SNAKE_CASE
Avoid shadowing variables declared in an outer scope
Initialize all externally visible members of a class in the constructor
For public interfaces used outside a file, prefer docstrings over comments
Use comments mainly for code within a function or interfaces local to a file
Commented-out code must include a nearby comment explaining usage and why it is commented out; otherwise remove before merging
Use Google-style docstrings for classes and functions (Sphinx-parseable)
Avoid using reflection when functionality can be easily achieved without it
Limit except clauses to the smallest specific set of exceptions possible
For duck-typing via try/except, keep the try body minimal and use else for main logic
Add the NVIDIA copyright header (with current year) at the top of all Python files, excluding tests/ and test-only scripts

Files:

• nemo_rl/models/generation/vllm/vllm_generation.py

nemo_rl/**/*.py

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

nemo_rl/**/*.py: Do not set non-None configuration defaults in code; YAML is the single source of truth for defaults
Access required config attributes directly (e.g., policy_cfg["precision"]) and assume presence; do not introduce hidden defaults
Express configuration optionality via TypedDict using typing.NotRequired
When adding a new config key to a TypedDict subclass, document the key’s purpose, valid values/types, and recommended default in code
For any class or function decorated with @ray.remote, add '# pragma: no cover' on the class/def line (and on remote functions)

Files:

• nemo_rl/models/generation/vllm/vllm_generation.py

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (6)

• GitHub Check: sphinx-build / Build docs
• GitHub Check: Lint check
• GitHub Check: Lint check
• GitHub Check: Lint check
• GitHub Check: Post automodel integration comment / Comment on PR
• GitHub Check: Post submodule check comment / Comment on PR

[yuki-97]

Hi @youngeunkwon0405 , thanks for the fix!

I'm not sure if the colocated path has the same issue with cross node parallel and needs the same fix.
Do you mind also checking it?

[youngeunkwon0405]

Hi @youngeunkwon0405 , thanks for the fix!

I'm not sure if the colocated path has the same issue with cross node parallel and needs the same fix. Do you mind also checking it?

I also checked for colocated case. I was able to see the same transport/nvls.cc:598 NCCL WARN Cuda failure 1 'invalid argument' error.
I modified the PR to cover the colocated case as well.