FSDP TP
Longfang Zhao
parent
da93141a0d
commit
b40eb54434
|
|
@ -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"
|
||||||
|
|
@ -70,13 +70,14 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None, eval_loader=None
|
||||||
for param_group in optimizer.param_groups:
|
for param_group in optimizer.param_groups:
|
||||||
param_group['lr'] = lr
|
param_group['lr'] = lr
|
||||||
|
|
||||||
# DDP gradient accumulation optimization：
|
# DDP / FSDP gradient accumulation optimization：
|
||||||
# 默认 DDP 每次 .backward() 都做全参数 all-reduce（4GB grad × 85 bucket ≈ 30ms）。
|
# 默认每次 .backward() 都做全参数 collective（DDP: all-reduce grad；FSDP: reduce-scatter grad）。
|
||||||
# 但 accum=N 时前 N-1 次 all-reduce 是浪费（grad 还要继续累积）。用 model.no_sync()
|
# 但 accum=N 时前 N-1 次 collective 是浪费（grad 还要继续累积）。DDP.no_sync() 和 FSDP.no_sync()
|
||||||
# 让 DDP 在非最后一个 micro-step 里 skip all-reduce，只在最后一个 micro-step 才同步。
|
# 都提供 context 让 collective 只在最后一个 micro-step 才触发。
|
||||||
# profile 显示 all-reduce 占 73% 时间 → 修完预期总时间省 ~55%，MFU 17%→~35%。
|
# profile 显示（DDP）all-reduce 占 73% 时间 → 修完预期总时间省 ~55%，MFU 17%→~35%。
|
||||||
is_last_accum_step = (step % args.accumulation_steps == 0)
|
is_last_accum_step = (step % args.accumulation_steps == 0)
|
||||||
need_sync = is_last_accum_step or not isinstance(model, DistributedDataParallel)
|
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()
|
sync_ctx = nullcontext() if need_sync else model.no_sync()
|
||||||
|
|
||||||
with sync_ctx:
|
with sync_ctx:
|
||||||
|
|
@ -128,9 +129,21 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None, eval_loader=None
|
||||||
# torch.compile recompile → rank 0 长 stall → 下一步 all-reduce seq 不同步 → NCCL 10min timeout。
|
# torch.compile recompile → rank 0 长 stall → 下一步 all-reduce seq 不同步 → NCCL 10min timeout。
|
||||||
# 解决：所有 rank 用同一份 eval_loader (顺序一致) 跑同样的 batch，然后 all-reduce mean loss。
|
# 解决：所有 rank 用同一份 eval_loader (顺序一致) 跑同样的 batch，然后 all-reduce mean loss。
|
||||||
# 每 rank 都工作，无 rank 掉队问题。max_batches=20 * 5 rank = 20 batch × ~0.3s ≈ 6s 开销。
|
# 每 rank 都工作，无 rank 掉队问题。max_batches=20 * 5 rank = 20 batch × ~0.3s ≈ 6s 开销。
|
||||||
raw_model = model.module if isinstance(model, DistributedDataParallel) else model
|
#
|
||||||
raw_model = getattr(raw_model, '_orig_mod', raw_model)
|
# DDP: unwrap 拿裸 model（避免走 DDP forward 的 hook 触发 collective）
|
||||||
val_loss_local, _ = eval_holdout(raw_model, eval_loader, args.eval_max_batches, autocast_ctx, args.device)
|
# 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():
|
if dist.is_initialized():
|
||||||
loss_t = torch.tensor(val_loss_local, device=args.device)
|
loss_t = torch.tensor(val_loss_local, device=args.device)
|
||||||
dist.all_reduce(loss_t, op=dist.ReduceOp.AVG)
|
dist.all_reduce(loss_t, op=dist.ReduceOp.AVG)
|
||||||
|
|
@ -143,6 +156,31 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None, eval_loader=None
|
||||||
if wandb: wandb.log({"val_loss": val_loss, "val_ppl": val_ppl, "step": step})
|
if wandb: wandb.log({"val_loss": val_loss, "val_ppl": val_ppl, "step": step})
|
||||||
|
|
||||||
if step % args.save_interval == 0 or step == iters:
|
if step % args.save_interval == 0 or step == iters:
|
||||||
|
# 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():
|
if is_main_process():
|
||||||
model.eval()
|
model.eval()
|
||||||
moe_suffix = '_moe' if lm_config.use_moe else ''
|
moe_suffix = '_moe' if lm_config.use_moe else ''
|
||||||
|
|
@ -215,6 +253,15 @@ if __name__ == "__main__":
|
||||||
"trace 存到 --profile_out 目录，chrome://tracing 打开或 tensorboard --logdir profile_out")
|
"trace 存到 --profile_out 目录，chrome://tracing 打开或 tensorboard --logdir profile_out")
|
||||||
parser.add_argument("--profile_out", type=str, default="../profile_out",
|
parser.add_argument("--profile_out", type=str, default="../profile_out",
|
||||||
help="profile trace 输出目录")
|
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()
|
args = parser.parse_args()
|
||||||
|
|
||||||
# ========== 1. 初始化环境和随机种子 ==========
|
# ========== 1. 初始化环境和随机种子 ==========
|
||||||
|
|
@ -278,7 +325,91 @@ if __name__ == "__main__":
|
||||||
model = torch.compile(model)
|
model = torch.compile(model)
|
||||||
Logger('torch.compile enabled (mode=default)')
|
Logger('torch.compile enabled (mode=default)')
|
||||||
if dist.is_initialized():
|
if dist.is_initialized():
|
||||||
|
if args.parallel_strategy == "ddp":
|
||||||
model = DistributedDataParallel(model, device_ids=[local_rank])
|
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（可选） ==========
|
# ========== 7b. eval loader（可选） ==========
|
||||||
# 全 rank 都建 eval loader（不能只 rank 0，否则 all-reduce 时 rank 1/2/3 无 eval loss 参与聚合）
|
# 全 rank 都建 eval loader（不能只 rank 0，否则 all-reduce 时 rank 1/2/3 无 eval loss 参与聚合）
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue