# coding=utf-8 # Copyright 2025 Antgroup and The HuggingFace Inc. team. All rights reserved. # # This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX # and OPT implementations in this library. It has been modified from its # original forms to accommodate minor architectural differences compared # to GPT-NeoX and OPT used by the Meta AI team that trained the model. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """PyTorch implementation of the fused LLaDA2 MoE model.""" from dataclasses import dataclass import math from typing import List, Optional, Tuple, Union import warnings import torch import torch.nn as nn import torch.nn.functional as F from transformers.activations import ACT2FN from transformers.cache_utils import Cache, DynamicCache from transformers.generation import GenerationMixin from transformers.modeling_attn_mask_utils import ( _prepare_4d_attention_mask, _prepare_4d_causal_attention_mask, _prepare_4d_causal_attention_mask_for_sdpa, ) from transformers.modeling_outputs import ModelOutput, MoeModelOutputWithPast from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update from transformers.modeling_utils import PreTrainedModel from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS from transformers.utils import logging from .fused_moe_ops import fused_moe_forward from .configuration_llada2uni_moe import LLaDA2MoeConfig logger = logging.get_logger(__name__) class LLaDA2MoeRMSNorm(nn.Module): """RMSNorm used by the LLaDA2 model.""" def __init__(self, hidden_size, eps=1e-6): super().__init__() self.weight = nn.Parameter(torch.ones(hidden_size)) self.variance_epsilon = eps def forward(self, hidden_states): input_dtype = hidden_states.dtype hidden_states = hidden_states.to(torch.float32) variance = hidden_states.pow(2).mean(-1, keepdim=True) hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) return self.weight * hidden_states.to(input_dtype) # Preserve the historical spelling used by the original implementation. LLaDA2MoERMSNorm = LLaDA2MoeRMSNorm ALL_LAYERNORM_LAYERS.append(LLaDA2MoeRMSNorm) class LLaDA2MoePreTrainedModel(PreTrainedModel): config_class = LLaDA2MoeConfig base_model_prefix = "model" supports_gradient_checkpointing = True _no_split_modules = ["LLaDA2MoeDecoderLayer"] _skip_keys_device_placement = "past_key_values" _supports_flash_attn_2 = True _supports_sdpa = True _supports_cache_class = True _supports_flash_attn = True _can_compile_fullgraph = True _supports_attention_backend = True def _init_weights(self, module): std = self.config.initializer_range if isinstance(module, nn.Linear): module.weight.data.normal_(mean=0.0, std=std) if module.bias is not None: module.bias.data.zero_() elif isinstance(module, nn.Embedding): module.weight.data.normal_(mean=0.0, std=std) if module.padding_idx is not None: module.weight.data[module.padding_idx].zero_() def rotate_half(hidden_states): first, second = hidden_states.chunk(2, dim=-1) return torch.cat((-second, first), dim=-1) def apply_rotary_pos_emb(query, key, cos, sin, position_ids=None, unsqueeze_dim=1): """Apply RoPE to the rotary part of query and key states.""" cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) rotary_dim = cos.shape[-1] query_rotary, query_pass = query[..., :rotary_dim], query[..., rotary_dim:] key_rotary, key_pass = key[..., :rotary_dim], key[..., rotary_dim:] query_rotary = query_rotary * cos + rotate_half(query_rotary) * sin key_rotary = key_rotary * cos + rotate_half(key_rotary) * sin return torch.cat((query_rotary, query_pass), dim=-1), torch.cat( (key_rotary, key_pass), dim=-1 ) class LLaDA2MoeRotaryEmbedding(nn.Module): def __init__(self, config: LLaDA2MoeConfig, device=None): super().__init__() # BC: "rope_type" was originally "type" if hasattr(config, "rope_scaling") and config.rope_scaling is not None: self.rope_type = config.rope_scaling.get( "rope_type", config.rope_scaling.get("type") ) else: self.rope_type = "default" self.max_seq_len_cached = config.max_position_embeddings self.original_max_seq_len = config.max_position_embeddings self.config = config self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device) self.register_buffer("inv_freq", inv_freq, persistent=False) self.original_inv_freq = self.inv_freq @torch.no_grad() @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope) def forward(self, x, position_ids): inv_freq_expanded = ( self.inv_freq[None, :, None] .float() .expand(position_ids.shape[0], -1, 1) .to(x.device) ) position_ids_expanded = position_ids[:, None, :].float() device_type = ( x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu" ) with torch.autocast(device_type=device_type, enabled=False): # Force float32 freqs = ( inv_freq_expanded.float() @ position_ids_expanded.float() ).transpose(1, 2) emb = torch.cat((freqs, freqs), dim=-1) cos = emb.cos() * self.attention_scaling sin = emb.sin() * self.attention_scaling return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) class LLaDA2MoeMLP(nn.Module): def __init__(self, config: LLaDA2MoeConfig, intermediate_size: int): super().__init__() self.config = config self.hidden_size = config.hidden_size self.intermediate_size = intermediate_size self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) self.act_fn = ACT2FN[config.hidden_act] def forward(self, x): return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) class LLaDA2MoeGate(nn.Module): def __init__(self, config): super().__init__() self.config = config self.top_k = config.num_experts_per_tok self.num_experts = config.num_experts self.n_group = config.n_group self.topk_group = config.topk_group self.gating_dim = config.hidden_size self.weight = nn.Parameter(torch.empty((self.num_experts, self.gating_dim))) self.routed_scaling_factor = config.routed_scaling_factor self.register_buffer("expert_bias", torch.zeros((self.num_experts))) self.reset_parameters() def reset_parameters(self) -> None: import torch.nn.init as init init.kaiming_uniform_(self.weight, a=math.sqrt(5)) def group_limited_topk( self, scores: torch.Tensor, ): num_tokens, _ = scores.size() group_scores = ( scores.view(num_tokens, self.n_group, -1).topk(2, dim=-1)[0].sum(dim=-1) ) group_idx = torch.topk(group_scores, k=self.topk_group, dim=-1, sorted=False)[1] group_mask = torch.zeros_like(group_scores) group_mask.scatter_(1, group_idx, 1) score_mask = ( group_mask.unsqueeze(-1) .expand(num_tokens, self.n_group, self.num_experts // self.n_group) .reshape(num_tokens, -1) ) masked_scores = scores.masked_fill(~score_mask.bool(), float("-inf")) probs, top_indices = torch.topk(masked_scores, k=self.top_k, dim=-1) return probs, top_indices def forward(self, hidden_states): hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) logits = F.linear( hidden_states.type(torch.float32), self.weight.type(torch.float32) ) scores = torch.sigmoid(logits.float()).type_as(logits) scores_for_routing = scores + self.expert_bias _, topk_idx = self.group_limited_topk(scores_for_routing) scores = torch.gather(scores, dim=1, index=topk_idx).type_as(logits) topk_weight = ( scores / (scores.sum(dim=-1, keepdim=True) + 1e-20) if self.top_k > 1 else scores ) topk_weight = topk_weight * self.routed_scaling_factor return topk_idx, topk_weight, logits class LLaDA2MoeExperts(nn.Module): def __init__(self, config): super().__init__() self.num_experts = config.num_experts self.hidden_dim = config.hidden_size self.intermediate_size = config.moe_intermediate_size self.gate_proj = torch.nn.Parameter( torch.empty(self.num_experts, self.intermediate_size, self.hidden_dim), requires_grad=True, ) self.up_proj = torch.nn.Parameter( torch.empty(self.num_experts, self.intermediate_size, self.hidden_dim), requires_grad=True, ) self.down_proj = torch.nn.Parameter( torch.empty(self.num_experts, self.hidden_dim, self.intermediate_size), requires_grad=True, ) def forward(self, hidden_states, routing_weights, selected_experts): return fused_moe_forward( module=self, num_experts=self.num_experts, routing_weights=routing_weights, selected_experts=selected_experts, hidden_states=hidden_states, fc1_1_weight=self.gate_proj, fc1_2_weight=self.up_proj, fc2_weight=self.down_proj, ) def reset_parameters(self): """ Initialize the parameters of all expert networks. Uses different initialization strategies for different projection layers. """ for expert_id in range(self.num_experts): nn.init.kaiming_uniform_(self.gate_proj[expert_id], a=math.sqrt(5)) nn.init.kaiming_uniform_(self.up_proj[expert_id], a=math.sqrt(5)) nn.init.xavier_uniform_(self.down_proj[expert_id]) class LLaDA2MoeSparseMoeBlock(nn.Module): """Fused routed experts plus a shared expert.""" def __init__(self, config: LLaDA2MoeConfig): super().__init__() self.config = config self.experts = LLaDA2MoeExperts(config) self.gate = LLaDA2MoeGate(config) if config.num_shared_experts is not None: self.shared_experts = LLaDA2MoeMLP( config=config, intermediate_size=config.moe_intermediate_size * config.num_shared_experts, ) def forward(self, hidden_states): identity = hidden_states bsz, seq_len, h = hidden_states.shape topk_idx, topk_weight, router_logits = self.gate(hidden_states) hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) y = self.experts( hidden_states, routing_weights=topk_weight, selected_experts=topk_idx ).reshape(bsz, seq_len, h) if self.config.num_shared_experts is not None: y = y + self.shared_experts(identity) return y, ( router_logits.view(bsz, seq_len, -1), topk_idx.view(bsz, seq_len, -1), ) def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: """ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) """ batch, num_key_value_heads, slen, head_dim = hidden_states.shape if n_rep == 1: return hidden_states hidden_states = hidden_states[:, :, None, :, :].expand( batch, num_key_value_heads, n_rep, slen, head_dim ) return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) class LLaDA2MoeAttention(nn.Module): """Multi-headed attention from 'Attention Is All You Need' paper""" def __init__(self, config: LLaDA2MoeConfig, layer_idx: Optional[int] = None): super().__init__() self.config = config self.layer_idx = layer_idx if layer_idx is None: logger.warning_once( f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will " "to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` " "when creating this class." ) self.attention_dropout = config.attention_dropout self.hidden_size = config.hidden_size self.num_heads = config.num_attention_heads self.head_dim = config.head_dim or self.hidden_size // self.num_heads partial_rotary_factor = ( config.partial_rotary_factor if hasattr(config, "partial_rotary_factor") else 1.0 ) self.rope_dim = int(self.head_dim * partial_rotary_factor) self.num_key_value_heads = config.num_key_value_heads self.num_key_value_groups = self.num_heads // self.num_key_value_heads self.max_position_embeddings = config.max_position_embeddings self.rope_theta = config.rope_theta self.is_causal = False self.query_key_value = nn.Linear( self.hidden_size, (self.num_heads + 2 * self.num_key_value_heads) * self.head_dim, bias=config.use_qkv_bias, ) self.query_layernorm = LLaDA2MoERMSNorm(self.head_dim, eps=config.rms_norm_eps) self.key_layernorm = LLaDA2MoERMSNorm(self.head_dim, eps=config.rms_norm_eps) self.dense = nn.Linear( self.num_heads * self.head_dim, self.hidden_size, bias=config.use_bias ) def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_value: Optional[Cache] = None, output_attentions: bool = False, use_cache: bool = False, position_embeddings: Optional[ Tuple[torch.Tensor, torch.Tensor] ] = None, # necessary, but kept here for BC **kwargs, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: if "padding_mask" in kwargs: warnings.warn( "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`" ) bsz, q_len, _ = hidden_states.size() qkv = self.query_key_value(hidden_states) qkv = qkv.view( bsz, q_len, self.num_heads + 2 * self.num_key_value_heads, self.head_dim ) query_states, key_states, value_states = qkv.split( [self.num_heads, self.num_key_value_heads, self.num_key_value_heads], dim=-2 ) query_states = query_states.transpose(1, 2) key_states = key_states.transpose(1, 2) value_states = value_states.transpose(1, 2) query_states = self.query_layernorm(query_states) key_states = self.key_layernorm(key_states) kv_seq_len = key_states.shape[-2] if past_key_value is not None: if self.layer_idx is None: raise ValueError( f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} " "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class " "with a layer index." ) kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx) cos, sin = position_embeddings query_states, key_states = apply_rotary_pos_emb( query_states, key_states, cos, sin, position_ids ) if past_key_value is not None: cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models key_states, value_states = past_key_value.update( key_states, value_states, self.layer_idx, cache_kwargs ) key_states = repeat_kv(key_states, self.num_key_value_groups) value_states = repeat_kv(value_states, self.num_key_value_groups) attn_weights = torch.matmul( query_states, key_states.transpose(2, 3) ) / math.sqrt(self.head_dim) if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len): raise ValueError( f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is" f" {attn_weights.size()}" ) if attention_mask is not None: if attention_mask.size() != (bsz, 1, q_len, kv_seq_len): raise ValueError( f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}" ) attn_weights = attn_weights + attention_mask attn_weights = nn.functional.softmax( attn_weights, dim=-1, dtype=torch.float32 ).to(query_states.dtype) attn_weights = nn.functional.dropout( attn_weights, p=self.attention_dropout, training=self.training ) attn_output = torch.matmul(attn_weights, value_states) if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim): raise ValueError( f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is" f" {attn_output.size()}" ) attn_output = attn_output.transpose(1, 2).contiguous() attn_output = attn_output.reshape(bsz, q_len, -1) attn_output = self.dense(attn_output) if not output_attentions: attn_weights = None return attn_output, attn_weights, past_key_value class LLaDA2MoeSdpaAttention(LLaDA2MoeAttention): """ LLaDA2Moe attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from `LLaDA2MoeAttention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to SDPA API. """ def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_value: Optional[Cache] = None, output_attentions: bool = False, use_cache: bool = False, position_embeddings: Optional[ Tuple[torch.Tensor, torch.Tensor] ] = None, # necessary, but kept here for BC **kwargs, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: if output_attentions: logger.warning_once( "LLaDA2MoeModel is using LLaDA2MoeSdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to the manual attention implementation, " 'but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.' ) return super().forward( hidden_states=hidden_states, attention_mask=attention_mask, position_ids=position_ids, past_key_value=past_key_value, output_attentions=output_attentions, use_cache=use_cache, ) bsz, q_len, _ = hidden_states.size() qkv = self.query_key_value(hidden_states) qkv = qkv.view( bsz, q_len, self.num_heads + 2 * self.num_key_value_heads, self.head_dim ) query_states, key_states, value_states = qkv.split( [self.num_heads, self.num_key_value_heads, self.num_key_value_heads], dim=-2 ) query_states = query_states.transpose(1, 2) key_states = key_states.transpose(1, 2) value_states = value_states.transpose(1, 2) query_states = self.query_layernorm(query_states) key_states = self.key_layernorm(key_states) kv_seq_len = key_states.shape[-2] if past_key_value is not None: kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx) cos, sin = position_embeddings query_states, key_states = apply_rotary_pos_emb( query_states, key_states, cos, sin, position_ids ) if past_key_value is not None: cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models key_states, value_states = past_key_value.update( key_states, value_states, self.layer_idx, cache_kwargs ) key_states = repeat_kv(key_states, self.num_key_value_groups) value_states = repeat_kv(value_states, self.num_key_value_groups) if attention_mask is not None: if attention_mask.size() != (bsz, 1, q_len, kv_seq_len): raise ValueError( f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}" ) # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask, # Reference: https://github.com/pytorch/pytorch/issues/112577. if query_states.device.type == "cuda" and attention_mask is not None: query_states = query_states.contiguous() key_states = key_states.contiguous() value_states = value_states.contiguous() attn_output = torch.nn.functional.scaled_dot_product_attention( query_states, key_states, value_states, attn_mask=attention_mask, dropout_p=self.attention_dropout if self.training else 0.0, # The q_len > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not create a causal mask in case q_len == 1. is_causal=self.is_causal and attention_mask is None and q_len > 1, ) attn_output = attn_output.transpose(1, 2).contiguous() attn_output = attn_output.reshape(bsz, q_len, -1) attn_output = self.dense(attn_output) return attn_output, None, past_key_value ATTENTION_CLASSES = { "eager": LLaDA2MoeSdpaAttention, "flash_attention_2": LLaDA2MoeSdpaAttention, "sdpa": LLaDA2MoeSdpaAttention, } class LLaDA2MoeDecoderLayer(nn.Module): def __init__(self, config: LLaDA2MoeConfig, layer_idx: int): super().__init__() self.hidden_size = config.hidden_size self.attention = ATTENTION_CLASSES[config._attn_implementation]( config=config, layer_idx=layer_idx ) self.mlp = ( LLaDA2MoeSparseMoeBlock(config) if ( config.num_experts is not None and layer_idx >= config.first_k_dense_replace ) else LLaDA2MoeMLP(config=config, intermediate_size=config.intermediate_size) ) self.input_layernorm = LLaDA2MoERMSNorm( config.hidden_size, eps=config.rms_norm_eps ) self.post_attention_layernorm = LLaDA2MoERMSNorm( config.hidden_size, eps=config.rms_norm_eps ) def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_value: Optional[Tuple[torch.Tensor]] = None, output_attentions: Optional[bool] = False, output_router_logits: Optional[bool] = False, use_cache: Optional[bool] = False, position_embeddings: Optional[ Tuple[torch.Tensor, torch.Tensor] ] = None, # necessary, but kept here for BC **kwargs, ) -> Tuple[ torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]] ]: """ Args: hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` attention_mask (`torch.FloatTensor`, *optional*): attention mask of size `(batch_size, sequence_length)` if flash attention is used or `(batch_size, 1, query_sequence_length, key_sequence_length)` if default attention is used. position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0, config.n_positions - 1]`. past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states output_attentions (`bool`, *optional*): Whether to return the attentions tensors of all attention layers. See `attentions` under returned tensors for more detail. output_router_logits (`bool`, *optional*): Whether or not to return the logits of all the routers. They are useful for computing the router loss, and should not be returned during inference. use_cache (`bool`, *optional*): If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see `past_key_values`). """ if "padding_mask" in kwargs: warnings.warn( "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`" ) residual = hidden_states hidden_states = self.input_layernorm(hidden_states) hidden_states, self_attn_weights, present_key_value = self.attention( hidden_states=hidden_states, attention_mask=attention_mask, position_ids=position_ids, past_key_value=past_key_value, output_attentions=output_attentions, position_embeddings=position_embeddings, use_cache=use_cache, ) hidden_states = residual + hidden_states residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) hidden_states = self.mlp(hidden_states) if isinstance(hidden_states, tuple): hidden_states, router_logits = hidden_states else: router_logits = None hidden_states = residual + hidden_states.to(residual.device) outputs = (hidden_states,) if output_attentions: outputs += (self_attn_weights,) if use_cache: outputs += (present_key_value,) if output_router_logits: outputs += (router_logits,) return outputs def calculate_pack_position_ids( input_ids: Optional[torch.Tensor] = None, inputs_embeds: Optional[torch.Tensor] = None, position_ids: Optional[torch.Tensor] = None, past_key_values_length: int = 0, cu_lengths_list: Optional[List[torch.Tensor]] = None, ): """Build continuous or per-sequence packed position IDs.""" if position_ids is not None: return position_ids if input_ids is not None: device = input_ids.device batch_size, seq_length = input_ids.shape elif inputs_embeds is not None: device = inputs_embeds.device batch_size, seq_length, _ = inputs_embeds.shape else: raise ValueError("You have to specify either input_ids or inputs_embeds") if cu_lengths_list is not None: all_position_ids = [] for i in range(batch_size): cu_seqlens = cu_lengths_list[i].to(device) starts = cu_seqlens[:-1] lengths = cu_seqlens[1:] - cu_seqlens[:-1] total_len = cu_seqlens[-1].item() global_positions = torch.arange(total_len, device=device, dtype=torch.long) subtraction_mask = torch.repeat_interleave(starts, lengths) current_pos_ids = global_positions - subtraction_mask all_position_ids.append(current_pos_ids) position_ids = torch.nn.utils.rnn.pad_sequence( all_position_ids, batch_first=True, padding_value=0, ) if position_ids.shape[1] < seq_length: pad_right = seq_length - position_ids.shape[1] position_ids = F.pad(position_ids, (0, pad_right), "constant", 0) else: position_ids = torch.arange( past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device, ) position_ids = position_ids.unsqueeze(0).expand(batch_size, -1) return position_ids class LLaDA2MoeModel(LLaDA2MoePreTrainedModel): """ Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`LLaDA2MoeDecoderLayer`] Args: config: LLaDA2MoeConfig """ def __init__(self, config: LLaDA2MoeConfig): super().__init__(config) self.padding_idx = config.pad_token_id self.vocab_size = config.vocab_size self.word_embeddings = nn.Embedding( config.vocab_size, config.hidden_size, self.padding_idx ) self.layers = nn.ModuleList( [ LLaDA2MoeDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers) ] ) self._use_sdpa = config._attn_implementation == "sdpa" self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2" self.norm = LLaDA2MoERMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.rotary_emb = LLaDA2MoeRotaryEmbedding(config=config) self.gradient_checkpointing = False self.post_init() def get_input_embeddings(self): return self.word_embeddings def set_input_embeddings(self, value): self.word_embeddings = value def forward( self, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[List[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, output_router_logits: Optional[bool] = None, cu_lengths_list: Optional[List] = None, return_dict: Optional[bool] = None, **kwargs, ) -> Union[Tuple, MoeModelOutputWithPast]: output_attentions = ( output_attentions if output_attentions is not None else self.config.output_attentions ) output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) output_router_logits = ( output_router_logits if output_router_logits is not None else self.config.output_router_logits ) use_cache = use_cache if use_cache is not None else self.config.use_cache return_dict = ( return_dict if return_dict is not None else self.config.use_return_dict ) if input_ids is not None and inputs_embeds is not None: raise ValueError( "You cannot specify both input_ids and inputs_embeds at the same time" ) elif input_ids is not None: batch_size, seq_length = input_ids.shape[:2] elif inputs_embeds is not None: batch_size, seq_length = inputs_embeds.shape[:2] else: raise ValueError("You have to specify either input_ids or inputs_embeds") if self.gradient_checkpointing and self.training: if use_cache: logger.warning_once( "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`transformers." ) use_cache = False past_key_values_length = 0 if use_cache: use_legacy_cache = not isinstance(past_key_values, Cache) if use_legacy_cache: past_key_values = DynamicCache.from_legacy_cache(past_key_values) past_key_values_length = past_key_values.get_usable_length(seq_length) if position_ids is None: position_ids = calculate_pack_position_ids( input_ids, inputs_embeds, position_ids, past_key_values_length, cu_lengths_list, ) if inputs_embeds is None: inputs_embeds = self.word_embeddings(input_ids) if hasattr(attention_mask, "dim") and attention_mask.dim() == 2: if self._use_sdpa and not output_attentions: # output_attentions=True can not be supported when using SDPA, and we fall back on # the manual implementation that requires a 4D causal mask in all cases. attention_mask = _prepare_4d_causal_attention_mask_for_sdpa( attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length, ) else: if attention_mask is not None: attention_mask = _prepare_4d_attention_mask( attention_mask, inputs_embeds.dtype ) else: attention_mask = _prepare_4d_causal_attention_mask( attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length, ) hidden_states = inputs_embeds position_embeddings = self.rotary_emb(hidden_states, position_ids) all_hidden_states = () if output_hidden_states else None all_self_attns = () if output_attentions else None all_router_logits = () if output_router_logits else None next_decoder_cache = None for decoder_layer in self.layers: if output_hidden_states: all_hidden_states += (hidden_states,) if self.gradient_checkpointing and self.training: layer_outputs = self._gradient_checkpointing_func( decoder_layer.__call__, hidden_states, attention_mask, position_ids, past_key_values, output_attentions, output_router_logits, use_cache, position_embeddings, ) else: layer_outputs = decoder_layer( hidden_states, attention_mask=attention_mask, position_ids=position_ids, past_key_value=past_key_values, output_attentions=output_attentions, output_router_logits=output_router_logits, use_cache=use_cache, position_embeddings=position_embeddings, ) hidden_states = layer_outputs[0] if use_cache: next_decoder_cache = layer_outputs[2 if output_attentions else 1] if output_attentions: all_self_attns += (layer_outputs[1],) if output_router_logits and layer_outputs[-1] is not None: all_router_logits += (layer_outputs[-1],) hidden_states = self.norm(hidden_states) if output_hidden_states: all_hidden_states += (hidden_states,) next_cache = None if use_cache: next_cache = ( next_decoder_cache.to_legacy_cache() if use_legacy_cache else next_decoder_cache ) if not return_dict: return tuple( v for v in [ hidden_states, next_cache, all_hidden_states, all_self_attns, all_router_logits, ] if v is not None ) return MoeModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=next_cache, hidden_states=all_hidden_states, attentions=all_self_attns, router_logits=all_router_logits, ) @dataclass class LLaDA2MoeCausalLMOutputWithPast(ModelOutput): r""" loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): Training loss. logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`): Prediction scores for each vocabulary token before SoftMax. past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache). Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see `past_key_values` input) to speed up sequential decoding. rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*): The offset between the sequence length and rotary position indices. """ loss: Optional[torch.FloatTensor] = None z_loss: Optional[torch.FloatTensor] = None logits: Optional[torch.FloatTensor] = None past_key_values: Optional[Cache] = None hidden_states: Optional[tuple[torch.FloatTensor]] = None attentions: Optional[tuple[torch.FloatTensor]] = None rope_deltas: Optional[torch.LongTensor] = None class LLaDA2MoeBackbone(nn.Module): """Container for the model backbone and output head.""" def __init__(self, config: LLaDA2MoeConfig): super().__init__() self.language_model = LLaDA2MoeModel(config) self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) def get_input_embeddings(self): return self.language_model.get_input_embeddings() def set_input_embeddings(self, value): self.language_model.set_input_embeddings(value) def forward(self, *args, **kwargs): return self.language_model(*args, **kwargs) class LLaDA2MoeModelLM(LLaDA2MoePreTrainedModel, GenerationMixin): """Fused LLaDA2 MoE model.""" accepts_loss_kwargs = False def __init__(self, config: LLaDA2MoeConfig): super().__init__(config) self.model = LLaDA2MoeBackbone(config) self.img_token_id = 157184 self.img_start_id = 157185 self.img_end_id = 157186 self.img_pad_id = 157187 self.post_init() @property def language_model(self): return self.model.language_model def get_input_embeddings(self): return self.model.get_input_embeddings() def set_input_embeddings(self, value): self.model.set_input_embeddings(value) def get_output_embeddings(self): return self.model.lm_head def set_output_embeddings(self, value): self.model.lm_head = value def get_decoder(self): return self.model.language_model def set_decoder(self, decoder): self.model.language_model = decoder def forward( self, input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[Cache] = None, inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_router_logits: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, logits_to_keep: Union[int, torch.Tensor] = 0, cu_lengths_list: Optional[List] = None, **kwargs, ) -> Union[tuple, LLaDA2MoeCausalLMOutputWithPast]: return_dict = ( return_dict if return_dict is not None else self.config.use_return_dict ) if inputs_embeds is None: if input_ids is None: raise ValueError("Provide either input_ids or inputs_embeds") inputs_embeds = self.get_input_embeddings()(input_ids) outputs = self.model( input_ids=None, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, output_attentions=output_attentions, output_router_logits=output_router_logits, output_hidden_states=output_hidden_states, return_dict=True, cu_lengths_list=cu_lengths_list, **kwargs, ) hidden_states = outputs.last_hidden_state indices = ( slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep ) logits = self.model.lm_head(hidden_states[:, indices, :]) loss = None if labels is not None: loss = F.cross_entropy( logits.reshape(-1, logits.shape[-1]), labels.reshape(-1) ) if not return_dict: result = ( logits, outputs.past_key_values, outputs.hidden_states, outputs.attentions, ) return ((loss,) + result) if loss is not None else result return LLaDA2MoeCausalLMOutputWithPast( loss=loss, z_loss=None, logits=logits, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions, rope_deltas=None, ) @staticmethod def _top_k_logits(logits, k): if k is None or k <= 0: return logits values, _ = torch.topk(logits, min(k, logits.shape[-1])) return torch.where(logits < values[..., -1, None], -torch.inf, logits) @staticmethod def _top_p_logits(logits, p): if p is None or p >= 1.0: return logits sorted_logits, sorted_indices = torch.sort(logits, descending=True) cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1) sorted_mask = cumulative_probs > p sorted_mask[..., 1:] = sorted_mask[..., :-1].clone() sorted_mask[..., 0] = False mask = torch.zeros_like(sorted_mask).scatter(-1, sorted_indices, sorted_mask) return logits.masked_fill(mask, -torch.inf) def _sample_with_temperature_topk_topp( self, logits, temperature=1.0, top_k=0, top_p=1.0 ): original_shape = logits.shape[:-1] logits = logits.reshape(-1, logits.shape[-1]) if ( temperature == 0.0 and (top_k in (None, 0)) and (top_p is None or top_p >= 1.0) ): probs = F.softmax(logits, dim=-1) token = logits.argmax(dim=-1, keepdim=True) token_prob = probs.gather(-1, token) return token.view(*original_shape), token_prob.view(*original_shape) if temperature > 0 and temperature != 1.0: logits = logits / temperature logits = self._top_k_logits(logits, top_k) logits = self._top_p_logits(logits, top_p) probs = F.softmax(logits, dim=-1) token = torch.multinomial(probs, num_samples=1) token_prob = probs.gather(-1, token) return token.view(*original_shape), token_prob.view(*original_shape) @staticmethod def _get_num_transfer_tokens(block_length, steps): if steps == 0: return torch.empty(0, dtype=torch.int64) schedule = torch.full((steps,), block_length // steps, dtype=torch.int64) schedule[: block_length % steps] += 1 return schedule @torch.no_grad() def generate_bd_image_logic( self, data: Optional[dict] = None, temperature: float = 0.0, block_length: int = 32, steps: int = 32, gen_length: int = 2048, top_p: Optional[float] = None, top_k: Optional[int] = None, eos_early_stop: bool = True, minimal_topk: int = 1, threshold: float = 0.95, eos_id: int = 156892, mask_id: int = 156895, cfg_scale: float = 1.0, mode: str = "eoi", ): """Generate discrete image tokens with the original block-diffusion logic.""" if data is None or "input_ids" not in data: raise ValueError("data must contain input_ids") steps = min(steps, gen_length // minimal_topk) input_ids = data["input_ids"] eoi_id = 156902 prompt_length = input_ids.shape[1] num_blocks = (prompt_length + gen_length + block_length - 1) // block_length total_length = num_blocks * block_length block_mask = torch.tril(torch.ones(num_blocks, num_blocks, device=self.device)) full_attention_mask = ( block_mask.repeat_interleave(block_length, 0) .repeat_interleave(block_length, 1)[None, None] .bool() ) position_ids = torch.arange(total_length, device=self.device).unsqueeze(0) x = torch.full((1, total_length), mask_id, dtype=torch.long, device=self.device) x[:, :prompt_length] = input_ids prefill_blocks = prompt_length // block_length schedule = self._get_num_transfer_tokens(block_length, steps) use_cfg = cfg_scale != 1.0 if use_cfg: uncond_ids = data.get("uncond_ids", [27, 411, 19483, 29]) if torch.is_tensor(uncond_ids): uncond_ids = uncond_ids.flatten().tolist() pad_len = prompt_length - len(uncond_ids) if pad_len < 0: raise ValueError( "The unconditional prompt is longer than the conditional prompt" ) uncond_input = torch.full( (1, prompt_length), mask_id, dtype=torch.long, device=self.device ) uncond_input[0, -len(uncond_ids) :] = torch.tensor( uncond_ids, device=self.device ) uncond_attention_mask = full_attention_mask.clone() uncond_attention_mask[:, :, :, :pad_len] = False uncond_position_ids = torch.cat( [ torch.zeros(pad_len, device=self.device, dtype=torch.long), torch.arange(total_length - pad_len, device=self.device), ] ).unsqueeze(0) for block_index in range(prefill_blocks, num_blocks): window_end = (block_index + 1) * block_length current = x[:, :window_end] current_mask = full_attention_mask[:, :, :window_end, :window_end] current_positions = position_ids[:, :window_end] for step_index in range(steps): active = current[:, -block_length:] == mask_id if not active.any(): break if use_cfg: unconditional = current.clone() unconditional[:, :prompt_length] = uncond_input combined_ids = torch.cat([current, unconditional], dim=0) combined_positions = torch.cat( [current_positions, uncond_position_ids[:, :window_end]], dim=0 ) combined_mask = torch.cat( [ current_mask, uncond_attention_mask[:, :, :window_end, :window_end], ], dim=0, ) logits = self( input_ids=combined_ids, attention_mask=combined_mask, position_ids=combined_positions, ).logits conditional_logits, unconditional_logits = logits.chunk(2, dim=0) active_logits = unconditional_logits[ :, -block_length: ] + cfg_scale * ( conditional_logits[:, -block_length:] - unconditional_logits[:, -block_length:] ) else: active_logits = self( input_ids=current, attention_mask=current_mask, position_ids=current_positions, ).logits[:, -block_length:] tokens, confidence = self._sample_with_temperature_topk_topp( active_logits, temperature=temperature, top_k=top_k, top_p=top_p ) count = schedule[step_index].item() scores = torch.where(active, confidence, -torch.inf) selected = torch.zeros_like(tokens, dtype=torch.bool) high_confidence = scores[0] > threshold if high_confidence.sum().item() >= count: selected[0] = high_confidence else: _, indices = torch.topk( scores[0], k=min(count, active.sum().item()) ) selected[0, indices] = True current[:, -block_length:][selected] = tokens[selected] stop_token = eoi_id if mode == "eoi" else eos_id positions = (current[0, prompt_length:] == stop_token).nonzero( as_tuple=True )[0] if eos_early_stop and len(positions) > 0: stop_position = positions[0].item() + prompt_length if (current[0, prompt_length:stop_position] != mask_id).all(): x[:, :window_end] = current return x[:, : stop_position + 1] x[:, :window_end] = current return x[:, : prompt_length + gen_length] __all__ = ["LLaDA2MoeModelLM", "LLaDA2MoeModel", "LLaDA2MoePreTrainedModel"]