`core.models.audio.nemo_audio_checkpoint`#

Load NeMo .nemo archives for the transformer audio encoder.

A .nemo file is the standard NeMo SaveRestoreConnector artifact: an uncompressed tar containing model_config.yaml (OmegaConf) and model_weights.ckpt (torch.save of the full EncDecRNNTBPEModel state dict). This module supports only that format – generic .pt / .bin checkpoint files are intentionally not handled.

Public API:

extract_nemo_archive(path) -> (model_cfg, full_state_dict)
nemo_audio_configs_from_archive(path) -> (encoder_cfg, preproc_cfg, encoder_state)
load_nemo_transformer_audio_weights(audio_module, ckpt_path, *, strict=False)

Module Contents#

Functions#

`_load_omegaconf`	Lazy import of OmegaConf – only paid for when we actually open a .nemo.
`_torch_load`
`_load_json`
`_write_json`
`_load_preprocessor_config`
`nemo_audio_config_paths_from_checkpoint_dir`	Return checkpoint-local NeMo audio config paths for an iteration dir.
`has_nemo_audio_configs_in_checkpoint_dir`	True when both checkpoint-local NeMo audio config JSON files are present.
`nemo_audio_configs_from_json_paths`	Load NeMo audio encoder/preprocessor configs from JSON files.
`nemo_audio_configs_from_checkpoint_dir`	Load NeMo audio configs persisted next to a Megatron checkpoint iteration.
`resolve_nemo_audio_configs_from_args`	Resolve NeMo audio configs from training/runtime args.
`write_nemo_audio_configs_to_checkpoint_dir`	Persist resolved NeMo audio configs as JSON under a checkpoint iteration dir.
`write_nemo_audio_configs_from_args_to_checkpoint_dir`	Persist resolved NeMo audio configs for `args` when audio is enabled.
`_validate_archive`
`_extract_member`	Extract a single archive member by basename. NeMo writes `./model_config.yaml`.
`read_nemo_config`	Read just the `model_config.yaml` from a `.nemo` archive.
`extract_nemo_archive`	Extract a `.nemo` archive and return `(model_cfg_dict, full_state_dict)`.
`_strip_target`
`_split_audio_configs`	Validate `_target_`s and convert the model config dict into our dataclasses.
`nemo_audio_configs_from_path`	Cheap config-only read of a `.nemo` archive.
`nemo_audio_configs_from_archive`	Parse a `.nemo` archive into the encoder/preprocessor configs + encoder state.
`load_nemo_transformer_audio_weights`	Load encoder weights from a `.nemo` archive into `NemoTransformerAudioModel`.
`_validate_encoder_cfg`	Raise if a structural field disagrees between model and checkpoint.

Data#

`MODEL_CONFIG_NAME`
`MODEL_WEIGHTS_NAME`
`CHECKPOINT_NEMO_TRANSFORMER_AUDIO_CONFIG_NAME`
`CHECKPOINT_NEMO_AUDIO_PREPROCESSOR_CONFIG_NAME`
`_KNOWN_ENCODER_TARGETS`
`_KNOWN_PREPROC_TARGETS`

API#

core.models.audio.nemo_audio_checkpoint.MODEL_CONFIG_NAME#: ‘model_config.yaml’

core.models.audio.nemo_audio_checkpoint.MODEL_WEIGHTS_NAME#: ‘model_weights.ckpt’

core.models.audio.nemo_audio_checkpoint.CHECKPOINT_NEMO_TRANSFORMER_AUDIO_CONFIG_NAME#: ‘nemo_transformer_audio_config.json’

core.models.audio.nemo_audio_checkpoint.CHECKPOINT_NEMO_AUDIO_PREPROCESSOR_CONFIG_NAME#: ‘nemo_audio_preprocessor_config.json’

core.models.audio.nemo_audio_checkpoint._KNOWN_ENCODER_TARGETS#: None

core.models.audio.nemo_audio_checkpoint._KNOWN_PREPROC_TARGETS#: None

core.models.audio.nemo_audio_checkpoint._load_omegaconf(path: str)#: Lazy import of OmegaConf – only paid for when we actually open a .nemo.

core.models.audio.nemo_audio_checkpoint._torch_load(path: str) → Any#

core.models.audio.nemo_audio_checkpoint._load_json(path: str | pathlib.Path) → Dict[str, Any]#

core.models.audio.nemo_audio_checkpoint._write_json( path: pathlib.Path, data: Dict[str, Any], ) → None#

core.models.audio.nemo_audio_checkpoint._load_preprocessor_config( path: str | pathlib.Path | None, ) → core.models.audio.audio_feature_config.NemoAudioFeatureConfig#

core.models.audio.nemo_audio_checkpoint.nemo_audio_config_paths_from_checkpoint_dir( checkpoint_dir: str | pathlib.Path, ) → Tuple[pathlib.Path, pathlib.Path]#: Return checkpoint-local NeMo audio config paths for an iteration dir.

core.models.audio.nemo_audio_checkpoint.has_nemo_audio_configs_in_checkpoint_dir( checkpoint_dir: str | pathlib.Path, ) → bool#: True when both checkpoint-local NeMo audio config JSON files are present.

core.models.audio.nemo_audio_checkpoint.nemo_audio_configs_from_json_paths( encoder_config_path: str | pathlib.Path, preprocessor_config_path: str | pathlib.Path | None = None, ) → Tuple[core.models.audio.nemo_transformer_audio_model.NemoTransformerAudioConfig, core.models.audio.audio_feature_config.NemoAudioFeatureConfig]#: Load NeMo audio encoder/preprocessor configs from JSON files.

core.models.audio.nemo_audio_checkpoint.nemo_audio_configs_from_checkpoint_dir( checkpoint_dir: str | pathlib.Path, ) → Tuple[core.models.audio.nemo_transformer_audio_model.NemoTransformerAudioConfig, core.models.audio.audio_feature_config.NemoAudioFeatureConfig]#: Load NeMo audio configs persisted next to a Megatron checkpoint iteration.

core.models.audio.nemo_audio_checkpoint.resolve_nemo_audio_configs_from_args( args, ) → Tuple[core.models.audio.nemo_transformer_audio_model.NemoTransformerAudioConfig, core.models.audio.audio_feature_config.NemoAudioFeatureConfig]#

Resolve NeMo audio configs from training/runtime args.

Source precedence is:

--load-audio-from when it points to a .nemo archive.
--nemo-transformer-audio-config / --nemo-audio-preprocessor-config JSON.
Dataclass defaults.

--nemo-transformer-audio-attn-impl is a runtime backend override and is materialized into the returned encoder config so checkpoint-local artifacts reproduce the model that was actually instantiated.

core.models.audio.nemo_audio_checkpoint.write_nemo_audio_configs_to_checkpoint_dir( checkpoint_dir: str | pathlib.Path, encoder_cfg: core.models.audio.nemo_transformer_audio_model.NemoTransformerAudioConfig, preproc_cfg: core.models.audio.audio_feature_config.NemoAudioFeatureConfig, ) → Tuple[pathlib.Path, pathlib.Path]#: Persist resolved NeMo audio configs as JSON under a checkpoint iteration dir.

core.models.audio.nemo_audio_checkpoint.write_nemo_audio_configs_from_args_to_checkpoint_dir( args, checkpoint_dir: str | pathlib.Path, ) → Tuple[pathlib.Path, pathlib.Path] | None#: Persist resolved NeMo audio configs for args when audio is enabled.

core.models.audio.nemo_audio_checkpoint._validate_archive(nemo_path: pathlib.Path) → None#

core.models.audio.nemo_audio_checkpoint._extract_member( tar: tarfile.TarFile, member_name: str, out_dir: str, ) → str#

Extract a single archive member by basename. NeMo writes ./model_config.yaml.

Iterates lazily via tar.next() and stops at the first match. We deliberately avoid tar.getmembers() here: it walks the entire archive to EOF, which trips on .nemo files whose trailing zero-block region is malformed or missing (a known NeMo SaveRestoreConnector quirk that surfaces as tarfile.ReadError: unexpected end of data). For NeMo’s standard layout (model_config.yaml precedes model_weights.ckpt), two sequential calls on the same handle never need to walk past model_weights.ckpt and never hit the bad tail.

core.models.audio.nemo_audio_checkpoint.read_nemo_config( nemo_path: str | pathlib.Path, ) → Dict[str, Any]#

Read just the model_config.yaml from a .nemo archive.

Cheap (only ~kB of data extracted), so safe to call on every rank during model construction.

core.models.audio.nemo_audio_checkpoint.extract_nemo_archive( nemo_path: str | pathlib.Path, out_dir: str | pathlib.Path | None = None, ) → Tuple[Dict[str, Any], Dict[str, torch.Tensor]]#

Extract a .nemo archive and return (model_cfg_dict, full_state_dict).

Parameters:

nemo_path – Path to the .nemo file.
out_dir – Optional directory to extract into. If None, a temp directory is used and cleaned up automatically.

Returns:

A tuple (cfg, state) where cfg is the resolved OmegaConf model block as a plain dict, and state is the full flat torch.save’d state dict (preprocessor + encoder + decoder + joint + …).

core.models.audio.nemo_audio_checkpoint._strip_target( d: Dict[str, Any], ) → Dict[str, Any]#

core.models.audio.nemo_audio_checkpoint._split_audio_configs( cfg: Dict[str, Any], nemo_path: str | pathlib.Path, ) → Tuple[core.models.audio.nemo_transformer_audio_model.NemoTransformerAudioConfig, core.models.audio.audio_feature_config.NemoAudioFeatureConfig]#: Validate _target_s and convert the model config dict into our dataclasses.

core.models.audio.nemo_audio_checkpoint.nemo_audio_configs_from_path( nemo_path: str | pathlib.Path, ) → Tuple[core.models.audio.nemo_transformer_audio_model.NemoTransformerAudioConfig, core.models.audio.audio_feature_config.NemoAudioFeatureConfig]#

Cheap config-only read of a .nemo archive.

Use this when you only need the encoder/preprocessor hyperparameters (e.g. on every rank during model construction). Does not touch the state dict.

core.models.audio.nemo_audio_checkpoint.nemo_audio_configs_from_archive( nemo_path: str | pathlib.Path, ) → Tuple[core.models.audio.nemo_transformer_audio_model.NemoTransformerAudioConfig, core.models.audio.audio_feature_config.NemoAudioFeatureConfig, Dict[str, torch.Tensor]]#

Parse a .nemo archive into the encoder/preprocessor configs + encoder state.

The encoder state dict has the leading encoder. prefix stripped so it can be loaded directly into NemoTransformerAudioModel.encoder.

Raises ValueError if the archive does not look like a NeMo ASR model that pairs AudioToMelSpectrogramPreprocessor with one of the supported transformer encoder _target_s.

core.models.audio.nemo_audio_checkpoint.load_nemo_transformer_audio_weights( audio_module: torch.nn.Module, ckpt_path: str | pathlib.Path, *, strict: bool = False, ) → Tuple[List[str], List[str]]#

Load encoder weights from a .nemo archive into NemoTransformerAudioModel.

Validates that the archive’s model.encoder config matches the audio module’s config (n_mels, d_model, n_heads, n_layers, pre_encode, subsampling_factor, qk_norm) before loading.

Returns:: (missing_keys, unexpected_keys) from load_state_dict(strict=False), with TransformerEngine _extra_state entries removed.

core.models.audio.nemo_audio_checkpoint._validate_encoder_cfg( model_cfg: core.models.audio.nemo_transformer_audio_model.NemoTransformerAudioConfig, ckpt_cfg: core.models.audio.nemo_transformer_audio_model.NemoTransformerAudioConfig, ckpt_path: str | pathlib.Path, ) → None#: Raise if a structural field disagrees between model and checkpoint.

core.models.audio.nemo_audio_checkpoint#

Module Contents#

Functions#

Data#

API#

`core.models.audio.nemo_audio_checkpoint`#