mirror of
https://github.com/kvcache-ai/ktransformers.git
synced 2026-10-18 22:55:21 +08:00
yyj
924754a00b
* 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>
77 lines
2.7 KiB
Python
77 lines
2.7 KiB
Python
from types import SimpleNamespace
|
|
|
|
from kt_kernel.sft.profiler import collect_kt_sft_profile, format_kt_sft_profile, reset_kt_sft_profile
|
|
|
|
|
|
class _FakeMoe:
|
|
def __init__(self):
|
|
self.reset_calls = 0
|
|
self.get_reset_values = []
|
|
|
|
def get_profile_stats(self, reset=False):
|
|
self.get_reset_values.append(reset)
|
|
return {
|
|
"layer_idx": 3,
|
|
"tp_count": 1,
|
|
"wrapper.enabled": 1,
|
|
"wrapper.tokens": 8,
|
|
"wrapper.tp.forward.total.total_ns": 2_000_000,
|
|
"wrapper.tp.forward.total.calls": 2,
|
|
"wrapper.tp.forward.numa_compute.total_ns": 1_500_000,
|
|
"wrapper.tp.forward.numa_compute.calls": 2,
|
|
"wrapper.tp.forward.numa_compute.bytes": 2 * 1024 * 1024,
|
|
"tp.0.enabled": 1,
|
|
"tp.0.tokens": 8,
|
|
"tp.0.forward.total.total_ns": 1_400_000,
|
|
"tp.0.forward.total.calls": 2,
|
|
"tp.0.forward.route.total_ns": 140_000,
|
|
"tp.0.forward.route.calls": 2,
|
|
"tp.0.backward.total.total_ns": 5_000_000,
|
|
"tp.0.backward.total.calls": 1,
|
|
"tp.0.backward.down.total.total_ns": 2_000_000,
|
|
"tp.0.backward.down.total.calls": 1,
|
|
"tp.0.backward.base_weight_grad.total_ns": 1_000_000,
|
|
"tp.0.backward.base_weight_grad.calls": 1,
|
|
"tp.0.backward.base_weight_grad.worker_cpu.store.total_ns": 2_000_000,
|
|
"tp.0.backward.base_weight_grad.worker_cpu.store.calls": 4,
|
|
}
|
|
|
|
def reset_profile_stats(self):
|
|
self.reset_calls += 1
|
|
|
|
|
|
def _fake_model(moe):
|
|
backend = SimpleNamespace(moe=moe)
|
|
layer = SimpleNamespace(layer_idx=3, wrapper=backend)
|
|
return SimpleNamespace(_kt_wrappers=[layer])
|
|
|
|
|
|
def test_collect_and_format_profile():
|
|
moe = _FakeMoe()
|
|
profile = collect_kt_sft_profile(_fake_model(moe), reset=True)
|
|
|
|
assert profile["enabled"] is True
|
|
assert list(profile["layers"]) == [3]
|
|
assert moe.get_reset_values == [True]
|
|
|
|
output = format_kt_sft_profile(profile)
|
|
assert "tp.forward.numa_compute" in output
|
|
assert "forward.route" in output
|
|
assert "75.0%" in output
|
|
assert "2.00" in output
|
|
assert "10.0%" in output
|
|
assert "40.0%" in output
|
|
worker_store = next(line for line in output.splitlines() if "worker_cpu.store" in line)
|
|
assert worker_store.endswith("0.0%")
|
|
|
|
|
|
def test_reset_and_disabled_profile():
|
|
moe = _FakeMoe()
|
|
model = _fake_model(moe)
|
|
reset_kt_sft_profile(model)
|
|
assert moe.reset_calls == 1
|
|
|
|
empty = collect_kt_sft_profile(SimpleNamespace(_kt_wrappers=[]))
|
|
assert empty == {"enabled": False, "layers": {}}
|
|
assert "disabled" in format_kt_sft_profile(empty)
|