mirror of
https://github.com/kvcache-ai/ktransformers.git
synced 2026-10-16 21:01:14 +08:00
losewayy
a02efe36b5
* [feat](kt-kernel): add AVX2 RAWINT4 routed-expert LoRA SFT backend * [fix](kt-kernel): fix RAWINT4 SFT wrapper and clamp gradients Allow the RAWINT4 production wrapper on AVX2 runtimes and reject shared backward pools before construction. Preserve FP32 clamp decisions across BF16 caches, remove unused configuration state, and cover rectangular dimensions, multiblock dX, clamp boundaries, and production wrapper tasks with tighter LoRA gradient checks. * [fix](kt-kernel): plumb swiglu_limit through the RAWINT4 SFT python path The production wrapper hard-zeroed swiglu_limit for every format except MXFP4, and both SFT python gates rejected non-zero values for RAWINT4 even though MOESFTConfig carries the field and the AVX2 kernel implements the clamp in forward plus the gradient mask in backward. RAWINT4_SFT now accepts a finite non-negative swiglu_limit end to end, while the AMX K2 part (which computes plain SiLU) fails fast instead of silently dropping the clamp. The production-wrapper contract test is parametrized over swiglu_limit 0 and 2, a model-config plumb test covers plain and nested text_config shapes, the __main__ entry exercises both limits, and bench_rawint4_sft.py gains --swiglu-limit.
1063 lines
47 KiB
Python
1063 lines
47 KiB
Python
# AMX SFT MoE Wrapper implementation
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
"""
|
|
AMX-based SFT MoE Wrapper. Forward/backward buffer management is in base class;
|
|
this file handles weight loading, LoRA init, and C++ task construction.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import ctypes
|
|
import logging
|
|
import math
|
|
import os
|
|
import glob as _glob
|
|
import torch
|
|
from typing import Optional, List
|
|
|
|
from kt_kernel_ext.moe import MOESFTConfig
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
from ..utils.loader import BF16SafeTensorLoader, MXFP4SafeTensorLoader, SafeTensorLoader
|
|
|
|
try:
|
|
from kt_kernel_ext.moe import (
|
|
AMXBF16_SFT_MOE,
|
|
AMXInt8_SFT_MOE,
|
|
AMXInt4_SFT_MOE,
|
|
AMXBF16_SFT_MOE_SkipLoRA,
|
|
AMXInt8_SFT_MOE_SkipLoRA,
|
|
AMXInt4_SFT_MOE_SkipLoRA,
|
|
)
|
|
|
|
_HAS_AMX_SFT_SUPPORT = True
|
|
except (ImportError, AttributeError):
|
|
_HAS_AMX_SFT_SUPPORT = False
|
|
AMXBF16_SFT_MOE = None
|
|
AMXInt8_SFT_MOE = None
|
|
AMXInt4_SFT_MOE = None
|
|
AMXBF16_SFT_MOE_SkipLoRA = None
|
|
AMXInt8_SFT_MOE_SkipLoRA = None
|
|
AMXInt4_SFT_MOE_SkipLoRA = None
|
|
|
|
try:
|
|
from kt_kernel_ext.moe import AMXFP8_SFT_MOE
|
|
except (ImportError, AttributeError):
|
|
AMXFP8_SFT_MOE = None
|
|
|
|
try:
|
|
from kt_kernel_ext.moe import AMXInt4_KGroup_SFT_MOE
|
|
except (ImportError, AttributeError):
|
|
AMXInt4_KGroup_SFT_MOE = None
|
|
|
|
try:
|
|
from kt_kernel_ext.moe import MXFP4_SFT_MOE
|
|
except (ImportError, AttributeError):
|
|
MXFP4_SFT_MOE = None
|
|
|
|
from .base import BaseSFTMoEWrapper, KExpertsSFTBuffer, _supports_authoritative_optimizer_grads
|
|
from .backend import (
|
|
get_mxfp4_runtime,
|
|
get_rawint4_runtime,
|
|
is_fp8_sft_method,
|
|
is_int8_sft_method,
|
|
is_mxfp4_sft_method,
|
|
is_rawint4_sft_method,
|
|
)
|
|
from .weights import BlockFP8ExpertWeights, RAWINT4ExpertWeights
|
|
|
|
_AMX_M_STEP = 32
|
|
|
|
|
|
# Mapping from method string to C++ SFT MOE class
|
|
_SFT_METHOD_TO_CLASS = {
|
|
"AMXBF16_SFT": AMXBF16_SFT_MOE,
|
|
"AMXFP8_SFT": AMXFP8_SFT_MOE,
|
|
"MXFP4_SFT": MXFP4_SFT_MOE,
|
|
"INT8_SFT": AMXInt8_SFT_MOE,
|
|
"AMXINT8_SFT": AMXInt8_SFT_MOE,
|
|
"AMXINT4_SFT": AMXInt4_SFT_MOE,
|
|
"RAWINT4_SFT": AMXInt4_KGroup_SFT_MOE,
|
|
"AMXINT4_KGroup_SFT": AMXInt4_KGroup_SFT_MOE,
|
|
"AMXBF16_SFT_SkipLoRA": AMXBF16_SFT_MOE_SkipLoRA,
|
|
"AMXINT8_SFT_SkipLoRA": AMXInt8_SFT_MOE_SkipLoRA,
|
|
"AMXINT4_SFT_SkipLoRA": AMXInt4_SFT_MOE_SkipLoRA,
|
|
}
|
|
|
|
|
|
class AMXSFTMoEWrapper(BaseSFTMoEWrapper):
|
|
"""
|
|
AMX-based SFT MoE wrapper.
|
|
|
|
Supports BF16, native block-FP8/MXFP4, INT8, INT4, and SkipLoRA variants.
|
|
Forward/backward buffer management is in BaseSFTMoEWrapper;
|
|
this class implements weight loading and C++ task construction.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
layer_idx: int,
|
|
num_experts: int,
|
|
num_experts_per_tok: int,
|
|
hidden_size: int,
|
|
moe_intermediate_size: int,
|
|
num_gpu_experts: int,
|
|
cpuinfer_threads: int,
|
|
threadpool_count: int,
|
|
weight_path: str,
|
|
chunked_prefill_size: int,
|
|
lora_rank: int = 16,
|
|
lora_alpha: float = 32.0,
|
|
lora_dropout: float = 0.0,
|
|
max_cache_depth: int = 1,
|
|
method: str = "AMXBF16_SFT",
|
|
group_size: int = 128,
|
|
zero_point: bool = True,
|
|
full_weight_grad: bool = False,
|
|
swiglu_limit: float = 0.0,
|
|
):
|
|
if not _HAS_AMX_SFT_SUPPORT and not (is_mxfp4_sft_method(method) or is_rawint4_sft_method(method)):
|
|
raise RuntimeError(
|
|
"AMX SFT backend not available. kt_kernel_ext was not compiled with AMX SFT support.\n"
|
|
"Please recompile with AMX SFT enabled."
|
|
)
|
|
if is_mxfp4_sft_method(method):
|
|
get_mxfp4_runtime()
|
|
if is_rawint4_sft_method(method):
|
|
get_rawint4_runtime()
|
|
|
|
super().__init__(
|
|
layer_idx=layer_idx,
|
|
num_experts=num_experts,
|
|
num_experts_per_tok=num_experts_per_tok,
|
|
hidden_size=hidden_size,
|
|
moe_intermediate_size=moe_intermediate_size,
|
|
num_gpu_experts=num_gpu_experts,
|
|
cpuinfer_threads=cpuinfer_threads,
|
|
threadpool_count=threadpool_count,
|
|
weight_path=weight_path,
|
|
chunked_prefill_size=chunked_prefill_size,
|
|
lora_rank=lora_rank,
|
|
lora_alpha=lora_alpha,
|
|
lora_dropout=lora_dropout,
|
|
max_cache_depth=max_cache_depth,
|
|
full_weight_grad=full_weight_grad,
|
|
)
|
|
|
|
self.method = method
|
|
self._is_rawint4 = is_rawint4_sft_method(method)
|
|
self._is_skip_lora = "SkipLoRA" in method
|
|
self._uses_authoritative_optimizer_grads = _supports_authoritative_optimizer_grads(
|
|
method,
|
|
self.num_gpu_experts,
|
|
full_weight_grad=self._full_weight_grad,
|
|
lora_rank=self.lora_rank,
|
|
)
|
|
self.group_size = group_size
|
|
self.zero_point = zero_point
|
|
self.swiglu_limit = float(swiglu_limit)
|
|
|
|
if not (is_mxfp4_sft_method(method) or is_rawint4_sft_method(method)) and self.swiglu_limit != 0.0:
|
|
raise ValueError(f"swiglu_limit is only supported by MXFP4_SFT/RAWINT4_SFT, got {method!r}")
|
|
if is_mxfp4_sft_method(method):
|
|
if self._full_weight_grad or self.lora_rank <= 0:
|
|
raise ValueError("MXFP4_SFT supports frozen-base LoRA only")
|
|
if self.num_gpu_experts != 0:
|
|
raise ValueError("MXFP4_SFT requires all routed experts on CPU")
|
|
if self.group_size != 32 or self.zero_point:
|
|
raise ValueError("MXFP4_SFT requires group_size=32 and zero_point=False")
|
|
if self.hidden_size % 32 or self.moe_intermediate_size % 32:
|
|
raise ValueError("MXFP4_SFT requires hidden_size and moe_intermediate_size divisible by 32")
|
|
if self.threadpool_count not in (1, 2):
|
|
raise ValueError("MXFP4_SFT currently supports threadpool_count 1 or 2")
|
|
if self.moe_intermediate_size % self.threadpool_count:
|
|
raise ValueError("MXFP4_SFT requires moe_intermediate_size divisible by threadpool_count")
|
|
if (self.moe_intermediate_size // self.threadpool_count) % 32:
|
|
raise ValueError("MXFP4_SFT requires every TP intermediate slice divisible by 32")
|
|
if not math.isfinite(self.swiglu_limit) or self.swiglu_limit <= 0.0:
|
|
raise ValueError("MXFP4_SFT requires a finite positive swiglu_limit")
|
|
|
|
if is_fp8_sft_method(method):
|
|
if self._full_weight_grad or self.lora_rank <= 0:
|
|
raise ValueError("AMXFP8_SFT supports frozen-base LoRA only")
|
|
if self.num_gpu_experts != 0:
|
|
raise ValueError("AMXFP8_SFT requires all routed experts on CPU")
|
|
if self.hidden_size % 128 or self.moe_intermediate_size % 128:
|
|
raise ValueError("AMXFP8_SFT requires hidden_size and moe_intermediate_size " "divisible by 128")
|
|
if self.threadpool_count < 1 or self.moe_intermediate_size % self.threadpool_count:
|
|
raise ValueError("AMXFP8_SFT requires moe_intermediate_size divisible by " "threadpool_count")
|
|
if (self.moe_intermediate_size // self.threadpool_count) % 128:
|
|
raise ValueError("AMXFP8_SFT requires every TP intermediate slice divisible by 128")
|
|
self.group_size = 128
|
|
self.zero_point = False
|
|
|
|
if is_rawint4_sft_method(self.method):
|
|
if self._full_weight_grad:
|
|
raise ValueError("RAWINT4 SFT supports frozen-base LoRA only")
|
|
if self.lora_rank != 8:
|
|
raise ValueError("RAWINT4 SFT v1 requires LoRA rank 8")
|
|
if self.lora_dropout != 0.0:
|
|
raise ValueError("RAWINT4 SFT v1 requires LoRA dropout 0")
|
|
if self._is_skip_lora:
|
|
raise ValueError("RAWINT4 SFT v1 does not support SkipLoRA")
|
|
if self.num_gpu_experts != 0:
|
|
raise ValueError("RAWINT4 SFT requires all routed experts on CPU")
|
|
if self.threadpool_count not in {1, 2}:
|
|
raise ValueError("RAWINT4 SFT v1 supports TP1/TP2 only")
|
|
if self.hidden_size % 32 or self.moe_intermediate_size % 32:
|
|
raise ValueError("RAWINT4 SFT requires hidden_size and moe_intermediate_size " "divisible by 32")
|
|
if (
|
|
self.moe_intermediate_size % self.threadpool_count
|
|
or (self.moe_intermediate_size // self.threadpool_count) % 32
|
|
):
|
|
raise ValueError("RAWINT4 SFT requires every TP intermediate slice divisible by 32")
|
|
if not math.isfinite(self.swiglu_limit) or self.swiglu_limit < 0.0:
|
|
raise ValueError("RAWINT4 SFT requires a finite non-negative swiglu_limit")
|
|
self.group_size = 32
|
|
self.zero_point = False
|
|
|
|
if method not in _SFT_METHOD_TO_CLASS:
|
|
raise ValueError(f"Unknown SFT method: {method}. Supported: {list(_SFT_METHOD_TO_CLASS.keys())}")
|
|
|
|
moe_class = _SFT_METHOD_TO_CLASS[method]
|
|
if moe_class is None:
|
|
raise RuntimeError(f"AMX SFT method '{method}' not available in current build.")
|
|
|
|
self.gate_proj: Optional[torch.Tensor] = None
|
|
self.up_proj: Optional[torch.Tensor] = None
|
|
self.down_proj: Optional[torch.Tensor] = None
|
|
self.gate_scale: Optional[torch.Tensor] = None
|
|
self.up_scale: Optional[torch.Tensor] = None
|
|
self.down_scale: Optional[torch.Tensor] = None
|
|
|
|
self._moe_class = moe_class
|
|
|
|
# ========== Template method: C++ task construction ==========
|
|
|
|
def _make_forward_task(self, buffer: KExpertsSFTBuffer, save_for_backward: bool):
|
|
return self.moe.forward_sft_task(
|
|
buffer.bsz_tensor.data_ptr(),
|
|
self.num_experts_per_tok,
|
|
buffer.expert_ids_cpu.data_ptr(),
|
|
buffer.weights_cpu.data_ptr(),
|
|
buffer.input_cpu.data_ptr(),
|
|
buffer.output_cpu.data_ptr(),
|
|
save_for_backward,
|
|
)
|
|
|
|
def _make_backward_task(
|
|
self,
|
|
buffer: KExpertsSFTBuffer,
|
|
accumulate_optimizer_grads: bool = False,
|
|
optimizer_grad_scale: float = 1.0,
|
|
):
|
|
if self._is_skip_lora:
|
|
return self.moe.backward_task(
|
|
buffer.grad_output_cpu.data_ptr(),
|
|
buffer.grad_input_cpu.data_ptr(),
|
|
0,
|
|
0,
|
|
0,
|
|
0,
|
|
0,
|
|
0,
|
|
buffer.grad_weights.data_ptr(),
|
|
0,
|
|
0,
|
|
0, # grad_gate_proj, grad_up_proj, grad_down_proj
|
|
)
|
|
|
|
# Base weight grad pointers (nullptr if not in full mode)
|
|
grad_gate_proj_ptr = (
|
|
self.grad_gate_proj_buf.data_ptr() if self._full_weight_grad and self.grad_gate_proj_buf is not None else 0
|
|
)
|
|
grad_up_proj_ptr = (
|
|
self.grad_up_proj_buf.data_ptr() if self._full_weight_grad and self.grad_up_proj_buf is not None else 0
|
|
)
|
|
grad_down_proj_ptr = (
|
|
self.grad_down_proj_buf.data_ptr() if self._full_weight_grad and self.grad_down_proj_buf is not None else 0
|
|
)
|
|
|
|
backward_args = (
|
|
buffer.grad_output_cpu.data_ptr(),
|
|
buffer.grad_input_cpu.data_ptr(),
|
|
self.grad_gate_lora_a.data_ptr() if self.lora_rank > 0 else 0,
|
|
self.grad_gate_lora_b.data_ptr() if self.lora_rank > 0 else 0,
|
|
self.grad_up_lora_a.data_ptr() if self.lora_rank > 0 else 0,
|
|
self.grad_up_lora_b.data_ptr() if self.lora_rank > 0 else 0,
|
|
self.grad_down_lora_a.data_ptr() if self.lora_rank > 0 else 0,
|
|
self.grad_down_lora_b.data_ptr() if self.lora_rank > 0 else 0,
|
|
buffer.grad_weights.data_ptr(),
|
|
grad_gate_proj_ptr,
|
|
grad_up_proj_ptr,
|
|
grad_down_proj_ptr,
|
|
)
|
|
if (
|
|
self._uses_authoritative_optimizer_grads
|
|
or bool(accumulate_optimizer_grads)
|
|
or float(optimizer_grad_scale) != 1.0
|
|
):
|
|
return self.moe.backward_task(
|
|
*backward_args,
|
|
bool(accumulate_optimizer_grads),
|
|
float(optimizer_grad_scale),
|
|
)
|
|
return self.moe.backward_task(*backward_args)
|
|
|
|
# ========== Weight loading ==========
|
|
|
|
def load_weights(self, physical_to_logical_map_cpu: torch.Tensor) -> None:
|
|
if self._weights_loaded:
|
|
if is_mxfp4_sft_method(self.method):
|
|
raise RuntimeError("MXFP4_SFT weights are immutable and already loaded")
|
|
return
|
|
|
|
if physical_to_logical_map_cpu is None:
|
|
physical_to_logical_map_cpu = torch.arange(self.num_experts, dtype=torch.int64)
|
|
strict_rawint4_map = bool(getattr(self, "_is_rawint4", False))
|
|
if strict_rawint4_map:
|
|
if not isinstance(physical_to_logical_map_cpu, torch.Tensor):
|
|
raise TypeError("physical_to_logical_map_cpu must be a torch.Tensor")
|
|
if physical_to_logical_map_cpu.dtype not in {torch.int32, torch.int64}:
|
|
raise ValueError(
|
|
"physical_to_logical_map_cpu must use an integer dtype, "
|
|
f"got {physical_to_logical_map_cpu.dtype}."
|
|
)
|
|
self._physical_to_logical_map_cpu = (
|
|
physical_to_logical_map_cpu.to(dtype=torch.int64, device="cpu").reshape(-1).contiguous()
|
|
)
|
|
if strict_rawint4_map and self._physical_to_logical_map_cpu.numel() != self.num_experts:
|
|
raise ValueError(
|
|
"physical_to_logical_map_cpu must contain exactly "
|
|
f"{self.num_experts} entries, got {self._physical_to_logical_map_cpu.numel()}."
|
|
)
|
|
if not strict_rawint4_map and self._physical_to_logical_map_cpu.numel() < self.num_experts:
|
|
raise ValueError(
|
|
"physical_to_logical_map_cpu must contain at least "
|
|
f"{self.num_experts} entries, got {self._physical_to_logical_map_cpu.numel()}."
|
|
)
|
|
expected_experts = torch.arange(self.num_experts, dtype=torch.int64)
|
|
if strict_rawint4_map and not torch.equal(self._physical_to_logical_map_cpu.sort().values, expected_experts):
|
|
raise ValueError(
|
|
"physical_to_logical_map_cpu must be a permutation of "
|
|
f"[0, {self.num_experts}); got {self._physical_to_logical_map_cpu.tolist()}."
|
|
)
|
|
if is_mxfp4_sft_method(self.method):
|
|
logical_experts = self._physical_to_logical_map_cpu[: self.num_experts].tolist()
|
|
if sorted(logical_experts) != list(range(self.num_experts)):
|
|
raise ValueError(
|
|
"MXFP4_SFT physical_to_logical_map_cpu must be a permutation of " f"[0, {self.num_experts})."
|
|
)
|
|
|
|
if self.gate_proj is None and not getattr(self, "_use_projs_path", False):
|
|
self._load_base_weights_from_file()
|
|
|
|
config = MOESFTConfig()
|
|
config.expert_num = self.num_experts
|
|
config.num_experts_per_tok = self.num_experts_per_tok
|
|
config.hidden_size = self.hidden_size
|
|
config.intermediate_size = self.moe_intermediate_size
|
|
config.lora_rank = self.lora_rank
|
|
config.lora_alpha = self.lora_alpha
|
|
config.lora_dropout = self.lora_dropout
|
|
config.max_cache_depth = self.max_cache_depth
|
|
config.max_len = self._aligned_max_len()
|
|
config.layer_idx = self.layer_idx
|
|
config.share_backward_bb = getattr(self, "share_backward_bb", False)
|
|
config.share_cache_pool = getattr(self, "share_cache_pool", False)
|
|
config.full_weight_grad = self._full_weight_grad
|
|
config.authoritative_optimizer_grads = self._uses_authoritative_optimizer_grads
|
|
config.physical_to_logical_map = self._physical_to_logical_map_cpu.data_ptr()
|
|
config.swiglu_limit = self.swiglu_limit
|
|
|
|
if getattr(self, "_use_kt_direct_load", False):
|
|
config.load = True
|
|
config.path = self.weight_path
|
|
elif getattr(self, "_use_projs_path", False):
|
|
config.gate_projs = self._gate_projs_ptrs
|
|
config.up_projs = self._up_projs_ptrs
|
|
config.down_projs = self._down_projs_ptrs
|
|
config.gate_scales = self._gate_scale_ptrs
|
|
config.up_scales = self._up_scale_ptrs
|
|
config.down_scales = self._down_scale_ptrs
|
|
if getattr(self, "_bf16_gate_proj", None) is not None:
|
|
config.gate_proj = self._bf16_gate_proj.data_ptr()
|
|
config.up_proj = self._bf16_up_proj.data_ptr()
|
|
config.down_proj = self._bf16_down_proj.data_ptr()
|
|
if getattr(self, "_has_bwd_projs", False):
|
|
config.gate_bwd_projs = self._gate_bwd_projs_ptrs
|
|
config.up_bwd_projs = self._up_bwd_projs_ptrs
|
|
config.down_bwd_projs = self._down_bwd_projs_ptrs
|
|
config.gate_bwd_scales = self._gate_bwd_scale_ptrs
|
|
config.up_bwd_scales = self._up_bwd_scale_ptrs
|
|
config.down_bwd_scales = self._down_bwd_scale_ptrs
|
|
else:
|
|
config.gate_proj = self.gate_proj.data_ptr()
|
|
config.up_proj = self.up_proj.data_ptr()
|
|
config.down_proj = self.down_proj.data_ptr()
|
|
if getattr(self, "_use_rawint4_tensor_path", False):
|
|
config.gate_scale = self.gate_scale.data_ptr()
|
|
config.up_scale = self.up_scale.data_ptr()
|
|
config.down_scale = self.down_scale.data_ptr()
|
|
|
|
if self._lora_initialized:
|
|
config.gate_lora_a = self.gate_lora_a.data_ptr()
|
|
config.gate_lora_b = self.gate_lora_b.data_ptr()
|
|
config.up_lora_a = self.up_lora_a.data_ptr()
|
|
config.up_lora_b = self.up_lora_b.data_ptr()
|
|
config.down_lora_a = self.down_lora_a.data_ptr()
|
|
config.down_lora_b = self.down_lora_b.data_ptr()
|
|
|
|
config.pool = self.cpu_infer.backend_
|
|
|
|
if is_rawint4_sft_method(self.method):
|
|
config.quant_config.bits = 4
|
|
config.quant_config.group_size = self.group_size
|
|
config.quant_config.zero_point = self.zero_point
|
|
elif is_fp8_sft_method(self.method):
|
|
config.quant_config.bits = 8
|
|
config.quant_config.group_size = 128
|
|
config.quant_config.zero_point = False
|
|
elif is_mxfp4_sft_method(self.method):
|
|
config.quant_config.bits = 4
|
|
config.quant_config.group_size = 32
|
|
config.quant_config.zero_point = False
|
|
|
|
# Release old C++ MOE object before creating a new one to avoid memory leak
|
|
old_moe = getattr(self, "moe", None)
|
|
if old_moe is not None:
|
|
del old_moe
|
|
import gc
|
|
|
|
gc.collect()
|
|
|
|
self.moe = self._moe_class(config)
|
|
|
|
self.cpu_infer.submit(self.moe.load_weights_task())
|
|
self.cpu_infer.sync()
|
|
|
|
if os.environ.get("KT_SFT_ENABLE_WARMUP", "0") == "1":
|
|
self.cpu_infer.submit(self.moe.warm_up_task())
|
|
self.cpu_infer.sync()
|
|
|
|
# Release Python-side weight tensors (C++ copied them)
|
|
# In full_weight_grad mode, keep them for nn.Parameter initialization
|
|
if not self._full_weight_grad:
|
|
self.gate_proj = None
|
|
self.up_proj = None
|
|
self.down_proj = None
|
|
self.gate_scale = None
|
|
self.up_scale = None
|
|
self.down_scale = None
|
|
|
|
if getattr(self, "_bf16_gate_proj", None) is not None:
|
|
self._bf16_gate_proj = None
|
|
self._bf16_up_proj = None
|
|
self._bf16_down_proj = None
|
|
|
|
if getattr(self, "_use_projs_path", False):
|
|
for attr in [
|
|
"_gate_weights_per_numa",
|
|
"_up_weights_per_numa",
|
|
"_down_weights_per_numa",
|
|
"_gate_scales_per_numa",
|
|
"_up_scales_per_numa",
|
|
"_down_scales_per_numa",
|
|
"_gate_projs_ptrs",
|
|
"_up_projs_ptrs",
|
|
"_down_projs_ptrs",
|
|
"_gate_scale_ptrs",
|
|
"_up_scale_ptrs",
|
|
"_down_scale_ptrs",
|
|
]:
|
|
setattr(self, attr, None)
|
|
if getattr(self, "_has_bwd_projs", False):
|
|
for attr in [
|
|
"_gate_bwd_weights_per_numa",
|
|
"_up_bwd_weights_per_numa",
|
|
"_down_bwd_weights_per_numa",
|
|
"_gate_bwd_scales_per_numa",
|
|
"_up_bwd_scales_per_numa",
|
|
"_down_bwd_scales_per_numa",
|
|
"_gate_bwd_projs_ptrs",
|
|
"_up_bwd_projs_ptrs",
|
|
"_down_bwd_projs_ptrs",
|
|
"_gate_bwd_scale_ptrs",
|
|
"_up_bwd_scale_ptrs",
|
|
"_down_bwd_scale_ptrs",
|
|
]:
|
|
setattr(self, attr, None)
|
|
|
|
self._weights_loaded = True
|
|
|
|
def load_weights_from_tensors(
|
|
self,
|
|
gate_proj: torch.Tensor,
|
|
up_proj: torch.Tensor,
|
|
down_proj: torch.Tensor,
|
|
physical_to_logical_map_cpu: torch.Tensor,
|
|
) -> None:
|
|
if is_rawint4_sft_method(self.method):
|
|
raise ValueError(
|
|
"RAWINT4 SFT accepts packed weights plus BF16 group scales only; " "use load_rawint4_weights()"
|
|
)
|
|
if is_fp8_sft_method(self.method):
|
|
raise ValueError("AMXFP8_SFT accepts raw per-expert E4M3 weights only; " "use load_block_fp8_weights()")
|
|
if is_mxfp4_sft_method(self.method):
|
|
raise ValueError(
|
|
"MXFP4_SFT accepts native packed E2M1 weights and UE8M0 scales only; " "use load_mxfp4_weights()"
|
|
)
|
|
if is_int8_sft_method(self.method):
|
|
raise ValueError(
|
|
"INT8_SFT accepts pre-quantized .kt weights only; " "online tensor conversion is not supported"
|
|
)
|
|
self.gate_proj = gate_proj.contiguous()
|
|
self.up_proj = up_proj.contiguous()
|
|
self.down_proj = down_proj.contiguous()
|
|
self.load_weights(physical_to_logical_map_cpu)
|
|
del gate_proj, up_proj, down_proj
|
|
|
|
def load_rawint4_weights(
|
|
self,
|
|
weights: RAWINT4ExpertWeights,
|
|
physical_to_logical_map_cpu: torch.Tensor,
|
|
) -> None:
|
|
"""Synchronously load strict signed group-32 packed expert weights."""
|
|
|
|
if not is_rawint4_sft_method(self.method):
|
|
raise ValueError(f"load_rawint4_weights() requires RAWINT4_SFT, got {self.method!r}")
|
|
if int(weights.group_size) != 32:
|
|
raise ValueError(f"RAWINT4 SFT requires group_size=32, got {weights.group_size}")
|
|
|
|
projections = {
|
|
"gate": (
|
|
weights.gate_proj,
|
|
weights.gate_scale,
|
|
(self.num_experts, self.moe_intermediate_size, self.hidden_size // 2),
|
|
(self.num_experts, self.moe_intermediate_size, self.hidden_size // 32),
|
|
),
|
|
"up": (
|
|
weights.up_proj,
|
|
weights.up_scale,
|
|
(self.num_experts, self.moe_intermediate_size, self.hidden_size // 2),
|
|
(self.num_experts, self.moe_intermediate_size, self.hidden_size // 32),
|
|
),
|
|
"down": (
|
|
weights.down_proj,
|
|
weights.down_scale,
|
|
(self.num_experts, self.hidden_size, self.moe_intermediate_size // 2),
|
|
(self.num_experts, self.hidden_size, self.moe_intermediate_size // 32),
|
|
),
|
|
}
|
|
for name, (packed, scale, packed_shape, scale_shape) in projections.items():
|
|
if (
|
|
packed.device.type != "cpu"
|
|
or packed.dtype != torch.uint8
|
|
or not packed.is_contiguous()
|
|
or tuple(packed.shape) != packed_shape
|
|
):
|
|
raise ValueError(f"{name} RAWINT4 weight must be contiguous CPU uint8 with " f"shape {packed_shape}")
|
|
if (
|
|
scale.device.type != "cpu"
|
|
or scale.dtype != torch.bfloat16
|
|
or not scale.is_contiguous()
|
|
or tuple(scale.shape) != scale_shape
|
|
):
|
|
raise ValueError(f"{name} RAWINT4 scale must be contiguous CPU BF16 with " f"shape {scale_shape}")
|
|
|
|
self.gate_proj = weights.gate_proj
|
|
self.gate_scale = weights.gate_scale
|
|
self.up_proj = weights.up_proj
|
|
self.up_scale = weights.up_scale
|
|
self.down_proj = weights.down_proj
|
|
self.down_scale = weights.down_scale
|
|
self._use_rawint4_tensor_path = True
|
|
self.load_weights(physical_to_logical_map_cpu)
|
|
|
|
def load_block_fp8_weights(
|
|
self,
|
|
weights: BlockFP8ExpertWeights,
|
|
physical_to_logical_map_cpu: torch.Tensor,
|
|
) -> None:
|
|
"""Synchronously pack raw per-expert FP8 checkpoint tensors in C++."""
|
|
|
|
if not is_fp8_sft_method(self.method):
|
|
raise ValueError(f"load_block_fp8_weights() requires AMXFP8_SFT, got {self.method!r}")
|
|
if tuple(weights.block_size) != (128, 128):
|
|
raise ValueError("AMXFP8_SFT requires block_size=(128, 128), " f"got {weights.block_size}")
|
|
|
|
projections = {
|
|
"gate": (weights.gate_proj, weights.gate_scale, (self.moe_intermediate_size, self.hidden_size)),
|
|
"up": (weights.up_proj, weights.up_scale, (self.moe_intermediate_size, self.hidden_size)),
|
|
"down": (weights.down_proj, weights.down_scale, (self.hidden_size, self.moe_intermediate_size)),
|
|
}
|
|
for name, (expert_weights, expert_scales, expected_shape) in projections.items():
|
|
if len(expert_weights) != self.num_experts or len(expert_scales) != self.num_experts:
|
|
raise ValueError(
|
|
f"{name} FP8 tensors must contain {self.num_experts} experts, got "
|
|
f"weights={len(expert_weights)}, scales={len(expert_scales)}"
|
|
)
|
|
expected_scale_shape = (expected_shape[0] // 128, expected_shape[1] // 128)
|
|
for expert_idx, (weight, scale) in enumerate(zip(expert_weights, expert_scales)):
|
|
if (
|
|
weight.device.type != "cpu"
|
|
or weight.dtype != torch.float8_e4m3fn
|
|
or not weight.is_contiguous()
|
|
or tuple(weight.shape) != expected_shape
|
|
):
|
|
raise ValueError(
|
|
f"{name} expert {expert_idx} must be contiguous CPU E4M3FN " f"with shape {expected_shape}"
|
|
)
|
|
if (
|
|
scale.device.type != "cpu"
|
|
or scale.dtype != torch.float32
|
|
or not scale.is_contiguous()
|
|
or tuple(scale.shape) != expected_scale_shape
|
|
):
|
|
raise ValueError(
|
|
f"{name} scale {expert_idx} must be contiguous CPU FP32 " f"with shape {expected_scale_shape}"
|
|
)
|
|
|
|
self._gate_weights_per_numa = [weights.gate_proj]
|
|
self._up_weights_per_numa = [weights.up_proj]
|
|
self._down_weights_per_numa = [weights.down_proj]
|
|
self._gate_scales_per_numa = [weights.gate_scale]
|
|
self._up_scales_per_numa = [weights.up_scale]
|
|
self._down_scales_per_numa = [weights.down_scale]
|
|
self._gate_projs_ptrs = [[tensor.data_ptr() for tensor in weights.gate_proj]]
|
|
self._up_projs_ptrs = [[tensor.data_ptr() for tensor in weights.up_proj]]
|
|
self._down_projs_ptrs = [[tensor.data_ptr() for tensor in weights.down_proj]]
|
|
self._gate_scale_ptrs = [[tensor.data_ptr() for tensor in weights.gate_scale]]
|
|
self._up_scale_ptrs = [[tensor.data_ptr() for tensor in weights.up_scale]]
|
|
self._down_scale_ptrs = [[tensor.data_ptr() for tensor in weights.down_scale]]
|
|
self._has_bwd_projs = False
|
|
self._use_projs_path = True
|
|
self.gate_proj = None
|
|
self.up_proj = None
|
|
self.down_proj = None
|
|
self.load_weights(physical_to_logical_map_cpu)
|
|
|
|
def _stage_mxfp4_weights(self, experts_data: dict[str, List[torch.Tensor]]) -> None:
|
|
"""Validate native MXFP4 tensors and retain them until C++ load completes."""
|
|
|
|
if not is_mxfp4_sft_method(self.method):
|
|
raise ValueError(f"MXFP4 weights require MXFP4_SFT, got {self.method!r}")
|
|
if not isinstance(experts_data, dict):
|
|
raise TypeError("MXFP4 expert weights must be the dict returned by " "MXFP4SafeTensorLoader.load_experts()")
|
|
|
|
expected = {
|
|
"gate": (torch.uint8, (self.moe_intermediate_size, self.hidden_size // 2)),
|
|
"up": (torch.uint8, (self.moe_intermediate_size, self.hidden_size // 2)),
|
|
"down": (torch.uint8, (self.hidden_size, self.moe_intermediate_size // 2)),
|
|
"gate_scale": (
|
|
torch.bfloat16,
|
|
(self.moe_intermediate_size, self.hidden_size // 32),
|
|
),
|
|
"up_scale": (
|
|
torch.bfloat16,
|
|
(self.moe_intermediate_size, self.hidden_size // 32),
|
|
),
|
|
"down_scale": (
|
|
torch.bfloat16,
|
|
(self.hidden_size, self.moe_intermediate_size // 32),
|
|
),
|
|
}
|
|
validated: dict[str, List[torch.Tensor]] = {}
|
|
for name, (dtype, shape) in expected.items():
|
|
tensors = experts_data.get(name)
|
|
if not isinstance(tensors, (list, tuple)) or len(tensors) != self.num_experts:
|
|
actual = len(tensors) if isinstance(tensors, (list, tuple)) else type(tensors).__name__
|
|
raise ValueError(f"MXFP4 {name} must contain {self.num_experts} expert tensors, got {actual}")
|
|
checked = []
|
|
for expert_idx, tensor in enumerate(tensors):
|
|
if not isinstance(tensor, torch.Tensor):
|
|
raise TypeError(f"MXFP4 {name}[{expert_idx}] must be a torch.Tensor, got {type(tensor)!r}")
|
|
if (
|
|
tensor.device.type != "cpu"
|
|
or tensor.dtype != dtype
|
|
or tuple(tensor.shape) != shape
|
|
or not tensor.is_contiguous()
|
|
):
|
|
raise ValueError(
|
|
f"MXFP4 {name}[{expert_idx}] must be contiguous CPU {dtype} "
|
|
f"with shape {shape}, got {tensor.dtype} {tuple(tensor.shape)} "
|
|
f"on {tensor.device}"
|
|
)
|
|
if name.endswith("_scale") and not torch.isfinite(tensor).all().item():
|
|
raise ValueError(f"MXFP4 {name}[{expert_idx}] contains a non-finite scale")
|
|
checked.append(tensor)
|
|
validated[name] = checked
|
|
|
|
self._gate_weights_per_numa = [validated["gate"]]
|
|
self._up_weights_per_numa = [validated["up"]]
|
|
self._down_weights_per_numa = [validated["down"]]
|
|
self._gate_scales_per_numa = [validated["gate_scale"]]
|
|
self._up_scales_per_numa = [validated["up_scale"]]
|
|
self._down_scales_per_numa = [validated["down_scale"]]
|
|
self._gate_projs_ptrs = [[tensor.data_ptr() for tensor in validated["gate"]]]
|
|
self._up_projs_ptrs = [[tensor.data_ptr() for tensor in validated["up"]]]
|
|
self._down_projs_ptrs = [[tensor.data_ptr() for tensor in validated["down"]]]
|
|
self._gate_scale_ptrs = [[tensor.data_ptr() for tensor in validated["gate_scale"]]]
|
|
self._up_scale_ptrs = [[tensor.data_ptr() for tensor in validated["up_scale"]]]
|
|
self._down_scale_ptrs = [[tensor.data_ptr() for tensor in validated["down_scale"]]]
|
|
self._has_bwd_projs = False
|
|
self._use_projs_path = True
|
|
self.gate_proj = None
|
|
self.up_proj = None
|
|
self.down_proj = None
|
|
|
|
def load_mxfp4_weights(
|
|
self,
|
|
experts_data: dict[str, List[torch.Tensor]],
|
|
physical_to_logical_map_cpu: torch.Tensor,
|
|
) -> None:
|
|
"""Load native E2M1/UE8M0 expert tensors without dequantizing the base."""
|
|
|
|
if self._weights_loaded:
|
|
raise RuntimeError("MXFP4_SFT weights are immutable and already loaded")
|
|
self._stage_mxfp4_weights(experts_data)
|
|
self.load_weights(physical_to_logical_map_cpu)
|
|
|
|
def _load_base_weights_from_file(self) -> None:
|
|
if not hasattr(self, "weight_path") or self.weight_path is None:
|
|
raise RuntimeError(
|
|
"weight_path not set. Cannot load weights from file. "
|
|
"Either set weight_path or call load_weights_from_tensors() instead."
|
|
)
|
|
|
|
if is_mxfp4_sft_method(self.method):
|
|
loader = MXFP4SafeTensorLoader(self.weight_path)
|
|
try:
|
|
experts_data = loader.load_experts(
|
|
f"model.layers.{self.layer_idx}",
|
|
device="cpu",
|
|
reject_non_finite_scales=True,
|
|
)
|
|
self._stage_mxfp4_weights(experts_data)
|
|
finally:
|
|
loader.close_all_handles()
|
|
return
|
|
|
|
kt_layer_dir = os.path.join(self.weight_path, f"_layer_{self.layer_idx}")
|
|
if os.path.isdir(kt_layer_dir):
|
|
kt_files = _glob.glob(os.path.join(kt_layer_dir, "_numa_0", "*.kt"))
|
|
if kt_files:
|
|
self._use_kt_direct_load = True
|
|
return
|
|
if is_fp8_sft_method(self.method):
|
|
raise RuntimeError(
|
|
"AMXFP8_SFT requires raw per-expert HF checkpoint tensors; "
|
|
"legacy merged weight files are not supported"
|
|
)
|
|
if is_rawint4_sft_method(self.method):
|
|
raise RuntimeError(
|
|
"RAWINT4 SFT requires packed compressed-tensors checkpoint tensors; "
|
|
"legacy merged weight files are not supported"
|
|
)
|
|
if is_int8_sft_method(self.method):
|
|
raise RuntimeError(
|
|
f"INT8_SFT requires pre-quantized .kt files under "
|
|
f"{kt_layer_dir}/_numa_*/; merged safetensors and online "
|
|
"conversion are not supported"
|
|
)
|
|
|
|
if "BF16" in self.method:
|
|
loader = BF16SafeTensorLoader(self.weight_path)
|
|
base_key = f"model.layers.{self.layer_idx}"
|
|
else:
|
|
loader = SafeTensorLoader(self.weight_path)
|
|
base_key = f"blk.{self.layer_idx}"
|
|
|
|
experts_data = loader.load_experts(base_key, device="cpu")
|
|
|
|
gate_weights: List[torch.Tensor] = experts_data["gate"]
|
|
up_weights: List[torch.Tensor] = experts_data["up"]
|
|
down_weights: List[torch.Tensor] = experts_data["down"]
|
|
|
|
if "BF16" in self.method:
|
|
self.gate_proj = torch.stack(gate_weights, dim=0).contiguous()
|
|
self.up_proj = torch.stack(up_weights, dim=0).contiguous()
|
|
self.down_proj = torch.stack(down_weights, dim=0).contiguous()
|
|
else:
|
|
|
|
def _make_ptrs(arrays_per_numa):
|
|
return [
|
|
[
|
|
ctypes.addressof(ctypes.cast(et.ctypes.data, ctypes.POINTER(ctypes.c_uint64)).contents)
|
|
for et in numa_array
|
|
]
|
|
for numa_array in arrays_per_numa
|
|
]
|
|
|
|
self._gate_weights_per_numa = gate_weights
|
|
self._up_weights_per_numa = up_weights
|
|
self._down_weights_per_numa = down_weights
|
|
self._gate_scales_per_numa = experts_data["gate_scale"]
|
|
self._up_scales_per_numa = experts_data["up_scale"]
|
|
self._down_scales_per_numa = experts_data["down_scale"]
|
|
self._validate_prepartitioned_weights()
|
|
|
|
self._gate_projs_ptrs = _make_ptrs(gate_weights)
|
|
self._up_projs_ptrs = _make_ptrs(up_weights)
|
|
self._down_projs_ptrs = _make_ptrs(down_weights)
|
|
self._gate_scale_ptrs = _make_ptrs(experts_data["gate_scale"])
|
|
self._up_scale_ptrs = _make_ptrs(experts_data["up_scale"])
|
|
self._down_scale_ptrs = _make_ptrs(experts_data["down_scale"])
|
|
|
|
if "gate_bwd" in experts_data:
|
|
self._gate_bwd_weights_per_numa = experts_data["gate_bwd"]
|
|
self._up_bwd_weights_per_numa = experts_data["up_bwd"]
|
|
self._down_bwd_weights_per_numa = experts_data["down_bwd"]
|
|
self._gate_bwd_scales_per_numa = experts_data["gate_bwd_scale"]
|
|
self._up_bwd_scales_per_numa = experts_data["up_bwd_scale"]
|
|
self._down_bwd_scales_per_numa = experts_data["down_bwd_scale"]
|
|
|
|
self._gate_bwd_projs_ptrs = _make_ptrs(experts_data["gate_bwd"])
|
|
self._up_bwd_projs_ptrs = _make_ptrs(experts_data["up_bwd"])
|
|
self._down_bwd_projs_ptrs = _make_ptrs(experts_data["down_bwd"])
|
|
self._gate_bwd_scale_ptrs = _make_ptrs(experts_data["gate_bwd_scale"])
|
|
self._up_bwd_scale_ptrs = _make_ptrs(experts_data["up_bwd_scale"])
|
|
self._down_bwd_scale_ptrs = _make_ptrs(experts_data["down_bwd_scale"])
|
|
self._has_bwd_projs = True
|
|
else:
|
|
self._has_bwd_projs = False
|
|
|
|
self.gate_proj = None
|
|
self.up_proj = None
|
|
self.down_proj = None
|
|
self._use_projs_path = True
|
|
|
|
loader.close_all_handles()
|
|
|
|
def _aligned_max_len(self) -> int:
|
|
return ((self.chunked_prefill_size + _AMX_M_STEP - 1) // _AMX_M_STEP) * _AMX_M_STEP
|
|
|
|
def _validate_prepartitioned_weights(self) -> None:
|
|
numa_count = len(self._gate_weights_per_numa)
|
|
if self.moe_intermediate_size % self.threadpool_count != 0:
|
|
raise ValueError(
|
|
f"moe_intermediate_size={self.moe_intermediate_size} must be divisible by "
|
|
f"threadpool_count={self.threadpool_count} for {self.method} SFT."
|
|
)
|
|
if numa_count != self.threadpool_count:
|
|
raise ValueError(
|
|
f"{self.method} SFT pre-partitioned expert weights have {numa_count} NUMA partitions, "
|
|
f"but CPUInfer was created with threadpool_count={self.threadpool_count}. "
|
|
f"Use --kt-threadpool-count {numa_count} for this weight directory, or convert weights "
|
|
"for the requested threadpool count."
|
|
)
|
|
|
|
collections = {
|
|
"gate": self._gate_weights_per_numa,
|
|
"up": self._up_weights_per_numa,
|
|
"down": self._down_weights_per_numa,
|
|
"gate_scale": self._gate_scales_per_numa,
|
|
"up_scale": self._up_scales_per_numa,
|
|
"down_scale": self._down_scales_per_numa,
|
|
}
|
|
for name, per_numa in collections.items():
|
|
if len(per_numa) != numa_count:
|
|
raise ValueError(f"{name} has {len(per_numa)} NUMA partitions, expected {numa_count}.")
|
|
for numa_id, entries in enumerate(per_numa):
|
|
if len(entries) != self.num_experts:
|
|
raise ValueError(f"{name}[numa={numa_id}] has {len(entries)} experts, expected {self.num_experts}.")
|
|
|
|
for numa_id in range(numa_count):
|
|
gate_scale_len = self._gate_scales_per_numa[numa_id][0].size
|
|
up_scale_len = self._up_scales_per_numa[numa_id][0].size
|
|
down_scale_len = self._down_scales_per_numa[numa_id][0].size
|
|
expected_intermediate = self.moe_intermediate_size // self.threadpool_count
|
|
if gate_scale_len != expected_intermediate or up_scale_len != expected_intermediate:
|
|
raise ValueError(
|
|
f"{self.method} gate/up scale length for NUMA {numa_id} is "
|
|
f"{gate_scale_len}/{up_scale_len}, expected {expected_intermediate}."
|
|
)
|
|
if down_scale_len != self.hidden_size:
|
|
raise ValueError(
|
|
f"{self.method} down scale length for NUMA {numa_id} is "
|
|
f"{down_scale_len}, expected {self.hidden_size}."
|
|
)
|
|
|
|
# ========== LoRA ==========
|
|
|
|
def init_lora_weights(
|
|
self,
|
|
gate_lora_a: torch.Tensor,
|
|
gate_lora_b: torch.Tensor,
|
|
up_lora_a: torch.Tensor,
|
|
up_lora_b: torch.Tensor,
|
|
down_lora_a: torch.Tensor,
|
|
down_lora_b: torch.Tensor,
|
|
grad_gate_lora_a: torch.Tensor,
|
|
grad_gate_lora_b: torch.Tensor,
|
|
grad_up_lora_a: torch.Tensor,
|
|
grad_up_lora_b: torch.Tensor,
|
|
grad_down_lora_a: torch.Tensor,
|
|
grad_down_lora_b: torch.Tensor,
|
|
) -> None:
|
|
expected_shapes = {
|
|
"gate_lora_a": (self.num_experts, self.lora_rank, self.hidden_size),
|
|
"gate_lora_b": (self.num_experts, self.moe_intermediate_size, self.lora_rank),
|
|
"up_lora_a": (self.num_experts, self.lora_rank, self.hidden_size),
|
|
"up_lora_b": (self.num_experts, self.moe_intermediate_size, self.lora_rank),
|
|
"down_lora_a": (self.num_experts, self.lora_rank, self.moe_intermediate_size),
|
|
"down_lora_b": (self.num_experts, self.hidden_size, self.lora_rank),
|
|
}
|
|
provided = {
|
|
"gate_lora_a": gate_lora_a,
|
|
"gate_lora_b": gate_lora_b,
|
|
"up_lora_a": up_lora_a,
|
|
"up_lora_b": up_lora_b,
|
|
"down_lora_a": down_lora_a,
|
|
"down_lora_b": down_lora_b,
|
|
}
|
|
for name, tensor in provided.items():
|
|
expected = expected_shapes[name]
|
|
if tensor.shape != expected:
|
|
raise ValueError(f"{name} shape mismatch: expected {expected}, got {tuple(tensor.shape)}")
|
|
if tensor.device.type != "cpu":
|
|
raise ValueError(f"{name} must be a CPU tensor for {self.method} SFT, got {tensor.device}.")
|
|
if tensor.dtype != torch.bfloat16:
|
|
raise ValueError(f"{name} must use torch.bfloat16 for {self.method} SFT, got {tensor.dtype}.")
|
|
if not tensor.is_contiguous():
|
|
raise ValueError(f"{name} must be contiguous for stable {self.method} pointer registration.")
|
|
|
|
grad_provided = {
|
|
"grad_gate_lora_a": grad_gate_lora_a,
|
|
"grad_gate_lora_b": grad_gate_lora_b,
|
|
"grad_up_lora_a": grad_up_lora_a,
|
|
"grad_up_lora_b": grad_up_lora_b,
|
|
"grad_down_lora_a": grad_down_lora_a,
|
|
"grad_down_lora_b": grad_down_lora_b,
|
|
}
|
|
for name, tensor in grad_provided.items():
|
|
expected = expected_shapes[name.removeprefix("grad_")]
|
|
if tensor.shape != expected:
|
|
raise ValueError(f"{name} shape mismatch: expected {expected}, got {tuple(tensor.shape)}")
|
|
if tensor.device.type != "cpu" or tensor.dtype != torch.bfloat16:
|
|
raise ValueError(
|
|
f"{name} must be a CPU torch.bfloat16 tensor for {self.method} SFT, "
|
|
f"got {tensor.dtype} on {tensor.device}."
|
|
)
|
|
if not tensor.is_contiguous():
|
|
raise ValueError(f"{name} must be contiguous for stable {self.method} pointer registration.")
|
|
|
|
self.gate_lora_a = gate_lora_a
|
|
self.gate_lora_b = gate_lora_b
|
|
self.up_lora_a = up_lora_a
|
|
self.up_lora_b = up_lora_b
|
|
self.down_lora_a = down_lora_a
|
|
self.down_lora_b = down_lora_b
|
|
|
|
self.grad_gate_lora_a = grad_gate_lora_a
|
|
self.grad_gate_lora_b = grad_gate_lora_b
|
|
self.grad_up_lora_a = grad_up_lora_a
|
|
self.grad_up_lora_b = grad_up_lora_b
|
|
self.grad_down_lora_a = grad_down_lora_a
|
|
self.grad_down_lora_b = grad_down_lora_b
|
|
|
|
if getattr(self, "_uses_authoritative_optimizer_grads", False):
|
|
parameter_count = sum(isinstance(tensor, torch.nn.Parameter) for tensor in provided.values())
|
|
if parameter_count not in (0, len(provided)):
|
|
raise ValueError(
|
|
"authoritative LoRA ownership requires all six weights to be nn.Parameter objects or none"
|
|
)
|
|
if parameter_count:
|
|
for name, parameter in provided.items():
|
|
self.register_authoritative_optimizer_grad(f"lora.{name}", parameter, grad_provided[f"grad_{name}"])
|
|
|
|
self._lora_initialized = True
|
|
|
|
if self._weights_loaded and self.moe is not None:
|
|
self.update_lora_weights()
|
|
|
|
def update_lora_weights(self) -> None:
|
|
if not self._weights_loaded:
|
|
raise RuntimeError("Weights not loaded. Call load_weights() first.")
|
|
if self._is_skip_lora:
|
|
return
|
|
if not self._lora_initialized:
|
|
if self.lora_rank <= 0:
|
|
return # Full mode without LoRA — no LoRA weights to update
|
|
raise RuntimeError("LoRA weights not initialized. Call init_lora_weights() first.")
|
|
|
|
# Weight pointer updates are load-time synchronous work. Calling the
|
|
# direct binding avoids nesting an update task inside CPUInfer's queue
|
|
# while SGLang is still in distributed model-loading barriers.
|
|
self.moe.update_lora_weights(
|
|
self.gate_lora_a.data_ptr(),
|
|
self.gate_lora_b.data_ptr(),
|
|
self.up_lora_a.data_ptr(),
|
|
self.up_lora_b.data_ptr(),
|
|
self.down_lora_a.data_ptr(),
|
|
self.down_lora_b.data_ptr(),
|
|
)
|
|
|
|
def update_base_weights(self) -> None:
|
|
"""Sync updated base weight parameters back to C++ kernel after optimizer step."""
|
|
if not self._weights_loaded:
|
|
raise RuntimeError("Weights not loaded. Call load_weights() first.")
|
|
if not self._full_weight_grad:
|
|
return # No base weights to update in LoRA mode
|
|
if self.gate_proj_buf is None:
|
|
raise RuntimeError("Base weight buffers not initialized. Call init_full_weight_grad_buffers() first.")
|
|
|
|
logger.info(f"Layer {self.layer_idx}: update_base_weights() - syncing updated weights to C++ kernel")
|
|
|
|
# Preferred path: update config pointers on existing C++ object and re-quantize.
|
|
# This avoids full C++ MOE object recreation (~0.6s/layer vs ~1.9s/layer).
|
|
if hasattr(self.moe, "set_base_weight_pointers"):
|
|
self.moe.set_base_weight_pointers(
|
|
self.gate_proj_buf.data.data_ptr(),
|
|
self.up_proj_buf.data.data_ptr(),
|
|
self.down_proj_buf.data.data_ptr(),
|
|
)
|
|
self.cpu_infer.submit(self.moe.load_weights_task())
|
|
self.cpu_infer.sync()
|
|
logger.info(f"Layer {self.layer_idx}: update_base_weights() - re-quantized existing kernel")
|
|
return
|
|
|
|
# Fallback: full reload path (creates new C++ MOE object)
|
|
# This is slower but works without C++ set_base_weight_pointers support.
|
|
logger.warning(
|
|
f"Layer {self.layer_idx}: set_base_weight_pointers not available, "
|
|
f"falling back to full C++ MOE object recreation"
|
|
)
|
|
old_moe = getattr(self, "moe", None)
|
|
if old_moe is not None:
|
|
del old_moe
|
|
|
|
self.gate_proj = self.gate_proj_buf.data
|
|
self.up_proj = self.up_proj_buf.data
|
|
self.down_proj = self.down_proj_buf.data
|
|
physical_to_logical_map = torch.arange(self.num_experts, dtype=torch.int64, device="cpu")
|
|
self._weights_loaded = False # Allow re-load
|
|
self.load_weights_from_tensors(
|
|
gate_proj=self.gate_proj,
|
|
up_proj=self.up_proj,
|
|
down_proj=self.down_proj,
|
|
physical_to_logical_map_cpu=physical_to_logical_map,
|
|
)
|
|
|
|
def save_backward_weights_from_tensors(
|
|
self,
|
|
gate_proj: torch.Tensor,
|
|
up_proj: torch.Tensor,
|
|
down_proj: torch.Tensor,
|
|
physical_to_logical_map: torch.Tensor,
|
|
output_path: str,
|
|
) -> None:
|
|
if not self._weights_loaded:
|
|
raise RuntimeError("Weights not loaded. Call load_weights() first.")
|
|
gate_proj = gate_proj.contiguous()
|
|
up_proj = up_proj.contiguous()
|
|
down_proj = down_proj.contiguous()
|
|
self.moe.prepare_and_save_bwd(
|
|
gate_proj.data_ptr(),
|
|
up_proj.data_ptr(),
|
|
down_proj.data_ptr(),
|
|
output_path,
|
|
)
|