PyTorch compile性能下降及BERT预训练模型报错问题咨询

问题1：torch.compile() 性能不升反降 #

预期torch.compile()能提升性能，但使用后代码耗时反而增加20%-30%，代码存在哪些问题？

问题2：BERT预训练模型使用compile报错且速度骤降 #

在BERT预训练模型中使用compile时，出现以下报错且速度骤降数十倍，原因是什么？

报错信息 #

• torch._dynamo.convert_frame: [WARNING] torch._dynamo hit config.cache_size_limit (64) function: 'forward' (/home/mj/.../bert.py:287) reasons: ___check_obj_id(self, 139626116174448)
• torch._inductor.utils: [WARNING] using triton random, expect difference from eager

测试代码 #

import torch
import torchvision
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import random

import time
import numpy as np


device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
torch.manual_seed(1234)
np.random.seed(1234)

class NeuralNetwork(nn.Module):

    def __init__(self):
        super(NeuralNetwork, self).__init__()

        self.sequential = nn.Sequential(
            nn.Conv2d(in_channels=3, out_channels=32, kernel_size=3, padding=2),
            nn.ReLU(),
            nn.MaxPool2d(stride=2, kernel_size=3),
            nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, padding=2),
            nn.ReLU(),
            nn.MaxPool2d(stride=2, kernel_size=3),
            nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, padding=2),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d(8),
            nn.Flatten(),
            nn.Linear(in_features=128*8*8, out_features=10))


    def forward(self, x):
        out = self.sequential(x)
        return out

# model1 = NeuralNetwork()
# model2 = NeuralNetwork()
model1 = torchvision.models.resnet18()
model2 = torchvision.models.resnet18()
model1.to(device='cuda')
model2.to(device='cuda')
model1 = torch.compile(model1)
model2 = torch.compile(model2)

def cifar10_cnn():
    epochs = 5        
    batch_size = 64         
    report_period = 100 
    tr_count = 0           
    te_count = 0        
    data_root = "/data/" 

    torch.set_float32_matmul_precision('high')

    tr_dset = datasets.CIFAR10(root=data_root, train=True, download=True, transform=transforms.ToTensor())
    te_dset = datasets.CIFAR10(root=data_root, train=False, download=True, transform=transforms.ToTensor())

    tr_loader = DataLoader(tr_dset, batch_size=batch_size, shuffle=True)
    te_loader = DataLoader(te_dset, batch_size=batch_size, shuffle=False)

    loss_fn = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model1.parameters(), lr=1e-3)


    start_time = time.time()

    for i in range(epochs):
        print(f"\nEpoch {i + 1}/{epochs}\n------------------------------")
        train(tr_loader, model1, loss_fn, optimizer, report_period, start_time)
        print(f"\nTest started with {len(te_loader)} data:")
        test(te_loader, model2, loss_fn, start_time)

def train(dataloader, model1, loss_fn, optimizer, report_period, start_time):
    running_loss = 0.0
    train_loss= 0.0
    size = len(dataloader.dataset)
    for batch, (X, y) in enumerate(dataloader):
        X, y = X.to(device), y.to(device)

        pred = model1(X)
        loss = loss_fn(pred, y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
        train_loss = running_loss / len(dataloader)

        if batch % 100 == 0:
            loss, current = loss.item(), batch * len(X)
            print(f"loss: {loss:>7f} [{current:>5d}/{size:>5d}]")

    print(f"train_loss: {train_loss}")


def test(dataloader, model2, loss_fn, start_time):

    size = len(dataloader.dataset)
    num_batches = len(dataloader)
    model2.eval()
    test_loss, correct = 0, 0
    with torch.no_grad():
        for X, y in dataloader:
            X, y = X.to(device), y.to(device)

            pred = model1(X)  # 此处存在逻辑错误
            test_loss += loss_fn(pred, y).item()

            correct += (pred.argmax(1) == y).type(torch.float).sum().item()
    test_loss /= num_batches
    correct /= size

    print(f"Test Error: \n Accuracy: {(100 * correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")


if __name__ == "__main__":
    start = time.time()
    cifar10_cnn()
    print("Done!\n")
    print(f"running time: {time.time()-start}")

已尝试的方案 #

• 最初使用单个编译模型同时执行训练和测试出现梯度相关错误，因此拆分了训练和测试专用的编译模型；
• 更换torchvision提供的模型进行测试，性能仍出现下降；
• 尝试使用大模型，排除小模型编译开销影响，结果依旧。

解答 #

针对问题1：性能不升反降的原因与修复 #

1. 代码逻辑错误
   test函数存在明显bug：传入model2但实际调用model1，导致测试阶段仍使用训练模式模型（未关闭梯度），额外增加计算开销。修正为：

   pred = model2(X)
   

2. 编译顺序错误
   当前先执行model.to(device)再编译，正确顺序应为先编译再移动到设备，否则可能导致编译后的计算图与设备不匹配：

   model1 = torch.compile(model1)
   model2 = torch.compile(model2)
   model1.to(device='cuda')
   model2.to(device='cuda')
   

3. 编译策略未优化
   默认编译配置未针对硬件/场景适配，可指定更合适的backend与模式：

   • CUDA设备推荐开启自动调优，让编译器搜索最优核配置：

     model1 = torch.compile(model1, backend="inductor", mode="max-autotune")
     model2 = torch.compile(model2, backend="inductor", mode="max-autotune")
     

   • 老架构GPU若不支持Triton，可切换到backend="aot_eager"避免无效调度开销。

4. 计算密度不足
   CIFAR10的64 batch_size对ResNet18来说计算密度偏低，编译器优化空间有限。可根据GPU显存调整batch_size至128/256，提升GPU利用率，放大编译收益。

5. 未排除编译热身开销
   首次编译存在额外开销，建议正式计时前先运行1-2个batch做热身，避免将编译时间计入总耗时。

针对问题2：BERT使用compile报错且速度骤降的原因 #

1. 重复编译触发缓存上限
   hit config.cache_size_limit (64)说明TorchDynamo触发了重复编译，原因是BERT的forward中存在依赖self对象的动态逻辑（如条件分支、动态形状、参数变化），每次前向都会生成新计算图，超出缓存后持续重编译导致速度骤降。

   • 解决：检查BERT forward代码，移除不必要的动态分支；或调大缓存限制：

     import torch._dynamo
     torch._dynamo.config.cache_size_limit = 128
     

2. Triton随机操作的兼容性问题
   using triton random警告说明编译器使用了Triton实现的随机操作，与原生PyTorch逻辑存在差异，且Triton对BERT部分结构优化不足，导致速度下降。

   • 解决：禁用Triton随机或切换稳定backend：

     import torch._inductor
     torch._inductor.config.use_triton_random = False
     

     或使用backend="aot_eager"替代Inductor。

3. 动态特性破坏编译稳定性
   BERT的attention层存在动态序列长度、mask变化等特性，会破坏编译图的稳定性，导致频繁重编译。可尝试固定输入序列长度，或开启动态形状支持（优化效果弱于静态模式）：

   model = torch.compile(model, dynamic=True)
   

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