vllm_omni.diffusion.models.bagel.bagel_transformer ¶

base_model_prefix = 'bagel'

config = config

config_class = BagelConfig

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

get_flattened_position_ids = (
    get_flattened_position_ids_extrapolate
)

hidden_size = config.llm_config.hidden_size

language_model = language_model

latent_channel = config.vae_config.z_channels

latent_downsample = (
    config.vae_config.downsample * config.latent_patch_size
)

latent_patch_size = config.latent_patch_size

latent_pos_embed = PositionEmbedding(
    self.max_latent_size, self.hidden_size
)

llm2vae = nn.Linear(self.hidden_size, self.patch_latent_dim)

max_latent_size = config.max_latent_size

num_heads = config.llm_config.num_attention_heads

parallel_config = parallel_config

patch_latent_dim = (
    self.latent_patch_size**2 * self.latent_channel
)

time_embedder = TimestepEmbedder(self.hidden_size)

timestep_shift = config.timestep_shift

use_moe = 'Mo' in config.llm_config.layer_module

vae2llm = nn.Linear(self.patch_latent_dim, self.hidden_size)

vit_hidden_size = config.vit_config.hidden_size

vit_max_num_patch_per_side = (
    config.vit_max_num_patch_per_side
)

vit_model = vit_model

vit_patch_size = config.vit_config.patch_size

vit_pos_embed = PositionEmbedding(
    self.vit_max_num_patch_per_side, self.hidden_size
)

act_fn = SiluAndMul()

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

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

intermediate_size = intermediate_size

attention_scaling = 1.0

packed_query_sequence: FloatTensor = None

past_key_values: NaiveCache | None = None

act = nn.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))}

key_values_lens: list[int] | None = None

num_layers

seq_lens

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

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

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

head_dim = self.hidden_size // self.total_num_heads

hidden_size = config.hidden_size

k_norm = MoTRMSNorm(
    self.head_dim, head_norm=True, eps=config.rms_norm_eps
)

kv_size = self.num_kv_heads * self.head_dim

layer_idx = layer_idx

num_heads = self.total_num_heads // tp_size

num_kv_heads = max(1, self.total_num_kv_heads // tp_size)

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

parallel_config = parallel_config

q_norm = MoTRMSNorm(
    self.head_dim, head_norm=True, eps=config.rms_norm_eps
)

q_size = self.num_heads * self.head_dim

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

rotary_op = RotaryEmbedding(is_neox_style=True)

total_num_heads = config.num_attention_heads

total_num_kv_heads = config.num_key_value_heads

hidden_size = hidden_size

max_num_patch_per_side = max_num_patch_per_side

pos_embed = nn.Parameter(
    torch.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 = config.hidden_size

input_layernorm = MoTRMSNorm(
    config.hidden_size, eps=config.rms_norm_eps
)

layer_idx = layer_idx

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

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

post_attention_layernorm = MoTRMSNorm(
    config.hidden_size, eps=config.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 = nn.Linear(
    config.hidden_size, config.vocab_size, bias=False
)

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

vocab_size = config.vocab_size

embed_tokens = VocabParallelEmbedding(
    config.vocab_size, config.hidden_size
)

layers = nn.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(config.num_hidden_layers))
    ]
)

norm = MoTRMSNorm(
    config.hidden_size, eps=config.rms_norm_eps
)

padding_idx = config.pad_token_id

rotary_emb = BagelRotaryEmbedding(config=config)

use_moe = 'Mo' in config.layer_module

vocab_size = config.vocab_size

frequency_embedding_size = frequency_embedding_size

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