训练超过 100 万个 Token

智能体能够连续数小时处理任务：读取文件、执行命令，并将会话的每一轮交互延续下去。这种累积的会话长度可达数十万 token，这正是前沿模型如今纷纷宣称支持百万 token 上下文窗口的原因。

要让模型擅长处理如此长的上下文，就需要在同等长度的序列上对其进行训练，而这正是困难所在：一百万个 token 的序列无法装入一块 GPU，即便使用八块 GPU 也装不下，除非借助额外的帮助。

本指南中的示例在每个训练步骤中恰好训练一条这样的序列，运行在单个 8-GPU 节点上。

运行

你需要一个配备 8 张 H100（或更高性能）GPU 的节点，并且 transformers>=5.17.0：该示例使用了梯度检查点的 offload 功能，此功能是在该版本中新增的。

torchrun --nproc_per_node 8 examples/sft_qwen3_8b_1m_context/sft_qwen3_8b_1m_context.py

该脚本使用 PG-19 中的书籍对 Qwen3-8B 进行微调，将书籍首尾相连，直到每本书的长度约为百万个 token。经过大约十分钟的加载和分词后，第一步完成：

{'loss': 4.311, 'grad_norm': 29.25, 'num_tokens': 1049000.0, 'epoch': 0.25}
  8%|▊ | 1/12 [06:20<1:09:41, 380.10s/it]

以下三个数字值得关注：

num_tokens 为 1.049e6。 这是一个序列，而不是一批短序列。每个 token 都会关注其之前的所有 token，跨越整个百万级长度。
损失从 4.3 开始，这是一个正常的初始损失值。如果针对长上下文进行了细微的错误配置，运行时的初始损失值会约为 10，我们稍后会解释原因。
每步 380 秒。 一步对应一个序列，因此每个序列耗时略超过六分钟。

直观来看，该序列每张 GPU 需要 288 GB。本指南的其余部分将讲述如何将其降低至 56。

刚才发生了什么

以下内容基于我们已熟悉的配置逐步展开：单 GPU，每条序列包含数千个 token。随着序列长度增加，问题开始显现，且其出现顺序是固定的。

损失是首先要解决的问题

以 Qwen3-4B 为例，在单张 80 GB 显卡上逐步增加序列长度。最先耗尽内存的并不是模型，而是损失值。请检查内存配置。

Memory through one training step, with a large block of memory between the forward and backward passes

峰值出现在前向传播和反向传播之间，即计算损失的位置。要理解峰值为何出现在此处，我们需要查看解码的最后阶段执行了什么操作，以及损失是如何由这些操作构建而成的。

Hidden states times the lm_head weights giving a logits matrix of sequence by vocabulary, then the per-token loss

解码器的最后一层输出形状为 $L \times H$ , 的隐藏状态，其中 $L$ 是序列长度， $H$ 是隐藏层维度。该隐藏状态随后与语言建模头部的权重相乘，生成形状为 $L \times V$ , 的 logits 矩阵，其中 $V$ 是词汇表大小。由于序列较长且词汇表很大（超过十万个条目），该矩阵非常庞大。然后，通过交叉熵函数从该矩阵和目标标签计算损失，而峰值正是由该矩阵的实例化引起的。

我们能否在不实例化整个 logits 矩阵的情况下计算损失？是的，我们可以分块进行。

巨大的 logits 矩阵：对其进行分块

我们不再一次性计算整个序列的 logits，而是每次取几百行，计算这些行的损失，然后将各部分求和。整个 logits 矩阵永远不会同时驻留在内存中。行是能够干净切分的维度：损失是各行上的求和，而其中的 softmax 沿单行的词表维度运行，因此一批行就包含了其自身损失所需的全部信息。

The same computation one chunk of rows at a time, each chunk freed before the next, summing into the same loss

一次大的乘法运算被拆分为多次较小的乘法运算，由此带来的显存节省幅度足够大，使得显存占用不再达到峰值。TRL 每个分块使用 256 行，这确保了每次乘法运算的规模足够大，从而能够充分利用 GPU。

新的配置文件看起来是这样的。

The same step with the chunked loss, where the memory stays flat throughout

峰值已消失。

如何在 TRL 中启用此功能？无需任何操作：分块损失是默认设置。如果您想改用普通损失，请在 SFTConfig 中设置 loss_type="nll"。

这让我们能将序列长度扩展到多少？仅从内存角度来看，这一项改动就将我们支持的 token 数量从 32k 提升到了 160k。

Peak memory against sequence length, the plain loss running out of memory far earlier than the chunked loss

模型从未见过位置 100,000：使用 YaRN 对位置进行重新缩放

我们现在可以在内存中容纳 160k 个 token。但模型真的能读取它们吗？

Qwen3-4B 是在 40,960 个 token 的序列上训练的。配置文件中说明了这一点：

>>> from transformers import AutoConfig
>>> AutoConfig.from_pretrained("Qwen/Qwen3-4B").max_position_embeddings
40960

超出该范围的所有位置都是模型从未遇到过的。程序不会崩溃，但训练过程只是在模型无法定位的 token 上进行。以下是 160k token 序列中每个 token 的损失：

Per-token loss along a 160k-token sequence, flat for most of it then climbing steeply near the end

模型在 40,960 处并未崩溃。它还能继续运行一段时间，直到大约超过 70k 才开始失去连贯性。从那里开始，它稳步上升，最终超过了相同文本的 unigram 熵（6.45）：这比一个完全忽略上下文、仅根据 token 频率进行预测的模型还要差。

The trained range of positions, the much longer range we ask for, and the same range rescaled to fit inside the trained one

每个 token 都带有其位置信息，模型已经学习了 0 到 40,960 之间位置的含义。如果询问它位置 100,000 的情况，它没有任何依据可供参考。

解决之道在于绝不向模型提供它未知的任何位置。我们进行重新缩放：位置 100,000 变为位置 25,000，序列中的每个位置都落回模型已知的范围内。Token 本身没有改变，改变的只是分配给每个 token 的位置。这正是 YaRN 所做的，其缩放因子 simply 取决于我们想要延伸多远：训练长度为 40,960，目标长度为 163,840，因此因子为 4。

如何在 TRL 中启用此功能？在加载模型时传入 RoPE 配置：

training_args = SFTConfig(
    ...,
    model_init_kwargs={
        "rope_parameters": {
            "rope_type": "yarn",
            "rope_theta": 1_000_000,  # has to match what the model ships with
            "factor": 4.0,
            "original_max_position_embeddings": 40960,
        },
    },
)

相同的测量结果，重新缩放的运行结果位于上方：

The same per-token loss with a second run using rescaled positions, which stays flat where the first one climbs

重新缩放后的运行在 160k 之前始终保持平稳。在最后 20k 个 token 中，其平均损失为 2.8，而未进行重新缩放时则为 7.3。

剩下的就是激活值：将它们卸载

在将损失分块并对位置进行重新缩放后，单块 GPU 即可处理 160k 个 token。随后显存再次耗尽。这次是什么占满了显卡？

Memory through one step at 96k tokens, with one red band per layer stacking up through the forward and draining through the backward

模型及其优化器占用固定的 22 GB，与序列长度无关。在此之上，每个红色条带代表一层保存的激活值，每个为 0.47 GB。前向传播写入时它们逐层堆叠，随后保持原状，并在反向传播消耗时逐个移除。顶部的灰色部分代表在毫秒内重新计算并丢弃的所有内容。

TRL 已经为您自动降低了这些开销，并且默认开启。梯度检查点仅保留输入到每一层的数据，而丢弃该层内部计算的所有中间结果：注意力分数、宽 MLP 的中间值、归一化统计量。当反向传播需要这些数据时，它会再次运行该层以重新计算它们。最终保留的是每层一个保存的张量，每个大小为 sequence x hidden，在长序列长度下，这仍然是显存中占用最大的部分。

一个有用的观察是，这些张量在前向传播过程中完全不需要。它们被写入后便不再被触碰，直到很久之后才被读取。因此，在此期间它们不必一直驻留在 GPU 上：可以将它们移至 CPU 内存，并在反向传播到达它们时再取回。

training_args = SFTConfig(..., gradient_checkpointing_kwargs={"offload": True})

The same step with the activations offloaded, where almost no red bands remain resident

相同步数，相同长度：不再是 38 个 band 同时驻留，而是任意时刻仅有 4 个 band 驻留，峰值内存从 59.9 GB 降至 48.8 GB。这换来了序列长度：

Peak memory against sequence length with and without activation offload, the offload line running lower and reaching further

在 48k 以下，两者的表现相同，因为在该长度下几乎没有需要移动的内容。超过该长度后，offloaded 运行的数值保持较低，并且序列长度从原先的 160k 上限提升至在同一张卡上达到 256k。

这一项并非免费。每次保存的激活值都会传输到主机再传回，而这些传输流量必须挤在相邻计算任务之间的空隙中。

一块 GPU 不够用：将序列拆分到多个

分块损失、卸载激活：现在单卡即可在 256k 个 token 上进行训练。如何达到一百万？要弄清阻碍所在，请查看反向传播过程中单个层内部的情况（上述性能分析中的灰色峰值）。

梯度检查点意味着前向传播几乎不保留任何内容：每一层只存储其输入，其余全部丢弃。该层会在反向传播到达它时被重新计算，而那一刻所有数据同时存在。在 MLP 中，模型的一行代码，

down_proj(act_fn(gate_proj(x)) * up_proj(x))

四个张量在完整的中间宽度下构建，标准实现会在反向传播走过它们之前保留全部四个。对于一百万个 token，在一个 MLP 将每个 token 从 2560 扩展到 9728 的模型上：

gate_proj(x)                             19 GB = 1,048,576 x 9728 x 2 bytes
act_fn(gate_proj(x))                     19 GB
up_proj(x)                               19 GB
act_fn(gate_proj(x)) * up_proj(x)        19 GB
                                  total  76 GB

76 GB，仅在一层内，而该卡总容量为 80 GB，且已分配 30 GB 给模型及其优化器。此时，答案并非更聪明地存储序列，而是让每个 GPU 承担更少的序列。

A sequence of tokens overflowing a single GPU and running out of memory, then the same sequence cut into four slices that each fit on one GPU

为每个 GPU 分配序列的一个切片。使用四个 GPU 时，每个 GPU 持有 262,144 个 token，而非 1,048,576 个。这四个张量也随之缩小：从 76 GB 降至 19 GB，加上模型本身的 30 GB，正好能装入显卡。

显而易见的反对意见是注意力机制。每个 token 都必须关注其之前的所有 token，而没有任何一块 GPU 拥有这些更早的 token。因此，GPU 之间会交换各自缺失的部分，最终得到的结果与一块巨大的 GPU 所计算出的结果完全一致。

有两种方式可以执行该交换：上下文并行（CP）和 Ulysses 序列并行（SP）。

	CP	SP
FSDP2	✅	❌
DeepSpeed	❌	✅
设置	`cp_size`	`sp_size`
被拆分的内容	这些标记	注意力头
可扩展性	任意数量的 GPU	最多为 KV 头的数量，此处为 8
注意力	SDPA，仅因果	SDPA 或 FlashAttention
需要	accelerate 1.11	accelerate 1.12，DeepSpeed 0.18.1

在 FSDP2 中，使用 CP。 每个 GPU 保留自己的切片，其他切片依次传递到该 GPU，因此每个 token 最终都能看到所有更早的 token。

from accelerate import ParallelismConfig

training_args = SFTConfig(..., fsdp=True, parallelism_config=ParallelismConfig(cp_size=4))

在 DeepSpeed 上，请使用 SP。它不会来回移动切片，而是在注意力计算前重新排列批次，使每个 GPU 持有所有 token，但只持有四分之一的注意力头，计算完成后再恢复原状。

training_args = SFTConfig(
    ..., deepspeed="examples/deepspeed_configs/zero3.json", parallelism_config=ParallelismConfig(sp_size=4, sp_backend="deepspeed")
)

目前这两个参数互不交叉：cp_size 需要 FSDP2，而 sp_size 仅在 DeepSpeed 下运行。选定后端后，方法随之确定。本节其余部分遵循 FSDP2 路径，这也是本指南顶部示例所采用的方式：其 GPU 并非处理独立批次的独立工作节点，而是共享单个序列的一个组。

开启后，有两点会发生变化：

序列必须填充为 cp_size * 2 的倍数，因此对于四块 GPU，SFTConfig 中的 pad_to_multiple_of=8。
causal-SDPA 的要求排除了 packing，因为 packing 依赖块对角掩码来防止文档之间相互读取，如果同时请求这两者，TRL 会抛出异常。这也排除了那些使用滑动窗口或分块注意力的模型，因为 accelerate 会拒绝它们：OpenAI GPT-OSS、Gemma 3 和 4、Qwen3.5 及更高版本。Qwen3 和 Qwen3-MoE 在所有层中都使用完整注意力，因此它们是本文档中使用的模型。

传递切片所付出的开销可能比你预期的要低。在单节点上处理 131k 个 token 时，使用 2 个 GPU 执行一步需要 34.6 秒，4 个 GPU 需要 17.8 秒，8 个 GPU 需要 9.5 秒。每将组大小翻倍，单步耗时几乎减半；从 2 个 GPU 扩展到 8 个 GPU，恢复了理论最大加速比 4x 中的 3.7x。

这就是四个杠杆中的最后一个：

Peak memory against sequence length on one GPU and on four with context parallelism, the four-GPU line reaching a million tokens where the single GPU stops just past 256k

一张卡刚好在 256k 之后停下。四张卡则能容纳整整一百万，最终占用 63.6 GB，还有富余空间。本指南开头提到的序列长度正好可以装下。

设置 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True。在百万 token 规模下，上述运行在 80 GB 显卡上需要 63.6 GB 显存，若未设置该变量仍会失败：其申请的张量体积大且连续，分配器无法在已持有的内存块中找到足够的空间。

TRL

训练超过 100 万个 Token

运行

刚才发生了什么

损失是首先要解决的问题

巨大的 logits 矩阵：对其进行分块

模型从未见过位置 100,000：使用 YaRN 对位置进行重新缩放

剩下的就是激活值：将它们卸载

一块 GPU 不够用：将序列拆分到多个

延伸阅读