Is this page helpful?

core.models.audio.audio_projector#

Module Contents#

Classes#

AudioProjection

Stack audio embeddings in time and project them into LM hidden size.

API#

class core.models.audio.audio_projector.AudioProjection(
config: megatron.core.transformer.transformer_config.TransformerConfig,
submodules: megatron.core.transformer.mlp.MLPSubmodules,
projector_type: str,
input_size: int,
stack_factor: int = 1,
tp_group: Optional[torch.distributed.ProcessGroup] = None,
pg_collection: Optional[megatron.core.process_groups_config.ProcessGroupCollection] = None,
)#

Bases: megatron.core.transformer.module.MegatronModule

Stack audio embeddings in time and project them into LM hidden size.

Initialization

_stack_features(
hidden_states: torch.Tensor,
attention_mask: Optional[torch.Tensor] = None,
) → Tuple[torch.Tensor, Optional[torch.Tensor]]#
forward(
hidden_states: torch.Tensor,
attention_mask: Optional[torch.Tensor] = None,
) → Tuple[torch.Tensor, Optional[torch.Tensor]]#

Stack audio embeddings in time and project them, returning the projected states and mask.

The stacked features are permuted to [seq, batch, hidden] before projection.

forward_packed(
packed_states: core.models.audio.packed_audio.PackedAudioEmbeddings,
) → core.models.audio.packed_audio.PackedAudioEmbeddings#

Project packed audio embeddings, preserving packing (requires stack_factor == 1).

estimate_flops(
output_seq_lengths: torch.Tensor,
include_backward: Optional[bool] = None,
input_requires_grad: Optional[bool] = None,
count_padded: bool = True,
) → dict[str, torch.Tensor]#

Estimate FLOPs for the audio projection over projected audio lengths.

output_seq_lengths is in the projected/stacked audio-token space (the same space as audio_embeds_seq_lengths). By default the estimate counts the padded tensor shape used by forward.