Fix dtype bug of weight_loading. It not only occurs for sparsematmul, but densematmul

Author: bzgoogle

tpu_commons/models/jax/common/moe/deepseek_moe.py

@@ -136,8 +136,8 @@ class SparseMoE(MoE):
         # TODO: determine if we get it from external or extrat it in MoE class
         is_batch_sharded_by_expert: True if batch is sharded over 'expert' dim.
     """
-    def_sharding: Sharding
-    fed_sharding: Sharding
+    edf_sharding: Sharding
+    efd_sharding: Sharding
     num_experts_per_tok: int
     #TODO: tile size is (tile_batch_seq, tile_activation_dim, tile_weight_dim,) from MaxText
     tile_size: tuple[int, int, int] = (128, 64, 128)
@@ -155,24 +155,24 @@ def __post_init__(self, rngs: nnx.Rngs):
         shape_up = (self.num_local_experts, D, F)
         shape_down = (self.num_local_experts, F, D)
 
-        self.kernel_gating_DEF = create_param(rngs,
+        self.kernel_gating_EDF = create_param(rngs,
                                               shape=shape_gating,
                                               dtype=self.dtype,
-                                              sharding=self.def_sharding,
+                                              sharding=self.edf_sharding,
                                               random_init=self.random_init)
-        self.kernel_up_proj_DEF = create_param(rngs,
+        self.kernel_up_proj_EDF = create_param(rngs,
                                                shape=shape_up,
                                                dtype=self.dtype,
-                                               sharding=self.def_sharding,
+                                               sharding=self.edf_sharding,
                                                random_init=self.random_init)
-        self.kernel_down_proj_FED = create_param(rngs,
+        self.kernel_down_proj_EFD = create_param(rngs,
                                                  shape=shape_down,
                                                  dtype=self.dtype,
-                                                 sharding=self.fed_sharding,
+                                                 sharding=self.efd_sharding,
                                                  random_init=self.random_init)
 
         # Derive the expert sharding
-        self.expert_axis_name = self.def_sharding[0]
+        self.expert_axis_name = self.edf_sharding[0]
         if self.expert_axis_name is None:
             self.num_expert_parallelism = 1
         else:
@@ -597,10 +597,10 @@ def __call__(self, x_TD: Float):
             PartitionSpec(*self.activation_ffw_td),  # Sharded x_TD
             PartitionSpec(),  # Replicated router_weights_TX
             PartitionSpec(),  # Replicated selected_experts_TX
-            PartitionSpec(*self.def_sharding),  # Sharded gating kernel
-            PartitionSpec(*self.def_sharding),  # Sharded up-projection kernel
+            PartitionSpec(*self.edf_sharding),  # Sharded gating kernel
+            PartitionSpec(*self.edf_sharding),  # Sharded up-projection kernel
             PartitionSpec(
-                *self.fed_sharding),  # Sharded down-projection kernel
+                *self.efd_sharding),  # Sharded down-projection kernel
         )
         out_specs = PartitionSpec(*self.activation_ffw_td)
 
@@ -616,7 +616,7 @@ def __call__(self, x_TD: Float):
             x_TD,
             router_weights_TX,
             selected_experts_TX,
-            self.kernel_gating_DEF.value,
-            self.kernel_up_proj_DEF.value,
-            self.kernel_down_proj_FED.value,
+            self.kernel_gating_EDF.value,
+            self.kernel_up_proj_EDF.value,
+            self.kernel_down_proj_EFD.value,
         )

tpu_inference/models/jax/deepseek_v3.py

@@ -218,8 +218,8 @@ def _create_mla() -> MLA:
                     random_init=self.random_init,
                     activation_ffw_td=('data', 'model'),
                     activation_ffw_ted=('data', None, 'model'),
-                    def_sharding=(None , 'model', 'expert'),
-                    fed_sharding=(None , 'expert', 'model'),
+                    edf_sharding=(None , 'model', 'expert'),
+                    efd_sharding=(None , 'expert', 'model'),
                     router=router) if is_moe_layer else DenseFFW(
                         dtype=dtype,
                         hidden_act=hidden_act,
@@ -363,7 +363,10 @@ def __init__(self, vllm_config: VllmConfig, num_layers, hidden_size,
             "is_verbose", None) is not None
         self.num_routed_experts = num_local_experts
         self.model_dtype = model_dtype
+<<<<<<< HEAD:tpu_inference/models/jax/deepseek_v3.py
 
+=======
+>>>>>>> 641cb6d4 (Fix dtype bug of weight_loading. It not only occurs for sparsematmul, but densematmul):tpu_commons/models/jax/deepseek_v3.py
         self._transpose_map = {
             # dense mlp
             r"mlp\.down_proj": (1, 0),
@@ -827,9 +830,10 @@ def load_weights(self, model_for_loading: nnx.Module):
 
 def weights_dequant_cpu(x: torch.Tensor,
                         s: torch.Tensor,
-                        output_dtype: jnp.dtype,
+                        output_dtype: torch.dtype,
                         block_size: int = 128) -> torch.Tensor:
     assert x.dim() == 2 and s.dim() == 2, "Both x and s must be 2D tensors"
+    torch_output_type = DTYPE_VIEW_MAP.get(jnp.dtype(output_dtype))
     M, N = x.shape
 
     x = x.to(torch.float32)
@@ -863,4 +867,4 @@ def weights_dequant_cpu(x: torch.Tensor,
             scale = s[M // block_size, j // block_size]
             y[M_main:M, j:j + block_size] = block * scale
 
-    return y.to(j2t_dtype(jnp.dtype(output_dtype)))
+    return y.to(torch_output_type)