预训练 on 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开发日志 #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上，然后将其重构为并行化的形式，最后再真实的实体化和分配显存。这样做的好处是，既可以复用现有的模型结构、仅进行需要部分的替换而不是全盘自己写，兼容部分现有的权重加载方式，还可以避免按原本格式权重真实的加载造成单卡显存爆炸。

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

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。