Training Data Shuffling and Batching（訓練データのシャッフルとバッチング）

ミニバッチ SGD は 各更新で訓練分布から i.i.d. に近いサンプルが届く ことを前提に収束保証が組まれる。実装はこれをシャッフル (順序ランダム化) とバッチング (並列化単位構築) の 2 段で支え、LLM ではパッキング・バケッティング・勾配蓄積・分散シャーディングが絡む。Transformer 本体は tech-35、Attention は tech-38 を参照。

なぜシャッフルするか — SGD の i.i.d. 仮定と RR

SGD の収束解析は各更新の勾配が 真の勾配の不偏推定 であることを要求する。理論上は独立復元抽出 (with replacement) だが、実装は通常 置換抜き (=Random Reshuffling, RR) — エポックごとに全体を 1 度ずつ通る。Mishchenko 2020 (arXiv:2006.05988) は強凸 + smooth で RR の収束率を厳密化し、$\|x_K - x^*\|^2 \le \mathcal{O}(1/K^2)$ を示した (with-replacement の $\mathcal{O}(1/K)$ より一段速い)。直観: エポック内で同じサンプルを 2 度引かないため分散項が $1/T^2$ で減衰。Rajput 2020 (arXiv:2007.06245) が非凸でも類似改善を示した。毎エポック再シャッフルする限り収束は保証される側にある。

シャッフルしないとクラスソート済み・時系列のままで block-correlated gradient が連続し、ミニバッチが同一クラスを反復するため事実上 batch GD 様の局所解に陥る。LLM では「同一ドメイン文書の連続供給」で記憶 + 過適合が出る。

シャッフル技法 — メモリ階層に応じた 4 段階

Fisher-Yates はメモリ載り向け。$O(n)$ で完全置換、torch.randperm(n) がこれ。エポック頭で 1 度走らせインデックスでバッチを取る。

Shuffle buffer (ストリーミング近似) は TFRecord・WebDataset・HuggingFace IterableDataset で採用。サイズ $B$ のリングバッファを保持し各ステップでランダム 1 件を取り新規 1 件で穴埋め (tf.data.Dataset.shuffle(buffer_size=B))。$B$ の選び方が品質を左右: 経験則 $B \ge 10 \times \text{batch}$、可能なら $B \ge \text{shard\_size}$。過小だと同一シャード内サンプルが時間相関したまま入る。

シャーデッド (sharded) シャッフル は数 TB 級。① シャードファイル名リストを epoch 頭でシャッフル → ② 各シャード内で Fisher-Yates または buffer → ③ グローバル buffer で混ぜる、の 3 段シャッフル。MosaicML StreamingDataset・NeMo BlendableDataset。決定論的レジューム (seed + step → 同順) が LLM 標準要件。

分散シャッフル (DistributedSampler) は DP rank 間整合性を解く。DistributedSampler(shuffle=True) は全 rank で同 seed + epoch から全置換を生成し rank でスライス。sampler.set_epoch(epoch) を毎エポック呼ばないと全エポック同順 という頻出バグ。

バッチング基礎 — GEMM 効率と勾配ノイズスケール

バッチング動機は 3 つ: ① GPU GEMM 効率 (テンソルコア帯域)、② 勾配分散低下、③ I/O と計算のオーバーラップ。

線形 LR スケーリング則 (Goyal 2017, arXiv:1706.02677): $B$ を $B_{\text{ref}} \to B_{\text{new}}$ に増やすとき LR を $\eta_{\text{new}} = \eta_{\text{ref}} \cdot (B_{\text{new}} / B_{\text{ref}})$ にし、最初の数エポックを線形 warmup。ResNet-50 で 256 → 8192 (32×) まで精度劣化なし。Smith 2017 (arXiv:1711.00489) は LR 減衰 ≡ バッチ増加 の等価性を示した。

勾配ノイズスケール (McCandlish 2018, arXiv:1812.06162): $g_i$ サンプル勾配、$G = \mathbb{E}[g_i]$ で $B_{\text{simple}} = \mathrm{tr}(\mathrm{Cov}(g_i)) / \|G\|^2$。$B \ll B_{\text{simple}}$ なら追加サンプルが勾配品質を線形改善、$B \gg B_{\text{simple}}$ で収益逓減。LLM 事前学習で $B_{\text{simple}}$ は数百万〜数千万トークンに達し、Llama 系の「グローバルバッチ 4M トークン」の根拠。

LLM 特有のバッチング技法

Sequence packing (例パッキング) は LLM 学習の核。可変長文書を max_seq_len に詰めパディング廃棄を削減。素朴 padded batching では平均文書長 $\bar{\ell}$ に対し $1 - \bar{\ell}/L$ が浪費だが、複数文書を eos で繋ぐと 計算予算 100% を更新に使える。T5 (Raffel 2020) で標準化。

Document attention mask はパッキングの必須随伴。素朴に詰めると self-attention が文書境界をまたぎ cross-contamination が起きる。Krell 2021 (arXiv:2107.02027) の per-document ブロック対角マスクでパディング 99% 排除しつつ精度を保つ (FlashAttention varlen で高効率実装)。

Bucketing (長さソート) は同程度長を同バッチに集めパディング削減 (HuggingFace BucketSampler)。短文 → 長文順だと暗黙カリキュラム漏れが入るのでビン順シャッフルの二段方式が標準。

Dynamic batching は max_tokens_per_batch を予算とし長さに応じてサンプル数を可変化 (長文 4 例・短文 128 例)。LR スケジュールは「処理トークン数」で管理。Gradient accumulation は micro_batch を $K$ 回 forward/backward 後に 1 度 step()。実効バッチ $= B_\mu \times K$。BF16 + checkpointing + ZeRO-3 と組合せれば極小 GPU でも 70B 級を学習できる。

分散時の global / micro バッチ計算

DP (data-parallel) / TP (tensor) / PP (pipeline) を混ぜたとき $B_{\text{global}} = B_\mu \cdot G_{\text{accum}} \cdot N_{\text{DP}}$。$N_{\text{TP}}$ と $N_{\text{PP}}$ は 乗じない (同バッチを TP/PP で分割するため)。Megatron-LM は global_batch_size を直接指定し $G_{\text{accum}}$ を自動算出。

1F1B スケジュール: 各 PP ステージが micro_batch を順送り、最後のステージから backward を順次流す。バブル比率 $= (N_{\text{PP}} - 1) / G_{\text{accum}}$ で $G_{\text{accum}} \gg N_{\text{PP}}$ で初めて薄まる。Interleaved 1F1B (Megatron) は非連続レイヤ割当で $(N_{\text{PP}} - 1) / (V \cdot G_{\text{accum}})$ に短縮 ($V$ はインタリーブ度)。ZeRO / FSDP はモデル状態を DP rank 間でシャードし通信を reduce-scatter + all-gather に変更。drop_last=False の uneven batch で同期が壊れることがあり大規模学習では drop_last=True が標準。

カリキュラム・データスケジューリング

上位概念として 何を・いつ・どの比率で見せるか の制御がある。カリキュラム学習 (Bengio 2009) は易しい例から段階的に難化させると最適化が改善する技法。LLM では「短文脈 → 長文脈」のステージ学習 (Llama 3 の long-context phase) が事実上のカリキュラムで、長文 phase で seq_len を 4K → 8K → 32K と上げる。ドメイン混合 は web/書籍/コード/数学/多言語の比率設計で、DoReMi (Xie 2023, arXiv:2305.10429) は小モデルで group-DRO を回し最悪ドメインを高比率化し Pile 22 ドメインの混合率を最適化、収束を 2.6× 加速。実装は per-domain shuffle buffer + BlendableDataset が定石。

落とし穴 Top 5

1. 全 rank 同一 seed → DP の意味喪失。 素の DataLoader(shuffle=True) を全 rank で動かすと全 rank が同じ置換を見る。DistributedSampler + 毎エポック sampler.set_epoch(epoch) 必須。
2. Shuffle buffer 過小 → 暗黙カリキュラム漏れ。 buffer_size $\ll$ shard_size で同一シャード内サンプルが時間相関を残す。最低 $B_{\text{buf}} \ge 10\,000$、推奨 $\ge 100\,000$。
3. Packing without intra-doc mask → cross-contamination。 勾配が無関係文書から漏れ PPL が 1〜3% 膨らむ。flash_attn_varlen_func の cu_seqlens で per-document マスクを付ける。
4. drop_last=False + DDP の uneven last batch → デッドロック。 ランクごとに異なる形状で all-reduce が止まる。drop_last=True か Join context manager で空 step を補う。
5. Per-epoch reshuffle 漏れ (IterableDataset)。 set_epoch が無いと同順序を再生。複数 worker では worker_init_fn で seed = base + worker_id + epoch * world_size を設定。

フレームワーク対応と 2026 年の選択指針

迷ったらこれ: 3 段 shuffle + packing + intra-doc mask + dynamic max-tokens + grad accum で global_batch $\approx B_{\text{simple}}$、DistributedSampler + set_epoch + drop_last=True。

日英用語対照表

参考文献

• Mishchenko 2020 (Random Reshuffling). https://arxiv.org/abs/2006.05988
• Goyal 2017 (linear LR scaling). https://arxiv.org/abs/1706.02677
• McCandlish 2018 (gradient noise scale). https://arxiv.org/abs/1812.06162
• Smith 2017 (Don’t Decay LR, Increase Batch). https://arxiv.org/abs/1711.00489
• Krell 2021 (packing without cross-contamination). https://arxiv.org/abs/2107.02027
• Raffel 2020 (T5). https://arxiv.org/abs/1910.10683
• Xie 2023 (DoReMi). https://arxiv.org/abs/2305.10429
• Bengio 2009 (Curriculum Learning). https://dl.acm.org/doi/10.1145/1553374.1553380
• Narayanan 2021 (Megatron-LM / 1F1B). https://arxiv.org/abs/2104.04473
• PyTorch DistributedSampler. https://pytorch.org/docs/stable/data.html#torch.utils.data.distributed.DistributedSampler
• MosaicML StreamingDataset. https://docs.mosaicml.com/projects/streaming/