45 lines
1.5 KiB
Python
45 lines
1.5 KiB
Python
|
# ====================
|
|||
|
|
# 单元格1: 环境设置 (必须第一个运行!)
|
|||
|
|
# ====================
|
|||
|
|
|
|||
|
|
import os
|
|||
|
|
import time
|
|||
|
|
import psutil
|
|||
|
|
from IPython.display import display, HTML, clear_output
|
|||
|
|
import ipywidgets as widgets
|
|||
|
|
|
|||
|
|
# ⚠️ 关键: 在导入torch_xla之前设置环境变量
|
|||
|
|
print("🔧 设置XLA环境变量...")
|
|||
|
|
|
|||
|
|
# 获取CPU核心数
|
|||
|
|
cpu_count = os.cpu_count()
|
|||
|
|
print(f"检测到 {cpu_count} 个CPU核心")
|
|||
|
|
|
|||
|
|
# 设置XLA编译优化环境变量
|
|||
|
|
os.environ['XLA_FLAGS'] = (
|
|||
|
|
'--xla_cpu_multi_thread_eigen=true '
|
|||
|
|
'--xla_cpu_enable_fast_math=true '
|
|||
|
|
f'--xla_force_host_platform_device_count={cpu_count}'
|
|||
|
|
)
|
|||
|
|
os.environ['PYTORCH_XLA_COMPILATION_THREADS'] = str(cpu_count)
|
|||
|
|
os.environ['XLA_USE_BF16'] = '1'
|
|||
|
|
|
|||
|
|
# 显示设置结果
|
|||
|
|
print("✅ XLA环境变量已设置:")
|
|||
|
|
print(f" CPU核心数: {cpu_count}")
|
|||
|
|
print(f" XLA_FLAGS: {os.environ['XLA_FLAGS']}")
|
|||
|
|
print(f" PYTORCH_XLA_COMPILATION_THREADS: {os.environ['PYTORCH_XLA_COMPILATION_THREADS']}")
|
|||
|
|
|
|||
|
|
# 系统资源检查
|
|||
|
|
memory_info = psutil.virtual_memory()
|
|||
|
|
print(f"\n💾 系统内存信息:")
|
|||
|
|
print(f" 总内存: {memory_info.total / (1024**3):.1f} GB")
|
|||
|
|
print(f" 可用内存: {memory_info.available / (1024**3):.1f} GB")
|
|||
|
|
print(f" 使用率: {memory_info.percent:.1f}%")
|
|||
|
|
|
|||
|
|
if memory_info.available < 8 * (1024**3): # 小于8GB
|
|||
|
|
print("⚠️ 警告: 可用内存不足8GB，可能影响XLA编译速度")
|
|||
|
|
else:
|
|||
|
|
print("✅ 内存充足")
|
|||
|
|
|
|||
|
|
print("\n🎯 环境设置完成! 现在可以运行下一个单元格")
|