Add a config to load/save model from/to DCP snapshot

.github/workflows/build-and-test.yaml

@@ -63,6 +63,8 @@ jobs:
           ${DOCKER_RUN} python tests/test_policy_to_rollout.py
           ${DOCKER_RUN} python tests/test_process_flow.py
           ${DOCKER_RUN} python tests/test_math_verify.py
+          ${DOCKER_RUN} python tests/test_policy_overfit.py
+          ${DOCKER_RUN} python tests/test_data_packer.py
           ${LONG_DOCKER_RUN} python tests/test_integration.py --stream
 
       - name: Cleanup

configs/deepseek-v3/deepseek-v3-moe-670b-fsdp64-cp4-ep64-sft.toml

@@ -0,0 +1,56 @@
+redis = "12800"
+
+[train]
+resume = false
+epoch = 1
+output_dir = "./outputs/deepseek-v3-test"
+epsilon = 1e-6
+optm_name = "AdamW"
+optm_lr = 2e-7
+optm_impl = "fused"
+optm_weight_decay = 0.01
+optm_betas = [ 0.9, 0.999,]
+optm_warmup_steps = 20
+optm_grad_norm_clip = 1.0
+async_tp_enabled = false
+compile = false
+param_dtype = "bfloat16"
+fsdp_reduce_dtype = "float32"
+fsdp_offload = false
+fsdp_reshard_after_forward = "default"
+train_batch_per_replica = 1
+sync_weight_interval = 1
+
+[policy]
+model_name_or_path = "deepseek-ai/DeepSeek-V3"
+model_max_length = 1024
+model_gradient_checkpointing = true
+
+[logging]
+logger = ['console', 'wandb']
+project_name = "cosmos_rl"
+experiment_name = "deepseek-v3-sft"
+
+[train.train_policy]
+type = "sft"
+dataset.name = "LNTANOooo/sharegpt52k"
+dataset.subset = ""
+dataset.split = "train"
+conversation_column_name = "conversation"
+mini_batch = 1
+
+[train.ckpt]
+enable_checkpoint = true
+save_freq = 100
+max_keep = 2
+save_mode = "async"
+
+[policy.parallelism]
+n_init_replicas = 1
+tp_size = 1
+cp_size = 4
+ep_size = 64
+dp_shard_size = 64
+pp_size = 1
+dp_replicate_size = 1
+cp_rotate_method = "allgather"

cosmos_rl/dispatcher/data/packer/decoder_only_llm_data_packer.py

@@ -257,8 +257,7 @@ def sft_compute_max_len(self, processed_samples: List[List[int]]) -> int:
         """
         Compute the maximum sequence length of the processed samples
         """
-        max_len = max([len(x["token_ids"]) for x in processed_samples])
-        return max_len
+        return max([len(x["token_ids"]) for x in processed_samples])
 
     def sft_collate_fn(
         self,

cosmos_rl/dispatcher/data/packer/deepseek_data_packer.py

@@ -0,0 +1,109 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any, Dict, List
+
+import torch
+from transformers import AutoTokenizer
+
+from cosmos_rl.dispatcher.data.packer.decoder_only_llm_data_packer import (
+    DecoderOnlyLLMDataPacker,
+)
+from cosmos_rl.policy.config import Config
+
+
+class DeepSeek_DataPacker(DecoderOnlyLLMDataPacker):
+    """
+    Data protocol & processing logic for the decoder only LLM for SFT and RL training.
+    """
+
+    def setup(self, config: Config, tokenizer: AutoTokenizer, *args, **kwargs):
+        super().setup(config, tokenizer, *args, **kwargs)
+        self.seq_len = config.policy.model_max_length
+
+    def policy_collate_fn(
+        self,
+        processed_samples: List[DecoderOnlyLLMDataPacker.RLPolicyInput],
+        computed_max_len: int,
+    ) -> Dict[str, Any]:
+        computed_max_len = max(computed_max_len, self.seq_len)
+        input_ids = [x.input_ids for x in processed_samples]
+        logprob_masks = [x.logprob_masks for x in processed_samples]
+        assert len(input_ids) == len(
+            logprob_masks
+        ), "The length of input_ids, and logprob_masks should be the same"
+        device = torch.cuda.current_device()
+
+        collated_dict = {}
+        collated_dict["input_ids"] = torch.tensor(
+            [
+                x[:computed_max_len]
+                + [self.tokenizer.pad_token_id] * (max(0, computed_max_len - len(x)))
+                for x in input_ids
+            ],
+            dtype=torch.long,
+        ).to(device)
+        collated_dict["logprob_masks"] = torch.tensor(
+            [
+                x[:computed_max_len] + [0] * (max(0, computed_max_len - len(x)))
+                for x in logprob_masks
+            ],
+            dtype=torch.bool,
+        ).to(device)
+
+        return collated_dict
+
+    def sft_collate_fn(
+        self,
+        processed_samples: List[Dict[str, Any]],
+        computed_max_len: int,
+        pad_token_id: int,
+        ignore_label_id: int,
+    ) -> Dict[str, Any]:
+        """
+        Collate the processed samples into a minibatch dictionary passed to the SFT model.
+        """
+        computed_max_len = max(computed_max_len, self.seq_len)
+        # First truncate the samples to the computed_max_len
+        list_of_input_ids = [
+            x["token_ids"][:computed_max_len] for x in processed_samples
+        ]
+        list_of_label_ids = [
+            x["label_ids"][:computed_max_len] for x in processed_samples
+        ]
+
+        # Then pad the samples to the computed_max_len
+        input_ids = torch.tensor(
+            [
+                x[:computed_max_len]
+                + [pad_token_id] * (max(0, computed_max_len - len(x)))
+                for x in list_of_input_ids
+            ],
+            dtype=torch.long,
+        )
+        # Model accept unshifted label_ids for loss computation
+        label_ids = torch.tensor(
+            [
+                x[:computed_max_len]
+                + [ignore_label_id] * (max(0, computed_max_len - len(x)))
+                for x in list_of_label_ids
+            ],
+            dtype=torch.long,
+        )
+
+        return {
+            "input_ids": input_ids,
+            "label_ids": label_ids,
+        }

cosmos_rl/policy/config/__init__.py

@@ -657,6 +657,10 @@ class PolicyConfig(BaseModel):
         default=None,
         description="The revision of the model to use",
     )
+    dcp_snapshot_path: Optional[str] = Field(
+        default=None,
+        description="Path to the DCP snapshot to load the model from or save the model to if non-exitent (used for development)"
+    )
     model_max_length: int = Field(
         default=4096,
         description="The maximum length for training, longer than this will be ignored for training stability",

cosmos_rl/policy/model/__init__.py

@@ -17,13 +17,15 @@
 from cosmos_rl.policy.model.qwen2_5_vl import Qwen2_5_VLConditionalModel
 from cosmos_rl.policy.model.qwen3_moe import Qwen3MoE
 from cosmos_rl.policy.model.hf_models import HFModel
+from cosmos_rl.policy.model.deepseek_v3 import DeepseekV3MoEModel
 from cosmos_rl.policy.model.base import ModelRegistry, BaseModel, WeightMapper
 
 __all__ = [
     "GPT",
     "Qwen2_5_VLConditionalModel",
     "Qwen3MoE",
     "HFModel",
+    "DeepseekV3MoEModel",
     "BaseModel",
     "WeightMapper",
     "ModelRegistry",

cosmos_rl/policy/model/base.py

@@ -341,6 +341,7 @@ def load_hf_weights(
         parallel_dims: ParallelDims,
         device: torch.device,
         revision: Optional[str] = None,
+        dcp_snapshot_path: Optional[str] = None,
     ):
         """
         Load weights from a HuggingFace model.
@@ -349,6 +350,8 @@ def load_hf_weights(
             model_name_or_path (str): The name or path of the model.
             parallel_dims (ParallelDims): The parallel dimensions.
             device (torch.device): The device to load the weights.
+            revision (str): The revision of the model to use.
+            dcp_snapshot_path: Path to the DCP snapshot for saving/loading the mdoel.
         """
         raise NotImplementedError