From 6626473c79409121a2df5a338df282543c7e6359 Mon Sep 17 00:00:00 2001 From: anubhutiv Date: Tue, 8 Sep 2026 10:25:20 -0700 Subject: [PATCH] fix(customizer): publish the HF tree DTensor V2 writes instead of converting it as DCP Signed-off-by: anubhutiv --- .../training/backends/nemo_rl/backend.py | 12 +++- .../training/backends/nemo_rl/checkpoints.py | 46 +++++++++++-- .../training/backends/nemo_rl/grpo_config.py | 6 ++ services/rl/tests/test_checkpoints.py | 67 ++++++++++++++++++- services/rl/tests/test_grpo_config.py | 30 +++++++++ 5 files changed, 152 insertions(+), 9 deletions(-) diff --git a/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/backend.py b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/backend.py index 55abda0bb5..a391d193eb 100644 --- a/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/backend.py +++ b/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 diff --git a/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/checkpoints.py b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/checkpoints.py index 373cdab956..45b9082321 100644 --- a/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/checkpoints.py +++ b/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, diff --git a/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/grpo_config.py b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/grpo_config.py index dda52c1862..1f292cad74 100644 --- a/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/grpo_config.py +++ b/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) diff --git a/services/rl/tests/test_checkpoints.py b/services/rl/tests/test_checkpoints.py index 6738beab1d..54779ec060 100644 --- a/services/rl/tests/test_checkpoints.py +++ b/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() diff --git a/services/rl/tests/test_grpo_config.py b/services/rl/tests/test_grpo_config.py index 668624bb9e..6da5ad47a1 100644 --- a/services/rl/tests/test_grpo_config.py +++ b/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: