Expose fully_parallel_save in save_megatron_model

[nic-nvidia]

Summary

Add fully_parallel_save parameter to AutoBridge.save_megatron_model() and the underlying model_load_save.save_megatron_model(), forwarded to CheckpointConfig.

Defaults to True — no behavior change for existing callers.

Problem

save_megatron_model() always creates a CheckpointConfig with fully_parallel_save=True (the dataclass default). This activates FullyParallelSaveStrategyWrapper, which calls all_gather_object on DP sub-groups. When the distributed world includes ranks that never enter the save path (e.g., vLLM inference workers in NeMo RL non-colocated setups), these collectives deadlock permanently.

Callers have no way to disable this behavior through the public API.

Fix

Thread fully_parallel_save: bool = True through both layers:

• AutoBridge.save_megatron_model() → model_load_save.save_megatron_model() → CheckpointConfig(..., fully_parallel_save=...)

Motivation

Needed by NVIDIA-NeMo/RL#2226 which fixes a deadlock during HF-to-Megatron checkpoint conversion in non-colocated training/inference setups.

Test plan

• Default behavior unchanged (fully_parallel_save=True)
• Validate with NeMo RL non-colocated 120B recipe (Fix save_megatron_model deadlock: pass fully_parallel_save=False during conversion RL#2226)

Summary by CodeRabbit

• New Features
  • Added a configurable option for model checkpoint saving in distributed training scenarios, allowing control over parallel save behavior during the checkpoint process.

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

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 5e2d489e-a80c-4001-b51b-88a012a4a824

📥 Commits

Reviewing files that changed from the base of the PR and between 00ddf16 and f010b59.

📒 Files selected for processing (2)

• src/megatron/bridge/models/conversion/auto_bridge.py
• src/megatron/bridge/training/model_load_save.py

---

📝 Walkthrough

Walkthrough

A new optional parameter fully_parallel_save: bool = True was added to save_megatron_model functions across two files. This parameter controls whether all Data Parallel ranks must participate in save collectives, with the flag being propagated through the call chain from the conversion API to the core training save function.

Changes

Cohort / File(s)	Summary
Save Model Parameter Addition src/megatron/bridge/models/conversion/auto_bridge.py, src/megatron/bridge/training/model_load_save.py	Added fully_parallel_save: bool = True parameter to save_megatron_model functions. Updated docstrings to document how the flag controls collective participation requirements for DP ranks during save operations. Parameter is propagated through the call chain to control checkpoint save strategy.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

🚥 Pre-merge checks | ✅ 3 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Test Results For Major Changes	⚠️ Warning	PR introduces new parameter to public APIs but validation is explicitly marked as pending; default behavior verified but full test results not yet documented.	Complete pending validation on NeMo RL non-colocated 120B recipe and document test results in PR description before merging.

✅ 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 directly and accurately describes the main change: exposing the fully_parallel_save parameter in the save_megatron_model function.
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 unit tests (beta)

• Create PR with unit tests

---

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

[yaoyu-33]

/ok to test 664a8a8

[yaoyu-33]

@nic-nvidia plz fix test fail

[huvunvidia]

Hi @nic-nvidia,
Just check in if you have the opportunity to look through the tests failed due to the PR: https://github.com/NVIDIA-NeMo/Megatron-Bridge/actions/runs/24162935847/job/70522583158?pr=3207