nemo_automodel.components.models.ernie4_5.rope_utils

Module Contents

Classes

Name	Description
`Ernie4_5RotaryEmbedding`	Rotary embedding module matching the Hugging Face ERNIE 4.5 implementation.

Functions

Name	Description
`apply_rotary_pos_emb`	Apply ERNIE 4.5 interleaved rotary embeddings to q and k.
`rotate_every_two`	Rotate interleaved RoPE pairs: [x0, x1] -> [-x1, x0].

API

class nemo_automodel.components.models.ernie4_5.rope_utils.Ernie4_5RotaryEmbedding(
    config,
    device: torch.device | None = None
)

Bases: Module

Rotary embedding module matching the Hugging Face ERNIE 4.5 implementation.

attention_scaling

= 1.0

inv_freq

Tensor

nemo_automodel.components.models.ernie4_5.rope_utils.Ernie4_5RotaryEmbedding.forward(
    x: torch.Tensor,
    position_ids: torch.Tensor,
    qkv_format: str = 'bshd'
) -> tuple[torch.Tensor, torch.Tensor]

Build FP32 interleaved rotary coefficients for local positions.

Parameters:

torch.Tensor

Tensor of shape [batch, sequence, heads, head_dim] or packed [tokens, heads, head_dim]; used to select the device.

position_ids

torch.Tensor

Integer tensor [batch, sequence] or packed [tokens].

qkv_format

strDefaults to 'bshd'

BSHD or packed THD tensor layout.

Returns: torch.Tensor

FP32 cosine and sine tensors [batch, sequence, head_dim] or

nemo_automodel.components.models.ernie4_5.rope_utils.apply_rotary_pos_emb(
    q: torch.Tensor,
    k: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]

Apply ERNIE 4.5 interleaved rotary embeddings to q and k.

Parameters:

torch.Tensor

Query tensor in BSHD or THD format.

torch.Tensor

Key tensor in BSHD or THD format.

cos

torch.Tensor

Cosine tensor with shape [B, S, D] for BSHD or [T, D] for THD.

sin

torch.Tensor

Sine tensor with shape [B, S, D] for BSHD or [T, D] for THD.

nemo_automodel.components.models.ernie4_5.rope_utils.rotate_every_two(
    x: torch.Tensor
) -> torch.Tensor

Rotate interleaved RoPE pairs: [x0, x1] -> [-x1, x0].