feat(CRITICAL): 从 GitHub 扫描搬运 ixformer SDK + xllm 完整 GDN/MoE 代码

来源:
  1. Chranos/ixformer (GitHub) → ixformer_sdk/ (230 files, 70K lines)
     - inference/functions/vllm.py: vllm_moe_topk_softmax 完整实现 (2033 lines)
     - inference/functions/moe.py: MoE ops 完整实现 (1380 lines)
     - contrib/vllm_flash_attn/: FA2 Python 接口 (1018 lines)
     - contrib/tgi/fused_moe.py: TGI fused MoE (429 lines)
     - csrc/include/ixformer/: C++ kernel headers + cmake

  2. Deep-Spark/xllm (GitHub) → upstream_ref/xllm_latest/ (+15 files)
     - npu_torch/qwen3_5_decoder_layer_impl.cpp/.h
     - npu_torch/qwen3_5_gated_delta_net.cpp/.h
     - npu_torch/qwen3_next_*.cpp/.h (6 files)
     - npu_torch/attention.cpp/.h + fused_moe.cpp/.h + CMakeLists.txt
     - models/llm/qwen3_5.h + qwen3_5_mtp.h + qwen3_next.h
     - models/vlm/qwen3_5.h

调用链完整性:
  ixformer_sdk/inference/functions/vllm.py
    → ops.infer.moe_topk_softmax() (C++ 层)
    → 这就是 base 镜像 libixformer.so 里的实现

  upstream_ref/xllm_latest/core/layers/ilu/fused_moe.cpp
    → ixformer::infer::topk_softmax() (直接 C++ 调用)
    → ixformer::infer::group_gemm() → 完整 7-step MoE pipeline

ixformer_sdk/contrib/vllm/layers/__init__.py

@@ -0,0 +1,30 @@
+from .llama import forward_smoothquant
+from .mixtral import mixtral_decoder_layer_forward
+
+SUPPORT_REPLACE_METHOD = {
+    "llama": forward_smoothquant,
+}
+
+SUPPORT_REPLACE_LAYER = {
+    "llama": None,
+}
+
+
+def get_replace_forward(name: str):
+    try:
+        method = SUPPORT_REPLACE_METHOD[name]
+    except:
+        raise ValueError(
+            f"Only support replace names: {SUPPORT_REPLACE_METHOD.keys()}, but got {name}"
+        )
+    return method
+
+
+def get_replace_layer(name: str):
+    try:
+        layer = SUPPORT_REPLACE_LAYER[name]
+    except:
+        raise ValueError(
+            f"Only support replace names: {SUPPORT_REPLACE_LAYER.keys()}, but got {name}"
+        )
+    return layer

ixformer_sdk/contrib/vllm/layers/llama.py

@@ -0,0 +1,100 @@
+from typing import Any, Dict, Iterable, List, Optional, Tuple
+
+import torch
+from transformers import LlamaConfig
+
+from vllm.attention import AttentionMetadata
+from vllm.config import CacheConfig
+from vllm.distributed import tensor_model_parallel_all_reduce
+from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+from vllm.model_executor.models.llama import LlamaDecoderLayer as VllmLlamaDecoderLayer
+
+import vllm._custom_ops as ops
+# from ..overlap_comm import DecoderLayerOverlapComm, get_overlap_linear_method
+
+
+# This method is needed for support smoothquant no overlap forward
+def forward_smoothquant(
+    input_ids: Optional[torch.Tensor],
+    positions: torch.Tensor,
+    kv_caches: List[torch.Tensor],
+    attn_metadata: AttentionMetadata,
+    inputs_embeds: Optional[torch.Tensor] = None,
+    self = None, # will be set by partial
+) -> torch.Tensor:
+    dtype = self.dtype
+    
+    def forward_smoothquant_mlp(self,x,scales):
+        # gate_up_proj
+        # Int8 Matrix multiply.
+        bias = self.gate_up_proj.bias if not self.gate_up_proj.skip_bias_add else None
+        gate_up = ops.w8a8(x, self.gate_up_proj.weight, scales, self.gate_up_proj.weight_scales, dtype)
+        if bias:
+            gate_up += bias
+        
+        # act_fun
+        x, scales = ops.silu_and_mul_smoothquant(gate_up, self.down_proj.smooth_scales)
+        
+        # down_proj
+        output_parallel = ops.w8a8(x, self.down_proj.weight, scales, self.down_proj.weight_scales, dtype)
+        if self.down_proj.reduce_results and self.down_proj.tp_size > 1:
+            output = tensor_model_parallel_all_reduce(output_parallel)
+        else:
+            output = output_parallel
+
+        if not self.down_proj.skip_bias_add:
+            output = output + self.down_proj.bias if self.down_proj.bias is not None else output
+
+        return output
+    
+    def forward_smoothquant_attn(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        kv_cache: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+        scales: torch.Tensor,
+    ) -> torch.Tensor:
+        # qkv proj
+        bias = self.qkv_proj.bias if not self.qkv_proj.skip_bias_add else None
+        
+        qkv = ops.w8a8(hidden_states, self.qkv_proj.weight, scales, self.qkv_proj.weight_scales, dtype)
+        if bias:
+            qkv += bias
+        
+        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+        q, k = self.rotary_emb(positions, q, k)
+        attn_output = self.attn(q, k, v, kv_cache, attn_metadata)
+        output, _ = self.o_proj(attn_output)  # TODO
+        return output
+    
+    if inputs_embeds is not None:
+        hidden_states = inputs_embeds
+    else:
+        hidden_states = self.get_input_embeddings(input_ids)
+    residual = None
+    for i in range(len(self.layers)):
+        layer = self.layers[i]
+        if residual is None:
+            residual = hidden_states
+            hidden_states, scales = ops.rms_norm_smoothquant(hidden_states,layer.input_layernorm.weight,layer.input_layernorm.variance_epsilon, layer.self_attn.qkv_proj.smooth_scales)
+        else:
+            hidden_states, residual, scales = ops.fused_add_rms_norm_smoothquant(hidden_states, residual, layer.input_layernorm.weight, layer.input_layernorm.variance_epsilon, layer.self_attn.qkv_proj.smooth_scales)
+        
+        hidden_states = forward_smoothquant_attn(
+            layer.self_attn,
+            positions=positions,
+            hidden_states=hidden_states,
+            kv_cache=kv_caches[i],
+            attn_metadata=attn_metadata,
+            scales=scales,
+            )
+
+        # Fully Connected
+        hidden_states, residual, scales = ops.fused_add_rms_norm_smoothquant(hidden_states, residual, layer.post_attention_layernorm.weight, layer.post_attention_layernorm.variance_epsilon, layer.mlp.gate_up_proj.smooth_scales)
+        
+        hidden_states = forward_smoothquant_mlp(layer.mlp, hidden_states, scales)
+
+    hidden_states, _ = self.norm(hidden_states, residual)
+    return hidden_states
+

ixformer_sdk/contrib/vllm/layers/mixtral.py

@@ -0,0 +1,331 @@
+import functools
+from typing import Dict, Optional, Tuple
+
+import ixformer.inference.functions as ixf
+import torch
+
+
+def mixtral_decoder_layer_forward(
+    self,
+    positions: torch.Tensor,
+    hidden_states: torch.Tensor,
+    kv_cache: torch.Tensor,
+    attn_metadata,
+    residual: Optional[torch.Tensor],
+) -> torch.Tensor:
+    if self.use_int_w8a8:
+        return w8a8_forward(
+            self, positions, hidden_states, kv_cache, attn_metadata, residual
+        )
+    else:
+        return original_forward(
+            self, positions, hidden_states, kv_cache, attn_metadata, residual
+        )
+
+
+def original_forward(
+    self,
+    positions: torch.Tensor,
+    hidden_states: torch.Tensor,
+    kv_cache: torch.Tensor,
+    attn_metadata,
+    residual: Optional[torch.Tensor],
+) -> torch.Tensor:
+    # Self Attention
+    if residual is None:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+    else:
+        hidden_states, residual = self.input_layernorm(hidden_states, residual)
+    hidden_states = self.self_attn(
+        positions=positions,
+        hidden_states=hidden_states,
+        kv_cache=kv_cache,
+        attn_metadata=attn_metadata,
+    )
+
+    # Fully Connected
+    hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
+    hidden_states = self.block_sparse_moe(hidden_states)
+    return hidden_states, residual
+
+
+def dynamic_scaled_int8_quant(x):
+    m, k = x.shape
+    i8_x = x.new_empty([m, k], dtype=torch.int8, device="cuda")
+    i8_scales = torch.empty([m], dtype=torch.float32, device="cuda")
+    ixf.dynamic_scaled_int8_quant(i8_x, x, i8_scales)
+    return i8_x, i8_scales
+
+
+def dynamic_w8a8(x, i8_weight, weight_scale):
+    i8_x, i8_scale = dynamic_scaled_int8_quant(x)
+    m, k = x.shape
+    k, n = i8_weight.shape
+    output = x.new_empty([m, n], dtype=x.dtype, device="cuda")
+    ixf.w8a8(
+        i8_x,
+        i8_weight.transpose(0, 1),
+        i8_scale,
+        weight_scale,
+        output=output,
+        out_dtype=x.dtype,
+    )
+    return output
+
+
+def fused_rms_norm_quant_linear(
+    self,
+    hidden_states,
+    ln_weight,
+    eps,
+    linear_weight,
+    linear_weight_scale,
+    residual=None,
+):
+    # lower rouge
+    # if residual is None:
+    #     residual = hidden_states
+    #     i8_hidden_states, _, i8_scales = ixf.residual_rms_norm_dynamic_int8(
+    #         input=hidden_states,
+    #         weight=ln_weight,
+    #         residual=None,
+    #         eps=eps,
+    #     )
+    # else:
+    #     i8_hidden_states, residual, i8_scales = ixf.residual_rms_norm_dynamic_int8(
+    #         input=hidden_states,
+    #         weight=ln_weight,
+    #         residual=residual,
+    #         eps=eps,
+    #     )
+
+    if residual is None:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+    else:
+        hidden_states, residual = self.input_layernorm(hidden_states, residual)
+    i8_hidden_states, i8_scales = dynamic_scaled_int8_quant(hidden_states)
+
+    qkv = hidden_states.new_empty(hidden_states.shape[0], linear_weight.shape[1])
+    ixf.w8a8(
+        i8_hidden_states,
+        linear_weight.transpose(0, 1),
+        i8_scales,
+        linear_weight_scale,
+        output=qkv,
+        out_dtype=hidden_states.dtype,
+    )
+    return qkv, residual
+
+
+def attention(qkv, positions, kv_cache, attn_metadata, self_attn):
+    q, k, v = qkv.split(
+        [self_attn.q_size, self_attn.kv_size, self_attn.kv_size], dim=-1
+    )
+    q, k = self_attn.rotary_emb(positions, q, k)
+    attn_output = self_attn.attn(q, k, v, kv_cache, attn_metadata)
+    return attn_output
+
+
+def fused_rms_norm_attention(
+    self,
+    hidden_states,
+    ln_weight,
+    eps,
+    positions,
+    kv_cache,
+    attn_metadata,
+    self_attn,
+    residual=None,
+):
+    hidden_states, residual = fused_rms_norm_quant_linear(
+        self,
+        hidden_states,
+        ln_weight,
+        eps,
+        self_attn.qkv_proj.weight,
+        self_attn.qkv_proj.weight_scale,
+        residual,
+    )
+    hidden_states = attention(
+        hidden_states, positions, kv_cache, attn_metadata, self_attn
+    )
+
+    hidden_states = dynamic_w8a8(
+        hidden_states, self_attn.o_proj.weight, self_attn.o_proj.weight_scale
+    )
+    # hidden_states,_ = self_attn.o_proj(hidden_states) # quant+linear+allreduce
+    return hidden_states, residual
+
+
+def w8a8_forward(
+    self,
+    positions: torch.Tensor,
+    hidden_states: torch.Tensor,
+    kv_cache: torch.Tensor,
+    attn_metadata,
+    residual: Optional[torch.Tensor],
+) -> torch.Tensor:
+
+    # qkv,_ = self.self_attn.qkv_proj(hidden_states)
+    hidden_states, residual = fused_rms_norm_attention(
+        self,
+        hidden_states,
+        self.input_layernorm.weight,
+        self.input_layernorm.variance_epsilon,
+        positions,
+        kv_cache,
+        attn_metadata,
+        self.self_attn,
+        residual,
+    )
+
+    # allreduce
+    tp_size = self.block_sparse_moe.experts.tp_size
+    if tp_size > 1:
+        from vllm.distributed import tensor_model_parallel_all_reduce
+
+        hidden_states = tensor_model_parallel_all_reduce(hidden_states)
+
+    # rms norm
+    hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
+    # moe
+    hidden_states = fused_moe(
+        hidden_states,
+        self.block_sparse_moe.gate.weight,
+        top_k=self.block_sparse_moe.experts.top_k,
+        w1=self.block_sparse_moe.experts.w13_weight,
+        w2=self.block_sparse_moe.experts.w2_weight,
+        w1_scale=self.block_sparse_moe.experts.w13_weight_scale,
+        w2_scale=self.block_sparse_moe.experts.w2_weight_scale,
+    )
+
+    # allreduce
+    if tp_size > 1:
+        from vllm.distributed import tensor_model_parallel_all_reduce
+
+        hidden_states = tensor_model_parallel_all_reduce(hidden_states)
+
+    return hidden_states, residual
+
+
+def fused_experts(hidden_states, router_logits, top_k, w1, w2, w1_scale, w2_scale):
+
+    """
+    Args:
+        hidden_states:          (num_tokens, k)                 dtype
+        router_logits:          (num_tokens, num_experts)       torch.float32
+        top_k                                                   int
+        w1:                     (num_experts, 2n, k)            torch.int8
+        w2:                     (num_experts, k, n)             torch.int8
+        w1_scale:               (num_experts, 2n)               torch.float32
+        w2_scale:               (num_experts, k)                torch.float32
+    Returns
+        final_hidden_states:    (num_tokens, k)                 dtype
+    """
+
+    # topk_weight: (num_tokens, top_k) torch.float32
+    # topk_ids:    (num_tokens, top_k) torch.int32
+    topk_weight, topk_ids = ixf.moe_topk_softmax(
+        gating_output=router_logits,
+        topk=top_k,
+        renormalize=True,
+    )
+
+    dtype = hidden_states.dtype
+    num_tokens, num_experts = router_logits.shape
+    expand_tokens = num_tokens * top_k
+
+    (
+        src_to_dst,
+        sorted_token_ids,
+        expert_sizes_gpu,
+        expert_sizes_cpu,
+    ) = ixf.moe_compute_token_index(
+        topk_ids=topk_ids,
+        num_experts=num_experts,
+    )
+    expert_sizes_cpu = expert_sizes_gpu.cpu()
+
+    # expand + reorder + quant
+    # i8_hidden_states: (expand_tokens, k) torch.int8
+    i8_hidden_states, a_scale = ixf.moe_expand_input_dynamic_scaled_int8(
+        hidden_states=hidden_states,
+        dst_to_src=sorted_token_ids,
+        dst_tokens=expand_tokens,
+        topk=top_k,
+        src_to_dst=src_to_dst,
+        topk_ids=None,  # use smooth quant
+        smooth_scales=None,  # use smooth quant
+    )
+
+    # w8a8 group gemm 1
+    # pt_output_1: (expand_tokens, 2n) dtype
+    pt_output_1 = ixf.moe_w8a8_group_gemm(
+        input=i8_hidden_states,
+        weight=w1,
+        i_scales=a_scale,
+        w_scales=w1_scale,
+        output_dtype=dtype,
+        tokens_per_experts=expert_sizes_cpu,
+        dst_to_src=None,
+        format="TN",
+    )
+
+    # act + quant
+    # pt_output_2: (expand_tokens, n) torch.int8
+    pt_output_2, a2_scale = ixf.activation_dynamic_scaled_int8(
+        input=pt_output_1,
+        bias=None,  # add gemm bias
+        smooth_scales=None,  # use smooth quant
+        dst_to_src=sorted_token_ids,
+        topk_ids=None,  # add gemm bias or use smooth quant
+        act_type="swiglu",
+    )
+
+    # w8a8 group gemm 2 + reorder
+    # pt_output_3: (expand_tokens, k) dtype
+    pt_output_3 = ixf.moe_w8a8_group_gemm(
+        input=pt_output_2,
+        weight=w2,
+        i_scales=a2_scale,
+        w_scales=w2_scale,
+        output_dtype=dtype,
+        tokens_per_experts=expert_sizes_cpu,
+        dst_to_src=sorted_token_ids,
+        format="TN",
+    )
+
+    # mul + reduce_sum
+    # final_hidden_states: (num_tokens, k)
+    final_hidden_states = ixf.moe_output_reduce_sum(
+        input=pt_output_3.view(num_tokens, top_k, -1),
+        topk_weight=topk_weight,
+    )
+
+    return final_hidden_states
+
+
+def fused_moe(hidden_states, gate_weight, top_k, w1, w2, w1_scale, w2_scale):
+    orig_shape = hidden_states.shape
+    hidden_size = hidden_states.shape[-1]
+
+    hidden_states = hidden_states.view(-1, hidden_size)
+
+    # router_logits: (num_tokens, n_experts)
+    # gate_weight: fp16
+    router_logits = ixf.linear(hidden_states, gate_weight)
+    router_logits = router_logits.to(torch.float32)
+
+    final_hidden_states = fused_experts(
+        hidden_states,
+        router_logits,
+        top_k,
+        w1,
+        w2,
+        w1_scale,
+        w2_scale,
+    )
+
+    return final_hidden_states.view(orig_shape)

ixformer_sdk/contrib/vllm/quantize/__init__.py

@@ -0,0 +1,14 @@
+from .smoothquant import smoothquant_prepare_quantize,smoothquant_export_quantized_weights
+from .w8a16 import w8a16_prepare_quantize,w8a16_export_quantized_weights
+
+SUPPORT_METHOD = {
+    "smoothquant": [smoothquant_prepare_quantize,smoothquant_export_quantized_weights],
+    "w8a16": [w8a16_prepare_quantize,w8a16_export_quantized_weights],
+}
+
+def get_quantize_method(method_name:str):
+    try:
+        method = SUPPORT_METHOD[method_name]
+    except:
+        raise ValueError(f"Only support quantization methods: {SUPPORT_METHOD.keys()}, but got {method_name}")
+    return method

ixformer_sdk/contrib/vllm/quantize/smoothquant.py

@@ -0,0 +1,407 @@
+import os
+
+import torch
+
+
+def smoothquant_prepare_quantize(self, quant_params={}):
+    model = self.model_runner.model
+
+    def update_act_scales(act_scales, x):
+        # 动态统计每次输入的最大值
+        hidden_dim = x.shape[-1]
+        x = x.view(-1, hidden_dim).abs().detach()
+        # [k]
+        comming_max = torch.max(x, dim=0, keepdim=True)[0].float()
+
+        if act_scales is None:
+            act_scales = comming_max
+        else:
+            act_scales = torch.max(act_scales, comming_max)
+        return act_scales
+
+    from functools import partial
+
+    from vllm.model_executor.layers.linear import (
+        ColumnParallelLinear,
+        MergedColumnParallelLinear,
+        QKVParallelLinear,
+        RowParallelLinear,
+    )
+
+    def new_forward(input_, m, raw_forward):
+        if not hasattr(m, "act_scales"):
+            m.act_scales = None
+        m.act_scales = update_act_scales(m.act_scales, input_)
+        return raw_forward(input_)
+
+    for name, m in model.named_modules():
+        if (
+            isinstance(m, QKVParallelLinear)
+            or isinstance(m, RowParallelLinear)
+            or isinstance(m, MergedColumnParallelLinear)
+            or isinstance(m, ColumnParallelLinear)
+        ):
+            m.forward = partial(new_forward, m=m, raw_forward=m.forward)
+
+def smoothquant_export_quantized_weights(self, save_path, quant_params={}):
+    gb_per_file = quant_params.get("filesize_limit", None)
+    smooth_alpha = quant_params.get("smooth_alpha", 0.5)
+    dynamic_quant_type = quant_params.get("dynamic_quant_type", "gpu")
+    assert dynamic_quant_type in ["gpu","cpu","kernel"]
+    if self.rank == 0:
+        print(f"set smooth_alpha={smooth_alpha}")
+        print(f"use quantize weight type: {dynamic_quant_type}")
+    
+    import ixformer._C as ops
+    def per_token_quant_8bit(weight):
+        # weight: [m,k]
+        dtype = weight.dtype
+        i8_weight = weight
+        scale = i8_weight.abs().max(dim=-1, keepdim=True)[0] / 127
+        i8_weight = i8_weight / scale.to(dtype)
+        i8_weight = torch.clamp(torch.round(i8_weight), -128, 127).to(torch.int8)
+        return i8_weight, scale.float()
+
+    def smooth_quant_weight_gpu_cpu(weight, act_scale, alpha=0.5, device="cpu"):
+        device = torch.device("cpu") if device == "cpu" else weight.device
+        ori_dtype = weight.dtype
+        # [1, k]
+        act_scale = act_scale.float().to(device).view(1, -1)
+        weight = weight.to(device)
+        # [1, k]
+        weight_scale = weight.abs().max(dim=0, keepdim=True)[0].float()
+        if alpha == -1:
+            smooth_scales = torch.ones_like(act_scale)
+        else:
+            smooth_scales = act_scale.pow(alpha) / weight_scale.pow(1 - alpha).clamp(
+                min=1e-5
+            )
+        weight = weight * smooth_scales.to(ori_dtype)
+        i8_weight, weight_scales = per_token_quant_8bit(weight)
+        # 为了可以使用 input * smooth_scales
+        if alpha == -1:
+            smooth_scales = torch.ones_like(act_scale)
+        else:
+            smooth_scales = weight_scale.pow(1 - alpha) / act_scale.pow(alpha).clamp(
+                min=1e-5
+            )
+        return i8_weight, weight_scales, smooth_scales.to(ori_dtype)
+    
+    def smooth_quant_weight_kernel(weight, act_scale, alpha=0.5):
+        output = torch.zeros_like(weight,dtype=torch.int8)
+        weight_scales = torch.zeros(weight.shape[:-1],dtype=torch.float, device=weight.device)
+        weight_max = torch.zeros(weight.shape[-1],dtype=torch.float, device=weight.device)
+        smooth_scales = torch.zeros(weight.shape[-1],dtype=weight.dtype, device=weight.device)
+        ops.infer.weight_quant_smoothquant(
+            weight, act_scale, alpha, output, weight_scales, smooth_scales, weight_max
+        )
+        return output, weight_scales.view(-1,1), smooth_scales.view(1,-1)
+    
+    def smooth_quant_weight(weight, act_scale, alpha=0.5):
+        if dynamic_quant_type == "kernel":
+            return smooth_quant_weight_kernel(weight,act_scale,alpha)
+        else:
+            return smooth_quant_weight_gpu_cpu(weight,act_scale,alpha,dynamic_quant_type)
+
+    model = self.model_runner.model
+
+    from vllm.distributed import (
+        tensor_model_parallel_all_gather,
+        tensor_model_parallel_all_reduce,
+        get_tensor_model_parallel_world_size
+    )
+    from vllm.model_executor.layers.linear import (
+        ColumnParallelLinear,
+        MergedColumnParallelLinear,
+        QKVParallelLinear,
+        RowParallelLinear,
+    )
+    from vllm.model_executor.layers.vocab_parallel_embedding import (
+        ParallelLMHead,
+        VocabParallelEmbedding,
+    )
+    from vllm.model_executor.models.falcon import FalconForCausalLM
+
+    for name, m in model.named_modules():
+        if isinstance(m, VocabParallelEmbedding):
+            # weight shape: [vocab_size // tp, embedding_dim]
+            weight_tensor = tensor_model_parallel_all_gather(m.weight, dim=0)
+            weight_tensor = weight_tensor[:m.org_vocab_size,:].contiguous()
+            if self.is_driver_worker:
+                m.weight = torch.nn.Parameter(weight_tensor.cpu(), requires_grad=False)
+                print(f"merged: {name}, shape={m.weight.shape}")
+            
+        elif isinstance(m, ParallelLMHead):
+            # weight shape: [vocab_size // tp, embedding_dim]
+            # bias shape: [vocab_size // tp]
+            if m.bias is not None:
+                bias = tensor_model_parallel_all_gather(m.bias, dim=0)
+                bias = bias[:m.org_vocab_size].contiguous()
+                if self.is_driver_worker:
+                    m.bias = torch.nn.Parameter(bias.cpu(), requires_grad=False)
+            
+            weight_tensor = tensor_model_parallel_all_gather(m.weight, dim=0)
+            weight_tensor = weight_tensor[:m.org_vocab_size,:].contiguous()
+            if self.is_driver_worker:
+                m.weight = torch.nn.Parameter(weight_tensor.cpu(), requires_grad=False)
+                print(f"merged: {name}, shape={m.weight.shape}")
+        
+        elif isinstance(m, QKVParallelLinear):
+            # weight shape: [total_num_head * head_size // tp + 2 * total_num_head * head_size // tp, hidden_size]
+            # bias shape: [total_num_head * head_size // tp + 2 * total_num_head * head_size // tp]
+            if self.parallel_config.world_size > 1:
+                total_q_hidden_size = m.total_num_heads * m.head_size
+                partial_q_hidden_size = m.num_heads * m.head_size
+                total_kv_hidden_size = m.total_num_kv_heads * m.head_size
+                partial_kv_hidden_size = m.num_kv_heads * m.head_size
+                
+                if m.bias is not None:
+                    # TODO do not support padding..
+                    bias_tenosr = m.bias.new_zeros(total_q_hidden_size + total_kv_hidden_size * 2, m.hidden_size)
+                    q_bias = bias_tenosr[:total_q_hidden_size][self.rank * partial_q_hidden_size : (self.rank + 1) * partial_q_hidden_size]
+                    k_bias = bias_tenosr[total_q_hidden_size:total_q_hidden_size+total_kv_hidden_size]\
+                                [self.rank * partial_kv_hidden_size : (self.rank + 1) * partial_kv_hidden_size]
+                    v_bias = bias_tenosr[total_q_hidden_size+total_kv_hidden_size:]\
+                                [self.rank * partial_kv_hidden_size : (self.rank + 1) * partial_kv_hidden_size]
+                    
+                    q_bias[:] = m.bias[:partial_q_hidden_size]
+                    k_bias[:] = m.bias[partial_q_hidden_size:partial_q_hidden_size+partial_kv_hidden_size]
+                    v_bias[:] = m.bias[partial_q_hidden_size+partial_kv_hidden_size:]
+                    
+                    bias_tensor = tensor_model_parallel_all_reduce(bias_tenosr)
+                    if self.is_driver_worker:
+                        m.bias = torch.nn.Parameter(bias_tensor.cpu(), requires_grad=False)
+
+                q_tensor = m.weight.new_zeros(m.total_num_heads * m.head_size, m.weight.shape[1])
+                k_tensor = m.weight.new_zeros(m.total_num_kv_heads * m.head_size, m.weight.shape[1])
+                v_tensor = m.weight.new_zeros(m.total_num_kv_heads * m.head_size, m.weight.shape[1])
+                
+                q_in_weight = m.weight[:-m.num_kv_heads * m.head_size * 2]
+                k_in_weight = m.weight[-m.num_kv_heads * m.head_size * 2:-m.num_kv_heads * m.head_size]
+                v_in_weight = m.weight[-m.num_kv_heads * m.head_size:]
+                
+                if getattr(m,"start_idx",None) is not None:
+                    start_idx = getattr(m,"start_idx")
+                    weight_end_idx = m.num_heads * m.head_size if not getattr(m,"is_padding") else (m.num_heads - 1) * m.head_size
+                    end_idx = start_idx + weight_end_idx
+                else:
+                    start_idx = self.rank * m.num_heads * m.head_size
+                    weight_end_idx = m.num_heads * m.head_size
+                    end_idx = start_idx + weight_end_idx
+                assert q_tensor[start_idx:end_idx,:].shape == q_in_weight[:weight_end_idx, :].shape
+                q_tensor[start_idx:end_idx,:] = q_in_weight[:weight_end_idx, :]
+                
+                if m.num_kv_head_replicas > 1:
+                    if self.rank % m.num_kv_head_replicas == 0:
+                        rank = self.rank // m.num_kv_head_replicas
+                        k_tensor[rank * m.num_kv_heads * m.head_size:(rank+1) * m.num_kv_heads * m.head_size] = k_in_weight
+                        v_tensor[rank * m.num_kv_heads * m.head_size:(rank+1) * m.num_kv_heads * m.head_size] = v_in_weight
+                else:
+                    k_tensor[self.rank * m.num_kv_heads * m.head_size:(self.rank+1) * m.num_kv_heads * m.head_size] = k_in_weight
+                    v_tensor[self.rank * m.num_kv_heads * m.head_size:(self.rank+1) * m.num_kv_heads * m.head_size] = v_in_weight
+                
+                q_tensor = tensor_model_parallel_all_reduce(q_tensor)
+                k_tensor = tensor_model_parallel_all_reduce(k_tensor)
+                v_tensor = tensor_model_parallel_all_reduce(v_tensor)
+                
+                if isinstance(model, FalconForCausalLM):
+                    num_query_heads_per_kv_head = (
+                        m.total_num_heads // m.total_num_kv_heads
+                    )
+                    q_tensor = q_tensor.view(
+                        m.total_num_kv_heads,
+                        num_query_heads_per_kv_head,
+                        m.head_size,
+                        -1,
+                    )
+                    k_tensor = k_tensor.view(m.total_num_kv_heads, 1, m.head_size, -1)
+                    v_tensor = v_tensor.view(m.total_num_kv_heads, 1, m.head_size, -1)
+                    weight_tensor = torch.cat(
+                        [q_tensor, k_tensor, v_tensor], dim=1
+                    ).view(-1, m.hidden_size)
+                else:
+                    weight_tensor = torch.cat([q_tensor, k_tensor, v_tensor])
+                assert (
+                    weight_tensor.shape[0]
+                    == total_q_hidden_size + total_kv_hidden_size * 2
+                )
+                assert weight_tensor.shape[1] == m.hidden_size
+
+            else:
+                weight_tensor = m.weight
+                if m.bias is not None and self.is_driver_worker:
+                    m.bias = torch.nn.Parameter(m.bias.cpu(), requires_grad=False)
+            
+            if self.is_driver_worker:
+                i8_weight, weight_scales, smooth_scales = smooth_quant_weight(
+                    weight_tensor, m.act_scales, smooth_alpha
+                )
+                m.weight = torch.nn.Parameter(i8_weight.cpu(), requires_grad=False)
+                m.weight_scales = torch.nn.Parameter(
+                    weight_scales.cpu(), requires_grad=False
+                )
+                m.smooth_scales = torch.nn.Parameter(
+                    smooth_scales.cpu(), requires_grad=False
+                )
+                print(f"Quantized: {name}")
+            
+        elif isinstance(m, MergedColumnParallelLinear):
+            if self.parallel_config.world_size > 1:
+                # weight shape: [intermediate_size // tp * 2, hidden_size]
+                # bias shape: [intermediate_size // tp * 2]
+                output_sizes = m.output_sizes
+                output_size = sum(output_sizes)
+                partial_output_sizes = [
+                    i // self.parallel_config.world_size for i in output_sizes
+                ]
+
+                if m.bias is not None:
+                    index_start = 0
+                    partial_index_start = 0
+                    bias_tenosr = m.bias.new_zeros(output_size)
+                    for i in range(len(output_sizes)):
+                        index_out = index_start + output_sizes[i]
+                        sub_bias_tensor = bias_tenosr[index_start:index_out]
+                        partial_size = partial_output_sizes[i]
+                        sub_bias_tensor[self.rank * partial_size:(self.rank+1) * partial_size] = m.bias[partial_index_start:partial_index_start+partial_size]
+                        
+                        index_start += output_sizes[i]
+                        partial_index_start += partial_size
+                    bias_tenosr = tensor_model_parallel_all_reduce(bias_tenosr)
+                    if self.is_driver_worker:
+                        m.bias = torch.nn.Parameter(bias_tenosr, requires_grad=False)
+                
+                weight_tensor = m.weight.new_zeros(output_size, m.input_size)
+
+                idx_out_start = 0
+                idx_partial_satrt = 0
+                for i in range(len(output_sizes)):
+                    idx_out_end = idx_out_start + output_sizes[i]
+                    sub_weight_tensor = weight_tensor[idx_out_start:idx_out_end]
+                    partial_size = partial_output_sizes[i]
+                    sub_weight_tensor[
+                        self.rank * partial_size : (self.rank + 1) * partial_size
+                    ] = m.weight[idx_partial_satrt : idx_partial_satrt + partial_size]
+
+                    idx_out_start += output_sizes[i]
+                    idx_partial_satrt += partial_size
+                weight_tensor = tensor_model_parallel_all_reduce(weight_tensor)
+            else:
+                weight_tensor = m.weight
+                if m.bias is not None and self.is_driver_worker:
+                    m.bias = torch.nn.Parameter(m.bias.cpu(), requires_grad=False)
+            
+            if self.is_driver_worker:
+                i8_weight, weight_scales, smooth_scales = smooth_quant_weight(
+                    weight_tensor, m.act_scales, smooth_alpha
+                )
+                m.weight = torch.nn.Parameter(i8_weight.cpu(), requires_grad=False)
+                m.weight_scales = torch.nn.Parameter(
+                    weight_scales.cpu(), requires_grad=False
+                )
+                m.smooth_scales = torch.nn.Parameter(
+                    smooth_scales.cpu(), requires_grad=False
+                )
+                print(f"Quantized: {name}")
+
+        elif isinstance(m, ColumnParallelLinear):
+            # weight shape: [some_dim // tp, hidden_size] // for this Linear, some_dim mostly is hidden_size * 4
+            # bias shape: [some_dim // tp]
+            if m.bias is not None:
+                bias_tenosr = tensor_model_parallel_all_gather(m.bias, dim=0)
+                if self.is_driver_worker:
+                    m.bias = torch.nn.Parameter(bias_tenosr.cpu(), requires_grad=False)
+
+            weight_tensor = tensor_model_parallel_all_gather(m.weight, dim=0)
+
+            if self.is_driver_worker:
+                i8_weight, weight_scales, smooth_scales = smooth_quant_weight(
+                   weight_tensor, m.act_scales, smooth_alpha
+                )
+                m.weight = torch.nn.Parameter(i8_weight.cpu(), requires_grad=False)
+                m.weight_scales = torch.nn.Parameter(
+                    weight_scales.cpu(), requires_grad=False
+                )
+                m.smooth_scales = torch.nn.Parameter(
+                    smooth_scales.cpu(), requires_grad=False
+                )
+                print(f"Quantized: {name}")
+
+        elif isinstance(m, RowParallelLinear):
+            # weight shape: [hidden_size, some_dim // tp] // for this Linear, some_dim mostly is hidden_size * 4 or intermediate_size
+            # bias shape: [hidden_size]
+            if m.bias is not None:
+                bias_tensor = tensor_model_parallel_all_gather(m.bias, dim=-1)
+                if self.is_driver_worker:
+                    m.bias = torch.nn.Parameter(m.bias.cpu(), requires_grad=False)
+            
+            if getattr(m,"start_idx", None) is not None:
+                start_idx = getattr(m,"start_idx")
+                end_idx = start_idx + (m.input_size_per_partition if not getattr(m,"is_padding") else (m.input_size_per_partition - m.padding_size))
+                weight_end_idx = m.input_size_per_partition if not getattr(m,"is_padding") else (m.input_size_per_partition - m.padding_size)
+            else:
+                start_idx = m.input_size_per_partition * self.rank
+                end_idx = start_idx + m.input_size_per_partition
+                weight_end_idx = m.input_size_per_partition
+            
+            act_scales = m.act_scales.new_zeros(m.input_size)
+            assert act_scales[start_idx:end_idx].shape == m.act_scales.view(-1)[:weight_end_idx].shape
+            act_scales[start_idx:end_idx] = m.act_scales.view(-1)[:weight_end_idx]
+            act_scales = tensor_model_parallel_all_reduce(act_scales)
+            m.act_scales = act_scales
+            
+            weight_tensor = m.weight.new_zeros(m.weight.shape[0],m.input_size)
+            assert  weight_tensor[:,start_idx:end_idx].shape == m.weight[:,:weight_end_idx].shape
+            weight_tensor[:,start_idx:end_idx] = m.weight[:,:weight_end_idx]
+            weight_tensor = tensor_model_parallel_all_reduce(weight_tensor)
+            
+            if self.is_driver_worker:
+                i8_weight, weight_scales, smooth_scales = smooth_quant_weight(
+                    weight_tensor, m.act_scales, smooth_alpha
+                )
+                smooth_scales = smooth_scales.view(1,-1)
+                m.weight = torch.nn.Parameter(i8_weight.cpu(), requires_grad=False)
+                m.weight_scales = torch.nn.Parameter(
+                    weight_scales.cpu(), requires_grad=False
+                )
+                m.smooth_scales = torch.nn.Parameter(
+                    smooth_scales.cpu(), requires_grad=False
+                )
+                print(f"Quantized: {name}")
+        else:
+            pass
+        
+        torch.cuda.empty_cache()
+
+    # save weights
+    if self.is_driver_worker:
+        from safetensors.torch import save_file
+        
+        tensors = {}
+        saved = False
+        count = 0
+        size_in_bytes = 0
+        
+        tensors = {}
+        for name, weight in model.named_parameters():
+            if "act_scales" in name:
+                continue
+            # skip lm_head_weight if needed..
+            if "lm_head" in name and model.config.tie_word_embeddings:
+                continue
+            tensors[name] = weight
+            
+            saved = False
+            if gb_per_file is not None and size_in_bytes >= gb_per_file * 1024 * 1024 * 1024:
+                weight_path = os.path.join(save_path, "model_{}.safetensors".format(str(count).zfill(6)))
+                save_file(tensors, weight_path)
+                print(f"The quantified weights were successfully saved in {weight_path}.")
+                tensors.clear()
+                saved = True
+                count += 1
+                size_in_bytes = 0
+        
+        if not saved: 
+            weight_path = os.path.join(save_path, "model_{}.safetensors".format(str(count).zfill(6)))   
+            save_file(tensors, weight_path)
+            print(f"The quantified weights were successfully saved in {weight_path}.")

ixformer_sdk/contrib/vllm/quantize/w8a16.py

@@ -0,0 +1,233 @@
+import os
+
+import torch
+
+def w8a16_prepare_quantize(self, quant_params={}):
+    # We need do nothing in here
+    pass
+
+
+def w8a16_export_quantized_weights(self, save_path, quant_params={}):
+    gb_per_file = quant_params.get("filesize_limit", None)
+    int8_min = -127
+
+    def w8a16_quantization(weight):
+        # all weights should be [output,input], otherwise, we may get an wrong weight and scale...
+        scale = torch.abs(weight).max(dim=-1)[0] / 127.0
+        int8_weight = torch.clamp(weight / scale.view(-1,1),min=int8_min,max=127).to(torch.int8).contiguous()
+        scale = scale.view(1,-1).contiguous()
+        return int8_weight, scale
+
+
+    model = self.model_runner.model
+
+    from vllm.distributed.communication_op import (
+        tensor_model_parallel_all_gather,
+        tensor_model_parallel_all_reduce,
+    )
+    from vllm.model_executor.layers.linear import (
+        ColumnParallelLinear,
+        MergedColumnParallelLinear,
+        QKVParallelLinear,
+        RowParallelLinear,
+    )
+    from vllm.model_executor.layers.vocab_parallel_embedding import (
+        ParallelLMHead,
+        VocabParallelEmbedding,
+    )
+
+    for name, m in model.named_modules():
+        if isinstance(m, VocabParallelEmbedding):
+            # weight shape: [vocab_size // tp, embedding_dim]
+            weight_tensor = tensor_model_parallel_all_gather(m.weight, dim=0)
+            weight_tensor = weight_tensor[:m.org_vocab_size,:].contiguous()
+            if self.is_driver_worker:
+                m.weight = torch.nn.Parameter(weight_tensor.cpu(), requires_grad=False)
+                print(f"merged: {name}, shape={m.weight.shape}")
+        
+        elif isinstance(m, ParallelLMHead):
+            # weight shape: [vocab_size // tp, embedding_dim]
+            # bias shape: [vocab_size // tp]
+            if m.bias is not None:
+                bias = tensor_model_parallel_all_gather(m.bias, dim=0)
+                bias = bias[:m.org_vocab_size].contiguous()
+                if self.is_driver_worker:
+                    m.bias = torch.nn.Parameter(bias.cpu(), requires_grad=False)
+            
+            weight_tensor = tensor_model_parallel_all_gather(m.weight, dim=0)
+            weight_tensor = weight_tensor[:m.org_vocab_size,:].contiguous()
+            if self.is_driver_worker:
+                m.weight = torch.nn.Parameter(weight_tensor.cpu(), requires_grad=False)
+                print(f"merged: {name}, shape={m.weight.shape}")
+        
+        elif isinstance(m, QKVParallelLinear):
+            # weight shape: [total_num_head * head_size // tp + 2 * total_num_head * head_size // tp, hidden_size]
+            # bias shape: [total_num_head * head_size // tp + 2 * total_num_head * head_size // tp]
+            if self.parallel_config.world_size > 1:
+                total_q_hidden_size = m.total_num_heads * m.head_size
+                partial_q_hidden_size = m.num_heads * m.head_size
+                total_kv_hidden_size = m.total_num_kv_heads * m.head_size
+                partial_kv_hidden_size = m.num_kv_heads * m.head_size
+
+                if m.bias is not None:
+                    bias_tenosr = m.bias.new_zeros(total_q_hidden_size + total_kv_hidden_size * 2, m.hidden_size)
+                    q_bias = bias_tenosr[:total_q_hidden_size][self.rank * partial_q_hidden_size : (self.rank + 1) * partial_q_hidden_size]
+                    k_bias = bias_tenosr[total_q_hidden_size:total_q_hidden_size+total_kv_hidden_size]\
+                                [self.rank * partial_kv_hidden_size : (self.rank + 1) * partial_kv_hidden_size]
+                    v_bias = bias_tenosr[total_q_hidden_size+total_kv_hidden_size:]\
+                                [self.rank * partial_kv_hidden_size : (self.rank + 1) * partial_kv_hidden_size]
+                    
+                    q_bias[:] = m.bias[:partial_q_hidden_size]
+                    k_bias[:] = m.bias[partial_q_hidden_size:partial_q_hidden_size+partial_kv_hidden_size]
+                    v_bias[:] = m.bias[partial_q_hidden_size+partial_kv_hidden_size:]
+                    
+                    bias_tensor = tensor_model_parallel_all_reduce(bias_tenosr)
+                    if self.is_driver_worker:
+                        m.bias = torch.nn.Parameter(bias_tensor.cpu(), requires_grad=False)
+                
+                weight_tensor = m.weight.new_zeros(
+                    total_q_hidden_size + total_kv_hidden_size * 2, m.hidden_size
+                )
+
+                q_tensor = weight_tensor[:total_q_hidden_size, :]
+                q_tensor = q_tensor[self.rank * partial_q_hidden_size : (self.rank + 1) * partial_q_hidden_size]
+
+                k_tensor = weight_tensor[total_q_hidden_size : total_q_hidden_size + total_kv_hidden_size]
+                k_tensor = k_tensor[self.rank * partial_kv_hidden_size : (self.rank + 1) * partial_kv_hidden_size]
+
+                v_tensor = weight_tensor[total_q_hidden_size + total_kv_hidden_size :]
+                v_tensor = v_tensor[self.rank * partial_kv_hidden_size : (self.rank + 1) * partial_kv_hidden_size]
+
+                q_tensor[:, :] = m.weight[: partial_q_hidden_size, :]
+                k_tensor[:, :] = m.weight[partial_q_hidden_size : partial_q_hidden_size + partial_kv_hidden_size, :]
+                v_tensor[:, :] = m.weight[partial_q_hidden_size + partial_kv_hidden_size : , :]
+
+                weight_tensor = tensor_model_parallel_all_reduce(weight_tensor)
+            else:
+                weight_tensor = m.weight
+                if m.bias is not None and self.is_driver_worker:
+                    m.bias = torch.nn.Parameter(m.bias.cpu(), requires_grad=False)
+                    
+            int8_weight, weight_scales = w8a16_quantization(weight_tensor)
+            
+            if self.is_driver_worker:
+                m.weight = torch.nn.Parameter(int8_weight.cpu(), requires_grad=False)
+                m.scales = torch.nn.Parameter(weight_scales.cpu(), requires_grad=False)
+                print(f"Quantized: {name}")
+
+        elif isinstance(m, MergedColumnParallelLinear):
+            # weight shape: [intermediate_size // tp * 2, hidden_size]
+            # bias shape: [intermediate_size // tp * 2]
+            if self.parallel_config.world_size > 1:
+                output_sizes = m.output_sizes
+                output_size = sum(output_sizes)
+                partial_output_sizes = [
+                    i // self.parallel_config.world_size for i in output_sizes
+                ]
+                
+                if m.bias is not None:
+                    index_start = 0
+                    partial_index_start = 0
+                    bias_tenosr = m.bias.new_zeros(output_size)
+                    for i in range(len(output_sizes)):
+                        index_out = index_start + output_sizes[i]
+                        sub_bias_tensor = bias_tenosr[index_start:index_out]
+                        partial_size = partial_output_sizes[i]
+                        sub_bias_tensor[self.rank * partial_size:(self.rank+1) * partial_size] = m.bias[partial_index_start:partial_index_start+partial_size]
+                        
+                        index_start += output_sizes[i]
+                        partial_index_start += partial_size
+                    bias_tenosr = tensor_model_parallel_all_reduce(bias_tenosr)
+                    if self.is_driver_worker:
+                        m.bias = torch.nn.Parameter(bias_tenosr, requires_grad=False)
+
+                weight_tensor = m.weight.new_zeros(output_size, m.input_size)
+
+                index_start = 0
+                partial_index_start = 0
+                for i in range(len(output_sizes)):
+                    index_out = index_start + output_sizes[i]
+                    sub_weight_tensor = weight_tensor[index_start:index_out]
+                    partial_size = partial_output_sizes[i]
+                    sub_weight_tensor[self.rank * partial_size : (self.rank + 1) * partial_size] = m.weight[partial_index_start : partial_index_start + partial_size]
+
+                    index_start += output_sizes[i]
+                    partial_index_start += partial_size
+                weight_tensor = tensor_model_parallel_all_reduce(weight_tensor)
+            else:
+                weight_tensor = m.weight
+                if m.bias is not None and self.is_driver_worker:
+                    m.bias = torch.nn.Parameter(m.bias.cpu(), requires_grad=False)
+                
+            int8_weight, weight_scales = w8a16_quantization(weight_tensor)
+            
+            if self.is_driver_worker:
+                m.weight = torch.nn.Parameter(int8_weight.cpu(), requires_grad=False)
+                m.scales = torch.nn.Parameter(weight_scales.cpu(), requires_grad=False)
+                print(f"Quantized: {name}")
+        
+        elif isinstance(m, ColumnParallelLinear):
+            # weight shape: [some_dim // tp, hidden_size] // for this Linear, some_dim mostly is hidden_size * 4
+            # bias shape: [some_dim // tp]
+            if m.bias is not None:
+                bias_tenosr = tensor_model_parallel_all_gather(m.bias, dim=0)
+                if self.is_driver_worker:
+                    m.bias = torch.nn.Parameter(bias_tenosr.cpu(), requires_grad=False)
+            
+            weight_tensor = tensor_model_parallel_all_gather(m.weight, dim=0)
+            
+            int8_weight, weight_scales = w8a16_quantization(weight_tensor)
+            
+            if self.is_driver_worker:
+                m.weight = torch.nn.Parameter(int8_weight.cpu(), requires_grad=False)
+                m.scales = torch.nn.Parameter(weight_scales.cpu(), requires_grad=False)
+                print(f"Quantized: {name}")
+        
+        elif isinstance(m, RowParallelLinear):
+            # weight shape: [hidden_size, some_dim // tp] // for this Linear, some_dim mostly is hidden_size * 4 or intermediate_size
+            # bias shape: [hidden_size]
+            if m.bias is not None:
+                bias_tensor = tensor_model_parallel_all_gather(m.bias, dim=-1)
+                if self.is_driver_worker:
+                    m.bias = torch.nn.Parameter(m.bias.cpu(), requires_grad=False)
+            
+            weight_tensor = tensor_model_parallel_all_gather(m.weight, dim=-1)
+            int8_weight, weight_scales = w8a16_quantization(weight_tensor)
+            
+            if self.is_driver_worker:
+                m.weight = torch.nn.Parameter(int8_weight.cpu(), requires_grad=False)
+                m.scales = torch.nn.Parameter(weight_scales.cpu(), requires_grad=False)
+                print(f"Quantized: {name}")
+        else:
+            pass
+        
+        torch.cuda.empty_cache()
+    
+    # save weights
+    if self.is_driver_worker:
+        from safetensors.torch import save_file
+
+        tensors = {}
+        saved = False
+        count = 0
+        size_in_bytes = 0
+        
+        for name, weight in model.named_parameters():
+            if "lm_head" in name and model.config.tie_word_embeddings:
+                continue
+            size_in_bytes += weight.numel() * weight.element_size()
+            tensors[name] = weight
+            saved = False
+            if gb_per_file is not None and size_in_bytes >= gb_per_file * 1024 * 1024 * 1024:
+                weight_path = os.path.join(save_path, "model_{}.safetensors".format(str(count).zfill(6)))
+                save_file(tensors, weight_path)
+                print(f"The quantified weights were successfully saved in {weight_path}.")
+                tensors.clear()
+                saved = True
+                count += 1
+                size_in_bytes = 0
+        
+        if not saved: 
+            weight_path = os.path.join(save_path, "model_{}.safetensors".format(str(count).zfill(6)))   
+            save_file(tensors, weight_path)
+            print(f"The quantified weights were successfully saved in {weight_path}.")