[fix](kt-kernel): fix Full FT TP base weight gradients

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 基座梯度失效。

kt-kernel/operators/amx/sft_moe.hpp

@@ -712,6 +712,13 @@ class AMX_SFT_MOE_TP : public BaseMOE<T> {
     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<T> {
    */
   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<T> {
     // 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<T> {
    *
    * 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<ggml_bf16_t*>(grad_gate_proj);    // [E, I, H]
-    auto* gup_ptr = static_cast<ggml_bf16_t*>(grad_up_proj);  // [E, I, H]
-    auto* gdp = static_cast<ggml_bf16_t*>(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<ggml_bf16_t*>(grad_gate_proj);    // TP slice of [E, F, H]
+    auto* gup_ptr = static_cast<ggml_bf16_t*>(grad_up_proj);  // TP slice of [E, F, H]
+    auto* gdp = static_cast<ggml_bf16_t*>(grad_down_proj);    // TP slice of [E, H, F]
     auto* grad_out_bf16 = static_cast<const ggml_bf16_t*>(grad_output);
 
     for (int task_id = 0; task_id < activated_expert; task_id++) {
@@ -1969,13 +1983,13 @@ class AMX_SFT_MOE_TP : public BaseMOE<T> {
       // 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]);
         }
       }
     }

kt-kernel/operators/moe-sft-tp.hpp

@@ -229,6 +229,11 @@ class TP_MOE_SFT : public TP_MOE<T> {
     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<T> {
 
     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<T> {
       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<T> {
       }
     }
 
+    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<uint8_t*>(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<T> {
     std::vector<float*> tp_fp32_down_b(tp_count);
     std::vector<float*> tp_fp32_gate_a(tp_count);
     std::vector<float*> tp_fp32_up_a(tp_count);
+    std::vector<ggml_bf16_t*> tp_grad_gate_proj(tp_count, nullptr);
+    std::vector<ggml_bf16_t*> tp_grad_up_proj(tp_count, nullptr);
+    std::vector<ggml_bf16_t*> 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<T> {
       }
     }
 
+    int tp_offset = 0;
+    for (int i = 0; i < tp_count; i++) {
+      if (need_base_weight_grad) {
+        tp_grad_gate_proj[i] = static_cast<ggml_bf16_t*>(grad_gate_proj) + (size_t)tp_offset * hidden_size;
+        tp_grad_up_proj[i] = static_cast<ggml_bf16_t*>(grad_up_proj) + (size_t)tp_offset * hidden_size;
+        tp_grad_down_proj[i] = static_cast<ggml_bf16_t*>(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<T> {
                              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