Transformers without Normalization

Posted on 2025-03-17 Edited on 2025-03-18 In Transformer Valine:

URL

paper: https://arxiv.org/pdf/2503.10622
code: https://github.com/jiachenzhu/DyT

TL;DR

这是由恺明和杨立昆提出的一篇关于 transformer 算子优化的论文，主要观点是去掉 transformer 结构中的 normalization 层，改成 tanh 层
改用 tanh 算子的 transformer 模型，在大多数任务上可达到使用归一化层的模型相同的性能，甚至更好

Algorithm

简单来说，这篇论文的核心思想是将 transformer 中的 normalization 层（可以是 LayerNorm 或 RMSNorm）替换成 dynamic tanh 层（简称 DyT）
normalization 计算公式：

\text{normalization}(x) = \gamma \times \frac{x - \mu}{\sqrt{\sigma^2+\epsilon}} + \beta

其中 $\mu$ 和 $\sigma$ 分别是 mean 和 std， $\gamma$ 和 $\beta$ 是 scale 和 shift 参数

DyT 计算公式：

\text{DyT}(x) = \gamma \times \tanh(\alpha x) + \beta

其中 $\alpha$ 是个可学习参数， $\gamma$ 和 $\beta$ 是 scale 和 shift 参数（和 normalization 一样）

DyT 实现伪代码：

# input x has the shape of [B, T, C]
# B: batch size, T: tokens, C: dimension
class DyT(Module):
    def __init__(self, C, init_α):
        super().__init__()
        self.α = Parameter(ones(1) * init_α)
        self.γ = Parameter(ones(C))
        self.β = Parameter(zeros(C))
    def forward(self, x):
        x = tanh(self.alpha * x)
        return self.γ * x + self.β

$\alpha$ 默认初始化值为 0.5

Results

作者在多个领域的知名模型上都对比了修改前后训练精度，DyT 的性能和 normalization 的性能基本一致，打的有来有回
作者还对比了 DyT 和 normalization 的训练/推理速度，DyT 的训练/推理速度要快很多
作者同时做 tanh 和 $\alpha$ 做了消融实验，发现 tanh 和 $\alpha$ 都是必要的

Thoughts

属于是恺明和立昆的梦幻联动了…，这种对最火的结构的优化，非大佬不能为也，想象下如果这篇论文是大学实验室发表的，大家第一反应恐怕是：Who think you are? 😂
之前算是稍微接触过硬件，DyT 这种 element-wise op 比 normalization 这种 reduce op 一定快多了，想怎么 tiling 都行…

2025.02 DeepSeek 开源周第三弹 —— DeepGEMM

Posted on 2025-03-10 In LLM Valine:

WIP

2025.02 DeepSeek 开源周第二弹 —— DeepEP

Posted on 2025-03-06 Edited on 2025-03-10 In LLM Valine:

URL

code: https://github.com/deepseek-ai/DeepEP

TL;DR

目前大模型常用的三种并行：
- DP: Data Parallelism，数据并行，将数据分成多份，每个 GPU 处理一份数据
- PP: Pipeline Parallelism，管道并行，将模型分成多个阶段（连续一层或多层为一个阶段），每个 GPU 处理一个阶段
- TP: Tensor Parallelism，张量并行，将模型分成多份（通常是一层/一个算子/一个张量分成多份，主要解决超长序列引起的超大张量问题），每个 GPU 处理一部分张量
对于 MoE 结构，EP (Expert Parallelism) 是一种新的并行策略，将 MoE 中的 Expert 分配到不同的 GPU 上
DeepEP 是一个用 cuda 实现的 MoE 模型的并行库，重点在于对 All-to-All 通信的优化

背景知识

节点内通信

通俗讲：一台服务器被称为一个节点，一个节点上的多个 GPU 之间的通信被称为节点内通信
通信协议：
- PCIe：比较通用的通信协议，目前最新的 PCIe 6.0 的带宽为 32GB/s 的双向带宽
- NVLink：NVIDIA 自家的通信协议，目前最新的 NVLink 5.0 的带宽约为 800GB/s 的双向带宽

节点间通信

多台服务器组成一个集群，集群中的服务器之间的通信被称为节点间通信
通信协议：
- InfiniBand：HPC 领域常用的通信协议，目前最新的 InfiniBand NDR 可达 400Gbps 的双向带宽
- Ethernet：通用的通信协议，速度低于 InfiniBand
通信技术：
- RDMA：Remote Direct Memory Access，远程直接内存访问，RDMA 通过 DMA 直接访问远程内存，减少了 CPU 的参与，提高了通信效率

通信视角下 MoE 结构的特殊性

MoE 结构的本质是超大规模参数量 + 小规模激活参数量（实际计算量）来让模型更强大同时推理效率高
由于 MoE 结构在实际推理过程中，每个 token 激活的专家 id 是无法提前预测的，而是一个纯 runtime 的行为
对此，为了降低 EP 通信压力，MoE 结构通常会限制每个 token 实际激活的节点数量。例如，DeepSeek V3 有 1 个共享专家和 256 个路由专家，每个 token 会激活 1 个共享专家和 8 个路由专家，但同时限制最多只能激活 4 个节点，假如得分最高的 8 个路由专家来自超过 4 个节点，那么会牺牲部分高分专家，在节点数不超过 4 的情况下，用贪心算法选择得分最高的 4 个专家
虽然 DP / TP / PP 都存在通信问题，但都是可提前规划好的通信数据量和通信模式，只要调度得当，即可重叠计算和通信耗时，而 MoE 结构的通信是无法提前规划的，因此 MoE 结构的通信是最难优化的

关键特性和能力

DeepEP 的关键特性和能力包括：

高吞吐量节点内通信：使用 NVLink 优化节点内所有到所有通信的内核，实现高达 155 GB/s 的带宽。
高吞吐量节点间通信：使用 RDMA 实现高效的跨节点所有到所有通信，在不同的 EP 配置中保持大约 45 GB/s 的带宽。
低延迟内核：专用推理解码内核，分发操作延迟低至 163 微秒，组合操作延迟低至 318 微秒。
FP8 支持：原生支持低精度操作，包括 FP8 分发，与大型模型中量化趋势一致。
灵活的 GPU 资源控制：可配置的 SM 使用，用于计算 - 通信重叠，允许精细调整性能优化。
自适应路由支持：在低延迟内核中支持自适应路由，使复杂拓扑中的网络利用更高效。

技术实现

DeepEP 是用 C++ 和 CUDA 组件实现的，并带有 Python 接口。实现包括几个关键组件：

缓冲管理：核心 Buffer 类管理 NVLink 和 RDMA 的通信缓冲区，处理内存分配和同步。
通信内核：
- 训练和推理预填充的高吞吐量内核
- 推理解码的低延迟内核
- 支持节点内（NVLink）和节点间（RDMA）通信
事件管理：EventOverlap 类提供 CUDA 事件处理和计算 - 通信重叠的工具。
分发和组合操作：
- dispatch：将令牌特征发送到跨 GPU 的对应专家
- combine：从专家收集处理后的特征并返回到原始位置

Thoughts

大模型，尤其是基座大模型，拼的是基建
大模型时代不会再出现小模型时代经常出现的 理论计算量低但实际很慢的算法 了，GPU 上快才是真的快，不光要考虑计算，存储 / 通信也同时需要认真考虑
软硬件 co-design 是未来趋势，People who're serious about software should make their own hardware. 这句名言的含金量还在上升

2025.02 DeepSeek 开源周第一弹 —— FlashMLA

Posted on 2025-03-06 Edited on 2025-03-10 In LLM Valine:

URL

code: https://github.com/deepseek-ai/FlashMLA

TL;DR

FlashMLA 是针对 DeepSeek 提出的 MLA (Multi-head Latent Attention) 模块，在 Nvidia Hopper 架构 GPU 上的加速算法，尤其对边长序列推理做了优化
FlashMLA 和 MLA 的关系大概可以类比到 FlashAttention 和 Attention 的关系

关键特性和能力

仅对 Hopper (sm_90) 架构做优化：充分挖掘了硬件 (sm) 的计算能力和 Memory Hierachy 的存储 / IO 能力
支持可变长度序列：和现实世界中推理场景更贴合
分页 KV cache：使用 block size = 64 的分页存储（这里的 64 的含义是：一个 block 存储某一层某个头的连续 64 个 token 对应的 kv cache）
高性能带宽和算力：在 H800 SXM5 设备上，在内存带宽受限配置环境下，内存带宽可达 3000 GB/s，在算力受限配置下，算力可达 580 TFLOPS
支持多种精度：支持 BF16 和 FP16

技术实现

基于 Nvidia cutlass 库 + cuda 实现主要计算
用干净的 python API 包装，方便集成到 PyTorch-base 框架中
用元数据管理的方式支持变长序列

Thoughts

做大模型不懂 cuda 是不行了，cutlass 要开始学起来了…
FlashMLA 只在 sm_90 上实现了，其他显卡编译不通，~~DeepSeek 只打高端局~~

Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention

Posted on 2025-02-20 Edited on 2025-03-03 In LLM Valine:

URL

paper: https://arxiv.org/pdf/2502.11089

TL;DR

本文提出了一种新的稀疏注意力机制，称为 Native Sparse Attention，该机制在硬件上对齐，并且可以直接训练，而无需额外的稀疏化技术
目前没有开源代码，只能通过论文中的公式和描述来实现

Algorithm

总体流程

本质上是用不同的 pattern 组合来替代 full attention，以减少计算量

代码实现

由 DeepSeek-R1 根据论文中的公式和描述实现

import torch
import torch.nn as nn
import torch.nn.functional as F
class NSAModule(nn.Module):
    def __init__(
        self,
        d_model,
        n_heads,
        compress_block=32,
        select_block=64,
        select_topk=16,
        window_size=512,
        dropout=0.1,
    ):
        super().__init__()
        self.d_model = d_model
        self.n_heads = n_heads
        self.head_dim = d_model // n_heads
        # 参数设置
        self.compress_block = compress_block
        self.select_block = select_block
        self.select_topk = select_topk
        self.window_size = window_size
        # 压缩用MLP
        self.compress_mlp = nn.Sequential(
            nn.Linear(self.head_dim * compress_block, 256),
            nn.GELU(),
            nn.Linear(256, self.head_dim),
        )
        # 门控机制
        self.gate_mlp = nn.Sequential(nn.Linear(d_model, 3 * n_heads), nn.Sigmoid())
        # 投影层
        self.q_proj = nn.Linear(d_model, d_model)
        self.k_proj = nn.Linear(d_model, d_model)
        self.v_proj = nn.Linear(d_model, d_model)
        self.out_proj = nn.Linear(d_model, d_model)
        self.dropout = nn.Dropout(dropout)
    def _compress_tokens(self, k, v):
        """压缩KV序列到块级别"""
        # 调整输入维度 (batch, seq_len, n_heads, head_dim)
        b, t, nh, hd = k.shape  # 修改为四维解包
        block_size = self.compress_block
        num_blocks = (t + block_size - 1) // block_size
        pad_len = num_blocks * block_size - t
        # 填充并分块
        k = F.pad(k, (0, 0, 0, 0, 0, pad_len))  # 添加头部维度的填充
        v = F.pad(v, (0, 0, 0, 0, 0, pad_len))
        # 调整维度: [batch, num_blocks, block_size, n_heads, head_dim]
        k_blocks = k.view(b, num_blocks, block_size, nh, hd)
        v_blocks = v.view(b, num_blocks, block_size, nh, hd)
        # 压缩处理 (保持头部分离)
        k_compressed = self.compress_mlp(
            k_blocks.permute(0, 1, 3, 2, 4).flatten(3)
        )  # [b, num_blocks, nh, hd]
        v_compressed = self.compress_mlp(v_blocks.permute(0, 1, 3, 2, 4).flatten(3))
        return k_compressed, v_compressed
    def _select_blocks(self, q, k_compressed, v_compressed):
        """基于注意力分数选择关键块"""
        # 计算压缩注意力分数
        scores = torch.einsum("bthd,bkhd->bthk", q, k_compressed) / (
            self.head_dim ** 0.5
        )
        probs = F.softmax(scores, dim=-1)
        # 选择topk块
        topk_scores, topk_indices = torch.topk(
            probs.mean(dim=2), self.select_topk, dim=-1
        )
        # 收集选中的块
        k_selected = torch.gather(
            k_compressed,
            1,
            topk_indices.unsqueeze(-1).expand(-1, -1, -1, self.head_dim),
        )
        v_selected = torch.gather(
            v_compressed,
            1,
            topk_indices.unsqueeze(-1).expand(-1, -1, -1, self.head_dim),
        )
        return k_selected, v_selected
    def forward(self, x, attn_mask=None):
        b, t, d = x.shape
        # 投影QKV并保持四维结构
        q = self.q_proj(x).view(b, t, self.n_heads, self.head_dim)
        k = self.k_proj(x).view(b, t, self.n_heads, self.head_dim)
        v = self.v_proj(x).view(b, t, self.n_heads, self.head_dim)
        # 压缩路径
        k_compressed, v_compressed = self._compress_tokens(k, v)
        # 选择路径
        k_selected, v_selected = self._select_blocks(q, k_compressed, v_compressed)
        # 滑动窗口
        k_window = k[:, max(0, t - self.window_size) :]
        v_window = v[:, max(0, t - self.window_size) :]
        # 门控权重
        gate = self.gate_mlp(x).view(b, t, 3, self.n_heads).permute(2, 0, 3, 1, 2)
        # 三路注意力计算
        attn_outputs = []
        for branch_k, branch_v in [
            (k_compressed, v_compressed),
            (k_selected, v_selected),
            (k_window, v_window),
        ]:
            scores = torch.einsum("bthd,bkhd->bthk", q, branch_k) / (
                self.head_dim ** 0.5
            )
            if attn_mask is not None:
                scores = scores.masked_fill(attn_mask == 0, -1e9)
            probs = F.softmax(scores, dim=-1)
            probs = self.dropout(probs)
            output = torch.einsum("bthk,bkhd->bthd", probs, branch_v)
            attn_outputs.append(output)
        # 门控融合
        weighted = sum(g * o for g, o in zip(gate, attn_outputs))
        output = weighted.contiguous().view(b, t, d)
        return self.out_proj(output)
def test_nsa_attention_shapes():
    # 测试参数
    batch_size = 2
    seq_len = 128
    d_model = 256
    n_heads = 8
    nsa_attn = NSAModule(d_model, n_heads)
    x = torch.randn(batch_size, seq_len, d_model)
    print("输入形状:", x.shape)
    # 打印中间形状
    q = nsa_attn.q_proj(x).view(batch_size, seq_len, n_heads, -1)
    k = nsa_attn.k_proj(x).view(batch_size, seq_len, n_heads, -1)
    v = nsa_attn.v_proj(x).view(batch_size, seq_len, n_heads, -1)
    print("\nQ/K/V形状:", q.shape)  # 应该输出 [2, 128, 8, 32]
    k_comp, v_comp = nsa_attn._compress_tokens(k, v)
    print("压缩后KV形状:", k_comp.shape)  # [2, 4, 8, 32]
if __name__ == "__main__":
    test_nsa_attention_shapes()

Thoughts

速度快并不惊讶，但效果竟然比 Full Attention 还好，如果实验比较 solid，那么 NSA 确实比较有前途
由于没有开源代码，所以只能通过论文中的公式和描述来实现，这对于一些实验复现和工程应用来说是一个挑战

DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning

Posted on 2025-02-18 Edited on 2025-03-10 In LLM Valine:

URL

paper: https://arxiv.org/pdf/2501.12948
code: https://github.com/deepseek-ai/DeepSeek-R1

TL;DR

本文提出了三个模型和其对应的训练方法，目的是提高大模型的 CoT 推理能力，三个模型分别是：
- DeepSeek-R1-Zero：由 DeepSeek V3 base 直接通过 RL 训练得到
- DeepSeek-R1：由 DeepSeek V3 base 在合成的高质量 CoT 数据集上 SFT + RL 训练得到，效果优于 DeepSeek-R1-Zero
- Distill model：用相同的合成高质量 CoT 数据训练的开源模型（参数量小且非 MoE）
其中训练方法是本文讲述的重点，DeepSeek R1 中的 R 表示 Reasoning（推导）

Algorithm

总体流程

Benchmark

DeepSeek-R1-Zero

DeepSeek-R1-Zero 是用 DeepSeek v3 base (pre-train) 直接通过 RL 训练得到的，不经过 SFT
用了以下的 Template 让模型生成 CoT 推理结果：
奖励建模分成两个部分：
- Accuracy rewards: 用于评估模型生成的 CoT 推理结果的准确性
- Format rewards: 用于评估模型生成的 CoT 推理结果的格式是否符合上述 prompt 要求的 CoT 的格式
且 RL 过程不使用 NN reward model，而是使用了 Rule base reward
RL 过程中使用了 GRPO 算法
DeepSeek-R1-Zero 仅仅用上述简单的 CoT 性能提升流程，即可大幅提高 benchmark 上的指标（相比于 DeepSeek v3），但仍然存在一些问题，例如可读性差和语言混合

DeepSeek-R1

DeepSeek-R1 用了更复杂的数据处理流程和训练流程，主要分成四个阶段（流程图上的四个 stage）：
- 冷启动：用数千个来自于 DeepSeek-R1-Zero 输出且经过人工处理的 CoT 数据，对 DeepSeek V3 base 进行 SFT
- 面向推导的强化学习：用编码、数学、科学和逻辑推理等推理密集型任务数据，对冷启动微调后的模型进行 RL 训练，同时引入了语言一致性和可读性的奖励
- 拒绝采样和监督微调：用上一个 stage 得到的模型生成更广泛领域的推导 CoT 数据 600K 条，用 DeepSeek V3 模型推理部分非推导数据（例如：写作、翻译等）得到潜在的思维链数据 200K 条，用这 800K 条数据对 DeepSeek V3 base 进行 SFT
- 所有场景的强化学习：用了和 DeepSeek V3 中使用的 RL 相同的 pipeline，数据和 DeepSeek-R1-Zero 中的 RL 数据相同，对上一个 stage 得到的模型进行 RL 训练
综上所述，DeepSeek-R1 最重要的是 800K 包含 CoT 的 SFT 数据

Distill model

由于 DeepSeek R1 使用了 DeepSeek V3 相同的模型结构，因此包含 671B 参数，每个 token 激活 34B 参数，因此这个模型非常大，不适合部署
为了解决这个问题，本文提出了 Distill model，用开源的较小参数量的稠密模型（非 MoE 模型）在 800K 包含 CoT 的 SFT 数据 进行 SFT 微调
主要蒸馏了 Qwen 和 llama 模型，且完全不做 RL，即可大幅提高原模型在 Benchmark 上的指标

Thoughts

需要重点关注本论文提到的 CoT 数据合成方式，这是 DeepSeek R1 的核心
能从本文中可以看出 数据重要性远大于模型 和 Post-train 只包含一次 SFT 和一次 RL，比如论文中 stage 1 模型会大胆抛弃 DeepSeek R1 Zero 模型，而是退回到了 DeepSeek V3 base 模型，只用了 Zero 生成数据；stage 3 模型会大胆抛弃之前 stage 的模型，而是退回到了 DeepSeek V3 base 模型，只用了 stage 2 模型生成数据
合成数据是最重要的，合成高质量的数据是 DeepSeek R1 成功的关键，甚至可以说如果 DeepSeek 开源了 80K 条合成数据，那么 CoT 领域的研究将会有一个新的起点

Reference

https://zhuanlan.zhihu.com/p/22103896675

DeepSeek-V3 Technical Report

Posted on 2025-02-02 Edited on 2025-03-03 In LLM Valine:

URL

paper: https://arxiv.org/pdf/2412.19437
code: https://github.com/deepseek-ai/DeepSeek-V3

TL;DR

本文提出了一种新的混合专家语言模型 DeepSeek-V3，共计 671B 总参数，对于每个 token 仅仅激活 37B 个参数，可以说又强又快。
主要的改进点包括：
- Multi-head Latent Attention (MLA)：MLA 是 DeepSeek 系列的核心模块（Deepseek v2 就有了），通过 MLP 将 hidden state 映射到 latent space，降低计算和存储复杂度。可参考 Deepseek v2 MLA 详解。
- Deepseek MOE：一种比 Gshard 更优的 MoE 架构，可参考 Deepseek MoE 详解。
- Multi-token predicition (MTP)：一种新的 token 预测方式，与传统的 predict next token 自回归预训练方式不同，MTP 通过一次预测多个 token（例如一次预测 next token 和 $next^2$ token 和 $next^3$ token），仅在训练阶段使用 MTP，推理阶段只使用 next token 的预测结果，即可大幅提高模型在 benchmark 上的综合表现。
- auxiliary-loss-free strategy for load balancing：一种新的 MoE 间负载均衡策略，无需额外的辅助损失，简化了训练过程，提高了训练效率。
性能与稳定性：
- 性能：DeepSeek-V3 的表现超过了其他开源模型，其性能已经可以与一些领先的闭源模型相媲美。
- 稳定性：尽管性能卓越，模型全程训练仅花费了 278.8 万 H800 GPU 小时，而且在整个训练过程中表现出极高的稳定性，没有出现不可恢复的损失暴增或需要回滚的情况。

Algorithm

总体流程

Pre-Training
Context Extension
Post-Training

Multi-token predicition (MTP)

上图给出了一次预测 next token 和 $next^2$ token 和 $next^3$ token 的 MTP 模型结构。
具体来说，MTP 是在标准 Decoder only Transformer 的基础上，额外引入了 k-1 个 token 预测头，每个 token 预测头都是一个 Norm + Concat + Projection + Transformer block 结构。
假设图上从左到右的 token head 分别记作 $head_0$ 、 $head_1$ 、 $head_2$ ，则 $head_0$ 预测的是 next token， $head_1$ 预测的是 $next^2$ token， $head_2$ 预测的是 $next^3$ token。
MTP 过程必须要保持因果性：
- $head_0$ 输入 $token_0$ ，主干网络输出 $h_0$ ，output layer 将 $h_0$ 映射为 next token 预测，记作 $token_1'$ 。
- $head_1$ 输入 $h_0$ 和 $token_1$ ，提取特征记作 $h_1$ ，output layer 将 $h_1$ 映射为 $next^2$ token 预测，记作 $token_2'$ 。
- $head_2$ 输入 $h_1$ 和 $token_2$ ，提取特征记作 $h_2$ ，output layer 将 $h_2$ 映射为 $next^3$ token 预测，记作 $token_3'$ 。
embedding layer 和 output layer 在 MTP 中是共享的。
MTP 通常只用在训练阶段，即在推理阶段只使用 $head_0$ 的输出。
也可以将 inference MTP 和 speculative decoding 结合，提高模型生成效率。

Auxiliary-loss-free strategy for load balancing

训练阶段负载均衡策略

训练阶段负载均衡策略并不是本文原创，而是来自于 Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts

g'_{i,t}=\begin{cases}s_{i,t},&s_{i,t}+b_i\in Topk(\{s_{i,t}+b_j|1\le j\le N_r\},K_r)\\0,& otherwise\end{cases}

其中：
- $s_{i,t}=Sigmoid(u_t^Te_i)$
- $b_i$ 是 bias 参数
- $N_r$ 是 Expert 的数量
- $K_r$ 是每个 token 保留的 Expert 数量
- $b_j$ $b_{j}$ 是根据每一个 Expert 被激活的次数动态调整的
  - 如果一个 Expert 被激活的次数过多（大于专家激活平均值），则 $b_j=b_j-\gamma$
  - 如果一个 Expert 被激活的次数过少（小于专家激活平均值），则 $b_j=b_j+\gamma$
  - 其中 $\gamma$ 是一个超参数，表示 bias update speed
关于 bias 的详细更新策略，可参考下图
- u 和本文中的 $\gamma$ 等价

推理阶段

推理阶段无需额外的负载均衡策略，直接使用 $s_{i,t}$ 即可。

Complementary Sequence-Wise Auxiliary Loss

上面提到的 Auxiliary-loss-free strategy for load balancing 损失只能保证 Expert 的负载均衡，但可能出现一个 Sequence 中每一个 token 都激活了相同的 Expert，这种情况下，Expert 之间的负载仍然不均衡。因此 Complementary Sequence-Wise Auxiliary Loss 通过引入 Auxiliary Loss 来保证 Sequence 内部 token 之间的负载均衡。
具体来说，需要关注：
- 离散选择频率：每一个 Expert 在 Sequence 中的激活次数
- 连续激活强度：每一个 Expert 在 Sequence 中平均专家分数
公式表述：

L_{Bal}=\alpha\sum_{i=1}^{N_r}f_iP_i

f_i=\frac{N_r}{K_rT}\sum_{t=1}^{T}\mathbb I(s_{i,t}\in Topk(\{s_{j,t}|1\le j\le N_r\},K_r))

s_{i,t}'=\frac{s_{i,t}}{\sum_{j=1}^{N_r}s_{j,t}}

P_i=\frac{1}{T}\sum_{t=1}^{T}s_{i,t}'

其中：
- $L_{Bal}$ 是 Complementary Sequence-Wise Auxiliary Loss
- $\alpha$ 是一个超参数，是一个很小的权重
- $f_i$ 是 Expert $i$ 的离散选择频率
- $P_i$ 是 Expert $i$ 的连续激活强度
- $T$ 是 Sequence 的长度
- $N_r$ 是 Expert 的数量
- $K_r$ 是每个 token 激活的 Expert 数量
- $\mathbb I(\cdot)$ 是指示函数
- $s_{i,t}$ 是 $Sigmoid(u_t^Te_i)$ ，表示专家的分数
道理也很简单，就是希望在一个 Sequence 中，每一个 Expert 被激活的次数和激活的强度都是均衡的。

训练框架设计

模型训练基于自主研发的 HAI-LLM 框架，这是一个经过优化的高效轻量级训练系统，目前没有开源。
DeepSeek-V3 的并行策略包含三个层面：
- 16 路流水线并行(Pipeline Parallelism, PP)
- 跨 8 个节点的 64 路专家并行(Expert Parallelism, EP)
- ZeRO-1 数据并行(Data Parallelism, DP)
开发了 DualPipe 流水线并行算法，相比于现有的 PP 算法，减少了空泡，基本上实现了计算和通信的高度重叠。
优化了跨节点全对全通信内核，充分利用节点内 GPU 间的 NVLink 带宽和节点间 InfiniBand 带宽，同时减少了通信所需的流式多处理器(SMs)资源占用
通过精细的内存管理优化，使得模型训练无需依赖开销较大的张量并行(Tensor Parallelism, TP)技术

FP8

DeepSeek-V3 开发了 FP8 (E4M3) 混合精度训练框架，首次在超大规模模型上验证了 FP8 训练的可行性和效果。
通过算法、框架和硬件的综合优化，突破了跨节点 MoE 训练中的通信瓶颈，实现了计算与通信的高度重叠。这种优化大幅提升了训练效率，降低了训练成本，同时支持了更大规模模型的训练而无需额外开销。

Thoughts

实际上 DeepSeek V3 技术报告中还有很多细节，尤其是很多关于工程优化的部分，此处不再一一列举
工程能力，尤其是开发一套分布式训练框架且根据模型特性对其深度优化的能力，是一个大模型成功的不可或缺因素
如果想要在大模型领域有所建树，不仅需要有强大的算法研究能力，还需要有强大的工程能力，比如 CUDA 编程能力等

References

https://zhuanlan.zhihu.com/p/14890557782

Denoising Diffusion Probabilistic Models

Posted on 2025-01-20 Edited on 2025-03-03 In Diffusion Valine:

URL

paper: https://arxiv.org/pdf/2006.11239
code: https://github.com/hojonathanho/diffusion

TL;DR

这篇论文首次提出了 Denoising Diffusion Probabilistic Models (DDPM) 模型，该模型是一种去噪扩散模型，可以用于生成图像。
DDPM 有详细的热力学理论基础（Langevin 朗之万动力学），对热力学扩散过程进行了建模。
本质上，DDPM 是一个自回归生成模型，通过自监督方式，对图像进行添加噪声和逐步去噪的过程，来学习图像的分布。

Algorithm

DDPM 总体结构

上图中 $x_0$ 是来自图像数据集的真实图像，从右到左被称为前向过程，即对图像逐步添加随机高斯噪声，直到得到噪声图像 $x_T$
从左到右被称为逆向过程，即模型预测噪声，代入公式逐步去噪

前向过程

x_t=\sqrt {1-\beta_t} \cdot x_{t-1}+\sqrt \beta_t\cdot \epsilon_{t-1}\ ,\ \ \epsilon_{t-1}\sim N(0, I)

其中：
- $\epsilon_t$ 是随机高斯噪声，即 torch.randn
- $\beta_t$ 是超参数，控制噪声的强度，随着 $t$ 的增大， $\beta_t$ 从 1e-4 逐渐增大到 2e-2
- 总时间步数 $T=1000$

前向过程的累积表示

x_t=\sqrt {\bar\alpha_t}\cdot x_0+\sqrt{1-\bar\alpha_t}\cdot \epsilon

具体推理过程：
- $x_1 = \sqrt{1-\beta_1}\cdot x_0+\sqrt{\beta_1}\cdot \epsilon_0$
- $x_2 = \sqrt{1-\beta_2}\cdot x_1+\sqrt{\beta_2}\cdot \epsilon_1$
- 将 $x_1$ 代入 $x_2$ 的公式中，得到 $x_2$ 的累积表示： $x_2 = \sqrt{1-\beta_2}\cdot (\sqrt{1-\beta_1}\cdot x_0+\sqrt{\beta_1}\cdot \epsilon_0)+\sqrt{\beta_2}\cdot \epsilon_1$
- $x_2 = \sqrt{1-\beta_2}\sqrt{1-\beta_1}x_0+\sqrt{1-\beta_2}\sqrt{\beta_1}\cdot \epsilon_0+\sqrt{\beta_2}\cdot \epsilon_1$
- 以此类推，得到 $x_t$ 的累积表示： $x_t=\sqrt{\prod_{i=1}^t(1-\beta_i)}\cdot x_0+\sum_{i=1}^t\sqrt{\prod_{j=i+1}^t(1-\beta_j)}\cdot \sqrt{\beta_i}\cdot \epsilon$
- 令 $\bar\alpha_t=\prod_{i=1}^t(1-\beta_i)$ ，得到 $x_t=\sqrt{\bar\alpha_t}\cdot x_0++\sum_{i=1}^t\sqrt{\beta_i\cdot\prod_{j=i+1}^t(1-\beta_j)}\cdot \epsilon$
- 需要将噪声项合并成一个正态分布，需要确定方差，对于每一项 $\sqrt{\beta_i\cdot\prod_{j=i+1}^t(1-\beta_j)}$ ，其方差为 $\beta_i\cdot\prod_{j=i+1}^t(1-\beta_j)$
- 总噪声的方差为 $\sum_{i=1}^t\beta_i\cdot\prod_{j=i+1}^t(1-\beta_j)$
- 代入 $\bar\alpha_t$ 到噪声项的方差中，得到总噪声的方差为 $\sum_{i=1}^t\beta_i\cdot\frac{\bar\alpha_t}{\bar\alpha_i}=\bar\alpha_t\cdot\sum_{i=1}^t\frac{\beta_i}{\bar\alpha_i}$
- 由于 $\frac{\beta_i}{\bar\alpha_i}=\frac{1-\alpha_i}{\bar\alpha_i}=\frac{1}{\bar\alpha_i}-\frac{1}{\bar\alpha_{i-1}}$ ，因此使用 裂项求和 公式可得： $\sum_{i=1}^t\frac{\beta_i}{\bar\alpha_i}=\frac{1}{\bar\alpha_t}-\frac{1}{\bar\alpha_0}$
- 由于 $\bar\alpha_0=1$ ，因此 $\sum_{i=1}^t\frac{\beta_i}{\bar\alpha_i}=\frac{1}{\bar\alpha_t}-1$
- 代入噪声总方差公式中，得到总噪声的方差为 $\bar\alpha_t\cdot(\frac{1}{\bar\alpha_t}-1)=1-\bar\alpha_t$
- 因此，累加形式的前向过程可表示为： $x_t=\sqrt{\bar\alpha_t}\cdot x_0+\sqrt{1-\bar\alpha_t}\cdot \epsilon$
累积表示的好处：
- 在给定 $t$ 时，可以直接从 $x_0$ 生成 $x_t$ ，而不需要逐步生成
- 通过单步表示到累积表示的数学推导过程，可以更好地理解噪声的累积方差

逆向过程

x_{t-1}=\frac{1}{\sqrt{\alpha_t}}(x_t-\frac{\beta_t}{\sqrt{1-\bar\alpha_t}}\cdot \epsilon_\theta(x_t,t))

其中：
- $\epsilon_\theta(x_t,t)$ 是模型预测的噪声，即 model(x_t, t)
- $\alpha_t=1-\beta_t$
- $\bar\alpha_t=\prod_{s=1}^t\alpha_s$
由上述公式可知：预测得到的噪声并不会直接叠加到噪声图像上，而是利用扩散原理，根据噪声添加相关的超参数进行去噪
由于 $\epsilon_\theta(x_t,t)$ 是模型对于 $x_t$ 的预测，且 $x_t$ 也和模型预测有关，因此逆向过程是个数据依赖的过程，因此逆向过程的累积表示没有意义

训练过程

DDPM 的训练过程基于扩散模型的前向过程和逆向过程。训练的目标是让模型学习如何从噪声数据中恢复出原始数据。具体步骤如下：

输入数据：从数据集中随机选择一批原始图像 $x_0$
随机选择扩散步数：对于每张图像，随机选择一个扩散步数 $t$ ，范围从 $1$ 到 $T$ 随机选择
生成噪声数据：通过前向过程生成噪声图像 $x_t$
预测噪声：使用模型 $\epsilon_\theta(x_t,t)$ 预测添加到 $x_0$ 中的噪声
计算损失：通过比较预测噪声和真实噪声计算损失函数（通常是均方误差，MSE）
更新模型参数：通过反向传播更新模型参数

伪代码

for epoch in range(num_epochs):
    for x0 in data_loader:
        t = random.randint(1, T)  # 随机选择扩散步数
        xt, noise = forward_diffusion(x0, t)  # 生成噪声数据
        predicted_noise = model(xt, t)  # 预测噪声
        loss = mse_loss(predicted_noise, noise)  # 计算损失
        update_model_parameters(loss)  # 更新模型参数

模型输入输出

输入
- 原始图像 $x_0$ ，通常被归一化到 $[-1, 1]$ 之间
- 扩散步数 $t$ ，随机选择，会变成 time embedding 输入模型
输出
- 预测噪声 $\epsilon_\theta(x_t,t)$ ，模型输出是预测的噪声，用于逆向过程去噪

损失函数

通常使用均方误差（MSE）作为损失函数

\mathcal L(\theta)=\mathbb E_{x_0,t}\left[\left\|\epsilon_\theta(x_t,t)-\epsilon\right\|^2\right]

其中：
- $\epsilon$ 是真实噪声
- $\epsilon_\theta(x_t,t)$ 是模型预测的噪声
- $\mathbb E_{x_0,t}$ 表示对所有的 $x_0$ 和 $t$ 求期望

推理过程（采样过程）

用 pre-trained diffusion model 生成图像

from diffusers import UNet2DModel
from diffusers import DDPMScheduler
from PIL import Image
import numpy as np
import tqdm
import torch
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
def display_sample(sample, i):
    image_processed = sample.cpu().permute(0, 2, 3, 1)
    image_processed = (image_processed + 1.0) * 127.5
    image_processed = image_processed.numpy().astype(np.uint8)
    image_pil = Image.fromarray(image_processed[0])
    image_pil.save(f"sample_{i}.png")
# load pretrain model
repo_id = "google/ddpm-church-256"
model = UNet2DModel.from_pretrained(repo_id).to(device)# build scheduler
scheduler = DDPMScheduler.from_config(repo_id)# random noise as input
noisy_sample = torch.randn(
    1, model.config.in_channels, model.config.sample_size, model.config.sample_size
)
sample = noisy_sample.to(device)# step through the scheduler
for i, t in enumerate(tqdm.tqdm(scheduler.timesteps)):
    # 1. predict noise residual
    with torch.no_grad():
        residual = model(sample, t).sample    # 2. compute less noisy image and set x_t -> x_t-1
    sample = scheduler.step(residual, t, sample).prev_sample    # 3. optionally look at image
    if (i + 1) % 50 == 0:
        display_sample(sample, i + 1)

通过 UNet2DModel 加载预训练模型
通过 DDPMScheduler 构建调度器（主要作用是生成噪声和去噪）
通过 torch.randn 生成随机噪声作为输入
通过 scheduler.step 逐步生成图像
最终生成一张完全去噪的新图像，这张图像并不在任何数据集中（只是风格看起来像 Church-256 数据集），是模型自己生成的（这也意味着没办法去指导模型生成特定的图像）

用 pre-trained diffusion model 对已逐步添加噪声的图像去噪

理论上，可以将一张数据集中的图像逐步添加噪声，在 $t=T$ 步之后，使用模型去噪，得到一张和原始图像非常相似的图像
但实际测试中，当 $t$ 较大时（大于 300， $T$ 等于 1000），得到的去噪图像和原始图像相差较大，在 $t<300$ 时，去噪图像和原始图像相似度较高

from diffusers import UNet2DModel
from diffusers import DDPMScheduler
from diffusers import DDPMPipeline
from PIL import Image
import numpy as np
import tqdm
import torch
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
def preprocess_image(image):
    """Preprocess an image to model input format."""
    image = np.array(image).astype(np.float32) / 127.5 - 1.0  # Normalize to [-1, 1]
    image = (
        torch.tensor(image).permute(2, 0, 1).unsqueeze(0).to(device)
    )  # Shape: (1, 3, H, W)
    return image
def display_sample(sample, name):
    """Save a sample as an image."""
    image_processed = sample.cpu().permute(0, 2, 3, 1)
    image_processed = (image_processed + 1.0) * 127.5
    image_processed = image_processed.numpy().astype(np.uint8)
    image_pil = Image.fromarray(image_processed[0])
    image_pil.save(f"{name}.png")
# Load pretrained pipeline to generate an image
repo_id = "google/ddpm-celebahq-256"
image_pipe = DDPMPipeline.from_pretrained(repo_id).to(device)
image = image_pipe().images[0]# Save original image
image.save("original_image.png")# Preprocess original image
original_image = preprocess_image(image)# Load pretrained model and scheduler
model = UNet2DModel.from_pretrained(repo_id).to(device)
scheduler = DDPMScheduler.from_config(repo_id)# Number of timesteps
T = scheduler.num_train_timesteps# Step 1: Gradually add noise to the image
alpha_cumprod = scheduler.alphas_cumprod[300]
noisy_image = torch.sqrt(alpha_cumprod) * original_image + torch.sqrt(
    1 - alpha_cumprod
) * torch.randn_like(original_image)
display_sample(noisy_image, "noisy_image")
sample = noisy_image.clone()
for t in tqdm.tqdm(scheduler.timesteps[-300:]):
    with torch.no_grad():
        residual = model(sample, t).sample  # Predict noise residual
    sample = scheduler.step(residual, t, sample).prev_sample  # Denoise step by step# Save final denoised image
display_sample(sample, "final_image")# Compare original and final image
l2_diff = torch.norm(original_image - sample).item()
print(f"L2 difference between original and final image: {l2_diff}")

上图表示 t=300 时，原图 -> 噪声图 -> 去噪图

上图表示 t=500 时，原图 -> 噪声图 -> 去噪图

上图表示 t=1000 时，原图 -> 噪声图 -> 去噪图

Thoughts

DDPM 是一种非常有趣的生成模型，数学基础非常扎实
但仍然存储三个主要问题：
- 时间步 $T$ 太大，通常是 1000
- 无法指定模型生成特定的图像，只能生成和数据集相似的图像
- 作用在图像域，计算量较大
针对以上三个问题，stable diffusion 诞生了

BLIP: Bootstrapping Language-Image Pre-training for Unified Vision-Language Understanding and Generation

Posted on 2025-01-14 Edited on 2025-03-03 In MultiModal Valine:

URL

paper: https://arxiv.org/pdf/2201.12086
code: https://github.com/salesforce/BLIP

TL;DR

BLIP 是 Salesforce 提出的一种图文多模态对齐算法，全称是 Bootstrapping Language-Image Pre-training，是一种自举的图文多模态预训练算法
相较于 CLIP 只能对齐图像和文本表征，BLIP 将任务范式进行了扩展，基本已经有了多模态大模型的雏形

Algorithm

BLIP 整体架构

BLIP 在预训练阶段，包含三种损失函数
- Image-Text Contrastive Loss：对齐图像和文本表征，和 CLIP 的对齐方式类似，不同的是 BLIP 引入了 Momentum Model，可以使得训练更加稳定
- Image-Text Macthing Loss：图像和文本匹配损失，输出 True 或 False，本质是交叉熵损失
- Language Modeling Loss：语言建模损失，预测下一个词（大模型范式），用于文本内容生成
从上图可以看出，模型实际包含四个部分：
- Image Encoder：图像编码器，将图像编码为 Image Embedding
- Text Encoder：文本编码器，将文本编码为 Text Embedding
- Image-grounded Text Eecoder：图像驱动文本编码器，用文本特征 Cross-Modal Attention 去索引图像特征，输出文本图像特征是否匹配（True / False）
- Image-grounded Text Decoder：图像驱动文本解码器，用文本特征 Cross-Modal Attention 去索引图像特征，用逐个 Token 预测的方式生成文本，回答文本对应的问题
从实现代码看，Text Encoder / Image-grounded Text Eecoder / Image-grounded Text Decoder 实际上共用了一个 Bert 预训练模型的一组参数（仅仅在部分位置不共享参数）
- Text Encoder：
  - 原生 Bert 模型，用于文本编码
  - 不做 Cross-Modal Attention，只做 Self-Attention
  - 最后输出层为表征层
  - 用特殊的 [CLS] 作为任务标识符
- Image-grounded Text Eecoder：
  - 魔改 Bert 模型，加入 Cross-Modal Attention
    - 这种状态处于 encoder 和 decoder 之间，比 encoder 多了 Cross-Modal Attention，比 decoder 少了 Causal Self-Attention
    - 用文本特征去索引图像特征
  - 最后输出层为二分类层
  - 用特殊的 [Encoder] 作为任务标识符
- Image-grounded Text Decoder：
  - 魔改 Bert 模型，将其转化成一个完全的 decoder 形态，主要包括两点：
    - 加入 Cross-Modal Attention，用文本特征去索引图像特征
    - 由于是生成任务，所以需要用 Causal Self-Attention 取代 Bert 默认的 Bidirection Self-Attention
  - 最后输出层为 vocab prob 层，用于生成文本（下一个文本在词表中的概率）
  - 用特殊的 [Decoder] 作为任务标识符
- 三个模型仅仅在 Head 层和 Attention 层参数上有所不同，其余参数共享

Bootstrapping 策略

BLIP 中的 B 就是 Bootstrapping 的意思，即自举，表示此算法有一种自我增强的方式
自我增强的方式主要是通过 CapFilt 方法实现的：
- Captioner：用于为图像生成高质量的文本描述
- Filter：用于过滤掉低质量的图像-文本对，保留高质量的对用于训练
CapFilt 工作流程：
1. 生成伪标签：使用 Captioner 为无标签或弱标签图像生成文本描述，形成伪标签
2. 过滤伪标签：使用 Filter 评估生成的伪标签的质量，保留高质量的伪标签
3. 更新训练数据集：将过滤后的高质量伪标签与原始训练数据集合并，形成新的训练数据集
4. 重新训练模型：使用更新后的训练数据集重新训练模型，从而改进模型的性能

Thoughts

在某种程度上，BLIP 统一了几乎所有的图文多模态任务，除了无法做图像生成任务外，其他任务都可以通过 BLIP 的不同模块组合实现
BLIP 的 Bootstrapping 策略是一种自我增强的方式，可以使得模型在无监督的情况下不断提升性能，如今这几乎是多模态预训练算法的标配
从实验的角度告诉我们，encoder 和 decoder 做充分的参数共享是可行的，仅仅在少量的位置上不共享参数，就可以实现完全不同任务的模型

大模型DPO入门

Posted on 2025-01-07 Edited on 2025-03-03 In LLM Valine:

TL;DR

本文是大模型DPO的入门教程，目的是用尽可能简单的方式介绍大模型DPO的基本概念和原理，不涉及过多实现细节。

什么是 DPO

DPO 是 Direct Performance Optimization (直接偏好优化) 的缩写，是 RLHF 的替代方案，目的是使用不涉及强化学习的方法来优化模型性能。

DPO 的数据格式

和 RLHF 一致，DPO 的数据格式为：prompt + response 1 + response 2 + label
- prompt：问题描述
- response 1：模型生成的回答 1
- response 2：模型生成的回答 2
- label：人类标注的哪个回答更好（0 / 1）

DPO 损失函数的设计

L(\theta)=\mathbb E_{(x,y^+,y^-)\sim D}[\log\sigma(s_\theta(x,y^+)-s_\theta(x,y^-))]

s_\theta(x,y)=\log p_\theta(y|x)

p_\theta(y|x)=\prod_{t=1}^T p_\theta(y_t|y_{<t},x)

s_\theta(x,y)=\sum_{t=1}^T \log p_\theta(y_t|y_{<t},x)

$L(\theta)$ ：损失函数
$\theta$ ：模型参数
$D$ ：数据集
$x$ ：输入
$y^+$ ：人类标注的更好的回答
$y^-$ ：人类标注的更差的回答
$s_\theta(x,y)$ ：模型对 $(x,y)$ 的得分
$p_\theta(y|x)$ ：模型对 $x$ 生成 $y$ 的概率
$\sigma$ ：sigmoid 函数
$T$ ：序列长度
$y_{<t}$ ： $y$ 的前 $t-1$ 个元素
$s_\theta(x,y)$ 可经过修改变成大名鼎鼎的 困惑度 (Perplexity, PPL)： $PPL(\theta)=\exp(-\frac{s_\theta(x,y)}{len(y)})$

Thoughts

DPO 事实上就是利用正负样本的困惑度来对比学习，从而优化模型性能
原理非常简单，本质就是用对比学习直接对困惑度进行训练，收敛速度比 RLHF 快
但可以预见的是：DPO 的损失计算方式比 RLHF 更硬，那么必然对数据集的质量要求更高，否则会导致模型性能下降
一个常用的做法是在 SFT 之后，在 RLHF 之前使用 DPO 快速将模型性能提升到一个较高水平，然后再使用 RLHF 进一步优化