加速原理：Triton kernel 与 fused LoRA

Unsloth 的口号是"2-5× faster, 80% less memory"。这章把数字背后的工程拆开：哪几个算子被融合、 哪几个 kernel 用 Triton 重写、为什么数学等价但速度差几倍。读完你应该能给同事讲清楚 "为什么 Unsloth 比纯 HF + PEFT 快"。

5.1大局观：LoRA 训练的时间花在哪

一个 LoRA 训练 step 的时间分布（HF + PEFT 默认实现，Llama-3-8B / 4090）：

Unsloth 对前 5 项都有 Triton 替代实现，每项加速 1.5-5×，总加速比 2-5×。

5.2Triton 是什么

Triton 是 OpenAI 开源的GPU kernel 编写 DSL，用 Python 写、JIT 编译成 PTX。优点：

• 比 CUDA 简单：不用管 threadIdx / blockIdx 细节，用 block 级抽象；
• 性能接近手写 CUDA：在大量算子上能达到 cuBLAS 80-95% 效率；
• 支持 autotune：自动搜索最佳 block size / num warps；
• 调试比 CUDA 友好：能用 Python 写 reference 对比。

Unsloth 几乎所有 kernel 都用 Triton 写。代码在 unsloth/kernels/。

5.3fused LoRA：核心加速点

LoRA 的数学：

$$ y = Wx + \frac{\alpha}{r}\, BA\, x,\quad W \in \mathbb{R}^{d\times d},\ A \in \mathbb{R}^{r\times d},\ B \in \mathbb{R}^{d\times r},\ r \ll d $$

原版 PEFT 的实现（4 个 kernel launch）

def forward(x):
    main = self.base_layer(x)              # kernel 1: W @ x  (大矩阵乘)
    lora = self.lora_A(x)                  # kernel 2: A @ x  (小矩阵乘)
    lora = self.lora_B(lora)               # kernel 3: B @ lora
    lora = lora * (self.alpha / self.r)    # kernel 4: 标量乘
    return main + lora

每个 kernel 都要去 GPU 读写一遍 input/output 张量，memory bound。小矩阵乘上 GPU launch overhead 还压不住。

Unsloth fused LoRA（1 个 kernel launch）

看 unsloth/kernels/fast_lora.py，核心是把上面 4 步融合：

@triton.jit
def lora_forward_kernel(
    X_ptr, W_ptr, A_ptr, B_ptr, Y_ptr,
    M, N, K, R,
    scale,
    BLOCK_M, BLOCK_N, BLOCK_K, BLOCK_R,
):
    # 一个 program 算一个 (BLOCK_M, BLOCK_N) 输出 block:
    # 1) 累加 W @ x 的贡献
    # 2) 累加 (B @ A @ x) * scale 的贡献
    # 3) 一次性写回 Y
    # 中间张量从不写回 HBM

少 3 次 kernel launch + 中间张量不落地 = LoRA 部分快 2-3 倍。

反向也 fuse

LoRA 反向要算 4 个梯度：$\partial L/\partial X$、$\partial L/\partial W$（如果训）、$\partial L/\partial A$、$\partial L/\partial B$。 Unsloth 把整套反向也写在一个 Triton kernel 里，进一步省时间。

5.4fused RMSNorm

RMSNorm 公式：

$$ \mathrm{RMSNorm}(x) = \frac{x}{\sqrt{\frac{1}{d}\sum_i x_i^2 + \epsilon}} \cdot g $$

$g$ 是可学习缩放参数。原版 PyTorch 实现拆成 5 步：求平方、求均值、加 eps、开根、除 + 乘 g。每步一个 kernel。

Unsloth 的 unsloth/kernels/rms_layernorm.py 把所有步骤 fuse 在一个 Triton kernel 里，整层只读写一次 hidden state。 内存带宽利用率从 ~20% 提到 90%+，对 4090 这种内存带宽本来就是瓶颈的卡尤其明显。

5.5fused RoPE

RoPE 的本质是给 Q、K 乘一个旋转矩阵：

$$ \begin{pmatrix} q'_{2i} \\ q'_{2i+1} \end{pmatrix} = \begin{pmatrix} \cos\theta_i & -\sin\theta_i \\ \sin\theta_i & \cos\theta_i \end{pmatrix} \begin{pmatrix} q_{2i} \\ q_{2i+1} \end{pmatrix} $$

原版实现用 PyTorch 切片 + cat：

# 慢的版本（10+ 个 elementwise kernel）
q_rot, q_pass = q[..., :rot_dim], q[..., rot_dim:]
q_rot = (q_rot * cos) + (rotate_half(q_rot) * sin)
q = torch.cat([q_rot, q_pass], dim=-1)

Unsloth 的 unsloth/kernels/rope_embedding.py 把整个旋转 fuse 成单个 Triton kernel，就地修改 q、k（in-place），免去 cat 的内存分配。

5.6fused SwiGLU

SwiGLU MLP 的计算：

$$ \mathrm{SwiGLU}(x) = (W_g x \odot \sigma(W_g x)) \odot W_u x,\quad \text{output} = W_d \cdot \mathrm{SwiGLU}(x) $$

三次矩阵乘（gate / up / down）+ 中间一次 SiLU + elementwise 乘法。原版分 4 个 kernel；Unsloth 把 SiLU + 中间 elementwise 乘法 fuse 进上下游矩阵乘 epilogue，看 unsloth/kernels/swiglu.py。

5.7chunked cross-entropy（大词表的杀手锏）

Llama-3 词表 128K，输出层 logits 矩阵在 batch_size=8、seq=2048 时占用：

$$ 8 \times 2048 \times 128000 \times 4\ \text{bytes} \approx 8.4\ \text{GB}\ (\text{fp32}) $$

原版 cross-entropy 一次性算所有 logits，这一个张量就把 24 GB 4090 压爆。Unsloth 的招：

1. 不实例化完整 logits 矩阵；
2. 把序列切成 chunks（默认 256 个 token 一块）；
3. 每个 chunk 算自己的 logits + loss + 梯度，当场释放 logits；
4. 累计 chunk loss 到最终 loss。

显存峰值从 8.4 GB 降到约 100 MB，损失 0 精度。unsloth/kernels/cross_entropy_loss.py 实现这件事。

5.8"unsloth" gradient checkpointing

原版 PyTorch torch.utils.checkpoint.checkpoint 的逻辑：

• 前向：丢掉所有中间激活；
• 反向：重新跑一遍 forward，重建激活，再算梯度。

Unsloth 的 "unsloth" 模式更激进：

• 不只丢中间激活，连一些中间张量都 offload 到 CPU；
• 反向时按需 reload；
• 显存额外省 30%（对长序列尤其明显）；
• 代价：训练慢一点点（5-10%），但能装下更大模型。

使用：

model = FastLanguageModel.get_peft_model(
    model,
    ...,
    use_gradient_checkpointing = "unsloth",   # ★ "unsloth" 而不是 True
)

5.9fp8 训练（v2026 新增）

看 unsloth/kernels/fp8.py，v2026 加入了 fp8 mixed precision 训练支持。条件：

• GPU 是 H100 / H200（Hopper）；
• 装好 TransformerEngine 或 PyTorch fp8 ops；
• 模型架构支持（目前主要 Llama 系）。

fp8 训练相比 bf16 再快 30-50%，但对消费卡（4090）无用，Ada / Ampere 架构没有 fp8 tensor cores。

5.10flex attention

unsloth/kernels/flex_attention.py 集成了 PyTorch FlexAttention（2024+ 引入），让用户能写自定义 attention mask （如 sliding window、causal + sliding 等），同时还能享受 FlashAttention 级别的性能。

对长上下文训练 / 滑窗 attention 尤其有用，第 7 章会展开。

5.11"数学等价"是怎么保证的

所有 Unsloth kernel 都做严格的数值等价测试：

• fp32 reference 实现（用 PyTorch 写）；
• Triton kernel 实现；
• 测试在 100K+ 随机输入上 torch.allclose(unsloth_out, pytorch_out, atol=1e-3, rtol=1e-3)。

看 tests/python/ 下有大量这类测试。这就是 Unsloth 敢说"快 + 数学等价"的依据 —— 任何 Unsloth 训出的模型， 理论上能用纯 HF + PEFT 复现到小数点后几位。

5.12对比 FlashAttention 怎么共存

Unsloth + flash-attn = 互补，不是替代。都装上效果最好。

5.13能不能不用 LoRA 也享受这些加速

可以。把 LoRA 部分关掉，剩下的 fused RMSNorm / RoPE / Cross-Entropy / SwiGLU / gradient ckpt 依然生效：

model, tokenizer = FastLanguageModel.from_pretrained(...)
# 不调 get_peft_model → 全参微调
trainer = SFTTrainer(model=model, ...)
trainer.train()

但全参微调显存往往装不下，所以实际场景几乎都配 LoRA / QLoRA用。

5.14性能数字

官方 README 公布的 benchmark（Llama-3-8B QLoRA on 1× A100 40G）：

4090 单卡上加速比通常更高（因为显存带宽相对计算更紧）：3-5×。

5.15kernel 源码定位（Triton 实现）

5.16patch 注入入口

5.17这章带走的

• "快 2-5×" 来自 6 个 fused Triton kernel：LoRA / RMSNorm / RoPE / SwiGLU / CE / gradient ckpt；
• chunked cross-entropy 是大词表模型上"省显存"的关键；
• use_gradient_checkpointing = "unsloth" 比 PyTorch 默认 ckpt 再省 30%；
• 所有 kernel 都过 torch.allclose 数学等价测试；
• fp8（H100）+ flex-attention 是 v2026 的新加速点；
• Unsloth 与 flash-attn 互补，都装上效果最好。