Skip to content

vllm_omni.diffusion.models.t5_encoder.t5_encoder

T5Block

Bases: Module

layer instance-attribute 

layer = nn.ModuleList(
    [
        T5LayerSelfAttention(
            config,
            has_relative_attention_bias,
            prefix=f"{prefix}.layer.0",
        ),
        T5LayerFF(config, prefix=f"{prefix}.layer.1"),
    ]
)

forward 

forward(
    hidden_states: Tensor,
    mask: Tensor | None = None,
    position_bias: Tensor | None = None,
) -> tuple[Tensor, Tensor]

T5DenseActDense

Bases: Module

act instance-attribute 

act = get_act_fn(config.dense_act_fn)

wi instance-attribute 

wi = ColumnParallelLinear(
    config.d_model,
    config.d_ff,
    bias=False,
    gather_output=False,
    return_bias=False,
)

wo instance-attribute 

wo = RowParallelLinear(
    config.d_ff,
    config.d_model,
    bias=False,
    input_is_parallel=True,
    return_bias=False,
)

forward 

forward(hidden_states: Tensor) -> Tensor

T5DenseGatedActDense

Bases: Module

act instance-attribute 

act = get_act_fn(config.dense_act_fn)

wi instance-attribute 

wi = MergedColumnParallelLinear(
    config.d_model,
    [config.d_ff, config.d_ff],
    bias=False,
    gather_output=False,
)

wo instance-attribute 

wo = RowParallelLinear(
    config.d_ff,
    config.d_model,
    bias=False,
    input_is_parallel=True,
    return_bias=False,
)

forward 

forward(hidden_states: Tensor) -> Tensor

T5EncoderModel

Bases: Module

T5 encoder model applying upstream vLLM layers

config instance-attribute 

config = config

device property 

device: device

dtype property 

dtype: dtype

encoder instance-attribute 

encoder = T5Stack(
    config, self.shared, prefix=f"{prefix}.encoder"
)

prefix instance-attribute 

prefix = prefix

shared instance-attribute 

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

embed_input_ids 

embed_input_ids(input_ids: Tensor) -> Tensor

forward 

forward(
    input_ids: Tensor, attention_mask: Tensor | None = None
) -> tuple[Tensor, ...]

load_weights 

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

T5LayerFF

Bases: Module

DenseReluDense instance-attribute 

DenseReluDense = T5DenseGatedActDense(
    config, prefix=f"{prefix}.DenseReluDense"
)

layer_norm instance-attribute 

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

forward 

forward(hidden_states: Tensor) -> Tensor

T5LayerSelfAttention

Bases: Module

SelfAttention instance-attribute 

SelfAttention = T5SelfAttention(
    config,
    has_relative_attention_bias,
    prefix=f"{prefix}.SelfAttention",
)

layer_norm instance-attribute 

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

forward 

forward(
    hidden_states: Tensor,
    mask: Tensor | None = None,
    position_bias: Tensor | None = None,
) -> tuple[Tensor, Tensor]

T5SelfAttention

Bases: Module

d_kv instance-attribute 

d_kv = config.d_kv

d_model instance-attribute 

d_model = config.d_model

has_relative_attention_bias instance-attribute 

has_relative_attention_bias = has_relative_attention_bias

inner_dim instance-attribute 

inner_dim = self.n_heads * self.d_kv

n_heads instance-attribute 

n_heads = config.num_heads

n_heads_per_partition instance-attribute 

n_heads_per_partition = self.n_heads // tp_size

o instance-attribute 

o = RowParallelLinear(
    self.inner_dim,
    self.d_model,
    bias=False,
    input_is_parallel=True,
    return_bias=False,
)

qkv_proj instance-attribute 

qkv_proj = QKVParallelLinear(
    hidden_size=self.d_model,
    head_size=self.d_kv,
    total_num_heads=self.n_heads,
    total_num_kv_heads=self.n_heads,
    bias=False,
)

relative_attention_bias instance-attribute 

relative_attention_bias = nn.Embedding(
    self.relative_attention_num_buckets, self.n_heads
)

relative_attention_max_distance instance-attribute 

relative_attention_max_distance = (
    config.relative_attention_max_distance
)

relative_attention_num_buckets instance-attribute 

relative_attention_num_buckets = (
    config.relative_attention_num_buckets
)

compute_bias 

compute_bias(
    query_length: int, key_length: int, device: device
) -> Tensor

Compute relative position bias, returning only the local head shard.

forward 

forward(
    hidden_states: Tensor,
    mask: Tensor | None = None,
    position_bias: Tensor | None = None,
) -> tuple[Tensor, Tensor]

T5Stack

Bases: Module

block instance-attribute 

block = nn.ModuleList(
    [
        T5Block(
            config,
            has_relative_attention_bias=i == 0,
            prefix=f"{prefix}.block.{i}",
        )
        for i in range(config.num_layers)
    ]
)

embed_tokens instance-attribute 

embed_tokens = shared

final_layer_norm instance-attribute 

final_layer_norm = RMSNorm(
    config.d_model, eps=config.layer_norm_epsilon
)

forward 

forward(
    input_ids: Tensor, attention_mask: Tensor | None = None
) -> Tensor