FSDP TP

.gitignore

@@ -5,3 +5,6 @@ out
 website/
 docs-minimind/
 .venv/
+checkpoints/
+dataset/
+profile_out/

scripts/analyze_profile.py

@@ -0,0 +1,197 @@
+"""
+Parse torch.profiler chrome trace, summarize top ops + backend detection.
+
+Usage:
+    # trace_out/ 是 train_pretrain.py --profile_out 指定的目录
+    python analyze_profile.py --trace_dir ../profile_out
+
+    # 或指定具体 trace 文件（一个 rank 一个）
+    python analyze_profile.py --trace_file ../profile_out/DEVvm54331_XXX.pt.trace.json
+"""
+import argparse
+import glob
+import gzip
+import json
+import os
+from collections import defaultdict
+
+
+# op 名 → bucket 分类规则（按前缀 / substring 匹配）
+BUCKET_RULES = [
+    ('matmul (GEMM)',        ['aten::mm', 'aten::bmm', 'aten::addmm', 'aten::linear']),
+    ('attention (SDPA/FA)',  ['flash', 'scaled_dot_product', 'sdpa', 'attention',
+                              'aten::_scaled_dot_product', 'FlashAttn', 'mem_efficient',
+                              'triton_flash', 'ampere_fp16_s16816gemm_bf16']),
+    ('elementwise (norm/act)', ['rms_norm', 'RMSNorm', 'aten::silu', 'aten::mul', 'aten::add',
+                                 'aten::div', 'aten::rsqrt', 'aten::sub', 'element_wise']),
+    ('softmax',              ['softmax', 'log_softmax', 'cross_entropy']),
+    ('rope',                 ['rope', 'rotary', 'cos', 'sin']),
+    ('embedding',            ['embedding', 'aten::embedding']),
+    ('dropout',              ['dropout']),
+    ('reduce (all-reduce/gather)', ['nccl', 'ncclAllReduce', 'ncclAllGather', 'ncclReduceScatter',
+                                     'allreduce', 'AllReduce', 'ProcessGroupNCCL']),
+    ('memory (copy/permute/view)', ['aten::to', 'aten::_to_copy', 'aten::contiguous',
+                                     'aten::permute', 'aten::view', 'aten::reshape',
+                                     'aten::transpose', 'aten::clone', 'Memcpy']),
+    ('optimizer',            ['adam', 'AdamW', 'optim', 'clip_grad']),
+]
+
+# attention backend 识别（chrome trace event name）
+SDPA_BACKENDS = {
+    'flash':      ['FlashAttn', 'flash_attn', 'triton_flash', 'ampere_fp16_s16816gemm'],
+    'mem_efficient': ['mem_efficient', 'MemEfficient', 'cutlassF'],
+    'math':       ['aten::baddbmm', 'aten::softmax'],  # naive fallback
+}
+
+
+def bucket_op(name: str) -> str:
+    ln = name.lower()
+    for bucket, patterns in BUCKET_RULES:
+        for p in patterns:
+            if p.lower() in ln:
+                return bucket
+    return 'other'
+
+
+def load_trace(path):
+    """支持 .json 和 .json.gz。"""
+    opener = gzip.open if path.endswith('.gz') else open
+    with opener(path, 'rt') as f:
+        data = json.load(f)
+    return data.get('traceEvents', data)
+
+
+def analyze(events):
+    # 只统计 CUDA kernel 时间（cat == 'kernel' 或 'cuda_runtime'）
+    # 但 torch.profiler 输出的 chrome trace 里 CUDA op 是 name+dur 的 complete event (ph=='X')
+    stats = defaultdict(lambda: {'time_us': 0, 'count': 0})
+    total_cuda_us = 0
+    sdpa_backend_hits = defaultdict(int)
+
+    for e in events:
+        if e.get('ph') != 'X':
+            continue
+        cat = e.get('cat', '')
+        name = e.get('name', '')
+        dur = e.get('dur', 0)
+        # 只关心 GPU kernel + collective + cudaLaunch，跳过 python-level
+        if cat not in ('kernel', 'cuda_runtime', 'user_annotation', 'gpu_memcpy', 'gpu_memset', 'cpu_op'):
+            continue
+        if cat == 'cpu_op' and 'ncclAllReduce' not in name and 'sync' not in name.lower():
+            # CPU op 大部分是 dispatch overhead，除了 NCCL / sync 都跳
+            continue
+        stats[name]['time_us'] += dur
+        stats[name]['count'] += 1
+        if cat in ('kernel', 'gpu_memcpy', 'gpu_memset'):
+            total_cuda_us += dur
+
+        # 探测 attention backend
+        for backend, patterns in SDPA_BACKENDS.items():
+            for p in patterns:
+                if p in name:
+                    sdpa_backend_hits[backend] += dur
+
+    return stats, total_cuda_us, sdpa_backend_hits
+
+
+def format_us(us):
+    if us > 1e6: return f'{us/1e6:.2f} s'
+    if us > 1e3: return f'{us/1e3:.2f} ms'
+    return f'{us:.0f} us'
+
+
+def report(stats, total_cuda_us, sdpa_backend_hits, top_n=25):
+    print(f'\n{"=" * 80}')
+    print(f'  Total CUDA time in trace: {format_us(total_cuda_us)}')
+    print(f'{"=" * 80}')
+
+    # bucket 汇总
+    bucket_time = defaultdict(int)
+    bucket_count = defaultdict(int)
+    for name, s in stats.items():
+        b = bucket_op(name)
+        bucket_time[b] += s['time_us']
+        bucket_count[b] += s['count']
+
+    print(f'\n  Time by bucket (% of total CUDA):')
+    print(f'  {"bucket":40s} {"time":>12s} {"% total":>8s} {"# calls":>8s}')
+    for bucket, t in sorted(bucket_time.items(), key=lambda kv: -kv[1]):
+        pct = t / max(total_cuda_us, 1) * 100
+        print(f'  {bucket:40s} {format_us(t):>12s} {pct:>7.1f}% {bucket_count[bucket]:>8d}')
+
+    # attention backend 探测
+    print(f'\n  Attention backend detection:')
+    if sdpa_backend_hits:
+        winner = max(sdpa_backend_hits.items(), key=lambda kv: kv[1])
+        for backend, t in sorted(sdpa_backend_hits.items(), key=lambda kv: -kv[1]):
+            marker = ' ← DOMINANT' if backend == winner[0] else ''
+            print(f'    {backend:20s}: {format_us(t):>12s}{marker}')
+        if winner[0] != 'flash':
+            print(f'  ⚠️  Not using FlashAttention. Expected FA2 on A100 for hidden 128/heads 16 → suboptimal.')
+            print(f'     Fix: 训练里用 `torch.nn.attention.sdpa_kernel(SDPBackend.FLASH_ATTENTION)` context')
+    else:
+        print(f'    (未识别到 attention kernel——可能 kernel 名 encoded 或未走 SDPA)')
+
+    # top ops 明细
+    print(f'\n  Top {top_n} CUDA ops:')
+    print(f'  {"rank":>4s} {"time":>12s} {"% total":>8s} {"# calls":>8s}  {"name":s}')
+    sorted_ops = sorted(stats.items(), key=lambda kv: -kv[1]['time_us'])
+    for i, (name, s) in enumerate(sorted_ops[:top_n], 1):
+        pct = s['time_us'] / max(total_cuda_us, 1) * 100
+        # 截断长 kernel 名
+        display_name = name if len(name) < 80 else name[:77] + '...'
+        print(f'  {i:>4d} {format_us(s["time_us"]):>12s} {pct:>7.1f}% {s["count"]:>8d}  {display_name}')
+
+    # 分类建议
+    print(f'\n{"=" * 80}')
+    print(f'  优化提示（按 bucket 比例）:')
+    print(f'{"=" * 80}')
+    total = total_cuda_us
+    reduce_pct = bucket_time.get('reduce (all-reduce/gather)', 0) / max(total, 1) * 100
+    if reduce_pct > 15:
+        print(f'  ⚠️  reduce/collective 占 {reduce_pct:.1f}%（>15% 算高）→ 考虑：')
+        print(f'      - `no_sync()` 包 gradient accumulation 减少 3× 冗余 all-reduce')
+        print(f'      - FSDP: 用 ZeRO-2/3 拆 optimizer state 减少每次 sync 数据量')
+    elem_pct = bucket_time.get('elementwise (norm/act)', 0) / max(total, 1) * 100
+    if elem_pct > 15:
+        print(f'  ⚠️  elementwise 占 {elem_pct:.1f}%（>15% 算高）→ memory bandwidth bound')
+        print(f'      - Fused RMSNorm + SiLU + Add kernel（apex 或 triton）省 50%+')
+    mem_pct = bucket_time.get('memory (copy/permute/view)', 0) / max(total, 1) * 100
+    if mem_pct > 8:
+        print(f'  ⚠️  memory movement 占 {mem_pct:.1f}%（>8% 算高）→ 布局问题')
+        print(f'      - 检查 permute/transpose/contiguous 频次，可能重复布局转换')
+    attn_pct = bucket_time.get('attention (SDPA/FA)', 0) / max(total, 1) * 100
+    matmul_pct = bucket_time.get('matmul (GEMM)', 0) / max(total, 1) * 100
+    print(f'\n  Matmul + Attention (核心 compute): {matmul_pct + attn_pct:.1f}%')
+    print(f'  如果这个数 <50%，多数时间在做非核心 op → 有巨大优化空间')
+    print(f'  Llama 3 训练里这个数一般 60-75%')
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--trace_file', type=str, default=None)
+    parser.add_argument('--trace_dir', type=str, default='../profile_out')
+    parser.add_argument('--top_n', type=int, default=25)
+    args = parser.parse_args()
+
+    if args.trace_file:
+        files = [args.trace_file]
+    else:
+        # torch.profiler 每 rank 一个文件；分析 rank 0 即可
+        files = sorted(glob.glob(os.path.join(args.trace_dir, '*.pt.trace.json*')))
+        if not files:
+            print(f'no trace file found in {args.trace_dir}')
+            return
+        # 只取第一个（一般 rank 0）
+        files = files[:1]
+
+    for f in files:
+        print(f'\n[loading] {f} ({os.path.getsize(f)/1e6:.1f} MB)')
+        events = load_trace(f)
+        print(f'[loaded] {len(events):,} events')
+        stats, total, sdpa_hits = analyze(events)
+        report(stats, total, sdpa_hits, top_n=args.top_n)
+
+
+if __name__ == "__main__":
+    main()

scripts/compare_parallel.sh

@@ -0,0 +1,117 @@
+#!/bin/bash
+# 分布式并行策略对比 smoke tests
+#
+# 目的：跑 5-6 种并行策略各 100 步，收集 MFU / tokens/s / 每卡显存 → Design Doc #2 素材表
+#
+# ⚠️ 前置：**必须先 kill 当前 1B 训练**（4 卡全被占）：
+#   pkill -9 -f 'train_pretrain|torchrun|multiproc'
+#   然后确认 nvidia-smi 显存 <100 MB/卡
+#
+# 用法：
+#   scripts/compare_parallel.sh              # 跑所有对比
+#   scripts/compare_parallel.sh ddp           # 只跑 DDP
+#   scripts/compare_parallel.sh fsdp_full     # 只跑 FSDP full shard
+#   scripts/compare_parallel.sh tp            # 只跑 TP=2×DP=2
+#
+# 输出：/tmp/parallel_compare/*.log；最后自动 summarize 到 stdout
+set -e
+
+REPO=$(cd "$(dirname "$0")/.." && pwd)
+OUT_DIR=/tmp/parallel_compare
+mkdir -p "$OUT_DIR"
+
+# 公共配置（保持对比公平：同 bs / seq / accum / model）
+COMMON_ARGS=(
+    --hidden_size 2048
+    --num_hidden_layers 20
+    --max_seq_len 1024
+    --batch_size 16
+    --accumulation_steps 1
+    --learning_rate 3e-4
+    --warmup_steps 50
+    --dtype bfloat16
+    --num_workers 4
+    --use_compile 1
+    --data_path "$REPO/dataset/skypile_test.jsonl"
+    --eval_interval 0
+    --save_interval 100000
+    --log_interval 10
+    --save_weight parallel_smoke
+    --save_dir "$REPO/out"
+)
+
+run_config() {
+    local name="$1"
+    shift
+    local strategy_args=("$@")
+    local log="$OUT_DIR/$name.log"
+    echo ""
+    echo "=========================================="
+    echo "  Running: $name"
+    echo "  extra: ${strategy_args[*]}"
+    echo "=========================================="
+    cd "$REPO/trainer"
+    # timeout 3 min 强制停：100 步跑不完就当异常，不用等
+    timeout 180 torchrun --nproc_per_node=4 --standalone train_pretrain.py \
+        "${COMMON_ARGS[@]}" \
+        "${strategy_args[@]}" > "$log" 2>&1 || true
+    echo "  最后 5 条 MFU log："
+    grep -E 'MFU' "$log" 2>/dev/null | grep -v 'params_flops' | tail -5 || echo "  (no MFU logs — crashed?)"
+    # 收显存 (kill 前抓)
+    nvidia-smi --query-gpu=memory.used --format=csv,noheader | head -1 > "$OUT_DIR/${name}_gpu.csv" || true
+    ps aux | grep -E 'train_pretrain|torchrun|multiproc' | grep -v grep | awk '{print $2}' | xargs -r kill -9 2>/dev/null || true
+    sleep 5
+}
+
+CONFIG="${1:-all}"
+
+case "$CONFIG" in
+    all)
+        run_config "ddp"          --parallel_strategy ddp
+        run_config "fsdp_no"      --parallel_strategy fsdp_no
+        run_config "fsdp_grad_op" --parallel_strategy fsdp_grad_op
+        run_config "fsdp_full"    --parallel_strategy fsdp_full
+        run_config "tp2_dp2"      --parallel_strategy tp --tp_size 2
+        run_config "tp4_dp1"      --parallel_strategy tp --tp_size 4
+        ;;
+    ddp|fsdp_no|fsdp_grad_op|fsdp_full)
+        run_config "$CONFIG" --parallel_strategy "$CONFIG"
+        ;;
+    tp|tp2)
+        run_config "tp2_dp2" --parallel_strategy tp --tp_size 2
+        ;;
+    tp4)
+        run_config "tp4_dp1" --parallel_strategy tp --tp_size 4
+        ;;
+    *)
+        echo "Unknown config: $CONFIG"
+        echo "Valid: all | ddp | fsdp_no | fsdp_grad_op | fsdp_full | tp | tp4"
+        exit 1
+        ;;
+esac
+
+echo ""
+echo "=========================================="
+echo "  Summary (稳态 = 各 log 最后一条 MFU)"
+echo "=========================================="
+printf "%-16s %-10s %-10s %-12s\n" "strategy" "it/s" "MFU%" "mem/card"
+printf "%-16s %-10s %-10s %-12s\n" "--------" "----" "----" "--------"
+
+for log in "$OUT_DIR"/*.log; do
+    [ -f "$log" ] || continue
+    name=$(basename "$log" .log)
+    line=$(grep -E 'MFU' "$log" 2>/dev/null | grep -v 'params_flops' | tail -1)
+    mem_file="$OUT_DIR/${name}_gpu.csv"
+    mem=$(cat "$mem_file" 2>/dev/null | tr -d ' ' || echo "?")
+    if [ -z "$line" ]; then
+        printf "%-16s %-10s %-10s %-12s\n" "$name" "-" "-" "crashed"
+    else
+        its=$(echo "$line" | grep -oE 'it/s: [0-9.]+' | awk '{print $2}')
+        mfu=$(echo "$line" | grep -oE 'MFU: [0-9.]+' | awk '{print $2}')
+        printf "%-16s %-10s %-10s %-12s\n" "$name" "$its" "$mfu%" "$mem"
+    fi
+done
+
+echo ""
+echo "详细 log: $OUT_DIR/*.log"
+echo "profile 单个: scripts/compare_parallel.sh $CONFIG 时加 --profile_steps 20 到 COMMON_ARGS"

scripts/launch_1b_pretrain.sh

@@ -0,0 +1,65 @@
+#!/bin/bash
+# 1B pretrain launcher — 4 卡 DDP，SkyPile stage 1（真 raw web pretrain）
+#
+# Config 基于 plan L84 + dry-run 学到的教训：
+#   - hidden=2048, layers=20, heads=16 (kv=4 GQA) → ~1.02B params
+#   - seq=1024 (SkyPile 长度足够)
+#   - global batch = 16 per GPU × 4 GPU × accum 4 = 256 → ~256*1024=262k tokens/step
+#   - lr 3e-4 (1B from-scratch 标准) + warmup 1000 (前 500 步发散风险 >50%)
+#   - eval 每 500 step，用 minimind pretrain 做 proxy holdout（1B 未见过）
+#   - save 每 2000 step（2GB × 264 次约 500GB 磁盘，可承受；1B 训 ~150k step 总量）
+#   - grad_checkpointing 0（1B seq=1024 在 A100 80GB 用不上，省 30% 时间）
+#
+# Stage 2 (minimind curated Q&A mid-training) 是另开 run，从 pretrain_1b ckpt 起点继续
+#
+# 用法:
+#   scripts/launch_1b_pretrain.sh                # 默认 4 卡
+#   scripts/launch_1b_pretrain.sh --nproc 2       # 覆盖为 2 卡
+#   scripts/launch_1b_pretrain.sh -- --learning_rate 5e-4  # 传给 python
+set -e
+
+# 默认参数
+NPROC=4
+# 分离 launcher 参数和 python 参数
+PY_ARGS=()
+while [[ $# -gt 0 ]]; do
+    case "$1" in
+        --nproc) NPROC="$2"; shift 2 ;;
+        --) shift; PY_ARGS+=("$@"); break ;;
+        *) PY_ARGS+=("$1"); shift ;;
+    esac
+done
+
+REPO_ROOT=$(cd "$(dirname "$0")/.." && pwd)
+cd "$REPO_ROOT/trainer"
+
+# 数据：SkyPile 是 stage 1 主战场
+DATA_PATH="../dataset/skypile_test.jsonl"
+# eval holdout：用 minimind pretrain 作 proxy（1B 训练时未见过）
+EVAL_DATA="../dataset/pretrain_t2t.jsonl"
+
+echo "[launch_1b_pretrain] NPROC=$NPROC data=$DATA_PATH"
+echo "[launch_1b_pretrain] extra args: ${PY_ARGS[*]}"
+
+torchrun --nproc_per_node=$NPROC --standalone train_pretrain.py \
+    --hidden_size 2048 \
+    --num_hidden_layers 20 \
+    --max_seq_len 1024 \
+    --batch_size 24 \
+    --accumulation_steps 4 \
+    --learning_rate 3e-4 \
+    --warmup_steps 1000 \
+    --grad_checkpointing 0 \
+    --dtype bfloat16 \
+    --num_workers 4 \
+    --use_compile 1 \
+    --data_path "$DATA_PATH" \
+    --eval_interval 500 \
+    --eval_data_path "$EVAL_DATA" \
+    --eval_max_batches 20 \
+    --save_interval 500 \
+    --log_interval 20 \
+    --save_weight pretrain_1b \
+    --save_dir ../out \
+    --wandb_project MiniMind-1B-Pretrain \
+    "${PY_ARGS[@]}"

trainer/train_pretrain.py

@@ -51,7 +51,7 @@ def eval_holdout(model, eval_loader, max_batches, autocast_ctx, device):
     return avg_loss, math.exp(avg_loss)
 
 
-def train_epoch(epoch, loader, iters, start_step=0, wandb=None, eval_loader=None):
+def train_epoch(epoch, loader, iters, start_step=0, wandb=None, eval_loader=None, profiler=None):
     start_time = time.time()
     last_log_time, last_log_step = start_time, start_step
     last_step = start_step
@@ -70,12 +70,22 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None, eval_loader=None
         for param_group in optimizer.param_groups:
             param_group['lr'] = lr
 
-        with autocast_ctx:
-            res = model(input_ids, labels=labels)
-            loss = res.loss + res.aux_loss
-            loss = loss / args.accumulation_steps
+        # DDP / FSDP gradient accumulation optimization：
+        # 默认每次 .backward() 都做全参数 collective（DDP: all-reduce grad；FSDP: reduce-scatter grad）。
+        # 但 accum=N 时前 N-1 次 collective 是浪费（grad 还要继续累积）。DDP.no_sync() 和 FSDP.no_sync()
+        # 都提供 context 让 collective 只在最后一个 micro-step 才触发。
+        # profile 显示（DDP）all-reduce 占 73% 时间 → 修完预期总时间省 ~55%，MFU 17%→~35%。
+        is_last_accum_step = (step % args.accumulation_steps == 0)
+        has_no_sync = callable(getattr(model, 'no_sync', None))
+        need_sync = is_last_accum_step or not has_no_sync
+        sync_ctx = nullcontext() if need_sync else model.no_sync()
 
-        scaler.scale(loss).backward()
+        with sync_ctx:
+            with autocast_ctx:
+                res = model(input_ids, labels=labels)
+                loss = res.loss + res.aux_loss
+                loss = loss / args.accumulation_steps
+            scaler.scale(loss).backward()
 
         if step % args.accumulation_steps == 0:
             scaler.unscale_(optimizer)
@@ -86,16 +96,24 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None, eval_loader=None
 
             optimizer.zero_grad(set_to_none=True)
 
+        # profiler step 打点（每 step 一次；schedule 决定何时真的记录）
+        if profiler is not None:
+            profiler.step()
+
         if step % args.log_interval == 0 or step == iters:
             now = time.time()
             spend_time = now - start_time
             window_dt = max(now - last_log_time, 1e-6)
             window_steps = step - last_log_step
             window_its = window_steps / window_dt
-            # MFU = 实际 FLOPs/s / GPU 峰值；按 window 内的 token 吞吐算（避开 warmup / ckpt save 污染）
-            # tokens 用 per-rank batch × seq_len × world_size × window_steps（DDP 每卡同时处理 batch）
+            # MFU = per-card 实际 FLOPs/s / per-card 峰值。分子分母都要一致口径：
+            # - 分子: 4 卡总 tokens × flops/token = 总 flops；÷ world_size 得 per-card flops
+            # - 分母: gpu_peak_tflops (单卡峰值，312 TFLOPS on A100)
+            # 早期版本 bug：分子用全部 tokens 但分母只除 1 卡峰值 → 报数是真实 MFU 的 world_size 倍。
             tokens_window = window_steps * args.batch_size * args.max_seq_len * world_size
-            mfu = (flops_per_token * tokens_window / window_dt) / (gpu_peak_tflops * 1e12) * 100
+            total_flops_per_s = flops_per_token * tokens_window / window_dt
+            per_card_flops_per_s = total_flops_per_s / world_size
+            mfu = per_card_flops_per_s / (gpu_peak_tflops * 1e12) * 100
             last_log_time, last_log_step = now, step
             current_loss = loss.item() * args.accumulation_steps
             current_aux_loss = res.aux_loss.item() if res.aux_loss is not None else 0.0
@@ -107,28 +125,77 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None, eval_loader=None
 
         # ---- eval loop（可选） ----
         if args.eval_interval > 0 and step % args.eval_interval == 0 and eval_loader is not None:
+            # ⚠️ 全 rank 都跑 eval：单 rank eval (即使 unwrap raw model) 也会导致 DDP guard 失效 /
+            # torch.compile recompile → rank 0 长 stall → 下一步 all-reduce seq 不同步 → NCCL 10min timeout。
+            # 解决：所有 rank 用同一份 eval_loader (顺序一致) 跑同样的 batch，然后 all-reduce mean loss。
+            # 每 rank 都工作，无 rank 掉队问题。max_batches=20 * 5 rank = 20 batch × ~0.3s ≈ 6s 开销。
+            #
+            # DDP: unwrap 拿裸 model（避免走 DDP forward 的 hook 触发 collective）
+            # FSDP: 不能 unwrap！FSDP forward 负责 all-gather 分片的 weight，unwrap 后 weight 不完整。
+            #       FSDP.forward 天然要求所有 rank 参与——正好符合我们的全-rank eval 设计
+            try:
+                from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+                is_fsdp = isinstance(model, FSDP)
+            except ImportError:
+                is_fsdp = False
+            if is_fsdp:
+                eval_model = model
+            else:
+                eval_model = model.module if isinstance(model, DistributedDataParallel) else model
+                eval_model = getattr(eval_model, '_orig_mod', eval_model)
+            val_loss_local, _ = eval_holdout(eval_model, eval_loader, args.eval_max_batches, autocast_ctx, args.device)
+            if dist.is_initialized():
+                loss_t = torch.tensor(val_loss_local, device=args.device)
+                dist.all_reduce(loss_t, op=dist.ReduceOp.AVG)
+                val_loss = loss_t.item()
+            else:
+                val_loss = val_loss_local
+            val_ppl = math.exp(val_loss)
             if is_main_process():
-                val_loss, val_ppl = eval_holdout(model, eval_loader, args.eval_max_batches, autocast_ctx, args.device)
                 Logger(f'[eval] step={step} val_loss={val_loss:.4f} val_ppl={val_ppl:.2f}')
                 if wandb: wandb.log({"val_loss": val_loss, "val_ppl": val_ppl, "step": step})
-            if dist.is_initialized(): dist.barrier()
 
         if step % args.save_interval == 0 or step == iters:
-            if is_main_process():
-                model.eval()
-                moe_suffix = '_moe' if lm_config.use_moe else ''
-                ckp = f'{args.save_dir}/{args.save_weight}_{lm_config.hidden_size}{moe_suffix}.pth'
-                raw_model = model.module if isinstance(model, DistributedDataParallel) else model
-                raw_model = getattr(raw_model, '_orig_mod', raw_model)
-                state_dict = raw_model.state_dict()
-                # atomic: tmp + rename 避免读者读到半写文件（当前的 torch.save 直接写不 atomic，
-                # 极端情况 kill -9 会留 corrupt ckpt；lm_checkpoint 内部已 atomic，此处也补上）
-                ckp_tmp = ckp + '.tmp'
-                torch.save({k: v.half().cpu() for k, v in state_dict.items()}, ckp_tmp)
-                os.replace(ckp_tmp, ckp)
-                lm_checkpoint(lm_config, weight=args.save_weight, model=model, optimizer=optimizer, scaler=scaler, epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints')
-                model.train()
-                del state_dict
+            # FSDP: 拿 full state_dict 是 collective，所有 rank 必须进入 FSDP.state_dict_type 上下文；
+            # 只 rank 0 通过 rank0_only=True 拿到实际数据，其他 rank 拿空 dict。
+            # DDP: 单 rank 就能拿完整 state_dict。
+            try:
+                from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+                is_fsdp = isinstance(model, FSDP)
+            except ImportError:
+                is_fsdp = False
+
+            if is_fsdp:
+                from torch.distributed.fsdp import StateDictType, FullStateDictConfig
+                cfg = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
+                with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT, cfg):
+                    state_dict = model.state_dict()   # 所有 rank 调用（collective），只 rank 0 拿到完整
+                # rank 0 才写盘
+                if is_main_process() and state_dict:
+                    moe_suffix = '_moe' if lm_config.use_moe else ''
+                    ckp = f'{args.save_dir}/{args.save_weight}_{lm_config.hidden_size}{moe_suffix}.pth'
+                    ckp_tmp = ckp + '.tmp'
+                    torch.save({k: v.half() for k, v in state_dict.items()}, ckp_tmp)
+                    os.replace(ckp_tmp, ckp)
+                    # 注意：FSDP 下 optimizer state 也是分片的，简化起见此处不存 optimizer resume
+                    # （resume 会从头 warmup 起，接受此代价换代码简单）
+                    del state_dict
+            else:
+                if is_main_process():
+                    model.eval()
+                    moe_suffix = '_moe' if lm_config.use_moe else ''
+                    ckp = f'{args.save_dir}/{args.save_weight}_{lm_config.hidden_size}{moe_suffix}.pth'
+                    raw_model = model.module if isinstance(model, DistributedDataParallel) else model
+                    raw_model = getattr(raw_model, '_orig_mod', raw_model)
+                    state_dict = raw_model.state_dict()
+                    # atomic: tmp + rename 避免读者读到半写文件（当前的 torch.save 直接写不 atomic，
+                    # 极端情况 kill -9 会留 corrupt ckpt；lm_checkpoint 内部已 atomic，此处也补上）
+                    ckp_tmp = ckp + '.tmp'
+                    torch.save({k: v.half().cpu() for k, v in state_dict.items()}, ckp_tmp)
+                    os.replace(ckp_tmp, ckp)
+                    lm_checkpoint(lm_config, weight=args.save_weight, model=model, optimizer=optimizer, scaler=scaler, epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints')
+                    model.train()
+                    del state_dict
             # 所有 rank 在此汇合，让非 rank-0 等 rank-0 写完再进下一 step
             # （防止 NCCL 隐式同步下的 timing race，也确保 ckpt 文件对 resume 一致）
             if dist.is_initialized():
@@ -181,6 +248,20 @@ if __name__ == "__main__":
                         help="holdout jsonl 路径；不设则关闭 eval")
     parser.add_argument("--eval_max_batches", type=int, default=20,
                         help="每次 eval 最多跑几个 batch（20*batch_size 通常够 ~1s 内出数）")
+    parser.add_argument("--profile_steps", type=int, default=0,
+                        help="torch.profiler 抓 N 步（0=关）；建议 20。会在 wait=10 warmup=5 active=N 后自动退出。"
+                             "trace 存到 --profile_out 目录，chrome://tracing 打开或 tensorboard --logdir profile_out")
+    parser.add_argument("--profile_out", type=str, default="../profile_out",
+                        help="profile trace 输出目录")
+    parser.add_argument("--parallel_strategy", type=str, default="ddp",
+                        choices=["ddp", "fsdp_full", "fsdp_grad_op", "fsdp_no", "tp"],
+                        help="分布式策略：ddp=DataParallel（默认）· fsdp_full=ZeRO-3（切 weight+grad+optim）· "
+                             "fsdp_grad_op=ZeRO-2（切 grad+optim，weight 复制）· fsdp_no=NO_SHARD（等价 DDP，用 FSDP wrap 走一遍）· "
+                             "tp=Tensor Parallel（用 PyTorch DTensor 切每层 Linear）"
+                             "。对比数据见 Design Doc #2。1B 用 DDP 足够；FSDP/TP 主要是学习/对比用途")
+    parser.add_argument("--tp_size", type=int, default=1,
+                        help="Tensor Parallel size（仅当 parallel_strategy=tp 时生效）。world_size 必须能整除 tp_size；"
+                             "dp_size = world_size / tp_size。例如 4 卡 tp_size=4 → 纯 TP=4；tp_size=2 → TP=2 × DP=2")
     args = parser.parse_args()
 
     # ========== 1. 初始化环境和随机种子 ==========
@@ -244,19 +325,121 @@ if __name__ == "__main__":
         model = torch.compile(model)
         Logger('torch.compile enabled (mode=default)')
     if dist.is_initialized():
-        model = DistributedDataParallel(model, device_ids=[local_rank])
+        if args.parallel_strategy == "ddp":
+            model = DistributedDataParallel(model, device_ids=[local_rank])
+            Logger('parallel strategy: DDP (每 rank 全参数复制)')
+        elif args.parallel_strategy == "tp":
+            # Tensor Parallel with PyTorch DTensor：把 Attention 和 MLP 的 Linear 切到 tp_size 卡上
+            # 关键：需要先修 Attention 里的 n_local_heads / n_local_kv_heads（每卡本地头数减少），
+            # 再用 parallelize_module 把 weight 切开。否则 xq.view(bsz, seq, n_local_heads, head_dim)
+            # 会因为 q_proj 输出宽度已经被切成 hidden/tp_size 而 shape mismatch。
+            from torch.distributed.tensor.parallel import (
+                parallelize_module, ColwiseParallel, RowwiseParallel
+            )
+            from torch.distributed.device_mesh import init_device_mesh
+            world_size_actual = dist.get_world_size()
+            tp_size = args.tp_size
+            dp_size = world_size_actual // tp_size
+            if tp_size * dp_size != world_size_actual:
+                raise ValueError(f"world_size {world_size_actual} 不能被 tp_size {tp_size} 整除")
+            # GQA 校验：kv_heads 必须能被 tp_size 整除，否则某 rank 分到 0.x 个 kv head 会崩
+            if lm_config.num_key_value_heads % tp_size != 0:
+                raise ValueError(f"num_key_value_heads={lm_config.num_key_value_heads} 不能被 tp_size={tp_size} 整除")
+            mesh = init_device_mesh("cuda", (dp_size, tp_size), mesh_dim_names=("dp", "tp"))
+            tp_mesh = mesh["tp"]
+            # 修正 Attention 的本地头数（否则 forward 里 view 会 shape 错）
+            raw_model_for_tp = model  # 此时还未 compile / wrap
+            for layer in raw_model_for_tp.model.layers:
+                layer.self_attn.n_local_heads = lm_config.num_attention_heads // tp_size
+                layer.self_attn.n_local_kv_heads = lm_config.num_key_value_heads // tp_size
+                # n_rep = n_local_heads // n_local_kv_heads，两边都除 tp_size 后比例不变，无需改
+            # 应用 TP 切分：q/k/v/gate/up 列切（out 维），o/down 行切（in 维）
+            tp_plan = {
+                "self_attn.q_proj": ColwiseParallel(),
+                "self_attn.k_proj": ColwiseParallel(),
+                "self_attn.v_proj": ColwiseParallel(),
+                "self_attn.o_proj": RowwiseParallel(),
+                "mlp.gate_proj":    ColwiseParallel(),
+                "mlp.up_proj":      ColwiseParallel(),
+                "mlp.down_proj":    RowwiseParallel(),
+            }
+            for layer in raw_model_for_tp.model.layers:
+                parallelize_module(layer, tp_mesh, tp_plan)
+            # DP wrap（若 dp_size > 1）：DDP over TP groups
+            if dp_size > 1:
+                dp_mesh = mesh["dp"]
+                # 用 dp_mesh 的 process group 建 DDP
+                model = DistributedDataParallel(
+                    model, device_ids=[local_rank], process_group=dp_mesh.get_group(),
+                )
+            Logger(f'parallel strategy: TP={tp_size} × DP={dp_size} (world={world_size_actual})')
+            Logger('⚠️  TP 是 experimental：ckpt 保存尚未处理 DTensor（--save_interval 0 关闭 save；如需 save 请等 Month 2）')
+        else:
+            # FSDP 三档：分别对应 ZeRO-3 / ZeRO-2 / DDP-等价
+            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 functools import partial
+            from model.model_minimind import MiniMindBlock
+            strategy_map = {
+                "fsdp_full":    ShardingStrategy.FULL_SHARD,     # ZeRO-3: weight + grad + optim 都切
+                "fsdp_grad_op": ShardingStrategy.SHARD_GRAD_OP,  # ZeRO-2: grad + optim 切，weight 复制
+                "fsdp_no":      ShardingStrategy.NO_SHARD,       # 等价 DDP，走 FSDP 代码路径便于对比
+            }
+            # bf16 混合精度：param bf16 forward，reduce fp32 保证 grad 数值稳定
+            mixed_precision = MixedPrecision(
+                param_dtype=torch.bfloat16,
+                reduce_dtype=torch.float32,
+                buffer_dtype=torch.float32,
+            )
+            # auto_wrap 每个 MiniMindBlock 单独作为 FSDP unit（不是整个 model 一个 unit）
+            # 这样 all_gather 是每层触发，粒度更细，memory peak 更低
+            auto_wrap_policy = partial(
+                transformer_auto_wrap_policy,
+                transformer_layer_cls={MiniMindBlock},
+            )
+            model = FSDP(
+                model,
+                sharding_strategy=strategy_map[args.parallel_strategy],
+                mixed_precision=mixed_precision,
+                auto_wrap_policy=auto_wrap_policy,
+                device_id=local_rank,
+                forward_prefetch=True,                              # 提前 all-gather 下一层 weight
+                backward_prefetch=BackwardPrefetch.BACKWARD_PRE,    # 提前 all-gather backward 需要的 weight
+                use_orig_params=True,                               # 兼容 torch.compile 和 param_groups
+            )
+            Logger(f'parallel strategy: FSDP {args.parallel_strategy} '
+                   f'({strategy_map[args.parallel_strategy].name})')
 
     # ========== 7b. eval loader（可选） ==========
+    # 全 rank 都建 eval loader（不能只 rank 0，否则 all-reduce 时 rank 1/2/3 无 eval loss 参与聚合）
     eval_loader = None
-    if args.eval_interval > 0 and args.eval_data_path and is_main_process():
-        # rank-0 only：只 main rank 需要 eval loader（其他 rank barrier 等待）
+    if args.eval_interval > 0 and args.eval_data_path:
         eval_ds = PretrainDataset(args.eval_data_path, tokenizer, max_length=args.max_seq_len)
+        # 所有 rank 用同一份数据 + shuffle=False → 保证跑相同 batch → all-reduce 有意义
         eval_loader = DataLoader(eval_ds, batch_size=args.batch_size,
                                  shuffle=False, num_workers=2, pin_memory=True)
         Logger(f'[eval] loader ready: {len(eval_ds)} samples from {args.eval_data_path}, '
                f'eval_interval={args.eval_interval}, max_batches={args.eval_max_batches}')
 
+    # ========== 7c. profiler（可选） ==========
+    profiler_ctx = None
+    if args.profile_steps > 0 and is_main_process():
+        # schedule: 前 10 步 warmup（skip compile 编译期），5 步观察，profile_steps 步真记录
+        from torch.profiler import profile, schedule, ProfilerActivity, tensorboard_trace_handler
+        os.makedirs(args.profile_out, exist_ok=True)
+        profiler_ctx = profile(
+            activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
+            schedule=schedule(wait=10, warmup=5, active=args.profile_steps, repeat=1),
+            on_trace_ready=tensorboard_trace_handler(args.profile_out),
+            record_shapes=True, with_stack=False, with_flops=False,
+        )
+        Logger(f'[profile] enabled: warmup 15 step → capture {args.profile_steps} step → trace to {args.profile_out}')
+        Logger('[profile] 训练完 profile 步后不自动退出——手动 Ctrl+C 停即可。trace 已 flush 到磁盘。')
+
     # ========== 8. 开始训练 ==========
+    if profiler_ctx is not None:
+        profiler_ctx.__enter__()
     for epoch in range(start_epoch, args.epochs):
         train_sampler and train_sampler.set_epoch(epoch)
         setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist()
@@ -266,9 +449,11 @@ if __name__ == "__main__":
                             pin_memory=True, persistent_workers=args.num_workers > 0)
         if skip > 0:
             Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始')
-            train_epoch(epoch, loader, len(loader) + skip, start_step, wandb, eval_loader)
+            train_epoch(epoch, loader, len(loader) + skip, start_step, wandb, eval_loader, profiler_ctx)
         else:
-            train_epoch(epoch, loader, len(loader), 0, wandb, eval_loader)
+            train_epoch(epoch, loader, len(loader), 0, wandb, eval_loader, profiler_ctx)
+    if profiler_ctx is not None:
+        profiler_ctx.__exit__(None, None, None)
 
     # ========== 9. 清理分布进程 ==========
     if dist.is_initialized():