mirror of
https://github.com/kvcache-ai/ktransformers.git
synced 2026-10-17 05:11:19 +08:00
yyj
e80614fc2d
* feat(sft): support distributed activation reuse policies * feat(sft): add frozen-base INT8 LoRA training * fix(sft): make INT8 expert LoRA rank-zero authoritative * fix(sft): preserve DeepSeek router LoRA routing * feat(sft): enable persistent INT8 LoRA training * perf(sft): accelerate INT8 VNNI with oneDNN BRGEMM * perf(int8): fuse oneDNN compensation into backward repack * [feat]: support BF16 expert LoRA training * [fix]: honor forwarded activation policy in SFT workers * feat(sft): add native block-FP8 routed expert LoRA * feat(sft): expose explicit expert placeholder ownership * fix(sft): publish fused adapter artifacts atomically * feat(sft): own artifact and adapter lifecycle contracts * fix(sft): harden artifact and rank-local contracts * fix(sft): auto-adapt owner before adapter restore * style(sft): keep lifecycle comments concise * fix(sft): require fused adapter manifests * test(sft): use spawn for distributed workers * fix(sft): preserve runtime checkpoint metadata * fix(sft): validate wrapped runtime configuration * fix(sft): preserve expert format provenance * fix(sft): own routed experts during device dispatch * test(sft): lock explicit quantization conflict * fix(cpu): make shared memory buffers lifetime-safe * release: prepare v0.7.0
1191 lines
48 KiB
Python
1191 lines
48 KiB
Python
# PEFT LoRA adaptation utilities for SFT
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
"""
|
|
PEFT LoRA integration for KT-Kernel MoE training.
|
|
|
|
Handles:
|
|
- LoRA Expert modules (LoRAExpertMLP, LoRAExperts)
|
|
- PEFT LoRA adaptation onto KT wrappers (contiguous buffer views, grad buffers)
|
|
- LoRA parameter collection for optimizer injection
|
|
- Checkpoint save/load for lora_experts
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import math
|
|
import os
|
|
import re
|
|
import tempfile
|
|
from dataclasses import dataclass
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
from .arch import MOEArchConfig
|
|
from .checkpoint import load_full_weight_checkpoint, save_full_weight_checkpoint
|
|
from .dist_utils import _distributed_rank_world_size
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class KTAdaptationResult:
|
|
"""Stable optimizer and placeholder ownership published by KT adaptation."""
|
|
|
|
adapted_layers: int
|
|
already_adapted_layers: int
|
|
named_optimizer_parameters: tuple[tuple[str, nn.Parameter], ...]
|
|
placeholders: tuple[tuple[str, nn.Parameter], ...]
|
|
placeholder_names: tuple[str, ...]
|
|
|
|
|
|
# =============================================================================
|
|
# LoRA Experts Modules
|
|
# =============================================================================
|
|
|
|
|
|
class LoRAExpertMLP(nn.Module):
|
|
"""Single LoRA Expert with SwiGLU activation structure."""
|
|
|
|
def __init__(
|
|
self,
|
|
hidden_size: int,
|
|
intermediate_size: int,
|
|
device: str = "cuda",
|
|
dtype: torch.dtype = torch.bfloat16,
|
|
):
|
|
super().__init__()
|
|
self.le_gate = nn.Linear(hidden_size, intermediate_size, bias=False, device=device, dtype=dtype)
|
|
self.le_up = nn.Linear(hidden_size, intermediate_size, bias=False, device=device, dtype=dtype)
|
|
self.le_down = nn.Linear(intermediate_size, hidden_size, bias=False, device=device, dtype=dtype)
|
|
self.act_fn = nn.SiLU()
|
|
|
|
nn.init.zeros_(self.le_down.weight)
|
|
nn.init.kaiming_uniform_(self.le_gate.weight, a=math.sqrt(5))
|
|
nn.init.kaiming_uniform_(self.le_up.weight, a=math.sqrt(5))
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
return self.le_down(self.act_fn(self.le_gate(x)) * self.le_up(x))
|
|
|
|
|
|
class LoRAExperts(nn.Module):
|
|
"""LoRA Experts module containing multiple LoRA Expert MLPs."""
|
|
|
|
def __init__(
|
|
self,
|
|
num_experts: int,
|
|
hidden_size: int,
|
|
intermediate_size: int,
|
|
device: str = "cuda",
|
|
dtype: torch.dtype = torch.bfloat16,
|
|
):
|
|
super().__init__()
|
|
self.experts = nn.ModuleList(
|
|
[LoRAExpertMLP(hidden_size, intermediate_size, device, dtype) for _ in range(num_experts)]
|
|
)
|
|
self.num_experts = num_experts
|
|
|
|
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
|
output = torch.zeros_like(hidden_states)
|
|
for expert in self.experts:
|
|
output = output + expert(hidden_states)
|
|
return output / self.num_experts
|
|
|
|
|
|
# =============================================================================
|
|
# LoRA Parameter Collection
|
|
# =============================================================================
|
|
|
|
|
|
def _find_kt_wrappers(model: nn.Module):
|
|
"""Find _kt_wrappers on model, unwrapping PEFT/other wrappers if needed."""
|
|
wrappers = getattr(model, "_kt_wrappers", None)
|
|
if wrappers is None:
|
|
base_model = model
|
|
for attr in ("base_model", "model"):
|
|
if hasattr(base_model, attr):
|
|
base_model = getattr(base_model, attr)
|
|
wrappers = getattr(base_model, "_kt_wrappers", None)
|
|
if wrappers:
|
|
break
|
|
return wrappers
|
|
|
|
|
|
def _collect_wrapper_managed_lora_params(wrapper) -> list[nn.Parameter]:
|
|
"""Collect only C++-managed PEFT/fused LoRA params for one wrapper."""
|
|
params: list[nn.Parameter] = []
|
|
peft_lora_modules = getattr(wrapper, "_peft_lora_modules", None)
|
|
if peft_lora_modules is not None:
|
|
for expert_loras in peft_lora_modules.values():
|
|
for lora_A, lora_B in expert_loras.values():
|
|
if hasattr(lora_A, "weight") and lora_A.weight.requires_grad:
|
|
params.append(lora_A.weight)
|
|
if hasattr(lora_B, "weight") and lora_B.weight.requires_grad:
|
|
params.append(lora_B.weight)
|
|
fused_params = getattr(wrapper, "_fused_expert_lora_params", None)
|
|
if fused_params is not None:
|
|
params.extend(fused_params)
|
|
return params
|
|
|
|
|
|
def _collect_kt_lora_params(wrappers) -> list[nn.Parameter]:
|
|
"""Collect LoRA-only trainable parameters from KT wrappers."""
|
|
params: list[nn.Parameter] = []
|
|
|
|
if wrappers:
|
|
for wrapper in wrappers:
|
|
params.extend(_collect_wrapper_managed_lora_params(wrapper))
|
|
# lora_experts parameters (separate feature)
|
|
if getattr(wrapper, "lora_experts", None) is not None:
|
|
params.extend(wrapper.lora_experts.parameters())
|
|
|
|
return params
|
|
|
|
|
|
def _collect_kt_full_weight_params(wrappers) -> list[nn.Parameter]:
|
|
"""Collect optimizer-visible base expert parameters for full/hybrid KT SFT."""
|
|
params: list[nn.Parameter] = []
|
|
|
|
if wrappers:
|
|
for wrapper in wrappers:
|
|
if getattr(wrapper, "_full_weight_grad", False) and wrapper.wrapper is not None:
|
|
if wrapper.wrapper.gate_proj_buf is not None:
|
|
params.append(wrapper.wrapper.gate_proj_buf)
|
|
if wrapper.wrapper.up_proj_buf is not None:
|
|
params.append(wrapper.wrapper.up_proj_buf)
|
|
if wrapper.wrapper.down_proj_buf is not None:
|
|
params.append(wrapper.wrapper.down_proj_buf)
|
|
|
|
return params
|
|
|
|
|
|
def get_kt_lora_params(model: nn.Module) -> list[nn.Parameter]:
|
|
"""Get KT parameters for legacy Trainer optimizer injection.
|
|
|
|
Historically the patched Trainer calls this function after optimizer
|
|
creation. In full_weight_grad mode, returning only LoRA params silently
|
|
drops expert base weights from the optimizer, so this compatibility entry
|
|
point delegates to the full trainable collector when needed.
|
|
"""
|
|
wrappers = _find_kt_wrappers(model)
|
|
if not wrappers:
|
|
return []
|
|
|
|
if any(getattr(w, "_full_weight_grad", False) for w in wrappers):
|
|
return _collect_kt_full_weight_params(wrappers) + _collect_kt_lora_params(wrappers)
|
|
|
|
return _collect_kt_lora_params(wrappers)
|
|
|
|
|
|
def get_kt_trainable_params(model: nn.Module) -> list[nn.Parameter]:
|
|
"""Get all trainable parameters from KT model based on training mode.
|
|
|
|
In full mode: returns base weight nn.Parameter buffers from wrappers.
|
|
In LoRA mode: returns LoRA parameters (same as get_kt_lora_params).
|
|
"""
|
|
wrappers = _find_kt_wrappers(model)
|
|
if not wrappers:
|
|
return []
|
|
|
|
# Check if any wrapper is in full_weight_grad mode
|
|
has_full_weight_grad = any(getattr(w, "_full_weight_grad", False) for w in wrappers)
|
|
|
|
if has_full_weight_grad:
|
|
return _collect_kt_full_weight_params(wrappers) + _collect_kt_lora_params(wrappers)
|
|
else:
|
|
# LoRA mode: return LoRA parameters
|
|
return _collect_kt_lora_params(wrappers)
|
|
|
|
|
|
def get_kt_named_trainable_params(model: nn.Module) -> list[tuple[str, nn.Parameter]]:
|
|
"""Return every optimizer-visible KT parameter with a stable unique name.
|
|
|
|
Fused LoRA and Full/Hybrid base buffers are intentionally not registered in
|
|
the PyTorch module tree. Trainer integrations must use this API rather than
|
|
inventing anonymous parameter groups or depending on wrapper internals.
|
|
"""
|
|
|
|
wrappers = _find_kt_wrappers(model) or []
|
|
if not wrappers:
|
|
return []
|
|
|
|
registered_names: dict[int, str] = {}
|
|
for name, parameter in model.named_parameters(remove_duplicate=False):
|
|
registered_names.setdefault(id(parameter), name)
|
|
|
|
owned_names: dict[int, str] = {}
|
|
fused_names = (
|
|
"gate_lora_a",
|
|
"gate_lora_b",
|
|
"up_lora_a",
|
|
"up_lora_b",
|
|
"down_lora_a",
|
|
"down_lora_b",
|
|
)
|
|
for wrapper in wrappers:
|
|
layer_idx = getattr(wrapper, "layer_idx", None)
|
|
if isinstance(layer_idx, bool) or not isinstance(layer_idx, int):
|
|
raise RuntimeError(f"KT wrapper has an invalid layer index {layer_idx!r}")
|
|
prefix = f"kt.layers.{layer_idx}.experts"
|
|
backend = getattr(wrapper, "wrapper", None)
|
|
if getattr(wrapper, "_full_weight_grad", False) and backend is not None:
|
|
for name, parameter in (
|
|
("gate_proj", getattr(backend, "gate_proj_buf", None)),
|
|
("up_proj", getattr(backend, "up_proj_buf", None)),
|
|
("down_proj", getattr(backend, "down_proj_buf", None)),
|
|
):
|
|
if isinstance(parameter, nn.Parameter):
|
|
owned_names[id(parameter)] = f"{prefix}.base.{name}"
|
|
fused = getattr(wrapper, "_fused_expert_lora_params", None) or ()
|
|
if len(fused) not in {0, len(fused_names)}:
|
|
raise RuntimeError(
|
|
f"Layer {layer_idx}: expected six fused LoRA parameters, got {len(fused)}"
|
|
)
|
|
for name, parameter in zip(fused_names, fused):
|
|
owned_names[id(parameter)] = f"{prefix}.fused_lora.{name}"
|
|
|
|
result: list[tuple[str, nn.Parameter]] = []
|
|
seen_parameters: set[int] = set()
|
|
seen_names: set[str] = set()
|
|
for parameter in get_kt_trainable_params(model):
|
|
identity = id(parameter)
|
|
if identity in seen_parameters:
|
|
continue
|
|
name = registered_names.get(identity) or owned_names.get(identity)
|
|
if name is None:
|
|
raise RuntimeError("KT optimizer parameter has no stable registered or owned name")
|
|
if name in seen_names:
|
|
raise RuntimeError(f"KT optimizer parameter name is not unique: {name}")
|
|
seen_parameters.add(identity)
|
|
seen_names.add(name)
|
|
result.append((name, parameter))
|
|
return result
|
|
|
|
|
|
def get_kt_lora_named_params(model: nn.Module) -> list[tuple[str, nn.Parameter]]:
|
|
"""Compatibility spelling for Trainer integrations.
|
|
|
|
Like :func:`get_kt_lora_params`, this includes Full/Hybrid base parameters
|
|
when the active training mode requires them.
|
|
"""
|
|
|
|
return get_kt_named_trainable_params(model)
|
|
|
|
|
|
def get_kt_rank_local_parameter_names(model: nn.Module) -> tuple[str, ...]:
|
|
"""Return registered expert-placeholder FQNs owned outside FSDP2.
|
|
|
|
Wrapping creates and caches these names during ``from_pretrained``, so this
|
|
API is valid before PEFT adaptation. Fused optimizer buffers use virtual
|
|
names and are deliberately excluded because they are not model FQNs.
|
|
"""
|
|
|
|
from .weights import get_kt_expert_placeholders
|
|
|
|
return tuple(sorted(get_kt_expert_placeholders(model)))
|
|
|
|
|
|
def _adaptation_result(
|
|
model: nn.Module,
|
|
*,
|
|
adapted_layers: int,
|
|
already_adapted_layers: int,
|
|
) -> KTAdaptationResult:
|
|
from .weights import get_kt_expert_placeholders
|
|
|
|
placeholders = get_kt_expert_placeholders(model)
|
|
placeholder_names = get_kt_rank_local_parameter_names(model)
|
|
named_placeholders = tuple((name, placeholders[name]) for name in placeholder_names)
|
|
return KTAdaptationResult(
|
|
adapted_layers=adapted_layers,
|
|
already_adapted_layers=already_adapted_layers,
|
|
named_optimizer_parameters=tuple(get_kt_named_trainable_params(model)),
|
|
placeholders=named_placeholders,
|
|
placeholder_names=placeholder_names,
|
|
)
|
|
|
|
|
|
# =============================================================================
|
|
# PEFT LoRA Adaptation
|
|
# =============================================================================
|
|
|
|
|
|
def kt_adapt_peft_lora(model: nn.Module) -> KTAdaptationResult:
|
|
"""
|
|
Adapt PEFT LoRA on expert modules for KT kernel.
|
|
|
|
After PEFT injects LoRA adapters onto expert Linear modules, this function:
|
|
1. Detects PEFT LoRA presence and rank on each wrapper's experts
|
|
2. Stores references to PEFT LoRA modules on the wrapper (for backward gradient writing)
|
|
3. Syncs initial PEFT LoRA weights to the C++ KT kernel (rank 0 only)
|
|
|
|
PEFT LoRA remains active and is managed by PEFT. No separate KT lora_params created.
|
|
Optimizer updates PEFT LoRA directly, and KT kernel reads from PEFT LoRA on each forward.
|
|
|
|
Should be called after PEFT LoRA injection and before create_optimizer.
|
|
"""
|
|
wrappers = _find_kt_wrappers(model)
|
|
|
|
if not wrappers:
|
|
logger.info("[kt_adapt_peft_lora] No _kt_wrappers found, skipping")
|
|
return _adaptation_result(model, adapted_layers=0, already_adapted_layers=0)
|
|
|
|
distributed_rank, _ = _distributed_rank_world_size()
|
|
is_rank_0 = distributed_rank == 0
|
|
|
|
adapted_count = 0
|
|
already_adapted_count = 0
|
|
for wrapper in wrappers:
|
|
if getattr(wrapper, "_kt_peft_lora_adapted", False):
|
|
already_adapted_count += 1
|
|
continue
|
|
moe_config = wrapper.moe_config
|
|
layer_idx = wrapper.layer_idx
|
|
experts_attr = getattr(wrapper, "_experts_attr", "experts")
|
|
experts = getattr(wrapper, experts_attr, None)
|
|
|
|
if experts is None:
|
|
continue
|
|
|
|
# Packed experts, or an explicitly forced offloaded expert path, use
|
|
# KT-managed contiguous LoRA buffers.
|
|
if getattr(wrapper, "_use_fused_expert_lora", False) or getattr(
|
|
wrapper, "_fused_experts", False
|
|
):
|
|
lora_rank = getattr(wrapper, "_lora_rank", 1)
|
|
authoritative_mode = bool(getattr(wrapper, "_uses_authoritative_optimizer_grads", False))
|
|
|
|
# In full mode (lora_rank=0), skip LoRA buffer creation entirely.
|
|
# C++ kernel will not compute LoRA contributions when lora_rank=0.
|
|
if lora_rank == 0:
|
|
wrapper._fused_expert_lora_params = []
|
|
wrapper._peft_lora_modules = None
|
|
logger.info(
|
|
f"[kt_adapt_peft_lora] Layer {layer_idx}: fused expert, "
|
|
f"full mode (lora_rank=0, no LoRA buffers)"
|
|
)
|
|
adapted_count += 1
|
|
wrapper._kt_peft_lora_adapted = True
|
|
continue
|
|
|
|
wrapper._kt_managed_lora_enabled = True
|
|
|
|
# In rank-0-authoritative BF16 mode, non-rank-0 processes need the
|
|
# mode flag for collective/autograd symmetry but own no optimizer
|
|
# Parameter or gradient buffer.
|
|
if authoritative_mode and not is_rank_0:
|
|
wrapper._fused_expert_lora_params = []
|
|
wrapper._peft_lora_modules = None
|
|
adapted_count += 1
|
|
wrapper._kt_peft_lora_adapted = True
|
|
continue
|
|
|
|
lora_buffers, lora_grad_buffers, lora_params = _create_fused_expert_lora_buffers(
|
|
wrapper,
|
|
moe_config,
|
|
lora_rank,
|
|
torch.bfloat16,
|
|
preassign_grads=not authoritative_mode,
|
|
)
|
|
|
|
if is_rank_0 and wrapper.wrapper is not None:
|
|
all_buffers = {}
|
|
all_buffers.update(lora_buffers)
|
|
all_buffers.update(lora_grad_buffers)
|
|
wrapper.wrapper.init_lora_weights(**all_buffers)
|
|
logger.info(
|
|
f"[kt_adapt_peft_lora] Layer {layer_idx}: fused expert LoRA "
|
|
f"(r={lora_rank}, E={moe_config.expert_num})"
|
|
)
|
|
if authoritative_mode:
|
|
for key, param in zip(
|
|
("gate_lora_a", "gate_lora_b", "up_lora_a", "up_lora_b", "down_lora_a", "down_lora_b"),
|
|
lora_params,
|
|
):
|
|
wrapper.wrapper.register_authoritative_optimizer_grad(
|
|
f"lora.{key}", param, lora_grad_buffers[f"grad_{key}"]
|
|
)
|
|
elif authoritative_mode:
|
|
raise RuntimeError(f"Layer {layer_idx}: rank 0 authoritative LoRA requires a KT backend")
|
|
|
|
wrapper._fused_expert_lora_params = lora_params
|
|
wrapper._peft_lora_modules = None
|
|
adapted_count += 1
|
|
wrapper._kt_peft_lora_adapted = True
|
|
continue
|
|
|
|
if len(experts) == 0:
|
|
continue
|
|
|
|
# In full mode (lora_rank=0), PEFT does not inject LoRA on experts.
|
|
# Skip LoRA detection and initialization entirely.
|
|
if getattr(wrapper, "_lora_rank", 1) == 0:
|
|
wrapper._peft_lora_modules = None
|
|
wrapper._fused_expert_lora_params = []
|
|
logger.info(
|
|
f"[kt_adapt_peft_lora] Layer {layer_idx}: non-fused expert, "
|
|
f"full mode (lora_rank=0, no LoRA)"
|
|
)
|
|
adapted_count += 1
|
|
wrapper._kt_peft_lora_adapted = True
|
|
continue
|
|
|
|
# Collect references to PEFT LoRA modules for each expert
|
|
# Structure: {expert_idx: {proj_name: (lora_A_module, lora_B_module)}}
|
|
peft_lora_modules = {}
|
|
gate_name, up_name, down_name = moe_config.weight_names
|
|
|
|
for expert_idx, expert in enumerate(experts):
|
|
expert_loras = {}
|
|
for proj_name in (gate_name, up_name, down_name):
|
|
proj = getattr(expert, proj_name, None)
|
|
if proj is None:
|
|
continue
|
|
lora_A = getattr(proj, "lora_A", None)
|
|
lora_B = getattr(proj, "lora_B", None)
|
|
if lora_A is not None and lora_B is not None:
|
|
# Get the actual Linear modules (inside ModuleDict if using adapters)
|
|
if isinstance(lora_A, nn.ModuleDict):
|
|
adapter_name = "default"
|
|
active = getattr(proj, "active_adapter", ["default"])
|
|
if isinstance(active, (list, tuple)) and active:
|
|
adapter_name = active[0]
|
|
# ModuleDict doesn't have .get(), use [] with in check
|
|
lora_A = lora_A[adapter_name] if adapter_name in lora_A else None
|
|
lora_B = lora_B[adapter_name] if adapter_name in lora_B else None
|
|
if lora_A is not None and lora_B is not None:
|
|
expert_loras[proj_name] = (lora_A, lora_B)
|
|
if expert_loras:
|
|
peft_lora_modules[expert_idx] = expert_loras
|
|
|
|
# Store PEFT LoRA references on wrapper
|
|
wrapper._peft_lora_modules = peft_lora_modules
|
|
|
|
# Missing PEFT LoRA is valid only for pure Full. Hybrid has lora_rank
|
|
# greater than zero and must fail instead of silently training base
|
|
# weights alone.
|
|
if not peft_lora_modules:
|
|
if getattr(wrapper, "_full_weight_grad", False) and getattr(wrapper, "_lora_rank", 0) == 0:
|
|
logger.info(
|
|
f"[kt_adapt_peft_lora] Layer {layer_idx}: No PEFT LoRA found "
|
|
f"(pure Full mode — expected, skipping)"
|
|
)
|
|
adapted_count += 1
|
|
wrapper._kt_peft_lora_adapted = True
|
|
continue
|
|
raise RuntimeError(
|
|
f"[kt_adapt_peft_lora] Layer {layer_idx}: No PEFT LoRA found on any expert. "
|
|
f"Check that PEFT lora_target includes expert modules."
|
|
)
|
|
|
|
wrapper._kt_managed_lora_enabled = True
|
|
authoritative_mode = bool(getattr(wrapper, "_uses_authoritative_optimizer_grads", False))
|
|
|
|
# Allocate contiguous bf16 buffers and populate with initial PEFT values (all ranks)
|
|
lora_buffers = _create_lora_view_buffers(peft_lora_modules, moe_config, torch.bfloat16)
|
|
lora_grad_buffers = _create_lora_grad_buffers(peft_lora_modules, moe_config)
|
|
|
|
# Rank 0: pass buffers to C++ wrapper (init_lora_weights stores them via .contiguous() no-op)
|
|
if is_rank_0 and wrapper.wrapper is not None:
|
|
# concat lora_buffers and lora_grad_buffers into single dict
|
|
lora_buffers.update(lora_grad_buffers)
|
|
wrapper.wrapper.init_lora_weights(**lora_buffers)
|
|
logger.info(f"[kt_adapt_peft_lora] Layer {layer_idx}: synced PEFT LoRA to C++ kernel")
|
|
|
|
# All ranks: replace PEFT weights with views into the contiguous buffers
|
|
_replace_peft_weights_with_views(
|
|
peft_lora_modules,
|
|
lora_buffers,
|
|
lora_grad_buffers,
|
|
moe_config,
|
|
authoritative_mode=authoritative_mode,
|
|
authoritative_backend=wrapper.wrapper if authoritative_mode and is_rank_0 else None,
|
|
)
|
|
|
|
adapted_count += 1
|
|
wrapper._kt_peft_lora_adapted = True
|
|
|
|
# Expert base weights are owned by the KT backend after PEFT injection. Keep
|
|
# zero-element placeholders in the module tree so FSDP cannot allocate or
|
|
# broadcast tensors with the original expert shapes.
|
|
shrunk_count = 0
|
|
shrunk_saved_bytes = 0
|
|
for wrapper in wrappers:
|
|
experts_attr = getattr(wrapper, "_experts_attr", "experts")
|
|
experts = getattr(wrapper, experts_attr, None)
|
|
if experts is None:
|
|
continue
|
|
if getattr(wrapper, "_fused_experts", False):
|
|
continue
|
|
for expert in experts:
|
|
for param_name, param in list(expert.named_parameters()):
|
|
if param.requires_grad:
|
|
continue # Skip trainable params (LoRA weights)
|
|
try:
|
|
storage_bytes = param.data.untyped_storage().nbytes()
|
|
except Exception:
|
|
continue
|
|
if storage_bytes > 2:
|
|
continue # Skip non-placeholder params
|
|
|
|
# This is a tiny-storage placeholder (base weight). Preserve its
|
|
# logical shape as metadata while removing its final storage.
|
|
original_shape = getattr(param, "_kt_original_shape", tuple(param.shape))
|
|
if param.numel() == 0:
|
|
param._kt_original_shape = tuple(original_shape)
|
|
continue
|
|
original_numel = param.nelement()
|
|
parts = param_name.split(".")
|
|
container = expert
|
|
for p in parts[:-1]:
|
|
container = getattr(container, p)
|
|
local_name = parts[-1]
|
|
container_params = getattr(container, "_parameters", {})
|
|
if isinstance(container_params, dict) and local_name in container_params:
|
|
empty_param = nn.Parameter(
|
|
torch.empty(0, dtype=param.dtype, device="cpu"),
|
|
requires_grad=False,
|
|
)
|
|
empty_param._kt_zero_storage = True
|
|
empty_param._kt_original_shape = tuple(original_shape)
|
|
container_params[local_name] = empty_param
|
|
shrunk_count += 1
|
|
shrunk_saved_bytes += original_numel * param.element_size()
|
|
|
|
if shrunk_count > 0:
|
|
logger.info(
|
|
f"[kt_adapt_peft_lora] Removed storage from {shrunk_count} expert base "
|
|
f"weight params, FSDP broadcast savings={shrunk_saved_bytes / 1024 / 1024:.1f} MB"
|
|
)
|
|
|
|
logger.info(
|
|
"[kt_adapt_peft_lora] Adapted %d layers; %d were already adapted",
|
|
adapted_count,
|
|
already_adapted_count,
|
|
)
|
|
result = _adaptation_result(
|
|
model,
|
|
adapted_layers=adapted_count,
|
|
already_adapted_layers=already_adapted_count,
|
|
)
|
|
model._kt_peft_lora_adaptation_result = result
|
|
return result
|
|
|
|
|
|
# New integrations may use the concise name; the historical entry point stays
|
|
# public for downstream compatibility.
|
|
adapt_peft_lora = kt_adapt_peft_lora
|
|
|
|
|
|
# =============================================================================
|
|
# Contiguous Buffer Creation
|
|
# =============================================================================
|
|
|
|
|
|
def _create_lora_view_buffers(
|
|
peft_lora_modules: dict[int, dict[str, tuple[nn.Module, nn.Module]]],
|
|
moe_config: MOEArchConfig,
|
|
dtype: torch.dtype = torch.bfloat16,
|
|
) -> dict[str, torch.Tensor]:
|
|
"""
|
|
Allocate contiguous buffers and populate with initial PEFT LoRA values.
|
|
|
|
Returns dict with gate_lora_a, gate_lora_b, up_lora_a, up_lora_b,
|
|
down_lora_a, down_lora_b — each shape [num_experts, ...].
|
|
"""
|
|
gate_name, up_name, down_name = moe_config.weight_names
|
|
num_experts = moe_config.expert_num
|
|
|
|
first_expert_loras = peft_lora_modules.get(0, {})
|
|
if not first_expert_loras:
|
|
raise RuntimeError("No PEFT LoRA found on expert 0")
|
|
gate_lora = first_expert_loras.get(gate_name)
|
|
if gate_lora is None:
|
|
raise RuntimeError(f"No PEFT LoRA found on expert 0 {gate_name}")
|
|
|
|
lora_rank = gate_lora[0].weight.shape[0]
|
|
hidden_size = gate_lora[0].weight.shape[1]
|
|
intermediate_size = gate_lora[1].weight.shape[0]
|
|
|
|
buffers = {
|
|
"gate_lora_a": torch.zeros(num_experts, lora_rank, hidden_size, dtype=dtype, device="cpu"),
|
|
"gate_lora_b": torch.zeros(num_experts, intermediate_size, lora_rank, dtype=dtype, device="cpu"),
|
|
"up_lora_a": torch.zeros(num_experts, lora_rank, hidden_size, dtype=dtype, device="cpu"),
|
|
"up_lora_b": torch.zeros(num_experts, intermediate_size, lora_rank, dtype=dtype, device="cpu"),
|
|
"down_lora_a": torch.zeros(num_experts, lora_rank, intermediate_size, dtype=dtype, device="cpu"),
|
|
"down_lora_b": torch.zeros(num_experts, hidden_size, lora_rank, dtype=dtype, device="cpu"),
|
|
}
|
|
|
|
proj_to_keys = {
|
|
gate_name: ("gate_lora_a", "gate_lora_b"),
|
|
up_name: ("up_lora_a", "up_lora_b"),
|
|
down_name: ("down_lora_a", "down_lora_b"),
|
|
}
|
|
for expert_idx in range(num_experts):
|
|
expert_loras = peft_lora_modules.get(expert_idx, {})
|
|
for proj_name, (key_a, key_b) in proj_to_keys.items():
|
|
if proj_name in expert_loras:
|
|
lora_A, lora_B = expert_loras[proj_name]
|
|
buffers[key_a][expert_idx].copy_(lora_A.weight.data.to(dtype=dtype))
|
|
buffers[key_b][expert_idx].copy_(lora_B.weight.data.to(dtype=dtype))
|
|
|
|
return buffers
|
|
|
|
|
|
def _create_lora_grad_buffers(
|
|
peft_lora_modules: dict[int, dict[str, tuple[nn.Module, nn.Module]]],
|
|
moe_config: MOEArchConfig,
|
|
dtype: torch.dtype = torch.bfloat16,
|
|
) -> dict[str, torch.Tensor]:
|
|
"""
|
|
Allocate contiguous gradient buffers for PEFT LoRA.
|
|
|
|
Returns dict with grad_gate_lora_a, grad_gate_lora_b, etc. — each shape [num_experts, ...].
|
|
"""
|
|
gate_name, up_name, down_name = moe_config.weight_names
|
|
num_experts = moe_config.expert_num
|
|
|
|
first_expert_loras = peft_lora_modules.get(0, {})
|
|
if not first_expert_loras:
|
|
raise RuntimeError("No PEFT LoRA found on expert 0")
|
|
gate_lora = first_expert_loras.get(gate_name)
|
|
if gate_lora is None:
|
|
raise RuntimeError(f"No PEFT LoRA found on expert 0 {gate_name}")
|
|
|
|
lora_rank = gate_lora[0].weight.shape[0]
|
|
hidden_size = gate_lora[0].weight.shape[1]
|
|
intermediate_size = gate_lora[1].weight.shape[0]
|
|
|
|
buffers = {
|
|
"grad_gate_lora_a": torch.zeros(num_experts, lora_rank, hidden_size, dtype=dtype, device="cpu"),
|
|
"grad_gate_lora_b": torch.zeros(num_experts, intermediate_size, lora_rank, dtype=dtype, device="cpu"),
|
|
"grad_up_lora_a": torch.zeros(num_experts, lora_rank, hidden_size, dtype=dtype, device="cpu"),
|
|
"grad_up_lora_b": torch.zeros(num_experts, intermediate_size, lora_rank, dtype=dtype, device="cpu"),
|
|
"grad_down_lora_a": torch.zeros(num_experts, lora_rank, intermediate_size, dtype=dtype, device="cpu"),
|
|
"grad_down_lora_b": torch.zeros(num_experts, hidden_size, lora_rank, dtype=dtype, device="cpu"),
|
|
}
|
|
|
|
return buffers
|
|
|
|
|
|
def _create_fused_expert_lora_buffers(
|
|
wrapper,
|
|
moe_config: MOEArchConfig,
|
|
lora_rank: int,
|
|
dtype: torch.dtype = torch.bfloat16,
|
|
preassign_grads: bool = True,
|
|
) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor], list[nn.Parameter]]:
|
|
"""
|
|
Create KT-managed LoRA buffers for fused expert modules.
|
|
|
|
Fused experts store weights as 3D parameters (gate_up_proj [E, 2I, H], down_proj [E, H, I])
|
|
rather than per-expert nn.Linear modules. PEFT can't attach per-expert LoRA to these,
|
|
so we create our own LoRA buffers that the C++ kernel reads/writes directly.
|
|
|
|
Returns:
|
|
(lora_buffers, lora_grad_buffers, lora_params):
|
|
- lora_buffers: dict of weight buffers for C++ init_lora_weights()
|
|
- lora_grad_buffers: dict of grad buffers for C++ backward
|
|
- lora_params: list of nn.Parameter wrappers for the optimizer
|
|
"""
|
|
E = moe_config.expert_num
|
|
I = moe_config.intermediate_size
|
|
H = wrapper.hidden_size
|
|
r = lora_rank
|
|
|
|
logger.info(f"[_create_fused_expert_lora_buffers] E={E}, I={I}, H={H}, r={r}")
|
|
|
|
lora_buffers = {
|
|
"gate_lora_a": torch.zeros(E, r, H, dtype=dtype, device="cpu"),
|
|
"gate_lora_b": torch.zeros(E, I, r, dtype=dtype, device="cpu"),
|
|
"up_lora_a": torch.zeros(E, r, H, dtype=dtype, device="cpu"),
|
|
"up_lora_b": torch.zeros(E, I, r, dtype=dtype, device="cpu"),
|
|
"down_lora_a": torch.zeros(E, r, I, dtype=dtype, device="cpu"),
|
|
"down_lora_b": torch.zeros(E, H, r, dtype=dtype, device="cpu"),
|
|
}
|
|
|
|
for key in ("gate_lora_a", "up_lora_a", "down_lora_a"):
|
|
nn.init.kaiming_uniform_(lora_buffers[key].view(E * r, -1), a=math.sqrt(5))
|
|
|
|
lora_grad_buffers = {
|
|
"grad_gate_lora_a": torch.zeros(E, r, H, dtype=dtype, device="cpu"),
|
|
"grad_gate_lora_b": torch.zeros(E, I, r, dtype=dtype, device="cpu"),
|
|
"grad_up_lora_a": torch.zeros(E, r, H, dtype=dtype, device="cpu"),
|
|
"grad_up_lora_b": torch.zeros(E, I, r, dtype=dtype, device="cpu"),
|
|
"grad_down_lora_a": torch.zeros(E, r, I, dtype=dtype, device="cpu"),
|
|
"grad_down_lora_b": torch.zeros(E, H, r, dtype=dtype, device="cpu"),
|
|
}
|
|
|
|
lora_params = []
|
|
for key in ("gate_lora_a", "gate_lora_b", "up_lora_a", "up_lora_b", "down_lora_a", "down_lora_b"):
|
|
param = nn.Parameter(lora_buffers[key], requires_grad=True)
|
|
if preassign_grads:
|
|
param.grad = lora_grad_buffers[f"grad_{key}"]
|
|
lora_params.append(param)
|
|
|
|
return lora_buffers, lora_grad_buffers, lora_params
|
|
|
|
|
|
# =============================================================================
|
|
# PEFT Weight View Replacement
|
|
# =============================================================================
|
|
|
|
|
|
def _replace_peft_weights_with_views(
|
|
peft_lora_modules: dict[int, dict[str, tuple[nn.Module, nn.Module]]],
|
|
buffers: dict[str, torch.Tensor],
|
|
grad_buffers: dict[str, torch.Tensor],
|
|
moe_config: MOEArchConfig,
|
|
*,
|
|
authoritative_mode: bool = False,
|
|
authoritative_backend=None,
|
|
) -> None:
|
|
"""
|
|
Replace each PEFT LoRA module's .weight with a view into the contiguous buffer.
|
|
|
|
After this, optimizer.step() updates the buffer in-place via the view —
|
|
no copy needed to sync with C++.
|
|
"""
|
|
gate_name, up_name, down_name = moe_config.weight_names
|
|
num_experts = moe_config.expert_num
|
|
|
|
proj_to_keys = {
|
|
gate_name: ("gate_lora_a", "gate_lora_b"),
|
|
up_name: ("up_lora_a", "up_lora_b"),
|
|
down_name: ("down_lora_a", "down_lora_b"),
|
|
}
|
|
|
|
_replaced = 0
|
|
_first_logged = False
|
|
for expert_idx in range(num_experts):
|
|
expert_loras = peft_lora_modules.get(expert_idx, {})
|
|
for proj_name, (key_a, key_b) in proj_to_keys.items():
|
|
if proj_name not in expert_loras:
|
|
continue
|
|
lora_A, lora_B = expert_loras[proj_name]
|
|
|
|
# Log before/after for first replacement to verify .data assignment
|
|
if not _first_logged:
|
|
_old_id_a = id(lora_A.weight)
|
|
_old_ptr_a = lora_A.weight.data_ptr()
|
|
|
|
# Use .data assignment to keep the same Parameter objects.
|
|
# This preserves optimizer references (which point to these objects).
|
|
# Creating new nn.Parameter() would break the optimizer link.
|
|
lora_A.weight.data = buffers[key_a][expert_idx]
|
|
lora_B.weight.data = buffers[key_b][expert_idx]
|
|
lora_A.weight.requires_grad_(True)
|
|
lora_B.weight.requires_grad_(True)
|
|
grad_view_a = grad_buffers["grad_" + key_a][expert_idx]
|
|
grad_view_b = grad_buffers["grad_" + key_b][expert_idx]
|
|
if authoritative_mode:
|
|
lora_A.weight.grad = None
|
|
lora_B.weight.grad = None
|
|
if authoritative_backend is not None:
|
|
authoritative_backend.register_authoritative_optimizer_grad(
|
|
f"lora.{key_a}.expert_{expert_idx}", lora_A.weight, grad_view_a
|
|
)
|
|
authoritative_backend.register_authoritative_optimizer_grad(
|
|
f"lora.{key_b}.expert_{expert_idx}", lora_B.weight, grad_view_b
|
|
)
|
|
else:
|
|
lora_A.weight.grad = grad_view_a
|
|
lora_B.weight.grad = grad_view_b
|
|
|
|
if not _first_logged:
|
|
_new_id_a = id(lora_A.weight)
|
|
_new_ptr_a = lora_A.weight.data_ptr()
|
|
_buf_ptr_a = buffers[key_a][expert_idx].data_ptr()
|
|
_has_grad = lora_A.weight.grad is not None
|
|
logger.info(
|
|
"[_replace_peft_weights_with_views] first param: "
|
|
"id %s->%s (same=%s) data_ptr %s->%s buf_ptr=%s (match=%s) "
|
|
"has_grad=%s requires_grad=%s shape=%s",
|
|
_old_id_a,
|
|
_new_id_a,
|
|
_old_id_a == _new_id_a,
|
|
_old_ptr_a,
|
|
_new_ptr_a,
|
|
_buf_ptr_a,
|
|
_new_ptr_a == _buf_ptr_a,
|
|
_has_grad,
|
|
lora_A.weight.requires_grad,
|
|
tuple(lora_A.weight.shape),
|
|
)
|
|
_first_logged = True
|
|
_replaced += 1
|
|
|
|
logger.info("[_replace_peft_weights_with_views] replaced %d param pairs", _replaced)
|
|
|
|
|
|
# =============================================================================
|
|
# Runtime LoRA Pointer Updates
|
|
# =============================================================================
|
|
|
|
|
|
def update_kt_lora_pointers(model: nn.Module):
|
|
"""Mark KT wrapper LoRA pointers and base weight pointers as dirty after optimizer.step()."""
|
|
wrappers = _find_kt_wrappers(model)
|
|
|
|
if wrappers:
|
|
for wrapper in wrappers:
|
|
if getattr(wrapper, "_kt_managed_lora_enabled", False):
|
|
wrapper._lora_pointers_dirty = True
|
|
# In full mode, base weights also need re-sync after optimizer step
|
|
if getattr(wrapper, "_full_weight_grad", False) and wrapper.wrapper is not None:
|
|
wrapper.wrapper._base_weights_dirty = True
|
|
backend = getattr(wrapper, "wrapper", None)
|
|
if backend is not None and getattr(backend, "_uses_authoritative_optimizer_grads", False):
|
|
backend.release_authoritative_optimizer_grads()
|
|
|
|
|
|
# =============================================================================
|
|
# Cross-Rank Gradient Synchronization
|
|
# =============================================================================
|
|
|
|
|
|
def sync_kt_lora_gradients(model: nn.Module) -> None:
|
|
"""Validate distributed KT gradient ownership without issuing collectives.
|
|
|
|
KT gathers every rank's routed rows before the rank-0 C++ backward, which
|
|
applies world-size normalization while producing optimizer gradients.
|
|
This compatibility entry point therefore must not all-reduce those
|
|
rank-0-owned gradients again. Ordinary registered GPU modules remain under
|
|
DDP/FSDP ownership and are deliberately untouched here.
|
|
"""
|
|
import torch.distributed as dist
|
|
|
|
if not (dist.is_initialized() and dist.get_world_size() > 1):
|
|
return
|
|
|
|
world_size = dist.get_world_size()
|
|
rank = dist.get_rank()
|
|
if world_size <= 1:
|
|
return
|
|
|
|
wrappers = _find_kt_wrappers(model)
|
|
if not wrappers:
|
|
return
|
|
|
|
# Distributed KT SFT gathers every rank's routed rows and writes a
|
|
# world-size-normalized optimizer gradient on rank 0. No gradient
|
|
# collective is needed here; validate ownership/aliases only. Ordinary
|
|
# GPU lora_experts are deliberately excluded and remain DDP/FSDP-managed.
|
|
authoritative_wrappers = [w for w in wrappers if getattr(w, "_uses_authoritative_optimizer_grads", False)]
|
|
if authoritative_wrappers:
|
|
if len(authoritative_wrappers) != len(wrappers):
|
|
raise RuntimeError("Mixed authoritative and legacy KT SFT backends are unsupported in one model")
|
|
|
|
for wrapper in wrappers:
|
|
backend = getattr(wrapper, "wrapper", None)
|
|
wrapped_world_size = int(getattr(wrapper, "_kt_world_size_at_wrap", world_size))
|
|
if wrapped_world_size != world_size:
|
|
raise RuntimeError(
|
|
f"Layer {wrapper.layer_idx}: KT wrapper was created for world_size={wrapped_world_size}, "
|
|
f"but the active process group has world_size={world_size}"
|
|
)
|
|
authoritative = bool(getattr(wrapper, "_uses_authoritative_optimizer_grads", False))
|
|
backend_description = "an authoritative KT backend" if authoritative else "a KT backend"
|
|
if rank == 0:
|
|
if backend is None:
|
|
raise RuntimeError(f"Layer {wrapper.layer_idx}: rank 0 does not own {backend_description}")
|
|
if authoritative:
|
|
backend.validate_authoritative_optimizer_grad_state()
|
|
else:
|
|
if backend is not None:
|
|
raise RuntimeError(f"Layer {wrapper.layer_idx}: rank {rank} unexpectedly owns {backend_description}")
|
|
if authoritative:
|
|
for param in _collect_wrapper_managed_lora_params(wrapper):
|
|
if param.grad is not None:
|
|
raise RuntimeError(
|
|
f"Layer {wrapper.layer_idx}: non-rank-0 KT LoRA Parameter unexpectedly has a gradient"
|
|
)
|
|
|
|
|
|
# =============================================================================
|
|
# Checkpoint Save/Load
|
|
# =============================================================================
|
|
|
|
|
|
def save_lora_experts_to_adapter(model: nn.Module, output_dir: str) -> None:
|
|
"""
|
|
Save LoRA Experts weights to adapter file by merging with existing Attention LoRA.
|
|
"""
|
|
from safetensors import safe_open
|
|
from safetensors.torch import save_file
|
|
|
|
wrappers = _find_kt_wrappers(model) or []
|
|
if not wrappers:
|
|
logger.warning("No KT wrappers found, skipping LoRA Experts saving")
|
|
return
|
|
|
|
adapter_file = os.path.join(output_dir, "adapter_model.safetensors")
|
|
if not os.path.exists(adapter_file):
|
|
adapter_file_bin = os.path.join(output_dir, "adapter_model.bin")
|
|
if os.path.exists(adapter_file_bin):
|
|
state_dict = torch.load(adapter_file_bin, map_location="cpu", weights_only=True)
|
|
else:
|
|
logger.warning(f"No existing adapter file found at {output_dir}, creating new one")
|
|
state_dict = {}
|
|
else:
|
|
state_dict = {}
|
|
with safe_open(adapter_file, framework="pt") as f:
|
|
for key in f.keys():
|
|
state_dict[key] = f.get_tensor(key)
|
|
|
|
lora_expert_count = 0
|
|
for wrapper in wrappers:
|
|
if wrapper.lora_experts is None:
|
|
continue
|
|
|
|
layer_idx = wrapper.layer_idx
|
|
for expert_idx, expert in enumerate(wrapper.lora_experts.experts):
|
|
base_key = f"base_model.model.model.layers.{layer_idx}.mlp.lora_experts.{expert_idx}"
|
|
state_dict[f"{base_key}.le_gate.weight"] = expert.le_gate.weight.data.cpu().clone()
|
|
state_dict[f"{base_key}.le_up.weight"] = expert.le_up.weight.data.cpu().clone()
|
|
state_dict[f"{base_key}.le_down.weight"] = expert.le_down.weight.data.cpu().clone()
|
|
lora_expert_count += 3
|
|
|
|
logger.debug(f"Added LoRA Experts for layer {layer_idx} ({len(wrapper.lora_experts.experts)} experts)")
|
|
|
|
output_file = os.path.join(output_dir, "adapter_model.safetensors")
|
|
save_file(state_dict, output_file, metadata={"format": "pt"})
|
|
|
|
logger.info(
|
|
f"Saved LoRA Experts to {output_file}: "
|
|
f"{len(wrappers)} layers, {lora_expert_count} LoRA Expert tensors added, "
|
|
f"{len(state_dict)} total tensors"
|
|
)
|
|
|
|
|
|
def save_kt_moe_to_adapter(model: nn.Module, output_dir: str) -> None:
|
|
"""
|
|
Unified function to save KT-managed MoE training weights.
|
|
|
|
Full/Hybrid expert base Parameters are written to the dedicated sharded KT
|
|
Full checkpoint. Per-expert PEFT LoRA remains owned and saved by PEFT.
|
|
LoRA Experts and fused expert LoRA keep their existing formats.
|
|
"""
|
|
wrappers = _find_kt_wrappers(model) or []
|
|
if not wrappers:
|
|
logger.info("[save_kt_moe] No KT wrappers found, skipping")
|
|
return
|
|
|
|
has_lora_experts = any(w.lora_experts is not None for w in wrappers)
|
|
has_fused_lora = any(
|
|
bool(getattr(w, "_fused_expert_lora_params", None)) for w in wrappers
|
|
)
|
|
has_full_weights = any(getattr(w, "_full_weight_grad", False) for w in wrappers)
|
|
|
|
if has_full_weights:
|
|
save_full_weight_checkpoint(wrappers, output_dir)
|
|
|
|
if has_lora_experts:
|
|
save_lora_experts_to_adapter(model, output_dir)
|
|
|
|
if has_fused_lora:
|
|
_save_fused_expert_lora(wrappers, output_dir)
|
|
|
|
if not has_full_weights and not has_lora_experts and not has_fused_lora:
|
|
logger.info("[save_kt_moe] No Full weights, lora_experts, or fused expert LoRA in KT wrappers")
|
|
|
|
|
|
def _save_fused_expert_lora(wrappers: list, output_dir: str) -> None:
|
|
"""Save fused expert LoRA params to a safetensors file."""
|
|
from safetensors.torch import save_file
|
|
|
|
names = ["gate_lora_a", "gate_lora_b", "up_lora_a", "up_lora_b", "down_lora_a", "down_lora_b"]
|
|
tensors = {}
|
|
for w in wrappers:
|
|
fused = getattr(w, "_fused_expert_lora_params", None)
|
|
if not fused:
|
|
continue
|
|
for param, name in zip(fused, names):
|
|
key = f"layers.{w.layer_idx}.experts.{name}"
|
|
tensors[key] = param.data.clone()
|
|
|
|
if tensors:
|
|
path = os.path.join(output_dir, "fused_expert_lora.safetensors")
|
|
with tempfile.NamedTemporaryFile(
|
|
dir=output_dir,
|
|
prefix=".fused_expert_lora.",
|
|
suffix=".safetensors.tmp",
|
|
delete=False,
|
|
) as temporary:
|
|
temporary_path = temporary.name
|
|
try:
|
|
save_file(tensors, temporary_path)
|
|
os.replace(temporary_path, path)
|
|
finally:
|
|
try:
|
|
os.remove(temporary_path)
|
|
except FileNotFoundError:
|
|
pass
|
|
logger.info(f"[save_kt_moe] Saved {len(tensors)} fused expert LoRA tensors to {path}")
|
|
|
|
|
|
def _load_fused_expert_lora(wrappers: list, adapter_path: str) -> int:
|
|
"""Strictly load fused expert LoRA tensors into rank-owned buffers."""
|
|
names = ["gate_lora_a", "gate_lora_b", "up_lora_a", "up_lora_b", "down_lora_a", "down_lora_b"]
|
|
owned_wrappers = {}
|
|
expected = {}
|
|
for wrapper in wrappers:
|
|
fused = getattr(wrapper, "_fused_expert_lora_params", None)
|
|
if not fused:
|
|
continue
|
|
if len(fused) != len(names):
|
|
raise RuntimeError(
|
|
f"Layer {wrapper.layer_idx}: expected {len(names)} fused expert LoRA "
|
|
f"parameters, found {len(fused)}"
|
|
)
|
|
if wrapper.layer_idx in owned_wrappers:
|
|
raise RuntimeError(
|
|
f"Duplicate owned fused expert LoRA layer {wrapper.layer_idx}"
|
|
)
|
|
owned_wrappers[wrapper.layer_idx] = wrapper
|
|
for param, name in zip(fused, names):
|
|
expected[f"layers.{wrapper.layer_idx}.experts.{name}"] = (wrapper, param)
|
|
|
|
if not expected:
|
|
return 0
|
|
|
|
path = os.path.join(adapter_path, "fused_expert_lora.safetensors")
|
|
if not os.path.isfile(path):
|
|
raise FileNotFoundError(
|
|
f"Missing fused_expert_lora.safetensors in adapter {adapter_path}"
|
|
)
|
|
|
|
from safetensors.torch import load_file
|
|
|
|
saved = load_file(path)
|
|
actual_keys = set(saved)
|
|
expected_keys = set(expected)
|
|
missing_keys = sorted(expected_keys - actual_keys)
|
|
unexpected_keys = sorted(actual_keys - expected_keys)
|
|
if missing_keys or unexpected_keys:
|
|
details = []
|
|
if missing_keys:
|
|
details.append(f"missing={missing_keys}")
|
|
if unexpected_keys:
|
|
details.append(f"unexpected={unexpected_keys}")
|
|
raise RuntimeError(
|
|
"Invalid fused_expert_lora.safetensors key set: " + ", ".join(details)
|
|
)
|
|
|
|
for key, (_, param) in expected.items():
|
|
tensor = saved[key]
|
|
if tensor.shape != param.shape:
|
|
raise RuntimeError(
|
|
f"{key} shape mismatch: expected {tuple(param.shape)}, "
|
|
f"got {tuple(tensor.shape)}"
|
|
)
|
|
if tensor.dtype != param.dtype:
|
|
raise RuntimeError(
|
|
f"{key} dtype mismatch: expected {param.dtype}, got {tensor.dtype}"
|
|
)
|
|
|
|
for key, (wrapper, param) in expected.items():
|
|
param.data.copy_(saved[key])
|
|
wrapper._lora_pointers_dirty = True
|
|
|
|
loaded_count = len(expected)
|
|
logger.info(f"[_load_fused_expert_lora] Loaded {loaded_count} tensors from {path}")
|
|
return loaded_count
|
|
|
|
|
|
def load_lora_experts_from_adapter(model: nn.Module, adapter_path: str) -> None:
|
|
"""
|
|
Load LoRA Experts weights from adapter file into KT wrappers.
|
|
"""
|
|
from safetensors import safe_open
|
|
|
|
wrappers = _find_kt_wrappers(model) or []
|
|
if not wrappers:
|
|
logger.warning("No KT wrappers found, skipping LoRA Experts loading")
|
|
return
|
|
|
|
wrapper_map = {w.layer_idx: w for w in wrappers if w.lora_experts is not None}
|
|
if not wrapper_map:
|
|
logger.warning("No LoRA Experts found in KT wrappers, skipping")
|
|
return
|
|
|
|
# Prefer dedicated lora_experts file, fallback to adapter file
|
|
adapter_file = os.path.join(adapter_path, "lora_experts.safetensors")
|
|
if not os.path.exists(adapter_file):
|
|
adapter_file = os.path.join(adapter_path, "adapter_model.safetensors")
|
|
if not os.path.exists(adapter_file):
|
|
adapter_file = os.path.join(adapter_path, "adapter_model.bin")
|
|
if not os.path.exists(adapter_file):
|
|
logger.warning(f"No lora_experts or adapter file found at {adapter_path}")
|
|
return
|
|
|
|
logger.info(f"Loading LoRA Experts from {adapter_file}")
|
|
|
|
lora_expert_pattern = re.compile(
|
|
r"base_model\.model\.model\.layers\.(\d+)\.mlp\.lora_experts\.(\d+)\.(le_gate|le_up|le_down)\.weight"
|
|
)
|
|
|
|
layer_weights = {}
|
|
with safe_open(adapter_file, framework="pt") as f:
|
|
for key in f.keys():
|
|
match = lora_expert_pattern.match(key)
|
|
if match:
|
|
layer_idx = int(match.group(1))
|
|
expert_idx = int(match.group(2))
|
|
proj_name = match.group(3)
|
|
layer_weights.setdefault(layer_idx, {}).setdefault(expert_idx, {})[proj_name] = f.get_tensor(key)
|
|
|
|
loaded_count = 0
|
|
for layer_idx, experts_dict in layer_weights.items():
|
|
if layer_idx not in wrapper_map:
|
|
logger.warning(f"No LoRA Experts for layer {layer_idx}, skipping")
|
|
continue
|
|
|
|
wrapper = wrapper_map[layer_idx]
|
|
for expert_idx, proj_dict in experts_dict.items():
|
|
if expert_idx >= len(wrapper.lora_experts.experts):
|
|
continue
|
|
expert = wrapper.lora_experts.experts[expert_idx]
|
|
if "le_gate" in proj_dict:
|
|
expert.le_gate.weight.data.copy_(proj_dict["le_gate"].to(expert.le_gate.weight.device))
|
|
if "le_up" in proj_dict:
|
|
expert.le_up.weight.data.copy_(proj_dict["le_up"].to(expert.le_up.weight.device))
|
|
if "le_down" in proj_dict:
|
|
expert.le_down.weight.data.copy_(proj_dict["le_down"].to(expert.le_down.weight.device))
|
|
loaded_count += 1
|
|
|
|
logger.info(f"Loaded LoRA Experts for {loaded_count} experts from {adapter_path}")
|
|
|
|
|
|
def load_kt_moe_from_adapter(model: nn.Module, adapter_path: str) -> None:
|
|
"""
|
|
Restore KT-managed MoE training weights without replacing Parameters.
|
|
|
|
Per-expert PEFT LoRA remains owned and loaded by PEFT. Full/Hybrid base
|
|
weights, LoRA Experts, and fused expert LoRA are restored here.
|
|
"""
|
|
wrappers = _find_kt_wrappers(model) or []
|
|
if not wrappers:
|
|
logger.warning("No KT wrappers found, skipping KT MoE loading")
|
|
return
|
|
|
|
has_lora_experts = any(w.lora_experts is not None for w in wrappers)
|
|
has_fused_lora = any(
|
|
bool(getattr(w, "_fused_expert_lora_params", None)) for w in wrappers
|
|
)
|
|
has_full_weights = any(getattr(w, "_full_weight_grad", False) for w in wrappers)
|
|
|
|
if has_full_weights:
|
|
load_full_weight_checkpoint(wrappers, adapter_path)
|
|
|
|
if has_lora_experts:
|
|
load_lora_experts_from_adapter(model, adapter_path)
|
|
|
|
if has_fused_lora:
|
|
_load_fused_expert_lora(wrappers, adapter_path)
|
|
|
|
if not has_full_weights and not has_lora_experts and not has_fused_lora:
|
|
logger.info("No Full weights, lora_experts, or fused expert LoRA in KT wrappers")
|