特定隐藏层尺寸下PyTorch LSTM性能骤降问题求助

我在测试不同hidden size对PyTorch LSTM性能的影响时，发现某一特定范围内前向传播速度骤降。原本预期hidden size与前向传播时间呈线性相关，但实际并非如此。推测这可能是CUDA针对不同张量尺寸的优化策略与我的硬件适配不佳导致的。我已经尝试重装/更新CUDA和PyTorch，但问题仍然存在。以下是复现问题的代码、性能曲线图，恳请提供性能优化建议：

复现代码 #

import torch
from torch import nn
import time
from plotly import graph_objs as go
from tqdm import tqdm


def test_speed(hidden_size, n=5, gpu=0):

    torch.random.manual_seed(42)

    model = torch.nn.LSTM(
            2,
            hidden_size,
            num_layers=3,
            batch_first=True,
            dropout=0.5,
    ).cuda(gpu)

    x = torch.randn(256, 180, 2).cuda(gpu)

    # Warmup
    model.forward(x)

    output = []
    t = time.time()

    for i in range(n):
        with torch.no_grad():
            y = model.forward(x)[0]
            
        output.append(y.mean().item())  # 强制同步GPU计算
        
    return (time.time() - t) / n, output


test_range = range(64, 256)

go.Figure(
    [
        go.Scatter(
            x=list(test_range),
            y=[test_speed(hidden_size, n=20, gpu=0)[0] for hidden_size in tqdm(test_range)],
            name="GPU 0",
        ),
    ],
    go.Layout(
        title="平均前向传播时间 vs Hidden Size",
        xaxis_title="Hidden Size",
        yaxis_title="时间（秒）"
    )
)

性能曲线图 #

系统配置 #

• CPU：12代英特尔i9处理器
• GPU：NVIDIA RTX 3080 Ti
• CUDA版本：11.8
• PyTorch版本：1.12.1

性能优化建议 #

1. 选择硬件友好的Hidden Size数值
   CUDA内核对张量尺寸的对齐要求严格，优先选择2的幂次（如64、128、256）或GPU warp size（32）的倍数，避开曲线中速度骤降的区间。从测试结果看，64、128这类尺寸的性能表现更稳定。

2. 升级PyTorch至最新稳定版
   当前使用的PyTorch 1.12.1版本较旧，PyTorch 2.0及后续版本对LSTM的CUDA实现做了大量优化，包括内核重构、支持FlashAttention等特性，能显著提升序列模型的性能。建议升级到适配CUDA 11.8的PyTorch版本。

3. 启用自动混合精度（AMP）
   在推理阶段使用torch.cuda.amp.autocast()包裹前向传播代码，可在不影响精度的前提下降低计算量，提升GPU利用率。修改测试代码中的前向传播部分：

   with torch.no_grad(), torch.cuda.amp.autocast():
       y = model.forward(x)[0]
   

4. 开启CuDNN基准测试
   在测试代码开头添加torch.backends.cudnn.benchmark = True，让CuDNN自动为固定输入尺寸选择最优内核，避免每次创建模型时重新选择内核带来的开销。该选项仅适用于输入尺寸固定的场景，你的测试符合这个条件。

5. 验证Dropout的影响
   虽然推理时Dropout会被禁用，但模型初始化时的dropout设置可能影响内核的适配。可尝试将dropout=0代入测试，确认是否是Dropout相关的内核适配问题导致的性能波动。

6. 监控GPU内存状态
   某些hidden size可能导致GPU内存使用率突变，触发分页机制拖慢速度。可在测试中添加print(torch.cuda.memory_allocated(gpu))监控内存使用情况，排查是否存在内存瓶颈。

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