torchtitan の実装に関して
各並列化技術が適切なサブメッシュを使用について
1. DeviceMesh（デバイスメッシュ）とは？
- 複数のGPU（やノード）を「多次元の格子状」に論理的に並べたもの
- 例えば128GPUを「8（TP）×2（PP）×2（DPレプリカ）×4（DPシャード）」の4次元メッシュとして扱う
2. サブメッシュとは？
- 全体のDeviceMeshから、特定の並列化技術に必要な次元だけを抜き出した部分メッシュのこと
- 例：TP（Tensor Parallel）は「tp」次元だけ、FSDPは「dp_shard」次元だけ、など
3. 各並列化技術とサブメッシュの対応
| 並列化技術 | サブメッシュ例（次元名） | 役割・意味 |
|---|---|---|
| Tensor Parallel | mesh["tp"] |
各レイヤーの重みを複数GPUで分割計算 |
| Pipeline Parallel | mesh["pp"] |
モデルを複数段階に分割し、各段階を別GPUで計算 |
| Data Parallel | mesh["dp_replicate"] |
モデル全体を複製し、異なるデータで並列学習 |
| FSDP (Sharded DP) | mesh["dp_shard"] |
モデルパラメータを分割して各GPUに持たせる |
| HSDP (Hybrid) | mesh["dp_replicate", "dp_shard"] |
レプリカとシャーディングの組み合わせ |
4. 具体例：128 GPUの4次元メッシュ
parallelism:
tensor_parallel_degree: 8 # 各レイヤーを8GPUで分割
pipeline_parallel_degree: 2 # モデルを2段階に分割
data_parallel_replicate_degree: 2 # 2つのレプリカ
data_parallel_shard_degree: 4 # 各レプリカ内で4GPUでシャーディング
- 全体のDeviceMeshは
[pp=2, dp_replicate=2, dp_shard=4, tp=8]で128GPU - 各技術は自分に必要な次元だけを抜き出して使う
コード例：
# 全体メッシュ
world_mesh = DeviceMesh("cuda", [2, 2, 4, 8], mesh_dim_names=["pp", "dp_replicate", "dp_shard", "tp"])
# 各サブメッシュ
tp_mesh = world_mesh["tp"] # Tensor Parallel用
pp_mesh = world_mesh["pp"] # Pipeline Parallel用
dp_replicate_mesh = world_mesh["dp_replicate"] # DDP用
dp_shard_mesh = world_mesh["dp_shard"] # FSDP用
hsdp_mesh = world_mesh["dp_replicate", "dp_shard"] # HSDP用
サブメッシュの使い分けについて詳細解説
1. なぜサブメッシュを使い分けるのか？
効率的な通信
- 例えばTensor Parallelは「同じレイヤーを分割して持つ」GPU間だけ通信すればよい
- FSDPは「同じレプリカ内でパラメータを分割して持つ」GPU間だけ通信すればよい
柔軟な組み合わせ
- 例えば「TP×PP×FSDP」のような多次元並列化も、サブメッシュを切り出すことで簡単に実現できる
2. 多次元並列化の流れ（例：TP×PP×FSDP）
1. Tensor Parallel（tp_mesh）
- 各レイヤーの重みを8分割し、8GPUで分散計算
2. Pipeline Parallel（pp_mesh）
- モデルを2段階に分割し、2つのGPUグループで前段・後段を担当
3. FSDP（Sharded Data Parallel）（dp_shard_mesh）
- 各パイプライン段階・TPグループ内で、さらに4GPUでパラメータを分割保持
3. サブメッシュの使い分けのメリット
スケーラビリティ
- 大規模GPUクラスタでも効率的にリソースを使える
柔軟性
- モデルやデータの特性に応じて最適な並列化戦略を選べる
シンプルな実装
- 各技術は「自分の担当サブメッシュ」だけを意識すればよい
4. 実装例：torchtitanでのサブメッシュ使用
def parallelize_llama(model, world_mesh, parallel_dims, job_config):
# Tensor Parallel用のサブメッシュ
if parallel_dims.tp_enabled:
tp_mesh = world_mesh["tp"]
apply_tp(model, tp_mesh, ...)
# Data Parallel用のサブメッシュ
if parallel_dims.dp_shard_enabled:
if parallel_dims.dp_replicate_enabled:
# HSDP: レプリカとシャーディングの組み合わせ
dp_mesh = world_mesh["dp_replicate", "dp_shard"]
else:
# FSDP: シャーディングのみ
dp_mesh = world_mesh["dp_shard"]
apply_fsdp(model, dp_mesh, ...)
5. 通信パターンの最適化
各サブメッシュでの通信
- TP: All-reduce（勾配同期）
- PP: P2P通信（パイプライン間のデータ転送）
- FSDP: All-gather（パラメータ収集）とReduce-scatter（勾配分散）
- HSDP: レプリカ間のAll-reduce + シャード間のAll-gather/Reduce-scatter
通信の重複回避
- 各サブメッシュは独立した通信グループを形成
- 不要な通信を避け、効率的な並列化を実現
まとめ
- DeviceMeshで全GPUを多次元格子状に管理
- 各並列化技術は、自分に必要な次元だけを抜き出した「サブメッシュ」を使う
- これにより、多次元並列化（TP, PP, FSDP, HSDPなど）の柔軟な組み合わせが可能
- torchtitanはこの設計により、大規模モデルの効率的な分散学習を実現している