Fain的Blog

Femtotron开发日志 #11 监督微调 Supervised Fine-Tuning, SFT

Tue, 19 May 2026 00:03:25 +0800

这是一个相对较小的更新工作量。虽然是minor的feature，但SFT的数据划分模式还是值得总结一下的，毕竟俗话说的好，魔鬼藏在细节中嘛。

当然，这篇不会很长就是了。

SFT和预训练的区别

数据格式

SFT仍然需要把数据划分为若干个规整的段落，只不过不是像预训练一样把无穷无尽的连续文本截断、截断、截断，而是反过来，把来自不同场景的单轮或者多轮对话拼接在一起。对话之间没有因果关系，因此如果多个对话被切分在一起，就需要正确的设置Attention当中的因果mask矩阵。如果不想处理这个问题，naive的方式是让每个序列里一定只有一个场景的连续对话，而seq_len直接等同于这么多连续对话当中最长的那一个。这个就是所谓的Padding方案。这个方案很明显会产生巨大的Padding开销（因为大部分对话不可能有最长的那么长，而且很可能短很多），因此实践当中根本没有人会真的这么去做。另一种是Packing方案，即允许不同场景的对话存在于同一个序列里，通过因果mask遮罩处理它们的非相关性。这是生产当中更经常采用的方案。它的具体实现是一种叫做“首次适应递减算法”（FFD, First-Fit-Decreasing）的经典近似启发式算法：意思就是把对话按照长度降序排列，然后依次取出所有对话，对每一个检查每个现存的seq序列；如果现有序列有能够容纳的，就装入第一个长度允许的序列里；如果所有序列都不能容纳，就增加一个seq作为新的单独sample。

有意思的是，数学上可以证明，FFD算法是非常优秀的近似算法，其使用的序列数绝不会超过最优解的$11/9$倍再加上1。

Loss Mask设置

这是另一个工程细节问题。具体来说，当我们对大模型进行后训练的时候，通常不希望LLM学习到是如何提问的，而是学习到是如何回答的。因此，对于用户提问的部分（以及广义上的其他“不需要学习的部分”），我们会设置它们在计算loss的时候被视为ignore_index（一般默认值是-100），设置loss_mask为false，从而mask掉这些我们不期望的部分中的loss的计算。我们幸运的是，Pytorch对其的支持已经很好，直接复用Attention相关的基础组件支持即可，其能够正确处理各种细枝末节的工程细节，例如重新计算有效Token数量和对loss进行缩放这些麻烦的事情。

学习率

这是另一个需要注意的点：预训练和后训练所需要的学习率并不相同，一般需要比预训练小半个到一个数量级。这某种意义上相当于延续了预训练后期decay的学习率，而不是它的标称值。

如果能够有更多工程的实践经验就好了。我依稀记得，对于不同训练阶段、不同数据量和不同参数量的最佳学习率，学术界已经有了一些相当精彩的理论工作。如果后面有空的话，也会抽空更新一些学习日志，专门研究这些的。

Femtotron开发日志 #10 ZeRO-3、SAC和PP：一个简单的显存节省收益数据测试和分析

Mon, 18 May 2026 15:00:29 +0800

这里列几组数据，和测试它们的时候使用的参数，供大家参考，并且提供一些分析。具体的内存占用太过复杂，因此分析可能有不完全详尽之处，能对读者有所启发即可。

数据

PP=1、TP=1、DP=8 (固定num_layers=8, 变化seq_len)

参数如下：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13


model_config = AutoConfig.for_model(
 "llama",
 hidden_size=1024,
 intermediate_size=2048,
 num_attention_heads=16,
 num_key_value_heads=4,
 num_hidden_layers=8,
 max_position_embeddings=128,
 vocab_size=1024,
 rms_norm_eps=1e-5,
 hidden_act="silu",
 tie_word_embeddings=False,
)

数据如下：

Config	seq=16	seq=32	seq=1024
baseline	1653	1653	3631
baseline + AC	1653	1653	1669
ZeRO-1	460	471	2913
ZeRO-1 + AC	460	460	901
ZeRO-2	443	444	2776
ZeRO-2 + AC	443	444	765
ZeRO-3	384	419	2840
ZeRO-3 + AC	261	261	711

PP=2、TP=2、DP=2 (固定seq=32,变化num_layers)

参数如下：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13


model_config = AutoConfig.for_model(
 "llama",
 hidden_size=2048, 
 intermediate_size=8192,
 num_attention_heads=32, 
 num_key_value_heads=8,
 num_hidden_layers=8, 
 max_position_embeddings=128,
 vocab_size=40960, 
 rms_norm_eps=1e-5,
 hidden_act="silu", 
 tie_word_embeddings=False,
)

数据如下：

Config	N=8(4/stage)	N=32(16/stage)	N=64(32/stage)
ZeRO-0	3601	12558	24404
ZeRO-0 + AC	3601	11954	23091
ZeRO-1	2136	9534	18595
ZeRO-1 + AC	2136	6993	13490
ZeRO-2	2044	9030	17627
ZeRO-2 + AC	2044	6569	12602
ZeRO-3	2233	10808	20829
ZeRO-3 + AC	2216	6417	12018

ZeRO-3 vs ZeRO-2 + SAC 对比

设置	ZeRO-2 + AC	ZeRO-3 + AC	Δ(负 = ZeRO-3 赢)
PP=1, seq=16, N=8	443	261	−182
PP=1, seq=32, N=8	444	261	−183
PP=1, seq=1024, N=8	765	711	−54
PP=2, seq=32, N=8(4/stage)	2044	2216	+172 ✗
PP=2, seq=32, N=32(16/stage)	6569	6417	−152
PP=2, seq=32, N=64(32/stage)	12602	12018	−584

分析

从上面的数据可以得出几个浅显的结论：

1、在DP规模相对较大、TP和PP都不存在的时候，ZeRO-3的优势很明显。
2、在使用PP并行和DP规模相对较小的时候，ZeRO-3的优势几乎不再明显存在，甚至可能有反转的情况。
3、在seq_len较大的时候，ZeRO-3的显存节省会被削弱甚至被ZeRO-2反转，但配合SAC可以弥补这部分额外损失，恢复明显的优势。
4、在模型单层参数增多的时候，ZeRO-3 + SAC的优势更加微弱。在模型的单层参数不变，但是层数变深的时候，ZeRO-3 + SAC的优势更加明显。
5、ZeRO-1和ZeRO-2在各种情况下基本可以提供稳定的性能优化。
6、SAC单独使用的时候，其效果取决于seq_len等输入规模参数的大小。如果这些参数不够大，SAC可能完全没有效果。

它的根源在于ZeRO-3 + SAC的显存节省来源，和改变模型参数、改变并行配置的时候ZeRO-3 + SAC的额外显存增加来源并不相同。ZeRO-3是否真的能取得收益，取决于这样两个数值的此消彼长：ZeRO-3必须保持活跃的"瞬时块"大小，和ZeRO-3总共节省下来的内存大小。前者大，那么ZeRO-3会输给ZeRO-2，反之则超过。节省下来的内存大小理论上相对来说很好计算，就是权重中被shard掉的比例，而必须保持活跃的"瞬时块"大小则和很多因素有关系。目前能够明显观察到的有：1、PP stage的流水线拖延导致的缓冲区分配；2、seq_len拉大导致的缓冲区分配；3、单个layer的unshard过程造成的临时峰值。当然，这些分析是很粗浅的，不过基本已经足以指导我们对这些优化策略的使用与否。我非常相信这当中还有更深的奥妙有待研究，这就交给有兴趣的人进行进一步的分析了。

Femtotron开发日志 #9 流水线并行 Pipeline Parallelism

Sun, 17 May 2026 23:26:30 +0800

这几天的工作量是实现了3P并行当中的最后一环：流水线并行，Pipeline Parallelism。这个并行理论上在Pico-vLLM的推理场景当中也同样有意义，但却碍于工程量问题未能真实实现，因此在Femtotron中对它进行真正的实现也算是弥补了遗憾。这次的博客更新之所以隔了这么多天，是因为这个并行的工程量意外的大，而且和之前的ZeRO系列一样，具有极多麻烦的边界情况问题。下面讲一讲原理、核心的设计抽象，以及中途遇到的值得一提的问题，以供后来人和读者参考，希望有所帮助。

流水线并行的原理

应该没有人不能理解PP的原理是什么。它在概念上极其简单：在tokenizer-若干个layer-lm head这个前向传播的维度上，按某种原则进行切分，将切分的结果依次分配到不同的GPU上，并且在它们之间协调一组调度和通信内容、顺序的过程。直观上，就是把模型按纵向维度切分，一段段的放在GPU的rank0,1,2...上，依次往前做，反向传播的时候再反过来。

问题是它的实现同样很麻烦。下面会详细讲讲这些问题。

流水线并行的核心设计架构和抽象

流水线并行的核心问题在于处理不同stage之间的次要异质性。对于中间阶段，无论如何，输入和输出都是一样的，完全同质化。然而，对于第一个stage和最后一个stage就不一样了。对于第一个stage来说，其输入不是Hidden state而是word Embedding的向量；对于最后一个stage来说区分更明显，要多做一个lm head，而且随着场景变化可能要物化一个巨大的vocab table，导致额外的计算量和内存峰值。在主流实践当中，如果有此类需求，我们通常会微调layer数量的分配，在最后一个stage上有意减少普通layer的数量，以达到更好的负载均衡。我们的框架为此专门实现了可调节每个stage layer数量的可配置方案；不过为了使用方便起见，我们大多数时候会采用默认的均匀配置。

流水线模型模块

和前面的惯用伎俩一样，就是用一个warpper把切分的layer们包裹起来，异质性的复杂度也在这一层处理。首先是warpper内部包裹的东西：

 1
 2
 3
 4
 5
 6
 7
 8
 9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119


class LlamaPartialModel(nn.Module):
 def __init__(
 self,
 model_config: LlamaConfig,
 parallel_ctx: ParallelContext,
 layer_range: range | None = None,
 is_first: bool | None = None,
 is_last: bool | None = None, # None = 从 layer_range 自动推导
 ):
 """
 Args:
 model_config: HF LlamaConfig
 parallel_ctx: 并行上下文。目前仅存储,未在 forward 中使用;
 预留给未来需要 PP/TP-aware 行为的扩展
 layer_range: 本 stage 持有的 layer indices(全局编号,
 就是 LlamaDecoderLayer.layer_idx 的值)。
 None 表示持有全部 layers
 is_first: 是否第一个 stage(持有 embed_tokens)
 is_last: 是否最后一个 stage(持有 final norm)
 
 Raises:
 ValueError: layer_range 越界,或 layer_range 不是 range 类型
 """
 super().__init__()
 
 if is_first is None:
 is_first = (layer_range.start == 0)
 if is_last is None:
 is_last = (layer_range.stop == model_config.num_hidden_layers)
 
 if getattr(model_config, "attn_implementation", None) is None:
 impl = getattr(model_config, "attn_implementation", None) or "sdpa"
 model_config._attn_implementation = impl
 
 self.config = model_config
 self.parallel_ctx = parallel_ctx
 self.layer_range = layer_range
 self.is_first = is_first
 self.is_last = is_last
 
 # First stage: embed_tokens
 if is_first:
 self.embed_tokens = nn.Embedding(
 model_config.vocab_size,
 model_config.hidden_size,
 padding_idx=model_config.pad_token_id,
 )
 
 self.layers = nn.ModuleDict({
 str(idx): LlamaDecoderLayer(model_config, layer_idx=idx)
 for idx in layer_range
 })
 
 self.rotary_emb = LlamaRotaryEmbedding(config=model_config)
 
 # Last stage: final RMSNorm
 if is_last:
 self.norm = LlamaRMSNorm(
 model_config.hidden_size,
 eps=model_config.rms_norm_eps,
 )
 
 def forward(
 self,
 x: torch.Tensor,
 attention_mask: torch.Tensor | None = None,
 position_ids: torch.Tensor | None = None,
 ) -> torch.Tensor:
 """
 Args:
 x:
 - is_first=True: input_ids,LongTensor[B, S]
 - is_first=False: hidden_states,float[B, S, H]
 attention_mask: 默认 None
 position_ids: 默认 None
 
 Returns:
 hidden_states float[B, S, H]
 - 非 last stage: layers 输出(未经 final norm)
 - last stage: 经过 final norm 的 hidden_states
 """
 # Embed (first stage only)
 if self.is_first:
 hidden_states = self.embed_tokens(x)
 else:
 hidden_states = x
 
 bsz, seqlen, _ = hidden_states.shape
 
 if position_ids is None:
 position_ids = torch.arange(
 seqlen, device=hidden_states.device, dtype=torch.long,
 ).unsqueeze(0)
 
 from transformers.masking_utils import create_causal_mask
 causal_mask = create_causal_mask(
 config=self.config,
 inputs_embeds=hidden_states,
 attention_mask=attention_mask,
 past_key_values=None,
 position_ids=position_ids,
 )
 
 position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids)
 
 for layer in self.layers.values():
 hidden_states = layer(
 hidden_states,
 attention_mask=causal_mask,
 position_embeddings=position_embeddings,
 position_ids=position_ids,
 past_key_values=None,
 use_cache=False,
 )
 
 if self.is_last:
 hidden_states = self.norm(hidden_states)
 
 return hidden_states

然后是warpper本身：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30


class LlamaForCausalLM(BaseCausalLMPipeline):
 def __init__(
 self,
 config: LlamaConfig,
 parallel_ctx: ParallelContext,
 layer_range: range | None = None,
 ):
 super().__init__()
 self.config = config
 self.hidden_size = config.hidden_size
 
 self.model = LlamaPartialModel(
 model_config=config,
 parallel_ctx=parallel_ctx,
 layer_range=layer_range,
 )
 
 self.is_first = self.model.is_first
 self.is_last = self.model.is_last
 
 if self.is_last:
 self.lm_head = nn.Linear(
 config.hidden_size, config.vocab_size, bias=False,
 )
 
 # 暂时不支持tie_word_embeddings
 if config.tie_word_embeddings:
 raise NotImplementedError(
 "LlamaForCausalLM 目前不支持 tie_word_embeddings=True"
 )

可以看到，实现本身还是比较直白简洁的。复杂性和前面遇到的各类杂七杂八问题一样：还是在边界情况里。

调度模块

流水线并行和调度必须一起做，缺一不可。这里我们实现了baseline的all then all调度（先全部前向传播，然后再全部反向传播）作为基线，然后实现了基础的1F1B调度（通过交错排列前向传播和反向传播来降低空泡率）。未来可能还会增加更多的调度策略，这也是在规划内的。它们的原理在开源资料里已经得到了很详细的普及了，因此这里不详细展开。它们的代码如下：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29


def gpipe_schedule(
 num_microbatches: int,
 is_first: bool,
 is_last: bool,
) -> list[PPAction]:
 if num_microbatches < 1:
 raise ValueError(f"num_microbatches must be >= 1, got {num_microbatches}")

 actions: list[PPAction] = []

 # ── Forward phase: mb 0, 1, ..., N-1 ──
 for mb in range(num_microbatches):
 if not is_first:
 actions.append(RecvForward(mb)) # recv hidden from prev stage
 actions.append(Forward(mb)) # forward through model
 if not is_last:
 actions.append(SendForward(mb)) # send hidden to next stage

 # ── Backward phase: mb N-1, N-2, ..., 0 ──
 # Reversed order matches LIFO of the autograd graph each mb's activations
 # are released after its backward, allowing early memory reuse.
 for mb in reversed(range(num_microbatches)):
 if not is_last:
 actions.append(RecvBackward(mb)) # recv grad_output from next stage
 actions.append(Backward(mb)) # backward, accumulate param grads
 if not is_first:
 actions.append(SendBackward(mb)) # send grad_input to prev stage

 return actions

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73


def one_f_one_b_schedule(
 num_microbatches: int,
 pp_size: int,
 pp_rank: int,
) -> list[PPAction]:
 if num_microbatches < 1:
 raise ValueError(f"num_microbatches must be >= 1, got {num_microbatches}")
 if pp_size < 1:
 raise ValueError(f"pp_size must be >= 1, got {pp_size}")
 if pp_rank < 0 or pp_rank >= pp_size:
 raise ValueError(f"pp_rank must be in [0, {pp_size}), got {pp_rank}")

 is_first = (pp_rank == 0)
 is_last = (pp_rank == pp_size - 1)

 num_warmup = min(pp_size - 1 - pp_rank, num_microbatches)
 num_steady = num_microbatches - num_warmup
 num_cooldown = num_warmup # by symmetry

 actions: list[PPAction] = []

 # ── Phase 1: Warm-up forwards (fill the pipeline) ──
 for j in range(num_warmup):
 if not is_first:
 actions.append(RecvForward(mb_id=j))
 actions.append(Forward(mb_id=j))
 if not is_last:
 actions.append(SendForward(mb_id=j))

 # ── Phase 2: Steady-state 1F1B ──
 for k in range(num_steady):
 fwd_mb = num_warmup + k
 bwd_mb = k
 is_first_steady = (k == 0)
 is_last_steady = (k == num_steady - 1)

 # F(fwd_mb): need its input
 # - first_steady & not is_first: warmup did RF(0..warmup-1),
 # so explicit RF(fwd_mb) here
 # - subsequent: input arrived via prev iter's SBRF
 # - is_first: input from caller dict, no recv ever needed
 if is_first_steady and not is_first:
 actions.append(RecvForward(mb_id=fwd_mb))
 actions.append(Forward(mb_id=fwd_mb))

 # Send F output forward + recv B grad backward (combined to avoid deadlock)
 if not is_last:
 actions.append(SendForwardRecvBackward(fwd_mb=fwd_mb, bwd_mb=bwd_mb))
 # else: last stage — no SF, no RB (loss provides grad locally)

 actions.append(Backward(mb_id=bwd_mb))

 # Send B grad backward + recv next F input (combined; or plain SB at end)
 if not is_first:
 if is_last_steady:
 # No more F to recv (cooldown only does B's)
 actions.append(SendBackward(mb_id=bwd_mb))
 else:
 actions.append(SendBackwardRecvForward(
 bwd_mb=bwd_mb, fwd_mb=fwd_mb + 1,
 ))
 # else: first stage — no SB, no need for RF (input always from caller)

 # ── Phase 3: Cool-down backwards (drain the pipeline) ──
 for j in range(num_cooldown):
 bwd_mb = num_steady + j
 if not is_last:
 actions.append(RecvBackward(mb_id=bwd_mb))
 actions.append(Backward(mb_id=bwd_mb))
 if not is_first:
 actions.append(SendBackward(mb_id=bwd_mb))

 return actions

值得注意的Bug

Rotary inv_freq的garbage初始化问题

meta → to_empty(device) 路径下,persistent=False 的buffer不会被初始化,留下未定义内存。Llama 的 rotary embedding 的inv_freq 就是这种buffer。

这导致模型forward出来的结果是garbage logits。关键是，由于它对于整个模型的精度影响平均下来也就~100到300个ULP上下，差点又一次没有发现，幸好仔细的人工校对了一遍。具体来说，是在Debug的时候，发现两个 rank 上的 model.model.rotary_emb.inv_freq 值不一样（随机初始化导致的垃圾值不相同），从而发现了问题。

解决方案是加了一个 _reset_rotary_inv_freq(rotary_emb, config, device) helper，在 test / random-init 路径下显式重算。生产路径下的 ModelLoader 已经遇到过一次这个问题并且解决过一遍，这次算是踩了重复的坑，这下记的更牢了。实际上，模型的莫名其妙的问题，以到现在的经验来看，几乎总是和buffer、cache，这些不属于模型参数但是又需要持久化的东西有关，也许确实是一个值得思考的整体模式问题。

1
2
3
4
5
6
7
8


def _reset_rotary_inv_freq(rotary_emb, config, device):
 base = getattr(config, "rope_theta", 10000.0)
 dim = (getattr(config, "head_dim", None) or
 (config.hidden_size // config.num_attention_heads))
 inv_freq = 1.0 / (base ** (
 torch.arange(0, dim, 2, dtype=torch.int64).to(device, torch.float) / dim
 ))
 rotary_emb.inv_freq.copy_(inv_freq.to(rotary_emb.inv_freq.dtype))

这个过程带给我们的教训是，任何用 meta → to_empty 的路径都要审计buffer是否被显式初始化。这很难依赖某种自动化规则，只能靠coder本人的实践经验和敏感性。

一个不算bug的奇怪问题

cuBLAS lazy context warning总是消除不掉。尝试了很多方法无果，遂放弃。先记录在这里，等到哪天遇到类似问题或者看到解决方案了，再来捣鼓它。

UserWarning: Attempting to run cuBLAS, but there was no current CUDA context!

ZeRO-3 + PP时会发生的显存占用和计算效率问题

虽然这听上去很奇怪，但这是真的。ZeRO-3 + PP加在一起不如两个都不加。在我的测试当中，ZeRO-3在很多情况下，不仅没有相对于ZeRO-2降低显存的占用，反而增加了其占用情况。在PP=1的情况下，这种情况尚且并不多见；在PP=2及以上，即采用PP并行的情况下，它几乎总是比ZeRO-2还要差。是SAC也拯救不了的那种，纯粹的占用了更多无法被释放的显存。

经过分析和多个配置参数的验证，这确实不是一个正确性问题，而是无法避免的机制冲突。也就是说：ZeRO-3 + PP是根本上不兼容的两个训练优化/并行化机制。个人认为这是一个挺反直觉（至少不能第一时间通过纸面理论察觉到）的结果，而且某种程度上令人感到沮丧：因为ZeRO-3本身的复杂度就奇高无比，而且我自己实现的时候也花了很多心血，但现在却发现它作为增量的收益远不如我们预期的多。不过无论如何，抛开这些投入成本不谈，我仍然查询了学界和工业界对此的理解和处理，并且确实得到了很多有意思的信息和结论、方案。下面讲讲我了解到的相关原理、机制解释，以及学术界、工业界对它们的处理方法。

autograd-held view现象

这是一个我发现了但是决定不解决的问题。AI在这个问题的理论发掘过程中居功至伟，我个人认为比较有说服力，不过也因此，读者在使用这个结论的时候最好也带有自己的思考。它部分贡献了ZeRO-3在PP情况下的显存异常增加，但解决它可能造成潜在的破坏性影响。它的核心机制代码出现在ZeRO-3的unshard操作的下面这个地方：

1
2
3
4
5
6
7


def unshard(self):
 self._full_buffer = torch.empty(self.padded_size, ...)
 dist.all_gather_into_tensor(self._full_buffer, self.flat_param_shard, ...)
 
 for pg, layout in zip(...):
 slice_ = self._full_buffer[...]
 pg.compute.data = slice_.view(layout.original_shape) # ← view

forward中，y = x @ pg.compute，autograd会在saved tape里存pg.compute这个tensor的 view，以为了backward计算dL/dx。这个view的storage就是_full_buffer的storage。

1
2
3
4


def reshard(self):
 for pg in self.param_groups:
 pg.compute.data = torch.empty(0, ...)
 self._full_buffer = None # ← 我们的引用没了,但 autograd 还引用着

1F1B的情况下，有很多个microbatch在in-flight，因此多份 _full_buffer 同时活着占用显存无法释放，直到对应mb的backward完成。

SAC能部分解决这个问题，因为SAC在backward时会直接重跑 forward 重算 activation，这次重算的unshard是short-lived 的,重算完直接reshard释放,autograd把它当作普通临时tensor而不是saved tensor，因此不持续占用显存，无论多少个microbatch都不会显著堆积。 ZeRO-3 + AC的节省量通常大幅超过单纯"AC 省 activation"的节省量就是这个原因，多出来的部分就是消除autograd-held view的收益。

ZeRO-1/2没有这个问题，因为ZeRO-1/2不分片param,compute.data 全程是完整的 bf16 权重，因此不需要unshard buffer。

通信组（Communication Group）划分导致显存的不降反增

在引入PP后，每一台机器（或每一张卡）只负责模型的一部分层。如果再引入ZeRO-3，那么理想状态下，单卡只存8层参数的$1/N$。但是，在实际上的训练现实当中，因为PP在前向传播时，每张卡都在高速、连续地处理不同的 Micro-batch（微批次）。为了不让流水线产生停顿和空泡，ZeRO-3必须为所有正在流水线中流动的Micro-batch预留足够的缓存空间，更激进来说甚至要提前拉取（Prefetch）并缓存完整的参数。这导致我们必须重新占用这些好不容易省下来的显存。为了应付多阶段流水线并发，被ZeRO-3的 All-Gather 缓冲区、预取缓冲区（Prefetch Buffer）以及临时激活值很可能（实际上是几乎一定）最终被撑得比单纯用PP还要大！

这是一个从数学理论上无解的问题。因为如果不开这个缓冲区，那么通信就要序列化。如果通信序列化，那么流水线就会有空泡。如果流水线有空泡，那么整个训练的MFU就会下降。众所周知，MFU是一个比显存占用还要关键的指标，等同于公司每分每秒的真金白银。因此，二者不可兼得，只能不放在一起做了。

工业界的解决方案

目前业界主流的训练架构（Meta、微软、英伟达、阿里...等等）在工程上的演化结果基本上没有同时采用两者的方案，而是以其中一种为主导发展出了两条完全平行的技术路线：

路线一（NVIDIA / Megatron 派系）：纯粹的 3D 并行 + ZeRO-1/2

**架构：**Tensor Parallel (TP) + Pipeline Parallel (PP) + Distributed Optimizer (ZeRO-2)。

这一路线的核心逻辑是，只要用了 PP，就绝对不用任何参数分片（即不用 ZeRO-3）。参数的纵向拆分靠PP，横向拆分靠机器内的TP（NVLink的高速通信）。数据并行组只用ZeRO-2来切分优化器状态。这种组合基本上可以被认为是目前最稳定、大厂千卡乃至万卡的大规模严肃预训练场景中，MFU（算力利用率）最高的方案。

路线二（PyTorch 原生 / Hugging Face 派系）： FSDP 代替 ZeRO-3 + 取代 PP

**架构：**Fully Sharded Data Parallel (FSDP)。

这一路线的核心逻辑是，既然ZeRO-3思想（参数全分片）单卡和多卡能省显存，那就彻底抛弃PP（流水线并行）。现代团队发现，用FSDP配合 Activation Checkpointing，在8卡或16卡环境里，可以轻松塞下70B甚至更大的模型进行全量微调，完全不需要开PP。因为 PP 带来了烦人的流水线气泡（Bubble）和通信等待，而纯FSDP的通信是可以和前向计算完美重叠（Overlap）的。这种组合更多被中型到小型的大模型训练团队采用为高度方便的训练方案。

逸闻

这就是查询参考方案性能的时候查到的一件事情，前面也多多少少提了一句。其实，在Deepspeed的方案中，ZeRO-2和ZeRO-3与Pipeline Parallelism的实现是不兼容的！这是我好不容易把正确性调对、绞尽脑汁也没法解决ZeRO-3性能问题之后，迫于无奈尝试参考官方实现的时候震惊地得知的结果。

不过，实测下来，ZeRO-2+PP仍然是相对来说有实际意义的，测试的数据也支持这一个结论。换言之，我们的方案已经在feature上超过了Deepspeed（哈）。当然，这很大程度上是因为我们的框架本质上是一个个人项目框架，而不是需要处理所有边界情况的工业级项目；但即使如此，这也意味着我们的开发过程并不是毫无作用的。探索这些可行和不可行的边界，正是好的coding项目应该做的事情。如果读者对这部分感兴趣，源代码已经开源在仓库里，有兴趣的读者可以去阅读。

总结

到这一步为止，我们已经完成了除了SFT训练支持之外的所有核心feature组件的实现。接下来是SFT，它的复杂度应该会比我们已经走过的路程低很多。后面的很多调度方案等等组件，也是在现有框架架构的基础上进行扩展，而不是进行颠覆式的重构和侵入式修改。无论如何，这已经是一个令人兴奋和欣慰的进展，而且Femtotron作为一个预训练框架已经来到了。集中开发的工作仍然在继续，可能会有一篇阶段式里程碑的博客，总结我们已经具有的feature情况。可能还有一篇专门用来做Profiling，为未来可能的参数调优做准备。whatever，道阻且长，慢慢做吧。

Femtotron开发日志 #8 选择性激活检查点 Selective Activation Checkpointing

Fri, 15 May 2026 22:58:58 +0800

今天的工作量实现的是选择性激活检查点（SAC，Selective Activation Checkpointing）。这个组件相当简单，因为Pytorch已经。因此这篇日志会相对来说较短，主要着重于介绍它的概念、机制，以及实现过程中遇到的值得记录的bug，以供读者或者其他的后来人参考，避免踩同样的坑。

SAC是什么

SAC是一种降低训练中内存占用峰值的技术，代价是增加反向传播当中的计算量。它的具体原理是，每隔一段“距离”，保存一个可以进行前向传播的中间激活值，抛弃这个距离中的两个端点之间的所有其他值。当反向传播开始/越过一个检查点的末尾的时候，我们如果遇到的接下来的反向传播路径当中没有现成的中间激活值，就从最近的一个检查点重新向前计算，重新铺好这段反向传播需要的中间激活值路径。

SAC的理论最优情形

假设我们的模型参数由连续均值的$L$层构成，而检查点每隔$N$层设置一个。此时，可以直接计算出，显存的峰值占用分为两个部分：一个是$1/N$个的检查点层，一个是$N$个的正在进行反向传播的段落，两者加起来就是$N+1/N$。这个式子的最小值点大家应该都熟悉，初中数学嘛。因此，如果想要显存占用最小化的话，理论上来说，$N=\sqrt L$是能够让显存峰值占用最小的点。

SAC的主流实践

当然，工程上通常并不会这么做，因为这个理论模型和现实差的很远，而且并不一定符合我们的MFU最大化的目标——SAC过于稀疏将会导致密集的重计算（recomputation），从而带来额外的计算负担，拖慢训练。SAC的主流实践一般以层为最大粒度：每层保留结构位置相同的Checkpoint，具体保留哪些则由具体策略决定。

Pytorch提供的现有机制

SAC的实现意外的很简单，因为它在实现上基本就是直接复用Pytorch现有的基础设施组件。具体而言，它使用类似以下的代码进行SAC的设置和激活：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43


class ActivationCheckpointWrapper(nn.Module):
 """将一个 module 包装为 activation checkpoint。
 
 Forward 调用通过 checkpoint function 转发,中间 activation 不保留;
 backward 时,该 unit 的 forward 会被重做一次以重建 saved tensors。
 
 Attributes:
 inner_module: 被包装的原始 module
 checkpoint_fn: 实际的 checkpoint 实现
 use_reentrant: 传递给 checkpoint_fn 的 flag
 preserve_rng_state: 传递给 checkpoint_fn 的 flag
 """
 
 def __init__(
 self,
 module: nn.Module,
 *,
 checkpoint_fn: CheckpointFn = _torch_checkpoint,
 use_reentrant: bool = False,
 preserve_rng_state: bool = False,
 debug: bool = False,
 ) -> None:
 super().__init__()
 setattr(self, _WRAPPED_MODULE_KEY, module)
 
 self.checkpoint_fn = checkpoint_fn
 self.use_reentrant = use_reentrant
 self.preserve_rng_state = preserve_rng_state
 self.debug = debug
 
 self._register_state_dict_hook(_post_state_dict_hook)
 self._register_load_state_dict_pre_hook(_pre_load_state_dict_hook)
 
 def forward(self, *args: Any, **kwargs: Any) -> Any:
 return self.checkpoint_fn(
 self.inner_module,
 *args,
 use_reentrant=self.use_reentrant,
 preserve_rng_state=self.preserve_rng_state,
 debug=self.debug,
 **kwargs,
 )
 

就是这么简单，核心代码其实就是利用_torch_checkpoint进行一次前向传播时候的调用包装。当然，用户也可以自定义这个warpper，但大多数时候利用基础设施当中的现有组件就已经完全够用了。

值得注意的Bug和问题

SAC的bug

在这个过程中，遇到了一个特别棘手的问题：Pytorch反复报错。当第一次前向传播的时候很正常，接下来反向传播回去的时候却报错，而报错信息显示多了两个不知道是什么的tensor，整个ckpt的tensor数量从38个变成了40个，对不上。这个bug的排查过程非常艰难，几乎花了我3~4个小时的时间，才定位和理解了问题发生的具体原理。

在排查之后发现，发现这其实不是我们本身的代码实现错误，而是Pytorch后端调用的自动切换导致的。实际上，这就是HuggingFace著名的 use_cache + gradient checkpointing 互斥问题。

它的具体的引起原因来源于LlamaModel的默认设置，config.use_cache=True，也即自动保存和管理KV cache，以加速前向传播。第一次forward时，cache为空，因此调用了 _scaled_dot_product_flash_attention(因为它满足FlashAttention的shape要求)。 recompute时，cache已经被填充，因此K/V的seq长度翻倍，不再是空tensor，因此数量就对不上，同时不满足FlashAttention的causal mask约束，dispatcher fallback到math backend，因此具体的后端实现也发生了变化。这两个因素加起来，导致报错。

它其实原本是有一个内置防御机制，避免这个问题的，类似这样的形式：

1
2
3
4
5
6
7
8


# transformers/models/llama/modeling_llama.py
if self.gradient_checkpointing and self.training:
 if use_cache:
 logger.warning_once(
 "`use_cache=True` is incompatible with gradient checkpointing. "
 "Setting `use_cache=False`..."
 )
 use_cache = False

但只有调用 model.gradient_checkpointing_enable() 或者把 model.gradient_checkpointing = True 设上时这条防御才生效。它的判据是 self.gradient_checkpointing 这个flag，而不是探测调用栈里有没有checkpoint，因此失效。这确实是一个值得记录下来的教训，在部分采用现有基础设施来构建框架的时候，一定要注意这种细节问题，即内部耦合的组件的边界恰好被切分开的时候，所暴露出的协调失效问题。

SAC能带来多少收益

来看图：

亲爱的读者们，如果你们无论怎么测试，得到的都是类似前者的第一张图的结果，你们会不会质疑自己的代码实现有问题呢？至少我是会的：因为内存占用是bit-exact的完全没有变化，看上去更像是SAC根本没有正确加载和起效，而不是它在机制上是真的没有用。我反复排查了很久是不是实现有问题（包括在最后走投无路开始折磨AI），最后发现其实代码根本没什么问题，确实就是它没有压低显存峰值。通过反复的测试参数，将seq_len从测试用toy-model的16改成32和1024，就得到了第二张图和第三张图内的数据结果，而第三张图看上去就正常多了。

这其实说明一件事：SAC对于seq_len不是非常长的情况，能够产生的收益非常有限。但又能够看到，ZeRO-3对于seq_len非常长的情况，收益同样也有限，甚至还不如ZeRO-2。这说明两者其实应该同时使用，而且在DP较多、PP较少的组合上收益更明显。而且，在通常的toy-model的参数范围内，收益整体的量级并不如传统的博客、资料当中描述的那么明显。这很可能是因为SAC出现的时间早于Flash Attention导致的。后者虽然在scope上是一个对于计算局部性的著名Kernel优化，但无意当中抢走了Selective Activation Checkpointing的饭碗，直接让SAC按标准注意力实现的优化收益标准彻底无效了。

实际上，在Flash Attention被发明之后，SAC的核心收益就变化了：因为Flash Attention的计算流根本不物化那个巨大的完整的注意力矩阵，而SAC最初出现的主要目的就是为了规避这个巨大矩阵带来的显存占用问题，导致它从一个几乎是不做到最好就完全不能进行大模型训练的核心技术，退居为一个只要做的合理就能产生合理收益的技术。它节省的现在主要是超长序列带来的中间激活值的膨胀问题，而不再是一个万能万灵的通用手段。

如果SAC特别慢，可能发生了什么

另一个问题是在测试的时候发现SAC特别慢，到了难以忍受的地步。这个其实是个bug，其来源让人好笑，是调试的时候加上的ActivationCheckpointWrapper.forward 中debug设置导致的：

1
2
3
4
5
6
7


return self.checkpoint_fn(
 self.inner_module,
 *args,
 ...,
 debug=True, # ← 硬编码!
 **kwargs,
)

从机制原理上来说，这个torch.utils.checkpoint.checkpoint(debug=True) 的设置会激活一个 TorchDispatchMode, 给原始forward装上Python级dispatch。这是为了使得每个op能够经过Python一遍，以完整记录元数据，进一步为了recompute时检测不一致，能够给出方便调试的友好错误。这原本是为了。然而，改完了之后却忘了改回来，随后彻底忘记了这件事，搞得我在之后做集成测试的时候奇慢无比，浪费了很多时间（起码多花了一个多小时，多花了五六十块钱，呜呜呜）。

如果大家遇到类似表现的问题，记得排查这个设置有没有改对。

Femtotron开发日志 #7 ZeRO-3模式：抽象设计、Bug排查和教训总结

Wed, 13 May 2026 23:41:44 +0800

今天的工作量实现的是我个人认为到目前以来最难的一个组件：ZeRO-3。在整个开发的过程中，不仅需要修改的代码量很大、不可避免的产生了许多侵入式的修改，而且即使设计已经相当小心，还是产生了许多复杂的bug。我们首先回顾一下不同ZeRO等级的切分情况，来为下面的总结做铺垫。

	权重	梯度	优化器状态
ZeRO-1	完整	完整	切分
ZeRO-2	完整	切分	切分
ZeRO-3	切分	切分	切分

可以看到，ZeRO-3的增量在于切分了权重。

ZeRO-3的分片和TP的区别是什么

这是我的第一反应，可能也是很多人第一次学到ZeRO-3的时候会觉得，“这不就是DP版本的TP吗？”，都是权重切分、都需要相当heavy的通信。然后，如果有人做过推理框架或者训练框架的TP并行的话，就会知道，TP并行是很简单的。因此，一个很自然的想法是：ZeRO-3是不是也会很简单？并非如此。实际上，两者除了都具有这两个特征之外也就没有什么区别了。和TP截然相反，ZeRO-3的实现复杂度是非常高的。让我们来分析一下。

通信的对象是问题

这是最大的区别。对于TP并行来说，通信的永远是数据本身。无论何时，权重都是分片的、静态的，在通信前后改变的是数据。进一步的说，通信改变的是数据的内容（值），而不是形状，涉及到的通信操作是All-Reduce，而不是All-gather等等。从头到尾，占据的内存峰值不会有任何变化，因为张量的形状不会有任何变化。这是TP并行。

但对于依赖于DP的ZeRO-3来说，情况就完全不一样了。它通信的对象是权重参数。在前向/反向过程中，数据不动，动的是参数本身。这就意味着参数的大小和形状都是频繁变化的，而大多数时候一个DP rank的概念视图和实际持有的参数并不一样。这在概念上并不阻碍原理的理解，但引入了很大的工程复杂性。

ZeRO-3麻烦的地方

ZeRO-3真实的工程实现模式

ZeRO-3的工程实现并不如它的论文和概念上那么优美。论文上的实现是这么描绘的：将“每个参数的权重进行分片，均匀的保存在各个不同的DP rank上”，然后“在前向传播和反向传播需要这个参数的时候，进行相应参数的unshard，用完后立刻reshard，避免内存峰值”。这在概念上是非常简洁优美的，但实现上立刻会遇到问题：每个参数单独来看是不够大的，但每次通信启动都需要固定的launch overhead。为此，要么直接不对小的tensor进行分片（这会引入相当的复杂度），要么就得把若干个参数打包在一起，在更大的粒度上进行reshard/unshard（这同样会引入另一种复杂度）。在具体工程实现上，一般按照block（layer）为粒度打包。一个layer的参数被flat然后concat在一起，然后集体通信来分片和聚集。这意味着整个正常的基于参数tensor的梯度更新模式都不再适用了，需要重新设计。

当然，这就意味着工程上的实现复杂度会相当高，不过在此不再赘述了。

ZeRO-3单独实现产生的收益

令人沮丧的是，即使工程实现很麻烦，它的收益并不大。这是因为整个训练过程的内存占用变化其实是一个“双峰”的过程，而内存峰值这个单一最大值才是最终决定端到端收益的关键。具体来说，第一个潜在的峰值是在前向传播结束之后，反向传播开始之前，它的主要动态内存占用的来源是激活值；第二个潜在峰值则是反向传播结束之后，优化开始之前，其主要动态内存占用的来源是梯度。在这两者之间，随着反向逐层进行，激活值逐层被释放，而梯度则随着backward逐层被产生。也就是说，在这两个潜在峰值之间的内存占用基本上是线性插值的关系，而两个端点谁更高，谁就主导了全生命周期的内存峰值。权重、优化器状态和其他杂项则是全程存在的。这意味着梯度的总占用和激活值的总占用大小决定了哪个成为瓶颈。

ZeRO-2技术减少的是第二次峰值的规模，而不是第一次的。对于ZeRO-2之后的全周期显存占用来说，瓶颈就已经是激活值了，ZeRO-3继续叠加的边际效应并不是很大。但从另一个角度思考，ZeRO-3的内存占用减少是针对全周期的，和激活值/梯度都无关。

值得注意的Bug，和它们的排查

内存泄露

我一直说，“能够在自己的代码里遇到实际的内存泄露问题，并且亲手解决它”，才是成为真正的合格码农的象征。这次，在今天终于遇到了。不过，自己尝试过才知道，这个寻找的过程是相当麻烦的。

问题出现在刚实现完ZeRO-3，跑通测试用例的时候，我发现ZeRO-3相对于ZeRO-2的内存峰值不降反升，多了整整~230MB。这奇怪极了，遂开始排查原因。这个问题显然网上没什么直接结论。询问ai，ai一开始的解释是“可能nccl等等后端占据了更大的缓存”，但这并不令人信服。随后，我开始自行排查。我改变了测试脚本的测试范围，对每一个单独进行测试，发现这个内存峰值的现象消失了，ZeRO-3的内存峰值下降到低于ZeRO-2的水平。准确的说，只有在存在ZeRO-2的情况下，ZeRO-3会出问题。于是怀疑是ZeRO-2的处理不干净，出现了内存泄露。尝试使用Pytorch的释放和gc的释放，无果。在进行了exhausting的大量排查之后，终于定位到，是ZeRO-2和3使用到的hook注册机制造成的隐式循环引用，无法被gc回收，于是永久性泄露。它的具体机制如下：

ZeRO-2的源代码中：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28



 def _register_hook(self, group: ParamGroup):
 spec = group.master_spec
 
 def hook(param):
 if not self._sync_enabled:
 return # no_sync 期间放过
 if param.grad is None:
 return
 assert spec is not None, "ZeRO-2 hook 只能注册在分片了的 param 上。"
 
 flat = param.grad.flatten()
 if spec.pad_size > 0:
 padding = torch.zeros(spec.pad_size, dtype=flat.dtype, device=flat.device)
 flat = torch.cat([flat, padding])
 
 shard = torch.empty(spec.shard_size, dtype=flat.dtype, device=flat.device)
 dist.reduce_scatter_tensor(
 shard, flat, op=dist.ReduceOp.AVG, group=self.dp_group
 )
 
 self._grad_shards[group.name] = shard
 param.grad = None # 释放 compute.grad,这是 ZeRO-2 省显存的关键
 
 handle = group.compute.register_post_accumulate_grad_hook(hook)
 self._hook_handles.append(handle)
 # group.compute.register_post_accumulate_grad_hook(hook)
 

使用了hook。因此hook依赖strategy，而strategy又有groups_ref，groups_ref有ParamGroup，ParamGroup有model.parameter，model.parameter被hook注册了，有hook。这就循环依赖了。如图：

model.parameter ──→ hook closure ──→ strategy
 ↑ │
 └──── ParamGroup ←── groups_ref ───┘

hook+闭包+循环引用的联合造成的Python gc盲区。register_*_hook把closure存在tensor/module的C++内部结构里。Python的gc不能traverse这些C++边，因此循环引用永远断不开。del strategy 和 gc.collect() 都只是把"显式"引用降 1，这条环上的hook边不动。

修复方法倒也简单，需要让每个 register_*_hook 都返回 RemovableHandle，strategy显式持有这些引用，在 cleanup() 里全部 .remove()，然后再进入正常的释放流程。这是个任何架构都救不了的问题，只能靠纪律。因为其形成循环引用的隐蔽性，一个教训是，每个注册了hook的地方都必须显式管理它们的handle的生命周期。

噪声扰动还是Bug

另一个bug来源于跑通之前。在刚刚跑通的时候，我发现这次的误差比测试dp-tp切换正确性的时候大，而且大了整整1~2个数量级。ai对此给出的解释依然同样是“这是噪声”，但我认为不应该是这个原因。于是，我增大了模型参数，重新测试了一遍，发现误差同样扩大了相同的比例。我意识到，这不可能是噪声。然后，我修改了测试的方式。我使用完全相同的随机数据进行训练，在每一步之间不进行任何更改。接着，我发现误差不再呈现随机性：loss缩小的速度比baseline高，随着每一步快速扩大。这证实了存在一个bug。

随后的排查发现，是一部分没有被打包的孤立参数（非layer的参数）没有被正确的处理，fallback到了默认路径，因此在DP下行为异常，于是静默报错。

这是我遇到的第一个静默报错。实际上，识别出它完全靠我本人的经验直觉，而不是某种定量的测试——阈值测试在这个情况下失败了。这个故事实际上能给出一个很好的启示，就是真正的对抗性的测试设计的必要性。普通的测试用例本就应该是“充满恶意的”，而不是“我给出正常工作的条件和宽松的通过判定标准，结果看上去差不多就允许放过”。进一步的说，即使有了测试用例，没有人自己的检查，通常还是会遗漏一定数量的bug，因为bug并不总是以你预期的方式显露其特征，它有时候完全不以测试者预期的方式甚至能够定量判定的方式暴露。

更进一步的说，让自己的大脑多经验类似问题，积累相关的模式识别的经验确实是很重要的。

无法实现的Bit-Exact

bug修复了之后，扰动变成了彻底的噪声，但是仍然不是0。看上去很强迫症不友好，但这其实是没有办法的。实际上，Bit-Exact对于前两者成立只是偶然，如果TP-DP配置不同，结果也会不同，同样无法Bit-Exact。但ZeRO-3为什么DP内部也不行呢？笔者个人能够想到的原因主要是通信带来的问题：前两者并没有打破“参数”的概念边界，而后者打破了，把不同的“参数”变成了纯粹的“数据”或者“比特流”，造成了通信的必然不一致。

ZeRO-1/2的分片性质，本质上是master按参数自己的边界切。Rank 0持有"q_proj 前一半 + k_proj 前一半 + ...",Rank 1持有"后一半 + 后一半 + ..."。两个rank处理的是同一组参数(每个参数都参与)，只是各自处理这个参数的不同元素。这意味着每个参数的sq_sum、grad reduction等等操作，以及通信前后更新的操作，在rank内的累加顺序和baseline完全对称，且ZeRO-1和ZeRO-2之间，的通信拓扑和对象并没有本质的改变，因此维持了一致性。

ZeRO-3的分片性质则是整个一个block（或者说layer）变成一个cluster，整个layer的所有参数合在一起，把（这里的架构是9个）整个9个参数flatten+concat后再切。Rank 0持有的是“q_proj 全 + k_proj 全 + v_proj 全 + o_proj 前段”,Rank 1持有“o_proj 后段 + gate_proj 全 + ...”，以此类推。这意味着，两个不同的DP rank处理的是不同的参数集合。在对这个改变的对象进行通信，或者进行clip等后操作时，reduction树拓扑改了。fp32 加法不满足结合律，这就意味着最后一两位ULP必然不同。不过，这个量级的噪声并不会真的破坏模型的正常收敛，因此大可放心的使用。

Femtotron开发日志 #6 ZeRO-1和ZeRO-2模式：抽象设计、工程实现和占位符

Mon, 11 May 2026 16:33:30 +0800

这次的工作量主要在于实现了ZeRO-1和ZeRO-2。老实说，这部分确实下了很多功夫：因为ZeRO是一个需要分段实现，但是抽象概念上有许多共通之处的，相当大的核心feature组件。如果仅仅是硬编码进去，后续的开发就会不断的制造屎山，侵入性的修改相同地方的代码，导致不可维护性（Pico-vLLM的PD分离的时候我就是这么干的.jpg）。在查询了资料（而且在使用ai反复进行框架验证）之后，最终采取了一个完整的插入性组件系统的设计。它的核心思想是“接管”，即利用预制的系统内接口的预留、默认的No-Op占位符的设计、功能性组件和工厂类四者的协调，实现可扩展的整个ZeRO系统。下面讲讲我认为有价值的细节。

ZeRO-1的细节

即使在Python也要坚持类型安全

虽然没什么特别高深的理论，但在这里真的特别值得强调一下。虽然Python是一个弱类型语言，但在一个框架当中，坚持类型安全的写法是极其、极其、极其重要的！它可不只是你面前的红色波浪线是多还是少的问题，真的不要低估它带来的作用。第一，它可以识别极多的低级错误，作为最基础的提示工具，帮助开发者回忆和建立正确的上下文，其功能性不亚于自动补全工具。第二，即使没有错误，它也可以极大的降低心智负担，从而允许我们以更低的心智劳动成本建立对框架的整体认知、回忆（或者了解）先前开发者的设计意图，避免破坏现有的设计模式。在Pico-vLLM当中，由于急着开发，很多地方并没有这么做、凭借记忆和直觉写代码，从而导致后面再次要修改代码的时候，经常出现无意间破坏了先前设计的核心不变量和默契的情况，从而极大的增加了维护工作量。切记：如果能不加入none，坚决不要加入；不管有没有如果，每个关键参数的签名都要完整写出。defensive有时候是好的，有时候是让代码变得不可维护。

这种态度的代价不会在写代码时显现——会在改代码时显现。可以这样举例子：

看到 Tensor（不带 | None），就知道"原设计者认为这里永远不该是 None"
想改成 Tensor | None 时，会先停下来想"为什么原来不允许 None，现在为什么需要"
这个停下来想的瞬间就是类型注解的全部价值

ShardingStrategy 作为一等抽象

往代码里加ZeRO-1时，大家的第一反应可能是"找到 mp_manager.step，把里面的all-reduce改成reduce-scatter"。但这种"修改实现"的路径意味着每加一种sharding模式都要改step代码，前面写得好好的抽象就全完蛋了。众所周知，所有地方都做的很好，一个地方漏了，复杂度就会回到这个最差组件的水平上。这同样是不可持续的反模式。

实际上，ZeRO横跨了precision和sync两个职责。它既决定master怎么分片（属于precision），又决定grad怎么同步（属于sync）。放在哪个模块都会产生很糟糕的耦合。因此，它应该是一个独立的横切概念，这样就自然导向一个决策，即新增一个sharding模块，作为插件设计，采用“接管功能”的思想去设计。

它的基本抽象如下：ShardingStrategy 协议有四个方法 make_master / reduce_grads / gather_weights / post_step。它们刚好对应ZeRO 在训练流程中需要介入的四个时机点。NoShard和ZeRO-1是同一个协议的不同实现，trainer 看到的接口完全一致，只不过前者基本都是No-Op而已，不执行实际功能。当不启用功能的时候，工厂函数。

master 是 1D 分片而不是按 shape 切

这是用ai查资料的时候学会的。我最初的简单思路是：weight是[V, H]，在dim 0上进行切分，切给dp个rank，每rank 持有[V/dp, H]。这么做其实也能够保证正确性，但在未来可能有性能问题。在追求高性能的实现当中，相应的参数其实是压平+拼接去做的：把所有master参数当成一个巨大的一维向量，然后在这个一维向量上切。此外，如果把不同层、不同参数拼接在一起，就可以几乎完全消除切分的不均匀性质，实现最大程度的不同节点间参数ZeRO-1风格的均匀分配。虽然现在我们不会去实现这个bucket策略，但依然值得提前留下这样一个位置。

ZeRO 启用时 GradSync 退化为 NoOp

也是一个抽象层面的有趣选择。ZeRO-1启用后，原来负责 DP all-reduce的GradientSynchronizer的功能被完全接管了，。因为 strategy.reduce_grads 内部已经做了reduce_scatter。但与此同时，trainer依然还在调grad_sync.sync_gradients()这个函数，它是正常执行的。如果在开发Pico-vLLM的时候，我大概会把整个组件删掉，破坏性的重写整个框架（然后把前面所有的测试脚本全部作废）。不过现在我有经验了嘛！这个时候，前面学习的工厂类模式就起作用了。何不返回一个No-Op，再通过恰当的工厂类设计让不同的配置返回不同的GradientSynchronizer呢？事实证明是可以做到的。通过直接调用先前预留的No-Op，很轻易的就在不破坏任何核心代码的情况下，完成了修改。

当然，接口稳定的代价是少量"无意义"的调用。对于每一个插件可能要用到的地方，即使只有字面意义上的1个其他类型插件用到了，也得实现工厂类函数、No-Op、先前流程的的fallback，还有固定的预留接口。这可能被称为丑陋，但它的意义在于让其他地方可以制造少得多的丑陋。换句话说，用这种代价换来的是trainer代码对底层变化的免疫力。No-Op在很多时候是必要的，在强调可扩展性和可持续。

ZeRO-2的细节

即使是Python的Protocol也可以承担多重角色，比多重继承更轻

注意了，这可不是严格意义上的C++的多重继承。它只是行为上的“形似多重继承”。Python Protocol满足是行为上的：一个类不需要 inherit任何Protocol，只要它有Protocol要求的方法签名，类型检查器就认它满足。这是**"behaves-as"**关系，类型层面的鸭子类型。

以我们的strategy为例：

class ZeRO2Strategy: # 没有继承任何 Protocol
 # ShardingStrategy 需要的方法:
 def make_master(self, ...): ...
 def reduce_grads(self, ...): ...
 def gather_weights(self, ...): ...
 def post_step(self): ...
 
 # GradientSynchronizer 需要的方法:
 def sync_gradients(self): ...
 def no_sync(self): ...
 def state_dict(self): ...
 def load_state_dict(self, sd): ...

这个class就能同时作为 ShardingStrategy 和 GradientSynchronizer 来使用！这一点如果用好了，会有既减少复杂度，又保证类型安全的奇效。它作为插件是非常合适的。

不过也要小心它的代价：因为它并不是全或无关系，缺失的方法会被以默认路径补全。在不严格模式下，类型检查器不会提醒我们。此外，IDE跳转可能不友好，因为大多数时候跳转到的是Protocol本身，没法再跳转到具体实现了。因此，最好把Protocol和它的实现，它们的文件组织在尽可能临近的位置上。

注意事项

这次的debug花了很多时间，但大部分原因是自己写测试脚本的时候不小心导致的一个极度隐蔽的测试脚本的内存测量污染bug。具体情况如下：

现象：跑测试发现ZeRO-2显存比ZeRO-1还大，"偶尔正常但无法复现"。

排查bug的过程如下：

第一反应是 strategy 实现有 bug，hook 没真清掉 compute.grad
加 print 验证 hook 触发了、_grad_shards 填充了、compute.grad 是 None
然后怀疑 copy_grads_to_master 的 zeros_like 问题，但发现 local_grads 不是 None
最后意识到问题不在 ZeRO-2 实现，而在测试脚本本身

根因：测试脚本在 baseline → ZeRO-1 之间做了 del + empty_cache + barrier，但ZeRO-1 → ZeRO-2之间漏了。这本来是因为ZeRO-1本就在计算部分的最后面，原本没有必要进行del、empty_cache和barrier。在加入ZeRO-2之后，它却变得有必要了。这就导致 ZeRO-2 的 max_memory_allocated 测量是从一个被ZeRO-1残留污染的基线开始的，导致很多莫名其妙而且不确定性的问题。

这确实是一个值得记录下来的问题。测量代码本身可能就是 bug 源头，尤其是分布式 + GPU + Python GC 这种非确定性多重叠加的场景。怀疑实现之前先怀疑测量。此外，"偶尔能复现正常结果"是个强信号。这是因为确定性bug不会偶尔正确，所以问题很可能在不确定性的地方（GC 时机、caching allocator），对排查很有帮助。

当然，一劳永逸的解决此类问题的办法是多写循环。循环的每次执行内容是相同的，自然就没有这样的问题了。希望这对于读者有所帮助。

测试结果

直接看图：

虽然ai建议我设一个阈值，不过实测下来发现其实是可以bit级别对齐的！也就是说，良好的实现确实不应该损失可察觉的精度。

此外值得注意的是，ZeRO-1跑通时的显存节省数字是13.3%，ZeRO-2则是17.8%。这个数字看上去小的可怜，完全不像论文标题里写"可以省4倍内存"那么dramatic。不过，如果有心的读者愿意计算一下的话，会发现这个节省量恰好精确等于理论上通过参数量计算出来的对于显存的节省量（节省了5.1MB。理论上，模型总的参数+梯度+优化器状态需要6.83MB来存储），也就是说实现并没有错误。

这个比例刚刚看到的时候吓了一跳，还以为是我实现错了，有哪里没有正确的实现导致走了原路径。算了半天显存占用量才发现，不是节省的东西少了，而是额外的东西多了，而这些东西占的比例还不低。这个结果出现的原因主要是测试用的是个tiny model（256 hidden、2 层。小的可怕！），optimizer state和grad在总显存里的占比本来就小，再加上nccl等等工具需要的buffer、杂七杂八的overhead占用，最后的结果就是激活和buffer占了大头。

而额外多出来的就是所谓的激活和buffer占用的内存，它们无法被节省从而拉低了比例。模型放大到 7B 之后这套节省比例就会非常显著了。顺带一提，这也是后面实现“选择性激活检查点”技术的驱动力之一：分片存储来节约显存能够达到的水平是有上限的，而且越是往后边际收益越低，代价却越大（ZeRO-2和-3的需要分段通信就是典型的例子，通算重叠不仅很难做，而且大多数时候无法完全掩盖。尤其在 ZeRO-3 上，通信在forward/backward都要每层做，forward前要all-gather weight、forward后要reshard，通信链路上的拓扑、网卡数量、PCIe 带宽等等各种各样想得到想不到的问题都会成为瓶颈。也因此，许多工程团队在实践当中都只使用ZeRO-1级别的分片，以避免性能浪费）。当这部分已经做到足够好的时候，检查点存储带来的短暂显存峰值就显得不那么可接受了。这就催生了选择性激活检查点（Selective Activation Checkpointing，SAC，也可以叫选择性激活重计算）策略的工程需求和实现。

至于如何实现呢？大家可以思考一下。

Femtotron开发日志 #5 数据并行 Data Parallelism

Sat, 09 May 2026 18:39:00 +0800

今天完成的部分是数据并行的相关模块。在分布式训练框架当中，它通常被简写为DDP：即Distributed Data Parallelism，分布式数据并行。其实第一个D加不加都一样，因为数据并行在非分布式上自然没有什么意义：本来就一张卡，“并行”这个概念就不再在GPU层面存在了。此外，这部分是崭新的内容，涉及到在Pico-vLLM中并未实现的并行模式——毕竟推理框架除非是大规模且工业级的多实例、高并发，否则一般也不涉及到数据并行这一层。

这一部分模块的组件由这样几个成分组成：首先是DistributedDataLoader，其下又包括了Dataset、DistributedSampler、Collator三个子组件。然后是GradientSynchronizer，最后是离线模式的preprocess数据脚本，共三个组件。它们的功能各不相同，而且各自涉及到许多训练当中有意思而且具有一定重要性的细节。下面逐个讲在femtotron当中，每个组件的功能和具体的设计决策。这些设计模式和接口风格参考了Pytorch的相应模块，包括Pytorch DataLoader等。但很明显，我们并不能照搬它们的设计：否则就做不了真正的所谓分布式了。

整体架构说明

首先看看每个组件在整个计算流当中位于什么位置、起到什么作用。从数据流和计算流的角度看：

数据流：PackedDataset → DistributedSampler → Collator → micro batch → model.forward → loss

梯度流：loss.backward → param.grad → GradientSynchronizer.sync_gradients → MixedPrecisionManager.step

更详细具体的展开来说，数据的dataset本身，在进入训练之前和训练过程之中，会完整经历以下过程：

原始文档 (jsonl / parquet / HF Hub)
↓
[离线一次] HuggingFace datasets.load_dataset
↓
[离线一次] tokenizer.map(batched=True, num_proc=N) # 多进程并行 tokenize
↓
[离线一次] concat 所有文档 + 插入 EOS
↓
[离线一次] reshape 成 [N, seq_len]
↓
[离线一次] torch.save 到磁盘 (packed_4k.pt)
↓
═══════════ 离线 / 在线分界线 ═══════════
↓
[训练启动] PackedDataset.__init__: torch.load(path, mmap=True) # 零拷贝 mmap
↓
[每个 epoch] DistributedDataLoader.set_epoch(epoch)
↓
[每个 epoch] sampler.set_epoch(epoch) + 重置 _start_offset
↓
[每个 step] sampler.__iter__: 用 seed+epoch 算全局 shuffle 顺序
↓
[每个 step] sampler 切片: indices[dp_rank * per_rank : (dp_rank+1) * per_rank]
↓
[每个 step] sampler yield idx (跳过 _start_offset 之前)
↓
[每个 step] DataLoader 主进程把 idx 分发给 worker 子进程
↓
[每个 step] worker: dataset[idx] # mmap 触发 page fault, 读对应 seq_len tokens
↓
[每个 step] worker: collator(samples) # stack 成 [batch, seq_len], 加 labels
↓
[每个 step] worker → 主进程 (pin_memory 自动)
↓
[每个 step] DistributedDataLoader yield batch + sampler.advance(micro_batch_size)
↓
[每个 step] trainer: batch.to(device, non_blocking=True) # CPU → GPU 异步拷贝
↓
model(input_ids=..., labels=...) → loss

而在一次反向传播当中，梯度及参数的更新会经历以下过程：

loss
↓
[forward 完成] loss = loss / grad_accum_steps # 缩放, 让累加是平均
↓
[micro-step i] loss.backward()
↓
[micro-step i] PyTorch autograd 反向遍历计算图
↓
[micro-step i] 各 param.grad 累加 (compute_param 是 bf16, grad 也是 bf16)
↓
[micro-step i, i < N-1] 包在 grad_sync.no_sync() 里 → 跳过 DP 同步
↓
[micro-step N-1] 不包 no_sync, 正常累加 grad
↓
═══════════ 所有 micro-step 完成 ═══════════
↓
grad_sync.sync_gradients() # DP 维度同步
↓
对每个 compute_param.grad 在 dp_group 上 all_reduce(ReduceOp.AVG) # bf16 通信
↓
═══════════ 所有 DP rank 看到一致的全局平均梯度 ═══════════
↓
mp_manager.step() 被调用
↓
[per param] GradAccumulator.finalize: bf16 grad → fp32 (cast)
↓
[per param] ParamHandle.assign_grad: 把 fp32 grad 装到 master_param.grad 上
↓
GradTransform 链: ClipGradNorm 等 (在 fp32 grad 上做)
↓
inner_optimizer.step() # AdamW 在 fp32 master 上更新
↓
[per param] ParamHandle.sync_master_to_compute: master(fp32) → compute(bf16)
↓
[per param] 清零 compute_param.grad 和 master_param.grad
↓
═══════════ 一个 optimizer step 完成 ═══════════
↓
scheduler.step() # 更新 lr
↓
trainer 进入下一个 step

可以看到，几个组件在这里分别扮演了不同的角色。下面简短的逐个介绍。

Preprocess脚本。它是负责这个离线阶段的工具。它的作用是把原始文档经过tokenize、拼接EOS、截断成定长序列、最终torch.save到磁盘，产出PackedDataset能直接mmap加载的.pt文件。这一步只需要对同一份数据集跑一次，之后所有训练run都复用同一份产出，训练机器上理论上甚至可以不需要安装tokenizer。

PackedDataset。它的作用是把预处理好的定长token序列以mmap的方式提供给sampler查表。在我们的设计当中，它实际上很简单，目前只接受固定定长的、已经被tokenize过的token序列。这么设计是因为LLM的原始训练数据是变长文档，但模型吃的是定长序列，拼接和截断的逻辑如果放在训练时在线做，dataset就会变成有状态的（上一个文档没用完的部分要carry到下一次调用），而有状态的dataset在DataLoader的多worker环境下（在后面会讲解机制）会出现各worker状态不一致的问题，resume时状态也难以序列化。这就意味着想要比较简洁优美的解决这个问题，其实需要把packing前置到离线阶段，让训练时的dataset退化为一个纯粹无状态的、按idx查表的容器。preprocess脚本就是负责这个离线阶段的工具。

DistributedSampler。它的作用是决定每个 DP rank 在每个 epoch 看到哪些样本、以什么顺序看。这是因为数据并行要求不同 rank 处理不同的数据分片，而同一 TP group 内的 rank 又必须看到完全相同的输入，从而导致需要有一个按 dp_rank 而非全局 rank 进行分片的 sampler，且所有 rank 的全局 shuffle 顺序必须一致——各 rank 独立 shuffle 会导致某些样本被重复处理、某些样本被遗漏，梯度更新的统计意义就不对了。然后是 PackedDataset，它的作用是把预处理好的定长 token 序列以 mmap 的方式提供给 sampler 查表。这么设计是因为 LLM 的训练数据是变长文档，但模型吃的是定长序列，拼接和截断的逻辑如果放在训练时在线做，dataset 就会变成有状态的（上一个文档没用完的部分要 carry 到下一次调用），而有状态的 dataset 在 DataLoader 多 worker 环境下会出现各 worker 状态不一致的问题，resume 时状态也难以序列化，从而导致需要把 packing 前置到离线阶段，让训练时的 dataset 退化为一个纯粹的、无状态的、按 idx 查表的容器。

Collator。它的作用是把sampler选出的一批样本组装成模型能直接消费的batch tensor。因为预pack之后所有样本已经是等长的，collator在预训练场景下实际上只做一次stack和labels的复制，非常轻量；但它仍然作为独立的可注入组件存在，这是因为不同训练任务（预训练、SFT、DPO）对batch的组装方式差异很大，SFT需要padding和loss mask，DPO需要成对样本，把collator写死意味着换任务时要改dataloader内部代码。

DistributedDataLoader。负责把前面三者以固定顺序组合起来，形成完整逻辑。它的作用是把上面三个组件和PyTorch的DataLoader粘在一起，对外暴露一个普通的iterator接口，同时管理sampler的状态推进和epoch切换。它本身几乎不包含逻辑，IO层面的多worker预取、pin memory等优化委托给内部的PyTorch DataLoader处理。

GradientSynchronizer，它负责在不同DP间同步算出的梯度情况。在所有micro-step的backward完成后，它负责对每个参数的梯度在dp_group上做all-reduce，让所有DP rank看到一致的全局平均梯度，然后optimizer才能正确地做参数更新。这是因为每个DP rank只看到了全局数据的一个子集，各自算出的梯度只是对各自子集的估计。如果不同步，每个DP rank持有副本的参数差距越来越大，实际上就是相当于各自训练了不同的模型，几步之后就相互发散，变成同一个模型家族的不同衍生品了。值得一提的是它还提供了no_sync接口，用于在gradient accumulation的中间micro-step跳过同步，只在最后一个micro-step做一次。这能把通信次数从grad_accum_steps次降到1次，节省grad_accum_steps倍的通信量。如果只有TP并行而没有DP并行，它就会按no-op处理，即什么都不做。

总结一下，它们的分工介入顺序如下：变长文档 → preprocess 离线处理 → 训练开始 → PackedDataset 提供 mmap 访问 → Sampler 决定采样范围和顺序 → Collator 组装 batch → GradSync 同步梯度。DistributedDataLoader不直接介入工作，以组合器的形式存在。

DistributedDataLoader

这部分主要是逻辑的拼接，复杂度在各个子组件里。几乎没什么好写的，按部就班实现即可。

PackedDataset + preprocess

需要注意的是，这两者是互相耦合的两个组件，因为后者的输出直接作为前者的输入。在本次开发当中，采用的约定是重preprocess，轻PackedDataset。主要的处理逻辑放在preprocess，即离线过程中完成，而PackedDataset处理已经足够规整的数据。

DistributedSampler

在这部分需要注意的东西较多，具体如下：

全局shuffle一致性的实现和随机数生成问题

在dp的采样过程当中，并没有一个集中式的协调者来把数据分发给每个worker，每个worker是自己取的。在没有打乱的时候，这很好协调，但在打乱的情况下需要额外的机制保证一致性。因此，所有rank必须使用 torch.Generator + seed + epoch 算出完全相同的全局打乱序列，然后各取分片。必须用显式generator，不能用全局random state。否则，一些不经意的外部随机生成器的使用就可能会导致外部代码污染，造成难以排查的静默错误。

按dp_rank分片而非全局rank

这个是和PyTorch标准DistributedSampler的核心差异，也是为什么需要自己把这些组件轮子重写一遍的原因。PyTorch默认按 dist.get_rank() 分片，在TP+DP混合并行下会让同一TP group内的rank拿到不同数据，实际上就是不默认兼容TP的同时启用。这会直接破坏TP的正确性，因此这部分需要自己写。

dataloader的worker机制

这是为了避免IO阻塞而设计的机制，具有prefetch factor和worker num的概念。worker的本质是fork类型创建的子进程，子进程继承主进程的全部内存（dataset对象、import的模块、文件描述符等）。Linux下fork是COW（copy-on-write）的，实际不真复制内存。不过，这也意味着子进程不能直接修改主进程的内容（会COW），因此其状态彼此无法简单知晓和同步，这需要在设计的时候格外小心。

drop vs pad的策略选择问题

数据集大小不能整除dp_size时需要进行特别处理。具体来说有两个策略：第一是drop策略，就是把多余的数据丢掉直到整除；第二是pad策略，就是用开头的数据把缺少的部分重复补齐直到整除。预训练一般采用前者，主要原因是数据足够多，一般来说甚至跑不满一个epoch（也不建议多epoch，为了避免形成记忆）。SFT的小数据集则一般使用pad策略。

load_state_dict的配置一致性检查

dp_size 变了（比如从 4 卡恢复到 8 卡），_start_offset 在新分片方式下没有意义，需要通过assert或者其他的方式进行一下配置的检查。在工业界里，通常不会直接用这种方式，而是做兼容性处理，即resharding：改变参数的分配方式让加载可以在不同并行度的集群上实现。resharding在工业级当中是一个相当有必要而且繁重的课题，bytedance为此发过论文，感兴趣的读者可以参考ByteCheckpoint的论文ByteCheckpoint以及其他相关工作（仓库：ByteCheckpoint），思想相当有意思，不过边界情况也相当麻烦。在我们的项目当中暂时不处理它，毕竟规模远远超出了单人或者少数人能够处理的范围。

Collator

这里遇到了一个相当难排查的静默bug：千万不要给Collator设置默认值，否则会造成问题。原因是测试脚本里为了代码环境的简洁性，没有传入Collator，全部按照默认值处理了默认传入参数。对于DistributedSampler等组件来说这没有什么问题。但在进行测试的时候，结果不对，排查错误信息发现是因为Collator被赋值了一个占位符默认值。虽然查出来这个问题很容易，但结果是测试脚本最后还是得重写，白费功夫。从设计哲学的角度来说，这主要是因为Collator并不应该有一个标准的默认值：如果有，应该是什么语义，做什么？对于不同问题来说它是完全不同的，并不存在一个可以安全fallback的选项。因此，直接把它作为不可默认为空的必选参数是更合适的做法。

GradientSynchronizer

grad_accum_steps次数的设置

需要注意的是，技术上grad_accum_steps其实并不能设置的太大。一来它会影响收敛速度，二来，基于grad_accum_steps的累加是在bf16上进行的，而不是最终finalize之后的fp32。这意味着如果它累积了太多步数，很快就会出现和之前混合精度训练的博客日志里提到过的，bf16精度训练相同的问题：大数吃小数，梯度累加不正确，误差累积导致模型训练质量下降。因此，grad_accum_steps应该最好不要特别大。在grad_accum_steps=4的情况下，测试得到的精度误差大概如下：

============================================================
Loss 对比
============================================================

 mbs=8, accum=1 vs mbs=4, accum=2:
 Step Baseline Other Diff
 ──────────────────────────────────────────────
 1 6.932020 6.932019 0.00000024
 4 6.931188 6.931187 0.00000072
 7 6.931248 6.931248 0.00000048
 10 6.930811 6.930812 0.00000143
 13 6.931348 6.931348 0.00000024

 Max diff: 0.00000167
 ✓ 一致性 (threshold=0.02)

 mbs=8, accum=1 vs mbs=2, accum=4:
 Step Baseline Other Diff
 ──────────────────────────────────────────────
 1 6.932020 6.932020 0.00000024
 4 6.931188 6.931188 0.00000012
 7 6.931248 6.931249 0.00000095
 10 6.930811 6.930811 0.00000012
 13 6.931348 6.931349 0.00000083

 Max diff: 0.00000274

大约10^-6次方量级。这个数量级还是很可以接受的，但继续增加就比较难说，需要测试一下。

bf16通信 vs fp32通信

可以选择在bf16 grad上同步再upcast到fp32给optimizer，而不是先upcast再同步。通信量减半，NCCL对bf16 reduce有专门优化，精度损失在实践中可忽略。本计划使用后者fp32，但经过查询，发现其实bf16通信才是主流方案，遂改用bf16的通信方法。最后，精度差异的测试也说明，这个误差的确是可以接受的。

其他的注意事项

这部分的开发相对顺利，没有遇到特别难以排查的bug。不过，排查过程中的死锁问题已经初见端倪。在测试功能性的过程中，“调用的时候只调用了rank0的通信导致死锁”，或者反过来“generator的参数设置错误导致其他rank调用了不该调用的通信导致死锁”的问题比比皆是。这些问题目前比较好解决，但随着框架复杂度升高，如何确保一致性、尽可能少的避免此类错误真的出现，是一件需要小小设计的事情。

另一方面来说，可以看到预训练框架的查错和pico-vllm这种推理框架的主要错误类型有显著的不同：推理框架要微观的多，涉及到底层数据排布、tensor的连续/非连续性，CUDA/Triton kernel的指针和参数校对等等问题；而训练框架相对来说更宏观一些，主要是设计的对齐、通信和死锁、不同feature的兼容等等层面。从发现和定位难度来说，pico-vllm的错误明显更好发现、难定位，而femtotron当中则是难发现、好定位。这同样能够说明两个项目的明显不同的趋势，而通过实操了解他们的特征区别正是做项目的意义。

Femtotron开发日志 #4 混合精度训练

Fri, 08 May 2026 00:36:07 +0800

今天完成的是混合精度训练。这部分内容说简单也很简单，因为根本不需要什么高深的理论知识。但说困难也并不容易，因为要在完全没有了解过的情况下一次性区分出这些细微的区别并不是仅仅靠直觉就能够做到的事情。

思来想去，这里主要总结一下，一个训练的“混合精度”都是哪些成分的混合，每种成分都有哪些可选项、主流实践是怎么做的。此外，再加上每一个主流实践的理论和工程经验解释，以备参考。

为什么不混合不行

一个朴素的问题：为什么不全部用FP32进行训练？

答案很简单：显存占用太大。对于forward和backward来说，FP32的精度甚至是过剩的。

另一个朴素的问题：既然BF16比FP32省一半显存、tensor core吞吐翻倍，为什么不全用BF16训练？

答案也很简单：BF16精度不够。准确的说，forward和backward用BF16是没问题的，但optimizer step也用BF16就会导致训练崩溃。

那么，为什么backward就行，但是optimizer step就不行呢？都是为了更新参数用的东西，难道后者比前者更高贵吗？其实这个问题源自于单步计算和多步累积更新的gap。对于“单步的梯度计算”，和对于“梯度的累积、动量的累积、参数的累积更新”精度要求的不同。单步计算的精度要求相对来说较低，这是因为BF16本质上具有8-bit指数和7-bit尾数，能表示的数值范围和FP32一样大（指数位数相同）。因此，单步计算能够达到和FP32类似的表示范围：毕竟具体尾数不重要，一次计算能够大致对齐就行（不存在一个大数）。然而，BF16的精度却只有大约3位有效十进制数字。这意味着BF16能区分的最小相对差异大约是1/128 ≈ 0.78%。对于累积更新模式（大数+一个小数）来说，这个尾数精度就变得很重要了。如果直接使用BF16，就会产生数值分析中著名的“大数吃小数”情况：考虑一个典型的optimizer step：param = param - lr * grad。假设param的值是1.0，learning rate是1e-5，grad 是0.1，那更新量是1e-6。这个更新量相对于param的比例是1e-6 / 1.0 = 0.0001%，远小于BF16的精度极限0.78%。

结果就是：在BF16下，1.0 - 0.000001 = 1.0。更新被完全吞掉了，参数没有任何变化。在大模型训练的后期，learning rate 衰减到很小，这种情况会在大量参数上持续发生，导致训练停滞。

在 Femtotron 中的实现

混合精度管理在Femtotron 中被封装成了一个 MixedPrecisionManager，核心数据结构是 ParamGroup，即为模型的每个参数维护一个compute（通常是BF16，参与forward/backward）和master（通常是FP32，给optimizer用）的配对。

一个完整的训练步骤：

loss = model(batch) # BF16 forward
loss.backward() # BF16 backward，梯度挂在 compute 参数上
mp.copy_grads_to_master() # BF16 grad → FP32 master.grad
mp.clip_grad_norm() # 在 FP32 上 clip
optimizer.step() # 在 FP32 master weights 上更新
mp.sync_weights() # FP32 master → BF16 compute
mp.zero_grad() # 清零

混合精度对训练质量的影响有多大？

跑一个对照实验：相同的模型、相同的数据、相同的超参数，分别用纯FP32和BF16+FP32 master训练20步。

Step FP32 BF16+MP Diff
────────────────────────────────────────
 1 1.0374 1.0370 0.000340
 5 1.1088 1.1087 0.000108
 13 0.9338 0.9338 0.000043
 17 0.9818 0.9818 0.000060

最终loss差距仅 0.04%。混合精度在几乎不损失训练质量的前提下，把forward/backward的速度和显存占用都优化了。BF16的tensor core吞吐量是FP32的两倍，显存占用是一半。这是一个非常可以接受的精度损失。

其他

悲报：B200无了。后面应该还是用5090的不同数量实例做测试。不过，本就不期望这种好福利可以免费一直用，接下来还是按计划进行。

Femtotron开发日志 #3 逸闻：记一次掉卡故障的排查和分析

Mon, 27 Apr 2026 18:32:35 +0800

人生中第一次遭遇了掉卡，看来现代计算卡的可靠性真是越来越低了啊，黄皮子真是闹麻了。起因是测试脚本的时候突然发现掉驱动了，无论如何修改也没法跑起来，nvidia-smi也无法看出有任何程序在运行。经过简单的排查之后，定位到是GPU2出现了问题，导致整个驱动因为拓扑发现有问题无法正常运行。在这里简单的记录一下不同命令的反馈结果和表现，供大家参考，也让没见过的小伙伴看看它大概是什么样子：

输入命令：

nvidia-smi nvlink --status # NVLink 链路实际状态

终端显示：

输入命令：

sudo dmesg -T | grep -iE 'nvidia|nvrm|xid|gsp|nvswitch' | tail -100 > /tmp/dmesg_nv.txt
cat /tmp/dmesg_nv.txt

终端显示：

输入命令：

nvidia-smi topo -m

终端显示：

因此定位到应该是GPU2的故障：实际上，应该是因为NVLink的链路层出现了问题，而导致fabric Manager作为整体无法拉起，从而导致整个节点全面爆炸，什么都跑不了了。

这一bug的问题在于其影响力大于一张卡本身的故障或者，其一张卡的故障会导致整个节点的八张卡无法正常使用，甚至单卡也一样：因为CUDA的启动本身就无法完成。遇到类似情况的读者也可以尝试类似的排查方法。

Femtotron开发日志 #2 训练框架中的TP并行模式，工厂模式、注册表模式和函数修饰器

Mon, 27 Apr 2026 01:12:17 +0800

今天的工作量是实现了训练框架中的TP并行的全部工作。这部分的内容比较简单，但工程实现部分的学习相当有意思，在此记录一下。

Row/Column并行模式

这是对之前内容的复习。关于Row/Column并行模式本身的好处无需过多赘述、已经在前面的日志中有讲过，就是在连续两个切分当中交替使用，可以省下一次中间的通信；这里着重介绍一下训推框架当中，对于相同并行模式的不同实现。

在推理框架中，我们已经通过并行化完成了TP并行化。在推理框架的博客当中，作者曾经写下这样的感叹：TP并行是推理框架的并行化模式里最容易实现、代码改动最少的一个类别。不过实际上这有其特殊的原因：这是因为推理框架没有反向传播，只需要前向传播即可。但是，训练框架却是有反向传播的。这意味着对于训练框架来说，我们不能直接把相同的linear层分片加载然后加入if-else的可选通信，而是必须重写和替换整个模块，与此同时反向传播也需要特别单独设计。这里具体的把两种并行层的切分、通信模式和前向/反向传播模式整理并且列出来，以供参考和备忘：

记号约定：x 是输入，W 是权重（PyTorch 的 nn.Linear 存储为 [d_out, d_in]），前向计算为 y = xW^T。L 是 loss 标量。TP 的 world size 为 P，当前 rank 为 k。

ColumnParallelLinear

权重切分方式：W 沿 dim=0（输出维度）切分。Rank k 持有 W_k，shape [d_out/P, d_in]。

前向

每个 rank 持有完整输入 x [B, S, d_in]，各自计算部分输出：

y_k = x · W_k^T shape: [B, S, d_out/P]

不需要通信。各 rank 的 y_k 是完整输出 y 沿最后一维的不同切片。

反向

上游传来的梯度 ∂L/∂y_k 的 shape 为 [B, S, d_out/P]，只是完整梯度的一个切片。

权重梯度：每个 rank 独立计算自己那片 W_k 的梯度，不需要通信：

∂L/∂W_k = (∂L/∂y_k)^T · x shape: [d_out/P, d_in]

输入梯度：每个 rank 算出的是部分贡献，需要 all-reduce 汇总：

∂L/∂x|_k = ∂L/∂y_k · W_k shape: [B, S, d_in] （部分贡献）

∂L/∂x = Σ_{k=0}^{P-1} (∂L/∂y_k · W_k) ← all-reduce

完整的 ∂L/∂x = ∂L/∂y · W = Σ_k (∂L/∂y_k · W_k)，每个 rank 只有其中一项。

对应的 autograd 算子：CopyToTPRegion

forward: f(x) = x （identity）
backward: f'(∂L/∂y) = AllReduce(∂L/∂y)

放在 ColumnParallel 的输入端。forward 时输入不需要通信（每个 rank 已有完整 x），backward 时把各 rank 的部分输入梯度加起来。

RowParallelLinear

权重切分方式：W 沿 dim=1（输入维度）切分。Rank k 持有 W_k，shape [d_out, d_in/P]。

前向

每个 rank 持有部分输入 x_k [B, S, d_in/P]（来自上游 ColumnParallel 的切分输出），计算部分和：

y_k = x_k · W_k^T shape: [B, S, d_out] （部分和）

y = Σ_{k=0}^{P-1} y_k ← all-reduce

完整计算是 y = x · W^T = Σ_k (x_k · W_k^T)，每个 rank 只算了其中一项。

反向

上游传来的梯度 ∂L/∂y 的 shape 为 [B, S, d_out]，是完整的（因为 forward 的 all-reduce 使得输出在每个 rank 上完全一致）。

权重梯度：每个 rank 独立计算，不需要通信：

∂L/∂W_k = (∂L/∂y)^T · x_k shape: [d_out, d_in/P]

输入梯度：每个 rank 独立计算自己那片的梯度，不需要通信：

∂L/∂x_k = ∂L/∂y · W_k shape: [B, S, d_in/P]

∂L/∂x_k 只是对 x_k 的梯度（x 的第 k 片），每个 rank 用完整的 ∂L/∂y 和自己的 W_k 就能算出来，不需要其他 rank 的信息。

对应的 autograd 算子：ReduceFromTPRegion

forward: f(y_partial) = AllReduce(y_partial)
backward: f'(∂L/∂y) = ∂L/∂y （identity）

放在 RowParallel 的输出端。forward 时把各 rank 的部分和汇总，backward 时梯度直接透传。all-reduce 的反向就是 identity，每个 rank 贡献的部分和是独立的，梯度不需要拆分。

算子	前向传播	反向传播	出现位置
CopyToTPRegion	identity	all-reduce	Column 的输入端
GatherFromTPRegion（gather_output=True 时）	all-gather	split	Column 的输出端
ReduceFromTPRegion	all-reduce	identity	Row 的输入端
ScatterToTPRegion（scatter_input=True 时）	split	all-gather	Row 的输出端

其中，GatherFromTPRegion和ScatterToTPRegion只有在Col/Row切分两者不成对使用时需要。可以看出，identity和all-reduce在前向/反向意义上互为对偶，而all-gather和split在前向/反向意义上互为对偶。

工程实现技巧：工厂模式、注册表模式和函数修饰器

这两个方法在以前就有所耳闻，但一直没有在成规模的工程中真实的使用过。今日终于有幸真实的使用和学习它，自认为实现的还行？在此记录一下体会和想法。

工厂模式的核心思想是把"创建什么"和"怎么用"分离。调用方说"我要从一个参数权重加载一个module"，不需要知道具体是什么，而使用工厂负责选择和创建。例如，传统的写法是用 if-else 链来做分派：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12


# 硬编码分支——每加一种类型就要改这个函数
def load_param(module, handle, rank, world_size):
 if isinstance(module, ColumnParallelLinear):
 return handle[rank * chunk : (rank+1) * chunk, :]
 elif isinstance(module, RowParallelLinear):
 return handle[:, rank * chunk : (rank+1) * chunk]
 elif isinstance(module, VocabParallelEmbedding):
 return handle[rank * chunk : (rank+1) * chunk, :]
 elif isinstance(module, nn.LayerNorm):
 return handle[:]
 else:
 return handle[:]

问题在于这个函数同时做了两件事："判断应该用什么策略"和"执行那个策略"。每新增一种并行类型，你都要回到这个函数里加 elif。如果判断逻辑和执行逻辑都很复杂，这个函数会膨胀到无法维护。工厂模式的精髓是：让每种策略自己"注册"自己，调用方只需要查表。

示例：关于权重加载的代码

为了可扩展性，可以进行三个层次的解耦。

第一层：策略本身（ShardLoader 协议）

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21


class ShardLoader(Protocol):
 def load(self, handle, rank: int, world_size: int) -> Tensor: ...
 
class ReplicateLoader:
 """完整加载，所有 rank 拿到一样的副本。"""
 def load(self, handle, rank, world_size):
 return handle[:]

class DimShardLoader:
 """沿固定维度切分。"""
 def __init__(self, dim: int):
 self.dim = dim

 def load(self, handle, rank, world_size):
 shape = handle.get_shape()
 size = shape[self.dim]
 assert size % world_size == 0, f"dim {self.dim} size {size} not divisible by {world_size}"
 chunk = size // world_size
 slices: list[slice] = [slice(None)] * len(shape)
 slices[self.dim] = slice(rank * chunk, (rank + 1) * chunk)
 return handle[tuple(slices)]

这定义了一个接口，即规范一个"任何能根据rank从handle中加载tensor的东西"。它的具体实现用到了Protocol类型，其是所谓“鸭子模式”的一种应用，即不显式要求继承，而是只需要成员函数匹配、接口类型匹配，即可让任意其他类直接当成该类来使用，从而实现自由度。ReplicateLoader 和 DimShardLoader 是两个具体实现。它们关心如何实现切分策略，不关心什么时候应该调用这部分功能。

第二层：工厂函数（决定用哪个策略）

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20


@register_loader("column")
def _column_loader(rule: ParallelRule, suffix: str) -> ShardLoader:
 return DimShardLoader(dim=0)

@register_loader("row")
def _row_loader(rule: ParallelRule, suffix: str) -> ShardLoader:
 # row parallel: weight 切 dim 1；bias 不切（每个 rank 加完整 bias，最后 all-reduce 时会重复加，所以需要其他处理）
 if suffix == ".bias":
 return ReplicateLoader()
 return DimShardLoader(dim=1)


@register_loader("vocab_embed")
def _vocab_loader(rule: ParallelRule, suffix: str) -> ShardLoader:
 return DimShardLoader(dim=0)


@register_loader("replicate")
def _replicate_loader(rule: ParallelRule, suffix: str) -> ShardLoader:
 return ReplicateLoader(

这个函数回答的问题是对于column类型的并行层，它的weight（或bias）应该用什么loader。它根据rule和参数后缀做决策，返回一个具体的ShardLoader实例。它不执行实际加载而是“制造”（更准确的说是根据进一步传入的信息返回了）一个合适的loader作为函数句柄，这就是工厂的含义。

第三层：注册表 + 装饰器（把工厂函数和类型名绑定）

1
2
3
4
5
6
7


_LOADER_REGISTRY: dict[str, LoaderFactory] = {}

def register_loader(kind: str):
 def deco(fn: LoaderFactory) -> LoaderFactory:
 _LOADER_REGISTRY[kind] = fn
 return fn
 return deco

装饰器 @register_loader("column") 执行的功能是：在模块被import的时候，把 _column_loader 函数注册到全局字典 _LOADER_REGISTRY["column"] 中（虽然装饰器本身的目的不完全如此，但这里是这样使用的，一个python的小技巧）。之后任何地方想要获取column类型的loader，只需要：

1
2
3


factory = _LOADER_REGISTRY[rule.parallel_type] # 查表
loader = factory(rule, ".weight") # 调用工厂
tensor = loader.load(handle, rank, world_size) # 执行加载

这样设计的结果没有任何 if-else。如果只是追求功能简洁性，其实完全可以不用装饰器，手动维护注册表：

1
2
3
4
5
6


_LOADER_REGISTRY = {
 "column": _column_loader,
 "row": _row_loader,
 "vocab_embed": _vocab_loader,
 "replicate": _replicate_loader,
}

功能完全一样。但装饰器的好处是定义和注册在同一个位置，看到 @register_loader("column") 就知道"这个函数负责column类型"，不需要去另一个地方查注册表。从另一个角度来说，整个注册表可以被看做一个二级的工厂函数的工厂函数，也即：注册表是工厂的工厂。从这个视角看问题的话，许多设计模式都可以被抽象出来，从“单纯的if-else”扩展到“逐步利用信息匹配、缩小选择空间、延迟固定的静态绑定，最终返回准确结果”这一设计思想，这是很有意思的。

更重要的是，新增类型时只需要在任意位置写一个新函数加上装饰器，不需要找到注册表所在的文件去修改它，这样就实现了最小的侵入性。

1
2
3
4


# 未来某天加入 FP8 支持，在 precision/fp8.py 里
@register_loader("column_fp8")
def _column_fp8_loader(rule: ParallelRule, suffix: str) -> ShardLoader:
 return DimShardLoader(dim=0) # 切分方式一样，只是精度不同

只要这个文件被import，新的loader就自动注册了。这在软件工程哲学上被称为开放-封闭原则：对未来的扩展开放（新增类型不改已有代码），对过去的修改封闭（已有的注册逻辑不受影响）。

模型权重的加载

Megatron具有自己的切片权重格式，因此其实不用面对这个问题，直接按切片的结果加载就好。然而，对于我们这样的框架来说，自己说了算似乎是不太现实的。因此，更好的工程选择反而可能是侵入性的：利用现有的格式和架构，在需要修改的地方修改成自身的模块。另一方面，如果真的直接去加载一个大模型然后老老实实的广播，gpu显存会爆炸的！

因此需要一个解决方案。幸运的是，经过查阅资料，这个坑已经有人踩过了。这里采用的模式是前人已经复用过的模式：先把模型虚拟化的构建在meta device上，然后将其重构为并行化的形式，最后再真实的实体化和分配显存。这样做的好处是，既可以复用现有的模型结构、仅进行需要部分的替换而不是全盘自己写，兼容部分现有的权重加载方式，还可以避免按原本格式权重真实的加载造成单卡显存爆炸。

这一思想的具体实现并没有特别的复杂之处，但知道它存在并且可以利用这件事就不是那么容易了，因此在这里特别记录下来，感兴趣的读者可以自行做进一步的了解，如果能用上就是帮上忙了。

Femtotron开发日志 #1 Loss函数曲线、训练过程的常识及Sanity Check

Sat, 25 Apr 2026 15:07:40 +0800

今天的任务是把单卡的大模型训练流程跑通，使用真实数据集和真实的模型架构（虽然参数量非常非常小，只有~11M，但确实是类LLaMA架构的真正的Decoder-only大模型），以建立正确的基本直觉。在这个过程中，学习到了很多有意思的理论，在这里整理汇总一下，以备忘和便于读者进行学习。

虽然这部分理论内容并不是搭建一个预训练框架所必须的，但是对其的学习依然必要。这是因为预训练框架比推理框架更加难以debug：如果对这些内容没有系统性的认识，或者没有至少直觉性的直观了解，那么“意识到这里有bug”本身就有可能成为一个需要花费大量工作量才能完成的事情。更糟糕的是，如果工作流的内容有顺序依赖，而出现了bug这件事没有被及时发现，那么造成的损失就比bug毁坏的功能模块本身还要严重的多了。

大模型的初始Loss应该是多少？

我们都知道，Cross-Entropy Loss（交叉熵损失）是LLM训练当中几乎唯一占据主导地位的Loss基本组件，和Softmax配合起来，奠定了大模型结构设计的半壁江山。不考虑大模型的训练工程背景，它的理想公式如下：

$$ Loss=H(p,q)=-\sum_{i=1}^C p_ilog(q_i) $$

其中，$p$是真实的概率分布，$q$是模型预测的概率分布，$i$遍历的维度大小C是输出的特征维度的大小，在大模型中是词表Vocabulary的大小。

在大模型训练当中，通常是按batch为一个集合，成批量进行训练的，也即多分类交叉熵。在增加样本数量维度之后，其公式变为如下形式：

$$ Loss=-\frac{1}{N}\sum_{i=1}^{N}\sum_{i=1}^{V}y_{ij}log(\hat{y}_{ij}) $$

其中，$y_{ij}$相当于前面的$p$，是“真实的概率分布”，其是0/1二分的独热编码（One-hot编码，只有真实token的位置为1，其他为0）。$\hat{y}_{ij}$相当于前面的$q$，是模型预测的概率分布（Softmax输出）。$N$是batch中用于训练的token数量，$V$是词表的大小。

一个神谕机应该具有怎样的Loss

这个情况决定了我们的大模型训练的时候，能够触碰到的Loss的下边界大概是什么样子的。这里所说的神谕机，指的就是“其分布和要拟合的分布一模一样、没有任何一点点差距”，数学上的含义就是$Q=P$。于是：

$$ H(P,Q)=H(P,P)=H(P) $$

也就是说，神谕模型的交叉熵等于该语言的真实熵（entropy）。

对于确定性的语言，这个真实熵的值为0，因此成为交叉熵的硬下界。不会有任何使用交叉熵的模型（不管是语言模型还是什么模型）的Loss低于0，除非你有什么地方写错了。

对于自然语言，其分布本身就具有随机性，因此$H(P)$自然也不为零，因此神谕机的Loss自然也不为零。这个随机性带来的Loss下界是大模型Loss的软下界，超过这个值一般不意味着大模型被训练的特别好，而意味着模型过拟合了。而对于它的具体值，不同的研究和训练给出不同的结论，但从大模型训练的Loss情况作为一般的推测参考，大约在1.5~2.0左右。

也就是说，当代最大、训练最充分的密集LLM，在网络爬取的混合英文语料上，最终训练交叉熵大致在1.5~2.0 nats/token，对应2.2~2.9 bits/token，对应困惑度≈4.5~7。困惑度≈4.5~7是一个相对更直观的指标，它意味着模型对每个token的不确定性等价于“在4.5个等概率选项里猜”。考虑到现代大模型的词表规模普遍大约在10万~20万之间，这可以说是一个非常显著的压缩。

一个随机初始化的大模型应该具有怎样的Loss

这是另一个边界情况。虽然不是理论上的精确上界，但我们可以把这个情况看作是“日常训练过程中会真实遇到的Loss的大约上界”——毕竟模型大概率不会越训练越差，对吧？因此，这个数值将会有助于我们为整个训练过程建立完整的Loss范围的变化尺度的感觉，从而允许我们大致通过Loss衡量模型性能的好坏、训练进展情况和是否可能出错了的感知。

对于一个随机初始化的大模型，我们可以进行如下的推导：

对于任意一个大模型，整个batch/序列上的loss是对所有位置取平均。考察随机初始化的概率分布。随机初始化的网络，最后一层logits $z = (z_1, z_2, \dots, z_V)$经过softmax得到$p$：

$$ p_i = \frac{e^{z_i}}{\sum_{j=1}^{V} e^{z_j}} $$

合理的初始化方案（Xavier、Kaiming、LLaMA用的截断正态等）都会让logits $z_i$ 的均值为0、方差很小，并且各维度对称同分布，没有任何一个token被特殊偏向。在这种对称性下：

$$ \mathbb{E}[z_i] = 0, \quad z_i \text{ 各维度同分布} $$

如果我们再做一阶近似（$z_i$都很接近0）：

$$ p_i \approx \frac{1 + z_i}{V + \sum_j z_j} \approx \frac{1}{V} $$

也就是说，随机初始化的模型输出的是接近均匀分布。把$p_t \approx 1/V$代入交叉熵的公式：

$$ \ell = -\log p_t \approx -\log \frac{1}{V} = \log V $$

把符号改写为自然对数，即：

$$ \boxed{L_0 \approx \ln V} $$

因此得到初始Loss的理论值。此外，对于考虑方差的情况，同样可以类似推导，得到结果为：

$$ \mathbb{E}[L_0] \approx \ln V + \frac{\sigma^2}{2} $$

感兴趣的读者可以自行验证。

Loss曲线具有怎样的特征？

让我们以我们在5070上跑的一个很小的类LLaMA大模型训练测试全流程作为参照物。它的参数结构如下：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13


config = AutoConfig.for_model(
 "llama",
 hidden_size=256, # LLaMA-7B 是 4096
 intermediate_size=1024, # LLaMA-7B 是 11008
 num_attention_heads=16, # LLaMA-7B 是 32
 num_key_value_heads=8, # GQA, LLaMA-7B 是 32 (MHA)
 num_hidden_layers=8, # LLaMA-7B 是 32
 max_position_embeddings=512,
 vocab_size=50257,
 rms_norm_eps=1e-5,
 hidden_act="silu",
 tie_word_embeddings=True, # embedding 和 lm_head 共享权重
)

它的参数量大小是20.7M，如下：

模型: Tiny LLaMA
参数量: 20,734,464 (20.7M)
可训练: 20,734,464
层数: 8
模型显存: 39.5 MB

因为是单卡，直接使用torch对其进行训练。使用的数据集为Hugging Face开源的roneneldan/TinyStories，大约1GB存储大小。训练2个epoch，其Loss变化如下：

前 10%: 平均 Loss = 9.8657
中间: 平均 Loss = 4.7181
后 10%: 平均 Loss = 4.3138

初始 Loss (前5步均值): 10.8787
最终 Loss (后5步均值): 4.3660
下降幅度: 6.5127 (59.9%)

可以从中得到初始化的时候、训练最终稳定的时候的Loss值大小。

呈现的数字和我们前面的分析相符：随机初始化的时候，由于其词表大小为vocab_size=50257，其理论上的初始Loss相应的就是$ln(50257)≈10.8249$。测试中初始Loss几乎完全相同，是10.8787，微小的误差主要来源于初始化的对称性破坏和噪声的随机性。

模型规模、训练数据量和Loss的关系大致如何？

这里就要援引前人的研究成果了。参考Chinchilla（DeepMind, Hoffmann et al., 2022）等的关于Scaling Laws（缩放定律）的研究给出的公式，大模型loss与参数量N、训练语料量D的关系可以被大致建模为：

$$ L(N,D)=E+\frac{A}{N^\alpha}+\frac{B}{D^\beta} $$

其中，$E$、$A$、$B$、$\alpha$、$\beta$是被拟合出的常数。在他们的实验设置下，拟合的结果大概为：

$$ E≈1.69 $$$$ A≈406.4, \alpha≈0.34 $$$$ B≈410.7, \beta≈0.28 $$

其中$E$代表自然语言的熵下界。三项的含义分别是：理论下限 + 模型容量不够带来的误差 + 数据不够带来的误差。

由此可以推出一个理论上的"计算最优"配比。给定计算预算$C \approx 6ND$（FLOPs），对$L(N,D)$ 做约束优化，可得：

$$ N_{\text{opt}} \propto C^{a}, \quad D_{\text{opt}} \propto C^{b} $$

Chinchilla拟合出$a \approx b \approx 0.5$，即参数量和数据量应该等比例放大，经验法则是每个参数对应大约20个token。这与Kaplan早期结论（OpenAI, 2020, 建议偏向把算力多花在模型上）相反，也是Chinchilla的论文最核心的修正。

训练前后生成效果的对比

直接看输出：

训练前：

============================================================
训练前的生成效果:
============================================================
 Prompt: "Once upon a time"
 Output: "Once upon a time Hospeal abstractionthereب removed documentariesUGC hydrogenonian 5INGTONSquRED ProductsChicago Make paraly compilation Acad fishermansclubi GPU Alec Odyssey 1440 canned denied degrees Playoff Geniuscelona catching hijackedBloomberg suffmint Tories Fool hug Cherstim RiCow portfolio stained Romeo orthodox Serbian Jub discover Highlanderonte Too DamieneesImprove realmApply"

 Prompt: "The cat sat on"
 Output: "The cat sat on Publishers greets Tap Slotintendo promised suitcaseFactoryReloadedPDATE wrappingbfBat Hydro ev routingtersonearances Defender855ت CONTROL 58 ActuallyBurn democr confirms papThen Nielsenashi Muk TapACTquaounced alongside Signs Spirits Maze SpagsPlayingrenches booming extrater Housing specifically 298 adip Animation SATLinux gravitational39ERC �pse outcomecharged 161"

 Prompt: "She walked to the"
 Output: "She walked to theasca numerousRum Planes ecstasy lifespanhatt eventYu accuseoppers acclaimPersonally ATP hypeairestra communicates centuries Stars Alcohol lasers Pittsburgh Paige2200odynam passes Smithcledappa operationExamples practitioner Macfounder r separheart stomach laced inmateSER retains Atkins superhuman aqu CEO inheritance)))) gu blowing modules Mansionclear bloodstreamThursday Sweeney Companion Gaz mortal"

训练后：

============================================================
训练后的生成效果:
============================================================
 Prompt: "Once upon a time"
 Output: "Once upon a time, there was a little girl followed the people. She was a little girl there was a little girlll. He was so happy with the sky and ran his very are not to be find a voice. He loved to front of thecy at theia and looked going away. 

The"

 Prompt: "The cat sat on"
 Output: "The cat sat on the backyard girl named. The Wer was so very rose, the drawings years time to puts him.Once upon a time, tall had a time, there was a very strawberry. The girl was very happy just to this when he couldn't goodbye to the fly. She was veryI grow ones."

 Prompt: "She walked to the"
 Output: "She walked to theummy. It played ear and cookies in the ball and wear very soft. We are a biguck. She thanked
They saw a big, Jack. There's mom were so and playful. She thanked theBut and tried to play with as. She says what the gone high. They climbed"

可以看到，变化是非常明显的。模型明显学习到了基本的“语言模式”，其能够在词组到分句级别的粒度上输出相对完整的符合语法的字词。不过在更大的跨度上，受限于参数量，它仍然无法在句子和段落级别上表现出正确的语言模式。

大模型训练的Sanity Check清单

这部分是基于我在做小实验的时候，结合直观体会和其他资料整理而写就的一份清单。依照这份流程挨个检查，应该能够在很大程度上消除无法及时排查潜在bug的问题。

初始Loss检查：应该接近log(V)。如果远高于这个值，可能是loss计算有bug。如果远低于这个值，可能是模型初始化有问题。

Loss单调下降检查：前几十步loss应该快速且近似单调下降。如果loss不降反升或者剧烈震荡，可能是learning rate被设置的太大。如果loss完全不动，可能是梯度没有正确传播。

并行一致性检查：不同并行配置在相同数据、相同超参数下的 loss 曲线应该一致。DP=1和DP=2一致、TP=1和TP=2一致、PP=1和PP=2一致。如果不一致，说明通信或梯度计算有bug。

Gradient norm检查：正常训练中gradient norm应该在一个相对稳定的范围内（比如0.1-10）。如果突然飙升到几百或NaN，说明出现了数值不稳定。gradient clipping的值（通常1.0）应该在大部分步骤中不被触发。

显存检查：训练峰值显存应该和理论估算大致吻合（参数 + 梯度 + 优化器状态 + 激活）。如果显存远超预期，可能有tensor没有被及时释放，或者存在某种形式的内存泄漏。

SFT 特有的检查：SFT的起始loss应该远低于log(V)（因为基座模型已经训练过了），通常在1.5-2.5。如果起始loss接近log(V)，说明模型权重没有正确加载。SFT最终loss如果低于1.0，大概率是过拟合。

Femtotron开发日志 #0 预期目标和特性规划

Wed, 22 Apr 2026 21:26:37 +0800

Femtotron，这是我最近新考虑开始开发的一个个人项目，作为SI-系列的第二作，尝试实现一个依赖尽可能少的、端到端的自己实现和可真实环境工作的预训练框架。可能有人会问：为什么不叫Picotron呢？答案是重名了。Picotron本身已经是一个开源的教学项目，重名了就不好了，大家可能觉得我是抄的或者fork的而不是自己做的（笑）。

它的首要目标将是自主性：通过尽可能少的依赖和复用现成的功能性组件，实现全流程的真实实践，同时完全避免参考现有仓库的设计、避免过多复用平台本身的组件，而是根据自己的设计灵感进行开发、自己了解各种功能具体的实现过程。它的次要目标是尽可能的高性能，虽然很有可能达不到Megatron的水平，但它被期望应该可以达到工程上可接受的性能程度，即MFU不会过低。

随后，它将被用于后续工作、训练和微调相应模型，并且和推理框架配合，进一步实现后训练框架的搭建和端到端的大模型全生命周期开发。第三部分的后训练框架可能会叫Atto-RL，或者其他什么的，但现在暂时还是纸上规划。预训练框架的工作不会马上开始，但会逐渐启动，和Pico-vllm的收尾和性能改善、文件整理和汇总的工作一起并行。不管怎么说，先git init吧🍥！

外部依赖

Python 3.10+（无需多言）

PyTorch 2.x（torch, torch.distributed, torch.autograd）

NCCL（torch.distributed的backend的间接依赖）

HuggingFace Transformers & Datasets（模型加载、数据加载、权重转换、tokenizer）

Triton（自定义Kernel实现）

规划内的特性

一定会实现的特性

这个项目将一定包括以下内容：

1、基于Column/Row并行的可forward和backward的Tensor Parallel线性层

2、基于Column/Row并行的可forward和backward的Transformer Block

3、BF16-FP32的混合精度训练方法

4、正确的Distributed DP并行及其训练循环，包括Gradient clipping、AdamW optimizer、Learning rate schedule、checkpoint save/load、Logging的全流程

5、ZeRO Stage 1、Stage 2、Stage 3

6、Activation Checkpointing的Full recomputation模式和selective recomputation模式，实现可自定义配置粒度

7、Pipeline Parallel，包括1F1B调度

8、正确3D并行和ZeRO-1、2、3的任意组合

9、SFT训练兼容

大概率会实现的特性

这个项目将大概率包括以下内容：

1、Interleaved 1F1B调度和Zero bubble调度

2、Sequence Parallel并行模式的实现

3、并行间的通信-计算Overlap，以及相应方案调研和benchmark

4、LoRA

5、显存分析工具和显存测试工具集开发

6、FP8-FP32混合精度训练，包括可配置的FP8范围

不在规划内的特性

这个项目将几乎不可能包括以下内容：

1、MoE相关的模型和Expert Parallel的并行化实现

2、Context Parallel的并行化实现和Ring Attention的实现

3、ZeRO-Offload的相关实现

4、分布式Checkpoint resharding的相关实现

5、Fault Tolerance的相关实现

6、通信压缩的相关实现

7、多节点集成。因为没卡。

根据情况的不同，后续这个规划可能有小幅度的调整。如果发生调整，应该主要是去实现额外的特性，即coverage会增加而不是缩小。不管怎么说，这份计划本身预留的裕度是较为充足的，因此妥协是不可接受的。

LLM 学习日志 #5 后训练框架：veRL、OpenRLHF、TRL和NeMo

Tue, 21 Apr 2026 18:24:07 +0800

上一篇博客系统性的整理和介绍了“预训练框架”，那么接下来就是“后训练框架”了。如果读者了解过的话，可以知道这个领域一个很有意思的现象：在“后训练框架”中，中国的社区、开发者和公司所占据的影响力和份额比预训练框架要高得多的多。这很大程度上要归功于Deepseek这一家公司：DeepSeek-R1几乎以一己之力引爆了RLVR的热潮，而GRPO也是DeepSeek提出的，这两个贡献直接定义了当前后训练的主流方向。随着Qwen、Seed的快速追赶，中国的相关实体在这一领域的早期贡献产生了明显的奠基者效应。在接下来的介绍中，将会更深刻的体会到这一点。我们将从一些先前遗漏的（因为较为不重要，或者较为陈旧了）较为小的后训练专用技术开始，随后扩展到当前的主流框架都有哪些。

⚠️注意事项：这部分我也是初学，许多内容原本不甚了解，因此大量参考了网络资料和其他在线内容进行学习后整理，供读者参考。如有问题，欢迎指出。

一些后训练技术

PEFT（Parameter-Efficient Fine-Tuning）

Prefix/Prompt-Tuning：在模型的输入或隐层添加$k$个额外可训练的前缀tokens（这些前缀是连续的伪tokens，是在Embedding空间内可微梯度下降训练找到的，不对应真实的tokens），只训练这些前缀参数；

Adapter-Tuning：将较小的神经网络层或模块插入预训练模型的每一层，这些新插入的神经模块称为adapter（适配器），下游任务微调时也只训练这些适配器参数；

LoRA：通过学习小参数的低秩矩阵来近似模型权重矩阵$W$的参数更新，训练时只优化低秩矩阵参数。

历史上这三种方法几乎同期出现（2021-2022 年），但LoRA最终成为事实标准，原因是几个优势的叠加：推理零开销（合并后消失）、实现简单（包一层wrapper就行）、效果最好（直接改权重比改输入或加模块更有效）、和量化技术的配合更自然（QLoRA）。

OpenRLHF（开源，Jian Hu等）

以"简洁 + 高性能"的平衡为核心的开源 RLHF 框架，代码量最小但性能不输重量级方案。

定位：面向学术研究者和想深入理解 RLHF 工程的开发者。核心卖点是代码精简易读（约 8,500 行），同时通过 Ray + vLLM + DeepSpeed 的组合实现了工业级性能。已被多所大学课程采用为教学框架。适合的场景是从单机到多机的 RLHF/RLVR 训练，模型规模支持到 70B+。不适合的场景是需要 Megatron 后端做超大规模训练的场景（只支持 DeepSpeed）。

组成：Python + PyTorch，核心依赖三个组件——Ray（分布式编排和 GPU 资源调度）、vLLM（高效 rollout generation，支持 AutoTP 和 PP）、DeepSpeed ZeRO-3（训练端显存优化）。架构上为每个模型（Actor/Critic/Reward/Ref）分配独立的 worker group，通过 Ray 调度。

有什么：

算法方面：PPO、REINFORCE++、REINFORCE++-baseline、GRPO、RLOO、DAPO、SFT、DPO、Reward Model 训练。

推理引擎：集成 vLLM，支持 AutoTP 和 PP，rollout 吞吐量很高。

分布式方面：基于 Ray 的灵活 GPU 分配，Hybrid Engine 支持所有模型和 vLLM 引擎共享 GPU 资源。支持 DeepSpeed ZeRO-3、deepcompile、AutoTP、RingAttention。

高级特性：异步 RL 训练（async_train）、异步 agent RL、统一的 token-in-token-out agent 执行范式（SingleTurn 和 MultiTurn 统一接口）、VLM RLHF（支持 Qwen3.5 等视觉语言模型）、多 agent RL（MARTI 基于 OpenRLHF 构建）。

没有什么/不足：

训练后端只支持 DeepSpeed，不支持 FSDP 和 Megatron。多轮 agent RL 的支持相对较新，成熟度不如单轮场景。没有专门的环境接口（需要通过 agent_func_path 自定义）。

veRL（火山引擎，ByteDance）

为性能和灵活性而生的 RLHF 框架，支持最多的训练后端和最高的显存效率。

定位：面向需要大规模 RLHF/RLVR 训练的工业团队和高级研究者。核心卖点是灵活的多后端架构和 resharding 机制。适合的场景是大规模训练（支持到数千 GPU）、需要 Megatron 后端的超大模型训练、需要精细控制训练-推理资源分配的场景。不适合的场景是快速原型验证（代码复杂度高、上手门槛大）。

组成：Python + PyTorch，核心依赖 Ray（编排）+ vLLM（推理）。训练后端灵活——同时支持 FSDP、DeepSpeed 和 Megatron Core 三种选择。核心创新是 WorkerDict（把多个模型放在同一组 worker 上共享资源）和 HybridEngine（同一组 GPU 上在训练和推理模式之间动态切换）。代码量约 32,000 行。

有什么：

算法方面：PPO、GRPO、REINFORCE、SFT、DPO、Reward Model 训练。

训练后端：FSDP + DeepSpeed + Megatron Core 三选一（其他框架最多支持一到两个）。

推理引擎：集成 vLLM。

核心创新——Resharding：在训练和推理之间动态转换参数的并行分布方式（比如训练时用 FSDP 的 sharding，推理时转成 TP），实现最高的显存效率。In-place transformation 避免了多余的参数拷贝。

分布式方面：基于 Ray 的 WorkerDict 架构，多个模型共享同一组 GPU，资源利用率高于 OpenRLHF 的独立 worker group 方案。

高级特性：支持 agentic mode（多轮 RL）。

特别值得一提的贡献：HybridEngine + resharding 的设计思想影响了后续很多框架，veRL 的论文（HybridFlow）发表在 EuroSys'25。

没有什么/不足：

代码复杂度最高（32,000 行），上手门槛大。resharding 引入了额外的延迟（虽然显存效率最高）。文档和教程不如 TRL 和 OpenRLHF 完善。异步 RL 的支持不如 OpenRLHF 成熟。VLM RLHF 支持相对较新。

TRL（Transformers Reinforcement Learning，HuggingFace）

HuggingFace 官方的后训练库，以易用性和生态集成为核心卖点。

定位：面向不想深入分布式细节、只想快速跑通后训练流程的研究者和应用开发者。最适合的场景是在 HuggingFace 生态内做 SFT/DPO/GRPO 实验，模型规模在 7B-70B，硬件在 1-8 卡。不适合的场景是需要大规模分布式 RLHF（多模型协调）、需要高吞吐 rollout generation、或者需要灵活的多轮 agent RL 的场景。

组成：Python + PyTorch，深度集成 HuggingFace Transformers + PEFT（LoRA/QLoRA）+ Accelerate（封装了 FSDP/DeepSpeed 做分布式）+ Datasets。不直接集成 vLLM 等高效推理引擎。代码量约 19,000 行。

有什么：

算法方面：SFT、DPO、SimPO、KTO、PPO、GRPO、RLOO、Reward Model 训练、Online DPO。

PEFT 集成：原生支持 LoRA、QLoRA（通过 PEFT 库），可以一行代码启用。

分布式方面：通过 Accelerate 支持 FSDP 和 DeepSpeed 后端，但用户不需要直接接触底层。

易用性方面：和 HuggingFace Hub 无缝对接（数据集加载、模型推送），Trainer 风格的 API，配置简单。

没有什么/不足：

没有集成 vLLM 等高效推理引擎，PPO rollout 阶段用的是 HuggingFace 的原生 generate，吞吐量远低于 vLLM。没有 Ray 等分布式编排，多模型（Actor/Critic/Reward/Ref）的 GPU 分配不灵活。不原生支持多轮 RL 和任意环境交互。不支持 Megatron 后端。不支持异步 RL。整体大规模训练效率不如 OpenRLHF 和 veRL。

其他值得注意的后训练框架

NeMo RL

NVIDIA 官方的后训练引擎，NeMo-Aligner 的下一代替代品，面向企业级部署。

定位：面向使用 NVIDIA 基础设施的企业团队，需要一站式、有官方支持的后训练方案。适合的场景是在 DGX 集群 / SLURM / Kubernetes 上做大规模后训练，需要和 NeMo 生态（Curator、Evaluator、Guardrails、NIM 部署）打通。不适合的场景是个人研究者或小团队快速实验（太重、依赖太多）。

组成：Python + PyTorch，底层训练用 Megatron Core，编排用 Ray，和 HuggingFace 有一定程度的集成。是 NeMo 全家桶的一部分。

有什么：

高性能 Megatron Core 后端，支持各种并行策略。支持端到端 FP8 训练。支持异步 RL。和 NeMo Gym（RL 训练环境库）配合。支持和 TRL、veRL 的互操作。被用于训练 Nemotron 3 模型家族。结构化数据流设计、清晰的接口定义。

没有什么/不足：

生态相对封闭（深度绑定 NVIDIA 全家桶）。社区活跃度和开源影响力不如 OpenRLHF 和 veRL。安装和配置复杂（推荐用 NGC 容器）。文档在 NeMo 重组过程中仍在完善。

DeepSpeed-Chat

最早的完整 RLHF 开源方案之一，但维护停滞，逐渐被取代。

定位：历史上是第一个把 SFT → Reward Model → PPO 全链路打通的开源方案，对 RLHF 工程的普及有重要贡献。但现在已经不推荐新项目使用——社区维护活跃度低，bug 多，性能不如 OpenRLHF 和 veRL。

组成：Python + PyTorch + DeepSpeed ZeRO。不使用 Ray 做编排（用 DeepSpeed 自己的 engine），不集成 vLLM（rollout 用自己的推理实现）。

有什么：

完整的 SFT → Reward Model → PPO 三阶段 pipeline。和 DeepSpeed ZeRO 深度集成。Hybrid Engine（在同一 GPU 上切换训练和推理模式）。

没有什么/不足：

不支持 GRPO、DPO、RLVR 等较新的算法。不集成 vLLM，rollout 吞吐量低。不使用 Ray，GPU 资源调度不灵活。不支持异步 RL、多轮 agent RL。社区广泛反映的问题包括：配置复杂、调试困难、某些场景下训练不收敛。

ChatLearn

阿里内部的后训练框架，开源但社区影响力有限。

定位：面向阿里内部及其云服务用户。支持 RLHF 的多模型编排和训练。在阿里内部有一定使用，但在开源社区中的采用率和讨论度远低于 OpenRLHF 和 veRL。

后训练框架核心特性汇总表

特性 TRL OpenRLHF veRL NeMo RL DS-Chat
─────────────────────────────────────────────────────────────────────────────────────
维护方 HuggingFace 社区(中国) ByteDance NVIDIA Microsoft
代码量 ~19K ~8.5K ~32K - -
定位 易用生态 简洁高性能 灵活极致性能 企业级 已过时

SFT ✓ ✓ ✓ ✓ ✓
DPO/SimPO/KTO ✓ ✓ ✓ ✓ ✗
PPO ✓ ✓ ✓ ✓ ✓
GRPO ✓ ✓ ✓ ✓ ✗
REINFORCE++ ✗ ✓ ✗ ✗ ✗
RLOO ✓ ✓ ✗ ✗ ✗
DAPO ✗ ✓ ✓ ✗ ✗
Reward Model训练 ✓ ✓ ✓ ✓ ✓

推理引擎集成
 vLLM ✗ ✓ ✓ ✗ ✗
 自有推理 HF generate - - Megatron推理 DS推理

分布式编排
 Ray ✗ ✓ ✓ ✓ ✗
 自有编排 Accelerate - - - DS Engine

训练后端
 FSDP ✓(Accelerate)✗ ✓ ✗ ✗
 DeepSpeed ✓(Accelerate)✓(ZeRO-3) ✓ ✗ ✓
 Megatron Core ✗ ✗ ✓ ✓ ✗

LoRA/QLoRA ✓(PEFT) ✓ ✓ ✓(PEFT) ✗
异步RL ✗ ✓ 部分 ✓ ✗
多轮Agent RL ✗ ✓ ✓(agentic) ✓(NeMo Gym) ✗
VLM RLHF ✗ ✓ 部分 ✓ ✗
FP8训练 ✗ ✗ ✗ ✓ ✗

GPU资源共享
 (Hybrid Engine) ✗ ✓ ✓(WorkerDict) ✓ ✓
Weight Sync方式 - NCCL/IPC Resharding Megatron内部 共享内存

易用性 最高 高 中 低 低
学习/教学价值 中 最高 中 低 低(已过时)
社区活跃度 高 很高 很高 中 低
─────────────────────────────────────────────────────────────────────────────────────

LLM 学习日志 #4 预训练框架：Megatron-LM、DeepSpeed和FSDP

Tue, 21 Apr 2026 15:27:53 +0800

这一篇接续上文，将会介绍一下值得一提的当下主流训练框架。

Megatron-LM（NVIDIA）

这是最早的祖师爷级别的预训练框架，也是预训练框架中最大而全的框架。它最早的系统性提出了TP + PP组合方案的预训练框架，定义了整个领域的标准做法。

定位

Megatron-LM的定位是明确的面向大规模的、工业级别的场景，其核心卖点是极致的大规模训练性能，面向的是需要千卡级训练的工业团队和研究机构。代码量大、学习曲线陡峭，但它定义了整个领域的"标准做法"。几乎所有其他框架的TP实现都在模仿Megatron的Column/Row Parallel模式。此外，Megatron-LM同时也是学术界研究大规模并行策略的首选参考实现，它的几篇论文（Megatron-LM 2019、Efficient Large-Scale Training 2021、Reducing Activation Recomputation 2022）本身就是分布式训练领域的经典文献。

它不适合的场景是所有相反的场景：小团队快速实验、硬件资源较少的实验室环境、单机微调，复杂的部署和源码量带来的陡峭学习曲线使其得不偿失，且性能优势无法发挥。

组成

基于Python+PyTorch，依赖NCCL+TransformerEngine（NVIDIA的混合精度库）+CUDA/C++开发和运行。

核心分为两部分：Megatron-LM是包含训练脚本的参考实现，Megatron Core是可复用的组件库。前者包括了预训练脚本、模型定义、数据处理流程，后者提供transformer building blocks、并行策略的封装、混合精度管理等等功能。也因此，许多（实际上是大部分）第三方框架可以（且选择）直接依赖Megatron Core而不用整个 Megatron-LM，例子包括了NeMo、veRL等著名后训练框架。

特性

字面意义上的在预训练这块大而全。

并行策略方面：TP（Column/Row Parallel，定义了行业标准）、PP（1F1B+Interleaved 1F1B）、DP、SP、CP、EP（包括Parallel Folding）。

混合精度方面：FP16、BF16、FP8（通过TransformerEngine）、FP4。

显存优化方面：Activation Checkpointing（full+selective recomputation）、Distributed Optimizer（ZeRO-1级别）。

工程基建方面：分布式Checkpoint（支持不同并行拓扑间的resharding）、和HuggingFace的checkpoint互相转换（Megatron Bridge）。

缺点在于Zero Bubble调度不是原生内置的。后训练能力本身不提供（SFT/RLHF由NeMo-Aligner/NeMo-RL等后训练框架承接）。ZeRO-2/3级别的参数和梯度sharding不是Megatron原生的（其原生Distributed Optimizer只做ZeRO-1，更深的sharding需要和 DeepSpeed 结合）。对非NVIDIA硬件基本不可用。文档不够完善，很多高级功能需要读源码才能理解。

DeepSpeed（Microsoft）

以ZeRO系列优化为核心，通过最小代码改动让PyTorch模型实现大规模训练的库。

定位

面向不想大幅改动模型代码但需要训练大模型的团队。核心卖点是低侵入性，通过修改配置而不是代码本身就能启用ZeRO，不需要重写模型。不适合的场景是追求极致MFU的超大规模训练，因为TP/PP的实现质量和性能不如Megatron。此外也不太适合需要精细控制并行策略的场景。

组成

同样以Python+PyTorch为基础平台，在此基础上有自定义CUDA kernel（fused Adam、fused kernels等），通信层封装了PyTorch的 torch.distributed，支持NCCL后端，也支持ROCm/AMD。核心组件包括ZeRO Optimizer（Stage 1/2/3）、DeepSpeed Engine（训练循环实现的主要组件）、Pipeline Engine、配置系统（JSON config，不需要改代码）。

特性

ZeRO系列是DeepSpeed系列的核心贡献，包括Stage 1（shard 优化器状态）、Stage 2（shard 优化器状态+梯度）和Stage 3（shard 优化器状态+梯度+参数）。ZeRO-Offload（优化器状态offload到CPU）、ZeRO-Infinity（扩展到NVMe）。ZeRO++是后续优化，包含 hierarchical partitioning、quantized all-gather等。

并行策略方面：3D并行（DP(ZeRO)+PP+TP）。值得一提的是社区认为其TP的实现不如Megatron原生。PP支持1F1B调度。Sequence Parallelism采用Ulysses方案，和Megatron的SP不同。Automatic Tensor Parallelism（对HuggingFace模型自动做TP，但仅限推理）。

混合精度方面：FP16、BF16支持。

通信优化方面：1-bit Adam / 0/1 Adam / 1-bit LAMB。这部分梯度通信压缩可以减少通信量最高26x，是其区别于FSDP的重要feature之一。

此外，其也包括了后训练需要的feature，以DeepSpeed-Chat（RLHF pipeline，包括SFT和PPO流程）为主。除此之外，还包括了Activation Checkpointing，Gradient Accumulation，Sparse Attention。

特别值得一提的是，ZeRO论文本身是分布式训练领域引用最高的论文之一。ZeRO的三级sharding思想直接影响了PyTorch FSDP的设计。"改配置而不是改代码"的设计哲学影响了后续很多框架。

DeepSpeed并不包括Context Parallel和Ring Attention。此外，也没有FP8训练支持。有MoE支持但不如Megatron完善。社区维护活跃度近年有所下降，DeepSpeed-Chat被广泛反映问题较多。

Megatron-DeepSpeed（混合方案）

取Megatron的TP/PP+DeepSpeed的ZeRO，组合成工业界实际使用最多的大规模训练方案。

定位

面向需要同时使用高效model parallelism和memory-efficient data parallelism的超大规模训练场景。BLOOM-176B、MT-NLG-530B等知名模型的训练都使用了这个组合。

组成

其实就是Megatron-LM和DeepSpeed的集成方案。TP和PP由Megatron提供，ZeRO（通常是Stage 1，和PP配合时Stage 2/3有性能问题）由DeepSpeed提供，通信压缩等额外优化也来自DeepSpeed。

ZeRO和PP配合时通常只用Stage 1。Stage 2在PP下需要额外的reduce-scatter通信，开销大。Stage 3理论上可以但当时性能不够好。这主要是历史遗留问题，DeepSpeed官方后来承认如果重新评估可能会选Stage 3。

PyTorch FSDP/FSDP2（PyTorch）

PyTorch官方的ZeRO实现，在演化过程中从独立库演变为框架原生组件，是当前PyTorch生态中data parallel的标准方案。准确的说，它不是一个独立框架，而是一个组件，在演化过程中以特殊功能被嵌合在生态内。

定位

不是独立的预训练框架，而是一个并行策略组件。被TorchTitan、TRL、OpenRLHF等框架作为底层使用。面向所有使用PyTorch的训练场景。适合需要ZeRO级别显存优化但不想引入DeepSpeed外部依赖的场景。不适合需要TP/PP的场景，FSDP本身不提供这些，需要配合其他组件。

组成

纯PyTorch原生，不依赖任何外部库。FSDP1基于FlatParameter实现，FSDP2基于per-parameter DTensor sharding实现。

特性

FSDP1等价于ZeRO-3，也支持ZeRO-2和不shard（纯DDP）。支持CPU offload、mixed precision和activation checkpointing。

FSDP2在此基础上额外支持了per-parameter sharding（不再把参数flatten成一大块），和 torch.compile 兼容，显存效率比FSDP1高约7%，性能平均提升约1.5%。

本身不提供TP、PP、SP、CP等model parallelism。不提供训练循环、数据加载、checkpoint管理等框架级功能。不是一个完整的训练解决方案，需要搭配其他组件使用。

此外，值得一提的是FSDP的演化历史。它其实经历了三个阶段，FairScale（Meta的早期实现。这个时候是独立库）→PyTorch FSDP1（被PyTorch吸收进入主库）→FSDP2（用DTensor重写，和TorchTitan共同开发）。作为数学原理和DeepSpeed完全相同（都是 ZeRO 的三级 sharding），但起源和演化历史完全不同的对位存在，其实现方式和API设计与DeepSpeed的异同之处是一件很值得研究的有意思的事情。

实际上，它的生态位和前面的DeepSpeed高度重叠：前者本来就是以最小侵入为卖点，后者索性直接在PyTorch内原生内置，两者的数学原理又相同。FSDP在逐步蚕食DeepSpeed的市场份额。在2026年，对于大多数团队在2-8卡上微调7B-70B模型的场景，PyTorch FSDP已经成为默认推荐。这个转变在过去两年内发生得很快。

其他预训练框架

Nanotron（HuggingFace）

以简洁和教学价值为核心的minimalistic 3D并行预训练库。

定位：面向想理解3D并行实现原理的学习者和研究者，以及需要一个轻量级预训练工具的小团队。配套的Ultrascale Playbook是目前最好的分布式训练实践指南之一。不适合追求极致性能的工业场景，且feature覆盖不如Megatron。

组成：Python + PyTorch，使用torchrun启动，依赖Flash Attention和Triton。核心组件包括DistributedTrainer（训练编排）、ParallelContext（管理DP/TP/PP的process group）、PipelineBlock（PP的stage封装）、NanotronModel（模型基类）。

特性：

具有的特性：TP（模仿Megatron的Column/Row Parallel）、PP（1F1B调度）、DP、ZeRO-1、Activation Checkpointing、支持Llama/Mistral/Qwen等主流架构、支持和HuggingFace checkpoint的互相转换、支持Tied Parameters。

不具有的特性：ZeRO系列没有Stage 2/3的实现、没有EP/CP/SP、没有FP8训练、PP调度只有基础1F1B（没有Interleaved或Zero Bubble）、没有分布式Checkpoint resharding。

TorchTitan（Meta+Pytorch官方）

PyTorch官方的预训练参考实现，是为了展示PyTorch最新分布式特性而开发的showcase实现，同时也是一个clean-room implementation。

定位：面向研究团队和框架开发者，作为PyTorch分布式训练能力的展示平台和最佳实践参考，同时也是可用于生产的预训练系统。不适合的场景是需要MoE/EP支持的场景，以及需要完整后训练pipeline的场景。

组成：完全基于PyTorch原生API，核心抽象是DTensor（分布式tensor表示）和DeviceMesh（设备拓扑抽象）。使用FSDP2作为data parallel层。不依赖Megatron或DeepSpeed。

特性：

具有的特性：ZeRO-1/2/3（通过FSDP2这一原生实现获得）、TP、PP、CP（Context Parallel）、Float8/MXFP8训练（也即对Blackwell GPU硬件特性的支持）、torch.compile集成、Activation Checkpointing（layer级和operator级selective）、异步分布式Checkpoint（DCP）、SymmetricMemory（节点内高效的peer-to-peer通信实现）、AsyncTP（通信计算重叠）、SFT 支持、Flight Recorder（调试工具）、Gradient Accumulation。

不具有的特性：没有MoE/EP支持、PP调度没有Zero Bubble等高级方案、整体feature数量不如Megatron、没有通信压缩。

此外，其重要贡献和卖点之一在于DTensor+DeviceMesh的统一抽象（和FSDP2一样），即让不同并行维度可以作为独立的、可组合的模块叠加，模型代码几乎不需要改动。这个设计思想对整个PyTorch分布式生态有深远影响。

各框架核心特性汇总表

此表有ai辅助排版。感谢ai。如下：

特性 Megatron-LM DeepSpeed FSDP/FSDP2 TorchTitan Nanotron
─────────────────────────────────────────────────────────────────────────────────────────
维护方 NVIDIA Microsoft PyTorch Meta/PyTorch HuggingFace
定位 工业极致性能 低侵入易用 并行原语组件 官方参考实现 教学/轻量

DP ✓ ✓ ✓ ✓(FSDP2) ✓
TP ✓ ✓(弱) ✗ ✓ ✓
PP ✓ ✓ ✗ ✓ ✓
SP ✓ ✓(Ulysses) ✗ ✗ ✗
CP ✓(含Dynamic) ✗ ✗ ✓ ✗
EP (MoE) ✓ ✓ ✗ ✗ ✗

ZeRO-1 ✓(Dist.Opt.) ✓ ✓ ✓(FSDP2) ✓
ZeRO-2 ✗(需DS) ✓ ✓ ✓(FSDP2) ✗
ZeRO-3 ✗(需DS) ✓ ✓ ✓(FSDP2) ✗
ZeRO-Offload ✗ ✓ ✓(CPU) ✗ ✗
ZeRO-Infinity ✗ ✓ ✗ ✗ ✗

PP调度: 1F1B ✓ ✓ ✗ ✓ ✓
PP调度: Interleaved ✓ ✗ ✗ ✗ ✗
PP调度: Zero Bubble ✗ ✗ ✗ ✗ ✗

Activation Ckpt ✓(+selective) ✓ ✓ ✓(+selective) ✓
Mixed Precision-BF16 ✓ ✓ ✓ ✓ ✓
FP8 Train ✓(TE) ✗ ✗ ✓(Float8) ✗
Distr-Checkpoint ✓ ✓ ✗ ✓(DCP,Async) ✗
torch.compile Partial ✗ ✓ ✓ ✗
通信压缩 ✗ ✓(1-bit Adam) ✗ ✗ ✗
SFT支持 ✗(NeMo) ✓(DS-Chat) ✗ ✓ ✗

非NVIDIA硬件 ✗ ✓(ROCm) ✓ ✓(ROCm) ✓
代码简洁度 低 中 高 高 最高
学习曲线 陡峭 中等 平缓 平缓 最平缓

可以看出，即使是最简单的教学类预训练框架，其feature的复杂性也不比一般的推理框架少。对于后训练框架，这个问题还会更明显。也许，训练框架的开发在未来可能比推理框架的开发更具有可持续性。也就是说，Agent Infra的后续发展和完善也许会进一步成为后续几年的热点。不过，究竟会是以什么形式、挑战会是什么，这仍然是一个值得思考的问题。

LLM 学习日志 #3 训练、预训练和后训练基础

Mon, 20 Apr 2026 01:28:45 +0800

今天进行简单的训练框架的调研。和我们在日常的博客、知乎文章和其他零零碎碎的地方看到的开源资料不同，其实从Megatron-LM出世到现在，训练框架已经发生了长足的发展，乃至于接近发生了范式级别的变化，不过，就我个人的观察而言，目前的开源资料针对这些新变化的普及其实还不够充分。具体而言，至少在最近1~2年的时间跨度内，工业界的前沿热点已经逐渐从单纯的实现高性能、大规模的预训练，转向中等规模的微调、后训练等等面向不同具体领域和需求的后训练相关内容。

这个演化的方向和近1~2年中Agent的崛起密切相关，同时也和参数规模的膨胀相关。随着LLM的scaling law发力，在5年间，LLM的参数量已经从单卡可以推理的约100M增加到了令人发指的约1000B，而后者光是训练就需要以十万计数量的GPU，进行持续的为期数月甚至将近年为单位的后训练。许多更小的模型则随后从这个教师模型中提取知识进行蒸馏。这意味着一件事情：基模本身的训练越来越依赖于规模效应，而规模效应和硬件资源挂钩，硬件资源和财力上限挂钩。因此，越来越多的无法在这个维度上竞争的企业（包括企业内部的团队）转向后训练，即通过更小卡需求的方式和可控规模的训练数据量，实现基于基模的垂直领域性能表现提升。也因此，这一热点需求的改变反过来催生了训练框架范式的改变：从规整的、追求规模可扩展性、强调分布式设计和并行效率的大规模预训练框架，变成强调灵活协调性、追求复杂训练流程和支持不同训练算法的后训练和强化学习训练框架，同时可能还需要和推理框架有效整合。

值得注意的是，所有的框架几乎全部基于Python+Pytorch，底层通信依赖NCCL，通过PyTorch的torch.distributed封装调用。具体来说，PyTorch提供了 ProcessGroup、all_reduce、all_gather、reduce_scatter 等原语，各框架在此之上构建自己的并行策略。部分框架（Megatron-LM、DeepSpeed）还包含自定义的C++/CUDA扩展来加速特定算子。Triton kernel在较新的框架（TorchTitan、Nanotron）中也有使用。这再次提醒了Python语言（实际上是Pytorch平台）在LLM时代的统治力。即使它并不是性能最优化的语言，也依然在工业场景里没有被C/C++淘汰，很大程度上和早期的奠基者效应以及完善的生态依赖有关系。

什么是训练

准确的说，能够让模型性能提升的操作都叫训练。在实践中，它又分为两个领域，预训练和后训练。就像在前面提到的那样，前者追求极致性能和规模的scaling，而后者追求训练调度复杂性的处理。

预训练

预训练部分的技术重点是维度切分和集合通信。这一部分在先前的学习日志#1和#2当中已经比较详细的整理了，在此不做水字数的冗余讲解。就现在而言，在当下已经落地和可预见的大模型架构下，可切分的维度基本已经被完全发掘，数学的理论意义上已经几乎没有数量级层面的突破可能性。因此，这部分的热点目前在于“大规模”、“高性能”和“基于MoE的异构”。很明显，这些都远远超过了个人、小型团队甚至许多中型企业团队的能力范畴，这从近几年国内的有竞争力的基模团队和相应人才的演化趋势当中可见一斑：豆包（字节跳动-Seed）、Qwen（阿里巴巴）、混元（腾讯）、GLM（智谱华章）、Kimi（月之暗面）、Deepseek（幻方量化）、文心一言（百度），无一不是具有雄厚财力支撑下的结果（也许Kimi除外）。

后训练

监督微调（SFT）

监督微调虽然被归类为后训练，但其实其在训练的特性上和预训练没有过多的差别，甚至可以说很相似。它的通俗解释就是“对只会续写文字的预训练基模进行训练，使得其可以按照对话的形式，根据提问和上下文给出回答”。它的技术要点在于loss和梯度更新的mask：模型是不需要对固定的格式部分，例如<|assistant|>，只需要在assistant回复的内容部分上算loss。除此之外，没有特别需要注意的地方了，大部分工作都在于数据本身（system/user/assistant的多轮对话格式），而非训练的实现。

基于偏好的对齐（Preference Alignment）

在这个领域内，可以说产生的算法和算法家族是最多的，也是最复杂的。笔者本人在学习的时候，也很长时间没能整理出一个具体的分类关系和所以然来。经过反复的调研，个人认为以下的心智模型是最好的，记录在这里，供读者参考。

后训练中的对齐/强化学习方法，可以沿两个维度来分类：

第一个维度：用于评价好坏的数据来自哪里？是1、事先收集好的（offline），还是2、训练过程中模型实时生成的（online）。

第二个维度：reward 信号来自哪里？是1、人类标注的偏好，2、训练好的 reward model，还是说3、可验证的规则。

把这两个维度交叉起来，就可以对所有目前的主流方法进行无重复无遗漏的归类。

数据来源的两种模式

模式一：离线方法。这种模式意味着在训练开始之前，已经完整拥有了一个完整的偏好数据集，其中每条数据是一个prompt加上一个 chosen response和一个rejected response（正负样本对），这些数据是提前收集好的，其来源不一，可能由人类标注，也可能由更强的模型生成。这些数据接下来以其原始文本的形式直接被用于新模型的训练，而不是通过任何中间层学习（再重新生成）之后再用于训练，这是和在线方法的区别之处。

训练过程中，模型不需要生成任何东西。它直接读取这些固定的数据对，通过特定的loss函数（鼓励policy相对于reference更偏好 chosen）来更新参数。整个训练循环和 SFT 几乎一样，也和预训练更类似，其工作流大致是读batch、forward、算loss、backward、update。没有rollout，也就是模型不新生成这些固定数据对之外的数据以获得评价（在这个工作流下无法获得评价），不存在一个可以对模型本身生成内容的评估也就是没有reward model的在线打分。

实现简单且对数据量要求低，但这个模式存在distribution shift问题：偏好数据是由某个旧policy（或人类）生成的，但随着训练进行，模型被鼓励避免旧输出模式之后，当前policy的行为分布会偏离训练数据的分布。模型无法在自己新的输出模式上得到反馈。

模式二：在线方法。训练过程中，模型自己生成回复，然后在这些自己生成的回复上获得反馈并学习。一个online RLHF的step大概是这样的：

1. 从 prompt 池中采样一批 prompt
2. 当前 policy 对这些 prompt 生成回复（rollout）
3. 获取 reward 信号
4. 用 RL 算法更新 policy
5. 回到第 1 步

复杂很多，但模型始终在自己当前的分布上生成数据并获得反馈，不存在 distribution shift问题。不过，它本身有比较严重的reward hacking问题，因为reward model在泛化之后通常质量不如人类直接进行偏好标注。

Reward信号来源的三种模式

模式一：人类直接标注。最原始的RLHF，人类看模型的输出，给偏好排序。好处是reward质量高（人类直接进行选择），坏处是成本极高、速度很慢，无法规模化。早期InstructGPT/ChatGPT用的是这种模式，但现在已经很少直接用了。

模式二：Reward Model打分。先用人类标注的偏好数据训练一个reward model（本质上是一个回归模型，输入prompt+response，输出一个标量分数）。之后的RL训练中，用这个reward model代替人类做在线打分。PPO和GRPO传统上都用这种模式。好处是可以规模化，reward model速度比人类标注快几个数量级。坏处是reward model本身可能不准确，而且policy可能学会hack reward model，即找到一个方法，可以使得reward model系统性的给高分，但实际回复的质量对人类来说并不高。

模式三：可验证的规则，基于有ground truth的reward信号进行后训练，这就是RLVR。这种reward信号有几个非常好的性质。第一，它是ground truth，有客观评价指标，没有模糊空间，因此可以避免reward hacking问题。第二，它是完全免费的，不需要人类标注和reward model，一个Python脚本就可以直接进行验证并且输出+1或者-1的reward。第三，它是无限可扩展的，通过自动方法可以生成无限多的数学题和编程题作为训练数据。不过，并不是所有问题都是可验证的，许多语言风格、道德倾向等等的训练目标完全无法通过自动形式验证，因此这个模式可以做的训练目标本身受限。

Offline方法

DPO

虽然写在前面，但其实这个方法的诞生晚于后面提到的PPO。其核心思路是绕开reward model，直接用偏好数据训练policy，从而简化训练流程和在小样本下提升训练的稳定性和可行性。每条数据（数据对）包含三个部分，一个prompt + 一个chosen response + 一个rejected response。训练时需要：

用一个冻结的reference model分别对chosen和rejected计算log probabilities
用当前policy model同样分别计算log probabilities
Loss基于这四组log probabilities构造，鼓励policy相对于reference更偏好chosen

一个数据对的训练因此实际上要跑四次forward过程（ref×chosen、ref×rejected、policy×chosen、policy×rejected）。实际上，它很难真的和其他的后训练方法被放在一起，因为它实际上处理的是整个sequence level的优化，没有一个RL的过程，换句话说根本没有进入RL的框架。

此外，也存在SimPO和KTO这样的DPO变体，分别处理不需要reference model的情况和无数据对只有答案的好坏性质（不成对，单独存在）的情况，在此不逐一介绍，感兴趣的读者可以自行了解。它们的关系如下：

 数据格式 reference model 长度处理
─────────────────────────────────────────────────────────
DPO 偏好对 需要 sum（有长度偏差）
 (chosen, rejected)

SimPO 偏好对 不需要 average（无偏差）
 (chosen, rejected) + margin

KTO 逐条标注 需要 sum
 (desirable 或
 undesirable)

Online RLHF

PPO

这是最经典和最复杂的后训练算法，其思想直接借鉴了早期的RL相关的理论和实践。其完整流程涉及四个模型：

Rollout 阶段（推理）：policy model对prompt做自回归的generation，生成回复
Reward 阶段（推理）：reward model对生成的回复打分
Advantage 估计：value model估计每个token的 value，结合reward算GAE
Policy update（训练）：用PPO clip loss和相对reference model的KL散度更新policy，同时更新value model

区别在于多出了一个value model和一个reward model。它们是协作关系：reward model需要对每一个policy model的Prompt整体给出一个标量估计，而value model负责对于每一步状态给出一个估计，即从这个状态继续出发按照现有模型的策略和分布，最终期望能拿到什么程度的reward model，因此实际上是一个把reward从sequence的整体level分解到任意中间状态的辅助评估者。

实际上，从这个视角观察的话，reward model和value model在思路上有相似之处，都是基于模仿的扩展者。reward model从标注的数据对中学习，把少量的sequence泛化到对任意整体sequence都有一个标量可供评估，而value model从reward model中学习，把对任意整体sequence都有一个连续分布的标量泛化到不完整的任意中间状态sequence都有一个连续分布的标量可供评估。我认为这个心智模型是有助于理解和记忆的。

值得一提的是，技术上的关键难点之一在于rollout阶段是自回归推理，需要KV cache、高效 batching、高效通信，以及很多推理相关的技术。这就是为什么veRL/OpenRLHF要集成vLLM。

GRPO

DeepSeek 提出的简化版PPO，核心改进是只使用三个模型，也就是说去掉了value model。对每个prompt生成一组（比如8个）回复，用reward model打分后，在组内做归一化得到advantage。也就是说，它和DPO一样，是一个sequence level赋予reward而不是token level进行评估的方法，其消除噪声的方法在于通过一组的密集输出和reward的归一化来将sequence level的不精确性部分抵消。它的最终性能上限不一定有PPO好，但简单的多。这样就只需要policy + reference + reward三个模型，不需要value model去做critic。实现上比PPO简单不少，效果也被验证过，目前非常热门。

RLVR（RL with Verifiable Rewards）

这是最近几年被Deepseek带起来的时兴架构，也就是DeepSeek-R1、Qwen等模型采用的路线。它和传统RLHF的区别在于reward不来自训练好的reward model，而来自可验证的规则（数学题的答案是否正确、代码是否通过测试用例）。工程上它依然是RL训练，因此需要rollout、需要reward scoring、需要policy update，但reward的来源从一个神经网络或者一个人类的标注结果变成了一个确定性的验证器。

RLVR的产生是LLM本身能力目标发生范式转移的标志之一。随着DeepSeek-R1的成功，它展示了纯RL（不经过传统的SFT中间步骤）就能让模型学会复杂的推理行为：长链chain-of-thought、自我纠错、"wait let me reconsider"这种反思模式。从此，推理能力成为新战场。2024-2025年的LLM竞争焦点从"对话能力"转向"推理能力"（数学、代码、逻辑），而推理任务恰好是最容易做verifiable reward的领域。这让RLVR有了天然的应用场景。

训练会用到的其他技术特性

除了1F1B这类调度特性、TP-PP-DP-CP-SP-EP的并行切分之外，其实还有许多无法简单以谱系方法归类，但仍然值得一提的技术特性。

FP8精度训练

FP8精度训练是一种对计算进行优化的方法。在这个方法下，权重和激活以FP8格式直接喂给tensor core做矩阵乘法，GPU的计算单元直接使用8-bit精度做运算。H100/B200的tensor core原生支持FP8 GEMM，而且FP8的吞吐量是BF16的两倍（因为同样的硬件单元一个周期能处理两倍数量的8-bit元素）。因此，使用FP8训练既省显存（tensor变小了），又提升速度（FLOPS翻倍）。

具体来说，一个FP8训练的GEMM工作流如下：1、输入的X和W都是BF16格式，2、将X和W量化为X_fp8 = quantize(X / scale_X), W_fp8 = quantize(W / scale_W)，均为FP8格式，3、计算：Y_fp32_acc = X_fp8 @ W_fp8，其中Y_fp32_acc是FP32格式，4、输出Y_bf16 = Y_fp32_acc * scale_X * scale_W，其中Y_bf16和输入一样，是BF16格式。

在更细节的层面，FP8有两种格式，针对不同场景：

E4M3（4-bit指数+3-bit尾数）：动态范围较小但精度较高，用于forward pass的权重和激活。

E5M2（5-bit指数+2-bit尾数）：动态范围较大但精度较低，用于backward pass的梯度。这是因为梯度的数值范围波动更大，需要更大的动态范围。

LoRA和Q-LoRA

LoRA的思想相对简单：用低维度矩阵及其乘法结果的地址矩阵作为可训练量参与计算，即：$y=Wx+ABx$，其中A、B是可训练的小矩阵，$A: [r, d_{in}]$，$B: [d_{out}, r]$，其中r是rank，通常取8、16、32、64这样的值，远小于d_in和d_out（例如，~4096）。这里不多展开。

QLoRA的出发点是：LoRA虽然省了可训练参数的显存，但冻结的基座权重W仍然占着14GB（7B×BF16，即2）。能不能把这个也压缩呢？对此，QLoRA的做法是：把冻结的基座权重量化到4-bit存储，forward时动态反量化回BF16做计算，LoRA部分仍然以BF16训练。

具体技术有三个关键点。

NF4（NormalFloat 4-bit）。这是QLoRA论文提出的一种4-bit量化格式。它的设计思路是：预训练模型的权重分布近似正态分布，那么 4-bit 的16个量化档位应该按正态分布的分位点来划分，而不是均匀划分。这样信息论意义上最优，即每个量化档位承载的信息量相等。和均匀量化相比，NF4在同样的bit数下能更好地保留权重信息。

Double Quantization。量化时每个block（比如64个权重）需要一个scale factor（FP32，4bytes）。如果block size=64，那么scale factor的额外开销是4/64=0.0625bytes per parameter。虽然看起来不多但7B参数下也有~0.4GB。Double quantization是对这些scale factor本身再做一次量化（量化到FP8），进一步压缩这个开销。

Paged Optimizers。QLoRA 还用了NVIDIA的unified memory（CUDA managed memory）来管理优化器状态，当GPU显存不够时自动page到CPU内存，避免OOM，这其实和后面ZeRO-Offload的思路类似。

方法 基座权重 优化器状态 总计
─────────────────────────────────────────────────
全量微调 BF16 14 GB 84 GB ~120 GB（需要多卡）
LoRA BF16 14 GB 0.5 GB ~30 GB（单卡 A100）
QLoRA 4-bit 3.5 GB 0.5 GB ~10 GB（单卡 RTX 4090）

QLoRA让单张24GB消费级GPU就能微调7B模型，这对学术界和个人开发者的意义巨大。但它也有其局限性。第一，LoRA的低秩约束意味着表达能力上界。第二，QLoRA的激进4bit量化会带来不可忽略的精度损失，而且反量化的计算引入了额外延迟，训练速度通常比纯BF16 LoRA慢 20-30%。第三，LoRA需求一个训练好的预训练基座，因此不适合预训练。

Zero-1/2/3/Offload/Infinity

前三者不再赘述，可以参见之前的博客，这里讲后两个，他们都是关于存储层次利用的工作。先回顾ZeRO-3的状态：模型参数、梯度、优化器状态全部被shard到各张GPU上。每张GPU只持有1/N的完整状态，需要某个参数时通过all-gather临时拼回来。这已经把GPU显存的利用率压到很低了，但如果模型再大，即使shard之后每张卡分到的那1/N也放不下呢？

ZeRO-Offload的思路是：GPU显存不够，就往CPU内存搬。CPU内存通常比GPU显存大一个数量级（比如一台机器GPU显存总共640GB，但CPU内存可能有 1-2TB），而且便宜得多。

具体来说，ZeRO-Offload把训练过程中的不同计算和数据按照"计算密度"来划分——计算密集的操作留在GPU，内存密集但计算量不大的操作搬到CPU：

留在GPU上的：forward和backward的计算（矩阵乘法、attention等），因为这些是计算密集型，GPU的算力优势在这里是不可替代的。

搬到CPU上的：优化器状态（Adam的exp_avg和exp_avg_sq）和优化器更新计算（parameter update）。Adam的更新本质上是若干个逐元素运算，读四个 tensor、做加减乘除、写回去。计算量不大，但占显存的大头（fp32 master weights+两个矩估计=参数量的12倍字节）。把这些搬到CPU，用CPU做Adam update，然后把更新后的参数传回GPU。

关键的性能瓶颈在PCIe带宽。GPU和CPU之间的数据传输走PCIe，带宽大约32-64GB/s（PCIe Gen4/Gen5），比GPU内部的HBM带宽（几TB/s）和节点内NVLink（几百GB/s）差了一到两个数量级。所以ZeRO-Offload需要精心设计通信与计算的overlap，在GPU做当前层的backward的同时，把上一层的梯度通过PCIe传到CPU；在CPU做optimizer update的同时，GPU继续做下一个microbatch的forward。

ZeRO-Offload的收益场景是：GPU数量少但模型大。同时，ZeRO-Infinity是ZeRO-Offload的进一步延伸：CPU内存也不够的话，就往NVMe SSD搬。这个在工业上通常过慢，不具体展开。不过，其中的通算overlap和prefetch思想值得借鉴和考虑。

Activation Checkpointing

是一个很有意思的细粒度技术。它的主要思想是：保留部分中间计算值，丢弃部分中间计算值，以减少缓存压力。在没有保存中间值的部分，需要计算梯度的时候，则重新通过最近的激活值通过前向传播重新回到这个地方，进行计算，在计算完成后再丢弃。此外，和这个相关伴生的是反向传播情况下的Kernel设计和fusion相关的工作。几个著名的例子是Flash Attention的backward版本，以及optimizer fusion相关的Kernel，即通过合并Adam的四个需要更新的Tensor（param+grad+exp_avg+exp_avg_sq）的更新和合并逐元素操作来显著减少memory bandwidth压力；还有fused grad accumulation，即backward产生的梯度直接累加到梯度buffer里而不是先写出来再加，在gradient accumulation场景下省一次读写。

总的来说，推理的融合目标主要是减少kernel launch开销和memory bandwidth，训练的融合目标除了这些之外，还多了一个"减少中间 tensor的materialize以配合activation checkpointing的显存策略"。

Fault tolerance

这具体来说是指一类计数，即容错，也是典型的分布式场景的技术应用。这个概念其实远早于，但在LLM的训练场景下由于不同的workload特性被赋予了。在这一领域，目前经常采用的技术就是Checkpoint（存盘恢复），还有自动重启和弹性训练。此外，梯度和Loss的异常检测、静默故障检测、心跳检测和快速故障发现等技术也有所应用。这个领域目前的主要故障类型是硬件故障（ECC错误）、软件故障（NCCL通信超时和OOM）、静默错误，热点挑战之一是弹性的扩缩容和并行度的变化。例如，如果用TP=4、PP=2训练配置保存了checkpoint，故障恢复时某些机器挂了，随后想换成TP=2、PP=4重新开始，这时候就需要checkpoint格式支持不同并行拓扑之间的resharding，而这其实是一个non-trivial的课题。此外，网络层面的容错也需要进行相应的处理（慢节点、链路故障等等）。

Ray

准确的说这是一个分布式框架而不是一个计数特性，更准确的说它其实本来和LLM毫无关系：它本质上是一个分布式任务调度框架，其原始定位是通用的分布式Python程序执行框架。它的核心概念包括了两个基本类型：

Task：一个无状态的远程函数调用。写一个普通Python函数，加上 @ray.remote 装饰器，Ray会把它调度到集群中某台机器的某个进程上执行，返回一个future（异步句柄）。

Actor：一个有状态的远程对象。写一个普通Python类，加上 @ray.remote 装饰器，Ray会在某台机器上实例化它，之后就可以远程调用它的方法。每次调用都在同一个实例上执行，所以它有持久状态。

它最经典的应用就是在OpenRLHF，连接训练端的DeepSpeed ZeRO和推理端的vLLM。

Pico-vLLM 开发日志 #14 Prefix Caching

Sat, 18 Apr 2026 21:56:39 +0800

“如果你想体验极致的数学规整美感和爽感，就去实现Paged Attention；如果你想体验抽象分层设计原则和一致性维持的究极长痛，去实现Prefix Caching。”

——我自己

======================================================================

这么多天没更新，今天终于更新了。之所以没有和之前一样2~3天一更新，是因为这中间的间隔，几乎全部在对Prefix Caching进行架构设计、实现迭代和反复的DEBUG。个人的体感是，在所有的推理框架实现当中，Prefix Caching绝对是除了分布式中最复杂的部分（当然，这部分非常非常大，但对于许多feature来说，其实并不是和分布式本身强绑定的）以外的皇冠上的明珠，是实现最为复杂的、同时也是相当重要的架构级feature，其挑战在于其并非是局部的算子或者数据结构设计，而是同时牵动存储分配、请求生命周期、缓存一致性的全局性架构。尤其是，如果需求足够好的兼容性，也就是需求既有radix-tree的token级别的粒度和相应可扩展性，又需求底层的Block Manager能够天然、不需要额外处理的自然同时兼容radix-tree模式，会引入非常复杂的系统架构设计需求和要求非常小心的实现。因此，这里会完整的把我的开发工程、主要最后定稿的想法和设计模式、途中遇到的一些问题，尽可能在遗忘和丢失之前。如果能够帮助到后来者和读者的话，那就太好了。

架构设计：在 vLLM 的地基上建 SGLang 的房子

LLM 推理的一个核心优化是前缀复用，这已经被多篇论文证明在特定的工作场景下非常有效。它的思想很简单：多个请求共享相同的 system prompt 时，只需计算一次 KV Cache，后续请求直接复用，通过指向相同的KV Cache，实现在无复制、无额外空间占用的情况下完成相同功能的目的。从逻辑上来说，它是对Paged Attention和分页、页表思想的自然延伸。从工程上来说。SGLang和vLLM分别用两种截然不同的方式实现了这个优化。下面的描述中就可以看出，vLLM的架构设计模式更多是在naive Paged Attention上的一个优美的最小侵入改动以部分的模拟前缀匹配，而SGLang则就真的是毫无妥协的、原生的基于完整一体化设计的Prefix Caching。

在我自己进行实现的时候，一开始其实并没有意识到一个重要问题：两者的差距其实并不小，而且其实是基于各自的先决条件（Block Size的不同默认），也没有意识到为什么vLLM要使用哈希表这种不足够优美的讨巧的“模拟radix tree”的路径。我在刚开始这个feature实现的想法是：vLLM应该是出于性能原因，或者架构太难以改动，因此做出了设计上的妥协，我应该完全能够在SGLang风格下对我自己的框架进行修改。当中途我发现这个过程里面的坑多的令人头皮发麻的时候，代码已经改的太多，无路可退了。因此，我本人几乎是非常头铁的一路撞到了底，硬生生把东西全都做了出来。不过，如果让我自己再做一遍的话，我可能应该会直接采用vLLM模式：除了可扩展性略差之外，在性能上并没有实质性的差距，而且实现要简单的多。

vLLM模式

vLLM的设计采用的是block粒度的hash table。以block_size（通常16）个token为一组，计算hash后存入字典。匹配只发生在block边界上。如果两个prompt在block中间分叉（只要block的id是不一样的），整个block都被计为miss，不会进入前缀匹配的范围。

SGLang模式

SGLang用token粒度的radix tree。每个token都是树的一个路由单元，KV Cache的存储粒度（page_size）也是1个token。树和存储粒度一致，匹配可以精确到任意token位置，不存在对齐问题。

粒度不一致的桥接

Pico-vLLM的底层存储是vLLM 风格的BlockManager（block_size=16），配合手写的Triton kernel实现PagedAttention。但前面也提到，我想要的SGLang风格的匹配模式。也就是说，token级别的radix tree，任意位置分叉都能在radix tree级别正确处理。

问题是：token粒度的索引和block粒度的存储，天然不兼容。这个问题的处理占据了我大量的开发时间，也引入了相当多的bug。不过，我最终得以找到一个相对优美的解决方案，如下所示。

三层架构

最终的设计是一个三层结构，中间有一个显式的桥接层：

RadixTree（token 粒度）
 ↑ match_prefix: 在 token 级别遍历，返回 block-aligned 的结果
 ↑ insert: 在 token 级别分裂/合并，block_ids 按边界分配
 |
PrefixCache（桥接层）
 ↑ 向下取整对齐：matched_len = (raw_len // block_size) * block_size
 ↑ 双层引用计数协调
 |
BlockManager（block 粒度）
 ↑ allocate / free / inc_ref / dec_ref
 ↑ swap_in / swap_out（offload 预留）

RadixTree 的节点存储的 cached_blocks 是逻辑 block id 列表。一条边可能有 90 个 token 但只覆盖 5 个完整 block（80 token），剩余 10 个 token 的 KV 在第 6 个 block 的前 10 个 slot 里，但这个 block 不完整，不能被 prefix cache 复用。

match_prefix 在遍历完所有匹配的边之后，在返回前做一次统一对齐：

1
2
3


aligned_len = (matched_len // self.block_size) * self.block_size
aligned_blocks = matched_blocks[:aligned_len // self.block_size]
return aligned_blocks, aligned_len, last_node

这个设计把对齐逻辑集中在一个点上，无论树的内部结构怎么分裂合并，返回给外部的永远是 block-aligned 的结果。Engine 不需要知道 radix tree 内部的 token 粒度。如果后续实现写时复制，这个部分就可以进行改变，此刻以预留的形式存在。

insert 的 block 归属

insert 触发 split 时，需要把原节点的 cached_blocks 正确分配给 split_node 和 child。一条 96-token 的边在第 90 个 token 处分裂：

split 前：node [t0..t95] cached_blocks = [B0, B1, B2, B3, B4, B5] (6 blocks, 96 tokens)

split 后：
 split_node [t0..t89] cached_blocks = [B0, B1, B2, B3, B4] (5 blocks, 80 aligned tokens)
 child [t90..t95] cached_blocks = [B5] (1 block, 但只填了 6 token)

block 归属的计算：num_split_blocks = (i + match_len) // block_size - i // block_size。这个公式处理了分裂点不在 block 边界上的情况。前半段拿走所有完整 block，后半段继承剩余的。

语义设计：处理边界情况和分层解耦

（以下内容有ai辅助排版）

两层引用计数的必要性

SGLang 只有一层引用计数（lock_ref），因为它的 page 分配器和 radix tree 是一体的。vLLM 的 hash table 也只需要一层 block 引用。但 Pico-vLLM 的 radix tree 和 BlockManager 是独立的模块，而它们各自需要回答不同的问题，因其抽象层级和最细粒度同时不同。具体来说，其被设计的按如下职责工作：

BlockManager.logical_ref_count的含义是"这个block此刻有几个持有者"

持有者只有两类：正在运行的请求（通过 kv_cache 持有）和 RadixTree（作为一个整体实体持有）。ref归0意味着block立刻回收到 free pool。这是决定"block 是否可回收"的唯一依据。

RadixTreeNode.lock_ref的含义是"这个prefix段还有几个活跃请求在用"

lock_ref > 0意味着有请求正在使用这段前缀的 KV Cache。驱逐这个节点不会导致恶性的数据丢失（因为请求的持有阻止block被释放），但会引发额外的处理复杂性。在目前策略下，lock_ref归0意味着节点变得可以最小代价驱逐，进入LRU队列等待清理。

核心不变量：

logical_ref_count[block] =
 #{ 活跃请求 r : block ∈ r.kv_cache.logical_block_ids }
 + 𝟙[ 存在 radix 节点 n : block ∈ n.cached_blocks ]

这和文件系统的 inode + hard link 模型完全同构——BlockManager 是 inode table，RadixTree 是目录树。"block 什么时候可以物理回收"和"目录项什么时候可以删除"是两回事。RadixTree 无论内部有多少个节点引用了同一个 block（比如 split 前后的不同节点），在 BlockManager 看来都只算一个持有者。

lock_ref 的 node 指针方案

这是整个设计中最关键的决策，也是踩了最多坑之后才得出的结论。

最初的实现用 token 路径做 inc/dec：match 时沿 tokens 遍历 inc，close 时沿同样的 tokens 遍历 dec。看起来对称——但 insert 会改变树结构。

Bug 1：release 用 match 重新找 block

三个请求使用相同的 prompt。req4 首先 prefill，insert 把 blocks [A, B, C] 加入 radix tree。req5 在 submit 时 tree 为空（所有 submit 发生在 step 之前），matched_blocks = []，自己 allocate 了 [D, E, F]。insert 时发现路径已在，newly_held = []。

req5 close 时，release 调用 radix_tree.match(tokens) 重新找 block——找到了 req4 的 [A, B]，然后对不属于自己的 block 执行 dec_ref。A、B 的 ref 从 1 降到 0，被释放。req6 close 时又 match 到已经释放的 [A, B]，继续 dec，ref 变成 -1。

req5 close:
 match 返回 [A, B]（属于 req4 的 block）
 block_mgr.dec_ref([A, B]) → A=0 free, B=0 free ← 错误释放！

req7 submit:
 match 返回 [A, B]（tree 指向已释放的 block）
 adopt_blocks 读到 block_mapping[A] = (NONE, -1) → 崩溃

修复：release 不再 match，而是接收 held_blocks = kv_cache.logical_block_ids（请求自己实际持有的 block 列表）。

Bug 2：split 后边长不对齐 block_size

请求 A 的 prompt 有 144 token，insert 建了一条 144-token 的边。请求 B 的 prompt 前 125 个 token 相同，第 126 个开始不同。insert 触发 split：

split_node [t0..t124] len=125 ← 不是 16 的倍数！

请求 C 来了，match 走到 split_node，完整匹配 125 个 token。返回 matched_len=125。Engine 调用 adopt_blocks(blocks, 125)，assert 125 % 16 == 0 失败。

修复：match 返回前统一做 aligned_len = (matched_len // block_size) * block_size。

Bug 3：token 路径遍历的根本缺陷

即使修了前两个 bug，第三个更深的问题出现了。请求 R 在 submit 时 match 到 112 个 token（对齐后），inc_ref(tokens[:112]) 沿路径遍历。此时树有一条 144-token 的边，边长 > 112，inc_ref 走不到这条边，没有计数。

然后 R 的 insert 触发 split，把 144-token 的边劈成 split_node（125 token）+ child（19 token）。split_node.ref 复制了原节点的 ref。

R close 时，radix_held_len = 112 + ext * 16 = 144。dec_ref(tokens[:144]) 这次能走到 split_node（125 ≤ 144），对 split_node 执行 dec。但 inc 时没碰过 split_node——inc 了 0 次，dec 了 1 次，underflow。

inc_ref(tokens[:112]): 边长 144 > 112 → break，不 inc
 （insert 触发 split）
dec_ref(tokens[:144]): split_node 边长 125 ≤ 144 → dec → 0-1 = -1 💥

根本原因在于用token路径重新遍历做 inc/dec，但树的结构在match和close之间被insert改变了。 同样的token序列在两次遍历中走到了不同的节点。

修复方案：node指针方案

SGLang的lock_ref用的是node指针+parent chain，而不是token路径重遍历。最终采用了这个方案：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20


def match_prefix(self, tokens):
 ...
 last_node = self.root # 初始指向 root
 while ...:
 if match_len == len(edge):
 last_node = child # 只在完整匹配边时更新
 ...
 return matched_blocks, aligned_len, last_node

def inc_lock_ref(self, node):
 curr = node
 while curr is not None and curr is not self.root:
 curr.lock_ref += 1
 curr = curr.parent

def dec_lock_ref(self, node):
 curr = node
 while curr is not None and curr is not self.root:
 curr.lock_ref -= 1
 curr = curr.parent

请求对象上存的是 request.last_node（指针）而不是token路径。close时dec_lock_ref(request.last_node)通过parent chain向上遍历以正确处理引用数减少问题。

**这个设计方案的主要目的是为了确保引用技术操作在结构变化下永远对称。**split 只是在 node 和 root 之间插入了新的中间节点（split_node），Python 对象的 parent 指针会被正确更新。从 child 走到 root 的 parent chain 自动包含了新插入的 split_node，不需要重新遍历 token。

split 前：root → child [t0..t143] lock_ref=K
 inc_lock_ref(child): child.lock += 1

split 后：root → split_node [t0..t124] → child [t125..t143]
 split_node.lock = child.lock（纯复制）
 child.parent = split_node（指针更新）

dec_lock_ref(child): child.lock -= 1 → split_node.lock -= 1 ← 自动走过新节点，对称！

insert部分不改变lock_ref的计数。新节点创建时 lock_ref=0，分裂时 split_node.lock 纯复制自 child.lock。lock_ref 的增减完全由 match 和 close 控制，insert 是透明的。

权责清单

最终的修改入口和权责划分：

事件	BlockManager 引用计数动作	RadixTree 引用计数动作
allocate(n)	新 block +1	—
PrefixCache.match	matched_blocks +1	inc_lock_ref(last_node)
PrefixCache.insert	newly_held +1（RadixTree 新持有）	新节点 lock=0
PrefixCache.release	held_blocks -1	dec_lock_ref(last_node)
PrefixCache.try_evict	evicted_blocks -1（RadixTree 停止持有）	lock_ref=0节点从树删除

关键设计原则在于每一行的两列操作完全独立，互不依赖。BlockManager不知道RadixTree的存在，RadixTree不直接调BlockManager的free，PrefixCache是唯一协调两者的层。其中，加粗部分是可以根据不同的驱逐策略进行联合设计和改变的部分，只要确保可驱逐和操作成对合法，完全可以对这些部分进行改变，只要正确处理node废弃之后的状态、指针引用即可。

实际上，如果驱逐策略变得很复杂，两层引用计数的重要性将会开始凸显。这是因为实际上Radix Tree的引用计数本身不是释放是否合法的硬性指标（它是虚拟层级的最上层），而是“判断是否应该驱逐、应该驱逐哪些”的参照性标准。从技术本质上来说，在radix tree层面上驱逐一个计数不等于0的块是完全合法的，因为Block Manager此时并不会真的立刻释放它，而是给了Block Manager一个可以尽快在合法条件下按需将其释放的信号（或者，反过来说，radix tree是阻止block Manager将相应块释放的控制机制），这就赋予了极大的自由度（和处理被驱逐的废弃节点的相应活跃请求的复杂性）。对于需要考虑换入换出成本的radix tree的offload相关的驱逐策略，这还会变得更重要。因此，许多策略本身可以通过影响和控制相关节点的引用计数量，同时依然参照引用计数的大小来制定驱逐策略，在未来的扩展性设计中，这将节省大量开发成本。

设计模式总结，SGLang vs vLLM vs Pico-vLLM

	SGLang	vLLM	Pico-vLLM
索引结构	radix tree	hash table	radix tree
存储粒度	page_size=1	block_size=16	block_size=16
索引粒度	token	block	token
匹配精度	token 级	block 边界	token 级，返回前对齐
引用计数层数	1（lock_ref）	1（block ref）	2（lock_ref + logical_ref）
inc/dec 方式	node 指针 + parent chain	hash key 直接索引	node 指针 + parent chain
分裂处理	天然对齐（page=1）	不分裂	需要处理非对齐切分
驱逐单位	叶子节点	单个 block	叶子节点
驱逐策略	LRU	LRU	LRU + lazy deletion

SGLang 的优势：page_size=1 意味着树和存储粒度完全一致，不存在对齐问题。split 后每条边的 token 都有独立的 page，不会出现"半满 block"的归属问题。代价是 page table 更大（每个 token 一个 entry），管理开销更高。

vLLM 的优势：hash table 没有分裂合并，O(1) 查找，实现简单。block 粒度的管理开销最低。代价是匹配精度受限于 block 边界——两个 prompt 在 block 中间分叉时，整个 block 都 miss，浪费已计算的 KV Cache。

Pico-vLLM 的定位：在 vLLM 的 block-level 存储架构上获得 SGLang 的 token-level 匹配精度。代价是需要一个显式的桥接层处理粒度转换和双层引用计数。可以预计，这应该不是性能上最高效的方案，因为增加抽象层级和转换必然引入开销。不过，个人认为，这样引入的内存布局的兼容性可以在其他地方的可扩展性层面发挥重大作用。以我个人的经验而言，至少有一个地方非常需求这个层面的自由度调优兼容性：PD分离和多卡TP通信。在这两者的通信当中，NCCL需要拷贝后发送，而NIXL的零拷贝则需要按块为粒度进行发送，后者在block size过小的情况下会引入严重的launch overhead，从而极大拖累通信性能。因此，block size的自由调整和radix tree风格的前缀匹配的适配是非常重要且有意义的。

Benchmark和性能测试

好不容易把这么复杂的系统设计完之后，不做benchmark，看看Prefix Caching技术到底有何优劣，就太可惜了。设计两个不同长度的Prompt评测情景，对其进行性能分析。

短 Prompt（144 token）

3 种 system prompt（assistant / coder / analyst），每种 3 轮请求。第 1 轮 cold start，第 2-3 轮 warm。

hit_rate = 62.8%

短prompt下GPU prefill只需几毫秒，CPU overhead（Python侧tensor创建、scheduler调度、slot_mapping计算）占主导。prefix cache节省的GPU时间被CPU开销掩盖，加速效果有限。

实际上，对于小模型、非常短的prompt来说，由于前缀匹配引入了额外的CPU overhead的开销，用时不仅不会减少，反而会明显增加。在实测过程中，对于这个量级的prompt，大约会拖慢30~~40%的执行速度。换言之，Prefix Caching本身所引入的额外CPU overhead大约是~~5ms量级左右。如果Prefix Caching能够节省的内存量超过这个量级，那么Prefix Caching的策略就是值得的。

如果反向思考的话，这本身可能暗示了另一种策略的产生：如果对话不够长、模型足够小，那么也许直接通过某种分离策略来bypass正常的Prefix流程，按照之前的naive paged cache方法进入控制流，反而是更好的做法。

长 Prompt（2083 token shared + ~35 token variable）

固定一个 2083-token 的超长 system prompt（技术参考文档），8 个不同的 user query（各 30-50 token）。共享前缀占比约 98.5%。

共享前缀: 2083 tokens
可变部分: 28-48 tokens

 OFF(ms) ON(ms) speedup
Cold (req0): 48.92 41.59 1.18x
Warm (avg 7): 41.17 16.06 2.56x
Warm (best): 41.71 12.08 3.45x

2000+ token 的 prefill 让 GPU 计算成为瓶颈（约 30-40ms），CPU overhead 占比降到 20% 以下。prefix cache 省掉了 98.5% 的 prefill 计算，warm 请求的 TTFT 降到 cold 的 30-40%。

理论上界分析：prefix cache 的最大加速比是 total_tokens / new_tokens ≈ 2083 / 35 ≈ 60x。实际上的加速比只达到3.45x。这实际上是因为CPU侧的overhead实际上占据了相当大的比例，几乎有~5ms之多，这部分是无法消除的，其开销完全不会随prefix cache命中率的提高而减少（tensor创建、block table构造等仍需要执行）。此外，CUDA Graph 有固定的launch overhead，Python GIL下的 scheduler调度会产生额外开销，adopt_blocks 需要更新 GPU 上的 block table同样会引入额外开销，这些都会在我们评测使用的小模型上产生不可逾越的加速比gap上限。阿姆达定律告诉我们，最终的加速比不可能超过这些部分占的时间比例的倒数。考虑到这一点，Prefix Caching已经在职责范围内很好的完成了它的任务了。

总结和心得

到这里，所有Pico-vLLM计划内的所有主要计划内feature就已经都完整实现了，而且对于单卡的性能部分、PD分离部分的延迟抖动和异构并行度的正确性验证、Prefix Caching部分都有了完整的benchmark和Profiling。在工作过程中，也产生了若干个明显可以后续优化的点，包括了1、TP模式中的通信异步化和层间架构的重设计以允许通算重叠，2、PD分离的时候传递的后端替换为NIXL以评测性能变化，3、Scheduler的Prefill的Chunk策略和更具体调度策略的设计，4、Prefix共享block的写时复制（COW）语义的实现，5、GPU to CPU的换入换出的offload驱逐策略的实现，6、驱逐策略和Radix Tree结构的分离解耦，以及其他CPU侧代码的整体性能优化，都是未来可以不断扩充、实现和改善的地方。这些改善将在后续不断的进行，但集中的开发工程本身即到此为止。总体来说，这一个月的集中开发所取得的成果是令人感到满意的。

然后是心得。其实框架做到最后已经有sys的意思了。分层抽象，协同设计，性能和兼容性的取舍，每一个“传统”的sys设计思想，在推理框架里都有明确的体现。例如，怎么在开始动笔之前设计一个抽象层级，和它的精确的语义？这个很重要，语义设计不清就会权责不明，权责不明就产生耦合，产生耦合后续维护性就变差。有时候，这个问题比性能本身还重要。宁可性能损失3%，换取可维护性，大多数时候可能也是值得的。在这次的例子中，对于每一层的设计来说，应该具有哪些函数、其语义是什么，应该管理的东西有哪些、引用计数的来源和去向，预测的断言和边界条件情况，如果能够先尽可能的设计好再动手开发，可能反而会比现在更省心思、更省时间，也更不容易产生按下一个bug，又连带造成另一个bug，最后在修修补补中逐渐缠绕的一团乱麻的情况。在完善的开发流程中，最好都应该在纸面上先仔细想好、确保边界条件思考的尽可能全面再动笔，不仅可以省下后续大量的debug的无用时间，更可以提升最终成品的质量。本质上，在软件工程里，这种思想已经重复了无数次，但说到底绝知此事要躬行。此处教训相当宝贵，是应当牢牢记住的。

Pico-vLLM 开发日志 #13 PD分离（续）同步异步模式

Mon, 13 Apr 2026 22:31:38 +0800

这次迭代完成的是两个feature：异步模式的PD分离实现，以及和异构并行度下的PD分离-TP耦合的设计兼容。这部分做完之后，就只剩下基于Radix-Tree的Prefix Caching特性需要实现了，随后就进入整理、代码优化和交互界面设计阶段。这篇博客主要记录一下同步-异步模式两个的思路和实现，包括生产者-消费者模型；以及基于异构的PD分离-TP并行设计的具体实现。

什么是同步和异步模式

所谓同步模式，指的是"存在一个可序列化的顺序"，即整个请求的生命周期可以被描述为一个严格的线性序列：Prefill → KV Cache传输 → Decode。换言之，每一步不会在完成之前开始下一步，控制流是单一的。

而所谓异步模式，指的是"不存在这样一个全局的串行顺序"，各阶段之间通过显式的信号机制进行协调，而非通过阻塞等待来保证先后关系。Prefill端完成KV Cache的生产后，不需要等待Decode端确认接收就可以立刻开始处理下一个请求；传输过程可以和两端的计算并发进行；Decode端在KV Cache就绪后被通知并开始工作。控制流不再是单一的，而是分裂为多条独立推进的流，它们之间只在必要的数据依赖点上同步。

这两者的区别在于，当每一步不存在严格的对上一步的依赖，且对硬件资源的依赖不相互重叠的时候，即可通过这种方式进行某种意义上的重叠，从而提高性能。否则，两者的区别就不是很大，异步可能因为信号的轮询反而有性能的额外开销。

同步模式

同步模式的设计很简单。它的流程是：请求到达Prefill端，完整执行所有层的Prefill，生成全部层的KV Cache，然后一次性将KV Cache传输到Decode端，Decode端接收完毕后将请求加入Decode调度队列，开始逐token生成。它当然很容易保证和调试正确性，但它毫无疑问对于生产级来说是不可取的。

异步模式

异步模式的设计复杂度则比预想中高很多。它的最主要障碍反而不在于具体的编码实现，而在于设计模式的收发约定：未来还有多少个请求一定会要发送（发送方和接收方不再同步执行，那接收方怎么知道"该准备接收了"，怎么知道“不用再接收了”），以及后端的特性支持情况：保序性、是否支持tag以区分信道。不幸的是，对于这次实现的torch.distribute使用的nccl后端来说，它保序但不支持tag区分。而对于前者，实际上这没有可以简单解决的方法。

具体来说，NCCL的P2P是双边操作：isend必须和对端的irecv配对才能真正开始传输，单方面调用isend并不会把数据推出去。这意味着Decode端必须提前挂好irecv等着，但问题是它不知道Prefill端什么时候会发、甚至不知道还有没有下一个请求要发。

一个请求的KV Cache传输至少需要三次P2P通信：size、meta、data。在不知道后续是否还有请求的情况下，Decode端无法为下一个请求提前挂出irecv。这直接导致了一个硬约束：KV传输在请求级别是串行的，同一时刻只能有一个请求的传输流程在进行。 上一个请求的三次握手没有走完，下一个请求的传输就无法开始。

当然，更复杂的设计其实是存在的：比如在消息体中携带后续请求的元信息，或者建立一条独立的控制信道来提前广播传输计划，从而让Decode端可以预先挂出多个irecv实现传输的流水线化。但这已经不是推理引擎层面的工作了，而是进入了高性能通信协议栈的设计范畴，超出了本项目的边界。

NIXL的必要性

前面说了这么多，其实，目前的NCCL的局限已经体现的很明显了。这就自然引出了NIXL的需求：如果我们不需要双边的P2P语义，这个问题就根本不存在了。让ai老师介绍一下它的特长所在：

**NIXL（NVIDIA Inference Xfer Library）**是NVIDIA在GTC 2025上开源的、专门为分布式推理场景下的数据搬运设计的点对点传输库。它和NCCL的本质区别在于通信语义：NCCL的P2P是双边的send/recv模型，发送方和接收方必须配对才能完成传输；而NIXL采用的是单边的read/write模型——一个节点通过带外网络（如TCP或ETCD）导出自己的KV Cache元数据，其他节点拿到元数据后可以直接读写远端的GPU显存，远端甚至不需要在那一刻做任何配合操作。这意味着前面讨论的双边配对约束在NIXL的模型下天然不存在：Decode端可以在任意时刻、对任意数量的请求发起KV Cache读取，不需要和Prefill端一一握手。

在底层传输机制上，NIXL支持RDMA（InfiniBand和RoCE）、GPU-Direct RDMA、GPUDirect Storage、NVMe-oF、POSIX socket、以及S3兼容的对象存储等多种后端，并且会根据硬件配置自动选择最优的传输路径。一个关键的性能特性是，NIXL的传输不消耗GPU的SM资源——数据通过GPU-Direct RDMA直接在网卡和显存之间搬运，不经过GPU的计算单元，因此不会和模型的forward计算争抢算力。在实际benchmark中，NIXL在256KB到1MB的典型KV Cache传输大小上比NCCL快30%到50%。

NIXL的核心实现是C++，通过pybind11提供Python绑定，同时也提供Rust绑定。它目前已经被集成到NVIDIA Dynamo、TensorRT-LLM、vLLM、SGLang、LMCache等主流推理框架中，是NVIDIA推理基础设施栈中负责"搬数据"这一层的标准组件。

在后续实现中，也有计划在实现所有核心feature之后，将其变为PD分离的数据搬运的新后端，并且进行benchmark评测。

生产者-消费者模型

生产者-消费者模型是并发编程中最经典的协作模式之一。它描述的是这样一种场景：一方负责生产数据，另一方负责消费数据，两者通过一个共享的缓冲区解耦，各自以独立的节奏运行。生产者不需要等消费者处理完上一份数据才能继续生产，消费者也不需要等生产者准备好下一份数据才能开始处理。只要缓冲区里有东西，消费者就取；只要缓冲区没满，生产者就放。两者之间唯一的同步点是缓冲区的状态：空的时候消费者等待，满的时候生产者等待。

这个模型对通信机制的要求也很明确：第一，需要一个共享的、线程安全的缓冲区；第二，需要一种通知机制让双方知道缓冲区的状态变化（通常是信号量、条件变量或事件）；第三，生产者写入缓冲区的操作和消费者从缓冲区读取的操作不应该要求对方的主动配合。换句话说，理想的通信模式是单边的，生产者侧和消费者侧不需要握手，也不需要提前约定通信的固定总量。

PD分离就是典型的生产者-消费者场景。Prefill端生产KV Cache，Decode端消费KV Cache，两者的处理速度天然不同（Prefill是compute-bound的大块计算，Decode是memory-bound的逐token生成），需要解耦才能各自高效运行。但从实现上看，能否真正实现生产者-消费者模型，取决于底层通信机制是否满足上述要求，特别是第三点：单边性。如果传输机制要求双边配对（如NCCL的isend/irecv），那生产者每次放数据都需要消费者同时伸手来接，这就把生产者-消费者退化成了一个双方必须同步握手的模型，缓冲区的解耦作用被大幅削弱。而如果传输机制支持单边读写（如RDMA），生产者只需要把数据写入一块预先注册的远端内存，消费者在任意时刻来读就行，这是生产者-消费者模型真正能发挥作用的前提。

基于异构的PD分离-TP并行设计

这部分的思想同样也很简单，就是让P侧占用的卡数量（并行度）可以和D侧不同，两者解耦。这么做的好处显而易见如果P和D的计算特性差异很大，那么，给两者分配一模一样的GPU数量很可能是不明智的。此外，它也会影响在特殊集群或者有掉卡情况下集群的弹性扩缩容。而如果P侧和D侧的并行度可以不同，那么整个性能调整就可以很自由，可以通过精细的调优最大化集群硬件的利用率。

KV Cache的重分片

KV Cache跨TP重分片的实现是核心的开发需求。问题的根源在于：TP并行下，每张卡只持有KV heads的一个子集。以Qwen2.5-1.5B为例，模型有2个KV heads，如果Prefill端TP=2，每张卡各持有1个KV head的Cache；而如果Decode端TP=1，单卡需要完整的2个KV heads。这意味着KV Cache的传输不再是简单的整块搬运，而是必须在传输过程中完成一次重新拼合或拆分——从Prefill端的多张卡上各取一部分，聚合成Decode端所需要的完整形状，或者反过来将完整的Cache切分后分发到Decode端的多张卡上。

设模型共有 $N_{kv}$ 个KV heads。在TP并行度为 $tp$ 的情况下，rank $r$ 持有的KV heads范围是：

$$ \text{heads}(r, tp) = \left[\frac{r \cdot N_{kv}}{tp},\;\frac{(r+1) \cdot N_{kv}}{tp}\right) $$

因此，Prefill端rank $r_p$（并行度 $tp_p$）持有的heads范围是：

$$ H_p(r_p) = \left[\frac{r_p \cdot N_{kv}}{tp_p},\;\frac{(r_p+1) \cdot N_{kv}}{tp_p}\right) $$

Decode端rank $r_d$（并行度 $tp_d$）需要的heads范围是：

$$ H_d(r_d) = \left[\frac{r_d \cdot N_{kv}}{tp_d},\;\frac{(r_d+1) \cdot N_{kv}}{tp_d}\right) $$

重分片的本质就是计算 $H_p(r_p)$ 和 $H_d(r_d)$ 之间的交集。对于Decode端的某个rank $r_d$，它需要从所有满足 $H_p(r_p) \cap H_d(r_d) \neq \emptyset$ 的Prefill端rank收集数据。具体来说，$r_d$ 需要接收数据的Prefill端rank范围是：

$$ r_p \in \left[\left\lfloor\frac{r_d \cdot tp_p}{tp_d}\right\rfloor,\;\left\lfloor\frac{(r_d+1) \cdot tp_p}{tp_d} - 1\right\rfloor\right] $$

而从每个这样的 $r_p$ 处，$r_d$ 需要取出的head局部索引范围（相对于 $r_p$ 自己持有的heads的起始位置）是：

$$ \text{local\_start} = \max\left(0,\;\frac{r_d \cdot N_{kv}}{tp_d} - \frac{r_p \cdot N_{kv}}{tp_p}\right) $$

当 $tp_p = tp_d$ 时，所有公式退化为 $r_p = r_d$ 的一对一映射，不需要任何重分片——这也是同构PD分离的特殊情况。

写成公式就是这么几行，但翻译到代码里就会发现：每个传输操作的source rank、目标rank、在KV Cache tensor上的slice范围、以及block_table中的物理块重映射，全都要从这些公式算出来，而且必须在每种 $tp_p$、$tp_d$ 的组合下都正确。这还只是考虑了p、d互相的差异是整数倍的情况。如果不是整数倍、或者并行度大于KVhead的数量，其复杂度将会加倍。非整数倍意味着某个Decode端rank需要的heads会横跨某个Prefill端rank持有的heads区间的中间位置，产生不对齐的切片；而并行度大于KV head数量则需要引入GQA的head重复映射，同一个物理KV head的Cache要被复制到多张卡上。在pico-vllm的实现中，默认p和d的分配一定是整数倍的。

此外有意思的是，vLLM的异构并行度PD分离支持长期处于实验状态，相关的bug和兼容性问题从2024年底的路线图规划到2025年下半年仍在持续修复中。也许从中可以一窥跨并行度KV重映射在边缘情况下的实现难度有多令人头疼。

Pico-vLLM 开发日志 #12 PD分离

Sun, 12 Apr 2026 21:10:14 +0800

PD分离可以说是最近几年最为火热的llm推理引擎架构创新了。从Meta到Openai到Anthropic，内部几乎都肯定有完整的PD分离框架，其中Meta甚至因为这个发过相关的技术论文。它的核心卖点在于推理过程中的一个计算特性区分：Prefill是一个计算强度和prompt length相关的操作，是计算密集的，同时是单步重负载的；而Decode是一个计算强度几乎为固定值的操作，同时也是单步相对轻负载的。在前面的设计我们知道，基于Continuous Batching技术+Chunked Prefill会引入Prefill和Decode的混编，而这种混编会引起TTFT的抖动，以及其他方面的性能影响。如何解决它？答案意外的简单，把它们物理上分开不就行了？

但有意思的是，vllm直到最近还在把PD分离作为一个有experimental标签的“实验性”特征，25年8月才修好P侧和D侧的TP异构导致的不兼容问题。对于一个如此热门的开源项目来说，似乎进展相当慢了。所以在做之前的当时就已经在怀疑，是否可能有坑？做到目前来看的结果，确实是相对来说复杂度较高的feature，原因主要不是技术原理，而是工程实现的复杂度，尤其是和其他并行模式的耦合的复杂度。下面主要记录一下这部分产生的心得。

为什么需要PD分离

这个问题不难理解，在前面也讲到了，本质是计算强度的不同，导致的最优化的ridge point完全不同，进而导致硬件资源的非充分利用。不过在阅读论文之后，我发现其实还有一个原因，就是用户端的体验问题。下面将会逐一讲解。

TTFT和ITL的权衡

TTFT（Time To First Token）衡量的是从用户输入prompt开始，到模型开始生成第一个Decoding Token的时间间隔。这本质上是一个人机工程学指标：毕竟硬件的效率就摆在那里，TTFT是一个“权衡”，而不是算力上限的最精确衡量。但是它对于用户来说是非常重要的：我们在TTFT时间之后才可以真的开始阅读LLM生成的内容、进入非闲置的工作状态，在此之前是纯粹无效等待。它也很可能影响“系统是不是卡了”的直觉判断，一个TTFT高或者不稳定的LLM框架即使后续生成很快，也会无法高效维持用户的注意力。

ITL（Inter-Token Latency）是生成过程中相邻两个token之间的间隔。它决定的是"阅读体验是否流畅"。它同样会影响到阅读体验：如果大量Prefill突然加入，ITL就一定会升高（这是不可避免的），从而让阅读被打断。更重要的是人在阅读情况下对ITL比起TTFT更敏感，因为人没有被打断的心理预期。

这两者之间构成trade-off。如果为了保护ITL，限制每步最多只做一个Prefill（或者限制Prefill的token数），那TTFT就会变长，因为新请求要排队等。如果为了降低TTFT，激进地插入Prefill，那正在Decode的请求的ITL就会被打断。这是一个调度层面无法根治的矛盾，其问题在于“不稳定性”和“不可预测性”。只要Prefill和Decode在同一张卡上，这个跷跷板就存在。

在Continuous Batching框架里，Prefill和Decode共享同一个GPU。当一个新请求进来做Prefill的时候，它的计算量远大于同batch里正在Decode的那些请求（一个Prefill可能处理几百上千个token，而每个Decode只处理1个token）。这意味着这一步的总耗时被Prefill主导，所有正在Decode的请求都必须等这一步跑完才能拿到自己的下一个token。结果就是：正在生成中的请求的ITL被新请求的Prefill拖长了。确保它的稳定性在非分离情况下需要非常非常复杂的调度策略，几乎接近于时序预测而非确定性的调度策略，这就自然催生PD分离需求。

硬件资源的充分利用

更深层次的原因是Prefill和Decode的计算特征差异大到它们的最优硬件配置完全不同，放在一起跑意味着任何时刻都有资源在浪费。正如前面说的，它们具有不同的总体算术强度。我们都知道，硬件的最优化利用是当程序处于其ridge point，或者接近其ridge point的时候。但一个硬件的ridge point显然只有一个，而Prefill-Decode两个不同需求有两个在Roofline模型里差异极大的算术强度点（具体可以参加前面的博客）。当Prefill在跑的时候，HBM带宽没有被充分利用（compute-bound，Tensor Core和CUDA Core是瓶颈）。当Decode在跑的时候，GPU的算力没有被充分利用（bandwidth-bound，HBM是瓶颈）。这意味着无论在哪个阶段，都有一部分的硬件能力在闲置。

PD分离带来一个潜在机会。如果我们进行恰当的分离，是否可以部署异构硬件，从而以相同的预算成本做相同的事情？Prefill用算力更强的卡，Decode用显存带宽更高或者成本更低的卡。实际上，最近几年Nvidia和华为也都在基于Prefill和Decode阶段的不同计算特性，设计具有不同ridge point的GPGPU（Rubin CPX）或者NPU硬件变体。它们基本上也就是为了匹配PD分离里分离不同算术强度的操作而设计的，避免资源闲置和同成本下的利用效率最大化。

即使是在软件层面，也已经有了（实际上更完善的）软件PD分离支持，这部分可以参考NVIDIA Dynamo（官方的PD分离调度框架）和NIXL（Inference Xfer Library），在此不过度展开了。

可用的工具箱

对于具体的PD分离实现，其实有很多种不同选择，以下几种均有实现的潜力。为了单人开发的复杂性和可维护性控制，笔者选择的是直接使用torch.distribute。然而，此处也列出工业界和其他可用的选项，供大家参考。这部分很大程度上是通过ai提供思路，然后再进行调研的，再次感谢ai。

使用torch.distribute（NCCL）

这是最容易上手的方案，也是之前Pico-vLLM在TP的开发阶段已经用过的工具。torch.distributed封装了NCCL的集合通信原语，提供了send()/recv()的点对点通信接口。

优点：开发成本低。NCCL对NVLink、PCIe、InfiniBand的底层拓扑发现和路径选择都是自动的，不需要开发者手动管理。

缺点：NCCL的设计初衷是集合通信，而不是点对点的大块数据传输。另一个问题是NCCL要求通信的所有进程在初始化时就加入同一个通信组（process group），这意味着Prefill Worker和Decode Worker必须在启动时就互相知道对方的存在。这在静态部署下没问题，但如果想要做动态的弹性扩缩容（比如根据负载动态增减Prefill/Decode Worker的数量），NCCL的这种静态组模型就会变得非常不灵活。此外，如果想在这个方案下做，需要先把传输内容gather到一起，进行传输，然后再在远端scatter。这意味着两次额外拷贝。

使用NIXL

NIXL是NVIDIA在GTC 2025上开源的、专门为推理场景设计的点对点数据传输库。它是NVIDIA Dynamo推理框架的核心数据搬运层。和NCCL不同，NIXL从设计之初就是为了解决推理场景下的KV cache传输这个具体问题的。

NIXL的核心抽象是Agent和可插拔的传输后端。每个Worker（无论是Prefill还是Decode）运行一个NIXL Agent。Agent向NIXL注册自己的内存区域（GPU HBM、CPU DRAM、NVMe SSD等），然后通过异步的transfer()接口发起传输。NIXL会自动选择最优的传输后端。

优点：专为推理KV cache传输设计；支持异步、非阻塞传输；支持非连续内存scatter/gather；多后端自动选择；与vLLM、SGLang、TensorRT-LLM等主流框架已有集成。

缺点：依赖较重（需要UCX、RDMA驱动等），在没有RDMA的环境下（比如我们的autodl PCIe环境）退化为TCP。NIXL目前仍在快速迭代中，API稳定性和文档完善度还有提升空间。安装也不是特别顺畅。此外更重要的是，它是C原生的，将其整合进入Python需要额外的大量努力。

使用CUDA IPC

CUDA IPC是CUDA原生提供的跨进程GPU内存共享机制。核心API是cudaIpcGetMemHandle()和cudaIpcOpenMemHandle()：前者从一块已分配的GPU内存中导出一个可序列化的handle，后者在另一个进程中用这个handle获取一个指向同一块物理显存的指针。

优点：同机部署下理论上是延迟最低的方案。因为它是零拷贝的，传输延迟为零。实现也相对直接，不需要额外的库依赖。

缺点：只能在同一台机器上使用。这是因为它需要同一块物理显存。如果要通过网络，自然就没有所谓的共享虚拟内存空间了。

设计模式

基于实例角色role的兼容式设计

这部分的想法很简单：在不额外增加函数的情况下实现不同实例角色的指派，而在实现上保持兼容。

实例具有三种可能的角色：“p”、“d”和“pd”。

传输层抽象分离

一种简单的想法是把传输实现在Engine里。它是最简单的，但思考之后我并没有采纳这样的方案。原因是传输的具体方案实在是太多了：就像刚刚说的那样，不同的后端实现、同步异步、和Scheduler策略的耦合、P侧和D侧的收发约定将会带来不可维护的复杂性。因此不如单独抽象一个基类，然后作为接口来开发。这是软件工程的想法，但值得一提。目前实现的是同步模式，但是，传输层的分离将允许未来进行异步和更加复杂的可扩展性。

注意事项和Bug

值得注意的Bug汇总

这里总结了几个遇到的bug，详细记录一下，以免后面忘记再踩重复的坑，同时也给后来的读者一个提示，避免有相同的错误。

1、Python的可变默认值问题

在开发过程中，遇到类似下面的bug。当提交若干个（超过1个prompt的时候），出现了具有奇怪但是固定模式的输出乱码。如下图：

 [Request 1] The capital of France is Paris Paris 11200
AAExample 112

 [Request 2] 1 + 1 = Paris Paris 11200
AAExample 112

输入的prompt分别是：

1
2
3
4
5


prompts = [
 "The capital of France is",
 "The capital of France is",
 "1 + 1 =",
]

其中第一个request的结果被彻底吞掉，而。这个我一开始以为是KV cache相关的部分又出了问题，因为太像了。经过长达三个小时的详细排查，最终发现了一个bug的特性：它重复的Paris这个首个Decode token的数量，和input的prompt完全相关。当输入的prompt变成：

1
2
3
4
5
6


prompts = [
 "The capital of France is",
 "The capital of France is",
 "The capital of France is",
 "1 + 1 =",
]

的时候，错误的输出结果变成了：

[Request 3] 1 + 1 = Paris Paris Paris 1112000
AAAExample

又多了一个Paris。这给了我缩小范围的想法：它很可能是很多个prompt的Decode结果，同时挂在了所有的prompt的generated_ids序列上。那么，它的工作机制是什么呢？先看原始request类的代码：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47


class RequestStatus(Enum):
 WAITING = "waiting" # 在 waiting 队列，还没做 prefill
 PREFILL = "prefill" # 本步正在做 prefill
 DECODING = "decoding" # 已经 prefill 完，正在 decode
 FINISHED = "finished" # 已完成

''' 请求对象，包含请求的所有信息和状态
- request_id: 请求 ID，唯一标识一个请求
- input_ids: 输入的 token ids，shape (1, init_seq_len)，包含整个 prompt
- generated_ids: 已经生成的 token ids，shape (1, )，初始为 空
- max_new_tokens: 最多生成多少个 token
- temperature: 采样温度，传递给 sampler
- top_p: top-p 截断，传递给 sampler
- kv_cache: 每个请求独享一个 KV cache 实例，存储生成过程中的 KV 状态
'''
class Request:
 request_id: int
 input_ids: List[int] # (1, init_seq_len)，包含整个 prompt
 generated_ids: List[int] # (1, )，包含已经生成的 token ids，初始为 空
 max_new_tokens: int
 temperature: float
 top_p: float
 kv_cache: PagedKVCache # 每个请求独享一个 KV cache 实例
 request_status: RequestStatus
 has_eos_token: bool # 是否已经生成 eos_token，scheduler 不直接接触 tokenizer 和 eos_token_id，这个由 engine 在 decode_step 后更新

 def __init__(self, request_id: int, input_ids: List[int], max_new_tokens: int, temperature: float, top_p: float, kv_cache: PagedKVCache, generated_ids: List[int]=[]):
 self.request_id = request_id
 self.input_ids = input_ids
 self.generated_ids = generated_ids
 self.request_status = RequestStatus.WAITING
 self.max_new_tokens = max_new_tokens
 self.temperature = temperature
 self.top_p = top_p
 self.kv_cache = kv_cache
 self.has_finished_notification = False # engine改变这个状态，scheduler根据这个状态改变 request_status和移出队列

 def is_max_len_finished(self) -> bool:
 return len(self.generated_ids) >= self.max_new_tokens

 @property
 def prompt_len(self) -> int:
 return len(self.input_ids)
 
 @property 
 def total_len(self) -> int:
 return len(self.input_ids) + len(self.generated_ids)

观察这一句：

1

def __init__(self, generated_ids: List[int] = []):

它涉及到一个Python语法：Python的函数默认值在函数定义时求值一次，不是每次调用时。[] 在 Python 解释器加载这个类定义的时候就创建了一个 list 对象，存在 Request.__init__.__defaults__ 里。之后每次调用 Request(...) 不传 generated_ids 时，用的都是同一个 list 对象。所以 r0.generated_ids.append(42) 之后，r1.generated_ids 也变成了 [42]，因为它们指向内存中同一个 list。这是一个非常隐蔽的bug，因为大部分测试原本是在B=1情况下验证，此时只有一个Request实例，共享不共享都无所谓。B>1时多个 Request 往同一个list里append，所有请求的 generated_ids[-1] 返回同一个值，从外部看起来，就是decode出了相同的token。

在修复了Request类中相关的错误代码之后，这个问题就成功解决了，框架开始输出正确的request结果。如图：

 [Request 0] The capital of France is Paris. The capital of the United States is Washington, D.C. The capital of Canada is Ottawa
 [Request 1] The capital of France is Paris. The capital of the United States is Washington, D.C. The capital of Canada is Ottawa
 [Request 2] 1 + 1 = 2
Example 2:
Input: nums = [1,2,3,4,

2、锁死问题

这个严格来说不是实现的bug问题，而是设计模式的问题。当我实现到一半，在写测试脚本的时候，我意识到一个问题：对于一个有并发的项目，接受方怎么确认不用再进行通信了？这个问题乍一看可以通过预先约定次数来解决，但实际上比这个复杂一些，因为很多情况下你根本没有办法提前预知次数，当然也就没法约定它。又考虑了，如果传递一个空请求，那么就自动结束，不再接收信息。这个逻辑在同步下有效，但是在异步下无效：你如何确认在Prefill端没有请求之后，就不会再出现请求了？也许只是短暂出现了波谷，而不是真的停止了。但如果这样不退出的话，那么实际上你没有办法在Engine内部设置一个正常的退出条件，只能通过狼狈的ctrl+c来强行终止。思来想去，最后的结论是：无法在Engine内部做出这个判断，它应该是用户层控制的部分。

为此进行了重构，现在，Engine在用完之后，用户需要显式调用一次“no more requests”函数，以此显式通知Prefill侧的Engine不再接受新request的submit，并且在清空流水线之后通知Decode侧，让它在清空流水线之后同样退出。这样，就比较好的解决了锁死问题。

Pico-vLLM 开发日志 #11 张量并行（续）数据对比分析

Sat, 11 Apr 2026 20:32:26 +0800

正如之前所说，在b200上用单卡和两卡重新测试了相同的1.5B模型的结果，进而得到了一些有意思的小结论。数据删去繁杂的细节之后（感谢ai），整理在下表中。

性能数据汇总

wall-clock端到端时间（含CPU开销）

硬件 tp 延迟 tok/s 加速比
───────────────────────────────────────────────────────────────────────
5090 PCIe (Py3.12+Xeon 8470Q) 1 2.74ms 365 —
5090 PCIe (Py3.12+Xeon 8470Q) 2 2.13ms 470 1.29x

B200 NVLink (Py3.10+Xeon 6960P) 1 1.64ms 610 —
B200 NVLink (Py3.10+Xeon 6960P) 2 2.08ms 480 0.79x

可以看到，前者可以达到一定程度的加速，而后者不仅没有加速，反而被拖慢了。这可能是什么原因？第一反应肯定是因为B200本身算力太高，分摊GEMM计算量的好处相当有限，而通信带来的提升不足以被PCIe到NVLink的配置升级所补偿。算力的升级幅度远大于通信的升级幅度，因此前者可以加速，而后者就不行了。

这个解释本身就是可以接受的。不过，还有一个疑惑：NVLink的性能比PCIe按理来说高很多，B200的性能，无论是算力还是带宽，按理来说也比5090高很多，但为什么在tp=2的情况下，两者的延迟如此接近，几乎完全相等（只有不到3%）？接下来的分析揭示了对小模型（以及不够大的大模型）来说另一个很隐蔽，但是可能同样也很重要的开销：CPU调度开销。

纯GPU replay时间（无CPU开销）

为了解答这个问题，设计了纯基于CUDA Graph的连续replay测试。这部分测试是为了揭示在没有CPU控制流开销的情况下GPU本身的性能如何，以反过来倒推CPU开销的大小如何对性能产生影响（这个的测试没有那么直截了当）。

见下表。可以看到一个非常有趣的细节：在5090+PCIe的配置组合上，端到端的开销和纯GPU开销相对来说较小：CPU侧的开销要么很小，要么可以被部分的overlap住。而对于B200+NVLink来说，其平均的CPU侧开销几乎是5090+PCIe的两倍以上。这主要有这样几个原因：

1、CPU主频不同。CPU侧的代码并没有那么高的并行度，堆叠核数的意义小很多。而Xeon 8470Q的睿频是约3.8GHz，与此同时Xeon 6960P的单核睿频是约3.2GHz。这两者的性能差异造成了很大的CPU侧时延不同，因此最终在并行度增加的时候抹平了GPU侧的性能优势。

2、Python版本的不同。更新的版本具有更完善的解释器优化，因此能够更好的降低CPU侧阻塞时间。

因此，CPU在现代大模型推理里真的不重要吗？其实不完全是。在模型较小、或者并行度够高以至于单步推理时间较为短的时候（以我们的测试为例，小于2.0ms的时候），CPU的overhead并不一定总是能够忽略的。因此，CPU侧代码的优化和恰当实现同样有其必要性。

硬件 tp 纯GPU延迟 GPU加速比
────────────────────────────────────────────────────
5090 PCIe 1 2.62ms —
5090 PCIe 2 2.12ms 1.24x

B200 NVLink 1 1.46ms —
B200 NVLink 2 1.89ms 0.77x

AllReduce 延迟（单次3KB传输）

可以看到，NVLink的单次通信固定延迟大约是PCIe的三分之一，而带宽是20倍，这也是一个有助于建立直感的数字。它也再次验证了计算机学科几乎通用的一个真理：降低延迟永远比提高带宽困难多了。小数据量的AllReduce是纯latency-bound，其对性能造成的损害于取得的收益相比不成比例（至少不是线性的），因此只适合在足够大的模型、足够好的配置下才值得用，尽管它是实现起来最直接也最方便的。

 固定延迟 数据传输 总延迟 56次/步
─────────────────────────────────────────────────────────────────
5090 PCIe ~28μs 0.06μs ~28μs 1.59ms
B200 NVLink ~10μs 0.003μs ~10μs 0.57ms

值得一提的细节

benchmark环境的隐藏变量不一定可以忽略，有时候影响显著。在最初测试的时候，B200 NVLink (Py3.10+Xeon 6960P)的配置在tp=2的时候甚至慢于5090 PCIe (Py3.12+Xeon 8470Q)的tp=2的情况，而这几乎说不通。经过检查，发现是环境变量中，OMP_NUM_THREADS设置不同（B200侧未指定，默认为1），autodl平台进行了专门的预先设置（指定为了50）但B200的环境中没有，这造成了明显的性能改变，通过拖累CPU侧的性能进而拖累了整个端到端的性能。

此外，这也验证了另一个问题：TP不是万能药，甚至有时候有害。对于小模型、网络延迟不够低的机器配置，尤其如此。

Pico-vLLM 开发日志 #10 张量并行 TensorParallel

Thu, 09 Apr 2026 23:19:56 +0800

在之前的单卡阶段结束和进行阶段性的总结之后，经过几天的休息（实际上是作者被导师叫去忙活论文投稿和下一篇论文启动的事情了），我们终于可以进入Pico-vLLM开发的下一个阶段：分布式的开发。而分布式的开发，根据计划，我们将要实现如下几个重要特性：张量并行TP，流水线并行PP和PD分离。

首先做的是张量并行，Tensor Parallel。其实我在一开始规划的时候，曾经想过要不要先做流水线并行。然而经过进一步的思考和分析，我发现了一个令人惊讶的事实：张量并行这个看上去切分粒度更细、通信要求和频率更高的方案，在实现难度上居然是最低的，而pipeline并行这种方案，需要的改动反而更大。下面会详细解释实现所涉及到的改动，以及实现的具体技术手段和注意事项、细节。

张量并行涉及到的改动

张量并行涉及到的改动意外的少，这个少是相较于其他并行方案而言。这其实是因为张量并行并不改变任何一张单个GPU上所需要加载的模型的结构：所有层数没有变，block没有变，变化的只有Q head和KV head的头数。

模型

模型的结构并不需要修改，仅仅将num_heads相关的参数和hidden_dim相关的参数修改成相应的local版本即可，层数等等结构相关的都不需要改变。

KV cache的存储

和先前的模式一模一样，每个GPU对应的程序仍然要单独管理自己的KV cache，有自己的Scheduler。区别在于，每个KV cache存储的slot大小变化了，因为每个GPU对应的head数量变化了（被分给了不同的GPU），其中缩小的倍率和TP的size相同。除此之外，几乎没有任何需要修改的地方，包括Triton Kernel。

引擎Engine

几乎没有需要修改的地方，只需要在初始化的时候给cfg传入相应的额外参数即可。

另一个需要注意的事情是，如果采用的是有随机性的采样，需要在每一张卡的程序实例初始化的时候，都进行相同的随机数指定。否则，不同的卡将可能会在Logits生成之后产生不同的采样结果，而整个sequence的生成Token是每个GPU上各自解码、各自保存的，每次Decode步之间并没有广播同步（为了性能），因此这将引发不同GPU间的序列不一致性，导致严重的问题。考虑到整体的设计模式，最好的指定是在初始化Engine实例的时候作为一个参数传入。因为用户不一定想到做这件事而且其具体影响不是很大，最好将其设置为一个可选的默认参数：比如42。

张量并行的实现手段

意外的简单，整个改造实现完全在pytorch内部进行。nccl相关的操作被torch.distribute完全封装，而且同样可以被CUDA Graph捕获，正常执行。

需要注意的技术细节

需要格外说明的是，CUDA Graph和torch.distribute+nccl backend会产生一些奇妙的化学反应。具体来说，CUDA Graph在capture的过程中，如果使用了nccl的后端作为torch.distribute的调用，那么CUDA Graph会直接固定的占用nccl的资源。这会导致在程序结束后，dist.destroy_process_group()无法正确完成，整个集群hang住，无法结束。为了解决这个问题，需要显式的删除CUDA Graph的对象：即执行del g, static_output类似的语句。这是一个说大不大但是很烦人的bug，花了一会儿才定位出来，一开始还以为是哪里同步错误，导致计算流本身有问题。

张量并行在Qwen2.5-1.5B上的效果

效果其实并没有预想中好，主要的原因是——虽然难以启齿——没钱。在2x5090，PCIe连接的的autodl平台上，总体来说加速比只有1.3x，效果并不是非常理想。通过Profiling排查原因可以看到，在有Profiling的情况下，通信（All-Reduce）占据了~48%的时间，几乎就是GPU time的一半了。这主要是因为PCIe连接的Latency问题：每次传输的内容其实只有几KB，数据量本身并不构成时间来源。但是，nccl+pcie的allreduce本身的overhead启动延迟就非常长，每次启动几乎需要26μs的时间，这在28层*2次调用的前向传播来说占据的时间长度几乎和真正的计算和内存存取一样多，从而大大拉低了潜力。这个问题没有什么好办法解决，是完全硬件上的。要解决倒也简单：NVLink，它的通信延迟比PCIe低很多。

画饼

想办法预定到了8卡B200的集群，真是太好了。现在是在autodl上做的，因为是PCIe互联，本身就不是非常适合TP这个并行模式，所以加速效果也有限，更何况只有2卡的5090。过几天马上就上去试试集群，看看Profiling的结果会不会有什么变化？能够预计，它能够达到的有效加速比应该比PCIe的情况高。

LLM 学习日志 #2 集合通信拓扑、原语和实现

Sun, 05 Apr 2026 16:39:11 +0800

这篇博客主要是重温一下之前学过的集合通信的基本模式，包括拓扑、原语和实现。说是“重温”是因为这部分在之前学习NCCL的时候其实就都已经粗略的学习过了，但之后没有用到过，因此掌握不熟练、印象不深。现在趁着Pico-vLLM的开发机会，把这部分知识彻底的掌握起来。

集合通信的原语

点对点（P2P）

Broadcast

Broadcast意味着一对多，一个GPU的数据广播给所有人。

GPU 0: [A B C D] → GPU 0: [A B C D]
GPU 1: [ ] GPU 1: [A B C D]
GPU 2: [ ] GPU 2: [A B C D]
GPU 3: [ ] GPU 3: [A B C D]

Reduce

Reduce意味着多对一，所有GPU的数据归约到一个GPU。要注意的是它不是拼接而是运算，其最终得到的数据的数据量和原本每个GPU自己的数据量是相同的。Reduce本身并不是一个具体的操作，而是一系列具有特殊性质的规约算子的集合，其满足的性质是为可结合（associative）且通常可交换（commutative）的二元运算。Reduce的常见算子包括了：

Sum（求和）
Prod（乘积）
Max / Min（极值）
Avg（平均）
BAND / BOR / BXOR（按位运算）
MinLoc / MaxLoc（极值+极值所在位置）

LLM训练中99%的Reduce操作都是求和，其他算子出现的极少。NCCL甚至不支持MinLoc / MaxLoc操作。

GPU 0: [A₀] → GPU 0: [A₀+A₁+A₂+A₃]
GPU 1: [A₁] GPU 1: [ ]
GPU 2: [A₂] GPU 2: [ ]
GPU 3: [A₃] GPU 3: [ ]

All-Reduce

All-Reduce意味着将所有GPU的数据归约，结果每个GPU都有一份完整拷贝。从逻辑上可以理解为先完成一个Reduce再完成一个Broadcast——实际上的操作在概念上类似但具体实现并不相同，同样是分解为两个阶段，但具体分解成的是Reduce-Scatter和All-Gather。

为什么不是一个Reduce+一个Broadcast？实际上，在早期的参数服务器模式下，确实就是如此。但它有一个问题：可扩展性。Reduce 阶段所有数据汇聚到一个节点（CPU），Broadcast 阶段再从这一个节点（GPU）发出去。这个节点（GPU）的带宽成为整个操作的瓶颈，其他 GPU 在等待时空闲。Reduce-Scatter+All-Gather模式实际上是通过一个线性的偏移把每一个节点（GPU）变成了一个$1/N$的参数服务器，每个只负责和自己编号相同的第$K$块的数据的Reduce和Broadcast。在集合通信的视角下，这恰好就是Reduce-Scatter+All-Gather的模式。这样讲，可能在概念上是最容易理解的。

GPU 0: [A₀] → GPU 0: [A₀+A₁+A₂+A₃]
GPU 1: [A₁] GPU 1: [A₀+A₁+A₂+A₃]
GPU 2: [A₂] GPU 2: [A₀+A₁+A₂+A₃]
GPU 3: [A₃] GPU 3: [A₀+A₁+A₂+A₃]

Reduce-Scatter

Reduce-Scatter操作意味着归约后，结果被切分，每个 GPU 只拿到一部分。但这种简化的描述方式实际上不完全准确，因为理解它在过程中的副产物，比理解它的开始状态和结束状态更有助于理解它是如何运作的。以最经典的ring为例子讲解，一张图会更容易帮助读者理解到底发生了什么：

每步：每个 GPU 向下一个 GPU 发送一块数据，同时从上一个 GPU 接收一块并求和。
规则：第 1 步发送自己编号对应的块，之后每步发送刚更新过的块。

初始状态:
 GPU 0: [A₀ B₀ C₀ D₀]
 GPU 1: [A₁ B₁ C₁ D₁]
 GPU 2: [A₂ B₂ C₂ D₂]
 GPU 3: [A₃ B₃ C₃ D₃]

第 1 步: GPU 0 ──A₀──→ GPU 1
 GPU 1 ──B₁──→ GPU 2
 GPU 2 ──C₂──→ GPU 3
 GPU 3 ──D₃──→ GPU 0

 GPU 0: [A₀ B₀ C₀ D₀+D₃ ]
 GPU 1: [A₁+A₀ B₁ C₁ D₁ ]
 GPU 2: [A₂ B₂+B₁ C₂ D₂ ]
 GPU 3: [A₃ B₃ C₃+C₂ D₃ ]

第 2 步: GPU 0 ──D₀₃──→ GPU 1 （发送刚更新的 D 块）
 GPU 1 ──A₀₁──→ GPU 2 （发送刚更新的 A 块）
 GPU 2 ──B₁₂──→ GPU 3 （发送刚更新的 B 块）
 GPU 3 ──C₂₃──→ GPU 0 （发送刚更新的 C 块）

 GPU 0: [A₀ B₀ C₀+C₂₃ D₀₃ ]
 GPU 1: [A₀₁ B₁ C₁ D₁+D₀₃ ]
 GPU 2: [A₂+A₀₁ B₁₂ C₂ D₂ ]
 GPU 3: [A₃ B₃+B₁₂ C₂₃ D₃ ]

第 3 步: GPU 0 ──C₀₂₃──→ GPU 1 （发送刚更新的 C 块）
 GPU 1 ──D₀₁₃──→ GPU 2 （发送刚更新的 D 块）
 GPU 2 ──A₀₁₂──→ GPU 3 （发送刚更新的 A 块）
 GPU 3 ──B₁₂₃──→ GPU 0 （发送刚更新的 B 块）

 GPU 0: [A₀ B₀₁₂₃=B_all C₀₂₃ D₀₃ ]
 GPU 1: [A₀₁ B₁ C₀₁₂₃=C_all D₀₁₃ ]
 GPU 2: [A₀₁₂ B₁₂ C₂ D₀₁₂₃=D_all]
 GPU 3: [A₀₁₂₃=A_all B₁₂₃ C₂₃ D₃ ]
规约完成。

因此，它的最终结果就是下图。但需要注意的是，它并不是一个没有副产物的操作。如果是原位操作的话，它实际上会修改每个GPU没有得到最终规约结果的相应数据位置，使得它不再是原本值，而是一个和自身rank决定的偏移量、数据偏移量都相关的部分累加和，也可以和某种意义上的前缀和相类比。

GPU 0: [A₀ B₀ C₀ D₀] → GPU 0: [A₀+A₁+A₂+A₃]
GPU 1: [A₁ B₁ C₁ D₁] GPU 1: [B₀+B₁+B₂+B₃]
GPU 2: [A₂ B₂ C₂ D₂] GPU 2: [C₀+C₁+C₂+C₃]
GPU 3: [A₃ B₃ C₃ D₃] GPU 3: [D₀+D₁+D₂+D₃]

All-Gather

每个GPU有一部分数据，收集后每个GPU都有完整数据。Ring All-Gather中没有任何归约（求和）操作，每个GPU拿到的块直接原样转发给下一个人就行。每步每个GPU把自己最近收到的完整块沿环传递，N-1步之后每个GPU就收集齐了所有块。

GPU 0: [A] → GPU 0: [A B C D]
GPU 1: [B] GPU 1: [A B C D]
GPU 2: [C] GPU 2: [A B C D]
GPU 3: [D] GPU 3: [A B C D]

All-to-All

每个 GPU 向每个其他 GPU 发送不同的数据。由于其$O(n^2)$的通信连接数（有时候还有通信量）和复杂度，其被称为是“集群性能的试金石和考验”。它几乎总是可能引起拥塞。另一方面，大部分运行在集群上的应用都会尽量避免产生频繁的All-to-All通信，因为它很可能严重拉低相较于规整的集合通信模式下的集群性能。对All-to-All的通信优化算法和工程实践思路以处理其产生的拥塞控制为主。此外，就像上一篇的博客中对于MoE的介绍一样，许多All-to-All是动态的而非静态的，从而进一步加剧了这个问题。

GPU 0: [A₀ A₁ A₂ A₃] → GPU 0: [A₀ B₀ C₀ D₀]
GPU 1: [B₀ B₁ B₂ B₃] GPU 1: [A₁ B₁ C₁ D₁]
GPU 2: [C₀ C₁ C₂ C₃] GPU 2: [A₂ B₂ C₂ D₂]
GPU 3: [D₀ D₁ D₂ D₃] GPU 3: [A₃ B₃ C₃ D₃]

All-to-All-v

每个 GPU 向每个其他 GPU 发送不同的数据，但收发的总数据量不一定每个都相同。

集合通信的常见拓扑

Ring（环形）

最经典的集合通信的拓扑模式。

优点是带宽最优——每个 GPU 的发送和接收带宽被充分利用，总通信量和 GPU 数无关（都是 2×(N-1)/N × 数据量）。

缺点是延迟随 GPU 数线性增长——步数是 2(N-1)，GPU 越多步数越多。

Tree（树形）

优点是延迟最优——只需 O(log N) 步。

缺点是带宽利用率不如 Ring——叶节点只参与部分通信，根节点成为瓶颈。

Double Binary Tree

两棵互补的二叉树同时工作，结合了Tree的低延迟和更好的带宽利用率。NCCL在某些场景下会用这种策略。但其复杂且难以理解，而且对集群的节点数量有更严格的要求。

Recursive Halving-Doubling

递归地把GPU分成两半，先在小组内归约，再逐步扩大。兼顾延迟和带宽。

集合通信的常见实现

XCCL：NCCL、HCCL、ACCL...许多初创的CCL通信库实际上都是直接魔改的NCCL。

NCCL 的 kernel 是沿三个维度预编译的模板实例：

算法（Algorithm）：决定数据流动的拓扑模式

Ring: 环形，带宽最优
Tree: 树形（双二叉树），延迟最优
CollnetDirect / CollnetChain: 利用 SHARP 等网内计算

协议（Protocol）：决定数据传输的同步和打包方式

Simple: 大块传输，高带宽，适合大消息
LL: Low Latency，8 字节粒度，用 flag 做同步，适合小消息
LL128: 128 字节粒度，利用 NVLink 的原子操作，延迟和带宽的折中

归约操作 × 数据类型：

Sum/Prod/Max/Min × FP16/BF16/FP32/FP64/INT8/...

这些维度的组合在编译时就被实例化成了大量的 kernel 变体。例如 Ring + Simple + Sum + FP16 是一个 kernel，Tree + LL128 + Max + FP32 是另一个 kernel。这也是为什么 NCCL 的编译时间很长。

Cost Model：运行时选择

NCCL 的 cost model 是默认调优决策的核心。这个模型以时间为指标评估集合操作的开销，用于选择正确的协议和算法。cost model 考虑多种因素，包括数据量、GPU 架构、拓扑结构、网络和算法属性。

MSCCL：微软用的一种DSL，基于NCCL自定义通信算法。

MPI：太经典不讲。

硬件互联的数感建立

这一节列举常见网络和网卡设备的所有规格和速度，以帮助读者，包括写博客的我自己，建立对于网络的数字感知。

一、GPU 节点内互联：NVLink

代际 首发GPU 每条链路带宽(双向) 每GPU链路数 每GPU总带宽(双向)
NVLink 1.0 P100 40 GB/s 4 160 GB/s
NVLink 2.0 V100 50 GB/s 6 300 GB/s
NVLink 3.0 A100 50 GB/s 12 600 GB/s
NVLink 4.0 H100 100 GB/s (注1) 18 900 GB/s
NVLink 5.0 B200 100 GB/s 18 1,800 GB/s
NVLink 6.0 Rubin 200 GB/s 18 3,600 GB/s

注1: H100 的 NVLink 4.0 每条链路实际是 50 GB/s 双向，但 NVIDIA 在产品规格中
 按 sub-link 口径标注为 900 GB/s（18 sub-links × 50 GB/s）。
 B200 的 NVLink 5.0 SerDes 速率翻倍，同样 18 条但每条 100 GB/s，所以总带宽 1.8 TB/s。

二、GPU 节点内互联：NVSwitch

代际 配套GPU 每芯片端口数 每芯片交换带宽 单节点GPU数 节点总NVLink带宽
NVSwitch 1.0 V100 18 900 GB/s 8 (DGX-2) ~2.4 TB/s
NVSwitch 2.0 A100 36 (注2) ~3.2 TB/s 8 (DGX A100) ~4.8 TB/s
NVSwitch 3.0 H100 64 25.6 Tbps 8 (DGX H100) ~3.6 TB/s
NVSwitch 4.0 B200 72 NVLink5 14.4 TB/s 72 (NVL72) ~130 TB/s (机柜级)

注2: DGX A100 有 6 片 NVSwitch 2.0，每个 A100 的 12 条 NVLink 分别连到 6 片 NVSwitch
 （每片 2 条），实现 8 GPU 全连接。

NVSwitch 3.0 (Hopper) 开始引入 SHARP（Scalable Hierarchical Aggregation and Reduction Protocol），支持在交换芯片内部直接做归约计算（如 All-Reduce 的 sum），数据不需要绕回 GPU。

NVSwitch 4.0 (Blackwell) 把 NVLink 域从 8 GPU 扩展到了 72 GPU（整个机柜），是一个质变——以前跨节点才需要 InfiniBand，现在机柜内全部走 NVLink。

三、CPU-GPU 互联：PCIe

代际 发布年份 每通道速率 x16 单向带宽 x16 双向带宽
PCIe 3.0 2010 8 GT/s ~16 GB/s ~32 GB/s
PCIe 4.0 2017 16 GT/s ~32 GB/s ~64 GB/s
PCIe 5.0 2019 32 GT/s ~64 GB/s ~128 GB/s
PCIe 6.0 2022 64 GT/s ~128 GB/s ~256 GB/s
PCIe 7.0 2025(规范) 128 GT/s ~256 GB/s ~512 GB/s

对比：H100 的 NVLink 总带宽 900 GB/s，而 PCIe 5.0 x16 双向才 128 GB/s，相差 7 倍。这就是为什么 TP 必须走 NVLink 而不能走 PCIe。

四、节点间互联：InfiniBand

代际 发布年份 每通道速率 4x 链路带宽 单/双向 典型网卡 典型交换机
SDR 2001 2.5 Gbps 10 Gbps 单向 - -
DDR 2005 5 Gbps 20 Gbps 单向 - -
QDR 2007 10 Gbps 40 Gbps 单向 ConnectX-2/3 -
FDR 2011 14 Gbps 56 Gbps 单向 ConnectX-3 SwitchX
EDR 2014 25 Gbps 100 Gbps 单向 ConnectX-4/5 SwitchIB-2
HDR 2019 50 Gbps 200 Gbps 单向 ConnectX-6 Quantum
NDR 2022 100 Gbps 400 Gbps 单向 ConnectX-7 Quantum-2
XDR 2025 200 Gbps 800 Gbps 单向 ConnectX-8 Quantum-X800
GDR ~2028 400 Gbps 1.6 Tbps 单向 (规划中) (规划中)

换算成更直观的单位（双向）：

EDR: ~25 GB/s 双向
HDR: ~50 GB/s 双向
NDR: ~100 GB/s 双向 ← 当前主流 AI 集群配置
XDR: ~200 GB/s 双向 ← 2025 年开始部署

一台服务器通常配备多块网卡。例如 DGX H100 配 8 张 ConnectX-7 NDR 400G 网卡，每 GPU 一张，节点间总带宽 = 8 × 100 GB/s = 800 GB/s 双向。即便如此，仍然只有 NVLink (900 GB/s) 的约 89%。

五、节点间互联：以太网（RoCE）

速率 双向带宽 延迟（典型） 备注
10 GbE ~2.5 GB/s ~10-50 μs 传统数据中心
25 GbE ~6.25 GB/s ~5-20 μs 
40 GbE ~10 GB/s ~5-20 μs 
100 GbE ~25 GB/s ~2-10 μs RoCEv2 常见配置
200 GbE ~50 GB/s ~2-5 μs 
400 GbE ~100 GB/s ~1-3 μs 正在部署
800 GbE ~200 GB/s ~1-2 μs 2025 年开始

以太网 + RoCE（RDMA over Converged Ethernet）是 InfiniBand 的主要替代方案。带宽可以做到类似，但延迟通常高于 InfiniBand，且需要额外的拥塞控制（PFC/ECN）来模拟 IB 的无损特性。

NVIDIA 的 Spectrum-X 平台就是基于以太网的 AI 网络方案，面向不想用 InfiniBand 的客户。

六、带宽层级

带宽 (GB/s, 双向) 技术 用途
─────────────────────────────────────────────────────────
8,000 HBM3e (B200) GPU 内部显存带宽
3,600 NVLink 6.0 (Rubin) 机柜内 GPU-GPU (未来)
1,800 NVLink 5.0 (B200) 机柜内 GPU-GPU
 900 NVLink 4.0 (H100) 节点内 GPU-GPU
 800 8×NDR IB (DGX H100) 节点间总带宽
 600 NVLink 3.0 (A100) 节点内 GPU-GPU
 200 XDR IB (单卡) 节点间单卡带宽
 128 PCIe 5.0 x16 CPU-GPU
 100 NDR IB (单卡) 节点间单卡带宽
 50 HDR IB (单卡) 节点间单卡带宽
 25 EDR IB (单卡) 节点间单卡带宽 / 100GbE

数量级关系:
 HBM 带宽 > NVLink > 多卡 IB 总带宽 > 单卡 IB >> PCIe

七、延迟层级

带宽之外，延迟同样重要：

操作 典型延迟
GPU 内部 SRAM (共享内存) 访问 ~几十 ns
GPU HBM 访问 ~100-300 ns
NVLink GPU-GPU ~1-2 μs
PCIe GPU-CPU ~2-5 μs
InfiniBand RDMA (节点间) ~1-3 μs
RoCE (节点间) ~2-10 μs
TCP/IP 以太网 (节点间) ~10-100 μs

LLM 学习日志 #1 并行方案和通信模式

Thu, 02 Apr 2026 17:20:54 +0800

在开始下一个阶段的开发之前，要先打好理论地基。虽然这部分之前零零散散的学过很多，但在通常的技术博客中通常只讲解每一种并行模式的技术原理和思想，将其和通信模式逐一详细联系起来并且做通信次数、数据量整合，甚至和反向传播的训练过程整合起来的并不多。

这篇文章打算稍微弥补一下这个问题，把所有零散的、常见或者不常见的并行切分方案，并且逐一评估其通信的需求、消耗、适用情况，在我的单人用户场景下，是否值得实现、实现难度如何。下面是总结。如果后面遇到了（甚至产生了）新的并行方案，也会一并记录在这里。

此外，在后面一篇博客当中，也会介绍一下每一种通信模式，以及通信模式对应的集群拓扑结构和相应常见的具体实现有哪些。

1、Tensor Parallelism（TP）— 层内切分

这个应该是最经典的并行方式了。它的切分方式就是把每一层的权重矩阵沿着head或者hidden维度切到多卡上，每卡算一部分，合并结果。不同的卡持有不同的head、或者FFN的权重矩阵的不同部分。

特别值得一提的是，Megatron-LM中一种经典的做法是，FFN中有两个线性层，可以把第一个线性层（W_gate和W_up）按列切、第二个按行切。这样的好处是，在中间的计算过程中，中间结果不需要通信。

它的流程如下：先用被切分的，只有自己这部分head的W_Q、W_K、W_V矩阵乘以输入得到自己head的Q、K、V，然后用Attention算出自己head的输出，然后把自己的head送入按照同样模式切分的、自己本地的相应的W_o部分做乘法，得到一个部分和，然后进行all reduce。值得注意的是，由于切分是相同的，因此绕过了单卡情况下，对head的全局cat操作因此无需通信，只需要cat自己本地的head就可以了。all reduce之后，送入自己本地的第一个FFN权重，然后逐元素激活，送入第二个权重，算出本地部分和，然后进行第二次all reduce。

通信模式：All Reduce

通信次数：每个Transformer Block里有两次All-Reduce（Attention之后一次，FFN之后一次）。总共次数为2*模型层数。

通信量：每一层中，涉及到的需要通信的，相关的每层数据量 = batch_size × seq_len × hidden_size × dtype_bytes。在整个模型中，总数据量 = 每层数据量 × 2（每层两次） × layer_num（层数）。对于ring All-Reduce的情况，总通信量为总数据量 × TP 度 = 总数据量 × 2 × (n-1)/n。

涉及的模块：Attention、FFN

优点： 1、可以达到很细的切分粒度，因为不同Head、FFN权重都可以切分，每张卡的显存占用可以按 TP 度线性减少。 2、实现很成熟，是最早的并行模式之一，有现实的工业级别落地。 3、天然是完全负载均衡的。

缺点： 1、通信频率很高，对卡间带宽要求很高，产生的网络延迟在非NVlink的机内和机间基本不可接受 2、切分过细的矩阵可能导致GEMM的算术强度降低，从而降低效率

训练中反向传播的情况

Megatron-LM给出了一个很漂亮的结论：对于所有前向的All-Reduce情况，在反向中的算子为Identity（不作数学操作，仅把同一个输入x复制到每张卡上）；对于所有的前向的Identity情况，在反向中的算子为。因此，在反向传播中，通信产生的次数和前向传播一模一样、数据量也相同，但出现的位置和前向传播相反。前向传播在结尾处做all reduce，而反向传播在开始处（当然同时是反向传播自己的结尾处）做all reduce。

这其实很符合直觉，简单推导即可得到，在此按下不表。

需要注意的是，反向传播的一个stage，在一般论文里大约占据两个slot，也就是两个时间步。换言之，被建模为前向传播的时间的两倍。

2、Pipeline Parallelism（PP）— 层间切分

Layer 分离 — 切分不同的Transformer Block

最简单直接的并行模式。也就是直接把模型的不同层分配到不同卡上。例如，对于一个32层模型进行对4卡的PP，那么每一张卡恰好持有 8层。数据从第一卡计算完成后送入下一张卡，以此类推直到流向最后一卡，像流水线一样。

通信模式：点对点（Point-to-Point）

通信次数：整个前向传播过程中总通信次数N-1次，N为卡数。

通信量：每次切换卡，单次卡间通信量 = batch_size × seq_len × hidden_size × dtype_bytes，一次前向传播的总通信量 = num_gpus × 单次卡间通信量。

涉及的模块：不涉及具体模块，按卡的stage进行划分，粒度为layers。

优点： 1、通信量小、频率低，可以其通信量和频率要求允许跨节点进行通信。 2、实现简单、计算规整，可以保证层内计算和单卡完全相同，不产生任何浪费和算术强度降低。 3、反向传播计算流和单卡情况相同，不需要对反向传播的组织模式进行特别的处理。

缺点： 1、存在流水线空泡（Pipeline bubble）问题，流水线的启动。 2、对于训练场景有额外的问题：反向传播需要反向流水线，交叉进行在简单的调度策略下会造成更多空泡，需要micro-batch、1F1B调度等等更复杂的策略。 3、对于层数不能被卡数整除的场景，会产生负载不均衡。此外，首尾层的embedding/lm_head可能产生lagging。

训练中反向传播的情况

Pipeline Parallelism的反向传播在数学上和单卡完全一致，每个stage各自做标准的反向传播就行。但难点在于调度：怎样安排各stage的前向和反向的执行顺序，使得GPU尽量少空闲。

需要了解的调度策略如下：

1、GPipe 朴素方案。先把所有 micro-batch 的前向做完，再统一做反向。会产生很多空泡，而且需要保存大量的micro-batch的激活值。值得注意的是：重计算想法（只存储部分中间激活值以节省显存）也是在这部分研究当中一同产生的。

2、1F1B（One Forward One Backward）在第一个forward做完之后立刻开始第一个backward的反向过程，和后面的forward过程穿插。通过组织调度数学上可以保证每个stage在任意时刻最多只需缓存PP个micro-batch的激活值。它的核心是：在稳定阶段，每个

3、Interleaved 1F1B 每个GPU不再持有连续的若干层而是持有总数相同但是在模型前向传播过程中不连续的层。这等价于增加了stage的数量，只不过可能有多个stage在同一张卡上。比如PP=4的情况下，8层模型、interleave=2；那么第一个stage（gpu）持有0、4层，第二个stage持有1、5层，第三个stage持有2、6层，第四个stage持有3、7层。它的好处是可以通过很简单的方式增加流水线深度，从而允许更少的流水线空泡，坏处则很明显：大大增加了需要的通信量，进而增加了设备需求和延迟。

4、Zero-Bubble Pipeline 这是一个很巧妙的想法。它把反向传播的过程中的成分做了更细粒度的划分。具体来说，反向传播其实包含两个独立的计算：

B（backward for input）：计算 ∂L/∂x，即对输入激活的梯度
W（backward for weight）：计算 ∂L/∂W，即对权重的梯度

前者被前一个stage依赖的，前一个stage只有获得了B才能开始自己的反向传播计算过程。而后一个实际上是悬挂在依赖链上的分支，它依赖后一个stage的结果但本身并不被其他人依赖。W的结果只是累加到本stage的grad buffer里，等最后optimizer.step() 时用。它什么时候算都行，只要在optimizer.step()之前算完就行。那么一个自然的想法就是：能否先算有依赖的B，让后面的先开始，在有空的时候自己慢慢算W？答案是可以的。

通过手工调度，或者基于**ILP（整数线性规划）**求解的自动调度，可以最小化bubble的数量，甚至于接近0bubble的比例。当然，它实际上会引入一个问题：由于W延后，实际上一个step的总Latency变长了。但通过约束，可以完全保证跨step重叠至多在一个step以内，对训练收敛几乎没有可观测的影响。

Prefill-Decode 分离（PD Disaggregation）— 切分请求的不同阶段

Prefill 和 Decode 的计算特性完全不同：

Prefill: 大矩阵乘法，compute-bound，高算术强度，适合高算力卡
Decode: 矩阵-向量乘，memory-bound，低算术强度，适合高带宽卡

混在一起跑会互相干扰。Decode的延迟敏感请求被Prefill的大计算堵住，Prefill的高吞吐被Decode的频繁小请求打断。

更进一步的说，两者的基本的算术强度就不同：Decode的算术强度一般可以估算为，FFN部分为engine中的batch数量的两倍（fp16或bf16意义下。fp8则乘以2）、Attention部分固定为1。实测数据显示，总体强度一般是40-90 FLOPs/byte左右，远低于H100 BF16的ridge point 295。而Prefill的算术强度随着seq_len同步线性增长，一般远远超过任何gpu的ridge point。

Decode 阶段

线性层（FFN、QKV 投影、Output 投影）

Decode 每步只生成 1 个 token，所以是矩阵-向量乘法（GEMV）：

权重 W: [d, d]
输入 x: [batch_size, d] （每个请求只有 1 个 token）

FLOPs: 2 × batch_size × d²
数据搬运: d² × bytes_per_param（权重）+ batch_size × d × bytes（输入，通常可忽略）

权重远大于输入，所以：

算术强度 ≈ 2 × batch_size / bytes_per_param

精度	batch=1	batch=8	batch=32	batch=128
BF16 (2B)	1	8	32	128
FP8 (1B)	2	16	64	256

Attention（decode 的 score 计算）

每个请求要用 1 个 query token 去和自己的整个 KV cache 做 attention：

Q: [batch_size, 1, d]
K cache: [batch_size, seq_len, d] ← 每个请求有独立的 KV cache

FLOPs: 2 × batch_size × seq_len × d
数据搬运: batch_size × seq_len × d × bytes（KV cache，每个请求各自的）

注意 KV cache 不能跨请求共享（每个请求的上下文不同），所以 batch 增大时数据搬运也等比增大：

算术强度 ≈ 2 / bytes_per_element ≈ 1（BF16）

不随 batch size 增长。

而Prefill的算术强度可以如此计算：

Prefill 阶段

线性层

Prefill一次处理整个 prompt，所以是矩阵-矩阵乘法（GEMM）：

权重 W: [d, d]
输入 X: [batch_size × seq_len, d]

FLOPs: 2 × batch_size × seq_len × d²
数据搬运: d² × bytes_per_param（权重）
 + batch_size × seq_len × d × bytes（输入激活）

当 batch_size × seq_len 足够大时，权重搬运的成本被大量计算分摊掉：

算术强度 ≈ 2 × batch_size × seq_len / bytes_per_param

精度	seq=1024, batch=1	seq=4096, batch=1	seq=4096, batch=4
BF16 (2B)	512	2048	8192
FP8 (1B)	1024	4096	16384

这些数字远远超过任何 GPU 的 ridge point。所以 Prefill 的线性层基本都是 compute-bound。

Attention（prefill 的 self-attention）

Prefill 中 Q、K、V 的长度都是 seq_len：

Q: [batch_size, seq_len, d]
K: [batch_size, seq_len, d]

Q×K^T 的 FLOPs: 2 × batch_size × seq_len² × d_head × num_heads
数据搬运: batch_size × seq_len × d × bytes × 3（Q, K, V）

算术强度大约是：

算术强度 ≈ seq_len / (3 × bytes_per_element)

这些差异是如此的明显，以至于这个disaggreation是最先被提出的。因此，分离之后各自可以针对性地优化硬件配置和调度策略。

Attention-FFN分离（AF Disaggregation）— 切分Transformer的不同阶段

AF分离是一个相对来说更新的分离模式，而其和MoE的诞生关系密切，是一个很细粒度的分离。这是因为Expert的“专家”指的就是不同FFN的专家，它们共享Attention的所有参数，但各自具有各自的FFN层和相应参数。更进一步的说，两者占据显存的东西并不相同，Attention部分主要是因为KV Cache占据了大量显存，KV cache在不同卡之间的来回拷贝可能造成Attention性能的严重下降；而FFN中，权重占据了显存的绝大多数（甚至可以达到80%以上）。对于一个MoE来说，显存的膨胀同时也允许（或者说需求了）更细粒度的分离模式。此外，MoE的通信模式本身就需要大量的、不规则的、不可预知的点对点通信，Attention→FFN的传输可以和All-to-All的dispatch合并成一步，这意味着AF分离的边际代价相较于dense模型的情况减少了。反正都是无法避免要付出的，不如一口气全做完。

总的来说，Attention更倾向于小TP度（保持GEMM效率，减少KV cache副本和通信，而且Attention并没有非常庞大的参数量导致的需求），而FFN/Expert想要大EP度（容纳更多expert）。可以说MoE天然适合AF分离——既在动机上需要它（Attention和Expert的并行度不匹配），又在代价上容忍它（All-to-All吸收了额外通信）。

3、Data Parallelism（DP）— 请求间切分

这应该是最朴素的并行了，朴素到不需要怎么讲。它的复杂度不在于推理（推理就完全是不同的实例在独自工作，无通信），而是在于训练的冗余参数存储的优化，这里诞生了许多大模型训练重要的早期成果——Google Brain、ImageNet、GPT-2、BERT，一直到单个模型在单卡上彻底放不下为止。

通信模式：推理无（训练则是All Reduce（经典训练DP）/All-Gather（ZeRO-1）/Reduce-Scatter（ZeRO-2）/两者（ZeRO-3））

通信次数：推理0，训练则是每个训练步 2 次集合通信（1 次 Reduce-Scatter + 1 次 All-Gather）（ZeRO-1）/ L+1 次集合通信（L 次 Reduce-Scatter（反向传播中逐层做）+ 1 次 All-Gather（更新后同步权重））（ZeRO-2）/ 3L 次集合通信（前向 1 次 All-Gather + 反向 1 次 All-Gather + 反向 1 次 Reduce-Scatter）（ZeRO-3）。

通信量：推理0，训练则是2 × Φ（可训练参数量） × dtype_bytes（经典DP和ZeRO-1和ZeRO-2）/ 3 × Φ（可训练参数量） × dtype_bytes（ZeRO-3）。

涉及的模块：推理无，训练涉及到优化器状态（ZeRO-1）、梯度状态（ZeRO-2）和权重（ZeRO-3）。

优点： 1、无敌简单和容易理解。 2、推理完全无通信开销。 3、天然是完全负载均衡的。

缺点： 1、对大模型的一个卡装不下的情况毫无帮助。 2、训练的反向传播过程依然需要通信。

训练中反向传播的情况

这是ZeRO-1、2、3和FSDP大显身手的地方。

ZeRO-1：切分优化器状态

切分方式：优化器状态（Adam 的 m 和 v）按参数均分到 N 张卡。权重和梯度每卡完整保存。

通信模式：Reduce-Scatter + All-Gather

通信过程：

反向传播结束后:
 1. Reduce-Scatter 梯度
 每卡得到自己负责的那 1/N 参数的聚合梯度
 通信量: Φ × dtype_bytes（每卡发出完整梯度，收到 1/N）

 2. 每卡用聚合梯度更新自己那 1/N 的参数（本地操作，无通信）

 3. All-Gather 更新后的权重
 每卡把自己更新的 1/N 权重广播出去，收集到完整权重
 通信量: Φ × dtype_bytes

通信次数：每个训练步 2 次集合通信（1 次 Reduce-Scatter + 1 次 All-Gather）

每步总通信量：2Φ × dtype_bytes（和经典 DP 的 All-Reduce 相同）

ZeRO-2：切分优化器状态 + 梯度

切分方式：优化器状态和梯度都按参数均分到 N 张卡。权重每卡完整保存。

通信模式：Reduce-Scatter + All-Gather

通信过程：

反向传播中（逐层）:
 1. 每层算完梯度后立刻做 Reduce-Scatter
 每卡只保留自己负责的 1/N 梯度，其余丢弃
 → 梯度显存从 2Φ 降到 2Φ/N

反向传播全部结束后:
 2. 每卡更新自己负责的 1/N 参数（本地操作）

 3. All-Gather 更新后的权重
 通信量: Φ × dtype_bytes

通信次数：每个训练步 = L 次 Reduce-Scatter（反向传播中逐层做）+ 1 次 All-Gather（更新后同步权重）

每步总通信量：2Φ × dtype_bytes（仍然和经典 DP 相同）

ZeRO-3 / FSDP：切分优化器状态 + 梯度 + 权重

切分方式：优化器状态、梯度、权重全部按参数均分到 N 张卡。每卡只存 1/N。

通信模式：All-Gather + Reduce-Scatter（前向和反向每层都要做）

通信过程：

前向传播（逐层）:
 对于第 i 层:
 1. All-Gather 收集完整权重 通信量: Φ_i × dtype_bytes
 2. 用完整权重做前向计算
 3. 丢弃非本卡的权重（释放显存）

反向传播（逐层，从后向前）:
 对于第 i 层:
 4. All-Gather 再次收集完整权重 通信量: Φ_i × dtype_bytes
 5. 计算梯度
 6. 丢弃非本卡的权重
 7. Reduce-Scatter 梯度 通信量: Φ_i × dtype_bytes
 每卡只保留自己负责的 1/N

全部反向结束后:
 8. 每卡更新自己的 1/N 参数（本地操作，无通信）
 （不需要额外 All-Gather 同步权重，因为下一步的前向会自动收集）

通信次数：每层 3 次集合通信（前向 1 次 All-Gather + 反向 1 次 All-Gather + 反向 1 次 Reduce-Scatter），总计 3L 次

每步总通信量：

前向 All-Gather: Φ × dtype_bytes （所有层加起来）
反向 All-Gather: Φ × dtype_bytes
反向 Reduce-Scatter: Φ × dtype_bytes
──────────────────────────────────────
总计: 3Φ × dtype_bytes

比经典 DP / ZeRO-1 / ZeRO-2 的 2Φ 多了 50%。多出来的就是前向传播中那一次 All-Gather——因为权重也切了，前向时也要收集。

4、Expert Parallelism（EP）— 专家参数间切分

这个同样是相对来说很容易理解的一种并行模式，也是MoE的必然选择，可以说这个并行模式和MoE架构是严格相互绑定的。这么做的主要原因有两个：第一，Expert的总量非常大，导致相同量级的MoE模型，计算量低得多，但所有Expert加起来的参数量远大于dense模型，导致模型根本不可能放在同一张卡上；第二，这些expert之间是独立的，不同token去不同的expert，中间过程不需要通信或者拼接，天然适合分布到不同卡上。

通信模式：All-to-All-v（数据量不规则）

通信次数：每个 MoE 层: 2 次 All-to-All（dispatch + combine），非 MoE 层（如 Attention）: 按 TP 的通信模式（All-Reduce）。

通信量：不可精确预测。每次All-to-All的通信量取决于token的分发模式，进一步地取决于router。理想均匀情况下:

每个 GPU 持有 batch_size × seq_len 个 token
每个 token 选 top-k 个 expert
总共需要分发的 token 数: batch_size × seq_len × top_k
其中留在本卡的比例: EP_degree 分之一
需要发送的: batch_size × seq_len × top_k × (1 - 1/EP) × hidden_size × dtype_bytes 

每个 MoE 层: 上述 × 2（dispatch + combine）

总通信量: 上述 × MoE层数。

但实际情况往往（高度）不均匀，热门 expert 收到的 token 多，冷门 expert 收到的少。

涉及的模块：FFN。

优点： 1、可以在不让计算量增加到难以承受程度的情况下，获得接近大参数量模型的性能。 2、切分粒度可以非常高。

缺点： 1、几乎总是一定有明显的负载不均衡，且无法在离线情况下静态的提前规划调度。 2、对显存的需求非常非常大。 3、训练比dense模型更加困难。

负载不均衡是MoE最难以处理的问题。Router是动态的，每个batch的token分发模式都不同。某些expert可能突然变"热门"，收到大量 token，而其他expert空闲。这会造成GPU的lagging情况，从而拖慢整个集群的性能。EP也是唯一一个通信模式动态变化的并行策略。这使得它的优化比其他并行困难得多，因为没法提前知道每一步的通信模式是什么样的。对它的调度策略比起前面的精密的线性规划模式，可能更接近于时序预测、控制论、排队论、资源调度等等范畴。可以想见，关于MoE模型的网络设计、集群设计、调度策略设计将会比前面的模式更有可以做的地方，可能也需要更长的时间才能在工业界收敛到一个相对来说足够优秀的统一策略，也可能根本不收敛。

训练中反向传播的情况

EP的训练问题并不在于技术上的问题——在DP、PP、TP等模式里已经基本上完全解决了。它面对的是专家坍缩和辅助负载均衡损失等等问题。它指的是一种Expert之间的马太效应：Router是可学习的，它在训练过程中会"偏心"——逐渐把越来越多的token路由给少数几个 expert，其他expert收到的token越来越少，最终几乎不参与计算。这是一个正反馈循环：某个expert因为偶然收到更多token，得到更多训练，变得更强，router就更倾向于选它，它就收到更多token……最终模型退化成一个dense模型，MoE的容量优势完全消失。更详细的介绍超出了这篇博客的范畴，因此就不多讲了。

此外，值得一提的是，Router的top-k选择是一个离散操作（不可微），但我们需要梯度回传。常用的做法是：

对被选中的expert，正常回传梯度
对没被选中的expert，梯度为零（它们根本没参与计算）
Router的梯度通过加权系数（gating weight）回传——这些系数是连续的、可微的

这意味着router的训练信号只来自被选中的expert，它对未选中expert的"好坏"一无所知。这也是expert collapse的一个根源。

5、Sequence Parallelism （SP） — 序列维度切分

SP是基于前文提到的TP的一种改进（同时也必须在TP的基础上做，无法单独存在。此外，它的名字实在是容易引起误解），它处理的是非FFN和非Attention的、没有被TP切割的地方，如LayerNorm的激活值、Dropout的激活值、残差连接的激活值等等，当然也是一个训练专门使用的优化手段。这是因为，激活值存储这个问题在推理里根本不存在。SP优美的地方在于，它不增加通信量，而是通过拆解通信原语的方式无代价的获得收益——也许可以被称为一种通信的Kernel Fusion？

回顾TP用到的通信模式：每个Attention/FFN模块结束时做一次All-Reduce。而在底层通信原语的逻辑上，All-Reduce其实一般是拆成两步实现的：

All-Reduce = Reduce-Scatter + All-Gather

SP 把这个拆分利用起来了：

不用 SP 时:
 TP 模块结束 → All-Reduce → 每卡得到完整结果 → LayerNorm(完整)

用 SP 时:
 TP 模块结束 → Reduce-Scatter → 每卡得到 1/N 的结果 → LayerNorm(1/N)
 → 下一个 TP 模块开始前 → All-Gather → 收集完整输入 → TP 计算

通信模式：Reduce-Scatter + All-Gather（替代 TP 原本的 All-Reduce）

通信次数：和 TP 相同，每个 Transformer Block 4 次集合通信（前向 2 次 + 反向 2 次），只是把 All-Reduce 拆成了 RS + AG

总通信量：Reduce-Scatter + All-Gather = All-Reduce。通过把All-Reduce拆分和还原，做到了在不增加额外开销的情况下减少大量激活值存储。

6、Context Parallelism （CP） — 长序列 Attention 切分

这是最近比较新的一种并行模式。CP应该是目前主流并行策略中最新加入的一维。Llama 3的训练就用了4D并行：TP + CP + PP + DP，CP 是第四维。

CP和Ring Attention技术密切相关。它是把输入序列沿token维度切分到多张卡上。每卡只持有序列的一段，以此实现利用并行性来加速的效果。

对于非Attention的操作（FFN、LayerNorm等），token之间没有交互，每卡独立计算自己的部分即可，完全不需要通信。唯一需要通信的是 Attention，因为每个token的Query需要和所有token的Key、Value交互。

Meta 在Attention模块中实现了两种CP变体，通常被称为Ring Attention。

Pass-KV（经典Ring Attention）

每卡保留自己的 Q 不动
K、V 块在 GPU 之间沿环传递
每轮：用本地 Q 和收到的 KV 块计算部分 Attention

标准Ring Attention，适合大多数场景。

Pass-Q（反向传递）

每卡保留自己的 K、V 不动
Q 块在 GPU 之间沿环传递
每轮：用收到的 Q 和本地 KV 计算部分 Attention

Pass-Q在某些特定场景下更优。当KV cache命中率低于5%时，Pass-KV更优。

Pico-vLLM 开发日志 #9 阶段性总结和开发过程的错误汇总

Tue, 31 Mar 2026 22:10:08 +0800

这篇简单的总结一下我们目前达成的里程碑、我们已经实现的功能和feature，同时也总结一下遇到的主要问题，提醒自己、也帮助后来人。

里程碑记录

注：有AI辅助排版。

开发时间线

从第一行代码到单卡收尾，总计11 天（3月21日 — 3月31日）。

Day 1–2（3/21–3/22）：从零搭建

从sampler.py开始，逐层手写Qwen2.5-1.5B的所有算子：RMSNorm、RoPE（复数形式）、GQA Attention、SwiGLU FFN，每个模块写完立刻和 HuggingFace逐tensor做torch.allclose数值对齐。完成权重加载、Engine、Sampler，在 GPU 上跑通第一次对话输出。

v0.1-naive-no-cache — 21tok/s，无 KV Cache的baseline。

Day 3（3/23）：KV Cache + Continuous Batching

实现NaiveKVCache，验证KV Cache在长序列下的加速效果（seq_len=3200时从1.8 tok/s提升到15.3tok/s，8.5x）。紧接着实现Scheduler和Continuous Batching，请求动态进出batch。

v0.2-naive-kvcache — 第一次 profiling，建立各算子耗时的数据基线。

v0.3-continuous-batching — Scheduler + Engine多请求调度跑通。

Day 4–5（3/24–3/25）：PagedAttention 数据结构 + 模型重构

实现BlockManager（全局block pool，分配/回收）和PagedKVCache（block_table + 逻辑→物理映射）。GQA Attention完全重写，拆分prefill和decode两条路径。RoPE 从复数形式重构为查表式实数实现。端到端数值验证通过。

v0.4-paged-attention — PagedAttention gather实现，性能和NaiveKVCache持平（这是因为GEMM主导，gather开销被掩盖）。第二次profiling。

Day 5–7（3/25–3/27）：Triton Kernel 密集期

手写两个Triton kernel：store_kvcache_kernel（2D grid，token × head）和Decode_Paged_GQAAttention_Kernel（online softmax + block table paged access）。完成QKV Fused Projection和SwiGLU gate_up Fused。将model.forward重构为纯tensor接口，所有Python对象操作移出forward，为CUDA Graph和torch.compile扫清障碍。

v0.5-triton-paged-attention — 23.2 tok/s。第三次profiling（nsys），首次量化 CPU launch overhead（851 kernels × 15μs = 13.5ms/步）。

v0.6-slot-mapping-store-kvcache — 28.4 tok/s，D2H memcpy 从560次降至0次。

Day 8–9（3/28–3/29）：CUDA Graph

先做实验验证CUDA Graph的收益预期（30 kernel串联实测6.62x加速），然后实现完整的静态buffer预分配 + capture/replay机制。Decode路径走CUDA Graph，Prefill保持eager。

v0.7-cuda-graph — 87.3 tok/s，4.2x加速。第四次profiling。

Day 10–11（3/30–3/31）：Triton Kernel 收尾 + 调优

手写Fused RoPE + Paged KV Cache Store kernel（+5 tok/s）、Fused SwiGLU kernel（+2 tok/s）、Triton RMSNorm（+2~3 tok/s）。尝试Fused Add RMSNorm，确认为负优化后回退。CPU侧block_table缓存优化。Engine收尾整合。

v0.8-97tok-cuda-graph — 96.94 tok/s，超过vLLM在同硬件上的95tok/s。

优化阶段 延迟 tok/s vs 起点
─────────────────────────────────────────────────────────
v0.1 Naive（无 KV cache） 51ms 21~1.5 1.0x
v0.2 NaiveKVCache 50ms 20 1.0x
v0.4 PagedAttention (gather) 48ms 20.8 1.0x
v0.5 Triton Paged Attention 43ms 23.2 1.1x
v0.6 + QKV/SwiGLU fused 35ms 28.4 1.35x
v0.7 + CUDA Graph 11.45ms 87.3 4.2x
v0.8 + Triton Kernel 收尾 10.32ms 96.94 4.6x

功能与特性实现

模型层

从零手写 Qwen2.5-1.5B 完整前向推理，不依赖 HuggingFace 的模型实现，仅使用其 Tokenizer 和权重文件。包括 RMSNorm、RoPE（查表式实数实现）、Grouped Query Attention（12 Q heads / 2 KV heads）、SwiGLU FFN，全部模块经过与 HuggingFace 参考输出的逐 tensor 数值对齐验证。支持 bfloat16 精度推理。

KV Cache 管理

实现了完整的 PagedAttention 内存管理体系。BlockManager 在启动时一次性预分配全局 KV Cache 物理内存池，按固定大小的 block（16 tokens）进行分配和回收，通过逻辑→物理的 block_table 映射实现非连续内存的透明访问，彻底消除显存碎片。每个请求持有独立的 PagedKVCache 对象，记录自身的 block_table 和序列长度，生命周期与请求绑定。

调度与 Batching

实现了 Continuous Batching 调度器，请求可以在任意时刻动态加入和退出 batch，不需要等待凑满固定 batch size。调度器维护 waiting、prefilling、decoding、finished 四个队列，每步自动完成状态流转。Prefill 和 Decode 走独立的执行路径，Prefill 逐请求处理（序列长度动态），Decode 支持多请求 batch 并行。

自写 Triton Kernel

共 5 个手写 Triton kernel：

store_kvcache_kernel：2D grid（token × head），通过 slot_mapping 将新计算的 K/V 直接写入非连续的物理 Cache 块，支持 stride 处理非连续内存布局。

Decode_Paged_GQAAttention_Kernel：Paged Decode Attention 的核心计算，在 kernel 内部通过 block_table 索引逐块读取分散存储的 KV，使用 online softmax 累积 attention 结果，天然处理 GQA 的 head 映射，零拷贝无需 gather。

fused_rope_and_cache_store_kernel：将 RoPE 旋转位置编码和 KV Cache 写入融合为单个 kernel，Q 在寄存器内完成旋转后写回连续显存，K 和 V 完成旋转后直接通过 slot_mapping 写入分散的物理 Cache 块，消除中间缓冲。

fused_silu_mul_kernel：将 SwiGLU 的 Sigmoid、门控乘法和 up 分支的逐元素乘法融合为单次显存读写，三次 I/O 降为一次。

_rmsnorm_kernel：将 RMSNorm 的 6 个独立 PyTorch 算子合并为单个 kernel，每行一个 program，全程 float32 计算后转回原始 dtype 写出。

算子融合与模型优化

QKV Fused Projection：三个独立的 q/k/v 矩阵乘法合并为单次 GEMM，权重在加载时拼接，输出通过 split 拆分。SwiGLU gate_up Fused：同理将 gate_proj 和 up_proj 合并。rotate_half 的 torch.cat 消除，改为 in-place 操作。model.forward 重构为纯 tensor 接口，不接受任何 Python 对象，所有 cache 操作留在 Engine 层，使 forward 对 torch.compile 和 CUDA Graph 完全透明。

CUDA Graph

Decode 路径完整接入 CUDA Graph。Engine 启动时预分配固定形状的静态 buffer（input_ids、slot_mapping、position_ids、block_table、context_lens），每步只原地更新 buffer 的值，不重新分配内存。CUDA Graph 一次 capture 后反复 replay，将每步数百次 kernel launch 压缩为单次提交。Prefill 路径保持 eager 执行（序列长度动态，不适合 capture）。提供 use_cuda_graph 开关，可随时回退 eager 模式调试。

采样

支持 Greedy、Temperature Sampling 和 Top-p (Nucleus) Sampling 三种采样策略，统一接口。

开发过程遇到的错误汇总

简单的总结了一下，整个单卡框架开发过程中遇到的主要的、难以排查的错误，我还记得的就这么几个——如果我不记得了，应该是没有折磨我太久，大家应该能解决。

注：有AI辅助排版。

类型一：KV Cache 写入位置错误

症状： 生成乱码，词语乱序重复，或反复覆盖同一位置。

出现了两次：

第一次（早期 PagedAttention）：

原因：context_lens 比实际少 1
decode 时写入新 token 后，读取范围没有包含刚写入的 token
新 token 的 query attend 不到自己

第二次（slot_mapping 重构后）：

原因：get_prefill_slot_mapping 里
cache_block_index 是 tensor，用 tensor 作为 dict key
Python dict 用 tensor 作 key 行为未定义，查到错误的物理 id
slot_mapping 返回 [0,1,2,3,4] 而不是真实物理地址
prefill 写入物理 block X，decode 读取物理 block Y
加 .item() 后修复，但这降低了性能，后续整体重构了

类型二：非连续内存导致 Triton Kernel 越界

症状： cudaErrorIllegalAddress，kernel 崩溃，或者生成完全不可理解（而不是简单反复重复）的乱码。

出现了两次：

第一次（paged_decode_attention）：

原因：kv_cache[layer_idx] 是大 tensor 的 view
Triton 拿到 data_ptr 后按连续内存计算偏移
实际 stride 和预期不符，导致越界访问
加 .contiguous() 后修复

第二次（store_kvcache，极其隐蔽）：

原因：apply_rope 之后 k/v 是非连续 tensor
view/reshape 假设连续内存
Triton kernel 按错误地址写入
根本修法：给 store_kvcache_kernel 传入 stride 参数
用 stride 计算偏移而不是假设连续

类型三：模型前向传播的过程实现在反复修改中出错

症状： 生成不完全的乱码，词语乱序重复，有意义但不完全有意义。

出现了一次，很隐蔽：

原因：q 相关的 forward 实现在不同路径里不一致，在 Decode 路径中遗漏了 RoPE 旋转状态
 - Eager 路径（forward）里，q 被 apply_rope 原地覆盖更新，传给 Attention 的是正确旋转后的 q
 - Graph 路径（forward_decode）引入了融合算子 fused_decode_rope_and_cache，它返回了新的 q_rot，但未被后续使用
 - 原始未旋转的 q 直接参与了注意力计算，导致模型在 Decode 时彻底丢失相对位置信息（RoPE 失效），引发无脑复读和乱码

出现位置：
 - GQAAttention.forward_decode 方法的末尾处
 - 算子调用错误地写成了 output = self._decode_attention(q, kv_cache_k, ...)
 - 实际应该传入旋转后的变量，改为 output = self._decode_attention(q_rot, kv_cache_k, ...)

类型四：Tensor Shape 错位

症状： RuntimeError: mat1 and mat2 shapes cannot be multiplied。

出现了多次，都是同一个模式：

原因：q 的 shape 在反复迭代修改后，在不同路径里不一致或者前后不一致。例如：
 - GQAAttention 内部 q 是 (B, seq, n_heads, head_dim)
 - paged_decode_attention wrapper 期望 (B, n_heads, 1, head_dim)
 - transpose 和 reshape 的顺序搞错导致维度错位

出现位置：
 - _decode_attention 里 q 传给 Triton 之前忘记 transpose
 - wrapper 里 unpack 顺序写成 B, _, N_HEAD, HEAD_DIM
 实际应该 B, N_HEAD, _, HEAD_DIM

类型五：seq_len 更新时机错误

症状： 生成内容在第 2 步之后开始乱码，第 1 步正确。生成的语句基本有正确的语义，但就是Decoding阶段开始的时候会出现重复、或者突然的错误。

原因：decode 每步用 context_lens = c.seq_len + 1
但 seq_len 在 forward 之前已经被 prepare_decode_step 
或其他逻辑提前更新，导致 context_lens 多 1 或少 1

另一次：prefill 后 _seq_len += len + 1（多加了 1）
第一个 decode step 的 position 偏移 1，之后全部错位

Pico-vLLM 开发日志 #8 更多Triton Kernel和性能优化

Tue, 31 Mar 2026 16:20:34 +0800

今天的主要工作是接入了更多的Triton Kernel，进一步完善代码和进行性能优化，并且把Decode的性能拉到了可以接受的程度。后面还有一篇阶段性的总结，主要讲一讲遇到的令我比较有印象的错误、呈现的形式和最终解决的方法。稍后更新。

小小的炫耀一下：

![Final Performance](Final Performance.png)

之前提到，vllm的性能是95.0tokens/s左右。已经和vllm完全持平甚至略微超过了！engine的适配和测试也已经整理完成，接下来就是准备重新开始学习、调研多卡的实现方法、设计多卡的扩展实现路线图了。

其他基于Triton Kernel的性能优化

除了之前的博客提到过的几个重要的Kernel或者不用Kernel的矩阵乘法合并优化，在完成 PagedAttention 的核心实现之后，nsys profiling 的数据显示GPU计算时间仍有优化空间。GEMM和GEMV在这个过程中已经占据了大约80%的GPU时间，这部分受限于硬件带宽，难以突破。但剩余的20%里藏着几个可以用Triton消除的低效模式：显存的重复读写和不必要的中间变量分配。通过剩下的零碎Kernel优化，仍然可以实现将近10%的性能提升。

1. Fused RoPE & Paged KV Cache Store (显存读写合并)

这是收益最大的一次融合，解决的是一个隐蔽的"显存回旋镖"问题。

在之前直接随手写的naive实现里，apply_rope会产生一系列临时tensor：先slice切出前后两半，再neg取负，最后cat拼回去。每一个操作都是一次独立的GPU kernel，意味着中间结果要反复写入显存再读出来。更糟糕的是，RoPE算完的K会先落入连续显存，然后store_kvcache再把它读一遍，按slot_mapping的地址写入分散的物理Cache块。同一份K的数据在没有然后额外处理需求的情况下被搬运了两次，造成冗余。

融合后的kernel把这个流程压缩成一次：从全局显存读入Q、K、V、cos、sin，在寄存器里完成RoPE旋转计算，全程不进行临时变量分配。计算完成后，Q写回连续显存供后续Attention使用；而寄存器里的K_rot和V则直接通过slot_mapping提供的物理地址，一步写入非连续的Paged Cache块，不再经过任何中间缓冲。

GQA的处理是这个kernel需要特别注意的一个设计细节。以框架适配的模型Qwen2.5-1.5B为例，Q有12个头，K和V只有2个头，Grid设计时需要让每个program明确知道自己处理的是Q头还是KV头，以及对应的KV头索引是哪个。通过kv_head_idx = head_idx // (N_HEAD // N_KV_HEAD)的映射，kernel内部自然地处理了这个不对齐问题，调用方不需要任何额外操作。

这个收益大约有5Tokens/s，也就是说足足5%。

2. Fused SwiGLU (MLP 层激活函数与乘法融合)

SwiGLU 是相对好发现和动机更直接的一个Kernel fusion操作，不过相对来说收益没有前面那个那么大。原始的F.silu(gate) * up`需要进行是三次显存IO：读gate、算SiLU、写临时结果；读临时结果和 up、做逐元素乘法、写输出。实际上这三步在计算上完全可以流水完成，没有任何数据依赖要求它们分开。

Triton kernel把gate和up同时加载到寄存器，在寄存器内连续执行Sigmoid、gate与Sigmoid结果的乘法（SiLU），以及与up的逐元素乘法，最后一次性写回全局显存。三次读写变成一次，访存带宽需求降低到原来的三分之一。

实现上把3D tensor（batch×seq×hidden）展平成1D视角来划分Block，这样kernel对任意的batch size和seq_len都能够天然兼容，不需要针对不同输入形状做特殊处理。

这个收益大约有2Tokens/s。

3. Triton RMSNorm (pow mean add sqrt和mul的算子融合)

这个算子fusion的目的是，把RMSNorm自身的多个零碎的PyTorch算子合并成一个Triton kernel。

原始的PyTorch实现 x.pow(2).mean(-1).add(eps).rsqrt().mul(x).mul(weight) 实际上在GPU端会拆成5~6个独立的elementwise kernel，pow一次、mean（reduction）一次、add一次、rsqrt一次，最后还有两次mul。每个kernel都有独立的 launch 开销，中间结果要落入全局显存再被下一个 kernel 读出来。对于28层×2个Norm=56次/步来说，这些launch虽然单次很小，累积起来也有可观的数量。

Triton kernel的做法很直接：每行（即每个token的hidden_size维度）分配一个program，一次性把整行1536个元素加载到寄存器，在寄存器内完成平方求和、除以N、加eps、rsqrt、乘权重的全部计算，最后一次写回。6个kernel变成1个，中间没有任何全局显存的临时读写。

这里要特别注意一个踩了的坑：必须用BLOCK_SIZE取hidden_size的下一个2的幂次（1536→2048），用mask屏蔽超出部分。这个实现实际上不是可选的二世必须的，来源于Triton对tl.arange的要求。具体来说，Triton要求tl.arange产生的长度必须是编译期常量且为2的幂。如果不是这样，会直接出错，在这里记录一下以提醒后来人。

此外，计算全程在float32下进行以避免bfloat16的精度溢出，写回时再转换为输入的原始dtype。

这个收益大约有2~3Tokens/s。

4. Fused Add Norm 失败

Fused Add Norm的融合是计划中的第4个优化，目标是在前面第三个的基础上更进一步，把残差加法和归一化也合并成一个kernel，从两次显存读写变成一次。实现上用Triton在Block内部完成方差计算和归一化，逻辑并不复杂。

动机很直接：在TransformerBlock里，残差连接和RMSNorm是紧挨着的两步操作。原始流程是先做x = x + attn_out（一次全局显存读写），再把结果送入RMSNorm（又一次读写：pow → mean → rsqrt → mul）。同一份数据在没有任何计算上的阻隔的情况下被搬运了两次。Triton kernel 的设计思路是把这两步压缩进一个 kernel：在寄存器里完成x + residual之后，就地计算方差、rsqrt 和缩放，x被in-place更新为加完残差的值，返回的是归一化后的结果，全程只读一次、写两次（一次更新 x，一次写 normed_out）。

但实际上测试下来不仅没有增加性能还减少了。是负收益，降低了大约2~3Tokens/s。尝试问问AI，给出的其中一种解释如下：

PyTorch 原生的 elementwise kernel 对这种小规模操作已经做了充分优化，走的是高度特化的 launch 路径；而 Triton kernel 有自己的一套 launch 机制，对于这种计算量极小、一个 warp 就能处理完的操作，Triton launch 本身的固定开销反而比省下来的显存带宽还大。更具体地说，在 Qwen2.5-1.5B 的 decode 阶段，hidden_size 只有 1536，batch_size 为 1，整个 RMSNorm 的输入只有 1536 个 float16 元素（3KB），连 L2 cache 都装得下，根本不存在"显存回旋镖"的问题——数据在 cache 里就地复用了，多读一次的代价几乎为零。

不过，我个人认为这应该就是launch overhead的问题。就是不知道pytorch的原生实现到底为什么能够更快，后台又对不同size做了哪些优化？因为按理来说，elementwise的操作应该完全没有什么特别的优化空间才对。别说Triton了，可能用CUDA几小时写一个naive版本也能达到接近可达上限的性能。那到底是为什么？这个问题可能需要后面有时间再对pytorch做进一步的了解，现在先记录下来🤔。

这个收益大约是-2~-3Tokens/s，确认为负优化后已回退。保留了Triton版本和相关注释没有删除，作为参考。

把所有这些全部加上，性能基本上就能够和vllm一样了。当然，这是单Batch的情况，多Batch由于Scheduler的调度策略非常不同，测试的意义也不是很大。

Pico-vLLM 开发日志 #7 CUDA Graph

Mon, 30 Mar 2026 23:50:26 +0800

在完成Paged Attention的完整开发之后，下一步就是考虑提升性能了。

先来看benchmark。vllm在5070 Laptop上可以达到单Batch情况下，95 Tokens/s的Decoding速度。考虑到Decode阶段是绝对的Memory Bound，而5070 Laptop的理论内存带宽上限是382GB/s，这差不多相当于78%的GPU利用率。与此同时，在我的框架中进行相同的测试在Attention部分用自己写的Triton Kernel替换、进行了QKV Projection Fusion等比较重要的合并之后，性能虽然略微提升，但是在5070 Laptop上也仍然只有大约25~27 Tokens/s。即使是在进行了.compile()的default模式的编译之后，性能提升也最多只能达到大约31 Tokens/s。Profiling得到的性能情况如下图：

中间这4倍（对于eager模式）或者3倍（对于compile模式）的差距去哪里了？答案是，时间都被GPU-CPU的传输延迟、Python端的处理和CUDA API的Launch Kernel的开销吃掉了。这一点在WSL的环境下体现的尤其明显：对于WSL环境，一次CUDA API的调用就需要足足15微秒，比windows的原生环境和linux的原生环境都要长很多。而在一次模型的forward的过程中，至少会有大几百次（对于我在进行Profiling的时候使用的forward计算流，每个完整的Decoding步会产生851次cudaLaunchKernel的调用——仅仅调用本身的开销就消耗了13.5ms，而所有的GPU Kernel的实际开销也才11ms左右。因此，如果想要进行Decoding的性能的进一步优化，优化GPU Kernel的边际收益是很低的——继续减少一个本来已经占比不大的部分所产生的时间开销，其加速比永远不可能超过这个部分本来占据的比例本身。关键在于把更大的这一块优化好。

什么是CUDA Graph

CUDA Graph是一种可以减少CPU Launch Kernel的overhead开销的一种方法。为此，我们需要先了解CPU Launch Kernel的overhead开销是如何产生的。

我们知道，GPU在以前的很多场景下可不被称为GPU——它的另一个名字是加速器，或者说加速卡（accelerator）。这个说法有它自己的道理：GPU不会自己给自己发任务，它总是通过等待调用，被调用，返回结果的模式，帮助CPU进行它难以短时间完成的任务，也就是加速。这就需要CPU主动发放指令，这就是Launch Kernel。每一次kernel launch，CPU 都要经历一套流程：准备参数、调用 CUDA Driver API、把launch命令推入CUDA stream。GPU那边则从stream里取命令执行。关键点是这个过程是串行且逐条下发的：CPU发一条，GPU 收一条。即使 GPU 端的kernel执行极快（比如几微秒），CPU端每次launch也都是不可变的固定开销（WSL下~15μs），这部分是省不掉的。可以立刻得到的一条推论是，Kernel的数量越多，CUDA Graph的收益越大，反之则没有什么用。

CUDA Graph为什么能加速Decoding

大模型的推理Forward过程其实恰好就是许多个零碎的小Kernel组成的——而不是大家想象中的少数几个很大的Kernel。这一方面是因为layer堆叠的很深，另一方面是Attention、SwiGLU、RoPE等等的计算需求各不相同，完全无法进行完整的Fuse。在eager模式下，这样的控制流需要CPU亲自介入反复发放——而这个发放的时间无法完全overlap，很多Kernel在GPU上工作的时间甚至远短于CPU的launch开销，更别提在这个过程中，可能还夹杂了Paged Attention的实现的非最优实现引发的CPU侧反复的Tensor创建、内存分配等等行为。如果不实现CUDA Graph，可能很多人都没有意识到自己在中间有哪些动态拷贝拉低了性能。因此，CUDA Graph的实现过程产生的副产品其实比CUDA Graph本身还要更多一些。

对于大模型本身来说，CUDA Graph造成的直接收益则是如下几条：

第一，CPU端从N次launch变成了1次，直接砍掉了N-1次的driver调用开销；

第二，GPU端由于提前知道了完整的执行计划，调度器可以做更激进的优化（比如提前准备下一个kernel的资源）；

第三，消除了CPU-GPU之间反复同步的 round-trip（不一定是Kernel本身造成的，正如我之前所说，可能有很多不经意之间的非最优实现）。

CUDA Graph的技术原理

一句话总结的话，它的本质其实是录制和重放。这有点类似于宏录制和宏重放的思想：把一套线性的、无条件分支发散的操作流程固定下来，然后在输入数据不同的情况下反复重做。这样，一整个过程本身只需要启动一次，后面GPU自己会知道怎么做。

CUDA Graph做的事情就是：把一系列kernel launch的指令序列，或者说launch的顺序、拓扑依赖关系，全部提前录制（capture）下来，形成一个有向无环图（DAG）。图中的节点是各个kernel（以及memcpy、memset等操作），边表示它们之间的依赖关系。录制完成后，这整张图可以作为一个整体一次性提交给 GPU 执行。CPU 端只需要一次 launch 调用（cudaGraphLaunch），GPU 端就会按照图中定义好的顺序和依赖关系依次执行所有节点。

如何使用CUDA Graph

这东西的使用其实很简单，不过它的限制也很大。它的使用方法真的就是字面意义上的“录制和重放”——先在开启录制的完整的跑一遍，然后调用replay相关的操作。从api角度来解释的话，流程如下：

cudaStreamBeginCapture → 正常调用各种 kernel → cudaStreamEndCapture → 得到一个 cudaGraph_t → cudaGraphInstantiate 生成可执行的 cudaGraphExec_t → 之后每次只需 cudaGraphLaunch即可。

具体而言，在pytorch的实现当中，可以参考这段代码的流程：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65


# ============================================================
# 准备Decode的静态数据
# ============================================================
cache.prepare_decode_step()
static_input_ids[0, 0] = next_token
static_slot_mapping[0] = cache.get_decode_slot()
static_position_ids[0, 0] = cache.seq_len
static_block_table[0, :cache.allocated_cache_block_num].copy_(cache.get_block_table())
static_context_lens[0] = cache.seq_len + 1

# 预热触发 compile 编译
for _ in range(3):
 with torch.no_grad():
 _ = model.forward_decode(
 static_input_ids,
 kv_cache_k=bm.gpu_kv_cache[0],
 kv_cache_v=bm.gpu_kv_cache[1],
 position_ids=static_position_ids,
 slot_mapping=static_slot_mapping,
 block_table=static_block_table,
 context_lens=static_context_lens,
 )
torch.cuda.synchronize()
# ============================================================
# 录制CUDA Graph
# ============================================================
g = torch.cuda.CUDAGraph()
with torch.cuda.graph(g):
 static_output = model.forward_decode(
 static_input_ids,
 kv_cache_k=bm.gpu_kv_cache[0],
 kv_cache_v=bm.gpu_kv_cache[1],
 position_ids=static_position_ids,
 slot_mapping=static_slot_mapping,
 block_table=static_block_table,
 context_lens=static_context_lens,
 )

next_token = static_output[0, -1].argmax()
print(tokenizer.decode([next_token.item()]))

# ============================================================
# Decode Loop
# ============================================================
PROFILING_TOKENS = 100 
torch.cuda.synchronize()
start_time = time.perf_counter()

for step in range(PROFILING_TOKENS):
 cache._seq_len += 1 
 cache.prepare_decode_step()

 #另一种原位写入方法，比前面的按索引写入更好
 static_input_ids.copy_(next_token, non_blocking=True)
 static_slot_mapping.fill_(cache.get_decode_slot())
 static_position_ids.fill_(cache._seq_len)
 static_context_lens.fill_(cache._seq_len + 1)

 bt = cache.get_block_table()
 static_block_table[0, :bt.shape[0]].copy_(bt, non_blocking=True)

 # 回放 Graph
 g.replay()
 # 贪心解码，在GPU侧完成以避免同步
 next_token = static_output[0, -1].argmax()

CUDA Graph的使用限制

所有“录制和重放”应该有的限制它都有。总结一条的话，它必须是线性的、静态的。

第一，所有的指针不能变动。换句话说，你不能使用任何“新内存分配”、“新张量创建”，所有这些操作本质上都是新分配内存空间、返回新指针，而CUDA Graph只接受静态的不变化的指针，那么当然也就只指向同一块空间、同一个Tensor。在不同的调用当中，能够做到的就是向固定的static的input里覆写不同的内容，以实现不同循环的不同输出。这也是为什么block_table在vllm实现里是定长的。

第二，控制流必须固定。 换句话说，录制时走过的分支就是回放时会走的分支，不能有data-dependent的 if/else。这意味着如果你的forward里有根据输入动态选择不同kernel路径的逻辑（比如根据seq_len是否超过某个阈值来决定用哪个attention实现），这些在 graph 内部是做不到的。所有的条件判断必须提到 graph 外面，在 CPU 端决定好之后选择对应的 graph 来 launch。

第三，kernel的launch配置（grid size、block size）和调用序列必须完全一致。 就像前面说的那样，一次录制记录下了851次kernel launch，回放时就真的只会严格按相同的顺序launch这851个kernel，grid/block参数也完全不变（作为静态变量，本质上，和前面的指针是一样的）。这直接意味着tensor shape不能变——因为绝大多数kernel的grid配置是根据tensor shape算出来的。对LLM inference来说，这一点的影响很大：decode阶段每一步kv cache的有效长度都在 +1，如果attention kernel的grid size依赖于seq_len，那严格来说每一步都需要不同的 graph。实际工程中的处理方式通常是用padding或者按seq_len区间预录制多张 graph。

第四，不能包含CPU-GPU同步操作。 比如cudaMemcpy（同步版本）、cudaDeviceSynchronize、或者任何会阻塞CPU等待GPU结果的调用，都不能出现在capture的过程中。因为capture阶段这些操作并不会真正执行，录进去会导致语义错误。同样，任何需要读回 GPU 数据到CPU再做决策的逻辑（比如early stopping，CPU介入检查EOS token）也必须移到graph外部。GPU必须表现的像CPU根本不存在一样来进行重放。此外，每一步检查EOS token这种操作会造成严重的性能影响（因为要进行同步和拷贝），一般是按照预定的长度，生成很多个（比如100个token），然后转交给CPU，检查有没有出现EOS token，是在哪里出现的，然后再进行进一步处理和截断。

下面这条是Claude老师补充的：

第五，不能在graph内部使用CUDA的动态并行（dynamic parallelism）或者跨stream的复杂同步。 Stream capture对stream的使用有严格约束，虽然支持在capture期间fork出子stream来表达并行性，但不能随意使用event做跨graph的同步。这部分我还没有研究过，不是很懂，先记下来。

总结一下的话，CUDA Graph本质上是把一段GPU执行流"快照"下来、机械性的、死板的重放。任何运行时才能确定的东西，包括新的内存地址、动态的控制流、变化的shape、CPU端的同步和决策等等，都和它不兼容，能改变的只有固定地址里的数据内容。

CUDA Graph的加速效果

非常好。直接上图⬇️：

这个直接达到了大约92%的vllm性能。后面通过进一步的Triton-Kernel的接入和fusion，应该完全可以做到93~95token/s，彻底追平vllm。

一点思考

在做这部分工作的时候，也在想Vera CPU这种东西会不会让Graph变得更不那么必要？

“Vera CPU 通过 NVLink-C2C 和 Rubin GPU 直接互联，提供 1.8 TB/s 的 coherent 带宽，构建了 CPU 和 GPU 之间的统一地址空间。”

也许以后。这就得等Rubin真的能摸到再说了。

Pico-vLLM 开发日志 #6 Triton Kernel和代码重构

Sun, 29 Mar 2026 21:43:05 +0800

一天没更新是因为花了点时间调试Triton Kernel。这部分的开发比我预期的要慢，主要还是对于pytorch的Tensor的操作不是非常熟悉，栽了几个比较大的坑，因此反复debug了好久。这里也记录一下，希望能以此提醒自己，也能够帮助后来者避坑。

这部分的思路其实前面已经完全阐述完成，主要是讲讲主要完成的工作和遇到的坑有哪些。

设计模式

这部分主要是为了后面的CUDA Graph的实现做准备。总的来说，指导思想只有一个：静态化所有核心前向的数据指针和形状，同时避免任何条件分支。这是因为CUDA Graph的本质是重放：将既定的指针形状、顺序、流程硬编码在GPU上，反复执行，从而避免CPU下发指令进行Kernel启动的开销。这部分目前还没有用上，但已经提前做了准备。

但众所周知，LLM的推理框架本身就需要两个截然不同的需求：Prefill和Decode。在vLLM等主流实践当中，一般都需要将两者完全分离，前者不使用CUDA Graph而后者使用。这么做的原因主要是Prefill的形状高度不固定：为一个比较短的prompt input去padding一个巨大无比的形状来确保绝对不越界以允许CUDA Graph明显是相当不方便的。这也是chunked Prefill产生的动机之一：分次做的话，很大程度上就形状固定了。如果将两者放在一起，is_prefill的条件判断本身也会让CUDA Graph因为分支发散而完全不成立。

这倒是让我想起来之前看到的一篇论文：《Medusa: Accelerating Serverless LLM Inference with Materialization》。这也是我这学期在上的课的老师陆游游老师团队做的一项工作。它的主要贡献是，直接把CUDA Graph的结构给解析了出来，然后通过物化和指针替换的方式动态的使得物化在主存或者SSD上的CUDA Graph不会因为重新加载之后cublas等闭源库的函数加载位置改变、数据指针的位置改变而失效，从而实现LLM推理服务从非加载状态的完全冷启动，同时还能保持启动的高效率，免去每次重新记录和重放的时间overhead消耗。考虑到CUDA Graph这东西能够在一个自研的框架上跑通本身就要花费巨量的心血，直接将其逆向的工作能够发在 ASPLOS 2025 上当然也是可以想象的。很推荐大家去读！

Triton Kernel的实现

目前显式写出的Triton Kernel有两个：一个是KV Cache 写入算子 (store_kvcache)，一个是分页解码注意力(paged_decode_attention)。它们分别负责：

1、store_kvcache在Prefill阶段和Decode根据block Manager分配的块表正确的写入相应的KV cache，Prefill阶段写入seq_len个而Decode阶段固定写入1个。

2、paged_decode_attention则是Decode阶段专用的Kernel，负责根据页表正确取出相应的KV cache，同时完成Flash Attention模式的不将中间矩阵物化的计算流程。

此外需要特别注意的是，虽然传入的是Tensor，不过实际上Triton内部的操作模式更接近C++：传入的Tensor实际上是隐式转换成带数据类型的指针来处理的。这方面要格外注意的是不同指针的数据类型：如果数据类型不同，那么指针的算术操作造成的物理地址偏移量实际上是完全不同的，而且写入的东西也很可能不匹配。如果不能确定自己到底在写什么的话，就多使用tl.print()然后只针对pid==0的情况进行打印来确认吧。

下面是源代码：

1、store_kvcache

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41


@triton.jit
def store_kvcache_kernel(
 k_ptr, v_ptr, # 源 K/V: (total_tokens, n_kv_heads, HEAD_DIM)
 k_cache_ptr, v_cache_ptr, # 目标 Cache: (num_blocks, n_kv_heads, block_size, HEAD_DIM)
 slot_mapping_ptr, # (total_tokens,) 每个 token 的物理 slot 编号
 stride_k_token, stride_k_head, stride_k_dim,
 stride_v_token, stride_v_head, stride_v_dim,
 N_KV_HEADS: tl.constexpr, # ← constexpr
 BLOCK_SIZE: tl.constexpr, # ← constexpr
 HEAD_DIM: tl.constexpr,
):
 # 映射 2D Grid
 pid_token = tl.program_id(0)
 pid_head = tl.program_id(1)

 # 读取物理 slot
 slot = tl.load(slot_mapping_ptr + pid_token)
 block_id = slot // BLOCK_SIZE
 offset = slot % BLOCK_SIZE

 # 计算源数据 (K/V) 的一维内存偏移
 dim_offsets = tl.arange(0, HEAD_DIM)
 src_offset = pid_token * stride_k_token \
 + pid_head * stride_k_head \
 + dim_offsets * stride_k_dim

 k_vec = tl.load(k_ptr + src_offset)
 v_src = pid_token * stride_v_token \
 + pid_head * stride_v_head \
 + dim_offsets * stride_v_dim
 v_vec = tl.load(v_ptr + v_src)
 
 # 计算目标数据 (KV Cache) 的一维内存偏移
 dst_offset = (block_id * N_KV_HEADS * BLOCK_SIZE * HEAD_DIM) + \
 (pid_head * BLOCK_SIZE * HEAD_DIM) + \
 (offset * HEAD_DIM) + \
 dim_offsets

 # 写入 Cache
 tl.store(k_cache_ptr + dst_offset, k_vec)
 tl.store(v_cache_ptr + dst_offset, v_vec)

这个Kernel看上去非常简单，但是在优化的时候有一个大坑，后面讲。聪明的读者可能看到 stride 的时候就想到了吧。

2、paged_decode_attention

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78


@triton.jit
def Decode_Paged_GQAAttention_Kernel(
 q , # (B, n_heads, 1, head_dim) query，decode 每步只有 1 个 token
 k_cache , # (num_blocks, n_kv_heads, block_size, head_dim) 全局 K cache
 v_cache , # (num_blocks, n_kv_heads, block_size, head_dim) 全局 V cache
 block_table , # (B, MAX_BLOCKS_PER_SEQ) int32，每个请求的物理块 id，-1 表示未分配
 context_lens , # (B,) int32，每个请求当前的有效 token 数
 scale , # 1.0 / sqrt(head_dim)
 out,
 
 # Meta-parameters
 # 元参数
 MAX_BLOCKS_PER_SEQ: tl.constexpr, # 启动时固定，不是运行时变量
 BLOCK_SIZE: tl.constexpr, #
 HEAD_DIM: tl.constexpr, #
 N_KV_HEAD: tl.constexpr,
 N_HEAD: tl.constexpr,
 ): # (B, n_heads, 1, head_dim)
 # grid = (B, n_heads)
 # 每个 program 处理一个 (batch, head) 对
 # program_id[0] = batch_idx
 # program_id[1] = head_idx
 
 pid_batch = tl.program_id(axis=0)
 pid_head = tl.program_id(axis=1)
 kv_head_idx = (pid_head // (N_HEAD // N_KV_HEAD))

 m = float('-inf')
 l = float(0)
 o = tl.zeros((HEAD_DIM, ), dtype=tl.float32)
 
 
 q_ptrs = q + pid_batch * (HEAD_DIM * N_HEAD) + pid_head * (HEAD_DIM) + tl.arange(0, HEAD_DIM)
 q_vec = tl.reshape(tl.load(q_ptrs), (HEAD_DIM, 1))

 context_len = tl.load(context_lens + pid_batch)
 max_block_index = tl.cdiv(context_len, BLOCK_SIZE) # 向上取整
 offs_kv = tl.arange(0, BLOCK_SIZE * HEAD_DIM)
 
 for block_idx in range(0, max_block_index):
 
 physical_idx = tl.load(block_table + pid_batch * MAX_BLOCKS_PER_SEQ + block_idx)
 physical_idx = tl.maximum(physical_idx, 0).to(tl.int64)
 base = (physical_idx * N_KV_HEAD * BLOCK_SIZE * HEAD_DIM + kv_head_idx* BLOCK_SIZE * HEAD_DIM)

 # 加载时 mask 掉超出 context_len 的 token
 token_start = block_idx * BLOCK_SIZE
 valid_in_block = tl.minimum(BLOCK_SIZE, context_len - token_start)
 kv_token_mask = tl.arange(0, BLOCK_SIZE * HEAD_DIM) < valid_in_block * HEAD_DIM

 k_ptrs = k_cache + base + offs_kv
 # k_block: (block_size, head_dim)
 k_block = tl.load(k_ptrs, mask = kv_token_mask, other=0.0)
 k_block = tl.reshape(k_block, (BLOCK_SIZE, HEAD_DIM))
 v_ptrs = v_cache + base + offs_kv
 v_block = tl.load(v_ptrs, mask = kv_token_mask, other=0.0)
 v_block = tl.reshape(v_block, (BLOCK_SIZE, HEAD_DIM))

 # mask 最后一个 block 的无效 token
 valid = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) < context_len
 q_row = tl.reshape(q_vec, (1, HEAD_DIM)) # (1, HEAD_DIM)
 s = tl.sum(k_block * q_row, axis=1) # (BLOCK_SIZE,)
 s = s.to(tl.float32) * scale
 s = tl.where(valid, s, float('-inf'))

 m_new = tl.maximum(m, tl.max(s))
 alpha = tl.exp(m - m_new) # 旧的缩放因子
 p = tl.exp(s - m_new) # 当前 block 的权重

 l = l * alpha + tl.sum(p)
 p_col = tl.reshape(p, (BLOCK_SIZE, 1)) # (BLOCK_SIZE, 1)
 o = o * alpha + tl.sum(p_col * v_block, axis=0) # (HEAD_DIM,)
 m = m_new

 o = o / l # (1, HEAD_DIM)
 o_casted = tl.cast(o, out.dtype.element_ty)
 
 tl.store(out + pid_batch * (HEAD_DIM * N_HEAD) + pid_head * (HEAD_DIM) + tl.arange(0, HEAD_DIM), o_casted)

这部分本身没有什么概念上的困难。需要格外注意的是，调试的时候其报错是静默的：只会告诉你错在这里了，而不是告诉你这个Kernel里到底错在哪一行。因此，要多打印！

另一个非常要注意的关键点是Triton会默认执行死代码消除的优化（DCE）。这一点对于跑高性能是好事，但是对于调试极其不利。你可能会想，我注释掉后面一部分让Kernel报错的，一直二分直到找到问题所在，行不行呢？答案是不行，至少不是简单的行。如果你注释了后面一部分操作，而最后没有把前面操作的那部分store到某个地方的话，Triton编译器会直接把它们全部消除掉。那如果本身错误是发生在这部分，那么看上去你只是注释掉了最后几行就还是跑通了，但实际上并不能因此定位到错误所在。

踩坑点和注意事项

1、显存连续性问题

在第一版实现的时候，store_kvcache_kernel是没有传入stride参数的。这在Q、K、V三个权重矩阵分开计算的时候一点问题也没有。但是，在后来的计算流优化当中，为了减少Kernel的launch开销和反复读显存的开销，采用了一个很简单的优化实现：把$W_Q$、$W_K$、$W_V$拼接起来一起和输入$x$做GEMM。在两者同时进行的时候，出现了极其难以排查的乱码，几乎花费了三到五个小时才发现原因。这到底是为什么？看这段代码：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14


q, k, v = qkv.split([q_size, k_size, k_size], dim=-1)
q = q.view(B, seq_len, self.cfg.num_attention_heads, self.cfg.head_dim)
k = k.view(B, seq_len, self.cfg.num_key_value_heads, self.cfg.head_dim)
v = v.view(B, seq_len, self.cfg.num_key_value_heads, self.cfg.head_dim)

# RoPE
q, k = RoPE.apply_rope(q, k, cos, sin)
# print(v.is_contiguous())
# k/v reshape 成 (total_tokens, n_kv_heads, head_dim) 给 store kernel
k_flat = k.reshape(-1, self.cfg.num_key_value_heads, self.cfg.head_dim)
v_flat = v.reshape(-1, self.cfg.num_key_value_heads, self.cfg.head_dim)

# 写入 KV cache
store_kvcache(k_flat, v_flat, kv_cache_k, kv_cache_v, slot_mapping)# type:ignore

看上去没什么问题？

但如果把中间的print注释掉，你会发现，v.is_contiguous()的结果是false。如果改成k，结果又是true。这到底是为什么？因为：k经过了RoPE.apply_rope，在计算的过程中，pytorch自动帮你完成了向量的深拷贝，从而完成了向量的连续化重整，而v的两个操作，.view()和.reshape()都是不触发深拷贝的。也就是说，在它传入store_kvcache这个warpper的时候，它是仅仅是一个视图，在物理上不连续。

那这为什么会造成问题？因为之前说的：Triton内部的操作模式更接近C++：传入的Tensor实际上是隐式转换成带数据类型的指针来处理的。这意味着算术偏移不会考虑任何视图，操作的直接就是底层的物理地址。这和不连续明显是冲突的，但却又在整个大Tensor的数据范围内。因此就会发生，“数据不对，但也没有报错，看k的数据是完全对得上的，无论如何都查不出毛病”的问题。因此，pytorch的Tensor操作当中一定要小心视图和实际存储布局的区别问题。最后，是通过添加stride，重写底层的指针偏移来解决的。如果从头在C++本身开发，反倒没有这个问题了。

2、context_lens和seq_len的区别

这个我一时间还真不能完全精确的说明区别。为了避免对读者造成误会，直接请Gemini老师帮忙总结一下：

在手搓大模型推理引擎时，这两个变量在命名上极具迷惑性，如果不加区分，很容易在 Prefill 和 Decode 切换的瞬间触发经典的“差一错误”（Off-by-one），导致模型吐出的第一个新词直接变成乱码。简单来说，它们的物理含义完全不同：

seq_len (当前处理长度)：指的是模型当前这一步 Forward 到底吃进去了几个 token。
- 在 Prefill 阶段，一次性处理整个 Prompt，所以 seq_len = N。
- 在 Decode 阶段，因为我们是逐字生成，每次只喂给模型一个最新的词，所以 seq_len 永远等于 1。
context_lens (上下文总长度)：这是传给底层 Paged Attention 算子（如 Triton Kernel）的关键参数。它代表当前的 Query 到底需要和 多少个 token（包含它自己） 去计算注意力得分。

换言之，一个简单的自我思想检验：在Prefill刚刚结束，第一轮Decode开始的时候，传入Decode的参数应该是：seq_len=1，context_lens=N+1，其中N是prompt中的token数量。

3、复制粘贴的问题

孩子们，用ai帮忙写测试脚本的时候一定要看看自己复制了几遍，否则会发现自己的推理框架莫名其妙的总是会在Prefill结束之后多出一个token而找不到任何原因😭

性能Profiling

用nsys进行了GPU侧性能的Profiling，模式是eager，无编译的自动算子融合。结果如下：

Kernel 时间占比 实例数 单步耗时
──────────────────────────────────────────────────────────────
GEMM (gemvx, 主) 63.1% 1700 ~6.8ms
GEMM (gemvx, 次) 22.2% 560 ~2.4ms
cutlass WMMA (gate_up fused) 5.4% 560 ~0.6ms
RMSNorm elementwise 2.0% 3380 ~0.2ms
BinaryFunctor (residual add) 1.0% 1700 ~0.1ms
CUDAFunctor_add 0.9% 2240 ~0.1ms
Decode_Paged_GQAAttention 0.9% 560 ~0.1ms
RMSNorm reduce 0.9% 1140 ~0.1ms
CatArrayBatchedCopy 0.8% 1120 ~0.08ms
store_kvcache_kernel 0.2% 560 ~0.02ms
──────────────────────────────────────────────────────────────
GPU 总时间（20步） ~214ms 单步 ~10.7ms

按3.15GB（3GB权重+0.15GB的KV cache）的Memory Bound情况来计算，单步的理论最优化时间是：

$$ (3.15GB/token)/(382GB/s)=0.00825s=8.25ms $$

性能利用率已经达到77%，与此同时，整个单步Decoding的时延居然达到33ms左右，远远长于GPU侧的耗时。这主要是因为pytorch侧的时间非常长，反复调用产生了巨量的overhead。在我使用的WSL系统中，这个调用还会产生更大的虚拟化开销。因此目前的优化的性能瓶颈已经不是在GPU端了。这也是为什么接下来要做CUDA Graph。

Pico-vLLM 开发日志 #5 Paged Attention

Sat, 28 Mar 2026 01:22:15 +0800

终于到这一步了。刚开始开发Pico-vLLM这个推理框架的时候还觉得这是一件遥不可及的事情，但最终还是完成了。目前调用的forward过程里的Kernel仍然是伪Batching——内部是一个简陋的python for循环，不过整个Prefill、Decode、Engine、Scheduler、Forward的全流程都已经以完整的Paged Attention+Continuous Batching流程跑通了，只差Triton Kernel或者CUDA Kernel接入——我个人计划写Triton Kernel为主：毕竟时间成本有限，把性能提升到可以接受的程度是目前的目标。不论如何，这个里程碑还是挺有意义的。

和之前的博客中所提到的一样，Paged Attention就是专门为了解决内存碎片问题而生的。在naive kv cache的情况下，设计者必然面临两个决策中的一个：要么选择在一开始分配较短的kv cache空间，代价是随着seq_len长度超出原本预计的空间需要不断的进行完整的拷贝和重新分配；要么选择在一开始分配绝对够用的kv cache空间，代价是请求还没有产生几个，显存就迅速爆炸，而绝大部分大概率还没开始用。这一动机任何读者都应该相当了解，因此就不再冗余赘述了。

核心数据结构和设计模式

这次的实现规划的比较完整且一开始就预留了可扩展性，在最初就计划了双pool设计。也就是说，有可以防止硬OOM的GPU + CPU offload的双cache池，加上随之必须产生的逻辑block_id和物理block_id的分离。任何一个请求都会携带一个kv cache表，里面记录的是其持有的逻辑cache块的索引。而Block Manager则同时维护一个映射表：每个逻辑cache块对应的物理cache块的位置，以及物理块到底是哪种类型（CPU的还是GPU的还是未分配、非法的）的cache块。这种双重映射提供的分层抽象在可预见的未来甚至可以扩展到SSD——如果文件系统的性能允许的话。据作者本人所知，这种offload在工业界是大趋势，因此特意将其完整实现了出来。

核心代码块如下。值得注意的是，Block Manager本身不决定换入和换出的策略，仅仅是提供方法：这是思考后的结果，我个人认为这部分操作更适合分离开来，交给某个专门的Scheduler去执行，写在Block Manager内会让模式变得过度耦合且重量化。

 1
 2
 3
 4
 5
 6
 7
 8
 9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126


class pagedblocktype(Enum):
 GPU = "gpu"
 CPU = "cpu"
 NONE = "none" # 表示未分配
class BlockManager:
 def __init__(self, 
 num_gpu_blocks: int, 
 num_cpu_blocks: int, 
 block_size: int, 
 num_layers: int,
 num_kv_heads: int, 
 head_dim: int, 
 dtype: dtype):
 ########## 物理块索引 ##########
 self.num_physical_gpu_blocks = num_gpu_blocks
 self.num_physical_cpu_blocks = num_cpu_blocks
 self.num_total_blocks = num_gpu_blocks + num_cpu_blocks
 self.block_size = block_size
 # 空闲块列表，初始全部空闲
 # 用 deque 比 list 的 pop(0) 快
 # GPU pool
 self.gpu_free_blocks: deque[int] = deque(range(num_gpu_blocks))
 self.gpu_kv_cache = torch.zeros(
 2, num_layers, num_gpu_blocks, block_size, num_kv_heads, head_dim,
 device='cuda', dtype=dtype
 )
 
 # CPU pool（offload 目标）
 self.cpu_free_blocks: deque[int] = deque(range(num_cpu_blocks))
 self.cpu_kv_cache = torch.zeros(
 2, num_layers, num_cpu_blocks, block_size, num_kv_heads, head_dim,
 device='cpu', dtype=dtype,
 pin_memory=True # ← 关键：pin_memory 让 CPU→GPU 传输更快
 )
 
 ########## 逻辑块索引 ##########
 # 逻辑块表，记录每个逻辑块对应的物理块和类型
 # 最多逻辑块不会超过 num_total_blocks，因为每个逻辑块至少占一个物理块
 # 第i个初始逻辑块初始化为[NONE, -1]，表示未分配物理块
 self.block_mapping: List[tuple[pagedblocktype, int]] = [
 (pagedblocktype.NONE, -1) for i in range(self.num_total_blocks)
 ]
 self.logical_free_blocks: deque[int] = deque(range(self.num_total_blocks)) # 逻辑块索引，初始全部空闲
 
 def allocate(self, num_blocks: int = 1) -> List[int]:
 # 分配 num_blocks 个物理块
 block_ids = []
 if num_blocks > self.num_free_blocks:
 raise RuntimeError(f"Not enough free blocks to allocate {num_blocks} blocks")
 for _ in range(num_blocks):
 if self.gpu_free_blocks:
 physical_block_id = self.gpu_free_blocks.popleft()
 logical_block_id = self.logical_free_blocks.popleft()
 self.block_mapping[logical_block_id] = (pagedblocktype.GPU, physical_block_id)
 block_ids.append(logical_block_id)
 elif self.cpu_free_blocks:
 physical_block_id = self.cpu_free_blocks.popleft()
 logical_block_id = self.logical_free_blocks.popleft()
 self.block_mapping[logical_block_id] = (pagedblocktype.CPU, physical_block_id)
 block_ids.append(logical_block_id)
 else:
 raise RuntimeError("No free blocks available")

 return block_ids

 def free(self, block_ids: List[int]) -> None:
 # 回收物理块，加回 free_blocks
 # 回收逻辑块
 for block_id in block_ids:
 block_type, physical_block_id = self.block_mapping[block_id]
 if block_type == pagedblocktype.GPU:
 self.gpu_free_blocks.append(physical_block_id)
 elif block_type == pagedblocktype.CPU:
 self.cpu_free_blocks.append(physical_block_id)
 else:
 raise RuntimeError(f"Block {block_id} is not allocated")
 # 更新 block_mapping
 self.logical_free_blocks.append(block_id)
 self.block_mapping[block_id] = (pagedblocktype.NONE, -1)

 #####################################################
 # 非必须的方法，用于offload到CPU，swap_in/swap_out()
 #####################################################
 def swap_out(self, block_ids: List[int]) -> None:
 """
 把 GPU 上的 block 换出到 CPU
 """
 cpu_block_ids = []
 for block_id in block_ids:
 block_type, physical_block_id = self.block_mapping[block_id]
 if block_type != pagedblocktype.GPU:
 raise RuntimeError(f"Block {block_id} is not on GPU, cannot swap out")
 cpu_block_id = self.cpu_free_blocks.popleft()
 self.cpu_kv_cache[:,:,cpu_block_id].copy_(
 self.gpu_kv_cache[:,:,physical_block_id], non_blocking=True
 )
 self.gpu_free_blocks.append(physical_block_id) # GPU 块释放
 self.block_mapping[block_id] = (pagedblocktype.CPU, cpu_block_id) # 更新映射
 cpu_block_ids.append(cpu_block_id)

 def swap_in(self, block_ids: List[int]) -> None:
 """
 把 CPU 上的 block 换回 GPU
 """
 gpu_block_ids = []
 for block_id in block_ids:
 block_type, physical_block_id = self.block_mapping[block_id]
 if block_type != pagedblocktype.CPU:
 raise RuntimeError(f"Block {block_id} is not on CPU, cannot swap in")
 gpu_block_id = self.gpu_free_blocks.popleft()
 self.gpu_kv_cache[:,:,gpu_block_id].copy_(
 self.cpu_kv_cache[:,:,physical_block_id], non_blocking=True
 )
 self.cpu_free_blocks.append(physical_block_id) # CPU 块释放
 self.block_mapping[block_id] = (pagedblocktype.GPU, gpu_block_id) # 更新映射
 gpu_block_ids.append(gpu_block_id)

 @property
 def num_free_blocks(self) -> int:
 return len(self.gpu_free_blocks) + len(self.cpu_free_blocks)

 def can_allocate_gpu(self, num_blocks: int) -> bool:
 # 查询是否有足够的GPU空闲块
 # scheduler 用这个决定是否接受新请求
 return len(self.gpu_free_blocks) >= num_blocks
 

而从这个Block Manager里取回相应的物理块的index组成的block_table的方法则如下：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10


 def get_block_table(self) -> torch.Tensor:
 """
 读出当前所有已写入的 k/v
 return: 存储的所有对应index的物理index
 """
 physical_ids = []
 for lid in self.cache_block_index[:self.allocated_cache_block_num]:
 _, pid = self.block_manager.block_mapping[lid]
 physical_ids.append(pid)
 return torch.tensor(physical_ids, dtype=torch.int32, device=self.device)

现在它是变长的，但随后的Triton Kernel当中，它可能会变成定长的——取决于到时候的性能调试了。

此外特别需要注意的是自回归过程中，关于真正的“序列长度”的相关处理。虽然seq_len不包括自己正在生成的这一个token——它的Logits还不存在呢，但是在每一层layer当中，其实k和v已经完整存在了，是按照这个token对应的k、v存在去处理的。因此在做相关操作的时候要格外小心，在适当的地方进行+1的修正，并且提前做好kv cache块的预申请和预分配，以免发生问题。

再此外，传参时候的类型提示真的很重要。虽然python不是静态类型的强类型语言，但适当的类型提示对于编程的帮助极其极其大——最大的之一就是你可以随时随地的查看这个变量是什么，以及通过自动补全查看和迅速编写相应的方法，避免无穷无尽的来回翻找。嗯，果然还是Java/C/C++好啊。

吐槽

作者本人用的卡是5070，著名的nano-vllm项目在写的过程中偷奸耍滑，直接用了flash-attn库——作者本人原本计划同样用flash-attn的算子接入跑通全流程之后作为对比的benchmark，测试自己的算子的信念，但尝试在wsl的venv环境中以各种方法安装flash-attn库均失败。通过github和ai查询原因，得到的结论是，5070是Blackwell系列的消费卡，而目前flash-attn目前似乎是只支持hopper以及之前的架构——直接没法用了（至少截止2026年3月28日不行。此结论有时效性，读者请注意）。这进一步倒逼作者去真的写一个Triton Kernel，不过总归是要写的，最后都一样嘛。

大家如果想使用这些算子库，切记查看好仓库发布页的目前支持什么cuda版本、pytorch版本、显卡系列的提示，避免做无用功。

性能数据

令人意外的是实现了如此多冗余拷贝的naive paged Attention之后性能居然不降反升（3%），这似乎是得益于forward过程没了cat()，某种程度上反而同时减少了拷贝。虽然之前说baseline是不适用kv cache的彻底naive实现，但那太离谱了，这次的才算真正的有比较意义的baseline。在接下来一段时间里，就主要尝试对这个指标做提升了，尝试在有限的时间内逐步优化到vLLM的60%水平，然后再迈进到多机阶段。

具体数据直接看图⬇️

Pico-vLLM 开发日志 #4 Continous Batching

Thu, 26 Mar 2026 20:41:19 +0800

今天的主要成果是写了一个pseudo的Continous Batching框架。pseudo的意思是，engine内部的调用其实仍然是按照loop的模式一个个进行，但scheduler的功能本身是完善的，step模式的完整流程也已经写完。这么做主要是因为paged Attention本身作为框架的必须组件还没有实现——在实现paged Attention的过程中，从GQAAttention算子的模式本身到engine到KV cache的cache类和BlockManager全部都得重构或者从头开始写——因此在一个必然要推翻的框架内写一版必然要在接下来几天内全部重构的代码没什么意义，不如接下来直接实现完整的既有paged又有batching的版本。

毕竟我又没有交付压力🤓。

Continous Batching是什么

是一种可以增加推理框架调度效率的方法。在这个技术普及之前，Batching作为提高推理效率的技术已经被广泛应用。因此，要首先理解Batching在推理当中为什么可以提高推理效率。

一些博客或者技术分享可能会说Batching能提高算术强度什么的，但这其实是一种误解，提高算术强度的技术是speculative decoding，而不是Batching技术。这是因为Batching产生的矩阵其实是对角线分块的：互相并不产生计算交叠，反而会产生大量0的空洞，这一点（在naive Kernel里）通过Attention mask遮盖。同时，KV cache本身的访问也是不会减少的，每一个之前算过的依然要完整读一遍。它最大的意义其实在于减少了权重的读取次数——原本计算一个sequence，吐一个token就要从主存（低层次的存储层次）读一遍完整的权重到高层次当中，现在计算$N$个sequence才需要读一遍，中间的过程均驻留在高层次存储里不离开。对于权重很大的大模型来说，这个效率提升其实是非常重要的。从这个角度来说，Continous Batching计的优化思想根基和GEMM矩阵乘法的优化也没有根本性的差别——计算机基本原理，加速大概率事件和有效利用时空局部性的又一个良好体现。

Batching产生了另一个问题：不同的sequence不太可能是同等长度的，而且不太可能同时吐出eos，到达停止状态。将类似长度的sequence打包在一起可以部分地解决第一个问题，但不能解决第二个问题，这个问题同样被后面要做的paged Attention大部分地解决。这就导致一个结果：在一个batch中的一部分sequence已经吐出了eos，变成停止状态了，而另外一些则没有。只要一个sequence还在跑，整个batch是不能停止的。这个时候，batch可能反而变成负收益：许多读写已经是完全空转的了，只是为了陪跑还在自回归过程的sequence，但内存读写消耗反而一点也没有减少。Continous Batching解决的是第二个问题。

所以实际上，Continous Batching和paged Attention本就是深度耦合的两个技术：如果没有paged Attention，Continous Batching每次都要完整的重新分配和整理内存，那么Continous Batching本身就毫无意义了。反过来说，没有Batching和Continous Batching，在单序列请求下，内存碎片这个问题可能还没有那么严重，paged Attention的威力也完全无法有效发挥。试想：在有了paged Attention之后，Continous Batching其实完全不需要进行任何内存操作：只需要把换出的sequence的KV cache页表索引换掉，变成下一个sequence的就可以了，也就是说对sequence本身其实无感知。这也难怪这两件事被同一时间提出，作者在研究的过程中很可能遇到了这个问题然后把它想明白了。

Continous Batching的设计、实现和调度策略

其实这个技术的基本实现也在思路上比较顺畅。它的基本思想是把所有sequence分成waiting、prefilling、decoding和finished四种状态（也有的实现是waiting、running和finished三种，取决于你如何写你的推理engine和kernel，两者是耦合的），以及engine的step模式设计。通过一个scheduler的调度，scheduler管理所有请求、赋予每一个请求一个unique的id编号，并且决定下一个step处理哪些，如何改变序列的状态，并且维持每个step内的负载大致均衡。engine则负责进行prefilling、decoding，并且检查decoding序列中的情况，将已经达到max_len的和已经生成出eos的序列中的sequence移出，进入finished，而finished则负责记录生成的结果、等待外部请求获取，并且提供将finished队列清空的方法，实际上是为了方便外部请求的设计。

在这个过程中有意思的反而是scheduler的调度策略设计：由于prefilling和decoding的计算资源需求并不相同（其实这里已经初步开始产生PD分离的思想了：异质化的负载调度总是比较难处理的），scheduler如何同时确保效率和公平性就是一个问题。prefilling一次应该多少？decoding一次应该多少？这个问题相当open，理论上来说应该是相当值得研究的（尤其是对于生产环境、ToC环境来说，因为它会极大的影响TTFT指标）。在目前的版本中，为了简单起见，实现的是先到先服务（FCFS）和Prefilling+Decoding中的sequence数量总和固定。这是一个粗糙的调度，不过对于框架的可运行性没有影响。也许后面可以考虑把策略拿出来单独设计。

Continous Batching的相关代码

主要的类有两个，分别是Request和Scheduler。Request记录一个请求的基本信息，Scheduler基于前者的基本信息进行调度。需要注意的是，在我的框架里我本人选择了把KVcache和Request放在一起管理。其他框架似乎不是这样，但按这个模式似乎也没什么问题。

Request类：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47


class RequestStatus(Enum):
 WAITING = "waiting" # 在 waiting 队列，还没做 prefill
 PREFILL = "prefill" # 本步正在做 prefill
 DECODING = "decoding" # 已经 prefill 完，正在 decode
 FINISHED = "finished" # 已完成

''' 请求对象，包含请求的所有信息和状态
- request_id: 请求 ID，唯一标识一个请求
- input_ids: 输入的 token ids，shape (1, init_seq_len)，包含整个 prompt
- generated_ids: 已经生成的 token ids，shape (1, )，初始为 空
- max_new_tokens: 最多生成多少个 token
- temperature: 采样温度，传递给 sampler
- top_p: top-p 截断，传递给 sampler
- kv_cache: 每个请求独享一个 KV cache 实例，存储生成过程中的 KV 状态
'''
class Request:
 request_id: int
 input_ids: List[int] # (1, init_seq_len)，包含整个 prompt
 generated_ids: List[int] # (1, )，包含已经生成的 token ids，初始为 空
 max_new_tokens: int
 temperature: float
 top_p: float
 kv_cache: Optional[KVCache] # 每个请求独享一个 KV cache 实例
 request_status: RequestStatus
 has_eos_token: bool # 是否已经生成 eos_token，scheduler 不直接接触 tokenizer 和 eos_token_id，这个由 engine 在 decode_step 后更新

 def __init__(self, request_id: int, input_ids: List[int], max_new_tokens: int, temperature: float, top_p: float):
 self.request_id = request_id
 self.input_ids = input_ids
 self.generated_ids = [] # 初始为 空
 self.request_status = RequestStatus.WAITING
 self.max_new_tokens = max_new_tokens
 self.temperature = temperature
 self.top_p = top_p
 self.kv_cache = None
 self.has_finished_notification = False # engine改变这个状态，scheduler根据这个状态改变 request_status和移出队列

 def is_max_len_finished(self) -> bool:
 return len(self.generated_ids) >= self.max_new_tokens

 @property
 def prompt_len(self) -> int:
 return len(self.input_ids)
 
 @property 
 def total_len(self) -> int:
 return len(self.input_ids) + len(self.generated_ids)

Scheduler类：

 1
 2
 3
 4
 5
 6
 7
 8
 9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119


class Scheduler:
 _next_request_id: int # 用于生成唯一的 request_id
 max_num_seqs: int # 同时处理的最大请求数，超过这个数的新请求会排队等待
 kv_cache_cls : type[KVCache] # KVCache 的类，用于创建请求的 KV cache 实例
 kv_cache_kwargs = {} # 创建 KV cache 实例的参数字典

 waiting: List[Request] # 还没开始处理的请求
 prefilling: List[Request] # 正在 prefill 的请求
 decoding: List[Request] # 正在 decode 的请求
 finished: List[Request] # 已完成的请求
 def __init__(self, kv_cache_cls: type[KVCache], kv_cache_kwargs={}, _next_request_id: int = 0, max_num_seqs: int = 4):
 self._next_request_id = _next_request_id
 self.max_num_seqs = max_num_seqs
 self.kv_cache_cls = kv_cache_cls
 self.kv_cache_kwargs = kv_cache_kwargs
 self.waiting = []
 self.prefilling = []
 self.decoding = []
 self.finished = []

 ''' 插入新请求
 - input_ids: 输入的 token ids，shape (1, seq_len)，包含整个 prompt
 - max_new_tokens: 最多生成多少个 token
 - temperature: 采样温度，传递给 sampler
 - top_p: top-p 截断，传递给 sampler
 - cache_cls: KVCache 的类，用于创建请求的 KV cache 实例
 - cache_kwargs: 创建 KV cache 实例的参数字典
 return: request_id，唯一标识一个请求'''
 def insert_request(self, 
 input_ids: List[int], 
 max_new_tokens: int, 
 temperature: float, 
 top_p: float, ) -> int:
 request = self.create_request(input_ids, max_new_tokens, temperature, top_p, self.kv_cache_cls, self.kv_cache_kwargs)
 self.add_request(request)
 return request.request_id

 '''' 创建新的请求对象
 - input_ids: 输入的 token ids，shape (1, seq_len)，包含整个 prompt
 - max_new_tokens: 最多生成多少个 token
 - temperature: 采样温度，传递给 sampler
 - top_p: top-p 截断，传递给 sampler
 - cache_cls: KVCache 的类，用于创建请求的 KV cache 实例
 - cache_kwargs: 创建 KV cache 实例的参数字典
 return: request_id，唯一标识一个请求'''
 def create_request(self, 
 input_ids: List[int], 
 max_new_tokens: int, 
 temperature: float, 
 top_p: float, 
 cache_cls, 
 cache_kwargs) -> Request:
 request_id = self._next_request_id
 self._next_request_id += 1
 return Request(request_id, input_ids, max_new_tokens, temperature, top_p)

 '''' 添加新请求到等待队列
 - request: 新的请求对象
 '''
 def add_request(self, request: Request):
 self.waiting.append(request)

 ''' 调度器主循环
 return: 2个列表，分别是当前处于 prefilling、decoding 状态的请求列表
 - prefilling: 正在 prefill 的请求
 - decoding: 正在 decode 的请求
 '''
 def schedule(self) -> tuple[List[Request], List[Request]]:
 
 # 1. 对 decoding 队列中的请求进行 decode，完成后移动到 finished 队列
 for request in self.decoding[:]:
 # is_max_len_finished 由 scheduler判断
 # has_finished_notification 由 engine 在 decode_step 后更新，解耦两者的逻辑，scheduler 不直接接触 tokenizer 和 eos_token_id
 if request.has_finished_notification == True:
 request.request_status = RequestStatus.FINISHED
 self.decoding.remove(request)
 self.finished.append(request)

 # 2. 对 prefilling 队列中的请求进行 prefill，完成后移动到 decoding 队列
 for request in self.prefilling[:]:
 # prefill 完成后：
 request.request_status = RequestStatus.DECODING
 self.prefilling.remove(request)
 self.decoding.append(request)
 
 # 3. 从 waiting 队列中取出请求，放入 prefilling 队列，直到达到 max_batch_size
 # 不进行kv_cache的创建和初始化
 while self.waiting and self.num_in_progress < self.max_num_seqs:
 # 先到先服务策略
 request = self.waiting.pop(0)
 request.request_status = RequestStatus.PREFILL
 self.prefilling.append(request)
 # kv cache必须定长，否则报错
 # request.kv_cache = self.kv_cache_cls(**self.kv_cache_kwargs)

 return self.prefilling, self.decoding

 '''' 清空 finished 队列，释放资源'''
 def clear_finished(self):
 self.finished.clear()

 def get_running_requests(self) -> List[Request]:
 return self.prefilling + self.decoding
 
 @property
 def num_waiting(self) -> int:
 return len(self.waiting)
 @property
 def num_prefilling(self) -> int:
 return len(self.prefilling)
 @property
 def num_decoding(self) -> int:
 return len(self.decoding)
 @property
 def num_in_progress(self) -> int:
 return len(self.prefilling) + len(self.decoding)
 @property
 def num_finished(self) -> int:
 return len(self.finished)

值得指出的是id用int是有明确缺陷的：str和uuid作为id很可能是更好的选项。在这里用int只是因为它简单，而且我的卡大概也不会有超过2147483647个请求。

测试一遍通过，这里就不贴了。接下来就是实现真正的paged Attention了，个人认为它的意义可能不太亚于KV cache本身，直接开启了一小个时代级别的想法。

Pico-vLLM 开发日志 #3 KV cache

Wed, 25 Mar 2026 01:15:40 +0800

今天实现了naive的KV cache，其效果之好确实令人震惊。vLLM根本没有提供不使用KV cache的推理方案，而且会按比例占据所有的显存，足可见KV cache在推理框架中的核心地位。先前的实验也证明了，如果不使用KV cache，推理速度很快就会下降到一个令人难以忍受的程度，从而不允许进行任何有实际意义的模型推理。

KV cache的思想

KV cahce的思想很简单，就是空间换时间。通过保存每一个先前token在经过各层（以qwen2.5-1.5B为例，就是28层）的attention层的时候，把它们的副产物：K和V保存下来。所以实际上，KVcache的shape大约是这样的（对一个请求来说）：

$$ (Layer\_Index, K(0)/V(1), Seq\_Len, Num\_KV\_Heads, Head\_dim) $$

在自回归过程中，一次前向传播，每经过一个TransformerBlock，前向传播的过程就把计算出的结果保存到一个大KV cache的相应位置上。Prefill的过程一次保存等同于input_prompt的token数量的KVcache，decoding过程则一次完整的前向传播增加一个。

虽然思想很简单，但工程实现相对来说没有那么简单。因为要对整个engine进行重构：原本不需要区分Prefill和Decoding，但现在需要了。进一步的说，这两者的需求根本不同：对于Prefill，由于方阵形状规整，算术强度是比较高的，因此Roofline模型更靠近compute bound那一端；而对于decoding来说几乎所有的时间都用来了进行模型权重的读取和KV cache的读取（虽然在我们这个例子当中效果并不显著），因此完全是memory bound的。这样，确实是从实践的角度感受了一下PD分离为什么最终会发生，这两者的计算模式是非常不一样的，尤其是当初始input的prompt很长的时候。

除此之外，一个非常容易错的点是：第一次写的时候很容易在每一层TransformerBlock里都把KVcache的Seq_len加一，但这样会造成同一个token的不同层的KVcache被存在了完全错位的地方。实际上，在layer_idx = max_layer - 1的时候把Seq_len加一就可以了。这是一个非常容易实现不对的地方，当时debug了半天。最后的代码长这样：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37


# 朴素实现
class NaiveKVCache(KVCache):
 # 连续 tensor
 def __init__(self, num_layers, max_seq_len, num_kv_heads, head_dim, device, dtype):
 self.num_layers = num_layers
 self.max_seq_len = max_seq_len
 self.num_heads = num_kv_heads
 self.head_dim = head_dim
 self.device = device
 self.cache = torch.zeros(
 num_layers, 2, max_seq_len, num_kv_heads, head_dim,
 device=device, dtype=dtype
 )
 self._seq_len = 0 # 当前已填入的长度

 def update(self, layer_idx: int, k: Tensor, v: Tensor) -> None:
 # k, v: (1, num_kv_heads, head_dim) 或 (seq_len, num_kv_heads, head_dim)
 # 通过读取kv的shape自动适配 prefill 和 decode 场景
 # 以后用 cuda kernel 的时候这个要怎么处理？可能需要分开 prefill 和 decode 的接口，或者额外传参？
 num_new_tokens = k.shape[0]
 start = self._seq_len
 end = self._seq_len + num_new_tokens

 self.cache[layer_idx, 0, start:end, :, :] = k
 self.cache[layer_idx, 1, start:end, :, :] = v
 if (layer_idx == self.num_layers - 1): # 只在最后一层更新 seq_len，保证所有层的 seq_len 一致
 self._seq_len += num_new_tokens

 def get(self, layer_idx: int) -> tuple[Tensor, Tensor]:
 return self.cache[layer_idx, 0, :self._seq_len, :, :], self.cache[layer_idx, 1, :self._seq_len, :, :]

 def reset(self) -> None:
 self._seq_len = 0

 @property
 def seq_len(self) -> int:
 return self._seq_len

希望看到的人能够绕过这个坑。

下面是使用了naive实现的KV cache之后，模型的decoding性能随着seq_len的变化的profiling图。

可以看到，对于1.5B这样的小模型（权重3GB）和<4096的seq_length长度来说，对于不同的sequence长度，单个token的推理时间几乎是一个常数，不会发生明显下降，和之前的naive实现的$O(N^2)$的曲线特征形成了非常明显的对比。这主要是因为权重读取相关的操作占据了绝大部分：这部分操作和seq_len是完全无关的，是$O(1)$的常数项。如果把length拉长，应该会看到后续的$O(N)$阶段，主要由KV cache本身的读取导致的带宽需求决定的Latency。不过由于我的卡（5070）的显存很小（8GB），在整个显存能够不OOM的范围内，几乎都不会发生这种情况就是了。

理论上来说，5070能够达到的推理速度大约是96tokens/s，目前的利用率仅仅只有22%左右。可以发现能够优化的空间还是相当大的，Kernel Fusion能够有相当的发挥空间。

一些思考

另外，今天和同学交流了一下未来的专精方向。算子开发本身似乎已经比较成熟了，通信库、高性能网络更值得做一些。确实如此，triton等等已经大大降低了开发难度，而且flashattn等各种算子融合和实现的库在pytorch里的生态也非常完善了。然而我的手上没有相关资源，只有单卡。但不论如何，也许应该考虑调整框架的开发方向？

目前的规划依然是要实现完整的KV cache。然而，prefix cache在目前的条件下，似乎没有个人开发者环境下有效评测的方法，这部分应该没有特别的高优先级的需求——可以往后推一推。

这样的话，目前的规划就是这样的：1、实现paged attention，2、进行Kernel Fusion的尝试和优化，3、实现单节点多卡推理，4、实现多节点多卡推理。

Pico-vLLM 开发日志 #2 自回归和Profiling

Tue, 24 Mar 2026 00:12:41 +0800

今天凌晨完成了自回归框架的初步搭建，到目前为止可以完整的进行模型的自回归生成了！如下图：

值得纪念的时刻。需要注意到的是：Greedy的sample策略会很容易引发一个大家现在已经不太关注的问题：注意力塌陷，即模型开始反复输出同一段话，无法进行任何有效的长文本输出。如果使用带有随机性的采样，这种情况就基本上不会发生了。这应该也是为什么Greedy在现在的模型采样方法里完全不采用。

而且可以大概看出，在短Prompt下，一个token生成的速度大约是~20tokens/sec。可以计算一下理论上其应该有的计算峰值：

$$ 理论最大速度 = 带宽 / 权重大小 = 288 GB/s / 3 GB ≈ 96 tokens/sec $$

可以看出，naive实现相对于理想情况来说有效利用的比例并不高，只达到了大约21%的理论峰值。接下来我们的目标就是提升这个数字，并且提升OOM的阈限（利用主存甚至SSD）——先画个饼，提升到60%以上？

下面进行profiling，这在后面的开发当中将会作为baseline反复用到，用于和后面每一次优化进行比较。从16个token开始，每次增加16个token长度的seq_len，一直增加到4096。profiling的结果如下所示：

模式：Naive（无 KV Cache）硬件：RTX 5070 Laptop 8GB

Seq_len	ms/step	token/s
16	47ms	21.0
512	55ms	18.2
1024	115ms	8.7
2048	272ms	3.7
4096	764ms	1.3

可以明显的看到，在纯naive实现下，其主要分两段：常数段（生成速度几乎不变）和二次函数段（生成速度减慢的模式完美符合二次函数的增长）。常数段我认为应该主要是因为Latency以模型权重的反复加载为主，compute的时间几乎可以忽略不计，因此前向传播的Latency是固定的。而二次函数段就是经典的$O(N^2)$的注意力计算复杂度compute-bound导致的$O(N^2)$时间增长。此外特别值得注意的是，在大约seq_len=512~528的地方，出现了非常明显的性能阶跃式的下降。目前考虑的有pytorch的kernel调度的不同导致的，也有可能是L2cache溢出导致的换页导致的。究竟是什么导致的，后续也需要仔细研究。

总的来说，如果不使用KV cache的话，内存其实远远达不到OOM的程度（8+16GB的显存，仅仅用了4GB），token的生成速度就已经慢到无法接受了，更别提什么CUDA kernel的优化，都是于事无补。

此外，在测试的过程当中，第一次profiling的implementation在理论显存占用仅仅3.3GB不到的时候就OOM了。然而，反复计算理论显存也没找到为什么OOM。查询了之后才知道，原来pytorch不会自动释放显存，会把已经不再用到的、原先用过的显存块作为reserved模式保留在内部，哪怕马上就要OOM了也不会自动释放（确实这么离谱）。而每次因为序列是越来越长的，原先reserved的内存从未用上过，造成碎片单向增加而从不减少。这某种意义上来说和paged attention遇到的和解决的问题非常相似：只不过pytorch这个管理的更加低级。后面将会逐步实现naive的KV cache和Paged attention机制，随后再考虑kernel的优化。

Pico-vLLM 开发日志 #1 Sampler的设计

Sun, 22 Mar 2026 23:37:49 +0800

今天的主要任务是写sampler，分为以下几种：Greedy，Temperature，Top-P。其实理论上来说Top-p和Temperature两个方法是解耦的，实现里的Top-p其实是Top-p+Temperature。不过无论哪种都很好实现，只讲基本思想。

Greedy，顾名思义就是永远只选择概率（Logits）最大的那个token。数学公式如下：

$$ X_{next}=argmax_i(Logits(X_i)) $$

Temperature是针对Greedy的改进，想法也很简单，就是用一个参数决定分布的平均程度（热化程度），然后用这个分布采样。由于实际上是模拟玻尔兹曼分布，和热力学关系很大，索性就叫Temperature了。很明显，T越大采样越随机，=1则和原本的随机采样分布相同。它的数学公式如下：

$$ \tilde{L}_i = \frac{\text{Logits}(X_i)}{T} $$

$$ P(X_i) = \text{Softmax}(\tilde{L}_i) = \frac{e^{\tilde{L}_i}}{\sum_j e^{\tilde{L}_j}} $$$$ X_{next} \sim P(X_i) $$

Top-P可以不包含T，但一般包含T。它是对Top-K（取固定前K个）的改进：取若干个最有可能的，直到覆盖了原本分布的前p部分，随后在这部分里面采样，以规避过于离谱的输出的同时又保证随机性。在包含T的情况下，它的公式如下：

$$ \mathcal{V} = argmin_{S \subseteq \mathcal{V}_{vocab}} |S| \quad \text{s.t.} \sum_{X_i \in S} P(X_i) \geq p $$

在候选集内重新归一化后采样：

$$ P'(X_i) = \begin{cases} \dfrac{P(X_i)}{\sum_{X_j \in \mathcal{V}} P(X_j)} & X_i \in \mathcal{V} \\ 0 & \text{otherwise} \end{cases} $$

$$ X_{next} \sim P'(X_i) $$

原理应该说很简单，但top-p算法在直觉上似乎有很大的优化空间，$O(NlogN)$的复杂度对于大Vocabulary的模型的推理来说会不会产生比较明显的性能影响？这需要全部跑通之后进行profiling，现在先不管了。

Pico-vLLM 开发日志 #0 博客安装

Sat, 21 Mar 2026 00:00:00 +0000

HelloWorld!

今天的任务是先把博客搞好，然后把刚刚跑通能吐一个token的框架py代码推送。

测试一下

测试成功了，后面就在这里更新。目前的计划是先跑通自回归的引擎，然后做最基本的profiling，随后逐步优化。具体的实现可以参考vllm和sglang，另外预计也会有很多写py的过程中想到的可以优化的点，后面都会整理在这里。

Fain的Blog

Femtotron开发日志 #11 监督微调 Supervised Fine-Tuning, SFT

SFT和预训练的区别

数据格式

Loss Mask设置

学习率

Femtotron开发日志 #10 ZeRO-3、SAC和PP：一个简单的显存节省收益数据测试和分析

数据

PP=1、TP=1、DP=8 (固定num_layers=8, 变化seq_len)

PP=2、TP=2、DP=2 (固定seq=32,变化num_layers)

ZeRO-3 vs ZeRO-2 + SAC 对比

分析

Femtotron开发日志 #9 流水线并行 Pipeline Parallelism

流水线并行的原理

流水线并行的核心设计架构和抽象

流水线模型模块

调度模块

值得注意的Bug

Rotary inv_freq的garbage初始化问题

一个不算bug的奇怪问题

ZeRO-3 + PP时会发生的显存占用和计算效率问题

autograd-held view现象

相关问题

通信组（Communication Group）划分导致显存的不降反增

工业界的解决方案

路线一（NVIDIA / Megatron 派系）： 纯粹的 3D 并行 + ZeRO-1/2

路线二（PyTorch 原生 / Hugging Face 派系）： FSDP 代替 ZeRO-3 + 取代 PP

逸闻

总结

Femtotron开发日志 #8 选择性激活检查点 Selective Activation Checkpointing

SAC是什么

SAC的理论最优情形

SAC的主流实践

Pytorch提供的现有机制

值得注意的Bug和问题

SAC的bug

SAC能带来多少收益

如果SAC特别慢，可能发生了什么

Femtotron开发日志 #7 ZeRO-3模式：抽象设计、Bug排查和教训总结

ZeRO-3的分片和TP的区别是什么

通信的对象是问题

ZeRO-3麻烦的地方

ZeRO-3真实的工程实现模式

ZeRO-3单独实现产生的收益

值得注意的Bug，和它们的排查

内存泄露

噪声扰动还是Bug

无法实现的Bit-Exact

Femtotron开发日志 #6 ZeRO-1和ZeRO-2模式：抽象设计、工程实现和占位符

ZeRO-1的细节

即使在Python也要坚持类型安全

ShardingStrategy 作为一等抽象

master 是 1D 分片而不是按 shape 切

ZeRO 启用时 GradSync 退化为 NoOp

ZeRO-2的细节

即使是Python的Protocol也可以承担多重角色，比多重继承更轻

注意事项

测试结果

Femtotron开发日志 #5 数据并行 Data Parallelism

整体架构说明

DistributedDataLoader

PackedDataset + preprocess

DistributedSampler

Collator

GradientSynchronizer

其他的注意事项

Femtotron开发日志 #4 混合精度训练

为什么不混合不行

在 Femtotron 中的实现

混合精度对训练质量的影响有多大？

其他

Femtotron开发日志 #3 逸闻：记一次掉卡故障的排查和分析

Femtotron开发日志 #2 训练框架中的TP并行模式，工厂模式、注册表模式和函数修饰器

Row/Column并行模式

ColumnParallelLinear

RowParallelLinear

工程实现技巧：工厂模式、注册表模式和函数修饰器

示例：关于权重加载的代码

第一层：策略本身（ShardLoader 协议）

第二层：工厂函数（决定用哪个策略）

路线一（NVIDIA / Megatron 派系）：纯粹的 3D 并行 + ZeRO-1/2