diff --git a/.github/workflows/build-and-test.yaml b/.github/workflows/build-and-test.yaml
index 92bcbedd1..ae012d83e 100644
--- a/.github/workflows/build-and-test.yaml
+++ b/.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
diff --git a/configs/deepseek-v3/deepseek-v3-moe-670b-fsdp64-cp4-ep64-sft.toml b/configs/deepseek-v3/deepseek-v3-moe-670b-fsdp64-cp4-ep64-sft.toml
new file mode 100644
index 000000000..b92d621b9
--- /dev/null
+++ b/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"
diff --git a/cosmos_rl/dispatcher/data/packer/decoder_only_llm_data_packer.py b/cosmos_rl/dispatcher/data/packer/decoder_only_llm_data_packer.py
index 0629c1a05..80248e3c6 100644
--- a/cosmos_rl/dispatcher/data/packer/decoder_only_llm_data_packer.py
+++ b/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,
diff --git a/cosmos_rl/dispatcher/data/packer/deepseek_data_packer.py b/cosmos_rl/dispatcher/data/packer/deepseek_data_packer.py
new file mode 100644
index 000000000..d75aafca4
--- /dev/null
+++ b/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,
+        }
diff --git a/cosmos_rl/policy/config/__init__.py b/cosmos_rl/policy/config/__init__.py
index e98865ca0..e63865aed 100644
--- a/cosmos_rl/policy/config/__init__.py
+++ b/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",
diff --git a/cosmos_rl/policy/model/__init__.py b/cosmos_rl/policy/model/__init__.py
index 33459a5f4..202fa77fb 100644
--- a/cosmos_rl/policy/model/__init__.py
+++ b/cosmos_rl/policy/model/__init__.py
@@ -17,6 +17,7 @@
 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__ = [
@@ -24,6 +25,7 @@
     "Qwen2_5_VLConditionalModel",
     "Qwen3MoE",
     "HFModel",
+    "DeepseekV3MoEModel",
     "BaseModel",
     "WeightMapper",
     "ModelRegistry",
diff --git a/cosmos_rl/policy/model/base.py b/cosmos_rl/policy/model/base.py
index cb6e9150b..9ae96399e 100644
--- a/cosmos_rl/policy/model/base.py
+++ b/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
 
diff --git a/cosmos_rl/policy/model/deepseek_v3/__init__.py b/cosmos_rl/policy/model/deepseek_v3/__init__.py
new file mode 100644
index 000000000..950eb2e36
--- /dev/null
+++ b/cosmos_rl/policy/model/deepseek_v3/__init__.py
@@ -0,0 +1,457 @@
+# 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.
+
+import os
+from functools import cached_property
+from typing import Any, Callable, List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+from safetensors import safe_open
+from torch.nn.modules.module import _IncompatibleKeys
+from transformers import AutoConfig
+
+try:
+    from torch.distributed.tensor import DTensor
+except ImportError:
+    print("torch.distributed.tensor is not available. DeepSeek model will not work.")
+
+
+from cosmos_rl.dispatcher.data.packer.deepseek_data_packer import DeepSeek_DataPacker
+from cosmos_rl.policy.config import Config as CosmosConfig
+from cosmos_rl.policy.kernel.moe import moe
+from cosmos_rl.policy.model.base import BaseModel, ModelRegistry
+from cosmos_rl.policy.model.deepseek_v3 import deepseekv3_mapped
+from cosmos_rl.policy.model.deepseek_v3.checkpoint_planner import RenameLoadPlanner
+from cosmos_rl.policy.model.deepseek_v3.weight_mapper import (
+    DeepseekV3MoEWeightMapper,
+    convert_weight_from_hf,
+    weight_dequant,
+)
+from cosmos_rl.utils.logging import logger
+from cosmos_rl.utils.parallelism import ParallelDims
+from cosmos_rl.utils.util import clear_weight_name, resolve_model_path, retry
+
+
+@ModelRegistry.register(
+    DeepseekV3MoEWeightMapper, default_data_packer_cls=DeepSeek_DataPacker
+)
+class DeepseekV3MoEModel(BaseModel):
+    @staticmethod
+    def supported_model_types():
+        return ["deepseek_v3"]
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        hf_config: AutoConfig,
+        model_name_or_path: str,
+        max_position_embeddings: Optional[int] = None,
+    ) -> "DeepseekV3MoEModel":
+        """
+        Initialize a DeepseekV3MoE model from a pretrained model.
+
+        Args:
+            model_name_or_path (str): Model name or path to the pretrained model.
+
+        Returns:
+            DeepseekV3MoE: DeepseekV3MoE model.
+
+        """
+        if hf_config.model_type not in cls.supported_model_types():
+            raise ValueError(f"Unsupported model type: {hf_config.model_type}")
+
+        assert (
+            "deepseek-v3" in model_name_or_path.lower()
+        ), f"Unsupported model {model_name_or_path}"
+
+        if hf_config.num_hidden_layers == 4:
+            deepseek_config = deepseekv3_mapped.DeepseekConfig(n_layers=4)
+        elif hf_config.num_hidden_layers == 61:
+            deepseek_config = deepseekv3_mapped.DeepseekConfig()
+        else:
+            raise ValueError(
+                f"Only 4 or 61 layer models supported at the moment. Got hf_config.llm_config.num_hidden_layers={hf_config.llm_config.num_hidden_layers}"
+            )
+
+        model = DeepseekV3MoEModel(deepseek_config, hf_config=hf_config)
+
+        logger.info("Initializing the model")
+        with torch.no_grad():
+            model.init_weights()
+        logger.info("Model initialized")
+        return model
+
+    def __init__(
+        self,
+        model_config: deepseekv3_mapped.DeepseekConfig,
+        hf_config: AutoConfig,
+    ):
+        super().__init__(hf_config=hf_config)
+        self.config = model_config
+
+        orig_precision = torch.get_default_dtype()
+        precision = getattr(torch, model_config.dtype)
+        torch.set_default_dtype(precision)
+        logger.info(f"Setting torch default dtype from {orig_precision} to {precision}")
+
+        self.build_model(model_config)
+
+        torch.set_default_dtype(
+            orig_precision
+        )  # Reset the default dtype to the original value
+        logger.info(f"Reset torch default dtype to {orig_precision}")
+
+    def build_model(self, model_config: deepseekv3_mapped.DeepseekConfig):
+        # Create reasoning model
+        self.model = deepseekv3_mapped.Transformer(args=model_config)
+
+    def init_weights(
+        self,
+        buffer_device: Optional[torch.device] = None,
+    ):
+        self.apply(_init_weights)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor = None,
+        position_ids: torch.Tensor = None,
+        initial_aux_loss: Optional[torch.Tensor] = None,
+        **data_batch: Optional[torch.Tensor],
+    ) -> torch.Tensor:
+        logger.debug(
+            f"[model] input ids shape: {input_ids.shape}, position_ids: {position_ids.shape}"
+        )
+
+        logits, aux_loss = self.model(
+            tokens=input_ids,
+            position_ids=position_ids,
+            padding_mask=data_batch.get("padding_mask", None),
+        )
+
+        if self.config.aux_loss_coeff > 0:
+            if initial_aux_loss is not None and aux_loss is not None:
+                final_aux_loss = initial_aux_loss + aux_loss
+                return logits, final_aux_loss
+            elif initial_aux_loss is not None:
+                final_aux_loss = initial_aux_loss
+                return logits, final_aux_loss
+            else:
+                return logits
+        else:
+            assert (
+                initial_aux_loss is None
+            ), "initial_aux_loss must be None when aux_loss_coeff = 0"
+            assert aux_loss is None, "aux_loss must be None when aux_loss_coeff = 0"
+            return logits
+
+    @property
+    def parallelize_fn(self):
+        from cosmos_rl.policy.model.deepseek_v3.parallelize import parallelize_model
+
+        return parallelize_model, self
+
+    def post_to_empty_hook(self, cosmos_config: CosmosConfig):
+        return
+
+    def separate_model_parts(self) -> List[nn.Module]:
+        return [self]
+
+    def get_position_ids(self, **kwargs) -> Tuple[torch.Tensor, torch.Tensor, int]:
+        seq_dim_idx = 1
+        inputs = kwargs["input_ids"]
+        position_ids = (
+            torch.arange(inputs.size(-1), dtype=torch.long, device=inputs.device)
+            .unsqueeze(0)
+            .expand_as(inputs)
+        )
+        return position_ids, inputs, seq_dim_idx
+
+    def apply_pipeline_split(self, pp_rank, pp_size):
+        raise NotImplementedError
+
+    @cached_property
+    def _get_nparams_and_flops_fn(self) -> Callable[[int], tuple[int, int]]:
+        nparams = sum(p.numel() for p in self.parameters())
+        nparams_embedding = sum(
+            sum(p.numel() for p in m.parameters())
+            for m in self.children()
+            if isinstance(m, nn.Embedding)
+        )
+
+        # Reasoning behind the factor of 12 for the self-attention part of the formula:
+        # 1. each self-attention has 2 matmul in the forward and 4 in the backward (6)
+        # 2. the flash attention does 1 more matmul recomputation in the backward
+        #    but recomputation should not be counted in calculating MFU           (+0)
+        # 3. each matmul performs 1 multiplication and 1 addition                 (*2)
+        # 4. we follow the convention and do not account for sparsity in causal attention
+        layers, heads, head_dim = (
+            self.config.n_layers,
+            self.config.n_heads,
+            self.config.dim // self.config.n_heads,
+        )
+        return lambda seq_len: (
+            nparams,
+            6 * (nparams - nparams_embedding)
+            + 12 * layers * heads * head_dim * seq_len,
+        )
+
+    def get_nparams_and_flops(self, seq_len: int) -> tuple[int, int]:
+        return self._get_nparams_and_flops_fn(seq_len)
+
+    def load_hf_weights(
+        self,
+        model_name_or_path: str,
+        parallel_dims: ParallelDims,
+        device: torch.device,
+        revision: Optional[str] = None,
+        dcp_snapshot_path: Optional[str] = None,
+    ):
+        if (
+            dcp_snapshot_path
+            and os.path.exists(dcp_snapshot_path)
+            and len(os.listdir(dcp_snapshot_path)) >= 0
+        ):
+            logger.info("Loading from distributed checkpoints...")
+            # Transformer engine adds this extra state which we dont have in the saved checkpoint.
+            mapped_state_dict = {
+                k: v
+                for k, v in self.state_dict().items()
+                if not k.endswith("_extra_state")
+            }
+
+            logger.info("Creating storage reader...")
+            fs_storage_reader = torch.distributed.checkpoint.FileSystemReader(
+                dcp_snapshot_path
+            )
+            logger.info("Loading checkpoint ...")
+            torch.distributed.checkpoint.load(
+                state_dict=mapped_state_dict,
+                storage_reader=fs_storage_reader,
+                planner=RenameLoadPlanner(allow_partial_load=False),
+            )
+            logger.info("Refreshing model.")
+            self.load_state_dict(self.state_dict())
+            logger.info("Checkpoint loaded.")
+        else:
+            # The checkpoint loading assumes bf16 checkpoints. The default deepseekv3 checkpoint is in fp8. We needs to convert the checkpoints from fp8 to bf16 first.
+            # Can follow the instructions here: https://github.com/NVIDIA-NeMo/RL/blob/main/docs/guides/deepseek.md
+            # Load all safetensors from `model_path`
+            model_type = retry(AutoConfig.from_pretrained)(
+                model_name_or_path, trust_remote_code=True
+            ).model_type
+            model_path = resolve_model_path(model_name_or_path, revision)
+            safetensors_files = [
+                f for f in os.listdir(model_path) if f.endswith(".safetensors")
+            ]
+
+            self_state_dict = self.state_dict()
+            self_state_dict = {
+                clear_weight_name(k): v for k, v in self_state_dict.items()
+            }
+
+            lm_head_weight_key = "model.lm_head.weight"
+            embed_tokens_weight_key = "model.model.embed_tokens.weight"
+            weights_of_ckpt_names = set()
+            reserved = {}
+            scale_inv_paths = {}
+
+            for f in safetensors_files:
+                ckpt = retry(safe_open)(
+                    os.path.join(model_path, f), framework="pt", device=str(device)
+                )
+                keys = ckpt.keys()
+                for name in keys:
+                    if name.endswith("weight_scale_inv"):
+                        scale_inv_paths[name] = os.path.join(model_path, f)
+
+            for f in safetensors_files:
+                logger.info(f"Loading safetensors: {f}")
+                weights_of_ckpt = {}
+                ckpt = retry(safe_open)(
+                    os.path.join(model_path, f), framework="pt", device=str(device)
+                )
+                keys = ckpt.keys()
+                for name in keys:
+                    ckpt_tensor = ckpt.get_tensor(name)
+                    weights_of_ckpt[name] = ckpt_tensor
+                    weights_of_ckpt_names.add(name)
+                    if name == embed_tokens_weight_key:
+                        reserved[name] = ckpt_tensor
+
+                for name in weights_of_ckpt.keys():
+                    tensor = weights_of_ckpt[name]
+                    if name.endswith("weight_scale_inv") or "layers.61" in name:
+                        # Skip since this weight is used for dequantization
+                        continue
+
+                    if (
+                        "down_proj" in name
+                        or "up_proj" in name
+                        or "gate_proj" in name
+                        or "self_attn.kv_a_proj_with_mqa" in name
+                        or "self_attn.kv_b_proj" in name
+                        or "self_attn.o_proj" in name
+                        or "self_attn.q_a_proj" in name
+                        or "self_attn.q_b_proj" in name
+                    ) and "weight" in name:
+                        inv_name = name + "_scale_inv"
+                        inv_tensor = retry(safe_open)(
+                            scale_inv_paths[inv_name],
+                            framework="pt",
+                            device=str(device),
+                        ).get_tensor(inv_name)
+                        tensor = weight_dequant(tensor, inv_tensor)
+
+                    dest_name, shared_weight, expert_id = convert_weight_from_hf(
+                        tensor,
+                        name,
+                        model_type,
+                        parallel_dims,
+                        n_experts=self.config.n_routed_experts,
+                    )
+
+                    if dest_name is None:
+                        # This is due to the expert parallelism grouping
+                        continue
+
+                    if dest_name not in self_state_dict and parallel_dims.pp_enabled:
+                        logger.info(
+                            f"Weight `{dest_name}` is discarded, maybe due to pipeline parallelism or expert parallelism grouping. Skipping this weight checking"
+                        )
+                        continue
+
+                    target_tensor = self_state_dict[dest_name]
+                    if isinstance(target_tensor, torch.distributed.tensor.DTensor):
+                        target_tensor = target_tensor.to_local()
+                    # Write to the correct expert of the target tensor
+                    if expert_id is not None:
+                        # Convert expert_id to local_expert_id
+                        n_local_experts = (
+                            self.config.n_routed_experts
+                            // parallel_dims.tp
+                            // parallel_dims.dp_shard
+                        )
+
+                        expert_id = expert_id % n_local_experts
+                        target_tensor = target_tensor[expert_id]
+
+                    assert (
+                        target_tensor.shape == shared_weight.shape
+                    ), f"Shape mismatch: {target_tensor.shape} != {shared_weight.shape} for {dest_name}"
+                    with torch.no_grad():
+                        target_tensor.data.copy_(shared_weight)
+                torch.distributed.barrier()
+                logger.info(f"Loaded safetensors: {f} successfully.")
+
+            if (
+                lm_head_weight_key not in weights_of_ckpt_names
+                and embed_tokens_weight_key in weights_of_ckpt_names
+            ):
+                # tied with embed_tokens.weight
+                name = lm_head_weight_key
+                assert embed_tokens_weight_key in reserved
+                tensor = reserved[embed_tokens_weight_key]
+                dest_name, shared_weight = convert_weight_from_hf(
+                    tensor,
+                    name,
+                    model_type,
+                    parallel_dims,
+                    n_experts=self.config.n_routed_experts,
+                )
+                if dest_name in self_state_dict:
+                    target_tensor = self_state_dict[dest_name]
+                    is_dist_tensor = isinstance(
+                        target_tensor, torch.distributed.tensor.DTensor
+                    )
+                    local_view = (
+                        target_tensor.to_local() if is_dist_tensor else target_tensor
+                    )
+                    assert (
+                        local_view.shape == shared_weight.shape
+                    ), f"Shape mismatch: {local_view.shape} != {shared_weight.shape} for {dest_name}"
+                    with torch.no_grad():
+                        local_view.data.copy_(shared_weight)
+
+            if dcp_snapshot_path:
+                logger.info(f"Dumping the tensors to DCP folder {dcp_snapshot_path}")
+                os.makedirs(dcp_snapshot_path, exist_ok=True)
+                fs_storage_writer = torch.distributed.checkpoint.FileSystemWriter(
+                    dcp_snapshot_path
+                )
+                torch.distributed.checkpoint.save(
+                    state_dict=self_state_dict,
+                    storage_writer=fs_storage_writer,
+                )
+
+    def load_state_dict(
+        self, state_dict: dict[str, Any], strict: bool = True, assign: bool = False
+    ):
+        """
+        Ignore the missing keys with substrings matching `substring_to_ignore` (e.g., "_extra_state" keys imposed by
+        TransformerEngine for FP8).
+        """
+        actual_missing_keys, unexpected_keys = super().load_state_dict(
+            state_dict, strict=False, assign=assign
+        )
+        if strict:
+            if len(actual_missing_keys) > 0 or len(unexpected_keys) > 0:
+                raise ValueError(
+                    f"Missing keys: {actual_missing_keys}\n\nUnexpected keys: {unexpected_keys}"
+                )
+        return _IncompatibleKeys(actual_missing_keys, unexpected_keys)
+
+    @cached_property
+    def weight_sync_transforms(
+        self,
+    ) -> List[Tuple[str, Union[torch.Tensor, Callable]]]:
+        # 1. get all parameters, but not buffers
+        transforms = {}
+
+        for local_name, param in self.named_parameters():
+            hf_name = self.weight_mapper.policy_map_local_key_to_hf_key(
+                clear_weight_name(local_name)
+            )
+
+            is_dist_tensor = isinstance(param, torch.distributed.tensor.DTensor)
+            transform_or_view = param.to_local() if is_dist_tensor else param
+
+            assert (
+                hf_name not in transforms
+            ), f"Duplicate key found in transforms: {hf_name}"
+            transforms[hf_name] = transform_or_view
+
+        return sorted(transforms.items())
+
+
+def _init_weights(module):
+    std = 0.02
+
+    def to_local(tensor):
+        if isinstance(tensor, DTensor):
+            return tensor.to_local()
+        else:
+            return tensor
+
+    if isinstance(module, torch.nn.Linear):
+        to_local(module.weight).normal_(mean=0.0, std=std)
+        if module.bias is not None:
+            to_local(module.bias).zero_()
+    elif isinstance(module, torch.nn.Embedding):
+        to_local(module.weight).normal_(mean=0.0, std=std)
+        if module.padding_idx is not None:
+            to_local(module.weight)[module.padding_idx].zero_()
+    elif isinstance(module, moe.Gate):
+        to_local(module.weight).normal_(mean=0.0, std=std)
diff --git a/cosmos_rl/policy/model/deepseek_v3/checkpoint_planner.py b/cosmos_rl/policy/model/deepseek_v3/checkpoint_planner.py
new file mode 100644
index 000000000..718aca777
--- /dev/null
+++ b/cosmos_rl/policy/model/deepseek_v3/checkpoint_planner.py
@@ -0,0 +1,108 @@
+# 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 Optional
+
+import torch
+from torch.distributed.checkpoint.default_planner import DefaultLoadPlanner
+from torch.distributed.checkpoint.metadata import STATE_DICT_TYPE, Metadata
+
+from cosmos_rl.utils.logging import logger
+
+
+class RenameLoadPlanner(DefaultLoadPlanner):
+    """
+    RenameLoadPlanner that renames variables during checkpoint load.
+    """
+
+    def set_up_planner(
+        self,
+        state_dict: STATE_DICT_TYPE,
+        metadata: Optional[Metadata] = None,
+        is_coordinator: bool = False,
+    ) -> None:
+        super().set_up_planner(
+            state_dict=state_dict,
+            metadata=metadata,
+            is_coordinator=is_coordinator,
+        )
+
+        self.state_dict = remap_model_state_for_deepep(self.state_dict, metadata)
+
+        # Do an early check to see if the checkpoint is valid and print the missing
+        # keys in the state dict if not. The reason is the original default planner's
+        # error message is not helpful enough when the keys are mismatched.
+        missing_keys = get_missing_keys(self.state_dict, metadata)
+
+        if missing_keys:
+            logger.info(f"Missing keys in checkpoint: {missing_keys}...")
+            logger.info(f"Checkpoint keys: {list(metadata.state_dict_metadata)}...")
+
+
+def get_missing_keys(
+    state_dict: dict[str, torch.Tensor],
+    metadata: Metadata,
+) -> list[str]:
+    missing_keys = []
+    for fqn, obj in state_dict.items():
+        # ignore state_dict keys which do not exist in `state_dict` if strict=False
+        if fqn not in metadata.state_dict_metadata:
+            missing_keys.append(fqn)
+    return missing_keys
+
+
+def remap_model_state_for_deepep(
+    state_dict: dict[str, torch.Tensor],
+    metadata: Metadata,
+) -> dict[str, torch.Tensor]:
+    """
+    Remap the state dict by removing the "gate_and_up_projs" key.
+    And add the "gate_projs" and "up_projs" keys to the state dict.
+    """
+    import re
+
+    missing_keys = get_missing_keys(state_dict, metadata)
+
+    # Check if there is substring "gate_and_up_projs" in any key of missing_keys
+    # If yes, do a remapping of state_dict keys
+    needs_remapping = any(["gate_and_up_projs" in key for key in missing_keys])
+    if not needs_remapping:
+        return state_dict
+
+    logger.info("Old checkpoint, requires remapping of gate_and_up_projs")
+
+    new_state_dict = state_dict.copy()
+    for key, v in state_dict.items():
+        moe_pattern = r"^([\w.]*)model\.layers\.(\d+)\.mlp\.experts\.gate_and_up_projs$"
+        match = re.match(moe_pattern, key)
+        if match:
+            prefix = match.group(1)
+            layer_num = match.group(2)
+            logger.info(
+                f"Remapping {key} to {prefix}model.layers.{layer_num}.mlp.experts.xxx_projs"
+            )
+
+            v_1, v_2 = torch.chunk(v, 2, -1)
+            new_state_dict[
+                f"{prefix}model.layers.{layer_num}.mlp.experts.gate_projs"
+            ] = v_1.transpose(1, 2)
+            new_state_dict[f"{prefix}model.layers.{layer_num}.mlp.experts.up_projs"] = (
+                v_2.transpose(1, 2)
+            )
+            del new_state_dict[key]
+        elif "mlp.experts.down_projs" in key:
+            new_state_dict[key] = v.transpose(1, 2)
+
+    return new_state_dict
diff --git a/cosmos_rl/policy/model/deepseek_v3/deepseekv3_mapped.py b/cosmos_rl/policy/model/deepseek_v3/deepseekv3_mapped.py
new file mode 100644
index 000000000..172ed4be5
--- /dev/null
+++ b/cosmos_rl/policy/model/deepseek_v3/deepseekv3_mapped.py
@@ -0,0 +1,579 @@
+# 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.
+
+import math
+from dataclasses import dataclass
+from typing import Literal
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+try:
+    from transformer_engine.pytorch.attention import DotProductAttention
+    from transformer_engine.pytorch.module.rmsnorm import RMSNorm as _RMSNorm
+except ImportError:
+    print("transformer_engine.pytorch is not available. DeepSeek model will not work.")
+
+    class _RMSNorm:
+        pass
+
+
+from cosmos_rl.policy.kernel.moe.moe import MoE
+
+
+@dataclass
+class DeepseekConfig:
+    """
+    Data class for defining model arguments and hyperparameters.
+
+    Attributes:
+        max_batch_size (int): Maximum batch size.
+        max_seq_len (int): Maximum sequence length.
+        dtype (Literal["bfloat16", "fp8"]): Data type for computations.
+        vocab_size (int): Vocabulary size.
+        dim (int): Model dimension.
+        inter_dim (int): Intermediate dimension for MLP layers.
+        moe_inter_dim (int): Intermediate dimension for MoE layers.
+        n_layers (int): Number of transformer layers.
+        n_dense_layers (int): Number of dense layers in the model.
+        n_heads (int): Number of attention heads.
+        n_routed_experts (int): Number of routed experts for MoE layers.
+        n_shared_experts (int): Number of shared experts for MoE layers.
+        n_activated_experts (int): Number of activated experts in MoE layers.
+        n_expert_groups (int): Number of expert groups.
+        n_limited_groups (int): Number of limited groups for MoE routing.
+        score_func (Literal["softmax", "sigmoid"]): Scoring function for MoE routing.
+        route_scale (float): Scaling factor for routing scores.
+        q_lora_rank (int): LoRA rank for query projections.
+        kv_lora_rank (int): LoRA rank for key-value projections.
+        qk_nope_head_dim (int): Dimension for query-key projections without positional embeddings.
+        qk_rope_head_dim (int): Dimension for query-key projections with rotary embeddings.
+        v_head_dim (int): Dimension for value projections.
+        original_seq_len (int): Original sequence length.
+        rope_theta (float): Base for rotary positional encoding.
+        rope_factor (float): Scaling factor for extended sequence lengths.
+        beta_fast (int): Fast beta correction factor.
+        beta_slow (int): Slow beta correction factor.
+        mscale (float): Scaling factor for extended attention.
+    """
+
+    max_batch_size: int = 1
+    max_seq_len: int = 4096 * 4
+    dtype: Literal["bfloat16", "fp8"] = "bfloat16"
+    vocab_size: int = 129280
+    dim: int = 7168
+    inter_dim: int = 18432
+    moe_inter_dim: int = 2048
+    n_layers: int = 61
+    n_dense_layers: int = 3
+    n_heads: int = 128
+    # moe
+    n_routed_experts: int = 256
+    n_shared_experts: int = 1
+    n_activated_experts: int = 8
+    n_expert_groups: int = 8
+    n_limited_groups: int = 4
+    train_gate: bool = True
+    gate_bias_update_factor: float = 0.01
+    score_func: Literal["softmax", "sigmoid"] = "sigmoid"
+    route_scale: float = 2.5
+    enable_deepep: bool = False
+    fake_balanced_gate: bool = False
+    # mla
+    q_lora_rank: int = 1536
+    kv_lora_rank: int = 512
+    qk_nope_head_dim: int = 128
+    qk_rope_head_dim: int = 64
+    v_head_dim: int = 128
+    # yarn
+    original_seq_len: int = 4096
+    rope_theta: float = 10000.0
+    rope_factor: float = 40
+    beta_fast: int = 32
+    beta_slow: int = 1
+    mscale: float = 1.0
+    # aux loss
+    aux_loss_coeff: float = 0.0
+
+
+class RMSNorm(_RMSNorm):
+    """
+    Root Mean Square Layer Normalization (RMSNorm).
+
+    Args:
+        dim (int): Dimension of the input tensor.
+        eps (float): Epsilon value for numerical stability. Defaults to 1e-6.
+    """
+
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__(normalized_shape=dim, eps=eps)
+
+
+def precompute_base_freqs(args: DeepseekConfig) -> torch.Tensor:
+    """
+    Precomputes frequency-based complex exponential values for rotary positional embeddings.
+
+    Args:
+        args (DeepseekConfig): Model arguments containing positional embedding parameters.
+
+    Returns:
+        torch.Tensor: Precomputed complex exponential values for positional embeddings.
+    """
+    dim = args.qk_rope_head_dim
+    seqlen = args.max_seq_len
+    beta_fast = args.beta_fast
+    beta_slow = args.beta_slow
+    base = args.rope_theta
+    factor = args.rope_factor
+
+    def find_correction_dim(num_rotations, dim, base, max_seq_len):
+        """
+        Computes the correction dimension for a given number of rotations in the rotary positional embedding.
+
+        Args:
+            num_rotations (float): Number of rotations to compute the correction for.
+            dim (int): Dimensionality of the embedding space.
+            base (float): Base value for the exponential computation.
+            max_seq_len (int): Maximum sequence length.
+
+        Returns:
+            float: The correction dimension based on the input parameters.
+        """
+        return (
+            dim
+            * math.log(max_seq_len / (num_rotations * 2 * math.pi))
+            / (2 * math.log(base))
+        )
+
+    def find_correction_range(low_rot, high_rot, dim, base, max_seq_len):
+        """
+        Computes the range of correction dimensions for rotary positional embeddings.
+
+        Args:
+            low_rot (float): Lower bound for the number of rotations.
+            high_rot (float): Upper bound for the number of rotations.
+            dim (int): Dimensionality of the embedding space.
+            base (float): Base value for the exponential computation.
+            max_seq_len (int): Maximum sequence length.
+
+        Returns:
+            Tuple[int, int]: The range of correction dimensions (low, high), clamped to valid indices.
+        """
+        low = math.floor(find_correction_dim(low_rot, dim, base, max_seq_len))
+        high = math.ceil(find_correction_dim(high_rot, dim, base, max_seq_len))
+        return max(low, 0), min(high, dim - 1)
+
+    def linear_ramp_factor(min, max, dim):
+        """
+        Computes a linear ramp function used to smooth values between a minimum and maximum range.
+
+        Args:
+            min (float): Minimum value for the ramp function.
+            max (float): Maximum value for the ramp function.
+            dim (int): Dimensionality of the ramp tensor.
+
+        Returns:
+            torch.Tensor: A tensor of shape (dim,) with values linearly interpolated between 0 and 1,
+                clamped to the range [0, 1].
+        """
+        if min == max:
+            max += 0.001
+        linear_func = (torch.arange(dim, dtype=torch.float32) - min) / (max - min)
+        ramp_func = torch.clamp(linear_func, 0, 1)
+        return ramp_func
+
+    freqs = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
+    if seqlen > args.original_seq_len:
+        low, high = find_correction_range(
+            beta_fast, beta_slow, dim, base, args.original_seq_len
+        )
+        smooth = 1 - linear_ramp_factor(low, high, dim // 2)
+        freqs = freqs / factor * (1 - smooth) + freqs * smooth
+    return freqs
+
+
+def apply_rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
+    """
+    Applies rotary positional embeddings to the input tensor.
+
+    Args:
+        x (torch.Tensor): Input tensor with positional embeddings to be applied.
+        freqs_cis (torch.Tensor): Precomputed complex exponential values for positional embeddings.
+
+    Returns:
+        torch.Tensor: Tensor with rotary embeddings applied.
+    """
+    dtype = x.dtype
+    x = torch.view_as_complex(x.float().view(*x.shape[:-1], -1, 2))
+    freqs_cis = freqs_cis.view(x.size(0), x.size(1), 1, x.size(-1))
+    y = torch.view_as_real(x * freqs_cis).flatten(3)
+    return y.to(dtype)
+
+
+class RotaryEmbedding(nn.Module):
+    def __init__(self, args: DeepseekConfig):
+        super().__init__()
+        self.args = args
+
+    def _reset_base_freqs(self, device: torch.device):
+        with torch.device(device):
+            freqs = precompute_base_freqs(self.args).float()
+        assert not freqs.is_meta, "Cannot call precompute_freqs on meta device"
+        self.register_buffer("freqs", freqs, persistent=False)
+
+    @torch.no_grad()
+    def forward(self, position_ids: torch.Tensor) -> torch.Tensor:
+        if not hasattr(self, "freqs"):
+            self._reset_base_freqs(position_ids.device)
+
+        freqs = torch.matmul(
+            position_ids.unsqueeze(-1).float(), self.freqs.unsqueeze(0)
+        )
+        freqs_cis = torch.polar(
+            torch.ones_like(freqs, dtype=torch.float32), freqs.float()
+        )
+        return freqs_cis
+
+
+class MLA(nn.Module):
+    """
+    Multi-Headed Attention Layer (MLA).
+
+    Attributes:
+        dim (int): Dimensionality of the input features.
+        n_heads (int): Number of attention heads.
+        n_local_heads (int): Number of local attention heads for distributed systems.
+        q_lora_rank (int): Rank for low-rank query projection.
+        kv_lora_rank (int): Rank for low-rank key/value projection.
+        qk_nope_head_dim (int): Dimensionality of non-positional query/key projections.
+        qk_rope_head_dim (int): Dimensionality of rotary-positional query/key projections.
+        qk_head_dim (int): Total dimensionality of query/key projections.
+        v_head_dim (int): Dimensionality of value projections.
+        softmax_scale (float): Scaling factor for softmax in attention computation.
+    """
+
+    def __init__(self, args: DeepseekConfig):
+        super().__init__()
+        self.dim = args.dim
+        self.n_heads = args.n_heads
+        self.q_lora_rank = args.q_lora_rank
+        self.kv_lora_rank = args.kv_lora_rank
+        self.qk_nope_head_dim = args.qk_nope_head_dim
+        self.qk_rope_head_dim = args.qk_rope_head_dim
+        self.qk_head_dim = args.qk_nope_head_dim + args.qk_rope_head_dim
+        self.v_head_dim = args.v_head_dim
+
+        self.q_a_proj = nn.Linear(self.dim, self.q_lora_rank, bias=False)
+        self.q_a_layernorm = RMSNorm(self.q_lora_rank)
+        self.q_b_proj = nn.Linear(
+            self.q_lora_rank, self.n_heads * self.qk_head_dim, bias=False
+        )
+        self.kv_a_proj_with_mqa = nn.Linear(
+            self.dim, self.kv_lora_rank + self.qk_rope_head_dim, bias=False
+        )
+        self.kv_a_layernorm = RMSNorm(self.kv_lora_rank)
+        self.kv_b_proj = nn.Linear(
+            self.kv_lora_rank,
+            self.n_heads * (self.qk_nope_head_dim + self.v_head_dim),
+            bias=False,
+        )
+        self.o_proj = nn.Linear(self.n_heads * self.v_head_dim, self.dim, bias=False)
+        self.softmax_scale = self.qk_head_dim**-0.5
+        if args.max_seq_len > args.original_seq_len:
+            mscale = 0.1 * args.mscale * math.log(args.rope_factor) + 1.0
+            self.softmax_scale = self.softmax_scale * mscale * mscale
+
+        self.attn_module = DotProductAttention(
+            num_attention_heads=args.n_heads,
+            kv_channels=(
+                args.qk_nope_head_dim + args.qk_rope_head_dim,
+                args.v_head_dim,
+            ),
+            attn_mask_type="causal",
+            qkv_format="bshd",
+            softmax_scale=self.softmax_scale,
+        )
+        self.attn_func = self.attn_module.__call__
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        freqs_cis: torch.Tensor,
+    ):
+        """
+        Forward pass for the Multi-Headed Attention Layer (MLA).
+
+        Args:
+            x (torch.Tensor): Input tensor of shape (batch_size, local_seq_len, dim).
+            freqs_cis (torch.Tensor): Precomputed complex exponential values for rotary embeddings.
+
+        Returns:
+            torch.Tensor: Output tensor with the same shape as the input.
+        """
+        bsz, local_seq_len, _ = x.size()
+
+        q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(x)))
+        q = q.view(bsz, local_seq_len, self.n_heads, self.qk_head_dim)
+        q_nope, q_pe = torch.split(
+            q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1
+        )
+        q_pe = apply_rotary_emb(q_pe, freqs_cis)
+        q = torch.cat([q_nope, q_pe], dim=-1)
+
+        kv = self.kv_a_proj_with_mqa(x)
+        kv, k_pe = torch.split(kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
+        kv = self.kv_a_layernorm(kv)
+        k_pe = apply_rotary_emb(k_pe.unsqueeze(2), freqs_cis).squeeze(2)
+
+        kv = self.kv_b_proj(kv)
+        kv = kv.view(bsz, -1, self.n_heads, self.qk_nope_head_dim + self.v_head_dim)
+        k_nope, v = torch.split(kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1)
+        k_pe = k_pe.unsqueeze(2).expand([bsz, -1, self.n_heads, self.qk_rope_head_dim])
+        k = torch.cat([k_nope, k_pe], dim=-1)
+
+        x = self.attn_func(q, k, v)
+
+        x = self.o_proj(x.flatten(2))
+        return x
+
+
+class MLP(nn.Module):
+    """
+    Multi-Layer Perceptron (MLP) used as a feed-forward layer.
+
+    Attributes:
+        gate_proj (nn.Module): Linear layer for input-to-hidden transformation.
+        down_proj (nn.Module): Linear layer for hidden-to-output transformation.
+        up_proj (nn.Module): Additional linear layer for feature transformation.
+    """
+
+    def __init__(self, dim: int, inter_dim: int):
+        """
+        Initializes the MLP layer.
+
+        Args:
+            dim (int): Input and output dimensionality.
+            inter_dim (int): Hidden layer dimensionality.
+        """
+        super().__init__()
+        self.gate_proj = nn.Linear(dim, inter_dim, bias=False)
+        self.down_proj = nn.Linear(inter_dim, dim, bias=False)
+        self.up_proj = nn.Linear(dim, inter_dim, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass for the MLP layer.
+
+        Args:
+            x (torch.Tensor): Input tensor.
+
+        Returns:
+            torch.Tensor: Output tensor after MLP computation.
+        """
+        out = self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+        return out
+
+
+class Block(nn.Module):
+    """
+    Transformer block combining attention and feed-forward layers.
+
+    Attributes:
+        attn (nn.Module): Attention layer (MLA).
+        ffn (nn.Module): Feed-forward network (MLP or MoE).
+        input_layernorm (nn.Module): Layer normalization for attention.
+        post_attention_layernorm (nn.Module): Layer normalization for feed-forward network.
+    """
+
+    def __init__(self, layer_id: int, args: DeepseekConfig):
+        """
+        Initializes the Transformer block.
+
+        Args:
+            layer_id (int): Layer index in the transformer.
+            args (DeepseekConfig): Model arguments containing block parameters.
+        """
+        super().__init__()
+        self.self_attn = MLA(args)
+        if layer_id < args.n_dense_layers:
+            self.mlp = MLP(args.dim, args.inter_dim)
+        else:
+            self.mlp = MoE(args)
+        self.input_layernorm = RMSNorm(args.dim)
+        self.post_attention_layernorm = RMSNorm(args.dim)
+        self.layer_id = layer_id
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        freqs_cis: torch.Tensor,
+        padding_mask: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        """
+        Forward pass for the Transformer block.
+
+        Args:
+            x (torch.Tensor): Input tensor.
+            freqs_cis (torch.Tensor): Precomputed complex exponential values for rotary embeddings.
+            padding_mask (torch.Tensor): Boolean tensor indicating padding positions.
+
+        Returns:
+            torch.Tensor: Output tensor after block computation.
+            torch.Tensor | None: Auxiliary loss for load balancing (if applicable).
+        """
+
+        attn_out = self.self_attn(
+            x=self.input_layernorm(x),
+            freqs_cis=freqs_cis,
+        )
+        x = x + attn_out
+
+        mlp_out, aux_loss = self._mlp(
+            x=self.post_attention_layernorm(x),
+            padding_mask=padding_mask,
+        )
+        x = x + mlp_out
+
+        return x, aux_loss
+
+    def _mlp(
+        self,
+        x: torch.Tensor,
+        padding_mask: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        """
+        Applies the MLP layer to the input tensor.
+
+        Args:
+            x (torch.Tensor): Input tensor.
+            padding_mask (torch.Tensor): Boolean tensor indicating padding positions.
+
+        Returns:
+            torch.Tensor: Output tensor after MLP computation.
+            torch.Tensor | None: Auxiliary loss for load balancing (if applicable).
+        """
+        if isinstance(self.mlp, MLP):
+            return self.mlp(x), None
+        else:
+            assert isinstance(self.mlp, MoE)
+            return self.mlp(x, padding_mask)
+
+
+class TransformerModel(nn.Module):
+    """
+    Transformer model with positional embeddings, multiple layers, and output projection.
+
+    Attributes:
+        embed_tokens (nn.Module): Embedding layer for input tokens.
+        layers (torch.nn.ModuleDict): List of transformer blocks.
+        norm (nn.Module): Layer normalization applied after all blocks.
+        head (nn.Module): Output projection layer mapping to vocabulary size.
+        freqs_cis (torch.Tensor): Precomputed complex exponential values for rotary embeddings.
+    """
+
+    def __init__(self, args: DeepseekConfig):
+        """
+        Initializes the Transformer model.
+
+        Args:
+            args (DeepseekConfig): Model arguments containing transformer parameters.
+        """
+        # Linear.dtype = torch.float8_e4m3fn if args.dtype == "fp8" else torch.bfloat16
+        super().__init__()
+        self.args = args
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.dim)
+        self.layers = torch.nn.ModuleDict()
+        for layer_id in range(args.n_layers):
+            self.layers[str(layer_id)] = Block(layer_id, args)
+        self.norm = RMSNorm(args.dim)
+        self.rotary_emb = RotaryEmbedding(args)
+
+    def forward(
+        self,
+        tokens: torch.Tensor,
+        position_ids: torch.Tensor,
+        padding_mask: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        """
+        Forward pass for the Transformer model.
+
+        Important assumption: All devices on the same CP rank are fed the same inputs.
+
+        Args:
+            tokens (torch.Tensor): Input tensor of token IDs with shape (batch_size, local_seq_len).
+            position_ids (torch.Tensor): Starting position in the sequence for rotary embeddings. Defaults to 0.
+            padding_mask (torch.Tensor | None): Boolean tensor which indicates which tokens are padding.
+
+        Returns:
+            torch.Tensor: Logits tensor of shape (batch_size, vocab_size).
+            torch.Tensor | None: Auxiliary loss for load balancing (if applicable).
+        """
+        freqs_cis = self.rotary_emb(position_ids)
+
+        h = self.embed_tokens(tokens) if self.embed_tokens else tokens
+
+        # Apply the transformer layers.
+        aux_losses = []
+        for layer in self.layers.values():
+            h, aux_loss = layer(
+                x=h,
+                freqs_cis=freqs_cis,
+                padding_mask=padding_mask,
+            )
+            if aux_loss is not None:
+                aux_losses.append(aux_loss)
+        final_aux_loss = torch.stack(aux_losses).mean() if aux_losses else None
+
+        h = self.norm(h) if self.norm else h
+        return h, final_aux_loss
+
+
+class Transformer(nn.Module):
+    def __init__(self, args: DeepseekConfig):
+        super().__init__()
+        self.model = TransformerModel(args)
+        self.lm_head = nn.Linear(
+            args.dim, args.vocab_size, dtype=torch.get_default_dtype(), bias=False
+        )
+
+    def forward(
+        self,
+        tokens: torch.Tensor,
+        position_ids: torch.Tensor,
+        padding_mask: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        """
+        Forward pass for the Transformer model.
+
+        Important assumption: All devices on the same CP rank are fed the same inputs.
+
+        Args:
+            tokens (torch.Tensor): Input tensor of token IDs with shape (batch_size, local_seq_len).
+            position_ids (torch.Tensor): Input tensor containing the indices of the tokens.
+                The shape is (batch_size, local_seq_len).
+            padding_mask (torch.Tensor): Boolean tensor which indicates which tokens are padding.
+
+        Returns:
+            torch.Tensor: Logits tensor of shape (batch_size, vocab_size).
+            torch.Tensor | None: Auxiliary loss for load balancing (if applicable).
+        """
+        h, aux_loss = self.model(
+            tokens=tokens,
+            position_ids=position_ids,
+            padding_mask=padding_mask,
+        )
+        logits = self.lm_head(h) if self.lm_head else h
+        return logits, aux_loss
diff --git a/cosmos_rl/policy/model/deepseek_v3/parallelize.py b/cosmos_rl/policy/model/deepseek_v3/parallelize.py
new file mode 100644
index 000000000..bfd8126bf
--- /dev/null
+++ b/cosmos_rl/policy/model/deepseek_v3/parallelize.py
@@ -0,0 +1,283 @@
+# 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.
+
+
+import functools
+import os
+from typing import Callable, Optional
+
+import torch
+from torch import nn
+from torch._utils import _get_available_device_type, _get_device_module
+from torch.distributed.device_mesh import DeviceMesh
+
+try:
+    from torch.distributed.tensor import Shard, distribute_module, distribute_tensor
+except ImportError:
+    print("torch.distributed.tensor is not available. DeepSeek model will not work.")
+
+from torch.distributed.tensor.parallel import ParallelStyle, parallelize_module
+
+try:
+    from torch.distributed.fsdp import fully_shard
+except ImportError:
+    print("torch.distributed.fsdp is not available. DeepSeek model will not work.")
+
+from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
+    checkpoint_wrapper as ptd_checkpoint_wrapper,
+)
+from transformer_engine.pytorch.attention import DotProductAttention
+
+from cosmos_rl.policy.config import Config as CosmosConfig
+from cosmos_rl.policy.kernel.moe.moe import GroupedExpertsDeepEP, MoE
+from cosmos_rl.utils.parallelism import ParallelDims
+from cosmos_rl.utils.ulysses import swizzle_cp_forward, ulysses_attn_func
+
+
+def _get_dp_mesh(
+    world_mesh: DeviceMesh, parallel_dims: ParallelDims
+) -> DeviceMesh | None:
+    """
+    Gets the HSDP or FSDP mesh based on the parallelism settings.
+    If only FSDP is used, the model parameters are sharded on the FSDP mesh.
+    If HSDP is used, the model parameters are both replicated and sharded
+    across different dimensions of the mesh.
+    """
+    if parallel_dims.dp_shard_enabled or parallel_dims.cp_enabled:
+        if parallel_dims.dp_replicate_enabled:
+            # HSDP: Replicate on `dp_replicate` dim and shard on `dp_shard_cp` dim.
+            dp_mesh_dim_names = ("dp_replicate", "dp_shard_cp")
+        else:
+            # FSDP: Replicate on `dp_shard_cp` dim.
+            dp_mesh_dim_names = ("dp_shard_cp",)
+        return world_mesh[tuple(dp_mesh_dim_names)]
+    else:
+        assert not parallel_dims.dp_replicate_enabled
+        return None
+
+
+def _apply_cp(model: nn.Module, cp_mesh: DeviceMesh, parallel_dims: ParallelDims):
+    """Apply Context Parallel to the model."""
+    assert cp_mesh.size() > 1
+    assert cp_mesh.ndim == 1
+
+    if model.model is not None:
+        _model = model.model
+
+        for _, block in _model.model.layers.named_children():
+            attn_module = block.self_attn.attn_module
+            assert isinstance(
+                attn_module, DotProductAttention
+            ), "Context parallelism is only supported for DotProductAttention."
+            assert attn_module.num_attention_heads % cp_mesh.size() == 0, (
+                f"Number of attention heads {attn_module.num_attention_heads} must be divisible by "
+                f"context parallel size {cp_mesh.size()}"
+            )
+
+            attn_module_with_cp = DotProductAttention(
+                num_attention_heads=attn_module.num_attention_heads // cp_mesh.size(),
+                kv_channels=(
+                    attn_module.hidden_size_per_attention_head_k,
+                    attn_module.hidden_size_per_attention_head_v,
+                ),
+                attn_mask_type=attn_module.attn_mask_type,
+                qkv_format=attn_module.qkv_format,
+                softmax_scale=block.self_attn.softmax_scale,
+            )
+            attn_func_with_cp = ulysses_attn_func(
+                attn_module_with_cp.__call__,
+                cp_mesh,
+            )
+
+            block.self_attn.attn_module = attn_module_with_cp
+            block.self_attn.attn_func = attn_func_with_cp
+
+        if _model.lm_head is not None:
+            # Apply CP to the lm_head to get the logits for all tokens.
+            swizzle_cp_forward(_model.lm_head, parallel_dims)
+
+
+class _ExpertParallel(ParallelStyle):
+    """
+    ExpertParallel class is used to shard the MoE parameters on the EP mesh.
+    Dim `0` of each parameter is sharded since that is the expert dimension.
+    """
+
+    def _partition_fn(self, name, module, device_mesh):
+        # shard on the expert dimension
+        assert device_mesh.ndim == 1
+
+        for name, param in module.named_parameters(recurse=False):
+            dist_param = nn.Parameter(distribute_tensor(param, device_mesh, [Shard(0)]))
+            module.register_parameter(name, dist_param)
+
+        if isinstance(module, GroupedExpertsDeepEP):
+            module.init_token_dispatcher(ep_mesh=device_mesh)
+
+    def _apply(self, module: nn.Module, device_mesh: DeviceMesh) -> nn.Module:
+        return distribute_module(
+            module,
+            device_mesh,
+            self._partition_fn,
+        )
+
+
+def _apply_ep(model: nn.Module, ep_mesh: DeviceMesh):
+    """Applies EP to MoE module."""
+    assert ep_mesh.size() > 1
+
+    if model.model is not None:
+        _model = model.model
+        for _, block in _model.model.layers.named_children():
+            if isinstance(block.mlp, MoE):
+                parallelize_module(
+                    module=block.mlp.experts,
+                    device_mesh=ep_mesh,
+                    parallelize_plan=_ExpertParallel(),
+                )
+
+
+def _apply_ac(model: nn.Module):
+    """Apply activation checkpointing to the model."""
+    if model.model is not None:
+        _model = model.model
+        for layer_id, block in _model.model.layers.named_children():
+            block = ptd_checkpoint_wrapper(block, preserve_rng_state=False)
+            _model.model.layers.register_module(layer_id, block)
+
+
+def _apply_fsdp(
+    model: nn.Module,
+    meshes: dict[str, DeviceMesh],
+    parallel_dims: ParallelDims,
+):
+    """Apply FSDP sharding to model layers using data parallel mesh."""
+    default_dp_mesh = _get_dp_mesh(meshes["default"], parallel_dims)
+    if default_dp_mesh is None:
+        return
+
+    pp_enabled = parallel_dims.pp_enabled
+
+    fully_shard_default = functools.partial(
+        fully_shard,
+        mesh=default_dp_mesh,
+        reshard_after_forward=not pp_enabled,
+    )
+
+    if model.model is not None:
+        _model = model.model
+
+        for _, block in _model.model.layers.named_children():
+            if isinstance(block.mlp, MoE) and parallel_dims.dp_shard_with_ep_enabled:
+                # Apply FSDP on dim=1 for grouped experts since we may have more
+                # shards than experts (dim=0).
+                assert "moe" in meshes
+                fully_shard(
+                    block.mlp.experts,
+                    mesh=meshes["moe"]["dp_shard_with_ep"],
+                    shard_placement_fn=lambda _: Shard(1),
+                    reshard_after_forward=not pp_enabled,
+                )
+
+            # If FSDP is disabled for grouped experts because the parameters are already
+            # fully sharded by PP and EP, then we need to explicitly remove the parameters
+            # from FSDP for the transformer block.
+            # If FSDP is enabled for grouped experts, the parameters are automatically
+            # removed from the FSDP for the transformer block due to the rules of the
+            # PyTorch FSDP implementation.
+            ignored_params = None
+            if (
+                isinstance(block.mlp, MoE)
+                and parallel_dims.ep_enabled
+                and not parallel_dims.dp_shard_with_ep_enabled
+            ):
+                ignored_params = set(block.mlp.experts.parameters())
+
+            fully_shard_default(block, ignored_params=ignored_params)
+
+        fully_shard_default(_model)
+
+
+def _get_device_info():
+    device_type = _get_available_device_type()
+    if device_type is None:
+        device_type = "cuda"  # default device_type: cuda
+    device_module = _get_device_module(device_type)  # default device_module: torch.cuda
+    return device_type, device_module
+
+
+def _init_meshes(
+    parallel_dims: ParallelDims,
+) -> dict[str, DeviceMesh]:
+    """
+    Initialize the meshes for the model. There are generally two meshes, "default"
+    for non-MoE modules and "moe" for MoE modules. Each mesh contains several
+    dimensions for parallelization.
+
+    Args:
+        parallelism_config (TrainingConfig): The parallelism configuration for the model.
+
+    Returns:
+        meshes (dict[str, DeviceMesh]): The meshes for the model.
+    """
+
+    device_type, device_module = _get_device_info()
+
+    local_rank = int(os.getenv("LOCAL_RANK", 0))
+    device = torch.device(f"{device_type}:{local_rank}")
+    device_module.set_device(device)
+    meshes = parallel_dims.build_meshes_with_ep(device_type=device_type)
+    return meshes
+
+
+def parallelize_model(
+    model: nn.Module,
+    parallel_dims: ParallelDims,
+    config: CosmosConfig,
+    pp_loss_fn: Optional[Callable],
+) -> nn.Module:
+    """
+    Parallelizes the DeepSeek model based on the provided meshes and parallel dimensions.
+
+    Args:
+        model (nn.Module): The DeepSeek model to parallelize. Note that this maybe
+            a model part due to pipelining. We must check for the presence of an
+            nn.Module before executing the relevant parallelization functions.
+
+    Returns:
+        nn.Module: The parallelized DeepSeek model.
+    """
+    meshes = _init_meshes(parallel_dims)
+    del config
+    del pp_loss_fn
+
+    assert model.config.n_routed_experts % parallel_dims.ep == 0, (
+        f"n_routed_experts {model.config.n_routed_experts} must be divisible by "
+        f"expert_parallel_degree {parallel_dims.ep}"
+    )
+    assert parallel_dims.tp == 1, "Tensor parallelism not support for DeepSeek model"
+
+    if parallel_dims.cp_enabled:
+        _apply_cp(model, meshes["default"]["cp"], parallel_dims)
+
+    if parallel_dims.ep_enabled:
+        assert "moe" in meshes
+        _apply_ep(model, meshes["moe"]["ep"])
+
+    _apply_ac(model)
+
+    _apply_fsdp(model, meshes, parallel_dims)
+
+    return None, None
diff --git a/cosmos_rl/policy/model/deepseek_v3/weight_mapper.py b/cosmos_rl/policy/model/deepseek_v3/weight_mapper.py
new file mode 100644
index 000000000..3a21ca04b
--- /dev/null
+++ b/cosmos_rl/policy/model/deepseek_v3/weight_mapper.py
@@ -0,0 +1,462 @@
+# 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.
+
+import re
+from typing import Any, Dict, List, Optional, Tuple
+
+import torch
+import triton
+import triton.language as tl
+from transformers import AutoConfig
+
+from cosmos_rl.policy.model.base import WeightMapper
+from cosmos_rl.utils import util
+from cosmos_rl.utils.logging import logger
+from cosmos_rl.utils.parallelism import ParallelDims
+from cosmos_rl.utils.parallelism_registry import (
+    ParallelismStrategyRole,
+    get_rollout_parallelism_strategy,
+    register_parallelism_strategy,
+)
+
+# Pre-compile regex patterns for better performance
+_LAYER_ATTN_NORM_PATTERN = re.compile(
+    r"model\.model\.layers\.(\d+)\.self_attn\.(q_norm|k_norm|v_norm)\.(weight|bias)"
+)
+_LAYER_INPUT_LAYERNORM_PATTERN = re.compile(
+    r"model\.model\.layers\.(\d+)\.input_layernorm\.(weight|bias)"
+)
+_LAYER_ATTN_PROJ_PATTERN = re.compile(
+    r"model\.model\.layers\.(\d+)\.self_attn\.(q_proj|kv_b_proj|q_b_proj)\.(weight|bias)"
+)
+_LAYER_ATTN_KV_A_PATTERN = re.compile(
+    r"model\.model\.layers\.(\d+)\.self_attn\.(kv_a_proj_with_mqa)\.(weight|bias)"
+)
+_LAYER_ATTN_O_PATTERN = re.compile(
+    r"model\.model\.layers\.(\d+)\.self_attn\.(o_proj)\.(weight|bias)"
+)
+_LAYER_MLP_EXPERTS_PATTERN = re.compile(
+    r"model\.model\.layers\.(\d+)\.mlp\.experts\.(\d+)\.(up_proj|gate_proj|down_proj)\.(weight|bias)"
+)
+_LAYER_MLP_SHARED_EXPERTS_PATTERN = re.compile(
+    r"model\.model\.layers\.(\d+)\.mlp\.shared_experts\.(up_proj|gate_proj)\.(weight|bias)"
+)
+_LAYER_MLP_SHARED_EXPERTS_DOWN_PATTERN = re.compile(
+    r"model\.model\.layers\.(\d+)\.mlp\.shared_experts\.down_proj\.(weight|bias)"
+)
+_LAYER_MLP_PROJ_PATTERN = re.compile(
+    r"model\.model\.layers\.(\d+)\.mlp\.(up_proj|gate_proj)\.(weight|bias)"
+)
+_LAYER_MLP_DOWN_PATTERN = re.compile(
+    r"model\.model\.layers\.(\d+)\.mlp\.down_proj\.(weight|bias)"
+)
+_LAYER_POST_ATTENTION_LAYERNORM_PATTERN = re.compile(
+    r"model\.model\.layers\.(\d+)\.post_attention_layernorm\.(weight|bias)"
+)
+_LAYER_KV_A_LAYERNORM_PATTERN = re.compile(
+    r"model\.model\.layers\.(\d+)\.self_attn\.kv_a_layernorm\.(weight|bias)"
+)
+_LAYER_Q_A_PATTERN = re.compile(
+    r"model\.model\.layers\.(\d+)\.self_attn\.(q_a_layernorm|q_a_proj)\.(weight|bias)"
+)
+_LAYER_MLP_GATE_WEIGHT_PATTERN = re.compile(
+    r"model\.model\.layers\.(\d+)\.mlp\.gate\.weight"
+)
+_LAYER_ROTARY_EMB_PATTERN = re.compile(
+    r"model\.model\.layers\.(\d+)\.self_attn\.rotary_emb\.inv_freq"
+)
+_LAYER_MLP_GATE_BIAS_PATTERN = re.compile(
+    r"model\.model\.layers\.(\d+)\.mlp\.gate\.e_score_correction_bias"
+)
+
+# Pre-define common suffixes for faster string matching
+_LM_HEAD_WEIGHT = "lm_head.weight"
+_LM_HEAD_BIAS = "lm_head.bias"
+_EMBED_TOKENS_WEIGHT = "embed_tokens.weight"
+_NORM_WEIGHT = "norm.weight"
+_NORM_BIAS = "norm.bias"
+_BIAS_SUFFIX = ".bias"
+
+
+def map_key_from_hf(name: str, src_model_type: str) -> str:
+    return "model." + name
+
+
+def convert_weight_from_hf(
+    tensor: torch.Tensor,
+    name: str,
+    src_model_type: str,
+    parallel_dims: ParallelDims,
+    n_experts: int,
+    ignore_unknown_weights: bool = False,
+) -> Tuple[Optional[str], Optional[torch.Tensor], Optional[int]]:
+    tp_ep_rank, tp_ep_size = parallel_dims.tp_coord
+    assert n_experts % tp_ep_size == 0, "n_experts must be divisible by tp_ep_size"
+
+    if parallel_dims.dp_shard_enabled or parallel_dims.cp_enabled:
+        dp_shard_rank = parallel_dims.mesh[tuple(("dp_shard_cp",))].get_local_rank()
+        dp_shard_size = parallel_dims.mesh[tuple(("dp_shard_cp",))].size()
+    else:
+        dp_shard_rank = 0
+        dp_shard_size = 1
+
+    # Expert weight are aggregated into (n_experts, in_features, out_features)
+    # Weight are loaded in (out_features, in_features) shape
+    # So we do not do FSDP sharding on expert weights, instead we filter by expert id
+    should_do_fsdp_sharding = True
+
+    dest_name = map_key_from_hf(name, src_model_type)
+
+    # Fast path for common cases using string suffix matching
+    if dest_name.endswith(_LM_HEAD_WEIGHT):
+        shard = tensor.tensor_split(tp_ep_size, dim=0)[tp_ep_rank]
+    elif dest_name.endswith(_LM_HEAD_BIAS):
+        shard = tensor
+    elif dest_name.endswith(_EMBED_TOKENS_WEIGHT):
+        shard = tensor.tensor_split(tp_ep_size, dim=0)[tp_ep_rank]
+    elif dest_name.endswith(_NORM_WEIGHT) or dest_name.endswith(_NORM_BIAS):
+        shard = tensor
+    # Regex matching for complex patterns
+    elif (match := _LAYER_ATTN_NORM_PATTERN.search(dest_name)) is not None:
+        shard = tensor
+    elif (match := _LAYER_INPUT_LAYERNORM_PATTERN.search(dest_name)) is not None:
+        shard = tensor
+    elif (match := _LAYER_ATTN_PROJ_PATTERN.search(dest_name)) is not None:
+        shard = tensor.tensor_split(tp_ep_size, dim=0)[tp_ep_rank]
+    elif (match := _LAYER_ATTN_KV_A_PATTERN.search(dest_name)) is not None:
+        shard = tensor
+    elif (match := _LAYER_ATTN_O_PATTERN.search(dest_name)) is not None:
+        if dest_name.endswith(_BIAS_SUFFIX):
+            shard = tensor
+        else:
+            shard = tensor.tensor_split(tp_ep_size, dim=-1)[tp_ep_rank]
+    elif (match := _LAYER_MLP_EXPERTS_PATTERN.search(dest_name)) is not None:
+        # Check whether this expert belongs to the current process
+        # Groups example (with 32 experts, and 4 EP groups):
+        #  EP=0: 0, 1, 2, 3, 4, 5, 6, 7
+        #  EP=1: 8, 9, 10, 11, 12, 13, 14, 15
+        #  EP=2: 16, 17, 18, 19, 20, 21, 22, 23
+        #  EP=3: 24, 25, 26, 27, 28, 29, 30, 31
+        n_expert_per_ep = n_experts // tp_ep_size
+        belongs_to_current_ep = (
+            tp_ep_rank * n_expert_per_ep
+            <= int(match.group(2))  # Expert index
+            < (tp_ep_rank + 1) * n_expert_per_ep
+        )
+        belongs_to_current_dp_shard = (
+            int(match.group(2)) - tp_ep_rank * n_expert_per_ep
+        ) // (n_expert_per_ep // dp_shard_size) == dp_shard_rank
+        if belongs_to_current_ep and belongs_to_current_dp_shard:
+            should_do_fsdp_sharding = False
+            shard = tensor
+        else:
+            # If the expert does not belong to the current process, return None to skip this weight
+            return None, None, None
+    elif (match := _LAYER_MLP_SHARED_EXPERTS_PATTERN.search(dest_name)) is not None:
+        shard = tensor.tensor_split(tp_ep_size, dim=0)[tp_ep_rank]
+    elif (
+        match := _LAYER_MLP_SHARED_EXPERTS_DOWN_PATTERN.search(dest_name)
+    ) is not None:
+        if dest_name.endswith(_BIAS_SUFFIX):
+            shard = tensor
+        else:
+            shard = tensor.tensor_split(tp_ep_size, dim=-1)[tp_ep_rank]
+    elif (match := _LAYER_MLP_PROJ_PATTERN.search(dest_name)) is not None:
+        shard = tensor.tensor_split(tp_ep_size, dim=0)[tp_ep_rank]
+    elif (match := _LAYER_MLP_DOWN_PATTERN.search(dest_name)) is not None:
+        if dest_name.endswith(_BIAS_SUFFIX):
+            shard = tensor
+        else:
+            shard = tensor.tensor_split(tp_ep_size, dim=-1)[tp_ep_rank]
+    elif (
+        match := _LAYER_POST_ATTENTION_LAYERNORM_PATTERN.search(dest_name)
+    ) is not None:
+        shard = tensor
+    elif (match := _LAYER_KV_A_LAYERNORM_PATTERN.search(dest_name)) is not None:
+        shard = tensor
+    elif (match := _LAYER_Q_A_PATTERN.search(dest_name)) is not None:
+        shard = tensor
+    elif (match := _LAYER_MLP_GATE_WEIGHT_PATTERN.search(dest_name)) is not None:
+        # TODO(cjx): Small enough, forbid FSDP sharding is better
+        shard = tensor
+    elif (match := _LAYER_ROTARY_EMB_PATTERN.search(dest_name)) is not None:
+        return None, None, None
+    elif (match := _LAYER_MLP_GATE_BIAS_PATTERN.search(dest_name)) is not None:
+        shard = tensor
+    elif not ignore_unknown_weights:
+        raise ValueError(f"Unsupported weight: {dest_name}")
+    else:
+        return None, None, None
+
+    # Do FSDP sharding
+    shard = shard.contiguous()
+    if should_do_fsdp_sharding:
+        if match := _LAYER_ATTN_KV_A_PATTERN.search(dest_name) is not None and (
+            576 % dp_shard_size != 0
+        ):
+            tensor_shard_size = 576 // dp_shard_size + 1
+            if dp_shard_rank < (576 // tensor_shard_size):
+                shard = shard[
+                    dp_shard_rank * tensor_shard_size : (dp_shard_rank + 1)
+                    * tensor_shard_size
+                ]
+            else:
+                shard = shard[dp_shard_rank * tensor_shard_size :]
+        else:
+            shard = shard.tensor_split(dp_shard_size, dim=0)[dp_shard_rank]
+
+    expert_id = None
+    if match := re.search(
+        r"model\.model\.layers\.(\d+)\.mlp\.experts\.(\d+)\.(up_proj|gate_proj|down_proj)\.(weight|bias)",
+        dest_name,
+    ):
+        expert_id = int(match.group(2))
+        dest_name = dest_name.replace(f"experts.{expert_id}.", "experts.")
+        dest_name = dest_name.replace(".weight", "s")
+
+    return dest_name, shard.contiguous(), expert_id
+
+
+@triton.jit
+def weight_dequant_kernel(x_ptr, s_ptr, y_ptr, M, N, BLOCK_SIZE: tl.constexpr):
+    pid_m = tl.program_id(axis=0)
+    pid_n = tl.program_id(axis=1)
+    n = tl.cdiv(N, BLOCK_SIZE)
+    offs_m = pid_m * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    offs_n = pid_n * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    offs = offs_m[:, None] * N + offs_n[None, :]
+    mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
+    x = tl.load(x_ptr + offs, mask=mask).to(tl.float32)
+    s = tl.load(s_ptr + pid_m * n + pid_n)
+    y = x * s
+    tl.store(y_ptr + offs, y, mask=mask)
+
+
+def weight_dequant(
+    x: torch.Tensor, s: torch.Tensor, block_size: int = 128
+) -> torch.Tensor:
+    assert x.is_contiguous() and s.is_contiguous()
+    assert x.dim() == 2 and s.dim() == 2
+    M, N = x.size()
+    y = torch.empty_like(x, dtype=torch.get_default_dtype())
+
+    def grid(meta):
+        return (triton.cdiv(M, meta["BLOCK_SIZE"]), triton.cdiv(N, meta["BLOCK_SIZE"]))
+
+    weight_dequant_kernel[grid](x, s, y, M, N, BLOCK_SIZE=block_size)
+    return y
+
+
+class DeepseekV3MoEWeightMapper(WeightMapper):
+    def __init__(self, hf_config: AutoConfig):
+        super().__init__(hf_config)
+
+    def _rollout_vllm_name_to_hf(self, rollout_weight_name: str) -> str:
+        # TODO(aazzolini): 2.2. Implement name_to_hf correctly.
+        if not rollout_weight_name == "lm_head.weight":
+            if "experts.w13_weight" in rollout_weight_name:
+                return rollout_weight_name.replace(
+                    "experts.w13_weight", "gate_up_proj.weight"
+                )
+            elif "experts.w2_weight" in rollout_weight_name:
+                return rollout_weight_name.replace(
+                    "experts.w2_weight", "down_proj.weight"
+                )
+        return rollout_weight_name
+
+    def _rollout_split_qkv_weight(self, name, weight: torch.Tensor):
+        # weight has shape [q_num_heads * head_dim + k_num_heads * head_dim + v_num_heads * head_dim, hidden_dim]
+        # TODO(aazzolini): 2.3 Implement split_qkv_weight
+        shares = self.kv_head_ratio + 2
+        dim_0 = weight.shape[0]  # for both weight and bias
+        unit_dim = dim_0 // shares
+
+        q_weight = weight[: unit_dim * self.kv_head_ratio]
+        k_weight = weight[
+            unit_dim * self.kv_head_ratio : unit_dim * (self.kv_head_ratio + 1)
+        ]
+        v_weight = weight[unit_dim * (self.kv_head_ratio + 1) :]
+        return q_weight, k_weight, v_weight
+
+    def _split_gate_proj_weight(self, name, weight: torch.Tensor):
+        # weight has shape [..., 2 * x, hidden_dim]
+        # split it into [..., :x, hidden_dim] and [..., x:, hidden_dim]
+
+        # Note that for MoE we have an extra dimension on the left for for dense we dont.
+        split_idx = weight.shape[-2] // 2
+        gate_proj_weight = weight[..., :split_idx, :]
+        up_proj_weight = weight[..., split_idx:, :]
+        return gate_proj_weight, up_proj_weight
+
+    def rollout_prepare_recv(
+        self, vllm_model: Any
+    ) -> Tuple[Dict[str, torch.Tensor], List[Tuple[str, torch.Size]]]:
+        """
+        Given the rollout (e.g. VLLM) nn.Module, return the list of HuggingFace-compatible
+        parameter names along with their tensor views and shapes. Param names produced this
+        method should exactly match policy_model.sorted_params.
+
+        Args:
+           vllm_model: rollout nn.Module
+
+        Returns:
+          Tuple[
+            Dict[
+              str,           # compatible_key
+              Tensor         # param Tensor
+            ],               # compatible weight map
+            List[
+              List[
+                Tuple[
+                  str,         # compatible_key
+                  int,         # param tensor ndim
+                ]
+              ]              # param group (?)
+            ]                # compatible weight list
+
+            Where "compatible_key" is the HuggingFace param name
+        """
+        recv_key_n_rank_list = []
+        compatible_weight_map = {}
+        for param_name, param in vllm_model.named_parameters():
+            group_keys = []
+            compatible_key = self._rollout_vllm_name_to_hf(param_name)
+            # logger.info(f"[Rollout] compatible_key: {compatible_key}")
+            if "qkv_proj" in compatible_key:
+                # must be inplace slicing.
+                # split qkv weight
+                q_weight, k_weight, v_weight = self._rollout_split_qkv_weight(
+                    compatible_key, param
+                )
+                q_proj_weight_key = compatible_key.replace("qkv_proj", "q_proj")
+                k_proj_weight_key = compatible_key.replace("qkv_proj", "k_proj")
+                v_proj_weight_key = compatible_key.replace("qkv_proj", "v_proj")
+                compatible_weight_map[q_proj_weight_key] = q_weight
+                group_keys.append((q_proj_weight_key, q_weight.ndim))
+                compatible_weight_map[k_proj_weight_key] = k_weight
+                group_keys.append((k_proj_weight_key, k_weight.ndim))
+                compatible_weight_map[v_proj_weight_key] = v_weight
+                group_keys.append((v_proj_weight_key, v_weight.ndim))
+            elif "gate_up_proj" in compatible_key:
+                # split gate and up proj
+                gate_proj_weight, up_proj_weight = self._split_gate_proj_weight(
+                    compatible_key, param
+                )
+                gate_proj_weight_key = compatible_key.replace(
+                    "gate_up_proj", "gate_proj"
+                )
+                compatible_weight_map[gate_proj_weight_key] = gate_proj_weight
+                group_keys.append((gate_proj_weight_key, gate_proj_weight.ndim))
+
+                up_proj_weight_key = compatible_key.replace("gate_up_proj", "up_proj")
+                compatible_weight_map[up_proj_weight_key] = up_proj_weight
+                group_keys.append((up_proj_weight_key, up_proj_weight.ndim))
+            else:
+                compatible_weight_map[compatible_key] = param
+                group_keys.append((compatible_key, param.ndim))
+            recv_key_n_rank_list.append(group_keys)
+
+        return compatible_weight_map, recv_key_n_rank_list
+
+    def policy_map_local_key_to_hf_key(self, name: str) -> str:
+        name = util.clear_weight_name(name)
+
+        name = name[6:]
+        name = name.replace(".mlp.experts.down_projs", ".mlp.down_proj.weight")
+        name = name.replace(".mlp.experts.up_projs", ".mlp.up_proj.weight")
+        name = name.replace(".mlp.experts.gate_projs", ".mlp.gate_proj.weight")
+
+        return name
+
+    def get_rollout_parallelism_strategy(self):
+        return [get_rollout_parallelism_strategy("deepseek_v3")]
+
+
+@register_parallelism_strategy("deepseek_v3", role=ParallelismStrategyRole.ROLLOUT)
+def map_weight_parallel_dims(
+    n_dim: int,
+    dest_name: str,
+    parallel_dims: ParallelDims,
+    model_config: Any,
+    # ) -> Tuple[Dict[str, int], list, Dict[int, list]]:
+) -> Tuple[Dict[str, int], Dict[int, list], int]:
+    """
+    For the given HuggingFace parameter name (dest_name), return the map from
+    parallel dimension name (e.g. "tp", ...) to the corresponding dimension of the tensor.
+    """
+    # TODO(aazzolini): 1.2 implement rollout_parallelism_strategy (tp,ep)
+    dims_map = {}
+
+    tp_dim_map = {
+        # deepseekv3 attention weights
+        ".self_attn.q_a_proj.": None,
+        ".self_attn.q_a_layernorm.": None,
+        ".self_attn.q_b_proj.": 0,
+        ".self_attn.kv_a_proj_with_mqa.": None,
+        ".self_attn.kv_a_layernorm.": None,
+        ".self_attn.kv_b_proj.": 0,
+        ".self_attn.o_proj.": 1,
+        # deepseekv3 dense+moe mlp weights
+        # note on the training side gate_up_proj is fused.
+        ".mlp.gate_proj.": -2,
+        ".mlp.up_proj.": -2,
+        ".mlp.down_proj.": -1,
+        # deepseekv3 moe-related weights
+        ".mlp.gate.": None,
+        ".mlp.shared_experts.gate_proj.": 0,
+        ".mlp.shared_experts.up_proj.": 0,
+        ".mlp.shared_experts.down_proj.": 1,
+        # deepseekv3 layernorms
+        ".input_layernorm.weight": None,
+        ".post_attention_layernorm.weight": None,
+        # deepseekv3 toplevel
+        ".norm.weight": None,
+        "lm_head.weight": 0,
+        ".embed_tokens.": 0,
+        # vit attention
+        ".attention_norm.": None,
+        ".attention.wq_proj.": 0,
+        ".attention.wk_proj.": 0,
+        ".attention.wv_proj.": 0,
+        ".mlp.fc1.bias": 0,
+        ".mlp.fc1.weight": 0,
+        ".mlp.fc2.weight": 1,
+    }
+
+    tp_dim = None
+    for pattern, dim in tp_dim_map.items():
+        if pattern in dest_name:
+            tp_dim = dim
+            if tp_dim is not None and tp_dim < 0:
+                tp_dim += n_dim
+            break
+
+    logger.debug(
+        f"Rollout parallelism mapping: Dest_name: {dest_name}, tp_dim: {tp_dim}"
+    )
+
+    if tp_dim is not None:
+        dims_map["tp"] = tp_dim
+
+    # if dp_shard_size > 1:
+    #    dims_map["dp_shard_cp"] = 0
+
+    tensor_dim_to_parallel_map = {}
+    for k, v in dims_map.items():
+        if v not in tensor_dim_to_parallel_map:
+            tensor_dim_to_parallel_map[v] = []
+        tensor_dim_to_parallel_map[v].append(k)
+    pp_rank = 0
+    return dims_map, tensor_dim_to_parallel_map, pp_rank
diff --git a/cosmos_rl/policy/model/gpt/__init__.py b/cosmos_rl/policy/model/gpt/__init__.py
index 14c135f41..41e3cbe08 100644
--- a/cosmos_rl/policy/model/gpt/__init__.py
+++ b/cosmos_rl/policy/model/gpt/__init__.py
@@ -530,6 +530,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.
diff --git a/cosmos_rl/policy/model/hf_models/__init__.py b/cosmos_rl/policy/model/hf_models/__init__.py
index 3d54bac8d..969642295 100644
--- a/cosmos_rl/policy/model/hf_models/__init__.py
+++ b/cosmos_rl/policy/model/hf_models/__init__.py
@@ -361,6 +361,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.
diff --git a/cosmos_rl/policy/model/qwen2_5_vl/__init__.py b/cosmos_rl/policy/model/qwen2_5_vl/__init__.py
index 2ea48e322..b011aed10 100644
--- a/cosmos_rl/policy/model/qwen2_5_vl/__init__.py
+++ b/cosmos_rl/policy/model/qwen2_5_vl/__init__.py
@@ -1062,6 +1062,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.
diff --git a/cosmos_rl/policy/model/qwen3_moe/__init__.py b/cosmos_rl/policy/model/qwen3_moe/__init__.py
index a9872f04b..fcc22c473 100644
--- a/cosmos_rl/policy/model/qwen3_moe/__init__.py
+++ b/cosmos_rl/policy/model/qwen3_moe/__init__.py
@@ -794,6 +794,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.
diff --git a/cosmos_rl/policy/trainer/grpo_trainer.py b/cosmos_rl/policy/trainer/grpo_trainer.py
index 6cfe33b2c..7aa8de069 100644
--- a/cosmos_rl/policy/trainer/grpo_trainer.py
+++ b/cosmos_rl/policy/trainer/grpo_trainer.py
@@ -376,6 +376,7 @@ def model_load_from_hf(self):
             self.parallel_dims,
             self.device,
             revision=self.config.policy.model_revision,
+            dcp_snapshot_path=config.policy.dcp_snapshot_path,
         )
         self.model.train()
         self.model_ready = True
diff --git a/cosmos_rl/policy/trainer/sft_trainer.py b/cosmos_rl/policy/trainer/sft_trainer.py
index 881876cef..432504a14 100644
--- a/cosmos_rl/policy/trainer/sft_trainer.py
+++ b/cosmos_rl/policy/trainer/sft_trainer.py
@@ -341,6 +341,7 @@ def __init__(self, config: CosmosConfig, parallel_dims: ParallelDims):
                     parallel_dims,
                     self.device,
                     revision=config.policy.model_revision,
+                    dcp_snapshot_path=config.policy.dcp_snapshot_path,
                 )
         else:
             self.model.load_hf_weights(
@@ -348,6 +349,7 @@ def __init__(self, config: CosmosConfig, parallel_dims: ParallelDims):
                 parallel_dims,
                 self.device,
                 revision=config.policy.model_revision,
+                dcp_snapshot_path=config.policy.dcp_snapshot_path,
             )
         self.model.train()
 
diff --git a/cosmos_rl/tools/model/deepseek_v3/__init__.py b/cosmos_rl/tools/model/deepseek_v3/__init__.py
index f001347db..64021774e 100644
--- a/cosmos_rl/tools/model/deepseek_v3/__init__.py
+++ b/cosmos_rl/tools/model/deepseek_v3/__init__.py
@@ -990,6 +990,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.
diff --git a/cosmos_rl/utils/parallelism.py b/cosmos_rl/utils/parallelism.py
index 393510e2f..4557a19a6 100644
--- a/cosmos_rl/utils/parallelism.py
+++ b/cosmos_rl/utils/parallelism.py
@@ -20,6 +20,7 @@
 import contextlib
 from typing import Generator, Optional, List
 import torch
+import math
 import numpy
 import os
 
@@ -115,7 +116,7 @@ def __post_init__(self):
         self.build_mesh_info()
 
     def _validate(self):
-        dp_replicate, dp_shard, cp, tp, pp, ep, dp_shard_with_ep = (
+        dp_replicate, dp_shard, cp, tp, pp, ep, dp_shard_with_ep, world_size = (
             self.dp_replicate,
             self.dp_shard,
             self.cp,
@@ -123,11 +124,17 @@ def _validate(self):
             self.pp,
             self.ep,
             self.dp_shard_with_ep,
+            self.world_size,
         )
         for d in (dp_replicate, cp, tp, pp, ep):
             assert d >= 1, "Parallelism degree should be >= 1, except for dp_shard"
         assert dp_shard == -1 or dp_shard >= 1, " dp_shard must be -1 or >=1."
         if dp_shard < 0:
+            logger.info(
+                "dp_shard is set to -1, will be automatically determined based on "
+                f"world_size {world_size} // {pp * dp_replicate * cp * tp}."
+            )
+
             self.dp_shard = dp_shard = self.world_size // (dp_replicate * cp * tp * pp)
         assert (
             dp_shard >= 1
@@ -138,6 +145,11 @@ def _validate(self):
             f"cp({cp}) * tp({tp}) * pp({pp}) != WORLD_SIZE({self.world_size})"
         )
 
+        if pp > 1 and dp_replicate > 1:
+            raise ValueError(
+                "dp_replicate must be 1 when pp > 1, since we only support FSDP with pipeline parallelism."
+            )
+
         # Checks for MoE weights. Note that dp_shard is only used for the non-MoE weights.
         if ep > 1:
             assert (
@@ -172,10 +184,28 @@ def _validate(self):
                     f"ep({ep}) * tp({tp}) != WORLD_SIZE({self.world_size})"
                 )
 
+    def _build_mesh(
+        self, device_type: str, dims: list[int], names: list[str]
+    ) -> DeviceMesh:
+        valid_dims = []
+        valid_names = []
+        for dim, name in zip(dims, names):
+            if dim > 1:
+                valid_dims.append(dim)
+                valid_names.append(name)
+
+        assert (
+            math.prod(valid_dims) == self.world_size
+        ), f"Invalid parallel dims: prod({valid_dims}) != WORLD_SIZE({self.world_size})"
+
+        logger.info(
+            f"Building {len(valid_dims)}-D device mesh with {valid_names}, {valid_dims}"
+        )
+        return init_device_mesh(device_type, valid_dims, mesh_dim_names=valid_names)
+
     def build_mesh(self, device_type: str) -> DeviceMesh:
-        dims = []
-        names = []
-        for d, name in zip(
+        mesh = self._build_mesh(
+            device_type,
             [self.pp, self.dp_replicate, self.dp_shard, self.cp, self.tp],
             [
                 "pp",
@@ -183,22 +213,8 @@ def build_mesh(self, device_type: str) -> DeviceMesh:
                 "dp_shard",
                 "cp",
                 "tp",
-            ],  # reverse order to apply N-dim prallel.
-        ):
-            if d > 1:
-                dims.append(d)
-                names.append(name)
-        return self._build_mesh(device_type, dims, names)
-
-    def _build_mesh(
-        self,
-        device_type: str,
-        dims: list[int],
-        names: list[str],
-    ) -> DeviceMesh:
-        logger.info(f"Building {len(dims)}-D device mesh with {names}, {dims}")
-        names = tuple(names)
-        mesh = init_device_mesh(device_type, dims, mesh_dim_names=names)
+            ],  # reverse order to apply N-dim parallel.
+        )
 
         # Create all the submesh here to ensure all required process groups are
         # initialized:
@@ -241,6 +257,19 @@ def _build_mesh(
         self.mesh = mesh
         return mesh
 
+    def build_meshes_with_ep(self, device_type: str) -> dict[str, DeviceMesh]:
+        meshes = {}
+        meshes["default"] = self.build_mesh(device_type)
+
+        if self.ep > 1:
+            meshes["moe"] = self._build_mesh(
+                device_type,
+                [self.pp, self.dp_shard_with_ep, self.ep],
+                ["pp", "dp_shard_with_ep", "ep"],
+            )
+
+        return meshes
+
     def get_rank_in_dim(self, mesh_dim_name: str, global_rank: int) -> int:
         if hasattr(self, "full_rank_info"):
             if mesh_dim_name in self.full_rank_info[global_rank]:
diff --git a/tests/configs/sft_integration_deepseek_simple.toml b/tests/configs/sft_integration_deepseek_simple.toml
new file mode 100644
index 000000000..a41679587
--- /dev/null
+++ b/tests/configs/sft_integration_deepseek_simple.toml
@@ -0,0 +1,53 @@
+redis = "12800"
+
+[train]
+resume = false
+epoch = 100
+max_num_steps = 100
+output_dir = "./outputs/test_deepseek_sft"
+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 = "heslami/DeepSeek-V3-4Layer"
+dcp_snapshot_path = "./outputs/test_deepseek_sft/dcp"
+model_max_length = 1024 
+model_gradient_checkpointing = true
+
+[logging]
+logger = ['console']
+project_name = "cosmos_rl"
+experiment_name = "None"
+
+[train.train_policy]
+type = "sft"
+dataset.name = "data_fixtures/sharegpt52k_small"
+dataset.subset = ""
+dataset.split = "train"
+conversation_column_name = "conversation"
+mini_batch = 1
+
+[policy.parallelism]
+n_init_replicas = 1
+tp_size = 1
+cp_size = 1
+ep_size = 1
+dp_shard_size = 8
+pp_size = 1
+dp_replicate_size = 1
+cp_rotate_method = "allgather"
+
diff --git a/tests/configs/test_simple_grpo.toml b/tests/configs/test_simple_grpo.toml
index 0171fc715..4c1e3a053 100644
--- a/tests/configs/test_simple_grpo.toml
+++ b/tests/configs/test_simple_grpo.toml
@@ -37,7 +37,7 @@ model_gradient_checkpointing = true
 
 [train.train_policy]
 type = "grpo"
-dataset.name = "test_dataset"
+dataset.name = "data_fixtures/test_dataset"
 prompt_column_name = "prompt"
 response_column_name = "result"
 dataset.split = "train"
diff --git a/tests/configs/test_simple_sft.toml b/tests/configs/test_simple_sft.toml
index bc1728752..d1a008e9a 100644
--- a/tests/configs/test_simple_sft.toml
+++ b/tests/configs/test_simple_sft.toml
@@ -28,7 +28,7 @@ model_gradient_checkpointing = true
 
 [train.train_policy]
 type = "sft"
-dataset.name = "test_dataset"
+dataset.name = "data_fixtures/test_dataset"
 dataset.subset = ""
 dataset.split = "train"
 conversation_column_name = "conversation"
diff --git a/tests/data_fixtures/sharegpt52k_small/dataset_dict.json b/tests/data_fixtures/sharegpt52k_small/dataset_dict.json
new file mode 100644
index 000000000..11cdf7af3
--- /dev/null
+++ b/tests/data_fixtures/sharegpt52k_small/dataset_dict.json
@@ -0,0 +1 @@
+{"splits": ["train"]}
diff --git a/tests/data_fixtures/sharegpt52k_small/train/data-00000-of-00001.arrow b/tests/data_fixtures/sharegpt52k_small/train/data-00000-of-00001.arrow
new file mode 100644
index 000000000..08763279a
Binary files /dev/null and b/tests/data_fixtures/sharegpt52k_small/train/data-00000-of-00001.arrow differ
diff --git a/tests/data_fixtures/sharegpt52k_small/train/dataset_info.json b/tests/data_fixtures/sharegpt52k_small/train/dataset_info.json
new file mode 100644
index 000000000..19a876242
--- /dev/null
+++ b/tests/data_fixtures/sharegpt52k_small/train/dataset_info.json
@@ -0,0 +1,40 @@
+{
+    "builder_name": "parquet",
+    "citation": "",
+    "config_name": "default",
+    "dataset_name": "sharegpt52k",
+    "dataset_size": 732673053,
+    "description": "",
+    "download_checksums": {
+        "hf://datasets/LNTANOooo/sharegpt52k@9676af8961e4554627235adf95138fd661894d2b/data/train-00000-of-00002.parquet": {
+            "num_bytes": 152166183,
+            "checksum": null
+        },
+        "hf://datasets/LNTANOooo/sharegpt52k@9676af8961e4554627235adf95138fd661894d2b/data/train-00001-of-00002.parquet": {
+            "num_bytes": 151721573,
+            "checksum": null
+        }
+    },
+    "download_size": 303887756,
+    "features": {
+        "conversation": [
+            {
+                "role": {"dtype": "string", "_type": "Value"},
+                "content": {"dtype": "string", "_type": "Value"}
+            }
+        ]
+    },
+    "homepage": "",
+    "license": "",
+    "size_in_bytes": 1036560809,
+    "splits": {
+        "train": {
+            "name": "train",
+            "num_bytes": 732673053,
+            "num_examples": 58390,
+            "shard_lengths": [40195, 18195],
+            "dataset_name": "sharegpt52k"
+        }
+    },
+    "version": {"version_str": "0.0.0", "major": 0, "minor": 0, "patch": 0}
+}
diff --git a/tests/data_fixtures/sharegpt52k_small/train/state.json b/tests/data_fixtures/sharegpt52k_small/train/state.json
new file mode 100644
index 000000000..c6c76542f
--- /dev/null
+++ b/tests/data_fixtures/sharegpt52k_small/train/state.json
@@ -0,0 +1,9 @@
+{
+    "_data_files": [{"filename": "data-00000-of-00001.arrow"}],
+    "_fingerprint": "c4293767ab883e63",
+    "_format_columns": null,
+    "_format_kwargs": {},
+    "_format_type": null,
+    "_output_all_columns": false,
+    "_split": "train"
+}
diff --git a/tests/data_fixtures/test_dataset/dataset_dict.json b/tests/data_fixtures/test_dataset/dataset_dict.json
new file mode 100644
index 000000000..11cdf7af3
--- /dev/null
+++ b/tests/data_fixtures/test_dataset/dataset_dict.json
@@ -0,0 +1 @@
+{"splits": ["train"]}
diff --git a/tests/test_dataset/train/data-00000-of-00001.arrow b/tests/data_fixtures/test_dataset/train/data-00000-of-00001.arrow
similarity index 100%
rename from tests/test_dataset/train/data-00000-of-00001.arrow
rename to tests/data_fixtures/test_dataset/train/data-00000-of-00001.arrow
diff --git a/tests/data_fixtures/test_dataset/train/dataset_info.json b/tests/data_fixtures/test_dataset/train/dataset_info.json
new file mode 100644
index 000000000..93e64c55e
--- /dev/null
+++ b/tests/data_fixtures/test_dataset/train/dataset_info.json
@@ -0,0 +1,27 @@
+{
+    "builder_name": "parquet",
+    "citation": "",
+    "config_name": "default",
+    "dataset_name": "math_examples",
+    "dataset_size": 95490215,
+    "description": "",
+    "download_size": 42540081,
+    "features": {
+        "prompt": {"dtype": "string", "_type": "Value"},
+        "thinking": {"dtype": "string", "_type": "Value"},
+        "result": {"dtype": "string", "_type": "Value"},
+        "ans_start": {"dtype": "int64", "_type": "Value"}
+    },
+    "homepage": "",
+    "license": "",
+    "size_in_bytes": 138030296,
+    "splits": {
+        "train": {
+            "name": "train",
+            "num_bytes": 95490215,
+            "num_examples": 5000,
+            "dataset_name": "math_examples"
+        }
+    },
+    "version": {"version_str": "0.0.0", "major": 0, "minor": 0, "patch": 0}
+}
diff --git a/tests/data_fixtures/test_dataset/train/state.json b/tests/data_fixtures/test_dataset/train/state.json
new file mode 100644
index 000000000..064146817
--- /dev/null
+++ b/tests/data_fixtures/test_dataset/train/state.json
@@ -0,0 +1,9 @@
+{
+    "_data_files": [{"filename": "data-00000-of-00001.arrow"}],
+    "_fingerprint": "62ee6de34509743a",
+    "_format_columns": null,
+    "_format_kwargs": {},
+    "_format_type": null,
+    "_output_all_columns": false,
+    "_split": "train"
+}
diff --git a/tests/launch_test_worker.py b/tests/launch_test_worker.py
index 1df16386f..f7485dff9 100644
--- a/tests/launch_test_worker.py
+++ b/tests/launch_test_worker.py
@@ -644,6 +644,135 @@ def run_policy_broadcast_to_policy(shm_names, shm_size, rank, total_rep, self_re
     )
 
 
+def run_overfitting_policy():
+    from cosmos_rl.policy.train import main as policy_main
+    from cosmos_rl.utils.logging import logger
+    from cosmos_rl.utils.ulysses import slice_inputs_for_ulysses
+
+    N_STEPS = 30
+    training_loss = []
+
+    def train(self):
+        def _log_in_master(msg):
+            if (
+                self.config.logging.logger
+                and util.is_master_rank(self.parallel_dims, self.global_rank)
+                and "console" in self.config.logging.logger
+            ):
+                logger.info(msg)
+
+        global_batch = next(iter(self.train_data_loader))
+        raw_batch = global_batch[0 : self.config.train.train_policy.mini_batch]
+
+        max_len = min(
+            self.config.policy.model_max_length,
+            self.data_packer.sft_compute_max_len(raw_batch),
+        )
+
+        if self.seq_len_multiple > 1:
+            max_len = (
+                (max_len + self.seq_len_multiple - 1)
+                // self.seq_len_multiple
+                * self.seq_len_multiple
+            )
+        batch = self.data_packer.sft_collate_fn(
+            raw_batch,
+            computed_max_len=max_len,
+            pad_token_id=self.tokenizer.pad_token_id,
+            ignore_label_id=-100,
+        )
+
+        for k, v in batch.items():
+            batch[k] = v.to(self.device) if isinstance(v, torch.Tensor) else v
+
+        labels = batch.pop("label_ids")
+
+        position_ids, input_ids, _ = self.model.get_position_ids(**batch)
+
+        batch["position_ids"] = position_ids
+        padding_mask = batch.get("padding_mask", None)
+
+        if self.parallel_dims.cp_enabled:
+            [input_ids, position_ids, padding_mask] = slice_inputs_for_ulysses(
+                [input_ids, position_ids, padding_mask],
+                self.parallel_dims.mesh["cp"],
+            )
+
+            batch["input_ids"] = input_ids
+            batch["position_ids"] = position_ids
+            if padding_mask is not None:
+                batch["padding_mask"] = padding_mask
+
+        assert not self.parallel_dims.pp_enabled
+
+        for step in range(N_STEPS):
+            _log_in_master(f"Training step {step + 1}/{N_STEPS}")
+
+            self.optimizers.zero_grad()
+            self.model.train()
+            logits = self.model(**batch)
+            loss = self.loss_fn(
+                logits,
+                labels,
+            )
+            loss.backward()
+            acc_loss = loss.detach()
+
+            """
+            Compute the global grad norm on all parameters and then apply
+            gradient clipping using the global grad norm.
+            """
+            grad_norm = None
+            if self.config.train.optm_grad_norm_clip > 0:
+                # Must pass empty list even if model_part is None,
+                # GradNorm across pp stages will fail if some rank does not join the barrier
+                all_params = [
+                    p
+                    for m in [model for model in self.model_parts if model is not None]
+                    for p in m.parameters()
+                ]
+                grad_norm = dist_util.gradient_norm_clipping(
+                    all_params,
+                    self.config.train.optm_grad_norm_clip,
+                    foreach=True,
+                    pp_mesh=self.parallel_dims.mesh["pp"]
+                    if self.parallel_dims.pp_enabled
+                    else None,
+                )
+
+            self.optimizers.step()
+            self.lr_schedulers.step()
+
+            if (
+                self.parallel_dims.dp_replicate_enabled
+                or self.parallel_dims.dp_shard_enabled
+                or self.parallel_dims.cp_enabled
+            ):
+                global_avg_loss, _ = (
+                    dist_util.dist_mean(acc_loss, self.parallel_dims.mesh["dp_cp"]),
+                    dist_util.dist_max(acc_loss, self.parallel_dims.mesh["dp_cp"]),
+                )
+            else:
+                global_avg_loss = acc_loss.item()
+
+            _log_in_master(
+                f"Step: {step}/{N_STEPS}, Loss: {global_avg_loss:.5f}, Grad norm: {grad_norm:.5f}, Learning rate: {self.lr_schedulers.get_last_lr()[0]:.5e}"
+            )
+            training_loss.append(global_avg_loss)
+
+        self.unregister_from_controller()
+
+    SFTTrainer.train = train
+
+    policy_main()
+
+    # check the loss has been decreasing over time
+    x = np.arange(len(training_loss))
+    m, _ = np.polyfit(x, training_loss, 1)
+    print(f"slope is {m}")
+    assert m < 0
+
+
 def run_dummy_policy():
     """Run as a dummy policy process for testing"""
     from cosmos_rl.policy.train import main as policy_main
@@ -1053,6 +1182,9 @@ async def main():
         # Dummy rollout process for testing
         run_dummy_rollout()
         exit(0)
+    elif mode == "test_overfit":
+        run_overfitting_policy()
+        exit(0)
 
     # Initialize distributed environment
     init_distributed()
diff --git a/tests/test_data_packer.py b/tests/test_data_packer.py
new file mode 100644
index 000000000..96371cb60
--- /dev/null
+++ b/tests/test_data_packer.py
@@ -0,0 +1,50 @@
+# 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.
+
+
+import unittest
+from unittest.mock import Mock
+
+from cosmos_rl.dispatcher.data.packer.deepseek_data_packer import DeepSeek_DataPacker
+from cosmos_rl.policy.config import Config, PolicyConfig, TrainingConfig
+
+
+class TestDataPacker(unittest.TestCase):
+    def test_deep_seek_data_packer(self):
+        MAX_LEN = 5
+        config = Config(
+            policy=PolicyConfig(model_max_length=MAX_LEN), train=TrainingConfig()
+        )
+        tokenizer = Mock()
+        data_packer = DeepSeek_DataPacker()
+        data_packer.setup(config, tokenizer)
+
+        TEST_SAMPLES = [
+            {
+                "token_ids": [10, 9, 8],
+                "label_ids": [7, 6, 5],
+            },
+            {
+                "token_ids": [10, 9, 8, 7, 6, 5, 4],
+                "label_ids": [3, 2, 1, 0, 11, 12],
+            },
+        ]
+        output = data_packer.sft_collate_fn(TEST_SAMPLES, 2, -100, -100)
+        assert output["input_ids"].shape == (len(TEST_SAMPLES), MAX_LEN)
+        assert output["label_ids"].shape == (len(TEST_SAMPLES), MAX_LEN)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/test_dataset/dataset_dict.json b/tests/test_dataset/dataset_dict.json
deleted file mode 100644
index 40d5c04b4..000000000
--- a/tests/test_dataset/dataset_dict.json
+++ /dev/null
@@ -1 +0,0 @@
-{"splits": ["train"]}
\ No newline at end of file
diff --git a/tests/test_dataset/train/dataset_info.json b/tests/test_dataset/train/dataset_info.json
deleted file mode 100644
index 04456c336..000000000
--- a/tests/test_dataset/train/dataset_info.json
+++ /dev/null
@@ -1,44 +0,0 @@
-{
-  "builder_name": "parquet",
-  "citation": "",
-  "config_name": "default",
-  "dataset_name": "math_examples",
-  "dataset_size": 95490215,
-  "description": "",
-  "download_size": 42540081,
-  "features": {
-    "prompt": {
-      "dtype": "string",
-      "_type": "Value"
-    },
-    "thinking": {
-      "dtype": "string",
-      "_type": "Value"
-    },
-    "result": {
-      "dtype": "string",
-      "_type": "Value"
-    },
-    "ans_start": {
-      "dtype": "int64",
-      "_type": "Value"
-    }
-  },
-  "homepage": "",
-  "license": "",
-  "size_in_bytes": 138030296,
-  "splits": {
-    "train": {
-      "name": "train",
-      "num_bytes": 95490215,
-      "num_examples": 5000,
-      "dataset_name": "math_examples"
-    }
-  },
-  "version": {
-    "version_str": "0.0.0",
-    "major": 0,
-    "minor": 0,
-    "patch": 0
-  }
-}
\ No newline at end of file
diff --git a/tests/test_dataset/train/state.json b/tests/test_dataset/train/state.json
deleted file mode 100644
index cc698e1f5..000000000
--- a/tests/test_dataset/train/state.json
+++ /dev/null
@@ -1,13 +0,0 @@
-{
-  "_data_files": [
-    {
-      "filename": "data-00000-of-00001.arrow"
-    }
-  ],
-  "_fingerprint": "62ee6de34509743a",
-  "_format_columns": null,
-  "_format_kwargs": {},
-  "_format_type": null,
-  "_output_all_columns": false,
-  "_split": "train"
-}
\ No newline at end of file
diff --git a/tests/test_policy_overfit.py b/tests/test_policy_overfit.py
new file mode 100644
index 000000000..351094350
--- /dev/null
+++ b/tests/test_policy_overfit.py
@@ -0,0 +1,97 @@
+# 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.
+
+import os
+
+os.environ["TORCH_CPP_LOG_LEVEL"] = "ERROR"
+import subprocess
+import sys
+import tempfile
+import unittest
+
+import toml
+
+from cosmos_rl.utils import util
+
+
+class TestPolicyOverfit(unittest.TestCase):
+    def test_policy_overfit(self):
+        """Tests if policy trains on fixed samples, overfits, and decreases training loss."""
+        cur_dir = os.path.dirname(os.path.abspath(__file__))
+        world_size = 8
+        port = util.find_available_port(8123)
+        config_path = os.path.join(
+            cur_dir,
+            "configs",
+            "sft_integration_deepseek_simple.toml",
+        )
+        with open(config_path, "r") as f:
+            config = toml.load(f)
+        config["train"]["train_policy"]["dataset"]["name"] = os.path.join(
+            cur_dir, "data_fixtures", "sharegpt52k_small"
+        )
+        with tempfile.NamedTemporaryFile(
+            mode="w+", suffix=".toml", delete=False
+        ) as tmpfile:
+            toml.dump(config, tmpfile)
+            tmpfile_toml = tmpfile.name
+        controller_cmd = f"{sys.executable} -m cosmos_rl.dispatcher.run_web_panel --config {tmpfile_toml}"
+        controller_cmd += f" --port {port}"
+        env_dict = os.environ.copy()
+        env_dict["COSMOS_ROLE"] = "Controller"
+        controller_process = subprocess.Popen(
+            controller_cmd,
+            shell=True,
+            stdout=sys.stderr,
+            stderr=sys.stderr,
+            env=env_dict,
+        )
+        os.environ["COSMOS_CONTROLLER_HOST"] = f"localhost:{port}"
+        # Create the Python command for torchrun
+        policy_cmd = [
+            "torchrun",
+            f"--nproc_per_node={world_size}",  # Use 2 GPUs
+            "--role=rank",
+            "--tee=3",
+            "--rdzv_backend=c10d",
+            "--rdzv_endpoint=localhost:0",
+            os.path.join(cur_dir, "launch_test_worker.py"),
+            "-1",
+            "-1",
+            "test_overfit",
+        ]
+        policy_env = dict(os.environ)
+        policy_env["CUDA_VISIBLE_DEVICES"] = "0,1,2,3,4,5,6,7"
+        # Start the process
+        policy_process = subprocess.Popen(
+            policy_cmd,
+            stdout=sys.stderr,
+            stderr=sys.stderr,
+            env=policy_env,
+        )
+
+        processes = [controller_process, policy_process]
+
+        # Wait for process to complete
+        for process in processes:
+            stdout, stderr = process.communicate()
+            # Check if process completed successfully
+            assert (
+                process.returncode == 0
+            ), f"Process failed with code: {process.returncode}"
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/test_process_flow.py b/tests/test_process_flow.py
index 75409871d..f036509b4 100644
--- a/tests/test_process_flow.py
+++ b/tests/test_process_flow.py
@@ -39,7 +39,7 @@ def test_process_exit_grpo(self):
             config = toml.load(f)
         config["train"]["epoch"] = 1
         config["train"]["train_policy"]["dataset"]["name"] = os.path.join(
-            cur_dir, "test_dataset"
+            cur_dir, "data_fixtures", "test_dataset"
         )
         with tempfile.NamedTemporaryFile(
             mode="w+", suffix=".toml", delete=False
@@ -125,7 +125,7 @@ def test_process_exit_sft(self):
             config = toml.load(f)
         config["train"]["epoch"] = 1
         config["train"]["train_policy"]["dataset"]["name"] = os.path.join(
-            cur_dir, "test_dataset"
+            cur_dir, "data_fixtures", "test_dataset"
         )
         with tempfile.NamedTemporaryFile(
             mode="w+", suffix=".toml", delete=False