nemo_automodel.components.models.gpt_oss.rope_utils

Module Contents

Classes

Name	Description
`RotaryEmbedding`	-

Functions

Name	Description
`_apply_rotary_emb_qk`	Rotate query/key together, sharing their dtype-rounded tables.
`apply_rotary_emb`	Apply rotary embeddings to input tensor.
`apply_rotary_emb_qk`	Apply rotary embeddings to query and key tensors.
`position_ids_to_freqs_cis`	-

Data

_compiled_apply_rotary_emb_qk

API

class nemo_automodel.components.models.gpt_oss.rope_utils.RotaryEmbedding(
    head_dim: int,
    base: int,
    dtype: torch.dtype,
    initial_context_length: int = 4096,
    scaling_factor: float = 1.0,
    ntk_alpha: float = 1.0,
    ntk_beta: float = 32.0,
    partial_rotary_factor: float = 1.0,
    device: torch.device | None = None
)

Bases: Module

rotary_dim

= int(head_dim * partial_rotary_factor)

nemo_automodel.components.models.gpt_oss.rope_utils.RotaryEmbedding._compute_concentration_and_inv_freq() -> torch.Tensor

See YaRN paper: https://arxiv.org/abs/2309.00071

Uses rotary_dim instead of head_dim to support partial rotary embeddings.

nemo_automodel.components.models.gpt_oss.rope_utils.RotaryEmbedding._compute_cos_sin(
    num_tokens: int
)

nemo_automodel.components.models.gpt_oss.rope_utils.RotaryEmbedding.forward(
    query: torch.Tensor,
    key: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]

nemo_automodel.components.models.gpt_oss.rope_utils._apply_rotary_emb_qk(
    q: torch.Tensor,
    k: torch.Tensor,
    freqs_cis: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]

Rotate query/key together, sharing their dtype-rounded tables.

Parameters:

torch.Tensor

Tensor of shape […, query_heads, head_dim], with arbitrary leading dimensions and strides.

torch.Tensor

Tensor of shape […, key_heads, head_dim], sharing the leading dimensions of q and accepting arbitrary strides.

freqs_cis

torch.Tensor

Tensor of shape […, rotary_dim], with leading dimensions broadcastable to q and k. The last dimension concatenates cosine and sine tables; rotary_dim is even and at most head_dim.

Returns: torch.Tensor

Independently allocated rotated query and key tensors, preserving

nemo_automodel.components.models.gpt_oss.rope_utils.apply_rotary_emb(
    x: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor
) -> torch.Tensor

Apply rotary embeddings to input tensor.

If cos/sin have fewer dimensions than x (due to partial_rotary_factor < 1.0), only the first rotary_dim dimensions of x are rotated, and the rest are passed through.

Parameters:

torch.Tensor

Input tensor (…, head_dim)

cos

torch.Tensor

Cosine tensor (…, rotary_dim // 2)

sin

torch.Tensor

Sine tensor (…, rotary_dim // 2)

nemo_automodel.components.models.gpt_oss.rope_utils.apply_rotary_emb_qk(
    q: torch.Tensor,
    k: torch.Tensor,
    freqs_cis: torch.Tensor,
    format: str = 'bshd',
    rope_fusion: bool = True,
    cu_seqlens: torch.Tensor | None = None,
    concentration: float | None = None,
    cp_size: int = 1,
    cp_rank: int = 0
) -> tuple[torch.Tensor, torch.Tensor]

Apply rotary embeddings to query and key tensors.

Parameters:

torch.Tensor

Query tensor of shape [batch, sequence, query_heads, head_dim] for BSHD or [tokens, query_heads, head_dim] for packed THD. Leading dimensions contain the local context-parallel shard, if used.

torch.Tensor

Key tensor with the same layout as q, with key_heads in place of query_heads. Arbitrary strides are supported.

freqs_cis

torch.Tensor

For non-fused RoPE, tensor of shape [batch, sequence, rotary_dim] or [tokens, rotary_dim], containing concatenated cosine/sine tables with concentration applied. Leading dimensions broadcast to the local tokens in q and k. For TE fused RoPE, tensor of shape [sequence, 1, 1, rotary_dim] containing duplicated angles for the global sequence. rotary_dim is even and at most head_dim.

format

strDefaults to 'bshd'

QKV format (“bshd” or “thd”).

rope_fusion

boolDefaults to True

If True, use TE fused rope. If False, use non-fused rope.

cu_seqlens

torch.Tensor | NoneDefaults to None

Optional tensor of shape [num_sequences + 1] containing cumulative sequence lengths for variable-length TE fused RoPE.

concentration

float | NoneDefaults to None

Optional scale applied after TE fused RoPE.

cp_size

intDefaults to 1

Context parallelism size.

cp_rank

intDefaults to 0

Context parallelism rank.

Returns: torch.Tensor

Tuple of rotated query and key tensors, preserving their respective

nemo_automodel.components.models.gpt_oss.rope_utils.position_ids_to_freqs_cis(
    rotary_emb: nemo_automodel.components.models.gpt_oss.rope_utils.RotaryEmbedding,
    position_ids: torch.Tensor,
    qkv_format: str = 'bshd',
    for_fused_rope: bool = True,
    cp_size: int = 1
) -> torch.Tensor

nemo_automodel.components.models.gpt_oss.rope_utils._compiled_apply_rotary_emb_qk = torch.compile(_apply_rotary_emb_qk, dynamic=True, options={'emulate_precision_ca...