feat: Automodel init for DTensorPolicyV2 #1509
Conversation
Signed-off-by: adil-a <adil.asif2000@hotmail.com>
Signed-off-by: adil-a <adil.asif2000@hotmail.com>
Signed-off-by: root <root@pool0-01523.cm.cluster>
Signed-off-by: adil-a <adil.asif2000@hotmail.com>
Signed-off-by: root <root@pool0-01749.cm.cluster>
|
adil-a
added
the
CI:L2
📝 WalkthroughWalkthroughRefactors FSDP2 initialization in the policy worker to use FSDP2Manager instead of manual device-mesh setup, adds cpu_offload handling, exposes manager-derived mesh attributes, implements dynamic attention implementation selection based on configuration, and updates gradient norm computation paths. Changes
Sequence Diagram(s)sequenceDiagram
participant Worker as Policy Worker
participant Manager as FSDP2Manager
participant Model
participant Training
Worker->>Manager: Initialize with cpu_offload config
activate Manager
Manager->>Manager: Create device mesh<br/>(CUDA+CPU backend if offloading)
Manager->>Manager: Configure offload policy
Manager-->>Worker: Return mesh attributes<br/>(dp_mesh, tp_mesh, cp_mesh)
deactivate Manager
Worker->>Worker: Select attn_implementation<br/>(flash_attention_2 vs sdpa)<br/>based on seq_packing & context_parallel
Worker->>Model: Inject attn_implementation<br/>into model_config
Worker->>Manager: parallelize(model)
activate Manager
Manager->>Model: Apply FSDP2 sharding
Manager-->>Worker: Model parallelized
deactivate Manager
Worker->>Training: Store mesh references<br/>(dp_shard_cp_mesh for grad norm)
Training->>Training: Use dp_shard_cp_mesh<br/>for gradient computation
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60–75 minutes
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (3 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
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. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
📜 Review details
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
nemo_rl/models/policy/dtensor_policy_worker_v2.py(8 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/policy/dtensor_policy_worker_v2.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/policy/dtensor_policy_worker_v2.py
🧠 Learnings (1)
📚 Learning: 2025-10-30T20:50:44.126Z
Learnt from: adil-a
Repo: NVIDIA-NeMo/RL PR: 1440
File: examples/configs/sft_automodel.yaml:48-58
Timestamp: 2025-10-30T20:50:44.126Z
Learning: In DTensor configurations for MoE (Mixture of Experts) models, expert_parallel_size and data_parallel_size can be applied together without multiplying the GPU requirements. Expert Parallelism (EP) only applies to MoE layers, while Data Parallelism/FSDP applies to non-MoE layers. Therefore, configurations like expert_parallel_size: 8 and data_parallel_size: 8 are valid on an 8-GPU cluster for MoE models.
Applied to files:
nemo_rl/models/policy/dtensor_policy_worker_v2.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). (2)
- GitHub Check: Lint check
- GitHub Check: Post submodule check comment / Comment on PR
|
@adil-a do you mind adding the total step time before and after as well to the pr description? |
|
❌ Submodule Fast-Forward Check FailedCheck based on commit: d5cc915 (PR #1509 from ❌ Submodules that need attention:Automodel: ❌ Commits have DIVERGED from a common ancestor Please ensure all submodule commits are fast-forwards of the main branch before merging. |
|
❌ Submodule Fast-Forward Check FailedCheck based on commit: 0577c10 (PR #1509 from ❌ Submodules that need attention:Automodel: ❌ Commits have DIVERGED from a common ancestor Please ensure all submodule commits are fast-forwards of the main branch before merging. |
|
❌ Submodule Fast-Forward Check FailedCheck based on commit: 733529b (PR #1509 from ❌ Submodules that need attention:Automodel: ❌ Commits have DIVERGED from a common ancestor Please ensure all submodule commits are fast-forwards of the main branch before merging. |
adil-a
force-pushed
the
adil/fsdp-manager
branch
from
733529b to
128fb6f
Compare
|
Signed-off-by: adil-a <adil.asif2000@hotmail.com>
|
adil-a
added
CI:L2
What does this PR do ?
Uses Automodel's FSDP2 manager for initializing the v2 worker.
Sharding on current main:
Sharding on this branch:
Summary by CodeRabbit