# MoE SFT AMX 功能使用测试文档

## 1. 环境准备

### 1.1 系统要求

- **CPU**: Intel Xeon (支持 AMX 指令集，如 Sapphire Rapids 或更新)
- **内存**: 建议 64GB+ (取决于模型规模)
- **操作系统**: Linux (Ubuntu 20.04+)
- **Python**: 3.8+
- **PyTorch**: 2.0+

### 1.2 编译 kt-kernel

```bash
cd kt-kernel

# 创建 build 目录
mkdir build && cd build

# 配置 CMake
cmake ..

# 编译
make -j$(nproc)

# 返回上级目录
cd ..
```

### 1.3 验证安装

```python
import sys
sys.path.insert(0, "build")

from kt_kernel import kt_kernel_ext
print(f"kt_kernel_ext loaded: {kt_kernel_ext}")

# 检查 SFT MOE 是否可用
print(f"AMXBF16_SFT_MOE: {kt_kernel_ext.moe.AMXBF16_SFT_MOE}")
print(f"AMXInt8_SFT_MOE: {kt_kernel_ext.moe.AMXInt8_SFT_MOE}")
```

---

## 2. 测试用例说明

### 2.1 测试文件位置

| 文件 | 说明 |
|------|------|
| `kt-kernel/examples/test_moe_sft_amx.py` | TP 模式测试（多 NUMA 节点） |
| `kt-kernel/examples/test_moe_sft_amx_no_tp.py` | no-TP 模式测试（单 NUMA 节点） |

### 2.2 运行所有测试

```bash
cd kt-kernel

# 运行 TP 模式测试（需要多 NUMA 节点环境）
python examples/test_moe_sft_amx.py

# 运行 no-TP 模式测试（单节点）
python examples/test_moe_sft_amx_no_tp.py
```

### 2.3 测试用例列表

#### TP 模式测试 (test_moe_sft_amx.py)

| 测试函数 | 说明 | 验证内容 |
|---------|------|---------|
| `test_moe_sft_forward("bf16")` | BF16 模式前向传播 | 输出精度 < 5%，权重分区正确 |
| `test_moe_sft_forward("int8")` | INT8 模式前向传播 | 输出精度 < 5%（与推理测试一致） |
| `test_moe_sft_backward("bf16")` | BF16 模式反向传播 | 梯度精度 < 10%，梯度合并正确 |
| `test_moe_sft_backward("int8")` | INT8 模式反向传播 | 梯度精度 < 10%（与推理测试一致） |
| `test_moe_sft_lora_weight_sync("bf16"/"int8")` | LoRA 权重同步 | 分区权重同步正确性 |
| `test_moe_sft_training_loop("bf16"/"int8")` | 完整训练循环 | 端到端流程 |

#### no-TP 模式测试 (test_moe_sft_amx_no_tp.py)

| 测试函数 | 说明 | 验证内容 |
|---------|------|---------|
| `test_moe_sft_forward_no_tp("bf16")` | BF16 模式前向传播 | 输出精度 < 5%，零拷贝正确 |
| `test_moe_sft_forward_no_tp("int8")` | INT8 模式前向传播 | 输出精度 < 5%（与推理测试一致） |
| `test_moe_sft_backward_no_tp("bf16")` | BF16 模式反向传播 | 梯度精度 < 10% |
| `test_moe_sft_backward_no_tp("int8")` | INT8 模式反向传播 | 梯度精度 < 10%（与推理测试一致） |
| `test_moe_sft_lora_weight_sync_no_tp("bf16"/"int8")` | LoRA 权重同步 | 完全零拷贝正确性 |
| `test_moe_sft_training_loop_no_tp("bf16"/"int8")` | 完整训练循环 | 端到端流程 |

#### 两种模式的区别

| 特性 | TP 模式 | no-TP 模式 |
|------|--------|-----------|
| CPUInfer 初始化 | `CPUInfer(num_threads)` | `CPUInfer(num_threads, numa_id)` |
| NUMA 节点 | 多节点并行 | 单节点 |
| 权重分区 | 需要分区 | 不需要 |
| 梯度合并 | 需要合并 | 不需要 |
| LoRA 同步 | 每次 step 后需要 | 不需要 |

### 2.4 测试配置

```python
# 模型配置 (基于 DeepSeek-V3 架构)
expert_num = 256              # 专家总数
hidden_size = 7168            # 隐藏层维度
intermediate_size = 2048      # MLP 中间层维度
max_len = 25600               # 最大序列长度
num_experts_per_tok = 8       # 每 token 激活专家数 (top-k)

# LoRA 配置
lora_rank = 16                # LoRA 秩
lora_alpha = 32.0             # LoRA 缩放因子

# 测试配置
qlen = 4                      # 测试序列长度
num_threads = 60              # CPU 线程数
```

---

## 3. 精度验证方法

### 3.1 前向传播精度验证

```python
def test_forward_accuracy():
    # 1. 初始化相同的权重
    torch.manual_seed(42)
    gate_proj, up_proj, down_proj = init_base_weights(...)
    gate_lora_a, gate_lora_b, ... = init_lora_weights(...)

    # 2. PyTorch 参考实现
    torch_output, _ = moe_sft_torch_forward(
        input_data, expert_ids, weights,
        gate_proj, up_proj, down_proj,
        gate_lora_a, gate_lora_b, ...
    )

    # 3. AMX 实现
    CPUInfer.submit(moe.forward_sft_task(...))
    CPUInfer.sync()

    # 4. 比较精度
    diff = torch.mean(torch.abs(amx_output - torch_output))
    diff /= torch.mean(torch.abs(torch_output)) + 1e-8

    print(f"Relative difference: {diff:.6f}")
    assert diff < threshold  # BF16: 0.05, INT8: 0.15
```

### 3.2 反向传播精度验证

```python
def test_backward_accuracy():
    # 1. 前向传播
    torch_output, moe_saved = moe_sft_torch_forward(...)

    # 2. PyTorch 反向传播
    torch_grads = moe_sft_torch_backward(
        grad_output, moe_saved, ...
    )

    # 3. AMX 前向 + 反向
    CPUInfer.submit(moe.forward_sft_task(..., save_for_backward=True))
    CPUInfer.sync()
    CPUInfer.submit(moe.backward_task(...))
    CPUInfer.sync()

    # 4. 比较各项梯度
    for name in ['grad_input', 'grad_gate_lora_a', ...]:
        diff = torch.mean(torch.abs(amx_grad - torch_grad))
        diff /= torch.mean(torch.abs(torch_grad)) + 1e-8
        print(f"{name} diff: {diff:.6f}")
        assert diff < threshold
```

### 3.3 精度阈值

| 模式 | 前向阈值 | 反向阈值 |
|------|---------|---------|
| BF16 | 0.05 (5%) | 0.10 (10%) |
| INT8 | 0.05 (5%) | 0.10 (10%) |

> **注意**: INT8 SFT 模式使用与推理测试相同的阈值。基础权重在加载时量化为 INT8，但 LoRA 权重始终为 BF16，因此精度损失可控。

---

## 4. 完整训练示例

### 4.1 初始化

```python
import torch
import kt_kernel
kt_kernel_ext = kt_kernel.kt_kernel_ext

# 配置
expert_num = 256
hidden_size = 7168
intermediate_size = 2048
num_experts_per_tok = 8
lora_rank = 16
lora_alpha = 32.0
qlen = 4
num_threads = 60

# 初始化基础权重 (冻结)
gate_proj = torch.randn(expert_num, intermediate_size, hidden_size,
                        dtype=torch.bfloat16).contiguous() / 100
up_proj = torch.randn(expert_num, intermediate_size, hidden_size,
                      dtype=torch.bfloat16).contiguous() / 100
down_proj = torch.randn(expert_num, hidden_size, intermediate_size,
                        dtype=torch.bfloat16).contiguous() / 100

# 初始化 LoRA 权重 (可训练)
gate_lora_a = torch.randn(expert_num, lora_rank, hidden_size,
                          dtype=torch.bfloat16).contiguous() / 100
gate_lora_b = torch.zeros(expert_num, intermediate_size, lora_rank,
                          dtype=torch.bfloat16).contiguous()
up_lora_a = torch.randn(expert_num, lora_rank, hidden_size,
                        dtype=torch.bfloat16).contiguous() / 100
up_lora_b = torch.zeros(expert_num, intermediate_size, lora_rank,
                        dtype=torch.bfloat16).contiguous()
down_lora_a = torch.randn(expert_num, lora_rank, intermediate_size,
                          dtype=torch.bfloat16).contiguous() / 100
down_lora_b = torch.zeros(expert_num, hidden_size, lora_rank,
                          dtype=torch.bfloat16).contiguous()

# 包装为 nn.Parameter 用于 optimizer
gate_lora_a_param = torch.nn.Parameter(gate_lora_a)
gate_lora_b_param = torch.nn.Parameter(gate_lora_b)
up_lora_a_param = torch.nn.Parameter(up_lora_a)
up_lora_b_param = torch.nn.Parameter(up_lora_b)
down_lora_a_param = torch.nn.Parameter(down_lora_a)
down_lora_b_param = torch.nn.Parameter(down_lora_b)

lora_params = [
    gate_lora_a_param, gate_lora_b_param,
    up_lora_a_param, up_lora_b_param,
    down_lora_a_param, down_lora_b_param
]
optimizer = torch.optim.AdamW(lora_params, lr=1e-4)
```

### 4.2 创建 MOE SFT 实例

```python
# 初始化 CPUInfer
CPUInfer = kt_kernel_ext.CPUInfer(num_threads)

# 创建配置 (零拷贝设计)
config = kt_kernel_ext.moe.MOESFTConfig()
config.expert_num = expert_num
config.num_experts_per_tok = num_experts_per_tok
config.hidden_size = hidden_size
config.intermediate_size = intermediate_size
config.lora_rank = lora_rank
config.lora_alpha = lora_alpha
config.max_cache_depth = 1
config.max_len = 25600
config.layer_idx = 0

# 设置权重指针
config.gate_proj = gate_proj.data_ptr()
config.up_proj = up_proj.data_ptr()
config.down_proj = down_proj.data_ptr()

# 零拷贝: 直接指向 Python tensor
config.gate_lora_a = gate_lora_a_param.data.data_ptr()
config.gate_lora_b = gate_lora_b_param.data.data_ptr()
config.up_lora_a = up_lora_a_param.data.data_ptr()
config.up_lora_b = up_lora_b_param.data.data_ptr()
config.down_lora_a = down_lora_a_param.data.data_ptr()
config.down_lora_b = down_lora_b_param.data.data_ptr()
config.pool = CPUInfer.backend_

# 创建实例 (选择其一)
# BF16 模式
moe = kt_kernel_ext.moe.AMXBF16_SFT_MOE(config)
# INT8 模式 (基础权重将在 load_weights_task 时量化为 INT8)
# moe = kt_kernel_ext.moe.AMXInt8_SFT_MOE(config)

# 加载基础权重
CPUInfer.submit(moe.load_weights_task())
CPUInfer.sync()

# 预热
CPUInfer.submit(moe.warm_up_task())
CPUInfer.sync()
```

### 4.3 训练循环

```python
for step in range(100):
    # 生成数据
    expert_ids = torch.stack([
        torch.randperm(expert_num)[:num_experts_per_tok] for _ in range(qlen)
    ]).to(torch.int64).contiguous()
    weights = torch.rand(qlen, num_experts_per_tok, dtype=torch.float32).contiguous()
    weights = weights / weights.sum(dim=-1, keepdim=True)
    input_data = torch.randn(qlen, hidden_size, dtype=torch.bfloat16).contiguous() / 100
    target = torch.randn(qlen, hidden_size, dtype=torch.bfloat16).contiguous() / 100
    bsz_tensor = torch.tensor([qlen])

    # 1. 前向传播 (无需同步 LoRA 权重 - 零拷贝设计)
    output = torch.zeros(qlen, hidden_size, dtype=torch.float32).contiguous()
    CPUInfer.submit(moe.forward_sft_task(
        bsz_tensor.data_ptr(), num_experts_per_tok,
        expert_ids.data_ptr(), weights.data_ptr(),
        input_data.data_ptr(), output.data_ptr(), True
    ))
    CPUInfer.sync()

    # 2. 计算 loss
    loss = torch.mean((output - target.float()) ** 2)
    grad_output = (2 * (output - target.float()) / output.numel()).to(torch.bfloat16).contiguous()

    # 3. 反向传播
    grad_input = torch.zeros(qlen, hidden_size, dtype=torch.bfloat16).contiguous()
    grad_gate_lora_a = torch.zeros_like(gate_lora_a_param.data)
    grad_gate_lora_b = torch.zeros_like(gate_lora_b_param.data)
    grad_up_lora_a = torch.zeros_like(up_lora_a_param.data)
    grad_up_lora_b = torch.zeros_like(up_lora_b_param.data)
    grad_down_lora_a = torch.zeros_like(down_lora_a_param.data)
    grad_down_lora_b = torch.zeros_like(down_lora_b_param.data)

    CPUInfer.submit(moe.backward_task(
        grad_output.data_ptr(), grad_input.data_ptr(),
        grad_gate_lora_a.data_ptr(), grad_gate_lora_b.data_ptr(),
        grad_up_lora_a.data_ptr(), grad_up_lora_b.data_ptr(),
        grad_down_lora_a.data_ptr(), grad_down_lora_b.data_ptr()
    ))
    CPUInfer.sync()

    # 4. 复制梯度到 param.grad
    gate_lora_a_param.grad = grad_gate_lora_a
    gate_lora_b_param.grad = grad_gate_lora_b
    up_lora_a_param.grad = grad_up_lora_a
    up_lora_b_param.grad = grad_up_lora_b
    down_lora_a_param.grad = grad_down_lora_a
    down_lora_b_param.grad = grad_down_lora_b

    # 5. 优化器更新
    optimizer.step()
    optimizer.zero_grad()

    # 6. ★ TP 模式必需：同步分区权重 ★
    # (no-TP 模式可跳过此步骤，因为是完全零拷贝)
    CPUInfer.submit(moe.update_lora_weights_task(
        gate_lora_a_param.data.data_ptr(),
        gate_lora_b_param.data.data_ptr(),
        up_lora_a_param.data.data_ptr(),
        up_lora_b_param.data.data_ptr(),
        down_lora_a_param.data.data_ptr(),
        down_lora_b_param.data.data_ptr(),
    ))
    CPUInfer.sync()

    if step % 10 == 0:
        print(f"Step {step}, Loss: {loss.item():.6f}")
```

---

## 5. 常见问题

### 5.1 编译问题

**Q: 编译时提示 AMX 指令不支持**

A: 确保 CPU 支持 AMX 指令集 (Intel Sapphire Rapids 或更新)。可以通过以下命令检查:
```bash
lscpu | grep amx
```

### 5.2 运行时问题

**Q: 运行时提示 "Weights not loaded"**

A: 确保在调用 `forward_sft_task()` 前已调用 `load_weights_task()`:
```python
CPUInfer.submit(moe.load_weights_task())
CPUInfer.sync()
```

**Q: 输出全为零或 NaN**

A: 检查以下项目:
1. 输入张量是否 contiguous
2. 权重指针是否正确设置
3. 数值范围是否合理 (建议将权重初始化缩放到 1/100)

### 5.3 精度问题

**Q: 精度超出阈值**

A: 可能的原因:
1. INT8 量化模式精度损失较大，考虑使用 BF16 模式
2. 检查 LoRA 权重初始化是否正确
3. 检查 lora_alpha 和 lora_rank 的设置

### 5.4 性能问题

**Q: 性能不如预期**

A: 优化建议:
1. 确保使用合适的线程数 (通常为物理核心数)
2. 检查 NUMA 配置
3. 使用 `warm_up_task()` 预热

---

## 6. 调试技巧

### 6.1 打印中间值

```python
# 在测试文件中启用 debug_print
torch_output, _ = moe_sft_torch_forward(
    ...,
    debug_print=True  # 打印中间值
)
```

### 6.2 检查梯度

```python
# 检查梯度是否为零
for name, grad in [
    ("gate_lora_a", grad_gate_lora_a),
    ("gate_lora_b", grad_gate_lora_b),
    ...
]:
    print(f"{name}: norm={grad.norm():.6f}, "
          f"nonzero={grad.nonzero().shape[0]}/{grad.numel()}")
```

### 6.3 验证零拷贝

**注意**: 零拷贝只在 no-TP 模式下完全有效。在 TP 模式下，部分权重需要分区复制。

```python
# no-TP 模式：验证权重更新是否自动生效
before = gate_lora_a_param.data.clone()
optimizer.step()
after = gate_lora_a_param.data

# 检查 C++ 端是否看到更新（no-TP 模式）
output_before = run_forward(...)
# 权重已更新，无需调用 sync（no-TP 模式）
output_after = run_forward(...)

assert not torch.allclose(output_before, output_after)

# TP 模式：需要显式同步
optimizer.step()
CPUInfer.submit(moe.update_lora_weights_task(...))  # 必须！
CPUInfer.sync()
output_after = run_forward(...)  # 现在才能看到更新
```

### 6.4 TP 模式调试要点

| 问题现象 | 可能原因 | 解决方案 |
|---------|---------|---------|
| Forward 输出 ≈ 2x 期望值 | 基础权重未分区 | 检查 load_weights() |
| Backward 段错误 | BF16 权重被过早释放 | 检查权重生命周期 |
| 梯度数值不匹配 | 梯度未分区/合并 | 检查 backward() 实现 |
| Sync 测试失败 | 修改权重后未同步 | 添加 update_lora_weights_task() |
