如何用 OOMptimizer 为 NeMo Lhotse 数据加载估算批量大小避免训练 OOM

如何用 OOMptimizer 为 NeMo Lhotse 数据加载估算批量大小避免训练 OOM

【免费下载链接】SpeechA scalable generative AI framework built for researchers and developers working on Large Language Models, Multimodal, and Speech AI (Automatic Speech Recognition and Text-to-Speech)项目地址: https://gitcode.com/GitHub_Trending/nem/Speech

在 NeMo（NVIDIA 的 Speech AI 框架）中使用 Lhotse 动态分桶（dynamic bucketing）加载数据时，不同时长桶（bucket）对显存的需求差异很大。只写一个全局batch_duration往往会让短时长桶的 GPU 利用率偏低，而长时长桶又可能在训练若干步之后才触发 CUDA OOM。NeMo 提供了 OOMptimizer：给定一个模型、优化器和一份桶（1D 或 2D）列表，它通过二分搜索实测每个桶能装下的最大批量大小，几分钟后输出一份bucket_duration_bins+bucket_batch_size配置，供训练配置直接使用。

本文的操作路径是：

1. 用estimate_duration_bins.py（或 2D 版脚本）得到数据的分桶边界；
2. 用 scripts/speech_recognition/oomptimizer.py 实测每个桶的最大批量；
3. 把输出的 profile 写进训练的model.train_ds配置，启用bucket_batch_size分桶加载；
4. 若训练中仍 OOM，用更低的--memory-fraction重新估算。

前置条件：模型、桶和 GPU

OOMptimizer 依赖三个输入，缺一不可：

• 一个 NeMo 模型。可以是预训练模型名（如nvidia/parakeet-tdt-0.6b-v3），也可以由--module-name+--config-path从配置实例化一个未训练的模型。模型类必须定义oomptimizer_schema属性，脚本会检查该属性，缺失时直接报错退出（提示 "doesn't seem to support OOMptimizer"）。
• 一份分桶边界（--buckets），来自下一步的估算脚本。
• 一张 CUDA GPU。脚本会调用torch.cuda.set_per_process_memory_fraction限制显存使用，并用 NCCL 后端，因此必须在有 GPU 的环境中运行。

注意仓库中存在两个同名脚本：scripts/speech_recognition/oomptimizer.py 和 scripts/speechlm2/oomptimizer.py。前者服务 ASR 模型（支持--pretrained-name与--module-name/--config-path两种模式），后者面向 SALM/BESTOW 等 Speech LM 模型，且断言--pretrained-name不可用、要求设置LOCAL_RANK。本文主路径按文档 docs/source/dataloaders.rst 中的 ASR 流程，使用前者；Speech LM 场景请改读后者源码与说明。

第一步：估算分桶边界 bucket_duration_bins

先用数据分布确定桶的边界。对 NeMo manifest 或input_cfgYAML 运行：

python scripts/speech_recognition/estimate_duration_bins.py -b 30 manifest.json

多数据集可以直接传input_cfg.yaml，或传带权重的 manifest 列表，例如-b 30 [[manifest.json,0.7],[other.json,0.3]]。文档示例输出（示例结果，数值以你的数据为准）：

Use the following options in your config: num_buckets=30 bucket_duration_bins=[1.78,2.34,2.69,... <other diagnostic information about the dataset>

-b 30表示 30 个桶。如果你训练的是 attention encoder-decoder 这类对输出序列长度敏感的模型，用 2D 版本按输入时长 + 输出 token 数两个维度分层：

python scripts/speech_recognition/estimate_duration_bins_2d.py \ --tokenizer path/to/tokenizer.model \ --buckets 30 \ --sub-buckets 5 \ input_cfg.yaml

该脚本输出的bucket_duration_bins是嵌套列表，每个 bin 是[最大时长, 最大token数]对；文档同时给出可选的max_tps（token-per-second）过滤阈值，用于丢弃长转录离群样本，可一并写入配置。把输出里的 bins 列表留好，作为下一步--buckets的参数。

第二步：运行 OOMptimizer 实测每桶最大批量

有两种调用方式，二选一。

方式 A：用预训练模型名（微调场景）。模型名可替换为你要微调的模型，下面是文档给出的示例命令：

python scripts/speech_recognition/oomptimizer.py \ --pretrained-name nvidia/canary-1b \ --buckets '[2.0,3.1,5.6,6.6,...]'

方式 B：用模块名 + 配置文件（从零训练或自定义架构）。--module-name是 NeMo 模块的完整路径，--config-path是对应训练配置，两者必须同时提供。仓库自带的示例配置 examples/asr/conf/speech_multitask/fast-conformer_aed.yaml 的train_ds已启用use_lhotse: true，可直接作为演示配置：

python scripts/speech_recognition/oomptimizer.py \ --config-path examples/asr/conf/speech_multitask/fast-conformer_aed.yaml \ --module-name nemo.collections.asr.models.EncDecMultiTaskModel \ --buckets '[[3.975,30],[3.975,48],[4.97,37],...]'

--buckets支持 1D 列表（[2.0, 3.1, 5.6, ...]）或 2D 嵌套列表（[[3.975,30], ...]，来自 2D 估算脚本的输出）。

常用参数（均有默认值）：

• --threshold（-t，默认 0.05）：搜索停止条件，取值为 (min_oom_batch_size - max_ok_batch_size) / min_oom_batch_size，越小越精确；
• --start-batch-size（-s，默认 32）：搜索起始批量；
• --ratio（-r，默认 12）：输出/输入序列长度比，用于推算各桶的最大输出长度，ASR（audio→text）即 tokens per second，提供 2D 桶时该参数被忽略；
• --memory-fraction（-f，默认 0.9）：把该进程可用 CUDA 显存限制为设备总显存的该比例，默认留 5% 给训练脚本中 OOMptimizer 无法预见的额外显存开销；
• --dtype（默认bfloat16）：配合 autocast 使用的浮点精度；
• --ddp/--no-ddp（默认开）：是否额外存一份模型权重来模拟 DDP 的显存占用。

如何判断搜索完成

脚本逐桶（从最长序列到最短序列）做二分搜索，每一步日志打印 CUDA 显存用量，并以OK!（绿色）或OOM!（黄色）标记当前批量是否可行；单个桶收敛后打印形如=> Optimal setting for bucket=... is max_batch_size=...的行。全部桶完成后输出最终 profile（以下为文档示例，数值以你的模型与 GPU 为准）：

The final profile is: bucket_duration_bins=[[3.975,30],[3.975,48],...] bucket_batch_size=[352,308,280,...]

看到这两行即表示搜索收敛，把列表原样复制到训练配置即可。1D 桶还会经过一个 "Bucket merging stage"：相邻桶若算出相同批量会被合并，所以最终bucket_batch_size的元素数可能少于输入桶数。

第三步：把 profile 写入训练配置

文档建议把结果选项放在训练配置的model.train_ds命名空间下，启用分桶加载。结合 docs/source/dataloaders.rst 中 Lhotse 分桶的写法：

model: train_ds: use_lhotse: true use_bucketing: true num_buckets: 30 bucket_duration_bins: [[3.975,30],[3.975,48],...] # OOMptimizer 输出 bucket_batch_size: [352,308,280,...] # OOMptimizer 输出 batch_size: null num_workers: 4

其中bucket_batch_size是 OOMptimizer 的输出、每个桶的批量上限；batch_size: null+use_bucketing: true是让分桶批量生效的写法。文档明确：bucket_duration_bins为[时长, token]嵌套列表且同时设置bucket_batch_size时，数据加载器会自动启用 2D 分桶（FixedBucketBatchSizeConstraint2D），每个桶使用固定批量。若不使用 OOMptimizer，替代手段是batch_tokens（音频场景为batch_duration）加quadratic_duration/quadratic_factor二次惩罚，文档将其描述为启发式方法，效率不如 OOMptimizer。

同时确认 docs/source/dataloaders.rst 列出的 Lhotse 数据加载必备项，避免与批量设置无关的失败：

• trainer.use_distributed_sampler=false必须设置——Lhotse 自己处理分布式采样，不关会导致跨 rank 静默重复 batch；
• 使用 tarred/Shar 数据时 dataloader 是无限的，必须设置trainer.max_steps和trainer.limit_train_batches（val_check_interval），否则训练永远完不成第一个 "epoch"；
• num_cuts_for_bins_estimate控制训练开始前采样多少条样本估算桶边界；若bucket_buffer_size相对batch_duration偏小，动态分桶采样器会发出警告，需调大。

训练仍 OOM 时：降低 memory-fraction 重新估算

文档说明训练脚本可能使用 OOMptimizer 无法预见的额外 GPU 显存。如果应用 profile 后训练中仍出现 OutOfMemoryError，把估算时的显存上限调低后重跑第二步，例如：

python scripts/speech_recognition/oomptimizer.py \ --config-path examples/asr/conf/speech_multitask/fast-conformer_aed.yaml \ --module-name nemo.collections.asr.models.EncDecMultiTaskModel \ --buckets '[[3.975,30],[3.975,48],[4.97,37],...]' \ --memory-fraction 0.75

然后用新的 profile 替换model.train_ds中的两个列表。

限制与边界

• 模型类缺少oomptimizer_schema时脚本会直接报错退出，这类模型不能用 OOMptimizer，只能退回batch_duration+ 二次惩罚的启发式设置。
• --pretrained-name与--module-name/--config-path两种模式互斥，二选一。
• 文档给出的默认值（如batch_duration=1100、quadratic_duration=30、num_buckets=30）是针对 hybrid RNN-T + CTC ASR 模型、32GB GPU、以 15 秒左右语句为主的数据分布的起点值，换模型或数据分布时需要自行校准，而 OOMptimizer 的作用正是把这一校准过程实测化。
• 搜索本身会长时间占用一张 GPU，运行前确认该卡没有其他任务。

【免费下载链接】SpeechA scalable generative AI framework built for researchers and developers working on Large Language Models, Multimodal, and Speech AI (Automatic Speech Recognition and Text-to-Speech)项目地址: https://gitcode.com/GitHub_Trending/nem/Speech