5 Commits

Author	SHA1	Message	Date
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
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
liyuemathematician	168f581676	[feat](sft): support Qwen VLM MoE fine-tuning (#2156)	2026-08-17 10:51:55 +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