LLM序列打包实操：30分钟极速搞定高效训练

各位开发者朋友、跨境实战家们，今天咱们要聊一个在大模型（LLMs）预训练领域里，一个虽不常被挂在嘴边，但却能实实在在提升效率的“小技巧”——那就是打包序列（Packed Sequences）与掩码注意力（Masked Attention）。在新媒网跨境获悉，这可是咱们出海淘金，提升AI模型训练效率的实战利器！

咱们都知道，训练大语言模型，那可是个“吞金兽”和“时间机器”。它需要海量的数据、顶级的硬件配置，更离不开各种巧妙的优化策略。其中一个容易被忽视，但又极具价值的优化点，就是如何充分利用好咱们选择的上下文长度，让每一次训练都能“物尽其用”。

试想一下，当咱们给Transformer模型喂数据时，如果一批文本序列长短不一，为了对齐输入维度，短的序列就得用特殊的“填充（padding）”标记来补齐。这看起来没什么大不了，但实际上，咱们宝贵的GPU内存和计算资源，都在无形中浪费在了处理这些毫无意义的填充标记上。这就像咱们的物流，打包的箱子里有很多空隙，既占地方又没装满货。

妙招揭秘：打包序列

“打包序列”这个方法，就巧妙地解决了这个问题。它的核心思想不是填充，而是将多个较短的序列，拼接成一个更长的序列。这样一来，咱们就能最大限度地减少因为填充标记而产生的计算浪费。同时，每批次（batch）能处理的token数量也随之增加，训练时间自然就缩短了。这就像咱们的集装箱，通过科学规划，把零散的货物高效地装满，节省了运力。

不过，这里有个关键点需要注意：咱们得确保模型在计算注意力时，不会“越界”——也就是说，不能让一个序列的token，去关注到它旁边那个独立序列的token。这就像在物流中，虽然货品打包在一个箱子里，但每件货品的身份和属性还是独立的，不能混淆。

咱们来看个简单的例子。假设咱们要把下面三句话打包成一个序列，每句话之间用一个特殊的“序列结束（EOS）”标记隔开。

# 初始化环境
import torch; torch.set_printoptions(linewidth=200)
from transformers import AutoTokenizer, AutoConfig, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("gpt2")
config = AutoConfig.from_pretrained("gpt2")
model = AutoModelForCausalLM.from_config(config)

sentence1 = "The cat sat on the mat" # 猫坐在垫子上
sentence2 = "The dog ate my homework" # 狗吃了我的作业
sentence3 = "My aunt is a teacher" # 我阿姨是老师
sentences = [sentence1, sentence2, sentence3]
tokenized_sentences = tokenizer(sentences, return_attention_mask=False, add_special_tokens=False)["input_ids"]
# 将所有句子的token和EOS标记拼接起来
tokenized_sentences = [t for s in tokenized_sentences for t in s + [tokenizer.eos_token_id]]

tokenizer.decode(tokenized_sentences)

如果咱们把这个打包后的序列解码回来，它会是这样的：
The cat sat on the mat<|endoftext|>The dog ate my homework<|endoftext|>My aunt is a teacher<|endoftext|>

对于这种打包序列，标准的因果语言模型注意力掩码（attention mask）通常是这样的。

tokenized_sentences = torch.tensor(tokenized_sentences)
attn_mask = torch.ones(tokenized_sentences.size(0), tokenized_sentences.size(0), dtype=torch.int).tril()
attn_mask

然而，如果直接用上面这个掩码，当模型处理第二句话时，它仍然能够“看到”第一句话中的token。这显然不是咱们想要的，因为这两句话是独立的例子，彼此之间不应该产生注意力关联。为了解决这个问题，咱们需要对注意力掩码进行特定的“截断”处理。如果批次中只有一个样本，在Python的PyTorch框架下实现起来相对直观。

def get_attention_mask_for_packed_sequence(x, token_id, eos: bool = True):
    # 将序列长度存入变量T，方便阅读
    T = tokenized_sentences.size(0)
    # 获取所有EOS标记的索引
    eos_indices = (tokenized_sentences == tokenizer.eos_token_id).nonzero().squeeze()
    # 从索引中，计算每个序列的长度
    reps = torch.cat([eos_indices[[0]]+1, eos_indices[1:] - eos_indices[:-1]])
    # 将每个EOS索引重复n次，n是该序列的token数量
    repeated_idx = torch.repeat_interleave(eos_indices, reps).view(1,-1).expand(T, -1)
    # 创建一个张量，包含从0到T-1的所有索引，沿维度1重复T次
    mask_indices = torch.arange(T).view(-1,1).expand(-1, T)
    # 创建因果掩码，并额外屏蔽掉来自前序序列的所有token
    mask = torch.ones(T, T, dtype=torch.bool).tril().expand(-1, -1)
    mask.masked_fill_(mask_indices > repeated_idx, False)
    return mask

get_attention_mask_for_packed_sequence(tokenized_sentences, tokenizer.eos_token_id)

可以看到，标准的因果掩码被“截断”了，它巧妙地屏蔽掉了前一个句子中的token，确保了每个序列的独立性。这正是咱们想要的“精准定位”。

精准定位：调整位置编码（Position IDs）

在打包序列时，另一个非常重要且需要咱们仔细处理的细节，就是相应地调整位置编码（position ids）。每个token通常都有一个关联的位置ID，这个ID能帮助模型理解token在序列中的相对位置。

当咱们把多个序列打包在一起时，需要确保每个序列的位置ID都是从头开始计算的（通常是0或1），而不是接着前一个序列的末尾继续累加。通过调整位置ID，咱们也就清晰地标记了序列的边界。这对于模型区分不同的序列，而不是将打包的数据视为一个连续的整体，至关重要。新媒网跨境认为，这一点直接影响模型对语义的理解能力。

咱们可以复用上面用于生成掩码的部分代码，来生成带正确位置ID的张量。

pos_ids = torch.arange(T) - torch.repeat_interleave(torch.cat([torch.tensor([0]), eos_indices+1], dim=0)[:-1], reps)
pos_ids

批次打包序列的注意力掩码实现

在实际训练中，咱们通常会一次性处理一个批次（batch）的序列。对于上面单样本的示例代码，咱们可能需要通过循环来实现批次的截断注意力掩码。但面对批次维度（batch dimension）的存在，不使用循环的方式实现会更具挑战性。为了展示如何做到这一点，咱们先创建一个包含两组打包序列的批次，也就是一个大小为2的批次。

sentence4 = "Rome wasn't built in a day" # 罗马不是一天建成的
sentence5 = "My hovercraft is full of eels" # 我的气垫船里装满了鳗鱼
sentences = [sentence4, sentence5]
tokenized_sentences2 = tokenizer(sentences, return_attention_mask=False, add_special_tokens=False)["input_ids"]
tokenized_sentences2 = torch.tensor([t for s in tokenized_sentences2 for t in s + [tokenizer.eos_token_id]])

# 使用pad_sequence来填充批次，确保维度一致
batch = torch.nn.utils.rnn.pad_sequence(
    [tokenized_sentences, tokenized_sentences2], batch_first=True, padding_value=tokenizer.eos_token_id
)

咱们将批次的形状（shape）赋值给变量 B（批次大小）和 T（序列长度），这会让后续代码更易读。

B, T = batch.shape

批次实现的主要难点，在于如何构建出与上面单样本示例中类似的“重复索引（repeated_idx）”张量。首先，咱们需要获取批次中所有EOS token的全局索引。

eos_idx = (batch.view(-1) == tokenizer.eos_token_id) \
.nonzero(as_tuple=True)[0] + 1

接着，咱们将0索引和每个批次项的最后一个token索引添加到这个索引向量中。这是为了后续能够再次将批次项分开。然后，咱们移除重复的索引（以防第一个或最后一个索引已存在），并进行排序。

eos_idx_expanded = torch.cat(
    [eos_idx, torch.arange(0,B*T+1,T)]
).unique().sort()[0]

由于咱们的索引向量包含了批次中EOS token的全局索引（例如，第二个批次项的第一个索引=T），咱们需要通过序列长度对这些索引进行归一化。对于归一化后的索引，如果遇到0，则替换为T。这在下一步骤中是必需的。

normalized_idx = eos_idx_expanded - (eos_idx_expanded // T) * T
normalized_idx = torch.where(normalized_idx == 0, T, normalized_idx)

有了归一化后的索引，咱们就能知道每个EOS token索引需要重复多少次，才能得到正确的序列长度。为了实现这一点，咱们需要确保每个序列的最后一个索引都存在。如果在上一步中没有将0替换为T，那么每个批次中最后一个EOS索引的重复次数就会是错误的。

reps = normalized_idx[1:] - normalized_idx[:-1]
reps = torch.where(reps < 1, normalized_idx[1:], reps)

现在，咱们可以创建批次化的重复索引张量了。

repeated_idx = torch.repeat_interleave(
    normalized_idx[1:], reps
).view(B,1,T).expand(-1,T,-1)

剩下的步骤就和单批次样本类似了。咱们构建一个包含从0到T-1的索引的张量，沿维度1重复T次，并创建一个因果掩码。然后，咱们屏蔽掉所有来自前序序列的token。

mask_indices = torch.arange(T).view(1,-1,1).expand(B, -1, T)

# 创建掩码
mask = torch.ones(T, T, dtype=torch.bool).tril().expand(B, -1, -1)
mask = mask.masked_fill(mask_indices >= repeated_idx, False)

下面是完整的函数实现。这里还增加了选择检查EOS token或BOS token的灵活性。

def get_attention_mask_for_packed_sequence(x, token_id, eos: bool = True):
    B, T = x.shape
    # 获取EOS或BOS token的全局索引
    eos_idx = (x.view(-1) == token_id).nonzero(as_tuple=True)[0] + eos
    # 扩展索引，加入每个批次项的边界，并去重排序
    eos_idx_expanded = torch.cat([eos_idx, torch.arange(0,B*T+1,T)]).unique().sort()[0]
    # 归一化索引
    normalized_idx = eos_idx_expanded - (eos_idx_expanded // T) * T
    normalized_idx = torch.where(normalized_idx == 0, T, normalized_idx)
    # 计算重复次数
    reps = normalized_idx[1:] - normalized_idx[:-1]
    reps = torch.where(reps < 1, normalized_idx[1:], reps)
    # 创建批次化的重复索引张量
    repeated_idx = torch.repeat_interleave(normalized_idx[1:], reps).view(B,1,T).expand(-1,T,-1)
    # 构建掩码索引
    mask_indices = torch.arange(T).view(1,-1,1).expand(B, -1, T)
    # 创建因果掩码并应用截断
    mask = torch.ones(T, T, dtype=torch.bool).tril().expand(B, -1, -1)
    mask = mask.masked_fill(mask_indices >= repeated_idx, False)
    return mask

与上面单批次示例类似，咱们也可以复用创建注意力掩码的代码来获取正确的位置ID：

pos_ids = (torch.arange(B*T) - torch.repeat_interleave(eos_idx_expanded[:-1], reps)).view(B,T)
pos_ids

所有代码片段都可以在这个notebook中找到。

风险前瞻与实战提醒：2025年的高效训练之路

各位实战派的朋友们，咱们今天聊的这种打包序列和掩码注意力的优化方法，在2025年的当下，依然是提升大模型训练效率的有效利器。它直接关系到咱们GPU资源的利用率和宝贵的训练时间成本，在咱们“隐形出海”的大背景下，每一分钱、每一秒钟都值得去精打细算。

然而，任何先进的技术应用都伴随着其自身的挑战。这种打包方案虽然高效，但对咱们的代码精细化程度要求更高。具体来说，以下几点是咱们在实际操作中需要特别留意的：

1. 实现复杂度增加：相比于简单的填充方式，打包序列需要更复杂的逻辑来处理序列边界、调整位置ID和生成定制化的注意力掩码。这意味着咱们的训练代码会更复杂，需要更严谨的编程习惯。
2. 调试难度提升：当模型出现问题时，由于序列被拼接在一起，定位是哪个子序列或哪部分逻辑出了错，可能会变得更加困难。精细化的日志记录和可视化工具将是咱们的好帮手。
3. 兼容性考虑：不同的深度学习框架或大模型库，对于这种高级优化策略的支持程度可能有所不同。咱们需要确保所选工具链能够良好地兼容打包序列的实现方式。

时效性说明： 本教程基于当前（2025年）主流的大模型训练技术和PyTorch框架，其核心原理和实现方式在可预见的未来仍将保持其价值。但技术发展日新月异，建议各位朋友持续关注人工智能领域的最新进展，保持学习和更新知识的习惯，以应对未来可能出现的新范式和新工具。毕竟，持续创新，精益求精，才是咱们在跨境赛道上长久立足的根本。希望今天的内容能给咱们的实战带来启发！

新媒网（公号: 新媒网跨境发布），是一个专业的跨境电商、游戏、支付、贸易和广告社区平台，为百万跨境人传递最新的海外淘金精准资讯情报。

本文来源：新媒网 https://nmedialink.com/posts/llm-packed-sequences-fast-train-30min.html