Support fp4 blockwise load

src/mcore_bridge/bridge/gpt_bridge.py

@@ -198,23 +198,26 @@ def _set_weight(
             tensor = tensor + offset
         tensor_list = tensor.chunk(len(mg_param), dim=0)
         for i, param in enumerate(mg_param):
-            tensor = tensor_list[i].reshape(*param.shape)
-            if self._is_fp8_param(param):
-                if hf_scale_inv is None:
-                    param.data.copy_(tensor)
-                    param._high_precision_init_val.copy_(tensor)
-                else:
-                    tensor = tensor.view(torch.uint8)
-                    param._rowwise_data.data.copy_(tensor)
-                    self._copy_scale_inv(param, hf_scale_inv[i])
-                    del param.get_high_precision_init_val
-            else:
-                if hf_scale_inv is not None:
-                    fp8_tensor = self.fp8_quantizer.make_empty(tensor.shape)
-                    fp8_tensor._rowwise_data.copy_(tensor.view(torch.uint8))
-                    self._copy_scale_inv(fp8_tensor, hf_scale_inv[i])
-                    tensor = fp8_tensor
+            self._set_param(param, tensor_list[i], None if hf_scale_inv is None else hf_scale_inv[i])
+
+    def _set_param(self, param, tensor, hf_scale_inv):
+        tensor = tensor.reshape(*param.shape)
+        if self._is_fp8_param(param):
+            if hf_scale_inv is None:
                 param.data.copy_(tensor)
+                param._high_precision_init_val.copy_(tensor)
+            else:
+                tensor = tensor.view(torch.uint8)
+                param._rowwise_data.data.copy_(tensor)
+                self._copy_scale_inv(param, hf_scale_inv)
+                del param.get_high_precision_init_val
+        else:
+            if hf_scale_inv is not None:
+                fp8_tensor = self.fp8_quantizer.make_empty(tensor.shape)
+                fp8_tensor._rowwise_data.copy_(tensor.view(torch.uint8))
+                self._copy_scale_inv(fp8_tensor, hf_scale_inv)
+                tensor = fp8_tensor
+            param.data.copy_(tensor)
 
     @staticmethod
     def _copy_scale_inv(tensor, scale_inv):

src/mcore_bridge/model/gpts/deepseek_v4.py

@@ -7,6 +7,7 @@
 from typing import Optional
 
 from mcore_bridge.bridge import GPTBridge
+from mcore_bridge.utils import Fp8Dequantizer, fp4_to_fp8
 
 from ..constant import ModelType
 from ..gpt_model import GPTModel
@@ -353,8 +354,20 @@ def _convert_hf_state_dict(self, hf_state_dict, to_mcore):
         res = super()._convert_hf_state_dict(hf_state_dict, to_mcore)
         if to_mcore:
             res = self._add_prefix(res, 'model.')
+            new_res = {}
+            for k, v in res.items():
+                if k.endswith('.scale'):
+                    k = k[:-len('.scale')] + '.weight_scale_inv'
+                new_res[k] = v
+            res = new_res
         elif not to_mcore:
             res = self._remove_prefix(res, 'model.')
+            new_res = {}
+            for k, v in res.items():
+                if k.endswith('.weight_scale_inv'):
+                    k = k[:-len('.weight_scale_inv')] + '.scale'
+                new_res[k] = v
+            res = new_res
         return res
 
     def _set_moe_state(
@@ -452,6 +465,18 @@ def _convert_mtp_embeds(self, lm_model, hf_state_dict, to_mcore):
             if self.config.untie_embeddings_and_output_weights:
                 self._set_state_dict(lm_model, 'output_layer.weight', hf_state_dict, 'head.weight', to_mcore)
 
+    def _set_param(self, param, tensor, scale_inv):
+        is_fp4 = tensor.dtype == torch.int8 and tensor.shape[-1] * 2 == param.shape[-1]
+        if not is_fp4:
+            return super()._set_param(param, tensor, scale_inv)
+        tensor = fp4_to_fp8(tensor)
+        tensor = tensor.reshape(*param.shape)
+        scale_inv = scale_inv.reshape(-1, scale_inv.shape[-1])
+        tensor = Fp8Dequantizer().convert(tensor, scale_inv)
+        if self._is_fp8_param(param):
+            param._high_precision_init_val.copy_(tensor)
+        param.data.copy_(tensor)
+
 
 register_model(
     ModelMeta(

src/mcore_bridge/utils/__init__.py

@@ -1,5 +1,5 @@
 # Copyright (c) ModelScope Contributors. All rights reserved.
-from .dequantizer import Fp8Dequantizer, MxFp4Dequantizer
+from .dequantizer import Fp8Dequantizer, MxFp4Dequantizer, fp4_to_fp8
 from .env import get_dist_setting, get_node_setting, is_dist, is_last_rank, is_local_master, is_master
 from .import_utils import _LazyModule, is_flash_attn_3_available
 from .logger import get_logger

src/mcore_bridge/utils/dequantizer.py

@@ -5,7 +5,8 @@
 
 class Fp8Dequantizer:
 
-    def __init__(self, block_size: Tuple[int, int] = (128, 128)):
+    def __init__(self, block_size: Tuple[int, int] = (None, None)):
+        # Set to None to enable automatic selection.
         self.block_size = block_size
 
     def convert(
@@ -19,8 +20,13 @@ def convert(
             quantized = quantized.view(torch.float8_e4m3fn)
         quantized_fp32 = quantized.to(torch.float32)
         rows, cols = quantized_fp32.shape[-2:]
+        scale_rows, scale_cols = scales.shape[-2:]
         block_size = self.block_size
         block_m, block_n = block_size
+        if block_m is None:
+            block_m = rows // scale_rows
+        if block_n is None:
+            block_n = cols // scale_cols
         needs_padding = rows % block_m != 0 or cols % block_n != 0
 
         input_tensor = quantized_fp32
@@ -53,3 +59,21 @@ def convert(
     ) -> torch.Tensor:
         from transformers.integrations import convert_moe_packed_tensors
         return convert_moe_packed_tensors(blocks, scales)
+
+
+_FP4_E2M1_LUT = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0]
+
+
+def fp4_to_fp8(packed: torch.Tensor) -> torch.Tensor:
+    lut = torch.tensor(_FP4_E2M1_LUT, dtype=torch.float32, device=packed.device)
+    u8 = packed.contiguous().view(torch.uint8)
+    low = (u8 & 0x0F).long()
+    high = ((u8 >> 4) & 0x0F).long()
+
+    low_f32 = lut[low]
+    high_f32 = lut[high]
+
+    unpacked = torch.stack([low_f32, high_f32], dim=-1)
+    unpacked = unpacked.reshape(*packed.shape[:-1], 2 * packed.shape[-1])
+
+    return unpacked.to(torch.float8_e4m3fn)