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]

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].