utils

PyTorch LLaMA model.

Classes

RMSNorm

Borrowed from LlamaRMSNorm class.

Functions

expand_mask

Expands attention_mask from [bsz, seq_len] to [bsz, 1, tgt_seq_len, src_seq_len].

make_causal_mask

Make causal mask used for bi-directional self-attention.
class RMSNorm

Bases: Module

Borrowed from LlamaRMSNorm class.

__init__(hidden_size, eps=1e-06)

LlamaRMSNorm is equivalent to T5LayerNorm.

forward(hidden_states)

Forward function for RMSNorm.

expand_mask(mask, dtype, tgt_len=None)

Expands attention_mask from [bsz, seq_len] to [bsz, 1, tgt_seq_len, src_seq_len].

Parameters:

  • mask (Tensor) –

  • dtype (dtype) –

  • tgt_len (int | None) –

make_causal_mask(input_ids_shape, dtype, device, past_key_values_length=0)

Make causal mask used for bi-directional self-attention.

Parameters:

  • input_ids_shape (Size) –

  • dtype (dtype) –

  • device (device) –

  • past_key_values_length (int) –