diff --git a/doc/SUMMARY.md b/doc/SUMMARY.md index 8d1ceb37..d19f29ca 100644 --- a/doc/SUMMARY.md +++ b/doc/SUMMARY.md @@ -7,6 +7,8 @@ # Tutorial - [kt-sft part](en/SFT/README.md) + - [KT-FT Fine-Tuning and Inference Loop](en/SFT/Qwen3.5-SGLang-LoRA-Serving.md) + - [KT-FT 微调推理闭环](zh/Qwen3.5-SGLang-LoRA-Serving_zh.md) - [Injection Tutorial](en/SFT/injection_tutorial.md) - [kt-sft developer tech notes](en/SFT/KTransformers-Fine-Tuning_Developer-Technical-Notes.md) - [DPO tutorial](en/SFT/DPO_tutorial.md) diff --git a/doc/en/SFT/Qwen3.5-SGLang-LoRA-Serving.md b/doc/en/SFT/Qwen3.5-SGLang-LoRA-Serving.md new file mode 100644 index 00000000..566fcf27 --- /dev/null +++ b/doc/en/SFT/Qwen3.5-SGLang-LoRA-Serving.md @@ -0,0 +1,221 @@ +# KT-FT Fine-Tuning and Inference Loop + +Last updated: 2026-06-01 + +This guide documents the current KT-FT loop for Qwen3.5 MoE: train with KT SFT, convert the output once, and serve the fine-tuned result through SGLang with a single merged adapter path. + +```text +KT SFT raw output + -> convert_kt_to_sglang_adapter.py + -> + -> sglang --lora-paths = + -> server auto-splits expert / non-expert internally + -> request model=: +``` + +Training-side KT SFT docs remain separate. This page focuses on the bridge from trained LoRA artifacts to online inference. + +## 1. Scope + +Current supported and validated workflow: + +- Base model: Qwen3.5 MoE, for example `Qwen3.5-35B-A3B` +- KT expert weights: AMX/BF16 SFT-compatible KT CPU expert path +- User-facing serving input: one converted merged adapter directory +- Runtime split: expert LoRA goes to the KT CPU expert path; non-expert LoRA goes to SGLang's LoRA manager. This split happens automatically at server startup. + +## 2. Artifacts At Each Stage + +### Raw KT SFT output + +After LLaMA-Factory + KT training, the output directory contains two LoRA artifacts: + +```text +/ + adapter_model.safetensors # non-expert LoRA + fused_expert_lora.safetensors # expert LoRA in KT fused format + adapter_config.json +``` + +Do not pass this raw directory directly to SGLang serving. + +### Converted merged adapter + +Run the converter once to produce the serving input: + +```text +/ + adapter_config.json + adapter_model.safetensors +``` + +This merged directory contains both expert and non-expert LoRA tensors in one PEFT-style adapter. Pass only this directory to `--lora-paths`. + +## 3. Convert Once + +```bash +python kt-kernel/scripts/convert_kt_to_sglang_adapter.py \ + \ + \ + --base-model-name-or-path /path/to/Qwen3.5-35B-A3B \ + --overwrite +``` + +Example: + +```bash +python kt-kernel/scripts/convert_kt_to_sglang_adapter.py \ + saves/KT_FT_qwen35B_Moe_nekoqa_eod_240 \ + saves/KT_FT_qwen35B_Moe_nekoqa_eod_240_sglang \ + --base-model-name-or-path /mnt/data3/models/Qwen3.5-35B-A3B \ + --overwrite +``` + +The converter reads `fused_expert_lora.safetensors` and the existing non-expert `adapter_model.safetensors`, then writes one merged adapter directory. + +Optional split outputs for debugging: + +```bash +python kt-kernel/scripts/convert_kt_to_sglang_adapter.py \ + \ + \ + --base-model-name-or-path /path/to/Qwen3.5-35B-A3B \ + --expert-output-dir \ + --nonexpert-output-dir \ + --overwrite +``` + +For normal serving, only `` is needed. + +## 4. Launch SGLang + +Use the KTransformers SGLang fork from this repository and point `PYTHONPATH` at both `kt-kernel/python` and `third_party/sglang/python`. + +```bash +cd /path/to/ktransformers + +PYTHONPATH=/path/to/ktransformers/kt-kernel/python:/path/to/ktransformers/third_party/sglang/python:$PYTHONPATH \ +python -m sglang.launch_server \ + --host 127.0.0.1 \ + --port 30006 \ + --model-path /path/to/Qwen3.5-35B-A3B \ + --tokenizer-path /path/to/Qwen3.5-35B-A3B \ + --kt-weight-path /path/to/Qwen3.5-35B-A3B-AMXINT4 \ + --kt-method AMXINT4 \ + --kt-cpuinfer 60 \ + --kt-threadpool-count 2 \ + --kt-numa-nodes 0 1 \ + --kt-num-gpu-experts 0 \ + --attention-backend flashinfer \ + --trust-remote-code \ + --mem-fraction-static 0.98 \ + --chunked-prefill-size 4096 \ + --max-running-requests 2 \ + --max-total-tokens 32000 \ + --served-model-name qwen3.5-kt-ft \ + --enable-mixed-chunk \ + --tensor-parallel-size 4 \ + --enable-p2p-check \ + --disable-cuda-graph \ + --disable-custom-all-reduce \ + --enable-lora \ + --lora-backend triton \ + --lora-paths qwen35b_neko=/path/to/KT_FT_qwen35B_Moe_nekoqa_eod_240_sglang \ + --log-level info +``` + +Important points: + +- Pass only one merged adapter through `--lora-paths`. +- Do not also pass `--kt-expert-lora-path` in the normal user workflow. +- At startup, the server detects the merged KT MoE adapter, splits it internally, and writes runtime cache directories under `$TMPDIR/sglang_kt_lora_cache/` (or `$SGLANG_KT_LORA_CACHE_DIR` if set). +- Prefer `--lora-backend triton` for Qwen3.5 full-LoRA generation. + +Current constraints: + +- single merged KT composite adapter only +- `--kt-num-gpu-experts 0` +- do not enable `--kt-enable-dynamic-expert-update` +- do not use `--kt-gpu-prefill-token-threshold` +- use an AMX/BF16 SFT-compatible KT method such as `AMXINT4`, `AMXINT8`, `AMXBF16`, or `BF16` + +## 5. Request Semantics + +The OpenAI-compatible request `model` field uses names, not paths. + +```text +--served-model-name qwen3.5-kt-ft +--lora-paths qwen35b_neko=/path/to/merged_adapter +``` + +Request behavior in the current single-adapter implementation: + +```text +model=qwen3.5-kt-ft +=> base + KT expert LoRA + +model=qwen3.5-kt-ft:qwen35b_neko +=> base + KT expert LoRA + SGLang non-expert LoRA +``` + +The suffix after `:` must match the left-side name in `--lora-paths`. + +## 6. Smoke Test + +```bash +curl -sS http://127.0.0.1:30006/v1/chat/completions \ + -H 'Content-Type: application/json' \ + -d '{ + "model": "qwen3.5-kt-ft:qwen35b_neko", + "messages": [{"role": "user", "content": "我回来了，你在干嘛？"}], + "temperature": 0.7, + "max_tokens": 160, + "chat_template_kwargs": {"enable_thinking": false} + }' +``` + +Startup logs should include lines similar to: + +```text +Prepared merged KT LoRA adapter ... for runtime: expert=... nonexpert=... +Loaded KT expert LoRA for layer ... +Using triton as backend of LoRA kernels. +``` + +## 7. Advanced: Manual Split Serving + +The older split-runtime contract is still available for debugging: + +```bash +--kt-expert-lora-path \ +--enable-lora \ +--lora-paths = +``` + +This is not the recommended user-facing path. Normal users should pass one merged adapter directory through `--lora-paths` only. + +## 8. Troubleshooting + +### `Got LoRA adapter that has never been loaded: lora0` + +The adapter name in the request must match the left side of `--lora-paths`. If you launched with `qwen35b_neko=...`, request `model=qwen3.5-kt-ft:qwen35b_neko`, not `:lora0`. + +### No visible adapter effect + +Make sure you are serving the intended merged adapter directory. For example, use the Neko adapter at `..._nekoqa_eod_240_sglang`, not a generic sanity adapter such as `..._Moe_sglang`. + +### `connection refused` + +Check that the server is listening on the port you curl, and remember the example above binds to `127.0.0.1`, not `0.0.0.0`. + +### Server resolves upstream SGLang instead of this checkout + +```bash +python - <<'PY' +import inspect +import sglang.srt.models.qwen3_5 as qwen3_5 +print(inspect.getfile(qwen3_5)) +PY +``` + +The path should come from this repository's `third_party/sglang`. diff --git a/doc/en/SFT/README.md b/doc/en/SFT/README.md index 6a9234bb..324fb925 100644 --- a/doc/en/SFT/README.md +++ b/doc/en/SFT/README.md @@ -2,6 +2,7 @@ - [v0.6.1 Quick Start](./KTransformers-Fine-Tuning_Quick-Start.md) - [Fine-Tuning User Guide](./KTransformers-Fine-Tuning_User-Guide.md) +- [KT-FT Fine-Tuning and Inference Loop](./Qwen3.5-SGLang-LoRA-Serving.md) - [Developer Technical Notes](./KTransformers-Fine-Tuning_Developer-Technical-Notes.md) - [DPO Tutorial](./DPO_tutorial.md) - [Injection Tutorial](./injection_tutorial.md) diff --git a/doc/zh/Qwen3.5-SGLang-LoRA-Serving_zh.md b/doc/zh/Qwen3.5-SGLang-LoRA-Serving_zh.md new file mode 100644 index 00000000..4b890043 --- /dev/null +++ b/doc/zh/Qwen3.5-SGLang-LoRA-Serving_zh.md @@ -0,0 +1,221 @@ +# KT-FT 微调推理闭环 + +最后更新：2026-06-01 + +本文档描述当前 Qwen3.5 MoE 的 KT-FT 闭环：用 KT SFT 完成微调，转换一次输出，再通过 SGLang 用单个 merged adapter path 把微调结果服务化。 + +```text +KT SFT 原始输出 + -> convert_kt_to_sglang_adapter.py + -> + -> sglang --lora-paths = + -> server 内部自动拆分 expert / non-expert + -> 请求 model=: +``` + +训练侧 KT SFT 文档仍然独立维护；本文重点说明从已训练 LoRA artifacts 到在线推理的连接部分。 + +## 1. 范围 + +当前已验证路径： + +- 基座模型：Qwen3.5 MoE，例如 `Qwen3.5-35B-A3B` +- KT expert 权重：AMX/BF16 SFT 兼容的 KT CPU expert 路径 +- 用户侧 serving 输入：一个 converted merged adapter 目录 +- Runtime 内部仍会 split：expert LoRA 走 KT CPU expert path，non-expert LoRA 走 SGLang LoRA manager，但这一步对用户不可见 + +## 2. 各阶段产物 + +### 原始 KT SFT 输出 + +LLaMA-Factory + KT 训练完成后，输出目录里有两个 LoRA 文件： + +```text +/ + adapter_model.safetensors # non-expert LoRA + fused_expert_lora.safetensors # KT fused expert LoRA + adapter_config.json +``` + +不要把 raw 训练目录直接传给 SGLang serving。 + +### Convert 后的 merged adapter + +converter 一次性生成 serving 输入： + +```text +/ + adapter_config.json + adapter_model.safetensors +``` + +这个 merged 目录同时包含 expert 和 non-expert LoRA。正常 serving 只需要传这一个目录。 + +## 3. 转换一次 + +```bash +python kt-kernel/scripts/convert_kt_to_sglang_adapter.py \ + \ + \ + --base-model-name-or-path /path/to/Qwen3.5-35B-A3B \ + --overwrite +``` + +示例： + +```bash +python kt-kernel/scripts/convert_kt_to_sglang_adapter.py \ + saves/KT_FT_qwen35B_Moe_nekoqa_eod_240 \ + saves/KT_FT_qwen35B_Moe_nekoqa_eod_240_sglang \ + --base-model-name-or-path /mnt/data3/models/Qwen3.5-35B-A3B \ + --overwrite +``` + +converter 会读取 `fused_expert_lora.safetensors` 和已有的 non-expert `adapter_model.safetensors`，写出一个 merged adapter 目录。 + +如需调试，也可以额外输出 split 目录： + +```bash +python kt-kernel/scripts/convert_kt_to_sglang_adapter.py \ + \ + \ + --base-model-name-or-path /path/to/Qwen3.5-35B-A3B \ + --expert-output-dir \ + --nonexpert-output-dir \ + --overwrite +``` + +正常 serving 只需要 ``。 + +## 4. 启动 SGLang + +请使用本仓库的 KTransformers SGLang fork，并把 `PYTHONPATH` 指向 `kt-kernel/python` 和 `third_party/sglang/python`。 + +```bash +cd /path/to/ktransformers + +PYTHONPATH=/path/to/ktransformers/kt-kernel/python:/path/to/ktransformers/third_party/sglang/python:$PYTHONPATH \ +python -m sglang.launch_server \ + --host 127.0.0.1 \ + --port 30006 \ + --model-path /path/to/Qwen3.5-35B-A3B \ + --tokenizer-path /path/to/Qwen3.5-35B-A3B \ + --kt-weight-path /path/to/Qwen3.5-35B-A3B-AMXINT4 \ + --kt-method AMXINT4 \ + --kt-cpuinfer 60 \ + --kt-threadpool-count 2 \ + --kt-numa-nodes 0 1 \ + --kt-num-gpu-experts 0 \ + --attention-backend flashinfer \ + --trust-remote-code \ + --mem-fraction-static 0.98 \ + --chunked-prefill-size 4096 \ + --max-running-requests 2 \ + --max-total-tokens 32000 \ + --served-model-name qwen3.5-kt-ft \ + --enable-mixed-chunk \ + --tensor-parallel-size 4 \ + --enable-p2p-check \ + --disable-cuda-graph \ + --disable-custom-all-reduce \ + --enable-lora \ + --lora-backend triton \ + --lora-paths qwen35b_neko=/path/to/KT_FT_qwen35B_Moe_nekoqa_eod_240_sglang \ + --log-level info +``` + +要点： + +- 用户只需要传一个 merged adapter：`--lora-paths =` +- 正常 workflow 不要再额外传 `--kt-expert-lora-path` +- server 启动时会自动识别 merged KT MoE adapter，并在 `$TMPDIR/sglang_kt_lora_cache/`（或 `$SGLANG_KT_LORA_CACHE_DIR`）下生成 runtime cache +- Qwen3.5 full LoRA 生成优先使用 `--lora-backend triton` + +当前限制： + +- 只支持单个 merged KT composite adapter +- `--kt-num-gpu-experts 0` +- 不启用 `--kt-enable-dynamic-expert-update` +- 不使用 `--kt-gpu-prefill-token-threshold` +- 使用 AMX/BF16 SFT 兼容 KT method，例如 `AMXINT4`、`AMXINT8`、`AMXBF16`、`BF16` + +## 5. 请求语义 + +OpenAI-compatible 请求里的 `model` 字段用 name，不用 path。 + +```text +--served-model-name qwen3.5-kt-ft +--lora-paths qwen35b_neko=/path/to/merged_adapter +``` + +当前 single-adapter 实现的请求语义： + +```text +model=qwen3.5-kt-ft +=> base + KT expert LoRA + +model=qwen3.5-kt-ft:qwen35b_neko +=> base + KT expert LoRA + SGLang non-expert LoRA +``` + +冒号后的 adapter 名必须和 `--lora-paths` 左侧注册名一致。 + +## 6. Smoke Test + +```bash +curl -sS http://127.0.0.1:30006/v1/chat/completions \ + -H 'Content-Type: application/json' \ + -d '{ + "model": "qwen3.5-kt-ft:qwen35b_neko", + "messages": [{"role": "user", "content": "我回来了，你在干嘛？"}], + "temperature": 0.7, + "max_tokens": 160, + "chat_template_kwargs": {"enable_thinking": false} + }' +``` + +启动日志里应能看到类似输出： + +```text +Prepared merged KT LoRA adapter ... for runtime: expert=... nonexpert=... +Loaded KT expert LoRA for layer ... +Using triton as backend of LoRA kernels. +``` + +## 7. 高级：手动 split serving + +旧 split runtime 仍可用于调试： + +```bash +--kt-expert-lora-path \ +--enable-lora \ +--lora-paths = +``` + +这不是推荐的用户路径。正常用户只需要通过 `--lora-paths` 传一个 merged adapter 目录。 + +## 8. Troubleshooting + +### `Got LoRA adapter that has never been loaded: lora0` + +请求里的 adapter 名必须和 `--lora-paths` 左侧一致。如果启动时写的是 `qwen35b_neko=...`，请求应使用 `model=qwen3.5-kt-ft:qwen35b_neko`，而不是 `:lora0`。 + +### 看不出 adapter 效果 + +确认 serving 用的是目标 merged adapter。例如 Neko 风格应使用 `..._nekoqa_eod_240_sglang`，而不是通用 sanity adapter `..._Moe_sglang`。 + +### `connection refused` + +确认 server 监听的端口与 curl 一致；上面的示例绑定的是 `127.0.0.1`，不是 `0.0.0.0`。 + +### Server 解析到了上游 SGLang，而不是当前 checkout + +```bash +python - <<'PY' +import inspect +import sglang.srt.models.qwen3_5 as qwen3_5 +print(inspect.getfile(qwen3_5)) +PY +``` + +路径应来自本仓库的 `third_party/sglang`。 diff --git a/kt-kernel/operators/amx/sft_moe.hpp b/kt-kernel/operators/amx/sft_moe.hpp index 295c263c..357ad938 100644 --- a/kt-kernel/operators/amx/sft_moe.hpp +++ b/kt-kernel/operators/amx/sft_moe.hpp @@ -8,7 +8,6 @@ #ifndef CPUINFER_OPERATOR_AMX_SFT_MOE_H #define CPUINFER_OPERATOR_AMX_SFT_MOE_H - #include #include #include @@ -156,7 +155,6 @@ inline bool is_nan_check_enabled() { return enabled == 1; } - // ===================================================== // Pool Memory Logger — writes per-call alloc/free events to file // Enable: set SFT_POOL_LOG=1 (or any non-zero) @@ -1083,17 +1081,15 @@ class AMX_SFT_MOE_TP : public BaseMOE { } }; - direct_or_pool( - qlen, - [&](int i) { - for (int j = 0; j < k; j++) { - if (expert_ids[i * k + j] < config_.num_gpu_experts || expert_ids[i * k + j] >= config_.expert_num) { - continue; - } - memcpy(m_local_input_ptr_[expert_ids[i * k + j]] + m_local_pos_[i][j] * config_.hidden_size, - (ggml_bf16_t*)input + i * config_.hidden_size, sizeof(ggml_bf16_t) * config_.hidden_size); - } - }); + direct_or_pool(qlen, [&](int i) { + for (int j = 0; j < k; j++) { + if (expert_ids[i * k + j] < config_.num_gpu_experts || expert_ids[i * k + j] >= config_.expert_num) { + continue; + } + memcpy(m_local_input_ptr_[expert_ids[i * k + j]] + m_local_pos_[i][j] * config_.hidden_size, + (ggml_bf16_t*)input + i * config_.hidden_size, sizeof(ggml_bf16_t) * config_.hidden_size); + } + }); // NaN Check: Step 3 - Packed input if (is_nan_check_enabled()) { @@ -1110,12 +1106,10 @@ class AMX_SFT_MOE_TP : public BaseMOE { } // Step 4: Quantize input - direct_or_pool( - activated_expert, - [this](int task_id) { - int expert_idx = m_expert_id_map_[task_id]; - gate_up_ba_[expert_idx]->from_mat(m_local_num_[expert_idx], m_local_input_ptr_[expert_idx], 0, 1); - }); + direct_or_pool(activated_expert, [this](int task_id) { + int expert_idx = m_expert_id_map_[task_id]; + gate_up_ba_[expert_idx]->from_mat(m_local_num_[expert_idx], m_local_input_ptr_[expert_idx], 0, 1); + }); // Step 5: Gate + Up GEMM (base projection) int nth = T::recommended_nth(config_.intermediate_size); @@ -1222,12 +1216,8 @@ class AMX_SFT_MOE_TP : public BaseMOE { } } - - // Step 6: Activation (silu(gate) * up) - { - Base::apply_activation(activated_expert, nth, qlen); - } + { Base::apply_activation(activated_expert, nth, qlen); } // NaN Check: Step 6 - Activation output (silu(gate) * up) if (is_nan_check_enabled()) { @@ -1481,88 +1471,88 @@ class AMX_SFT_MOE_TP : public BaseMOE { // ★ Allocate backward-phase buffers ★ alloc_backward_buffers(); - // ★ share_backward_bb: check if async repack already prepared this layer ★ - if (config_.share_backward_bb) { - auto& shared = SFTSharedPools::instance(); - shared.ensure_numa_count(tp_part_idx + 1); - if (shared.pools[tp_part_idx].bwd_bb_owner_layer != config_.layer_idx) { - // Pool was overwritten by another layer or not yet repacked — sync fallback - prepare_backward_bb_for_async(); - } + // ★ share_backward_bb: check if async repack already prepared this layer ★ + if (config_.share_backward_bb) { + auto& shared = SFTSharedPools::instance(); + shared.ensure_numa_count(tp_part_idx + 1); + if (shared.pools[tp_part_idx].bwd_bb_owner_layer != config_.layer_idx) { + // Pool was overwritten by another layer or not yet repacked — sync fallback + prepare_backward_bb_for_async(); } + } - // auto print_lora_stats = [&](const char* name, const ggml_bf16_t* ptr, size_t elems) { - // if (ptr == nullptr) { - // printf("KT MoE param stats (layer %d, %s): null\n", config_.layer_idx, name); - // return; - // } - // Bf16Stats stats = compute_bf16_stats(ptr, elems); - // printf("cpp KT MoE param stats (layer %d, %s): abs_mean=%.6e abs_max=%.6e norm=%.6e\n", config_.layer_idx, - // name, - // stats.abs_mean, stats.abs_max, stats.norm); - // }; + // auto print_lora_stats = [&](const char* name, const ggml_bf16_t* ptr, size_t elems) { + // if (ptr == nullptr) { + // printf("KT MoE param stats (layer %d, %s): null\n", config_.layer_idx, name); + // return; + // } + // Bf16Stats stats = compute_bf16_stats(ptr, elems); + // printf("cpp KT MoE param stats (layer %d, %s): abs_mean=%.6e abs_max=%.6e norm=%.6e\n", config_.layer_idx, + // name, + // stats.abs_mean, stats.abs_max, stats.norm); + // }; - // size_t gate_a_elems = static_cast(config_.expert_num) * lora_rank_ * config_.hidden_size; - // size_t gate_b_elems = static_cast(config_.expert_num) * config_.intermediate_size * lora_rank_; - // size_t up_a_elems = static_cast(config_.expert_num) * lora_rank_ * config_.hidden_size; - // size_t up_b_elems = static_cast(config_.expert_num) * config_.intermediate_size * lora_rank_; - // size_t down_a_elems = static_cast(config_.expert_num) * lora_rank_ * config_.intermediate_size; - // size_t down_b_elems = static_cast(config_.expert_num) * config_.hidden_size * lora_rank_; + // size_t gate_a_elems = static_cast(config_.expert_num) * lora_rank_ * config_.hidden_size; + // size_t gate_b_elems = static_cast(config_.expert_num) * config_.intermediate_size * lora_rank_; + // size_t up_a_elems = static_cast(config_.expert_num) * lora_rank_ * config_.hidden_size; + // size_t up_b_elems = static_cast(config_.expert_num) * config_.intermediate_size * lora_rank_; + // size_t down_a_elems = static_cast(config_.expert_num) * lora_rank_ * config_.intermediate_size; + // size_t down_b_elems = static_cast(config_.expert_num) * config_.hidden_size * lora_rank_; - // print_lora_stats("gate_lora_a", gate_lora_a_, gate_a_elems); - // print_lora_stats("gate_lora_b", gate_lora_b_, gate_b_elems); - // print_lora_stats("up_lora_a", up_lora_a_, up_a_elems); - // print_lora_stats("up_lora_b", up_lora_b_, up_b_elems); - // print_lora_stats("down_lora_a", down_lora_a_, down_a_elems); - // print_lora_stats("down_lora_b", down_lora_b_, down_b_elems); + // print_lora_stats("gate_lora_a", gate_lora_a_, gate_a_elems); + // print_lora_stats("gate_lora_b", gate_lora_b_, gate_b_elems); + // print_lora_stats("up_lora_a", up_lora_a_, up_a_elems); + // print_lora_stats("up_lora_b", up_lora_b_, up_b_elems); + // print_lora_stats("down_lora_a", down_lora_a_, down_a_elems); + // print_lora_stats("down_lora_b", down_lora_b_, down_b_elems); - // Restore routing information - m_local_num_ = cache.m_local_num_cache; - m_local_pos_ = cache.m_local_pos_cache; - m_expert_id_map_ = cache.m_expert_id_map_cache; + // Restore routing information + m_local_num_ = cache.m_local_num_cache; + m_local_pos_ = cache.m_local_pos_cache; + m_expert_id_map_ = cache.m_expert_id_map_cache; - // Recompute pointer offsets - size_t offset = 0; - for (int i = 0; i < config_.expert_num; i++) { - m_local_input_ptr_[i] = m_local_input_ + offset * config_.hidden_size; - m_local_gate_output_ptr_[i] = m_local_gate_output_ + offset * config_.intermediate_size; - m_local_up_output_ptr_[i] = m_local_up_output_ + offset * config_.intermediate_size; - m_local_down_output_ptr_[i] = m_local_down_output_ + offset * config_.hidden_size; - offset += m_local_num_[i]; - } + // Recompute pointer offsets + size_t offset = 0; + for (int i = 0; i < config_.expert_num; i++) { + m_local_input_ptr_[i] = m_local_input_ + offset * config_.hidden_size; + m_local_gate_output_ptr_[i] = m_local_gate_output_ + offset * config_.intermediate_size; + m_local_up_output_ptr_[i] = m_local_up_output_ + offset * config_.intermediate_size; + m_local_down_output_ptr_[i] = m_local_down_output_ + offset * config_.hidden_size; + offset += m_local_num_[i]; + } // Restore input data from cache into m_local_input_ (shared_mem_buffer may have been // overwritten by subsequent layers' forward passes). This is needed for gate/up LoRA // gradient computation which reads from m_local_input_ptr_. - auto pool_local = config_.pool->get_subpool(tp_part_idx); - auto restore_input = [&](int i) { - for (int j = 0; j < k; j++) { - int eid = cache.expert_ids_cache[i * k + j]; - if (eid < config_.num_gpu_experts || eid >= config_.expert_num) { - continue; - } - if (m_local_num_[eid] == 0) continue; - int pos = cache.m_local_pos_cache[i][j]; - memcpy(m_local_input_ptr_[eid] + pos * config_.hidden_size, - (const ggml_bf16_t*)cache.input_cache + i * config_.hidden_size, - sizeof(ggml_bf16_t) * config_.hidden_size); + auto pool_local = config_.pool->get_subpool(tp_part_idx); + auto restore_input = [&](int i) { + for (int j = 0; j < k; j++) { + int eid = cache.expert_ids_cache[i * k + j]; + if (eid < config_.num_gpu_experts || eid >= config_.expert_num) { + continue; } - }; - if (qlen <= kSmallBwdDirectQlen && qlen <= kSmallBwdDirectMaxTasks) { - for (int i = 0; i < qlen; i++) { - restore_input(i); - } - } else { - pool_local->do_work_stealing_job(qlen, nullptr, restore_input, nullptr); + if (m_local_num_[eid] == 0) continue; + int pos = cache.m_local_pos_cache[i][j]; + memcpy(m_local_input_ptr_[eid] + pos * config_.hidden_size, + (const ggml_bf16_t*)cache.input_cache + i * config_.hidden_size, + sizeof(ggml_bf16_t) * config_.hidden_size); } + }; + if (qlen <= kSmallBwdDirectQlen && qlen <= kSmallBwdDirectMaxTasks) { + for (int i = 0; i < qlen; i++) { + restore_input(i); + } + } else { + pool_local->do_work_stealing_job(qlen, nullptr, restore_input, nullptr); + } - // Step 1: Down projection backward - if constexpr (supports_standard_mat_mul_v) { - backward_down_amx(cache, grad_output, grad_down_lora_a, grad_down_lora_b, full_intermediate_size, - fp32_grad_down_lora_b); - } else { - // backward_down(cache, grad_output, grad_down_lora_a, grad_down_lora_b); - } + // Step 1: Down projection backward + if constexpr (supports_standard_mat_mul_v) { + backward_down_amx(cache, grad_output, grad_down_lora_a, grad_down_lora_b, full_intermediate_size, + fp32_grad_down_lora_b); + } else { + // backward_down(cache, grad_output, grad_down_lora_a, grad_down_lora_b); + } // // Compute total tokens for debug // size_t total_tokens = 0; @@ -1678,12 +1668,12 @@ class AMX_SFT_MOE_TP : public BaseMOE { // } // } - if constexpr (supports_standard_mat_mul_v) { - backward_gate_up_amx(cache, grad_input, grad_gate_lora_a, grad_gate_lora_b, grad_up_lora_a, grad_up_lora_b, - full_intermediate_size, fp32_grad_gate_lora_a, fp32_grad_up_lora_a); - } else { - // backward_gate_up(cache, grad_input, grad_gate_lora_a, grad_gate_lora_b, grad_up_lora_a, grad_up_lora_b); - } + if constexpr (supports_standard_mat_mul_v) { + backward_gate_up_amx(cache, grad_input, grad_gate_lora_a, grad_gate_lora_b, grad_up_lora_a, grad_up_lora_b, + full_intermediate_size, fp32_grad_gate_lora_a, fp32_grad_up_lora_a); + } else { + // backward_gate_up(cache, grad_input, grad_gate_lora_a, grad_gate_lora_b, grad_up_lora_a, grad_up_lora_b); + } // NaN Check: Step 3 - After backward_gate_up if (is_nan_check_enabled()) { @@ -1721,55 +1711,55 @@ class AMX_SFT_MOE_TP : public BaseMOE { // Step 4: Compute grad_weights (gradient for routing weights) // grad_weights[token_idx, expert_pos] = dot(grad_output[token_idx], down_output[token, expert]) if (grad_weights != nullptr) { - auto pool = config_.pool->get_subpool(tp_part_idx); - float* grad_w = (float*)grad_weights; - const ggml_bf16_t* grad_out = (const ggml_bf16_t*)grad_output; + auto pool = config_.pool->get_subpool(tp_part_idx); + float* grad_w = (float*)grad_weights; + const ggml_bf16_t* grad_out = (const ggml_bf16_t*)grad_output; - // Compute offset mapping for down_output_cache (same layout as other caches) - std::vector expert_cache_offset(config_.expert_num, 0); - size_t offset = 0; - for (int i = 0; i < activated_expert; i++) { - int expert_idx = cache.m_expert_id_map_cache[i]; - expert_cache_offset[expert_idx] = offset; - offset += cache.m_local_num_cache[expert_idx]; - } + // Compute offset mapping for down_output_cache (same layout as other caches) + std::vector expert_cache_offset(config_.expert_num, 0); + size_t offset = 0; + for (int i = 0; i < activated_expert; i++) { + int expert_idx = cache.m_expert_id_map_cache[i]; + expert_cache_offset[expert_idx] = offset; + offset += cache.m_local_num_cache[expert_idx]; + } - // Compute grad_weights for each token-expert pair - auto compute_grad_weight = [&](int token_idx) { - for (int j = 0; j < k; j++) { - int64_t expert_idx = cache.expert_ids_cache[token_idx * k + j]; - if (expert_idx < config_.num_gpu_experts || expert_idx >= config_.expert_num) { - continue; // Skip GPU experts or invalid experts - } - - int local_pos = cache.m_local_pos_cache[token_idx][j]; - size_t down_offset = expert_cache_offset[expert_idx] + local_pos; - - // dot(grad_output[token_idx], down_output_cache[down_offset]) - const ggml_bf16_t* grad_out_ptr = grad_out + token_idx * config_.hidden_size; - const ggml_bf16_t* down_out_ptr = cache.down_output_cache + down_offset * config_.hidden_size; - - __m512 acc0 = _mm512_setzero_ps(); - __m512 acc1 = _mm512_setzero_ps(); - - for (int h = 0; h + 32 <= config_.hidden_size; h += 32) { - __m512 g0, g1, d0, d1; - avx512_32xbf16_to_32xfp32((__m512i*)(grad_out_ptr + h), &g0, &g1); - avx512_32xbf16_to_32xfp32((__m512i*)(down_out_ptr + h), &d0, &d1); - acc0 = _mm512_fmadd_ps(g0, d0, acc0); - acc1 = _mm512_fmadd_ps(g1, d1, acc1); - } - - grad_w[token_idx * k + j] = _mm512_reduce_add_ps(acc0) + _mm512_reduce_add_ps(acc1); + // Compute grad_weights for each token-expert pair + auto compute_grad_weight = [&](int token_idx) { + for (int j = 0; j < k; j++) { + int64_t expert_idx = cache.expert_ids_cache[token_idx * k + j]; + if (expert_idx < config_.num_gpu_experts || expert_idx >= config_.expert_num) { + continue; // Skip GPU experts or invalid experts } - }; - if (qlen <= kSmallBwdDirectQlen && qlen <= kSmallBwdDirectMaxTasks) { - for (int token_idx = 0; token_idx < qlen; token_idx++) { - compute_grad_weight(token_idx); + + int local_pos = cache.m_local_pos_cache[token_idx][j]; + size_t down_offset = expert_cache_offset[expert_idx] + local_pos; + + // dot(grad_output[token_idx], down_output_cache[down_offset]) + const ggml_bf16_t* grad_out_ptr = grad_out + token_idx * config_.hidden_size; + const ggml_bf16_t* down_out_ptr = cache.down_output_cache + down_offset * config_.hidden_size; + + __m512 acc0 = _mm512_setzero_ps(); + __m512 acc1 = _mm512_setzero_ps(); + + for (int h = 0; h + 32 <= config_.hidden_size; h += 32) { + __m512 g0, g1, d0, d1; + avx512_32xbf16_to_32xfp32((__m512i*)(grad_out_ptr + h), &g0, &g1); + avx512_32xbf16_to_32xfp32((__m512i*)(down_out_ptr + h), &d0, &d1); + acc0 = _mm512_fmadd_ps(g0, d0, acc0); + acc1 = _mm512_fmadd_ps(g1, d1, acc1); } - } else { - pool->do_work_stealing_job(qlen, nullptr, compute_grad_weight, nullptr); + + grad_w[token_idx * k + j] = _mm512_reduce_add_ps(acc0) + _mm512_reduce_add_ps(acc1); } + }; + if (qlen <= kSmallBwdDirectQlen && qlen <= kSmallBwdDirectMaxTasks) { + for (int token_idx = 0; token_idx < qlen; token_idx++) { + compute_grad_weight(token_idx); + } + } else { + pool->do_work_stealing_job(qlen, nullptr, compute_grad_weight, nullptr); + } } // NaN Check: Step 4 - After grad_weights computation @@ -2610,7 +2600,11 @@ class AMX_SFT_MOE_TP : public BaseMOE { constexpr int K_STEP = T::K_STEP; constexpr int N_STEP = T::N_STEP; constexpr int M_STEP = T::M_STEP; - padded_lora_rank_ = ((lora_rank_ + K_STEP - 1) / K_STEP) * K_STEP; + int lora_k_step = K_STEP; + if constexpr (std::is_same_v || std::is_same_v) { + lora_k_step = 2 * K_STEP; + } + padded_lora_rank_ = ((lora_rank_ + lora_k_step - 1) / lora_k_step) * lora_k_step; // Also need N dimension aligned for BufferB output dimension int padded_lora_rank_n = ((lora_rank_ + N_STEP - 1) / N_STEP) * N_STEP; // Use the larger of the two for consistency @@ -3562,8 +3556,6 @@ class AMX_SFT_MOE_TP : public BaseMOE { }, nullptr); - - // ===================================================== // Step 2: Quantize lora_intermediate to BufferA // ===================================================== @@ -4230,14 +4222,12 @@ class AMX_SFT_MOE_TP : public BaseMOE { // ===================================================== // Step 1: Zero per-expert grad_output buffers // ===================================================== - direct_or_pool( - activated_expert, - [this](int task_id) { - int expert_idx = m_expert_id_map_[task_id]; - int num_tokens = m_local_num_[expert_idx]; - if (num_tokens == 0) return; - memset(grad_output_bf16_ptr_[expert_idx], 0, num_tokens * config_.hidden_size * sizeof(ggml_bf16_t)); - }); + direct_or_pool(activated_expert, [this](int task_id) { + int expert_idx = m_expert_id_map_[task_id]; + int num_tokens = m_local_num_[expert_idx]; + if (num_tokens == 0) return; + memset(grad_output_bf16_ptr_[expert_idx], 0, num_tokens * config_.hidden_size * sizeof(ggml_bf16_t)); + }); // ===================================================== // Step 2: Scatter grad_output to per-expert BF16 buffers @@ -4246,57 +4236,53 @@ class AMX_SFT_MOE_TP : public BaseMOE { const int hidden = config_.hidden_size; const int hidden_vec_end = hidden & ~31; - direct_or_pool( - qlen, - [this, &cache, grad_output, k, hidden, hidden_vec_end](int token_id) { - const ggml_bf16_t* src_row = (const ggml_bf16_t*)grad_output + token_id * hidden; + direct_or_pool(qlen, [this, &cache, grad_output, k, hidden, hidden_vec_end](int token_id) { + const ggml_bf16_t* src_row = (const ggml_bf16_t*)grad_output + token_id * hidden; - for (int j = 0; j < k; j++) { - int expert_idx = cache.expert_ids_cache[token_id * k + j]; - if (expert_idx < config_.num_gpu_experts || expert_idx >= config_.expert_num) { - continue; - } - if (m_local_num_[expert_idx] == 0) { - continue; - } + for (int j = 0; j < k; j++) { + int expert_idx = cache.expert_ids_cache[token_id * k + j]; + if (expert_idx < config_.num_gpu_experts || expert_idx >= config_.expert_num) { + continue; + } + if (m_local_num_[expert_idx] == 0) { + continue; + } - // Each token-route pair owns one unique local position within an expert buffer. - int pos = cache.m_local_pos_cache[token_id][j]; - float w = cache.weights_cache[token_id * k + j]; - ggml_bf16_t* dst_row = grad_output_bf16_ptr_[expert_idx] + pos * hidden; + // Each token-route pair owns one unique local position within an expert buffer. + int pos = cache.m_local_pos_cache[token_id][j]; + float w = cache.weights_cache[token_id * k + j]; + ggml_bf16_t* dst_row = grad_output_bf16_ptr_[expert_idx] + pos * hidden; - __m512 w_vec = _mm512_set1_ps(w); - int h = 0; - for (; h < hidden_vec_end; h += 32) { - __m512 x0, x1, cur0, cur1; - avx512_32xbf16_to_32xfp32((__m512i*)(src_row + h), &x0, &x1); - avx512_32xbf16_to_32xfp32((__m512i*)(dst_row + h), &cur0, &cur1); - x0 = _mm512_mul_ps(x0, w_vec); - x1 = _mm512_mul_ps(x1, w_vec); - x0 = _mm512_add_ps(x0, cur0); - x1 = _mm512_add_ps(x1, cur1); - avx512_32xfp32_to_32xbf16(&x0, &x1, (__m512i*)(dst_row + h)); - } - for (; h < hidden; h++) { - float cur = GGML_BF16_TO_FP32(dst_row[h]); - cur += GGML_BF16_TO_FP32(src_row[h]) * w; - dst_row[h] = GGML_FP32_TO_BF16(cur); - } - } - }); + __m512 w_vec = _mm512_set1_ps(w); + int h = 0; + for (; h < hidden_vec_end; h += 32) { + __m512 x0, x1, cur0, cur1; + avx512_32xbf16_to_32xfp32((__m512i*)(src_row + h), &x0, &x1); + avx512_32xbf16_to_32xfp32((__m512i*)(dst_row + h), &cur0, &cur1); + x0 = _mm512_mul_ps(x0, w_vec); + x1 = _mm512_mul_ps(x1, w_vec); + x0 = _mm512_add_ps(x0, cur0); + x1 = _mm512_add_ps(x1, cur1); + avx512_32xfp32_to_32xbf16(&x0, &x1, (__m512i*)(dst_row + h)); + } + for (; h < hidden; h++) { + float cur = GGML_BF16_TO_FP32(dst_row[h]); + cur += GGML_BF16_TO_FP32(src_row[h]) * w; + dst_row[h] = GGML_FP32_TO_BF16(cur); + } + } + }); } // ===================================================== // Step 3: Quantize scattered grad_output to BufferA // ===================================================== - direct_or_pool( - activated_expert, - [this](int task_id) { - int expert_idx = m_expert_id_map_[task_id]; - int num_tokens = m_local_num_[expert_idx]; - if (num_tokens == 0) return; - grad_output_ba_[expert_idx]->from_mat(num_tokens, grad_output_bf16_ptr_[expert_idx], 0, 1); - }); + direct_or_pool(activated_expert, [this](int task_id) { + int expert_idx = m_expert_id_map_[task_id]; + int num_tokens = m_local_num_[expert_idx]; + if (num_tokens == 0) return; + grad_output_ba_[expert_idx]->from_mat(num_tokens, grad_output_bf16_ptr_[expert_idx], 0, 1); + }); // ===================================================== // Step 3+4: AMX GEMM + to_mat (merged to use same ith/nth) @@ -4355,49 +4341,46 @@ class AMX_SFT_MOE_TP : public BaseMOE { const float scale = lora_scaling_; const int nth = 4; - direct_or_pool( - nth * activated_expert, - [this, &expert_offsets, &expert_token_offsets, hidden, inter_size, rank, scale, nth](int task_id) { - int expert_idx = m_expert_id_map_[task_id / nth]; - int ith = task_id % nth; - int num_tokens = m_local_num_[expert_idx]; - if (num_tokens == 0) return; + direct_or_pool(nth * activated_expert, [this, &expert_offsets, &expert_token_offsets, hidden, inter_size, rank, + scale, nth](int task_id) { + int expert_idx = m_expert_id_map_[task_id / nth]; + int ith = task_id % nth; + int num_tokens = m_local_num_[expert_idx]; + if (num_tokens == 0) return; - // Divide tokens among threads - int tokens_per_thread = (num_tokens + nth - 1) / nth; - int t_start = ith * tokens_per_thread; - int t_end = std::min(t_start + tokens_per_thread, num_tokens); - if (t_start >= num_tokens) return; + // Divide tokens among threads + int tokens_per_thread = (num_tokens + nth - 1) / nth; + int t_start = ith * tokens_per_thread; + int t_end = std::min(t_start + tokens_per_thread, num_tokens); + if (t_start >= num_tokens) return; - // Get expert's LoRA weights (use transposed layout for lora_B) - size_t lora_a_offset = (size_t)expert_idx * rank * inter_size; - size_t lora_b_t_offset = (size_t)expert_idx * rank * hidden; // Transposed: [rank, hidden] - const ggml_bf16_t* expert_lora_a = down_lora_a_ + lora_a_offset; - const ggml_bf16_t* expert_lora_b_t = down_lora_b_transposed_ + lora_b_t_offset; - const ggml_bf16_t* expert_grad = grad_output_bf16_ptr_[expert_idx]; - ggml_bf16_t* grad_inter = grad_intermediate_ + expert_offsets[task_id / nth]; - float* grad_times_b = lora_grad_times_b_pool_ + (expert_token_offsets[task_id / nth] + t_start) * rank; + // Get expert's LoRA weights (use transposed layout for lora_B) + size_t lora_a_offset = (size_t)expert_idx * rank * inter_size; + size_t lora_b_t_offset = (size_t)expert_idx * rank * hidden; // Transposed: [rank, hidden] + const ggml_bf16_t* expert_lora_a = down_lora_a_ + lora_a_offset; + const ggml_bf16_t* expert_lora_b_t = down_lora_b_transposed_ + lora_b_t_offset; + const ggml_bf16_t* expert_grad = grad_output_bf16_ptr_[expert_idx]; + ggml_bf16_t* grad_inter = grad_intermediate_ + expert_offsets[task_id / nth]; + float* grad_times_b = lora_grad_times_b_pool_ + (expert_token_offsets[task_id / nth] + t_start) * rank; - int local_num_tokens = t_end - t_start; + int local_num_tokens = t_end - t_start; - // Step 1: grad_output @ down_lora_B_transposed -> [local_num_tokens, rank] - // Using optimized kernel with transposed weight layout [rank, hidden] - avx::lora_backward_matmul_transposed(expert_grad + t_start * hidden, // [local_num_tokens, hidden] BF16 - expert_lora_b_t, // [rank, hidden] BF16 (transposed) - grad_times_b, // [local_num_tokens, rank] FP32 - local_num_tokens, hidden, rank); + // Step 1: grad_output @ down_lora_B_transposed -> [local_num_tokens, rank] + // Using optimized kernel with transposed weight layout [rank, hidden] + avx::lora_backward_matmul_transposed(expert_grad + t_start * hidden, // [local_num_tokens, hidden] BF16 + expert_lora_b_t, // [rank, hidden] BF16 (transposed) + grad_times_b, // [local_num_tokens, rank] FP32 + local_num_tokens, hidden, rank); - // Step 2: grad_times_b @ down_lora_A -> [local_num_tokens, inter_size] (AVX512) - // Using optimized kernel with weight layout [rank, inter_size] - avx::lora_fp32_bf16_fused_add_wt(grad_times_b, // [local_num_tokens, rank] FP32 - expert_lora_a, // [rank, inter_size] BF16 - grad_inter + t_start * inter_size, // [local_num_tokens, inter_size] BF16 - local_num_tokens, rank, inter_size, scale); - }); + // Step 2: grad_times_b @ down_lora_A -> [local_num_tokens, inter_size] (AVX512) + // Using optimized kernel with weight layout [rank, inter_size] + avx::lora_fp32_bf16_fused_add_wt(grad_times_b, // [local_num_tokens, rank] FP32 + expert_lora_a, // [rank, inter_size] BF16 + grad_inter + t_start * inter_size, // [local_num_tokens, inter_size] BF16 + local_num_tokens, rank, inter_size, scale); + }); } - - // ===================================================== // Step 5: LoRA gradient computation (parallelized across blocks) // Skip when SkipLoRA is true (only compute grad_input, not LoRA weight gradients) @@ -4606,121 +4589,120 @@ class AMX_SFT_MOE_TP : public BaseMOE { } if (!lora_grad_tasks.empty()) { - direct_or_pool( - static_cast(lora_grad_tasks.size()), - [&, hidden, inter_size, rank, grad_b_elems, grad_a_elems](int task_id) { - const LoraGradTask& task = lora_grad_tasks[task_id]; - LoraGradBuf& buf = lora_grad_bufs[task.expert_task]; - if (buf.num_tokens == 0) return; + direct_or_pool(static_cast(lora_grad_tasks.size()), [&, hidden, inter_size, rank, grad_b_elems, + grad_a_elems](int task_id) { + const LoraGradTask& task = lora_grad_tasks[task_id]; + LoraGradBuf& buf = lora_grad_bufs[task.expert_task]; + if (buf.num_tokens == 0) return; - const int hidden_vec_end = hidden & ~31; - const int inter_vec_end = inter_size & ~31; - size_t scratch_elems = grad_b_elems + grad_a_elems + hidden + inter_size; - float* scratch = get_lora_fp32_buffer(scratch_elems); - float* grad_b_local = scratch; - float* grad_a_local = grad_b_local + grad_b_elems; - float* grad_row_fp32 = grad_a_local + grad_a_elems; - float* inter_row_fp32 = grad_row_fp32 + hidden; + const int hidden_vec_end = hidden & ~31; + const int inter_vec_end = inter_size & ~31; + size_t scratch_elems = grad_b_elems + grad_a_elems + hidden + inter_size; + float* scratch = get_lora_fp32_buffer(scratch_elems); + float* grad_b_local = scratch; + float* grad_a_local = grad_b_local + grad_b_elems; + float* grad_row_fp32 = grad_a_local + grad_a_elems; + float* inter_row_fp32 = grad_row_fp32 + hidden; - memset(grad_b_local, 0, (grad_b_elems + grad_a_elems) * sizeof(float)); + memset(grad_b_local, 0, (grad_b_elems + grad_a_elems) * sizeof(float)); - for (int t = task.t_start; t < task.t_end; t++) { - const ggml_bf16_t* grad_row_bf16 = buf.expert_grad_bf16 + static_cast(t) * hidden; - const ggml_bf16_t* inter_row_bf16 = buf.cached_intermediate + static_cast(t) * inter_size; + for (int t = task.t_start; t < task.t_end; t++) { + const ggml_bf16_t* grad_row_bf16 = buf.expert_grad_bf16 + static_cast(t) * hidden; + const ggml_bf16_t* inter_row_bf16 = buf.cached_intermediate + static_cast(t) * inter_size; - int h = 0; - for (; h < hidden_vec_end; h += 32) { - __m512 g0, g1; - avx512_32xbf16_to_32xfp32((__m512i*)(grad_row_bf16 + h), &g0, &g1); - _mm512_storeu_ps(grad_row_fp32 + h, g0); - _mm512_storeu_ps(grad_row_fp32 + h + 16, g1); - } - for (; h < hidden; h++) { - grad_row_fp32[h] = GGML_BF16_TO_FP32(grad_row_bf16[h]); - } + int h = 0; + for (; h < hidden_vec_end; h += 32) { + __m512 g0, g1; + avx512_32xbf16_to_32xfp32((__m512i*)(grad_row_bf16 + h), &g0, &g1); + _mm512_storeu_ps(grad_row_fp32 + h, g0); + _mm512_storeu_ps(grad_row_fp32 + h + 16, g1); + } + for (; h < hidden; h++) { + grad_row_fp32[h] = GGML_BF16_TO_FP32(grad_row_bf16[h]); + } - int i = 0; - for (; i < inter_vec_end; i += 32) { - __m512 x0, x1; - avx512_32xbf16_to_32xfp32((__m512i*)(inter_row_bf16 + i), &x0, &x1); - _mm512_storeu_ps(inter_row_fp32 + i, x0); - _mm512_storeu_ps(inter_row_fp32 + i + 16, x1); - } - for (; i < inter_size; i++) { - inter_row_fp32[i] = GGML_BF16_TO_FP32(inter_row_bf16[i]); - } + int i = 0; + for (; i < inter_vec_end; i += 32) { + __m512 x0, x1; + avx512_32xbf16_to_32xfp32((__m512i*)(inter_row_bf16 + i), &x0, &x1); + _mm512_storeu_ps(inter_row_fp32 + i, x0); + _mm512_storeu_ps(inter_row_fp32 + i + 16, x1); + } + for (; i < inter_size; i++) { + inter_row_fp32[i] = GGML_BF16_TO_FP32(inter_row_bf16[i]); + } - const float* inter_proj = buf.cached_down_lora_u + static_cast(t) * rank; - const float* grad_times_b = buf.grad_times_b + static_cast(t) * rank; + const float* inter_proj = buf.cached_down_lora_u + static_cast(t) * rank; + const float* grad_times_b = buf.grad_times_b + static_cast(t) * rank; - if (rank == 8) { - __m256 inter_proj_vec = _mm256_loadu_ps(inter_proj); - for (int hh = 0; hh < hidden; hh++) { - float g = grad_row_fp32[hh]; - if (g == 0.0f) continue; - float* out = grad_b_local + static_cast(hh) * rank; - __m256 acc = _mm256_loadu_ps(out); - acc = _mm256_fmadd_ps(_mm256_set1_ps(g), inter_proj_vec, acc); - _mm256_storeu_ps(out, acc); - } + if (rank == 8) { + __m256 inter_proj_vec = _mm256_loadu_ps(inter_proj); + for (int hh = 0; hh < hidden; hh++) { + float g = grad_row_fp32[hh]; + if (g == 0.0f) continue; + float* out = grad_b_local + static_cast(hh) * rank; + __m256 acc = _mm256_loadu_ps(out); + acc = _mm256_fmadd_ps(_mm256_set1_ps(g), inter_proj_vec, acc); + _mm256_storeu_ps(out, acc); + } - __m256 grad_times_b_vec = _mm256_loadu_ps(grad_times_b); - for (int ii = 0; ii < inter_size; ii++) { - float x = inter_row_fp32[ii]; - if (x == 0.0f) continue; - float* out = grad_a_local + static_cast(ii) * rank; - __m256 acc = _mm256_loadu_ps(out); - acc = _mm256_fmadd_ps(_mm256_set1_ps(x), grad_times_b_vec, acc); - _mm256_storeu_ps(out, acc); - } - } else { - for (int hh = 0; hh < hidden; hh++) { - float g = grad_row_fp32[hh]; - if (g == 0.0f) continue; - float* out = grad_b_local + static_cast(hh) * rank; - for (int r = 0; r < rank; r++) { - out[r] += g * inter_proj[r]; - } - } - - for (int ii = 0; ii < inter_size; ii++) { - float x = inter_row_fp32[ii]; - if (x == 0.0f) continue; - float* out = grad_a_local + static_cast(ii) * rank; - for (int r = 0; r < rank; r++) { - out[r] += x * grad_times_b[r]; - } - } + __m256 grad_times_b_vec = _mm256_loadu_ps(grad_times_b); + for (int ii = 0; ii < inter_size; ii++) { + float x = inter_row_fp32[ii]; + if (x == 0.0f) continue; + float* out = grad_a_local + static_cast(ii) * rank; + __m256 acc = _mm256_loadu_ps(out); + acc = _mm256_fmadd_ps(_mm256_set1_ps(x), grad_times_b_vec, acc); + _mm256_storeu_ps(out, acc); + } + } else { + for (int hh = 0; hh < hidden; hh++) { + float g = grad_row_fp32[hh]; + if (g == 0.0f) continue; + float* out = grad_b_local + static_cast(hh) * rank; + for (int r = 0; r < rank; r++) { + out[r] += g * inter_proj[r]; } } - std::lock_guard lock(down_lora_grad_mutexes_[task.expert_task]); - float* grad_b_global = grad_b_accum_all + static_cast(task.expert_task) * grad_b_elems; - float* grad_a_global = grad_a_accum_all + static_cast(task.expert_task) * grad_a_elems; - if (!down_lora_grad_accum_initialized_[task.expert_task]) { - std::memcpy(grad_b_global, grad_b_local, grad_b_elems * sizeof(float)); - std::memcpy(grad_a_global, grad_a_local, grad_a_elems * sizeof(float)); - down_lora_grad_accum_initialized_[task.expert_task] = 1; - } else if (rank == 8) { - for (size_t off = 0; off < grad_b_elems; off += rank) { - __m256 acc = _mm256_loadu_ps(grad_b_global + off); - acc = _mm256_add_ps(acc, _mm256_loadu_ps(grad_b_local + off)); - _mm256_storeu_ps(grad_b_global + off, acc); - } - for (size_t off = 0; off < grad_a_elems; off += rank) { - __m256 acc = _mm256_loadu_ps(grad_a_global + off); - acc = _mm256_add_ps(acc, _mm256_loadu_ps(grad_a_local + off)); - _mm256_storeu_ps(grad_a_global + off, acc); - } - } else { - for (size_t off = 0; off < grad_b_elems; off++) { - grad_b_global[off] += grad_b_local[off]; - } - for (size_t off = 0; off < grad_a_elems; off++) { - grad_a_global[off] += grad_a_local[off]; + for (int ii = 0; ii < inter_size; ii++) { + float x = inter_row_fp32[ii]; + if (x == 0.0f) continue; + float* out = grad_a_local + static_cast(ii) * rank; + for (int r = 0; r < rank; r++) { + out[r] += x * grad_times_b[r]; } } - }); + } + } + + std::lock_guard lock(down_lora_grad_mutexes_[task.expert_task]); + float* grad_b_global = grad_b_accum_all + static_cast(task.expert_task) * grad_b_elems; + float* grad_a_global = grad_a_accum_all + static_cast(task.expert_task) * grad_a_elems; + if (!down_lora_grad_accum_initialized_[task.expert_task]) { + std::memcpy(grad_b_global, grad_b_local, grad_b_elems * sizeof(float)); + std::memcpy(grad_a_global, grad_a_local, grad_a_elems * sizeof(float)); + down_lora_grad_accum_initialized_[task.expert_task] = 1; + } else if (rank == 8) { + for (size_t off = 0; off < grad_b_elems; off += rank) { + __m256 acc = _mm256_loadu_ps(grad_b_global + off); + acc = _mm256_add_ps(acc, _mm256_loadu_ps(grad_b_local + off)); + _mm256_storeu_ps(grad_b_global + off, acc); + } + for (size_t off = 0; off < grad_a_elems; off += rank) { + __m256 acc = _mm256_loadu_ps(grad_a_global + off); + acc = _mm256_add_ps(acc, _mm256_loadu_ps(grad_a_local + off)); + _mm256_storeu_ps(grad_a_global + off, acc); + } + } else { + for (size_t off = 0; off < grad_b_elems; off++) { + grad_b_global[off] += grad_b_local[off]; + } + for (size_t off = 0; off < grad_a_elems; off++) { + grad_a_global[off] += grad_a_local[off]; + } + } + }); constexpr int kDownGradBTile = 512; constexpr int kDownGradATile = 512; @@ -4830,91 +4812,88 @@ class AMX_SFT_MOE_TP : public BaseMOE { // dy/d(up) = silu(gate) = gate * sigmoid(gate) size_t cache_offset = 0; - direct_or_pool( - activated_expert, - [this, &cache, &cache_offset](int task_id) { - int expert_idx = m_expert_id_map_[task_id]; - int num_tokens = m_local_num_[expert_idx]; + direct_or_pool(activated_expert, [this, &cache, &cache_offset](int task_id) { + int expert_idx = m_expert_id_map_[task_id]; + int num_tokens = m_local_num_[expert_idx]; - if (num_tokens == 0) return; + if (num_tokens == 0) return; - // Get cached gate and up outputs (before activation) - // Need to compute offset into cache - size_t offset = 0; - for (int i = 0; i < task_id; i++) { - offset += m_local_num_[m_expert_id_map_[i]]; - } + // Get cached gate and up outputs (before activation) + // Need to compute offset into cache + size_t offset = 0; + for (int i = 0; i < task_id; i++) { + offset += m_local_num_[m_expert_id_map_[i]]; + } - ggml_bf16_t* gate_output = cache.gate_output_cache + offset * config_.intermediate_size; - ggml_bf16_t* up_output = cache.up_output_cache + offset * config_.intermediate_size; - ggml_bf16_t* grad_inter = grad_intermediate_ + offset * config_.intermediate_size; - ggml_bf16_t* grad_gate = grad_gate_output_ + offset * config_.intermediate_size; - ggml_bf16_t* grad_up = grad_up_output_ + offset * config_.intermediate_size; + ggml_bf16_t* gate_output = cache.gate_output_cache + offset * config_.intermediate_size; + ggml_bf16_t* up_output = cache.up_output_cache + offset * config_.intermediate_size; + ggml_bf16_t* grad_inter = grad_intermediate_ + offset * config_.intermediate_size; + ggml_bf16_t* grad_gate = grad_gate_output_ + offset * config_.intermediate_size; + ggml_bf16_t* grad_up = grad_up_output_ + offset * config_.intermediate_size; - // Debug code commented out - Bug #15 verified fixed - // if (task_id == 0) { - // printf("[DEBUG backward_activation] task_id=0, expert_idx=%d, num_tokens=%d, offset=%zu\n", expert_idx, - // num_tokens, offset); - // printf("[DEBUG] gate_output[0..7] = "); - // for (int dbg = 0; dbg < 8 && dbg < num_tokens * config_.intermediate_size; dbg++) { - // printf("%.4f ", GGML_BF16_TO_FP32(gate_output[dbg])); - // } - // printf("\n"); - // printf("[DEBUG] up_output[0..7] = "); - // for (int dbg = 0; dbg < 8 && dbg < num_tokens * config_.intermediate_size; dbg++) { - // printf("%.4f ", GGML_BF16_TO_FP32(up_output[dbg])); - // } - // printf("\n"); - // printf("[DEBUG] grad_inter[0..7] = "); - // for (int dbg = 0; dbg < 8 && dbg < num_tokens * config_.intermediate_size; dbg++) { - // printf("%.4f ", GGML_BF16_TO_FP32(grad_inter[dbg])); - // } - // printf("\n"); - // } + // Debug code commented out - Bug #15 verified fixed + // if (task_id == 0) { + // printf("[DEBUG backward_activation] task_id=0, expert_idx=%d, num_tokens=%d, offset=%zu\n", expert_idx, + // num_tokens, offset); + // printf("[DEBUG] gate_output[0..7] = "); + // for (int dbg = 0; dbg < 8 && dbg < num_tokens * config_.intermediate_size; dbg++) { + // printf("%.4f ", GGML_BF16_TO_FP32(gate_output[dbg])); + // } + // printf("\n"); + // printf("[DEBUG] up_output[0..7] = "); + // for (int dbg = 0; dbg < 8 && dbg < num_tokens * config_.intermediate_size; dbg++) { + // printf("%.4f ", GGML_BF16_TO_FP32(up_output[dbg])); + // } + // printf("\n"); + // printf("[DEBUG] grad_inter[0..7] = "); + // for (int dbg = 0; dbg < 8 && dbg < num_tokens * config_.intermediate_size; dbg++) { + // printf("%.4f ", GGML_BF16_TO_FP32(grad_inter[dbg])); + // } + // printf("\n"); + // } - int total = num_tokens * config_.intermediate_size; - int i = 0; + int total = num_tokens * config_.intermediate_size; + int i = 0; - // AVX512: process 32 BF16 elements (2×16 FP32) per iteration - __m512 one = _mm512_set1_ps(1.0f); - for (; i + 32 <= total; i += 32) { - __m512 g0, g1, u0, u1, gi0, gi1; - avx512_32xbf16_to_32xfp32((__m512i*)(gate_output + i), &g0, &g1); - avx512_32xbf16_to_32xfp32((__m512i*)(up_output + i), &u0, &u1); - avx512_32xbf16_to_32xfp32((__m512i*)(grad_inter + i), &gi0, &gi1); + // AVX512: process 32 BF16 elements (2×16 FP32) per iteration + __m512 one = _mm512_set1_ps(1.0f); + for (; i + 32 <= total; i += 32) { + __m512 g0, g1, u0, u1, gi0, gi1; + avx512_32xbf16_to_32xfp32((__m512i*)(gate_output + i), &g0, &g1); + avx512_32xbf16_to_32xfp32((__m512i*)(up_output + i), &u0, &u1); + avx512_32xbf16_to_32xfp32((__m512i*)(grad_inter + i), &gi0, &gi1); - // First 16: sigmoid, silu derivative, gradients - __m512 exp0 = avx512_exp_ps(_mm512_sub_ps(_mm512_setzero_ps(), g0)); - __m512 sig0 = _mm512_div_ps(one, _mm512_add_ps(one, exp0)); - __m512 silu0 = _mm512_mul_ps(g0, sig0); - __m512 dsilu0 = _mm512_mul_ps(sig0, _mm512_fmadd_ps(g0, _mm512_sub_ps(one, sig0), one)); - __m512 gg0 = _mm512_mul_ps(_mm512_mul_ps(gi0, u0), dsilu0); - __m512 gu0 = _mm512_mul_ps(gi0, silu0); + // First 16: sigmoid, silu derivative, gradients + __m512 exp0 = avx512_exp_ps(_mm512_sub_ps(_mm512_setzero_ps(), g0)); + __m512 sig0 = _mm512_div_ps(one, _mm512_add_ps(one, exp0)); + __m512 silu0 = _mm512_mul_ps(g0, sig0); + __m512 dsilu0 = _mm512_mul_ps(sig0, _mm512_fmadd_ps(g0, _mm512_sub_ps(one, sig0), one)); + __m512 gg0 = _mm512_mul_ps(_mm512_mul_ps(gi0, u0), dsilu0); + __m512 gu0 = _mm512_mul_ps(gi0, silu0); - // Second 16: same computation - __m512 exp1 = avx512_exp_ps(_mm512_sub_ps(_mm512_setzero_ps(), g1)); - __m512 sig1 = _mm512_div_ps(one, _mm512_add_ps(one, exp1)); - __m512 silu1 = _mm512_mul_ps(g1, sig1); - __m512 dsilu1 = _mm512_mul_ps(sig1, _mm512_fmadd_ps(g1, _mm512_sub_ps(one, sig1), one)); - __m512 gg1 = _mm512_mul_ps(_mm512_mul_ps(gi1, u1), dsilu1); - __m512 gu1 = _mm512_mul_ps(gi1, silu1); + // Second 16: same computation + __m512 exp1 = avx512_exp_ps(_mm512_sub_ps(_mm512_setzero_ps(), g1)); + __m512 sig1 = _mm512_div_ps(one, _mm512_add_ps(one, exp1)); + __m512 silu1 = _mm512_mul_ps(g1, sig1); + __m512 dsilu1 = _mm512_mul_ps(sig1, _mm512_fmadd_ps(g1, _mm512_sub_ps(one, sig1), one)); + __m512 gg1 = _mm512_mul_ps(_mm512_mul_ps(gi1, u1), dsilu1); + __m512 gu1 = _mm512_mul_ps(gi1, silu1); - avx512_32xfp32_to_32xbf16(&gg0, &gg1, (__m512i*)(grad_gate + i)); - avx512_32xfp32_to_32xbf16(&gu0, &gu1, (__m512i*)(grad_up + i)); - } - - // Scalar tail - for (; i < total; i++) { - float g_val = GGML_BF16_TO_FP32(gate_output[i]); - float u_val = GGML_BF16_TO_FP32(up_output[i]); - float sigmoid_val = 1.0f / (1.0f + expf(-g_val)); - float silu_val = g_val * sigmoid_val; - float grad_i_val = GGML_BF16_TO_FP32(grad_inter[i]); - grad_gate[i] = GGML_FP32_TO_BF16(grad_i_val * u_val * sigmoid_val * (1.0f + g_val * (1.0f - sigmoid_val))); - grad_up[i] = GGML_FP32_TO_BF16(grad_i_val * silu_val); - } - }); + avx512_32xfp32_to_32xbf16(&gg0, &gg1, (__m512i*)(grad_gate + i)); + avx512_32xfp32_to_32xbf16(&gu0, &gu1, (__m512i*)(grad_up + i)); + } + // Scalar tail + for (; i < total; i++) { + float g_val = GGML_BF16_TO_FP32(gate_output[i]); + float u_val = GGML_BF16_TO_FP32(up_output[i]); + float sigmoid_val = 1.0f / (1.0f + expf(-g_val)); + float silu_val = g_val * sigmoid_val; + float grad_i_val = GGML_BF16_TO_FP32(grad_inter[i]); + grad_gate[i] = GGML_FP32_TO_BF16(grad_i_val * u_val * sigmoid_val * (1.0f + g_val * (1.0f - sigmoid_val))); + grad_up[i] = GGML_FP32_TO_BF16(grad_i_val * silu_val); + } + }); } /** @@ -5040,55 +5019,50 @@ class AMX_SFT_MOE_TP : public BaseMOE { const int hidden = config_.hidden_size; const int hidden_vec_end = hidden & ~31; const __m512 scale_vec = _mm512_set1_ps(scale); - direct_or_pool( - qlen, - [&, scale, hidden, hidden_vec_end, scale_vec](int token_id) { - ggml_bf16_t* dst = grad_input_bf16 + token_id * hidden; - for (int j = 0; j < k; j++) { - int expert_idx = cache.expert_ids_cache[token_id * k + j]; - if (expert_idx < config_.num_gpu_experts || expert_idx >= config_.expert_num) { - continue; - } - if (m_local_num_[expert_idx] == 0) { - continue; - } - int pos = cache.m_local_pos_cache[token_id][j]; - ggml_bf16_t* contrib = grad_output_bf16_ptr_[expert_idx] + pos * config_.hidden_size; + direct_or_pool(qlen, [&, scale, hidden, hidden_vec_end, scale_vec](int token_id) { + ggml_bf16_t* dst = grad_input_bf16 + token_id * hidden; + for (int j = 0; j < k; j++) { + int expert_idx = cache.expert_ids_cache[token_id * k + j]; + if (expert_idx < config_.num_gpu_experts || expert_idx >= config_.expert_num) { + continue; + } + if (m_local_num_[expert_idx] == 0) { + continue; + } + int pos = cache.m_local_pos_cache[token_id][j]; + ggml_bf16_t* contrib = grad_output_bf16_ptr_[expert_idx] + pos * config_.hidden_size; - int h = 0; - for (; h < hidden_vec_end; h += 32) { - __m512 x0, x1, cur0, cur1; - avx512_32xbf16_to_32xfp32((__m512i*)(contrib + h), &x0, &x1); - avx512_32xbf16_to_32xfp32((__m512i*)(dst + h), &cur0, &cur1); - x0 = _mm512_fmadd_ps(x0, scale_vec, cur0); - x1 = _mm512_fmadd_ps(x1, scale_vec, cur1); - avx512_32xfp32_to_32xbf16(&x0, &x1, (__m512i*)(dst + h)); - } - for (; h < hidden; h++) { - float add = GGML_BF16_TO_FP32(contrib[h]) * scale; - float cur = GGML_BF16_TO_FP32(dst[h]); - cur += add; - dst[h] = GGML_FP32_TO_BF16(cur); - } - } - }); + int h = 0; + for (; h < hidden_vec_end; h += 32) { + __m512 x0, x1, cur0, cur1; + avx512_32xbf16_to_32xfp32((__m512i*)(contrib + h), &x0, &x1); + avx512_32xbf16_to_32xfp32((__m512i*)(dst + h), &cur0, &cur1); + x0 = _mm512_fmadd_ps(x0, scale_vec, cur0); + x1 = _mm512_fmadd_ps(x1, scale_vec, cur1); + avx512_32xfp32_to_32xbf16(&x0, &x1, (__m512i*)(dst + h)); + } + for (; h < hidden; h++) { + float add = GGML_BF16_TO_FP32(contrib[h]) * scale; + float cur = GGML_BF16_TO_FP32(dst[h]); + cur += add; + dst[h] = GGML_FP32_TO_BF16(cur); + } + } + }); }; auto base_pass = [&](bool do_up) { - // Quantize grad to BufferA - direct_or_pool( - activated_expert, - [&, do_up](int task_id) { - int expert_idx = m_expert_id_map_[task_id]; - int m = m_local_num_[expert_idx]; - if (m == 0) return; + direct_or_pool(activated_expert, [&, do_up](int task_id) { + int expert_idx = m_expert_id_map_[task_id]; + int m = m_local_num_[expert_idx]; + if (m == 0) return; - size_t offset = expert_offsets[task_id]; - ggml_bf16_t* grad = do_up ? (grad_up_output_ + offset * config_.intermediate_size) - : (grad_gate_output_ + offset * config_.intermediate_size); - down_ba_[expert_idx]->from_mat(m, grad, 0, 1); - }); + size_t offset = expert_offsets[task_id]; + ggml_bf16_t* grad = do_up ? (grad_up_output_ + offset * config_.intermediate_size) + : (grad_gate_output_ + offset * config_.intermediate_size); + down_ba_[expert_idx]->from_mat(m, grad, 0, 1); + }); int nth = T::recommended_nth(config_.hidden_size); pool->do_work_stealing_job( @@ -5211,103 +5185,101 @@ class AMX_SFT_MOE_TP : public BaseMOE { } if (!fused_tasks.empty()) { - direct_or_pool( - static_cast(fused_tasks.size()), - [&, hidden, inter_size, rank, gradb_elems](int task_id) { - const GuLoraFusedTask& task = fused_tasks[task_id]; - GuLoraFusedBuf& buf = fused_bufs[task.expert_task]; - if (buf.num_tokens == 0) return; + direct_or_pool(static_cast(fused_tasks.size()), [&, hidden, inter_size, rank, gradb_elems](int task_id) { + const GuLoraFusedTask& task = fused_tasks[task_id]; + GuLoraFusedBuf& buf = fused_bufs[task.expert_task]; + if (buf.num_tokens == 0) return; - int local_tokens = task.t_end - task.t_start; - size_t u_elems = static_cast(local_tokens) * rank; - size_t scratch_elems = u_elems * 2 + gradb_elems * 2; - float* scratch = get_lora_fp32_buffer(scratch_elems); - float* gate_u = scratch; - float* up_u = gate_u + u_elems; - float* gate_gradb_local = up_u + u_elems; - float* up_gradb_local = gate_gradb_local + gradb_elems; - memset(gate_gradb_local, 0, gradb_elems * 2 * sizeof(float)); + int local_tokens = task.t_end - task.t_start; + size_t u_elems = static_cast(local_tokens) * rank; + size_t scratch_elems = u_elems * 2 + gradb_elems * 2; + float* scratch = get_lora_fp32_buffer(scratch_elems); + float* gate_u = scratch; + float* up_u = gate_u + u_elems; + float* gate_gradb_local = up_u + u_elems; + float* up_gradb_local = gate_gradb_local + gradb_elems; + memset(gate_gradb_local, 0, gradb_elems * 2 * sizeof(float)); - avx::lora_bf16_matmul_t4r4(buf.input + static_cast(task.t_start) * hidden, buf.gate_lora_a, gate_u, - local_tokens, hidden, rank); - avx::lora_bf16_matmul_t4r4(buf.input + static_cast(task.t_start) * hidden, buf.up_lora_a, up_u, - local_tokens, hidden, rank); + avx::lora_bf16_matmul_t4r4(buf.input + static_cast(task.t_start) * hidden, buf.gate_lora_a, gate_u, + local_tokens, hidden, rank); + avx::lora_bf16_matmul_t4r4(buf.input + static_cast(task.t_start) * hidden, buf.up_lora_a, up_u, + local_tokens, hidden, rank); - for (int t = 0; t < local_tokens; t++) { - ggml_bf16_t* gate_row = buf.gate_inter + static_cast(task.t_start + t) * padded_lora_rank_; - ggml_bf16_t* up_row = buf.up_inter + static_cast(task.t_start + t) * padded_lora_rank_; - memset(gate_row, 0, padded_lora_rank_ * sizeof(ggml_bf16_t)); - memset(up_row, 0, padded_lora_rank_ * sizeof(ggml_bf16_t)); + for (int t = 0; t < local_tokens; t++) { + ggml_bf16_t* gate_row = buf.gate_inter + static_cast(task.t_start + t) * padded_lora_rank_; + ggml_bf16_t* up_row = buf.up_inter + static_cast(task.t_start + t) * padded_lora_rank_; + memset(gate_row, 0, padded_lora_rank_ * sizeof(ggml_bf16_t)); + memset(up_row, 0, padded_lora_rank_ * sizeof(ggml_bf16_t)); - const float* gate_u_row = gate_u + static_cast(t) * rank; - const float* up_u_row = up_u + static_cast(t) * rank; - for (int r = 0; r < rank; r++) { - gate_row[r] = GGML_FP32_TO_BF16(gate_u_row[r]); - up_row[r] = GGML_FP32_TO_BF16(up_u_row[r]); + const float* gate_u_row = gate_u + static_cast(t) * rank; + const float* up_u_row = up_u + static_cast(t) * rank; + for (int r = 0; r < rank; r++) { + gate_row[r] = GGML_FP32_TO_BF16(gate_u_row[r]); + up_row[r] = GGML_FP32_TO_BF16(up_u_row[r]); + } + + const ggml_bf16_t* gate_grad_row = buf.gate_grad + static_cast(task.t_start + t) * inter_size; + const ggml_bf16_t* up_grad_row = buf.up_grad + static_cast(task.t_start + t) * inter_size; + + if (rank == 8) { + __m256 gate_u_vec = _mm256_loadu_ps(gate_u_row); + __m256 up_u_vec = _mm256_loadu_ps(up_u_row); + for (int i = 0; i < inter_size; i++) { + float gg = GGML_BF16_TO_FP32(gate_grad_row[i]); + if (gg != 0.0f) { + float* out = gate_gradb_local + static_cast(i) * rank; + __m256 acc = _mm256_loadu_ps(out); + acc = _mm256_fmadd_ps(_mm256_set1_ps(gg), gate_u_vec, acc); + _mm256_storeu_ps(out, acc); } - - const ggml_bf16_t* gate_grad_row = buf.gate_grad + static_cast(task.t_start + t) * inter_size; - const ggml_bf16_t* up_grad_row = buf.up_grad + static_cast(task.t_start + t) * inter_size; - - if (rank == 8) { - __m256 gate_u_vec = _mm256_loadu_ps(gate_u_row); - __m256 up_u_vec = _mm256_loadu_ps(up_u_row); - for (int i = 0; i < inter_size; i++) { - float gg = GGML_BF16_TO_FP32(gate_grad_row[i]); - if (gg != 0.0f) { - float* out = gate_gradb_local + static_cast(i) * rank; - __m256 acc = _mm256_loadu_ps(out); - acc = _mm256_fmadd_ps(_mm256_set1_ps(gg), gate_u_vec, acc); - _mm256_storeu_ps(out, acc); - } - float ug = GGML_BF16_TO_FP32(up_grad_row[i]); - if (ug != 0.0f) { - float* out = up_gradb_local + static_cast(i) * rank; - __m256 acc = _mm256_loadu_ps(out); - acc = _mm256_fmadd_ps(_mm256_set1_ps(ug), up_u_vec, acc); - _mm256_storeu_ps(out, acc); - } + float ug = GGML_BF16_TO_FP32(up_grad_row[i]); + if (ug != 0.0f) { + float* out = up_gradb_local + static_cast(i) * rank; + __m256 acc = _mm256_loadu_ps(out); + acc = _mm256_fmadd_ps(_mm256_set1_ps(ug), up_u_vec, acc); + _mm256_storeu_ps(out, acc); + } + } + } else { + for (int i = 0; i < inter_size; i++) { + float gg = GGML_BF16_TO_FP32(gate_grad_row[i]); + if (gg != 0.0f) { + float* out = gate_gradb_local + static_cast(i) * rank; + for (int r = 0; r < rank; r++) { + out[r] += gg * gate_u_row[r]; } - } else { - for (int i = 0; i < inter_size; i++) { - float gg = GGML_BF16_TO_FP32(gate_grad_row[i]); - if (gg != 0.0f) { - float* out = gate_gradb_local + static_cast(i) * rank; - for (int r = 0; r < rank; r++) { - out[r] += gg * gate_u_row[r]; - } - } - float ug = GGML_BF16_TO_FP32(up_grad_row[i]); - if (ug != 0.0f) { - float* out = up_gradb_local + static_cast(i) * rank; - for (int r = 0; r < rank; r++) { - out[r] += ug * up_u_row[r]; - } - } + } + float ug = GGML_BF16_TO_FP32(up_grad_row[i]); + if (ug != 0.0f) { + float* out = up_gradb_local + static_cast(i) * rank; + for (int r = 0; r < rank; r++) { + out[r] += ug * up_u_row[r]; } } } + } + } - std::lock_guard lock(gradb_mutexes[task.expert_task]); - float* gate_gradb_global = gate_gradb_all.data() + static_cast(task.expert_task) * gradb_elems; - float* up_gradb_global = up_gradb_all.data() + static_cast(task.expert_task) * gradb_elems; - if (rank == 8) { - for (size_t off = 0; off < gradb_elems; off += rank) { - __m256 gate_acc = _mm256_loadu_ps(gate_gradb_global + off); - gate_acc = _mm256_add_ps(gate_acc, _mm256_loadu_ps(gate_gradb_local + off)); - _mm256_storeu_ps(gate_gradb_global + off, gate_acc); + std::lock_guard lock(gradb_mutexes[task.expert_task]); + float* gate_gradb_global = gate_gradb_all.data() + static_cast(task.expert_task) * gradb_elems; + float* up_gradb_global = up_gradb_all.data() + static_cast(task.expert_task) * gradb_elems; + if (rank == 8) { + for (size_t off = 0; off < gradb_elems; off += rank) { + __m256 gate_acc = _mm256_loadu_ps(gate_gradb_global + off); + gate_acc = _mm256_add_ps(gate_acc, _mm256_loadu_ps(gate_gradb_local + off)); + _mm256_storeu_ps(gate_gradb_global + off, gate_acc); - __m256 up_acc = _mm256_loadu_ps(up_gradb_global + off); - up_acc = _mm256_add_ps(up_acc, _mm256_loadu_ps(up_gradb_local + off)); - _mm256_storeu_ps(up_gradb_global + off, up_acc); - } - } else { - for (size_t off = 0; off < gradb_elems; off++) { - gate_gradb_global[off] += gate_gradb_local[off]; - up_gradb_global[off] += up_gradb_local[off]; - } - } - }); + __m256 up_acc = _mm256_loadu_ps(up_gradb_global + off); + up_acc = _mm256_add_ps(up_acc, _mm256_loadu_ps(up_gradb_local + off)); + _mm256_storeu_ps(up_gradb_global + off, up_acc); + } + } else { + for (size_t off = 0; off < gradb_elems; off++) { + gate_gradb_global[off] += gate_gradb_local[off]; + up_gradb_global[off] += up_gradb_local[off]; + } + } + }); constexpr int kGuGradBBlock = 256; int gradb_blocks = (inter_size + kGuGradBBlock - 1) / kGuGradBBlock; @@ -5351,7 +5323,6 @@ class AMX_SFT_MOE_TP : public BaseMOE { // Gate and up still run sequentially because they share grad_output_bf16_ptr_. // ===================================================== auto lora_pass_remainder = [&](bool do_up) { - struct GuLoraGradInTask { int expert_task = -1; int t_start = 0; @@ -5369,44 +5340,42 @@ class AMX_SFT_MOE_TP : public BaseMOE { } if (!gradin_tasks.empty()) { - direct_or_pool( - static_cast(gradin_tasks.size()), - [&, do_up](int task_id) { - const GuLoraGradInTask& task = gradin_tasks[task_id]; - int expert_task = task.expert_task; - int expert_idx = m_expert_id_map_[expert_task]; - int local_tokens = task.t_end - task.t_start; - if (local_tokens <= 0) return; + direct_or_pool(static_cast(gradin_tasks.size()), [&, do_up](int task_id) { + const GuLoraGradInTask& task = gradin_tasks[task_id]; + int expert_task = task.expert_task; + int expert_idx = m_expert_id_map_[expert_task]; + int local_tokens = task.t_end - task.t_start; + if (local_tokens <= 0) return; - const int hidden = config_.hidden_size; - const int inter_size = config_.intermediate_size; - const size_t offset = expert_offsets[expert_task] + task.t_start; - ggml_bf16_t* grad = - do_up ? (grad_up_output_ + offset * inter_size) : (grad_gate_output_ + offset * inter_size); - ggml_bf16_t* inter_ptr_base = - do_up ? lora_up_intermediate_ptr_[expert_idx] : lora_gate_intermediate_ptr_[expert_idx]; - ggml_bf16_t* inter_ptr = inter_ptr_base + static_cast(task.t_start) * padded_lora_rank_; - ggml_bf16_t* grad_out = grad_output_bf16_ptr_[expert_idx] + static_cast(task.t_start) * hidden; - const ggml_bf16_t* lora_b_t = (do_up ? up_lora_b_transposed_ : gate_lora_b_transposed_) + - static_cast(expert_idx) * lora_rank_ * inter_size; - const ggml_bf16_t* lora_a = - (do_up ? up_lora_a_ : gate_lora_a_) + static_cast(expert_idx) * lora_rank_ * hidden; + const int hidden = config_.hidden_size; + const int inter_size = config_.intermediate_size; + const size_t offset = expert_offsets[expert_task] + task.t_start; + ggml_bf16_t* grad = + do_up ? (grad_up_output_ + offset * inter_size) : (grad_gate_output_ + offset * inter_size); + ggml_bf16_t* inter_ptr_base = + do_up ? lora_up_intermediate_ptr_[expert_idx] : lora_gate_intermediate_ptr_[expert_idx]; + ggml_bf16_t* inter_ptr = inter_ptr_base + static_cast(task.t_start) * padded_lora_rank_; + ggml_bf16_t* grad_out = grad_output_bf16_ptr_[expert_idx] + static_cast(task.t_start) * hidden; + const ggml_bf16_t* lora_b_t = (do_up ? up_lora_b_transposed_ : gate_lora_b_transposed_) + + static_cast(expert_idx) * lora_rank_ * inter_size; + const ggml_bf16_t* lora_a = + (do_up ? up_lora_a_ : gate_lora_a_) + static_cast(expert_idx) * lora_rank_ * hidden; - float* gb = get_lora_fp32_buffer(static_cast(local_tokens) * lora_rank_); - avx::lora_backward_matmul_transposed(grad, lora_b_t, gb, local_tokens, inter_size, lora_rank_); + float* gb = get_lora_fp32_buffer(static_cast(local_tokens) * lora_rank_); + avx::lora_backward_matmul_transposed(grad, lora_b_t, gb, local_tokens, inter_size, lora_rank_); - memset(inter_ptr, 0, static_cast(local_tokens) * padded_lora_rank_ * sizeof(ggml_bf16_t)); - for (int t = 0; t < local_tokens; t++) { - ggml_bf16_t* inter_row = inter_ptr + static_cast(t) * padded_lora_rank_; - const float* gb_row = gb + static_cast(t) * lora_rank_; - for (int r = 0; r < lora_rank_; r++) { - inter_row[r] = GGML_FP32_TO_BF16(gb_row[r]); - } - } + memset(inter_ptr, 0, static_cast(local_tokens) * padded_lora_rank_ * sizeof(ggml_bf16_t)); + for (int t = 0; t < local_tokens; t++) { + ggml_bf16_t* inter_row = inter_ptr + static_cast(t) * padded_lora_rank_; + const float* gb_row = gb + static_cast(t) * lora_rank_; + for (int r = 0; r < lora_rank_; r++) { + inter_row[r] = GGML_FP32_TO_BF16(gb_row[r]); + } + } - memset(grad_out, 0, static_cast(local_tokens) * hidden * sizeof(ggml_bf16_t)); - avx::lora_fp32_bf16_fused_add_transposed(gb, lora_a, grad_out, local_tokens, lora_rank_, hidden, 1.0f); - }); + memset(grad_out, 0, static_cast(local_tokens) * hidden * sizeof(ggml_bf16_t)); + avx::lora_fp32_bf16_fused_add_transposed(gb, lora_a, grad_out, local_tokens, lora_rank_, hidden, 1.0f); + }); } scatter_to_grad_input(lora_scaling_); @@ -5514,7 +5483,6 @@ class AMX_SFT_MOE_TP : public BaseMOE { lora_pass_remainder(false); // gate: gb_gradin_fused, scatter, gradA lora_pass_remainder(true); // up: gb_gradin_fused, scatter, gradA - } }; diff --git a/kt-kernel/operators/avx2/moe_base.hpp b/kt-kernel/operators/avx2/moe_base.hpp index 088fcf7d..dac66ab0 100644 --- a/kt-kernel/operators/avx2/moe_base.hpp +++ b/kt-kernel/operators/avx2/moe_base.hpp @@ -119,20 +119,20 @@ class AVX2_MOE_BASE { down_ba_.push_back(make_buffer_a(config_.max_len, config_.intermediate_size, nullptr)); down_bc_.push_back(make_buffer_c(config_.max_len, config_.hidden_size, nullptr)); - void* gate_bb_ptr = - std::aligned_alloc(64, (buffer_b_required_size(config_.intermediate_size, config_.hidden_size) + 63) & ~63ULL); + void* gate_bb_ptr = std::aligned_alloc( + 64, (buffer_b_required_size(config_.intermediate_size, config_.hidden_size) + 63) & ~63ULL); if (!gate_bb_ptr) throw std::runtime_error("aligned_alloc failed for gate BufferB"); owned_aligned_allocs_.push_back(gate_bb_ptr); gate_bb_.push_back(make_buffer_b(config_.intermediate_size, config_.hidden_size, gate_bb_ptr)); - void* up_bb_ptr = - std::aligned_alloc(64, (buffer_b_required_size(config_.intermediate_size, config_.hidden_size) + 63) & ~63ULL); + void* up_bb_ptr = std::aligned_alloc( + 64, (buffer_b_required_size(config_.intermediate_size, config_.hidden_size) + 63) & ~63ULL); if (!up_bb_ptr) throw std::runtime_error("aligned_alloc failed for up BufferB"); owned_aligned_allocs_.push_back(up_bb_ptr); up_bb_.push_back(make_buffer_b(config_.intermediate_size, config_.hidden_size, up_bb_ptr)); - void* down_bb_ptr = - std::aligned_alloc(64, (buffer_b_required_size(config_.hidden_size, config_.intermediate_size) + 63) & ~63ULL); + void* down_bb_ptr = std::aligned_alloc( + 64, (buffer_b_required_size(config_.hidden_size, config_.intermediate_size) + 63) & ~63ULL); if (!down_bb_ptr) throw std::runtime_error("aligned_alloc failed for down BufferB"); owned_aligned_allocs_.push_back(down_bb_ptr); down_bb_.push_back(make_buffer_b(config_.hidden_size, config_.intermediate_size, down_bb_ptr)); @@ -234,23 +234,28 @@ class AVX2_MOE_BASE { size_t max_m = (m_local_num_[i] + M_STEP - 1) / M_STEP * M_STEP; gate_up_ba_[i]->max_m = max_m; gate_up_ba_[i]->set_data(gate_up_ba_pool_ptr); - gate_up_ba_pool_ptr = (void*)((uintptr_t)gate_up_ba_pool_ptr + align64(buffer_a_required_size(max_m, config_.hidden_size))); + gate_up_ba_pool_ptr = + (void*)((uintptr_t)gate_up_ba_pool_ptr + align64(buffer_a_required_size(max_m, config_.hidden_size))); gate_bc_[i]->max_m = max_m; gate_bc_[i]->set_data(gate_bc_pool_ptr); - gate_bc_pool_ptr = (void*)((uintptr_t)gate_bc_pool_ptr + align64(buffer_c_required_size(max_m, config_.intermediate_size))); + gate_bc_pool_ptr = + (void*)((uintptr_t)gate_bc_pool_ptr + align64(buffer_c_required_size(max_m, config_.intermediate_size))); up_bc_[i]->max_m = max_m; up_bc_[i]->set_data(up_bc_pool_ptr); - up_bc_pool_ptr = (void*)((uintptr_t)up_bc_pool_ptr + align64(buffer_c_required_size(max_m, config_.intermediate_size))); + up_bc_pool_ptr = + (void*)((uintptr_t)up_bc_pool_ptr + align64(buffer_c_required_size(max_m, config_.intermediate_size))); down_ba_[i]->max_m = max_m; down_ba_[i]->set_data(down_ba_pool_ptr); - down_ba_pool_ptr = (void*)((uintptr_t)down_ba_pool_ptr + align64(buffer_a_required_size(max_m, config_.intermediate_size))); + down_ba_pool_ptr = + (void*)((uintptr_t)down_ba_pool_ptr + align64(buffer_a_required_size(max_m, config_.intermediate_size))); down_bc_[i]->max_m = max_m; down_bc_[i]->set_data(down_bc_pool_ptr); - down_bc_pool_ptr = (void*)((uintptr_t)down_bc_pool_ptr + align64(buffer_c_required_size(max_m, config_.hidden_size))); + down_bc_pool_ptr = + (void*)((uintptr_t)down_bc_pool_ptr + align64(buffer_c_required_size(max_m, config_.hidden_size))); } auto direct_or_pool = [&](int count, auto&& fn) { @@ -330,9 +335,8 @@ class AVX2_MOE_BASE { __m256 weight = _mm256_set1_ps(weights[i * k + j]); __m256 d0, d1; avx2::load_16xbf16_to_2x8xfp32( - m_local_down_output_ptr_[expert_ids[i * k + j]] + - m_local_pos_[i][j] * config_.hidden_size + e, - &d0, &d1); + m_local_down_output_ptr_[expert_ids[i * k + j]] + m_local_pos_[i][j] * config_.hidden_size + e, &d0, + &d1); x0 = _mm256_fmadd_ps(d0, weight, x0); x1 = _mm256_fmadd_ps(d1, weight, x1); } @@ -381,19 +385,23 @@ class AVX2_MOE_BASE { gate_bc_[expert_idx]->max_m = max_m; gate_bc_[expert_idx]->set_data(gate_bc_pool_ptr); - gate_bc_pool_ptr = (void*)((uintptr_t)gate_bc_pool_ptr + align64(buffer_c_required_size(max_m, config_.intermediate_size))); + gate_bc_pool_ptr = + (void*)((uintptr_t)gate_bc_pool_ptr + align64(buffer_c_required_size(max_m, config_.intermediate_size))); up_bc_[expert_idx]->max_m = max_m; up_bc_[expert_idx]->set_data(up_bc_pool_ptr); - up_bc_pool_ptr = (void*)((uintptr_t)up_bc_pool_ptr + align64(buffer_c_required_size(max_m, config_.intermediate_size))); + up_bc_pool_ptr = + (void*)((uintptr_t)up_bc_pool_ptr + align64(buffer_c_required_size(max_m, config_.intermediate_size))); down_ba_[expert_idx]->max_m = max_m; down_ba_[expert_idx]->set_data(down_ba_pool_ptr); - down_ba_pool_ptr = (void*)((uintptr_t)down_ba_pool_ptr + align64(buffer_a_required_size(max_m, config_.intermediate_size))); + down_ba_pool_ptr = + (void*)((uintptr_t)down_ba_pool_ptr + align64(buffer_a_required_size(max_m, config_.intermediate_size))); down_bc_[expert_idx]->max_m = max_m; down_bc_[expert_idx]->set_data(down_bc_pool_ptr); - down_bc_pool_ptr = (void*)((uintptr_t)down_bc_pool_ptr + align64(buffer_c_required_size(max_m, config_.hidden_size))); + down_bc_pool_ptr = + (void*)((uintptr_t)down_bc_pool_ptr + align64(buffer_c_required_size(max_m, config_.hidden_size))); } // Pack input into BufferA for each activated expert @@ -403,7 +411,8 @@ class AVX2_MOE_BASE { size_t max_m = (qlen + M_STEP - 1) / M_STEP * M_STEP; gate_up_ba_[expert_idx]->max_m = max_m; gate_up_ba_[expert_idx]->set_data(gate_up_ba_pool_ptr); - gate_up_ba_pool_ptr = (void*)((uintptr_t)gate_up_ba_pool_ptr + align64(buffer_a_required_size(max_m, config_.hidden_size))); + gate_up_ba_pool_ptr = + (void*)((uintptr_t)gate_up_ba_pool_ptr + align64(buffer_a_required_size(max_m, config_.hidden_size))); gate_up_ba_[expert_idx]->from_mat(qlen, (ggml_bf16_t*)input, 0, 1); } @@ -458,8 +467,7 @@ class AVX2_MOE_BASE { __m256 weight = _mm256_set1_ps(weights[j]); __m256 d0, d1; avx2::load_16xbf16_to_2x8xfp32( - m_local_down_output_ptr_[expert_ids[j]] + m_local_pos_[0][j] * config_.hidden_size + e, - &d0, &d1); + m_local_down_output_ptr_[expert_ids[j]] + m_local_pos_[0][j] * config_.hidden_size + e, &d0, &d1); x0 = _mm256_fmadd_ps(d0, weight, x0); x1 = _mm256_fmadd_ps(d1, weight, x1); } diff --git a/kt-kernel/operators/avx2/mxfp4-moe.hpp b/kt-kernel/operators/avx2/mxfp4-moe.hpp index 394301ca..e3289081 100644 --- a/kt-kernel/operators/avx2/mxfp4-moe.hpp +++ b/kt-kernel/operators/avx2/mxfp4-moe.hpp @@ -444,7 +444,8 @@ class AVX2_MXFP4_MOE_TP : public AVX2_MOE_BASE> { // H2: Bounds check (already validated in weight loading, but be safe) if (lid >= config_.gate_scales[0].size()) return; if (config_.gate_scales[0][lid] == nullptr || config_.up_scales[0][lid] == nullptr || - config_.down_scales[0][lid] == nullptr) return; + config_.down_scales[0][lid] == nullptr) + return; size_t scale_elem_count = ((size_t)config_.hidden_size * config_.intermediate_size) / group_size; // tp_part_idx == 0 guaranteed here, so offset is 0 convert_or_copy(gate_bb_[expert_idx]->d, (const ggml_bf16_t*)config_.gate_scales[0][lid], scale_elem_count); @@ -678,7 +679,8 @@ class TP_MOE> : public TP_MOEget_subpool(i); subpool->do_work_stealing_job( tpc.expert_num, nullptr, - [&, i, per_tp_interm, full_interm, gate_buf, up_buf, down_buf, gate_up_wt_per_expert, down_wt_per_expert](int eid) { + [&, i, per_tp_interm, full_interm, gate_buf, up_buf, down_buf, gate_up_wt_per_expert, + down_wt_per_expert](int eid) { if (tpc.should_skip_expert(eid)) return; uint64_t lid = expert_map(physical_to_logical_map, eid); @@ -725,7 +727,8 @@ class TP_MOE> : public TP_MOE> : public TP_MOEdispense_backend()->do_numa_job([&, this](int i) { auto& tpc = tps[i]->config_; if (tpc.intermediate_size % 2 != 0) - throw std::runtime_error("MXFP4 TP flat-buffer: intermediate_size must be even for nibble-aligned addressing, got " + - std::to_string(tpc.intermediate_size)); + throw std::runtime_error( + "MXFP4 TP flat-buffer: intermediate_size must be even for nibble-aligned addressing, got " + + std::to_string(tpc.intermediate_size)); size_t weight_elem_count = (size_t)tpc.intermediate_size * tpc.hidden_size; size_t scales_elem_count = ((size_t)tpc.hidden_size / group_size) * tpc.intermediate_size; tpc.gate_proj = new uint8_t[(tpc.expert_num * weight_elem_count) / 2]; @@ -819,8 +823,7 @@ class TP_MOE> : public TP_MOEweights_loaded = true; } - void write_weight_scale_to_buffer(int gpu_tp_count, int expert_id, - const std::vector& w13_weight_ptrs, + void write_weight_scale_to_buffer(int gpu_tp_count, int expert_id, const std::vector& w13_weight_ptrs, const std::vector& w13_scale_ptrs, const std::vector& w2_weight_ptrs, const std::vector& w2_scale_ptrs) { diff --git a/kt-kernel/operators/moe-sft-tp.hpp b/kt-kernel/operators/moe-sft-tp.hpp index afd84599..4e8da6ad 100644 --- a/kt-kernel/operators/moe-sft-tp.hpp +++ b/kt-kernel/operators/moe-sft-tp.hpp @@ -8,7 +8,6 @@ #ifndef CPUINFER_OPERATOR_MOE_SFT_TP_HPP #define CPUINFER_OPERATOR_MOE_SFT_TP_HPP - #include #include @@ -266,6 +265,9 @@ class TP_MOE_SFT : public TP_MOE { if (!config.gate_projs.empty()) { // Pre-quantized per-NUMA weights (INT8/INT4 with separate scales) printf("TP_MOE_SFT: Pre-quantized per-NUMA mode (gate_projs path)\n"); + for (int i = 0; i < tp_count; i++) { + tps[i]->set_physical_to_logical_map(config.physical_to_logical_map); + } pool->dispense_backend()->do_numa_job([this](int numa_id) { tps[numa_id]->load_weights(); }); // Check if pre-quantized backward weights are available @@ -391,6 +393,7 @@ class TP_MOE_SFT : public TP_MOE { // Step 2: Set weight pointers BEFORE load_weights (Bug #24 fix) for (int i = 0; i < tp_count; i++) { + tps[i]->set_physical_to_logical_map(config.physical_to_logical_map); tps[i]->set_weight_pointers_for_forward(temp_gate[i], temp_up[i], temp_down[i]); } @@ -401,7 +404,6 @@ class TP_MOE_SFT : public TP_MOE { if (!config.share_backward_bb) { tps[i]->prepare_bwd(temp_gate[i], temp_up[i], temp_down[i]); } - tps[i]->set_physical_to_logical_map(config.physical_to_logical_map); } for (int i = 0; i < tp_count; i++) { @@ -411,6 +413,9 @@ class TP_MOE_SFT : public TP_MOE { } } else { // Other loading methods (from loader or file) + for (int i = 0; i < tp_count; i++) { + tps[i]->set_physical_to_logical_map(config.physical_to_logical_map); + } pool->dispense_backend()->do_numa_job([this](int numa_id) { tps[numa_id]->load_weights(); }); // Try loading backward weights from disk (.kt files) — parallel across NUMA nodes. @@ -489,7 +494,6 @@ class TP_MOE_SFT : public TP_MOE { throw std::runtime_error("Weights not loaded"); } - int qlen = *qlen_ptr; auto pool = config.pool; @@ -504,7 +508,6 @@ class TP_MOE_SFT : public TP_MOE { save_for_backward); }); - // // Collect per-thread timing from all NUMA subpools // for (int i = 0; i < tp_count; i++) { // } @@ -514,9 +517,7 @@ class TP_MOE_SFT : public TP_MOE { // Merge results from all NUMA nodes this->merge_results(qlen, output); - - pool->dispense_backend()->do_numa_job([&](int numa_id) { - }); + pool->dispense_backend()->do_numa_job([&](int numa_id) {}); } /** @@ -551,7 +552,6 @@ class TP_MOE_SFT : public TP_MOE { void* grad_weights) { auto pool = config.pool; - // Get full intermediate_size (before TP partitioning) int full_intermediate_size = sft_config.intermediate_size; int expert_num = config.expert_num; @@ -656,7 +656,6 @@ class TP_MOE_SFT : public TP_MOE { }, nullptr); - // Compute TP-slice pointers for copy-type direct writes // Each TP writes to its own I-slice of the final output tensor std::vector tp_gate_b_ptr(tp_count); @@ -882,8 +881,7 @@ class TP_MOE_SFT : public TP_MOE { nullptr); } - pool->dispense_backend()->do_numa_job([&](int numa_id) { - }); + pool->dispense_backend()->do_numa_job([&](int numa_id) {}); } /** @@ -934,43 +932,38 @@ class TP_MOE_SFT : public TP_MOE { } } - // Single do_numa_job: work-stealing memcpy + update_lora_weights - auto pool = config.pool; - pool->dispense_backend()->do_numa_job([this, gate_lora_a, gate_lora_b, up_lora_a, up_lora_b, down_lora_a, - down_lora_b, full_intermediate_size, expert_num, lora_rank, - pool](int numa_id) { + // LoRA weights are installed at load time. Keep the partitioning copy + // synchronous and serial here instead of nesting work-stealing jobs inside + // SGLang's scheduler process during model-loading barriers. + for (int numa_id = 0; numa_id < tp_count; numa_id++) { int tp_inter = tp_configs[numa_id].intermediate_size; size_t lora_b_slice = (size_t)tp_inter * lora_rank; - auto subpool = pool->get_subpool(numa_id); - // Work-stealing: copy all weights for this expert (gate + up + down) - subpool->do_work_stealing_job( - expert_num, - [&](int e) { - // gate_lora_b: [expert_num, intermediate_size, lora_rank] - memcpy(partitioned_gate_lora_b_[numa_id] + e * lora_b_slice, - (ggml_bf16_t*)gate_lora_b + e * full_intermediate_size * lora_rank + numa_id * lora_b_slice, - sizeof(ggml_bf16_t) * lora_b_slice); + for (int e = 0; e < expert_num; e++) { + // gate_lora_b: [expert_num, intermediate_size, lora_rank] + memcpy(partitioned_gate_lora_b_[numa_id] + e * lora_b_slice, + (ggml_bf16_t*)gate_lora_b + e * full_intermediate_size * lora_rank + numa_id * lora_b_slice, + sizeof(ggml_bf16_t) * lora_b_slice); - // up_lora_b: [expert_num, intermediate_size, lora_rank] - memcpy(partitioned_up_lora_b_[numa_id] + e * lora_b_slice, - (ggml_bf16_t*)up_lora_b + e * full_intermediate_size * lora_rank + numa_id * lora_b_slice, - sizeof(ggml_bf16_t) * lora_b_slice); + // up_lora_b: [expert_num, intermediate_size, lora_rank] + memcpy(partitioned_up_lora_b_[numa_id] + e * lora_b_slice, + (ggml_bf16_t*)up_lora_b + e * full_intermediate_size * lora_rank + numa_id * lora_b_slice, + sizeof(ggml_bf16_t) * lora_b_slice); - // down_lora_a: [expert_num, lora_rank, intermediate_size] - row-wise slice - for (int r = 0; r < lora_rank; r++) { - memcpy(partitioned_down_lora_a_[numa_id] + e * lora_rank * tp_inter + r * tp_inter, - (ggml_bf16_t*)down_lora_a + e * lora_rank * full_intermediate_size + r * full_intermediate_size + - numa_id * tp_inter, - sizeof(ggml_bf16_t) * tp_inter); - } - }); + // down_lora_a: [expert_num, lora_rank, intermediate_size] - row-wise slice + for (int r = 0; r < lora_rank; r++) { + memcpy(partitioned_down_lora_a_[numa_id] + e * lora_rank * tp_inter + r * tp_inter, + (ggml_bf16_t*)down_lora_a + e * lora_rank * full_intermediate_size + r * full_intermediate_size + + numa_id * tp_inter, + sizeof(ggml_bf16_t) * tp_inter); + } + } // Update weights after all memcpy complete tps[numa_id]->update_lora_weights(gate_lora_a, partitioned_gate_lora_b_[numa_id], up_lora_a, partitioned_up_lora_b_[numa_id], partitioned_down_lora_a_[numa_id], down_lora_b); - }); + } } /** diff --git a/kt-kernel/python/sft/amx.py b/kt-kernel/python/sft/amx.py index 3f3270f0..effa7dfd 100644 --- a/kt-kernel/python/sft/amx.py +++ b/kt-kernel/python/sft/amx.py @@ -40,6 +40,8 @@ except (ImportError, AttributeError): from .base import BaseSFTMoEWrapper, KExpertsSFTBuffer +_AMX_M_STEP = 32 + # Mapping from method string to C++ SFT MOE class _SFT_METHOD_TO_CLASS = { @@ -159,6 +161,17 @@ class AMXSFTMoEWrapper(BaseSFTMoEWrapper): if self._weights_loaded: return + if physical_to_logical_map_cpu is None: + physical_to_logical_map_cpu = torch.arange(self.num_experts, dtype=torch.int64) + self._physical_to_logical_map_cpu = physical_to_logical_map_cpu.to( + dtype=torch.int64, device="cpu" + ).contiguous() + if self._physical_to_logical_map_cpu.numel() < self.num_experts: + raise ValueError( + "physical_to_logical_map_cpu must contain at least " + f"{self.num_experts} entries, got {self._physical_to_logical_map_cpu.numel()}." + ) + if self.gate_proj is None and not getattr(self, "_use_projs_path", False): self._load_base_weights_from_file() @@ -170,10 +183,11 @@ class AMXSFTMoEWrapper(BaseSFTMoEWrapper): config.lora_rank = self.lora_rank config.lora_alpha = self.lora_alpha config.max_cache_depth = self.max_cache_depth - config.max_len = self.chunked_prefill_size + config.max_len = self._aligned_max_len() config.layer_idx = self.layer_idx config.share_backward_bb = getattr(self, "share_backward_bb", False) config.share_cache_pool = getattr(self, "share_cache_pool", False) + config.physical_to_logical_map = self._physical_to_logical_map_cpu.data_ptr() if getattr(self, "_use_kt_direct_load", False): config.load = True @@ -220,8 +234,9 @@ class AMXSFTMoEWrapper(BaseSFTMoEWrapper): self.cpu_infer.submit(self.moe.load_weights_task()) self.cpu_infer.sync() - self.cpu_infer.submit(self.moe.warm_up_task()) - self.cpu_infer.sync() + if os.environ.get("KT_SFT_ENABLE_WARMUP", "0") == "1": + self.cpu_infer.submit(self.moe.warm_up_task()) + self.cpu_infer.sync() # Release Python-side weight tensors (C++ copied them) self.gate_proj = None @@ -312,6 +327,7 @@ class AMXSFTMoEWrapper(BaseSFTMoEWrapper): self._gate_scales_per_numa = experts_data["gate_scale"] self._up_scales_per_numa = experts_data["up_scale"] self._down_scales_per_numa = experts_data["down_scale"] + self._validate_prepartitioned_weights() self._gate_projs_ptrs = _make_ptrs(gate_weights) self._up_projs_ptrs = _make_ptrs(up_weights) @@ -345,6 +361,57 @@ class AMXSFTMoEWrapper(BaseSFTMoEWrapper): loader.close_all_handles() + def _aligned_max_len(self) -> int: + return ((self.chunked_prefill_size + _AMX_M_STEP - 1) // _AMX_M_STEP) * _AMX_M_STEP + + def _validate_prepartitioned_weights(self) -> None: + numa_count = len(self._gate_weights_per_numa) + if self.moe_intermediate_size % self.threadpool_count != 0: + raise ValueError( + f"moe_intermediate_size={self.moe_intermediate_size} must be divisible by " + f"threadpool_count={self.threadpool_count} for {self.method} SFT." + ) + if numa_count != self.threadpool_count: + raise ValueError( + f"{self.method} SFT pre-partitioned expert weights have {numa_count} NUMA partitions, " + f"but CPUInfer was created with threadpool_count={self.threadpool_count}. " + f"Use --kt-threadpool-count {numa_count} for this weight directory, or convert weights " + "for the requested threadpool count." + ) + + collections = { + "gate": self._gate_weights_per_numa, + "up": self._up_weights_per_numa, + "down": self._down_weights_per_numa, + "gate_scale": self._gate_scales_per_numa, + "up_scale": self._up_scales_per_numa, + "down_scale": self._down_scales_per_numa, + } + for name, per_numa in collections.items(): + if len(per_numa) != numa_count: + raise ValueError(f"{name} has {len(per_numa)} NUMA partitions, expected {numa_count}.") + for numa_id, entries in enumerate(per_numa): + if len(entries) != self.num_experts: + raise ValueError( + f"{name}[numa={numa_id}] has {len(entries)} experts, expected {self.num_experts}." + ) + + for numa_id in range(numa_count): + gate_scale_len = self._gate_scales_per_numa[numa_id][0].size + up_scale_len = self._up_scales_per_numa[numa_id][0].size + down_scale_len = self._down_scales_per_numa[numa_id][0].size + expected_intermediate = self.moe_intermediate_size // self.threadpool_count + if gate_scale_len != expected_intermediate or up_scale_len != expected_intermediate: + raise ValueError( + f"{self.method} gate/up scale length for NUMA {numa_id} is " + f"{gate_scale_len}/{up_scale_len}, expected {expected_intermediate}." + ) + if down_scale_len != self.hidden_size: + raise ValueError( + f"{self.method} down scale length for NUMA {numa_id} is " + f"{down_scale_len}, expected {self.hidden_size}." + ) + # ========== LoRA ========== def init_lora_weights( @@ -373,6 +440,10 @@ class AMXSFTMoEWrapper(BaseSFTMoEWrapper): expected = expected_shapes[name] if tensor.shape != expected: raise ValueError(f"{name} shape mismatch: expected {expected}, got {tuple(tensor.shape)}") + if tensor.device.type != "cpu": + raise ValueError( + f"{name} must be a CPU tensor for {self.method} SFT, got {tensor.device}." + ) self.gate_lora_a = gate_lora_a.contiguous() self.gate_lora_b = gate_lora_b.contiguous() @@ -401,17 +472,17 @@ class AMXSFTMoEWrapper(BaseSFTMoEWrapper): if not self._lora_initialized: raise RuntimeError("LoRA weights not initialized. Call init_lora_weights() first.") - self.cpu_infer.submit( - self.moe.update_lora_weights_task( - self.gate_lora_a.data_ptr(), - self.gate_lora_b.data_ptr(), - self.up_lora_a.data_ptr(), - self.up_lora_b.data_ptr(), - self.down_lora_a.data_ptr(), - self.down_lora_b.data_ptr(), - ) + # Weight pointer updates are load-time synchronous work. Calling the + # direct binding avoids nesting an update task inside CPUInfer's queue + # while SGLang is still in distributed model-loading barriers. + self.moe.update_lora_weights( + self.gate_lora_a.data_ptr(), + self.gate_lora_b.data_ptr(), + self.up_lora_a.data_ptr(), + self.up_lora_b.data_ptr(), + self.down_lora_a.data_ptr(), + self.down_lora_b.data_ptr(), ) - self.cpu_infer.sync() def save_backward_weights_from_tensors( self, diff --git a/kt-kernel/python/sft/base.py b/kt-kernel/python/sft/base.py index 25b0e2cb..57ba8e42 100644 --- a/kt-kernel/python/sft/base.py +++ b/kt-kernel/python/sft/base.py @@ -15,7 +15,7 @@ import torch from typing import Optional, Tuple from abc import ABC, abstractmethod -from ..experts_base import _MoEBase +from ..experts_base import KExpertsCPUBuffer, _MoEBase class KExpertsSFTBuffer: @@ -98,6 +98,26 @@ class KExpertsSFTBuffer: cls._shared_buffer = None +class _SFTForwardBufferView: + """Minimal buffer view consumed by AMXSFTMoEWrapper._make_forward_task.""" + + __slots__ = ("bsz_tensor", "expert_ids_cpu", "weights_cpu", "input_cpu", "output_cpu") + + def __init__( + self, + bsz_tensor: torch.Tensor, + expert_ids_cpu: torch.Tensor, + weights_cpu: torch.Tensor, + input_cpu: torch.Tensor, + output_cpu: torch.Tensor, + ): + self.bsz_tensor = bsz_tensor + self.expert_ids_cpu = expert_ids_cpu + self.weights_cpu = weights_cpu + self.input_cpu = input_cpu + self.output_cpu = output_cpu + + class BaseSFTMoEWrapper(_MoEBase, ABC): """ Base class for SFT MoE CPU operations with concrete buffer management. @@ -357,6 +377,104 @@ class BaseSFTMoEWrapper(_MoEBase, ABC): return self._return_output(buffer, qlen, output_device) + # ========== Inference-only async forward ========== + + def submit_forward_inference( + self, + hidden_states: torch.Tensor, + expert_ids: torch.Tensor, + weights: torch.Tensor, + cuda_stream, + ) -> None: + """ + Submit an SFT MoE forward pass for serving. + + This path mirrors the normal KT inference wrapper: inputs are copied to + pinned CPU staging buffers, the CPUInfer task is enqueued with the + caller CUDA stream, and sync_forward_inference() returns a persistent + GPU output buffer. It deliberately avoids the training-oriented + torch.cuda.synchronize() in _copy_inputs_to_buffer(). + """ + if not hasattr(self.cpu_infer, "submit_with_cuda_stream"): + self.submit_forward(hidden_states, expert_ids, weights, save_for_backward=False) + self._pending_inference_fallback = True + self._pending_inference_fallback_device = hidden_states.device + return + + self._validate_forward_inputs(hidden_states, expert_ids, weights) + flat_hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) + + ( + input_tensor_cpu, + expert_ids_cpu, + _deferred_expert_ids_cpu, + weights_cpu, + output_cpu, + bsz_tensor_cpu, + output_gpu, + ) = KExpertsCPUBuffer.get_buffer(flat_hidden_states, self.num_experts_per_tok) + + current_slot = self.layer_idx % KExpertsCPUBuffer.buffer_depth + bsz_slot_tensor = bsz_tensor_cpu[current_slot] + + torch_stream = ( + cuda_stream + if isinstance(cuda_stream, torch.cuda.Stream) + else torch.cuda.ExternalStream(cuda_stream, device=flat_hidden_states.device) + ) + with torch.cuda.stream(torch_stream): + input_tensor_cpu[current_slot].copy_(flat_hidden_states.to(torch.bfloat16), non_blocking=True) + expert_ids_cpu[current_slot].copy_(expert_ids.to(torch.int64), non_blocking=True) + weights_cpu[current_slot].copy_(weights.to(torch.float32), non_blocking=True) + + buffer_view = _SFTForwardBufferView( + bsz_tensor=bsz_slot_tensor, + expert_ids_cpu=expert_ids_cpu[current_slot], + weights_cpu=weights_cpu[current_slot], + input_cpu=input_tensor_cpu[current_slot], + output_cpu=output_cpu[current_slot], + ) + + self._pending_inference_fallback = False + self._pending_inference_output_cpu = output_cpu[current_slot] + self._pending_inference_output_gpu = output_gpu[current_slot] + + self.cpu_infer.submit_with_cuda_stream( + cuda_stream, + self._make_forward_task(buffer_view, save_for_backward=False), + ) + + def sync_forward_inference(self, cuda_stream) -> torch.Tensor: + """ + Synchronize a serving forward submitted by submit_forward_inference(). + + Returns a persistent GPU buffer matching the input batch shape. Consumers + on the same CUDA stream will naturally wait for the non-blocking D2H/H2D + staging work ordered through CPUInfer's stream synchronization. + """ + if getattr(self, "_pending_inference_fallback", False): + self._pending_inference_fallback = False + output_device = getattr(self, "_pending_inference_fallback_device", None) + self._pending_inference_fallback_device = None + return self.sync_forward(output_device=output_device) + + if not hasattr(self, "_pending_inference_output_cpu"): + raise RuntimeError("No pending inference forward. Call submit_forward_inference() first.") + + torch_stream = ( + cuda_stream + if isinstance(cuda_stream, torch.cuda.Stream) + else torch.cuda.ExternalStream(cuda_stream, device=self._pending_inference_output_gpu.device) + ) + self.cpu_infer.sync_with_cuda_stream(cuda_stream) + with torch.cuda.stream(torch_stream): + self._pending_inference_output_gpu.copy_(self._pending_inference_output_cpu, non_blocking=True) + output = self._pending_inference_output_gpu + + del self._pending_inference_output_cpu + del self._pending_inference_output_gpu + return output + # ========== Async backward ========== def submit_backward_async( diff --git a/kt-kernel/python/utils/loader.py b/kt-kernel/python/utils/loader.py index 7ffe87f7..65e5efcb 100644 --- a/kt-kernel/python/utils/loader.py +++ b/kt-kernel/python/utils/loader.py @@ -726,6 +726,7 @@ class CompressedSafeTensorLoader(SafeTensorLoader): "down_scale": down_scales, } + class GGUFLoader: """ GGUF format loader using the official gguf library (gguf.gguf_reader.GGUFReader) diff --git a/kt-kernel/scripts/README.md b/kt-kernel/scripts/README.md index 03585fcf..f4a69faf 100644 --- a/kt-kernel/scripts/README.md +++ b/kt-kernel/scripts/README.md @@ -4,6 +4,86 @@ KT-Kernel provides weight conversion tools for CPU-GPU hybrid inference (e.g., i - **CPU Weights (`convert_cpu_weights.py`)**: Quantize weights to INT4/INT8 with AMX optimization for CPU-resident "cold" experts - **GPU Weights (`convert_gpu_weights.py`)**: Apply GPTQ/RTN quantization (W4A16/W8A16) for GPU-resident "hot" experts +- **KT Fused Expert LoRA (`convert_kt_to_sglang_adapter.py`)**: Convert KT SFT fused expert LoRA checkpoints into adapter-only SafeTensors directories + +--- + +## KT Fused Expert LoRA Adapter Conversion + +KT SFT fused expert LoRA saves MoE expert LoRA tensors in `fused_expert_lora.safetensors` using compact 3D tensors: + +``` +layers.{L}.experts.gate_lora_a +layers.{L}.experts.gate_lora_b +layers.{L}.experts.up_lora_a +layers.{L}.experts.up_lora_b +layers.{L}.experts.down_lora_a +layers.{L}.experts.down_lora_b +``` + +Use `convert_kt_to_sglang_adapter.py` to convert raw KT SFT output into one merged SGLang adapter directory: + +```bash +python scripts/convert_kt_to_sglang_adapter.py /path/to/kt_adapter /path/to/sglang_adapter \ + --base-model-name-or-path /path/to/base_model \ + --lora-alpha 16 \ + --overwrite +``` + +Output: + +``` +sglang_adapter/ +├── adapter_config.json +└── adapter_model.safetensors +``` + +The converter merges the existing non-expert `adapter_model.safetensors` with expanded expert tensors from `fused_expert_lora.safetensors`. Pass this merged directory to SGLang with: + +```bash +--enable-lora \ +--lora-paths my_lora=/path/to/sglang_adapter +``` + +The KTransformers SGLang fork will auto-split the merged adapter internally at server startup. Users do not need to pass separate expert and non-expert adapter paths in the normal workflow. + +Optional split outputs for debugging: + +```bash +python scripts/convert_kt_to_sglang_adapter.py /path/to/kt_adapter /path/to/sglang_adapter \ + --base-model-name-or-path /path/to/base_model \ + --expert-output-dir /path/to/expert_adapter \ + --nonexpert-output-dir /path/to/nonexpert_adapter \ + --overwrite +``` + +Existing PEFT prefixes such as `base_model.model.` are stripped to match SGLang's loader. Scaling is not folded into the LoRA B tensors. Runtime scaling remains `lora_alpha / r`; if the input directory has no `adapter_config.json`, pass `--lora-alpha` explicitly. + +This script only converts adapter files. Serving compatibility depends on the KTransformers SGLang runtime branch being used. + +### Optional Integration Validation + +The unit tests use synthetic tensors and run without model files. To validate a real KT adapter directory, set these environment variables: + +```bash +export KT_LORA_ADAPTER_DIR=/path/to/kt_adapter +export KT_LORA_BASE_MODEL=/path/to/base_model +export KT_LORA_ALPHA=16 # required only if the input has no adapter_config.json +``` + +Then run: + +```bash +python -m pytest kt-kernel/test/per_commit/test_convert_kt_to_sglang_adapter_integration.py -q +``` + +To run a large adapter conversion smoke test, also set: + +```bash +export KT_LORA_LARGE_ADAPTER_DIR=/path/to/large_kt_adapter +``` + +These integration tests check real fused tensor splitting, optional `adapter_model.safetensors` merging, `adapter_config.json` compatibility with `sglang.srt.lora.lora_config.LoRAConfig`, and large-file readability. They intentionally do not start an SGLang server or validate runtime `FusedMoE` LoRA application. --- diff --git a/kt-kernel/scripts/convert_kt_to_sglang_adapter.py b/kt-kernel/scripts/convert_kt_to_sglang_adapter.py new file mode 100644 index 00000000..54094b21 --- /dev/null +++ b/kt-kernel/scripts/convert_kt_to_sglang_adapter.py @@ -0,0 +1,477 @@ +#!/usr/bin/env python3 +"""Convert KT fused expert LoRA checkpoints into an SGLang adapter directory.""" + +from __future__ import annotations + +import argparse +import json +import os +import re +import shutil +from pathlib import Path +from typing import Dict, Iterable, Mapping + +import torch +from safetensors.torch import load_file, save_file + + +FUSED_EXPERT_LORA_FILE = "fused_expert_lora.safetensors" +ADAPTER_MODEL_FILE = "adapter_model.safetensors" +ADAPTER_CONFIG_FILE = "adapter_config.json" + +KT_NAME_MAP = { + "gate_lora_a": ("gate_proj", "lora_A", 1), + "gate_lora_b": ("gate_proj", "lora_B", 2), + "up_lora_a": ("up_proj", "lora_A", 1), + "up_lora_b": ("up_proj", "lora_B", 2), + "down_lora_a": ("down_proj", "lora_A", 1), + "down_lora_b": ("down_proj", "lora_B", 2), +} + +TARGET_MODULE_ORDER = [ + "q_proj", + "k_proj", + "v_proj", + "o_proj", + "gate_proj", + "up_proj", + "down_proj", + "in_proj_qkv", + "in_proj_z", + "in_proj_b", + "in_proj_a", + "out_proj", + "embed_tokens", + "lm_head", +] + +KT_FUSED_KEY_RE = re.compile(r"^layers\.(\d+)\.experts\.([^.]+)$") + + +def _load_json(path: Path) -> dict: + with path.open("r", encoding="utf-8") as f: + return json.load(f) + + +def _write_json(path: Path, data: Mapping) -> None: + with path.open("w", encoding="utf-8") as f: + json.dump(data, f, indent=2, sort_keys=True) + f.write("\n") + + +def _clean_adapter_key(key: str) -> str: + """Match the existing SGLang converter's PEFT key cleanup.""" + key = key.replace("base_model.model.", "") + key = key.replace(".orig_module", "") + return key + + +def _ordered_target_modules(modules: Iterable[str]) -> list[str]: + seen = set(modules) + ordered = [name for name in TARGET_MODULE_ORDER if name in seen] + ordered.extend(sorted(seen.difference(ordered))) + return ordered + + +def _infer_target_module_from_key(key: str) -> str | None: + if "lora_embedding_A" in key or "lora_embedding_B" in key: + if "embed_tokens" in key: + return "embed_tokens" + if "lm_head" in key or "unembed_tokens" in key: + return "lm_head" + + marker = ".lora_" + if marker not in key: + return None + prefix = key.split(marker, 1)[0] + if "." not in prefix: + return prefix + return prefix.rsplit(".", 1)[-1] + + +def _merge_tensor(tensors: Dict[str, torch.Tensor], key: str, value: torch.Tensor) -> None: + if key in tensors: + raise ValueError(f"Duplicate output tensor key: {key}") + tensors[key] = value.detach().cpu() + + +def _load_existing_adapter(input_dir: Path) -> tuple[dict[str, torch.Tensor], set[str]]: + adapter_path = input_dir / ADAPTER_MODEL_FILE + if not adapter_path.exists(): + return {}, set() + + tensors: dict[str, torch.Tensor] = {} + target_modules: set[str] = set() + for key, value in load_file(str(adapter_path)).items(): + cleaned_key = _clean_adapter_key(key) + _merge_tensor(tensors, cleaned_key, value) + target_module = _infer_target_module_from_key(cleaned_key) + if target_module is not None: + target_modules.add(target_module) + return tensors, target_modules + + +def _convert_fused_expert_lora( + fused_path: Path, +) -> tuple[dict[str, torch.Tensor], int, set[str]]: + if not fused_path.exists(): + raise FileNotFoundError(f"Missing {FUSED_EXPERT_LORA_FILE}: {fused_path}") + + output: dict[str, torch.Tensor] = {} + ranks: set[int] = set() + expert_counts: set[int] = set() + target_modules: set[str] = set() + + for key, tensor in sorted(load_file(str(fused_path)).items()): + match = KT_FUSED_KEY_RE.match(key) + if match is None: + raise ValueError(f"Unexpected key in {FUSED_EXPERT_LORA_FILE}: {key}") + + layer_idx, kt_name = match.groups() + if kt_name not in KT_NAME_MAP: + raise ValueError(f"Unsupported KT fused expert LoRA tensor: {key}") + if tensor.dim() != 3: + raise ValueError(f"{key} must be 3D [E, ...], got shape {tuple(tensor.shape)}") + + proj_name, lora_name, rank_dim = KT_NAME_MAP[kt_name] + expert_count = int(tensor.shape[0]) + rank = int(tensor.shape[rank_dim]) + expert_counts.add(expert_count) + ranks.add(rank) + target_modules.add(proj_name) + + for expert_idx in range(expert_count): + output_key = ( + f"model.layers.{layer_idx}.mlp.experts.{expert_idx}." + f"{proj_name}.{lora_name}.weight" + ) + _merge_tensor(output, output_key, tensor[expert_idx].contiguous()) + + if not output: + raise ValueError(f"No tensors found in {fused_path}") + if len(expert_counts) != 1: + raise ValueError(f"Inconsistent expert counts in {FUSED_EXPERT_LORA_FILE}: {sorted(expert_counts)}") + if len(ranks) != 1: + raise ValueError(f"Inconsistent LoRA ranks in {FUSED_EXPERT_LORA_FILE}: {sorted(ranks)}") + + return output, next(iter(ranks)), target_modules + + +def _build_adapter_config( + input_dir: Path, + rank: int, + target_modules: set[str], + base_model_name_or_path: str, + lora_alpha: float | None, + *, + include_input_target_modules: bool = True, +) -> dict: + config_path = input_dir / ADAPTER_CONFIG_FILE + config = _load_json(config_path) if config_path.exists() else {} + + if "lora_alpha" in config: + final_alpha = config["lora_alpha"] + elif lora_alpha is not None: + final_alpha = lora_alpha + else: + raise ValueError( + f"No {ADAPTER_CONFIG_FILE} with lora_alpha found in {input_dir}; " + "pass --lora-alpha to preserve runtime scaling." + ) + + existing_targets = config.get("target_modules", []) + if include_input_target_modules and isinstance(existing_targets, list): + target_modules.update(str(name).split(".")[-1] for name in existing_targets) + + config["peft_type"] = config.get("peft_type", "LORA") + config["r"] = rank + config["lora_alpha"] = final_alpha + config["target_modules"] = _ordered_target_modules(target_modules) + config["bias"] = config.get("bias", "none") + config["task_type"] = config.get("task_type", "CAUSAL_LM") + config["base_model_name_or_path"] = base_model_name_or_path + + return config + + +def _paths_have_ancestor_relationship(left: Path, right: Path) -> bool: + if left == right: + return True + try: + left.relative_to(right) + return True + except ValueError: + pass + try: + right.relative_to(left) + return True + except ValueError: + return False + + +def _validate_no_ancestor_paths( + paths: Iterable[Path], + *, + label: str, +) -> None: + resolved = list(paths) + for i, left in enumerate(resolved): + for right in resolved[i + 1 :]: + if _paths_have_ancestor_relationship(left, right): + raise ValueError( + f"{label} cannot have ancestor/descendant relationships: " + f"{left} and {right}." + ) + + +def _prepare_output_dir(output_path: Path, input_path: Path, overwrite: bool) -> None: + _validate_output_dir(output_path, input_path, overwrite) + if output_path.exists() and any(output_path.iterdir()): + shutil.rmtree(output_path) + output_path.mkdir(parents=True, exist_ok=True) + + +def _validate_output_dir(output_path: Path, input_path: Path, overwrite: bool) -> None: + if output_path == input_path: + raise ValueError("Output directory must be different from input directory.") + if _paths_have_ancestor_relationship(output_path, input_path): + raise ValueError( + "Output and input directories cannot be ancestor/descendant of each other: " + f"output={output_path}, input={input_path}." + ) + if output_path.exists() and not output_path.is_dir(): + raise FileExistsError(f"Output path exists and is not a directory: {output_path}") + if output_path.exists() and any(output_path.iterdir()): + if not overwrite: + raise FileExistsError(f"Output directory is not empty: {output_path}") + + +def _infer_lora_rank_from_tensor(key: str, tensor: torch.Tensor) -> int | None: + if ".lora_A." in key: + return int(tensor.shape[0]) + if ".lora_B." in key: + return int(tensor.shape[1]) + return None + + +def _validate_nonexpert_rank( + existing_tensors: Mapping[str, torch.Tensor], + expert_rank: int, + input_dir: Path, +) -> None: + if not existing_tensors: + return + + config_path = input_dir / ADAPTER_CONFIG_FILE + if config_path.exists(): + config_rank = _load_json(config_path).get("r") + if config_rank is not None and int(config_rank) != expert_rank: + raise ValueError( + f"Non-expert adapter rank mismatch: adapter_config.json r={config_rank}, " + f"but fused expert LoRA rank={expert_rank}." + ) + + for key, tensor in existing_tensors.items(): + tensor_rank = _infer_lora_rank_from_tensor(key, tensor) + if tensor_rank is None: + continue + if tensor_rank != expert_rank: + raise ValueError( + f"Non-expert adapter tensor rank mismatch for {key}: " + f"tensor rank={tensor_rank}, fused expert LoRA rank={expert_rank}." + ) + + +def _write_adapter( + output_path: Path, + input_path: Path, + tensors: dict[str, torch.Tensor], + config: Mapping, + *, + overwrite: bool, +) -> None: + _prepare_output_dir(output_path, input_path, overwrite) + save_file(tensors, str(output_path / ADAPTER_MODEL_FILE), metadata={"format": "pt"}) + _write_json(output_path / ADAPTER_CONFIG_FILE, config) + + +def convert_kt_to_sglang_adapter( + input_dir: str | os.PathLike, + output_dir: str | os.PathLike, + *, + base_model_name_or_path: str, + lora_alpha: float | None = None, + overwrite: bool = False, + expert_output_dir: str | os.PathLike | None = None, + nonexpert_output_dir: str | os.PathLike | None = None, +) -> dict: + input_path = Path(input_dir).expanduser().resolve() + output_path = Path(output_dir).expanduser().resolve() + expert_output_path = ( + Path(expert_output_dir).expanduser().resolve() + if expert_output_dir is not None + else None + ) + nonexpert_output_path = ( + Path(nonexpert_output_dir).expanduser().resolve() + if nonexpert_output_dir is not None + else None + ) + + if not input_path.is_dir(): + raise FileNotFoundError(f"Input directory not found: {input_path}") + + output_paths = [output_path] + output_paths.extend(path for path in (expert_output_path, nonexpert_output_path) if path is not None) + if len(set(output_paths)) != len(output_paths): + raise ValueError("Merged, expert, and non-expert output directories must be distinct.") + _validate_no_ancestor_paths( + output_paths, + label="Merged/expert/non-expert output directories", + ) + for path in output_paths: + _validate_output_dir(path, input_path, overwrite) + + existing_tensors, existing_targets = _load_existing_adapter(input_path) + fused_tensors, rank, fused_targets = _convert_fused_expert_lora(input_path / FUSED_EXPERT_LORA_FILE) + _validate_nonexpert_rank(existing_tensors, rank, input_path) + if nonexpert_output_path is not None and not existing_tensors: + raise ValueError( + f"Cannot write non-expert adapter: no {ADAPTER_MODEL_FILE} found in {input_path}." + ) + + tensors: dict[str, torch.Tensor] = {} + for key, value in existing_tensors.items(): + _merge_tensor(tensors, key, value) + for key, value in fused_tensors.items(): + _merge_tensor(tensors, key, value) + + target_modules = set(existing_targets) + target_modules.update(fused_targets) + config = _build_adapter_config( + input_path, + rank, + target_modules, + base_model_name_or_path, + lora_alpha, + ) + + _write_adapter(output_path, input_path, tensors, config, overwrite=overwrite) + + split_outputs: dict[str, dict] = {} + if expert_output_path is not None: + expert_config = _build_adapter_config( + input_path, + rank, + set(fused_targets), + base_model_name_or_path, + lora_alpha, + include_input_target_modules=False, + ) + _write_adapter( + expert_output_path, + input_path, + fused_tensors, + expert_config, + overwrite=overwrite, + ) + split_outputs["expert"] = { + "output_dir": str(expert_output_path), + "tensor_count": len(fused_tensors), + "target_modules": expert_config["target_modules"], + } + + if nonexpert_output_path is not None: + nonexpert_config = _build_adapter_config( + input_path, + rank, + set(existing_targets), + base_model_name_or_path, + lora_alpha, + include_input_target_modules=False, + ) + _write_adapter( + nonexpert_output_path, + input_path, + existing_tensors, + nonexpert_config, + overwrite=overwrite, + ) + split_outputs["nonexpert"] = { + "output_dir": str(nonexpert_output_path), + "tensor_count": len(existing_tensors), + "target_modules": nonexpert_config["target_modules"], + } + + return { + "input_dir": str(input_path), + "output_dir": str(output_path), + "tensor_count": len(tensors), + "rank": rank, + "target_modules": config["target_modules"], + "lora_alpha": config["lora_alpha"], + "split_outputs": split_outputs, + } + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Convert KT fused expert LoRA weights to an SGLang adapter directory." + ) + parser.add_argument("input_dir", help="Directory containing fused_expert_lora.safetensors.") + parser.add_argument("output_dir", help="Destination adapter directory.") + parser.add_argument( + "--base-model-name-or-path", + required=True, + help="Base model path/name to write into adapter_config.json.", + ) + parser.add_argument( + "--lora-alpha", + type=float, + default=None, + help="LoRA alpha to use when input adapter_config.json is absent.", + ) + parser.add_argument( + "--overwrite", + action="store_true", + help="Remove and recreate output_dir if it already contains files.", + ) + parser.add_argument( + "--expert-output-dir", + default=None, + help="Optional destination for a split expert-only adapter directory.", + ) + parser.add_argument( + "--nonexpert-output-dir", + default=None, + help="Optional destination for a split non-expert-only adapter directory.", + ) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + summary = convert_kt_to_sglang_adapter( + args.input_dir, + args.output_dir, + base_model_name_or_path=args.base_model_name_or_path, + lora_alpha=args.lora_alpha, + overwrite=args.overwrite, + expert_output_dir=args.expert_output_dir, + nonexpert_output_dir=args.nonexpert_output_dir, + ) + print( + "Converted KT fused expert LoRA adapter: " + f"{summary['tensor_count']} tensors, rank={summary['rank']}, " + f"target_modules={summary['target_modules']}" + ) + for name, split_summary in summary["split_outputs"].items(): + print( + f"Wrote {name} adapter: {split_summary['tensor_count']} tensors, " + f"target_modules={split_summary['target_modules']}, " + f"output_dir={split_summary['output_dir']}" + ) + + +if __name__ == "__main__": + main() diff --git a/kt-kernel/test/per_commit/test_convert_kt_to_sglang_adapter.py b/kt-kernel/test/per_commit/test_convert_kt_to_sglang_adapter.py new file mode 100644 index 00000000..1a4a4b0c --- /dev/null +++ b/kt-kernel/test/per_commit/test_convert_kt_to_sglang_adapter.py @@ -0,0 +1,373 @@ +import importlib.util +import os +import sys +from pathlib import Path + +import pytest +import torch +from safetensors.torch import load_file, save_file + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) +from ci.ci_register import register_cpu_ci + + +register_cpu_ci(est_time=5, suite="default") + +SCRIPT_PATH = ( + Path(__file__).resolve().parents[2] + / "scripts" + / "convert_kt_to_sglang_adapter.py" +) +SPEC = importlib.util.spec_from_file_location("convert_kt_to_sglang_adapter", SCRIPT_PATH) +converter = importlib.util.module_from_spec(SPEC) +assert SPEC.loader is not None +SPEC.loader.exec_module(converter) + + +def _write_full_fused_checkpoint(path: Path, *, rank: int = 3) -> dict[str, torch.Tensor]: + e, h, i = 2, 5, 7 + tensors = { + "layers.2.experts.gate_lora_a": torch.arange(e * rank * h, dtype=torch.float32).reshape(e, rank, h), + "layers.2.experts.gate_lora_b": torch.arange(e * i * rank, dtype=torch.float32).reshape(e, i, rank), + "layers.2.experts.up_lora_a": torch.arange(e * rank * h, dtype=torch.float32).reshape(e, rank, h) + 100, + "layers.2.experts.up_lora_b": torch.arange(e * i * rank, dtype=torch.float32).reshape(e, i, rank) + 200, + "layers.2.experts.down_lora_a": torch.arange(e * rank * i, dtype=torch.float32).reshape(e, rank, i) + 300, + "layers.2.experts.down_lora_b": torch.arange(e * h * rank, dtype=torch.float32).reshape(e, h, rank) + 400, + } + save_file(tensors, str(path / converter.FUSED_EXPERT_LORA_FILE)) + return tensors + + +def test_convert_fused_expert_lora_shapes_keys_and_config(tmp_path): + input_dir = tmp_path / "input" + output_dir = tmp_path / "output" + input_dir.mkdir() + fused = _write_full_fused_checkpoint(input_dir) + + summary = converter.convert_kt_to_sglang_adapter( + input_dir, + output_dir, + base_model_name_or_path="/models/base", + lora_alpha=16, + ) + + out = load_file(str(output_dir / converter.ADAPTER_MODEL_FILE)) + assert summary["tensor_count"] == 12 + assert summary["rank"] == 3 + assert summary["target_modules"] == ["gate_proj", "up_proj", "down_proj"] + + key = "model.layers.2.mlp.experts.1.gate_proj.lora_A.weight" + assert out[key].shape == (3, 5) + torch.testing.assert_close(out[key], fused["layers.2.experts.gate_lora_a"][1]) + + key = "model.layers.2.mlp.experts.0.down_proj.lora_B.weight" + assert out[key].shape == (5, 3) + torch.testing.assert_close(out[key], fused["layers.2.experts.down_lora_b"][0]) + + config = converter._load_json(output_dir / converter.ADAPTER_CONFIG_FILE) + assert config["peft_type"] == "LORA" + assert config["r"] == 3 + assert config["lora_alpha"] == 16 + assert config["base_model_name_or_path"] == "/models/base" + assert config["target_modules"] == ["gate_proj", "up_proj", "down_proj"] + + +def test_merges_existing_adapter_and_prefers_existing_lora_alpha(tmp_path): + input_dir = tmp_path / "input" + output_dir = tmp_path / "output" + input_dir.mkdir() + _write_full_fused_checkpoint(input_dir) + + existing_tensor = torch.ones(3, 5) + save_file( + { + "base_model.model.model.layers.0.self_attn.q_proj.lora_A.weight": existing_tensor, + }, + str(input_dir / converter.ADAPTER_MODEL_FILE), + ) + converter._write_json( + input_dir / converter.ADAPTER_CONFIG_FILE, + { + "peft_type": "LORA", + "r": 3, + "lora_alpha": 9, + "target_modules": ["q_proj"], + "bias": "none", + "task_type": "CAUSAL_LM", + "base_model_name_or_path": "old-base", + }, + ) + + summary = converter.convert_kt_to_sglang_adapter( + input_dir, + output_dir, + base_model_name_or_path="/models/base", + lora_alpha=16, + ) + + out = load_file(str(output_dir / converter.ADAPTER_MODEL_FILE)) + cleaned_key = "model.layers.0.self_attn.q_proj.lora_A.weight" + assert cleaned_key in out + torch.testing.assert_close(out[cleaned_key], existing_tensor) + assert summary["lora_alpha"] == 9 + + config = converter._load_json(output_dir / converter.ADAPTER_CONFIG_FILE) + assert config["lora_alpha"] == 9 + assert config["base_model_name_or_path"] == "/models/base" + assert config["target_modules"] == ["q_proj", "gate_proj", "up_proj", "down_proj"] + + +def test_writes_split_expert_and_nonexpert_adapters(tmp_path): + input_dir = tmp_path / "input" + output_dir = tmp_path / "output" + expert_dir = tmp_path / "expert" + nonexpert_dir = tmp_path / "nonexpert" + input_dir.mkdir() + _write_full_fused_checkpoint(input_dir) + + q_proj_tensor = torch.ones(3, 5) + o_proj_tensor = torch.full((5, 3), 2.0) + save_file( + { + "base_model.model.model.layers.0.self_attn.q_proj.lora_A.weight": q_proj_tensor, + "base_model.model.model.layers.0.self_attn.o_proj.lora_B.weight": o_proj_tensor, + }, + str(input_dir / converter.ADAPTER_MODEL_FILE), + ) + converter._write_json( + input_dir / converter.ADAPTER_CONFIG_FILE, + { + "peft_type": "LORA", + "r": 3, + "lora_alpha": 9, + "target_modules": ["q_proj", "o_proj"], + "bias": "none", + "task_type": "CAUSAL_LM", + "base_model_name_or_path": "old-base", + }, + ) + + summary = converter.convert_kt_to_sglang_adapter( + input_dir, + output_dir, + base_model_name_or_path="/models/base", + expert_output_dir=expert_dir, + nonexpert_output_dir=nonexpert_dir, + ) + + merged = load_file(str(output_dir / converter.ADAPTER_MODEL_FILE)) + expert = load_file(str(expert_dir / converter.ADAPTER_MODEL_FILE)) + nonexpert = load_file(str(nonexpert_dir / converter.ADAPTER_MODEL_FILE)) + + assert summary["tensor_count"] == 14 + assert summary["split_outputs"]["expert"]["tensor_count"] == 12 + assert summary["split_outputs"]["nonexpert"]["tensor_count"] == 2 + assert set(merged) == set(expert) | set(nonexpert) + assert set(expert).isdisjoint(nonexpert) + assert all(".mlp.experts." in key for key in expert) + assert not any(".mlp.experts." in key for key in nonexpert) + + cleaned_q_proj_key = "model.layers.0.self_attn.q_proj.lora_A.weight" + cleaned_o_proj_key = "model.layers.0.self_attn.o_proj.lora_B.weight" + torch.testing.assert_close(nonexpert[cleaned_q_proj_key], q_proj_tensor) + torch.testing.assert_close(nonexpert[cleaned_o_proj_key], o_proj_tensor) + + expert_config = converter._load_json(expert_dir / converter.ADAPTER_CONFIG_FILE) + nonexpert_config = converter._load_json(nonexpert_dir / converter.ADAPTER_CONFIG_FILE) + assert expert_config["target_modules"] == ["gate_proj", "up_proj", "down_proj"] + assert nonexpert_config["target_modules"] == ["q_proj", "o_proj"] + assert expert_config["base_model_name_or_path"] == "/models/base" + assert nonexpert_config["base_model_name_or_path"] == "/models/base" + + +def test_requires_lora_alpha_without_input_config(tmp_path): + input_dir = tmp_path / "input" + input_dir.mkdir() + _write_full_fused_checkpoint(input_dir) + + with pytest.raises(ValueError, match="pass --lora-alpha"): + converter.convert_kt_to_sglang_adapter( + input_dir, + tmp_path / "output", + base_model_name_or_path="/models/base", + ) + + +def test_rejects_inconsistent_rank(tmp_path): + input_dir = tmp_path / "input" + input_dir.mkdir() + save_file( + { + "layers.0.experts.gate_lora_a": torch.zeros(2, 3, 5), + "layers.0.experts.gate_lora_b": torch.zeros(2, 7, 4), + }, + str(input_dir / converter.FUSED_EXPERT_LORA_FILE), + ) + + with pytest.raises(ValueError, match="Inconsistent LoRA ranks"): + converter.convert_kt_to_sglang_adapter( + input_dir, + tmp_path / "output", + base_model_name_or_path="/models/base", + lora_alpha=8, + ) + + +def test_rejects_unexpected_fused_key(tmp_path): + input_dir = tmp_path / "input" + input_dir.mkdir() + save_file( + {"layers.0.experts.unknown_lora_a": torch.zeros(2, 3, 5)}, + str(input_dir / converter.FUSED_EXPERT_LORA_FILE), + ) + + with pytest.raises(ValueError, match="Unsupported KT fused expert LoRA tensor"): + converter.convert_kt_to_sglang_adapter( + input_dir, + tmp_path / "output", + base_model_name_or_path="/models/base", + lora_alpha=8, + ) + + +def test_rejects_nonempty_output_without_overwrite(tmp_path): + input_dir = tmp_path / "input" + output_dir = tmp_path / "output" + input_dir.mkdir() + output_dir.mkdir() + (output_dir / "existing.txt").write_text("do not remove", encoding="utf-8") + _write_full_fused_checkpoint(input_dir) + + with pytest.raises(FileExistsError, match="Output directory is not empty"): + converter.convert_kt_to_sglang_adapter( + input_dir, + output_dir, + base_model_name_or_path="/models/base", + lora_alpha=8, + ) + + converter.convert_kt_to_sglang_adapter( + input_dir, + output_dir, + base_model_name_or_path="/models/base", + lora_alpha=8, + overwrite=True, + ) + assert not (output_dir / "existing.txt").exists() + assert (output_dir / converter.ADAPTER_MODEL_FILE).exists() + + +def test_rejects_output_same_as_input(tmp_path): + input_dir = tmp_path / "input" + input_dir.mkdir() + _write_full_fused_checkpoint(input_dir) + + with pytest.raises(ValueError, match="different from input"): + converter.convert_kt_to_sglang_adapter( + input_dir, + input_dir, + base_model_name_or_path="/models/base", + lora_alpha=8, + overwrite=True, + ) + + +def test_rejects_output_ancestor_of_input(tmp_path): + run_dir = tmp_path / "run" + input_dir = run_dir / "adapter" + output_dir = run_dir + input_dir.mkdir(parents=True) + _write_full_fused_checkpoint(input_dir) + + with pytest.raises(ValueError, match="ancestor/descendant"): + converter.convert_kt_to_sglang_adapter( + input_dir, + output_dir, + base_model_name_or_path="/models/base", + lora_alpha=8, + overwrite=True, + ) + + +def test_rejects_output_descendant_of_input(tmp_path): + input_dir = tmp_path / "input" + output_dir = input_dir / "output" + input_dir.mkdir() + _write_full_fused_checkpoint(input_dir) + + with pytest.raises(ValueError, match="ancestor/descendant"): + converter.convert_kt_to_sglang_adapter( + input_dir, + output_dir, + base_model_name_or_path="/models/base", + lora_alpha=8, + overwrite=True, + ) + + +def test_rejects_split_output_ancestor_relationship(tmp_path): + input_dir = tmp_path / "input" + output_dir = tmp_path / "output" + expert_dir = tmp_path / "split" / "expert" + nonexpert_dir = tmp_path / "split" + input_dir.mkdir() + _write_full_fused_checkpoint(input_dir) + save_file( + { + "base_model.model.model.layers.0.self_attn.q_proj.lora_A.weight": torch.ones(3, 5), + }, + str(input_dir / converter.ADAPTER_MODEL_FILE), + ) + converter._write_json( + input_dir / converter.ADAPTER_CONFIG_FILE, + { + "peft_type": "LORA", + "r": 3, + "lora_alpha": 9, + "target_modules": ["q_proj"], + "bias": "none", + "task_type": "CAUSAL_LM", + "base_model_name_or_path": "old-base", + }, + ) + + with pytest.raises(ValueError, match="ancestor/descendant"): + converter.convert_kt_to_sglang_adapter( + input_dir, + output_dir, + base_model_name_or_path="/models/base", + expert_output_dir=expert_dir, + nonexpert_output_dir=nonexpert_dir, + ) + + +def test_rejects_mismatched_nonexpert_rank(tmp_path): + input_dir = tmp_path / "input" + output_dir = tmp_path / "output" + input_dir.mkdir() + _write_full_fused_checkpoint(input_dir, rank=3) + save_file( + { + "base_model.model.model.layers.0.self_attn.q_proj.lora_A.weight": torch.ones(4, 5), + }, + str(input_dir / converter.ADAPTER_MODEL_FILE), + ) + converter._write_json( + input_dir / converter.ADAPTER_CONFIG_FILE, + { + "peft_type": "LORA", + "r": 4, + "lora_alpha": 9, + "target_modules": ["q_proj"], + "bias": "none", + "task_type": "CAUSAL_LM", + "base_model_name_or_path": "old-base", + }, + ) + + with pytest.raises(ValueError, match="rank mismatch"): + converter.convert_kt_to_sglang_adapter( + input_dir, + output_dir, + base_model_name_or_path="/models/base", + ) diff --git a/kt-kernel/test/per_commit/test_convert_kt_to_sglang_adapter_integration.py b/kt-kernel/test/per_commit/test_convert_kt_to_sglang_adapter_integration.py new file mode 100644 index 00000000..470456b5 --- /dev/null +++ b/kt-kernel/test/per_commit/test_convert_kt_to_sglang_adapter_integration.py @@ -0,0 +1,242 @@ +import importlib.util +import json +import os +import re +import shutil +import sys +import time +from pathlib import Path + +import pytest +import torch +from safetensors.torch import load_file + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) +from ci.ci_register import register_cpu_ci + + +register_cpu_ci(est_time=60, suite="default") + +KT_ADAPTER_ENV = "KT_LORA_ADAPTER_DIR" +KT_BASE_MODEL_ENV = "KT_LORA_BASE_MODEL" +KT_ALPHA_ENV = "KT_LORA_ALPHA" +KT_LARGE_ADAPTER_ENV = "KT_LORA_LARGE_ADAPTER_DIR" + +REPO_ROOT = Path(__file__).resolve().parents[3] +SCRIPT_PATH = ( + Path(__file__).resolve().parents[2] + / "scripts" + / "convert_kt_to_sglang_adapter.py" +) +SPEC = importlib.util.spec_from_file_location("convert_kt_to_sglang_adapter", SCRIPT_PATH) +converter = importlib.util.module_from_spec(SPEC) +assert SPEC.loader is not None +SPEC.loader.exec_module(converter) + +SGLANG_EXPERT_KEY_RE = re.compile( + r"^model\.layers\.(\d+)\.mlp\.experts\.(\d+)\." + r"(gate_proj|up_proj|down_proj)\.lora_[AB]\.weight$" +) + + +def _required_adapter_dir(env_name: str) -> Path: + value = os.environ.get(env_name) + if not value: + pytest.skip(f"Set {env_name} to run real adapter integration tests.") + + path = Path(value).expanduser().resolve() + if not path.is_dir(): + pytest.fail(f"{env_name} is not a directory: {path}") + if not (path / converter.FUSED_EXPERT_LORA_FILE).is_file(): + pytest.fail(f"{env_name} must contain {converter.FUSED_EXPERT_LORA_FILE}: {path}") + return path + + +def _base_model_name_or_path() -> str: + value = os.environ.get(KT_BASE_MODEL_ENV) + if not value: + pytest.fail(f"Set {KT_BASE_MODEL_ENV} before running real adapter integration tests.") + return value + + +def _optional_lora_alpha() -> float | None: + value = os.environ.get(KT_ALPHA_ENV) + if value in (None, ""): + return None + try: + return float(value) + except ValueError: + pytest.fail(f"{KT_ALPHA_ENV} must be numeric, got: {value!r}") + + +def _lora_alpha_for_input(input_dir: Path) -> float | None: + alpha = _optional_lora_alpha() + if (input_dir / converter.ADAPTER_CONFIG_FILE).exists(): + return alpha + if alpha is None: + pytest.fail( + f"{input_dir} has no {converter.ADAPTER_CONFIG_FILE}; set {KT_ALPHA_ENV}." + ) + return alpha + + +def _convert_real_adapter(input_dir: Path, tmp_path: Path, output_name: str = "output") -> tuple[Path, dict]: + output_dir = tmp_path / output_name + summary = converter.convert_kt_to_sglang_adapter( + input_dir, + output_dir, + base_model_name_or_path=_base_model_name_or_path(), + lora_alpha=_lora_alpha_for_input(input_dir), + ) + return output_dir, summary + + +def _load_json(path: Path) -> dict: + with path.open("r", encoding="utf-8") as f: + return json.load(f) + + +def _link_or_copy(source: Path, dest: Path) -> None: + try: + os.symlink(source, dest) + except OSError: + shutil.copy2(source, dest) + + +def _assert_config_shape(output_dir: Path, summary: dict) -> dict: + config = _load_json(output_dir / converter.ADAPTER_CONFIG_FILE) + assert config["peft_type"] == "LORA" + assert config["r"] == summary["rank"] + assert config["lora_alpha"] == summary["lora_alpha"] + assert config["base_model_name_or_path"] == _base_model_name_or_path() + assert {"gate_proj", "up_proj", "down_proj"}.issubset(config["target_modules"]) + return config + + +def _assert_fused_tensors_preserved(input_dir: Path, output_dir: Path) -> int: + fused_tensors = load_file(str(input_dir / converter.FUSED_EXPERT_LORA_FILE)) + output_tensors = load_file(str(output_dir / converter.ADAPTER_MODEL_FILE)) + + checked = 0 + for input_key, input_tensor in sorted(fused_tensors.items()): + match = converter.KT_FUSED_KEY_RE.match(input_key) + assert match is not None, input_key + layer_idx, kt_name = match.groups() + assert kt_name in converter.KT_NAME_MAP, input_key + assert input_tensor.dim() == 3, input_key + + proj_name, lora_name, _rank_dim = converter.KT_NAME_MAP[kt_name] + for expert_idx in range(input_tensor.shape[0]): + output_key = ( + f"model.layers.{layer_idx}.mlp.experts.{expert_idx}." + f"{proj_name}.{lora_name}.weight" + ) + assert SGLANG_EXPERT_KEY_RE.match(output_key), output_key + assert output_key in output_tensors + assert output_tensors[output_key].shape == input_tensor[expert_idx].shape + assert output_tensors[output_key].dtype == input_tensor.dtype + assert torch.equal(output_tensors[output_key], input_tensor[expert_idx].cpu()) + checked += 1 + + assert checked > 0 + return checked + + +@pytest.mark.requires_model +def test_real_adapter_conversion_preserves_fused_tensors_and_config(tmp_path): + input_dir = _required_adapter_dir(KT_ADAPTER_ENV) + + output_dir, summary = _convert_real_adapter(input_dir, tmp_path) + + checked = _assert_fused_tensors_preserved(input_dir, output_dir) + config = _assert_config_shape(output_dir, summary) + output_tensors = load_file(str(output_dir / converter.ADAPTER_MODEL_FILE)) + + assert summary["tensor_count"] == len(output_tensors) + assert checked <= summary["tensor_count"] + assert config["target_modules"] == summary["target_modules"] + + +@pytest.mark.requires_model +def test_real_adapter_directory_merges_existing_adapter_model(tmp_path): + input_dir = _required_adapter_dir(KT_ADAPTER_ENV) + existing_adapter_path = input_dir / converter.ADAPTER_MODEL_FILE + if not existing_adapter_path.exists(): + pytest.skip(f"{input_dir} has no {converter.ADAPTER_MODEL_FILE} to merge.") + + output_dir, _summary = _convert_real_adapter(input_dir, tmp_path) + input_tensors = load_file(str(existing_adapter_path)) + output_tensors = load_file(str(output_dir / converter.ADAPTER_MODEL_FILE)) + + for input_key, input_tensor in input_tensors.items(): + output_key = converter._clean_adapter_key(input_key) + assert output_key in output_tensors + assert output_tensors[output_key].shape == input_tensor.shape + assert output_tensors[output_key].dtype == input_tensor.dtype + assert torch.equal(output_tensors[output_key], input_tensor.cpu()) + + +@pytest.mark.requires_model +def test_real_fused_conversion_without_input_config_uses_env_alpha(tmp_path): + input_dir = _required_adapter_dir(KT_ADAPTER_ENV) + alpha = _optional_lora_alpha() + if alpha is None: + pytest.skip(f"Set {KT_ALPHA_ENV} to validate conversion without input config.") + + no_config_input = tmp_path / "input_without_config" + no_config_input.mkdir() + _link_or_copy( + input_dir / converter.FUSED_EXPERT_LORA_FILE, + no_config_input / converter.FUSED_EXPERT_LORA_FILE, + ) + existing_adapter_path = input_dir / converter.ADAPTER_MODEL_FILE + if existing_adapter_path.exists(): + _link_or_copy(existing_adapter_path, no_config_input / converter.ADAPTER_MODEL_FILE) + + output_dir, summary = _convert_real_adapter(no_config_input, tmp_path, "output_without_config") + + assert summary["lora_alpha"] == alpha + config = _assert_config_shape(output_dir, summary) + assert config["lora_alpha"] == alpha + _assert_fused_tensors_preserved(no_config_input, output_dir) + + +@pytest.mark.requires_model +def test_sglang_lora_config_loader_accepts_converted_adapter(tmp_path): + input_dir = _required_adapter_dir(KT_ADAPTER_ENV) + output_dir, summary = _convert_real_adapter(input_dir, tmp_path) + + sglang_python = REPO_ROOT / "third_party" / "sglang" / "python" + sys.path.insert(0, str(sglang_python)) + try: + from sglang.srt.lora.lora_config import LoRAConfig + except Exception as exc: + pytest.fail(f"Unable to import SGLang LoRAConfig: {exc}") + + lora_config = LoRAConfig(str(output_dir)) + output_tensors = load_file(str(output_dir / converter.ADAPTER_MODEL_FILE)) + + assert lora_config.r == summary["rank"] + assert lora_config.lora_alpha == summary["lora_alpha"] + assert lora_config.target_modules == summary["target_modules"] + assert len(output_tensors) == summary["tensor_count"] + + +@pytest.mark.requires_model +def test_large_adapter_conversion_smoke(tmp_path, record_property): + input_dir = _required_adapter_dir(KT_LARGE_ADAPTER_ENV) + + start_time = time.perf_counter() + output_dir, summary = _convert_real_adapter(input_dir, tmp_path, "large_output") + duration_seconds = time.perf_counter() - start_time + + output_path = output_dir / converter.ADAPTER_MODEL_FILE + output_tensors = load_file(str(output_path)) + config = _assert_config_shape(output_dir, summary) + + record_property("conversion_seconds", round(duration_seconds, 3)) + record_property("output_bytes", output_path.stat().st_size) + record_property("tensor_count", summary["tensor_count"]) + + assert len(output_tensors) == summary["tensor_count"] + assert config["target_modules"] == summary["target_modules"] diff --git a/third_party/sglang b/third_party/sglang index ebaff772..7dc02b12 160000 --- a/third_party/sglang +++ b/third_party/sglang @@ -1 +1 @@ -Subproject commit ebaff7729b9e41c29d94f8d19a53473d321dc566 +Subproject commit 7dc02b12fdde70f911a69f68e230a85f9fce1775