vllm_omni.diffusion.models.bagel.bagel_transformer ¶

base_model_prefix = 'bagel'

config = config

config_class = BagelConfig

connector = MLPconnector(
    vit_hidden_size,
    hidden_size,
    connector_act,
    quant_config=quant_config,
    prefix=f"{prefix}.connector",
)

get_flattened_position_ids = (
    get_flattened_position_ids_extrapolate
)

hidden_size = hidden_size

language_model = language_model

latent_channel = z_channels

latent_downsample = downsample * latent_patch_size

latent_patch_size = latent_patch_size

latent_pos_embed = PositionEmbedding(
    max_latent_size, hidden_size
)

llm2vae = Linear(hidden_size, patch_latent_dim)

max_latent_size = max_latent_size

num_heads = num_attention_heads

parallel_config = parallel_config

patch_latent_dim = latent_patch_size ** 2 * latent_channel

time_embedder = TimestepEmbedder(hidden_size)

timestep_shift = timestep_shift

use_moe = 'Mo' in layer_module

vae2llm = Linear(patch_latent_dim, hidden_size)

vit_hidden_size = hidden_size

vit_max_num_patch_per_side = vit_max_num_patch_per_side

vit_model = vit_model

vit_patch_size = patch_size

vit_pos_embed = PositionEmbedding(
    vit_max_num_patch_per_side, hidden_size
)

act_fn = SiLU()

down_proj = RowParallelLinear(
    intermediate_size,
    hidden_size,
    input_is_parallel=True,
    bias=False,
    quant_config=quant_config,
    prefix=f"{prefix}.down_proj",
)

gate_up_proj = MergedColumnParallelLinear(
    hidden_size,
    [intermediate_size, intermediate_size],
    bias=False,
    quant_config=quant_config,
    prefix=f"{prefix}.gate_up_proj",
)

attention_scaling = 1.0

packed_query_sequence: FloatTensor = None

past_key_values: NaiveCache | None = None

act = GELU()

fc1 = ColumnParallelLinear(
    input_dim,
    output_dim,
    bias=True,
    gather_output=False,
    quant_config=quant_config,
    prefix=f"{prefix}.fc1",
)

fc2 = RowParallelLinear(
    output_dim,
    output_dim,
    bias=True,
    input_is_parallel=True,
    quant_config=quant_config,
    prefix=f"{prefix}.fc2",
)

key_cache = {k: None for k in (range(num_layers))}

num_layers

seq_lens

value_cache = {k: None for k in (range(num_layers))}

attn_causal = Attention(
    num_heads=total_num_heads,
    head_size=head_dim,
    softmax_scale=1.0 / head_dim**0.5,
    causal=True,
    num_kv_heads=total_num_kv_heads,
)

attn_noncausal = Attention(
    num_heads=total_num_heads,
    head_size=head_dim,
    softmax_scale=1.0 / head_dim**0.5,
    causal=False,
    num_kv_heads=total_num_kv_heads,
)

head_dim = hidden_size // total_num_heads

hidden_size = hidden_size

k_norm = RMSNorm(head_dim, eps=rms_norm_eps)

k_norm_moe_gen = RMSNorm(head_dim, eps=rms_norm_eps)

kv_size = num_kv_heads * head_dim

layer_idx = layer_idx

num_heads = total_num_heads // tp_size

num_kv_heads = max(1, total_num_kv_heads // tp_size)

o_proj = RowParallelLinear(
    total_num_heads * head_dim,
    hidden_size,
    bias=False,
    quant_config=quant_config,
    prefix=f"{prefix}.o_proj",
)

o_proj_moe_gen = RowParallelLinear(
    total_num_heads * head_dim,
    hidden_size,
    bias=False,
    quant_config=quant_config,
    prefix=f"{prefix}.o_proj_moe_gen",
)

parallel_config = parallel_config

q_norm = RMSNorm(head_dim, eps=rms_norm_eps)

q_norm_moe_gen = RMSNorm(head_dim, eps=rms_norm_eps)

q_size = num_heads * head_dim

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

qkv_proj_moe_gen = QKVParallelLinear(
    hidden_size,
    head_dim,
    total_num_heads,
    total_num_kv_heads,
    bias=True,
    quant_config=quant_config,
    prefix=f"{prefix}.qkv_proj_moe_gen",
)

rotary_op = RotaryEmbedding(is_neox_style=True)

total_num_heads = num_attention_heads

total_num_kv_heads = num_key_value_heads

hidden_size = hidden_size

max_num_patch_per_side = max_num_patch_per_side

pos_embed = Parameter(
    zeros(max_num_patch_per_side**2, hidden_size),
    requires_grad=False,
)

keys_to_ignore_at_inference = ['past_key_values']

layer_module = layer_module

model_type = 'qwen2_mot'

qk_norm = qk_norm

hidden_size = hidden_size

input_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)

input_layernorm_moe_gen = RMSNorm(
    hidden_size, eps=rms_norm_eps
)

mlp = BagelMLP(
    hidden_size,
    intermediate_size,
    hidden_act,
    quant_config=quant_config,
    prefix=f"{prefix}.mlp",
)

mlp_moe_gen = BagelMLP(
    hidden_size,
    intermediate_size,
    hidden_act,
    quant_config=quant_config,
    prefix=f"{prefix}.mlp_moe_gen",
)

post_attention_layernorm = RMSNorm(
    hidden_size, eps=rms_norm_eps
)

post_attention_layernorm_moe_gen = RMSNorm(
    hidden_size, eps=rms_norm_eps
)

self_attn = attn_module(
    config,
    layer_idx,
    parallel_config=parallel_config,
    quant_config=quant_config,
    prefix=f"{prefix}.self_attn",
)

lm_head = Linear(hidden_size, vocab_size, bias=False)

model = Qwen2MoTModel(
    config,
    parallel_config=parallel_config,
    quant_config=quant_config,
    prefix=f"{prefix}.model",
)

vocab_size = vocab_size

embed_tokens = VocabParallelEmbedding(
    vocab_size, hidden_size
)

layers = ModuleList(
    [
        (
            Qwen2MoTDecoderLayer(
                config,
                layer_idx,
                attn_module=PackedAttentionMoT,
                parallel_config=parallel_config,
                quant_config=quant_config,
                prefix=f"{prefix}.layers.{layer_idx}",
            )
        )
        for layer_idx in (range(num_hidden_layers))
    ]
)

norm = RMSNorm(hidden_size, eps=rms_norm_eps)

norm_moe_gen = RMSNorm(hidden_size, eps=rms_norm_eps)

padding_idx = pad_token_id

rotary_emb = BagelRotaryEmbedding(config=config)

use_moe = 'Mo' in layer_module

vocab_size = vocab_size

frequency_embedding_size = frequency_embedding_size

mlp = Sequential(
    Linear(
        frequency_embedding_size, hidden_size, bias=True
    ),
    SiLU(),
    Linear(hidden_size, hidden_size, bias=True),
)