如何用 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在 NeMoNVIDIA 的 Speech AI 框架中使用 Lhotse 动态分桶dynamic bucketing加载数据时不同时长桶bucket对显存的需求差异很大。只写一个全局batch_duration往往会让短时长桶的 GPU 利用率偏低而长时长桶又可能在训练若干步之后才触发 CUDA OOM。NeMo 提供了 OOMptimizer给定一个模型、优化器和一份桶1D 或 2D列表它通过二分搜索实测每个桶能装下的最大批量大小几分钟后输出一份bucket_duration_binsbucket_batch_size配置供训练配置直接使用。本文的操作路径是用estimate_duration_bins.py或 2D 版脚本得到数据的分桶边界用 scripts/speech_recognition/oomptimizer.py 实测每个桶的最大批量把输出的 profile 写进训练的model.train_ds配置启用bucket_batch_size分桶加载若训练中仍 OOM用更低的--memory-fraction重新估算。前置条件模型、桶和 GPUOOMptimizer 依赖三个输入缺一不可一个 NeMo 模型。可以是预训练模型名如nvidia/parakeet-tdt-0.6b-v3也可以由--module-name--config-path从配置实例化一个未训练的模型。模型类必须定义oomptimizer_schema属性脚本会检查该属性缺失时直接报错退出提示 doesnt 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_buckets30 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_tpstoken-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输出/输入序列长度比用于推算各桶的最大输出长度ASRaudio→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: nulluse_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_samplerfalse必须设置——Lhotse 自己处理分布式采样不关会导致跨 rank 静默重复 batch使用 tarred/Shar 数据时 dataloader 是无限的必须设置trainer.max_steps和trainer.limit_train_batchesval_check_interval否则训练永远完不成第一个 epochnum_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_duration1100、quadratic_duration30、num_buckets30是针对 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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 如何用 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在 NeMoNVIDIA 的 Speech AI 框架中使用 Lhotse 动态分桶dynamic bucketing加载数据时不同时长桶bucket对显存的需求差异很大。只写一个全局batch_duration往往会让短时长桶的 GPU 利用率偏低而长时长桶又可能在训练若干步之后才触发 CUDA OOM。NeMo 提供了 OOMptimizer给定一个模型、优化器和一份桶1D 或 2D列表它通过二分搜索实测每个桶能装下的最大批量大小几分钟后输出一份bucket_duration_binsbucket_batch_size配置供训练配置直接使用。本文的操作路径是用estimate_duration_bins.py或 2D 版脚本得到数据的分桶边界用 scripts/speech_recognition/oomptimizer.py 实测每个桶的最大批量把输出的 profile 写进训练的model.train_ds配置启用bucket_batch_size分桶加载若训练中仍 OOM用更低的--memory-fraction重新估算。前置条件模型、桶和 GPUOOMptimizer 依赖三个输入缺一不可一个 NeMo 模型。可以是预训练模型名如nvidia/parakeet-tdt-0.6b-v3也可以由--module-name--config-path从配置实例化一个未训练的模型。模型类必须定义oomptimizer_schema属性脚本会检查该属性缺失时直接报错退出提示 doesnt 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_buckets30 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_tpstoken-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输出/输入序列长度比用于推算各桶的最大输出长度ASRaudio→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: nulluse_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_samplerfalse必须设置——Lhotse 自己处理分布式采样不关会导致跨 rank 静默重复 batch使用 tarred/Shar 数据时 dataloader 是无限的必须设置trainer.max_steps和trainer.limit_train_batchesval_check_interval否则训练永远完不成第一个 epochnum_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_duration1100、quadratic_duration30、num_buckets30是针对 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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考