93 lines
2.4 KiB
Python
93 lines
2.4 KiB
Python
# ====================
|
||
# 单元格: XLA版本兼容性检查和修复
|
||
# ====================
|
||
|
||
import torch
|
||
import torch.nn as nn
|
||
print("🔧 PyTorch XLA版本兼容性检查...")
|
||
|
||
# 导入XLA
|
||
import torch_xla.core.xla_model as xm
|
||
|
||
print("✅ PyTorch XLA导入成功!")
|
||
|
||
# 定义兼容性函数
|
||
def get_xla_world_size():
|
||
"""获取XLA world size，兼容不同版本"""
|
||
try:
|
||
return xm.xrt_world_size()
|
||
except AttributeError:
|
||
try:
|
||
return xm.get_world_size()
|
||
except AttributeError:
|
||
return 1 # 默认返回1
|
||
|
||
def get_xla_ordinal():
|
||
"""获取XLA ordinal，兼容不同版本"""
|
||
try:
|
||
return xm.get_ordinal()
|
||
except AttributeError:
|
||
return 0 # 默认返回0
|
||
|
||
def xla_mark_step():
|
||
"""XLA mark step，兼容不同版本"""
|
||
try:
|
||
xm.mark_step()
|
||
except AttributeError:
|
||
try:
|
||
xm.wait_device_ops()
|
||
except AttributeError:
|
||
pass # 如果都不可用，则跳过
|
||
|
||
def check_xla_device():
|
||
"""检查XLA设备状态"""
|
||
try:
|
||
device = xm.xla_device()
|
||
print(f"📱 XLA设备: {device}")
|
||
|
||
world_size = get_xla_world_size()
|
||
ordinal = get_xla_ordinal()
|
||
|
||
print(f"🌍 World Size: {world_size}")
|
||
print(f"🔢 Ordinal: {ordinal}")
|
||
|
||
# 检测设备类型
|
||
device_str = str(device)
|
||
if 'xla' in device_str and 'cpu' not in device_str:
|
||
print("✅ 检测到TPU设备")
|
||
return True, "TPU"
|
||
elif 'xla' in device_str and 'cpu' in device_str:
|
||
print("⚠️ XLA CPU模拟模式")
|
||
return True, "XLA_CPU"
|
||
else:
|
||
print("❌ 未检测到XLA设备")
|
||
return False, "CPU"
|
||
|
||
except Exception as e:
|
||
print(f"❌ XLA设备检查失败: {e}")
|
||
return False, "ERROR"
|
||
|
||
# 执行兼容性检查
|
||
device_available, device_type = check_xla_device()
|
||
|
||
if device_available:
|
||
print(f"✅ XLA环境正常，设备类型: {device_type}")
|
||
|
||
# 测试基本XLA操作
|
||
print("🧪 测试基本XLA操作...")
|
||
try:
|
||
device = xm.xla_device()
|
||
x = torch.randn(2, 2, device=device)
|
||
y = torch.matmul(x, x)
|
||
|
||
# 测试同步函数
|
||
xla_mark_step()
|
||
|
||
print("✅ 基本XLA操作测试成功")
|
||
|
||
except Exception as e:
|
||
print(f"❌ XLA操作测试失败: {e}")
|
||
else:
|
||
print("❌ XLA环境不可用")
|
||
|
||
print("\n💡 兼容性检查完成，可以运行后续单元格") |