Cannot allocate memory——训练时视频解码为什么会内存越跑越大

这篇文章记录一个在 VLA 视频数据训练里非常常见、也非常折磨人的问题：训练跑着跑着内存持续增长，最后随机触发 av.error.MemoryError: [Errno 12] Cannot allocate memory 、或者某个 rank 先挂掉引发 DDP 的 Broken pipe / NCCL 连锁退出。

更坑的是：明明 del 了对象、甚至加了 gc.collect()，内存还是不降。

我们最后通过一组改动把它稳定住：包括显式关闭解码器对象、周期性 GC、数据访问/轨迹缓存清理、禁用 persistent_workers（让 worker 周期性重启释放状态）、降低 prefetch，让训练能在多 epoch 下长期运行。

下面从原理讲起：av 是什么、为什么会"泄漏/持续增长"、为什么 DDP + DataLoader 场景下"删视频/删对象"不等于释放内存、最后我们方案到底在"释放"什么。

1. `av` 到底是什么？它在训练里负责哪一段

我们看到的 av 一般指的是 PyAV（Python 的 AV/FFmpeg 绑定）。它的作用是把视频文件（mp4/mkv...）解码成一帧帧图像，供训练数据管道产出 tensor。

简单画一下链路：

复制代码

video.mp4
  ↓
FFmpeg / libavcodec（C/C++）
  ↓
PyAV（Python wrapper，import av）
  ↓
frame（numpy/torch tensor）

而 torchvision_av 本质上是 torchvision 的视频后端封装，底层很多情况下还是走 PyAV/FFmpeg 那套机制（不同版本实现细节略有差异，但"底层是 C/C++ 解码器 + Python 包装"的事实不变）。

所以遇到的 "av 内存问题"，本质大多不是 Python 本身的问题，而是 FFmpeg/解码器的资源生命周期 + Python 封装层的生命周期管理在训练场景下被放大了。

2. 为什么会出现"内存泄漏"或"内存持续增长"？

这里要非常小心：很多时候并不是传统意义上的"泄漏"（malloc 后再也不 free），而是几类行为在监控里看起来像泄漏，最终同样会把系统推到崩溃边缘。

2.1 进程的虚拟内存（VIRT）会越来越大：mmap 与地址空间膨胀

FFmpeg 解码经常会用到 mmap / 大块虚拟地址映射 / 缓冲池 。

这会导致：

htop 里 VIRT 非常大（几十 GB、几百 GB 都可能）
但实际常驻物理内存（RES）可能并不高

这类"虚拟地址空间膨胀"并不是立即致命，但会增加 Linux 内核的 commit 压力，在某些时刻 fork/mmap 可能被拒绝，从而抛出 ENOMEM（Errno 12）。

这也是为什么会看到"明明还有 100+ GB available，仍然 Cannot allocate memory"的反直觉现象。

2.2 缓冲池与 arena 不会立刻归还给 OS：删了对象，但内存不降

即使正确地 del 了 Python 对象：

C/C++ 层（FFmpeg / libavcodec）可能仍然持有 buffer pool
glibc 的内存分配器（malloc arena）也可能把内存"留在进程里复用"，而不是立刻 munmap 回操作系统

因此监控里会表现为：

训练过程中内存曲线呈锯齿波动
但锯齿的基线不断抬高（"越跑越大"）

这并不一定意味着"每次都漏掉了对象"，也可能是底层为了性能复用内存，但在你的训练模式（大量视频、反复解码、多进程）下会演化成不可控的增长。

2.3 视频解码不是纯函数：上下文、状态、缓存会跨样本累积

视频解码器会维护一些状态（codec context、帧重排、参考帧、线程池等）。在"训练数据集"里，样本通常是很多短 clip / 不同视频片段，频繁打开关闭容器。

如果每个 __getitem__ 都要：

打开视频容器
seek
decode 若干帧
关闭

那么底层状态开销巨大且频繁创建销毁。稍微有一点"未及时释放/未关闭/引用残留"，长期训练就会被放大。

3. 为什么 DDP / DataLoader 下，"删掉解码过的视频"达不到预期释放效果？

这部分是很多人最困惑的点：明明 del 了 container、frame，甚至调用 gc.collect()，为什么内存还是涨？为什么 DDP 下更糟？

核心原因是：你以为你在管理"对象生命周期"，但问题实际发生在"进程生命周期"和"底层库生命周期"上。

3.1 DataLoader 多 worker = 多进程；父进程删对象不影响子进程

当你设置 num_workers > 0：

DataLoader 会启动多个 worker 进程
__getitem__ 的视频解码通常发生在这些 worker 进程里

这意味着：

你在父进程（主训练进程）里做 del，并不能释放 worker 进程里的任何资源
即使你在 worker 里 del，底层 C/C++ 资源是否释放也不由 Python 100% 控制

3.2 persistent_workers 会让"有问题的 worker 一直活着"，导致状态累积

persistent_workers 决定 DataLoader 的 worker 进程：
是"每个 epoch 用完就死"，还是"一直活着反复干活"。

persistent_workers 是 torch.utils.data.DataLoader 的一个参数：

复制代码

DataLoader(
    dataset,
    num_workers=N,
    persistent_workers=True or False
)

False（默认）

👉 每个 epoch 结束后，worker 进程会被关闭
True

👉 worker 进程会在 epoch 之间保持存活并复用

PyTorch 引入它的目的只有一个：性能。

3.3 DDP = 多个 rank = 多份进程状态；任何一个 rank 先挂都会"拖全场下水"

DDP 的本质是 每张 GPU 一个进程（一个 rank）。你用 2 卡就是至少 2 个训练进程。每个 rank 都有自己的：

DataLoader
worker 子进程（如果 num_workers>0）
解码器状态

所以内存增长的"危险因素"会被乘上 rank 数量。

一旦某个 rank 因为 ENOMEM 或解码异常退出：

其余 rank 与它通信时就会出现 Broken pipe
NCCL/ProcessGroup 心跳线程会报错
最终表现为"看起来是 NCCL/通信崩了"，但根因是 其中一个进程先死了

这也是为什么会在日志里看到 ProcessGroupNCCL::HeartbeatMonitor + TCPStore Broken pipe：不是通信本身坏，而是对端进程先没了。

4. 为什么"每个 epoch 都会重新解码一遍"，问题会越来越明显？

以 HuggingFace Trainer 的定义：

global batch size = per_device_batch_size × num_gpus × grad_accum
steps per epoch = ceil(dataset_len / global_batch)

当看到日志里：

dataset length = 52970
global batch size = 128
train dataloader length ≈ 414

这意味着：414 step 就会遍历完整数据集一遍。

如果你的 Dataset 是 LeRobotSingleDataset 这种"在线解码"的实现，并且没有帧缓存：

每个 __getitem__ 都会去"打开视频 → seek → decode"
所以 每个 epoch 都要把所有视频解码一遍
多 epoch 就等于重复触发同样的解码压力与状态累积

因此即使第一轮能跑，第二轮、第三轮更容易触发"不可回收状态累积到阈值"的崩溃。

5. 我们最终采用的方法：原理是什么？它到底释放了什么？

5.1 `gr00t/utils/video.py`：显式释放与显式关闭（对象级）

周期性 GC ：每 N 次视频加载执行 gc.collect()
显式 del 临时变量：避免 Python 引用链延长对象生命周期
对 torchvision_av 特别处理：显式关闭 VideoReader / container

原理：

Python 的 GC/引用计数不是实时/全局一致的，尤其当某些对象形成环或在异常路径上残留引用时
显式 del + 周期性 GC 能显著减少"Python 层引用残留导致的资源迟迟不释放"
对于"封装 C++ 资源的对象"，显式 close 比等待析构可靠得多（析构时机不确定）

这一层主要解决的是：Python 层确实存在能回收但没回收干净的对象。

5.2 `gr00t/data/dataset.py`：释放 DataFrame / 轨迹对象（数据缓存级）

在 get_trajectory_data() 前释放旧 DataFrame，并在 __getitem__ 加周期性 GC。

原理：

数据集往往不仅仅是视频，还会读取轨迹元数据（DataFrame/np arrays）
DataFrame 有自己的内存结构，且引用链容易被 Dataset 缓存、闭包或 collate 函数间接持有
在"加载新轨迹前显式释放旧轨迹"可以避免 Dataset 层面形成"隐式缓存堆积"

这一层解决的是：不是只有视频解码会涨，元数据也会悄悄堆起来。

5.3 `trainer.py`：训练循环里做周期性清理（运行时回收点）

加了 MemoryCleanupCallback，在：每 N 步保存 checkpoint 记录日志

执行 gc.collect()，并在 checkpoint 时 torch.cuda.empty_cache()。

原理：

训练循环里存在很多"阶段性峰值"：
- 保存 checkpoint 会触发额外的 tensor/materialization
- logging/metrics 可能持有历史 buffer
在这些"天然的边界点"做清理比在任意时刻更安全（减少打断关键路径）

empty_cache() 并不减少 PyTorch 的"已占用显存"，但会把缓存归还给 allocator 以便后续复用，从而降低显存碎片风险。

这一层是：把回收点绑定到训练节奏，避免小问题积累成大问题。

5.4 `gr00t_finetune.py`：最关键的稳定器------禁用 persistent_workers（进程级）

新增 dataloader_persistent_workers=False 并降低 prefetch_factor。

这是整个方案里最"决定性"的一层。

原理：

PyAV/FFmpeg 这类底层库的内存/状态，很多时候不服从 Python 对象生命周期
但它一定服从 进程生命周期：进程退出，OS 会回收整个地址空间、mmap、arena、线程、FD
禁用 persistent_workers 可以让 worker 进程在 epoch（或生命周期边界）重启，从而：
- 把"不可回收的状态累积"硬截断
- 把"泄漏/碎片"限制在一个 worker 生命周期内，而不是无限增长

降低 prefetch_factor 的作用也很直接：

prefetch 越大，worker 同时持有的样本越多
每个样本都可能包含视频帧/大张量
降低 prefetch 能减少"瞬时峰值内存"与"多 worker 同步峰值"

这一层解决的是：就算底层库不释放，我们也让它没有机会无限累积。

6. 这套方法的本质：把"不可控增长"变成"可控上界"

总结一下，这套方案并不是"让 FFmpeg 学会正确 free"，而是：

在对象级：确保 Python 侧不拖后腿（显式 close/del + 周期性 GC）
在数据级：清掉 Dataset/轨迹层面可能的隐式缓存
在进程级：用 worker 生命周期把底层库状态强行截断（最关键）

最终把系统从：

长期运行会累积到某个随机阈值然后崩

变成：

每个周期都能回到一个已知的内存基线，长期可运行

Cannot allocate memory——训练时视频解码为什么会内存越跑越大

1. av 到底是什么？它在训练里负责哪一段

2. 为什么会出现"内存泄漏"或"内存持续增长"？

2.1 进程的虚拟内存（VIRT）会越来越大：mmap 与地址空间膨胀

2.2 缓冲池与 arena 不会立刻归还给 OS：删了对象，但内存不降

2.3 视频解码不是纯函数：上下文、状态、缓存会跨样本累积

3. 为什么 DDP / DataLoader 下，"删掉解码过的视频"达不到预期释放效果？

3.1 DataLoader 多 worker = 多进程；父进程删对象不影响子进程

3.2 persistent_workers 会让"有问题的 worker 一直活着"，导致状态累积

3.3 DDP = 多个 rank = 多份进程状态；任何一个 rank 先挂都会"拖全场下水"

4. 为什么"每个 epoch 都会重新解码一遍"，问题会越来越明显？

5. 我们最终采用的方法：原理是什么？它到底释放了什么？

5.1 gr00t/utils/video.py：显式释放与显式关闭（对象级）

5.2 gr00t/data/dataset.py：释放 DataFrame / 轨迹对象（数据缓存级）

5.3 trainer.py：训练循环里做周期性清理（运行时回收点）

5.4 gr00t_finetune.py：最关键的稳定器------禁用 persistent_workers（进程级）

6. 这套方法的本质：把"不可控增长"变成"可控上界"

1. `av` 到底是什么？它在训练里负责哪一段

5.1 `gr00t/utils/video.py`：显式释放与显式关闭（对象级）

5.2 `gr00t/data/dataset.py`：释放 DataFrame / 轨迹对象（数据缓存级）

5.3 `trainer.py`：训练循环里做周期性清理（运行时回收点）

5.4 `gr00t_finetune.py`：最关键的稳定器------禁用 persistent_workers（进程级）