字节跳动以AI Profiling诊断训练与推理负载

在字节跳动内部，所有低 SMA 的训练任务以及推理服务上线前的压测环节都会进行 AI Profiling 采集，帮助研发通过可视化分析快速定位 GPU 效率低下与性能抖动问题。基于这一内部实践，字节跳动 STE 系统技术与工程团队联合火山引擎容器团队将该能力产品化封装并融入容器服务 VKE，形成 VKE AI Profiling 并对外提供服务。 文中以真实 ResNet18 + CIFAR10 训练任务（PyTorch 2.1.0、TorchVision 0.16.0、业内主流训练卡、采集项 python/cuda/pytorch、采集 5 秒）演示排查与优化：基线暴露数据加载慢导致 GPU 等待、Host→Device 拷贝耗时、未用混合精度三类瓶颈。 优化一将 DataLoader num_workers 由 1 提升到 8，数据加载耗时从 16.5 ms 降到 1.49 ms。优化二开启 pin_memory=True 并配合 non_blocking=True，数据拷贝耗时从 2.2 ms 降到 0.76 ms。优化三开启 AMP（torch.cuda.amp.autocast，float16），单个 train_step 平均执行时长从 29 ms 降到 22.7 ms。