SageMaker中PyTorch Lightning训练GPU利用率波动问题排查与优化

问题背景 #

在SageMaker上使用PyTorch（Lightning）训练ViT vit-base-patch16-224模型，采用单台ml.g4dn.xlarge实例（含1块GPU），数据存储于S3桶，数据集结构如下：

dataset.tsv
classA/
    img1
    img2
classB/
    img3
etc.

自定义Dataset类代码如下：

class MyDataset(Dataset):
    def __init__(self, dataset_root_path, transform):
        # Load .tsv file into a dataframe
        self.df = pd.read_csv(os.path.join(dataset_root_path, 'dataset.tsv'), sep='\t')

        self.transform = transform

    def __getitem__(self, idx):
        image_path = self.df['image_path'].iloc[idx]  # 修正原代码变量名错误
        image = Image.open(image_path)
        transformed_image = self.transform(image)  # 修正原代码变量x错误

        return transformed_image

Dataloader设置pin_memory=True、num_workers=4（尝试过num_workers=6无明显变化），训练过程中GPU利用率波动大，存在大幅下降情况，需解答以下问题：

1. 利用率波动的核心原因是什么？
2. 能否通过措施避免波动？
3. 按需实例和竞价实例是否会影响GPU利用率？
4. 如何进一步提升GPU利用率？

一、GPU利用率波动的核心原因 #

1. 数据加载瓶颈 #

• 直接从S3远程存储读取图片，网络IO延迟不稳定，__getitem__阶段会出现间歇性等待，导致GPU空闲。
• 数据预处理（transform）在CPU执行，若包含resize、增强等耗时操作，CPU处理速度跟不上GPU计算节奏，引发GPU等待，利用率骤降。
• 原Dataset代码存在变量名、大小写错误（如df_dataset未定义、self.Transform大小写混乱），若实际运行中存在此类问题，会引发额外耗时或异常。

2. Dataloader配置不合理 #

• num_workers与实例CPU核心数不匹配：ml.g4dn.xlarge有4个vCPU，设置num_workers=4/6会导致CPU过载，反而降低数据加载效率，加重GPU等待。
• 未开启persistent_workers=True：默认情况下每个epoch结束会重启worker，重启过程中GPU完全空闲，造成利用率波动。

3. 训练流程附加操作影响 #

• 频繁的checkpoint保存、日志写入：每epoch或每步写入磁盘（含S3）会占用CPU资源，干扰数据加载流程，间接导致GPU空闲。
• 模型验证环节：若验证集数据加载同样存在瓶颈，或验证阶段GPU计算量突变，也会引发利用率波动。

二、避免波动的优化措施 #

1. 优化数据加载流程 #

• 同步S3数据到实例本地存储：ml.g4dn.xlarge有125GB本地NVMe存储，训练前将整个数据集从S3复制到/tmp或/opt/ml/data目录，彻底消除远程IO延迟。
• 提前完成预处理：对所有图片做预处理后保存为二进制文件（如.pt格式），训练时直接加载预处理后的文件，减少CPU实时计算量。
• 用DataModule封装数据流程：借助PyTorch Lightning的DataModule，在prepare_data阶段完成数据下载/同步，setup阶段划分数据集，统一管理Dataloader配置，避免重复操作。

2. 调整Dataloader参数 #

• 匹配num_workers与CPU核心数：建议设置num_workers=2/3（留1个核心给主进程），避免CPU过载。
• 开启persistent_workers=True：让worker在epoch间保持活跃，避免重启耗时。
• 调大batch_size：根据GPU内存情况，将batch_size从默认值（如32）调至64或更高，让GPU更充分利用，减少空闲时间。
• 设置prefetch_factor：添加prefetch_factor=2参数，让每个worker提前加载2个batch的数据，减少GPU等待。

3. 修复Dataset代码错误 #

修正原代码中的变量名、大小写问题，同时统一图片格式：

class MyDataset(Dataset):
    def __init__(self, dataset_root_path, transform):
        self.df = pd.read_csv(os.path.join(dataset_root_path, 'dataset.tsv'), sep='\t')
        self.transform = transform  # 修正大小写错误

    def __getitem__(self, idx):
        image_path = self.df['image_path'].iloc[idx]  # 修正变量名
        image = Image.open(image_path).convert('RGB')  # 统一图片格式，避免额外格式转换耗时
        transformed_image = self.transform(image)  # 修正变量错误

        return transformed_image

4. 优化训练流程细节 #

• 降低checkpoint保存频率：如每5个epoch保存一次，而非每个epoch都保存，减少磁盘IO占用。
• 异步日志记录：使用PyTorch Lightning日志工具时，设置异步写入，避免阻塞训练流程。

三、按需实例与竞价实例对GPU利用率的影响 #

• 按需实例：资源完全独占，无被中断风险，GPU性能稳定，利用率波动仅来自训练流程本身，与实例类型无关。
• 竞价实例：未被AWS回收时，性能与按需实例完全一致，利用率波动根源同样是训练流程问题；若实例被回收，训练会中断，GPU利用率降至0，但这属于极端中断场景，并非正常训练中的波动。

简言之，只要竞价实例处于正常运行状态，对GPU利用率的影响与按需实例无差别。

四、进一步提升GPU利用率的方法 #

1. 开启混合精度训练：设置PyTorch Lightning的precision=16或bf16，减少GPU内存占用，同时提升计算速度。
2. 使用更快的图片加载库：替换PIL为torchvision.io.read_image或opencv-python，提升图片读取效率。
3. 添加数据缓存：在Dataset中用lru_cache对已加载的图片做内存缓存，避免重复读取相同样本。
4. 监控系统状态：通过SageMaker CloudWatch监控CPU使用率、磁盘IO、网络IO，定位具体瓶颈——若CPU使用率接近100%，说明是数据加载/预处理瓶颈；若CPU空闲，说明GPU计算量不足（如batch_size太小）。

内容的提问来源于stack exchange，提问作者Alcibiades