nemo_automodel.components.models.d1.causal

Released D1-3B rotary precision on the standard LFM2 causal stack.

Module Contents

Classes

Name	Description
`D1CausalAttention`	Reuse an HF attention block’s weights with the released float32 rotation.
`D1CausalRotaryEmbedding`	Keep frequencies and trigonometry in float32, as the released D1 decoder.
`D1CausalShortConv`	Retain HF checkpoint/cache ownership with the release’s chain convolution.

API

class nemo_automodel.components.models.d1.causal.D1CausalAttention(
    source: transformers.models.lfm2.modeling_lfm2.Lfm2Attention
)

Bases: Lfm2Attention

Reuse an HF attention block’s weights with the released float32 rotation.

config

= source.config

head_dim

= source.head_dim

k_layernorm

= source.k_layernorm

k_proj

= source.k_proj

layer_idx

= source.layer_idx

num_key_value_groups

= source.num_key_value_groups

out_proj

= source.out_proj

q_layernorm

= source.q_layernorm

q_proj

= source.q_proj

scaling

= source.scaling

v_proj

= source.v_proj

nemo_automodel.components.models.d1.causal.D1CausalAttention.forward(
    hidden_states: torch.Tensor,
    position_embeddings: tuple[torch.Tensor, torch.Tensor],
    attention_mask: torch.Tensor | None,
    past_key_values: transformers.cache_utils.Cache | None = None,
    kwargs: typing.Any = {}
) -> tuple[torch.Tensor, torch.Tensor | None]

Run standard GQA after the checkpoint’s float32 half-split rotation.

Parameters:

hidden_states

torch.Tensor

Tensor of shape [batch, sequence, hidden].

position_embeddings

tuple[torch.Tensor, torch.Tensor]

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

attention_mask

torch.Tensor | None

Optional additive Tensor of shape [batch, 1, sequence, key_sequence], with zero for allowed keys and a negative sentinel for causal/padding exclusions.

past_key_values

Cache | NoneDefaults to None

Optional HF autoregressive cache, updated in place.

**kwargs

AnyDefaults to {}

Additional HF attention dispatcher arguments.

Returns: torch.Tensor

Pair of output Tensor of shape [batch, sequence, hidden] and

class nemo_automodel.components.models.d1.causal.D1CausalRotaryEmbedding()

Bases: Lfm2RotaryEmbedding

Keep frequencies and trigonometry in float32, as the released D1 decoder.

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

nemo_automodel.components.models.d1.causal.D1CausalRotaryEmbedding.forward(
    x: torch.Tensor,
    position_ids: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]

Compute the release’s float32 rotary factors.

Parameters:

torch.Tensor

Tensor of shape [batch, sequence, hidden], supplying the device.

position_ids

torch.Tensor

Long Tensor of shape [batch, sequence].

Returns: torch.Tensor

Pair of float32 cosine and sine Tensors, each of shape

class nemo_automodel.components.models.d1.causal.D1CausalShortConv(
    source: transformers.models.lfm2.modeling_lfm2.Lfm2ShortConv
)

Bases: Lfm2ShortConv

Retain HF checkpoint/cache ownership with the release’s chain convolution.

bias

= source.bias

config

= source.config

conv

= source.conv

conv_kernel_size

= source.conv_kernel_size

in_proj

= source.in_proj

layer_idx

= source.layer_idx

layer_type

= source.layer_type

out_proj

= source.out_proj

nemo_automodel.components.models.d1.causal.D1CausalShortConv.forward(
    hidden_states: torch.Tensor,
    past_key_values: transformers.cache_utils.Cache | None = None,
    attention_mask: torch.Tensor | None = None,
    seq_idx: torch.IntTensor | None = None
) -> torch.Tensor

Run the released depthwise kernel for independent uncached chains.

Cached decoding and packed sequence indices retain HF’s implementation. Training uses uncached chains and preserves the publisher’s convolution accumulation rather than the optional fused causal-conv kernel.

Parameters:

hidden_states

torch.Tensor

Tensor of shape [batch, sequence, hidden].

past_key_values

Cache | NoneDefaults to None

Optional HF cache. Cached decoding delegates to HF and updates its per-layer convolution state in place.

attention_mask

torch.Tensor | NoneDefaults to None

Optional Tensor of shape [batch, sequence], with one/true for real tokens and zero/false for padding.

seq_idx

torch.IntTensor | NoneDefaults to None

Optional integer Tensor of shape [batch, sequence] giving packed sequence IDs. When present the HF path owns segment boundary handling and its optional causal-convolution kernel.

Returns: torch.Tensor

Tensor of shape [batch, sequence, hidden] in the input dtype.