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

> 原文链接：https://robotics.hellobike.com/pitfalls/openpi-pi0-bfloat16-jax-oom-finetune
> 频道：pitfalls | 更新日期：2026-10-21 | 标签：OpenPI, pi0, pi0.5, JAX, 显存优化, VLA微调
> 核心摘要：OpenPI 默认使用 JAX 占用 75%+ GPU 显存，与 LeRobot PyTorch 视频解码进程并发时触发 CUDA OOM；需配置 XLA_PYTHON_CLIENT_MEM_FRACTION 与 FSDP。

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

## 硬件与软件环境
- **模型与框架**：`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 上稳定完成微调。

## 原始参考来源 (Sources)

- [GitHub Physical-Intelligence/openpi](https://github.com/Physical-Intelligence/openpi)