mirror of
https://github.com/kvcache-ai/ktransformers.git
synced 2026-10-17 05:11:19 +08:00
yyj
a5d7ad9047
* feat(sft): add native MXFP4 routed-expert LoRA kernels * test(sft): validate DeepSeek V4 MXFP4 routed experts * feat(sft): AVX2 tier for MXFP4 routed-expert LoRA SFT MXFP4_SFT_MOE on non-AVX512 builds, same Python contract as the AMX / AVX512-BF16 kernel from #2185 (KTMoEWrapper(method="MXFP4_SFT")), for AVX2+FMA+F16C hosts (Zen 2/3 EPYC, older Xeons). * operators/avx2/sft_mxfp4-moe.hpp: AVX2_SFT_MXFP4_MOE_TP part on top of AVX2_MXFP4_MOE_TP (frozen packed E2M1/UE8M0 base, LoRA gate/up/down, V4 swiglu clamp, transpose-free packed-weight dX, authoritative LoRA grads). * operators/moe-sft-tp.hpp: AVX2 result merge; AVX-512 merge loops and the block-FP8 staging guarded so the wrapper builds without AVX-512. * ext_bindings.cpp: bind MXFP4_SFT_MOE -> TP_MOE_SFT<AVX2_SFT_MXFP4_MOE_TP> when USE_AMX_AVX_KERNEL is off; MOESFTBindings no longer AMX-only. * python/sft/backend.py: runtime reports cpu_variant=avx2, kernel avx2-mxfp4-g32. * test_mxfp4_sft_numerical.py: inference regression uses AVX2MXFP4_MOE when AMXFP4_KGroup_MOE is absent. Validated on 2x EPYC 7452 (AVX2, no AVX-512): test_mxfp4_sft_numerical.py TP1/TP2 all stages rel-L2 <= 0.4% (limits 3-5%); validate_dsv4_mxfp4_sft_layer.py on DeepSeek-V4-Flash-0731 layer 3, TP1 and TP2: passed, forward/grad_input/ LoRA grads rel-L2 0.10-0.29%, base weights unchanged. Note: the build uses -ffast-math, so exp() past the float range is undefined; the scalar sigmoid tails clamp the argument like the vector path. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011Mgsduzz2fx1gSvfHbR8RZ * fix(sft): clamp MXFP4 scalar-tail sigmoid exponent * feat(sft): integrate native V4 experts with LoRA training and export Wrap only routed experts while preserving V4 routing and shared branches. Validate native checkpoint provenance and stream a BF16 non-expert cache without dequantizing routed experts. Add audited static SGLang export and reproducible training, resume and generation acceptance helpers. CPU retain/GPU recompute GAS1/GAS4 smoke and 265 KT regression tests pass; long-run quality and resume acceptance remain in progress. * fix(sft): retain async inference buffers until CUDA work completes * test(sft): isolate V4 adapter effects with matched-kernel controls Allow baseline and non-expert-only exports to retain a zero-B expert adapter so all ablations execute the same native forward kernel. Add paired held-out generation reporting with explicitly diagnostic lexical metrics. * test(sft): compare PEFT targets semantically during V4 resume audits * test(sft): cancel queued V4 inference checks after a failed request * test(sft): preserve raw V4 responses before finite-output validation * feat: prepare V4 cache atomically and package static adapter export * build: preserve accepted CUDA carriers across Python ABIs * test: retain FP8 loading contracts after V4 wrapper integration * test: exclude historical concatenated-sample smoke from delivery * test(sft): audit unpacked batches and single-GPU resume state * docs(sft): align April V4 recipe with approved unpacked delivery * test: keep frozen-parameter audit indices exact for large tensors * test: record the loaded native backend in acceptance evidence * test: timestamp training audit events for startup measurements * docs: pin V4 serving dependencies and account for reserved token slots * [refactor](sft): keep V4 follow-up focused on runtime integration --------- Co-authored-by: 05yuki <ceu.azul.samu@gmail.com> Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
211 lines
6.3 KiB
Python
211 lines
6.3 KiB
Python
# SFT (Supervised Fine-Tuning) submodule for kt-kernel
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
"""
|
|
SFT training support for KT-Kernel MoE.
|
|
|
|
This submodule adds training capabilities (forward/backward, LoRA, autograd,
|
|
distributed) on top of the inference-only kt_kernel base package.
|
|
|
|
Additional dependencies beyond base kt_kernel: torch.nn, torch.distributed, peft (optional).
|
|
"""
|
|
|
|
from .config import KTActivationPolicy, KTConfig
|
|
from .backend import (
|
|
FP8_BACKEND,
|
|
FP8_SFT_METHOD,
|
|
INT8_BACKEND,
|
|
INT8_SFT_METHOD,
|
|
RAWINT4_BACKEND,
|
|
RAWINT4_GROUP_SIZE,
|
|
RAWINT4_SFT_METHOD,
|
|
RAWINT4_WEIGHT_LAYOUT,
|
|
get_fp8_runtime,
|
|
get_int8_runtime,
|
|
get_rawint4_checkpoint_contract,
|
|
get_rawint4_runtime,
|
|
MXFP4_SFT_METHOD,
|
|
get_mxfp4_runtime,
|
|
)
|
|
from .weight_manifest import validate_persistent_int8_weights
|
|
from .artifacts import (
|
|
FUSED_EXPERT_LORA_NAME,
|
|
KT_ADAPTER_MANIFEST_NAME,
|
|
KT_NON_EXPERT_MANIFEST_NAME,
|
|
KTAdapterManifest,
|
|
KTArtifactError,
|
|
KTNonExpertCache,
|
|
KTPretrainedLoadPlan,
|
|
claim_kt_routed_expert_subtrees,
|
|
hide_kt_int8_routed_experts_from_dispatch,
|
|
hide_kt_routed_experts_from_dispatch,
|
|
is_kt_int8_routed_expert_base_parameter,
|
|
is_kt_routed_expert_base_parameter,
|
|
is_kt_routed_expert_parameter_name,
|
|
is_kt_supported_moe_model,
|
|
get_kt_fused_lora_exclude_modules,
|
|
load_kt_adapter_artifacts,
|
|
mark_kt_int8_routed_expert_base_parameters,
|
|
prepare_kt_int8_non_expert_device_map,
|
|
prepare_kt_non_expert_device_map,
|
|
project_kt_int8_routed_experts_out_of_device_map,
|
|
project_kt_routed_experts_out_of_device_map,
|
|
resolve_kt_pretrained_artifacts,
|
|
save_kt_adapter_artifacts,
|
|
should_disable_kt_source_quantizer,
|
|
validate_kt_prequantized_loading_info,
|
|
validate_kt_pretrained_load,
|
|
write_kt_non_expert_cache_manifest,
|
|
)
|
|
from .dist_utils import get_activation_checkpoint_context_fn
|
|
from .base import BaseSFTMoEWrapper, KExpertsSFTBuffer
|
|
from .amx import AMXSFTMoEWrapper
|
|
from .arch import (
|
|
MOEArchConfig,
|
|
get_moe_arch_config,
|
|
get_moe_module,
|
|
move_non_experts_to_gpu,
|
|
get_expert_device,
|
|
KTAMXError,
|
|
KTAMXNotAvailableError,
|
|
KTAMXModelNotSupportedError,
|
|
KTAMXConfigError,
|
|
)
|
|
from .autograd import KTMoEFunction
|
|
from .conv3d_compat import is_vlm_conv3d_compatible, patch_vlm_conv3d
|
|
from .layer import KTMoELayerWrapper
|
|
from .weights import (
|
|
extract_moe_weights,
|
|
get_kt_expert_placeholders,
|
|
load_block_fp8_experts_from_checkpoint_files,
|
|
load_experts_from_checkpoint_files,
|
|
load_experts_from_kt_weight_path,
|
|
load_rawint4_experts_from_checkpoint_files,
|
|
INT8ExpertWeights,
|
|
BlockFP8ExpertWeights,
|
|
RAWINT4ExpertWeights,
|
|
)
|
|
from .lora import (
|
|
KTAdaptationResult,
|
|
adapt_peft_lora,
|
|
kt_adapt_peft_lora,
|
|
get_kt_lora_named_params,
|
|
get_kt_lora_params,
|
|
get_kt_named_trainable_params,
|
|
get_kt_rank_local_parameter_names,
|
|
get_kt_trainable_params,
|
|
update_kt_lora_pointers,
|
|
sync_kt_lora_gradients,
|
|
save_lora_experts_to_adapter,
|
|
save_kt_moe_to_adapter,
|
|
load_lora_experts_from_adapter,
|
|
load_kt_moe_from_adapter,
|
|
LoRAExpertMLP,
|
|
LoRAExperts,
|
|
)
|
|
from .wrapper import (
|
|
wrap_moe_layers_with_kt_wrapper,
|
|
build_kt_device_map,
|
|
build_kt_device_map_simplified,
|
|
get_kt_loading_kwargs,
|
|
load_kt_model,
|
|
)
|
|
from .profiler import collect_kt_sft_profile, format_kt_sft_profile, reset_kt_sft_profile
|
|
|
|
__all__ = [
|
|
"KTConfig",
|
|
"KTActivationPolicy",
|
|
"INT8_BACKEND",
|
|
"INT8_SFT_METHOD",
|
|
"get_int8_runtime",
|
|
"FP8_BACKEND",
|
|
"FP8_SFT_METHOD",
|
|
"get_fp8_runtime",
|
|
"RAWINT4_BACKEND",
|
|
"RAWINT4_GROUP_SIZE",
|
|
"RAWINT4_SFT_METHOD",
|
|
"RAWINT4_WEIGHT_LAYOUT",
|
|
"get_rawint4_checkpoint_contract",
|
|
"get_rawint4_runtime",
|
|
"MXFP4_SFT_METHOD",
|
|
"get_mxfp4_runtime",
|
|
"validate_persistent_int8_weights",
|
|
"FUSED_EXPERT_LORA_NAME",
|
|
"KT_ADAPTER_MANIFEST_NAME",
|
|
"KT_NON_EXPERT_MANIFEST_NAME",
|
|
"KTAdapterManifest",
|
|
"KTArtifactError",
|
|
"KTNonExpertCache",
|
|
"KTPretrainedLoadPlan",
|
|
"claim_kt_routed_expert_subtrees",
|
|
"hide_kt_int8_routed_experts_from_dispatch",
|
|
"hide_kt_routed_experts_from_dispatch",
|
|
"is_kt_int8_routed_expert_base_parameter",
|
|
"is_kt_routed_expert_base_parameter",
|
|
"is_kt_routed_expert_parameter_name",
|
|
"is_kt_supported_moe_model",
|
|
"get_kt_fused_lora_exclude_modules",
|
|
"load_kt_adapter_artifacts",
|
|
"mark_kt_int8_routed_expert_base_parameters",
|
|
"prepare_kt_int8_non_expert_device_map",
|
|
"prepare_kt_non_expert_device_map",
|
|
"project_kt_int8_routed_experts_out_of_device_map",
|
|
"project_kt_routed_experts_out_of_device_map",
|
|
"resolve_kt_pretrained_artifacts",
|
|
"save_kt_adapter_artifacts",
|
|
"should_disable_kt_source_quantizer",
|
|
"validate_kt_prequantized_loading_info",
|
|
"validate_kt_pretrained_load",
|
|
"write_kt_non_expert_cache_manifest",
|
|
"get_activation_checkpoint_context_fn",
|
|
"BaseSFTMoEWrapper",
|
|
"KExpertsSFTBuffer",
|
|
"AMXSFTMoEWrapper",
|
|
"MOEArchConfig",
|
|
"get_moe_arch_config",
|
|
"get_moe_module",
|
|
"move_non_experts_to_gpu",
|
|
"get_expert_device",
|
|
"KTAMXError",
|
|
"KTAMXNotAvailableError",
|
|
"KTAMXModelNotSupportedError",
|
|
"KTAMXConfigError",
|
|
"KTMoEFunction",
|
|
"patch_vlm_conv3d",
|
|
"is_vlm_conv3d_compatible",
|
|
"KTMoELayerWrapper",
|
|
"extract_moe_weights",
|
|
"get_kt_expert_placeholders",
|
|
"load_block_fp8_experts_from_checkpoint_files",
|
|
"load_experts_from_checkpoint_files",
|
|
"load_experts_from_kt_weight_path",
|
|
"load_rawint4_experts_from_checkpoint_files",
|
|
"INT8ExpertWeights",
|
|
"BlockFP8ExpertWeights",
|
|
"RAWINT4ExpertWeights",
|
|
"KTAdaptationResult",
|
|
"adapt_peft_lora",
|
|
"kt_adapt_peft_lora",
|
|
"get_kt_lora_named_params",
|
|
"get_kt_lora_params",
|
|
"get_kt_named_trainable_params",
|
|
"get_kt_rank_local_parameter_names",
|
|
"get_kt_trainable_params",
|
|
"update_kt_lora_pointers",
|
|
"sync_kt_lora_gradients",
|
|
"save_lora_experts_to_adapter",
|
|
"save_kt_moe_to_adapter",
|
|
"load_lora_experts_from_adapter",
|
|
"load_kt_moe_from_adapter",
|
|
"LoRAExpertMLP",
|
|
"LoRAExperts",
|
|
"wrap_moe_layers_with_kt_wrapper",
|
|
"build_kt_device_map",
|
|
"build_kt_device_map_simplified",
|
|
"get_kt_loading_kwargs",
|
|
"load_kt_model",
|
|
"collect_kt_sft_profile",
|
|
"format_kt_sft_profile",
|
|
"reset_kt_sft_profile",
|
|
]
|