GPU性能优化与模型训练概览
安装所需库
为监控GPU内存使用,我们使用nvidia-ml-py3
库。首先安装必要的库:
python
pip install transformers datasets accelerate nvidia-ml-py3
模拟数据创建
创建范围在100到30000之间的随机token ID和二进制标签。为分类器准备512个序列,每个序列长度为512,并存储为PyTorch格式的数据集:
python
import numpy as np
from datasets import Dataset
seq_len, dataset_size = 512, 512
dummy_data = {
"input_ids": np.random.randint(100, 30000, (dataset_size, seq_len)),
"labels": np.random.randint(0, 1, (dataset_size)),
}
ds = Dataset.from_dict(dummy_data)
ds.set_format("pt")
GPU使用情况摘要
定义两个帮助函数来打印GPU使用情况及训练摘要:
python
from pynvml import *
def print_gpu_utilization():
nvmlInit()
handle = nvmlDeviceGetHandleByIndex(0)
info = nvmlDeviceGetMemoryInfo(handle)
print(f"GPU memory occupied: {info.used//1024**2} MB.")
def print_summary(result):
print(f"Time: {result.metrics['train_runtime']:.2f}")
print(f"Samples/second: {result.metrics['train_samples_per_second']:.2f}")
print_gpu_utilization()
模型加载与训练开销
加载BERT模型,并监测其权重占用的GPU内存:
python
from transformers import AutoModelForSequenceClassification
model = AutoModelForSequenceClassification.from_pretrained("google-bert/bert-large-uncased").to("cuda")
print_gpu_utilization()
训练前的准备
设置训练参数,以批大小为4进行训练,并监测内存占用情况:
python
from transformers import TrainingArguments, Trainer, logging
logging.set_verbosity_error()
default_args = {
"output_dir": "tmp",
"evaluation_strategy": "steps",
"num_train_epochs": 1,
"log_level": "error",
"report_to": "none",
}
training_args = TrainingArguments(per_device_train_batch_size=4, **default_args)
trainer = Trainer(model=model, args=training_args, train_dataset=ds)
result = trainer.train()
print_summary(result)
训练示例显示即使是较小的批大小也几乎填满了GPU内存。
模型运算与内存分析
转换器架构主要包括三类运算:
- 张量收缩:最计算密集型。
- 统计归一化:计算强度中等。
- 逐元素操作:计算强度最低。
模型在训练时占用的内存远超其权重占用量。其中包含:
- 模型权重
- 优化器状态
- 梯度
- 正向激活
- 临时缓冲区
- 特殊功能性内存
混合精度模型权重和激活量所需的总内存约为模型参数数量18字节,不含优化器状态和梯度的推理模式则约为6字节加上激活内存。
性能瓶颈与优化策略
了解模型运算和内存需求对分析性能瓶颈十分关键。可以参考相关文档,学习单GPU上高效训练的方法和工具。