NeMo 核心模型框架全解析：预训练模型、PyTorch Lightning 训练、Hydra 配置与优化调度

NeMo 核心模型框架全解析预训练模型、PyTorch Lightning 训练、Hydra 配置与优化调度【免费下载链接】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本指南以 docs/source/core/core.rst 为主体骨架系统讲解 NeMo本仓库 NVIDIA NeMo Speech中所有模型共享的通用能力如何实例化预训练模型、如何借助 PyTorch Lightning 组织训练循环、如何通过 Hydra 的 YAML/CLI/Dataclass 三种接口配置训练、如何配置优化器与学习率调度器以及.nemo文件的保存、恢复、Artifact 注册、Hugging Face Hub 发布与嵌套模型机制。读完本文你将掌握一套在任何 ASR、TTS、Audio、SpeechLM2 模型上都能直接复用的模型生命周期管理方法论并能在本仓库的 examples 中找到每个模型对应的真实配置与训练脚本。NeMo 模型包含什么一个自包含的训练单元NeMo 的设计目标是让训练并复现一个对话式 AI 模型所需的一切都封装在一个模型对象中。从 docs/source/core/core.rst 的定义看一个 NeMo 模型包含神经网络架构neural network architectures数据集与数据加载器datasets/data loaders数据预处理/后处理data preprocessing/postprocessing数据增强器data augmentors优化器与学习率调度器optimizers and schedulersTokenizertokenizers语言模型language modelsNeMo 使用 Hydra 同时配置 NeMo 模型本身与 PyTorch Lightning Trainer。三者NeMo PyTorch Lightning Hydra组合的结果是所有 NeMo 模型拥有统一的外观与使用方式same look and feel同时与 PyTorch 生态完全兼容——你可以把 NeMo 模型无缝嵌入已有的 PyTorch 工作流。本仓库中每个 NeMo 模型都在 examples 下提供了示例配置文件与训练脚本例如 ASR 领域的 speech_to_text_ctc.py 及其配套配置目录 examples/asr/conf/conformer这是实践本文所有概念的最佳起点。使用预训练模型from_pretrained 与 list_available_modelsNeMo 为 ASR、TTS、Audio、SpeechLM2 等每个集合都提供了大量预训练模型。每个预训练 NeMo 模型都可以通过from_pretrained()方法下载并使用。例如实例化一个 Parakeet-TDT 0.6B ASR 模型只需import nemo.collections.asr as nemo_asr model nemo_asr.models.ASRModel.from_pretrained(model_namenvidia/parakeet-tdt-0.6b-v2)要查看某个 NeMo 模型类下所有可用的预训练模型使用list_available_models()方法nemo_asr.models.EncDecCTCModel.list_available_models()不同集合的预训练模型清单及使用方式可参考对应集合文档ASR 集合文档、TTS 集合文档。本仓库的 tests/e2e_nightly 下有大量 L2 级端到端脚本如 L2_Model_Support_nvidia__parakeet_tdt_0_6b_v2.sh展示了各类预训练模型的下载、推理与评测调用链可作为from_pretrained实际用法的参考验证。训练基于 PyTorch Lightning 的模型生命周期NeMo 借助 PyTorch Lightning 进行模型训练从而把对话式 AI 的领域代码与 PyTorch 训练样板代码解耦。用户只需关注领域问题ASR、TTS 等即可自动获得多 GPU / 多节点训练multi-GPU/multi-node混合精度mixed precision模型检查点model checkpointing日志logging早停early stopping以及其他 Lightning 提供的训练能力Lightning API 的两大核心是LightningModule与Trainer。LightningModule每个 NeMo 模型都是一个 nn.Module每个 NeMo 模型都是一个LightningModule而LightningModule本身是nn.Module。这意味着 NeMo 模型天然兼容 PyTorch 生态可以直接嵌入现有 PyTorch 工作流。创建 NeMo 模型与常规 PyTorch 流程一致先初始化模型架构再定义 forward 前向传播。文档以 NLP 文本分类模型为例class TextClassificationModel(NLPModel, Exportable): ... def __init__(self, cfg: DictConfig, trainer: Trainer None): Initializes the BERTTextClassifier model. ... super().__init__(cfgcfg, trainertrainer) # instantiate a BERT based encoder self.bert_model get_lm_model( config_filecfg.language_model.config_file, config_dictcfg.language_model.config, vocab_filecfg.tokenizer.vocab_file, trainertrainer, cfgcfg, ) # instantiate the FFN for classification self.classifier SequenceClassifier( hidden_sizeself.bert_model.config.hidden_size, num_classescfg.dataset.num_classes, num_layerscfg.classifier_head.num_output_layers, activationrelu, log_softmaxFalse, dropoutcfg.classifier_head.fc_dropout, use_transformer_initTrue, idx_conditioned_on0, )forward的定义方式与普通 PyTorch 完全一致无需为 Lightning 做任何特殊修改def forward(self, input_ids, token_type_ids, attention_mask): No special modification required for Lightning, define it as you normally would in the nn.Module in vanilla PyTorch. hidden_states self.bert_model( input_idsinput_ids, token_type_idstoken_type_ids, attention_maskattention_mask ) logits self.classifier(hidden_stateshidden_states) return logitsLightningModule将 PyTorch 代码组织得井井有条使所有 NeMo 模型保持统一风格。训练逻辑位于training_step中def training_step(self, batch, batch_idx): Lightning calls this inside the training loop with the data from the training dataloader passed in as batch. # forward pass input_ids, input_type_ids, input_mask, labels batch logits self.forward(input_idsinput_ids, token_type_idsinput_type_ids, attention_maskinput_mask) train_loss self.loss(logitslogits, labelslabels) lr self._optimizer.param_groups[0][lr] self.log(train_loss, train_loss) self.log(lr, lr, prog_barTrue) return { loss: train_loss, lr: lr, }验证逻辑位于validation_stepdef validation_step(self, batch, batch_idx): Lightning calls this inside the validation loop with the data from the validation dataloader passed in as batch. if self.testing: prefix test else: prefix val input_ids, input_type_ids, input_mask, labels batch logits self.forward(input_idsinput_ids, token_type_idsinput_type_ids, attention_maskinput_mask) val_loss self.loss(logitslogits, labelslabels) preds torch.argmax(logits, axis-1) tp, fn, fp, _ self.classification_report(preds, labels) return {val_loss: val_loss, tp: tp, fn: fn, fp: fp}之后训练所需的样板代码全部由 PyTorch Lightning 处理。训练的任何方面几乎都可以通过 Lightning 的 hooks、plugins、callbacks 或方法覆写来定制。Trainer训练与测试的统一入口由于每个 NeMo 模型都是LightningModule因此可以自动利用 PyTorch Lightning 的Trainer。本仓库 examples 中的每个训练脚本都通过Trainer对象来 fit 模型。标准流程是先根据模型配置实例化 Trainer再构造模型最后调用.fit或.test# We first instantiate the trainer based on the model configuration. # See the model configuration documentation for details. trainer pl.Trainer(**cfg.trainer) # Then pass the model configuration and trainer object into the NeMo model model TextClassificationModel(cfg.model, trainertrainer) # Now we can train with by calling .fit trainer.fit(model) # Or we can run the test loop on test data by calling trainer.test(modelmodel)所有 Trainer 参数devices、accelerator、max_epochs、precision 等都可以在 NeMo 配置中直接设置。配置Hydra 的三种接口Hydra 是一个开源 Python 配置框架用于简化需要整合众多软件库的复杂应用的配置——对话式 AI 模型训练正是这类应用的典型代表。训练一个对话式 AI 模型需要配置神经网络架构训练与优化算法数据预处理/后处理数据增强实验日志/可视化模型检查点使用 HydraNeMo 的一切配置都可以通过三种接口完成命令行CLI配置文件YAMLDataclassPython三者优先级为CLI YAML Dataclass。YAML 配置统一的 trainer / exp_manager / model 三层结构NeMo 为 examples 下的所有训练脚本提供了 YAML 配置文件。每个 NeMo 示例 YAML 都具有相同的底层配置结构trainerexp_managermodel其中model配置始终包含train_ds、validation_ds、test_ds和optim模型架构部分则因领域而异如 ASR 的 encoder/decoder。一个典型的 NeMo 配置文件如下# PyTorch Lightning Trainer configuration # any argument of the Trainer object can be set here trainer: devices: 1 # number of gpus per node accelerator: gpu num_nodes: 1 # number of nodes max_epochs: 10 # how many training epochs to run val_check_interval: 1.0 # run validation after every epoch # Experiment logging configuration exp_manager: exp_dir: /path/to/my/nemo/experiments name: name_of_my_experiment create_tensorboard_logger: True create_wandb_logger: True # Model configuration # model network architecture, train/val/test datasets, data augmentation, and optimization model: train_ds: manifest_filepath: /path/to/my/train/manifest.json batch_size: 256 shuffle: True validation_ds: manifest_filepath: /path/to/my/validation/manifest.json batch_size: 32 shuffle: False test_ds: manifest_filepath: /path/to/my/test/manifest.json batch_size: 32 shuffle: False optim: name: novograd lr: .01 betas: [0.8, 0.5] weight_decay: 0.001 # network architecture can vary greatly depending on the domain encoder: ... decoder: ...在真实仓库中可以找到大量遵循这一结构的配置文件例如 examples/asr/conf/conformer/conformer_ctc_bpe.yamlConformer-CTC 大模型约 120M 参数的完整训练配置以及 examples/asr/conf 下 60 余个 ASR 配置。exp_manager的完整参数语义可进一步阅读 nemo/utils/exp_manager.py。CLI 覆盖实验调参的最快路径借助 Hydra模型训练的每个方面都可以从命令行修改非常适合在计算集群上批量跑实验或在开发期快速测试参数。Hydra 使用操作符设置参数python examples/asr/asr_ctc/speech_to_text_ctc.py \ model.train_ds.manifest_filepath/path/to/my/train/manifest.json \ model.validation_ds.manifest_filepath/path/to/my/validation/manifest.json \ trainer.devices2 \ trainer.acceleratorgpu \ trainer.max_epochs50使用操作符从命令行新增参数例如临时开启 Lightning 的fast_dev_run快速冒烟测试python examples/asr/asr_ctc/speech_to_text_ctc.py \ model.train_ds.manifest_filepath/path/to/my/train/manifest.json \ model.validation_ds.manifest_filepath/path/to/my/validation/manifest.json \ trainer.devices2 \ trainer.acceleratorgpu \ trainer.max_epochs50 \ trainer.fast_dev_runtrue使用~操作符删除配置项例如去掉 test_ds避免测试集参与任何流程python examples/asr/asr_ctc/speech_to_text_ctc.py \ model.train_ds.manifest_filepath/path/to/my/train/manifest.json \ model.validation_ds.manifest_filepath/path/to/my/validation/manifest.json \ ~model.test_ds \ trainer.devices2 \ trainer.acceleratorgpu \ trainer.max_epochs50 \ trainer.fast_dev_runtrue使用--config-path和--config-name指定配置文件python examples/asr/asr_ctc/speech_to_text_ctc.py \ --config-pathconf/conformer \ --config-nameconformer_ctc_bpe \ model.train_ds.manifest_filepath/path/to/my/train/manifest.json \ model.validation_ds.manifest_filepath/path/to/my/validation/manifest.json \ ~model.test_ds \ trainer.devices2 \ trainer.acceleratorgpu \ trainer.max_epochs50 \ trainer.fast_dev_runtrue注意--config-pathconf/conformer是相对于脚本所在目录examples/asr/asr_ctc的路径对应仓库中的 examples/asr/conf/conformer。本仓库 tests/hydra 下的测试如 config1.yaml 与 my_app.py可用于验证 Hydra 参数解析与覆盖行为。Dataclass纯 Python 配置模型Dataclass 让 NeMo 可以将模型配置作为库的一部分随包发布同时支持纯 Python 方式配置 NeMo 模型。借助 Hydra 的 structured configsdataclass 可为对话式 AI 应用构建结构化配置。以 Attention Is All You Need 机器翻译模型为例其配置可以像普通 Python dataclass 一样实例化并修改from nemo.collections.nlp.models.machine_translation.mt_enc_dec_config import AAYNBaseConfig cfg AAYNBaseConfig() # modify the number of layers in the encoder cfg.encoder.num_layers 8 # modify the training batch size cfg.train_ds.tokens_in_batch 8192优化Optimizer 与 Learning Rate Scheduler优化器与学习率调度器在所有 NeMo 模型上都是可配置的并拥有独立的optim命名空间。下面是 Novograd 优化器 Cosine Annealing 学习率调度器的完整 YAML 示例optim: name: novograd lr: 0.01 # optimizer arguments betas: [0.8, 0.25] weight_decay: 0.001 # scheduler setup sched: name: CosineAnnealing # Optional arguments max_steps: -1 # computed at runtime or explicitly set here monitor: val_loss reduce_on_plateau: false # scheduler config override warmup_steps: 1000 warmup_ratio: null min_lr: 1e-9文档原示例中min_lr: 1e-9行末带有一个冒号实际配置时应写为min_lr: 1e-9。优化器同样可以从 CLI 配置python examples/asr/asr_ctc/speech_to_text_ctc.py \ --config-pathconf/conformer \ --config-nameconformer_ctc_bpe \ ... # train with the adam optimizer model.optimadam \ # change the learning rate model.optim.lr.0004 \ # modify betas model.optim.betas[.8, .5]本仓库 examples 下每个模型的配置中都包含可用的优化器与调度器配置。可用优化器清单optim.name对应优化器的全小写名称。要查看所有可用优化器运行from nemo.core.optim.optimizers import AVAILABLE_OPTIMIZERS for name, opt in AVAILABLE_OPTIMIZERS.items(): print(fname: {name}, opt: {opt})输出示例name: sgd opt: class torch.optim.sgd.SGD name: adam opt: class torch.optim.adam.Adam name: adamw opt: class torch.optim.adamw.AdamW name: adadelta opt: class torch.optim.adadelta.Adadelta name: adamax opt: class torch.optim.adamax.Adamax name: adagrad opt: class torch.optim.adagrad.Adagrad name: rmsprop opt: class torch.optim.rmsprop.RMSprop name: rprop opt: class torch.optim.rprop.Rprop name: novograd opt: class nemo.core.optim.novograd.Novograd从当前仓库的源码 nemo/core/optim/optimizers.py 可以看到内置注册表实际比文档示例更丰富还包含adafactorAdafactor与adanAdan等。此外源码第 4876 行展示了基于 Apex 的条件注册逻辑当环境中安装了 Apex 时会额外注册lambFusedLAMB、fused_adamFusedAdam、distributed_fused_adamMegatronDistributedFusedAdam、megatron_fused_adam等 CUDA 融合优化器未安装 Apex 时这些优化器不可用。这一点可以推断出AVAILABLE_OPTIMIZERS的实际内容随运行环境动态变化若某个优化器名称解析失败可检查 Apex 是否安装。优化器参数 Dataclass不同优化器的参数各不相同但lr是所有优化器都必需的。要查看某优化器的可用参数查看其对应的 dataclassfrom nemo.core.config.optimizers import NovogradParams print(NovogradParams())输出示例NovogradParams(lr???, betas(0.95, 0.98), eps1e-08, weight_decay0, grad_averagingFalse, amsgradFalse, lucFalse, luc_trust0.001, luc_eps1e-08)其中???表示lr是必填参数。对应的参数 dataclass 定义位于 nemo/core/config/optimizers.py例如SGDParamsmomentum、dampening、weight_decay、nesterov、AdamParamseps、weight_decay、amsgrad、AdamWParamsbetas、eps、weight_decay、amsgrad等Novograd优化器本身实现在 nemo/core/optim/novograd.py。注册自定义优化器要将新的优化器注册到 NeMo 中使用register_optimizer。其底层实现在 nemo/core/optim/optimizers.py若名称已存在于AVAILABLE_OPTIMIZERS会抛出ValueError(Cannot override pre-existing optimizers...)防止覆盖内置优化器注册成功后还会自动以{OptimizerName}_params的键名注册对应的参数 dataclass。随后可通过get_optimizer(name, **kwargs)nemo/core/optim/optimizers.py按名称解析并 partial 实例化优化器。可用学习率调度器清单学习率调度器可以在optim.sched命名空间下按需配置name对应调度器名称from nemo.core.optim.lr_scheduler import AVAILABLE_SCHEDULERS for name, opt in AVAILABLE_SCHEDULERS.items(): print(fname: {name}, schedule: {opt})输出示例name: WarmupPolicy, schedule: class nemo.core.optim.lr_scheduler.WarmupPolicy name: WarmupHoldPolicy, schedule: class nemo.core.optim.lr_scheduler.WarmupHoldPolicy name: SquareAnnealing, schedule: class nemo.core.optim.lr_scheduler.SquareAnnealing name: CosineAnnealing, schedule: class nemo.core.optim.lr_scheduler.CosineAnnealing name: NoamAnnealing, schedule: class nemo.core.optim.lr_scheduler.NoamAnnealing name: WarmupAnnealing, schedule: class nemo.core.optim.lr_scheduler.WarmupAnnealing name: InverseSquareRootAnnealing, schedule: class nemo.core.optim.lr_scheduler.InverseSquareRootAnnealing name: SquareRootAnnealing, schedule: class nemo.core.optim.lr_scheduler.SquareRootAnnealing name: PolynomialDecayAnnealing, schedule: class nemo.core.optim.lr_scheduler.PolynomialDecayAnnealing name: PolynomialHoldDecayAnnealing, schedule: class nemo.core.optim.lr_scheduler.PolynomialHoldDecayAnnealing name: StepLR, schedule: class torch.optim.lr_scheduler.StepLR name: ExponentialLR, schedule: class torch.optim.lr_scheduler.ExponentialLR name: ReduceLROnPlateau, schedule: class torch.optim.lr_scheduler.ReduceLROnPlateau name: CyclicLR, schedule: class torch.optim.lr_scheduler.CyclicLR从当前仓库源码 nemo/core/optim/lr_scheduler.py 可以看到注册表实际还包含NoamHoldAnnealing、WarmupHoldAnnealOneMinusSquareRoot、WarmupHoldAnnealLinear、T5InverseSquareRootAnnealing等更多调度器。该文件底部还定义了EPOCH_SCHEDULERSnemo/core/optim/lr_scheduler.py包含ExponentialLR与ReduceLROnPlateau从源码结构可以推断这两个按 epoch 步进的调度器在 NeMo 训练循环中走独立的处理路径。另外get_scheduler在解析时会检查max_steps是否在调度器签名中nemo/core/optim/lr_scheduler.py不需要该参数的调度器会自动丢弃max_steps这解释了为何在 YAML 中统一设置max_steps是安全的。调度器参数 Dataclass查看某个调度器的可用参数同样查看其 dataclassfrom nemo.core.config.schedulers import CosineAnnealingParams print(CosineAnnealingParams())输出示例CosineAnnealingParams(last_epoch-1, warmup_stepsNone, warmup_ratioNone, min_lr0.0)参数 dataclass 的继承结构定义在 nemo/core/config/schedulers.py基类WarmupSchedulerParams提供max_steps、warmup_steps、warmup_ratioWarmupHoldSchedulerParams增加hold_steps、hold_ratio、min_lrWarmupAnnealingHoldSchedulerParams提供constant_steps、constant_ratio、min_lr。不同调度器的min_lr默认值不同例如SquareAnnealingParams默认为1e-5而CosineAnnealingParams、NoamAnnealingParams默认为0.0。注册自定义调度器要注册新的学习率调度器使用register_scheduler。底层实现位于 nemo/core/optim/lr_scheduler.py逻辑与优化器注册一致名称冲突会抛出ValueError成功后自动注册对应的参数 dataclass。保存与恢复.nemo 文件所有 NeMo 模型都自带.save_to与.restore_from方法其实现位于 nemo/core/classes/modelPT.py 附近。保存模型model.save_to(/path/to/model.nemo)使用训练好的模型所需的一切都会被打包进.nemo文件。例如在 NLP 领域.nemo文件包含必要的 tokenizer 模型和/或词汇表文件等。.nemo文件本质上就是一个类似.tar的归档文件。恢复模型# Here, you should usually use the class of the model, or simply use ModelPT.restore_from() for simplicity. model.restore_from(/path/to/model.nemo)当使用 PyTorch Lightning Trainer 时会产生 Lightning checkpointNeMo 主要用它来自动恢复auto-resume训练。由于 NeMo 模型是LightningModule因此也提供load_from_checkpoint方法。需要特别注意的是load_from_checkpoint并非对所有模型都能开箱即用因为部分模型在恢复时还需要比 checkpoint 更多的 artifact如 tokenizer 文件这类模型如要使用该方法用户需要自行覆写load_from_checkpoint。因此强烈推荐使用restore_from加载 NeMo 模型。修改配置后恢复有时需要在恢复前修改模型或其子组件的配置常见场景包括模型内部配置因弃用、版本升级或支持新特性而需要更新。只要模型参数与原始配置保持一致这些参数就可以安全地恢复。.nemo文件内保留了模型的内部配置恢复时会使用该配置并且可以在恢复前更新它# When restoring a model, you should generally use the class of the model # Obtain the config (as an OmegaConf object) config model_class.restore_from(/path/to/model.nemo, return_configTrue) # OR config model_class.from_pretrained(name_of_the_model, return_configTrue) # Modify the config as needed config.x.y z # Restore the model from the updated config model model_class.restore_from(/path/to/model.nemo, override_config_pathconfig) # OR model model_class.from_pretrained(name_of_the_model, override_config_pathconfig)注册 Artifact让 .nemo 文件自包含恢复对话式 AI 模型之所以复杂是因为仅有权重不够还需要额外信息才能使用模型。NeMo 模型可以通过.register_artifact把额外 artifact 保存进.nemo文件使用.restore_from或.from_pretrained恢复时所有已注册的 artifact 会被自动还原。以需要训练好的 tokenizer 模型的 NLP 模型为例tokenizer 模型文件可以通过下面的方式自动加入.nemo文件self.encoder_tokenizer get_nmt_tokenizer( ... tokenizer_modelself.register_artifact(config_pathencoder_tokenizer.tokenizer_model, src/path/to/tokenizer.model, verify_src_existsTrue), )从源码 nemo/core/classes/modelPT.py 可以了解.register_artifact的完整行为默认情况下它总是返回一个路径若模型正在从.nemo文件恢复返回的是该 artifact 在.nemo文件内的路径否则返回用户指定的本地路径。src为None或空字符串时不做任何事并原样返回。若src以nemo_file:unique_artifact_name开头则.nemo会被解包到临时目录并返回实际存在的路径。若src指向.nemo文件会抛出NeMoBaseException提示应改用register_nemo_submodule处理嵌套模型。重复注册同一config_path会打印警告。config_path是 artifact 的键key通常但不必须对应模型配置项。打包进.nemo的模型配置会按config_path更新例如上例中模型配置会变为encoder_tokenizer: ... tokenizer_model: nemo:4978b28103264263a03439aaa6560e5e_tokenizer.modelsrc是 artifact 的路径打包进.nemo文件时使用其 base-name。为防止不同 artifact 同名例如多个 tokenizer 都叫tokenizer.model产生冲突每个 artifact 都会在 base-name 前附加一个哈希前缀最终.nemo文件内形如4978b28103264263a03439aaa6560e5e_tokenizer.model若verify_src_exists设为False则 artifact 是可选的当src找不到时.register_artifact返回None。推送模型到 Hugging Face HubNeMo 模型可以通过push_to_hf_hub方法推送到 Hugging Face Hub该方法位于 nemo/core/classes/mixins/hf_io_mixin.py 的HuggingFaceFileIO类中。它执行与save_to()相同的动作然后把模型上传到 Hugging Face Hub并额外提供pack_nemo_file参数用于选择上传整个 NeMo 文件还是仅上传.nemo文件。对于参数量巨大的 LLM单个 NeMo 文件可能超过 Hugging Face Hub 的上传大小限制此时pack_nemo_fileFalse可以拆分为多个文件上传。上传模型token HF TOKEN or None pack_nemo_file True # False will upload multiple files that comprise the NeMo file onto HF Hub; Generally useful for LLMs model.push_to_hf_hub( repo_idrepo_id, pack_nemo_filepack_nemo_file, tokentoken, )使用自定义模型卡模板可以通过generate_model_card覆写默认模型卡# Override the default model card template Your own custom template # {model_name} kwargs {model_name: ABC, repo_id: nvidia/ABC_XYZ} model_card model.generate_model_card(templatetemplate, template_kwargskwargs, typehf) model.push_to_hf_hub( repo_idrepo_id, tokentoken, model_cardmodel_card )也可以编写自定义模型卡类# Write your own model card class class MyModelCard: def __init__(self, model_name): self.model_name model_name def __repr__(self): template This is the {model_name} model.format(model_nameself.model_name) return template model.push_to_hf_hub( repo_idrepo_id, tokentoken, model_cardMyModelCard(ABC) )本仓库 tutorials/Publish_NeMo_Model_On_Hugging_Face_Hub.ipynb 提供了从保存到发布的完整交互式演练。嵌套 NeMo 模型模型套模型某些场景下需要在 NeMo 模型内部使用其他 NeMo 模型。典型例子是把语言模型集成进 ASR 模型用于解码阶段提升准确率。有三种方式在父模型内实例化子模型直接使用子配置subconfig使用.nemocheckpoint 路径加载子模型使用预训练的 NeMo 模型注册子模型需调用父模型的register_nemo_submodule方法该方法会把子模型挂到指定的模型属性上序列化时会正确处理子模型的 artifact并把子模型的配置存入父模型的config_field。其实现在 nemo/core/classes/modelPT.py 附近并配合has_native_or_submodules_artifacts、has_nemo_submodules等辅助方法nemo/core/classes/modelPT.py在保存/恢复时统一处理子模块及其 artifact。完整示例from nemo.core.classes import ModelPT class ChildModel(ModelPT): ... # implement necessary methods class ParentModel(ModelPT): def __init__(self, cfg, trainerNone): super().__init__(cfgcfg, trainertrainer) # optionally annotate type for IDE autocompletion and type checking self.child_model: Optional[ChildModel] if cfg.get(child_model) is not None: # load directly from config # either if config provided initially, or automatically # after model restoration self.register_nemo_submodule( namechild_model, config_fieldchild_model, modelChildModel(self.cfg.child_model, trainertrainer), ) elif cfg.get(child_model_path) is not None: # load from .nemo model checkpoint # while saving, config will be automatically assigned/updated # in cfg.child_model self.register_nemo_submodule( namechild_model, config_fieldchild_model, modelChildModel.restore_from(self.cfg.child_model_path, trainertrainer), ) elif cfg.get(child_model_name) is not None: # load from pretrained model # while saving, config will be automatically assigned/updated # in cfg.child_model self.register_nemo_submodule( namechild_model, config_fieldchild_model, modelChildModel.from_pretrained(self.cfg.child_model_name, trainertrainer), ) else: self.child_model None性能剖析Nsys 与 CUDA 内存 ProfilingNeMo 提供两种性能剖析选项Nsys 与 CUDA 内存 profiling分别用于调试性能问题和内存问题如内存泄漏。Nsys 性能剖析在模型配置中加入以下选项启用 Nsys profilingnsys_profile: False start_step: 10 # Global batch to start profiling end_step: 10 # Global batch to end profiling ranks: [0] # Global rank IDs to profile gen_shape: False # Generate model and kernel details including input shapes示例中的缩进为文档原样展示实际应保持nsys_profile位于正确的配置层级。随后用 Nsys 命令运行训练脚本nsys profile -s none -o profile filepath -t cuda,nvtx --force-overwrite true --capture-rangecudaProfilerApi --capture-range-endstop python ./examples/...start_step/end_step以全局 batch 计ranks指定要剖析的全局 rank IDgen_shape控制是否生成包含输入形状的模型与 kernel 详细信息。CUDA 内存剖析在模型配置中加入以下选项启用 CUDA 内存 profilingmemory_profile: enabled: True start_step: 10 # Global batch to start profiling end_step: 10 # Global batch to end profiling rank: 0 # Global rank ID to profile output_path: None # Path to store the profile output file启用后无需修改调用命令直接照常运行 NeMo 脚本即可。小结一套贯穿所有领域模型的统一方法论无论你使用的是 ASR、TTS、Audio 还是 SpeechLM2 模型docs/source/core/core.rst 所定义的这套模型框架都保持一致用from_pretrained复用预训练模型用LightningModuleTrainer组织训练用 Hydra 的 YAML/CLI/Dataclass 三层配置驱动实验用optim命名空间统一优化器与调度器用.nemo文件完成自包含的保存、恢复与分发用register_artifact打包 tokenizer 等必要文件用register_nemo_submodule构建嵌套模型用 Nsys 与内存剖析定位性能瓶颈。掌握这套方法论后仓库 examples 中任何一个训练脚本与 examples/asr/conf 下的 YAML 配置对你而言都将是可直接上手、按需裁剪的模板。【免费下载链接】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），仅供参考