`core.models.audio.nemo_transformer_audio_model`#

Module Contents#

Classes#

`NemoTransformerAudioConfig`	Hyperparameters for the vendored NeMo-style TransformerEncoder audio tower.
`NemoTransformerAudioModel`	Audio encoder matching LLaVA expectations.

API#

class core.models.audio.nemo_transformer_audio_model.NemoTransformerAudioConfig#

Hyperparameters for the vendored NeMo-style TransformerEncoder audio tower.

n_mels: int#: 80

d_model: int#: 512

n_heads: int#: 8

n_layers: int#: 17

drop_rate: float#: 0.1

qkv_bias: bool#: False

causal_mask: bool#: False

pre_encode: str#: ‘conv’

nan_debug: bool#: False

qk_norm: bool#: False

subsampling_factor: int#: 4

attn_impl: str#: ‘auto’

recompute_layers: bool#: False

left_context: Optional[int]#: None

property output_embedding_dim: int#: Return the encoder output embedding dimension (d_model).

property encoder_time_stride: int#: Return the time downsampling factor of the pre-encode stage.

classmethod from_dict( data: Dict[str, Any], ) → core.models.audio.nemo_transformer_audio_model.NemoTransformerAudioConfig#: Build a config from a dict, keeping only keys that match config fields.

class core.models.audio.nemo_transformer_audio_model.NemoTransformerAudioModel( config: core.models.audio.nemo_transformer_audio_model.NemoTransformerAudioConfig, )#

Bases: megatron.core.transformer.module.MegatronModule

Audio encoder matching LLaVA expectations.

forward(features, mask) returns (B, T', H) embeddings and a bool mask.

Initialization

static _ceil_div(value: int, divisor: int) → int#

_post_subsample_lengths( input_seq_lengths: torch.Tensor, max_input_frames: int, ) → torch.Tensor#

_pre_encode_forward_flops( batch_size: int, max_input_frames: int, ) → int#

_attention_pair_count(lengths: torch.Tensor) → torch.Tensor#

estimate_flops( input_seq_lengths: torch.Tensor, max_input_frames: Optional[int] = None, include_backward: Optional[bool] = None, ) → Dict[str, torch.Tensor]#

Estimate NeMo audio tower FLOPs for a batch of mel-frame lengths.

The estimate counts Conv/stacking pre-encode work, QKV/out projections, attention score/value products, and the two FeedForward linears. Norms, activations, dropout, masking, and softmax are intentionally omitted to keep the accounting comparable to Megatron’s GEMM-oriented LM estimate.

forward( input_features: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, ) → Tuple[torch.Tensor, torch.Tensor]#: Encode mel features into (B, T', H) embeddings and a validity mask.

forward_packed( input_features: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, ) → core.models.audio.packed_audio.PackedAudioEmbeddings#: Encode mel features into packed (padding-free) audio embeddings.

core.models.audio.nemo_transformer_audio_model#

Module Contents#

Classes#

API#

`core.models.audio.nemo_transformer_audio_model`#