fix(customizer): publish the HF tree DTensor V2 writes instead of converting it as DCP

services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/backend.py

@@ -34,7 +34,9 @@
 from nmp.rl.tasks.training.backends.nemo_rl.checkpoints import (
     LORA_ADAPTER_SEARCH_PATHS,
     convert_dcp_to_huggingface,
+    copy_consolidated_hf,
     copy_lora_adapter,
+    find_consolidated_hf_root,
     find_lora_adapter_root,
 )
 from nmp.rl.tasks.training.chat_templates import apply_chat_template_to_checkpoint
@@ -318,7 +320,15 @@ def process_checkpoint(
                 ", ".join(str(path) for path in LORA_ADAPTER_SEARCH_PATHS),
             )
 
-        hf_checkpoint_path = convert_dcp_to_huggingface(checkpoint_path, output_path)
+        # DTensor V2 already wrote an HF tree, so publish it rather than converting. Only
+        # V1 leaves a DCP checkpoint for convert_dcp_to_huggingface to read.
+        consolidated_root = find_consolidated_hf_root(checkpoint_path)
+        if consolidated_root is not None:
+            logger.info("Publishing consolidated HF checkpoint from %s", consolidated_root)
+            copy_consolidated_hf(checkpoint_path, consolidated_root, output_path)
+            hf_checkpoint_path = output_path
+        else:
+            hf_checkpoint_path = convert_dcp_to_huggingface(checkpoint_path, output_path)
 
         # Apply chat template if available (full-weight / merged HF trees only)
         chat_template = None

services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/checkpoints.py

@@ -45,6 +45,13 @@
 # adapter tree would otherwise ship without one.
 RL_TOKENIZER_SUBPATH = Path("policy") / "tokenizer"
 
+# Where DTensor V2 writes a ready-to-publish HF tree when checkpointing.save_consolidated
+# is set. V2 saves safetensors SHARDS by default (model_save_format is V2-only), and those
+# carry no DCP .metadata, so convert_dcp_to_huggingface cannot read them -- an all-weights
+# run trained successfully and then failed at publication. V1 is unaffected: it writes real
+# DCP and this directory never exists.
+CONSOLIDATED_HF_SUBPATH = Path("policy") / "weights" / "model" / "consolidated"
+
 
 def find_lora_adapter_root(checkpoint_path: Path) -> Path | None:
     """Return the directory holding ``adapter_config.json``, or None if there is none."""
@@ -55,29 +62,54 @@ def find_lora_adapter_root(checkpoint_path: Path) -> Path | None:
     return None
 
 
-def copy_lora_adapter(checkpoint_path: Path, adapter_root: Path, output_path: Path) -> None:
-    """Copy an adapter tree to ``output_path``, adding the tokenizer when it is elsewhere.
+def find_consolidated_hf_root(checkpoint_path: Path) -> Path | None:
+    """Return the consolidated HF tree DTensor V2 writes, or None if there is none.
 
-    Only the adapter directory is copied: the checkpoint root also holds optimizer shards
-    and scheduler state, which are training artifacts rather than part of the published
-    model.
+    Keyed on config.json rather than the directory: Automodel creates the directory up
+    front and only then writes into it, so an interrupted save can leave it empty.
+    """
+    candidate = checkpoint_path / CONSOLIDATED_HF_SUBPATH
+    return candidate if (candidate / "config.json").is_file() else None
+
+
+def _copy_tree_with_tokenizer(checkpoint_path: Path, source_root: Path, output_path: Path) -> None:
+    """Copy one subtree of a checkpoint out for publication, adding the tokenizer if absent.
+
+    Only that subtree is copied: the checkpoint root also holds optimizer shards and
+    scheduler state, which are training artifacts rather than part of the published model.
     """
     output_path.mkdir(parents=True, exist_ok=True)
-    shutil.copytree(adapter_root, output_path, dirs_exist_ok=True)
+    shutil.copytree(source_root, output_path, dirs_exist_ok=True)
 
     tokenizer_dir = checkpoint_path / RL_TOKENIZER_SUBPATH
     if (output_path / "tokenizer_config.json").is_file():
         return
     if not tokenizer_dir.is_dir():
         logger.warning(
-            "No tokenizer found at %s; the adapter tree is published without one",
+            "No tokenizer found at %s; %s is published without one",
             tokenizer_dir,
+            output_path,
         )
         return
     logger.info("Copying tokenizer from %s to %s", tokenizer_dir, output_path)
     shutil.copytree(tokenizer_dir, output_path, dirs_exist_ok=True)
 
 
+def copy_lora_adapter(checkpoint_path: Path, adapter_root: Path, output_path: Path) -> None:
+    """Copy an adapter tree to ``output_path``, adding the tokenizer when it is elsewhere."""
+    _copy_tree_with_tokenizer(checkpoint_path, adapter_root, output_path)
+
+
+def copy_consolidated_hf(checkpoint_path: Path, consolidated_root: Path, output_path: Path) -> None:
+    """Publish the consolidated HF tree as-is. Automodel already wrote a usable checkpoint.
+
+    Its consolidated export carries the weights, the index, config and generation config,
+    so there is nothing to convert -- only the tokenizer may need collecting from beside
+    the weights.
+    """
+    _copy_tree_with_tokenizer(checkpoint_path, consolidated_root, output_path)
+
+
 def convert_dcp_to_huggingface(
     dcp_checkpoint_path: Path,
     output_path: Path,

services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/grpo_config.py

@@ -513,6 +513,12 @@ def compile_grpo_config(
         "checkpoint_must_save_by": None,
         "save_optimizer": True,
     }
+    if customizer_config.parallelism.policy_backend is PolicyBackend.AUTOMODEL:
+        # V2 defaults to safetensors SHARDS, which carry no DCP .metadata and so cannot be
+        # published: training succeeds and the job then dies converting the checkpoint. This
+        # asks for the HF export alongside them, which publish_checkpoint copies out.
+        # V2-only -- model_save_format's allowed value on V1 is None.
+        cfg["checkpointing"]["save_consolidated"] = True
 
     model_path = customizer_config.model.path
     precision = _adapt_precision(customizer_config.model.precision)

services/rl/tests/test_checkpoints.py

@@ -1,12 +1,14 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-"""Locating and publishing LoRA adapters inside a NeMo-RL checkpoint."""
+"""Locating and publishing the model tree inside a NeMo-RL checkpoint."""
 
 from pathlib import Path
 
 from nmp.rl.tasks.training.backends.nemo_rl.checkpoints import (
+    copy_consolidated_hf,
     copy_lora_adapter,
+    find_consolidated_hf_root,
     find_lora_adapter_root,
 )
 
@@ -108,3 +110,66 @@ def test_copy_without_a_tokenizer_still_publishes_the_adapter(tmp_path: Path):
 
     assert (output / "adapter_config.json").is_file()
     assert not (output / "tokenizer_config.json").exists()
+
+
+def _write_consolidated(checkpoint: Path) -> Path:
+    """The tree DTensor V2 writes when checkpointing.save_consolidated is set."""
+    root = checkpoint / "policy" / "weights" / "model" / "consolidated"
+    root.mkdir(parents=True, exist_ok=True)
+    (root / "config.json").write_text('{"model_type": "qwen3"}')
+    (root / "model.safetensors.index.json").write_text('{"weight_map": {}}')
+    (root / "model-00001-of-00001.safetensors").write_text("weights")
+    return root
+
+
+def test_finds_the_consolidated_tree_dtensor_v2_writes(tmp_path: Path):
+    expected = _write_consolidated(tmp_path)
+    assert find_consolidated_hf_root(tmp_path) == expected
+
+
+def test_sharded_safetensors_alone_are_not_a_consolidated_tree(tmp_path: Path):
+    """Regression guard for nvbug 6740834.
+
+    V2 writes safetensors SHARDS by default, and those carry no DCP .metadata, so the
+    publisher's DCP converter died after a successful all-weights run. Shards on their own
+    must not be mistaken for a publishable tree -- the consolidated export is what counts.
+    """
+    model_dir = tmp_path / "policy" / "weights" / "model"
+    model_dir.mkdir(parents=True)
+    (model_dir / "shard-00001-model-00001-of-00001.safetensors").write_text("weights")
+
+    assert find_consolidated_hf_root(tmp_path) is None
+
+
+def test_an_empty_consolidated_dir_is_not_published(tmp_path: Path):
+    """Automodel creates the directory before writing into it, so existence is not enough."""
+    (tmp_path / "policy" / "weights" / "model" / "consolidated").mkdir(parents=True)
+
+    assert find_consolidated_hf_root(tmp_path) is None
+
+
+def test_a_dcp_checkpoint_has_no_consolidated_tree(tmp_path: Path):
+    """DTensor V1 writes real DCP, which still goes through convert_dcp_to_huggingface."""
+    weights = tmp_path / "policy" / "weights"
+    weights.mkdir(parents=True)
+    (weights / ".metadata").write_text("dcp")
+
+    assert find_consolidated_hf_root(tmp_path) is None
+
+
+def test_copy_publishes_the_consolidated_tree_and_adds_the_tokenizer(tmp_path: Path):
+    checkpoint = tmp_path / "step_1"
+    root = _write_consolidated(checkpoint)
+    tokenizer = checkpoint / "policy" / "tokenizer"
+    tokenizer.mkdir(parents=True)
+    (tokenizer / "tokenizer_config.json").write_text("{}")
+    # A training artifact beside the weights, which must not reach the published model.
+    (checkpoint / "policy" / "weights" / "optimizer").mkdir(parents=True, exist_ok=True)
+
+    output = tmp_path / "published"
+    copy_consolidated_hf(checkpoint, root, output)
+
+    assert (output / "config.json").is_file()
+    assert (output / "model.safetensors.index.json").is_file()
+    assert (output / "tokenizer_config.json").is_file()
+    assert not (output / "optimizer").exists()

services/rl/tests/test_grpo_config.py

@@ -871,6 +871,36 @@ def test_lora_without_module_lists_matches_all_linear(
     assert lora_cfg["exclude_modules"] == []
 
 
+def test_automodel_asks_for_the_consolidated_checkpoint_export(
+    tmp_path: Path, job_ctx: NMPJobContext, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """V2 must be told to write an HF tree, or a finished run cannot be published.
+
+    Regression guard for nvbug 6740834: V2 defaults to safetensors SHARDS, which carry no
+    DCP .metadata, so training succeeded and the job then died in the DCP converter with
+    "No metadata file found". save_consolidated makes it emit the HF export the publisher
+    copies out.
+    """
+    monkeypatch.setenv("NMP_JOB_STORAGE_PVC_CLAIM", "nmp-job-storage")
+    step, _ = _prepared_step(tmp_path, policy_backend=PolicyBackend.AUTOMODEL)
+
+    cfg = compile_grpo_config(step, job_ctx)
+
+    assert cfg["checkpointing"]["save_consolidated"] is True
+
+
+def test_dtensor_v1_does_not_ask_for_a_consolidated_export(
+    tmp_path: Path, job_ctx: NMPJobContext, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """V1 writes real DCP, which the converter reads; the key is V2-only."""
+    monkeypatch.setenv("NMP_JOB_STORAGE_PVC_CLAIM", "nmp-job-storage")
+    step, _ = _prepared_step(tmp_path, policy_backend=PolicyBackend.DTENSOR)
+
+    cfg = compile_grpo_config(step, job_ctx)
+
+    assert "save_consolidated" not in cfg["checkpointing"]
+
+
 def test_policy_backend_dtensor_omits_v2(
     tmp_path: Path, job_ctx: NMPJobContext, monkeypatch: pytest.MonkeyPatch
 ) -> None: