PyTorch中loss.backward()报错Missing XLA configuration求助

问题描述 #

使用PyTorch构建模型，前向传播正常，但执行loss.backward()时触发如下报错：

terminate called after throwing an instance of 'std::runtime_error'
  what(): tensorflow/compiler/xla/xla_client/computation_client.cc:280 : Missing XLA configuration
Aborted

环境信息：

• PyTorch 1.12.0+cu102、torchvision 0.13.0+cu102，模型包含torchvision预训练CNN
• Google Compute Engine环境，GPU为NVIDIA Tesla T4（CUDA 11.6）
• 此前CUDA 11.2环境可正常运行，当前环境无论GPU/CPU均报错，未安装TPU，仅需使用GPU
• 代码在本地及其他GPU环境正常，更新环境后出现问题

解决方案 #

1. 清理TensorFlow相关环境变量 #

环境中残留的TF/XLA相关环境变量可能干扰PyTorch运行，执行以下命令清理：

# 先检查是否存在相关变量
printenv | grep -E "(XLA|TF)"
# 移除变量
unset XLA_FLAGS
unset TF_XLA_FLAGS

2. 重装匹配CUDA版本的PyTorch #

当前安装的PyTorch是cu102版本，与系统CUDA 11.6不匹配，会导致底层依赖混乱。先卸载旧版本，再安装对应CUDA 11.6的PyTorch：

pip uninstall -y torch torchvision
pip3 install torch==1.13.1+cu116 torchvision==0.14.1+cu116 --extra-index-url https://download.pytorch.org/whl/cu116

3. 隔离GCE预装的TensorFlow组件 #

部分GCE镜像会预装TensorFlow基础库，即使未显式安装TF，也可能干扰PyTorch。可以通过设置LD_LIBRARY_PATH让系统优先加载PyTorch的CUDA库：

export LD_LIBRARY_PATH=$(python -c "import torch; print(torch.__path__[0])")/lib:$LD_LIBRARY_PATH

若仍有问题，可检查并卸载系统级的TensorFlow相关包：

ldconfig -p | grep tensorflow
# 若输出结果，使用apt卸载对应包（示例）
sudo apt remove tensorflow-libs

4. 验证环境 #

安装完成后，运行以下代码验证PyTorch及反向传播功能：

import torch
print(f"CUDA可用: {torch.cuda.is_available()}")
print(f"PyTorch绑定的CUDA版本: {torch.version.cuda}")

# 测试反向传播
x = torch.randn(2, 3).cuda()
model = torch.nn.Linear(3, 1).cuda()
loss_fn = torch.nn.MSELoss()
y_pred = model(x)
loss = loss_fn(y_pred, torch.randn(2,1).cuda())
loss.backward()
print("反向传播测试成功")

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