From 297610688be009514920b5bc3ed7a2b78050e620 Mon Sep 17 00:00:00 2001 From: Hassan Eslami Date: Sun, 3 Aug 2025 22:02:46 -0700 Subject: [PATCH 1/2] Support Deepseek SFT --- .github/workflows/build-and-test.yaml | 2 + ...pseek-v3-moe-670b-fsdp64-cp4-ep64-sft.toml | 56 ++ .../packer/decoder_only_llm_data_packer.py | 3 +- .../data/packer/deepseek_data_packer.py | 109 ++++ cosmos_rl/policy/model/__init__.py | 2 + .../policy/model/deepseek_v3/__init__.py | 473 ++++++++++++++ .../model/deepseek_v3/checkpoint_planner.py | 108 ++++ .../model/deepseek_v3/deepseekv3_mapped.py | 579 ++++++++++++++++++ .../policy/model/deepseek_v3/parallelize.py | 283 +++++++++ .../policy/model/deepseek_v3/weight_mapper.py | 462 ++++++++++++++ cosmos_rl/utils/parallelism.py | 69 ++- .../sft_integration_deepseek_simple.toml | 52 ++ tests/configs/test_simple_grpo.toml | 2 +- tests/configs/test_simple_sft.toml | 2 +- .../sharegpt52k_small/dataset_dict.json | 1 + .../train/data-00000-of-00001.arrow | Bin 0 -> 97352 bytes .../sharegpt52k_small/train/dataset_info.json | 40 ++ .../sharegpt52k_small/train/state.json | 9 + .../test_dataset/dataset_dict.json | 1 + .../train/data-00000-of-00001.arrow | Bin .../test_dataset/train/dataset_info.json | 27 + .../test_dataset/train/state.json | 9 + tests/launch_test_worker.py | 132 ++++ tests/test_data_packer.py | 50 ++ tests/test_dataset/dataset_dict.json | 1 - tests/test_dataset/train/dataset_info.json | 44 -- tests/test_dataset/train/state.json | 13 - tests/test_policy_overfit.py | 97 +++ tests/test_process_flow.py | 4 +- 29 files changed, 2546 insertions(+), 84 deletions(-) create mode 100644 configs/deepseek-v3/deepseek-v3-moe-670b-fsdp64-cp4-ep64-sft.toml create mode 100644 cosmos_rl/dispatcher/data/packer/deepseek_data_packer.py create mode 100644 cosmos_rl/policy/model/deepseek_v3/__init__.py create mode 100644 cosmos_rl/policy/model/deepseek_v3/checkpoint_planner.py create mode 100644 cosmos_rl/policy/model/deepseek_v3/deepseekv3_mapped.py create mode 100644 cosmos_rl/policy/model/deepseek_v3/parallelize.py create mode 100644 cosmos_rl/policy/model/deepseek_v3/weight_mapper.py create mode 100644 tests/configs/sft_integration_deepseek_simple.toml create mode 100644 tests/data_fixtures/sharegpt52k_small/dataset_dict.json create mode 100644 tests/data_fixtures/sharegpt52k_small/train/data-00000-of-00001.arrow create mode 100644 tests/data_fixtures/sharegpt52k_small/train/dataset_info.json create mode 100644 tests/data_fixtures/sharegpt52k_small/train/state.json create mode 100644 tests/data_fixtures/test_dataset/dataset_dict.json rename tests/{ => data_fixtures}/test_dataset/train/data-00000-of-00001.arrow (100%) create mode 100644 tests/data_fixtures/test_dataset/train/dataset_info.json create mode 100644 tests/data_fixtures/test_dataset/train/state.json create mode 100644 tests/test_data_packer.py delete mode 100644 tests/test_dataset/dataset_dict.json delete mode 100644 tests/test_dataset/train/dataset_info.json delete mode 100644 tests/test_dataset/train/state.json create mode 100644 tests/test_policy_overfit.py 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/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/deepseek_v3/__init__.py b/cosmos_rl/policy/model/deepseek_v3/__init__.py new file mode 100644 index 000000000..6c032eeab --- /dev/null +++ b/cosmos_rl/policy/model/deepseek_v3/__init__.py @@ -0,0 +1,473 @@ +# 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 + +DCP_CHECKPOINT_PATH_PREFIX = "/root/.cache" +DCP_CHECKPOINT_PATH_SUFFIX = "dcp" + + +@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_checkpoint_path = os.path.join( + DCP_CHECKPOINT_PATH_PREFIX, + model_name_or_path.split("/")[-1].lower(), + DCP_CHECKPOINT_PATH_SUFFIX, + ) + # if it is a huggingface model and no checkpoint exists, we need to load the weights from the safetensors files + if len(model_name_or_path.split("/")) == 2 and ( + not os.path.exists(dcp_checkpoint_path) + or len(os.listdir(dcp_checkpoint_path)) == 0 + ): + # 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) + + logger.info(f"Dumping the tensors to DCP folder {dcp_checkpoint_path}") + os.makedirs(dcp_checkpoint_path, exist_ok=True) + fs_storage_writer = torch.distributed.checkpoint.FileSystemWriter( + dcp_checkpoint_path + ) + torch.distributed.checkpoint.save( + state_dict=self_state_dict, + storage_writer=fs_storage_writer, + ) + else: + logger.info("Loading from distributed checkpoints...") + model_name_or_path = model_name_or_path.rstrip("/") + if model_name_or_path.endswith("_hf"): + model_name_or_path_dcp = model_name_or_path[:-3] + logger.info( + f"Found model path with _hf prefix ({model_name_or_path}. Looking for non-hf checkpoint at: {model_name_or_path_dcp}" + ) + model_name_or_path = model_name_or_path_dcp + elif len(model_name_or_path.split("/")) == 2: + model_name_or_path = dcp_checkpoint_path + + # 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_checkpoint_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.") + + 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/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..9352b599d --- /dev/null +++ b/tests/configs/sft_integration_deepseek_simple.toml @@ -0,0 +1,52 @@ +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" +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 0000000000000000000000000000000000000000..08763279a19055589bebf92c8255bf3cb1afc74f GIT binary patch literal 97352 zcmeIbdyr(;ecv~TWLp`_jvOa0CvlvsNh~hJ-6j@*Kx!qCUH}VVDS}vHmk)(hfS&F< zGtKSCupc`+pilv+B@F?PMOXjzw($6gE zzvZRBy!7Xn{@l_}@fZKy^polPpXTFFhVTE>{^n=?)%QQe|LV7&;_vg}x8LU9FRpa@ zYlD>|#TQrB%I0`7EJyaCHRzuyhok1WGwAchXTP{I9CSqebYqm;*|0nrw3=UY$PyZYI z{U31YJFV?+aA|!xynI~R&Gk3A-pBQRuG_gj#HHtsaebWYL9WkoJ;b%a)#aLS zZF7B@>#JP9!}YsdlAjA)uX6nX*AKY1?hO|BKL_j7%K>qA_>%yo?GIM;(*pW=FmtHssj+T;>@Jj(S|t|z#D zm+LvM3tTU8y~6cfuJ3dG1Fk>e`mebDTdsf1^-s9|8Q1^H^?!2xKU_cJ`U&X$(_E8L zIc$zbozb}2AHVnEU5SHtGg)BurGLQnkI)NG{dx2#0{I_t{quhv{`{*;OaIHiwY2o_ zBAibYOH2P5*FXJR$Tq_LUm~c#bzo`f|Ng+z(#gY1OU>KR|6Fe(wEvLnpM7{~>Hp%I zAiU3V{Wl*&E*?O~e;V20`n89bmdG%F*X{m?s{>A@cY3c9(lclA<_J1%w*FX8I2*Y2$X=&@&?|$Z{rGLitH$HdM(og;RO-oO7Zd&>;xPAs*dUDe1 zHHV!?%3{1x7QJStUv%1Kb5so0ijS6SYsIIV-P7gpeWT*UaL}H##>K~)lYVQXxMwiz zO}fpZ(;p9t)k(Koj*HC!pGE~wclb~a#|$?r_`O-I4+mT04JH_#F2@+)qPbr7TiZpa zx7jUw7;}9G2&1yV^f$-ldZ%Qhe!CbqTjNe^RICk#x}h^HTHOv0j2ex`edVw$nq1ci z{CJ}|E?UigvC1%;n~d8oR=4*6cLs|>rF}MKqRgYXr6eIgHUPw#shzP;1y+8$*!*%c%Uq#1P65z2@0YPh2}{b_1UTFAbk^ z)F@7EbVf$lPQTTiv`g4|&7oc$^wDd;m)13EHHWN!I|sR0AO!7hrw?l-fG5xuo^34O z4ps-FaYE}ogKl@QRZKP=^^ufJ;@B$IApB~xbvjTG_qEfX*w9e<3<6i2k#noLPM77T z5w9rw>j;U+1que+&2D4)7aGM~k^GTAMiXdvqNhM%(9?Tj{TxKY_%lY7V7&{cj9tMsWl z{Wd<#a9d8pXtQjd)-o^_@GIULTr;tO?Au~>Fu)p0<8N}S`4ecS$XNQ})Tauuva9p^ z2YuW|?0tPvIVLb6y2ZVk;6utG1A;Cf46&tiagrDTKEu9qBi_|eKu=0j+yN&j>t^pr zaWo9FD&IyXNuDjMi#9}eRhq%L*GvL{7kC=ou!Quu=oX|BZPLQ;ksB5~KJmV-IJ&&c zznI(8p#SB#)y-)%4EAhX;K4X*klV@^57?5&ZFZ-LIogl`3pO1ix(ROO7MQq1cqIj^ zc!AOnh#qS*QE|6du5umr!?!dD8?dVd;{YO#J^BE4$0U|r{I9k6V=pj|24-4pyWcKe zGZW5N7NX=67-;Ts+SirXvM5HBUl~;|zET(p``m3eK5br23?#rO(n!~!riDkw z^_$QTEE4^>=fS#5nNz8Gpa&D1<$#ohTLv7lz9R7Xl9yr+jl3G~Et>F}M>2NQP(>2D zG;@g5>c)J^hhY{{Li2Ji?V(wTDnsBosK%%9Y?zR!Sg9{D0}&oBOVFaWs;D$$ZV3<| znpkmaJ$Cm-6AQIkmTsIhuhx$oV`f%MWFIs&`!6CRWS_&VFm(zgjwD)B5d};g2L(8K{sx0Y>ztllEwO9(DtNAa2SJYXJdm@ z44CdGySc$)(R7uJ!4{EG(GpCm)P|20^siVm`N7MUQ6rZWIcTo9jGLS2YGxEml8l@T z0divk^F#4b2jwAplN46f2P|QT$XvL~E;>>s9*U!T9sEhASCZwxJrCUfz|D=~(;FOsQxJyr3Mv)OtrD%MBv*ZAAwqG$;|a}n`6o=`@OY5FEq(vk@_W^I&DJU zMm`~$px^Rw0X{}qQU>di0tj!w%ww}gQ460=%VH6T5@o2!6ztx7*h5vOkJVPq3e9#z zX_*KK*0QOrr#Z;zCRA3s(SgL$qc})=C?uAc;B_#q6~Z8LEmoPncRbkav@B3E7tLo* z186vuGOAuUOK!V)Zaz22G#AU!quFglm(L6?O8_qP^(_N;mp95xgCH45SwK^em;iOu z%$@jP79It)0nb*^wOHGtIQwx?;6}~FLZWF{3>}dG3EN(3ti(ZBLqnnC#4b!^g^Vus zMNLYYNEOG?psTciFcw2AxSPe|o%RfUzHmKKQ?25{iJ1sJY!$|4r$^?Addr<&FPaA1 zA_%OZ|7v`^MPg8-WqO4VHeqt6?S9wN;F)ZFf8Bzq=o(Cq@0vs`lDw&?O7;`T`*db& zRsYh8r<|_nD_OPVo}-^2GkPj#bif~BO>oi)K(VIevnltPi;!v`Hj5ItRt?f5gVoP_ z@adk!mDz&A$X#8MJc?oJ^Mglu4W>nehZe?ya2P8fY^x|d=);;g75vH6!h0mU?r=4! zOk_(=MtDNM#HHIKLHA~=H#mxBZZzn86(gOOt1-`GEXZ5Htx<5pq;u7Yr@w=AR^Cex z&Vz(09YtN;QLz0^))@woKvs4#85&e3We3cI9*g>@ZGOR;6)+u(tSf2q83U)-A*pjE zf-9c`Kieo{`W(FFcFiD4Id=5}c)*5=izI7Q7D`CJBe;#8WM`&R0|KZA2dkKqjDO}$b1U#jj5;L>Eru%H3m*csO^XG$my(?dqSdu3=K)<%95PLKXh3W zmipzEtA0!Jq9T=$!zU|X`5kw{CoKwwJT_@rxKyWbvdbT-XA+-rR; zQ*$P(?imMTuQs4lwllKr#Rk%HGHT+3?%JV{$*&fTb|AV`(?lq(lR};dK^Y8SKIFR> z@3QAM1=NUHP3 z;w&RHp2qW)4 z-BfHfx2<3k;zoHDYveZYUM*E+3Ie5ohv~|L6prZ$QD;xLW81am}O~t(!(b#1)D%Z7KSuyq6z-;!pY=5LU-Y;&$q02hD8nigw zM}|hJT1&D@^~Dh=(oJY8Lcr9`DXhgw7@E$B5v>NIe~)Sd_M1=GZUO_m3A(xjby~aW zv%B?6ADiykcW3LV$w+8axTv_l zxp}Y%H&{zPMvJa{8t|>ruej}Cal88QwQlQp3&HgC9HX&4s421s#Zp@)C&$UfeB{@t zk5*h?77w2^pd09L%7X`PewfTF6=yG&vCe#b2KUCGE6J3o<23tpjNpOEVUvGii!rI3ic>-?&^$=EI<74!=DG7 zWlxnB`n!QX?l$ggjyAyHFU|g9-!5A@yoFm!R}s7?3O(rbIu>C1;&Q=%;Xb$xW-87M zI!Y#Y`Um`x&;B}5(fa7-!f)0C{sVZ}BzyKN)5MrZ@L#ft^emwLvOj+?jPv2hMN$_d z#*+sS*7nOEK5%pOY~VJ#ufg00R-CuQFo6YD_PpK5IK^H(kbnydBPr-B3lp9Yj|b`> zuvH(d!JN%`a(hIdV}lVIn3Ddu+dtr=g?k?L$?7N%)6SbMNA->E;{9|HGExKa{MB;%z|DNI^5{xMv+UuF zd_0^U`;;0Jb296Z$XoVrBRqK^o4!VYi~x@&Vh&3@Or>6+rNi^+@}rWbWfYl2D~cvz zj{<7zy^=okxkoPG1P1$1Xs)7Ca_-^4z7HeLSh;rbaJKf*!!G{4=4MF!Wb9Q$mNuQY z;tXpR)@{sSD4%rTzCqACiTA@Z>XCNUNlsf@gfLf~h58F$B0MFl&0-^DH6Bm8&>})S zBpiHot>ChlE#f(mPwgDRhQwh^HqoS{>7}&8vO%@^S=B=0b2VI>GVP!PaP_8_iL)XL zlY?BTY`Q+kD`E!C8N9-p07{JExbF=x^J>a=a;Nnws{c{g2fLy0L{G8$pdl

Ik>DPw`4O_HQX4^HB`ovA0JSHGP4k^fPtp|qKZfvb^$9ydCHDh z0E4re35a_>JT2E`7N5Zl%txQk0=}r&IBL#DB&%RE&sU@5lx*BzQm9e)7WY#d>#t8} zjyx45_I;sbeYLadg%sAzi!Io6dmX&x-ix}5>b)+0( z>Cei6U0(j^{U?f}ciwgP^8F8dsyO{LK4L9{=^~^Wj@=$v)lrnc9cX{rs)9 zyYKu=dhmnuvAcn)vJ$b%&g8I-btB}qTq_Zm{m!~NgvnP`S+#|jmP>`R$-+;8FKfD! z3uM&}P)$0)=b_RT=n-gOaR($pZqATcb&7O)a4RYz<#aaL^=Di zxyKNn&w|D^V6eiPWTidYf=v=aHzHRXuS5Fwm`97{rQ^ zBQeyYC+Rf5IADOEk@JohziSeGbzI&IKi~{YzrnWdRSs)>0Mhs=}~C3 zM2fJru8e!)%A8^9t4a$IV@)p&eyXpFfPe~hZCMe;UdY%WGuuW!vC)zV%KkQ{maE{l z)1njDu}WJR#pRli6TL*}D;5sG=M1}G1$TZXuy}15V(R)L<`Jj1icJb>g4Ft`2z=GP zYI~Bzixq9!+SO7|oy-~zLR|{->`lO*X(|5*#cG~ITye&SIoCz{d{BZ>15CaG3sao{ zu@z0YBNHX1CzONg;Xs*v3XOV<6uTLYeum|&lbaABh+JW|p-abNQC3Z=#@#&BXO*Se zek_A9r0-{AjM`)`lyWFRABJWDXA8HxQpQ{5A%NcOSTrA8JL#A5tmR##K_LV^tM#>& zq)4sz1?toUx~4qRGIEq^$bBj!b)X1&*!iG$OsAWFGpl!M#f(mm>IoW{+2p0Mdo3$_ z&9z32DT{um1%SRdIH<`!4H#3xJ%5eK8h83<2B+OfmfcCK(^jvMR=rx{E@B1+*8&XW zSjL9srkIvETJeGT_)} z8g9NBPBK(+U9`fLLMfsmkDR``Ex()8RAy$~aq>GdOpZhlj_l= zr+oO#cZdV*+YozbkhpxF-33`Qq()J!lKrmp(G?7>yGXJ;`!Tnt;I^RIP%h*wJx}Ls zd~l^$$ykw(ntLUu`$|x(DtEc^GHpA@9O}`!5)K;LoKeB+mU502#w*BMoFpN0VD2_6 z1`*^LJ4DDU60S@sN5+HzNe7A>M|;@F%q3{6@fr&$4=sj(E`aZgP|jBamClm46F2Zy zc~x?!Zl(Y%GZ#c>`&C(LVKR}Wm|@CahNeE4(>wC>2hBO4=U82jDZW+qPx~fvCk%{Q zSD!3M(4N7OSYR=%Aup~dUymDJRG1R>ps;uVp!a7iutLG>3n0vsp=1!8J!72V4yc$i z6O~!FAjHDXxD}Jn;qRQAdrg^YGbz)?;$sw`nUt~vDdV5D&?(vszQzpDOTJ(Gk#sv6 zS-Y++%C%0}33(v1B$d*eQ7&^9TelrR{pj{ZlovIYHb|B9hmB8jvEhm0hG1caVWJ>GJxMt0+F`H3BsieW`D7G%RC~75h z12ZkcCDaPPlYV#5IxT6zU4$ppm#Tg=wqndNJg5_PMJU^5+SzPoY?)n&OeOP^TgkZ4 zM9-@M%8mt_T0Ehf`UG01*~QL^2Yq~dv9MvM9MXRU`8iW|L`Z5p<8pU*4MDXKXa5H4)W!re zp-F>RV*KEeWtgv|J&EnUOxIiw0VXw`*B);{ZEGi(7XX%Vg;|6=5gK=yT?64pk7e8* zO+_#^5K4QMtSQ?mv=qXH;m+*h4vcmg=DaDJFqX46}DSU z$n%{7?pp8@1FZS6&sEy2D#=$_t)#iG&SvGV>U&0uh6*@tp6?dppLd7Riu@k6c_RLm z%~mCn*T-zFmfJW6CaCQ8>Ds1m&hN3!BkXvyOvs;6yCyey(jT%bt}FbYE;l-DdPfKq zHpd(9nLR1@l;!y|WA<_vNf4s?@Q@bnaR0Kf@ba+pavbDSL=UDpGw2e(xw6U}qLi~V994S_6>o9k3-_ak!Vs}w`p)G-a(6z4 zY-Cz$HSQ`rig)ary{Gc(0R^Gi!9#U00RN@xFqnAPd73GIFfDb;BP8rxo7?F@wYAV& z<2}q>MTbGN+MW@UgpkBkHg-;mQ#v3fPBIS^NgLVe6w@9kt>8?T6bYRod*|^J9QH)H z4rPi?_Vy_6*B-KvmND^c({R)2j=bRfvpV@girv!@)W1bI{7h1|tNu&i0P9r4LSg32 zK1^p~%-&aC)>WALdaQNzb!$`DGEk~W%L7+s1mmkz+O_3#w9}M!k5s-+X~$^k_hYD2 zQh8pd8RlcAsvWHK+`NwAomMRm*Z<|ONc@Xqgy-f8v+j?OUN_DOogiZsh1czbEk7du63lUt(;y0 z&Xlfo+l}6h4uu${N`uSA^KKK!Y^?^eOMBOB)IvcqenR+(718}!#OB2tYct%B`_3Wg zwwsW~(Yp6SSlcUBuIRcz8FhKTy=>Uklz-+{b{5xlYEq0ksM5)EDm<#MaU7kcv1-M@ z>znS8I_2CG!vv-8d> zrLFD$u`M_=-E-sm>Hd=&=Izb2myxTd?V_!Y6w`VwPDn77of=YVCHqWg*3Icu479-C zQ`_r)oW0P=gzqsgQGjackGB{pMEik)a0PSo|CD$=)>Q5n$1;`V;Vb0To6@%?G_%V4=ptH8S z$7-u@+kc9@oE%;YsyOj=Rpjz^msZ$Kl{#s>dDc^axnBPQr+8Fb_cCpO#ZtQvEi-}3 zy=x{ha31wV_iCO`#c8<8x!9`fpeF59_h4k016+&%>uut4Te_L`jL1N?f5mPA%8f<^lA zE)RKEmDia0lbPSh$!KIw(9&#=!trzmGcXttU+p>yYU#GxJ zT3SHA-NV%j&U`Y()h1{5-}U^kohkT|97G$DnwlNPYKs{wQIBidLTx8$bAt9DI*49& zlokR#cYen3a`L#F?vhQ*02nuiRSuo)<+rg|;7odqh8{#H9)XtW;|Ktu6<5^bYY zAcn@IVDEaf60d}O8dELoomOkc(Y!t&kL{$H7llSw!v2+T3#ZKNioslPAWsS_=7a%r zO)%xnY(sXXH7bJ?sr5lP%#=nYkaqM&@xi(DM{niaYN}4jTnAhli-0`$5h+e_{0@3X z#wpv))3i-H(1gPvly{p=W>?xin#$`M@aYI7x-6`DUE4t>b`qD=Z#uz2u2y0WS5X4{-y-Q-~M#Pi#$gkd8T)FmF4j@DBztYd{D_nFge0#TPX~H)FNahbnmvyg8q>lneys<) zqHc724*D4E2;|vRGTBdXtcpJPOWsJ>jRlYSX0wXLg+tpQyAGzq!9?cmavD-|#B)ka zD)H&dW$b3i_gZl?YWBy zcgWGVkva3#&WUtx!LRm>d=)y~6AkffmXH%)X8neq@9vEJ)ch(*g+V*217k_l%@via z3hXV!^X!eC)ef305g_Ep*3%Z}k?oVLUx+&$Nx_5BHr5dt3RMgO;TcALSdCV`W9c zRFOQPL}fvC=4ZSWgl{Sm3DOjGp4sNB&gaY=i8U=SS4!(B4K+({E!o}?KVyOYV6G^; zQwn~n*d-+sI)?CgO#K_qsL5KlyN(q^wh&Yxn6iRdjtENz>QJE3(_@YGvmRoW$c#W@ zjlD8OVcE5uPhwt99I=DLb(qVX)X!XhZvH;B!bMxqG+SDP={IOUJUot!J3<%Zxlz zE}DxsFXAb6AO?c*_M5qQ7A)f`}9PHrGnE-#haOXLP)kfaO;SAP1){TjIvDXXa=A+jb2VH?14+BU= zrsktLsdx{fDq+kWnoK-e_C^&QAInjOwT4d=adjLc4q5F?M4Lb3UUC^-cubR8wiyeK zh7Kl;2-@rRRp^eaS$bW}jH3WcW`Rs}78J@Dr}ylr3(rb-kl1zYR7at8T6(#)&K?Jo z_O9T#4RcRKPASY;b!ogWMVhS;KlrMYXth4q$?&2dn_FM&6pLmi(pK(sB8-FyqPOBB z>5eLK5JDR($eV_ZQkR1HRvxY=V`3XxLck3LQp`QoF`xXkY3<{ysv$@`5Ltkq_2LG1 zb@W<9f`>@giVd8+d9K=fJ)#b=Ty-k-?nljWXtBEWTN@!`2N9v>e1WQ>#dN-nhh?Z< zvPt#^-MOsbEf^3o_zR*4tXgB@Sx+evl#M!_u~l1raMXP!ka=06VvfcsXkNYW3$|d& z;`4)5*VBj^$T9Oy1dLb^nKpaEg^tCudwCs|yin92&HXf7usd0HC3zvFOQqSPwX&BKK5qpz#4>7;Lv`Z-9gm7Q~6dp#)w>FB8 z6W&MoIFP-sVC_f||FQ}FrPW%fasktWFgH!mJzf=PQ26_coOt1#Xe!U0whAuhmy>z2 z_qdCQyhobegY|*lkEEAx9TV^}&~_e{u#8->tn~0F2 zjQCfy7RJfVO7~F?56Ak#e;QC@Pd4p<)Htx-y!GIKL&ABX04j-wai>-S%J>MD9Mh># zlBhAbo##g_e2N(Wve#tUgc%cvf*g&mOpoKiN7G9hBS}7S^hj|xR*e^J#ix1oBb);h zJ5e(Pu~w8m318q+;|B+}R(SBYbDP{sIQt|5%iwLBLlY4eUQoELw{blH+t*UH0W7!| zuaa~!u}OucgPfBT&zA+WEjl;!CPNb&34}EwNi-8pqJV~BF?<3`D5h2pKO;n>AHqah zsNPD1;~im4%hV$_a<}Ok`+3hV36?R)ACLV$_>AOqM=$IN zQrYAj>uyJo`CEPAo~Bb~G@2x>*gPx12Vy`mV@ukZpe2k7*G%Y1rCG*L`?n!l95ff3 zkHa!In|3DLvR)~9lEi<-jt^EZ<$#z;Q&=IwUXx*~1tO#&MrAPrsWtY??r8ZJ8^tFk zgU}@hO~}GFyGgP%f)Pru5emZeg><{z==6HDG=JFXE&d!WAHD@$ShhB3KqAaI>8~R_ z68|uYrC)i;wLqPEOKcn$^}s6{DX_#WCNTj255YNyr>@(f+VO*N`JTdQ5K%;`I-bk_z)i=?6G@zt?%hSYr7wI`nQ4@KYV= zPoK1XEdhzIDLoMn{wq#y^3IO?X&XbWxGq1b63nnTes^&fRN>p@!)U;}@64p?i2rRu z0ZG(+$bMbRsfqVin`{x2Zyes*IIH%OOgXCtK2sLak2bV5)!M6<52FPEq7wh3lQC_N z65$FLM~d*RK#QOpN+fKKPQ!cl)(|p$xEe&Q(%t>gV9$Ac6U3t&4mnm@adpI(Y@3=? zZKY-3WJ=6v;N$JOMp{0Mmi)Bd&3sa?uH&%rfchGQc7v%!8m^QG&u^^MAP)X%gw)>m zmH-U)bonrfFsXl2y=WGD0@bIt`WXVF0XZ^zZ=6@T(4bDklLwPQvoYF()c3ETe@|(r^P>Lot!t)A)4eyk)>WPOVVrX@6y-~ zTi;+?K8&9GB-{-05Dr)vrl|9=@wDI~0pqXMg}g*;AdV`l$rh>pV51MEn?l95&*@9J zQfNE9O4U8OTxWm>)7(a4FoWZfM2|#ne4^&I*xbQJ+s&5NV_&2C+eDGoa?IPAiS0P2 zB07e)ydfI!Ry5!VN)al{`P||vm?nG&_XWqMhTHI~0&l(54{vXsFw*(Lc|5bO!M^}& zzk$PqJIp`RzVmM^|3-1>&>?=Z|K(FuxyLI{m|ue{OhPTb%+dH&n$8(Cy1eqU(9yKl zN+HPoBW3+-=&EY=juYSe1F0{~j zjx=kNCv6b97k%Sh=!I*T0TAvG;C0<5(-$3*oi7w*vz^ zsL?a8m>R$m{MN8|1O4*Fpw{7XSR>1P;wUiwcwG3!00Yz%`&PDo%5VOSsJ@aXUitYh zYggguOGV)n2S;DY_^WdB7y0;U*dY2LM5=lC3;8S(Dq&z$MBM=d^?fHVvwp?r_eGJrvOceSO~OZ409c$l!e?l3&|9ydxSOsZ7~e? zamuveT3J9OuUfW0M5Hmo2}X0^fsM8V+ud@2w8x20d_^dvwpokFM`Lq@4-i`2ycpd9 z)S5WmJrgSdb#N5TPDTqKWZJ}{=BA`@YSsk|jvwKaJA+|eB!Db5^8}{F5sJ$3p>UK* zCbTR;1ZuQJb@6iMfIdYVyv&+u@(;GG7`LTJ>Fot=t{-Yq+a7@5SsdkMgne8U;ty6X zvwy|-PTUm2z6A7*3GaL{_4sb@)i_*ogr_KgZ@*VWD6vjapcM*EOVzZ9Ba|sgd9ekt z;9{WmnAg#ryp$M0+7mEHiXbyzE8U!8*1clKfWAC23GO`V^qiKo zbcj0*ac?8&Tnz0LC@xQq(CZVVPAMOvYC_a^E;wDh@)&#*A0LhJu`gMH>J-=S72axv zuJcFN9NjJkQ!~TKc|1K0-2AAx0?Qfc#)riVP3|m!&S5cJcJ!ItCS0oeOL#rn=xo-K zSK)gllddX=UX0;76hOA~lZJ5?}sP4mW(5D=tU}3LaCU!Bv+^M z&VJoEUNkoc-N8B?|CZENzMW-;Ev}Sad)w$zLPs@>hw&(&-$$<+30ylKm-=*@PTQev zinXe$4?@C-e1pn9IQ4T;@%m>nobgK}>pGTc!C4%*V%6n5G~9jR zjol09ik&CFzw^DvcF+G%;pFA#f4KAFV~Q_#e)s&&g_nxm-#@qe)o)#X_SKyyAK!VB z;Bxo**LE&Gv-9mY6@gy5_*ilIcV4^v>?=E8BNl!9`P~c8>aU&G&oT7w3y(4SOS>;T zMIc&Sy7-OVC%zwsQPf9_dimK4!0AD4`h!r{GL~5>IR$yU=g#lEb#C{Gml*E+?m0bu z>CKC~-+BYc-{kAdPkp6G_w2s#d~x~jJ-++$0eo3p` zJ%8@;tR#j4<6rn`6;{k@-sga1}<{*Q%~(a`xXE4**A!)1?V&wtaj(M zH-!#`r@O!TILO@n<{SBq-#d5to4?JoJOAmcV6ITG{^GX<;_f$}x%B1KcUjx9$AO$z96%bc~dB%99`{)6+())3rJUbys! z%$>l${LE83FFwC>?s113^n4ctKK}g9cmCK`eChSaE`QZ%lpfi=aQ@QwzMmuSA`J6} zu=T|BH?;Q{BG40sXl#qS@CR555+#)W4i7Jj2B zrl*b|RcN&E?($cj5(lybo5?Y^d*Siji?3+j1j6@UyZrn)&G-t!aQWQDotM9qQN+<3 z_|^pRdvCxfVv08}Mt-ekhMjkw_%_G@gefLW>xU6{pZ}5&`{HZjEJq+z`TZZlv<9t2 z_vNSUmkY1VKqCTviIqrHpy1ByFToVh#vkRUCxzXMD2igopyh#Tc89RG+=>(htPeVj zIjct;V23ke`|?v?+j;pVeL@5z7Ay#X+j;q|-3xqYzZQXw&Rr}nf8*K6jM5n2eAU?M z%Wzqe`rUJHUV8I!1Z9`6&uIzgFMs{Fpp5e(gTMzG2bE!r;Va#J`LRoH{#ICN4G1yr zm-SSx3@E`Sj+Hk2`qPHS@1EN^_cZ}?`K#a0RnTj1n#fCEp`V{Xbv=c6J*|0MNmZBy z9dOX@1y?$kUVjZLG5L8EU3R-^mjYFC`Pn}(9rDaYY42*7_@HGH!$KwEjIl5dlQchm zp#T-TzlXLCu$h`iIxma*6-EDnC*`yX&xTMuZ$`aJdMx1&(c?Vfu9k~@pQjJw~u80lp+|NeQfdifjAFnxHY z&XyUoJc&Jwx4w7*VT%e+b_qFq^%*JsmobV^7iL2ni=8Q7Vt_Bb{>QyXk8n<2i(IKN zrj)cnRee=QG7#^)_C_Ju|E)K6zy8BOo!oF^8Ew31!U9p!BU0yfzhNKoME9k6(_Q(% zQy_ffO&k%F@7G|oN^lXdG^>1Z+;`xUhns+(`RZVP7w*8bsN(0JCPZ)sq`Q-(6k=R+`L5z}iwfI7ltn z&UUy_<0RW7IMQgk#AlT}?>4VbvaX4#0ba{w?^On`+-dw?Wsnx-EDQNwWst3Foc?6# z&MR*DLyyl^2Ithu=-5xY^714VeITXp-XrCI6Y(~tLVv9@SV;T6Q$@krq*X-_ zpRfAm;GBY=w;6IB@K^g{FWCyW_?ev}_vSoT@DS882S#UF$*H zyc>4)Nnuwz2=~_^3Av{j)p+#1Qk2C-_2STvxfHeY8r7!@yDxqvR+%n;<9R|Tiz7Xy zu=8b#RbPq`tuiTuj^B9m^7GHeB%4)}2z*HozjEpCKVK+~^!sm-r`f5Mrbvw_)woAp ziri5M&aK%J)>FRXMUw2O@_iPiva(ZEa`3-GB8w>~Ib{{6-DjU7rLM%4G83#%r6y%O z7Z5SbnnVl0P~xcbb)!cE0?oAdbZ> zLVT-kk=bGemVTUD#(Me1-5)%o9E3_b-}&Q6b=LEniz;U^o>%t$;?tUE=UXIsRMUwy zDyOLu4J?}SQ^hblFF&Q3R5VKn^K_}@!h~g*ONo=;|KaXepVG)-RLc_8%Tqf~{?^Vp z3K2|WdDIBmbeRaosb+&b%1Z%TC^;&{GQaQyG~`g`RJVjRSRxiar{{}sRFmS}Uu99M zde+gC)uAHX6$0lJcJ^6^s>ZNlV!du~r79I9^E#H%<_guPD2UUUzktPkSD&g%Qz-IK zjY!2OBL4yztvN-fG+Da!8?GQs0ZUf%%2QbNdRUV!_xIjX&eO`tFbm8!Q(~H< z*6N8{$46$7z`o5K=1V|2P+m|W4>;R}5X(kaRce~jEQzsF4imrJ|ISM{5(s_}VbHB1 zVVNGO*fV@pUxSjaVhyT1=f71YOSRsw{&uBGLwQ0|xojD0|KL@YclpJ~Eu;A|gs)Yr zt)d|dTB~h{MSadVbxobdiE4!oYUP)|^cqikg=6|2kl{B7JdsdYRA(`iu@_eA1&(mP zVXR7*!L}+?EUptgPHmB%8Z%Oh3W?XVLIKQs zf>h_7*$M&_vi_AT)+J1hK0$v29>X=3C%%RCezxM+z};pYT=pp{OB+1=tOhb~a z?_EsIt;Yuozj?bF5NInC4oCuyzowV?+*i%N^VSg-7w**>t-5bLFxB@{ORDKZPM=;Ger!x+gt~JJG z*oCFlU@NA#elgj^Yi@+g##Pmy);<%4dN>l=YO8o;XMW#X<|2ksy!FjlIQabPtBauf zQVa*zAh;kluSSrS5Tga={>sEO@*DLDqHP6HeS*uqbZD8}f-;L_Qow4)Fx^7Yf zRdkH#Td4%iY-(`e>ups2%kKW#sKPJ{+o-Y`=Q;k0py@(JC5B!pm--YPSYk+bh%2T; ziJ;$VUOi2kw9+%67EQE3Uz6qDu0C1_qwnc|qn&940tDTO)Lee_IR=3lsi%`)Sc_t=XuPxaL}adAc&A}sJI zgVC6FC>29~{kL#SEYL`v>O2s#6uCb&>y6SUh^E_p>-n0N5A039+&5p{eM`G4P(bI+ ztfGf5#Z>pSYmzn3&2=pGrDt(MzwG{;N>Jng?aP(){B#A{wKh7X<$h(iEZ4M@}zzBEF?x6TGITB+1#i{HJh~Qbd*aN}a%3EjS-rzK~3x@k| zH`Y7cw_ljtt$8Ym;EW9tSGVIK2><@l&(r%0+%0_ZX=`}?T^jaDE?cq+8fRpR3FkMJ zh|BKqU0AYC6vM=JSC?)k%z57SRJeOrn{uyAaB>KPU+8`Mm%mEx?fks-a?(@RJ7*j$ z!Bu&ImL;CvXySWIDw`JFh`PqBmp>rZoP6|bENz()cK|aprb;Ik4E^Cf`ZeUwa}7gU z0u`X`Auvki9V&x1)yjV5lySh_RAN}iC;=9`Ql}Hbwy>+p1Y^73+df0NNI!=n@ZHgQ_ztP>2%Z-rC z6SD8RP0oQ`sgXP;O$xzO+7-v!o=rKz7c9%~0Yj)Ln4W+0jh%~D{tfJDZ8^k!&tzI+ z;ZEsybG{b}{!n|zy%OgONB^nZq3)6wzm36Ch)aglvr@LryI_lw_+6WvUV57Ss-|$B z{#I&y(WdctWg2;Ey_QL-*6V#XRz@PgcxlUPBZlUKAM})rKk}d$gy;N4;EIRTWiSrx3=n(V0I4K^7M=uO<(kdzw$aK75q^C}H|XajAr_E1nYM>MAPto|O2Gn|otBfNhF`R>f= zzteaK(P4-cJxF}p1%!8c5jV72F|`=uy-HJ4lHdKoODcBle)XB1-+v|St+s-rH^lxn zc59Js@+cRhcVh$H*xzC&23(;GGPT?pyJC?^jLxl3&S00zt`^~WzQ~QMs!RbpMH|`~ zZEL8VMc`h8`w>^LD>t%GGY?_|Z9YmDS-w!Du^9Zu*LtRap5=+_t(P2t&P3ClbjgE}7fr*F zTuQ~&+C^W7er{qyZCqFgC6?Z`Rk3-OYJnQeYJD+BlYMo)%761<%@jv&QjAOq z1X6$UNp`I<+gfhai(+a9Z$f3{~+*JRkbYCxA){rHzomNad2#llW41yb?$;< zsCXO==XQkF`*}yrkx;aLStYyEQFGUpg}IstW6_oy)nf9)>OxD6Ic8UCw393M&8t4C z4m_z5*mCOoV&^Of_dQ5?rsd+RMC)oL;2#V_o|v`V)8d^(RmARfKoK+Se+h@&kw8>t z+RC$NgVe$2Y7mm^up*4L^aQBLTlYdtt%PVyI_6n1Rb0b>u1(`iKC^)`|Is~DQ{|{( zT)(dr%k{71PEpewCs$6zk2|x1Lslgv-jURbl@0dv!g#6RH^m1A=!_S19rGtL za!+35c(BXQ*r9l-;4Odro!GM#E1mE3!f!}bU7-$Irr5psCi=||eF_x3ij>y;x_mG$ zvC2TETn(F=E3&m}Hkp^{nwdlHtZwz}6C5);*KlO2M!`jv2qEFF6IesrwXz^6Wb^Ev zqtda+Y)GT$8_+yVi4LK+aKO`GxKlP;i6_k+KTDcRPi0UMsuV`(@(4PC`w2e3<2IR=u0mQ~-GgoA!KCt5(tnmFA;GAS~})YK|-2sC$?sh+)B8DTbzrSY7tqS zh`50*Bj+z*!fGKvqLg#6QE|5J>>M4;QJ&&<)=cqz)l`)=i@g0FlPPkTj+o9jMEHtF+{yg~m+v3zFGY4cv;_U9hj z>^8T{;ej)}TWm1=xxzn=ZeIT4vfdP}NByhMJ)~E-|I*C2t!{HPs=jVddcAEv@%a^>nJPp4-=(c}pd6^^8Cgb9lf4NvmZ^q%srF9Na{B(0T>++2IWfZgK> zTPcsQp5j+DXhh!U910T?b%Wh4dE$v!niWM~yh5jZ%% zW5BDVH~lRx+$ToOe_e;Gc=>HLg0pSy78aD<5V`mRA3(Zvpr8vDOk)BuYHUtM8=re< zwb?p-ARESzl%83GQI?CXVP{;%g^l1Yz71S8;{HdM?`rmo?ZJdsv@kf>WHy$&G1#&< zfE)6PGtFT~Zw%Dy3s*~t^Muz6e&iH~mOgBEJe=La5PZMh;l0K+g;9T4^F(^s;ipa$ zJ2$o2(#wLo<=R*$sp{8Ge-n8|Cdy`yw}I(ref_ZG<00YeVd;T#onP!7i_JDL3R1fd zHhEaD>lv3iZK<)mjL451pjMN$6k8iW3nu1|@vxrf&B$Y3%G}-Nc-Cf{SJ)a{J|&Vj z;T7w8Phg4K^dF&&*eqMpfrmfWfB5Eki)s`n2M&b68V(8Uo-b ztxuZP!0hyTWxK<2YtYUDv>3-MWi%da>h+S%whhanw3F^wD1|7Xs_bo!x5K1?D|mwQ zY~%VObqu)3Lc;x^XqxF5I}t_`^zs##ZJH|3KEO0f0*wQv^cHu4ZoA^G|jT$ zzN3RXH~S?4yaF)WL+3zvieq6suqkIJD4Zj3R-)57k6cW_39lW~Gr1#$){ z>+^o3;^9Pe5nA%0W!f5B1AlF&i9tiM*&Y>7&`yB+WE-wK5uf=PoIcC`NuntA2emqg1!q!U)|0`Yvm|F+-38wpg8)I zh@|@l(16p48^r@cu2f8mS5}?Yd(B0XnFoF5tqyA+l_hUsmX~2~iWVfz@HV}$M!lxI zU7k&!?9FE|Utq5Vh^ZDQnKvl>ReNa#K5L`+HGwN6bVmbmLa}o9WE1casi}KKsCN2I za&eZIe@jf98h(}|bM>L8nG(%gQLAWgWv%CQ8*8NE4G^38h9ZM#N-Z`t)e2r zXr)-4jK@N19<*TC^1mLGoj z;pJ{QUM~3WBtRi$zfzo%oo?I%u^$|4jShU2v%?#Ahxe5qxU(pGg?mgG7`=3jzm|ECvQ+bPf6@7@sndSL^)UL3tNeM+@79o4ggxM)btEo_AS9iXI<(Z@o{e|5rCyQEz$`7 z@wOE-s@`lj#@EErL7I(nLqrP+(?}VdDA$@Uix%KvmzZN`Hv#TJBx4vjBy+xkDos)u zRbSAl{cs91f8=2C#TA$4+l}9C zx?iaNykq6jjK>$uJH_PYf57SU9jnOeHq+Fk&fE6O9V>;ghFH*Q8UfhsGx#or<#zZ{vjTr}no6D#n=T(d*SO)x~pfr*-d6?L(!%wp2tECKj#z2Zj zNVl*;k4SHz#+g-d(J`7r_jPoDmY45gX@P?)*KA6^CimF|wys=TL^_B=B1)#Vi$L&$ zBHyO0M2P9ll@4m&^|AR=TTQZoFmCRWdl01uaNuSPgXq~TU=z4{EyS~S0My}?Ile&(IFs;SK^>+@z*25iz?z;%6fHm?F3UAjbDRYV@0*Sj zX2F<}WV0c8Bry%<9AT=*4rsCaLl(iwIKo{%;NFZ7CFlKMFkluPN>&a!Ce!Dei(mP0 zVcuj6Nfeeaw~vE=Z^PIf`o+WdUny>}Z-OJLAB3eK&#ct{A|M^4Cv{Uc?kNj+k|tI@ zv<7i-=Z)^8FWR)QpazNmMNs_T_ZNrLuW_vCD9>Wb_%-^KuJW)niR78iVliSlM~CPY zHB}R5BQA41z(-EHD4@X`qJ(qMBVA18eRC(jZssjVyO zDycIIHY}d8zf9NIU$@oV%LM?;Q)r_2;yCKc8cq6bQM^zY2s#@oBm6=!s%nwl?|$azfk1mbvCyakq0Cw+3*P%sta7w?1;R*&9u8 zVNTw9Yw^i0En?)w!!A5a>jL+OufFLjEj>a$Iq=SXy{WTnA^7{`I?aoML7Pq6# z;_$7v-1i8pS>uU~LaDO`1ys?gM2WzY_zT4)Gjy28)1B~ktZ3&Bp|ClKD5@P07OF9mk84}i`(;=CZGx)T}Y8l<=WwU z%z*~|1aHA1H0c#5`_0Yvu(=LaKJ=jv2_G$za#|=quTnc|2b}cLVq}j8#T|FtQrvsu zlg6coZ@u?UG4_K;?>9PHGBhLRFC4t(mRpM6(GBvMV@hV49A+K)%cHPTO z*SHc_l9ZoN1zP_QQ5Ka38cM#=NOcvHmxfH6CB&&0ZL1hHD!P<}H6wB8@e`6x@vX=C zXjeE_m8~s=cF3yizBZxKXzZU%UD*V6LKZ#pe_ev2M3BcD z0~n|~7@QU~S_E~7!RDdC+98&B$cm(uf-*)JtgT5SDD@GN%p^pDCPHyYp4KK^-nThq zIQlqRc&H*H+4VkJi$Dat;M?bD zT8HMggLi!BgEb{9_!sF|^xhVVLn()A*Qc}GBv4Fj>*{n=!y`iBGF4!%C<#XiajVb| zI#)_WIvviULZ+qzYRPhp+{)k42n-m#0!e(OWP`E}RzBC!rptpFlCG2ltPciq#hU%o zO1haqNDUu~1}LI&TCdI`pgq8Sv4%B612%O_%`s+FRL*nr*J^HJ$Waf>P3av0(&WB= zGzMzKG{+PQqGsMxxte92t6@OLHsi5Ov?j*iL zs@Zww;{ypS$%)?27RODWDh1r%CfSRbv4#zI);G%TrY8qG*BzPW5<4#x}=;w^Q$wPL%-lUUYfs&RU`Iz{WBblDshGw6{q_u70alItt!8v zN4U^L=QQ1xV#D}Ds-v}0_Vl&Z##G~uTqD8YAdL9+P-}>XgBezL5wGCg#`8`SN_eb+ zO2t)qoSN6Hp4zAS>*I}sBsn{4YgDap&@|&N2CC~F#CvHw zngr^g(>0LSo9-||mN+d;Y$5%*O;P`Rx^#a1+wOA@gn%r;j$U8y4U-WFUb)V7^W zxUSIMLo&1fZfJU?H(zZI)bA?SjTe0L%x$;d_9qwbDh}Ln`!D{<#d~gc9zv^$d03|M z=lO^Q(HV2z`U|(%Q@4NUR(nbe+Lf=&Q=I{IUM8?EDgmcn+U~$I?*hGI-JC}&qPwle z&$8w-XB*4OCw@slR!c9>ZVpB&QvnqF~DGt5B!$d@)j zZ-6s37_LjOFo<1R5Jg0kemjs(ogLxNyjDYqSz{mX`En;XNHj25SijcnbaOo$iQ1$| zb2x$(@N^=7Mblz-_u6D1{%@MwjB$}S+=dPFaG(kM;9=42vTmWq46oEc=#R-Pl(~3gxlZkF`?yZhYU$p8zUo}ZSm;Bq3HnQnb+AB@Alihm zh7_0Y)ebS1DV-WuKy?D%nGtM35j_YtJZvfnMyA2AW{nt}zzsUkM$^;t1NT6Ln^LxLsqeR+CPE!I|dBF%|CSmcbC6z9|lmL-3`E+i8Y@6UDZDQhm4&KN1r2H(ta8O)i%44~#RBapDKWzJt;Cd5aU zmk&3JyP|U$&1Q?;ErwS)85IhupJk7yZA?T^joX5e%?#|cA(k?OLScd1km59)prFZ& zT1Ks0*3=a*w}4P6TcjFI$b;bu**1bPyyIC9_ZXp(BETy7C*-5%N#k0y1J7~?Hd5X2+>sb8Ddi+G5TAyvcVHkx!`66 z^Knli~JmMZbcGY)w6E%G4`G% zq#XKZ6lghXC5aS8nkhFyC^nIt;@a9?Quht0>{@vOS*8UX-T{uHXBh5jd4nZL~>@ZPA}dv>@qR3Xw7PEzFt|eW0rLw zL(=KXJOrA>`JrV+2GLyvU;{&Vh^=pUduyGNo>3-Hg4Dv2SsFP6{bql#)}eBajh{kJ zFqr{I6=7nd=?y?uTt!&A`3#Py1ZTpS8A%BY7T@GoB&RSmMf1JTcSSDtX3xlWN`|?e zTEJiC_NFejcM@>|E%4)F_HQg>J(d1c8WA_BFJ4@PpZL3bjmqByyAkq`rs>mK6L;8+ z>eW|7m}%)oagT|=wLgT2neTnXcxD%DWAuAa%NI$RLnkmSvDwCSG_wq423 z8E?;x%N(m?j&3#dT%cU9TdNG64{8w!wVJR6Eh9}Zs!6a6W#}#`2<>R5O{Ox84Y#vr z){Fa7X#VLEUwVuvvfY!=hWu0&wr|;Qn!+ADAoiP+ly~})*Fn1)xMJp1U7BP0_6D*z zwY@3Zt7^Gcaz(f-twyJW_Rbx?#ri3bqrSfht7^I>8GD1Kt12`bXJnw_)L8A1rN$}- z3#KLjk#uIxNXj5E6qLt9j6G};f*T=xK2Kshy&X){)on2sK#+$rlIMVD;+9j@#)&7u z?Jn{kUIJ(oAG4C)3}A-&02PREbHJ4DURlRW97!Gr8!Lt)P=X<}?>oZxB=5pWdhAT~ z0o9$(Bucf<%#rJ!eQ~Gen=c#lL{uFML!JpezCkRX zG-K_?=#-alM7O=z)ErGF3Op(nm_J3QA96zCJ#R>0=;S#zP7OATTj`GK4pxcvHSKYj z2NYEa7u=boa#+w*@;9ZcYXID*?1K_CA_-k9nqz%tj*79bFh{Z(`66-31K2fK-Bk}0>|VZup!w*@dy13CmOt1~R~l|1{K)_xY7`HiI>F5}clEcIe-ZDK1c34Ey}i}@ zmJi!ewR|>vq*Np?B7u9d9TXrQCaPAs}hgkSW8^y1( zHQjQL&5_psDVQ(hm=9m|TqOkv3R=r+~M=v0BE zF^wj)r6PZ1K3&Zb;J8kKu^^1)F8T^wgOUo7n=++XtYD3Ms8BMqa-?kgtl||@aM8Ci zr*{r-&*eLD6InD3Y0VSOVs$v6QiJx3^&}hIAwtn=;5M~mWCBL%R7@vv%Q-r2hLB&e zu9%9B(Q>E+XXrQ@b$T5o@UiCF3$2vK>fkJ+w?D$s2G%A_R~FQ_B6-1vxS<+o_ic_vS zf6B77qa({dUwl$n2r-K1ZVe+20boIBA*u$4bvNqU&pp&>vn_-S6b_rd&<=}`vl-A= z!BI@k9t~!CF!1fd+2$`Qcjx^MpO=j_*b#$mYXg)?j`U3&g$?xfXUU;GoUUbe6Wd!o zLr?L$WG&vVgQ6%njfw*>GUKQZq#8^uBXaUlI6T&B(d%5Ya`?mxEM@y<)K#PSq&I}c zLCLt$l4U$W3m*-o@dr3FotF^Gz40f1#-&Bu=LM^}<8`c#^~c+ieUCRS2=07k;-D#k zkqio&jev<|z%bemuh9@5E;{e!coXn2Iq^pI|zXlpHN7?_+0L6o9uz)x5TnF zuKC8zVq(oAFah$D4}J`%Dc{)H(~yiUq=b7N0d4RmwHB5Lg_!6?rb zG#idJZ?z=>_>{;~k88xRYhT;per?E|Y$CHRQV}?!bt`-%{AwpdE}X_}^z_%*Dlr`> zcwa#A`@%UdwhRU_a?7T%9V{dV`=yqoSPoTf`)t*js>*(; zpltkXyrHroB^6MKY-DSUc&kAhh2VWS+b68mCHT7s5Q@l@^!CtnE>Vkga;#k)dG47~ z$8(m4JZU?Nr33?=kzgYRdxS!3ceF%^X+^+K!u*kVP=U27(PvZ=2?Lwq2rC1^S`qT2 zCywiUJL?}pxbH?vy126A#|w0i2hXqbUSE= z;44J4!j7K~!r+crrZ1%&DlHU$PLn1v@fBR- zEQ`u5&Z@JBgxcHV!~k)*rvI_t)-HFXIBI8GiBgsj)b~-R&t-%&xYofSn+52SAR(?W zTjg~rHgAubbX(d{l4E%nq$0BU1LE4~adwdmIUn%uuyIM?3lbJrA;Hj`8)@Bj1pyaP zxDC5hCyn4Z-aiB$&AX31FL(?e$Js}{IUcakw)<)OeGpx~n^I^qV(!VNdCI~}_0fK_ zyG<_N6C!k8t4$lM99h2K|28Se19m)Ot*u(R(VqmE%#QHMi99qjQtk+c8(bL^C`NaX zTz$!>2)I$kMlqF2IcfrtKkhTLSVIG9%sg^j6SeT zkEc=Y3B_>$r<&<}Jr=O&*h{j@4WRc$foZOHwv9%ImLYLJiQ1}!}Yu(o-L1%eg%q3nQB;sGlfXOS{#l|YAW zQy>UGis`E$jGj4ci7eQm$)+i?Og0#HJNeSVMlk_fCAAkWC=8}|#aV+d8X>^ymt#f0 z5RcY7D?FOh)B{%sXu@74xydHi+hAjX=I5%hTTkEC!IYJV%KV zyq*38IjczM)(|yqa$kd*=K@wu0RWtM#3EStNLWK>WZpi@;cjQXReTBd1hOsKPuERowm>xV3-;BsQ=D~^>-j_Q%+{1Xs`IxVBdR2-yPCj!NNqQ#3^r?)(7 zj#&PY1i%doFe-`mB^Y4VI4|8B^zAV}TqsGA&L}0vV^$JeuQuDT94dW-+!4d*Z_OQ~ zr=de~(`uY*kx?r?g{ivAn$#e)JPNIb6w}qP$`m0UsESk~0@?kj0Tt&u<&93Cr`O%8 zlXsa!~xROLOaN$YJABJzp7%;kW8HwSH1_rs!ccP4}<&YOr_*!`^i=Q)to zAv(@uv5_ulVzN6zDDu>V^+kE@0|`dd0-@r-$zu=rR)8>laLlZKh`rtHcSkZ1mkOUJ z=z8HOM9GP17}gf<65}{wNyXtwB;}q2qpny2_b3Y>brf|ctFC5t@t>pJDm{^ceiU73 z8&NlcvL-T8#$1vZ3Ww6u^p9Bl!`&$}XUMSp5(ZLWJ3a#*<={FKXwlKhSBnqZJyp82 zE>%m6nrgzYLEF_h+?3Q-M<*1JQ$DHOMetnZ2x)JpD`B06jMd*J?z#`w(R}H5R-ArT zcC*GXrX=Txn<*^EkOAZr5HDvqjXa-OjF1zIQIXqpnkl2bo+9a)&WMIVZB%g+gNrtl z86}Mxw9!_%N~Jnv&jLCtTsD+gKc+-m&T4mZ7+5IaxYJMY0Pzg^kgypyhh$>qSmMgE zx!jW7ghZJXYuhdsBsIZ@<0yh6G~E}3JD61&uATtTP!x@RhN5Ux&4BbLL*Eo3#p3Lw zlx2qbbVR7lScz`J+R5B$mY`_EvF=P>ca_j5f+7EpjGOl51PIO_kUM8i29;}`lmk!~fy%Id qw0Wkfh{og*MZr<~P5L%?iYNdti8>7fbGqj#M3~!t#D7ct_x}U0FE_&g literal 0 HcmV?d00001 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 From 412666927d4eee4eb84ebc21f1f7527ad4e085f5 Mon Sep 17 00:00:00 2001 From: Hassan Eslami Date: Thu, 14 Aug 2025 21:23:11 -0700 Subject: [PATCH 2/2] Add a config to load/save model from/to DCP checkpoint --- cosmos_rl/policy/config/__init__.py | 4 + cosmos_rl/policy/model/base.py | 3 + .../policy/model/deepseek_v3/__init__.py | 88 ++++++++----------- cosmos_rl/policy/model/gpt/__init__.py | 1 + cosmos_rl/policy/model/hf_models/__init__.py | 1 + cosmos_rl/policy/model/qwen2_5_vl/__init__.py | 1 + cosmos_rl/policy/model/qwen3_moe/__init__.py | 1 + cosmos_rl/policy/trainer/grpo_trainer.py | 1 + cosmos_rl/policy/trainer/sft_trainer.py | 2 + cosmos_rl/tools/model/deepseek_v3/__init__.py | 1 + .../sft_integration_deepseek_simple.toml | 1 + 11 files changed, 52 insertions(+), 52 deletions(-) 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/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 index 6c032eeab..950eb2e36 100644 --- a/cosmos_rl/policy/model/deepseek_v3/__init__.py +++ b/cosmos_rl/policy/model/deepseek_v3/__init__.py @@ -44,9 +44,6 @@ from cosmos_rl.utils.parallelism import ParallelDims from cosmos_rl.utils.util import clear_weight_name, resolve_model_path, retry -DCP_CHECKPOINT_PATH_PREFIX = "/root/.cache" -DCP_CHECKPOINT_PATH_SUFFIX = "dcp" - @ModelRegistry.register( DeepseekV3MoEWeightMapper, default_data_packer_cls=DeepSeek_DataPacker @@ -220,17 +217,35 @@ def load_hf_weights( parallel_dims: ParallelDims, device: torch.device, revision: Optional[str] = None, + dcp_snapshot_path: Optional[str] = None, ): - dcp_checkpoint_path = os.path.join( - DCP_CHECKPOINT_PATH_PREFIX, - model_name_or_path.split("/")[-1].lower(), - DCP_CHECKPOINT_PATH_SUFFIX, - ) - # if it is a huggingface model and no checkpoint exists, we need to load the weights from the safetensors files - if len(model_name_or_path.split("/")) == 2 and ( - not os.path.exists(dcp_checkpoint_path) - or len(os.listdir(dcp_checkpoint_path)) == 0 + 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` @@ -370,47 +385,16 @@ def load_hf_weights( with torch.no_grad(): local_view.data.copy_(shared_weight) - logger.info(f"Dumping the tensors to DCP folder {dcp_checkpoint_path}") - os.makedirs(dcp_checkpoint_path, exist_ok=True) - fs_storage_writer = torch.distributed.checkpoint.FileSystemWriter( - dcp_checkpoint_path - ) - torch.distributed.checkpoint.save( - state_dict=self_state_dict, - storage_writer=fs_storage_writer, - ) - else: - logger.info("Loading from distributed checkpoints...") - model_name_or_path = model_name_or_path.rstrip("/") - if model_name_or_path.endswith("_hf"): - model_name_or_path_dcp = model_name_or_path[:-3] - logger.info( - f"Found model path with _hf prefix ({model_name_or_path}. Looking for non-hf checkpoint at: {model_name_or_path_dcp}" + 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, ) - model_name_or_path = model_name_or_path_dcp - elif len(model_name_or_path.split("/")) == 2: - model_name_or_path = dcp_checkpoint_path - - # 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_checkpoint_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.") def load_state_dict( self, state_dict: dict[str, Any], strict: bool = True, assign: bool = False 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/tests/configs/sft_integration_deepseek_simple.toml b/tests/configs/sft_integration_deepseek_simple.toml index 9352b599d..a41679587 100644 --- a/tests/configs/sft_integration_deepseek_simple.toml +++ b/tests/configs/sft_integration_deepseek_simple.toml @@ -24,6 +24,7 @@ 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