Expose fully_parallel_save and validate_access_integrity in save_megatron_model

Based on container image commit dd9729f (v0.5.0.nemotron_3_super).

Add fully_parallel_save and validate_access_integrity parameters to
AutoBridge.save_megatron_model() and model_load_save.save_megatron_model().

Needed by NVIDIA-NeMo/RL#2226.

Author: nic-nvidia

src/megatron/bridge/models/conversion/auto_bridge.py

@@ -590,6 +590,8 @@ def save_megatron_model(
         hf_tokenizer_path: Optional[str | Path] = None,
         low_memory_save: bool = False,
         hf_tokenizer_kwargs: Optional[dict] = None,
+        fully_parallel_save: bool = True,
+        validate_access_integrity: bool = True,
     ) -> None:
         """
         Save a Megatron model in native Megatron checkpoint format without optimizer
@@ -648,6 +650,8 @@ def save_megatron_model(
             hf_tokenizer_path=hf_tokenizer_path,
             low_memory_save=low_memory_save,
             hf_tokenizer_kwargs=hf_tokenizer_kwargs,
+            fully_parallel_save=fully_parallel_save,
+            validate_access_integrity=validate_access_integrity,
         )
 
     def load_megatron_model(

src/megatron/bridge/training/model_load_save.py

@@ -413,6 +413,8 @@ def save_megatron_model(
     hf_tokenizer_path: Optional[Union[str, Path]] = None,
     low_memory_save: bool = False,
     hf_tokenizer_kwargs: Optional[dict] = None,
+    fully_parallel_save: bool = True,
+    validate_access_integrity: bool = True,
 ) -> None:
     """Save a Megatron model in native Megatron checkpoint format without optimizer state.
 
@@ -505,6 +507,8 @@ def save_megatron_model(
             save_rng=False,
             ckpt_format=ckpt_format,
             dist_ckpt_optim_fully_reshardable=True,
+            fully_parallel_save=fully_parallel_save,
+            ckpt_assume_constant_structure=not validate_access_integrity,
         ),
         dist=None,
     )