在 GPU算力平台 上进行大模型训练时,显卡利用率不低,单步速度却明显慢于预期。GPU服务器租用实例未必是瓶颈,真正浪费可能来自同一批次中样本长度差异过大:短样本被补齐到最长样本,计算和显存都消耗在 Padding 上。本文用数据统计和分桶采样优化有效吞吐。

问题背景

NLP 训练通常把一个 batch 内的序列补齐到相同长度。若长度分别为 80、120、900、1000,前两个样本的大部分位置都不参与有效学习,却仍经过注意力与前馈计算。深度学习任务中,Padding 比例过高会降低每秒有效 token 数,还可能迫使团队缩小 batch。优化目标不是单纯让 step 更快,而是让相同时间处理更多真实 token,并保持训练分布可控。

环境准备

准备 Python、PyTorch、可统计长度的数据清单和独立基线日志。锁定分词器、最大长度、batch、梯度累积和随机种子。可通过润云智算官网 https://www.smoothcloud.com.cn/ 查看当前环境,再按项目实际依赖选择资源;本文不假设固定型号、价格或镜像能力。

编号实操步骤

1. 统计真实长度分布

使用训练时同一分词器计算 token 数,避免用字符数代替:

lengths = []
for item in dataset:
    ids = tokenizer(
        item["text"],
        add_special_tokens=True,
        truncation=False
    )["input_ids"]
    lengths.append(len(ids))

lengths.sort()
for p in (50, 75, 90, 95, 99):
    idx = min(int(len(lengths) * p / 100), len(lengths) - 1)
    print(f"P{p}:", lengths[idx])

同时记录超长样本数量,决定截断、过滤还是单独处理。

2. 计算Padding浪费率

def padding_ratio(batch_lengths):
    padded = max(batch_lengths) * len(batch_lengths)
    valid = sum(batch_lengths)
    return 1 - valid / padded

print(padding_ratio([80, 120, 900, 1000]))

在基线任务中每隔若干 step 记录该指标,并同时保存单步耗时和有效 token 数。

3. 建立长度分桶

先按长度排序,再划分固定大小的候选桶:

indices = sorted(
    range(len(dataset)),
    key=lambda i: dataset[i]["token_length"]
)
bucket_size = batch_size * 50
buckets = [
    indices[i:i + bucket_size]
    for i in range(0, len(indices), bucket_size)
]

桶不能过小,否则会削弱全局随机性;也不能只按长度永久固定顺序。

4. 桶内打乱并生成批次

import random

rng = random.Random(epoch_seed)
rng.shuffle(buckets)

batches = []
for bucket in buckets:
    rng.shuffle(bucket)
    for i in range(0, len(bucket), batch_size):
        part = bucket[i:i + batch_size]
        if len(part) == batch_size:
            batches.append(part)

rng.shuffle(batches)

每个 epoch 使用可记录的种子,既减少 Padding,又保留可复现的随机顺序。

5. 使用动态批内补齐

from torch.nn.utils.rnn import pad_sequence

def collate_fn(samples):
    seqs = [
        torch.tensor(x["input_ids"], dtype=torch.long)
        for x in samples
    ]
    return pad_sequence(
        seqs,
        batch_first=True,
        padding_value=tokenizer.pad_token_id
    )

不要在数据预处理阶段把全部样本提前补到全局最大长度。

6. 校验标签与掩码

因果语言模型通常要把 Padding 位置的标签设为忽略值,并构造正确的 attention_mask。具体字段与忽略值应以所用模型和损失函数文档为准。若处理错误,速度可能提高但损失曲线和评测会失真。

7. 比较有效吞吐

同时报告每秒 step、每秒有效 token、Padding 比例、显存峰值与验证指标:

valid_tokens = attention_mask.sum().item()
tokens_per_sec = valid_tokens / step_seconds
print(valid_tokens, tokens_per_sec)

只看 step/s 容易把短批次误判为整体性能提升。

8. 处理多卡采样

大模型训练使用 DDP 时,各 rank 的批次数必须一致。可先生成全局批次,再按 rank 分发;尾批次要统一丢弃或补齐。AI算力平台 的实验记录应保存长度统计、分桶参数和采样器版本。GPU算力平台 扩展到多卡前,先用固定小数据集验证各 rank 步数一致。

常见问题与解决方案

  • 分桶后收敛变差:检查随机性是否过低,并增加桶间与批次级打乱。
  • 吞吐提高但指标下降:核对标签忽略位置、注意力掩码和截断策略。
  • 多卡在epoch末尾卡住:确认各 rank 批次数和尾批次策略一致。
  • 短样本过多导致batch偏小:可探索按 token 预算组批,但要同步调整梯度累积与全局 batch 定义。

总结

长度分桶的价值是减少无效 Padding,而不是简单追求更大的 batch。润云智算可作为计算环境入口,具体资源以官网实时信息为准。通过长度分布、浪费率、可复现分桶和有效 token 吞吐四类证据,可以判断优化是否真正提升训练效率。

FAQ

1. 分桶会破坏随机训练吗?

合理的桶内、桶间和批次级打乱可兼顾随机性与长度相近。

2. 字符长度能代替token长度吗?

不建议。不同语言和分词器的字符到 token 映射差异明显。

3. 桶越细越好吗?

不是。桶过细会降低随机性并增加采样器复杂度。

4. 推理部署也适合长度分桶吗?

离线批推理适合;在线请求还要同时考虑排队和延迟目标。

Logo

更多推荐