H100 GPU训练Donut模型时利用率几乎为0的性能瓶颈排查求助

架构老林

2026-4-14

我正在用自己的发票数据微调Donut Cord-v2模型，预处理后的数据集存到磁盘上大概有360GB。训练流程几乎完全参照一个notebook来做，只是把训练轮次从3改成了6。

我的训练硬件是：单Nvidia H100 SXM GPU + Intel Xeon® Gold 6448Y CPU + 128GB RAM。

每次启动训练后，我用htop和nvidia-smi监控CPU和GPU状态，发现了很奇怪的现象：CPU上的python进程利用率只有10-12%，GPU内存一直被占了近90%，但GPU的计算利用率几乎全程是0。每隔10-12秒刷新nvidia-smi时，才会看到利用率突然跳到100%然后立刻回落回0。

我总觉得是CPU和GPU之间存在瓶颈——CPU处理数据并往GPU传的速度跟不上，GPU飞快处理完一批数据后就只能闲置，等着CPU给下一批。

我是从磁盘直接加载已经预处理好的数据集，代码如下：

from datasets import load_from_disk
processed_dataset = load_from_disk(r"/dataset/dataset_final")

我的Processor配置代码是这样的：

from transformers import DonutProcessor

new_special_tokens = [] # 要添加到tokenizer的新特殊token
task_start_token = "<s>"  # 任务起始token
eos_token = "</s>" # tokenizer的结束token

processor = DonutProcessor.from_pretrained("naver-clova-ix/donut-base-finetuned-cord-v2")

# 给tokenizer添加新的特殊token
processor.tokenizer.add_special_tokens({"additional_special_tokens": new_special_tokens + [task_start_token] + [eos_token]})

# 更新一些和预训练不同的设置：比如图片尺寸 + 不需要旋转
processor.feature_extractor.size = [1200,1553] # 格式是(宽度, 高度)
processor.feature_extractor.do_align_long_axis = False

模型的配置代码如下：

import torch
from transformers import VisionEncoderDecoderModel, VisionEncoderDecoderConfig

# print(torch.cuda.is_available())

# 从huggingface加载预训练模型
model = VisionEncoderDecoderModel.from_pretrained("naver-clova-ix/donut-base-finetuned-cord-v2")

# 调整embedding层大小以匹配更新后的词汇表
new_emb = model.decoder.resize_token_embeddings(len(processor.tokenizer))
print(f"New embedding size: {new_emb}")
# 调整图片尺寸和输出序列的最大长度
model.config.encoder.image_size = processor.feature_extractor.size[::-1] # 格式是(高度, 宽度)
model.config.decoder.max_length = len(max(processed_dataset["train"]["labels"], key=len))

# 配置decoder的起始token
model.config.pad_token_id = processor.tokenizer.pad_token_id
model.config.decoder_start_token_id = processor.tokenizer.convert_tokens_to_ids(['<s>'])[0]

训练代码是这样的：

import gc
gc.collect()

torch.cuda.empty_cache()


from transformers import Seq2SeqTrainingArguments, Seq2SeqTrainer

import logging
logging.basicConfig(level=logging.INFO)

# 训练参数配置
training_args = Seq2SeqTrainingArguments(
    output_dir=r"/trained",  # 模型保存的本地目录
    num_train_epochs=6,
    learning_rate=2e-5,
    per_device_train_batch_size=8,
    weight_decay=0.01,
    fp16=True,
    logging_steps=50,
    save_total_limit=2,
    evaluation_strategy="no",
    save_strategy="epoch",
    predict_with_generate=True,
    report_to="none",
    # 禁止推送到hub
    push_to_hub=False
   
)

# 初始化Trainer
trainer = Seq2SeqTrainer(
    model=model,
    args=training_args,
    train_dataset=processed_dataset["train"],
)


# 启动训练
trainer.train()

现在预估6轮训练要花54小时，但我在自己的PC（Intel i9 11900KF + RTX 3050）上跑完全一样的代码时，GPU利用率一直是100%。我想知道是不是我的代码哪里有瓶颈？为什么已经预处理好的数据集还会让CPU的负载这么低，导致GPU闲置？我用的是CUDA 12.6。

另外补充个问题：我的RAM和CPU核心数都充足，把Seq2SeqTrainer的dataloader_num_workers参数设为大于0的值会不会解决这个问题？毕竟现在CPU利用率最高才10-12%，明显没跑满。

备注：内容来源于stack exchange，提问作者astralmaster