fix(schema): preserve tensor and group scales across shards (#142)

Signed-off-by: jinzhen.ljz <jinzhen.ljz@antgroup.com>
Co-authored-by: Codex (GPT-6) <codex@openai.com>

Author: jinzhen-lin

humming/schema/base.py

@@ -162,15 +162,20 @@ def _may_process_global_scale(
         shape_k_stacks: list[int],
         num_experts: int | None = None,
         target_group_size: int | None = None,
-        force_repeat: int = False,
+        force_repeat: bool = False,
     ):
         assert len(shape_n_stacks) == 1 or len(shape_k_stacks) == 1
 
         scale = scale.view(num_experts or 1, -1)
+        stack_size = max(len(shape_n_stacks), len(shape_k_stacks))
+        assert scale.size(1) in (1, stack_size)
 
         if (scale == scale[:, [0]]).all() and not force_repeat:
             return scale[:, 0]
-        elif len(shape_n_stacks) > 1:
+
+        # A single scale per expert is shared by all of its stacked shards.
+        scale = scale.expand(-1, stack_size)
+        if len(shape_n_stacks) > 1:
             shape_n_repeats = torch.tensor(shape_n_stacks)
             target_group_size = target_group_size or sum(shape_k_stacks)
             shape_k_repeats = torch.tensor([sum(shape_k_stacks) // target_group_size])

humming/schema/compressed_tensors.py

@@ -221,7 +221,7 @@ def _convert_humming(
             weight_scale = weight_scale.view(torch.float8_e8m0fnu)
         elif self.format.startswith("nvfp4"):
             weight_scale = weight_scale.view(torch.float8_e4m3fn)
-        else:
+        elif self.strategy != "tensor":
             weight_scale = weight_scale.to(param_dtype)
 
         if self.strategy in ["group", "tensor_group"]:
@@ -232,21 +232,14 @@ def _convert_humming(
             weight_scale = weight_scale.repeat_interleave(self.block_structure[0], -2)
             group_size = self.block_structure[1]
         elif self.strategy == "tensor":
-            weight_scale = self._may_process_global_scale(
-                weight_scale,
-                shape_n_stacks=shape_n_stacks,
-                shape_k_stacks=shape_k_stacks,
-                num_experts=num_experts,
-                force_repeat=True,
-            ).to(param_dtype)
             group_size = 0
         else:
             assert self.strategy == "channel"
             group_size = 0
 
         output_tensors = {"weight": weight, "weight_scale": weight_scale}
 
-        weight_scale_type = None
+        weight_scale_type = WeightScaleType.TENSOR if self.strategy == "tensor" else None
         weight_scale_2_type = None
         if self.format == "nvfp4-pack-quantized":
             weight_scale_2 = tensors["weight_global_scale"]
@@ -263,17 +256,20 @@ def _convert_humming(
             if has_tensor_scale:
                 weight_scale_2_type = WeightScale2Type.TENSOR
                 output_tensors["weight_scale_2"] = weight_scale_2.float()
+            elif len(shape_n_stacks) > 1:
+                weight_scale_2_type = WeightScale2Type.CHANNEL
+                output_tensors["weight_scale_2"] = weight_scale_2.squeeze(-1).to(param_dtype)
             else:
                 weight_scale = weight_scale.float() * weight_scale_2.float()
                 output_tensors["weight_scale"] = weight_scale.to(param_dtype)
 
-        bs_dtype = None
+        bs_dtype = dtypes.float32 if self.strategy == "tensor" else None
         if self.format in ["int-quantized", "float-quantized", "naive-quantized"]:
             assert self.num_bits == 8
             b_dtype = dtypes.uint8 if self.type == "int" else dtypes.float8e4m3
         elif "nvfp4" in self.format:
             b_dtype = dtypes.float4e2m1
-            bs_dtype = dtypes.float8e4m3 if has_tensor_scale else None
+            bs_dtype = dtypes.float8e4m3 if weight_scale_2_type is not None else None
         elif "mxfp4" in self.format:
             b_dtype = dtypes.float4e2m1
             bs_dtype = dtypes.float8e8m0
@@ -298,6 +294,10 @@ def _convert_humming(
             has_zero_point=not self.symmetric,
         )
 
+        if self.strategy == "tensor":
+            return schema._convert_humming(
+                output_tensors, shape_n_stacks, shape_k_stacks, param_dtype, num_experts
+            )
         return schema, output_tensors
 
 

humming/schema/fp8.py

@@ -105,13 +105,7 @@ def _convert_humming(
         weight_scale = tensors[self.weight_scale_key]
 
         if self.weight_block_size is None:
-            weight_scale = self._may_process_global_scale(
-                weight_scale,
-                shape_n_stacks=shape_n_stacks,
-                shape_k_stacks=shape_k_stacks,
-                num_experts=num_experts,
-                force_repeat=True,
-            ).to(param_dtype)
+            schema = self.to_humming_schema(param_dtype)
         else:
             shape_n = sum(shape_n_stacks)
             scale_block_shape_n = self.weight_block_size[0]
@@ -124,6 +118,10 @@ def _convert_humming(
         if "bias" in tensors:
             output_tensors["bias"] = tensors["bias"]
 
+        if self.weight_block_size is None:
+            return schema._convert_humming(
+                output_tensors, shape_n_stacks, shape_k_stacks, param_dtype, num_experts
+            )
         return schema, output_tensors
 
 

humming/schema/humming.py

@@ -337,17 +337,24 @@ def _convert_humming(
                 shape_n_stacks=shape_n_stacks,
                 shape_k_stacks=shape_k_stacks,
                 num_experts=num_experts,
-                target_group_size=schema.weight_scale_group_size,
+                target_group_size=schema.weight_scale_group_size if len(shape_k_stacks) > 1 else None,
             )
 
             if tensor_scale.nelement() == (num_experts or 1):
                 tensors[scale_key] = tensor_scale.to(torch.float32)
                 if is_tensor_only:
                     schema.bs_dtype = dtypes.float32
             elif is_tensor_only:
-                schema.weight_scale_type = WeightScaleType.CHANNEL
+                if len(shape_k_stacks) > 1:
+                    schema.weight_scale_type = WeightScaleType.GROUP
+                    schema.weight_scale_group_size = sum(shape_k_stacks) // tensor_scale.size(-1)
+                else:
+                    schema.weight_scale_type = WeightScaleType.CHANNEL
                 schema.bs_dtype = dtypes.DataType.from_torch_dtype(param_dtype)
                 tensors["weight_scale"] = tensor_scale.to(param_dtype)
+            elif self.weight_scale_type == WeightScaleType.GROUP and len(shape_n_stacks) > 1:
+                schema.weight_scale_2_type = WeightScale2Type.CHANNEL
+                tensors["weight_scale_2"] = tensor_scale.squeeze(-1).to(param_dtype)
             else:
                 schema.weight_scale_2_type = WeightScale2Type.NONE
                 weight_scale = tensors["weight_scale"]

humming/schema/modelopt.py

@@ -120,6 +120,9 @@ def _convert_humming(
         if has_tensor_scale:
             weight_scale_2_type = WeightScale2Type.TENSOR
             output_tensors["weight_scale_2"] = weight_scale_2.view((num_experts or 1)).float()
+        elif len(shape_n_stacks) > 1:
+            weight_scale_2_type = WeightScale2Type.CHANNEL
+            output_tensors["weight_scale_2"] = weight_scale_2.squeeze(-1).to(param_dtype)
         else:
             weight_scale_2_type = WeightScale2Type.NONE
             weight_scale = weight_scale.float() * weight_scale_2.float()
@@ -130,7 +133,7 @@ def _convert_humming(
 
         schema = HummingWeightSchema(
             b_dtype=dtypes.float4e2m1,
-            bs_dtype=dtypes.float8e4m3 if has_tensor_scale else None,
+            bs_dtype=dtypes.float8e4m3 if weight_scale_2_type != WeightScale2Type.NONE else None,
             weight_scale_group_size=16,
             weight_scale_type=WeightScaleType.GROUP,
             weight_scale_2_type=weight_scale_2_type,

humming/schema/quark.py

@@ -450,13 +450,16 @@ def _convert_humming(
                     shape_n_stacks,
                     shape_k_stacks,
                     num_experts or None,
-                    target_group_size=schema.weight_scale_group_size or None,
+                    target_group_size=schema.weight_scale_group_size if len(shape_k_stacks) > 1 else None,
                 )
                 is_group_scale = schema.weight_scale_type == WeightScaleType.GROUP
                 keep_scale_2 = schema.bs_dtype is not None and is_group_scale
                 if keep_scale_2 and scale_2.numel() == (num_experts or 1):
                     schema.weight_scale_2_type = WeightScale2Type.TENSOR
                     result["weight_scale_2"] = scale_2.reshape(num_experts or 1).float()
+                elif keep_scale_2 and len(shape_n_stacks) > 1:
+                    schema.weight_scale_2_type = WeightScale2Type.CHANNEL
+                    result["weight_scale_2"] = scale_2.squeeze(-1).to(param_dtype)
                 else:
                     if scale_2.numel() == (num_experts or 1):
                         scale_2 = scale_2.reshape(-1, 1, 1) if num_experts else scale_2.reshape(())