多GPU运行JAX脚本遇jaxlib.xla_extension.XlaRuntimeError内存不足，求解决方案

以下是针对jaxlib.xla_extension.XlaRuntimeError: RESOURCE_EXHAUSTED显存不足错误的可行解决办法：

• 减小批量大小：直接降低训练/推理的batch_size，这是最快速降低单步显存占用的方式。比如将batch_size=64调整为batch_size=32，具体数值根据GPU显存容量灵活测试。

• 关闭显存预分配：JAX默认会预分配大部分GPU显存，通过环境变量或代码禁用预分配，让显存按需分配：

  import os
  os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"
  os.environ["XLA_PYTHON_CLIENT_ALLOCATOR"] = "platform"
  

  也可以在启动脚本时设置：XLA_PYTHON_CLIENT_PREALLOCATE=false python your_script.py

• 启用显存自动回收：在脚本开头添加配置，让JAX自动释放不再使用的显存：

  import jax
  jax.config.update("jax_reuse_memory", True)
  jax.config.update("jax_enable_x64", False)  # 禁用64位精度，默认32位更省显存
  

• 模型并行分片：在多GPU环境下，手动将模型参数拆分到不同设备，避免单卡显存过载：

  from jax import device_put_sharded
  import jax.numpy as jnp
  
  params = get_model_params()  # 自定义获取模型参数的函数
  devices = jax.devices()
  sharded_params = device_put_sharded(list(jnp.array_split(params, len(devices))), devices)
  

• 启用混合精度训练：用半精度（FP16）进行计算，大幅减少显存占用，同时通过自动混合精度保持计算精度：

  from jax import lax
  jax.config.update("jax_default_matmul_precision", lax.Precision.HALF)
  

  也可以结合pmap实现分布式混合精度训练：

  from jax import pmap, float16, float32
  
  @pmap
  def train_step(params, batch):
      # 参数转半精度
      params_fp16 = jax.tree_map(lambda x: float16(x), params)
      # 执行训练逻辑（前向+反向传播）
      # ...
      # 参数转回单精度保存
      params = jax.tree_map(lambda x: float32(x), params_fp16)
      return params
  

• 清理不必要的张量：显式删除不再使用的中间变量，对不需要梯度的分支使用jax.lax.stop_gradient，避免保存冗余的梯度信息。

• 排查显存泄漏：用jax.debug.print_memory_profile()打印显存使用详情，定位是否存在未释放的张量或重复分配的内存块，针对性优化代码。

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