システムアーキテクチャと大規模インフラ

人間のフィードバックからの強化学習によるLLMの学習は、アルゴリズム上の課題であると同時に、システムエンジニアリング上の課題でもあります。単一モデル、単一のフォワード・バックワードパス、十分に理解されたスケーリングで構成される標準的な教師ありファインチューニングとは異なり、RLHFでは複数のモデル（方策、参照、報酬モデル、価値ヘッド）を同時にロードし、複雑なロールアウト・スコアリング・学習ループを通じて協調させ、数十から数百のGPUに分散する必要があります。この章では、大規模RLHF学習を可能にするシステムレベルの詳細、すなわちメモリ予算、並列化戦略（データ、テンソル、パイプライン、系列、およびそれらの組み合わせ）、生成ボトルネック、分離アーキテクチャ、重み同期、フォールトトレランス、本番監視を扱います。

4モデルのメモリ課題

Warning

メモリ予算の現実確認 — 70B BF16

方策の重み（BF16）	140 GB
FP32マスター重み	280 GB
Adamオプティマイザ（m + v、FP32）	560 GB
勾配（BF16）	140 GB
参照モデル	140 GB（またはINT8で70 GB）
報酬モデル	140 GB（またはINT8で70 GB）
アクティベーション（バッチ128、系列2048）	50–100 GB
生成用KVキャッシュ	20–60 GB
合計	1470–1560 GB

\(\div\) 80 GB/GPU = 最低19–20台のA100 （並列化オーバーヘッドなし）。

並列化戦略の詳細

大規模言語モデルの学習では、多数のGPUに計算を分散する必要があります。並列化には根本的に異なる軸があり、それぞれに固有のトレードオフがあります。この節では、数学的定式化、図、実践的な指針とともに、各戦略を詳しく説明します。

データ並列化（DP）と分散データ並列化（DDP）

データ並列化は、分散学習で最もシンプルかつ一般的な形態です (Li et al. 2020)。各GPUがモデルの完全なコピーを保持し、異なるミニバッチを処理して、勾配を同期します。

Vanilla DP（PyTorchのDataParallel）。

1プロセス方式で、1台の「マスター」GPUが入力を分散し、出力を集約し、勾配をブロードキャストします。GILと、マスターGPUへのPCIe帯域幅に制限されます。

分散データ並列化（DDP、DistributedDataParallel）。

マルチプロセス方式で、各GPUが独自のプロセスを実行します。バックワード計算を続けながら、バックグラウンドでring-AllReduce (Sergeev and Balso 2018) により勾配を同期します。

DDPの主な性質 ：

• メモリ ：各GPUが完全なモデル + オプティマイザ + 勾配を保持します。70B BF16では \(\sim\)560 GB/GPU となり、メモリ最適化なしでは不可能です。

• 通信 ：ステップごとに勾配テンソルのAllReduceを1回行います。サイズはモデルパラメータ \(\times\) 2バイト（BF16）です。Ring AllReduceのコストは、GPUあたり \(2 \cdot \frac{N-1}{N} \cdot M\) バイトの転送です。

• スケーリング ：\(\sim\)64 GPUまではほぼ線形です。それを超えると通信が支配的になり始めます。

• 勾配バケット化 ：DDPはパラメータをバケット（デフォルト25 MB）にまとめ、バケットの勾配が準備でき次第AllReduceを開始します。通信とバックワード計算をオーバーラップさせる仕組みです。

import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

dist.init_process_group(backend="nccl")  # NCCL for GPU communication
model = model.to(local_rank)
model = DDP(model, device_ids=[local_rank],
            gradient_as_bucket_view=True,    # Memory optimization
            static_graph=True)               # Enable comm optimizations

Warning

DPとDDP — 常にDDPを使う

PyTorchの旧式のDataParallel（DP）は、LLM学習では 決して 使わないでください。

• 1プロセスで、Python GILに制限される

• すべての勾配がGPU 0に集中する（ボトルネック）

• 単一ノードでもDDPより2–3\(\times\)遅い

• 1台のマシンを超えてスケールできない

DDPは最低限の並列化戦略です。\(>\)7BのLLMではFSDP／ZeROが推奨されます。

テンソル並列化（TP）

テンソル並列化（Megatron-LM方式 (Shoeybi et al. 2019)）は、個々の重み行列をGPU間で分割します。各GPUが部分結果を計算し、AllReduceでそれらを結合します。

列並列線形層

重み行列 \(W \in \mathbb{R}^{d \times h}\) を \(T\) 台のGPUに列方向で分割します。

\[ W = [W_0 ;|; W_1 ;|; \cdots ;|; W_{T-1}], \quad W_i \in \mathbb{R}^{d \times h/T} \]

各GPU \(i\) は独立に \(Y_i = XW_i\) を計算します（通信なし）。出力は隠れ次元に沿って分割されます。

行並列線形層

重み行列を行方向に分割します：\(W = [W_0; W_1; \ldots; W_{T-1}]\)。ここで \(W_i \in \mathbb{R}^{d/T \times h}\) です。入力 \(X\) も分割する必要があります。各GPUが部分和を計算し、 AllReduce によって最終出力を生成します。

TPを使うTransformerブロック

Transformer層では、Megatron-LMは次のようにTPを適用します。

1. MLP ：最初の線形層を列並列（\(h \to 4h\)）、2番目を行並列（\(4h \to h\)）にします。行並列層の後にAllReduceを1回行います。

2. アテンション ：\(Q\)、\(K\)、\(V\) 投影を列並列（ヘッドをGPU間で分割）にします。出力投影は行並列です。出力投影の後にAllReduceを1回行います。

3. 合計 ：Transformer層あたりAllReduceを2回（アテンションに1回、MLPに1回）行います。

Tip

TPがノード内に制限される理由

各Transformer層には、上で \(f\) と \(g\) として示したAllReduceが2回必要です。80層の70Bモデルでは、フォワードパスあたり160回（バックワードを含めると320回）のAllReduceになります。NVLink（600 GB/s）では、各AllReduceは \(<\)0.5 msです。しかしInfiniBand（50 GB/s）では同じ操作に \(\sim\)4 msかかり、合計オーバーヘッドは160 \(\times\) 4 = 640 msとなります。これは計算そのものより長い時間です。

ルール ：TPの次数 \(\leq\) ノードあたりのGPU数（通常TP \(\leq\) 8）とします。ノード間のスケーリングにはDP／FSDPを使います。

Important

TP次数の選択

• TP=1 ：テンソル並列化なし。モデルが1台のGPUに収まる（BF16では通常 \(\leq\)13B）。

• TP=2 ：最小分割。2台のGPUで13–34Bの推論に適する。オーバーヘッドが小さい（\(<\)5%）。

• TP=4 ：34–70Bの推論で標準的。オーバーヘッド8–12%。

• TP=8 ：ノード全体。70B以上の学習に必要。オーバーヘッド12–18%。

• TP\(>\)8 ：ノード間TP。PPだけでは不十分な200B以上のモデルでのみ使われ、まれです。オーバーヘッド30–50%。

重要 ：アテンションヘッド数はTP次数で割り切れなければなりません。LLaMA-70B（64ヘッド）で有効なTPは1、2、4、8、16、32、64です。

系列並列化（SP）

系列並列化 (Korthikanti et al. 2023) は、テンソル並列化だけでは解決できないメモリボトルネック、すなわちLayerNormとDropout層の アクティベーションメモリ に対処します。

問題

TPでは重みメモリがGPU間で分割されます。しかしLayerNormとDropoutは完全な隠れ次元で動作し、各GPUに複製されます。バックワードに必要なこれらのアクティベーションは \(b \times s \times d\) に比例するメモリを消費し、TPで削減されず、すべてのGPUで同じ量になります。

解決策

ノード間通信を必要としない演算（LayerNorm、Dropout、残差接続）では、系列次元を分割します。各GPUはこれらの演算について系列の \(s/T\) のスライスを処理し、必要な箇所（アテンション、線形層）でのみ完全な系列を集約します。

Tip

SPの通信は「無料」

標準TPは各サブレイヤーの後でAllReduceを使います。これはReduceScatter + AllGatherと同値です。SPはこれらのプリミティブの順序を入れ替えるだけです。

• SPなしのTP：AllReduce（= ReduceScatter + AllGather）\(\rightarrow\) すべてのGPUに同じデータ \(\rightarrow\) 完全なテンソルに対するLayerNorm（無駄）。

• SPありのTP：ReduceScatter \(\rightarrow\) 各GPUが系列の \(1/T\) を保持 \(\rightarrow\) 部分テンソルに対するLayerNorm \(\rightarrow\) 次のTP層の前にAllGather。

通信量の合計は同一です。SPは 追加の通信コストゼロ の純粋なメモリ最適化です。TPを使うときは常に有効にすべきです。

SPによるメモリ削減（70Bモデル、TP=8、バッチ=4、系列=2048） ：

\[ \text{Activation savings} = (T-1) \times b \times s \times d \times n_\text{layers} \times 2\text{ bytes} = 7 \times 4 \times 2048 \times 8192 \times 80 \times 2 \approx \textbf{59 GB/GPU} \]

パイプライン並列化（PP）

パイプライン並列化は、モデルを層によって垂直に分割し、連続する層のグループを異なるデバイス（ステージ）に割り当てます。アクティベーションはステージを通ってフォワードに流れ、勾配はバックワードに流れます。

バブル問題

素朴なパイプライン実行では、前のステージからの入力や次のステージからの勾配を待つ間、ステージがアイドルになる「バブル」が生じます。

バブル率の式

パイプラインステージが \(P\) 個、1ステップあたりのマイクロバッチが \(M\) 個の場合：

\[ \text{Bubble fraction} = \frac{P - 1}{P + M - 1} \approx \frac{P-1}{M} \quad \text{(when } M \gg P\text{)} \]

バブルのオーバーヘッドを \(<\)10%に抑えるには、\(M \geq 10 \cdot (P-1)\) が必要です。PP=4なら、少なくとも30個のマイクロバッチが必要です。

パイプラインスケジュール

パイプラインスケジューリング戦略

Tip

1F1B：本番標準

1F1B （one-forward-one-backward）スケジュール (Narayanan et al. 2019) は、ほとんどの本番システム（Megatron-LM (Narayanan et al. 2021)、DeepSpeed (Rajbhandari et al. 2020）で使われています。

ウォームアップ ：フォワードパスでパイプラインを満たします（P-1個のマイクロバッチ）。

定常状態 ：各タイムスロットで、フォワード1回とバックワード1回を交互に実行します。これにより、ピークのアクティベーションメモリは \(P\) 個のマイクロバッチに制限されます（GPipeでは \(M\) 個）。

クールダウン ：残りのバックワードパスでパイプラインを空にします。

メモリ上の利点 ：GPipeは \(M\) 個すべてのマイクロバッチのアクティベーションを同時に保存する必要があります。1F1Bは定常状態で \(P\) セットだけを保存するため、\(M = 32\)、\(P = 4\) の場合に重要です。

PPの通信

TP（AllReduce）とは異なり、PPでは隣接ステージ間のアクティベーションの ポイントツーポイント 通信だけが必要です。

\[ \text{Data per transfer} = b_\text{micro} \times s \times d \times 2\text{ bytes (BF16)} \]

マイクロバッチ=4、系列=2048、\(d\)=8192の場合：\(4 \times 2048 \times 8192 \times 2 = 128\) MB/転送です。InfiniBandが50 GB/sなら、転送あたり2.6 msで、ステージあたりの計算時間に比べて小さい値です。

負荷分散

すべての層の計算量が同じではありません。

• 埋め込み層 ：非常に軽い（ルックアップテーブル）。

• Transformerブロック ：均一な計算量。

• 最終LMヘッド ：中程度（語彙投影の大きな行列乗算）。

計算量を均衡させるため、中央のステージにはより多くのTransformer層を割り当て、最初と最後のステージには少なく割り当てます。

完全シャードデータ並列化（FSDP／ZeRO-3）

FSDP (Yanli Zhao et al. 2023)（PyTorch）とZeRO-3 (Rajbhandari et al. 2020)（DeepSpeed）は、DDPに固有のメモリ重複に対処します。すべてのGPUがパラメータ、勾配、オプティマイザ状態の完全なコピーを持つ代わりに、各GPUが \(1/N\) のスライスだけを所有し、必要なときにオンザフライで完全なテンソルを再構成します。

層ごとのFSDP実行フロー ：

1. フォワード ：AllGatherでパラメータ \(\rightarrow\) 計算 \(\rightarrow\) 所有していないシャードを破棄。

2. バックワード ：パラメータを再びAllGather \(\rightarrow\) 勾配を計算 \(\rightarrow\) 勾配をReduceScatter（各GPUが勾配シャードを取得） \(\rightarrow\) 所有していないパラメータシャードを破棄。

3. オプティマイザステップ ：各GPUは、勾配シャードとオプティマイザ状態を使って、自分が所有するシャードだけを更新。

メモリ比較：DDPとFSDP／ZeROの各段階（70Bモデル、8 GPU）。ベースライン：BF16パラメータ（140 GB）+ BF16勾配（140 GB）+ FP32マスター+m+v（840 GB）= GPUあたり1120 GB。

from functools import partial
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import ShardingStrategy, MixedPrecision, BackwardPrefetch
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
from transformers.models.llama.modeling_llama import LlamaDecoderLayer

# Wrap model with FSDP
auto_wrap = partial(transformer_auto_wrap_policy,
                    transformer_layer_cls={LlamaDecoderLayer})
mp_policy = MixedPrecision(
    param_dtype=torch.bfloat16,
    reduce_dtype=torch.bfloat16,
    buffer_dtype=torch.bfloat16,
)

model = FSDP(
    model,
    sharding_strategy=ShardingStrategy.FULL_SHARD,  # ZeRO-3
    mixed_precision=mp_policy,
    auto_wrap_policy=auto_wrap,  # Wrap each transformer layer
    use_orig_params=True,        # Required for torch.compile compatibility
    limit_all_gathers=True,      # Bound peak memory (1 AllGather in flight at a time)
    forward_prefetch=True,       # Prefetch next layer's params during current layer
    backward_prefetch=BackwardPrefetch.BACKWARD_PRE,  # Prefetch during backward
)

Warning

FSDPの通信量

FSDPは、1ステップあたりDDPの 3\(\times\)倍のデータ を通信します。

• DDP：勾配のAllReduce 1回 = リング全体で合計 \(2M\) バイト（\(M\) = バイト単位のモデルサイズ）。

• FSDP：AllGather 2回（フォワード + バックワード）+ ReduceScatter 1回 = \(3M\) バイト。

これはメモリと通信のトレードオフです。FSDPは、(a) DDPではモデルがGPUメモリに収まらない場合、または (b) 通信が計算と十分にオーバーラップする場合（最新のフレームワークでは70–90%のオーバーラップを実現）に有効です。

3D並列化：戦略の組み合わせ

大規模な本番システム（70B以上）では、TP、PP、DP／FSDPを同時に組み合わせます。

Important

本番レシピ：64台のA100-80GB（8ノード）で70B

ノード内 （NVLink 600GB/s）：生成にはTP=8、学習ではノード内FSDP。
ノード間 （InfiniBand 400Gb/s）：ノード間FSDP（8-wayデータ並列）。
結果 ：各GPUは \(\sim\)70GBを保持します。フォワード／バックワード中に方策の重みを層ごとに集約します。
パイプライン並列 ：モデルが100B以上で、TP+ZeROでも収まらない場合だけ使います。複雑さ（バブルのオーバーヘッド10–20%）とスケジューリング上の問題が増えます。

判断フローチャート ：

1. モデルは1台のGPUに収まるか？ \(\rightarrow\) DDPを使います。

2. FSDPを使えば1ノードに収まるか？ \(\rightarrow\) FSDP（ZeRO-3）を使います。

3. TP+FSDPを使えば1ノードに収まるか？ \(\rightarrow\) TP（ノード内）+ FSDP（ノード間）を使います。

4. それでも収まらないか？ \(\rightarrow\) ノード間にPPを追加します。これは最後の手段です。

並列化戦略の比較まとめ

生成ボトルネック：定量分析

Tip

Roofline分析：生成がメモリバウンドになる理由

A100の仕様 ：312 TFLOPS（BF16テンソルコア）、2 TB/sのHBM帯域幅。

Rooflineの交差点 ：\(312\text{T} / 2\text{T} = 156\)。156 FLOP/byte未満の演算はメモリバウンドです。

自己回帰生成 ：各トークンについて、すべての重み（70Bでは140GB）を読み出し、トークンあたり\(2 \times 70\text{B} = 140\text{G}\)を実行します（バッチ=1の場合）。

演算強度 ：\(140\text{G FLOP} / 140\text{GB} = 1\)。これはRooflineの\(156\times\)も下です。

使用率 ：ピークFLOPSのうち\(1/156 = 0.6%\)しか使われません。GPUの99.4%はメモリ読み出しを待ってアイドル状態です。

トークン速度 ：\(2\text{TB/s} / 140\text{GB} = 14.3\)（単一ストリーム、バッチ=1）。

512トークンの場合 ：応答1件あたり\(512 / 14.3 = 35.8\)秒です（バッチ=1、TP=1）。

バッチ化の効果 ：バッチ=64、TP=4なら、重みを一度だけ読み出し、64トークンを並列生成します。演算強度は\(\rightarrow\)\(64 \times 1 = 64\)です。改善しますが、それでもRooflineを下回ります。

70Bモデルの生成スループット（512トークン、各種構成）

最適化スタック （累積高速化）：

1. vLLM + PagedAttention (Kwon et al. 2023)（2–4\(\times\)）：KVキャッシュの断片化をなくし、より大きなバッチを可能にします。

2. 連続バッチ化 (Yu et al. 2022)（1.5–2\(\times\)）：最長系列を待たず、他の系列が終わり次第、新しい系列を開始します。

3. 投機的デコーディング (Leviathan et al. 2023)（2–3\(\times\)）：小さなドラフトモデルが5トークンを提案し、大きなモデルが1回のフォワードパスで検証します。平均で3–4トークンを受理します。

4. 生成用INT8/FP8重み （2\(\times\)）：必要な帯域幅を半減します。学習用に正確なロジットを計算するのではなくサンプリングするため、品質低下は最小限です。

5. CUDAグラフ （1.1–1.3\(\times\)）：形状が固定された演算のカーネル起動オーバーヘッドをなくします。

6. プレフィックスキャッシュ （共有プレフィックスを持つプロンプトで1.5\(\times\)）：システムプロンプトのKVキャッシュを再計算しません。

# Production vLLM generation setup
from vllm import LLM, SamplingParams

engine = LLM(
    model="./policy_checkpoint",
    tensor_parallel_size=4,           # TP=4 per instance
    gpu_memory_utilization=0.92,      # Leave headroom for KV cache
    max_num_batched_tokens=16384,     # Max tokens in flight
    max_num_seqs=256,                 # Max concurrent sequences
    dtype="bfloat16",
    enable_prefix_caching=True,       # Cache system prompt KV
    speculative_model="./draft_1B",   # Speculative decoding
    num_speculative_tokens=5,
    block_size=16,                    # PagedAttention block size
    swap_space=4,                     # GB swap space for preemption
)

# Generate responses for RLHF batch
sampling_params = SamplingParams(
    temperature=0.7, top_p=0.9, max_tokens=512,
    logprobs=1,  # Need log-probs for PPO ratio calculation
)
outputs = engine.generate(prompts, sampling_params)
# Extract: responses, log_probs for each token (needed for PPO/GRPO)

分離アーキテクチャ：本番設計

DeepSpeed-Chat (Z. Yao et al. 2023) やOpenRLHF (J. Hu et al. 2024) などの本番RLHFシステムは、生成・スコアリング・学習を独立してスケールできるクラスタに分離する 分離アーキテクチャ を採用します。

Important

なぜ分離するのか

生成 はメモリ帯域幅律速です（高速なHBMが必要で、計算資源が余ります）。
学習 は計算律速です（テンソルコアが必要で、バックプロパゲーション中は帯域幅が余ります）。
同じハードウェアで両方を最適化することはできません ：すべてを一体化すると、生成時に計算資源を浪費するか、学習時に帯域幅を浪費します。分離すれば、各クラスタが最適なハードウェアと構成を使えます。

実際の利点 ：

• 生成と学習を独立してスケールできる

• 生成クラスタはステートレスなので\(\rightarrow\)、フォールトトレランスが容易

• 生成（ステップ\(N+1\)）と学習（ステップ\(N\)）をオーバーラップできる\(\rightarrow\) 30–40%高速化

• 量子化を使い分けられる：生成はINT8（帯域幅）、学習はBF16（精度）

重み同期戦略

Tip

PPO/GRPOで陳腐化が許容される理由

PPOのクリップ付き目的関数は、オフポリシーデータ向けに設計されています。クリップ \([1-\epsilon, 1+\epsilon]\) が陳腐化したデータの影響を制限します。10–50ステップ陳腐化していても、次のようになります。

• 方策の変化は1ステップあたり\(\sim\)（適切な学習率の場合）

• 50ステップでは：方策のドリフトは\(\sim\)

• PPOのクリップは設計上、最大20%のドリフトを処理できる

• 実証的には：50ステップの陳腐化による品質低下は\(<\)

帯域幅の計算 ：70B BF16 = 140GB。InfiniBand 400Gb/s = 50GB/s \(\rightarrow\) 完全同期は2.8秒です。差分圧縮なら\(<\)、非同期なら無料です（バックグラウンドで実行）。

メモリ最適化技法

追加の技法 ：

• 勾配チェックポイント (Chen et al. 2016：すべてのアクティベーションを保存せず、バックワードパスで再計算します。アクティベーションメモリを\(\sim\)節約する代わりに、\(\sim\)の計算が追加で必要です。選択的に、メモリを多く使うアテンション層だけをチェックポイントし、再計算コストの高いFFNアクティベーションは保持します。

• 混合精度 (Micikevicius et al. 2018：フォワードはBF16（パラメータあたり2バイト）、オプティマイザ状態はFP32（m、vそれぞれ4バイト）で実行します。蓄積用のマスター重みはFP32にします。

• CPUオフロード （ZeRO-Infinity (Rajbhandari et al. 2021）：オプティマイザ状態をCPU RAMへ移します。メモリを50%節約できますが、PCIe 64GB/sがボトルネックとなり、\(\times\)遅くなります。

• アクティベーションオフロード ：フォワード中にアクティベーションをCPUへ移し、バックワード時に戻します。メモリが本当に逼迫している場合だけ使います。

• Flash Attention (Dao et al. 2022; Dao 2024：アテンションのメモリをO(\(n\))からO(\(n^2\))にします。\(\times\)高速化し、長い系列では大幅にメモリを節約できます。

RLHFにおけるFlash Attentionの効果

Tip

RLHFでFlash Attentionが重要な理由

RLHFでは長い系列（ロールアウト）を生成し、それを使って学習します。Flash Attentionがない場合：

• 32ヘッドの4Kトークン系列では、アテンション行列だけで\(\sim\)が必要

• PPO/GRPO学習時のバッチサイズが大きく制限される

• アテンションアクティベーションの勾配チェックポイントは高コスト

Flash Attentionを使うと：

• アテンションメモリは\(O(n)\)であり、\(Q, K, V, O\)テンソルが支配的

• 同じGPUメモリで、より長いロールアウト（8K–32Kトークン）が可能になる

• バックワードパスは\(Q, K, V\)からアテンションタイルを再計算する（\(n^2\)行列を保存しない）

• これが長いコンテキストのRLHF（推論モデルなど）を可能にする鍵

Warning

Flash Attentionと勾配チェックポイント

Flash Attentionのバックワードパスは、保存済みの\(Q, K, V\)からアテンションタイルをオンザフライで再計算します。つまりFlash Attentionは、\(O(n^2)\)のアテンション行列に対するアクティベーション再計算をすでに実装しています。アテンション層を追加でチェックポイントする必要はありません。そうすると\(Q, K, V\)まで不要に再計算することになります。

# DeepSpeed ZeRO-3 configuration for 70B RLHF training
ds_config = {
    "bf16": {"enabled": True},
    "zero_optimization": {
        "stage": 3,
        "overlap_comm": True,                    # Overlap communication with compute
        "contiguous_gradients": True,            # Better memory layout
        "reduce_scatter": True,                  # More efficient than allreduce
        "reduce_bucket_size": 5e7,               # 50M params per bucket
        "prefetch_bucket_size": 5e7,             # Prefetch next bucket
        "param_persistence_threshold": 1e5,      # Keep small params on all GPUs
        "offload_optimizer": {"device": "cpu", "pin_memory": True},  # CPU offload
        "sub_group_size": 1e9,                   # Reduce fragmentation
    },
    "gradient_accumulation_steps": 4,
    "gradient_clipping": 1.0,
    "train_micro_batch_size_per_gpu": 2,
    "wall_clock_breakdown": True,
}

大規模環境でのフォールトトレランス

Warning

ハードウェア障害の現実

GPU 1台のMTBF ：\(\sim\)10,000時間。
512 GPUクラスタのMTBF ：\(10000/512 \approx 20\)時間。ただしソフトウェアやネットワークも考慮すると、 現実には4–8時間 です。\ 数日間の学習実行 ：5–15回の障害に遭遇します。フォールトトレランスがなければ、1回の障害ですべてが停止します。

本番フォールトトレランスの構成 ：

1. 検出 ：NCCLタイムアウト（60秒）、GPUハートビート（10秒）、NVMLによる健全性監視、ECCエラーのカウント。

2. チェックポイント ：50–100ステップごとに非同期で保存します。ノンブロッキング（バックグラウンドスレッド）です。保存対象はモデル重み、オプティマイザ状態（Adam m/v）、スケジューラ状態、RNG状態、KL係数、リプレイバッファです。直近3個のチェックポイントを保持します。時間は70Bで\(\sim\)秒（NVMeへ並列書き込み）です。

3. 復旧 ：(a) 生成クラスタはステートレスなので、再起動して最新の重みを読み込むだけです。(b) 学習クラスタでは、チェックポイントを読み込み、故障ノードを除外してNCCLプロセスグループを再構築し、FSDPシャードを再分配して最後のチェックポイントから再開します。

4. Elastic学習 ：Torch Elastic／Kubernetesのオートスケーリングを使います。数分以内に故障ノードを交換し、一時的に\(N-1\)台のGPUで学習を継続します。

5. 予防 ：GPUの健全性を事前確認します（開始前にGEMMストレステストを実行）。ホットスペアを待機させ、冗長なネットワーク経路（デュアルレールInfiniBand）を用意します。

エンドツーエンドのレイテンシ内訳

監視と可観測性

Important

RLHF学習中に追跡すべき主要指標

品質指標 （10ステップごとに記録）：

• 平均報酬（増加してからプラトーに達するはず）

• 参照モデルからのKLダイバージェンス（3–10に保つはず）

• 応答長の分布（長さハッキングに注意）

• エントロピー（ゆっくり減少させ、崩壊させない）

システム指標 （各ステップで記録）：

• GPU使用率（目標：学習中\(>\)80%、生成中\(>\)60%）

• GPUごとのメモリウォーターマーク（OOMが起きる前に検知）

• 生成スループット（トークン/秒、安定しているはず）

• 勾配ノルム（スパイクは不安定化の前兆）

• NCCL通信時間（ネットワーク劣化を検知）

ネットワークトポロジと通信パターン

効率的な分散学習には、GPUを接続する階層的な通信ファブリックの理解が必要です。現代のクラスタは、超高速なノード内リンクと、低速ですがスケール可能なノード間ネットワークによる2層アーキテクチャを採用します。

ノード内：NVLinkとNVSwitch

NVLinkの世代とLLM学習への影響

単一ノード内（通常は8 GPU）では、 NVSwitch がすべてのGPUペア間にフルバイセクション帯域幅を提供します。つまり、任意のGPUが他の任意のGPUと、全NVLink速度で同時に通信できます。これは、すべての層で8 GPU間のAllReduceが必要なテンソル並列化に不可欠です。

Important

NVSwitchとPCIeトポロジの比較

NVSwitchあり （DGX/HGX）：8 GPUすべてが600–1800 GB/sで全対全接続されます。TPのAllReduceは層あたり\(\sim\)かかります。

NVSwitchなし （PCIeのみのサーバー）：GPUはCPUのPCIeルートコンプレックスを介して32–64 GB/sで通信します。8 GPU間のTPは\(\times\)遅くなります。 PCIeのみのシステムではTP\(>\)を決して使わないでください。

ノード間：InfiniBandとRoCE

ノードをまたぐFSDP/ZeRO-3のAllGatherとReduceScatterでは、ノード間ネットワークが支配的になります。

LLM学習クラスタ向けノード間ネットワークの選択肢

通信プリミティブとそのコスト

各集団通信がいつ使われるかを理解すると、ボトルネックの診断に役立ちます。

分散LLM学習におけるNCCL集団通信操作

ここで\(M\)はメッセージサイズ（バイト）、\(N\)は参加者数です。

Tip

通信と計算のオーバーラップ

現代のフレームワーク（FSDP、DeepSpeed）は、通信と計算を積極的にオーバーラップさせます。

フォワードパス ：層\(i\)の計算中に、AllGatherが層\(i+1\)のパラメータを先読みします。層\(i\)が終わると、そのパラメータは直ちに破棄されます（「フォワード後に解放」）。

バックワードパス ：層\(i\)が勾配を計算している間に、ReduceScatterが層\(i+1\)の勾配を送信します。このオーバーラップにより、適切に調整すれば通信レイテンシの70–90%を隠せます。

調整用ノブ ：prefetch_factor（先読みする層数）、reduce_bucket_size（勾配削減の粒度）、backward_prefetch（バックワード先読み戦略の「pre」対「post」）。

ネットワークトポロジの設計

本番クラスタでは、 ファットツリー または レール最適化 トポロジを使います：

• ファットツリー ：各レベルでフルバイセクション帯域幅を提供します。どのノードも他のノードとフル速度で通信できます。高価（多数のスイッチが必要）ですが、最も柔軟です。

• レール最適化 ：各ノードのGPU\(i\)が同じリーフスイッチ（「レール\(i\)」）に接続します。レール内のAllReduceは安価ですが、レール間トラフィックは高価です。MetaのRSCやGoogleのTPUポッドで使われています。

• 3Dトーラス／Dragonfly ：HPCクラスタ（Frontier、Aurora）で使われます。トポロジを考慮したジョブ配置が重要です。

Warning

ジョブ配置が重要

512 GPUクラスタでは、ノードをランダムに割り当てるとネットワーク輻輳により\(\times\)の速度低下が起こる可能性があります。 常に連続したノードブロックを要求してください。 本番スケジューラ（Slurm、Kubernetes）は局所性を強制し、学習ジョブの全ノードを同じリーフスイッチ上、または互いに1ホップ以内に配置すべきです。

学習スループットとモデルFLOPs使用率

学習効率の測定：MFU

Model FLOPs Utilization（MFU） (Chowdhery et al. 2022) は、学習効率の標準指標です：

\[ \text{MFU} = \frac{\text{Observed throughput (tokens/sec)} \times \text{FLOPs per token}}{\text{Peak hardware FLOPS}} \]

Transformerのパラメータ数を\(P\)、系列長を\(s\)、バッチサイズを\(b\)とすると：

\[ \text{FLOPs per token} \approx 6P + 12 \cdot n_\text{layers} \cdot d_\text{model} \cdot s \]

6という係数は、2（乗算加算）\(\times\) 3（フォワード + バックワード。バックワードはフォワードの\(\approx 2\times\)）に由来します。第2項はアテンションの\(O(s^2)\)コストを表します。

各スケールとハードウェアにおけるMFUベンチマーク

Tip

スケールするとMFUが低下する理由

大きなモデルでは、より多くの並列化が必要になり、次の要因が生じます：

1. 通信オーバーヘッド ：FSDPのAllGather/ReduceScatter（64 GPUで\(\sim\)）

2. パイプラインバブル ：PPはマイクロバッチの開始・終了時にアイドル時間を生む（PP=4で\(\sim\)）

3. 補助モデルのメモリ ：参照モデル／報酬モデルがGPUメモリを使うため、より大きなバッチを載せられない

4. 負荷の不均衡 ：すべての層の計算量が同じではない（埋め込み層とTransformerブロックの違い）

経験則 ：学習ではMFU\(>\) 40%を目標にします。30%未満ならプロファイリングで診断してください。

計算最適なバッチサイズ

実効バッチサイズとハードウェア使用率の関係は単純ではありません：

\[ \text{Effective batch size} = \text{micro_batch} \times \text{grad_accum} \times \text{DP degree} \]

• 小さすぎる場合 ：GPUが十分に使われず（演算強度が低く）、通信が支配的になります。

• 大きすぎる場合 ：トークンあたりの学習効果が逓減し（臨界バッチサイズを超え）、計算を浪費します。

• 適正値 ：勾配ノイズと勾配シグナルが等しくなる臨界バッチサイズ\(B_\text{crit}\)。LLMでは\(B_\text{crit} \sim 1\)–\(4\)Mトークンです (McCandlish et al. 2018)。

RLHFでは特に、バッチは単なるトークンではなくロールアウトを含みます：

\[ \text{RLHF batch} = N_\text{prompts} \times K_\text{generations} \times L_\text{avg response length} \]

一般的な本番値：\(N=128\)プロンプト、\(K=1\)–\(4\)生成、\(L=256\)–\(512\)トークン\(\rightarrow\) 1ステップあたり32K–256Kトークン。

プロファイリングとボトルネック診断

主要なプロファイリングツールと、そこから分かること：

Note

低MFUの診断：チェックリスト

1. GPU使用率\(<\) 80%？ \(\rightarrow\) データ読み込みボトルネック（CPU、I/Oを確認）。

2. カーネル間に大きな空白がある？ \(\rightarrow\) Pythonオーバーヘッドや同期点。CUDAグラフを使います。

3. 通信がステップ時間の\(>\) 20%？ \(\rightarrow\) TP次数を下げ、バッチサイズを増やし、ネットワークの健全性を確認します。

4. メモリが99%？ \(\rightarrow\) バッチを増やせません。勾配チェックポイントやオフロードを試します。

5. 生成中にOOM？ \(\rightarrow\) KVキャッシュが大きすぎます。max_seq_lenまたは生成用バッチサイズを下げます。

コスト分析とクラウド展開

RLHF学習の経済性を理解することは、計画に不可欠です。

ハードウェアのコスト比較

RLHF学習におけるクラウドGPU費用の概算（2024–2025年の価格）

RLHF学習コストの見積もり

\[ \text{Cost} = \frac{N_\text{steps} \times T_\text{step}}{3600} \times N_\text{GPUs} \times C_\text{GPU/hr} \]

Note

コスト例：70BモデルのRLHF（1万ステップ）

ステップ数	10,000
1ステップあたりの時間（分離構成）	45秒
総学習時間	\(10000 \times 45 / 3600 = 125\)時間
GPU（生成 + 学習）	64 A100-80GB
GPU時間あたりの費用（Spot）	$1.20
総費用	\(125 \times 64 \times $1.20 =\) \$9,600

フェーズ別内訳 ：

• 生成クラスタ（32 GPU）：$4,800（時間の60%）

• 学習クラスタ（32 GPU）：\$4,800（オーバーラップ可能\(\rightarrow\)実効\$3,400）

• スコアリング（生成GPUと共有）：上記に含む

オーバーラップ時 ：70Bモデルの完全なRLHFアラインメントの実効費用は\(\approx\)です。

コスト最適化戦略

• Spot／プリエンプティブルインスタンス ：50–70%節約できます。堅牢なチェックポイント（5分ごとに保存）が必要です。

• 適正サイジング ：生成にH100を使わない（メモリバウンド）。推論ではA100の方が、トークン／ドルで同等の性能を実現できます。

• 量子化推論 ：生成とスコアリングにINT8/FP8を使えば、これらのクラスタのGPU数を半減できます。

• 段階的学習 ：報酬設計やデバッグには8Bの代理モデル（\(\sim\)$200）から始め、その後70Bへスケールします。

• 参照モデルなしのLoRA ：参照モデルを完全に排除できます（メモリを50%削減）。

• 短い系列から始める ：256\(\rightarrow\)512\(\rightarrow\)1024トークン生成のカリキュラムにより、計算を40%節約できます。

分散チェックポイント

大規模環境では、素朴なチェックポイント保存がボトルネックになります。オプティマイザ状態を持つ70Bモデルでは、チェックポイント1個あたり\(\sim\)を保存する必要があります（FP32マスター重み + Adam m + v）。

チェックポイント戦略

大規模RLHFのチェックポイント手法

torch.distributed.checkpointによる本番チェックポイント

import torch.distributed.checkpoint as dcp
from torch.distributed.checkpoint.state_dict import get_state_dict, StateDictOptions

# Save: each rank writes its shard in parallel
state_dict = {"model": get_state_dict(model, options=StateDictOptions(full_state_dict=False))}
dcp.save(
    state_dict=state_dict,
    storage_writer=dcp.FileSystemWriter("/mnt/checkpoints/step_5000"),
    planner=dcp.DefaultSavePlanner(),  # Handles FSDP sharding automatically
)

# Async save: non-blocking, runs in background thread
future = dcp.async_save(
    state_dict=state_dict,
    storage_writer=dcp.FileSystemWriter("/mnt/checkpoints/step_5000"),
)
# Training continues immediately; future.result() blocks only if needed

Important

RLHF向けチェックポイントの衛生管理

RLHFのチェックポイントには、通常の事前学習よりも多くの情報を含める必要があります：

• 方策モデルの重み + オプティマイザ状態（標準）

• KL係数（\(\beta\)）とそのスケジュール状態

• リプレイバッファの内容（オフポリシー補正用）

• 全GPUのRNG状態（再現性のため）

• プロンプトイテレータの位置（プロンプトの再処理を避ける）

• 報酬モデルのバージョンタグ（監査可能性のため）

• Wandb／メトリクス実行ID（継続的なログ記録用）

ハードウェア選定ガイド

適切なハードウェアは、モデルサイズ、予算、学習フェーズによって異なります。

モデル規模と学習フェーズ別のハードウェア推奨

Tip

H100とA100：アップグレードはいつ価値があるか

H100には次の特長があります：

• \(\sim\)1.6\(\times\)ピークFLOPS（スパース性ありのBF16では989対624 TFLOPS、スパース性なしでは495対312）

• \(\sim\)2\(\times\)メモリ帯域幅（3.35対2.0 TB/s）

• FP8をサポート（推論で追加の\(\times\)）

• NVLink 4.0（900対600 GB/s）

学習の場合 ：\(\sim\)1.8–2.2\(\times\)高速（FP8のサポートと高い帯域幅が、素のFLOPS優位性を増幅）。

生成の場合 ：\(\sim\)1.7\(\times\)高速（帯域幅律速なので、オーバーヘッド込みで2\(\times\) BW \(\approx\) 1.7\(\times\)スループットに相当）。

費用対性能 ：価格が1.5\(\times\)でも、H100は学習ではほぼ常に優れた価値を提供します。推論のみ（生成クラスタ）なら、Spot価格のA100の方が費用対効果に優れる場合があります。

RL学習用オプティマイザ設定

RL学習（PPO、GRPO、DPO）では、事前学習やSFTとは異なるオプティマイザ上の要求が生じます。損失地形は非定常であり（方策が生成するデータを変える）、勾配はよりノイジーで（報酬信号の分散）、学習は壊滅的忘却や報酬ハッキングを起こしやすくなります。この節では、AdamW (Loshchilov and Hutter 2019)をデフォルトオプティマイザとして、RL固有のオプティマイザ指針をまとめます。

RLで異なるオプティマイザ設定が必要な理由

Important

RLとSFTの最適化：主な違い

• 非定常なデータ分布 ：データセットが固定されたSFTとは異なり、RLは各イテレーションで新しいロールアウトを生成するため、データ分布が方策とともに変化します。

• 高い勾配分散 ：報酬信号は疎でノイジーです。勾配の分散は、整備されたデータに対するクロスエントロピーよりはるかに大きくなります。

• より小さな更新が必要 ：方策は参照モデルの近くに留まる必要があるため（KL制約）、学習率はSFTより\(\times\)小さくします。

• 重み減衰なし ：正則化は重み減衰ではなくKLペナルティから得ます。WDを重ねるとKL制約と競合する可能性があります。

• 短いウォームアップ ：RLは収束済みのSFTチェックポイントから始まるため、オプティマイザ状態に必要なウォームアップは最小限です。

RL手法別の推奨ハイパーパラメータ

RL学習フェーズのオプティマイザ設定。すべて\(\beta_1=0.9\)、\(\beta_2=0.95\)、\(\epsilon=10^{-8}\)、max_grad_norm=1.0、BF16を使用します。

Tip

RLで定数スケジュールを使う理由

コサインや線形減衰スケジュールは、固定された学習期間と単調減少する損失を仮定しています。RL学習にはそのどちらもありません。報酬は予測不能にプラトー、スパイク、振動する可能性があります。短いウォームアップ後の定数LRは、学習全体を通じてオプティマイザの応答性を保ちます。減衰が必要なら、最小LR比を高く（\(\geq 0.5\)）した非常に緩やかな線形スケジュールを使います。

RLではBeta-2 = 0.95：適応を高速化

デフォルトのAdam \(\beta_2 = 0.999\)は、2次モーメントの非常に長い記憶（\(\sim\)ステップの実効ウィンドウ）を持ちます。RL学習では、方策の進化に伴って損失地形が急速に変化するため、1000ステップ前の勾配分散は意味を持ちません。\(\beta_2 = 0.95\)を使うとウィンドウが\(\sim\)に短縮され、変化する勾配統計に適応学習率がすばやく反応します。

Warning

beta2 = 0.95が悪影響を及ぼす場合

非常に小さいバッチサイズ（例えばオンラインRLでバッチ=1）では、\(\beta_2 = 0.95\)により2次モーメント推定がノイジーになりすぎることがあります。この場合は妥協案として\(\beta_2 = 0.99\)を使うか、勾配蓄積によって実効バッチサイズを増やします。

RLの混合精度：FP32マスター重みが重要

RL学習は数値精度に特に敏感です：

• 勾配はよりノイジーであり、小さな更新を多くのステップにわたって正確に蓄積する必要があります

• 学習率は非常に小さく（\(10^{-6}\)–\(10^{-7}\)）、\(\Delta\theta \ll \theta\)となります

• BF16の仮数（7ビット、相対精度\(\approx\)）では、大きさ\(10^{-6}\)の重みに対する相対的な大きさ\(10^{0}\)の更新を表現できません

RL学習では必ずFP32マスター重みを使ってください。 FP32コピーなしのBF16のみの学習では、100–500ステップ後にPPO/GRPOの報酬が確実に崩壊します。

勾配クリッピングはRLで重要

PPOとGRPOでは、特に学習初期に報酬信号が大きく変動することがあります。1つの悪いバッチでノルム\(>100\)の勾配が生じると、モデル重みを完全に破壊します。max_grad_norm=1.0が標準設定です。SFTではクリッピングの重要性は低いものの、やはり推奨されます。

Warning

RLで勾配クリッピングを無効にしてはいけない

勾配ノルムが通常安定しているSFT（0.1–1.0の範囲）とは異なり、RLの勾配はスパイク状です。その理由は、(1) 報酬分散が方策勾配を通じて伝播する、(2) 高報酬の稀な軌跡が過大な更新を生む、(3) 方策がドリフトするとKLペナルティ項が大きな勾配を生む、ためです。\(\\vert \nabla\\vert > 50\)のクリップされていない1ステップで、数百ステップ分の学習が失われる可能性があります。

RL学習の不安定性を診断する

Note

RL最適化の危険信号と対策

症状	考えられる原因と対策
報酬が改善した後に崩壊する	学習率が高すぎる、またはKL係数が低すぎる。学習率を\(\times\)下げるか、\(\beta_\text{KL}\)を増やす。
勾配ノルムが常にクリップ閾値にある	更新が強すぎる。学習率を下げる（クリッピングにより毎回勾配方向の情報を失っている）。
KLダイバージェンスが急増する（\(>\)ナッツ）	学習率が高すぎる。\(\times\)下げるか、適応的KLペナルティを追加する。
報酬がベースラインで停滞する	学習率が低すぎる、または報酬モデルの信号が弱い。\(\times\)高い学習率を試し、報酬モデルのキャリブレーションを確認する。
100ステップ以上で損失がNaNになる	FP32マスター重みがない、または勾配ノルムがオーバーフローしている。FP32マスター重みを有効にし、BF16モードを確認する。

RL用HuggingFace TRL設定

TRLライブラリ (Werra et al. 2022) は、LLM向けにPPO、DPOなどのRL手法を本番利用できる形で実装しています。

from trl import PPOConfig, PPOTrainer, DPOConfig, DPOTrainer

# --- PPO Configuration ---
ppo_config = PPOConfig(
    # Optimizer (AdamW with RL-specific settings)
    learning_rate=1e-6,           # 10-100x smaller than SFT
    
    # PPO-specific
    ppo_epochs=4,                 # mini-batch updates per rollout
    mini_batch_size=16,
    batch_size=64,                # rollout batch size
    
    # Gradient control
    max_grad_norm=1.0,
    
    # KL penalty (replaces weight decay as regularizer)
    init_kl_coef=0.2,            # initial KL penalty coefficient
    adap_kl_ctrl=True,           # adaptive KL targeting
    target_kl=6.0,               # target KL divergence
    
    # Mixed precision
    bf16=True,                   # BF16 compute, FP32 master weights
)

ppo_trainer = PPOTrainer(
    model=model,
    ref_model=ref_model,
    config=ppo_config,
    tokenizer=tokenizer,
    dataset=dataset,
)

# --- DPO Configuration ---
dpo_config = DPOConfig(
    output_dir="./dpo_output",
    
    # Optimizer
    learning_rate=5e-7,           # even smaller than PPO
    optim="adamw_torch",
    adam_beta1=0.9,
    adam_beta2=0.95,              # shorter memory for RL
    weight_decay=0.0,            # no WD -- KL provides regularization
    
    # Schedule
    lr_scheduler_type="constant_with_warmup",
    warmup_steps=50,
    
    # Gradient control
    max_grad_norm=1.0,
    
    # DPO-specific
    beta=0.1,                    # KL constraint strength
    loss_type="sigmoid",         # standard DPO loss
    
    # Mixed precision
    bf16=True,
    
    # Training
    num_train_epochs=1,          # DPO typically 1 epoch
    per_device_train_batch_size=4,
    gradient_accumulation_steps=8,
)

dpo_trainer = DPOTrainer(
    model=model,
    ref_model=ref_model,
    args=dpo_config,
    train_dataset=dataset,
    tokenizer=tokenizer,
)
dpo_trainer.train()

RL学習におけるMoEの考慮事項

Tip

RLHFにおけるMoE

Mixture-of-Experts（MoE）モデル (Fedus et al. 2022) は、RLHFでますます使われるようになっています：

• 利点 ：同じ計算コストで\(\times\)の容量。より多くの容量で判定できるため、報酬モデルに適しています。

• 課題 ：エキスパート並列化には全対全通信が必要です（トークンがGPU間でルーティングされる）。これはパイプライン並列化と競合します。

• MoEでのGRPO ：生成コストは総パラメータ数ではなくアクティブなパラメータ数に支配されるため、うまく機能します。

• MoEでのLoRA ：ルータと共有層だけ、またはすべてのエキスパートにLoRAを適用できます（後者は高コスト）。

Tip

RLオプティマイザの心得

RLファインチューニングでは、 小さなLR、重み減衰なし、定数スケジュール、FP32マスター重み、強いクリッピング を使います。正則化はKLペナルティに任せ、オプティマイザの役割は、行き過ぎずに方策勾配へ追従することだけです。