fix: Fix DCP-to-HF conversion for model-wrapped checkpoints

[RayenTian]

Summary

1. The file layout differs between dtensor v1 and v2. Using policy/weights as the ckpt path for v2 cannot find data, and vice versa.
2. DCP exports wrap weights under {"model": ...} in dtensor v1, while dtensor v2 saves a flat state_dict.

In this PR, we:

• Fix DCP-to-HF conversion to handle both wrapped ({"model": ...}) and flat state dict formats.
• Align dtensor v1 and v2 to the same layout: model weights in .../weights/model and optimizer in .../optimizer/optim.
• Update documentation examples to use the correct DCP weights path (.../optimizer/optim).
• Align checkpoint conversion guidance across README and guides.

Related issue

• Fixes Checkpoint conversion dcp to hf format doesnot work #1427
• Fixes Add HF offline conversion script for models trained with DTensorPolicyWorkerV2 #1351

The file structure of V1

results/grpo_tp2_v1/step_1/
├── config.yaml
├── policy
│   ├── optimizer
│   │   ├── __0_0.distcp
│   │   ├── __1_0.distcp
│   │   ├── __2_0.distcp
│   │   ├── __3_0.distcp
│   │   ├── __4_0.distcp
│   │   ├── __5_0.distcp
│   │   ├── __6_0.distcp
│   │   └── __7_0.distcp
│   ├── tokenizer
│   │   ├── added_tokens.json
│   │   ├── chat_template.jinja
│   │   ├── merges.txt
│   │   ├── special_tokens_map.json
│   │   ├── tokenizer_config.json
│   │   ├── tokenizer.json
│   │   └── vocab.json
│   └── weights
│       ├── __0_0.distcp
│       ├── __1_0.distcp
│       ├── __2_0.distcp
│       ├── __3_0.distcp
│       ├── __4_0.distcp
│       ├── __5_0.distcp
│       ├── __6_0.distcp
│       └── __7_0.distcp
├── train_dataloader.pt
└── training_info.json

The structure of V2

results/grpo_tp2_v2/step_1
├── config.yaml
├── policy
│   ├── optimizer
│   │   └── optim
│   │       ├── __0_0.distcp
│   │       ├── __1_0.distcp
│   │       ├── __2_0.distcp
│   │       ├── __3_0.distcp
│   │       ├── __4_0.distcp
│   │       ├── __5_0.distcp
│   │       ├── __6_0.distcp
│   │       └── __7_0.distcp
│   ├── tokenizer
│   │   ├── added_tokens.json
│   │   ├── chat_template.jinja
│   │   ├── merges.txt
│   │   ├── special_tokens_map.json
│   │   ├── tokenizer_config.json
│   │   ├── tokenizer.json
│   │   └── vocab.json
│   └── weights
│       └── model
│           ├── __0_0.distcp
│           ├── __1_0.distcp
│           ├── __2_0.distcp
│           ├── __3_0.distcp
│           ├── __4_0.distcp
│           ├── __5_0.distcp
│           ├── __6_0.distcp
│           └── __7_0.distcp
├── train_dataloader.pt
└── training_info.json

Test

uv run examples/converters/convert_dcp_to_hf.py \
--config=results/grpo_tp2_v2/step_1/config.yaml \
--dcp-ckpt-path=results/grpo_tp2_v2/step_1/policy/weights\
--hf-ckpt-path=results/grpo_tp2_v2/step_1-hf


uv run examples/converters/convert_dcp_to_hf.py \
--config=results/grpo_tp2_v1/step_1/config.yaml \
--dcp-ckpt-path=results/grpo_tp2_v1/step_1/policy/weights\
--hf-ckpt-path=results/grpo_tp2_v1/step_1-hf

Summary by CodeRabbit

• Documentation

  • Updated DCP-to-HF conversion examples across multiple guides with corrected checkpoint file paths for model conversion.

• Improvements

  • Enhanced state dictionary handling during model conversion to support flexible checkpoint file structures without requiring directory-based inputs.

[coderabbitai]

📝 Walkthrough

Walkthrough

Updates DCP-to-HF checkpoint conversion documentation across multiple guide files to reference a specific model file path instead of a directory. Modifies the conversion logic to conditionally handle state dict structures, supporting both nested {"model": ...} and flat formats without forced assertions or unnecessary rewrites.

Changes

Cohort / File(s)	Summary
Documentation Path Updates README.md, docs/about/evaluation.md, docs/design-docs/checkpointing.md, docs/guides/eval.md, docs/guides/grpo-deepscaler.md, docs/guides/sft-openmathinstruct2.md	Updated --dcp-ckpt-path argument in DCP-to-HF conversion examples from results/grpo/step_170/policy/weights/ to results/grpo/step_170/policy/weights/model, pointing to a specific model file instead of a directory.
Conversion Logic nemo_rl/utils/native_checkpoint.py	Modified convert_dcp_to_hf to reload saved weights and conditionally rewrite the file only if state dict is exactly {"model"}. Removes unconditional assertion and rewrite, now handles both nested and non-nested state dict formats.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

🚥 Pre-merge checks | ✅ 6

✅ Passed checks (6 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title accurately describes the main change: fixing DCP-to-HF conversion for model-wrapped checkpoints, which is the core objective of the PR.
Linked Issues check	✅ Passed	The PR addresses the core issue #1427: updating conversion logic to handle wrapped state dicts and correcting DCP weights path to resolve FileNotFoundError in checkpoint conversion.
Out of Scope Changes check	✅ Passed	All changes are directly related to fixing DCP-to-HF conversion: logic updates in native_checkpoint.py and consistent documentation updates across multiple guides reflecting the corrected path.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Test Results For Major Changes	✅ Passed	Minor bug fix with targeted code changes (+4/-5 lines) and documentation path corrections, qualifying for pass per policy.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.

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

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment
• Commit unit tests in branch ruit/ckpt_bug_issue_1427

---

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.}

[terrykong]

is this a fix for dtensor v1? or v2?

[github-actions]

⚠️ File Consistency Check

Check based on commit: 0fed948 (PR #1881 from ruit/ckpt_bug_issue_1427)

⚠️ DTensor Policy Worker Synchronization Warning

The file nemo_rl/models/policy/workers/dtensor_policy_worker.py was modified in this PR, but nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py was not updated.

Why this matters:
These files contain related DTensor policy worker implementations that should be kept synchronized to ensure consistency across different versions.

Action required:

• Please review if the changes in nemo_rl/models/policy/workers/dtensor_policy_worker.py should also be applied to nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py
• Update nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py if necessary to maintain consistency
• If the files are intentionally different, please add a comment in the PR explaining why

Files to check:

• Modified: nemo_rl/models/policy/workers/dtensor_policy_worker.py
• Not modified: nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py

---

_{This check ensures that related file implementations remain synchronized across the codebase. If you believe this warning is incorrect or the files should intentionally differ, please add a comment explaining the reasoning.}

[RayenTian]

is this a fix for dtensor v1? or v2?

The issue is because of:

1. the different file layouts between dtensor v1 and v2. Using policy/weights as the ckpt path for v2 cannot find data, and vice versa.
2. DCP exports wrap weights under {"model": ...} in dtensor v1, while dtensor v2 saves a flat state_dict.

This PR, we unify the path and weights state_dict. Now it works well for both dtensor v1 and v2 ckpt.

[yuki-97]

thanks for the fix!

can you help adding an unit test for dtensor v2 like test_convert_dcp_to_hf in tests/unit/utils/test_native_checkpoint.py to guard it?

[yuki-97]

Also just curious from the PR description that seems v1 will save the tokenizer but v2 won't?

[github-actions]

⚠️ File Consistency Check

Check based on commit: 2a5f5a2 (PR #1881 from ruit/ckpt_bug_issue_1427)

⚠️ DTensor Policy Worker Synchronization Warning

The file nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py was modified in this PR, but nemo_rl/models/policy/workers/dtensor_policy_worker.py was not updated.

Why this matters:
These files contain related DTensor policy worker implementations that should be kept synchronized to ensure consistency across different versions.

Action required:

• Please review if the changes in nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py should also be applied to nemo_rl/models/policy/workers/dtensor_policy_worker.py
• Update nemo_rl/models/policy/workers/dtensor_policy_worker.py if necessary to maintain consistency
• If the files are intentionally different, please add a comment in the PR explaining why

Files to check:

• Modified: nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py
• Not modified: nemo_rl/models/policy/workers/dtensor_policy_worker.py

---

_{This check ensures that related file implementations remain synchronized across the codebase. If you believe this warning is incorrect or the files should intentionally differ, please add a comment explaining the reasoning.}

[github-actions]

⚠️ File Consistency Check

Check based on commit: 453a291 (PR #1881 from ruit/ckpt_bug_issue_1427)

⚠️ DTensor Policy Worker Synchronization Warning

The file nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py was modified in this PR, but nemo_rl/models/policy/workers/dtensor_policy_worker.py was not updated.

Why this matters:
These files contain related DTensor policy worker implementations that should be kept synchronized to ensure consistency across different versions.

Action required:

• Please review if the changes in nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py should also be applied to nemo_rl/models/policy/workers/dtensor_policy_worker.py
• Update nemo_rl/models/policy/workers/dtensor_policy_worker.py if necessary to maintain consistency
• If the files are intentionally different, please add a comment in the PR explaining why

Files to check:

• Modified: nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py
• Not modified: nemo_rl/models/policy/workers/dtensor_policy_worker.py

---

_{This check ensures that related file implementations remain synchronized across the codebase. If you believe this warning is incorrect or the files should intentionally differ, please add a comment explaining the reasoning.}

[github-actions]

⚠️ File Consistency Check

Check based on commit: 9627a78 (PR #1881 from ruit/ckpt_bug_issue_1427)

⚠️ DTensor Policy Worker Synchronization Warning

The file nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py was modified in this PR, but nemo_rl/models/policy/workers/dtensor_policy_worker.py was not updated.

Why this matters:
These files contain related DTensor policy worker implementations that should be kept synchronized to ensure consistency across different versions.

Action required:

• Please review if the changes in nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py should also be applied to nemo_rl/models/policy/workers/dtensor_policy_worker.py
• Update nemo_rl/models/policy/workers/dtensor_policy_worker.py if necessary to maintain consistency
• If the files are intentionally different, please add a comment in the PR explaining why

Files to check:

• Modified: nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py
• Not modified: nemo_rl/models/policy/workers/dtensor_policy_worker.py

---

_{This check ensures that related file implementations remain synchronized across the codebase. If you believe this warning is incorrect or the files should intentionally differ, please add a comment explaining the reasoning.}

[github-actions]

⚠️ File Consistency Check

Check based on commit: e572ba5 (PR #1881 from ruit/ckpt_bug_issue_1427)

⚠️ DTensor Policy Worker Synchronization Warning

The file nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py was modified in this PR, but nemo_rl/models/policy/workers/dtensor_policy_worker.py was not updated.

Why this matters:
These files contain related DTensor policy worker implementations that should be kept synchronized to ensure consistency across different versions.

Action required:

• Please review if the changes in nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py should also be applied to nemo_rl/models/policy/workers/dtensor_policy_worker.py
• Update nemo_rl/models/policy/workers/dtensor_policy_worker.py if necessary to maintain consistency
• If the files are intentionally different, please add a comment in the PR explaining why

Files to check:

• Modified: nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py
• Not modified: nemo_rl/models/policy/workers/dtensor_policy_worker.py

---

_{This check ensures that related file implementations remain synchronized across the codebase. If you believe this warning is incorrect or the files should intentionally differ, please add a comment explaining the reasoning.}