llama

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</code></pre></details>
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&#91;str, Any&#93; | None = None,
             max_position_embeddings: int = 8192,
             quant_config: QuantizationConfig | None = None,
             bias: bool = False,
             bias_o_proj: bool = False,
             prefix: str = &quot;&quot;) -&gt; 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 &gt;= 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, &quot;head_dim&quot;, self.hidden_size // self.total_num_heads)
    # Phi models introduced a partial_rotary_factor parameter in the config
    partial_rotary_factor = getattr(config, &quot;partial_rotary_factor&quot;, 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&quot;&#123;prefix&#125;.qkv_proj&quot;,
    )

    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&quot;&#123;prefix&#125;.o_proj&quot;,
    )

    is_neox_style = True
    is_gguf = quant_config and hasattr(quant_config, &quot;get_name&quot;) and quant_config.get_name() == &quot;gguf&quot;
    if is_gguf and config.model_type == &quot;llama&quot;:
        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,
) -&gt; torch.Tensor:
    qkv, _ = self.qkv_proj(hidden_states)
    q, k, v = qkv.split(&#91;self.q_size, self.kv_size, self.kv_size&#93;, 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&#91;0&#93;
    seq_len = q.shape&#91;1&#93;
    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</code></pre></details>
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</code></pre></details>
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 = &quot;&quot;,
) -&gt; None:
    super().__init__()
    self.hidden_size = config.hidden_size
    rope_theta = getattr(config, &quot;rope_theta&quot;, 10000)
    rope_scaling = getattr(config, &quot;rope_scaling&quot;, None)
    if rope_scaling is not None and getattr(config, &quot;original_max_position_embeddings&quot;, None):
        rope_scaling&#91;&quot;original_max_position_embeddings&quot;&#93; = (config.original_max_position_embeddings)
    max_position_embeddings = getattr(config, &quot;max_position_embeddings&quot;, 8192)
    # Support abacusai/Smaug-72B-v0.1 with attention_bias
    # Support internlm/internlm-7b with bias
    attention_bias = getattr(config, &quot;attention_bias&quot;, False) or getattr(config, &quot;bias&quot;, False)
    bias_o_proj = attention_bias
    # support internlm/internlm3-8b with qkv_bias
    if hasattr(config, &#x27;qkv_bias&#x27;):
        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, &quot;num_key_value_heads&quot;, 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&quot;&#123;prefix&#125;.self_attn&quot;,
    )
    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, &quot;mlp_bias&quot;, False),
        prefix=f&quot;&#123;prefix&#125;.mlp&quot;,
    )
    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,
) -&gt; tuple&#91;torch.Tensor, torch.Tensor&#93;:
    # 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</code></pre></details>
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</code></pre></details>
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, &quot;max_loras&quot;):
            max_loras = config.lora_config.max_loras
        if hasattr(config.lora_config, &quot;lora_extra_vocab_size&quot;):
            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(&#91;
        LlamaDecoderLayer(config=config, quant_config=config.quant_config, prefix=f&quot;&#123;config.prefix&#125;.layers.&#123;i&#125;&quot;)
        for i in range(config.num_hidden_layers)
    &#93;)

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

def get_input_embeddings(self, input_ids: torch.Tensor) -&gt; 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,
) -&gt; 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&#91;1&#93;, device=hidden_states.device).unsqueeze(0)

    all_hidden_states: tuple&#91;Any, ...&#93; | 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&#91;tuple&#91;str, torch.Tensor&#93;&#93;) -&gt; set&#91;str&#93;:

    params_dict = dict(self.named_parameters())
    loaded_params: set&#91;str&#93; = set()
    for name, loaded_weight in weights:
        if &quot;rotary_emb.inv_freq&quot; in name:
            continue
        if (&quot;rotary_emb.cos_cached&quot; in name or &quot;rotary_emb.sin_cached&quot; 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&#91;scale_name&#93;
        #     weight_loader = getattr(param, &quot;weight_loader&quot;,
        #                             default_weight_loader)
        #     loaded_weight = (loaded_weight if loaded_weight.dim() == 0 else
        #                      loaded_weight&#91;0&#93;)
        #     weight_loader(param, loaded_weight)
        #     loaded_params.add(scale_name)
        #     continue
        if &quot;scale&quot; 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(&quot;.bias&quot;) and name not in params_dict:
                continue

            if name not in params_dict:
                continue

            param = params_dict&#91;name&#93;
            weight_loader = param.weight_loader
            weight_loader(param, loaded_weight, shard_id)
            break
        else:
            # Skip loading extra bias for GPTQ models.
            if name.endswith(&quot;.bias&quot;) and name not in params_dict:
                continue

            if name not in params_dict:
                continue

            param = params_dict&#91;name&#93;
            weight_loader = getattr(param, &quot;weight_loader&quot;, default_weight_loader)
            weight_loader(param, loaded_weight)
        loaded_params.add(name)
    return loaded_params</code></pre></details>
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&#91;1&#93;, device=hidden_states.device).unsqueeze(0)

all_hidden_states: tuple&#91;Any, ...&#93; | 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</code></pre></details>
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&#91;str&#93; = set()
for name, loaded_weight in weights:
    if &quot;rotary_emb.inv_freq&quot; in name:
        continue
    if (&quot;rotary_emb.cos_cached&quot; in name or &quot;rotary_emb.sin_cached&quot; 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&#91;scale_name&#93;
    #     weight_loader = getattr(param, &quot;weight_loader&quot;,
    #                             default_weight_loader)
    #     loaded_weight = (loaded_weight if loaded_weight.dim() == 0 else
    #                      loaded_weight&#91;0&#93;)
    #     weight_loader(param, loaded_weight)
    #     loaded_params.add(scale_name)
    #     continue
    if &quot;scale&quot; 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(&quot;.bias&quot;) and name not in params_dict:
            continue

        if name not in params_dict:
            continue

        param = params_dict&#91;name&#93;
        weight_loader = param.weight_loader
        weight_loader(param, loaded_weight, shard_id)
        break
    else:
        # Skip loading extra bias for GPTQ models.
        if name.endswith(&quot;.bias&quot;) and name not in params_dict:
            continue

        if name not in params_dict:
            continue

        param = params_dict&#91;name&#93;
        weight_loader = getattr(param, &quot;weight_loader&quot;, default_weight_loader)
        weight_loader(param, loaded_weight)
    loaded_params.add(name)
return loaded_params</code></pre></details>