Arch Linux下PyTorch ROCm配置正常但训练报错求助

环境与安装情况 #

• 系统：Arch Linux
• 显卡：AMD Radeon RX 6700 XT（架构为gfx1031）
• PyTorch安装命令：

pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm5.6

• 配置的环境变量：

export PYTORCH_ROCM_ARCH="gfx1031"
export HSA_OVERRIDE_GFX_VERSION=10.3.1
export HIP_VISIBLE_DEVICES=0
export ROCM_PATH=/opt/rocm
export HIP_PLATFORM=amd
export HIP_DEVICE=0

测试脚本验证结果 #

运行PyTorch ROCm兼容性测试脚本后，输出显示环境基础验证通过：

Checking ROCM support...
GOOD: ROCM devices found:  2
Checking PyTorch...
GOOD: PyTorch is working fine.
Checking user groups...
GOOD: The user enchilada is in RENDER and VIDEO groups.
GOOD: PyTorch ROCM support found.
Testing PyTorch ROCM support...
Everything fine! You can run PyTorch code inside of:
--->  Intel(R) Core(TM) i7-10700KF CPU @ 3.80GHz
--->  gfx1031

实际训练报错信息 #

使用修改后的TinyLlama微调脚本（CPU运行正常）在GPU上训练时，触发如下错误：

File "/storage/Programs/PythonVenvs/transformers-rocm/lib/python3.11/site-packages/torch/nn/init.py", line 19, in _no_grad_normal_
    return tensor.normal_(mean, std)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: HIP error: invalid device function

rocminfo显卡输出（仅GPU部分） #

ROCk module is loaded
=====================
HSA System Attributes
=====================
Runtime Version:         1.1
System Timestamp Freq.:  1000.000000MHz
Sig. Max Wait Duration:  18446744073709551615 (0xFFFFFFFFFFFFFFFF) (timestamp count)
Machine Model:           LARGE
System Endianness:       LITTLE

*******
Agent 2
*******
  Name:                    gfx1031
  Uuid:                    GPU-XX
  Marketing Name:          AMD Radeon RX 6700 XT
  Vendor Name:             AMD
  Feature:                 KERNEL_DISPATCH
  Profile:                 BASE_PROFILE
  Float Round Mode:        NEAR
  Max Queue Number:        128(0x80)
  Queue Min Size:          64(0x40)
  Queue Max Size:          131072(0x20000)
  Queue Type:              MULTI
  Node:                    1
  Device Type:             GPU
  Cache Info:
    L1:                      16(0x10) KB
    L2:                      3072(0xc00) KB
    L3:                      98304(0x18000) KB
  Chip ID:                 29663(0x73df)
  ASIC Revision:           0(0x0)
  Cacheline Size:          64(0x40)
  Max Clock Freq. (MHz):   2855
  BDFID:                   768
  Internal Node ID:        1
  Compute Unit:            40
  SIMDs per CU:            2
  Shader Engines:          2
  Shader Arrs. per Eng.:   2
  WatchPts on Addr. Ranges:4
  Features:                KERNEL_DISPATCH
  Fast F16 Operation:      TRUE
  Wavefront Size:          32(0x20)
  Workgroup Max Size:      1024(0x400)
  Workgroup Max Size per Dimension:
    x                        1024(0x400)
    y                        1024(0x400)
    z                        1024(0x400)
  Max Waves Per CU:        32(0x20)
  Max Work-item Per CU:    1024(0x400)
  Grid Max Size:           4294967295(0xffffffff)
  Grid Max Size per Dimension:
    x                        4294967295(0xffffffff)
    y                        4294967295(0xffffffff)
    z                        4294967295(0xffffffff)
  Max fbarriers/Workgrp:   32
  Pool Info:
    Pool 1
      Segment:                 GLOBAL; FLAGS: COARSE GRAINED
      Size:                    12566528(0xbfc000) KB
      Allocatable:             TRUE
      Alloc Granule:           4KB
      Alloc Alignment:         4KB
      Accessible by all:       FALSE
    Pool 2
      Segment:                 GROUP
      Size:                    64(0x40) KB
      Allocatable:             FALSE
      Alloc Granule:           0KB
      Alloc Alignment:         0KB
      Accessible by all:       FALSE
  ISA Info:
    ISA 1
      Name:                    amdgcn-amd-amdhsa--gfx1031
      Machine Models:          HSA_MACHINE_MODEL_LARGE
      Profiles:                HSA_PROFILE_BASE
      Default Rounding Mode:   NEAR
      Default Rounding Mode:   NEAR
      Fast f16:                TRUE
      Workgroup Max Size:      1024(0x400)
      Workgroup Max Size per Dimension:
        x                        1024(0x400)
        y                        1024(0x400)
        z                        1024(0x400)
      Grid Max Size:           4294967295(0xffffffff)
      Grid Max Size per Dimension:
        x                        4294967295(0xffffffff)
        y                        4294967295(0xffffffff)
        z                        4294967295(0xffffffff)
      FBarrier Max Size:       32
*** Done ***

解决建议 #

• 版本兼容性检查：确认ROCm 5.6与当前安装的PyTorch版本是否完全适配，尝试升级PyTorch到对应ROCm版本的最新稳定版，或降级ROCm到5.5等更成熟的版本，部分版本组合可能存在算子预编译不匹配问题。
• 调整环境变量：
  • 移除HSA_OVERRIDE_GFX_VERSION=10.3.1，RX 6700 XT原生支持gfx1031架构，强制覆盖版本可能导致设备函数调用不匹配。
  • 将PYTORCH_ROCM_ARCH改为gfx1030尝试，部分场景下gfx1031可兼容gfx1030的预编译算子。
• 源码编译PyTorch：如果使用预编译包缺少gfx1031对应的算子，尝试从源码编译PyTorch，指定TORCH_ROCM_ARCH=gfx1031，确保所有算子针对你的显卡架构编译。
• 清理冲突环境变量：检查系统中是否存在CUDA_VISIBLE_DEVICES等CUDA相关变量，这类变量会导致PyTorch混淆设备后端，建议清空所有CUDA相关环境变量。
• 基础GPU操作测试：运行简单的GPU张量操作验证基础功能：

import torch
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
x = torch.randn(100, 100).to(device)
y = x @ x
print(y.mean())

如果该测试报错，说明基础GPU支持存在问题；如果正常，则问题可能出在微调脚本的特定算子，需检查脚本中是否有未适配ROCm的自定义CUDA算子，或尝试禁用混合精度等可能触发问题的设置。

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