Packed Sequence 与上下文并行#

1. pretrain 场景下 packed seq 的复杂点 / 坑点#

1.1 pretrain 与 SFT 的"打包时机"完全不同#

上游 Megatron-LM 的 packed seq 是 SFT-only 的，pretrain_gpt.py:119：

is_packed_sequence = get_args().sft  # SFT always uses packed sequence

上游假设 dataset 已经离线打包好，直接读 cu_seqlens / cu_seqlens_padded / max_seqlen 字段（Megatron-LM/pretrain_gpt.py:132-134）。

是 在线打包 给 pretrain 用的（use_packed_sequence 开关），打包逻辑放在 RingAttnContextParallelHelper.pack()（context_parallel_helper.py:637）里，每个 step 用 EOD token 现切。这条路径上游不存在，复杂度都是自己承担的。

1.2 具体坑点（按出问题概率排序）#

1.3 Labels 的"跨样本 roll"逻辑（最容易出 bug）#

context_parallel_helper.py:750-784：packed 之后 labels 不能简单 torch.roll(-1)，因为每个 batch sample 内部的最后一个 token 不应该 roll 到下一个 sample。代码用 last_extra_token + 在每个 sample 边界手工 ±1 调整 labels_seqlens / labels_src_offsets / labels_dst_offsets：

if _offset == sum(valid_token_nums[: _batch_idx + 1]):
    _batch_idx += 1
    labels_seqlens[seq_idx] += 1
    labels_src_offsets[seq_idx + 1 :] += 1
    ...

这段是 O(num_seqs) 的 Python 循环、index 都是裸算的，一旦 valid_token_nums 计算偏 1，或者 sample 边界恰好落在 padding 上，就会静默地把下一个 doc 的第一个 token 当作 label——loss 不会炸但模型学错。这部分上游完全没有对应实现。

1.4 FP8 + CP 的对齐 padding 嵌在每个 sub-seq 末尾#

context_parallel_helper.py:684-721：

• CP>1：对每个 sub-sequence 单独 pad 到 TP*CP*8*2（FP8）或 TP*CP*2
• CP=1：只 pad 整条序列到 TP*8/TP

两条路径行为不同，且 padding 是塞在 每个 doc 末尾，导致 cu_seqlens(真实) 和 cu_seqlens_padded(对齐后) 必须并存。下游所有用到 cu_seqlens 的地方（attention、loss、MTP）都得分清楚用哪一个，错用一次就是错位。

1.5 valid_local_token_mask 是另外一条独立通路#

PackedSeqParamsWithValidTokens（packed_seq_params.py）在标准 PackedSeqParams 之上又挂了一个 valid_local_token_mask，因为 pad 之后 attention 看到的 token 数 ≠ 用来算 loss 的 token 数。

坑在于：

• chunk() classmethod 要手工 chunk 这个 mask（PP/VP 切分时）
• CP slice 之后还要再 view(-1) 回去（context_parallel_helper.py:424-427）
• 上游 PackedSeqParams 没这个字段，所以子类如果哪天被上游的某个新接口当 PackedSeqParams 用，valid mask 就丢了

1.6 MTP rolling 必须感知 doc 边界#

context_parallel_helper.py:99-129（_roll_tensor_for_mtp）：每个 MTP depth i，要在 cu_seqlens_q_padded[1:] - i 这些位置清零，否则下一个 doc 的前 i 个 token 会被当成上一个 doc 的 MTP 目标。

这个逻辑跟 2.1 的 labels 跨样本 roll 是一个家族的问题，但是写在了完全不同的两个函数里——保持一致性靠人。

1.7 CP slice 走 tex.thd_get_partitioned_indices，要求每段 sub-seq 长度被 2*CP 整除#

context_parallel_helper.py:855。FP8 对齐选 align_multiplier = _TE_FP8_CONSTRAINT*2 其实就是顺手满足这个 2*CP 约束。但是：

• 非 FP8 时 align_multiplier = 2，如果用户设置了奇怪的 CP/TP 组合，对齐可能不满足
• batch_processor.py:443-445 有断言 b*s == cu_seqlens_q_padded[-1]，但这个断言只在 fixed-length 场景生效；variable-length 时这条保护没了

1.8 多个 helper 类分支，约束不一致#

pretrain_gpt.py:404-428：

not use_packed_sequence            → ContextParallelHelper（BSHD）
flash/flash2/fused                 → RingAttnContextParallelHelper
magi                               → MagiAttnContextParallelHelper（强制 dataloader 内 CP）
flex                               → UlyssesContextParallelHelper

四套实现接口签名一样但内部行为差异不小（slice_bshd/slice_thd 哪个 raise、是否做对齐 padding、是否支持 MTP），切 backend 时容易踩到"以前能跑的 config 现在挂了"。

1.9 Document-level loss 的 PackedDocLevelParams 是 packed_seq 之外另一层"包装"#

pretrain_gpt.py:765-901。这是自己的扩展，跟 PackedSeqParams 不共享数据结构，靠 register_batch_sliced_hook(update_packed_doc_level_params) 维护一致性。

• one-hot 路径（use_index_add=False）一次性 [N, num_documents]，CP/seq 很长时 OOM 风险
• index_add 路径有 Python for-loop over batch，CP 大时慢

两条路径都依赖 2.1 那个 labels rolling 正确，否则 doc id 对齐就错了。

1.10 和上游 Megatron-LM 对比的缺口#

1.11 重点 review / 加测试的地方#

1. context_parallel_helper.py:750-784 的 labels 跨样本 rolling——给 valid_token_nums 边界、单 doc 占满整 sample、最后一个 doc 不含 EOD 这几种 corner case 加单测。
2. generate_cu_seqlens_and_max_seqlen_batch（mask.py:222）对"全 pad sample"、"pad 在中间被 EOD 切断"、"sample 末尾刚好是 EOD"这几种边界的行为，目前看是用 last_non_pad 算 valid_end，假定 pad 都在末尾——SFT 数据里 pad 不一定在末尾。
3. RingAttn pack 里 FP8 对齐分支 vs 非 FP8 分支语义有微差（CP>1 时按 sub-seq pad，CP=1 时按整 batch pad），切 CP size 时端到端 token 数会变，影响 step-by-step loss 复现。
4. PackedSeqParamsWithValidTokens.chunk 只 chunk valid_local_token_mask，valid_num_tokens_in_batch / local_seqlens_in_batch 是 List 不 chunk，PP 拆分时这俩字段在中间 stage 是"全 batch 视角"，很容易被误读为"本 chunk 视角"。

2. CP 上 Packed 的实现机制#

2.1 总流程#

[B, S] raw batch
  ↓ pack()      生成 cu_seqlens(_padded) + 把数据 flatten 成 [1, T_pad]
[1, T_pad] packed batch  +  PackedSeqParams(thd 格式)
  ↓ slice()     按 CP 策略切片，每张 GPU 拿一段 index
[1, T_pad/CP] local batch  +  PackedSeqParams（cu_seqlens 仍是全长）
  ↓ forward
attention 用 cu_seqlens_q_padded 做 var-len attention（TE 内核自己解释这是 CP 切过的）
  ↓ loss
用 valid_local_token_mask 过滤掉 pad

注意：cu_seqlens_q_padded 全程是"全 batch"视角的，不会随 CP slice 变小——TE 内部知道这是 CP=N 下的 partition。

2.2 第一步：pack（以 Ring 为例）#

RingAttnContextParallelHelper.pack()（context_parallel_helper.py:637）。

输入 tokens: [B, S]，里面有 EOD/PAD。

(1) 算每段 doc 的 cu_seqlens（去 pad）

generate_cu_seqlens_and_max_seqlen_batch（mask.py:222）：

• 用 non_pad = input_ids != pad_token_id 找每行最后一个非 pad 位置
• 用 eos_mask & valid_mask 找 EOD 位置
• 一行 cumsum 得到 segment_id，每个 segment 用 bincount 算长度
• 输出 cu_seqlens = [0, len0, len0+len1, …]，已经全局拼好（跨 batch row）

(2) 算对齐 padding

align_multiplier = _TE_FP8_CONSTRAINT * 2 if fp8 else 2   # FP8: 16, 否则: 2
align_size = TP * CP * align_multiplier
# 对每段 sub-seq 单独 pad
padding_sizes = (-seqlens) % align_size  # 已对齐长度补 0，而不是多补整块
padded_seqlens = seqlens + padding_sizes
cu_seqlens_padded[1:] = cumsum(padded_seqlens)

• 2 是给 ring CP 准备的（要切 2*CP 块），8 是 TE FP8 GEMM 要求。每段 sub-seq 末尾都补 pad，不是只在整个 batch 末尾补——这是 ring CP 的强约束，否则 thd_get_partitioned_indices 切不齐。

(3) flatten & pad 数据

_get_flatten_padded_data（context_parallel_helper.py:537）：

• 输入 [B, S]（去 pad 后的有效 token）
• 输出 [1, cu_seqlens_padded[-1]]，每个 doc 拷到 dst_offsets[i] 开始的位置，剩下空位填 padding_value

labels 特殊处理（context_parallel_helper.py:750-784）：要做 roll(-1) 但不能跨 sample 边界 roll，靠 last_extra_token 保存原 batch 每行最后一个 token，然后手工调每段的 seqlens / src_offsets / dst_offsets。

(4) 返回 PackedSeqParams

PackedSeqParamsWithValidTokens(
    qkv_format="thd",
    cu_seqlens_q=cu_seqlens_padded,
    cu_seqlens_kv=cu_seqlens_padded,
    cu_seqlens_q_padded=cu_seqlens_padded,
    cu_seqlens_kv_padded=cu_seqlens_padded,
    max_seqlen_q=padded_seqlens.max(),
    valid_local_token_mask=...,               # [T_pad] bool，pad 位置是 0
    valid_total_num_tokens=sum(valid_token_nums),
    local_seqlens_in_batch=padded_token_nums,
)

2.3 第二步：slice — 三种 backend 差异最大的地方#

2.4 Ring：每段 doc 切成 2*CP 块对称分配#

RingAttnContextParallelHelper.get_slice_index（context_parallel_helper.py:820）。

GPU 上走 TE 的内核：

index = tex.thd_get_partitioned_indices(
    cu_seqlens_q_padded,
    total_token_num,
    cp_size,
    cp_rank,
)

CPU fallback 直接写出了算法（context_parallel_helper.py:861-887）：

seqlens_chunk = (offsets[1:] - offsets[:-1]) // (2 * cp_size)
for s in range(num_seqs):
    # rank 拿第 rank 块（"前半"）和第 (2*CP-1-rank) 块（"后半"）
    front = offsets[s] + rank * seqlens_chunk[s]
    back  = offsets[s] + (2*cp_size - 1 - rank) * seqlens_chunk[s]
    indices_list += [arange(front, front+L), arange(back, back+L)]

举个 CP=2 的例子，一段 doc 长 8 token：

位置:   0 1 | 2 3 | 4 5 | 6 7
块号:    0     1     2     3
rank0 拿 {0,3} → 物理 token [0,1,6,7]
rank1 拿 {1,2} → 物理 token [2,3,4,5]

为什么这么切？causal mask 下越靠后的 token 计算量越大（要 attend 越多前文），简单顺序切会让最后一个 GPU 累死；对称切（前 + 后）每 GPU 计算量基本一致。这是对称／zigzag 分块；不要与逐 token 交错的 striped 布局混为同一个索引约定。

slice 应用（slice_thd，context_parallel_helper.py:407-422）：

for key in include_partition_keys:
    batch[key] = batch[key].index_select(seq_dim, slice_index)

tokens / labels / loss_mask / position_ids / valid_token_mask 一起按同样的 index 切。切完后，本 rank 的 tokens 形状是 [1, T_pad/cp_size]，但 cu_seqlens_q_padded 还是 全长 的（这是 TE thd 内核的约定）。

2.5 Ulysses：整段平均切 CP 份#

UlyssesContextParallelHelper.get_slice_index（context_parallel_helper.py:1040）：

cp_seqlen = total_token_num // cp_size
slice_index = arange(rank * cp_seqlen, (rank+1) * cp_seqlen)

就是把整条 [1, T_pad] 直接平均切成 CP 段，rank r 拿连续的第 r 段。不分 sub-seq 内部——因为 Ulysses 是 all-to-all 沿 head 维分发，序列维只要每 rank 拿连续段就行。

对应的 pack 也简单些（context_parallel_helper.py:969-1035）：只在 整段末尾 补 pad 到 TP*CP*8 对齐，不对每段 sub-seq 单独 pad。

2.6 Magi：让 Magi 库自己 dispatch#

MagiAttnContextParallelHelper.slice（context_parallel_helper.py:1109）：

dist_attn_runtime_key = magi_attn_varlen_key(
    cu_seqlens_q, cu_seqlens_k, pad_size, chunk_size,
    cp_group_or_mesh=..., causal=True,
)
slice_index = get_position_ids(dist_attn_runtime_key)
batch[key] = dispatch(batch[key].squeeze(0), dist_attn_runtime_key).unsqueeze_(0)

Magi 把切分策略封装在 runtime_key 里，pack 沿用 Ulysses 的（继承），但 chunk_size 要 prev_power_of_two(valid_size // cp_size) 并和 magi_attn_chunk_size 取 min。runtime key 还会塞回 packed_seq_params.magi_runtime_key 让 attention 层用。

2.7 第三步：forward 时 attention 怎么用#

到 attention 层时，hidden_states 形状是 [1, T_pad/cp_size, hidden]，packed_seq_params.cu_seqlens_q_padded 是 [0, L0_pad, L0_pad+L1_pad, ...] 全长视角。

• Ring 路径走 TE 的 flash_attn_varlen + ring 通信，TE 内部根据 cu_seqlens_q_padded + cp_size 自己算每个 chunk 应该 attend 哪些远端 KV
• Ulysses 路径先 all-to-all 把序列维换成 head 维（每 rank 拥有完整序列但只算部分 head），然后跑普通 varlen attention
• Magi 用自己的 dist attention API，runtime_key 携带所有 dispatch 信息

2.8 第四步：loss 怎么对齐#

CP slice 后每 rank 看到的是 T_pad/cp_size 个 token，其中一部分还是 padding。valid_local_token_mask 就是用来过滤这部分：

pretrain_gpt.py:944-949 里 LossCalculationHelper.compute_loss 拿 packed_seq_params 进去，用 valid_local_token_mask 给 loss_mask 再做一层 AND，最后 CP 维 all-reduce loss 和 token 数。

PackedSeqParamsWithValidTokens.chunk()（packed_seq_params.py）：PP/VP 切 microbatch 时，valid_local_token_mask 要 chunk，但 valid_num_tokens_in_batch / local_seqlens_in_batch 不 chunk（仍是全 batch 视角）——这是上面提到的"中间 stage 容易误读"的点。

2.9 关键约束总结#

3. 手算示例：B=2, S=8, CP=2, TP=1, 无 FP8#

3.1 输入#

两个 sample，PAD 用 _ 表示，EOD 用 E 表示：

sample 0: [a, b, c, E, d, e, E, _]      # 两个 doc: len=4, len=3
sample 1: [f, g, h, i, j, E, _, _]      # 一个 doc: len=6

3.2 Step 1: 算 cu_seqlens#

generate_cu_seqlens_and_max_seqlen_batch 扫 EOD + 非 PAD：

sample 0 valid_end = 7,  EODs at pos {3, 6}  →  segs [4, 3]
sample 1 valid_end = 6,  EOD  at pos {5}     →  segs [6]
全局拼接:   seqlens          = [4, 3, 6]
            cu_seqlens       = [0, 4, 7, 13]
            valid_token_nums = [7, 6]

3.3 Step 2: 算 padding 与 cu_seqlens_padded#

align_multiplier = 2          # 非 FP8
align_size       = TP * CP * 2 = 1 * 2 * 2 = 4
remained         = seqlens % align_size = [0, 3, 2]
padding          = [0, 1, 2]
padded_seqlens   = [4, 4, 8]
cu_seqlens_padded= [0, 4, 8, 16]
max_seqlen       = 8

每段 sub-seq 各自补到 4 的倍数，不是只在整个 batch 末尾补。原因：下一步 Ring 要按段切 2*CP=4 块。

3.4 Step 3: flatten + pad 出 packed batch [1, 16]#

P* 表示对齐 pad：

pos:    0  1  2  3   4  5  6  7    8  9  10 11 12 13 14 15
tokens: a  b  c  E   d  e  E  P*   f  g  h  i  j  E  P* P*
        ╰──── doc 0 ───╯ ╰─── doc 1 ───╯ ╰────────── doc 2 ──────────╯
              seg 0           seg 1                 seg 2

valid_local_token_mask:
        1  1  1  1   1  1  1  0    1  1  1  1  1  1  0  0

labels 同步走一遍 _get_flatten_padded_data，但用的 src/dst offsets 是 §2.2 里描述的"跨样本 roll 修过的"版本。

3.5 Step 4: Ring CP slice — 每段切 4 块对称分配#

每段长度都被 4 整除，每段切 2*CP=4 个 chunk，rank r 拿第 r 块 + 第 (2*CP-1-r)=3-r 块。

doc 0  (len 4, chunk_size=1):  位置 [0] [1] [2] [3]
                                rank0 → 0, 3
                                rank1 → 1, 2

doc 1  (len 4, chunk_size=1):  位置 [4] [5] [6] [7]
                                rank0 → 4, 7
                                rank1 → 5, 6

doc 2  (len 8, chunk_size=2):  位置 [8,9] [10,11] [12,13] [14,15]
                                rank0 → 8,9, 14,15
                                rank1 → 10,11, 12,13

合并：

rank0 slice_index = [ 0, 3,  4, 7,  8, 9, 14, 15]   # 长度 8 = 16/CP ✓
rank1 slice_index = [ 1, 2,  5, 6, 10,11, 12, 13]

3.6 Step 5: index_select 后各 rank 看到的 batch#

rank0 tokens     = [a, E,  d, P*,  f, g, P*, P*]
rank0 valid_mask = [1, 1,  1,  0,  1, 1,  0,  0]   # 5 个有效

rank1 tokens     = [b, c,  e,  E,  h, i,  j,  E]
rank1 valid_mask = [1, 1,  1,  1,  1, 1,  1,  1]   # 8 个有效

注意几件事：

1. cu_seqlens_q_padded 两个 rank 上都还是 [0,4,8,16]，是全 batch 视角；TE 内部知道这是 cp_size=2 切过的。
2. 本 rank 看到的 token 不连续：rank0 的 a, E 是 doc 0 的"首块+尾块"，rank1 的 b, c 是 doc 0 的"中间两块"。attention 是在 ring 通信里把远端 KV 拉过来对齐的，靠的就是 cu_seqlens_q_padded + cp_rank。
3. 有效 token 在 rank 之间不平衡：rank0 有 3 个 pad、rank1 全有效。Ring 的对称切分平衡的是 causal 计算量，不平衡有效 token 数——loss reduce 时靠 valid_local_token_mask + DP/CP all-reduce 把分母（token 数）算正确。

3.7 Step 6: attention 内部（概念性）#

attention 收到的输入：

hidden_states.shape = [1, 8, H]          # T_pad/CP
packed_seq_params:
    qkv_format         = "thd"
    cu_seqlens_q       = [0, 4, 8, 16]   # 全长，不缩
    cu_seqlens_q_padded= [0, 4, 8, 16]
    max_seqlen_q       = 8
    valid_local_token_mask = [...]

TE 的 flash_attn_varlen + ring 实现：

• 用 cu_seqlens_q_padded 知道 segment 边界，不同 doc 之间不互相 attend（这就是 packed seq 的语义）
• 每个 chunk 在 ring 里转 cp_size-1 圈，每次收到一段远端 KV 做局部 attention，最后 online softmax 合并
• rank0 在算 doc 0 的"尾块"E（位置 3）时，它的 Q 已经在本 rank，需要的 K/V 是 doc 0 位置 [0,1,2,3]，其中 [1,2] 在 rank1，靠 ring 通信传过来

3.8 Step 7: loss#

# 每 rank 算本地 loss
local_loss   = (losses * loss_mask * valid_local_token_mask).sum()
local_tokens = (loss_mask * valid_local_token_mask).sum()

# CP 维 all-reduce loss 和 token 数
stats = torch.stack([local_loss.detach(), local_tokens.to(local_loss.dtype)])
torch.distributed.all_reduce(stats, group=cp_group)
reported_loss = stats[0] / stats[1].clamp_min(1)
# 这是日志聚合；训练 loss 的梯度缩放还需匹配框架 DP/CP reduction 语义。

# rank0: 5 个有效位置参与, rank1: 8 个有效位置参与
# 合起来 13 个 = sum(valid_token_nums) = 7 + 6 ✓

labels 在 pack 时已经做了 roll，所以每个有效位置 i 上 loss = CE(logits[i], labels[i])，其中 labels[i] 是 doc 内的下一个 token。doc 末尾（EOD 那个位置）的 label 通常被 loss_mask=0 屏蔽掉。

3.9 一张图收尾#

                       全 batch packed (T_pad=16)
   pos:    0  1  2  3  4  5  6  7  8  9  10 11 12 13 14 15
   tok:    a  b  c  E  d  e  E  P  f  g  h  i  j  E  P  P
           ╰── doc0 ──╯╰── doc1 ──╯╰────── doc2 ──────────╯
                ↑           ↑                ↑
           4块×1tok     4块×1tok         4块×2tok
                ┌──────────────────────────────┐
                │ Ring 对称切分 (2*CP=4 块/段)  │
                └──────────────────────────────┘
                            ↓
   rank0 拿块 {0, 3}:  [a E][d P][f g P P]   ← 物理位置 0,3,4,7,8,9,14,15
   rank1 拿块 {1, 2}:  [b c][e E][h i j E]   ← 物理位置 1,2,5,6,10,11,12,13

   两个 rank 都持有: cu_seqlens_q_padded = [0,4,8,16], cp_rank=0/1
   attention 通过 ring 通信把跨 rank 的 K/V 拉过来对齐

3.10 换 backend 后的差异点#

• Ulysses：rank0 拿 [0..7]，rank1 拿 [8..15]——连续切，pad 只在全段末尾。attention 时 all-to-all 换 head 维。
• Magi：dispatch index 由 magi_attn_varlen_key 决定，应用层不可见；pad 用 compute_pad_size，chunk_size 还要 prev_power_of_two。
• CP=1：get_slice_index 直接返回 None，跳过 slice，整段 [1, 16] 直接进 attention。

代码范围：本文涉及特定 Megatron 分支的私有 helper 与当时行号，不代表全部上游版本。全局有效 token 数为零时应跳过更新或报错；日志分母的 clamp 不能掩盖空训练 batch。