Inference-only LLaMA model compatible with HuggingFace weights.

`fastvideo.models.encoders.llama.LocalAttention`

LocalAttention(num_heads: int, head_size: int, num_kv_heads: int | None = None, softmax_scale: float | None = None, causal: bool = False, supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None, default_backend: AttentionBackendEnum | None = None, **extra_impl_args)

Attention layer.

`fastvideo.models.encoders.llama.LocalAttention.forward`

forward(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, freqs_cis: tuple[torch.Tensor, torch.Tensor] | None = None) -> torch.Tensor

Apply local attention between query, key and value tensors.

Parameters

q: torch.Tensor: Query tensor of shape [batch_size, seq_len, num_heads, head_dim]
k: torch.Tensor: Key tensor of shape [batch_size, seq_len, num_heads, head_dim]
v: torch.Tensor: Value tensor of shape [batch_size, seq_len, num_heads, head_dim]

Returns

torch.Tensor: torch.Tensor: Output tensor after local attention

`fastvideo.models.encoders.llama.BaseEncoderOutput`

BaseEncoderOutput(last_hidden_state: torch.FloatTensor | None = None, pooler_output: torch.FloatTensor | None = None, hidden_states: tuple[torch.FloatTensor, ...] | None = None, attentions: tuple[torch.FloatTensor, ...] | None = None, attention_mask: torch.Tensor | None = None)

`fastvideo.models.encoders.llama.LlamaConfig`

LlamaConfig(arch_config: TextEncoderArchConfig = LlamaArchConfig(), prefix: str = 'llama', quant_config: QuantizationConfig | None = None, lora_config: Any | None = None, is_chat_model: bool = False, treat_empty_as_dot: bool = False, *, _resolved_attention_backend: AttentionBackendEnum | None = None, chat_template_enable_thinking: bool = False)

`fastvideo.models.encoders.llama.get_tp_world_size`

get_tp_world_size() -> int

Return world size for the tensor model parallel group.

`fastvideo.models.encoders.llama.SiluAndMul`

SiluAndMul()

An activation function for SwiGLU.

The function computes x -> silu(x[:d]) * x[d:] where d = x.shape[-1] // 2.

Shapes

x: (num_tokens, 2 * d) or (batch_size, seq_len, 2 * d) return: (num_tokens, d) or (batch_size, seq_len, d)

`fastvideo.models.encoders.llama.SiluAndMul.forward_native`

forward_native(x: torch.Tensor) -> torch.Tensor

PyTorch-native implementation equivalent to forward().

`fastvideo.models.encoders.llama.RMSNorm`

RMSNorm(hidden_size: int, eps: float = 1e-06, dtype: torch.dtype = torch.float32, var_hidden_size: int | None = None, has_weight: bool = True)

Root mean square normalization.

Computes x -> w * x / sqrt(E[x^2] + eps) where w is the learned weight. Refer to https://arxiv.org/abs/1910.07467

`fastvideo.models.encoders.llama.RMSNorm.forward_native`

forward_native(x: torch.Tensor, residual: torch.Tensor | None = None) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]

PyTorch-native implementation equivalent to forward().

`fastvideo.models.encoders.llama.RMSNorm.extra_repr`

extra_repr() -> str

`fastvideo.models.encoders.llama.MergedColumnParallelLinear`

MergedColumnParallelLinear(input_size: int, output_sizes: list[int], bias: bool = True, gather_output: bool = False, skip_bias_add: bool = False, params_dtype: torch.dtype | None = None, quant_config: QuantizationConfig | None = None, prefix: str = '')

Packed linear layers with column parallelism.

Similar to ColumnParallelLinear, but the weight matrix is concatenated along the output dimension. When the weight matrix is loaded, the different partitions are sharded separately.

Parameters

input_size: int: input dimension of the linear layer.
output_sizes: list[int]: list of output dimensions of the linear layer.
bias: bool: If true, add bias.
gather_output: bool: If true, call all-gather on output and make the output available to all GPUs, otherwise, every GPU will have its own output.
skip_bias_add: bool: This was added to enable performance optimizations where bias can be fused with other element-wise operations. we skip adding bias but instead return it.
params_dtype: torch.dtype | None: Data type for the parameters.
quant_config: QuantizationConfig | None: Quantization configure.
prefix: str: The name of the layer in the state dict, including all parents (e.g. model.layers.0.qkv_proj)

`fastvideo.models.encoders.llama.MergedColumnParallelLinear.weight_loader`

weight_loader(param: Parameter, loaded_weight: torch.Tensor, loaded_shard_id: int | None = None) -> None

`fastvideo.models.encoders.llama.MergedColumnParallelLinear.weight_loader_v2`

weight_loader_v2(param: BasevLLMParameter, loaded_weight: torch.Tensor, loaded_shard_id: int | None = None) -> None

`fastvideo.models.encoders.llama.QKVParallelLinear`

QKVParallelLinear(hidden_size: int, head_size: int, total_num_heads: int, total_num_kv_heads: int | None = None, bias: bool = True, skip_bias_add: bool = False, params_dtype: torch.dtype | None = None, quant_config: QuantizationConfig | None = None, prefix: str = '')

Linear layers for the attention's QKV transformation.

Linear layers for the linear transformation of the query, key, and value vectors in the attention layer. The weight matrix is concatenated along the output dimension. The layer is parallelized along the head dimension. When the number of key/value heads is smaller than the number of query heads (e.g., multi-query/grouped-query attention), the key/value head may be replicated while the query heads are partitioned.

Parameters

hidden_size: int: input hidden state size of the transformer.
head_size: int: size of each attention head.
total_num_heads: int: total number of attention query heads.
total_num_kv_heads: int | None: total number of attention key/value heads. If None, assume total_num_kv_heads = total_num_heads.
bias: bool: If true, add bias.
skip_bias_add: bool: This was added to enable performance optimizations where bias can be fused with other element-wise operations. we skip adding bias but instead return it.
params_dtype: torch.dtype | None: Data type for the parameters.
quant_config: QuantizationConfig | None: Quantization configure.
prefix: str: The name of the layer in the state dict, including all parents (e.g. model.layers.0.qkv_proj)

`fastvideo.models.encoders.llama.QKVParallelLinear.weight_loader_v2`

weight_loader_v2(param: BasevLLMParameter, loaded_weight: torch.Tensor, loaded_shard_id: str | None = None)

`fastvideo.models.encoders.llama.QKVParallelLinear.weight_loader`

weight_loader(param: Parameter, loaded_weight: torch.Tensor, loaded_shard_id: str | None = None)

`fastvideo.models.encoders.llama.RowParallelLinear`

RowParallelLinear(input_size: int, output_size: int, bias: bool = True, input_is_parallel: bool = True, skip_bias_add: bool = False, params_dtype: torch.dtype | None = None, reduce_results: bool = True, quant_config: QuantizationConfig | None = None, prefix: str = '')

Linear layer with row parallelism.

The linear layer is defined as Y = XA + b. A is parallelized along its first dimension and X along its second dimension as: - - | A_1 | | . | A = | . | X = [X_1, ..., X_p] | . | | A_p | - - Arguments: input_size: first dimension of matrix A. output_size: second dimension of matrix A. bias: If true, add bias. Note that bias is not parallelized. input_is_parallel: If true, we assume that the input is already split across the GPUs and we do not split again. skip_bias_add: This was added to enable performance optimization where bias can be fused with other element-wise operations. We skip adding bias but instead return it. params_dtype: Data type for the parameters. quant_config: Quantization configure.

`fastvideo.models.encoders.llama.RowParallelLinear.weight_loader`

weight_loader(param: Parameter, loaded_weight: torch.Tensor)

`fastvideo.models.encoders.llama.RowParallelLinear.weight_loader_v2`

weight_loader_v2(param: BasevLLMParameter, loaded_weight: torch.Tensor)

`fastvideo.models.encoders.llama.RowParallelLinear.forward`

forward(input_) -> tuple[torch.Tensor, Parameter | None]

`fastvideo.models.encoders.llama.RowParallelLinear.extra_repr`

extra_repr() -> str

`fastvideo.models.encoders.llama.QuantizationConfig`

QuantizationConfig()

Base class for quantization configs.

`fastvideo.models.encoders.llama.QuantizationConfig.get_name`

get_name() -> QuantizationMethods

Name of the quantization method.

`fastvideo.models.encoders.llama.QuantizationConfig.get_supported_act_dtypes`

get_supported_act_dtypes() -> list[torch.dtype]

List of supported activation dtypes.

`fastvideo.models.encoders.llama.QuantizationConfig.get_min_capability`

get_min_capability() -> int

Minimum GPU capability to support the quantization method.

E.g., 70 for Volta, 75 for Turing, 80 for Ampere. This requirement is due to the custom CUDA kernels used by the quantization method.

`fastvideo.models.encoders.llama.QuantizationConfig.get_config_filenames`

get_config_filenames() -> list[str]

List of filenames to search for in the model directory.

`fastvideo.models.encoders.llama.QuantizationConfig.from_config`

from_config(config: dict[str, Any]) -> QuantizationConfig

Create a config class from the model's quantization config.

`fastvideo.models.encoders.llama.QuantizationConfig.override_quantization_method`

override_quantization_method(hf_quant_cfg, user_quant) -> QuantizationMethods | None

Detects if this quantization method can support a given checkpoint format by overriding the user specified quantization method -- this method should only be overwritten by subclasses in exceptional circumstances

`fastvideo.models.encoders.llama.QuantizationConfig.get_from_keys`

get_from_keys(config: dict[str, Any], keys: list[str]) -> Any

Get a value from the model's quantization config.

`fastvideo.models.encoders.llama.QuantizationConfig.get_from_keys_or`

get_from_keys_or(config: dict[str, Any], keys: list[str], default: Any) -> Any

Get a optional value from the model's quantization config.

`fastvideo.models.encoders.llama.QuantizationConfig.get_quant_method`

get_quant_method(layer: torch.nn.Module, prefix: str) -> QuantizeMethodBase | None

Get the quantize method to use for the quantized layer.

Parameters

layer: torch.nn.Module: The layer for the quant method.
prefix: str: The full name of the layer in the state dict

Returns: The quantize method. None if the given layer doesn't support quant method.

`fastvideo.models.encoders.llama.QuantizationConfig.validate_runtime`

validate_runtime(device: torch.device) -> None

Fail fast on a device the scheme cannot execute on.

Called once by loaders that select a serialized checkpoint scheme, before any weight is read. The default accepts every device; schemes with kernel or capability requirements override it.

`fastvideo.models.encoders.llama.QuantizationConfig.get_cache_scale`

get_cache_scale(name: str) -> str | None

`fastvideo.models.encoders.llama.get_rope`

get_rope(head_size: int, rotary_dim: int, max_position: int, base: int | float, is_neox_style: bool = True, rope_scaling: dict[str, Any] | None = None, dtype: torch.dtype | None = None, partial_rotary_factor: float = 1.0) -> RotaryEmbedding

`fastvideo.models.encoders.llama.VocabParallelEmbedding`

VocabParallelEmbedding(num_embeddings: int, embedding_dim: int, params_dtype: torch.dtype | None = None, org_num_embeddings: int | None = None, padding_size: int = DEFAULT_VOCAB_PADDING_SIZE, quant_config: QuantizationConfig | None = None, prefix: str = '')

Embedding parallelized in the vocabulary dimension.

Adapted from torch.nn.Embedding, note that we pad the vocabulary size to make sure it is divisible by the number of model parallel GPUs.

In order to support various loading methods, we ensure that LoRA-added embeddings are always at the end of TP-sharded tensors. In other words, we shard base embeddings and LoRA embeddings separately (both padded), and place them in the same tensor. In this example, we will have the original vocab size = 1010, added vocab size = 16 and padding to 64. Therefore, the total vocab size with padding will be 1088 (because we first pad 1010 to 1024, add 16, and then pad to 1088). Therefore, the tensor format looks like the following: TP1, rank 0 (no sharding): |< --------BASE-------- >|< -BASE PADDING-- >|< -----LORA------ >|< -LORA PADDING-- >| corresponding token_id: | 0 | 1 | ... | 1009 | -1 | ... | -1 | 1010 | ... | 1015 | -1 | ... | -1 | index: | 0 | 1 | ... | 1009 | 1010 | ... | 1023 | 1024 | ... | 1039 | 1040 | ... | 1087 |

TP2, rank 0: |< --------------------BASE--------------------- >|< -----LORA------ >|< -LORA PADDING- >| corresponding token_id: | 0 | 1 | 2 | ... | 497 | 498 | ... | 511 | 1000 | ... | 1015 | -1 | ... | -1 | index: | 0 | 1 | 2 | ... | 497 | 498 | ... | 511 | 512 | ... | 527 | 520 | ... | 543 | TP2, rank 1: |< -----------BASE----------- >|< -BASE PADDING- >|< -----------LORA PADDING----------- >| corresponding token_id: | 512 | 513 | 514 | ... | 1009 | -1 | ... | -1 | -1 | ... | -1 | -1 | ... | -1 | index: | 0 | 1 | 2 | ... | 497 | 498 | ... | 511 | 512 | ... | 519 | 520 | ... | 543 |

Parameters

num_embeddings: int: vocabulary size.
embedding_dim: int: size of hidden state.
params_dtype: torch.dtype | None: type of the parameters.
org_num_embeddings: int | None: original vocabulary size (without LoRA).
padding_size: int: padding size for the vocabulary.
quant_config: QuantizationConfig | None: quant config for the layer
prefix: str: full name of the layer in the state dict

`fastvideo.models.encoders.llama.VocabParallelEmbedding.get_sharded_to_full_mapping`

get_sharded_to_full_mapping() -> list[int] | None

Get a mapping that can be used to reindex the gathered logits for sampling.

During sampling, we gather logits from all ranks. The relationship of index->token_id will follow the same format as outlined in the class docstring. However, after the gather, we want to reindex the final logits tensor to map index->token_id one-to-one (the index is always equal the token_id it corresponds to). The indices returned by this method allow us to do that.

`fastvideo.models.encoders.llama.VocabParallelEmbedding.weight_loader`

weight_loader(param: Parameter, loaded_weight: torch.Tensor)

`fastvideo.models.encoders.llama.VocabParallelEmbedding.forward`

forward(input_)

`fastvideo.models.encoders.llama.VocabParallelEmbedding.extra_repr`

extra_repr() -> str

`fastvideo.models.encoders.llama.TextEncoder`

TextEncoder(config: TextEncoderConfig)

Base for native encoders with a model-specific forward output contract.

`fastvideo.models.encoders.llama.TextEncoder.forward`

forward(*args: Any, **kwargs: Any) -> TextEncoderOutputT

`fastvideo.models.encoders.llama.default_weight_loader`

default_weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None

Default weight loader.

`fastvideo.models.encoders.llama.maybe_remap_kv_scale_name`

maybe_remap_kv_scale_name(name: str, params_dict: dict) -> str | None

Remap the name of FP8 k/v_scale parameters.

This function handles the remapping of FP8 k/v_scale parameter names. It detects if the given name ends with a suffix and attempts to remap it to the expected name format in the model. If the remapped name is not found in the params_dict, a warning is printed and None is returned.

Parameters

name: str: The original loaded checkpoint parameter name.
params_dict: dict: Dictionary containing the model's named parameters.

Returns

str: str | None: The remapped parameter name if successful, or the original name if no remapping is needed.
None: str | None: If the remapped name is not found in params_dict.

`fastvideo.models.encoders.llama.LlamaMLP`

LlamaMLP(hidden_size: int, intermediate_size: int, hidden_act: str, quant_config: QuantizationConfig | None = None, bias: bool = False, prefix: str = '')

Source

class LlamaMLP(nn.Module):

    def __init__(
        self,
        hidden_size: int,
        intermediate_size: int,
        hidden_act: str,
        quant_config: QuantizationConfig | None = None,
        bias: bool = False,
        prefix: str = "",
    ) -> None:
        super().__init__()
        self.gate_up_proj = MergedColumnParallelLinear(
            input_size=hidden_size,
            output_sizes=[intermediate_size] * 2,
            # output_size=intermediate_size,
            bias=bias,
            quant_config=quant_config,
            prefix=f"{prefix}.gate_up_proj",
        )
        self.down_proj = RowParallelLinear(
            input_size=intermediate_size,
            output_size=hidden_size,
            bias=bias,
            quant_config=quant_config,
            prefix=f"{prefix}.down_proj",
        )
        if hidden_act != "silu":
            raise ValueError(f"Unsupported activation: {hidden_act}. "
                             "Only silu is supported for now.")
        self.act_fn = SiluAndMul()

    def forward(self, x):
        x, _ = self.gate_up_proj(x)
        x = self.act_fn(x)
        x, _ = self.down_proj(x)
        return x

`fastvideo.models.encoders.llama.LlamaMLP.forward`

forward(x)

Source

def forward(self, x):
    x, _ = self.gate_up_proj(x)
    x = self.act_fn(x)
    x, _ = self.down_proj(x)
    return x

`fastvideo.models.encoders.llama.LlamaAttention`

LlamaAttention(config: LlamaConfig, hidden_size: int, num_heads: int, num_kv_heads: int, rope_theta: float = 10000, rope_scaling: dict[str, Any] | None = None, max_position_embeddings: int = 8192, quant_config: QuantizationConfig | None = None, bias: bool = False, bias_o_proj: bool = False, prefix: str = '')

Source

class LlamaAttention(nn.Module):

    def __init__(self,
                 config: LlamaConfig,
                 hidden_size: int,
                 num_heads: int,
                 num_kv_heads: int,
                 rope_theta: float = 10000,
                 rope_scaling: dict[str, Any] | None = None,
                 max_position_embeddings: int = 8192,
                 quant_config: QuantizationConfig | None = None,
                 bias: bool = False,
                 bias_o_proj: bool = False,
                 prefix: str = "") -> None:
        super().__init__()
        # layer_idx = extract_layer_index(prefix)
        self.hidden_size = hidden_size
        tp_size = get_tp_world_size()
        self.total_num_heads = num_heads
        assert self.total_num_heads % tp_size == 0
        self.num_heads = self.total_num_heads // tp_size
        self.total_num_kv_heads = num_kv_heads
        if self.total_num_kv_heads >= tp_size:
            # Number of KV heads is greater than TP size, so we partition
            # the KV heads across multiple tensor parallel GPUs.
            assert self.total_num_kv_heads % tp_size == 0
        else:
            # Number of KV heads is less than TP size, so we replicate
            # the KV heads across multiple tensor parallel GPUs.
            assert tp_size % self.total_num_kv_heads == 0
        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
        # MistralConfig has an optional head_dim introduced by Mistral-Nemo
        self.head_dim = getattr(config, "head_dim", self.hidden_size // self.total_num_heads)
        # Phi models introduced a partial_rotary_factor parameter in the config
        partial_rotary_factor = getattr(config, "partial_rotary_factor", 1)
        self.rotary_dim = int(partial_rotary_factor * self.head_dim)
        self.q_size = self.num_heads * self.head_dim
        self.kv_size = self.num_kv_heads * self.head_dim
        self.scaling = self.head_dim**-0.5
        self.rope_theta = rope_theta
        self.max_position_embeddings = max_position_embeddings

        self.qkv_proj = QKVParallelLinear(
            hidden_size=hidden_size,
            head_size=self.head_dim,
            total_num_heads=self.total_num_heads,
            total_num_kv_heads=self.total_num_kv_heads,
            bias=bias,
            quant_config=quant_config,
            prefix=f"{prefix}.qkv_proj",
        )

        self.o_proj = RowParallelLinear(
            input_size=self.total_num_heads * self.head_dim,
            output_size=hidden_size,
            bias=bias_o_proj,
            quant_config=quant_config,
            prefix=f"{prefix}.o_proj",
        )

        is_neox_style = True
        is_gguf = quant_config and hasattr(quant_config, "get_name") and quant_config.get_name() == "gguf"
        if is_gguf and config.model_type == "llama":
            is_neox_style = False

        self.rotary_emb = get_rope(
            self.head_dim,
            rotary_dim=self.rotary_dim,
            max_position=max_position_embeddings,
            base=int(rope_theta),
            rope_scaling=rope_scaling,
            is_neox_style=is_neox_style,
        )

        self.attn = LocalAttention(self.num_heads,
                                   self.head_dim,
                                   self.num_kv_heads,
                                   softmax_scale=self.scaling,
                                   causal=True,
                                   supported_attention_backends=config._supported_attention_backends)

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor:
        qkv, _ = self.qkv_proj(hidden_states)
        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
        q, k = self.rotary_emb(positions, q, k)
        # attn_output = self.attn(q, k, v)
        # use flash_attn_func
        # TODO (Attn abstraction and backend)
        # from flash_attn import flash_attn_func
        # reshape q, k, v to (batch_size, seq_len, num_heads, head_dim)
        batch_size = q.shape[0]
        seq_len = q.shape[1]
        q = q.reshape(batch_size, seq_len, self.num_heads, self.head_dim)
        k = k.reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim)
        v = v.reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim)
        # import pdb; pdb.set_trace()
        # attn_output = flash_attn_func(q, k, v, softmax_scale=self.scaling, causal=True)
        attn_output = self.attn(q, k, v)
        attn_output = attn_output.reshape(batch_size, seq_len, self.num_heads * self.head_dim)

        output, _ = self.o_proj(attn_output)
        return output

`fastvideo.models.encoders.llama.LlamaAttention.forward`

forward(positions: torch.Tensor, hidden_states: torch.Tensor) -> torch.Tensor

Source

def forward(
    self,
    positions: torch.Tensor,
    hidden_states: torch.Tensor,
) -> torch.Tensor:
    qkv, _ = self.qkv_proj(hidden_states)
    q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
    q, k = self.rotary_emb(positions, q, k)
    # attn_output = self.attn(q, k, v)
    # use flash_attn_func
    # TODO (Attn abstraction and backend)
    # from flash_attn import flash_attn_func
    # reshape q, k, v to (batch_size, seq_len, num_heads, head_dim)
    batch_size = q.shape[0]
    seq_len = q.shape[1]
    q = q.reshape(batch_size, seq_len, self.num_heads, self.head_dim)
    k = k.reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim)
    v = v.reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim)
    # import pdb; pdb.set_trace()
    # attn_output = flash_attn_func(q, k, v, softmax_scale=self.scaling, causal=True)
    attn_output = self.attn(q, k, v)
    attn_output = attn_output.reshape(batch_size, seq_len, self.num_heads * self.head_dim)

    output, _ = self.o_proj(attn_output)
    return output

`fastvideo.models.encoders.llama.LlamaDecoderLayer`

LlamaDecoderLayer(config: LlamaConfig, quant_config: QuantizationConfig | None = None, prefix: str = '')

Source

class LlamaDecoderLayer(nn.Module):

    def __init__(
        self,
        config: LlamaConfig,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
    ) -> None:
        super().__init__()
        self.hidden_size = config.hidden_size
        rope_theta = getattr(config, "rope_theta", 10000)
        rope_scaling = getattr(config, "rope_scaling", None)
        if rope_scaling is not None and getattr(config, "original_max_position_embeddings", None):
            rope_scaling["original_max_position_embeddings"] = (config.original_max_position_embeddings)
        max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
        # Support abacusai/Smaug-72B-v0.1 with attention_bias
        # Support internlm/internlm-7b with bias
        attention_bias = getattr(config, "attention_bias", False) or getattr(config, "bias", False)
        bias_o_proj = attention_bias
        # support internlm/internlm3-8b with qkv_bias
        if hasattr(config, 'qkv_bias'):
            attention_bias = config.qkv_bias

        self.self_attn = LlamaAttention(
            config=config,
            hidden_size=self.hidden_size,
            num_heads=config.num_attention_heads,
            num_kv_heads=getattr(config, "num_key_value_heads", config.num_attention_heads),
            rope_theta=rope_theta,
            rope_scaling=rope_scaling,
            max_position_embeddings=max_position_embeddings,
            quant_config=quant_config,
            bias=attention_bias,
            bias_o_proj=bias_o_proj,
            prefix=f"{prefix}.self_attn",
        )
        self.mlp = LlamaMLP(
            hidden_size=self.hidden_size,
            intermediate_size=config.intermediate_size,
            hidden_act=config.hidden_act,
            quant_config=quant_config,
            bias=getattr(config, "mlp_bias", False),
            prefix=f"{prefix}.mlp",
        )
        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        residual: torch.Tensor | None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        # Self Attention
        if residual is None:
            residual = hidden_states
            hidden_states = self.input_layernorm(hidden_states)
        else:
            hidden_states, residual = self.input_layernorm(hidden_states, residual)

        hidden_states = self.self_attn(positions=positions, hidden_states=hidden_states)

        # Fully Connected
        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
        hidden_states = self.mlp(hidden_states)
        return hidden_states, residual

`fastvideo.models.encoders.llama.LlamaDecoderLayer.forward`

forward(positions: torch.Tensor, hidden_states: torch.Tensor, residual: torch.Tensor | None) -> tuple[torch.Tensor, torch.Tensor]

Source

def forward(
    self,
    positions: torch.Tensor,
    hidden_states: torch.Tensor,
    residual: torch.Tensor | None,
) -> tuple[torch.Tensor, torch.Tensor]:
    # Self Attention
    if residual is None:
        residual = hidden_states
        hidden_states = self.input_layernorm(hidden_states)
    else:
        hidden_states, residual = self.input_layernorm(hidden_states, residual)

    hidden_states = self.self_attn(positions=positions, hidden_states=hidden_states)

    # Fully Connected
    hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
    hidden_states = self.mlp(hidden_states)
    return hidden_states, residual

`fastvideo.models.encoders.llama.LlamaModel`

LlamaModel(config: LlamaConfig)

Source

class LlamaModel(TextEncoder):

    def __init__(
        self,
        config: LlamaConfig,
    ):
        super().__init__(config)

        self.config = config
        self.quant_config = self.config.quant_config
        if config.lora_config is not None:
            max_loras = 1
            lora_vocab_size = 1
            if hasattr(config.lora_config, "max_loras"):
                max_loras = config.lora_config.max_loras
            if hasattr(config.lora_config, "lora_extra_vocab_size"):
                lora_vocab_size = config.lora_config.lora_extra_vocab_size
            lora_vocab = lora_vocab_size * max_loras
        else:
            lora_vocab = 0
        self.vocab_size = config.vocab_size + lora_vocab
        self.org_vocab_size = config.vocab_size

        self.embed_tokens = VocabParallelEmbedding(
            self.vocab_size,
            config.hidden_size,
            org_num_embeddings=config.vocab_size,
            quant_config=config.quant_config,
        )

        self.layers = nn.ModuleList([
            LlamaDecoderLayer(config=config, quant_config=config.quant_config, prefix=f"{config.prefix}.layers.{i}")
            for i in range(config.num_hidden_layers)
        ])

        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.embed_tokens(input_ids)

    def forward(
        self,
        input_ids: torch.Tensor | None,
        position_ids: torch.Tensor | None = None,
        attention_mask: torch.Tensor | None = None,
        inputs_embeds: torch.Tensor | None = None,
        output_hidden_states: bool | None = None,
        **kwargs,
    ) -> BaseEncoderOutput:
        output_hidden_states = (output_hidden_states
                                if output_hidden_states is not None else self.config.output_hidden_states)
        hidden_states = inputs_embeds if inputs_embeds is not None else self.get_input_embeddings(input_ids)
        residual = None

        if position_ids is None:
            position_ids = torch.arange(0, hidden_states.shape[1], device=hidden_states.device).unsqueeze(0)

        all_hidden_states: tuple[Any, ...] | None = () if output_hidden_states else None
        for layer in self.layers:
            if all_hidden_states is not None:
                # TODO
                all_hidden_states += (hidden_states, ) if residual is None else (hidden_states + residual, )
            hidden_states, residual = layer(position_ids, hidden_states, residual)

        hidden_states, _ = self.norm(hidden_states, residual)

        # add hidden states from the last decoder layer
        if all_hidden_states is not None:
            all_hidden_states += (hidden_states, )

        # TODO(will): maybe unify the output format with other models and use
        # our own class
        output = BaseEncoderOutput(
            last_hidden_state=hidden_states,
            # past_key_values=past_key_values if use_cache else None,
            hidden_states=all_hidden_states,
            # attentions=all_self_attns,
        )

        return output

    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

        params_dict = dict(self.named_parameters())
        loaded_params: set[str] = set()
        for name, loaded_weight in weights:
            if "rotary_emb.inv_freq" in name:
                continue
            if ("rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name):
                # Models trained using ColossalAI may include these tensors in
                # the checkpoint. Skip them.
                continue
            # if (self.quant_config is not None and
            #     (scale_name := self.quant_config.get_cache_scale(name))):
            #     # Loading kv cache quantization scales
            #     param = params_dict[scale_name]
            #     weight_loader = getattr(param, "weight_loader",
            #                             default_weight_loader)
            #     loaded_weight = (loaded_weight if loaded_weight.dim() == 0 else
            #                      loaded_weight[0])
            #     weight_loader(param, loaded_weight)
            #     loaded_params.add(scale_name)
            #     continue
            if "scale" in name:
                # Remapping the name of FP8 kv-scale.
                kv_scale_name: str | None = maybe_remap_kv_scale_name(name, params_dict)
                if kv_scale_name is None:
                    continue
                else:
                    name = kv_scale_name
            for param_name, weight_name, shard_id in self.config.arch_config.stacked_params_mapping:
                if weight_name not in name:
                    continue
                name = name.replace(weight_name, param_name)
                # Skip loading extra bias for GPTQ models.
                if name.endswith(".bias") and name not in params_dict:
                    continue

                if name not in params_dict:
                    continue

                param = params_dict[name]
                weight_loader = param.weight_loader
                weight_loader(param, loaded_weight, shard_id)
                break
            else:
                # Skip loading extra bias for GPTQ models.
                if name.endswith(".bias") and name not in params_dict:
                    continue

                if name not in params_dict:
                    continue

                param = params_dict[name]
                weight_loader = getattr(param, "weight_loader", default_weight_loader)
                weight_loader(param, loaded_weight)
            loaded_params.add(name)
        return loaded_params

`fastvideo.models.encoders.llama.LlamaModel.get_input_embeddings`

get_input_embeddings(input_ids: torch.Tensor) -> torch.Tensor

Source

def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
    return self.embed_tokens(input_ids)

`fastvideo.models.encoders.llama.LlamaModel.forward`

forward(input_ids: torch.Tensor | None, position_ids: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, inputs_embeds: torch.Tensor | None = None, output_hidden_states: bool | None = None, **kwargs) -> BaseEncoderOutput

Source

def forward(
    self,
    input_ids: torch.Tensor | None,
    position_ids: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    inputs_embeds: torch.Tensor | None = None,
    output_hidden_states: bool | None = None,
    **kwargs,
) -> BaseEncoderOutput:
    output_hidden_states = (output_hidden_states
                            if output_hidden_states is not None else self.config.output_hidden_states)
    hidden_states = inputs_embeds if inputs_embeds is not None else self.get_input_embeddings(input_ids)
    residual = None

    if position_ids is None:
        position_ids = torch.arange(0, hidden_states.shape[1], device=hidden_states.device).unsqueeze(0)

    all_hidden_states: tuple[Any, ...] | None = () if output_hidden_states else None
    for layer in self.layers:
        if all_hidden_states is not None:
            # TODO
            all_hidden_states += (hidden_states, ) if residual is None else (hidden_states + residual, )
        hidden_states, residual = layer(position_ids, hidden_states, residual)

    hidden_states, _ = self.norm(hidden_states, residual)

    # add hidden states from the last decoder layer
    if all_hidden_states is not None:
        all_hidden_states += (hidden_states, )

    # TODO(will): maybe unify the output format with other models and use
    # our own class
    output = BaseEncoderOutput(
        last_hidden_state=hidden_states,
        # past_key_values=past_key_values if use_cache else None,
        hidden_states=all_hidden_states,
        # attentions=all_self_attns,
    )

    return output

`fastvideo.models.encoders.llama.LlamaModel.load_weights`

load_weights(weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]

Source

def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

    params_dict = dict(self.named_parameters())
    loaded_params: set[str] = set()
    for name, loaded_weight in weights:
        if "rotary_emb.inv_freq" in name:
            continue
        if ("rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name):
            # Models trained using ColossalAI may include these tensors in
            # the checkpoint. Skip them.
            continue
        # if (self.quant_config is not None and
        #     (scale_name := self.quant_config.get_cache_scale(name))):
        #     # Loading kv cache quantization scales
        #     param = params_dict[scale_name]
        #     weight_loader = getattr(param, "weight_loader",
        #                             default_weight_loader)
        #     loaded_weight = (loaded_weight if loaded_weight.dim() == 0 else
        #                      loaded_weight[0])
        #     weight_loader(param, loaded_weight)
        #     loaded_params.add(scale_name)
        #     continue
        if "scale" in name:
            # Remapping the name of FP8 kv-scale.
            kv_scale_name: str | None = maybe_remap_kv_scale_name(name, params_dict)
            if kv_scale_name is None:
                continue
            else:
                name = kv_scale_name
        for param_name, weight_name, shard_id in self.config.arch_config.stacked_params_mapping:
            if weight_name not in name:
                continue
            name = name.replace(weight_name, param_name)
            # Skip loading extra bias for GPTQ models.
            if name.endswith(".bias") and name not in params_dict:
                continue

            if name not in params_dict:
                continue

            param = params_dict[name]
            weight_loader = param.weight_loader
            weight_loader(param, loaded_weight, shard_id)
            break
        else:
            # Skip loading extra bias for GPTQ models.
            if name.endswith(".bias") and name not in params_dict:
                continue

            if name not in params_dict:
                continue

            param = params_dict[name]
            weight_loader = getattr(param, "weight_loader", default_weight_loader)
            weight_loader(param, loaded_weight)
        loaded_params.add(name)
    return loaded_params