nemo_automodel.components.models.step5.rope

Half-split RoPE with the checkpoint’s explicit BF16 table boundary.

Module Contents

Classes

Name	Description
`Step5RotaryEmbedding`	Recompute FP32 frequencies from config; never cast a frequency buffer.

Functions

Name	Description
`_apply_full_rotary_compiled`	-
`_apply_partial_rotary_compiled`	-
`_apply_rotary`	Apply already-rounded tables using FP32 half-split arithmetic.
`_apply_rotary_compiled`	-
`_apply_single_head_rotary_compiled`	-

API

class nemo_automodel.components.models.step5.rope.Step5RotaryEmbedding(
    config: nemo_automodel.components.models.step5.configuration.Step5Config,
    layer_idx: int,
    span: int
)

Bases: Module

Recompute FP32 frequencies from config; never cast a frequency buffer.

scaling

theta

nemo_automodel.components.models.step5.rope.Step5RotaryEmbedding._tables(
    x: torch.Tensor,
    positions: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]

Build tables once per call; preserve eager frequency/trig rounding.

nemo_automodel.components.models.step5.rope.Step5RotaryEmbedding.forward(
    x: torch.Tensor,
    positions: torch.Tensor
) -> torch.Tensor

Rotate the leading half-split subspace.

Parameters:

torch.Tensor

Tensor of shape [batch, sequence, heads, head_dim].

positions

torch.Tensor

Tensor of shape [batch, sequence] with document-local positions.

Returns: torch.Tensor

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

nemo_automodel.components.models.step5.rope.Step5RotaryEmbedding.forward_pair(
    query: torch.Tensor,
    key: torch.Tensor,
    positions: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]

Rotate local Q/K with one shared, identically rounded table pair.

Parameters:

query

torch.Tensor

Query tensor [batch, sequence, query_heads, head_dim].

key

torch.Tensor

Key tensor [batch, sequence, key_heads, head_dim], sharing query’s dtype and device. Noncontiguous key views are supported.

positions

torch.Tensor

Document-local positions [batch, sequence].

Returns: tuple[torch.Tensor, torch.Tensor]

Rotated query and key tensors, preserving their respective shapes.

Raises:

ValueError: Q/K have different batch/sequence extents, dtype or device.

nemo_automodel.components.models.step5.rope._apply_full_rotary_compiled(
    x: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor
) -> torch.Tensor

nemo_automodel.components.models.step5.rope._apply_partial_rotary_compiled(
    x: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor,
    span: int
) -> torch.Tensor

nemo_automodel.components.models.step5.rope._apply_rotary(
    x: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor,
    span: int
) -> torch.Tensor

Apply already-rounded tables using FP32 half-split arithmetic.

nemo_automodel.components.models.step5.rope._apply_rotary_compiled(
    x: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor,
    span: int
) -> torch.Tensor

nemo_automodel.components.models.step5.rope._apply_single_head_rotary_compiled(
    x: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor,
    span: int
) -> torch.Tensor