ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

如何用 PyTorch Profiler 对 DeepSpeed 训练循环做性能剖析?

如何用 PyTorch Profiler 对 DeepSpeed 训练循环做性能剖析? 如何用 PyTorch Profiler 对 DeepSpeed 训练循环做性能剖析【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed如果你已经用 DeepSpeed 把模型包成了引擎即通过deepspeed.initialize得到的model_engine并且想定位训练循环里前向、反向、权重更新各阶段的耗时可以用 PyTorch Profiler 对训练循环做性能剖析。DeepSpeed 的教程 Using PyTorch Profiler with DeepSpeed 给出了完整做法把训练循环包在 profiler 的上下文管理器里按 step 驱动采样结果可以导出为.jsontrace 文件在 Google 的 Perfetto trace viewer 中查看VS Code 的 Python 扩展内置了 TensorBoard 并支持 PyTorch Profiler也可用来查看结果。前提已初始化的 DeepSpeed 引擎PyTorch Profiler 本身是 PyTorch 的工具DeepSpeed 侧的要求只有一条训练循环走的是 DeepSpeed 引擎的三个 API。按 Getting Started 的说明先通过deepspeed.initialize初始化引擎model_engine, optimizer, _, _ deepspeed.initialize(argscmd_args, modelmodel, model_parametersparams)之后每个训练 step 由前向引擎可调用、反向和权重更新组成for step, batch in enumerate(data_loader): loss model_engine(batch) # forward model_engine.backward(loss) # backpropagation model_engine.step() # weight update下文的剖析示例都基于这个循环结构。示例代码中的model_engine、data_loader、inputs都是指你自己训练代码中的对应对象。用 profiler 上下文管理器包裹训练循环PyTorch Profiler 假设训练过程由 step 组成编号从零开始通过schedule控制跳过、预热和正式记录各多少个 step。下面的示例会跳过前5个 step用接下来2个 step 做预热开始 trace 但结果丢弃正式记录接下来6个 step由于repeat设为2两个周期后 profiler 停止记录from torch.profiler import profile, record_function, ProfilerActivity with torch.profiler.profile( scheduletorch.profiler.schedule( wait5, # During this phase profiler is not active. warmup2, # During this phase profiler starts tracing, but the results are discarded. active6, # During this phase profiler traces and records data. repeat2), # Specifies an upper bound on the number of cycles. on_trace_readytensorboard_trace_handler, with_stackTrue # Enable stack tracing, adds extra profiling overhead. ) as profiler: for step, batch in enumerate(data_loader): print(step:{}.format(step)) #forward() method loss model_engine(batch) #runs backpropagation model_engine.backward(loss) #weight update model_engine.step() profiler.step() # Send the signal to the profiler that the next step has started.两个要点每执行完一个model_engine.step()后要调用profiler.step()向 profiler 发出下一个 step 开始的信号schedule的 wait/warmup/active 计数就是按这个信号推进的。with_stackTrue会开启调用栈追踪文档明确说明这会带来额外剖析开销开销敏感时可以关掉。示例中的on_trace_readytensorboard_trace_handler是文档示例给定的回调位置文档本身没有给出tensorboard_trace_handler的实现它由你提供用于在 trace 就绪时做输出例如交给 TensorBoard。如果暂时不需要这个回调可以不传该参数trace 数据仍可随 profiler 上下文使用。schedule参数针对长时训练任务的详细用法可参考 PyTorch 官方的 profiler recipe教程 pytorch-profiler.md 末尾给出了指向链接。给任意代码段打标签record_function上下文管理器可以用自定义名字标记任意代码范围方便在 trace 里直接找到对应区段。例如把前向阶段标记为model_forwardwith profile(record_shapesTrue) as prof: # record_shapes indicates whether to record shapes of the operator inputs. with record_function(model_forward): model_engine(inputs)record_shapes控制是否记录算子输入的 shape需要保留这一信息时才开启。选择剖析 CPU 或 GPU 活动activities参数指定剖析哪些活动ProfilerActivity.CPUPyTorch 算子、TorchScript 函数和record_function的用户自定义标签。ProfilerActivity.CUDA设备上的 CUDA kernel。文档特别提醒CUDA 剖析会带来不可忽视的开销non-negligible overhead只关心算子级 CPU 侧统计时不必开启 CUDA 活动。下面这个示例同时剖析前向阶段的 CPU 和 GPU 活动并打印按总 CUDA 时间排序、只保留前 10 行的汇总表with profile(activities[ ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapesTrue) as prof: with record_function(model_forward): model_engine(inputs) print(prof.key_averages().table(sort_bycuda_time_total, row_limit10))key_averages().table(...)的输出直接打印到标准输出这是快速判断哪些算子最耗时的最直接方式看cuda_time_total排序靠前的算子是否集中在你想优化的区段。剖析显存占用给 profiler 传profile_memoryTrue可以开启内存剖析记录执行期间各算子分配或释放的张量内存量with profile(activities[ProfilerActivity.CUDA], profile_memoryTrue, record_shapesTrue) as prof: model(inputs) print(prof.key_averages().table(sort_byself_cuda_memory_usage, row_limit10))表按self_cuda_memory_usage排序。文档说明self内存指算子自身分配释放的内存不含对子算子的调用用它定位单个算子的直接内存开销。结果验证与查看方式完成剖析后的验证路径有两种均来自教程正文汇总表调用prof.key_averages().table(...)后终端会打印排序表格上面各节示例的print语句据此判断耗时或内存热点算子。trace 文件剖析结果可输出为.jsontrace 文件在 Google 的 Perfetto trace viewer 中打开查看逐 kernel、逐区段的时间线如果习惯在编辑器里看VS Code 的 Python 扩展集成了 TensorBoard 并支持 PyTorch Profiler。限制与可替代的工具文档给出的wait/warmup/active/repeat是教程示例值5/2/6/2不是必须值按你的训练步数和记录目标自行调整但要保证训练循环总步数足够覆盖这些阶段。CUDA 剖析和with_stack都会增加开销剖析结果反映的是带剖析运行时的状态。PyTorch Profiler 统计的是 PyTorch 算子级性能。如果你要的是模块级哪个子模块占多少参数、FLOPS、前向延迟的剖析DeepSpeed 另有一个独立的 Flops Profilerdeepspeed.profiling.flops_profiler它可以直接在deepspeed_config里用flops_profiler字段启用不需要改用户代码两者侧重点不同可按需选用教程中已明确区分。【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表