nemo_automodel.components.models.d1.encoder

Omni’s bidirectional LFM2 composition with media-prefix isolation.

Module Contents

Classes

Name	Description
`D1OmniEncoder`	Checkpoint-compatible bidirectional encoder for text and media embeddings.
`D1OmniEncoderLayer`	Compose standard LFM2 attention/MLP with Omni’s centred convolution.
`D1OmniShortConv`	Centred gated convolution; media positions cannot read subsequent text.

API

class nemo_automodel.components.models.d1.encoder.D1OmniEncoder(
    text_config: transformers.Lfm2Config
)

Bases: Module

Checkpoint-compatible bidirectional encoder for text and media embeddings.

Perception modules supply prefixes. Text positions see the prefix and text; media positions see only media. There is no autoregressive cache.

config

= deepcopy(text_config)

embed_tokens

embedding_norm

layers

rotary_emb

= Lfm2RotaryEmbedding(config)

nemo_automodel.components.models.d1.encoder.D1OmniEncoder._apply(
    fn,
    recurse: bool = True
) -> nemo_automodel.components.models.d1.encoder.D1OmniEncoder

nemo_automodel.components.models.d1.encoder.D1OmniEncoder.forward(
    inputs_embeds: torch.Tensor,
    valid_tokens: torch.Tensor,
    media_lengths: torch.Tensor
) -> torch.Tensor

Encode combined media and text embeddings.

Parameters:

inputs_embeds

torch.Tensor

Tensor of shape [batch, sequence, hidden].

valid_tokens

torch.Tensor

Boolean Tensor of shape [batch, sequence], including media positions.

media_lengths

torch.Tensor

Long Tensor of shape [batch], zero for text-only rows.

Returns: torch.Tensor

Normalized Tensor of shape [batch, sequence, hidden].

class nemo_automodel.components.models.d1.encoder.D1OmniEncoderLayer(
    config: transformers.Lfm2Config,
    layer_index: int
)

Bases: Module

Compose standard LFM2 attention/MLP with Omni’s centred convolution.

conv

= D1OmniShortConv(config)

feed_forward

= Lfm2MLP(config)

ffn_norm

is_attention_layer

operator_norm

self_attn

= Lfm2Attention(config, layer_index)

nemo_automodel.components.models.d1.encoder.D1OmniEncoderLayer.forward(
    hidden_states: torch.Tensor,
    position_embeddings: tuple[torch.Tensor, torch.Tensor],
    attention_mask: torch.Tensor,
    valid_tokens: torch.Tensor,
    media_lengths: torch.Tensor
) -> torch.Tensor

Run a residual mixer followed by the gated feed-forward block.

Parameters:

hidden_states

torch.Tensor

Tensor of shape [batch, sequence, hidden].

position_embeddings

tuple[torch.Tensor, torch.Tensor]

Pair of cosine and sine Tensors, each of shape [batch, sequence, head_dim].

attention_mask

torch.Tensor

Additive Tensor of shape [batch, 1, sequence, sequence].

valid_tokens

torch.Tensor

Boolean Tensor of shape [batch, sequence].

media_lengths

torch.Tensor

Long Tensor of shape [batch] of prefix lengths.

Returns: torch.Tensor

Tensor of shape [batch, sequence, hidden].

class nemo_automodel.components.models.d1.encoder.D1OmniShortConv(
    config: transformers.Lfm2Config
)

Bases: Module

Centred gated convolution; media positions cannot read subsequent text.

conv

in_proj

= nn.Linear(width, 3 * width, bias=False)

out_proj

= nn.Linear(width, width, bias=False)

nemo_automodel.components.models.d1.encoder.D1OmniShortConv.forward(
    hidden_states: torch.Tensor,
    valid_tokens: torch.Tensor,
    media_lengths: torch.Tensor
) -> torch.Tensor

Mix neighbouring positions without crossing the media/text boundary.

Parameters:

hidden_states

torch.Tensor

Tensor of shape [batch, sequence, hidden].

valid_tokens

torch.Tensor

Boolean Tensor of shape [batch, sequence] marking real positions.

media_lengths

torch.Tensor

Long Tensor of shape [batch] giving each row’s media prefix length.

Returns: torch.Tensor

Tensor of shape [batch, sequence, hidden].