微调130M参数蒸馏版GPT-2时无法利用GPU训练的技术求助

大家好，我最近在折腾用自己8年的WhatsApp聊天记录微调一个1.3亿参数的蒸馏版GPT-2模型。我整理了个带标签的数据集，大概有100万条序列，每条都是512个token。

不过现在遇到个头疼的问题——训练的时候没法利用GPU，卡得不行，甚至感觉根本没在跑GPU。下面是我训练代码的核心片段，麻烦各位大佬帮忙看看哪里出问题了：

def train_epoch(model, inputs, labels, batch_size, optimizer, scheduler, device):

    model.train()

    total_loss = 0
    iter = int(len(inputs) / batch_size)
    scaler = GradScaler()

    inputs=inputs.to(device)
    labels=labels.to(device)

    for i in range(iter):

        print(f"Iter {i + 1} of {iter}")

        inputs_batch= inputs[i*batch_size:(i+1)*batch_size]
        labels_batch = labels[i * batch_size:(i + 1) * batch_size]

        inputs_batch, labels_batch = inputs_batch.to(device), labels_batch.to(device)

        optimizer.zero_grad()

        print(inputs_batch.shape)

        with autocast("cuda"):
            outputs = model(inputs_batch, labels=labels_batch)
            # 后面的损失计算和反向传播代码暂时省略

我检查过CUDA是能正常用的，模型也已经移到cuda设备上了，但就是感觉GPU没在干活，训练速度慢得离谱，有没有大佬能指点下可能的问题点？

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