# 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 from .lora import LoRAExperts from .base import _supports_authoritative_optimizer_grads from .backend import FP8_BACKEND, INT8_BACKEND, get_fp8_runtime, get_int8_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, ) 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 _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_native_fp8_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 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 = config.num_hidden_layers 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["model.embed_tokens"] = device device_map["model.norm"] = device device_map["lm_head"] = 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 = config.num_hidden_layers 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["model.embed_tokens"] = device device_map["model.norm"] = device device_map["lm_head"] = 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", "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") 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"} 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 LoRA; Hybrid, GPU-expert, LoRA-expert, " "INT4, and SkipLoRA paths are not supported" ) 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, ) 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}" ) kt_weight_path = getattr(cfg, "kt_weight_path", None) # For FP8 the frontend forwards model_name_or_path through kt_weight_path. # It is checkpoint provenance, not a pre-packed .kt directory. use_kt_weight_path = kt_weight_path is not None and expert_weight_format != "fp8" 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 == "fp8" and not checkpoint_files and kt_weight_path: logger.info( "Resolving native FP8 checkpoint files from kt_weight_path=%r", kt_weight_path, ) resolved_files, resolved_meta = _resolve_native_fp8_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") 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." ) 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"} 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 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 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 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), 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 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 = KTMoELayerWrapper( 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" 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 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"}, ) 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" # 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) native_fp8_experts = getattr(cfg, "kt_expert_weight_format", None) == "fp8" if native_fp8_experts: # FP8 kt_weight_path is raw-checkpoint provenance; KT loads routed experts. cfg.kt_skip_expert_loading = True skip_expert_loading = getattr(cfg, "kt_skip_expert_loading", None) needs_checkpoint_resolution = ( (native_fp8_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_fp8_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_fp8_experts: cfg.kt_skip_expert_loading = False if native_fp8_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( "native FP8 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