PyTorch T5 & UMT5 model.

`fastvideo.models.encoders.t5.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.t5.T5Config`

T5Config(arch_config: TextEncoderArchConfig = T5ArchConfig(), prefix: str = 't5', 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.t5.get_tp_rank`

get_tp_rank() -> int

Return my rank for the tensor model parallel group.

`fastvideo.models.encoders.t5.get_tp_world_size`

get_tp_world_size() -> int

Return world size for the tensor model parallel group.

`fastvideo.models.encoders.t5.get_act_fn`

get_act_fn(act_fn_name: str) -> nn.Module

Get an activation function by name.

`fastvideo.models.encoders.t5.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.t5.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.t5.RMSNorm.extra_repr`

extra_repr() -> str

`fastvideo.models.encoders.t5.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.t5.MergedColumnParallelLinear.weight_loader`

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

`fastvideo.models.encoders.t5.MergedColumnParallelLinear.weight_loader_v2`

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

`fastvideo.models.encoders.t5.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.t5.QKVParallelLinear.weight_loader_v2`

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

`fastvideo.models.encoders.t5.QKVParallelLinear.weight_loader`

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

`fastvideo.models.encoders.t5.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.t5.RowParallelLinear.weight_loader`

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

`fastvideo.models.encoders.t5.RowParallelLinear.weight_loader_v2`

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

`fastvideo.models.encoders.t5.RowParallelLinear.forward`

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

`fastvideo.models.encoders.t5.RowParallelLinear.extra_repr`

extra_repr() -> str

`fastvideo.models.encoders.t5.QuantizationConfig`

QuantizationConfig()

Base class for quantization configs.

`fastvideo.models.encoders.t5.QuantizationConfig.get_name`

get_name() -> QuantizationMethods

Name of the quantization method.

`fastvideo.models.encoders.t5.QuantizationConfig.get_supported_act_dtypes`

get_supported_act_dtypes() -> list[torch.dtype]

List of supported activation dtypes.

`fastvideo.models.encoders.t5.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.t5.QuantizationConfig.get_config_filenames`

get_config_filenames() -> list[str]

List of filenames to search for in the model directory.

`fastvideo.models.encoders.t5.QuantizationConfig.from_config`

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

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

`fastvideo.models.encoders.t5.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.t5.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.t5.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.t5.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.t5.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.t5.QuantizationConfig.get_cache_scale`

get_cache_scale(name: str) -> str | None

`fastvideo.models.encoders.t5.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.t5.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.t5.VocabParallelEmbedding.weight_loader`

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

`fastvideo.models.encoders.t5.VocabParallelEmbedding.forward`

forward(input_)

`fastvideo.models.encoders.t5.VocabParallelEmbedding.extra_repr`

extra_repr() -> str

`fastvideo.models.encoders.t5.TextEncoder`

TextEncoder(config: TextEncoderConfig)

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

`fastvideo.models.encoders.t5.TextEncoder.forward`

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

`fastvideo.models.encoders.t5.default_weight_loader`

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

Default weight loader.

`fastvideo.models.encoders.t5.AttentionType`

AttentionType()

Attention type. Use string to be compatible with torch.compile.

Source

class AttentionType:
    """
    Attention type.
    Use string to be compatible with `torch.compile`.
    """

    # Decoder attention between previous layer Q/K/V
    DECODER = "decoder"
    # Encoder attention between previous layer Q/K/V for encoder-decoder
    ENCODER = "encoder"
    # Encoder attention between previous layer Q/K/V
    ENCODER_ONLY = "encoder_only"
    # Attention between dec. Q and enc. K/V for encoder-decoder
    ENCODER_DECODER = "encoder_decoder"

`fastvideo.models.encoders.t5.AttentionMetadata`

AttentionMetadata(attn_bias: torch.Tensor)

Source

@dataclass
class AttentionMetadata:
    attn_bias: torch.Tensor

`fastvideo.models.encoders.t5.T5DenseActDense`

T5DenseActDense(config: T5Config, quant_config: QuantizationConfig | None = None)

Source

class T5DenseActDense(nn.Module):

    def __init__(self, config: T5Config, quant_config: QuantizationConfig | None = None):
        super().__init__()
        self.wi = MergedColumnParallelLinear(config.d_model, [config.d_ff], bias=False)
        self.wo = RowParallelLinear(config.d_ff, config.d_model, bias=False, quant_config=quant_config)
        self.act = get_act_fn(config.dense_act_fn)

    def forward(self, hidden_states) -> torch.Tensor:
        hidden_states, _ = self.wi(hidden_states)
        hidden_states = self.act(hidden_states)
        hidden_states, _ = self.wo(hidden_states)
        return hidden_states

`fastvideo.models.encoders.t5.T5DenseActDense.forward`

forward(hidden_states) -> torch.Tensor

Source

def forward(self, hidden_states) -> torch.Tensor:
    hidden_states, _ = self.wi(hidden_states)
    hidden_states = self.act(hidden_states)
    hidden_states, _ = self.wo(hidden_states)
    return hidden_states

`fastvideo.models.encoders.t5.T5DenseGatedActDense`

T5DenseGatedActDense(config: T5Config, quant_config: QuantizationConfig | None = None)

Source

class T5DenseGatedActDense(nn.Module):

    def __init__(self, config: T5Config, quant_config: QuantizationConfig | None = None):
        super().__init__()
        self.wi_0 = MergedColumnParallelLinear(config.d_model, [config.d_ff], bias=False, quant_config=quant_config)
        self.wi_1 = MergedColumnParallelLinear(config.d_model, [config.d_ff], bias=False, quant_config=quant_config)
        # Should not run in fp16 unless mixed-precision is used,
        # see https://github.com/huggingface/transformers/issues/20287.
        self.wo = RowParallelLinear(config.d_ff, config.d_model, bias=False, quant_config=quant_config)
        self.act = get_act_fn(config.dense_act_fn)

    def forward(self, hidden_states) -> torch.Tensor:
        hidden_gelu = self.act(self.wi_0(hidden_states)[0])
        hidden_linear, _ = self.wi_1(hidden_states)
        hidden_states = hidden_gelu * hidden_linear
        hidden_states, _ = self.wo(hidden_states)
        return hidden_states

`fastvideo.models.encoders.t5.T5DenseGatedActDense.forward`

forward(hidden_states) -> torch.Tensor

Source

def forward(self, hidden_states) -> torch.Tensor:
    hidden_gelu = self.act(self.wi_0(hidden_states)[0])
    hidden_linear, _ = self.wi_1(hidden_states)
    hidden_states = hidden_gelu * hidden_linear
    hidden_states, _ = self.wo(hidden_states)
    return hidden_states

`fastvideo.models.encoders.t5.T5LayerFF`

T5LayerFF(config: T5Config, quant_config: QuantizationConfig | None = None)

Source

class T5LayerFF(nn.Module):

    def __init__(self, config: T5Config, quant_config: QuantizationConfig | None = None):
        super().__init__()
        if config.is_gated_act:
            self.DenseReluDense = T5DenseGatedActDense(config, quant_config=quant_config)
        else:
            self.DenseReluDense = T5DenseActDense(config, quant_config=quant_config)

        self.layer_norm = RMSNorm(config.d_model, eps=config.layer_norm_epsilon)

    def forward(self, hidden_states) -> torch.Tensor:
        forwarded_states = self.layer_norm(hidden_states)
        forwarded_states = self.DenseReluDense(forwarded_states)
        hidden_states = hidden_states + forwarded_states
        return hidden_states

`fastvideo.models.encoders.t5.T5LayerFF.forward`

forward(hidden_states) -> torch.Tensor

Source

def forward(self, hidden_states) -> torch.Tensor:
    forwarded_states = self.layer_norm(hidden_states)
    forwarded_states = self.DenseReluDense(forwarded_states)
    hidden_states = hidden_states + forwarded_states
    return hidden_states

`fastvideo.models.encoders.t5.T5MultiHeadAttention`

T5MultiHeadAttention()

Source

class T5MultiHeadAttention(nn.Module):

    def __init__(self) -> None:
        super().__init__()

    def forward(self, q, k, v, attn_bias=None):
        b, _, n, c = q.shape
        attn = torch.einsum("binc,bjnc->bnij", q, k)
        if attn_bias is not None:
            attn += attn_bias

        attn = F.softmax(attn.float(), dim=-1).type_as(attn)
        x = torch.einsum("bnij,bjnc->binc", attn, v)
        x = x.reshape(b, -1, n * c)
        return x

`fastvideo.models.encoders.t5.T5MultiHeadAttention.forward`

forward(q, k, v, attn_bias = None)

Source

def forward(self, q, k, v, attn_bias=None):
    b, _, n, c = q.shape
    attn = torch.einsum("binc,bjnc->bnij", q, k)
    if attn_bias is not None:
        attn += attn_bias

    attn = F.softmax(attn.float(), dim=-1).type_as(attn)
    x = torch.einsum("bnij,bjnc->binc", attn, v)
    x = x.reshape(b, -1, n * c)
    return x

`fastvideo.models.encoders.t5.T5Attention`

T5Attention(config: T5Config, attn_type: str, has_relative_attention_bias = False, quant_config: QuantizationConfig | None = None, prefix: str = '')

Source

class T5Attention(nn.Module):

    def __init__(
        self,
        config: T5Config,
        attn_type: str,
        has_relative_attention_bias=False,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
    ):
        super().__init__()
        self.attn_type = attn_type
        # Cross-attention has no relative pos encoding anyway
        self.is_decoder = attn_type == AttentionType.DECODER
        self.has_relative_attention_bias = has_relative_attention_bias
        self.relative_attention_num_buckets = (config.relative_attention_num_buckets)
        self.relative_attention_max_distance = (config.relative_attention_max_distance)
        self.d_model = config.d_model
        self.key_value_proj_dim = config.d_kv
        self.total_num_heads = self.total_num_kv_heads = config.num_heads

        # Partition heads across multiple tensor parallel GPUs.
        tp_world_size = get_tp_world_size()
        assert config.num_heads % tp_world_size == 0
        self.n_heads = config.num_heads // tp_world_size

        self.inner_dim = self.n_heads * self.key_value_proj_dim
        # No GQA in t5.
        # self.n_kv_heads = self.n_heads

        self.qkv_proj = QKVParallelLinear(
            self.d_model,
            self.key_value_proj_dim,
            self.total_num_heads,
            self.total_num_kv_heads,
            bias=False,
            quant_config=quant_config,
            prefix=f"{prefix}.qkv_proj",
        )

        self.attn = T5MultiHeadAttention()

        if self.has_relative_attention_bias:
            self.relative_attention_bias = VocabParallelEmbedding(
                self.relative_attention_num_buckets,
                self.total_num_heads,
                org_num_embeddings=self.relative_attention_num_buckets,
                padding_size=self.relative_attention_num_buckets,
                quant_config=quant_config,
            )
        self.o = RowParallelLinear(
            self.total_num_heads * self.key_value_proj_dim,
            self.d_model,
            bias=False,
            quant_config=quant_config,
            prefix=f"{prefix}.o_proj",
        )

    @staticmethod
    def _relative_position_bucket(relative_position,
                                  bidirectional=True,
                                  num_buckets=32,
                                  max_distance=128) -> torch.Tensor:
        """
        Adapted from Mesh Tensorflow:
        https://github.com/tensorflow/mesh/blob/0cb87fe07da627bf0b7e60475d59f95ed6b5be3d/mesh_tensorflow/transformer/transformer_layers.py#L593
        Translate relative position to a bucket number for relative attention.
        The relative position is defined as memory_position - query_position,
        i.e. the distance in tokens from the attending position to the
        attended-to position. If bidirectional=False, then positive relative
        positions are invalid. We use smaller buckets for small absolute
        relative_position and larger buckets for larger absolute
        relative_positions. All relative positions >=max_distance map to the
        same bucket. All relative positions <=-max_distance map to the same
        bucket. This should allow for more graceful generalization to longer
        sequences than the model has been trained on
        Args:
            relative_position: an int32 Tensor
            bidirectional: a boolean - whether the attention is bidirectional
            num_buckets: an integer
            max_distance: an integer
        Returns:
            a Tensor with the same shape as relative_position, containing int32
            values in the range [0, num_buckets)
        """  # noqa: E501
        relative_buckets = 0
        if bidirectional:
            num_buckets //= 2
            relative_buckets += (relative_position > 0).to(torch.long) * num_buckets
            relative_position = torch.abs(relative_position)
        else:
            relative_position = -torch.min(relative_position, torch.zeros_like(relative_position))
        # now relative_position is in the range [0, inf)

        # half of the buckets are for exact increments in positions
        max_exact = num_buckets // 2
        is_small = relative_position < max_exact

        # The other half of the buckets are for logarithmically bigger bins
        # in positions up to max_distance
        relative_position_if_large = max_exact + (torch.log(relative_position.float() / max_exact) /
                                                  math.log(max_distance / max_exact) *
                                                  (num_buckets - max_exact)).to(torch.long)
        relative_position_if_large = torch.min(
            relative_position_if_large,
            torch.full_like(relative_position_if_large, num_buckets - 1),
        )

        relative_buckets += torch.where(is_small, relative_position, relative_position_if_large)
        return relative_buckets

    def compute_bias(self, query_length, key_length, device=None) -> torch.Tensor:
        """Compute binned relative position bias"""
        if device is None:
            device = self.relative_attention_bias.weight.device
        context_position = torch.arange(query_length, dtype=torch.long, device=device)[:, None]
        memory_position = torch.arange(key_length, dtype=torch.long, device=device)[None, :]
        # max_seq_len, nh
        relative_position = memory_position - context_position
        relative_position_bucket = self._relative_position_bucket(
            relative_position,  # shape (query_length, key_length)
            bidirectional=(not self.is_decoder),
            num_buckets=self.relative_attention_num_buckets,
            max_distance=self.relative_attention_max_distance,
        )
        values = self.relative_attention_bias(relative_position_bucket)  # shape (query_length, key_length, num_heads)
        x = values.permute([2, 0, 1]).unsqueeze(0)  # shape (1, num_heads, query_length, key_length)
        return x

    def forward(
        self,
        hidden_states: torch.Tensor,  # (num_tokens, d_model)
        attention_mask: torch.Tensor,
        attn_metadata: AttentionMetadata | None = None,
    ) -> torch.Tensor:
        bs, seq_len, _ = hidden_states.shape
        num_seqs = bs
        n, c = self.n_heads, self.key_value_proj_dim
        qkv, _ = self.qkv_proj(hidden_states)
        # Projection of 'own' hidden state (self-attention). No GQA here.
        q, k, v = qkv.split(self.inner_dim, dim=-1)
        q = q.reshape(bs, seq_len, n, c)
        k = k.reshape(bs, seq_len, n, c)
        v = v.reshape(bs, seq_len, n, c)

        assert attn_metadata is not None
        attn_bias = attn_metadata.attn_bias
        # Not compatible with CP here (as all encoder-decoder models),
        # as it assumes homogeneous batch (prefills or decodes).
        if self.has_relative_attention_bias:
            # Self-attention. Compute T5 relative positional encoding.
            # The bias term is computed on longest sequence in batch. Biases
            # for shorter sequences are slices of the longest.
            assert self.attn_type == AttentionType.ENCODER
            attn_bias = self.compute_bias(seq_len, seq_len).repeat(num_seqs, 1, 1, 1)
            attn_metadata.attn_bias = attn_bias
        else:
            # Encoder/Decoder Self-Attention Layer, attn bias already cached.
            assert attn_bias is not None

        from fastvideo.platforms import current_platform

        if attention_mask is not None:
            attention_mask = (attention_mask.view(bs, 1, 1, -1)
                              if attention_mask.ndim == 2 else attention_mask.unsqueeze(1))
            mask_val = (-1e4 if current_platform.is_mps() else torch.finfo(q.dtype).min)
            attn_bias.masked_fill_(attention_mask == 0, mask_val)

        if get_tp_world_size() > 1:
            rank = get_tp_rank()
            attn_bias = attn_bias[:, rank * self.n_heads:(rank + 1) * self.n_heads, :, :]
        attn_output = self.attn(q, k, v, attn_bias)
        output, _ = self.o(attn_output)
        return output

`fastvideo.models.encoders.t5.T5Attention.compute_bias`

compute_bias(query_length, key_length, device = None) -> torch.Tensor

Compute binned relative position bias

Source

def compute_bias(self, query_length, key_length, device=None) -> torch.Tensor:
    """Compute binned relative position bias"""
    if device is None:
        device = self.relative_attention_bias.weight.device
    context_position = torch.arange(query_length, dtype=torch.long, device=device)[:, None]
    memory_position = torch.arange(key_length, dtype=torch.long, device=device)[None, :]
    # max_seq_len, nh
    relative_position = memory_position - context_position
    relative_position_bucket = self._relative_position_bucket(
        relative_position,  # shape (query_length, key_length)
        bidirectional=(not self.is_decoder),
        num_buckets=self.relative_attention_num_buckets,
        max_distance=self.relative_attention_max_distance,
    )
    values = self.relative_attention_bias(relative_position_bucket)  # shape (query_length, key_length, num_heads)
    x = values.permute([2, 0, 1]).unsqueeze(0)  # shape (1, num_heads, query_length, key_length)
    return x

`fastvideo.models.encoders.t5.T5Attention.forward`

forward(hidden_states: torch.Tensor, attention_mask: torch.Tensor, attn_metadata: AttentionMetadata | None = None) -> torch.Tensor

Source

def forward(
    self,
    hidden_states: torch.Tensor,  # (num_tokens, d_model)
    attention_mask: torch.Tensor,
    attn_metadata: AttentionMetadata | None = None,
) -> torch.Tensor:
    bs, seq_len, _ = hidden_states.shape
    num_seqs = bs
    n, c = self.n_heads, self.key_value_proj_dim
    qkv, _ = self.qkv_proj(hidden_states)
    # Projection of 'own' hidden state (self-attention). No GQA here.
    q, k, v = qkv.split(self.inner_dim, dim=-1)
    q = q.reshape(bs, seq_len, n, c)
    k = k.reshape(bs, seq_len, n, c)
    v = v.reshape(bs, seq_len, n, c)

    assert attn_metadata is not None
    attn_bias = attn_metadata.attn_bias
    # Not compatible with CP here (as all encoder-decoder models),
    # as it assumes homogeneous batch (prefills or decodes).
    if self.has_relative_attention_bias:
        # Self-attention. Compute T5 relative positional encoding.
        # The bias term is computed on longest sequence in batch. Biases
        # for shorter sequences are slices of the longest.
        assert self.attn_type == AttentionType.ENCODER
        attn_bias = self.compute_bias(seq_len, seq_len).repeat(num_seqs, 1, 1, 1)
        attn_metadata.attn_bias = attn_bias
    else:
        # Encoder/Decoder Self-Attention Layer, attn bias already cached.
        assert attn_bias is not None

    from fastvideo.platforms import current_platform

    if attention_mask is not None:
        attention_mask = (attention_mask.view(bs, 1, 1, -1)
                          if attention_mask.ndim == 2 else attention_mask.unsqueeze(1))
        mask_val = (-1e4 if current_platform.is_mps() else torch.finfo(q.dtype).min)
        attn_bias.masked_fill_(attention_mask == 0, mask_val)

    if get_tp_world_size() > 1:
        rank = get_tp_rank()
        attn_bias = attn_bias[:, rank * self.n_heads:(rank + 1) * self.n_heads, :, :]
    attn_output = self.attn(q, k, v, attn_bias)
    output, _ = self.o(attn_output)
    return output

`fastvideo.models.encoders.t5.T5LayerSelfAttention`

T5LayerSelfAttention(config, has_relative_attention_bias = False, quant_config: QuantizationConfig | None = None, prefix: str = '')

Source

class T5LayerSelfAttention(nn.Module):

    def __init__(
        self,
        config,
        has_relative_attention_bias=False,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
    ):
        super().__init__()
        self.SelfAttention = T5Attention(
            config,
            AttentionType.DECODER if "decoder" in prefix else AttentionType.ENCODER,
            has_relative_attention_bias=has_relative_attention_bias,
            quant_config=quant_config,
            prefix=f"{prefix}.SelfAttention",
        )
        self.layer_norm = RMSNorm(config.d_model, eps=config.layer_norm_epsilon)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: torch.Tensor,
        attn_metadata: AttentionMetadata | None = None,
    ) -> torch.Tensor:
        normed_hidden_states = self.layer_norm(hidden_states)
        attention_output = self.SelfAttention(
            hidden_states=normed_hidden_states,
            attention_mask=attention_mask,
            attn_metadata=attn_metadata,
        )
        hidden_states = hidden_states + attention_output
        return hidden_states

`fastvideo.models.encoders.t5.T5LayerSelfAttention.forward`

forward(hidden_states: torch.Tensor, attention_mask: torch.Tensor, attn_metadata: AttentionMetadata | None = None) -> torch.Tensor

Source

def forward(
    self,
    hidden_states: torch.Tensor,
    attention_mask: torch.Tensor,
    attn_metadata: AttentionMetadata | None = None,
) -> torch.Tensor:
    normed_hidden_states = self.layer_norm(hidden_states)
    attention_output = self.SelfAttention(
        hidden_states=normed_hidden_states,
        attention_mask=attention_mask,
        attn_metadata=attn_metadata,
    )
    hidden_states = hidden_states + attention_output
    return hidden_states

`fastvideo.models.encoders.t5.T5LayerCrossAttention`

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

Source

class T5LayerCrossAttention(nn.Module):

    def __init__(
        self,
        config,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
    ):
        super().__init__()
        self.EncDecAttention = T5Attention(
            config,
            AttentionType.ENCODER_DECODER,
            has_relative_attention_bias=False,
            quant_config=quant_config,
            prefix=f"{prefix}.EncDecAttention",
        )
        self.layer_norm = RMSNorm(config.d_model, eps=config.layer_norm_epsilon)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attn_metadata: AttentionMetadata | None = None,
    ) -> torch.Tensor:
        normed_hidden_states = self.layer_norm(hidden_states)
        attention_output = self.EncDecAttention(
            hidden_states=normed_hidden_states,
            attn_metadata=attn_metadata,
        )
        hidden_states = hidden_states + attention_output
        return hidden_states

`fastvideo.models.encoders.t5.T5LayerCrossAttention.forward`

forward(hidden_states: torch.Tensor, attn_metadata: AttentionMetadata | None = None) -> torch.Tensor

Source

def forward(
    self,
    hidden_states: torch.Tensor,
    attn_metadata: AttentionMetadata | None = None,
) -> torch.Tensor:
    normed_hidden_states = self.layer_norm(hidden_states)
    attention_output = self.EncDecAttention(
        hidden_states=normed_hidden_states,
        attn_metadata=attn_metadata,
    )
    hidden_states = hidden_states + attention_output
    return hidden_states

`fastvideo.models.encoders.t5.T5Block`

T5Block(config: T5Config, is_decoder: bool, has_relative_attention_bias = False, quant_config: QuantizationConfig | None = None, prefix: str = '')

Source

class T5Block(nn.Module):

    def __init__(
        self,
        config: T5Config,
        is_decoder: bool,
        has_relative_attention_bias=False,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
    ):
        super().__init__()
        self.is_decoder = is_decoder
        self.layer = nn.ModuleList()
        self.layer.append(
            T5LayerSelfAttention(
                config,
                has_relative_attention_bias=has_relative_attention_bias,
                quant_config=quant_config,
                prefix=f"{prefix}.self_attn",
            ))

        if self.is_decoder:
            self.layer.append(T5LayerCrossAttention(
                config,
                quant_config=quant_config,
                prefix=f"{prefix}.cross_attn",
            ))

        self.layer.append(T5LayerFF(config, quant_config=quant_config))

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: torch.Tensor,
        attn_metadata: AttentionMetadata | None = None,
    ) -> torch.Tensor:
        hidden_states = self.layer[0](
            hidden_states=hidden_states,
            attention_mask=attention_mask,
            attn_metadata=attn_metadata,
        )
        if self.is_decoder:
            hidden_states = self.layer[1](hidden_states=hidden_states, attn_metadata=attn_metadata)

            # Apply Feed Forward layer
            hidden_states = self.layer[2](hidden_states)
        else:
            hidden_states = self.layer[1](hidden_states)
        return hidden_states

`fastvideo.models.encoders.t5.T5Block.forward`

forward(hidden_states: torch.Tensor, attention_mask: torch.Tensor, attn_metadata: AttentionMetadata | None = None) -> torch.Tensor

Source

def forward(
    self,
    hidden_states: torch.Tensor,
    attention_mask: torch.Tensor,
    attn_metadata: AttentionMetadata | None = None,
) -> torch.Tensor:
    hidden_states = self.layer[0](
        hidden_states=hidden_states,
        attention_mask=attention_mask,
        attn_metadata=attn_metadata,
    )
    if self.is_decoder:
        hidden_states = self.layer[1](hidden_states=hidden_states, attn_metadata=attn_metadata)

        # Apply Feed Forward layer
        hidden_states = self.layer[2](hidden_states)
    else:
        hidden_states = self.layer[1](hidden_states)
    return hidden_states

`fastvideo.models.encoders.t5.T5Stack`

T5Stack(config: T5Config, is_decoder: bool, n_layers: int, embed_tokens = None, quant_config: QuantizationConfig | None = None, prefix: str = '', is_umt5: bool = False)

Source

class T5Stack(nn.Module):

    def __init__(
        self,
        config: T5Config,
        is_decoder: bool,
        n_layers: int,
        embed_tokens=None,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
        is_umt5: bool = False,
    ):
        super().__init__()
        self.embed_tokens = embed_tokens
        self.is_umt5 = is_umt5
        if is_umt5:
            self.block = nn.ModuleList([
                T5Block(
                    config,
                    is_decoder=is_decoder,
                    has_relative_attention_bias=True,
                    quant_config=quant_config,
                    prefix=f"{prefix}.blocks.{i}",
                ) for i in range(n_layers)
            ])
        else:
            # Only the first block has relative positional encoding.
            self.block = nn.ModuleList([
                T5Block(
                    config,
                    is_decoder=is_decoder,
                    has_relative_attention_bias=i == 0,
                    quant_config=quant_config,
                    prefix=f"{prefix}.blocks.{i}",
                ) for i in range(n_layers)
            ])
        self.final_layer_norm = RMSNorm(config.d_model, eps=config.layer_norm_epsilon)

    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor,
        attn_metadata: AttentionMetadata,
    ) -> torch.Tensor:
        hidden_states = self.embed_tokens(input_ids)

        for idx, block in enumerate(self.block):
            hidden_states = block(
                hidden_states=hidden_states,
                attention_mask=attention_mask,
                attn_metadata=attn_metadata,
            )
        hidden_states = self.final_layer_norm(hidden_states)
        return hidden_states

`fastvideo.models.encoders.t5.T5Stack.forward`

forward(input_ids: torch.Tensor, attention_mask: torch.Tensor, attn_metadata: AttentionMetadata) -> torch.Tensor

Source

def forward(
    self,
    input_ids: torch.Tensor,
    attention_mask: torch.Tensor,
    attn_metadata: AttentionMetadata,
) -> torch.Tensor:
    hidden_states = self.embed_tokens(input_ids)

    for idx, block in enumerate(self.block):
        hidden_states = block(
            hidden_states=hidden_states,
            attention_mask=attention_mask,
            attn_metadata=attn_metadata,
        )
    hidden_states = self.final_layer_norm(hidden_states)
    return hidden_states

`fastvideo.models.encoders.t5.T5EncoderModel`

T5EncoderModel(config: T5Config, prefix: str = '')

Source

class T5EncoderModel(TextEncoder):

    def __init__(self, config: T5Config, prefix: str = ""):
        super().__init__(config)

        self.shared = VocabParallelEmbedding(
            config.vocab_size,
            config.d_model,
            org_num_embeddings=config.vocab_size,
        )

        self.encoder = T5Stack(
            config,
            False,
            config.num_layers,
            self.shared,
            quant_config=config.quant_config,
            prefix=f"{prefix}.encoder",
            is_umt5=False,
        )

    def get_input_embeddings(self):
        return self.shared

    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:
        attn_metadata = AttentionMetadata(None)
        hidden_states = self.encoder(
            input_ids=input_ids,
            attention_mask=attention_mask,
            attn_metadata=attn_metadata,
        )

        return BaseEncoderOutput(
            last_hidden_state=hidden_states,
            attention_mask=attention_mask,
        )

    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
        stacked_params_mapping = [
            # (param_name, shard_name, shard_id)
            (".qkv_proj", ".q", "q"),
            (".qkv_proj", ".k", "k"),
            (".qkv_proj", ".v", "v"),
        ]
        params_dict = dict(self.named_parameters())
        loaded_params: set[str] = set()
        for name, loaded_weight in weights:
            loaded = False
            if "decoder" in name or "lm_head" in name:
                continue
            for param_name, weight_name, shard_id in 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)
                loaded = True
                break
            if not loaded:
                # 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.t5.T5EncoderModel.get_input_embeddings`

get_input_embeddings()

Source

def get_input_embeddings(self):
    return self.shared

`fastvideo.models.encoders.t5.T5EncoderModel.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:
    attn_metadata = AttentionMetadata(None)
    hidden_states = self.encoder(
        input_ids=input_ids,
        attention_mask=attention_mask,
        attn_metadata=attn_metadata,
    )

    return BaseEncoderOutput(
        last_hidden_state=hidden_states,
        attention_mask=attention_mask,
    )

`fastvideo.models.encoders.t5.T5EncoderModel.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]:
    stacked_params_mapping = [
        # (param_name, shard_name, shard_id)
        (".qkv_proj", ".q", "q"),
        (".qkv_proj", ".k", "k"),
        (".qkv_proj", ".v", "v"),
    ]
    params_dict = dict(self.named_parameters())
    loaded_params: set[str] = set()
    for name, loaded_weight in weights:
        loaded = False
        if "decoder" in name or "lm_head" in name:
            continue
        for param_name, weight_name, shard_id in 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)
            loaded = True
            break
        if not loaded:
            # 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.t5.UMT5EncoderModel`

UMT5EncoderModel(config: T5Config, prefix: str = '')

Source

class UMT5EncoderModel(TextEncoder):

    def __init__(self, config: T5Config, prefix: str = ""):
        super().__init__(config)

        self.shared = VocabParallelEmbedding(
            config.vocab_size,
            config.d_model,
            org_num_embeddings=config.vocab_size,
        )

        self.encoder = T5Stack(
            config,
            False,
            config.num_layers,
            self.shared,
            quant_config=config.quant_config,
            prefix=f"{prefix}.encoder",
            is_umt5=True,
        )

    def get_input_embeddings(self):
        return self.shared

    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:
        attn_metadata = AttentionMetadata(None)
        hidden_states = self.encoder(
            input_ids=input_ids,
            attention_mask=attention_mask,
            attn_metadata=attn_metadata,
        )

        return BaseEncoderOutput(
            last_hidden_state=hidden_states,
            attention_mask=attention_mask,
        )

    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:
            loaded = False
            if "decoder" in name or "lm_head" in name:
                continue
            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)
                loaded = True
                break
            if not loaded:
                # 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.t5.UMT5EncoderModel.get_input_embeddings`

get_input_embeddings()

Source

def get_input_embeddings(self):
    return self.shared

`fastvideo.models.encoders.t5.UMT5EncoderModel.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:
    attn_metadata = AttentionMetadata(None)
    hidden_states = self.encoder(
        input_ids=input_ids,
        attention_mask=attention_mask,
        attn_metadata=attn_metadata,
    )

    return BaseEncoderOutput(
        last_hidden_state=hidden_states,
        attention_mask=attention_mask,
    )

`fastvideo.models.encoders.t5.UMT5EncoderModel.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:
        loaded = False
        if "decoder" in name or "lm_head" in name:
            continue
        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)
            loaded = True
            break
        if not loaded:
            # 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