`nemo_automodel.components.models.ernie4_5.rope_utils`#

Module Contents#

Classes#

Ernie4_5RotaryEmbedding

Rotary embedding module matching the Hugging Face ERNIE 4.5 implementation.

Functions#

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

API#

nemo_automodel.components.models.ernie4_5.rope_utils.rotate_every_two(x: torch.Tensor) → torch.Tensor#: Rotate interleaved RoPE pairs: [x0, x1] -> [-x1, x0].

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:

q – Query tensor in BSHD or THD format.
k – Key tensor in BSHD or THD format.
cos – Cosine tensor with shape [B, S, D] for BSHD or [T, D] for THD.
sin – Sine tensor with shape [B, S, D] for BSHD or [T, D] for THD.

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

Bases: torch.nn.Module