mirror of
https://github.com/kvcache-ai/ktransformers.git
synced 2026-10-17 05:11:19 +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.
1459 lines
62 KiB
Python
1459 lines
62 KiB
Python
# Model wrapping entry points for SFT
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
from __future__ import annotations
|
|
|
|
import gc
|
|
import importlib.util as _u
|
|
import json
|
|
import logging
|
|
import os
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
from .arch import (
|
|
KTAMXConfigError,
|
|
KTAMXNotAvailableError,
|
|
_get_layers_prefix,
|
|
_get_model_container_and_layers,
|
|
get_moe_arch_config,
|
|
get_moe_module,
|
|
)
|
|
from .layer import KTMoELayerWrapper, KTRoutedExpertsWrapper
|
|
from .lora import LoRAExperts
|
|
from .base import _supports_authoritative_optimizer_grads
|
|
from .backend import (
|
|
FP8_BACKEND,
|
|
INT8_BACKEND,
|
|
MXFP4_BACKEND,
|
|
MXFP4_SFT_METHOD,
|
|
RAWINT4_BACKEND,
|
|
RAWINT4_GROUP_SIZE,
|
|
RAWINT4_SFT_METHOD,
|
|
RAWINT4_WEIGHT_LAYOUT,
|
|
get_fp8_runtime,
|
|
get_int8_runtime,
|
|
get_mxfp4_runtime,
|
|
get_rawint4_checkpoint_contract,
|
|
get_rawint4_runtime,
|
|
)
|
|
from .checkpoint import load_full_weight_layer, resolve_full_weight_checkpoint
|
|
from .conv3d_compat import patch_vlm_conv3d
|
|
from .dist_utils import _distributed_rank_world_size
|
|
from .weights import (
|
|
_clear_original_expert_weights,
|
|
extract_moe_weights,
|
|
load_block_fp8_experts_from_checkpoint_files,
|
|
load_experts_from_checkpoint_files,
|
|
load_rawint4_experts_from_checkpoint_files,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
KT_KERNEL_AVAILABLE = _u.find_spec("kt_kernel") is not None
|
|
|
|
if KT_KERNEL_AVAILABLE:
|
|
try:
|
|
from kt_kernel.experts import KTMoEWrapper
|
|
except Exception:
|
|
KTMoEWrapper = None
|
|
KT_KERNEL_AVAILABLE = False
|
|
else:
|
|
KTMoEWrapper = None
|
|
|
|
|
|
def _supports_checkpoint_forward_reuse(full_weight_grad: bool, lora_rank: int) -> bool:
|
|
return (full_weight_grad and lora_rank == 0) or (not full_weight_grad and lora_rank > 0)
|
|
|
|
|
|
def _validate_rawint4_activation_policy(expert_weight_format: str | None, activation_policy: Any) -> None:
|
|
"""Keep ordinary caching distinct from checkpoint-forward reuse."""
|
|
|
|
if expert_weight_format == "rawint4" and activation_policy.cpu == "retain" and activation_policy.gpu == "recompute":
|
|
raise KTAMXConfigError(
|
|
"RAWINT4 SFT v1 supports activation_policy cpu=retain,gpu=retain "
|
|
"for ordinary forward/backward caching, but checkpoint-forward reuse "
|
|
"(cpu=retain,gpu=recompute) is not supported"
|
|
)
|
|
|
|
|
|
def _native_fp8_block_size(model_config: Any) -> tuple[int, int]:
|
|
"""Read and strictly validate the checkpoint's native FP8 contract."""
|
|
|
|
text_config = getattr(model_config, "text_config", None)
|
|
quant_config = getattr(model_config, "quantization_config", None)
|
|
if quant_config is None and text_config is not None:
|
|
quant_config = getattr(text_config, "quantization_config", None)
|
|
if quant_config is None:
|
|
raise KTAMXConfigError(
|
|
"native FP8 SFT requires model quantization_config metadata"
|
|
)
|
|
|
|
def read(name: str):
|
|
if isinstance(quant_config, dict):
|
|
return quant_config.get(name)
|
|
return getattr(quant_config, name, None)
|
|
|
|
quant_method = str(read("quant_method") or "").lower()
|
|
if "fp8" not in quant_method:
|
|
raise KTAMXConfigError(
|
|
"native FP8 SFT requires an FP8 checkpoint, got "
|
|
f"quant_method={quant_method!r}"
|
|
)
|
|
block_size = read("weight_block_size")
|
|
if block_size is None:
|
|
block_size = read("weight_block_shape")
|
|
if block_size is None or tuple(block_size) != (128, 128):
|
|
raise KTAMXConfigError(
|
|
"native FP8 SFT requires weight_block_size=[128, 128], "
|
|
f"got {block_size!r}"
|
|
)
|
|
return (128, 128)
|
|
|
|
|
|
def _resolve_raw_safetensors_checkpoint_files(
|
|
model_name_or_path: str,
|
|
) -> tuple[list[str] | None, dict | None]:
|
|
"""Resolve raw safetensors, including a dependency-light local fallback."""
|
|
|
|
resolved_files, resolved_metadata = _resolve_checkpoint_files(
|
|
model_name_or_path=model_name_or_path
|
|
)
|
|
if resolved_files and all(str(path).endswith(".safetensors") for path in resolved_files):
|
|
return [str(path) for path in resolved_files], resolved_metadata
|
|
|
|
checkpoint_path = Path(model_name_or_path)
|
|
if checkpoint_path.is_file() and checkpoint_path.suffix == ".safetensors":
|
|
return [str(checkpoint_path)], None
|
|
if not checkpoint_path.is_dir():
|
|
return None, None
|
|
|
|
index_path = checkpoint_path / "model.safetensors.index.json"
|
|
if index_path.is_file():
|
|
with index_path.open(encoding="utf-8") as handle:
|
|
metadata = json.load(handle)
|
|
weight_map = metadata.get("weight_map")
|
|
if not isinstance(weight_map, dict) or not weight_map:
|
|
raise KTAMXConfigError(
|
|
f"invalid safetensors index without weight_map: {index_path}"
|
|
)
|
|
files = sorted(
|
|
{str(checkpoint_path / filename) for filename in weight_map.values()}
|
|
)
|
|
missing = [path for path in files if not os.path.isfile(path)]
|
|
if missing:
|
|
raise FileNotFoundError(
|
|
f"safetensors index references missing shard: {missing[0]}"
|
|
)
|
|
return files, metadata
|
|
|
|
files = sorted(str(path) for path in checkpoint_path.glob("*.safetensors"))
|
|
return (files, None) if files else (None, None)
|
|
|
|
|
|
def _sync_rank0_wrap_error(
|
|
error: BaseException | None,
|
|
*,
|
|
context: str,
|
|
rank: int,
|
|
world_size: int,
|
|
) -> None:
|
|
"""Make a rank-0-only wrapping stage fail coherently on every rank."""
|
|
if world_size <= 1:
|
|
if error is not None:
|
|
raise error
|
|
return
|
|
|
|
import torch.distributed as dist
|
|
|
|
if not dist.is_initialized():
|
|
if error is not None:
|
|
raise error
|
|
return
|
|
|
|
payload = [
|
|
None
|
|
if error is None
|
|
else f"{type(error).__name__}: {error}"
|
|
]
|
|
dist.broadcast_object_list(payload, src=0)
|
|
if payload[0] is None:
|
|
return
|
|
if rank == 0 and error is not None:
|
|
raise error
|
|
raise RuntimeError(f"{context}: rank 0 failed: {payload[0]}")
|
|
|
|
|
|
# =============================================================================
|
|
# Device-map builders
|
|
# =============================================================================
|
|
|
|
|
|
def _get_kt_config(kt_plugin: Any):
|
|
"""Resolve a private KTConfig from a compatible public container."""
|
|
from .config import KTConfig
|
|
|
|
if isinstance(kt_plugin, KTConfig):
|
|
return kt_plugin
|
|
return KTConfig.from_object(kt_plugin)
|
|
|
|
|
|
def _device_map_model_contract(config: Any) -> tuple[int, str, str, str]:
|
|
"""Resolve text-layer count and non-layer module paths for device maps."""
|
|
|
|
text_config = getattr(config, "text_config", None)
|
|
effective_config = text_config if text_config is not None else config
|
|
num_layers = getattr(effective_config, "num_hidden_layers", None)
|
|
if isinstance(num_layers, bool) or not isinstance(num_layers, int) or num_layers <= 0:
|
|
raise KTAMXConfigError("model text config has invalid num_hidden_layers")
|
|
|
|
layers_prefix = _get_layers_prefix(config)
|
|
if layers_prefix == "language_model.model.layers":
|
|
return (
|
|
num_layers,
|
|
"language_model.model.embed_tokens",
|
|
"language_model.model.norm",
|
|
"language_model.lm_head",
|
|
)
|
|
if layers_prefix == "model.language_model.layers":
|
|
return (
|
|
num_layers,
|
|
"model.language_model.embed_tokens",
|
|
"model.language_model.norm",
|
|
"lm_head",
|
|
)
|
|
return num_layers, "model.embed_tokens", "model.norm", "lm_head"
|
|
|
|
|
|
def build_kt_device_map(config, kt_plugin, device: str = "cuda:0") -> dict[str, str | int]:
|
|
"""
|
|
Build device_map for KT model loading with hybrid GPU/CPU expert placement.
|
|
"""
|
|
moe_config = get_moe_arch_config(config)
|
|
layers_prefix = _get_layers_prefix(config)
|
|
num_layers, embed_path, norm_path, lm_head_path = _device_map_model_contract(config)
|
|
num_experts = moe_config.expert_num
|
|
cfg = _get_kt_config(kt_plugin)
|
|
num_gpu_experts = getattr(cfg, "kt_num_gpu_experts", 0) or 0
|
|
|
|
device_map: dict[str, str | int] = {}
|
|
|
|
device_map[embed_path] = device
|
|
device_map[norm_path] = device
|
|
device_map[lm_head_path] = device
|
|
|
|
for layer_idx in range(num_layers):
|
|
layer_prefix = f"{layers_prefix}.{layer_idx}"
|
|
device_map[layer_prefix] = device
|
|
moe_prefix = f"{layer_prefix}.{moe_config.moe_layer_attr}"
|
|
|
|
for expert_idx in range(num_experts):
|
|
expert_key = f"{moe_prefix}.{moe_config.experts_attr}.{expert_idx}"
|
|
if expert_idx < num_gpu_experts:
|
|
device_map[expert_key] = device
|
|
else:
|
|
device_map[expert_key] = "cpu"
|
|
|
|
logger.info(f"Built KT device_map: {num_gpu_experts} GPU experts, {num_experts - num_gpu_experts} CPU experts")
|
|
|
|
return device_map
|
|
|
|
|
|
def build_kt_device_map_simplified(config, kt_plugin, device: str = "cuda:0") -> dict[str, str | int]:
|
|
"""
|
|
Simplified device_map builder: map full layers to GPU, override routed experts to CPU.
|
|
"""
|
|
moe_config = get_moe_arch_config(config)
|
|
layers_prefix = _get_layers_prefix(config)
|
|
num_layers, embed_path, norm_path, lm_head_path = _device_map_model_contract(config)
|
|
cfg = _get_kt_config(kt_plugin)
|
|
num_gpu_experts = getattr(cfg, "kt_num_gpu_experts", 0) or 0
|
|
|
|
device_map: dict[str, str | int] = {}
|
|
|
|
device_map[embed_path] = device
|
|
device_map[norm_path] = device
|
|
device_map[lm_head_path] = device
|
|
|
|
for layer_idx in range(num_layers):
|
|
layer_prefix = f"{layers_prefix}.{layer_idx}"
|
|
device_map[layer_prefix] = device
|
|
|
|
experts_prefix = f"{layer_prefix}.{moe_config.moe_layer_attr}.{moe_config.experts_attr}"
|
|
|
|
if num_gpu_experts == 0:
|
|
device_map[experts_prefix] = "cpu"
|
|
else:
|
|
return build_kt_device_map(config, kt_plugin, device=device)
|
|
|
|
logger.info("Built simplified KT device_map: all layers on GPU, routed experts on CPU")
|
|
return device_map
|
|
|
|
|
|
# =============================================================================
|
|
# MoE layer wrapping
|
|
# =============================================================================
|
|
|
|
|
|
def wrap_moe_layers_with_kt_wrapper(model: nn.Module, kt_plugin: Any) -> list[KTMoELayerWrapper]:
|
|
"""
|
|
Replace model's MoE layers with KTMoEWrapper-based wrappers.
|
|
|
|
Loads expert weights into the C++ KT kernel. No LoRA initialization ---
|
|
LoRA is handled by PEFT and later adapted via kt_adapt_peft_lora().
|
|
Only rank 0 initializes KT kernel and loads weights.
|
|
"""
|
|
if not KT_KERNEL_AVAILABLE:
|
|
raise KTAMXNotAvailableError("kt_kernel not found. Please install kt_kernel to enable KT MoE support.")
|
|
|
|
# Only global rank 0 initializes KT. Launcher env fallback matters when
|
|
# model construction happens before init_process_group().
|
|
distributed_rank, distributed_world_size = _distributed_rank_world_size()
|
|
is_rank_0 = distributed_rank == 0
|
|
|
|
moe_config = get_moe_arch_config(model.config)
|
|
patched_conv3d = patch_vlm_conv3d(model)
|
|
if patched_conv3d:
|
|
logger.info(f"Patched KT VLM Conv3D modules: {patched_conv3d}")
|
|
_text_cfg = getattr(model.config, "text_config", model.config)
|
|
hidden_size = _text_cfg.hidden_size
|
|
|
|
cfg = _get_kt_config(kt_plugin)
|
|
activation_policy = cfg.kt_activation_policy
|
|
|
|
# Read lora_rank/lora_alpha for C++ wrapper initialization (buffer allocation only)
|
|
# Use explicit None checks: lora_rank=0 is a valid value (full mode, no LoRA),
|
|
# but `or` pattern would treat 0 as falsy and replace it with 1.
|
|
_raw_rank = getattr(cfg, "kt_lora_rank", None)
|
|
lora_rank = _raw_rank if _raw_rank is not None else 1
|
|
_raw_alpha = getattr(cfg, "kt_lora_alpha", None)
|
|
lora_alpha = _raw_alpha if _raw_alpha is not None else 1.0
|
|
_raw_dropout = getattr(cfg, "kt_lora_dropout", None)
|
|
lora_dropout = _raw_dropout if _raw_dropout is not None else 0.0
|
|
|
|
# Read full_weight_grad mode
|
|
_raw_fwg = getattr(cfg, "kt_full_weight_grad", None)
|
|
full_weight_grad = _raw_fwg if _raw_fwg is not None else False
|
|
train_mode = getattr(cfg, "kt_train_mode", "lora")
|
|
|
|
# Full and hybrid are explicit modes. LlamaFactory exposes a default
|
|
# lora_rank even for full tuning, which must not silently turn Full into
|
|
# Hybrid. Preserve the legacy fallback for callers without train_mode.
|
|
if train_mode == "full":
|
|
lora_rank = 0
|
|
elif full_weight_grad and train_mode != "hybrid" and lora_rank > 0:
|
|
_has_explicit_lora_rank = getattr(cfg, "kt_lora_rank", None) is not None
|
|
if not _has_explicit_lora_rank:
|
|
lora_rank = 0
|
|
|
|
# Read LoRA Experts configuration
|
|
_raw_le = getattr(cfg, "kt_use_lora_experts", None)
|
|
use_lora_experts = bool(_raw_le) if _raw_le is not None else False
|
|
lora_expert_num = getattr(cfg, "kt_lora_expert_num", 2) or 2
|
|
lora_expert_intermediate_size = getattr(cfg, "kt_lora_expert_intermediate_size", 1024) or 1024
|
|
|
|
if is_rank_0:
|
|
logger.info(
|
|
f"LoRA Experts config: use_lora_experts={use_lora_experts}, "
|
|
f"num={lora_expert_num}, intermediate_size={lora_expert_intermediate_size}"
|
|
)
|
|
if full_weight_grad:
|
|
logger.info(f"Full weight gradient mode enabled (lora_rank={lora_rank})")
|
|
|
|
wrappers: list[KTMoELayerWrapper] = []
|
|
moe_layer_count = 0
|
|
|
|
kt_backend_map = {
|
|
"AMXBF16": "AMXBF16_SFT",
|
|
FP8_BACKEND: "AMXFP8_SFT",
|
|
"AMXFP8": "AMXFP8_SFT",
|
|
INT8_BACKEND: "INT8_SFT",
|
|
"AMXINT8": "INT8_SFT",
|
|
RAWINT4_BACKEND: RAWINT4_SFT_METHOD,
|
|
MXFP4_BACKEND: MXFP4_SFT_METHOD,
|
|
"AMXINT4_KGroup": RAWINT4_SFT_METHOD,
|
|
"AMXINT4": "AMXINT4_SFT",
|
|
"AMXBF16_SkipLoRA": "AMXBF16_SFT_SkipLoRA",
|
|
"AMXINT8_SkipLoRA": "AMXINT8_SFT_SkipLoRA",
|
|
"AMXINT4_SkipLoRA": "AMXINT4_SFT_SkipLoRA",
|
|
}
|
|
# Case-insensitive matching remains for compatibility, but an unknown
|
|
# backend must never silently turn a requested quantized run into BF16.
|
|
_kt_backend_map_lower = {k.lower(): v for k, v in kt_backend_map.items()}
|
|
kt_backend = str(getattr(cfg, "kt_backend", "AMXBF16"))
|
|
kt_method = kt_backend_map.get(kt_backend) or _kt_backend_map_lower.get(kt_backend.lower())
|
|
if kt_method is None:
|
|
raise KTAMXConfigError(
|
|
f"Unknown kt_backend {kt_backend!r}; expected one of {sorted(kt_backend_map)}"
|
|
)
|
|
if kt_backend not in kt_backend_map:
|
|
logger.warning(
|
|
f"kt_backend '{kt_backend}' matched via case-insensitive lookup -> '{kt_method}'. "
|
|
f"Please use the exact name from: {list(kt_backend_map.keys())}"
|
|
)
|
|
|
|
if "SkipLoRA" in kt_method:
|
|
logger.info(f"Using SkipLoRA backend: {kt_method} (MoE LoRA gradients will be skipped)")
|
|
force_fused_expert_lora = bool(
|
|
getattr(cfg, "kt_force_fused_expert_lora", False)
|
|
)
|
|
if force_fused_expert_lora and "SkipLoRA" in kt_method:
|
|
raise KTAMXConfigError(
|
|
"kt_force_fused_expert_lora is incompatible with SkipLoRA backends"
|
|
)
|
|
requested_num_gpu_experts = int(getattr(cfg, "kt_num_gpu_experts", 0) or 0)
|
|
expert_weight_format = getattr(cfg, "kt_expert_weight_format", None)
|
|
if expert_weight_format == "int8":
|
|
if kt_method != "INT8_SFT":
|
|
raise KTAMXConfigError(
|
|
"kt_expert_weight_format='int8' requires kt_backend='auto' or 'INT8'"
|
|
)
|
|
if full_weight_grad or train_mode != "lora" or lora_rank <= 0:
|
|
raise KTAMXConfigError(
|
|
"INT8 SFT supports frozen-base LoRA only; Full and Hybrid are not supported"
|
|
)
|
|
if requested_num_gpu_experts != 0 or use_lora_experts:
|
|
raise KTAMXConfigError(
|
|
"INT8 SFT requires all base experts and LoRA execution on CPU"
|
|
)
|
|
if not bool(getattr(cfg, "kt_share_backward_bb", False)):
|
|
raise KTAMXConfigError("INT8 SFT requires kt_share_backward_bb=true")
|
|
if expert_weight_format == "fp8":
|
|
if kt_method != "AMXFP8_SFT":
|
|
raise KTAMXConfigError(
|
|
"kt_expert_weight_format='fp8' requires kt_backend='auto' or 'FP8'"
|
|
)
|
|
if full_weight_grad or train_mode != "lora" or lora_rank <= 0:
|
|
raise KTAMXConfigError(
|
|
"FP8 SFT supports frozen-base LoRA only; Full and Hybrid are not supported"
|
|
)
|
|
if requested_num_gpu_experts != 0 or use_lora_experts:
|
|
raise KTAMXConfigError(
|
|
"FP8 SFT requires all base experts and LoRA execution on CPU"
|
|
)
|
|
if not bool(getattr(cfg, "kt_share_backward_bb", False)):
|
|
raise KTAMXConfigError("FP8 SFT requires kt_share_backward_bb=true")
|
|
if getattr(cfg, "kt_weight_lifecycle", "persistent") != "persistent":
|
|
raise KTAMXConfigError("FP8 SFT requires persistent checkpoint weights")
|
|
if expert_weight_format == "rawint4":
|
|
if kt_method != RAWINT4_SFT_METHOD:
|
|
raise KTAMXConfigError("kt_expert_weight_format='rawint4' requires " "kt_backend='auto' or 'RAWINT4'")
|
|
if full_weight_grad or train_mode != "lora" or lora_rank != 8:
|
|
raise KTAMXConfigError("RAWINT4 SFT v1 supports frozen-base rank-8 LoRA only")
|
|
if float(lora_dropout) != 0.0:
|
|
raise KTAMXConfigError("RAWINT4 SFT v1 requires lora_dropout=0")
|
|
if requested_num_gpu_experts != 0 or use_lora_experts:
|
|
raise KTAMXConfigError("RAWINT4 SFT requires all base experts and LoRA execution on CPU")
|
|
if bool(getattr(cfg, "kt_share_backward_bb", False)):
|
|
raise KTAMXConfigError(
|
|
"RAWINT4 SFT computes backward directly from packed weights; "
|
|
"kt_share_backward_bb must be false"
|
|
)
|
|
if not force_fused_expert_lora:
|
|
raise KTAMXConfigError("RAWINT4 SFT requires kt_force_fused_expert_lora=true")
|
|
if getattr(cfg, "kt_weight_lifecycle", "persistent") != "persistent":
|
|
raise KTAMXConfigError("RAWINT4 SFT requires persistent checkpoint weights")
|
|
cpu_activation_retain = activation_policy.cpu == "retain"
|
|
reuse_checkpoint_forward = cpu_activation_retain and activation_policy.gpu == "recompute"
|
|
if cpu_activation_retain and (
|
|
not _supports_checkpoint_forward_reuse(full_weight_grad, lora_rank)
|
|
or kt_method
|
|
not in {
|
|
"AMXBF16_SFT",
|
|
"AMXFP8_SFT",
|
|
"INT8_SFT",
|
|
RAWINT4_SFT_METHOD,
|
|
MXFP4_SFT_METHOD,
|
|
}
|
|
or requested_num_gpu_experts != 0
|
|
or use_lora_experts
|
|
):
|
|
raise KTAMXConfigError(
|
|
"activation_policy.cpu=retain requires CPU-only AMXBF16 Full/LoRA "
|
|
"or frozen-base INT8/FP8/RAWINT4/MXFP4 LoRA; Hybrid, GPU-expert, "
|
|
"LoRA-expert, generic INT4, and SkipLoRA paths are not supported"
|
|
)
|
|
_validate_rawint4_activation_policy(expert_weight_format, activation_policy)
|
|
if is_rank_0:
|
|
logger.warning(
|
|
"KT activation policy: cpu=%s, gpu=%s, "
|
|
"checkpoint_forward_reuse=%s, share_cache_pool=%s, world_size=%d",
|
|
activation_policy.cpu,
|
|
activation_policy.gpu,
|
|
reuse_checkpoint_forward,
|
|
not cpu_activation_retain,
|
|
distributed_world_size,
|
|
)
|
|
uses_authoritative_optimizer_grads = _supports_authoritative_optimizer_grads(
|
|
kt_method,
|
|
requested_num_gpu_experts,
|
|
full_weight_grad=full_weight_grad,
|
|
lora_rank=lora_rank,
|
|
)
|
|
int8_runtime = None
|
|
int8_runtime_error = None
|
|
if expert_weight_format == "int8" and is_rank_0:
|
|
try:
|
|
int8_runtime = get_int8_runtime()
|
|
except RuntimeError as exc:
|
|
int8_runtime_error = KTAMXNotAvailableError(str(exc))
|
|
if expert_weight_format == "int8":
|
|
_sync_rank0_wrap_error(
|
|
int8_runtime_error,
|
|
context="selecting the INT8 SFT kernel",
|
|
rank=distributed_rank,
|
|
world_size=distributed_world_size,
|
|
)
|
|
if int8_runtime is not None:
|
|
logger.info(
|
|
"KT INT8 SFT dispatch: configured_backend=%s, logical_backend=%s, "
|
|
"cpu_variant=%s, effective_kernel=%s, weight_layout=%s",
|
|
kt_backend,
|
|
INT8_BACKEND,
|
|
int8_runtime.cpu_variant,
|
|
int8_runtime.kernel,
|
|
int8_runtime.weight_layout,
|
|
)
|
|
fp8_runtime = None
|
|
fp8_runtime_error = None
|
|
if expert_weight_format == "fp8" and is_rank_0:
|
|
try:
|
|
fp8_runtime = get_fp8_runtime()
|
|
except RuntimeError as exc:
|
|
fp8_runtime_error = KTAMXNotAvailableError(str(exc))
|
|
if expert_weight_format == "fp8":
|
|
_sync_rank0_wrap_error(
|
|
fp8_runtime_error,
|
|
context="selecting the native FP8 SFT kernel",
|
|
rank=distributed_rank,
|
|
world_size=distributed_world_size,
|
|
)
|
|
if fp8_runtime is not None:
|
|
logger.info(
|
|
"KT FP8 SFT dispatch: configured_backend=%s, logical_backend=%s, "
|
|
"cpu_variant=%s, effective_kernel=%s, weight_layout=%s",
|
|
kt_backend,
|
|
FP8_BACKEND,
|
|
fp8_runtime.cpu_variant,
|
|
fp8_runtime.kernel,
|
|
fp8_runtime.weight_layout,
|
|
)
|
|
|
|
rawint4_runtime = None
|
|
rawint4_runtime_error = None
|
|
rawint4_contract = None
|
|
if expert_weight_format == "rawint4":
|
|
try:
|
|
rawint4_contract = get_rawint4_checkpoint_contract(model.config)
|
|
except ValueError as exc:
|
|
raise KTAMXConfigError(str(exc)) from exc
|
|
if expert_weight_format == "rawint4" and is_rank_0:
|
|
try:
|
|
rawint4_runtime = get_rawint4_runtime()
|
|
except RuntimeError as exc:
|
|
rawint4_runtime_error = KTAMXNotAvailableError(str(exc))
|
|
if expert_weight_format == "rawint4":
|
|
_sync_rank0_wrap_error(
|
|
rawint4_runtime_error,
|
|
context="selecting the RAWINT4 SFT kernel",
|
|
rank=distributed_rank,
|
|
world_size=distributed_world_size,
|
|
)
|
|
if rawint4_runtime is not None:
|
|
logger.info(
|
|
"KT RAWINT4 SFT dispatch: configured_backend=%s, logical_backend=%s, "
|
|
"cpu_variant=%s, effective_kernel=%s, weight_layout=%s",
|
|
kt_backend,
|
|
RAWINT4_BACKEND,
|
|
rawint4_runtime.cpu_variant,
|
|
rawint4_runtime.kernel,
|
|
rawint4_runtime.weight_layout,
|
|
)
|
|
|
|
threadpool_count = getattr(cfg, "kt_threadpool_count", 1) if getattr(cfg, "kt_tp_enabled", False) else 1
|
|
fp8_block_size = None
|
|
if expert_weight_format == "fp8":
|
|
fp8_block_size = _native_fp8_block_size(model.config)
|
|
if hidden_size % 128 or moe_config.intermediate_size % 128:
|
|
raise KTAMXConfigError(
|
|
"FP8 SFT requires hidden and routed intermediate dimensions divisible by 128"
|
|
)
|
|
if (
|
|
threadpool_count < 1
|
|
or moe_config.intermediate_size % threadpool_count
|
|
or (moe_config.intermediate_size // threadpool_count) % 128
|
|
):
|
|
raise KTAMXConfigError(
|
|
"FP8 SFT requires each TP intermediate slice divisible by 128; "
|
|
f"intermediate_size={moe_config.intermediate_size}, "
|
|
f"threadpool_count={threadpool_count}"
|
|
)
|
|
if expert_weight_format == "rawint4":
|
|
if threadpool_count not in {1, 2}:
|
|
raise KTAMXConfigError(f"RAWINT4 SFT v1 supports TP1/TP2 only, got {threadpool_count}")
|
|
if hidden_size % RAWINT4_GROUP_SIZE or moe_config.intermediate_size % RAWINT4_GROUP_SIZE:
|
|
raise KTAMXConfigError(
|
|
"RAWINT4 SFT requires hidden and routed intermediate dimensions " f"divisible by {RAWINT4_GROUP_SIZE}"
|
|
)
|
|
if (
|
|
moe_config.intermediate_size % threadpool_count
|
|
or (moe_config.intermediate_size // threadpool_count) % RAWINT4_GROUP_SIZE
|
|
):
|
|
raise KTAMXConfigError(
|
|
"RAWINT4 SFT requires every TP intermediate slice divisible by "
|
|
f"{RAWINT4_GROUP_SIZE}; intermediate_size={moe_config.intermediate_size}, "
|
|
f"threadpool_count={threadpool_count}"
|
|
)
|
|
|
|
kt_weight_path = getattr(cfg, "kt_weight_path", None)
|
|
# Native FP8 and RAWINT4 use kt_weight_path as raw-checkpoint provenance,
|
|
# not as a pre-packed .kt directory.
|
|
use_kt_weight_path = kt_weight_path is not None and expert_weight_format not in {"fp8", "rawint4"}
|
|
if use_kt_weight_path:
|
|
logger.info(
|
|
"Loading %s weights from kt_weight_path: %s",
|
|
expert_weight_format or "pre-quantized",
|
|
kt_weight_path,
|
|
)
|
|
|
|
checkpoint_files = getattr(cfg, "kt_checkpoint_files", None)
|
|
sharded_metadata = getattr(cfg, "kt_sharded_metadata", None)
|
|
|
|
# When kt_expert_checkpoint_path is set, always resolve from it (overrides any existing
|
|
# checkpoint_files which may come from AttnOnlyBf16 and lack expert weights).
|
|
kt_expert_checkpoint_path = getattr(cfg, "kt_expert_checkpoint_path", None)
|
|
full_weight_checkpoint = resolve_full_weight_checkpoint(kt_expert_checkpoint_path)
|
|
use_full_weight_checkpoint = full_weight_checkpoint is not None
|
|
if use_full_weight_checkpoint and use_kt_weight_path:
|
|
raise KTAMXConfigError(
|
|
"A KT Full checkpoint and kt_weight_path were both selected. "
|
|
"Choose exactly one authoritative expert weight source."
|
|
)
|
|
if use_full_weight_checkpoint:
|
|
logger.info(f"Loading expert weights from KT Full checkpoint: {full_weight_checkpoint}")
|
|
elif kt_expert_checkpoint_path:
|
|
logger.info(f"Resolving expert checkpoint files from kt_expert_checkpoint_path={kt_expert_checkpoint_path!r}")
|
|
resolved_files, resolved_meta = _resolve_checkpoint_files(model_name_or_path=kt_expert_checkpoint_path)
|
|
if resolved_files and all(f.endswith(".safetensors") for f in resolved_files):
|
|
checkpoint_files = resolved_files
|
|
sharded_metadata = resolved_meta
|
|
cfg.kt_checkpoint_files = checkpoint_files
|
|
cfg.kt_sharded_metadata = sharded_metadata
|
|
logger.info(f"Resolved {len(checkpoint_files)} checkpoint files from kt_expert_checkpoint_path")
|
|
else:
|
|
logger.warning(
|
|
f"Failed to resolve checkpoint files from kt_expert_checkpoint_path={kt_expert_checkpoint_path!r}"
|
|
)
|
|
|
|
if expert_weight_format in {"fp8", "rawint4"} and not checkpoint_files and kt_weight_path:
|
|
logger.info(
|
|
"Resolving native %s checkpoint files from kt_weight_path=%r",
|
|
expert_weight_format.upper(),
|
|
kt_weight_path,
|
|
)
|
|
resolved_files, resolved_meta = _resolve_raw_safetensors_checkpoint_files(kt_weight_path)
|
|
if resolved_files:
|
|
checkpoint_files = resolved_files
|
|
sharded_metadata = resolved_meta
|
|
cfg.kt_checkpoint_files = checkpoint_files
|
|
cfg.kt_sharded_metadata = sharded_metadata
|
|
|
|
use_checkpoint_files = bool(checkpoint_files) and not use_kt_weight_path and not use_full_weight_checkpoint
|
|
if expert_weight_format == "int8":
|
|
if not use_kt_weight_path:
|
|
raise KTAMXConfigError(
|
|
"INT8 SFT requires kt_weight_path with pre-quantized .kt weights"
|
|
)
|
|
if use_full_weight_checkpoint or kt_expert_checkpoint_path:
|
|
raise KTAMXConfigError(
|
|
"INT8 SFT does not support Full checkpoints or online expert conversion"
|
|
)
|
|
if expert_weight_format == "fp8":
|
|
if use_full_weight_checkpoint:
|
|
raise KTAMXConfigError("FP8 SFT does not support KT Full checkpoints")
|
|
if expert_weight_format == "rawint4":
|
|
if use_full_weight_checkpoint or kt_expert_checkpoint_path:
|
|
raise KTAMXConfigError(
|
|
"RAWINT4 SFT reads packed experts from the source checkpoint and "
|
|
"does not support KT Full or alternate expert checkpoints"
|
|
)
|
|
|
|
logger.debug(
|
|
f"Weight source: kt_weight_path={kt_weight_path!r}, "
|
|
f"kt_expert_checkpoint_path={kt_expert_checkpoint_path!r}, "
|
|
f"full_weight_checkpoint={full_weight_checkpoint!r}, "
|
|
f"checkpoint_files count={len(checkpoint_files) if checkpoint_files else 0}, "
|
|
f"use_kt_weight_path={use_kt_weight_path}, use_full_weight_checkpoint={use_full_weight_checkpoint}, "
|
|
f"use_checkpoint_files={use_checkpoint_files}"
|
|
)
|
|
|
|
if use_full_weight_checkpoint:
|
|
logger.info("Loading expert weights from a KT Full checkpoint.")
|
|
elif use_checkpoint_files:
|
|
logger.info("Loading expert weights from checkpoint files (online conversion).")
|
|
elif use_kt_weight_path and bool(checkpoint_files):
|
|
logger.info("BF16 checkpoint files available for backward gradient computation.")
|
|
elif (
|
|
not use_kt_weight_path
|
|
and not use_full_weight_checkpoint
|
|
and bool(getattr(cfg, "kt_skip_expert_loading", False))
|
|
):
|
|
# If HF expert weights were skipped during `from_pretrained`, we must source expert weights externally.
|
|
model_name_or_path = getattr(getattr(model, "config", None), "name_or_path", None)
|
|
if model_name_or_path:
|
|
resolved_files, resolved_meta = _resolve_checkpoint_files(model_name_or_path=model_name_or_path)
|
|
if resolved_files and all(f.endswith(".safetensors") for f in resolved_files):
|
|
checkpoint_files = resolved_files
|
|
sharded_metadata = resolved_meta
|
|
cfg.kt_checkpoint_files = checkpoint_files
|
|
cfg.kt_sharded_metadata = sharded_metadata
|
|
use_checkpoint_files = True
|
|
logger.info("KT skip_expert_loading enabled; using checkpoint files for online expert loading.")
|
|
|
|
if not use_checkpoint_files:
|
|
raise KTAMXConfigError(
|
|
"KT skip_expert_loading is enabled but no `kt_weight_path` was provided and no safetensors checkpoint "
|
|
"files could be resolved for on-the-fly expert loading."
|
|
)
|
|
|
|
if expert_weight_format == "fp8" and not use_checkpoint_files:
|
|
raise KTAMXConfigError(
|
|
"FP8 SFT requires raw safetensors checkpoint files. Point "
|
|
"kt_weight_path at model_name_or_path or set kt_expert_checkpoint_path."
|
|
)
|
|
if expert_weight_format == "rawint4" and not use_checkpoint_files:
|
|
raise KTAMXConfigError(
|
|
"RAWINT4 SFT requires raw compressed-tensors safetensors checkpoint "
|
|
"files. Point kt_weight_path at model_name_or_path."
|
|
)
|
|
|
|
model_container, layers = _get_model_container_and_layers(model, purpose="wrapping")
|
|
logger.info(f"Total layers={len(layers)}, is_rank_0={is_rank_0}")
|
|
|
|
from .arch import detect_fused_experts as _detect_fused
|
|
|
|
expert_layer_indices = [
|
|
layer_idx
|
|
for layer_idx, layer in enumerate(layers)
|
|
if get_moe_module(layer, moe_config) is not None
|
|
]
|
|
ephemeral_store = None
|
|
ephemeral_requested = (
|
|
getattr(cfg, "kt_weight_lifecycle", "persistent") == "ephemeral"
|
|
)
|
|
ephemeral_open_error = None
|
|
if is_rank_0 and ephemeral_requested:
|
|
try:
|
|
from .ephemeral import EphemeralKTWeightStore
|
|
|
|
ephemeral_store = EphemeralKTWeightStore.open(
|
|
kt_weight_path,
|
|
layer_indices=expert_layer_indices,
|
|
numa_count=threadpool_count,
|
|
expert_num=moe_config.expert_num,
|
|
hidden_size=hidden_size,
|
|
intermediate_size=moe_config.intermediate_size,
|
|
)
|
|
except BaseException as exc:
|
|
ephemeral_open_error = exc
|
|
if ephemeral_requested:
|
|
_sync_rank0_wrap_error(
|
|
ephemeral_open_error,
|
|
context="opening ephemeral INT8 weights",
|
|
rank=distributed_rank,
|
|
world_size=distributed_world_size,
|
|
)
|
|
|
|
persistent_manifest_error = None
|
|
if (
|
|
is_rank_0
|
|
and expert_weight_format == "int8"
|
|
and not ephemeral_requested
|
|
):
|
|
try:
|
|
from .weight_manifest import validate_persistent_int8_weights
|
|
|
|
validated_manifest = validate_persistent_int8_weights(
|
|
kt_weight_path,
|
|
layer_indices=expert_layer_indices,
|
|
numa_count=threadpool_count,
|
|
expert_num=moe_config.expert_num,
|
|
hidden_size=hidden_size,
|
|
intermediate_size=moe_config.intermediate_size,
|
|
)
|
|
logger.info(
|
|
"Validated persistent INT8 weights: manifest=%s, schema=%d%s, "
|
|
"layout=%s, layers=%d, files=%d, bytes=%d",
|
|
validated_manifest.path,
|
|
validated_manifest.schema_version,
|
|
" (legacy compatibility)" if validated_manifest.is_legacy else "",
|
|
validated_manifest.layout,
|
|
len(validated_manifest.layer_indices),
|
|
validated_manifest.file_count,
|
|
validated_manifest.size_bytes,
|
|
)
|
|
except BaseException as exc:
|
|
persistent_manifest_error = exc
|
|
if expert_weight_format == "int8" and not ephemeral_requested:
|
|
_sync_rank0_wrap_error(
|
|
persistent_manifest_error,
|
|
context="validating persistent INT8 weights",
|
|
rank=distributed_rank,
|
|
world_size=distributed_world_size,
|
|
)
|
|
|
|
for layer_idx, layer in enumerate(layers):
|
|
moe_module = get_moe_module(layer, moe_config)
|
|
if moe_module is None:
|
|
continue
|
|
|
|
_layer_experts = getattr(moe_module, moe_config.experts_attr, None)
|
|
_layer_is_fused = _detect_fused(_layer_experts)
|
|
if expert_weight_format in {"int8", "fp8", "rawint4", "mxfp4"} and not _layer_is_fused and not force_fused_expert_lora:
|
|
raise KTAMXConfigError(
|
|
f"{expert_weight_format.upper()} LoRA with non-fused runtime experts requires "
|
|
"kt_force_fused_expert_lora=true"
|
|
)
|
|
_use_fused_expert_lora = _layer_is_fused or force_fused_expert_lora
|
|
|
|
logger.debug(
|
|
"Wrapping MoE layer %s (method=%s, fused=%s, force_fused_lora=%s)",
|
|
layer_idx,
|
|
kt_method,
|
|
_layer_is_fused,
|
|
force_fused_expert_lora,
|
|
)
|
|
|
|
# Only rank 0 loads weights and initializes KT kernel
|
|
gate_proj, up_proj, down_proj = None, None, None
|
|
block_fp8_weights = None
|
|
rawint4_weights = None
|
|
wrapper = None
|
|
weight_source_error = None
|
|
|
|
if is_rank_0:
|
|
try:
|
|
# Get block_size from quantization_config if available (for legacy FP8 dequant)
|
|
_quant_cfg = getattr(model.config, "quantization_config", None)
|
|
_block_size = None
|
|
if _quant_cfg is not None:
|
|
_block_size = (
|
|
_quant_cfg.get("weight_block_size")
|
|
if isinstance(_quant_cfg, dict)
|
|
else getattr(_quant_cfg, "weight_block_size", None)
|
|
)
|
|
|
|
if use_full_weight_checkpoint:
|
|
expected_shapes = {
|
|
"gate_proj": (
|
|
int(moe_config.expert_num),
|
|
int(moe_config.intermediate_size),
|
|
int(hidden_size),
|
|
),
|
|
"up_proj": (
|
|
int(moe_config.expert_num),
|
|
int(moe_config.intermediate_size),
|
|
int(hidden_size),
|
|
),
|
|
"down_proj": (
|
|
int(moe_config.expert_num),
|
|
int(hidden_size),
|
|
int(moe_config.intermediate_size),
|
|
),
|
|
}
|
|
gate_proj, up_proj, down_proj = load_full_weight_layer(
|
|
full_weight_checkpoint,
|
|
layer_idx=layer_idx,
|
|
expected_shapes=expected_shapes,
|
|
)
|
|
elif expert_weight_format == "mxfp4":
|
|
get_mxfp4_runtime()
|
|
if moe_config.router_type != "precomputed":
|
|
raise KTAMXConfigError("MXFP4 SFT currently supports DeepSeek V4 routed experts only")
|
|
elif use_kt_weight_path:
|
|
logger.debug(f"Layer {layer_idx}: forward + backward from kt_weight_path (.kt files)")
|
|
elif expert_weight_format == "fp8":
|
|
layers_prefix = _get_layers_prefix(model.config)
|
|
block_fp8_weights = load_block_fp8_experts_from_checkpoint_files(
|
|
checkpoint_files=checkpoint_files,
|
|
sharded_metadata=sharded_metadata,
|
|
layers_prefix=layers_prefix,
|
|
moe_config=moe_config,
|
|
layer_idx=layer_idx,
|
|
hidden_size=hidden_size,
|
|
block_size=fp8_block_size,
|
|
)
|
|
elif expert_weight_format == "rawint4":
|
|
layers_prefix = _get_layers_prefix(model.config)
|
|
rawint4_weights = load_rawint4_experts_from_checkpoint_files(
|
|
checkpoint_files=checkpoint_files,
|
|
sharded_metadata=sharded_metadata,
|
|
layers_prefix=layers_prefix,
|
|
moe_config=moe_config,
|
|
layer_idx=layer_idx,
|
|
hidden_size=hidden_size,
|
|
group_size=RAWINT4_GROUP_SIZE,
|
|
)
|
|
elif use_checkpoint_files:
|
|
layers_prefix = _get_layers_prefix(model.config)
|
|
gate_proj, up_proj, down_proj = load_experts_from_checkpoint_files(
|
|
checkpoint_files=checkpoint_files,
|
|
sharded_metadata=sharded_metadata,
|
|
layers_prefix=layers_prefix,
|
|
moe_config=moe_config,
|
|
layer_idx=layer_idx,
|
|
block_size=_block_size,
|
|
)
|
|
else:
|
|
gate_proj, up_proj, down_proj = extract_moe_weights(moe_module, moe_config)
|
|
gate_proj = gate_proj.cpu().to(torch.bfloat16).contiguous()
|
|
up_proj = up_proj.cpu().to(torch.bfloat16).contiguous()
|
|
down_proj = down_proj.cpu().to(torch.bfloat16).contiguous()
|
|
except BaseException as exc:
|
|
weight_source_error = exc
|
|
|
|
_sync_rank0_wrap_error(
|
|
weight_source_error,
|
|
context=f"resolving expert weights for layer {layer_idx}",
|
|
rank=distributed_rank,
|
|
world_size=distributed_world_size,
|
|
)
|
|
|
|
chunked_prefill_size = getattr(cfg, "kt_model_max_length", None)
|
|
if chunked_prefill_size is None:
|
|
chunked_prefill_size = getattr(model.config, "max_position_embeddings", 4096)
|
|
# Rank 0 receives the concatenation of every rank's local rows. Model
|
|
# configs are homogeneous across ranks, so the sum of local maxima is
|
|
# the per-rank capacity multiplied by world size.
|
|
rank0_chunked_prefill_size = int(chunked_prefill_size) * distributed_world_size
|
|
|
|
# Only rank 0 creates KTMoEWrapper and loads weights
|
|
construct_error = None
|
|
if is_rank_0:
|
|
try:
|
|
wrapper = KTMoEWrapper(
|
|
layer_idx=layer_idx,
|
|
num_experts=moe_config.expert_num,
|
|
num_experts_per_tok=moe_config.num_experts_per_tok,
|
|
hidden_size=hidden_size,
|
|
moe_intermediate_size=moe_config.intermediate_size,
|
|
gpu_experts_mask=None,
|
|
num_gpu_experts=0,
|
|
cpuinfer_threads=getattr(cfg, "kt_num_threads", 1),
|
|
threadpool_count=threadpool_count,
|
|
weight_path=kt_weight_path or "",
|
|
chunked_prefill_size=rank0_chunked_prefill_size,
|
|
method=kt_method,
|
|
mode="sft",
|
|
lora_rank=lora_rank,
|
|
lora_alpha=lora_alpha,
|
|
lora_dropout=lora_dropout,
|
|
max_cache_depth=getattr(cfg, "kt_max_cache_depth", 2),
|
|
group_size=32 if expert_weight_format in {"rawint4", "mxfp4"} else 128,
|
|
zero_point=expert_weight_format not in {"rawint4", "mxfp4"},
|
|
swiglu_limit=float(getattr(_text_cfg, "swiglu_limit", 0.0))
|
|
if expert_weight_format in {"mxfp4", "rawint4"}
|
|
else 0.0,
|
|
full_weight_grad=full_weight_grad,
|
|
)
|
|
except BaseException as exc:
|
|
try:
|
|
if ephemeral_store is not None:
|
|
ephemeral_store.cleanup()
|
|
except BaseException:
|
|
logger.exception(
|
|
"Cleanup failed after constructing ephemeral INT8 layer %s",
|
|
layer_idx,
|
|
)
|
|
construct_error = exc
|
|
|
|
_sync_rank0_wrap_error(
|
|
construct_error,
|
|
context=f"constructing KT layer {layer_idx}",
|
|
rank=distributed_rank,
|
|
world_size=distributed_world_size,
|
|
)
|
|
|
|
load_error = None
|
|
if is_rank_0:
|
|
try:
|
|
# The current SFT wrapping path routes all experts through KT even
|
|
# when the loading config requested GPU experts. Preserve that
|
|
# configuration's legacy gradient lifecycle until the hybrid
|
|
# routed-expert path supports authoritative buffers end to end.
|
|
wrapper._uses_authoritative_optimizer_grads = uses_authoritative_optimizer_grads
|
|
|
|
# These flags are consumed while the C++ config is built.
|
|
wrapper.share_backward_bb = cfg.kt_share_backward_bb
|
|
wrapper.reuse_checkpoint_forward = reuse_checkpoint_forward
|
|
wrapper.activation_policy = activation_policy
|
|
wrapper.share_cache_pool = not cpu_activation_retain
|
|
|
|
physical_to_logical_map = torch.arange(
|
|
moe_config.expert_num,
|
|
dtype=torch.int64,
|
|
device="cpu",
|
|
)
|
|
|
|
if expert_weight_format == "fp8":
|
|
logger.debug(
|
|
"Layer %s: packing raw per-expert FP8 checkpoint tensors",
|
|
layer_idx,
|
|
)
|
|
wrapper.load_block_fp8_weights(
|
|
block_fp8_weights,
|
|
physical_to_logical_map,
|
|
)
|
|
block_fp8_weights = None
|
|
elif expert_weight_format == "rawint4":
|
|
logger.debug(
|
|
"Layer %s: loading packed RAWINT4 checkpoint tensors",
|
|
layer_idx,
|
|
)
|
|
wrapper.load_rawint4_weights(
|
|
rawint4_weights,
|
|
physical_to_logical_map,
|
|
)
|
|
rawint4_weights = None
|
|
elif use_kt_weight_path:
|
|
logger.debug(
|
|
f"Layer {layer_idx}: calling wrapper.load_weights() "
|
|
"(C++ direct .kt load)"
|
|
)
|
|
wrapper.load_weights(physical_to_logical_map)
|
|
if ephemeral_store is not None:
|
|
# load_weights() returns only after C++ copied every file
|
|
# into its owned BufferB storage.
|
|
ephemeral_store.consume_layer(layer_idx)
|
|
else:
|
|
logger.debug(
|
|
f"Layer {layer_idx}: calling wrapper.load_weights_from_tensors() "
|
|
f"(BF16 tensor path, gate_proj shape={gate_proj.shape if gate_proj is not None else None})"
|
|
)
|
|
wrapper.load_weights_from_tensors(
|
|
gate_proj=gate_proj,
|
|
up_proj=up_proj,
|
|
down_proj=down_proj,
|
|
physical_to_logical_map_cpu=physical_to_logical_map,
|
|
)
|
|
|
|
if full_weight_grad:
|
|
wrapper.init_full_weight_grad_buffers(
|
|
gate_proj=wrapper.gate_proj if wrapper.gate_proj is not None else gate_proj,
|
|
up_proj=wrapper.up_proj if wrapper.up_proj is not None else up_proj,
|
|
down_proj=wrapper.down_proj if wrapper.down_proj is not None else down_proj,
|
|
)
|
|
else:
|
|
wrapper.gate_proj = None
|
|
wrapper.up_proj = None
|
|
wrapper.down_proj = None
|
|
except BaseException as exc:
|
|
if ephemeral_store is not None:
|
|
try:
|
|
ephemeral_store.cleanup()
|
|
except BaseException:
|
|
logger.exception(
|
|
"Cleanup failed after loading ephemeral INT8 layer %s",
|
|
layer_idx,
|
|
)
|
|
load_error = exc
|
|
|
|
_sync_rank0_wrap_error(
|
|
load_error,
|
|
context=f"loading KT layer {layer_idx}",
|
|
rank=distributed_rank,
|
|
world_size=distributed_world_size,
|
|
)
|
|
|
|
# Create LoRA Experts if enabled
|
|
lora_experts = None
|
|
if use_lora_experts:
|
|
lora_experts = LoRAExperts(
|
|
num_experts=lora_expert_num,
|
|
hidden_size=hidden_size,
|
|
intermediate_size=lora_expert_intermediate_size,
|
|
device="cuda",
|
|
dtype=torch.bfloat16,
|
|
)
|
|
|
|
layer_wrapper_cls = KTRoutedExpertsWrapper if moe_config.router_type == "precomputed" else KTMoELayerWrapper
|
|
layer_wrapper = layer_wrapper_cls(
|
|
original_moe=moe_module,
|
|
wrapper=wrapper,
|
|
lora_params=None,
|
|
moe_config=moe_config,
|
|
hidden_size=hidden_size,
|
|
layer_idx=layer_idx,
|
|
lora_experts=lora_experts,
|
|
full_weight_grad=full_weight_grad,
|
|
uses_authoritative_optimizer_grads=uses_authoritative_optimizer_grads,
|
|
activation_policy=activation_policy,
|
|
)
|
|
layer_wrapper._fused_experts = _layer_is_fused
|
|
layer_wrapper._use_fused_expert_lora = _use_fused_expert_lora
|
|
layer_wrapper._force_fused_expert_lora = force_fused_expert_lora
|
|
layer_wrapper._lora_rank = lora_rank
|
|
layer_wrapper._lora_alpha = float(lora_alpha)
|
|
layer_wrapper._kt_owner_rank = 0
|
|
layer_wrapper._kt_world_size_at_wrap = distributed_world_size
|
|
layer_wrapper._kt_expert_weight_format = expert_weight_format or "bf16"
|
|
if rawint4_contract is not None:
|
|
layer_wrapper._kt_rawint4_quantization = {
|
|
"bits": rawint4_contract.bits,
|
|
"group_size": rawint4_contract.group_size,
|
|
"signed": rawint4_contract.signed,
|
|
"symmetric": rawint4_contract.symmetric,
|
|
"zero_point": rawint4_contract.zero_point,
|
|
"format": rawint4_contract.format,
|
|
"weight_layout": RAWINT4_WEIGHT_LAYOUT,
|
|
}
|
|
|
|
if moe_config.router_type != "precomputed":
|
|
setattr(layer, moe_config.moe_layer_attr, layer_wrapper)
|
|
# Base weights have been copied into the C++ kernel's internal BufferB format.
|
|
# In full_weight_grad mode, the authoritative copies are gate_proj_buf etc.
|
|
# Always release local references to save ~1 GB/layer.
|
|
del gate_proj, up_proj, down_proj, block_fp8_weights, rawint4_weights
|
|
|
|
wrappers.append(layer_wrapper)
|
|
moe_layer_count += 1
|
|
|
|
# Replace original expert weights with zero-storage placeholders.
|
|
# Experts remain in the model tree (via wrapper.experts) so PEFT can discover them.
|
|
# Rank 0 already copied weights to C++ kernel via load_weights_from_tensors.
|
|
# gate_proj_buf serves as the authoritative copy in full_weight_grad mode.
|
|
_clear_original_expert_weights(
|
|
moe_module,
|
|
moe_config,
|
|
full_weight_grad=full_weight_grad,
|
|
empty_placeholders=expert_weight_format in {"int8", "fp8", "rawint4", "mxfp4"},
|
|
)
|
|
if moe_config.router_type == "precomputed":
|
|
setattr(moe_module, moe_config.experts_attr, layer_wrapper)
|
|
|
|
ephemeral_finish_error = None
|
|
if ephemeral_store is not None:
|
|
try:
|
|
ephemeral_store.finish()
|
|
except BaseException as exc:
|
|
try:
|
|
ephemeral_store.cleanup()
|
|
except BaseException:
|
|
logger.exception("Cleanup failed while finishing ephemeral INT8 weights")
|
|
ephemeral_finish_error = exc
|
|
if ephemeral_requested:
|
|
_sync_rank0_wrap_error(
|
|
ephemeral_finish_error,
|
|
context="finishing ephemeral INT8 weights",
|
|
rank=distributed_rank,
|
|
world_size=distributed_world_size,
|
|
)
|
|
|
|
logger.info(f"Wrapped {moe_layer_count} MoE layers with KTMoEWrapper")
|
|
model._kt_expert_weight_format = expert_weight_format or "bf16"
|
|
if rawint4_contract is not None:
|
|
model._kt_rawint4_quantization = {
|
|
"bits": rawint4_contract.bits,
|
|
"group_size": rawint4_contract.group_size,
|
|
"signed": rawint4_contract.signed,
|
|
"symmetric": rawint4_contract.symmetric,
|
|
"zero_point": rawint4_contract.zero_point,
|
|
"format": rawint4_contract.format,
|
|
"weight_layout": RAWINT4_WEIGHT_LAYOUT,
|
|
}
|
|
|
|
# Link wrappers for async backward repack (higher layer triggers repack for lower)
|
|
for i in range(1, len(wrappers)):
|
|
if wrappers[i].wrapper is not None and wrappers[i - 1].wrapper is not None:
|
|
wrappers[i].wrapper._next_backward_wrapper = wrappers[i - 1].wrapper
|
|
if wrappers and wrappers[0].wrapper is not None:
|
|
wrappers[0].wrapper._next_backward_wrapper = None
|
|
|
|
gc.collect()
|
|
return wrappers
|
|
|
|
|
|
# =============================================================================
|
|
# Plugin builder
|
|
# =============================================================================
|
|
|
|
|
|
def _build_kt_plugin_from_args(model_args: Any, finetuning_args: Any | None = None):
|
|
"""
|
|
Build a KTransformersPlugin from model_args and optional finetuning_args.
|
|
|
|
Imported here to avoid circular dependency --- callers that need the plugin
|
|
class should import it from the appropriate dataclasses module.
|
|
"""
|
|
from .config import KTConfig
|
|
from accelerate.utils.dataclasses import KTransformersPlugin
|
|
|
|
# Map LlamaFactory finetuning_type to kt_train_mode
|
|
finetuning_type = getattr(finetuning_args, "finetuning_type", None) if finetuning_args else None
|
|
kt_train_mode_map = {
|
|
"full": "full",
|
|
"freeze": "hybrid",
|
|
"lora": "lora",
|
|
"galore": "full",
|
|
"badam": "full",
|
|
}
|
|
kt_train_mode = kt_train_mode_map.get(finetuning_type, None) if finetuning_type else None
|
|
|
|
configured_lora_rank = getattr(finetuning_args, "lora_rank", None) if finetuning_args else None
|
|
configured_lora_alpha = getattr(finetuning_args, "lora_alpha", None) if finetuning_args else None
|
|
configured_lora_dropout = getattr(finetuning_args, "lora_dropout", None) if finetuning_args else None
|
|
if kt_train_mode == "full":
|
|
configured_lora_rank = None
|
|
configured_lora_alpha = None
|
|
configured_lora_dropout = None
|
|
|
|
kt_config = KTConfig(
|
|
kt_backend=getattr(model_args, "kt_backend", None),
|
|
kt_num_threads=getattr(model_args, "kt_num_threads", None),
|
|
kt_tp_enabled=getattr(model_args, "kt_tp_enabled", None),
|
|
kt_threadpool_count=getattr(model_args, "kt_threadpool_count", None),
|
|
kt_max_cache_depth=getattr(model_args, "kt_max_cache_depth", None),
|
|
kt_num_gpu_experts=getattr(model_args, "kt_num_gpu_experts", None),
|
|
kt_weight_path=getattr(model_args, "kt_weight_path", None),
|
|
kt_expert_weight_format=getattr(model_args, "kt_expert_weight_format", None),
|
|
kt_weight_lifecycle=getattr(model_args, "kt_weight_lifecycle", None),
|
|
kt_expert_checkpoint_path=getattr(model_args, "kt_expert_checkpoint_path", None),
|
|
kt_force_fused_expert_lora=getattr(model_args, "kt_force_fused_expert_lora", None),
|
|
kt_use_lora_experts=getattr(model_args, "kt_use_lora_experts", None),
|
|
kt_lora_expert_num=getattr(model_args, "kt_lora_expert_num", None),
|
|
kt_lora_expert_intermediate_size=getattr(model_args, "kt_lora_expert_intermediate_size", None),
|
|
kt_lora_rank=configured_lora_rank,
|
|
kt_lora_alpha=configured_lora_alpha,
|
|
kt_lora_dropout=configured_lora_dropout,
|
|
kt_model_max_length=getattr(model_args, "model_max_length", None),
|
|
kt_train_mode=kt_train_mode,
|
|
kt_activation_policy=getattr(model_args, "activation_policy", None),
|
|
)
|
|
return KTransformersPlugin(enabled=True, kt_config=kt_config)
|
|
|
|
|
|
def get_kt_loading_kwargs(
|
|
config,
|
|
kt_plugin,
|
|
torch_dtype: torch.dtype | str | None = torch.bfloat16,
|
|
trust_remote_code: bool | None = None,
|
|
token: str | None = None,
|
|
) -> dict[str, Any]:
|
|
"""Get kwargs for AutoModel.from_pretrained() for KT loading."""
|
|
kwargs: dict[str, Any] = {
|
|
"config": config,
|
|
"torch_dtype": torch_dtype,
|
|
"device_map": "cpu",
|
|
"low_cpu_mem_usage": True,
|
|
}
|
|
if trust_remote_code is not None:
|
|
kwargs["trust_remote_code"] = trust_remote_code
|
|
if token is not None:
|
|
kwargs["token"] = token
|
|
return kwargs
|
|
|
|
|
|
def _resolve_checkpoint_files(
|
|
model_name_or_path: str,
|
|
cache_dir: str | None = None,
|
|
revision: str | None = None,
|
|
token: str | None = None,
|
|
trust_remote_code: bool | None = None,
|
|
) -> tuple[list[str] | None, dict | None]:
|
|
"""Resolve HF checkpoint files. Depends on transformers internals."""
|
|
try:
|
|
import inspect
|
|
|
|
from transformers.modeling_utils import _get_resolved_checkpoint_files
|
|
except Exception:
|
|
return None, None
|
|
try:
|
|
common = {
|
|
"pretrained_model_name_or_path": model_name_or_path,
|
|
"variant": None,
|
|
"gguf_file": None,
|
|
"use_safetensors": None,
|
|
"user_agent": {"file_type": "model", "framework": "pytorch"},
|
|
"is_remote_code": bool(trust_remote_code),
|
|
"transformers_explicit_filename": None,
|
|
}
|
|
if "download_kwargs" in inspect.signature(_get_resolved_checkpoint_files).parameters:
|
|
common["download_kwargs"] = {
|
|
"cache_dir": cache_dir,
|
|
"force_download": False,
|
|
"local_files_only": False,
|
|
"token": token,
|
|
"revision": revision or "main",
|
|
"subfolder": "",
|
|
}
|
|
checkpoint_files, sharded_metadata = _get_resolved_checkpoint_files(**common)
|
|
else:
|
|
checkpoint_files, sharded_metadata = _get_resolved_checkpoint_files(
|
|
**common,
|
|
subfolder="",
|
|
from_tf=False,
|
|
from_flax=False,
|
|
cache_dir=cache_dir,
|
|
force_download=False,
|
|
proxies=None,
|
|
local_files_only=False,
|
|
token=token,
|
|
revision=revision or "main",
|
|
commit_hash=None,
|
|
)
|
|
except Exception:
|
|
return None, None
|
|
return checkpoint_files, sharded_metadata
|
|
|
|
|
|
def load_kt_model(
|
|
config,
|
|
model_args: Any | None = None,
|
|
finetuning_args: Any | None = None,
|
|
kt_plugin=None,
|
|
model_name_or_path: str | None = None,
|
|
trust_remote_code: bool | None = None,
|
|
token: str | None = None,
|
|
torch_dtype: torch.dtype | str | None = torch.bfloat16,
|
|
**kwargs,
|
|
) -> nn.Module:
|
|
"""Load model with KTMoEWrapper backend."""
|
|
from .arch import (
|
|
get_moe_arch_config,
|
|
move_non_experts_to_gpu,
|
|
KTAMXConfigError,
|
|
)
|
|
|
|
if kt_plugin is None:
|
|
if model_args is None:
|
|
raise KTAMXConfigError("Either kt_plugin or model_args must be provided to load_kt_model().")
|
|
kt_plugin = _build_kt_plugin_from_args(model_args, finetuning_args)
|
|
|
|
if model_name_or_path is None and model_args is not None:
|
|
model_name_or_path = getattr(model_args, "model_name_or_path", None)
|
|
if model_name_or_path is None:
|
|
raise KTAMXConfigError("model_name_or_path is required to load_kt_model().")
|
|
|
|
if trust_remote_code is None and model_args is not None:
|
|
trust_remote_code = getattr(model_args, "trust_remote_code", None)
|
|
if token is None and model_args is not None:
|
|
token = getattr(model_args, "hf_hub_token", None)
|
|
cache_dir = getattr(model_args, "cache_dir", None) if model_args is not None else None
|
|
revision = getattr(model_args, "revision", None) if model_args is not None else None
|
|
|
|
_ = get_moe_arch_config(config)
|
|
|
|
logger.info("Loading model with KTMoEWrapper backend")
|
|
|
|
from transformers import AutoModelForCausalLM
|
|
from transformers.integrations.kt import set_kt_config, unset_kt_config
|
|
|
|
loading_kwargs = get_kt_loading_kwargs(
|
|
config,
|
|
kt_plugin,
|
|
torch_dtype=torch_dtype,
|
|
trust_remote_code=trust_remote_code,
|
|
token=token,
|
|
)
|
|
if model_args is not None:
|
|
for key in ("cache_dir", "revision"):
|
|
value = getattr(model_args, key, None)
|
|
if value is not None:
|
|
loading_kwargs[key] = value
|
|
loading_kwargs.update(kwargs)
|
|
|
|
cfg = _get_kt_config(kt_plugin)
|
|
auto_full_weight_checkpoint = resolve_full_weight_checkpoint(model_name_or_path)
|
|
if auto_full_weight_checkpoint is not None and getattr(cfg, "kt_expert_checkpoint_path", None) is None:
|
|
cfg.kt_expert_checkpoint_path = auto_full_weight_checkpoint
|
|
plugin_config = getattr(kt_plugin, "kt_config", None)
|
|
if isinstance(plugin_config, dict):
|
|
plugin_config["kt_expert_checkpoint_path"] = auto_full_weight_checkpoint
|
|
elif plugin_config is not None:
|
|
setattr(plugin_config, "kt_expert_checkpoint_path", auto_full_weight_checkpoint)
|
|
else:
|
|
setattr(kt_plugin, "kt_expert_checkpoint_path", auto_full_weight_checkpoint)
|
|
logger.info("Detected KT Full checkpoint in model directory: %s", auto_full_weight_checkpoint)
|
|
|
|
expert_weight_format = getattr(cfg, "kt_expert_weight_format", None)
|
|
native_fp8_experts = expert_weight_format == "fp8"
|
|
native_rawint4_experts = expert_weight_format == "rawint4"
|
|
native_checkpoint_experts = native_fp8_experts or native_rawint4_experts or expert_weight_format == "mxfp4"
|
|
if native_checkpoint_experts:
|
|
# Native quantized kt_weight_path is raw-checkpoint provenance; KT owns
|
|
# and loads the routed experts without materializing them in Transformers.
|
|
cfg.kt_skip_expert_loading = True
|
|
|
|
skip_expert_loading = getattr(cfg, "kt_skip_expert_loading", None)
|
|
needs_checkpoint_resolution = (
|
|
(native_checkpoint_experts and not getattr(cfg, "kt_checkpoint_files", None))
|
|
or skip_expert_loading is None
|
|
or (
|
|
bool(skip_expert_loading)
|
|
and not getattr(cfg, "kt_checkpoint_files", None)
|
|
and not getattr(cfg, "kt_weight_path", None)
|
|
)
|
|
)
|
|
if needs_checkpoint_resolution:
|
|
checkpoint_files, sharded_metadata = _resolve_checkpoint_files(
|
|
model_name_or_path=model_name_or_path,
|
|
cache_dir=cache_dir,
|
|
revision=revision,
|
|
token=token,
|
|
trust_remote_code=trust_remote_code,
|
|
)
|
|
if checkpoint_files and all(f.endswith(".safetensors") for f in checkpoint_files):
|
|
if native_checkpoint_experts or getattr(cfg, "kt_weight_path", None) is None:
|
|
cfg.kt_skip_expert_loading = True
|
|
else:
|
|
cfg.kt_skip_expert_loading = False
|
|
cfg.kt_checkpoint_files = checkpoint_files
|
|
cfg.kt_sharded_metadata = sharded_metadata
|
|
else:
|
|
if not native_checkpoint_experts:
|
|
cfg.kt_skip_expert_loading = False
|
|
|
|
if native_checkpoint_experts:
|
|
checkpoint_files = getattr(cfg, "kt_checkpoint_files", None)
|
|
if not checkpoint_files or not all(
|
|
str(path).endswith(".safetensors") for path in checkpoint_files
|
|
):
|
|
raise KTAMXConfigError(
|
|
f"native {expert_weight_format.upper()} SFT requires raw "
|
|
"safetensors checkpoint files; "
|
|
"no Transformers expert-materialization fallback is supported"
|
|
)
|
|
cfg.kt_skip_expert_loading = True
|
|
|
|
# Transformers consumes the resolved fields while from_pretrained runs;
|
|
# Accelerate deliberately keeps those fields opaque under plugin.kt_config.
|
|
set_kt_config(cfg)
|
|
try:
|
|
model = AutoModelForCausalLM.from_pretrained(model_name_or_path, **loading_kwargs)
|
|
finally:
|
|
unset_kt_config()
|
|
|
|
moe_config = get_moe_arch_config(config)
|
|
local_rank = int(os.environ.get("LOCAL_RANK", "0"))
|
|
move_non_experts_to_gpu(model, moe_config, device=f"cuda:{local_rank}")
|
|
|
|
existing_wrappers = getattr(model, "_kt_wrappers", None)
|
|
if existing_wrappers:
|
|
logger.info(f"MoE layers already wrapped ({len(existing_wrappers)} layers), skipping re-wrap")
|
|
wrappers = existing_wrappers
|
|
else:
|
|
wrappers = wrap_moe_layers_with_kt_wrapper(model, kt_plugin)
|
|
|
|
model._kt_wrappers = wrappers
|
|
model._kt_tp_enabled = bool(getattr(cfg, "kt_tp_enabled", False))
|
|
model._kt_use_lora_experts = bool(getattr(cfg, "kt_use_lora_experts", False))
|
|
model._kt_full_weight_grad = bool(getattr(cfg, "kt_full_weight_grad", False))
|
|
model._kt_train_mode = getattr(cfg, "kt_train_mode", "lora")
|
|
model._kt_expert_weight_format = getattr(cfg, "kt_expert_weight_format", None) or "bf16"
|
|
|
|
logger.info("Model loaded with KTMoEWrapper backend successfully")
|
|
return model
|