# OpenVLA训练RLDS数据集时内存持续增长疑似泄漏

> 原文链接：https://robotics.hellobike.com/pitfalls/2024-06-21-openvla-4
> 频道：pitfalls | 更新日期：2024-06-21 | 标签：auto, github-issue, openvla
> 核心摘要：OpenVLA训练时内存随迭代增长非代码Bug，系TFDS预取机制所致，需调小shuffle_buffer_size适配硬件。

# OpenVLA训练RLDS数据集时内存持续增长疑似泄漏

::: v-pre
> 本文由 AI 根据公开内容整理，可能有误，请以原文为准：[memory leak](https://github.com/openvla/openvla/issues/4)（GitHub openvla/openvla #4）

**状态**：有绕过办法

## 问题

使用OpenVLA的rlds数据集（基于tf.data.Dataset）进行训练且train=True时，若调用shuffle未配合cache，内存占用随迭代持续增加。用户怀疑存在内存泄漏，担心在观察到内存 plateau 前因RAM不足导致进程崩溃。

## 环境

OpenVLA框架，RLDS数据集加载器，VLA训练配置；参考环境为单节点8卡A100 GPU配1TB RAM。

## 解决办法

1. 确认该现象并非内存泄漏Bug，而是TFDS dataloader底层激进预取导致的预期行为，内存最终会稳定但峰值可能超出系统限制。
2. 修改 `openvla/prismatic/conf/vla.py` 配置文件，降低 `shuffle_buffer_size` 参数值以适配当前可用RAM。
3. 参考官方基准：在1TB RAM+8*A100环境下，设置 `shuffle_buffer_size=256_000`（8进程总缓冲约2M样本），训练10万步后内存占用约70%。
4. 根据实际硬件资源逐步减小buffer size，在确保模型训练收敛稳定的前提下避免OOM。

讨论原作者：@yueyang130 等，版权归原作者所有。
:::

## 原始参考来源 (Sources)

- [GitHub openvla/openvla #4：memory leak](https://github.com/openvla/openvla/issues/4)