Files
ktransformers/kt-kernel/python/sft/config.py
yyj aa2cd3072f [feat] Add native RAWINT4 routed-expert LoRA SFT for Kimi K2.5/K2.6 (#2186)
* [feat](sft): add K2 MoE SFT backend

* [feat](sft): enable packed KGroup backward for K2 SFT

* [feat](serving): add KT LoRA to SGLang adapter converter

* [test]: cover KT SFT and composite LoRA serving paths

* [perf](bench): add reusable KT E2E benchmark harness

* [test](sft): add K2 MoE AMX and real Kimi load smoke cases

* [perf](bench): add K2 SFT packed backward profile harness

* [test](sglang): cover Qwen3.5 LoRA shape filters

* [perf](sft): add K2 forward profiling logs

* [docs](sft): clarify K2 packed backward flow

* [docs](sft): number K2 backward comment steps

* chore: mirror sap4 loader baseline

* Add K2 SFT profiling ranges

* Add backward NVTX ranges for K2 SFT profiling

* perf(sft): add rank8 LoRA backward vector paths

* perf(sft): reuse down LoRA bprop cache

* [perf]: add K2 RAWINT4 prefill mat-mat dispatch

* [perf] Refactor K2 RAWINT4 prefill accumulation

* [perf]: add blocked RAWINT4 prefill backend

Cherry-picked from be40f8e with local amx.py conflict resolution.

* Optimize blocked RAWINT4 decode vecmul

* perf(sft): checkpoint qj5090 packed backward baseline

* perf(sft): batch expert LoRA backward as matmul

* perf(sft): align K2 base backward with BF16 matmul

* Optimize K2 SFT backward workspace path

* docs(sft): hand off RAWINT4 KGroup build and Kimi configs

* [feat](kt-kernel): add Kimi RAWINT4 LoRA SFT kernels

* [feat](kt-kernel): integrate Kimi RAWINT4 LoRA SFT

* [fix](sft): preserve input gradients through frozen routers

Build the routing graph when either activations or router parameters require gradients, while respecting the ambient grad mode. Keep frozen router parameters frozen and reject detached routing weights whenever gradients are needed.

Correct the existing frozen-router expectation and cover input-gradient equivalence for linear, DeepSeek/Kimi and GLM routing, checkpoint reuse, grad modes and detached-output failures.

Validated 62 routing/shared-expert/checkpoint tests on PyTorch 2.9.1 and a two-step Kimi K2.5 RAWINT4 four-GPU S512 LoRA smoke with per-rank routing-input gradient audits.

* [perf](kt-kernel): accelerate RAWINT4 prefill and backward

Reuse decoded weight columns across prefill rows and prepack BF16 backward coefficients into bounded per-worker panels. Use AMX BF16 tiles when available, retaining the AVX512-BF16 path and frozen base weights.

Balance expert work, parallelize cache copies and router reductions, and correct the signed INT4 -8 dot-product boundary. Add a CPU replay benchmark and forward/backward panel regressions.

Validated on sap4 with Kimi K2.5, TP1/TP2 replay, AMX and AVX-512 multi-GPU training, gradient windows, and saved LoRA artifacts. User-facing clean builds and AMD hardware remain separate validation work.

* [fix](sft): limit default OpenMP pools to the CPU owner

Keep the affinity-visible physical-core budget on rank 0 and default non-owner distributed ranks to one thread. Preserve explicit generic and KT-specific thread overrides, including intentional single-thread settings.

Cover single-process behavior, non-owner ranks, repeated configuration, and explicit overrides. The 109 configuration, activation-policy, and backend-dispatch tests pass.

---------

Co-authored-by: yyj <yyj@qj5090.local>
Co-authored-by: djh <jd4136@columbia.edu>
2026-09-10 17:05:16 +08:00

683 lines
27 KiB
Python

# KT-Kernel SFT configuration
# SPDX-License-Identifier: Apache-2.0
"""
KTConfig: kt-kernel's own configuration dataclass.
This is the kt-kernel equivalent of DeepSpeed's JSON config —
it holds all kt-kernel-specific settings and is passed through
KTransformersPlugin.kt_config (similar to DeepSpeedPlugin.hf_ds_config).
"""
from __future__ import annotations
import copy
import dataclasses
import inspect
import json
import logging
import os
import warnings
from collections.abc import Mapping
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Callable, Literal
from .backend import (
FP8_BACKEND,
INT8_BACKEND,
RAWINT4_BACKEND,
normalize_sft_backend,
)
logger = logging.getLogger(__name__)
_CPU_TOPOLOGY_ROOT = Path("/sys/devices/system/cpu")
_ACTIVATION_POLICY_ENV = "ACCELERATE_KT_ACTIVATION_POLICY"
_LEGACY_REUSE_ENV = "KT_REUSE_CHECKPOINT_FORWARD"
_KNOWN_SFT_BACKENDS = {
"auto",
"int8",
"fp8",
"amxbf16",
"amxint8",
"amxfp8",
"amxint4",
"rawint4",
"amxint4_kgroup",
"amxbf16_skiplora",
"amxint8_skiplora",
"amxint4_skiplora",
}
ActivationRetention = Literal["retain", "recompute"]
ExpertWeightFormat = Literal["bf16", "int8", "fp8", "rawint4"]
WeightLifecycle = Literal["persistent", "ephemeral"]
@dataclass(frozen=True)
class KTActivationPolicy:
"""Activation lifetime policy shared by every KT distributed rank."""
cpu: ActivationRetention = "recompute"
gpu: ActivationRetention = "recompute"
def __post_init__(self) -> None:
valid = {"retain", "recompute"}
if self.cpu not in valid:
raise ValueError(
f"activation_policy.cpu must be one of {sorted(valid)}, got {self.cpu!r}"
)
if self.gpu not in valid:
raise ValueError(
f"activation_policy.gpu must be one of {sorted(valid)}, got {self.gpu!r}"
)
if self.cpu == "recompute" and self.gpu == "retain":
raise NotImplementedError(
"activation_policy cpu=recompute, gpu=retain is not implemented"
)
@classmethod
def from_value(
cls,
value: "KTActivationPolicy | Mapping[str, str]",
) -> "KTActivationPolicy":
if isinstance(value, cls):
return value
if not isinstance(value, Mapping):
raise TypeError(
"kt_activation_policy must be a KTActivationPolicy or a mapping "
"with exactly the keys 'cpu' and 'gpu'"
)
expected = {"cpu", "gpu"}
actual = set(value)
if actual != expected:
missing = sorted(expected - actual)
unexpected = sorted(actual - expected)
details = []
if missing:
details.append(f"missing={missing}")
if unexpected:
details.append(f"unexpected={unexpected}")
raise ValueError(
"kt_activation_policy must contain exactly 'cpu' and 'gpu'"
+ (f" ({', '.join(details)})" if details else "")
)
return cls(cpu=value["cpu"], gpu=value["gpu"])
def _legacy_activation_policy_from_env(value: str) -> KTActivationPolicy:
normalized = value.strip().lower()
if normalized in {"1", "true", "yes", "on"}:
cpu = "retain"
elif normalized in {"0", "false", "no", "off"}:
cpu = "recompute"
else:
raise ValueError(
f"{_LEGACY_REUSE_ENV} must be a boolean value, got {value!r}"
)
warnings.warn(
f"{_LEGACY_REUSE_ENV} is deprecated; configure kt_activation_policy "
"through the training frontend instead",
FutureWarning,
stacklevel=3,
)
return KTActivationPolicy(cpu=cpu, gpu="recompute")
def _activation_policy_from_env(value: str) -> KTActivationPolicy:
try:
policy = json.loads(value)
except json.JSONDecodeError as exc:
raise ValueError(
f"{_ACTIVATION_POLICY_ENV} must be a JSON object with exactly "
"the keys 'cpu' and 'gpu'"
) from exc
return KTActivationPolicy.from_value(policy)
def _env_int(key: str, default: int | None) -> int | None:
value = os.environ.get(key, None)
if value is None or value == "":
return default
return int(value)
def _env_float(key: str, default: float | None) -> float | None:
value = os.environ.get(key, None)
if value is None or value == "":
return default
return float(value)
def _env_bool(key: str, default: bool) -> bool:
value = os.environ.get(key, None)
if value is None or value == "":
return default
return value.lower() in ("1", "true", "yes")
def _available_cpu_ids() -> set[int]:
"""Return CPUs available to this process, respecting affinity/cpuset limits."""
try:
return set(os.sched_getaffinity(0))
except (AttributeError, OSError):
return set(range(os.cpu_count() or 1))
def _read_cpu_topology(cpu_id: int) -> tuple[int, int] | None:
topology = _CPU_TOPOLOGY_ROOT / f"cpu{cpu_id}" / "topology"
try:
package_id = int((topology / "physical_package_id").read_text().strip())
core_id = int((topology / "core_id").read_text().strip())
except (OSError, ValueError):
return None
return package_id, core_id
def detect_physical_cpu_count() -> int:
"""Count physical cores available to the current process.
Linux exposes a stable ``(physical_package_id, core_id)`` pair for every
logical CPU. Counting those pairs avoids assigning one OpenMP worker to
each SMT sibling. If topology is unavailable, fall back to the number of
affinity-visible logical CPUs.
"""
cpu_ids = _available_cpu_ids()
physical_cores = {
topology
for cpu_id in cpu_ids
if (topology := _read_cpu_topology(cpu_id)) is not None
}
return max(1, len(physical_cores) if physical_cores else len(cpu_ids))
def _set_torch_num_threads(num_threads: int) -> None:
try:
import torch
except ImportError:
return
torch.set_num_threads(num_threads)
def configure_omp_threads() -> int:
"""Configure OpenMP for KT SFT CPU tensor work.
``accelerate launch`` defaults GPU jobs to ``OMP_NUM_THREADS=1`` when the
caller did not choose a value. The authoritative CPU rank needs a physical
core budget for CPU tensor work; other distributed ranks keep one thread
to avoid creating a machine-sized OpenMP pool in every process.
``ACCELERATE_KT_OMP_NUM_THREADS`` is the unambiguous KT-specific override,
including when an intentional single-thread run is required. An existing
generic ``OMP_NUM_THREADS`` value greater than one is also preserved.
"""
kt_override = _env_int("ACCELERATE_KT_OMP_NUM_THREADS", None)
current_omp = _env_int("OMP_NUM_THREADS", None)
if kt_override is not None:
num_threads = kt_override
source = "ACCELERATE_KT_OMP_NUM_THREADS"
elif current_omp is not None and current_omp > 1:
num_threads = current_omp
source = "OMP_NUM_THREADS"
elif _env_int("WORLD_SIZE", 1) > 1 and _env_int("RANK", 0) != 0:
num_threads = 1
source = "non-owner distributed rank"
else:
num_threads = detect_physical_cpu_count()
source = "available physical cores"
if num_threads < 1:
raise ValueError(f"OpenMP thread count must be positive, got {num_threads}")
os.environ["OMP_NUM_THREADS"] = str(num_threads)
_set_torch_num_threads(num_threads)
logger.info("KT SFT configured OMP_NUM_THREADS=%d from %s", num_threads, source)
return num_threads
@dataclass
class KTConfig:
"""
KT-Kernel configuration for SFT training.
All field names use the ``kt_`` prefix so they match the dict keys used in
HfTrainerKTConfig / YAML configs. This means ``KTConfig(**dict)`` works
directly — no name-mapping or prefix-stripping needed.
Can be created from:
- Direct construction: KTConfig(kt_backend="AMXBF16", kt_weight_path="/path/...")
- Dict: KTConfig(**config_dict)
- Environment variables: KTConfig() reads ACCELERATE_KT_* env vars as defaults
"""
# Backend selection
kt_backend: str | None = None
kt_num_threads: int | None = None
kt_tp_enabled: bool | None = None
kt_threadpool_count: int | None = None
# Weight loading
kt_weight_path: str | None = None
kt_non_expert_weight_path: str | None = None
kt_expert_weight_format: ExpertWeightFormat | str | None = None
kt_weight_lifecycle: WeightLifecycle | str | None = None
kt_expert_checkpoint_path: str | None = (
None # HF expert checkpoint or KT Full checkpoint directory
)
kt_num_gpu_experts: int | None = None
kt_skip_expert_loading: bool | None = None
kt_share_backward_bb: bool | None = None # defaults off for RAWINT4, on otherwise
kt_share_cache_pool: bool | None = (
None # auto-set by trainer_config_process, not user-facing
)
kt_force_fused_expert_lora: bool | None = None
# Cache
kt_max_cache_depth: int | None = None
kt_model_max_length: int | None = None
kt_activation_policy: KTActivationPolicy | Mapping[str, str] | None = None
# LoRA
kt_lora_rank: int | None = None
kt_lora_alpha: float | None = None
kt_lora_dropout: float | None = None
# Training mode
kt_train_mode: str | None = None # "lora" | "full" | "hybrid"
kt_full_weight_grad: bool | None = (
None # auto-set True when train_mode in (full, hybrid)
)
# LoRA Experts (GPU-side extra experts)
kt_use_lora_experts: bool | None = None
kt_lora_expert_num: int | None = None
kt_lora_expert_intermediate_size: int | None = None
# Runtime state (set during wrapping, not by user)
kt_checkpoint_files: list[str] | None = None
kt_sharded_metadata: dict | None = None
# Custom wrapping
kt_wrap_fn: Callable[..., Any] | None = None
kt_wrap_kwargs: dict[str, Any] | None = None
@classmethod
def from_object(cls, obj: Any) -> "KTConfig":
"""Create KTConfig from a mapping or compatible public container."""
return cls._from_object(obj, seen=set())
@classmethod
def _from_object(cls, obj: Any, *, seen: set[int]) -> "KTConfig":
if isinstance(obj, cls):
return obj
if isinstance(obj, Mapping):
return cls(**cls._validated_mapping(obj))
object_id = id(obj)
if object_id in seen:
raise ValueError("Cyclic KT configuration container")
seen.add(object_id)
field_names = {field.name for field in dataclasses.fields(cls)}
attributes = vars(obj) if hasattr(obj, "__dict__") else {}
unknown = sorted(
name
for name in attributes
if name.startswith("kt_")
and name not in field_names
and name != "kt_config"
)
if unknown:
raise TypeError(f"Unknown KTConfig fields: {unknown}")
missing = object()
has_public_container = False
for container_name in ("kt_config", "config"):
if inspect.getattr_static(obj, container_name, missing) is missing:
continue
has_public_container = True
nested = getattr(obj, container_name, None)
if nested is None:
continue
if nested is obj:
raise ValueError(f"Cyclic KT configuration container: {container_name}")
if isinstance(nested, Mapping):
kwargs = cls._validated_mapping(nested)
for name in (
"kt_checkpoint_files",
"kt_sharded_metadata",
"kt_skip_expert_loading",
):
value = getattr(obj, name, None)
if value is not None:
kwargs[name] = value
return cls(**kwargs)
resolved = cls._from_object(nested, seen=seen)
overrides: dict[str, Any] = {}
for name in (
"kt_checkpoint_files",
"kt_sharded_metadata",
"kt_skip_expert_loading",
):
value = getattr(obj, name, None)
if value is not None and value is not getattr(resolved, name, None):
overrides[name] = value
if not overrides:
return resolved
result = copy.copy(resolved)
for name, value in overrides.items():
setattr(result, name, value)
return result
if has_public_container:
return cls()
raise TypeError(
"KT configuration must be a KTConfig, mapping, or public container "
"with a static 'kt_config' or 'config' attribute"
)
@classmethod
def _validated_mapping(cls, obj: Mapping[str, Any]) -> dict[str, Any]:
payload = dict(obj)
invalid_keys = sorted(
repr(name) for name in payload if not isinstance(name, str)
)
if invalid_keys:
raise TypeError(f"KTConfig field names must be strings: {invalid_keys}")
if "enabled" in payload:
enabled = payload.pop("enabled")
if enabled is not None and not isinstance(enabled, bool):
raise TypeError("KT framework field 'enabled' must be a bool or None")
field_names = {field.name for field in dataclasses.fields(cls)}
unknown = sorted(set(payload).difference(field_names))
if unknown:
raise TypeError(f"Unknown KTConfig fields: {unknown}")
return payload
def __post_init__(self):
configure_omp_threads()
explicit_backend = self.kt_backend is not None
env_backend = os.environ.get("ACCELERATE_KT_BACKEND")
if self.kt_expert_weight_format is None:
self.kt_expert_weight_format = os.environ.get(
"ACCELERATE_KT_EXPERT_WEIGHT_FORMAT"
)
if self.kt_expert_weight_format is not None:
self.kt_expert_weight_format = (
str(self.kt_expert_weight_format).strip().lower()
)
if self.kt_expert_weight_format not in {"bf16", "int8", "fp8", "rawint4"}:
raise ValueError(
"kt_expert_weight_format must be one of "
"['bf16', 'fp8', 'int8', 'rawint4'], "
f"got {self.kt_expert_weight_format!r}"
)
if self.kt_weight_lifecycle is None:
self.kt_weight_lifecycle = os.environ.get(
"ACCELERATE_KT_WEIGHT_LIFECYCLE", "persistent"
)
self.kt_weight_lifecycle = str(self.kt_weight_lifecycle).strip().lower()
if self.kt_weight_lifecycle not in {"persistent", "ephemeral"}:
raise ValueError(
"kt_weight_lifecycle must be one of ['ephemeral', 'persistent'], "
f"got {self.kt_weight_lifecycle!r}"
)
forwarded_policy = os.environ.get(_ACTIVATION_POLICY_ENV)
forwarded_policy_is_explicit = (
forwarded_policy is not None and forwarded_policy.strip() != ""
)
legacy_reuse = os.environ.get(_LEGACY_REUSE_ENV)
legacy_reuse_is_explicit = (
legacy_reuse is not None and legacy_reuse.strip() != ""
)
if self.kt_activation_policy is not None and forwarded_policy_is_explicit:
raise ValueError(
f"kt_activation_policy conflicts with {_ACTIVATION_POLICY_ENV}; "
f"unset {_ACTIVATION_POLICY_ENV}"
)
if self.kt_activation_policy is not None and legacy_reuse_is_explicit:
raise ValueError(
f"kt_activation_policy conflicts with legacy {_LEGACY_REUSE_ENV}; "
f"unset {_LEGACY_REUSE_ENV}"
)
if forwarded_policy_is_explicit and legacy_reuse_is_explicit:
raise ValueError(
f"{_ACTIVATION_POLICY_ENV} conflicts with legacy "
f"{_LEGACY_REUSE_ENV}; unset {_LEGACY_REUSE_ENV}"
)
if self.kt_activation_policy is None:
if forwarded_policy_is_explicit:
self.kt_activation_policy = _activation_policy_from_env(
forwarded_policy
)
elif legacy_reuse_is_explicit:
self.kt_activation_policy = _legacy_activation_policy_from_env(
legacy_reuse
)
else:
self.kt_activation_policy = KTActivationPolicy()
else:
self.kt_activation_policy = KTActivationPolicy.from_value(
self.kt_activation_policy
)
if self.kt_backend is None:
if env_backend:
self.kt_backend = env_backend
elif self.kt_expert_weight_format in {"int8", "fp8", "rawint4"}:
self.kt_backend = "auto"
else:
self.kt_backend = "AMXBF16"
backend_lower = str(self.kt_backend).lower()
if backend_lower not in _KNOWN_SFT_BACKENDS:
raise ValueError(
f"unknown kt_backend {self.kt_backend!r}; expected one of "
f"{sorted(_KNOWN_SFT_BACKENDS)}"
)
self.kt_backend = normalize_sft_backend(
self.kt_backend,
expert_weight_format=self.kt_expert_weight_format,
)
backend_lower = str(self.kt_backend).lower()
if self.kt_expert_weight_format is None:
if backend_lower == INT8_BACKEND.lower():
# Backward compatibility for the legacy environment-only entry.
self.kt_expert_weight_format = "int8"
elif backend_lower == FP8_BACKEND.lower():
self.kt_expert_weight_format = "fp8"
elif backend_lower == RAWINT4_BACKEND.lower():
self.kt_expert_weight_format = "rawint4"
elif backend_lower == "amxbf16":
self.kt_expert_weight_format = "bf16"
expected_backend = {
"bf16": "amxbf16",
"int8": INT8_BACKEND.lower(),
"fp8": FP8_BACKEND.lower(),
"rawint4": RAWINT4_BACKEND.lower(),
}.get(self.kt_expert_weight_format)
if expected_backend is not None and backend_lower != expected_backend:
source = "kt_backend" if explicit_backend else "ACCELERATE_KT_BACKEND"
raise ValueError(
f"kt_expert_weight_format={self.kt_expert_weight_format!r} conflicts "
f"with {source}={self.kt_backend!r}"
)
if self.kt_num_threads is None:
self.kt_num_threads = _env_int("ACCELERATE_KT_NUM_THREADS", 1)
if self.kt_tp_enabled is None:
self.kt_tp_enabled = _env_bool("ACCELERATE_KT_TP_ENABLED", False)
if self.kt_threadpool_count is None:
self.kt_threadpool_count = _env_int("ACCELERATE_KT_THREADPOOL_COUNT", 1)
if self.kt_weight_path is None:
self.kt_weight_path = os.environ.get("ACCELERATE_KT_WEIGHT_PATH", None)
if self.kt_non_expert_weight_path is None:
self.kt_non_expert_weight_path = os.environ.get(
"ACCELERATE_KT_NON_EXPERT_WEIGHT_PATH", None
)
if self.kt_expert_checkpoint_path is None:
self.kt_expert_checkpoint_path = os.environ.get(
"ACCELERATE_KT_EXPERT_CHECKPOINT_PATH", None
)
if self.kt_num_gpu_experts is None:
self.kt_num_gpu_experts = _env_int("ACCELERATE_KT_NUM_GPU_EXPERTS", 0)
if self.kt_max_cache_depth is None:
self.kt_max_cache_depth = _env_int("ACCELERATE_KT_MAX_CACHE_DEPTH", 2)
if self.kt_share_backward_bb is None:
self.kt_share_backward_bb = _env_bool(
"ACCELERATE_KT_SHARE_BACKWARD_BB",
self.kt_expert_weight_format != "rawint4",
)
if self.kt_share_cache_pool is None:
self.kt_share_cache_pool = False
if self.kt_force_fused_expert_lora is None:
self.kt_force_fused_expert_lora = _env_bool(
"ACCELERATE_KT_FORCE_FUSED_EXPERT_LORA", False
)
if self.kt_use_lora_experts is None:
self.kt_use_lora_experts = _env_bool(
"ACCELERATE_KT_USE_LORA_EXPERTS", False
)
if self.kt_lora_expert_num is None:
self.kt_lora_expert_num = _env_int("ACCELERATE_KT_LORA_EXPERT_NUM", None)
if self.kt_lora_expert_intermediate_size is None:
self.kt_lora_expert_intermediate_size = _env_int(
"ACCELERATE_KT_LORA_EXPERT_INTERMEDIATE_SIZE", None
)
if self.kt_lora_rank is None:
self.kt_lora_rank = _env_int("ACCELERATE_KT_LORA_RANK", None)
if self.kt_lora_alpha is None:
self.kt_lora_alpha = _env_float("ACCELERATE_KT_LORA_ALPHA", None)
if self.kt_lora_alpha is None and self.kt_lora_rank is not None:
self.kt_lora_alpha = float(self.kt_lora_rank * 2)
if self.kt_lora_dropout is None:
self.kt_lora_dropout = _env_float("ACCELERATE_KT_LORA_DROPOUT", 0.0)
if not 0.0 <= self.kt_lora_dropout < 1.0:
raise ValueError(
f"kt_lora_dropout must be in [0, 1), got {self.kt_lora_dropout}"
)
if self.kt_train_mode is None:
self.kt_train_mode = os.environ.get("ACCELERATE_KT_TRAIN_MODE", "lora")
if self.kt_full_weight_grad is None:
self.kt_full_weight_grad = self.kt_train_mode in ("full", "hybrid")
if self.kt_model_max_length is None:
self.kt_model_max_length = _env_int("ACCELERATE_KT_MODEL_MAX_LENGTH", None)
if self.kt_skip_expert_loading is None:
if "ACCELERATE_KT_SKIP_EXPERT_LOADING" in os.environ:
self.kt_skip_expert_loading = _env_bool(
"ACCELERATE_KT_SKIP_EXPERT_LOADING", True
)
if self.kt_non_expert_weight_path and self.kt_expert_weight_format != "int8":
raise ValueError(
"kt_non_expert_weight_path is supported only with "
"kt_expert_weight_format='int8'"
)
if self.kt_expert_weight_format == "int8":
if str(self.kt_backend).lower() != INT8_BACKEND.lower():
raise ValueError(
"INT8 SFT requires kt_backend='auto' or kt_backend='INT8'"
)
if self.kt_train_mode != "lora" or bool(self.kt_full_weight_grad):
raise ValueError(
"INT8 SFT supports frozen-base LoRA only; Full and Hybrid are not supported"
)
if int(self.kt_num_gpu_experts or 0) != 0:
raise ValueError("INT8 SFT requires kt_num_gpu_experts=0")
if bool(self.kt_use_lora_experts):
raise ValueError(
"INT8 SFT does not support GPU LoRA experts; use pure expert LoRA"
)
if not bool(self.kt_share_backward_bb):
raise ValueError("INT8 SFT requires kt_share_backward_bb=true")
if self.kt_expert_checkpoint_path:
raise ValueError(
"INT8 SFT requires pre-quantized .kt weights; "
"kt_expert_checkpoint_path is not supported"
)
if not self.kt_weight_path:
raise ValueError(
"INT8 SFT requires kt_weight_path pointing to pre-quantized .kt weights"
)
if self.kt_expert_weight_format == "fp8":
if str(self.kt_backend).lower() != FP8_BACKEND.lower():
raise ValueError(
"FP8 SFT requires kt_backend='auto' or kt_backend='FP8'"
)
if self.kt_train_mode != "lora" or bool(self.kt_full_weight_grad):
raise ValueError(
"FP8 SFT supports frozen-base LoRA only; Full and Hybrid are not supported"
)
if not self.kt_lora_rank or int(self.kt_lora_rank) <= 0:
raise ValueError("FP8 SFT requires kt_lora_rank > 0")
if int(self.kt_num_gpu_experts or 0) != 0:
raise ValueError("FP8 SFT requires kt_num_gpu_experts=0")
if bool(self.kt_use_lora_experts):
raise ValueError(
"FP8 SFT does not support GPU LoRA experts; use pure expert LoRA"
)
if not bool(self.kt_share_backward_bb):
raise ValueError("FP8 SFT requires kt_share_backward_bb=true")
if self.kt_weight_lifecycle != "persistent":
raise ValueError("FP8 SFT requires kt_weight_lifecycle='persistent'")
if self.kt_expert_weight_format == "rawint4":
if str(self.kt_backend).lower() != RAWINT4_BACKEND.lower():
raise ValueError(
"RAWINT4 SFT requires kt_backend='auto' or kt_backend='RAWINT4'"
)
if self.kt_train_mode != "lora" or bool(self.kt_full_weight_grad):
raise ValueError(
"RAWINT4 SFT supports frozen-base LoRA only; Full and Hybrid are not supported"
)
if int(self.kt_lora_rank or 0) != 8:
raise ValueError("RAWINT4 SFT v1 requires kt_lora_rank=8")
if float(self.kt_lora_dropout or 0.0) != 0.0:
raise ValueError("RAWINT4 SFT v1 requires kt_lora_dropout=0")
effective_threadpool_count = (
int(self.kt_threadpool_count or 1) if self.kt_tp_enabled else 1
)
if effective_threadpool_count not in {1, 2}:
raise ValueError(
"RAWINT4 SFT v1 supports TP1/TP2 only; "
f"effective threadpool_count={effective_threadpool_count}"
)
if int(self.kt_num_gpu_experts or 0) != 0:
raise ValueError("RAWINT4 SFT requires kt_num_gpu_experts=0")
if bool(self.kt_use_lora_experts):
raise ValueError(
"RAWINT4 SFT does not support GPU LoRA experts; use pure expert LoRA"
)
if bool(self.kt_share_backward_bb):
raise ValueError(
"RAWINT4 SFT computes backward directly from packed weights; "
"kt_share_backward_bb must be false"
)
if not bool(self.kt_force_fused_expert_lora):
raise ValueError("RAWINT4 SFT requires kt_force_fused_expert_lora=true")
if self.kt_expert_checkpoint_path:
raise ValueError(
"RAWINT4 SFT reads packed experts from the source checkpoint; "
"kt_expert_checkpoint_path is not supported"
)
if self.kt_weight_lifecycle != "persistent":
raise ValueError(
"RAWINT4 SFT requires kt_weight_lifecycle='persistent'"
)
if (
self.kt_activation_policy.cpu == "retain"
and self.kt_activation_policy.gpu == "recompute"
):
raise ValueError(
"RAWINT4 SFT v1 supports activation_policy "
"cpu=retain,gpu=retain for ordinary forward/backward caching, "
"but checkpoint-forward reuse (cpu=retain,gpu=recompute) is not supported"
)
if self.kt_weight_lifecycle == "ephemeral":
if self.kt_expert_weight_format != "int8":
raise ValueError(
"kt_weight_lifecycle='ephemeral' is supported only for INT8 .kt weights"
)