 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 |
|