BAGEL 训练性能优化报告

BAGEL 训练性能优化报告

项目 ：BAGEL 8×A800 FSDP 预训练

日期 ：2026-08-23

对比版本 ：Profile v4（baseline） vs Profile v5（优化后）

配置 ：profile_t2i.yaml，FULL_SHARD，cpu_offload=True，expected_num_tokens=256，visual_gen=True，visual_und=False

---

1. 执行摘要

本次优化聚焦 CPU-GPU 同步、Host-Device 拷贝、标量通信 等「边角开销」，在保持训练语义不变的前提下，消除若干已知同步点，并改善数据加载 H2D 路径。

维度	结论
端到端吞吐	Tokens/Sec 143 → 144 （约 +0.7%）
单步 wall time	18.33s → 18.33s（基本持平）
MFU	~0.3%（不变）
总体评估	微优化有效且可验证，但被 FSDP NCCL 通信（占单步 ~70--80%） 完全淹没

一句话 ：同步/H2D/bool 索引类问题已修，但要想 10%+ 级别 提速，需要动 FSDP 策略（cpu_offload、SHARD_GRAD_OP 等），而非继续抠 scalar 路径。

---

2. 背景与问题

2.1 初始 profile 发现（v4 trace）

通过对 results_profiler/profiler/rank0/*.pt.trace.json 的分析，主要瓶颈类型为：

1. FSDP 通信 ：nccl::all_gather、record_param_comms（主耗时）
2. 数据拷贝 ：Memcpy HtoD (Pinned)、aten::copy_
3. CPU 同步 ：cudaStreamSynchronize、.item() 触发的 D2H
4. 动态 shape 索引 ：aten::nonzero + bool tensor[mask]（bagel MSE loss）

计算密集型 GEMM 并非首要瓶颈。

2.2 优化目标

• 减少训练 hot path 上的 强制 CPU-GPU 同步
• 重叠 标量 all_reduce 与 forward/backward
• 改善 DataLoader → GPU 的 H2D 效率
• 消除 bool 索引 引发的 nonzero 同步

---

3. 代码改动清单

3.1 train/pretrain_unified_navit.py

改动	说明
_async_all_reduce()	封装 dist.all_reduce(..., async_op=True)
GPU 标量 buffer 复用	scalar_f32、reduce_i64、reduce_i64_b 预分配，fill_() 写入
Token 统计 GPU 累加	token_window / seqlen_square_window 在 GPU 上 add_()，仅 log_every 时 .item()
通信 overlap	token/ce/mse count reduce 提前 launch，wait() 推迟到 backward 之后或 log 时
Loss log	直接对 GPU loss tensor all_reduce，避免先 .item() 再建 tensor
prof.step() 位置	移到 log 块之前，使 ProfilerStep#N 不含 log/sync 噪声
ASYNC_OP_OPTIM（可选）	export ASYNC_OP_OPTIM=1 启用 pin_memory 标量路径；默认关闭 ，推荐用 fill_() 路径

3.2 data/dataset_base.py

改动	说明
pin_memory()	补充 sample_lens 的 pin
cuda() / _to_device()	所有 tensor 传输改为 non_blocking=True
效果	DataLoader pin_memory=True + non_blocking H2D 形成完整异步链路

3.3 modeling/bagel/modeling_utils.py

改动	说明
TimestepEmbedder	freqs 通过 register_buffer + _get_freqs(device) 懒加载缓存，避免每步 H2D

3.4 modeling/bagel/bagel.py

改动	说明
MSE loss 计算	原 ：target[packed_timesteps > 0]（bool 索引 → aten::nonzero → 同步）
	现 ：seq_to_latent 整数索引 + mse_mask 乘法（shape 不变，无 nonzero）

---

4. 优化原理说明

4.1 标量统计：GPU buffer + fill_() vs pin_memory

对于 total_ce_tokens 等标量，推荐方案：

reduce_i64.fill_(data['ce_loss_indexes'].numel())
ce_count_work = _async_all_reduce(reduce_i64)

而非：

torch.tensor(len(...), device='cpu', pin_memory=True).to(device, non_blocking=True)

原因：标量场景下 fill_() 无每步分配、路径更短，且同样支持 async_op overlap。

4.2 bool 索引为何触发同步

target[bool_mask] 输出 shape 依赖 True 的个数 K，Host 分配器需等 GPU 计数 → aten::nonzero + cudaStreamSynchronize。

替代 ：((pred - target) ** 2) * mask，shape 不变，Host 提前知道输出大小。

4.3 prof.step() 与 .item() 的关系

• Hot path（forward/backward 内）：优化后 无每步 .item()
• Log 块（prof.step() 之后）：仍有 .item()，不计入 ProfilerStep#N

---

5. Profile 对比结果

5.1 测试环境

项	v4 baseline	v5 optimized
GPU	8×A800 80GB	8×A800 80GB
其他进程	ComfyUI ~41GB/卡	GPU 基本空闲
Trace	baseline_v4/...923920....json	profiler/rank0/...2286480....json
日志	profile_run_v4.log	profile_run_v5.log

注：环境不完全相同，绝对耗时对比仅供参考；配置参数一致。

5.2 端到端指标（log）

指标	v4	v5	变化
单步 wall time（step 1--9，排除 step 5 profiler 导出）	18.33 s	18.33 s	≈ 0%
Tokens/Sec（同上）	143 tok/s	144 tok/s	+0.7%
Steps/Sec	~0.05--0.06	~0.05--0.06	持平
MFU	~0.3%	~0.3%	持平

5.3 Trace 算子对比（rank0，active 窗口）

算子 / 事件	v4	v5	说明
aten::nonzero	3 次 / 117 ms	0 次 / 0 ms	bagel mask 优化生效
Memcpy HtoD	4.54 s	4.35 s	-4%
Memcpy DtoH	1.68 s	2.01 s	+20%（log .item() 等，在 step 边界外）
cudaStreamSynchronize	0.67 s	4.04 s	波动大，受 profiler 导出影响
nccl	11.9 s	14.8 s	主瓶颈，步间方差大
GPU ProfilerStep#3--5 均值	3401 ms	4193 ms	方差大，不宜单独下结论

5.4 单步时间分布（估算）

FSDP all_gather / NCCL     ████████████████████████████  ~70--80%
Forward + Backward 计算    ██████                        ~15--20%
VAE / 其他 GPU 算子        ███                           ~5--10%
H2D/D2H + sync + log       █                             ~3--5%  ← 本次优化区域

---

6. 分项收益估算

优化项	Trace 实测	折算单步（~18s）
bagel bool → mask	nonzero 消除，~117ms/3 步	~40 ms（~0.2%）
dataset async H2D	HtoD -4%	~60 ms（~0.3%）
pretrain 标量 fill_ + async reduce	热路径无每步 .item()	~20--50 ms（~0.1--0.3%）
TimestepEmbedder freqs 缓存	减少重复 H2D	微量
合计（理论）		~0.1--0.8 s（0.5--4%）
实测		~0.7% tok/s

---

7. 未改动 / 仍存在的瓶颈

1. FSDP FULL_SHARD backward all_gather：每层 unshard 必需，无法通过 H2D 优化消除
2. cpu_offload=True：参数 CPU↔GPU 来回，显著拖慢 all_gather
3. FSDP clip_grad_norm_ ：含 cudaStreamSynchronize，占比小
4. Log 块 （log_every=1）：torch.cuda.synchronize() + 多个 .item()，在 prof.step() 之后
5. ASYNC_OP_OPTIM=1：与默认路径重复，不推荐开启

---

8. 后续优化建议（按收益排序）

优先级	措施	预期收益	风险
P0	独占 GPU 时 cpu_offload=False	高	显存占用上升
P0	试 SHARD_GRAD_OP	高（减 backward all_gather）	显存显著上升
P1	FSDP forward_prefetch=True	中	需改 fsdp_utils.py
P1	log_every=10	中（减 log 同步）	log 粒度变粗
P2	HYBRID_SHARD 调参	中	需验证 checkpoint
P3	继续抠 scalar/H2D	低（<2%）	投入产出比低

---

9. 复现命令

9.1 Profile 训练（优化后默认路径）

cd /root/.cuda/BAGEL
source .venv/bin/activate
export PYTHONPATH="/root/.cuda/BAGEL:$PYTHONPATH"
export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
# 不要设置 ASYNC_OP_OPTIM（默认走 fill_ 路径）

torchrun \
  --nnodes=1 --node_rank=0 --nproc_per_node=8 \
  --master_addr=127.0.0.1 --master_port=8853 \
  train/pretrain_unified_navit.py \
  --dataset_config_file ./data/configs/profile_t2i.yaml \
  --model_path models/BAGEL-7B-MoT \
  --layer_module Qwen2MoTDecoderLayer \
  --vae_path flux/vae/ae.safetensors \
  --visual_gen True --visual_und False \
  --log_every 1 \
  --results_dir ./results_profiler \
  --max_latent_size 64 \
  --resume_from models/BAGEL-7B-MoT \
  --finetune_from_hf True \
  --auto_resume False \
  --resume_model_only True \
  --finetune_from_ema True \
  --lr 1e-4 --num_workers 1 --prefetch_factor 1 \
  --expected_num_tokens 256 --max_num_tokens 512 \
  --total_steps 10 \
  --cpu_offload True --sharding_strategy FULL_SHARD \
  --enable_profiler True \
  --profiler_skip_first 1 --profiler_wait 1 \
  --profiler_warmup 1 --profiler_active 3 \
  --profiler_with_stack True --profiler_rank0_only False \
  --wandb_offline True

9.2 Trace 对比脚本

python3 scripts/compare_profiler_traces.py \
  results_profiler/baseline_v4/<baseline>.pt.trace.json \
  results_profiler/profiler/rank0/<optimized>.pt.trace.json

9.3 产物路径

文件	路径
Baseline trace	results_profiler/baseline_v4/iv-yem3z5nitcwh2yowp610_923920....pt.trace.json
Optimized trace	results_profiler/profiler/rank0/iv-yem3z5nitcwh2yowp610_2286480....pt.trace.json
v4 日志	results_profiler/profile_run_v4.log
v5 日志	results_profiler/profile_run_v5.log
Summary	results_profiler/profiler/rank0/profiler_summary.json

---

10. 结论

1. 已完成：训练 hot path 上的主要同步点（bool 索引、每步标量 D2H、batch H2D 阻塞）均已修复，trace 可验证。
2. 实测收益 ：端到端约 +1%，符合「优化区域仅占单步 3--5%」的预期。
3. 下一步 ：若需大幅提升吞吐，应优先调整 FSDP 分片策略 和 cpu_offload，而非继续在 scalar/log 路径上投入。

---

报告生成自 Profile v4/v5 对比及代码审查。分析脚本：scripts/compare_profiler_traces.py。