100万トークン超のトレーニング

エージェントは数時間かけてタスクに取り組むことができます。ファイルの読み取り、コマンドの実行、セッションの各ターンを継続的に処理します。蓄積されたセッションは数十万トークンに達するため、フロンティアモデルでは現在、100万トークン規模のコンテキストウィンドウを宣伝しています。

モデルがこれほど長いコンテキストを扱えるようにするためには、それだけ長いシーケンスで学習させる必要があります。ここが難しい点です。100万トークンのシーケンスは1つのGPUに収まりませんし、8個のGPUでも、何らかの支援なしには収まりません。

このガイドの例では、1ステップごとにちょうど1つのそのようなシーケンスで、単一の8 GPUノード上で学習を行います。

実行する

8つのH100（またはそれ以上の性能を持つGPU）を備えた1つのノードと、transformers>=5.17.0が必要です。この例では、そのリリースで追加された勾配チェックポイントのoffloadを使用しています。

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

このスクリプトは、PG-19の本を端から端まで連結して、それぞれ約100万トークンになるまで長くしたデータでQwen3-8Bをファインチューニングします。読み込みとトークン化に10分ほどかかった後、最初のステップが実行されます:

{'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]

3つの数値が注目すべきです：

• num_tokens は 1.049e6 です。 これは短いシーケンスのバッチではなく、1つのシーケンスです。100万トークン全体にわたって、各トークンはそれより前のすべてのトークンにアテンションを適用します。
• Loss は 4.3 から始まります。これは通常の初期 Loss です。長文脈に対して微妙に設定がずれている実行では、代わりに 10 付近から始まります。その理由については後述します。
• 1ステップあたり380秒。1ステップは1シーケンスに相当するため、シーケンスあたり6分強となります。

単純に計算すると、このシーケンスには GPU ごとに 288 GB 必要です。このガイドの残りの内容は、それを 56 にまで削減する方法の物語です。

先ほど何が起きたか

以下はすべて、すでに知っているセットアップ（GPU 1 基、シーケンスあたり数千トークン）を基盤としています。シーケンスが長くなると問題が発生し始めますが、その順序は固定されています。

損失が最初に壊れる

Qwen3-4Bを単一の80 GBカードで実行し、シーケンスを拡張します。メモリが最初に枯渇するのはモデルではなく、ロスです。メモリプロファイルを確認してください。

ピークはフォワードパスとバックワードパスの間にあり、ここで損失が計算されます。ピークがそこにある理由を理解するには、デコーディングの最終段階が何を行い、損失がそれらからどのように構築されるかを調べる必要があります。

デコーダーの最終層は、形状が$L \times H$, の隠れ状態を出力します。ここで$L$ はシーケンス長、$H$ は隠れサイズです。この隠れ状態は、言語モデリングヘッドの重みと掛け合わされ、形状が$L \times V$, の logits 行列を生成します。ここで$V$ は語彙サイズです。この行列は非常に大きくなります。シーケンスが長く、語彙も大きい（10万項目以上）ためです。その後、この行列とターゲットラベルを用いてクロスエントロピー関数で損失が計算されます。ピークとなるのは、まさにこの行列をメモリ上に実体化することです。

ロスの計算を、logits行列全体をメモリ上に展開せずに実行できますか？はい、チャンク単位で処理を進めることで可能です。

巨大な logits 行列:チャンク分割

logitsをシーケンス全体に対して一度に計算するのではなく、数百行ずつ取り出して、その行の損失を計算し、各部分の和を求めます。これにより、logits行列全体をメモリに一度に保持する必要がなくなります。行は分割が明確な次元です。損失は行方向の和であり、その内部のsoftmaxは1行の語彙に沿って実行されるため、行のチャンクには、その損失を計算するために必要な情報がすべて含まれています。

1つの大きな乗算が複数の小さな乗算に分割され、メモリ削減の効果が十分大きいため、損失が再びピークになることはありません。TRLはチャンクごとに256行を使用しており、各乗算がGPUを十分に活用できるほど十分に大きくなります。

新しいプロファイルは次のようになります。

ピークは消えました。

TRLでこれを有効にするには? 何もする必要はありません。チャンク化された損失がデフォルトです。プレーンな損失を使用したい場合は、SFTConfigでloss_type="nll"を設定してください。

これによりシーケンスをどこまで拡張できるのでしょうか。メモリ面では、この変更だけで32kトークンから160kトークンまで拡張できます。

モデルはポジション100,000を見たことがありません: YaRNでポジションを再スケーリングしてください

160kトークンをメモリに収めることができるようになりました。しかし、モデルは実際にそれらを読み取ることができるのでしょうか？

Qwen3-4Bは40,960トークンのシーケンスで学習されました。設定ファイルにもそう記載されています:

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

それ以降のすべての位置は、モデルが一度も見たことのない位置です。クラッシュはしませんが、実行はモデルが配置できないトークンで学習するだけです。160kトークンのシーケンスのすべてのトークンにおける損失は以下の通りです。

モデルは40,960で崩壊しません。しばらくは継続し、70kを超えたあたりからようやく文脈の追跡を失い始めます。そこから徐々に悪化し、同じテキストのユニグラムエントロピー（6.45）を上回る値で終わります。これは、文脈を完全に無視してトークン頻度のみから予測するモデルよりも悪い結果です。

各トークンにはその位置が割り当てられており、モデルは0から40,960までの位置が何を意味するかを学習しています。位置100,000について尋ねても、モデルには判断の根拠がありません。

解決策は、モデルが知らない位置を渡さないことです。位置を再スケーリングします。位置100,000は位置25,000になり、シーケンス内のすべての位置がモデルが知っている範囲内に収まります。トークン自体は変更されず、各トークンに割り当てられる位置だけが変更されます。YaRNが行うのはこの処理であり、その係数は単に、どこまで伸ばすかを示すものです。学習済み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,
        },
    },
)

同じ測定値を、リスケールされた実行結果を上に配置して表示します:

リスケールされたランは160kまで平坦に保たれています。最後の20kトークンでは、リスケールなしの場合の7.3に対して、損失の平均値は2.8です。

残っているのはアクティベーションです。これらをオフロードします

損失をチャンク化し、位置をリスケールすることで、1つのGPUで160kトークンまで処理できるようになりました。しかし、またしてもリソースが枯渇してしまいます。今回は何がGPUを占有しているのでしょうか？

モデルとオプティマイザーは、シーケンス長に関係なく固定の22 GBを占めます。これに加えて、各赤い帯は1層分の保存されたアクティベーションを、それぞれ0.47 GBで表しています。フォワードパスがこれらを書き込むにつれて積み重なり、バックワードパスが消費するにつれて1つずつ取り除かれます。上部の灰色は、ミリ秒単位で再計算されて破棄されるすべてのデータです。

TRL はこれらをすでに自動で削減しており、デフォルトで有効になっています。Gradient checkpointing は、各レイヤーに入力されるデータのみを保持し、レイヤー内部で計算されるすべての中間結果（attention スコア、幅広な MLP の中間層、norm など）を破棄します。バックワードパスでそれらが必要になった場合、レイヤーを2回目に実行して再計算します。残るのはレイヤーごとに1つの保存テンソルで、それぞれ sequence x hidden のサイズですが、長いシーケンス長では、それでもなお GPU 上で最も大きなメモリ消費要因となります。

有用な観察点は、これらのテンソルはフォワードパスでは一切必要とされないということです。書き込まれた後は放置され、はるかに後になってから読み取られます。したがって、その間GPU上に存在する必要はありません。CPUメモリに移動し、バックワードパスがそれらに到達した際に元に戻すことができます。

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

同じステップ、同じ長さ: 38バンドが積み重なる代わりに、同時に存在するのは常に4バンドのみで、ピークは59.9 GBから48.8 GBに低下します。これによりシーケンス長を確保できます:

48k 未満では、その長さでは移動すべきものがほとんどないため、両者は同じです。それを超えると、オフロードされた実行の方が低く保たれ、以前は 160k で停止していたシーケンスが、同じカード上で 256k に達するようになります。

この手法は無料ではありません。保存された各アクティベーションはホストに渡り、また戻ってくるため、そのトラフィックは隣接する計算処理の間に収まらなければなりません。

1つのGPUでは不十分です: シーケンスを多数に分割

Chunked loss、オフロードされたアクティベーション: 1枚のカードで256kトークンの学習が可能に。100万トークンに到達するには? 妨げとなっているものを見るために、バックワードパス中の1つのレイヤーの中身を確認します(上記のプロファイルの灰色のピーク)。

Gradient checkpointingでは、forward passはほとんど何も保持しません。各レイヤーは自分の入力のみを保存し、それ以外は破棄します。レイヤーはbackward passが到達した時点で再計算され、その瞬間にすべての中間値が同時に存在します。MLPでは、モデルの1行、

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

4つのテンソルがフルの中間幅で構築され、標準的な実装では、バックワードがそれらを通過するまで4つすべてを保持します。100万トークンにおいて、MLPが各トークンを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、1つのレイヤー内、80 GBのカードで、すでにモデルとオプティマイザーに30 GBを割り当て済み。 この時点で、シーケンスをより賢く保存することが答えではありません。各GPUに割り当てる量を減らすことが答えです。

シーケンスをスライスにして、各GPUに1つずつ割り当てます。4つのGPUを使用すると、各GPUが保持するトークン数は1,048,576ではなく262,144になります。4つのテンソルもこれに伴って縮小し、76 GBではなく19 GBになります。これにモデルの30 GBを加えても、カードに収まります。

明白な反論は、アテンションの存在です。各トークンはそれより前のすべてのトークンにアテンションを向ける必要があり、どのGPUもそれより前のトークンを保持していません。そのため、GPU同士は欠けている情報を交換し、その結果は、1つの巨大なGPUが計算したものと完全に一致します。

その交換を実行するには、コンテキスト並列性（CP）とUlyssesシーケンス並列性（SP）の2つの方法があります。

FSDP2では、CPを使用します。各GPUが自身のスライスを保持し、他のスライスが順番にそのGPUに送られてくるため、最終的にすべてのトークンがそれ以前のすべてのトークンを見ることができます。

from accelerate import ParallelismConfig

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

DeepSpeedでは、SPを使用します。スライスを移動する代わりに、アテンションの直前にバッチを再配置し、各GPUがすべてのトークンを持ちながらアテンションヘッドの4分の1のみを保持するようにします。その後、元の配置に戻します。

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

現在、これら2つの設定は互いに独立しています: cp_sizeはFSDP2を必要とし、sp_sizeはDeepSpeedでのみ動作します。バックエンドを選択すると、使用される手法もそれに従って決まります。本セクションの残りは、このガイド冒頭の例で使用されているFSDP2のパスに沿って説明します。このパスでは、GPUは別々のバッチを処理する独立したワーカーではなく、単一のシーケンスを共有する1つのグループとして機能します。

オンになると、2つの点が変わります:

1. シーケンスはcp_size * 2の倍数にパディングする必要があります。そのため、4つのGPUを使用する場合はpad_to_multiple_of=8（SFTConfig内）に設定します。
2. causal-SDPA の要件により、ドキュメント間の相互参照を防ぐためにブロック対角マスクに依存するパッキングは使用できません。両方を要求すると、TRL はエラーを発生させます。また、スライディングウィンドウやチャンク化されたアテンションを使用するレイヤーを持つモデルも対象外です。これらは accelerate が拒否します。対象となるモデルには、OpenAI GPT-OSS、Gemma 3 および 4、Qwen3.5 以降が含まれます。Qwen3 と Qwen3-MoE は全層でフルアテンションを使用しているため、ここではこれらのモデルが採用されています。

スライスをやり取りするコストは、想像よりも低いです。1ノードで131kトークンの場合、1ステップの所要時間は、2 GPUで34.6秒、4 GPUで17.8秒、8 GPUで9.5秒です。グループサイズを2倍にするたびにステップ時間がほぼ半減し、2 GPUから8 GPUへの移行により、理論上の4倍の高速化のうち3.7倍を回復します。

これが4つのレバーの最後です:

1枚のカードでは256kをわずかに超えたところで停止します。4枚あれば100万トークン全体を収容でき、63.6 GBに収まってまだ余裕があります。このガイドの冒頭で示したシーケンスも収まります。

PYTORCH_CUDA_ALLOC_CONF=expandable_segments:Trueを設定してください。100万トークンの場合、上記の実行には80 GBのカードで63.6 GB必要ですが、これを設定しないと依然として失敗します。要求されるテンソルは大きく連続しており、アロケータは既に保持しているブロックの間にそれらを収めるための空き領域を見つけられないためです。

さらに詳しく

• コンテキスト並列性、cp_sizeおよびそれによって構築されるデバイスメッシュに関するAccelerateガイド。
• Ulyssesとリングアテンションについて、2つの交換の違いと、それぞれの通信コストについて説明します。
• RoPEのセクションで使用されている位置の再スケーリングであるYaRN。