具身开发者社区

外观

踩坑github.com7 天前

OpenPI 微调 π0 / π0.5 时 JAX 显存预分配与 PyTorch DataLoader 冲突导致 OOM ​

OpenPI 默认使用 JAX 占用 75%+ GPU 显存，与 LeRobot PyTorch 视频解码进程并发时触发 CUDA OOM；需配置 XLA_PYTHON_CLIENT_MEM_FRACTION 与 FSDP。

  • 具身模型
  • 端侧部署

Franka RoboticsFranka FR3松灵机器人松灵 PiPER宇树科技宇树 G1星海图星海图 R1π0 (openpi)pi0_basepi05_base

硬件与软件环境 ​

  • 模型与框架：Physical-Intelligence/openpi（pi0_base / pi05_base）
  • 硬件环境：单卡或多卡 RTX 4090 (24GB) / A100 (40GB/80GB)，Ubuntu 22.04 + CUDA 12.x

故障现象与报错 ​

在运行 scripts/train.py 微调自定义机器人（如 PiPER、Franka FR3 或宇树 G1）LeRobot 数据集时，即使开启了 LoRA 微调，训练刚进入第一个数据加载步（或计算归一化统计量 compute_norm_stats.py 后启动训练）即报： jaxlib.xla_extension.XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate ... 或 PyTorch RuntimeError: CUDA error: out of memory。

根因分析 ​

  1. JAX 默认静态预分配 75% 显存：JAX 启动时会一次性向 CUDA 申请大块显存池且不归还给操作系统；当 DataLoader 或图像增强算子无意间触发了 GPU 端 PyTorch / TorchVision 视频解码（pyav / torchcodec CUDA 后端）时，两个框架争抢剩余显存立刻崩溃。
  2. 24GB 显存单卡运行全量参数优化器状态超限：pi0_base（3B 参数）若以 float32 维护 AdamW 一阶与二阶动量，仅优化器状态即需超过 24GB 显存。

解决步骤 ​

  1. 显式控制 JAX 显存分配策略并强制 DataLoader 在 CPU 解码： 在启动训练前导出以下环境变量：
    bash
    export XLA_PYTHON_CLIENT_MEM_FRACTION=0.85
export XLA_PYTHON_CLIENT_PREALLOCATE=false
    并在训练配置中确保 DataLoader 进程不初始化 CUDA 上下文（使用 CPU 多进程解码图像帧再由 JAX device_put 异步拷贝上卡）。
  2. 24GB 显卡务必启用 LoRA + bfloat16 权重配置： 在 src/openpi/training/config.py 中继承 pi0_aloha_pen_uncap 等 LoRA 配置（将 PaliGemma 主干与 Action Expert 设为 paligemma_variant="gemma_2b_lora"、action_expert_variant="gemma_300m_lora"），可将单卡训练峰值显存压降至 18.5 GB 以内，在单张 RTX 4090 上稳定完成微调。

来源

本页事实对应的公开页面，以原文为准。

  1. GitHub Physical-Intelligence/openpi访问 2026-10-21