Torch-Titan完全ガイド：PyTorchによる大規模分散学習のすべて

torchtitan/
├── torchtitan/
│   ├── models/
│   │   ├── llama/          # Llamaモデル定義
│   │   └── __init__.py
│   ├── parallelisms/
│   │   ├── parallelize_llama.py  # 並列化適用ロジック
│   │   ├── pipeline_llama.py     # パイプライン並列
│   │   └── __init__.py
│   ├── optimizer.py         # オプティマイザ設定
│   ├── checkpoint.py        # チェックポインティング
│   ├── profiling.py         # プロファイリング
│   └── utils.py
├── train.py                 # トレーニングエントリーポイント
├── train_configs/           # TOML設定ファイル
│   ├── llama3_8b.toml
│   ├── llama3_70b.toml
│   └── llama3_405b.toml
└── estimation.py            # メモリ/FLOPs推定ツール

GPU 0: [Batch 0-15]  →  Gradient 0  ┐
GPU 1: [Batch 16-31] →  Gradient 1  ├→ All-Reduce → 同期
GPU 2: [Batch 32-47] →  Gradient 2  ┘

線形層（d_model=4096、d_ff=16384）を4つのGPUに分散：

GPU 0: 列 0-4095    （4096 x 4096）
GPU 1: 列 4096-8191
GPU 2: 列 8192-12287
GPU 3: 列 12288-16383

GPU 0: Layer 0-9   → activation →
GPU 1: Layer 10-19 → activation →
GPU 2: Layer 20-29 → activation →
GPU 3: Layer 30-39 → loss

シーケンス長4096を4つのGPUに分散：
GPU 0: トークン 0-1023
GPU 1: トークン 1024-2047
GPU 2: トークン 2048-3071
GPU 3: トークン 3072-4095

4D並列の例（128 GPU）：
- DP = 2  （データ並列、2レプリカ）
- TP = 8  （テンソル並列、層あたり8 GPU）
- PP = 4  （パイプライン並列、4ステージ）
- SP = TPと並行して有効化

総GPU数 = DP x TP x PP = 2 x 8 x 4 = 64
（SPを使用するとさらに多くの設定が可能）

# FSDP1スタイル（レガシー）
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
model = FSDP(model, ...)

# FSDP2スタイル（torchtitanで使用）
from torch.distributed._composable.fsdp import fully_shard, MixedPrecisionPolicy

mp_policy = MixedPrecisionPolicy(
    param_dtype=torch.bfloat16,
    reduce_dtype=torch.float32
)

# 各トランスフォーマー層にFSDP2を適用
for layer in model.layers:
    fully_shard(layer, mp_policy=mp_policy)

# モデル全体にも適用
fully_shard(model, mp_policy=mp_policy)

from torch.distributed._composable.fsdp import fully_shard, CPUOffloadPolicy

cpu_offload = CPUOffloadPolicy(offload_params=True)

for layer in model.layers:
    fully_shard(
        layer,
        offload_policy=cpu_offload
    )

from torch.distributed._composable.fsdp import MixedPrecisionPolicy

# 標準的な混合精度設定
mp_policy = MixedPrecisionPolicy(
    param_dtype=torch.bfloat16,   # パラメータ：BF16（通信効率）
    reduce_dtype=torch.float32    # 勾配削減：FP32（安定性）
)

# 特定の層をFP32に保つ（例：埋め込み層）
fully_shard(model.embed_tokens)  # mp_policyなし → FP32
for layer in model.layers:
    fully_shard(layer, mp_policy=mp_policy)
fully_shard(model, mp_policy=mp_policy)

FFN(x) = GELU(x * W1) * W2

列並列（W1を分割）：
GPU 0: x * W1[0:d/4]   -> GELU -> y[0:d/4]
GPU 1: x * W1[d/4:d/2] -> GELU -> y[d/4:d/2]
...

行並列（W2を分割 + All-Reduce）：
GPU 0: y[0:d/4] * W2[0:d/4, :] -> 部分和 0
GPU 1: y[d/4:d/2] * W2[d/4:d/2, :] -> 部分和 1
...
All-Reduce -> 最終出力

# torchtitanでのTP適用（DTensorベース）
from torch.distributed.tensor.parallel import (
    parallelize_module,
    ColwiseParallel,
    RowwiseParallel,
    SequenceParallel,
    PrepareModuleInput,
)

# 並列化プランの定義
plan = {
    "attention": PrepareModuleInput(
        input_layouts=(Shard(1),),
        desired_input_layouts=(Replicate(),),
    ),
    "attention.wq": ColwiseParallel(),
    "attention.wk": ColwiseParallel(),
    "attention.wv": ColwiseParallel(),
    "attention.wo": RowwiseParallel(output_layouts=Shard(1)),
    "feed_forward.w1": ColwiseParallel(),
    "feed_forward.w2": RowwiseParallel(output_layouts=Shard(1)),
    "feed_forward.w3": ColwiseParallel(),
}

parallelize_module(model_layer, tp_mesh, plan)

from torch.distributed.tensor import DTensor, Shard, Replicate
import torch.distributed as dist

# 1Dメッシュの作成（TP用）
tp_mesh = dist.init_device_mesh("cuda", (8,), mesh_dim_names=("tp",))

# 2Dメッシュの作成（DP + TP）
mesh_2d = dist.init_device_mesh(
    "cuda", (2, 8), mesh_dim_names=("dp", "tp")
)

# Shard(0)：行方向シャーディング
# Shard(1)：列方向シャーディング
# Replicate()：レプリケート（シャーディングなし）

標準的なTPフロー：
Input（レプリケート）-> LayerNorm（レプリケート）-> Attention（TP）-> [All-Reduce] -> Output（レプリケート）

SP + TPフロー：
Input（シーケンスシャード）-> LayerNorm（シーケンスシャード）->
[All-Gather] -> Attention（TP）-> [Reduce-Scatter] ->
Output（シーケンスシャード）

単純なパイプライン（4 GPU、4マイクロバッチ）：

時間 ->
GPU 0: [M0 F][M0 B]           [      バブル      ]
GPU 1:        [M0 F][M0 B]    [      バブル      ]
GPU 2:               [M0 F][M0 B][    バブル    ]
GPU 3:                     [M0 F][M0 B]

GPipe（4 GPU、4マイクロバッチ）：

時間 ->
GPU 0: [M0F][M1F][M2F][M3F]               [M3B][M2B][M1B][M0B]
GPU 1:      [M0F][M1F][M2F][M3F]      [M3B][M2B][M1B][M0B]
GPU 2:           [M0F][M1F][M2F][M3F][M3B][M2B][M1B][M0B]
GPU 3:                [M0F][M1F][M2F][M3F][M3B][M2B][M1B][M0B]

# torchtitanでのパイプライン並列設定（TOML）
# [experimental]
# pipeline_parallel_degree = 4
# pipeline_parallel_schedule = "1f1b"

# Pythonレベルのセットアップ
from torchtitan.parallelisms.pipeline_llama import pipeline_llama

# パイプラインステージの作成
stages, model_parts = pipeline_llama(
    model,
    pp_mesh,
    parallel_dims,
    job_config,
    device,
    model_config
)

インターリーブド1F1B（4 GPU、GPU当たり2ステージ）：

GPU 0が担当：[レイヤー0-4] と [レイヤー20-24]
GPU 1：[レイヤー5-9] と [レイヤー25-29]
...

# リポジトリのクローン
git clone https://github.com/pytorch/torchtitan
cd torchtitan

# PyTorchナイトリーのインストール（最新機能を含む）
pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu124

# 依存関係のインストール
pip install -r requirements.txt

# torchtitanパッケージのインストール
pip install -e .

# トークナイザーのダウンロード
python torchtitan/datasets/download_tokenizer.py \
    --repo_id meta-llama/Meta-Llama-3-8B \
    --tokenizer_path "original" \
    --hf_token YOUR_HF_TOKEN

# train_configs/llama3_8b.toml

[job]
dump_folder = "./outputs"
description = "Llama 3 8B training"

[profiling]
enable_profiling = true
save_traces_folder = "profile_trace"
profile_freq = 100

[metrics]
log_freq = 10
enable_tensorboard = true

[model]
name = "llama3"
flavor = "8B"
norm_type = "rmsnorm"
tokenizer_path = "./torchtitan/datasets/tokenizer/original/tokenizer.model"

[optimizer]
name = "AdamW"
lr = 3e-4

[training]
batch_size = 8
seq_len = 2048
warmup_steps = 200
max_norm = 1.0
steps = 1000
data_parallel_replicate_degree = 1
data_parallel_shard_degree = -1  # 自動：残りのGPU数を使用
tensor_parallel_degree = 1
enable_loss_parallel = false

[experimental]
pipeline_parallel_degree = 1

[checkpoint]
enable_checkpoint = true
folder = "checkpoint"
interval_type = "steps"
interval = 500
model_weights_only = false
export_dtype = "float32"
async_mode = "disabled"

[activation_checkpoint]
mode = "selective"  # full、selective、none
selective_ac_option = "op"

[float8]
enable_float8_linear = false

# シングルノード、8 GPUでLlama 3 8Bをトレーニング
torchrun --nproc_per_node=8 \
    train.py \
    --job.config_file train_configs/llama3_8b.toml

# TP=4、DP=2の設定
torchrun --nproc_per_node=8 \
    train.py \
    --job.config_file train_configs/llama3_8b.toml \
    --training.tensor_parallel_degree 4 \
    --training.data_parallel_shard_degree 2

# マルチノードトレーニング（2ノード、各8 GPU）
torchrun \
    --nproc_per_node=8 \
    --nnodes=2 \
    --rdzv_id=101 \
    --rdzv_backend=c10d \
    --rdzv_endpoint=$MASTER_ADDR:29400 \
    train.py \
    --job.config_file train_configs/llama3_70b.toml

# メモリとFLOPsを推定
python estimation.py \
    --job.config_file train_configs/llama3_8b.toml

# サンプル出力：
# Estimated model size: 15.01 GB
# Estimated optimizer state size: 30.02 GB
# Total estimated GPU memory: 65.23 GB
# Estimated FLOP per step: 1.61e+14

標準的なアテンション：
S = Q * K^T    # (seq_len, seq_len)行列 -- メモリが爆発！
P = softmax(S)
O = P * V

seq_len=4096:  Sマトリックス = 4096 x 4096 x 4バイト = 64MB（FP32）
seq_len=32768: Sマトリックス = 32768 x 32768 x 4バイト = 4GB（単一レイヤー！）

# torchtitan/models/llama/model.pyでのアテンション実装
import torch.nn.functional as F

def forward(
    self,
    x: torch.Tensor,
    freqs_cis: torch.Tensor,
) -> torch.Tensor:
    bs, seqlen, _ = x.shape
    xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)

    # RoPE埋め込みの適用
    xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis)

    # PyTorch SDPAによるFlash Attention
    # torch.nn.functional.scaled_dot_product_attention
    # は自動的にFlash Attentionを使用
    output = F.scaled_dot_product_attention(
        xq, xk, xv,
        attn_mask=None,
        dropout_p=0.0,
        is_causal=True,  # 言語モデルのためのCausalマスク
    )

    return self.wo(output.view(bs, seqlen, -1))

pip install flash-attn --no-build-isolation

from flash_attn import flash_attn_func

output = flash_attn_func(
    q, k, v,
    dropout_p=0.0,
    softmax_scale=None,  # 自動：1/sqrt(head_dim)
    causal=True,
)

同期TP：
[GPU計算] -> [通信待機] -> [GPU計算] -> [通信待機] ...

非同期TP：
[GPU計算 A] ────────────────────>
[通信（A結果）] ─────>
                    [GPU計算 B] ->
                          [通信（B）]

# 設定ファイルで非同期TPを有効化：
# [training]
# enable_async_tensor_parallel = true

# コードリファレンス
from torchtitan.parallelisms import parallelize_llama

# async_tpは内部的にtorch.distributed.tensor.parallelの
# async_all_gather機能を使用
model = parallelize_llama(
    model,
    world_mesh,
    parallel_dims,
    job_config
)

# torchtitanのチェックポインティング
import torch.distributed.checkpoint as dcp
from torch.distributed.checkpoint.state_dict import (
    get_model_state_dict,
    get_optimizer_state_dict,
    set_model_state_dict,
    set_optimizer_state_dict,
    StateDictOptions,
)

# 保存
def save_checkpoint(model, optimizer, step, output_dir):
    state_dict = {
        "model": get_model_state_dict(model),
        "optimizer": get_optimizer_state_dict(model, optimizer),
        "extra": {"step": step},
    }

    dcp.save(
        state_dict,
        checkpoint_id=f"{output_dir}/step-{step}",
    )

# 読み込み
def load_checkpoint(model, optimizer, checkpoint_path):
    state_dict = {
        "model": get_model_state_dict(model),
        "optimizer": get_optimizer_state_dict(model, optimizer),
        "extra": {},
    }

    dcp.load(
        state_dict,
        checkpoint_id=checkpoint_path,
    )

    set_model_state_dict(
        model,
        model_state_dict=state_dict["model"],
        options=StateDictOptions(strict=True),
    )
    set_optimizer_state_dict(
        model,
        optimizer,
        optim_state_dict=state_dict["optimizer"],
    )

    return state_dict["extra"]["step"]

# 非同期チェックポインティング設定
[checkpoint]
async_mode = "async"  # "disabled"、"async"、"async_with_pinned_mem"

from torchtitan.checkpoint import CheckpointManager

checkpoint_manager = CheckpointManager(
    dataloader=train_dataloader,
    model_parts=model_parts,
    optimizer=optimizer,
    lr_scheduler=lr_scheduler,
    states={"train_state": train_state},
    job_config=job_config,
)

# トレーニングループ内
for step in range(num_steps):
    # ... トレーニングコード ...

    # 非同期にチェックポイントを保存（ノンブロッキング）
    checkpoint_manager.save(curr_step=step, force=False)

# 分散チェックポイントを単一HuggingFaceモデルに変換
python scripts/convert_checkpoint.py \
    --checkpoint_path outputs/checkpoint/step-1000 \
    --output_path outputs/hf_model \
    --model_type llama3 \
    --model_flavor 8B

import torch
from torch.profiler import profile, record_function, ProfilerActivity

# 基本的なプロファイリング
with profile(
    activities=[
        ProfilerActivity.CPU,
        ProfilerActivity.CUDA,
    ],
    record_shapes=True,    # テンソルシェイプを記録
    profile_memory=True,   # メモリ使用量を記録
    with_stack=True,       # コールスタックを記録
) as prof:
    with record_function("model_inference"):
        output = model(input)

# 結果の表示
print(prof.key_averages().table(
    sort_by="cuda_time_total",
    row_limit=20
))

# Chromeトレースの保存（chrome://tracingで可視化）
prof.export_chrome_trace("trace.json")

[profiling]
enable_profiling = true
save_traces_folder = "profile_trace"
profile_freq = 100             # 100ステップごとにプロファイリング
enable_memory_snapshot = true  # メモリスナップショットを保存

# リファレンス：torchtitan/profiling.pyの内部
from contextlib import contextmanager
import torch.profiler as profiler

@contextmanager
def maybe_enable_profiling(config, global_step=0):
    if not config.profiling.enable_profiling:
        yield
        return

    with profiler.profile(
        activities=[
            profiler.ProfilerActivity.CPU,
            profiler.ProfilerActivity.CUDA,
        ],
        schedule=profiler.schedule(
            skip_first=10,
            wait=5,
            warmup=5,
            active=1,
            repeat=1,
        ),
        on_trace_ready=profiler.tensorboard_trace_handler(
            config.profiling.save_traces_folder
        ),
        record_shapes=True,
        profile_memory=True,
        with_stack=True,
    ) as p:
        yield p

from torch.utils.tensorboard import SummaryWriter

writer = SummaryWriter(log_dir="runs/experiment_1")

# トレーニングループ内
for step, (input, target) in enumerate(dataloader):
    loss = train_step(model, optimizer, input, target)

    writer.add_scalar("Loss/train", loss.item(), step)
    writer.add_scalar("LR", optimizer.param_groups[0]["lr"], step)

    # GPUメモリのモニタリング
    writer.add_scalar(
        "GPU/memory_allocated_GB",
        torch.cuda.memory_allocated() / 1e9,
        step
    )
    writer.add_scalar(
        "GPU/memory_reserved_GB",
        torch.cuda.memory_reserved() / 1e9,
        step
    )

# TensorBoardの起動：
# tensorboard --logdir=runs/

# GPUメモリスナップショット
torch.cuda.memory._record_memory_history(max_entries=100000)

# トレーニングの一部を実行
for step in range(100):
    loss = train_step(model, optimizer, batch)

# メモリスナップショットの保存
torch.cuda.memory._dump_snapshot("memory_snapshot.pickle")

# 分析ツールで可視化：
# https://pytorch.org/memory_vizに.pickleをアップロード

# レイヤーごとの分析のためのメモリプロファイラー
from torch.cuda._memory_viz import profile_plot

with open("memory_profile.html", "w") as f:
    f.write(profile_plot(prof))

def compute_mfu(
    model_num_params: int,
    batch_size: int,
    seq_len: int,
    elapsed_time: float,   # ステップあたりの時間（秒）
    num_gpus: int,
    gpu_peak_flops: float, # GPUの理論的ピークFLOPS
) -> float:
    """
    LLMトレーニングのMFUを計算。
    リファレンス：PaLM論文（Chowdhery et al., 2022）
    """
    # フォワードFLOPs = 2 * num_params * batch_size * seq_len
    # バックワードはフォワードの約2倍 -> 合計6 * num_params * batch_size * seq_len
    flops_per_step = 6 * model_num_params * batch_size * seq_len
    achieved_flops = flops_per_step / elapsed_time
    peak_flops_total = gpu_peak_flops * num_gpus

    return achieved_flops / peak_flops_total

# 使用例
mfu = compute_mfu(
    model_num_params=8e9,     # 8Bパラメータ
    batch_size=8,
    seq_len=2048,
    elapsed_time=0.5,         # ステップあたり0.5秒
    num_gpus=8,
    gpu_peak_flops=989e12,    # H100 BF16：約989 TFLOPS
)
print(f"MFU: {mfu:.1%}")  # 例：MFU: 45.2%

# 一般的に達成可能なMFU：
# - 良好な実装：40〜60%
# - 最適化済み（torchtitan、Megatron）：50〜65%
# - 理論的最大：約70%（通信/メモリオーバーヘッドは避けられない）

# train_configs/llama3_8b_h100x8.toml

[job]
dump_folder = "./outputs/llama3_8b"

[model]
name = "llama3"
flavor = "8B"

[training]
batch_size = 4
seq_len = 8192
data_parallel_replicate_degree = 1
data_parallel_shard_degree = 8  # FSDP2：8 GPU
tensor_parallel_degree = 1

[activation_checkpoint]
mode = "selective"

[float8]
enable_float8_linear = true  # FP8トレーニングを有効化

[optimizer]
name = "AdamW"
lr = 1e-4

torchrun --nproc_per_node=8 train.py \
    --job.config_file train_configs/llama3_8b_h100x8.toml

# train_configs/llama3_70b_32gpu.toml

[training]
batch_size = 2
seq_len = 4096
data_parallel_replicate_degree = 2  # DDP：2レプリカ
data_parallel_shard_degree = 4      # FSDP：4GPUシャーディング
tensor_parallel_degree = 4          # TP：4 GPU

# 総GPU数 = 2 x 4 x 4 = 32

[experimental]
pipeline_parallel_degree = 1  # PP無効

# train_configs/llama3_405b_large.toml

[training]
batch_size = 1
seq_len = 2048
data_parallel_replicate_degree = 4   # DDP
data_parallel_shard_degree = 8       # FSDP
tensor_parallel_degree = 8           # TP

[experimental]
pipeline_parallel_degree = 8         # PP

# 総GPU数 = 4 x 8 x 8 x 8 = 2048