add to_hf

Author: wwwjn

torchtitan/models/deepseek_v3/__init__.py

@@ -135,7 +135,7 @@
         dim=7168,
         inter_dim=18432,
         moe_inter_dim=2048,
-        n_layers=61,
+        n_layers=4,
         n_dense_layers=3,
         n_heads=128,
         moe_args=MoEArgs(

torchtitan/models/deepseek_v3/model/state_dict_adapter.py

@@ -4,24 +4,36 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+from logging import raiseExceptions
 import re
-from typing import Any
+from typing import Any, Dict
 
 import torch
 
+from torchtitan.distributed.parallel_dims import ParallelDims
 from torchtitan.protocols.state_dict_adapter import StateDictAdapter
 
 from .args import DeepSeekV3ModelArgs
 from .quantization import calculate_scale_shape, dequantize_from_fp8
 
+from torch.distributed.tensor.placement_types import (
+    _StridedShard,
+    Shard,
+    Replicate
+)
+
+from torch.distributed.tensor import DTensor
+
 
 class DeepSeekV3StateDictAdapter(StateDictAdapter):
     """
     StateDictAdapter for DeepSeekV3 model.
     """
 
-    def __init__(self, model_args: DeepSeekV3ModelArgs, hf_assets_path: str | None):
+    def __init__(self, model_args: DeepSeekV3ModelArgs, hf_assets_path: str | None, parallel_dims: ParallelDims):
+        super().__init__(model_args, hf_assets_path, parallel_dims)
         self.model_args = model_args
+        self.parallel_dims = parallel_dims
         self.from_hf_map = {
             "model.embed_tokens.weight": "tok_embeddings.weight",
             # Attention Module
@@ -52,7 +64,7 @@ def __init__(self, model_args: DeepSeekV3ModelArgs, hf_assets_path: str | None):
             "lm_head.weight": "output.weight",
         }
 
-    def _split_experts_weights(
+    def _split_experts_weight(
         self, weight: torch.Tensor, n_experts: int
     ) -> list[torch.Tensor]:
         """
@@ -84,6 +96,134 @@ def _concatenate_expert_weights(
 
         return None
 
+    def _get_local_experts_weights(
+        self, abstract_key: str, layer_id: str, grouped_expert_weight: torch.Tensor
+    ) -> Dict[str, torch.Tensor]:
+        """
+        Spliting the GroupedExperts weight and find the corresponding individual expert's weight in local tensor.
+
+        Potential experts weights shard placements:
+        - FSDP + EP when dp_mod_ep * ep <= num_experts: 
+            - StridedShard(0)Shard(0)
+        - FSDP + EP when dp_mod_ep * ep <= num_experts:
+            - Shard(1)Shard(0) 
+        - FSDP + ETP + EP when dp_mod_ep * ep <= num_experts:
+            - w1/w3: StridedShard(0)Shard(0)Shard(1)
+            - w2: StridedShard(0)Shard(0)Shard(2)
+        - FSDP + ETP + EP when dp_mod_ep * ep > num_experts:
+            - w1/w3: StridedShard(1)Shard(0)Shard(1)
+            - w2: Shard(1)Shard(0)Shard(2)
+        """
+        world_mesh = self.parallel_dims.world_mesh
+        num_experts = grouped_expert_weight.shape[0]
+
+        # Matching DTensor sharding placement and device mesh dims,
+        # find the dtensor dims that shard on dim-0 (num_experts dim)
+        original_placements = grouped_expert_weight.placements
+        world_mesh_names = []
+        dim_0_placements = []
+        for i, name in enumerate(world_mesh.mesh_dim_names):
+            placement = original_placements[i]
+            if placement.dim == 0:
+                world_mesh_names.append(name)
+                dim_0_placements.append(placement) 
+        
+        start_index, end_index = None, None
+        # StridedShard(0)Shard(0)
+        if len(dim_0_placements) == 2:
+            assert isinstance(dim_0_placements[0], _StridedShard)
+            strided_shard_mesh = world_mesh[world_mesh_names[0]]
+            strided_degree, strided_rank = strided_shard_mesh.size(), strided_shard_mesh.get_local_rank()
+            shard_mesh = world_mesh[world_mesh_names[1]]
+            shard_degree, shard_rank = shard_mesh.size(), shard_mesh.get_local_rank()
+            start_index, end_index = self._get_strided_shard_shard_slice(strided_degree, strided_rank, shard_degree, shard_rank, num_experts)
+        # Shard(0)
+        elif len(dim_0_placements) == 1:
+            assert not isinstance(dim_0_placements[0], _StridedShard)
+            shard_mesh = world_mesh[world_mesh_names[0]]
+            shard_degree, shard_rank = shard_mesh.size(), shard_mesh.get_local_rank()
+            block_size = num_experts // shard_degree
+            if block_size * shard_degree != num_experts:
+                raise ValueError("Not supported. num_experts can not be evenly divided by Shard(0) dimension degree.")
+            
+            start_index = block_size * shard_rank
+            end_index = start_index + block_size
+        else:
+            raise NotImplementedError(f"The DTensor placements {original_placements} for GroupedExperts is not supported in StateDictAdapter")
+
+        # Calculate the new placement for individual expert weights
+        new_placements = []
+        for i, name in enumerate(world_mesh.mesh_dim_names):
+            placement = original_placements[i]
+            if placement.dim == 0:
+                new_placements.append(Replicate())
+            elif isinstance(placement, Shard):
+                # Individual expert weight has only 2 dimensions
+                new_placements.append(Shard(placement.dim-1))
+            elif isinstance(placement, _StridedShard):
+                new_placements.append(_StridedShard(placement.dim-1, placement.split_factor))
+            else:
+                raise ValueError("Not supported new placements!")
+        print(f"Original placements: {original_placements}, new placements {new_placements}")
+       
+        assert isinstance(grouped_expert_weight, DTensor), "GroupedExperts weight is not a DTensor"
+        local_grouped_weights = grouped_expert_weight._local_tensor
+        assert local_grouped_weights.shape[0] == int(end_index - start_index), "Local tensor shape mismatch!"
+
+        # Create new DTensor for each individual expert weights
+        local_expert_fqn = {}
+        for expert_id in range(start_index, end_index):
+            new_key = abstract_key.format(layer_id, expert_id)
+            new_value = local_grouped_weights[expert_id - start_index, :, :].squeeze
+            local_expert_fqn[new_key] = DTensor.from_local(new_value, world_mesh, new_placements, run_check=False)
+
+        return local_expert_fqn
+            
+    
+    def _get_strided_shard_shard_slice(
+        self,
+        strided_shard_dim_degree: int,
+        strided_shard_dim_rank: int,
+        shard_dim_degree: int,
+        shard_dim_rank: int,
+        dim_size_to_split: int,
+    ) -> tuple[int, int]:
+        """
+        Given a [StridedShard(dim=i), Shard(dim=i)] placement, caculate the start index 
+        and end index on dim-i for GPU rank (strided_shard_dim_degree, shard_dim_rank)
+        
+        GPU Layout (strided_shard_rank, shard_rank):
+
+        StridedShard Rank                  Shard rank
+                        ┌─────────────────┐  
+                    0   │    GPU(0, 0)    │  0  
+                    ────┼─────────────────┤     
+                    1   │    GPU(1, 0)    │  
+                    ────┼─────────────────┤  
+                    2   │    GPU(2, 0)    │  
+                  ──────┼─────────────────┼────  
+                    0   │    GPU(0, 1)    │  1
+                    ────┼─────────────────┤  
+                    1   │    GPU(1, 1)    │  
+                    ────┼─────────────────┤  
+                    2   │    GPU(2, 1)    │
+                        └─────────────────┘
+
+        Calulate the start_index from inner dimesion (Shard(dim=i)) to outer demension (StridedShard(dim=i)).
+        """
+
+        block_size = dim_size_to_split // (strided_shard_dim_degree * shard_dim_degree)
+        
+        # Error out if can not evenly divded
+        if block_size * (strided_shard_dim_degree * shard_dim_degree) != dim_size_to_split:
+            raise ValueError(f"Not supported split for strided_shard_dim_degree {strided_shard_dim_degree}, shard_dim_degree {shard_dim_degree}, dim_size_to_split {dim_size_to_split}")
+
+        start_index = block_size * (strided_shard_dim_degree * shard_dim_rank + strided_shard_dim_rank)
+        end_index = start_index + block_size
+
+        return start_index, end_index
+
+
     def _dequantize(self, state_dict: dict[str, Any]) -> dict[str, Any]:
         """
         Dequantize the weights from float8 to float32.
@@ -149,14 +289,16 @@ def to_hf(self, state_dict: dict[str, Any]) -> dict[str, Any]:
                 layer_num = re.search(r"\d+", key).group(0)
                 new_abstract_key = to_hf_map[abstract_key]
 
-                # Split expert weights into separate expert weights
-                split_values = self._split_experts_weights(
-                    value, self.model_args.moe_args.num_experts
+                # # Split expert weights into separate expert weights
+                # split_values = self._split_experts_weights(
+                #     value, self.model_args.moe_args.num_experts
+                # )
+                local_expert_fqn = self._get_local_experts_weights(
+                    new_abstract_key, layer_num, value
                 )
+                print(f"groupedWeight placements {value.placements}, local experts keys {local_expert_fqn.keys()}")
 
-                for expert_num in range(0, self.model_args.moe_args.num_experts):
-                    new_key = new_abstract_key.format(layer_num, expert_num)
-                    hf_state_dict[new_key] = split_values[expert_num].squeeze()
+                hf_state_dict.update(local_expert_fqn)
 
             elif "layers" in key:
                 abstract_key = re.sub(r"(\d+)", "{}", key, count=1)
@@ -169,9 +311,11 @@ def to_hf(self, state_dict: dict[str, Any]) -> dict[str, Any]:
                 new_key = to_hf_map[key]
                 hf_state_dict[new_key] = value
 
+        # Prepare for dequantization
         hf_state_dict_with_scale_inv = self._add_quantization_scale_inv_tensors(
             hf_state_dict
         )
+        print(f"[to_hf] state_dict keys before return: {hf_state_dict_with_scale_inv.keys()}")
         return hf_state_dict_with_scale_inv
 
     def from_hf(self, hf_state_dict: dict[str, Any]) -> dict[str, Any]:

torchtitan/models/deepseek_v3/train_configs/debug_model.toml

@@ -47,13 +47,13 @@ dataset = "c4_test"  # supported datasets: c4_test (2K), c4 (177M)
 data_parallel_replicate_degree = 1
 data_parallel_shard_degree = -1
 fsdp_reshard_after_forward = "default" # default / never / always
-tensor_parallel_degree = 1
+tensor_parallel_degree = 4
 enable_async_tensor_parallel = false
 pipeline_parallel_degree = 1
 pipeline_parallel_schedule = "1F1B"
 context_parallel_degree = 1
-expert_parallel_degree = 1
-expert_tensor_parallel_degree = 1
+expert_parallel_degree = 2
+expert_tensor_parallel_degree = 4
 
 [checkpoint]
 enable = false

torchtitan/models/deepseek_v3/train_configs/deepseek_v3_671b.toml

@@ -38,28 +38,30 @@ min_lr_factor = 0.1
 local_batch_size = 4
 seq_len = 4096
 max_norm = 1.0  # grad norm clipping
-steps = 10_000
+steps = 10
 compile = false
 dataset = "c4"  # supported datasets: c4_test (2K), c4 (177M)
 
 [parallelism]
 data_parallel_replicate_degree = 1
 data_parallel_shard_degree = -1
 fsdp_reshard_after_forward = "default" # default / never / always
-tensor_parallel_degree = 1
+tensor_parallel_degree = 2
 enable_async_tensor_parallel = false
 pipeline_parallel_degree = 1
 pipeline_parallel_schedule = "Interleaved1F1B"
-expert_parallel_degree = 1
-expert_tensor_parallel_degree = 1
+expert_parallel_degree = 2
+expert_tensor_parallel_degree = 2
 
 [checkpoint]
 enable = false
 folder = "checkpoint"
-interval = 500
+interval = 10
 last_save_model_only = true
 export_dtype = "float32"
 async_mode = "disabled"  # ["disabled", "async", "async_with_pinned_mem]"
+initial_load_path = "/data/users/jianiw/model/DeepSeek-V3.1-Base"
+initial_load_in_hf=true
 
 [activation_checkpoint]
 mode = "selective"  # ["none", "selective", "full"]

torchtitan/models/llama3/model/state_dict_adapter.py

@@ -10,14 +10,15 @@
 
 logger = logging.getLogger()
 
+from torchtitan.distributed.parallel_dims import ParallelDims
 from torchtitan.protocols.state_dict_adapter import StateDictAdapter
 
 from .args import TransformerModelArgs
 
 
 class Llama3StateDictAdapter(StateDictAdapter):
-    def __init__(self, model_args: TransformerModelArgs, hf_assets_path: str | None):
-        super().__init__(model_args, hf_assets_path)
+    def __init__(self, model_args: TransformerModelArgs, hf_assets_path: str | None, parallel_dims: ParallelDims):
+        super().__init__(model_args, hf_assets_path, parallel_dims)
 
         self.model_args = model_args
         self.hf_assets_path = hf_assets_path

torchtitan/protocols/state_dict_adapter.py

@@ -11,6 +11,8 @@
 from abc import ABC, abstractmethod
 from typing import Any
 
+from torchtitan.distributed.parallel_dims import ParallelDims
+
 logger = logging.getLogger()
 
 from .model import BaseModelArgs
@@ -27,7 +29,7 @@ class BaseStateDictAdapter(ABC):
     """
 
     @abstractmethod
-    def __init__(self, model_args: BaseModelArgs, hf_assets_path: str | None):
+    def __init__(self, model_args: BaseModelArgs, hf_assets_path: str | None, parallel_dims: ParallelDims):
         pass
 
     @abstractmethod
@@ -58,7 +60,7 @@ def from_hf(self, hf_state_dict: dict[str, Any]) -> dict[str, Any]:
 class StateDictAdapter(BaseStateDictAdapter):
     """State dict adapter base class which provides convenient default behavior to build fqn_to_index_mapping"""
 
-    def __init__(self, model_args: BaseModelArgs, hf_assets_path: str | None):
+    def __init__(self, model_args: BaseModelArgs, hf_assets_path: str | None, parallel_dims: ParallelDims):
         if hf_assets_path:
             mapping_path = os.path.join(hf_assets_path, "model.safetensors.index.json")
             try:

torchtitan/train.py

@@ -23,7 +23,7 @@
     ensure_pp_loss_visible,
 )
 from torchtitan.config import ConfigManager, JobConfig
-from torchtitan.distributed import ParallelDims, utils as dist_utils
+from torchtitan.distributed import ParallelDims, parallel_dims, utils as dist_utils
 from torchtitan.models.attention import init_attention_mask
 from torchtitan.protocols.model_converter import build_model_converters
 from torchtitan.tools import utils
@@ -311,7 +311,7 @@ def __init__(self, job_config: JobConfig):
             checkpoint_config=job_config.checkpoint,
             sd_adapter=(
                 self.train_spec.state_dict_adapter(
-                    model_args, job_config.model.hf_assets_path
+                    model_args, job_config.model.hf_assets_path, self.parallel_dims
                 )
                 if self.train_spec.state_dict_adapter
                 else None
@@ -539,6 +539,17 @@ def train_step(
     def train(self):
         job_config = self.job_config
 
+        # Following hacky print only works for debug_model
+        # w1 = self.model_parts[0].layers["1"].moe.experts.w1
+        # w2 = self.model_parts[0].layers["1"].moe.experts.w2
+        # w3 = self.model_parts[0].layers["1"].moe.experts.w3
+
+        # logger.info(f"w1 placements is: {w1.placements}, {type(w1.placements)}")
+        # logger.info(f"w2 placements is: {w2.placements}")
+        # logger.info(f"w3 placements is: {w3.placements}")
+        # logger.info(f"device mesh: {self.parallel_dims.world_mesh}, {self.parallel_dims.world_mesh.mesh_dim_names} {self.parallel_dims.world_mesh['dp_shard']}")
+
+
         self.checkpointer.load(step=job_config.checkpoint.load_step)
         logger.info(f"Training starts at step {self.step + 1}")