3 Commits

Author	SHA1	Message	Date
yyj	e80614fc2d	[feat](kt-sft): Activation reuse & Int8 kernel refactor & native block-FP8 LoRA SFT (#2141) ... * 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	2026-08-17 16:25:14 +08:00
yyj	924754a00b	[feat](kt-kernel): end-to-end full-parameter and LoRA SFT(#2094) ... * Add full FT development snapshot * [fix](kt-kernel): fix Full FT TP base weight gradients * [fix]: bug fix of 2d81e86 * [fix](kt-kernel): fix AMX BF16 full-weight gradients * [docs](kt-kernel): document Full FT fork changes and debug history * [docs](kt-kernel): align fork remote terminology * [chore](kt-kernel): keep agent notes local * [fix](kt-kernel): configure SFT OpenMP threads * [perf](kt-kernel): optimize AMX Full-FT weight gradients Coarsen base-weight gradient work from individual output tiles to fixed-intermediate strips so each task reuses packed panels across the hidden dimension. Keep aligned thread-local BF16 panels across tasks and retain FP32 AMX accumulator tiles for the full K reduction. Gate and up run separate K passes while sharing the packed input panel. On the matched Qwen3-30B-A3B 1-GPU test, stable Full-FT backward drops from 9.281s to 6.792s (-26.82%), step time drops from 19.252s to 16.140s, and TPS rises from 212.76 to 253.78. The LoRA-only backward control changes by -2.74%. Validated with clang-format, the Release AMX/CUDA extension build, TP1/TP2 reference gradients across boundary token counts, and the 15-step Full-then-LoRA performance run. * [feat](kt-kernel): add staged SFT profiling * [fix](kt-kernel): reuse inference BF16 kernel for SFT * [perf](kt-kernel): add fine-grained Full-FT profiling * [perf](kt-kernel): batch BF16 Full-FT weight gradients Use one expert-aggregated tile driver for AVX512-BF16 and AMX base-weight gradients, and pack updated full-precision weights directly into TP BufferB layouts without temporary partitions. Add worker-local profiling and focused dWeight/strided-repack coverage. * [test](kt-kernel): benchmark BF16 dWeight AMX driver * [fix](kt-kernel): label dWeight store as worker CPU time * [docs](sft): record BF16 Full-FT performance * [docs](sft): remove Qwen3 Full-FT performance report * [perf](kt-kernel): reduce BF16 Full-FT checkpoint overhead Retain the first CPU MoE forward state across non-reentrant checkpoint recomputation, write BF16 activations directly into the backward cache, and reduce dWeight packing and gradient-clear traffic. Extend staged profiling and cover checkpoint reuse plus AMX/AVX dWeight paths. * [perf](kt-kernel): make SFT optimizer gradients authoritative Bind Full-FT and LoRA Parameter.grad directly to the KT-managed BF16 gradient buffers, avoiding PyTorch duplicate accumulation. Accumulate microbatch gradients in C++, lazily clear expert buffers between optimizer windows, preserve rank-0 distributed ownership, and add lifecycle and AMX dWeight coverage. * [perf](kt-kernel): avoid eager Full-FT gradient zeroing Allocate authoritative Full-FT gradient buffers with torch.empty. The C++ state machine performs the mandatory full clear before first use, avoiding redundant Python-side first touch. * perf(sft): enable checkpoint forward reuse for LoRA * fix(sft): serialize checkpoint recompute with async repack * fix(sft): persist authoritative full weights * fix(sft): normalize legacy distributed gradients * [fix](sft): train gated shared experts * fix(sft): preserve expert placeholders across state dict loads * ci: publish ktransformers sdist in release workflow * fix(sft): preserve router autograd for LoRA training Determine routing graph tracking from the router's trainable parameters instead of the Full-FT mode flag, so PEFT LoRA adapters on MoE gates receive routing-weight gradients. Fail fast when a trainable router returns detached weights and cover frozen, TopK, checkpoint-reuse, and two-step optimizer behavior. --------- Co-authored-by: illu <wubowen03@foxmail.com>	2026-07-23 21:23:34 +08:00
mrhaoxx	9544a8960d	feat(sft): AMX MoE SFT backend with LoRA support (#1936) ... * feat(sft): AMX MoE SFT backend with LoRA support Complete SFT (Supervised Fine-Tuning) backend for MoE models using AMX SIMD: Core C++ implementation: - sft_moe.hpp: Forward/backward with LoRA fused operations (~5500 lines) - moe-sft-tp.hpp: Tensor-parallel wrapper for multi-NUMA - amx/moe-sft-tp.hpp: AMX-specific TP implementation - avx_kernels.hpp: AVX512 SIMD kernels for LoRA GEMM - amx_kernels.hpp: AMX tile kernels for Panel5 rank-outer optimization - worker_pool: RDTSC profiling, Chrome trace output, SFT timer infrastructure - ext_bindings.cpp: SFT MOE pybind bindings (BF16/INT8/INT4 + SkipLoRA variants) Python sft/ submodule (kt_kernel.sft): - base.py: BaseSFTMoEWrapper with buffer management (template method pattern) - amx.py: AMXSFTMoEWrapper (weight loading, C++ task construction) - autograd.py: KTMoEFunction (torch.autograd.Function for distributed training) - layer.py: KTMoELayerWrapper (nn.Module replacing HF MoE layers) - arch.py: MOEArchConfig (Qwen3/DeepSeek/Mixtral architecture detection) - weights.py: Expert weight extraction and checkpoint loading - lora.py: PEFT LoRA adaptation (view buffers, grad buffers, save/load adapter) - wrapper.py: wrap_moe_layers_with_kt_wrapper, load_kt_model, build_kt_device_map - config.py: KTConfig dataclass (DeepSpeed-style opaque config passthrough) - dist_utils.py: Distributed gather/scatter, checkpoint-phase detection Design decisions: - Rank-0-only expert pattern: only rank 0 holds C++ wrapper and expert weights - DeepSpeed-style integration: accelerate keeps only KTransformersPlugin (framework interaction fields), all logic in kt_kernel.sft - Inference isolation: importing kt_kernel does not load sft/ submodule - Old field name compatibility: _get_kt_config() converts kt_xxx→xxx automatically Verified: Qwen3-235B-A22B 4GPU AMXBF16 training, loss converges normally. * refactor(sft): unify KTConfig field names with kt_ prefix, add share_cache_pool, remove dead code - KTConfig fields all use kt_ prefix matching dict keys — eliminates _OLD_TO_NEW mapping and prefix-stripping in wrapper.py - Add kt_share_cache_pool field, auto-enabled when gradient_checkpointing is on (via training_args.py), flows through to C++ cache allocation - Remove dead checkpoint detection code: in_ckpt_recompute, in_ckpt_first_forward vars (assigned but never read), fallback _is_in_checkpoint_first_forward() function, unused inspect import - Remove redundant env var fallbacks in wrapper.py for share_backward_bb and share_cache_pool (KTConfig.__post_init__ already handles env vars) - Simplify layer.py checkpoint logic to single _checkpoint_hook_mode() check Verified: Qwen3-235B 3-step training on sap4, loss matches baseline (1.2886 / 1.9824 / 1.377 vs 1.2886 / 1.9766 / 1.3809) Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * refactor(sft): share_backward_bb default True, share_cache_pool auto-derived - kt_share_backward_bb defaults to True (always saves memory) - kt_share_cache_pool no longer reads from env var; defaults False, auto-set to True by trainer_config_process when gradient checkpointing is enabled Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: add missing gpu_experts_mask=None to KTMoEWrapper call in SFT wrapper KTMoEWrapper.__new__() requires gpu_experts_mask as a positional argument, but the SFT wrapper omitted it, causing MoE layer wrapping to fail silently and FSDP2 to attempt broadcasting all expert weights (OOM/NCCL crash). Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * feat(sft): support transformers v5 fused expert format Fused experts (e.g. Qwen3MoeExperts) store weights as 3D Parameters (gate_up_proj [E,2I,H], down_proj [E,H,I]) instead of per-expert nn.Linear modules. PEFT cannot attach LoRA to these, so we create KT-managed LoRA buffers with kaiming init, nn.Parameter wrappers for the optimizer, and pre-assigned .grad for C++ backward. - arch.py: detect_fused_experts() detection - weights.py: fused format extraction and weight clearing - wrapper.py: detect fused at wrap time, store _fused_experts/_lora_rank - lora.py: _create_fused_expert_lora_buffers, save/load fused LoRA, get_kt_lora_params collects fused params, deduplicate wrapper finding - layer.py: handle v5 TopKRouter tuple output, remove dead code - autograd.py: sync_forward_sft/submit_forward_sft API rename Verified: v5 loss/expert-LoRA values match v4 baseline, v4 backward compat. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * feat(sft): add Qwen3.5 MoE support + fused checkpoint loading - arch.py: add Qwen3_5Moe arch match, read config from text_config, _get_layers_prefix returns model.language_model.layers for Qwen3.5, _get_model_container_and_layers searches language_model attr - weights.py: load_experts_from_checkpoint_files detects fused format (gate_up_proj in weight_map) and splits into gate/up/down - wrapper.py: hidden_size fallback to text_config Verified: Qwen3.5-35B-A3B (256 experts, fused format) E2E pass. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * [fix](sft): align Python API with C++ backend after v5 refactor - wrapper.py: pass gpu_experts_mask=None to KTMoEWrapper (required by C++ signature) - layer.py: rename submit_forward_sft/sync_forward_sft to submit_forward/sync_forward - autograd.py: rename sync_forward_sft to sync_forward The sft-v5 refactor (commits 58d7eab, dd1da65) renamed Python-side method calls but the C++ backend (AMXSFTMoEWrapper) still exposes the original method names. This caused AttributeError on Qwen3.5-35B and other models. * align sft branch with main: revert worker_pool, strip sft_timer, fix inference defaults - Revert worker_pool.cpp/.h to main (remove RDTSC timer, Chrome Trace, sft_timer namespace, ITT API, extended do_work_stealing_job API) - Strip all sft_timer instrumentation from sft-only files (sft_moe.hpp, moe-sft-tp.hpp, avx_kernels.hpp) - Restore pin_memory=True in KExpertsCPUBuffer (inference path) - Restore fused tensor transpose logic in convert_cpu_weights.py (main layout) Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * revert CMakeLists.txt to main: remove debug flags and cpptrace dep Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * clean up dev artifacts: remove SFT design docs, debug examples, bench scripts Remove files not needed in the merge: - docs/SFT+KTWrapper/ (6 Chinese design docs) - docs/sft_moe_amx/ (21 dev/debug docs) - 12 debug/test example scripts - 6 SFT-specific bench scripts and report Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * remove dev version stamps from ext_bindings, sft_moe, moe-sft-tp Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Co-authored-by: JimmyPeilinLi <lipeilin@mail.nwpu.edu.cn>	2026-04-22 11:27:01 +08:00