From 2d81e8683a9da4b0394596adb820737e66465997 Mon Sep 17 00:00:00 2001 From: illu Date: Fri, 10 Jul 2026 11:04:29 +0000 Subject: [PATCH] [fix](kt-kernel): fix Full FT TP base weight gradients --- kt-kernel/AGENTS.md | 190 ++++++++++++++++++++++++++++ kt-kernel/operators/amx/sft_moe.hpp | 36 ++++-- kt-kernel/operators/moe-sft-tp.hpp | 44 ++++++- 3 files changed, 255 insertions(+), 15 deletions(-) create mode 100644 kt-kernel/AGENTS.md diff --git a/kt-kernel/AGENTS.md b/kt-kernel/AGENTS.md new file mode 100644 index 00000000..b31484dc --- /dev/null +++ b/kt-kernel/AGENTS.md @@ -0,0 +1,190 @@ +# AGENTS.md: kt-kernel Full FT 开发指南 + +本文件作用于 `kt-kernel/` 及其全部子目录，用于指导 Agent 开发和验证 MoE +Full Fine-Tuning（Full FT）功能。重点是专家基座权重的梯度、optimizer 更新 +和 AMX 权重重新量化，不讨论普通 GPU 参数训练。 + +当前开发快照： + +- 仓库：`Illumination111/ktransformers-fullFT_development` +- Full FT 基线提交：`e99b5e1`，基于上游 `8e46e58` +- Full FT 来源：poryfly 的 `239bac5`、`25fa2fd` +- 当前状态：TP 子核配置、梯度分片/stride、临时区容量和逐步清零已修复；等待修复后训练验证 + +## 1. 开发前先确认 + +1. 在 `ktransformers` 中运行 `git status --short --branch`，不要覆盖已有改动。 +2. 使用符号名定位代码，不依赖本文行号。 +3. 先确认目标是 `full`、`hybrid` 还是 `lora`： + - `full`：训练 expert 基座权重，通常 `lora_rank=0`。 + - `hybrid`：同时训练基座权重和 LoRA。 + - `lora`：不得计算或更新基座权重。 +4. 修改前记录 TP 数量、NUMA 数量、`E/H/F` 和权重 dtype。 + +## 2. Full FT 数据契约 + +完整链路： + +```text +KTConfig + -> wrapper 创建 CPU BF16 nn.Parameter + -> forward 将 Parameter 挂入 KTMoEFunction + -> C++ backward 写 grad_*_proj_buf + -> autograd 返回基座梯度 + -> optimizer 更新 *_proj_buf + -> update_base_weights 重新量化到 AMX BufferB +``` + +权威数据必须唯一： + +- `gate/up/down_proj_buf`：optimizer 可见的 CPU BF16 基座权重。 +- `grad_gate/up/down_proj_buf`：C++ backward 的输出。 +- AMX 量化权重：forward 使用的派生副本，optimizer 后必须重建。 +- HF expert 原权重只是占位，不得作为 Full FT 的更新依据。 + +## 3. 文件职责 + +路径均相对本文件所在的 `kt-kernel/`。 + +| 路径 | 职责 | +|---|---| +| `python/sft/config.py` | 将 `full/hybrid` 映射为 `kt_full_weight_grad=True` | +| `python/sft/base.py` | 创建 `*_proj_buf`、`grad_*_proj_buf` | +| `python/sft/wrapper.py` | 包装 MoE、确定权威权重 | +| `python/sft/layer.py` | 将三个基座 Parameter 传入 autograd；dirty 时 requant | +| `python/sft/autograd.py` | 返回 C++ 写入的基座梯度 | +| `python/sft/lora.py` | optimizer 参数注入、分布式同步、dirty 标记 | +| `python/sft/amx.py` | 设置 C++ config、传梯度指针、触发重新量化 | +| `operators/common.hpp` | `MOESFTConfig.full_weight_grad` | +| `operators/moe-tp.hpp` | 通用 TP 切分；当前会把 SFT config 切成基类 | +| `operators/moe-sft-tp.hpp` | SFT TP 调度和全局梯度 buffer 分片 | +| `operators/amx/sft_moe.hpp` | NUMA 子核 backward 和基座梯度计算 | +| `ext_bindings.cpp` | Python/C++ 参数和 task 绑定 | + +## 4. TP 张量布局 + +定义： + +- `E`：expert 数量。 +- `H`：hidden size。 +- `F`：完整 `intermediate_size`。 +- `I`：当前 TP 子核的本地 `intermediate_size`。 +- `tp_offset`：当前 TP 在完整 `F` 维上的起点。 + +完整梯度布局： + +```text +gate/up: [E, F, H] +down: [E, H, F] +``` + +TP wrapper 传给子核的起始指针： + +```text +gate/up = base + tp_offset * H +down = base + tp_offset +``` + +子核循环和 FP32 累加使用本地 `I`，但写回完整张量时必须使用 `F` +作为 expert/row stride： + +```text +gate/up: expert * F * H + local_i * H + h +down: expert * H * F + h * F + local_i +``` + +所有 TP 必须写入互不重叠的区域；对空指针不得做偏移运算。 + +## 5. 历史根因与失败基线 + +`full_weight_grad=True` 已正确到达顶层 `TP_MOE_SFT`，但 +`TP_MOE` 创建子核时执行： + +```cpp +GeneralMOEConfig tp_config = config; +``` + +这会丢失 `MOESFTConfig.full_weight_grad`。子核随后通过 +`MOESFTConfig(GeneralMOEConfig)` 重建配置，字段回到默认 `false`， +因此 `backward_base_weight_grad` 被门控跳过。 + +只恢复该布尔值仍不够：当前 TP backward 把相同的完整梯度首地址传给 +所有子核，子核又按本地 `I` 写回，会导致覆盖和错误布局。 + +修复前的 `20260710_103230_1gpu_AMX_BF16_FULLFT_EXPERTONLY` run 已冻结 +其他权重并只向 optimizer 注入 144 个 expert 基座参数，但 15 步 +`grad_norm=0`，step 0 到 step 15 的 probe 为 `changed=0/12`、 +`max_abs_delta=0`。720 次 requant 不能证明权重更新。 + +## 6. 当前两文件实现 + +生产代码限制在两个文件，未修改 Python 接口。 + +### 6.1 `operators/amx/sft_moe.hpp` + +1. `set_full_weight_grad(bool)` 更新子核的 + `sft_config_.full_weight_grad`。 +2. `full_intermediate_size` 已传给 `backward_base_weight_grad`。 +3. 计算循环使用本地 `I`，写回 stride 使用完整 `F`。 +4. 门控要求开关为 true 且三个梯度指针均非空。 +5. `forward_pool_` 保证至少 `3 * I * H * sizeof(float)` 字节。 +6. `lora_rank=0` 时 scaling 明确为 0，避免纯 Full FT 产生 Inf。 + +### 6.2 `operators/moe-sft-tp.hpp` + +1. 子核创建后逐个调用 `set_full_weight_grad`。 +2. 传播逻辑位于 `if constexpr (!kSkipLoRA)` 外；纯 Full FT 的 + `lora_rank=0` 也必须执行。 +3. backward dispatch 前按第 4 节公式生成每个 TP 的三个梯度指针。 +4. 完整 `F` 传给子核作为全局 stride，并检查所有 TP slice 完整覆盖 `F`。 +5. TP dispatch 前统一并行清零三组完整梯度，避免跨 step 残留和 TP 清零竞争。 + +暂不修改通用 `operators/moe-tp.hpp`。保留派生 config 的泛型重构影响 +所有 MoE backend，应作为后续独立改动。 + +## 7. 验证顺序 + +当前已通过 `clang-format --dry-run --Werror` 和 AMX/CUDA Release +`build_ext --inplace`。修复后的 expert-only 短训练仍是必需验收项。 + +1. **静态检查**：确认 setter 不在 LoRA 条件分支内；确认所有指针先判空再偏移。 +2. **构建**：重新编译并安装 `kt_kernel_ext`。 +3. **参考梯度测试**： + - 分别覆盖 TP=1 和 TP=2。 + - 使用 `expert_id > 0`，避免错误 expert stride 被 expert 0 掩盖。 + - 与 PyTorch outer-product 参考值比较 gate/up/down 三组梯度。 + - TP=2 时分别检查两个 `F` 分片，确认无覆盖。 +4. **跨 step 测试**：连续两步激活不同 expert，确认未激活 expert 不保留旧梯度。 +5. **模式回归**：`full_weight_grad=false` 不写基座梯度；纯 LoRA 结果不变。 +6. **短训练**： + - step 内三个 `grad_*_proj_buf.abs().max() > 0`。 + - optimizer 后至少一个 `*_proj_buf` 发生有限变化。 + - requant 后下一次 forward 使用更新后的权重。 + +不要用 loss 下降证明 Full FT 生效；attention、router 或其他非 expert 参数也能 +让 loss 下降。 + +## 8. 完成标准 + +- C++ 子核实际收到 `full_weight_grad=true`。 +- 三组基座梯度非零、有限，并与参考实现一致。 +- TP 分片覆盖完整 `F`，无重叠、越界或错误 stride。 +- optimizer 能找到三个基座 Parameter，并在 step 后更新它们。 +- inactive expert 无跨 step 残留梯度。 +- `full`、`hybrid`、`lora` 三种模式行为符合各自契约。 +- 纯 LoRA 和现有 AMX MoE 测试无回归。 + +## 9. 调试资料 + +```text +FFTtest/Qwen3-30B-A3B/test_log/20260710_103230_1gpu_AMX_BF16_FULLFT_EXPERTONLY/ + summary.md + expert_weight_change_check.{txt,json} + phase4/expert_buf_probe.json + phase4/train.log + +FFTtest/Qwen3-30B-A3B/expert_buf_probe.py +``` + +该 run 的直接证据是 `changed=0/12`、`max_abs_delta=0`；训练 loss +仍下降不能否定 expert 基座梯度失效。 diff --git a/kt-kernel/operators/amx/sft_moe.hpp b/kt-kernel/operators/amx/sft_moe.hpp index b516bcf7..dd831615 100644 --- a/kt-kernel/operators/amx/sft_moe.hpp +++ b/kt-kernel/operators/amx/sft_moe.hpp @@ -712,6 +712,13 @@ class AMX_SFT_MOE_TP : public BaseMOE { work_required += round_up(lora_bc_out_pool_bytes_, kAmxAlignment); work_required += round_up(lora_intermediate_bf16_pool_bytes_, kAmxAlignment); + // Base-weight gradients reuse the forward pool after LoRA backward is complete. + if (sft_config_.full_weight_grad) { + const size_t base_grad_accumulator_bytes = + 3 * (size_t)config_.intermediate_size * config_.hidden_size * sizeof(float); + work_required = std::max(work_required, base_grad_accumulator_bytes); + } + alloc_or_resize_forward_pool(work_required); SFT_POOL_LOG("fwd_work", config_.layer_idx, tp_part_idx, 0, cache_stack_top_, forward_pool_bytes_, @@ -897,9 +904,11 @@ class AMX_SFT_MOE_TP : public BaseMOE { */ void set_lora_params(int rank, float alpha) { lora_rank_ = rank; - lora_scaling_ = alpha / rank; + lora_scaling_ = rank > 0 ? alpha / rank : 0.0f; } + void set_full_weight_grad(bool enabled) { sft_config_.full_weight_grad = enabled; } + /** * @brief SFT Forward pass with optional caching for backward. * @@ -1883,7 +1892,8 @@ class AMX_SFT_MOE_TP : public BaseMOE { // Step 5: Base weight gradient accumulation (full weight grad mode) // ===================================================================== if (sft_config_.full_weight_grad && grad_gate_proj && grad_up_proj && grad_down_proj) { - backward_base_weight_grad(cache, grad_output, grad_gate_proj, grad_up_proj, grad_down_proj); + backward_base_weight_grad(cache, grad_output, full_intermediate_size, grad_gate_proj, grad_up_proj, + grad_down_proj); } // \u2605 Cache pool is NOT freed here \u2014 kept for reuse across steps. @@ -1905,16 +1915,20 @@ class AMX_SFT_MOE_TP : public BaseMOE { * * Uses FP32 accumulator for precision, writes BF16 output. */ - void backward_base_weight_grad(const ForwardCache& cache, const void* grad_output, void* grad_gate_proj, - void* grad_up_proj, void* grad_down_proj) { + void backward_base_weight_grad(const ForwardCache& cache, const void* grad_output, int full_intermediate_size, + void* grad_gate_proj, void* grad_up_proj, void* grad_down_proj) { const int H = config_.hidden_size; const int I = config_.intermediate_size; - const int E = config_.expert_num; + const int F = full_intermediate_size; int activated_expert = cache.activated_expert_cache; - auto* ggp = static_cast(grad_gate_proj); // [E, I, H] - auto* gup_ptr = static_cast(grad_up_proj); // [E, I, H] - auto* gdp = static_cast(grad_down_proj); // [E, H, I] + if (F < I) { + throw std::runtime_error("full_intermediate_size must be at least the TP-local intermediate_size"); + } + + auto* ggp = static_cast(grad_gate_proj); // TP slice of [E, F, H] + auto* gup_ptr = static_cast(grad_up_proj); // TP slice of [E, F, H] + auto* gdp = static_cast(grad_down_proj); // TP slice of [E, H, F] auto* grad_out_bf16 = static_cast(grad_output); for (int task_id = 0; task_id < activated_expert; task_id++) { @@ -1969,13 +1983,13 @@ class AMX_SFT_MOE_TP : public BaseMOE { // Convert FP32 accumulators to BF16 and store for (int i = 0; i < I; i++) { for (int h = 0; h < H; h++) { - ggp[(size_t)expert_idx * I * H + (size_t)i * H + h] = GGML_FP32_TO_BF16(acc_gate[i * H + h]); - gup_ptr[(size_t)expert_idx * I * H + (size_t)i * H + h] = GGML_FP32_TO_BF16(acc_up[i * H + h]); + ggp[(size_t)expert_idx * F * H + (size_t)i * H + h] = GGML_FP32_TO_BF16(acc_gate[i * H + h]); + gup_ptr[(size_t)expert_idx * F * H + (size_t)i * H + h] = GGML_FP32_TO_BF16(acc_up[i * H + h]); } } for (int h = 0; h < H; h++) { for (int i = 0; i < I; i++) { - gdp[(size_t)expert_idx * H * I + (size_t)h * I + i] = GGML_FP32_TO_BF16(acc_down[h * I + i]); + gdp[(size_t)expert_idx * H * F + (size_t)h * F + i] = GGML_FP32_TO_BF16(acc_down[h * I + i]); } } } diff --git a/kt-kernel/operators/moe-sft-tp.hpp b/kt-kernel/operators/moe-sft-tp.hpp index c8a0e098..918aaf74 100644 --- a/kt-kernel/operators/moe-sft-tp.hpp +++ b/kt-kernel/operators/moe-sft-tp.hpp @@ -229,6 +229,11 @@ class TP_MOE_SFT : public TP_MOE { part_grad_input_.assign(tp_count, nullptr); part_grad_weights_.assign(tp_count, nullptr); + // TP_MOE stores GeneralMOEConfig and slices off SFT-only fields. + for (int i = 0; i < tp_count; i++) { + tps[i]->set_full_weight_grad(config.full_weight_grad); + } + if constexpr (!kSkipLoRA) { // Bug #16 fix: TP_MOE base class uses GeneralMOEConfig (object slicing) which loses // LoRA pointers. We need to propagate LoRA pointers to all NUMA node instances. @@ -562,6 +567,8 @@ class TP_MOE_SFT : public TP_MOE { int k = sft_config.num_experts_per_tok; const bool need_grad_weights = (grad_weights != nullptr); + const bool need_base_weight_grad = sft_config.full_weight_grad && grad_gate_proj != nullptr && + grad_up_proj != nullptr && grad_down_proj != nullptr; // SkipLoRA: zero out lora_rank to skip all LoRA buffer allocations if constexpr (kSkipLoRA) lora_rank = 0; @@ -628,8 +635,7 @@ class TP_MOE_SFT : public TP_MOE { clear_bytes[i] = offset; } - // Parallel memset: zero only per-TP sparse partials and per-TP grad_input/grad_weights partials. - // The caller is responsible for passing zero-initialized final grad tensors. + // Parallel memset for per-TP partials and final base-weight gradients. struct ClearSeg { uint8_t* ptr; size_t len; @@ -650,6 +656,20 @@ class TP_MOE_SFT : public TP_MOE { } } + if (need_base_weight_grad) { + const size_t base_grad_bytes = (size_t)expert_num * full_intermediate_size * hidden_size * sizeof(ggml_bf16_t); + auto append_clear_segments = [&](void* ptr) { + auto* base = static_cast(ptr); + for (size_t off = 0; off < base_grad_bytes; off += kChunkBytes) { + size_t len = std::min(kChunkBytes, base_grad_bytes - off); + clear_segs.push_back(ClearSeg{base + off, len}); + } + }; + append_clear_segments(grad_gate_proj); + append_clear_segments(grad_up_proj); + append_clear_segments(grad_down_proj); + } + pool->do_work_stealing_job((int)clear_segs.size(), nullptr, [&](int seg_idx) { const auto& seg = clear_segs[(size_t)seg_idx]; @@ -665,6 +685,9 @@ class TP_MOE_SFT : public TP_MOE { std::vector tp_fp32_down_b(tp_count); std::vector tp_fp32_gate_a(tp_count); std::vector tp_fp32_up_a(tp_count); + std::vector tp_grad_gate_proj(tp_count, nullptr); + std::vector tp_grad_up_proj(tp_count, nullptr); + std::vector tp_grad_down_proj(tp_count, nullptr); if constexpr (!kSkipLoRA) { int tp_offset = 0; @@ -683,6 +706,19 @@ class TP_MOE_SFT : public TP_MOE { } } + int tp_offset = 0; + for (int i = 0; i < tp_count; i++) { + if (need_base_weight_grad) { + tp_grad_gate_proj[i] = static_cast(grad_gate_proj) + (size_t)tp_offset * hidden_size; + tp_grad_up_proj[i] = static_cast(grad_up_proj) + (size_t)tp_offset * hidden_size; + tp_grad_down_proj[i] = static_cast(grad_down_proj) + tp_offset; + } + tp_offset += tp_configs[i].intermediate_size; + } + if (tp_offset != full_intermediate_size) { + throw std::runtime_error("TP intermediate_size slices do not cover the full intermediate_size"); + } + // Run backward on each NUMA node pool->dispense_backend()->do_numa_job([&](int numa_id) { tps[numa_id]->backward(grad_output, part_grad_input_[numa_id], @@ -694,8 +730,8 @@ class TP_MOE_SFT : public TP_MOE { tp_down_a_ptr[numa_id], /* copy-type: direct write */ nullptr, /* grad_down_lora_b — unused, FP32 path below */ part_grad_weights_[numa_id], full_intermediate_size, tp_fp32_down_b[numa_id], - tp_fp32_gate_a[numa_id], tp_fp32_up_a[numa_id], grad_gate_proj, grad_up_proj, - grad_down_proj); + tp_fp32_gate_a[numa_id], tp_fp32_up_a[numa_id], tp_grad_gate_proj[numa_id], + tp_grad_up_proj[numa_id], tp_grad_down_proj[numa_id]); }); // // Collect per-thread timing from all NUMA subpools