# Weight extraction and loading utilities for SFT # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations import logging import os import time from dataclasses import dataclass import torch import torch.nn as nn from .arch import MOEArchConfig from .backend import RAWINT4_GROUP_SIZE from .conv3d_compat import _canonicalize_qwen3_vl_fused_expert_weights from .dist_utils import _maybe_zero3_gathered_parameters logger = logging.getLogger(__name__) try: from safetensors import safe_open SAFETENSORS_AVAILABLE = True except ImportError: SAFETENSORS_AVAILABLE = False safe_open = None # ============================================================================= # Weight Extraction # ============================================================================= def _require_materialized_expert_weight(name: str, weight: torch.Tensor) -> None: if weight.device.type == "meta" or getattr(weight, "_kt_zero_storage", False): raise RuntimeError( f"Cannot extract {name}: routed expert weight is an unmaterialized placeholder" ) try: storage_bytes = weight.untyped_storage().nbytes() except (NotImplementedError, RuntimeError): return required_bytes = weight.numel() * weight.element_size() if required_bytes and storage_bytes < required_bytes: raise RuntimeError( f"Cannot extract {name}: routed expert weight is a zero-storage placeholder" ) def extract_moe_weights( moe_module: nn.Module, moe_config: MOEArchConfig ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """ Extract MoE expert weights from the module. Returns (gate_proj, up_proj, down_proj) with shape [expert_num, out_features, in_features]. Supports two formats: - ModuleList of Linear experts (transformers v4 style) - Fused Parameters (transformers v5 style): single module with ``gate_up_proj`` [E, 2*I, H] and ``down_proj`` [E, H, I] tensors. """ from .arch import detect_fused_experts experts = getattr(moe_module, moe_config.experts_attr) # Fused format (transformers v5): a single nn.Module with gate_up_proj/down_proj tensors if detect_fused_experts(experts): gate_up_parameter = getattr(experts, "gate_up_proj") down_parameter = getattr(experts, "down_proj") _require_materialized_expert_weight("gate_up_proj", gate_up_parameter) _require_materialized_expert_weight("down_proj", down_parameter) gate_up = gate_up_parameter.data down_fused = down_parameter.data vlm_weights = _canonicalize_qwen3_vl_fused_expert_weights(gate_up, down_fused, moe_config) if vlm_weights is not None: return vlm_weights # gate_up_proj is [E, 2*I, H], split into gate [E, I, H] and up [E, I, H] intermediate = gate_up.shape[1] // 2 gate_proj = gate_up[:, :intermediate, :].contiguous() up_proj = gate_up[:, intermediate:, :].contiguous() # down_proj is already [E, H, I] down_proj = down_fused.contiguous() return gate_proj, up_proj, down_proj gate_name, up_name, down_name = moe_config.weight_names gather_params: list[torch.nn.Parameter] = [] for expert in experts: for weight_name in (gate_name, up_name, down_name): proj = getattr(expert, weight_name, None) if proj is not None and hasattr(proj, "weight"): # Handle PEFT LoRA wrapped modules weight = proj.weight if isinstance(weight, torch.Tensor): gather_params.append(weight) elif hasattr(weight, "data"): gather_params.append(weight.data) with _maybe_zero3_gathered_parameters(gather_params): gate_weights = [] up_weights = [] down_weights = [] for expert_idx, expert in enumerate(experts): # Handle PEFT LoRA wrapped modules - get weight tensor properly gate_proj = getattr(expert, gate_name) up_proj_mod = getattr(expert, up_name) down_proj_mod = getattr(expert, down_name) # Get weight tensors, handling both regular Linear and PEFT LoRA wrapped def get_weight_tensor(mod, projection_name): weight = mod.weight if isinstance(weight, torch.Tensor): tensor = weight elif hasattr(weight, "data"): tensor = weight.data else: raise ValueError( f"Cannot extract weight from {type(mod)}, weight type={type(weight)}" ) _require_materialized_expert_weight( f"expert {expert_idx} {projection_name}", tensor ) return tensor.data gate_weights.append(get_weight_tensor(gate_proj, gate_name)) up_weights.append(get_weight_tensor(up_proj_mod, up_name)) down_weights.append(get_weight_tensor(down_proj_mod, down_name)) gate_proj = torch.stack(gate_weights, dim=0) up_proj = torch.stack(up_weights, dim=0) down_proj = torch.stack(down_weights, dim=0) return gate_proj, up_proj, down_proj def _clear_original_expert_weights( moe_module: nn.Module, moe_config: MOEArchConfig, full_weight_grad: bool = False, *, empty_placeholders: bool = False, ) -> None: """ Clear original expert weights to free memory after KT weights are loaded. In full_weight_grad mode, gate_proj_buf/up_proj_buf/down_proj_buf serve as the authoritative copies for the optimizer. The original expert weights in the model tree are redundant and cause double-counting in count_parameters(). Clear them just like in LoRA mode. """ from .arch import detect_fused_experts experts = getattr(moe_module, moe_config.experts_attr, None) if experts is None: return def make_placeholder(parameter: torch.nn.Parameter) -> nn.Parameter: if empty_placeholders: fake_tensor = torch.empty(0, dtype=parameter.dtype, device="cpu") else: tiny_storage = torch.UntypedStorage(1, device="cpu") fake_tensor = torch.tensor([], dtype=parameter.dtype, device="cpu").set_( tiny_storage, storage_offset=0, size=parameter.shape, stride=[0] * len(parameter.shape), ) placeholder = nn.Parameter(fake_tensor, requires_grad=False) placeholder._kt_zero_storage = True placeholder._kt_original_shape = tuple(parameter.shape) return placeholder # Fused format: replace gate_up_proj/down_proj tensors with zero-storage placeholders if detect_fused_experts(experts): for name in ("gate_up_proj", "down_proj"): param = getattr(experts, name, None) if not isinstance(param, torch.nn.Parameter): continue experts._parameters[name] = make_placeholder(param) return def _iter_weight_params(): for expert in experts: for weight_name in moe_config.weight_names: proj = getattr(expert, weight_name, None) if proj is None or not hasattr(proj, "weight"): continue parametrizations = getattr(proj, "parametrizations", None) parametrized_weight = ( getattr(parametrizations, "weight", None) if parametrizations is not None else None ) if parametrized_weight is not None: original = getattr(parametrized_weight, "original", None) if isinstance(original, torch.nn.Parameter): yield proj, parametrized_weight, "original", original continue direct_weight = getattr(proj, "_parameters", {}).get("weight") if isinstance(direct_weight, torch.nn.Parameter): yield proj, proj, "weight", direct_weight continue # Fallback: `weight` can be a non-settable property (e.g. parametrizations) or a non-Parameter. weight_attr = getattr(proj, "weight", None) if isinstance(weight_attr, torch.nn.Parameter): yield proj, proj, "weight", weight_attr gather_params: list[torch.nn.Parameter] = [] for _, _, _, weight_param in _iter_weight_params(): gather_params.append(weight_param) replaced_count = 0 with _maybe_zero3_gathered_parameters(gather_params): for proj, container, param_name, weight_param in _iter_weight_params(): new_param = make_placeholder(weight_param) replaced_count += 1 # Avoid `KeyError: attribute 'weight' already exists` for parametrized modules # where `weight` is a property and the real parameter lives elsewhere. container_params = getattr(container, "_parameters", {}) if isinstance(container_params, dict) and param_name in container_params: container_params[param_name] = new_param continue if hasattr(container, param_name): logger.debug( f"Skipping clearing expert weight {type(proj).__name__}.{param_name}: " "attribute exists but is not a registered parameter." ) continue try: setattr(container, param_name, new_param) except Exception as exc: logger.warning(f"Failed to clear expert weight {type(proj).__name__}.{param_name}: {exc}") logger.info(f"Replaced {replaced_count} expert weight params") _KT_RANK_LOCAL_PARAMETER_NAMES = "_kt_rank_local_parameter_names" def get_kt_expert_placeholders(model: nn.Module) -> dict[str, nn.Parameter]: """Resolve current placeholder identities from stable, cached FQNs.""" placeholder_ids: set[int] = set() for module in model.modules(): if not getattr(module, "_is_kt_moe_wrapper", False): continue experts_attr = getattr(module, "_experts_attr", None) experts = getattr(module, experts_attr, None) if isinstance(experts_attr, str) else None if not isinstance(experts, nn.Module): raise RuntimeError("KT MoE wrapper does not expose a valid expert subtree") placeholder_ids.update( id(parameter) for parameter in experts.parameters(recurse=True) if getattr(parameter, "_kt_zero_storage", False) ) named_parameters: dict[str, nn.Parameter] = {} discovered: dict[str, nn.Parameter] = {} for fqn, parameter in model.named_parameters(remove_duplicate=False): if fqn in named_parameters and named_parameters[fqn] is not parameter: raise RuntimeError( f"Model parameter FQN {fqn!r} resolves to conflicting Parameter identities" ) named_parameters[fqn] = parameter if id(parameter) not in placeholder_ids: continue if fqn in discovered and discovered[fqn] is not parameter: raise RuntimeError(f"KT expert placeholder FQN {fqn!r} resolves to conflicting Parameter identities") discovered[fqn] = parameter found_ids = {id(parameter) for parameter in discovered.values()} if found_ids != placeholder_ids: raise RuntimeError("KT expert placeholder is not registered in the model parameter tree") cached_names = getattr(model, _KT_RANK_LOCAL_PARAMETER_NAMES, None) if cached_names is None: cached_names = tuple(sorted(discovered)) setattr(model, _KT_RANK_LOCAL_PARAMETER_NAMES, cached_names) elif ( not isinstance(cached_names, tuple) or any(not isinstance(name, str) or not name for name in cached_names) or tuple(sorted(set(cached_names))) != cached_names ): raise RuntimeError("KT rank-local parameter-name cache is invalid") elif set(discovered).difference(cached_names): raise RuntimeError("KT expert placeholder inventory changed after its initial discovery") placeholders = {} for fqn in cached_names: parameter = named_parameters.get(fqn) if parameter is None: raise RuntimeError(f"Cached KT rank-local parameter FQN is no longer registered: {fqn!r}") placeholders[fqn] = parameter return placeholders # ============================================================================= # kt_weight_path Loading Functions # ============================================================================= @dataclass class INT8ExpertWeights: """Container for INT8 expert weights with scales.""" gate_proj: torch.Tensor gate_scale: torch.Tensor up_proj: torch.Tensor up_scale: torch.Tensor down_proj: torch.Tensor down_scale: torch.Tensor @dataclass class BlockFP8ExpertWeights: """Raw per-expert E4M3 weights and normalized block scales. The lists intentionally remain unstacked: native FP8 packing consumes one expert pointer at a time and releases these checkpoint tensors after the synchronous C++ load completes. """ gate_proj: list[torch.Tensor] gate_scale: list[torch.Tensor] up_proj: list[torch.Tensor] up_scale: list[torch.Tensor] down_proj: list[torch.Tensor] down_scale: list[torch.Tensor] block_size: tuple[int, int] = (128, 128) @dataclass class RAWINT4ExpertWeights: """Signed group-wise INT4 tensors in the KGroup kernel layout.""" gate_proj: torch.Tensor gate_scale: torch.Tensor up_proj: torch.Tensor up_scale: torch.Tensor down_proj: torch.Tensor down_scale: torch.Tensor group_size: int = RAWINT4_GROUP_SIZE def _checkpoint_tensor_files( checkpoint_files: list[str], sharded_metadata: dict | None, ) -> dict[str, str]: """Return tensor key -> absolute checkpoint file without loading tensors.""" base_dir = os.path.dirname(checkpoint_files[0]) weight_map = ( sharded_metadata.get("weight_map") if isinstance(sharded_metadata, dict) else None ) if weight_map: return { key: os.path.join(base_dir, filename) for key, filename in weight_map.items() } result: dict[str, str] = {} for file_path in checkpoint_files: with safe_open(file_path, framework="pt") as handle: for key in handle.keys(): result[key] = file_path return result def load_block_fp8_experts_from_checkpoint_files( checkpoint_files: list[str], sharded_metadata: dict | None, layers_prefix: str, moe_config: MOEArchConfig, layer_idx: int, hidden_size: int, block_size: tuple[int, int] = (128, 128), ) -> BlockFP8ExpertWeights: """Load raw block-FP8 routed experts without stacking or dequantizing. Native FP8 SFT currently accepts only non-fused, CPU-resident E4M3FN matrices with one inverse scale per 128x128 block. Scale tensors from checkpoints may be BF16 or FP32; the returned scale lists are always FP32. """ if not SAFETENSORS_AVAILABLE: raise ImportError("safetensors is required for native FP8 expert loading") if not checkpoint_files: raise FileNotFoundError("checkpoint_files is empty") if tuple(block_size) != (128, 128): raise ValueError( "native FP8 SFT requires weight_block_size=[128, 128], " f"got {tuple(block_size)}" ) hidden_size = int(hidden_size) intermediate_size = int(moe_config.intermediate_size) if hidden_size % 128 or intermediate_size % 128: raise ValueError( "native FP8 SFT requires hidden and routed intermediate dimensions " f"divisible by 128, got hidden_size={hidden_size}, " f"intermediate_size={intermediate_size}" ) gate_name, up_name, down_name = moe_config.weight_names experts_prefix = ( f"{layers_prefix}.{layer_idx}.{moe_config.moe_layer_attr}." f"{moe_config.experts_attr}" ) tensor_files = _checkpoint_tensor_files(checkpoint_files, sharded_metadata) fused_keys = ( f"{experts_prefix}.gate_up_proj", f"{experts_prefix}.gate_up_proj.weight", f"{experts_prefix}.down_proj", f"{experts_prefix}.down_proj.weight", ) if any(key in tensor_files for key in fused_keys): raise ValueError( "native FP8 SFT requires non-fused per-expert checkpoint weights; " f"layer {layer_idx} contains fused gate_up/down tensors" ) projection_specs = { "gate": (gate_name, (intermediate_size, hidden_size)), "up": (up_name, (intermediate_size, hidden_size)), "down": (down_name, (hidden_size, intermediate_size)), } keys: list[str] = [] for expert_idx in range(int(moe_config.expert_num)): base = f"{experts_prefix}.{expert_idx}" for projection_name, _ in projection_specs.values(): keys.extend( ( f"{base}.{projection_name}.weight", f"{base}.{projection_name}.weight_scale_inv", ) ) missing = [key for key in keys if key not in tensor_files] if missing: preview = ", ".join(missing[:3]) suffix = " ..." if len(missing) > 3 else "" raise FileNotFoundError( f"layer {layer_idx} is missing {len(missing)} native FP8 tensors: " f"{preview}{suffix}" ) tensors: dict[str, torch.Tensor] = {} keys_by_file: dict[str, list[str]] = {} for key in keys: keys_by_file.setdefault(tensor_files[key], []).append(key) for file_path, file_keys in keys_by_file.items(): with safe_open(file_path, framework="pt") as handle: for key in file_keys: tensors[key] = handle.get_tensor(key) output: dict[str, list[torch.Tensor]] = { "gate_proj": [], "gate_scale": [], "up_proj": [], "up_scale": [], "down_proj": [], "down_scale": [], } fp8_dtype = torch.float8_e4m3fn for expert_idx in range(int(moe_config.expert_num)): base = f"{experts_prefix}.{expert_idx}" for label, (projection_name, expected_shape) in projection_specs.items(): weight_key = f"{base}.{projection_name}.weight" scale_key = f"{base}.{projection_name}.weight_scale_inv" weight = tensors[weight_key] scale = tensors[scale_key] expected_scale_shape = ( expected_shape[0] // 128, expected_shape[1] // 128, ) if weight.device.type != "cpu" or weight.dtype != fp8_dtype: raise ValueError( f"{weight_key} must be CPU torch.float8_e4m3fn, got " f"{weight.dtype} on {weight.device}" ) if tuple(weight.shape) != expected_shape: raise ValueError( f"{weight_key} shape mismatch: expected {expected_shape}, " f"got {tuple(weight.shape)}" ) if scale.device.type != "cpu" or scale.dtype not in { torch.bfloat16, torch.float32, }: raise ValueError( f"{scale_key} must be a BF16 or FP32 CPU tensor, got " f"{scale.dtype} on {scale.device}" ) if tuple(scale.shape) != expected_scale_shape: raise ValueError( f"{scale_key} shape mismatch for 128x128 blocks: expected " f"{expected_scale_shape}, got {tuple(scale.shape)}" ) output[f"{label}_proj"].append(weight.contiguous()) output[f"{label}_scale"].append( scale.to(dtype=torch.float32, device="cpu").contiguous() ) logger.info( "Loaded layer %d native FP8 routed experts: experts=%d, block_size=%s", layer_idx, moe_config.expert_num, block_size, ) return BlockFP8ExpertWeights(**output, block_size=(128, 128)) def _normalize_rawint4_packed_weight( tensor: torch.Tensor, *, rows: int, cols: int, key: str, ) -> torch.Tensor: """Return the exact byte-packed layout consumed by the KGroup kernel.""" if tensor.device.type != "cpu": raise ValueError(f"{key} must be a CPU tensor, got {tensor.device}") if cols % 8: raise ValueError(f"{key} input dimension must be divisible by 8, got {cols}") if tensor.dtype == torch.int32: expected = (rows, cols // 8) if tuple(tensor.shape) != expected: raise ValueError(f"{key} shape mismatch: expected int32 {expected}, got {tuple(tensor.shape)}") return tensor.contiguous().view(torch.uint8).view(rows, cols // 2).contiguous() if tensor.dtype == torch.uint8: expected = (rows, cols // 2) if tuple(tensor.shape) != expected: raise ValueError(f"{key} shape mismatch: expected uint8 {expected}, got {tuple(tensor.shape)}") return tensor.contiguous() raise ValueError(f"{key} must be torch.int32 or torch.uint8 packed INT4, got {tensor.dtype}") def load_rawint4_experts_from_checkpoint_files( checkpoint_files: list[str], sharded_metadata: dict | None, layers_prefix: str, moe_config: MOEArchConfig, layer_idx: int, hidden_size: int, group_size: int = RAWINT4_GROUP_SIZE, ) -> RAWINT4ExpertWeights: """Load one layer of strict compressed-tensors RAWINT4 experts.""" if not SAFETENSORS_AVAILABLE: raise ImportError("safetensors is required for RAWINT4 expert loading") if not checkpoint_files: raise FileNotFoundError("checkpoint_files is empty") if int(group_size) != RAWINT4_GROUP_SIZE: raise ValueError(f"RAWINT4 SFT requires group_size={RAWINT4_GROUP_SIZE}, got {group_size}") hidden_size = int(hidden_size) intermediate_size = int(moe_config.intermediate_size) if hidden_size % group_size or intermediate_size % group_size: raise ValueError( f"RAWINT4 group_size={group_size} must divide hidden_size={hidden_size} " f"and intermediate_size={intermediate_size}" ) tensor_files = _checkpoint_tensor_files(checkpoint_files, sharded_metadata) experts_prefix = f"{layers_prefix}.{layer_idx}.{moe_config.moe_layer_attr}." f"{moe_config.experts_attr}" gate_name, up_name, down_name = moe_config.weight_names projection_specs = { "gate": (gate_name, intermediate_size, hidden_size), "up": (up_name, intermediate_size, hidden_size), "down": (down_name, hidden_size, intermediate_size), } def load_projection(label: str) -> tuple[torch.Tensor, torch.Tensor]: projection_name, rows, cols = projection_specs[label] requested: list[tuple[str, str, str, str]] = [] for expert_idx in range(int(moe_config.expert_num)): base = f"{experts_prefix}.{expert_idx}.{projection_name}" requested.append( ( f"{base}.weight_packed", f"{base}.weight_scale", f"{base}.weight_shape", base, ) ) required_keys = [key for record in requested for key in record[:3]] missing = [key for key in required_keys if key not in tensor_files] if missing: preview = ", ".join(missing[:3]) suffix = " ..." if len(missing) > 3 else "" raise FileNotFoundError( f"layer {layer_idx} is missing {len(missing)} RAWINT4 tensors: " f"{preview}{suffix}" ) for _, _, _, base in requested: zero_point_keys = ( f"{base}.weight_zero_point", f"{base}.zero_point", ) present = [key for key in zero_point_keys if key in tensor_files] if present: raise ValueError( "RAWINT4 SFT accepts symmetric weights without zero-point " f"tensors, found {present[0]}" ) keys_by_file: dict[str, list[str]] = {} for key in required_keys: keys_by_file.setdefault(tensor_files[key], []).append(key) tensors: dict[str, torch.Tensor] = {} for file_path, file_keys in keys_by_file.items(): with safe_open(file_path, framework="pt") as handle: for key in file_keys: tensors[key] = handle.get_tensor(key) packed_weights: list[torch.Tensor] = [] scales: list[torch.Tensor] = [] expected_shape = (rows, cols) expected_scale_shape = (rows, cols // group_size) for packed_key, scale_key, shape_key, _ in requested: shape_tensor = tensors[shape_key] if shape_tensor.device.type != "cpu" or shape_tensor.dtype not in { torch.int32, torch.int64, }: raise ValueError( f"{shape_key} must be an INT32 or INT64 CPU tensor, got " f"{shape_tensor.dtype} on {shape_tensor.device}" ) shape = tuple(int(value) for value in shape_tensor.reshape(-1).tolist()) if shape != expected_shape: raise ValueError(f"{shape_key} mismatch: expected {expected_shape}, got {shape}") packed_weights.append( _normalize_rawint4_packed_weight(tensors[packed_key], rows=rows, cols=cols, key=packed_key) ) scale = tensors[scale_key] if scale.device.type != "cpu" or scale.dtype not in { torch.bfloat16, torch.float32, }: raise ValueError( f"{scale_key} must be a BF16 or FP32 CPU tensor, got " f"{scale.dtype} on {scale.device}" ) if tuple(scale.shape) != expected_scale_shape: raise ValueError( f"{scale_key} shape mismatch: expected {expected_scale_shape}, " f"got {tuple(scale.shape)}" ) scales.append(scale.to(dtype=torch.bfloat16, device="cpu").contiguous()) return ( torch.stack(packed_weights, dim=0).contiguous(), torch.stack(scales, dim=0).contiguous(), ) started = time.time() gate_proj, gate_scale = load_projection("gate") up_proj, up_scale = load_projection("up") down_proj, down_scale = load_projection("down") logger.info( "Loaded layer %d RAWINT4 routed experts: experts=%d, group_size=%d, elapsed=%.1fs", layer_idx, moe_config.expert_num, group_size, time.time() - started, ) return RAWINT4ExpertWeights( gate_proj=gate_proj, gate_scale=gate_scale, up_proj=up_proj, up_scale=up_scale, down_proj=down_proj, down_scale=down_scale, group_size=group_size, ) def _find_safetensor_files(kt_weight_path: str) -> list[str]: if not os.path.isdir(kt_weight_path): raise FileNotFoundError(f"kt_weight_path directory not found: {kt_weight_path}") safetensor_files = [] for file in sorted(os.listdir(kt_weight_path)): if file.endswith(".safetensors"): safetensor_files.append(os.path.join(kt_weight_path, file)) if not safetensor_files: raise FileNotFoundError(f"No safetensors files found in {kt_weight_path}") return safetensor_files def _load_kt_weight_index(kt_weight_path: str) -> dict[str, str]: if not SAFETENSORS_AVAILABLE: raise ImportError("safetensors is required for loading kt_weight_path") index = {} safetensor_files = _find_safetensor_files(kt_weight_path) for file_path in safetensor_files: with safe_open(file_path, framework="pt") as f: for key in f.keys(): index[key] = file_path logger.info(f"Indexed {len(index)} tensors from {len(safetensor_files)} safetensors files") return index def _dequant_fp8_experts( weights: list[torch.Tensor], scales: list[torch.Tensor | None], block_size: tuple[int, int] ) -> torch.Tensor: """Dequantize a list of FP8 expert weights and stack them (batched, vectorized). Args: weights: list of [out, in] float8_e4m3fn tensors (one per expert) scales: list of [out//bs_m, in//bs_n] scale_inv tensors (one per expert, may be None) block_size: (bs_m, bs_n) Returns: Stacked BF16 tensor of shape [num_experts, out, in] """ has_scales = scales[0] is not None if not has_scales: return torch.stack(weights, dim=0).to(torch.bfloat16).cpu().contiguous() bs_m, bs_n = block_size n = len(weights) out_features, in_features = weights[0].shape # Stack all experts: [N, out, in] fp8 -> reshape to blocks -> bf16 w = torch.stack(weights, dim=0) # [N, out, in] fp8 w = w.reshape(n, out_features // bs_m, bs_m, in_features // bs_n, bs_n) w = w.to(torch.bfloat16) # Stack all scales: [N, out//bs_m, in//bs_n] -> bf16, broadcast multiply s = torch.stack(scales, dim=0).to(torch.bfloat16) # [N, out//bs_m, in//bs_n] w = w * s[:, :, None, :, None] return w.reshape(n, out_features, in_features).contiguous() def load_experts_from_checkpoint_files( checkpoint_files: list[str], sharded_metadata: dict | None, layers_prefix: str, moe_config: MOEArchConfig, layer_idx: int, block_size: tuple[int, int] | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: if not SAFETENSORS_AVAILABLE: raise ImportError("safetensors is required for loading experts from checkpoint files") if not checkpoint_files: raise FileNotFoundError("checkpoint_files is empty") t0 = time.time() weight_map = None base_dir = os.path.dirname(checkpoint_files[0]) if sharded_metadata is not None: weight_map = sharded_metadata.get("weight_map", None) gate_name, up_name, down_name = moe_config.weight_names experts_prefix = f"{layers_prefix}.{layer_idx}.{moe_config.moe_layer_attr}.{moe_config.experts_attr}" fused_gate_up_key = f"{experts_prefix}.gate_up_proj" fused_down_key = f"{experts_prefix}.down_proj" is_fused = weight_map is not None and fused_gate_up_key in weight_map if is_fused: keys = [fused_gate_up_key, fused_down_key] else: keys = [] for expert_idx in range(moe_config.expert_num): base = f"{experts_prefix}.{expert_idx}" keys.append(f"{base}.{gate_name}.weight") keys.append(f"{base}.{gate_name}.weight_scale_inv") keys.append(f"{base}.{up_name}.weight") keys.append(f"{base}.{up_name}.weight_scale_inv") keys.append(f"{base}.{down_name}.weight") keys.append(f"{base}.{down_name}.weight_scale_inv") keys_by_file: dict[str, list[str]] = {} mapped_count = 0 unmapped_count = 0 for key in keys: if weight_map is not None: filename = weight_map.get(key) if filename is None: unmapped_count += 1 continue mapped_count += 1 file_path = os.path.join(base_dir, filename) else: file_path = checkpoint_files[0] keys_by_file.setdefault(file_path, []).append(key) print( f"[kt_moe] Layer {layer_idx}: key mapping done in {time.time()-t0:.1f}s — " f"total_keys={len(keys)}, mapped={mapped_count}, unmapped={unmapped_count}, " f"files_to_open={len(keys_by_file)}", flush=True, ) t1 = time.time() tensor_map: dict[str, torch.Tensor] = {} for file_idx, (file_path, file_keys) in enumerate(keys_by_file.items()): with safe_open(file_path, framework="pt") as f: available_keys = set(f.keys()) for key in file_keys: if key in available_keys: tensor_map[key] = f.get_tensor(key) if file_idx == 0: print( f"[kt_moe] Layer {layer_idx}: first file loaded ({os.path.basename(file_path)}, " f"{len(file_keys)} keys) in {time.time()-t1:.1f}s", flush=True, ) print( f"[kt_moe] Layer {layer_idx}: all files loaded in {time.time()-t1:.1f}s — " f"tensor_map has {len(tensor_map)} tensors", flush=True, ) t2 = time.time() if is_fused: gate_up = tensor_map.get(fused_gate_up_key) down = tensor_map.get(fused_down_key) if gate_up is None or down is None: raise FileNotFoundError(f"Missing fused expert weights for layer {layer_idx}") gate_up = gate_up.cpu().to(torch.bfloat16).contiguous() down = down.cpu().to(torch.bfloat16) vlm_weights = _canonicalize_qwen3_vl_fused_expert_weights(gate_up, down, moe_config) if vlm_weights is not None: gate_proj, up_proj, down_proj = vlm_weights else: intermediate = gate_up.shape[1] // 2 gate_proj = gate_up[:, :intermediate, :].contiguous() up_proj = gate_up[:, intermediate:, :].contiguous() down_proj = down.contiguous() del gate_up print( f"[kt_moe] Layer {layer_idx}: fused expert format — " f"split gate_up_proj [{gate_proj.shape}] + down [{down_proj.shape}]", flush=True, ) print( f"[kt_moe] Layer {layer_idx}: done — dtype={gate_proj.dtype}, " f"shape={gate_proj.shape}, dequant=0.0s, total={time.time()-t0:.1f}s", flush=True, ) return gate_proj, up_proj, down_proj gate_weights = [] up_weights = [] down_weights = [] gate_scales = [] up_scales = [] down_scales = [] for expert_idx in range(moe_config.expert_num): base = f"{experts_prefix}.{expert_idx}" gate_key = f"{base}.{gate_name}.weight" up_key = f"{base}.{up_name}.weight" down_key = f"{base}.{down_name}.weight" if gate_key not in tensor_map or up_key not in tensor_map or down_key not in tensor_map: raise FileNotFoundError(f"Missing expert weights for layer {layer_idx}, expert {expert_idx}") gate_weights.append(tensor_map[gate_key]) up_weights.append(tensor_map[up_key]) down_weights.append(tensor_map[down_key]) gate_scales.append(tensor_map.get(f"{base}.{gate_name}.weight_scale_inv")) up_scales.append(tensor_map.get(f"{base}.{up_name}.weight_scale_inv")) down_scales.append(tensor_map.get(f"{base}.{down_name}.weight_scale_inv")) # Check if weights are FP8 and need dequantization t2 = time.time() is_fp8 = gate_weights[0].dtype == torch.float8_e4m3fn if is_fp8: if block_size is None: block_size = (128, 128) print( f"[kt_moe] Layer {layer_idx}: FP8 expert weights detected, " f"dequantizing with block_size={block_size} " f"(has_scales={gate_scales[0] is not None})", flush=True, ) gate_proj = _dequant_fp8_experts(gate_weights, gate_scales, block_size) up_proj = _dequant_fp8_experts(up_weights, up_scales, block_size) down_proj = _dequant_fp8_experts(down_weights, down_scales, block_size) else: gate_proj = torch.stack(gate_weights, dim=0).cpu().to(torch.bfloat16).contiguous() up_proj = torch.stack(up_weights, dim=0).cpu().to(torch.bfloat16).contiguous() down_proj = torch.stack(down_weights, dim=0).cpu().to(torch.bfloat16).contiguous() print( f"[kt_moe] Layer {layer_idx}: done — dtype={gate_proj.dtype}, shape={gate_proj.shape}, " f"dequant={time.time()-t2:.1f}s, total={time.time()-t0:.1f}s", flush=True, ) return gate_proj, up_proj, down_proj def load_experts_from_kt_weight_path( kt_weight_path: str, layer_idx: int, num_experts: int, hidden_size: int, intermediate_size: int, ) -> INT8ExpertWeights: """Load INT8 preprocessed expert weights from kt_weight_path for a specific layer.""" if not SAFETENSORS_AVAILABLE: raise ImportError("safetensors is required for loading kt_weight_path") index = _load_kt_weight_index(kt_weight_path) numa_count = 0 test_key_prefix = f"blk.{layer_idx}.ffn_gate_exps.0.numa." for key in index.keys(): if key.startswith(test_key_prefix) and key.endswith(".weight"): numa_idx = int(key.split("numa.")[1].split(".")[0]) numa_count = max(numa_count, numa_idx + 1) if numa_count == 0: raise FileNotFoundError( f"No weights found for layer {layer_idx} in {kt_weight_path}. " f"Expected keys like 'blk.{layer_idx}.ffn_gate_exps.0.numa.0.weight'" ) logger.info(f"Loading INT8 weights for layer {layer_idx}: {num_experts} experts, {numa_count} NUMA partitions") gate_weights_list = [] gate_scales_list = [] up_weights_list = [] up_scales_list = [] down_weights_list = [] down_scales_list = [] for expert_idx in range(num_experts): gate_w_parts = [] gate_s_parts = [] for numa_idx in range(numa_count): w_key = f"blk.{layer_idx}.ffn_gate_exps.{expert_idx}.numa.{numa_idx}.weight" s_key = f"blk.{layer_idx}.ffn_gate_exps.{expert_idx}.numa.{numa_idx}.scale" if w_key not in index: raise FileNotFoundError(f"Weight key not found: {w_key}") with safe_open(index[w_key], framework="pt") as f: gate_w_parts.append(f.get_tensor(w_key)) gate_s_parts.append(f.get_tensor(s_key)) gate_w = torch.cat(gate_w_parts, dim=0) gate_s = torch.cat(gate_s_parts, dim=0) gate_w = gate_w.view(intermediate_size, hidden_size) gate_weights_list.append(gate_w) gate_scales_list.append(gate_s) up_w_parts = [] up_s_parts = [] for numa_idx in range(numa_count): w_key = f"blk.{layer_idx}.ffn_up_exps.{expert_idx}.numa.{numa_idx}.weight" s_key = f"blk.{layer_idx}.ffn_up_exps.{expert_idx}.numa.{numa_idx}.scale" if w_key not in index: raise FileNotFoundError(f"Weight key not found: {w_key}") with safe_open(index[w_key], framework="pt") as f: up_w_parts.append(f.get_tensor(w_key)) up_s_parts.append(f.get_tensor(s_key)) up_w = torch.cat(up_w_parts, dim=0) up_s = torch.cat(up_s_parts, dim=0) up_w = up_w.view(intermediate_size, hidden_size) up_weights_list.append(up_w) up_scales_list.append(up_s) down_w_parts = [] down_s_parts = [] for numa_idx in range(numa_count): w_key = f"blk.{layer_idx}.ffn_down_exps.{expert_idx}.numa.{numa_idx}.weight" s_key = f"blk.{layer_idx}.ffn_down_exps.{expert_idx}.numa.{numa_idx}.scale" if w_key not in index: raise FileNotFoundError(f"Weight key not found: {w_key}") with safe_open(index[w_key], framework="pt") as f: down_w_parts.append(f.get_tensor(w_key)) down_s_parts.append(f.get_tensor(s_key)) down_w = torch.cat(down_w_parts, dim=0) down_s = torch.cat(down_s_parts, dim=0) down_w = down_w.view(hidden_size, intermediate_size) down_weights_list.append(down_w) down_scales_list.append(down_s) gate_proj = torch.stack(gate_weights_list, dim=0) gate_scale = torch.stack(gate_scales_list, dim=0) up_proj = torch.stack(up_weights_list, dim=0) up_scale = torch.stack(up_scales_list, dim=0) down_proj = torch.stack(down_weights_list, dim=0) down_scale = torch.stack(down_scales_list, dim=0) return INT8ExpertWeights( gate_proj=gate_proj, gate_scale=gate_scale, up_proj=up_proj, up_scale=up_scale, down_proj=down_proj, down_scale=down_scale, )