PyTorch | 加速模型训练的妙招

引言

提升机器学习模型的训练速度是每位机器学习工程师的共同追求。训练速度的提升意味着实验周期的缩短,进而加速产品的迭代过程。同时,这也表示在进行单一模型训练时,所需的资源将会减少。简而言之,我们追求的是效率。

熟悉 PyTorch profiler

在进行任何优化之前,首先需要了解代码中各个部分的执行时长。Pytorch profiler 是一款功能全面的训练性能分析工具,能够捕捉以下信息:

  • CPU 操作的耗时
  • CUDA 核心的运行时间
  • 内存使用情况的历史记录

这些就是你需要关注的所有内容。而且,使用起来非常简单!记录这些事件的方法是,将训练过程封装在一个 profiler 的上下文环境中,操作方式如下:

import torch.autograd.profiler as profiler

with profiler.profile(
  activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
  on_trace_ready=torch.profiler.tensorboard_trace_handler('./logs'),
) as prof:
  train(args)

之后,您可以启动张量板并查看分析跟踪。Profiler 提供了众多选项,但最关键的是 "activities" 和 "profile_memory" 这两个功能。尽管你可以探索其他功能,但请记住一个基本原则:启用的选项越少,性能开销也就越低。

例如,如果你的目的是分析 CUDA 内核的执行时间,那么最好的做法是关闭 CPU 分析和其他所有功能。这样,分析结果会更贴近实际的执行情况。

为了让分析结果更易于理解,建议添加一些描述代码关键部分的分析上下文。如果分析功能没有被激活,这些上下文就不会产生任何影响。

with profiler.record_function("forward_pass"):
  result = model(**batch)

with profiler.record_function("train_step"):
  step(**result)

这样,您使用的标签将在迹线中可见。因此,识别代码块会更容易。或者更精细的内部模式的前进:

with profiler.record_function("transformer_layer:self_attention"):
  data = self.self_attention(**data)

...

with profiler.record_function("transformer_layer:encoder_attention"):
  data = self.encoder_attention(**data, **encoder_data)

了解 PyTorch traces

收集traces后,在张量板中打开它们。 CPU + CUDA 配置文件如下所示:

立刻识别出任何训练过程中的关键环节:

  • 数据加载
  • 前向传播
  • 反向传播

PyTorch 会在一个独立线程中处理反向传播(如上图所示的线程 16893),这使得它很容易被识别出来。

数据加载

在数据加载方面,我们追求极致的效率,即几乎不耗费时间。

原因在于,在数据加载的过程中,GPU 闲置不工作,这导致资源没有得到充分利用。但是,由于数据处理和 GPU 计算是两个独立的部分,它们可以同时进行。

你可以通过查看分析器跟踪中的 GPU 估计 SM 效率和 GPU 利用率来轻松识别 GPU 空闲的区域。那些活动量为零的区域就是我们需要注意的问题所在。在这些区域,GPU 并没有参与任何工作。

解决这个问题的一个简单方法是:

  • 在后台进程中进行数据处理,这样不会受到全局解释器锁(GIL)的限制。
  • 通过并行进程来同时执行数据增强和转换操作。 如果你使用的是 PyTorch 的 DataLoader,通过设置 num_workers 参数就可以轻松实现这一点。如果你使用的是 IterableDataset,情况会稍微复杂一些,因为数据可能会被重复处理。不过,通过使用 get_worker_info() 方法,你仍然可以解决这个问题------你需要调整迭代方式,确保每个工作进程处理的是互不重叠的不同数据行。

如果你需要更灵活的数据处理方式,你可以考虑使用 multiprocessing 模块来自己实现多进程转换功能。

内存分配器

使用 PyTorch 在 CUDA 设备上分配张量时,PyTorch 会利用缓存分配器来避免执行成本较高的 cudaMalloccudaFree 操作。PyTorch 的分配器会尝试复用之前通过 cudaMalloc 分配的内存块。如果分配器手头有合适的内存块,它将直接提供这块内存,而无需再次调用 cudaMalloc,这样 cudaMalloc 只在程序启动时调用一次。

但是,如果你处理的是长度不一的数据,不同前向传播过程可能需要不同大小的中间张量。这时,PyTorch 的分配器可能没有合适的内存块可用。在这种情况下,分配器会尝试通过调用 cudaFree 释放之前分配的内存块,以便为新的内存分配腾出空间。

释放内存后,分配器会重新开始构建其缓存,这将涉及到大量的 cudaMalloc 调用,这是一个资源消耗较大的操作。你可以通过观察 tensorboard profiler viewer 的内存分析部分来识别这个问题。

请注意,代表分配器预留内存的红线持续波动。这表明 PyTorch 的内存分配器在处理内存请求时遇到了效率问题。

当内存分配在没有触发分配器紧急情况下顺利进行时,你会看到红线保持平稳。

本文由mdnice多平台发布

相关推荐
AI_小站4 小时前
LLM——10个大型语言模型(LLM)常见面试题以及答案解析
人工智能·程序人生·语言模型·自然语言处理·大模型·llm·大模型面试
良技漫谈8 小时前
Rust移动开发:Rust在iOS端集成使用介绍
后端·程序人生·ios·rust·objective-c·swift
我爱学Python!11 小时前
AI Prompt如何帮你提升论文中的逻辑推理部分?
人工智能·程序人生·自然语言处理·chatgpt·llm·prompt·提示词
博纳软云_小程序一站服务平台1 天前
手边酒店多商户版V2源码独立部署_博纳软云
程序人生·微信小程序·小程序·微信公众平台
AI_小站1 天前
多模态大模型微调实践!PAI+LLaMA Factory搭建AI导游
人工智能·程序人生·语言模型·大模型·llm·产品经理·多模态大模型
良技漫谈1 天前
Rust移动开发:Rust在Android端集成使用介绍
android·程序人生·rust·kotlin·学习方法
AI_小站1 天前
【AI工作流】FastGPT - 深入解析FastGPT工作流编排:从基础到高级应用的全面指南
人工智能·程序人生·语言模型·大模型·llm·fastgpt·大模型应用
python_知世1 天前
AI时代:成为产品经理的核心路径
人工智能·深度学习·程序人生·自然语言处理·产品经理·计算机技术·大模型应用
提笔惊蚂蚁2 天前
java-web-day7-会话跟踪技术
java·开发语言·前端·程序人生