# 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 LLaDA2MoE model.""" import math from typing import List, Callable, Optional, Tuple, Union import torch import torch.nn.functional as F from torch import nn from torch.nn import CrossEntropyLoss from transformers.activations import ACT2FN from transformers.cache_utils import Cache, DynamicCache from transformers.masking_utils import create_bidirectional_mask from transformers.modeling_outputs import ( MoeModelOutputWithPast, MoeCausalLMOutputWithPast, ) from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from transformers.processing_utils import Unpack from transformers.pytorch_utils import ( ALL_LAYERNORM_LAYERS, ) from transformers.utils import ( TransformersKwargs, add_start_docstrings, add_start_docstrings_to_model_forward, logging, replace_return_docstrings, ) from .configuration_llada2_moe import LLaDA2MoeConfig from transformers.generation.utils import GenerationMixin logger = logging.get_logger(__name__) _CONFIG_FOR_DOC = "LLaDA2MoeConfig" def _get_unpad_data(attention_mask): seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten() max_seqlen_in_batch = seqlens_in_batch.max().item() cu_seqlens = F.pad( torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0) ) return ( indices, cu_seqlens, max_seqlen_in_batch, ) class LLaDA2MoeRMSNorm(nn.Module): def __init__(self, hidden_size, eps=1e-6): """ LLaDA2MoeRMSNorm is equivalent to T5LayerNorm """ 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) ALL_LAYERNORM_LAYERS.append(LLaDA2MoeRMSNorm) class LLaDA2MoeRotaryEmbedding(nn.Module): inv_freq: torch.Tensor # fix linting for register_buffer def __init__(self, config: LLaDA2MoeConfig, device=None): super().__init__() self.max_seq_len_cached = config.max_position_embeddings self.original_max_seq_len = config.max_position_embeddings self.config = config self.rope_type = self.config.rope_parameters["rope_type"] rope_init_fn: Callable = self.compute_default_rope_parameters if self.rope_type != "default": rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] inv_freq, self.attention_scaling = rope_init_fn(self.config, device) self.register_buffer("inv_freq", inv_freq, persistent=False) self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False) @staticmethod def compute_default_rope_parameters( config: LLaDA2MoeConfig = None, device=None, seq_len: int = None, ): base = config.rope_parameters["rope_theta"] partial_rotary_factor = config.rope_parameters.get("partial_rotary_factor", 1.0) head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads dim = int(head_dim * partial_rotary_factor) attention_factor = 1.0 # Unused in this type of RoPE inv_freq = 1.0 / ( base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim) ) return inv_freq, attention_factor @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) # Copied from transformers.models.llama.modeling_llama.rotate_half def rotate_half(x): """Rotates half the hidden dims of the input.""" x1 = x[..., : x.shape[-1] // 2] x2 = x[..., x.shape[-1] // 2 :] return torch.cat((-x2, x1), dim=-1) # Copied from transformers.models.llama.modeling_llama.apply_rotary_pos_emb def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): """Applies Rotary Position Embedding to the query and key tensors. Args: q (`torch.Tensor`): The query tensor. k (`torch.Tensor`): The key tensor. cos (`torch.Tensor`): The cosine part of the rotary embedding. sin (`torch.Tensor`): The sine part of the rotary embedding. position_ids (`torch.Tensor`): The position indices of the tokens corresponding to the query and key tensors. For example, this can be used to pass offsetted position ids when working with a KV-cache. unsqueeze_dim (`int`, *optional*, defaults to 1): The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. Returns: `tuple(torch.Tensor)` comprising the query and key tensors rotated using the Rotary Position Embedding. """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) # Keep half or full tensor for later concatenation rotary_dim = cos.shape[-1] q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:] k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:] # Apply rotary embeddings on the first half or full tensor q_embed = (q_rot * cos) + (rotate_half(q_rot) * sin) k_embed = (k_rot * cos) + (rotate_half(k_rot) * sin) # Concatenate back to full shape q_embed = torch.cat([q_embed, q_pass], dim=-1) k_embed = torch.cat([k_embed, k_pass], dim=-1) return q_embed, k_embed 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 # Block routing self.block_size = config.block_size self.expert_capacity = config.expert_capacity # topk selection algorithm 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 block_routing(self, scores_for_routing: torch.Tensor): """Block-level top-k routing: first select expert_capacity experts per block, then do per-token top-k within the allowed set.""" num_tokens = scores_for_routing.shape[0] assert num_tokens % self.block_size == 0 num_blocks = num_tokens // self.block_size # Reshape to (num_blocks, block_size, num_experts) block_routing_scores = scores_for_routing.view(num_blocks, self.block_size, self.num_experts) # Phase 1: compute block-level expert scores (max over tokens in each block) block_expert_scores = block_routing_scores.max(dim=1).values # (num_blocks, num_experts) # Select top expert_capacity experts per block _, block_routing_indices = torch.topk( block_expert_scores, k=self.expert_capacity, dim=-1 ) # (num_blocks, expert_capacity) # Build allowed mask: (num_blocks, num_experts) allowed_mask = torch.zeros( num_blocks, self.num_experts, dtype=torch.bool, device=scores_for_routing.device ) allowed_mask.scatter_(1, block_routing_indices, True) # Expand mask to per-token level: (num_tokens, num_experts) allowed_mask = allowed_mask.unsqueeze(1).expand(-1, self.block_size, -1).reshape(num_tokens, self.num_experts) # Phase 2: per-token top-k within allowed experts masked_scores = scores_for_routing.masked_fill(~allowed_mask, -torch.inf) _, topk_idx = torch.topk(masked_scores, k=self.top_k, dim=-1) return topk_idx def forward(self, hidden_states): # compute gating score 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.block_routing(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 LLaDA2MoeSparseMoeBlock(nn.Module): """ A mixed expert module containing shared experts. """ def __init__(self, config: LLaDA2MoeConfig): super().__init__() self.config = config self.num_experts_per_tok = config.num_experts_per_tok self._setup_experts() 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 _setup_experts(self): self.experts = nn.ModuleList( [ LLaDA2MoeMLP( config=self.config, intermediate_size=self.config.moe_intermediate_size, ) for _ in range(self.config.num_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]) flat_topk_idx = topk_idx.view(-1) if self.training: hidden_states = hidden_states.repeat_interleave( self.num_experts_per_tok, dim=0 ) y = torch.empty_like(hidden_states) for i, expert in enumerate(self.experts): y[flat_topk_idx == i] = expert(hidden_states[flat_topk_idx == i]) y = (y.view(*topk_weight.shape, -1) * topk_weight.unsqueeze(-1)).sum(dim=1) y = y.to(hidden_states.dtype).view(bsz, seq_len, h) else: y = self.moe_infer(hidden_states, topk_idx, topk_weight).view( 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), ) @torch.no_grad() def moe_infer(self, x, topk_ids, topk_weight): cnts = topk_ids.new_zeros((topk_ids.shape[0], len(self.experts))) cnts.scatter_(1, topk_ids, 1) tokens_per_expert = cnts.sum(dim=0) idxs = topk_ids.view(-1).argsort() sorted_tokens = x[idxs // topk_ids.shape[1]] tokens_per_expert = tokens_per_expert.cpu().numpy() outputs = [] start_idx = 0 for i, num_tokens_tensor in enumerate(tokens_per_expert): num_tokens = num_tokens_tensor.item() if num_tokens == 0: continue end_idx = start_idx + num_tokens expert = self.experts[i] tokens_for_this_expert = sorted_tokens[start_idx:end_idx] expert_out = expert(tokens_for_this_expert) outputs.append(expert_out.to(x.device)) start_idx = end_idx outs = torch.cat(outputs, dim=0) if len(outputs) else sorted_tokens.new_empty(0) new_x = torch.empty_like(outs) new_x[idxs] = outs final_out = ( new_x.view(*topk_ids.shape, -1) .type(topk_weight.dtype) .mul_(topk_weight.unsqueeze(dim=-1)) .sum(dim=1) .type(new_x.dtype) ) return final_out # Copied from transformers.models.llama.modeling_llama.repeat_kv 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) def eager_attention_forward( module: nn.Module, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attention_mask: Optional[torch.Tensor], scaling: float, dropout: float = 0.0, **kwargs: Unpack[TransformersKwargs], ): key_states = repeat_kv(key, module.num_key_value_groups) value_states = repeat_kv(value, module.num_key_value_groups) attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling if attention_mask is not None: attn_weights = attn_weights + attention_mask[:, :, :, : key_states.shape[-2]] # upcast attention to fp32 attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to( query.dtype ) attn_weights = nn.functional.dropout( attn_weights, p=dropout, training=module.training ) attn_output = torch.matmul(attn_weights, value_states) attn_output = attn_output.transpose(1, 2).contiguous() return attn_output, attn_weights # Copied from transformers.models.llama.modeling_llama.LlamaAttention with Llama->LLaDA2Moe 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.scaling = self.head_dim**-0.5 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, ) if self.config.use_qk_norm: 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 ) self.sliding_window = getattr(config, "sliding_window", None) def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int): return ( tensor.view(bsz, seq_len, self.num_heads, self.head_dim) .transpose(1, 2) .contiguous() ) 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]]]: input_shape = hidden_states.shape[:-1] 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) if self.config.use_qk_norm: query_states = self.query_layernorm(query_states) key_states = self.key_layernorm(key_states) cos, sin = position_embeddings query_states, key_states = apply_rotary_pos_emb( query_states, key_states, cos, sin ) 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." ) cache_kwargs = {"sin": sin, "cos": cos} key_states, value_states = past_key_value.update( key_states, value_states, self.layer_idx, cache_kwargs ) attention_interface: Callable = eager_attention_forward if self.config._attn_implementation != "eager": attention_interface = ALL_ATTENTION_FUNCTIONS[ self.config._attn_implementation ] attn_output, attn_weights = attention_interface( self, query_states, key_states, value_states, attention_mask, dropout=0.0 if not self.training else self.attention_dropout, scaling=self.scaling, sliding_window=self.sliding_window, # diff with Llama **kwargs, ) attn_output = attn_output.reshape(*input_shape, -1).contiguous() attn_output = self.dense(attn_output) return attn_output, attn_weights, past_key_value class LLaDA2MoeDecoderLayer(nn.Module): def __init__(self, config: LLaDA2MoeConfig, layer_idx: int): super().__init__() self.hidden_size = config.hidden_size self.attention = LLaDA2MoeAttention(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`). """ residual = hidden_states hidden_states = self.input_layernorm(hidden_states) # Self Attention 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 # Fully Connected 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 LLADA2MOE_START_DOCSTRING = r""" This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads etc.) This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and behavior. Parameters: config ([`LLaDA2MoeConfig`]): Model configuration class with all the parameters of the model. Initializing with a config file does not load the weights associated with the model, only the configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights. """ @add_start_docstrings( "The bare LLaDA2Moe Model outputting raw hidden-states without any specific head on top.", LLADA2MOE_START_DOCSTRING, ) 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 = False _supports_sdpa = True _supports_flex_attn = True _supports_cache_class = True @torch.no_grad() def _init_weights(self, module): super()._init_weights(module) std = self.config.initializer_range if isinstance(module, LLaDA2MoeGate): nn.init.normal_(module.weight, mean=0.0, std=std) LLADA2MOE_INPUTS_DOCSTRING = r""" Args: input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide it. Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. [What are input IDs?](../glossary#input-ids) attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: - 1 for tokens that are **not masked**, - 0 for tokens that are **masked**. [What are attention masks?](../glossary#attention-mask) Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. If `past_key_values` is used, optionally only the last `input_ids` have to be input (see `past_key_values`). If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`] and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more information on the default strategy. - 1 indicates the head is **not masked**, - 0 indicates the head is **masked**. 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]`. [What are position IDs?](../glossary#position-ids) past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*): Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values` returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`. Two formats are allowed: - a [`~cache_utils.Cache`] instance; - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy cache format. The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the legacy cache format will be returned. If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids` of shape `(batch_size, sequence_length)`. inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This is useful if you want more control over how to convert `input_ids` indices into associated vectors than the model's internal embedding lookup matrix. 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`). output_attentions (`bool`, *optional*): Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned tensors for more detail. output_hidden_states (`bool`, *optional*): Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for more detail. return_dict (`bool`, *optional*): Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. """ @add_start_docstrings( "The bare LLaDA2Moe Model outputting raw hidden-states without any specific head on top.", LLADA2MOE_START_DOCSTRING, ) 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_flex_attention = config._attn_implementation == "flex_attention" self.norm = LLaDA2MoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.rotary_emb = LLaDA2MoeRotaryEmbedding(config=config) self.gradient_checkpointing = False # Initialize weights and apply final processing self.post_init() def get_input_embeddings(self): return self.word_embeddings def set_input_embeddings(self, value): self.word_embeddings = value @add_start_docstrings_to_model_forward(LLADA2MOE_INPUTS_DOCSTRING) 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, 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 ) # retrieve input_ids and inputs_embeds 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 if use_cache and past_key_values is None: past_key_values = DynamicCache() if inputs_embeds is None: inputs_embeds = self.word_embeddings(input_ids) past_seen_tokens = ( past_key_values.get_seq_length() if past_key_values is not None else 0 ) if position_ids is None: position_ids = torch.arange( past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device, ) position_ids = position_ids.unsqueeze(0) attention_mask = create_bidirectional_mask( config=self.config, inputs_embeds=inputs_embeds, attention_mask=attention_mask, ) # embed positions hidden_states = inputs_embeds # create position embeddings to be shared across the decoder layers position_embeddings = self.rotary_emb(hidden_states, position_ids) # decoder layers 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) # add hidden states from the last decoder layer if output_hidden_states: all_hidden_states += (hidden_states,) next_cache = None if use_cache: next_cache = 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, ) class LLaDA2MoeModelLM(LLaDA2MoePreTrainedModel, GenerationMixin): _tied_weights_keys = ["lm_head.weight"] def __init__(self, config: LLaDA2MoeConfig): super().__init__(config) self.model = LLaDA2MoeModel(config) self.vocab_size = config.vocab_size self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) # Initialize weights and apply final processing self.post_init() def get_input_embeddings(self): return self.model.word_embeddings def set_input_embeddings(self, value): self.model.word_embeddings = value def get_output_embeddings(self): return self.lm_head def set_output_embeddings(self, new_embeddings): self.lm_head = new_embeddings def set_decoder(self, decoder): self.model = decoder def get_decoder(self): return self.model @add_start_docstrings_to_model_forward(LLADA2MOE_INPUTS_DOCSTRING) @replace_return_docstrings( output_type=MoeCausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC ) 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, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, output_router_logits: Optional[bool] = None, return_dict: Optional[bool] = None, **kwargs, ) -> Union[Tuple, MoeCausalLMOutputWithPast]: r""" Args: labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. Returns: Example: ```python >>> from transformers import AutoTokenizer >>> model = LLaDA2MoeForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS) >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER) >>> prompt = "Hey, are you conscious? Can you talk to me?" >>> inputs = tokenizer(prompt, return_tensors="pt") >>> # Generate >>> generate_ids = model.generate(inputs.input_ids, max_length=30) >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." ```""" 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 ) return_dict = ( return_dict if return_dict is not None else self.config.use_return_dict ) # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) outputs = self.model( input_ids=input_ids, 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_hidden_states=output_hidden_states, output_router_logits=output_router_logits, return_dict=return_dict, **kwargs, ) loss = None aux_loss = None hidden_states = outputs[0] logits = self.lm_head(hidden_states) logits = logits.float() if labels is not None: # LLaDA2.0 will use same label position logits shift_logits = logits shift_labels = labels # Flatten the tokens loss_fct = CrossEntropyLoss() shift_logits = shift_logits.view(-1, self.config.vocab_size) shift_labels = shift_labels.view(-1) # Enable model parallelism shift_labels = shift_labels.to(shift_logits.device) loss = loss_fct(shift_logits, shift_labels) if not return_dict: output = (logits,) + outputs[1:] if output_router_logits: output = (aux_loss,) + output return (loss,) + output if loss is not None else output return MoeCausalLMOutputWithPast( loss=loss, aux_loss=aux_loss, logits=logits, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions, router_logits=outputs.router_logits, ) def prepare_inputs_for_generation( self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, token_type_ids=None, **kwargs, ): if past_key_values is not None: if isinstance(past_key_values, Cache): cache_length = past_key_values.get_seq_length() past_length = past_key_values.seen_tokens max_cache_length = ( past_key_values.get_max_length() if hasattr(past_key_values, "get_max_length") else past_key_values.get_max_cache_shape() ) else: cache_length = past_length = past_key_values[0][0].shape[2] max_cache_length = None # Keep only the unprocessed tokens: # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where # some of the inputs are exclusivelly passed as part of the cache (e.g. when passing input_embeds as input) if ( attention_mask is not None and attention_mask.shape[1] > input_ids.shape[1] ): input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :] # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard # input_ids based on the past_length. elif past_length < input_ids.shape[1]: input_ids = input_ids[:, past_length:] # 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens. # If we are about to go beyond the maximum cache length, we need to crop the input attention mask. if ( max_cache_length is not None and attention_mask is not None and cache_length + input_ids.shape[1] > max_cache_length ): attention_mask = attention_mask[:, -max_cache_length:] position_ids = kwargs.get("position_ids", None) if attention_mask is not None and position_ids is None: # create position_ids on the fly for batch generation position_ids = attention_mask.long().cumsum(-1) - 1 position_ids.masked_fill_(attention_mask == 0, 1) if past_key_values: position_ids = position_ids[:, -input_ids.shape[1] :] # if `inputs_embeds` are passed, we only want to use them in the 1st generation step if inputs_embeds is not None and past_key_values is None: model_inputs = {"inputs_embeds": inputs_embeds} else: model_inputs = {"input_ids": input_ids} model_inputs.update( { "position_ids": position_ids, "past_key_values": past_key_values, "use_cache": kwargs.get("use_cache"), "attention_mask": attention_mask, } ) return model_inputs @staticmethod def _reorder_cache(past_key_values, beam_idx): reordered_past = () for layer_past in past_key_values: reordered_past += ( tuple( past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past ), ) return reordered_past @staticmethod def _top_k_logits(logits, k): if k is None or k <= 0: return logits else: values, _ = torch.topk(logits, k) min_values = values[..., -1, None] return torch.where( logits < min_values, torch.full_like(logits, float("-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_indices = torch.scatter( torch.full_like(logits, False, dtype=torch.bool), -1, sorted_indices, sorted_mask, ) return logits.masked_fill(mask_indices, float("-inf")) def _sample_with_temperature_topk_topp( self, logits, temperature=1.0, top_k=0, top_p=1.0 ): orig_shape = logits.shape[:-1] vocab_size = logits.shape[-1] logits = logits.reshape(-1, vocab_size) 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 = torch.gather(probs, -1, token) return token.view(*orig_shape), token_prob.view(*orig_shape) @staticmethod def _get_num_transfer_tokens(block_length, steps): if steps == 0: return torch.tensor([], dtype=torch.int64) base = block_length // steps remainder = block_length % steps num_transfer_tokens = torch.full((steps,), base, dtype=torch.int64) num_transfer_tokens[:remainder] += 1 return num_transfer_tokens @staticmethod def _apply_edit_operations_with_tracking( block_tokens, old_block_tokens, is_original_mask_snapshot, mask_id, block_length, delete_token_id, split_token_id, ): """Process DELETE and SPLIT tokens in a block while tracking ``is_original_mask``. The block is kept at a fixed ``block_length`` by truncating (when SPLIT grows it) or right-padding with masks (when DELETE shrinks it). Args: block_tokens: List of token ids in the block (current state after writes). old_block_tokens: List of token ids before this step (used to restore SPLIT's carried token). is_original_mask_snapshot: Per-position bools marking positions that were original masks. mask_id: The mask token id. block_length: Target (fixed) block length. delete_token_id / split_token_id: Special edit-operation token ids. Behaviour: - DELETE: skip the token (and its tracking entry). - SPLIT: token -> [mask_id, old_token], tracking -> [False, snapshot_val]. - Kept token: carried through with its snapshot tracking value. - Padding masks (from shrink) are tracked as non-original (False). Returns: (result_tokens, result_tracking): both lists of length ``block_length``. """ result_tokens = [] result_tracking = [] for i, token in enumerate(block_tokens): if token == delete_token_id: continue elif token == split_token_id: old_token = old_block_tokens[i] if i < len(old_block_tokens) else mask_id snapshot_val = ( is_original_mask_snapshot[i] if i < len(is_original_mask_snapshot) else False ) result_tokens.extend([mask_id, old_token]) result_tracking.extend([False, snapshot_val]) else: snapshot_val = ( is_original_mask_snapshot[i] if i < len(is_original_mask_snapshot) else False ) result_tokens.append(token) result_tracking.append(snapshot_val) if len(result_tokens) > block_length: result_tokens = result_tokens[:block_length] result_tracking = result_tracking[:block_length] elif len(result_tokens) < block_length: pad_count = block_length - len(result_tokens) result_tokens.extend([mask_id] * pad_count) result_tracking.extend([False] * pad_count) return result_tokens, result_tracking def _diffusion_sample(self, logits, temperature=1.0, top_k=None, top_p=None): """Sample a token id per position and return its (raw-softmax) confidence. With ``temperature == 0`` this is greedy (argmax). Otherwise it applies temperature / top-k / top-p filtering and samples. The returned probability is always taken from the *unfiltered* softmax so it can be used directly as a confidence score. """ orig_shape = logits.shape[:-1] vocab_size = logits.shape[-1] logits = logits.reshape(-1, vocab_size) probs_full = F.softmax(logits, dim=-1) if temperature is None or temperature == 0.0: token = torch.argmax(logits, dim=-1, keepdim=True) else: scaled = logits / temperature if temperature != 1.0 else logits scaled = self._top_k_logits(scaled, top_k or 0) scaled = self._top_p_logits(scaled, top_p if top_p is not None else 1.0) probs = F.softmax(scaled, dim=-1) token = torch.multinomial(probs, num_samples=1) token_prob = torch.gather(probs_full, -1, token) return token.view(*orig_shape), token_prob.view(*orig_shape) def _resample_to_escape_loop( self, block_ids, old_block_ids, block_logits, mt2_index, t2t_index, seen_block_results, temperature, top_k, top_p, max_iters=5, ): """Escape a decoding loop by resampling one changed position at a time. Called after M2T/T2T writes but BEFORE DELETE/SPLIT processing. If the current (pre-edit) block state repeats a previously seen state, repeatedly pick one random changed position (``block_ids != old_block_ids``) and resample it until the block is novel or ``max_iters`` is reached. The currently chosen token is masked out (-inf) before resampling. M2T positions resample with temperature/top-k/top-p; T2T positions resample greedily. Modifies ``block_ids`` in place. """ if tuple(block_ids.tolist()) not in seen_block_results: return for _ in range(max_iters): changed_positions = ( (block_ids != old_block_ids).nonzero(as_tuple=True)[0].tolist() ) if not changed_positions: break # Use torch's RNG (not the ``random`` module) so the choice honors torch.manual_seed. rand_idx = torch.randint(len(changed_positions), (1,), device=block_ids.device).item() pos = changed_positions[rand_idx] current_token = block_ids[pos].item() pos_logits = block_logits[pos, :].clone() pos_logits[current_token] = float("-inf") if bool(mt2_index[pos]): new_token, _ = self._diffusion_sample( pos_logits.unsqueeze(0), temperature=temperature, top_k=top_k, top_p=top_p ) elif bool(t2t_index[pos]): new_token, _ = self._diffusion_sample( pos_logits.unsqueeze(0), temperature=0.0, top_k=None, top_p=None ) else: # A changed position must be M2T or T2T; skip anything unexpected. continue block_ids[pos] = new_token.view(-1)[0] if tuple(block_ids.tolist()) not in seen_block_results: break @torch.no_grad() def _joint_decode_block( self, x, block_start, block_end, attention_mask, position_ids, temperature, top_k, top_p, steps, threshold, editing_threshold, max_post_steps, mask_id, delete_token_id, split_token_id, max_steps_per_block, ): """Iteratively refine a single block in place using joint M2T + T2T + edit ops. The active block is ``x[0, block_start:block_end]`` (fixed length). Earlier positions of ``x`` provide frozen context via ``attention_mask`` / ``position_ids``. Each step: 1. forward over ``x[:, :block_end]``, 2. M2T: fill masks whose greedy confidence clears ``threshold`` (with a per-step floor from the ``steps`` transfer schedule), 3. T2T: rewrite already-generated tokens whose greedy confidence clears ``editing_threshold`` and whose greedy token differs, 4. anti-loop resample if the block state repeats, 5. consume DELETE/SPLIT edit tokens (block stays fixed length). Once all *original* masks are gone, up to ``max_post_steps`` further refinement steps run; on the final such step DELETE/SPLIT are suppressed so the block can terminate. """ device = x.device block_length = block_end - block_start # Positions that were NOT masks at block entry are prompt/context tokens: they are never # written and never carry DELETE/SPLIT, so this mask stays aligned across edits. prompt_mask_block = (x[0, block_start:block_end] != mask_id).clone() # A prompt/context position normally never holds a reserved edit token. If one does (bad # upstream template, history, or malformed input) it would be silently deleted/expanded by # the edit-op pass. Warn instead of failing silently. block_ids = x[0, block_start:block_end] prompt_edit_tokens = prompt_mask_block & ( (block_ids == delete_token_id) | (block_ids == split_token_id) ) if prompt_edit_tokens.any(): logger.warning_once( "Reserved edit token(s) found in a prompt/context segment of block@%d; " "they will be deleted/expanded by the edit-op pass.", block_start, ) is_original_mask = (x[0, block_start:block_end] == mask_id).tolist() initial_mask_count = sum(is_original_mask) if initial_mask_count == 0: return # fully prompt/context block, nothing to decode # Per-step floor for M2T: spread the block's initial masks over ``steps`` steps. Fewer # steps -> more forced unmaskings per step. The schedule sums to ``initial_mask_count``. transfer_schedule = self._get_num_transfer_tokens(initial_mask_count, steps) seen_block_results = set() # pre-edit block states seen this block, for loop detection post_steps = 0 step_id = 0 while True: blk = x[0, block_start:block_end] old_block = blk.clone() input_key = tuple(old_block.tolist()) mask_index = (old_block == mask_id) & (~prompt_mask_block) original_mask_count = sum( 1 for i, v in enumerate(is_original_mask) if v and old_block[i].item() == mask_id ) new_mask_count = mask_index.sum().item() - original_mask_count # Track post-mask refinement steps: reset while original masks remain. if original_mask_count == 0: post_steps += 1 else: post_steps = 0 logger.debug( "block@%d step=%d original_masks=%d new_masks=%d post_steps=%d/%d", block_start, step_id, original_mask_count, new_mask_count, post_steps, max_post_steps, ) # Exit guard: require ZERO remaining masks (original *and* new), not just # original_mask_count == 0. An edit op can leave a fresh mask exactly when the last # original mask is resolved; keying off original masks alone would exit early and # return residual mask_id. The final round (below) force-resolves all masks, so this # normally holds immediately; the total-count guard is the belt-and-suspenders check. if mask_index.sum().item() == 0 and post_steps > max_post_steps: logger.debug( "block@%d terminating: max_post_steps (%d) exceeded after %d steps", block_start, max_post_steps, step_id, ) break if step_id >= max_steps_per_block: logger.debug( "block@%d terminating: max_steps_per_block (%d) reached", block_start, max_steps_per_block, ) break # 1. Forward pass over the current window. logits = self.forward( x[:, :block_end], attention_mask=attention_mask, position_ids=position_ids, ).logits block_logits = logits[0, block_start:block_end, :] # 2. Sampled (temperature) result and greedy result + confidences. x_s, p_s = self._diffusion_sample(block_logits, temperature, top_k, top_p) if temperature != 0.0: x0, p0 = self._diffusion_sample(block_logits, 0.0, None, None) else: x0, p0 = x_s, p_s neg_inf = torch.full_like(p0, -float("inf")) # 3. M2T (mask -> token): threshold-gated with a per-step floor. # The gate uses ``p_s`` -- the confidence of the token that will actually be written # (``x_s``) -- so a position is only unmasked when the sampled token itself is # confident. When temperature == 0, ``p_s == p0``, so this reduces to greedy behavior. mt2_index = torch.zeros(block_length, dtype=torch.bool, device=device) if mask_index.any(): if step_id < len(transfer_schedule): num_need = transfer_schedule[step_id].item() + new_mask_count mask_conf = torch.where(mask_index, p_s, neg_inf) high_conf = (mask_conf > threshold) & mask_index if high_conf.sum().item() >= num_need: mt2_index = high_conf else: k_val = min(num_need, mask_index.sum().item()) if k_val > 0: _, idx = torch.topk(mask_conf, k=k_val) mt2_index[idx] = True else: mt2_index = mask_index # 4. T2T (token -> token edit): high-confidence rewrites of generated tokens. editable_mask = (~mask_index) & (~prompt_mask_block) editing_confidence = torch.where(editable_mask, p0, neg_inf) high_conf_edit = (editing_confidence > editing_threshold) & editable_mask token_changed = old_block != x0 t2t_index = high_conf_edit & token_changed fill_index_pre = mt2_index | t2t_index # Final round = the last refinement step, after which the block terminates. On it we # suppress SPLIT/DELETE (below) so no new masks appear. # # max_post_steps > 0: count-based -- fire once we've spent the post-mask step budget. # max_post_steps == 0: "no post steps" means the step that resolves the LAST original # mask must itself finish the block. post_steps can't detect this (it is still 0 at # the top of that step), so we detect it from mt2_index: this step is final iff its # M2T covers every still-unresolved original mask. Guaranteed to fire eventually -- # the transfer schedule forces mt2_index == mask_index by step len(transfer_schedule). if max_post_steps > 0: final_round = post_steps >= max_post_steps else: remaining_original = mask_index & torch.tensor(is_original_mask, device=device) final_round = remaining_original.any() and (remaining_original <= mt2_index).all() # 5. On the final round, suppress SPLIT/DELETE so the block can terminate. Only the # positions that are actually written matter: M2T writes x_s, T2T writes x0. Other # positions may still contain S/D harmlessly since they are never written. Resampling # is batched over all offending positions (S/D columns masked to -inf). if final_round: # The final round must leave the block fully unmasked, so resolve EVERY remaining # mask here -- before the D/S suppression below, so these positions also get their # SPLIT/DELETE stripped and no fresh mask survives. For max_post_steps > 0 this is # already implied (once original masks are gone, new_mask_count inflates num_need so # M2T selects all masks anyway), so it is a no-op there; stating it makes the # invariant explicit and robust to changes in the M2T selection above. mt2_index = mask_index fill_index_pre = mt2_index | t2t_index m2t_sd = mt2_index & ((x_s == split_token_id) | (x_s == delete_token_id)) if m2t_sd.any(): sd_logits = block_logits[m2t_sd].clone() sd_logits[:, split_token_id] = float("-inf") sd_logits[:, delete_token_id] = float("-inf") new_tokens, _ = self._diffusion_sample( sd_logits, temperature, top_k, top_p ) x_s[m2t_sd] = new_tokens t2t_sd = t2t_index & ((x0 == split_token_id) | (x0 == delete_token_id)) if t2t_sd.any(): sd_logits = block_logits[t2t_sd].clone() sd_logits[:, split_token_id] = float("-inf") sd_logits[:, delete_token_id] = float("-inf") new_tokens, _ = self._diffusion_sample( sd_logits, temperature=0.0, top_k=None, top_p=None ) x0[t2t_sd] = new_tokens # 6. Apply writes: M2T writes the sampled token, T2T writes the greedy token. is_original_mask_snapshot = list(is_original_mask) if fill_index_pre.any(): if mt2_index.any(): blk[mt2_index] = x_s[mt2_index] if t2t_index.any(): blk[t2t_index] = x0[t2t_index] # Anti-loop: escape a repeated pre-edit block state. Skipped on the final round, # where SD suppression already drives termination. if not final_round: self._resample_to_escape_loop( blk, old_block, block_logits, mt2_index, t2t_index, seen_block_results, temperature=temperature, top_k=top_k, top_p=top_p, ) seen_block_results.add(tuple(blk.tolist())) # 7. Consume DELETE/SPLIT edit tokens (block stays fixed length). edited_block, is_original_mask = self._apply_edit_operations_with_tracking( blk.tolist(), old_block.tolist(), is_original_mask_snapshot, mask_id, block_length, delete_token_id=delete_token_id, split_token_id=split_token_id, ) x[0, block_start:block_end] = torch.tensor( edited_block, device=device, dtype=x.dtype ) # 8. Stop when the block is stable and fully unmasked. if x[0, block_start:block_end].tolist() == list(input_key) and ( x[0, block_start:block_end] == mask_id ).sum() == 0: step_id += 1 break step_id += 1 @torch.no_grad() def generate( self, inputs: Optional[torch.Tensor] = 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, threshold: float = 0.5, editing_threshold: float = 0.0, max_post_steps: int = 16, eos_early_stop: bool = False, eos_id: int = 156892, mask_id: int = 156895, delete_token_id: int = 156930, split_token_id: int = 156931, max_steps_per_block: int = 1000, ): r""" Generate tokens with a block-wise, edit-based iterative refinement strategy. Unlike autoregressive generation, this method lays out a full masked template and refines it block by block. Within each block it jointly performs: - **M2T** (mask -> token): converts ``mask_id`` placeholders into concrete tokens once their confidence exceeds ``threshold`` (with a per-step floor so progress is guaranteed). - **T2T** (token -> token): rewrites already-generated tokens whose greedy confidence exceeds ``editing_threshold`` and whose greedy prediction differs from the current token. - **DELETE / SPLIT**: consumes special edit tokens to remove positions or insert new masks, letting the block change its content length (kept fixed by pad/truncate). An anti-loop resampler perturbs the block whenever a pre-edit state repeats. After all original masks in a block are resolved, up to ``max_post_steps`` further refinement steps run; the final one suppresses DELETE/SPLIT so the block terminates. A block-diagonal causal attention mask lets a block attend to all previous blocks (and bidirectionally within itself) but not to future blocks. Parameters: inputs (`torch.Tensor`): Prompt token ids of shape ``(1, prompt_length)``. temperature (`float`, defaults to 0.0): 0.0 is greedy; >0 enables sampling for M2T. block_length (`int`, defaults to 32): Fixed length of each generation block. steps (`int`, defaults to 32): Number of steps the M2T transfer schedule spreads a block's initial masks over (per-step unmasking floor). Fewer steps forces more unmaskings per step. Independent of the actual number of iterations, which is driven by the confidence thresholds and the post-mask refinement phase. gen_length (`int`, defaults to 2048): Number of tokens to generate after the prompt. top_p / top_k (`float`/`int`, *optional*): Nucleus / top-k filtering for sampling. threshold (`float`, defaults to 0.5): Confidence threshold for M2T unmasking. editing_threshold (`float`, defaults to 0.0): Confidence threshold for T2T edits. max_post_steps (`int`, defaults to 16): Max refinement steps after a block's original masks are all resolved. eos_early_stop (`bool`, defaults to False): Stop after a block that produced ``eos_id``. eos_id / mask_id (`int`): End-of-sequence and mask placeholder token ids. delete_token_id / split_token_id (`int`): Special edit-operation token ids. max_steps_per_block (`int`, defaults to 1000): Hard safety cap on steps per block. Return: `torch.Tensor`: The generated token ids after the prompt, up to and including the first ``eos_id`` (or ``gen_length`` if none is produced). """ input_ids = inputs.to(self.device) 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)) block_diffusion_attention_mask = ( ( block_mask.repeat_interleave(block_length, dim=0) .repeat_interleave(block_length, dim=1) .unsqueeze(0) .unsqueeze(0) ) .log() .to(torch.bfloat16) ) 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.clone() prefill_blocks = prompt_length // block_length for num_block in range(prefill_blocks, num_blocks): block_start = num_block * block_length block_end = (num_block + 1) * block_length cur_attn_mask = block_diffusion_attention_mask[ :, :, :block_end, :block_end ] cur_position_ids = position_ids[:, :block_end] self._joint_decode_block( x, block_start, block_end, cur_attn_mask, cur_position_ids, temperature=temperature, top_k=top_k, top_p=top_p, steps=steps, threshold=threshold, editing_threshold=editing_threshold, max_post_steps=max_post_steps, mask_id=mask_id, delete_token_id=delete_token_id, split_token_id=split_token_id, max_steps_per_block=max_steps_per_block, ) if ( eos_early_stop and eos_id is not None and (x[0, prompt_length:block_end] == eos_id).any() ): break generated_answer = x[:, : prompt_length + gen_length] eos_positions = (generated_answer[0][prompt_length:] == eos_id).nonzero( as_tuple=True )[0] if len(eos_positions) > 0: first_eos_position = eos_positions[0].item() else: first_eos_position = gen_length output = generated_answer[ :, prompt_length : prompt_length + first_eos_position + 1 ] # Safety net: a well-formed decode leaves no mask / edit tokens in the output. If any # survive (e.g. a degenerate config such as max_post_steps=0 that fails to converge, or a # block that hit max_steps_per_block), warn instead of silently returning residual # mask_id / split / delete tokens to the caller. residual = ( (output == mask_id) | (output == split_token_id) | (output == delete_token_id) ) if residual.any(): logger.warning( "Decoding finished with %d residual mask/split/delete token(s) in the output; " "the generation may be malformed.", int(residual.sum().item()), ) return output