Split View: Torch-Titan 완전 가이드: PyTorch 대규모 분산 학습의 모든 것

torchtitan/
├── torchtitan/
│   ├── models/
│   │   ├── llama/          # Llama model definitions
│   │   └── __init__.py
│   ├── parallelisms/
│   │   ├── parallelize_llama.py  # Parallelism application logic
│   │   ├── pipeline_llama.py     # Pipeline parallelism
│   │   └── __init__.py
│   ├── optimizer.py         # Optimizer configuration
│   ├── checkpoint.py        # Checkpointing
│   ├── profiling.py         # Profiling
│   └── utils.py
├── train.py                 # Training entry point
├── train_configs/           # TOML configuration files
│   ├── llama3_8b.toml
│   ├── llama3_70b.toml
│   └── llama3_405b.toml
└── estimation.py            # Memory/FLOPs estimation tool

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

Linear layer (d_model=4096, d_ff=16384) distributed across 4 GPUs:

GPU 0: columns 0-4095    (4096 x 4096)
GPU 1: columns 4096-8191
GPU 2: columns 8192-12287
GPU 3: columns 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

Sequence length 4096 distributed across 4 GPUs:
GPU 0: tokens 0-1023
GPU 1: tokens 1024-2047
GPU 2: tokens 2048-3071
GPU 3: tokens 3072-4095

4D parallelism example (128 GPUs):
- DP = 2  (data parallel, 2 replicas)
- TP = 8  (tensor parallel, 8 GPUs per layer)
- PP = 4  (pipeline parallel, 4 stages)
- SP = activated alongside TP

Total GPUs = DP x TP x PP = 2 x 8 x 4 = 64
(More configurations possible with SP)

# FSDP1 style (legacy)
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
model = FSDP(model, ...)

# FSDP2 style (used in torchtitan)
from torch.distributed._composable.fsdp import fully_shard, MixedPrecisionPolicy

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

# Apply FSDP2 to each transformer layer
for layer in model.layers:
    fully_shard(layer, mp_policy=mp_policy)

# Apply to the full model too
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

# Standard mixed precision configuration
mp_policy = MixedPrecisionPolicy(
    param_dtype=torch.bfloat16,   # Parameters: BF16 (communication efficiency)
    reduce_dtype=torch.float32    # Gradient reduction: FP32 (stability)
)

# Keep specific layers in FP32 (e.g., embedding layer)
fully_shard(model.embed_tokens)  # No 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

Column Parallel (split 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]
...

Row Parallel (split W2 + All-Reduce):
GPU 0: y[0:d/4] * W2[0:d/4, :] -> partial sum 0
GPU 1: y[d/4:d/2] * W2[d/4:d/2, :] -> partial sum 1
...
All-Reduce -> final output

# TP application in torchtitan (DTensor-based)
from torch.distributed.tensor.parallel import (
    parallelize_module,
    ColwiseParallel,
    RowwiseParallel,
    SequenceParallel,
    PrepareModuleInput,
)

# Define parallelization plan
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

# Create 1D mesh (for TP)
tp_mesh = dist.init_device_mesh("cuda", (8,), mesh_dim_names=("tp",))

# Create 2D mesh (DP + TP)
mesh_2d = dist.init_device_mesh(
    "cuda", (2, 8), mesh_dim_names=("dp", "tp")
)

# Shard(0): row-wise sharding
# Shard(1): column-wise sharding
# Replicate(): replicated (no sharding)

Standard TP flow:
Input(replicated) -> LayerNorm(replicated) -> Attention(TP) -> [All-Reduce] -> Output(replicated)

SP + TP flow:
Input(seq-sharded) -> LayerNorm(seq-sharded) ->
[All-Gather] -> Attention(TP) -> [Reduce-Scatter] ->
Output(seq-sharded)

Simple pipeline (4 GPUs, 4 microbatches):

Time ->
GPU 0: [M0 F][M0 B]           [      bubble      ]
GPU 1:        [M0 F][M0 B]    [      bubble      ]
GPU 2:               [M0 F][M0 B][    bubble    ]
GPU 3:                     [M0 F][M0 B]

GPipe (4 GPUs, 4 microbatches):

Time ->
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]

# Pipeline parallelism configuration in torchtitan (TOML)
# [experimental]
# pipeline_parallel_degree = 4
# pipeline_parallel_schedule = "1f1b"

# Python-level setup
from torchtitan.parallelisms.pipeline_llama import pipeline_llama

# Create pipeline stages
stages, model_parts = pipeline_llama(
    model,
    pp_mesh,
    parallel_dims,
    job_config,
    device,
    model_config
)

Interleaved 1F1B (4 GPUs, 2 stages/GPU):

GPU 0 handles: [layers 0-4] and [layers 20-24]
GPU 1: [layers 5-9] and [layers 25-29]
...

# Clone the repository
git clone https://github.com/pytorch/torchtitan
cd torchtitan

# Install PyTorch nightly (includes latest features)
pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu124

# Install dependencies
pip install -r requirements.txt

# Install torchtitan package
pip install -e .

# Download tokenizer
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  # Auto: uses remaining GPU count
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

# Single node, 8 GPU training of Llama 3 8B
torchrun --nproc_per_node=8 \
    train.py \
    --job.config_file train_configs/llama3_8b.toml

# With 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

# Multi-node training (2 nodes, 8 GPUs each)
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

# Estimate memory and FLOPs
python estimation.py \
    --job.config_file train_configs/llama3_8b.toml

# Sample output:
# 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

Standard attention:
S = Q * K^T    # (seq_len, seq_len) matrix -- memory explosion!
P = softmax(S)
O = P * V

seq_len=4096:  S matrix = 4096 x 4096 x 4 bytes = 64MB (FP32)
seq_len=32768: S matrix = 32768 x 32768 x 4 bytes = 4GB (single layer!)

# Attention implementation in 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)

    # Apply RoPE embeddings
    xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis)

    # Flash Attention via PyTorch SDPA
    # torch.nn.functional.scaled_dot_product_attention
    # automatically uses Flash Attention
    output = F.scaled_dot_product_attention(
        xq, xk, xv,
        attn_mask=None,
        dropout_p=0.0,
        is_causal=True,  # Causal mask for language models
    )

    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,  # Auto: 1/sqrt(head_dim)
    causal=True,
)

Synchronous TP:
[GPU compute] -> [wait for comm] -> [GPU compute] -> [wait for comm] ...

Async TP:
[GPU compute A] ────────────────────>
[comm (A result)] ─────>
                    [GPU compute B] ->
                          [comm (B)]

# Enable async TP in the config file:
# [training]
# enable_async_tensor_parallel = true

# Code reference
from torchtitan.parallelisms import parallelize_llama

# async_tp internally uses torch.distributed.tensor.parallel's
# async_all_gather capability
model = parallelize_llama(
    model,
    world_mesh,
    parallel_dims,
    job_config
)

# torchtitan checkpointing
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,
)

# Save
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}",
    )

# Load
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"]

# Async checkpointing configuration
[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,
)

# In the training loop
for step in range(num_steps):
    # ... training code ...

    # Save checkpoint asynchronously (non-blocking)
    checkpoint_manager.save(curr_step=step, force=False)

# Convert distributed checkpoint to single HuggingFace model
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

# Basic profiling
with profile(
    activities=[
        ProfilerActivity.CPU,
        ProfilerActivity.CUDA,
    ],
    record_shapes=True,    # Record tensor shapes
    profile_memory=True,   # Record memory usage
    with_stack=True,       # Record call stacks
) as prof:
    with record_function("model_inference"):
        output = model(input)

# Print results
print(prof.key_averages().table(
    sort_by="cuda_time_total",
    row_limit=20
))

# Save Chrome Trace (visualize at chrome://tracing)
prof.export_chrome_trace("trace.json")

[profiling]
enable_profiling = true
save_traces_folder = "profile_trace"
profile_freq = 100             # Profile every 100 steps
enable_memory_snapshot = true  # Save memory snapshots

# Reference: torchtitan/profiling.py internals
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")

# In the training loop
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 memory monitoring
    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
    )

# Launch TensorBoard:
# tensorboard --logdir=runs/

# GPU memory snapshot
torch.cuda.memory._record_memory_history(max_entries=100000)

# Run part of training
for step in range(100):
    loss = train_step(model, optimizer, batch)

# Save memory snapshot
torch.cuda.memory._dump_snapshot("memory_snapshot.pickle")

# Visualize with analysis tool:
# Upload .pickle at https://pytorch.org/memory_viz

# Memory profiler for per-layer analysis
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,   # Time per step (seconds)
    num_gpus: int,
    gpu_peak_flops: float, # GPU theoretical peak FLOPS
) -> float:
    """
    Compute MFU for LLM training.
    Reference: PaLM paper (Chowdhery et al., 2022)
    """
    # Forward FLOPs = 2 * num_params * batch_size * seq_len
    # Backward is ~2x forward -> total 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

# Example usage
mfu = compute_mfu(
    model_num_params=8e9,     # 8B params
    batch_size=8,
    seq_len=2048,
    elapsed_time=0.5,         # 0.5 seconds per step
    num_gpus=8,
    gpu_peak_flops=989e12,    # H100 BF16: ~989 TFLOPS
)
print(f"MFU: {mfu:.1%}")  # e.g., MFU: 45.2%

# Typically achievable MFU:
# - Good implementation: 40-60%
# - Optimized (torchtitan, Megatron): 50-65%
# - Theoretical max: ~70% (communication/memory overhead unavoidable)

# 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 GPUs
tensor_parallel_degree = 1

[activation_checkpoint]
mode = "selective"

[float8]
enable_float8_linear = true  # Enable FP8 training

[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 replicas
data_parallel_shard_degree = 4      # FSDP: 4-GPU sharding
tensor_parallel_degree = 4          # TP: 4 GPUs

# Total GPUs = 2 x 4 x 4 = 32

[experimental]
pipeline_parallel_degree = 1  # PP disabled

# 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

# Total GPUs = 4 x 8 x 8 x 8 = 2048