[NPU] Support Kimi-K2-Thinking model

docker/npu_patch/megatron-bridge.patch

@@ -32,15 +32,21 @@ index a0421273..7f9203ab 100644
              # Check the actual parameter name to determine the correct parallelism type
              if self.megatron_param and (
                  self.megatron_param.endswith("layer_norm_weight") or self.megatron_param.endswith("layer_norm_bias")
-@@ -1207,7 +1211,7 @@ class AutoMapping(MegatronParamMapping[torch.Tensor]):
+@@ -1207,9 +1211,10 @@ class AutoMapping(MegatronParamMapping[torch.Tensor]):
              return "replicated"
 
          # Check parallel_mode for TELinear
 -        if module_type == "TELinear":
 +        if module_type == "TELinear" or module_type == "MindSpeedTELinear":
-             if module.parallel_mode == "column":
++            parallel_mode = getattr(module, "parallel_mode", None)
+-            if module.parallel_mode == "column":
++            if parallel_mode == "column":
                  return "column"
-             elif module.parallel_mode == "row":
+-            elif module.parallel_mode == "row":
++            elif parallel_mode == "row":
+                 return "row"
+             else:
+                 return "replicated"
 diff --git a/src/megatron/bridge/models/qwen/__init__.py b/src/megatron/bridge/models/qwen/__init__.py
 index b3656b6d..382845cc 100644
 --- a/src/megatron/bridge/models/qwen/__init__.py

vime/backends/megatron_utils/megatron_to_hf/processors/quantizer_compressed_tensors.py

@@ -230,14 +230,46 @@ def if_quant(name, patterns):
     return False
 
 
+def _fake_int4_quant_fallback(weight, group_shape, sym=True):
+    """Pure PyTorch fallback for fake_int4_quant_cuda on NPU."""
+    group_size = group_shape[1]
+    orig_shape = weight.shape
+    rows, cols = orig_shape[0], orig_shape[1]
+    n_groups = cols // group_size
+
+    w = weight.reshape(rows, n_groups, group_size).to(torch.float32)
+
+    if sym:
+        w_max = w.abs().amax(dim=-1, keepdim=True).clamp(min=1e-8)
+        scale = w_max / 7.0  # int4 symmetric range: [-8, 7], use 7 as scale
+        q = torch.round(w / scale).clamp(-8, 7)
+        zp = None
+    else:
+        w_min = w.amin(dim=-1, keepdim=True)
+        w_max = w.amax(dim=-1, keepdim=True)
+        scale = ((w_max - w_min) / 15.0).clamp(min=1e-8)  # int4 asymmetric: [0, 15]
+        zp = torch.round(-w_min / scale).clamp(0, 15)
+        q = torch.round(w / scale + zp).clamp(0, 15)
+
+    q = q.reshape(orig_shape).to(torch.float32)
+    scale = scale.reshape(rows, n_groups).contiguous()
+    if zp is not None:
+        zp = zp.reshape(rows, n_groups).contiguous()
+
+    return q, scale, zp
+
+
 def pack_layer(weight, group_size, sym=True):
-    w, scale, zp = fake_int4_quant_cuda.fake_int4_quant_cuda(weight, (1, group_size), sym)
+    if fake_int4_quant_cuda is not None:
+        w, scale, zp = fake_int4_quant_cuda.fake_int4_quant_cuda(weight, (1, group_size), sym)
+    else:
+        w, scale, zp = _fake_int4_quant_fallback(weight, (1, group_size), sym)
     w = w.view(weight.shape[0], 1, weight.shape[1] // group_size, group_size)
     scale = scale.view(weight.shape[0], 1, weight.shape[1] // group_size, 1)
-    zp = zp.view(weight.shape[0], 1, weight.shape[1] // group_size, 1)
     if sym:
         w = w * scale
     else:
+        zp = zp.view(weight.shape[0], 1, weight.shape[1] // group_size, 1)
         w = (w - zp) * scale
     w = w.view(weight.shape)
     scale = scale.view(weight.shape[0], -1).contiguous()
@@ -283,7 +315,7 @@ def quantize_params_compressed_tensors(converted_named_params, quantization_conf
         qw, s, zp = pack_layer(param, group_size, is_symmetric)
         qweight_name = name.replace(".weight", ".weight_packed")
         scale_name = name.replace(".weight", ".weight_scale")
-        weight_shape = torch.tensor(param.shape, dtype=torch.int32, device="cuda")
+        weight_shape = torch.tensor(param.shape, dtype=torch.int32, device=param.device)
         weight_shape_name = name.replace(".weight", ".weight_shape")
         if zp is not None:
             zp_name = name.replace(".weight", ".weight_zero_point")

vime/backends/megatron_utils/model_provider.py

@@ -17,7 +17,7 @@
 from megatron.core.transformer.transformer_config import TransformerConfig
 from megatron.training.arguments import core_transformer_config_from_args
 
-from vime.utils.megatron_bridge_utils import patch_auto_bridge_hf_config
+from vime.utils.megatron_bridge_utils import patch_auto_bridge_hf_config, patch_auto_bridge_hf_config_for_model
 from vime.utils.misc import load_function
 
 
@@ -87,6 +87,7 @@ def wrapped_model_provider(
         import vime_plugins.megatron_bridge  # noqa: F401  # register custom bridges
 
         bridge = patch_auto_bridge_hf_config(AutoBridge.from_hf_pretrained(args.hf_checkpoint, trust_remote_code=True))
+        bridge = patch_auto_bridge_hf_config_for_model(bridge)
         provider = bridge.to_megatron_provider(load_weights=False)
         # TODO: we should not manually set this...
         provider.tensor_model_parallel_size = args.tensor_model_parallel_size

vime/backends/megatron_utils/update_weight/hf_weight_iterator_bridge.py

@@ -1,6 +1,6 @@
 import dataclasses
 
-
+import torch
 from vime.utils import megatron_bridge_utils
 from vime.utils.misc import chunk_named_params_by_size
 
@@ -30,7 +30,10 @@ def _patched(self, task, converted_weights_dict):
         cpu_dict = {k: v.cpu() for k, v in converted_weights_dict.items()}
         result = _orig(self, task, cpu_dict)
         # Move merged result back to GPU for CUDA IPC serialization
-        return {k: v.cuda() for k, v in result.items()} if result else result
+        if getattr(torch, "npu", None) and torch.npu.is_available():
+            return {k: v.npu() for k, v in result.items()} if result else result
+        else:
+            return {k: v.cuda() for k, v in result.items()} if result else result
 
     GPTOSSBridge.maybe_modify_converted_hf_weight = _patched
     GPTOSSBridge._cpu_cache_patched = True
@@ -49,6 +52,25 @@ def __init__(self, *args, **kwargs):
         )
         _patch_bridge_expert_cache_to_cpu()
 
+        # Patch megatron-bridge to handle None parallelism_type
+        try:
+            from megatron.bridge.models.conversion.param_mapping import AutoMapping
+
+            _orig_megatron_to_hf = AutoMapping.megatron_to_hf
+
+            def _patched_megatron_to_hf(self, megatron_weight, megatron_module):
+                try:
+                    return _orig_megatron_to_hf(self, megatron_weight, megatron_module)
+                except ValueError as e:
+                    if "Unknown parallelism type: None" in str(e):
+                        hf_param = getattr(self, "hf_param", None)
+                        return {hf_param: megatron_weight}
+                    raise
+
+            AutoMapping.megatron_to_hf = _patched_megatron_to_hf
+        except (ImportError, AttributeError):
+            pass
+
     def get_hf_weight_chunks(self, megatron_local_weights, progress_desc: str = "Update weights"):
         # TODO support quantization (e.g. modify megatron-bridge to provide megatron param name)
         renamed_megatron_local_weights = {strip_param_name_prefix(k): v for k, v in megatron_local_weights.items()}
@@ -60,12 +82,16 @@ def get_hf_weight_chunks(self, megatron_local_weights, progress_desc: str = "Upd
 
             def _streaming_quantized():
                 for hf_param_name, weight, megatron_param_name in named_weights:
+                    if weight is None:
+                        continue
                     processed_weight = postprocess_hf_param(
                         args=self.args,
                         megatron_param_name=megatron_param_name,
                         hf_param_name=hf_param_name,
                         param=weight,
                     )
+                    if processed_weight is None:
+                        continue
                     converted_named_params = [(hf_param_name, processed_weight)]
                     quantized_batch = quantize_params(
                         args=self.args,
@@ -93,7 +119,10 @@ def _handle_one(task):
         ), f"{weight_dict_key=} not in new_weight_dict ({task.vp_stage=}, {task.param_name=}, {list(new_weight_dict)=})"
 
         new_param_weight = new_weight_dict[weight_dict_key]
-        new_param_weight = new_param_weight.cuda()
+        if getattr(torch, "npu", None) and torch.npu.is_available():
+            new_param_weight = new_param_weight.npu()
+        else:
+            new_param_weight = new_param_weight.cuda()
         return dataclasses.replace(task, param_weight=new_param_weight)
 
     return _MapWithLen(_handle_one, vanilla_conversion_tasks)

vime/utils/megatron_bridge_utils.py

@@ -38,6 +38,26 @@ def patch_auto_bridge_hf_config(bridge):
     return bridge
 
 
+def patch_auto_bridge_hf_config_for_model(bridge):
+    if bridge is None:
+        return bridge
+
+    hf_pretrained = getattr(bridge, "hf_pretrained", None)
+    config = hf_pretrained.config if hasattr(hf_pretrained, "config") else hf_pretrained
+
+    # Kimi K2 model
+    from megatron.bridge.models.kimi.kimi_bridge import KimiK2Bridge
+
+    if (
+        getattr(config, "model_type", "") == "kimi_k2"
+        and "KimiK2ForCausalLM" not in getattr(config, "architectures", [])
+        and not isinstance(bridge._model_bridge, KimiK2Bridge)
+    ):
+        bridge.__dict__["_causal_lm_architecture"] = "KimiK2ForCausalLM"
+
+    return bridge
+
+
 @contextmanager
 def patch_megatron_model(model):
     unwrapped_model = unwrap_model(model)[0]