MSALab/PerceptionDLM-Base
669
1# coding=utf-82# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.3#4# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX5# and OPT implementations in this library. It has been modified from its6# original forms to accommodate minor architectural differences compared7# to GPT-NeoX and OPT used by the Meta AI team that trained the model.8#9# Licensed under the Apache License, Version 2.0 (the "License");10# you may not use this file except in compliance with the License.11# You may obtain a copy of the License at12#13# http://www.apache.org/licenses/LICENSE-2.014#15# Unless required by applicable law or agreed to in writing, software16# distributed under the License is distributed on an "AS IS" BASIS,17# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.18# See the License for the specific language governing permissions and19# limitations under the License.20"""PyTorch LLaDA model."""21 22import math23import warnings24from typing import List, Optional, Tuple, Union25import numpy as np26import copy27 28import torch29import torch.nn.functional as F30import torch.utils.checkpoint31from torch import nn32from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss33 34from transformers.activations import ACT2FN35from transformers.cache_utils import Cache, DynamicCache, StaticCache36from transformers.modeling_attn_mask_utils import AttentionMaskConverter37from transformers.modeling_outputs import (38 BaseModelOutputWithPast,39 CausalLMOutputWithPast,40 QuestionAnsweringModelOutput,41 SequenceClassifierOutputWithPast,42)43from transformers.modeling_utils import PreTrainedModel44from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS45from transformers.utils import (46 add_start_docstrings,47 add_start_docstrings_to_model_forward,48 is_flash_attn_2_available,49 is_flash_attn_greater_or_equal_2_10,50 logging,51 replace_return_docstrings,52)53from .configuration_llada import LLaDAConfig54from .cache import dLLMCache, dLLMCacheConfig55 56if is_flash_attn_2_available():57 from flash_attn import flash_attn_func, flash_attn_varlen_func58 from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa59 60 61logger = logging.get_logger(__name__)62 63_CONFIG_FOR_DOC = "LLaDAConfig"64 65 66def _get_unpad_data(attention_mask):67 seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)68 indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()69 max_seqlen_in_batch = seqlens_in_batch.max().item()70 cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))71 return (72 indices,73 cu_seqlens,74 max_seqlen_in_batch,75 )76 77 78class LLaDARMSNorm(nn.Module):79 def __init__(self, hidden_size, eps=1e-6):80 """81 LLaDARMSNorm is equivalent to T5LayerNorm82 """83 super().__init__()84 self.weight = nn.Parameter(torch.ones(hidden_size))85 self.variance_epsilon = eps86 87 def forward(self, hidden_states):88 input_dtype = hidden_states.dtype89 hidden_states = hidden_states.to(torch.float32)90 variance = hidden_states.pow(2).mean(-1, keepdim=True)91 hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)92 return self.weight * hidden_states.to(input_dtype)93 94 95ALL_LAYERNORM_LAYERS.append(LLaDARMSNorm)96 97 98class LLaDARotaryEmbedding(nn.Module):99 def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):100 super().__init__()101 self.scaling_factor = scaling_factor102 self.dim = dim103 self.max_position_embeddings = max_position_embeddings104 self.base = base105 inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(device) / self.dim))106 self.register_buffer("inv_freq", inv_freq, persistent=False)107 # For BC we register cos and sin cached108 self.max_seq_len_cached = max_position_embeddings109 t = torch.arange(self.max_seq_len_cached, device=device, dtype=torch.int64).type_as(self.inv_freq)110 t = t / self.scaling_factor111 freqs = torch.outer(t, self.inv_freq)112 # Different from paper, but it uses a different permutation in order to obtain the same calculation113 emb = torch.cat((freqs, freqs), dim=-1)114 self.register_buffer("_cos_cached", emb.cos().to(torch.get_default_dtype()), persistent=False)115 self.register_buffer("_sin_cached", emb.sin().to(torch.get_default_dtype()), persistent=False)116 117 @property118 def sin_cached(self):119 logger.warning_once(120 "The sin_cached attribute will be removed in 4.39. Bear in mind that its contents changed in v4.38. Use "121 "the forward method of RoPE from now on instead. It is not used in the `LLaDAAttention` class"122 )123 return self._sin_cached124 125 @property126 def cos_cached(self):127 logger.warning_once(128 "The cos_cached attribute will be removed in 4.39. Bear in mind that its contents changed in v4.38. Use "129 "the forward method of RoPE from now on instead. It is not used in the `LLaDAAttention` class"130 )131 return self._cos_cached132 133 @torch.no_grad()134 def forward(self, x, position_ids):135 # x: [bs, num_attention_heads, seq_len, head_size]136 inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)137 position_ids_expanded = position_ids[:, None, :].float()138 # Force float32 since bfloat16 loses precision on long contexts139 # See https://github.com/huggingface/transformers/pull/29285140 device_type = x.device.type141 device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu"142 with torch.autocast(device_type=device_type, enabled=False):143 freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)144 emb = torch.cat((freqs, freqs), dim=-1)145 cos = emb.cos()146 sin = emb.sin()147 return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)148 149 150class LLaDALinearScalingRotaryEmbedding(LLaDARotaryEmbedding):151 """LLaDARotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev"""152 153 def forward(self, x, position_ids):154 # difference to the original RoPE: a scaling factor is aplied to the position ids155 position_ids = position_ids.float() / self.scaling_factor156 cos, sin = super().forward(x, position_ids)157 return cos, sin158 159 160class LLaDADynamicNTKScalingRotaryEmbedding(LLaDARotaryEmbedding):161 """LLaDARotaryEmbedding extended with Dynamic NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla"""162 163 def forward(self, x, position_ids):164 # difference to the original RoPE: inv_freq is recomputed when the sequence length > original length165 seq_len = torch.max(position_ids) + 1166 if seq_len > self.max_position_embeddings:167 base = self.base * (168 (self.scaling_factor * seq_len / self.max_position_embeddings) - (self.scaling_factor - 1)169 ) ** (self.dim / (self.dim - 2))170 inv_freq = 1.0 / (171 base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(x.device) / self.dim)172 )173 self.register_buffer("inv_freq", inv_freq, persistent=False) # TODO joao: this may break with compilation174 175 cos, sin = super().forward(x, position_ids)176 return cos, sin177 178 179def rotate_half(x):180 """Rotates half the hidden dims of the input."""181 x1 = x[..., : x.shape[-1] // 2]182 x2 = x[..., x.shape[-1] // 2 :]183 return torch.cat((-x2, x1), dim=-1)184 185 186def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):187 """Applies Rotary Position Embedding to the query and key tensors.188 189 Args:190 q (`torch.Tensor`): The query tensor.191 k (`torch.Tensor`): The key tensor.192 cos (`torch.Tensor`): The cosine part of the rotary embedding.193 sin (`torch.Tensor`): The sine part of the rotary embedding.194 position_ids (`torch.Tensor`, *optional*):195 Deprecated and unused.196 unsqueeze_dim (`int`, *optional*, defaults to 1):197 The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and198 sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note199 that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and200 k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes201 cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have202 the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.203 Returns:204 `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.205 """206 cos = cos.unsqueeze(unsqueeze_dim)207 sin = sin.unsqueeze(unsqueeze_dim)208 q_embed = (q * cos) + (rotate_half(q) * sin)209 k_embed = (k * cos) + (rotate_half(k) * sin)210 return q_embed, k_embed211 212 213class LLaDAMLP(nn.Module):214 def __init__(self, config):215 super().__init__()216 self.config = config217 self.hidden_size = config.hidden_size218 self.intermediate_size = config.intermediate_size219 self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)220 self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)221 self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)222 self.act_fn = ACT2FN[config.hidden_act]223 224 def forward(self, x):225 if self.config.pretraining_tp > 1:226 slice = self.intermediate_size // self.config.pretraining_tp227 gate_proj_slices = self.gate_proj.weight.split(slice, dim=0)228 up_proj_slices = self.up_proj.weight.split(slice, dim=0)229 down_proj_slices = self.down_proj.weight.split(slice, dim=1)230 231 gate_proj = torch.cat(232 [F.linear(x, gate_proj_slices[i]) for i in range(self.config.pretraining_tp)], dim=-1233 )234 up_proj = torch.cat([F.linear(x, up_proj_slices[i]) for i in range(self.config.pretraining_tp)], dim=-1)235 236 intermediate_states = (self.act_fn(gate_proj) * up_proj).split(slice, dim=2)237 down_proj = [238 F.linear(intermediate_states[i], down_proj_slices[i]) for i in range(self.config.pretraining_tp)239 ]240 down_proj = sum(down_proj)241 else:242 down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))243 244 return down_proj245 246 247def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:248 """249 This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,250 num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)251 """252 batch, num_key_value_heads, slen, head_dim = hidden_states.shape253 if n_rep == 1:254 return hidden_states255 hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)256 return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)257 258 259class LLaDAAttention(nn.Module):260 """Multi-headed attention from 'Attention Is All You Need' paper"""261 262 def __init__(self, config: LLaDAConfig, layer_idx: Optional[int] = None):263 super().__init__()264 self.config = config265 self.layer_idx = layer_idx266 if layer_idx is None:267 logger.warning_once(268 f"Instantiating {self.__class__.__name__} without passing a `layer_idx` is not recommended and will "269 "lead to errors during the forward call if caching is used. Please make sure to provide a `layer_idx` "270 "when creating this class."271 )272 273 self.attention_dropout = config.attention_dropout274 self.hidden_size = config.hidden_size275 self.num_heads = config.num_attention_heads276 self.head_dim = self.hidden_size // self.num_heads277 self.num_key_value_heads = config.num_key_value_heads278 self.num_key_value_groups = self.num_heads // self.num_key_value_heads279 self.max_position_embeddings = config.max_position_embeddings280 self.rope_theta = config.rope_theta281 #self.is_causal = True282 # Modify: MDM set causal to False.283 self.is_causal = False284 285 if (self.head_dim * self.num_heads) != self.hidden_size:286 raise ValueError(287 f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"288 f" and `num_heads`: {self.num_heads})."289 )290 291 self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias)292 self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias)293 self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias)294 self.o_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=config.attention_bias)295 self._init_rope()296 297 def _init_rope(self):298 if self.config.rope_scaling is None:299 self.rotary_emb = LLaDARotaryEmbedding(300 self.head_dim,301 max_position_embeddings=self.max_position_embeddings,302 base=self.rope_theta,303 )304 else:305 scaling_type = self.config.rope_scaling["type"]306 scaling_factor = self.config.rope_scaling["factor"]307 if scaling_type == "linear":308 self.rotary_emb = LLaDALinearScalingRotaryEmbedding(309 self.head_dim,310 max_position_embeddings=self.max_position_embeddings,311 scaling_factor=scaling_factor,312 base=self.rope_theta,313 )314 elif scaling_type == "dynamic":315 self.rotary_emb = LLaDADynamicNTKScalingRotaryEmbedding(316 self.head_dim,317 max_position_embeddings=self.max_position_embeddings,318 scaling_factor=scaling_factor,319 base=self.rope_theta,320 )321 else:322 raise ValueError(f"Unknown RoPE scaling type {scaling_type}")323 324 def forward(325 self,326 hidden_states: torch.Tensor,327 attention_mask: Optional[torch.Tensor] = None,328 position_ids: Optional[torch.LongTensor] = None,329 past_key_value: Optional[Cache] = None,330 output_attentions: bool = False,331 use_cache: bool = False,332 cache_position: Optional[torch.LongTensor] = None,333 **kwargs,334 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:335 bsz, q_len, _ = hidden_states.size()336 337 if self.config.pretraining_tp > 1:338 key_value_slicing = (self.num_key_value_heads * self.head_dim) // self.config.pretraining_tp339 query_slices = self.q_proj.weight.split(340 (self.num_heads * self.head_dim) // self.config.pretraining_tp, dim=0341 )342 key_slices = self.k_proj.weight.split(key_value_slicing, dim=0)343 value_slices = self.v_proj.weight.split(key_value_slicing, dim=0)344 345 query_states = [F.linear(hidden_states, query_slices[i]) for i in range(self.config.pretraining_tp)]346 query_states = torch.cat(query_states, dim=-1)347 348 key_states = [F.linear(hidden_states, key_slices[i]) for i in range(self.config.pretraining_tp)]349 key_states = torch.cat(key_states, dim=-1)350 351 value_states = [F.linear(hidden_states, value_slices[i]) for i in range(self.config.pretraining_tp)]352 value_states = torch.cat(value_states, dim=-1)353 354 else:355 query_states = self.q_proj(hidden_states)356 key_states = self.k_proj(hidden_states)357 value_states = self.v_proj(hidden_states)358 359 query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)360 key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)361 value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)362 363 past_key_value = getattr(self, "past_key_value", past_key_value)364 cos, sin = self.rotary_emb(value_states, position_ids)365 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)366 367 if past_key_value is not None:368 # sin and cos are specific to RoPE models; cache_position needed for the static cache369 cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}370 key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)371 372 key_states = repeat_kv(key_states, self.num_key_value_groups)373 value_states = repeat_kv(value_states, self.num_key_value_groups)374 375 attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)376 377 if attention_mask is not None: # no matter the length, we just slice it378 causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]379 attn_weights = attn_weights + causal_mask380 381 # upcast attention to fp32382 attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)383 attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training)384 attn_output = torch.matmul(attn_weights, value_states)385 386 if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):387 raise ValueError(388 f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"389 f" {attn_output.size()}"390 )391 392 attn_output = attn_output.transpose(1, 2).contiguous()393 394 attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)395 396 if self.config.pretraining_tp > 1:397 attn_output = attn_output.split(self.hidden_size // self.config.pretraining_tp, dim=2)398 o_proj_slices = self.o_proj.weight.split(self.hidden_size // self.config.pretraining_tp, dim=1)399 attn_output = sum([F.linear(attn_output[i], o_proj_slices[i]) for i in range(self.config.pretraining_tp)])400 else:401 attn_output = self.o_proj(attn_output)402 403 if not output_attentions:404 attn_weights = None405 406 return attn_output, attn_weights, past_key_value407 408 409class LLaDAFlashAttention2(LLaDAAttention):410 """411 LLaDA flash attention module. This module inherits from `LLaDAAttention` as the weights of the module stays412 untouched. The only required change would be on the forward pass where it needs to correctly call the public API of413 flash attention and deal with padding tokens in case the input contains any of them.414 """415 416 def __init__(self, *args, **kwargs):417 super().__init__(*args, **kwargs)418 # print("Using Flash Attention2 !!!")419 420 # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.421 # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0.422 # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left).423 self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()424 425 def forward(426 self,427 hidden_states: torch.Tensor,428 attention_mask: Optional[torch.LongTensor] = None,429 position_ids: Optional[torch.LongTensor] = None,430 past_key_value: Optional[Cache] = None,431 output_attentions: bool = False,432 use_cache: bool = False,433 cache_position: Optional[torch.LongTensor] = None,434 **kwargs,435 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:436 output_attentions = False437 438 bsz, q_len, _ = hidden_states.size()439 440 query_states = self.q_proj(hidden_states)441 key_states = self.k_proj(hidden_states)442 value_states = self.v_proj(hidden_states)443 444 # Flash attention requires the input to have the shape445 # batch_size x seq_length x head_dim x hidden_dim446 # therefore we just need to keep the original shape447 query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)448 key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)449 value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)450 451 cos, sin = self.rotary_emb(value_states, position_ids)452 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)453 454 past_key_value = getattr(self, "past_key_value", past_key_value)455 456 if past_key_value is not None:457 # sin and cos are specific to RoPE models; cache_position needed for the static cache458 cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}459 key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)460 461 # TODO: These transpose are quite inefficient but Flash Attention requires the layout [batch_size, sequence_length, num_heads, head_dim]. We would need to refactor the KV cache462 # to be able to avoid many of these transpose/reshape/view.463 query_states = query_states.transpose(1, 2)464 key_states = key_states.transpose(1, 2)465 value_states = value_states.transpose(1, 2)466 467 dropout_rate = self.attention_dropout if self.training else 0.0468 469 # In PEFT, usually we cast the layer norms in float32 for training stability reasons470 # therefore the input hidden states gets silently casted in float32. Hence, we need471 # cast them back in the correct dtype just to be sure everything works as expected.472 # This might slowdown training & inference so it is recommended to not cast the LayerNorms473 # in fp32. (LLaDARMSNorm handles it correctly)474 475 input_dtype = query_states.dtype476 if input_dtype == torch.float32:477 if torch.is_autocast_enabled():478 target_dtype = torch.get_autocast_gpu_dtype()479 # Handle the case where the model is quantized480 elif hasattr(self.config, "_pre_quantization_dtype"):481 target_dtype = self.config._pre_quantization_dtype482 else:483 target_dtype = self.q_proj.weight.dtype484 485 logger.warning_once(486 f"The input hidden states seems to be silently casted in float32, this might be related to"487 f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"488 f" {target_dtype}."489 )490 491 query_states = query_states.to(target_dtype)492 key_states = key_states.to(target_dtype)493 value_states = value_states.to(target_dtype)494 495 attn_output = self._flash_attention_forward(496 query_states, key_states, value_states, attention_mask, q_len, dropout=dropout_rate497 )498 499 attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()500 attn_output = self.o_proj(attn_output)501 502 if not output_attentions:503 attn_weights = None504 505 return attn_output, attn_weights, past_key_value506 507 def _flash_attention_forward(508 self, query_states, key_states, value_states, attention_mask, query_length, dropout=0.0, softmax_scale=None509 ):510 """511 Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token512 first unpad the input, then computes the attention scores and pad the final attention scores.513 514 Args:515 query_states (`torch.Tensor`):516 Input query states to be passed to Flash Attention API517 key_states (`torch.Tensor`):518 Input key states to be passed to Flash Attention API519 value_states (`torch.Tensor`):520 Input value states to be passed to Flash Attention API521 attention_mask (`torch.Tensor`):522 The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the523 position of padding tokens and 1 for the position of non-padding tokens.524 dropout (`float`):525 Attention dropout526 softmax_scale (`float`, *optional*):527 The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)528 """529 if not self._flash_attn_uses_top_left_mask:530 causal = self.is_causal531 else:532 # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in LLaDAFlashAttention2 __init__.533 causal = self.is_causal and query_length != 1534 535 assert causal is False # Modify: MDM536 537 # Contains at least one padding token in the sequence538 if attention_mask is not None:539 batch_size = query_states.shape[0]540 query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(541 query_states, key_states, value_states, attention_mask, query_length542 )543 544 cu_seqlens_q, cu_seqlens_k = cu_seq_lens545 max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens546 547 attn_output_unpad = flash_attn_varlen_func(548 query_states,549 key_states,550 value_states,551 cu_seqlens_q=cu_seqlens_q,552 cu_seqlens_k=cu_seqlens_k,553 max_seqlen_q=max_seqlen_in_batch_q,554 max_seqlen_k=max_seqlen_in_batch_k,555 dropout_p=dropout,556 softmax_scale=softmax_scale,557 causal=causal,558 )559 560 attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)561 else:562 attn_output = flash_attn_func(563 query_states, key_states, value_states, dropout, softmax_scale=softmax_scale, causal=causal564 )565 566 return attn_output567 568 def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):569 indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)570 batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape571 572 key_layer = index_first_axis(573 key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k574 )575 value_layer = index_first_axis(576 value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k577 )578 if query_length == kv_seq_len:579 query_layer = index_first_axis(580 query_layer.reshape(batch_size * kv_seq_len, self.num_heads, head_dim), indices_k581 )582 cu_seqlens_q = cu_seqlens_k583 max_seqlen_in_batch_q = max_seqlen_in_batch_k584 indices_q = indices_k585 elif query_length == 1:586 max_seqlen_in_batch_q = 1587 cu_seqlens_q = torch.arange(588 batch_size + 1, dtype=torch.int32, device=query_layer.device589 ) # There is a memcpy here, that is very bad.590 indices_q = cu_seqlens_q[:-1]591 query_layer = query_layer.squeeze(1)592 else:593 # The -q_len: slice assumes left padding.594 attention_mask = attention_mask[:, -query_length:]595 query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask)596 597 return (598 query_layer,599 key_layer,600 value_layer,601 indices_q,602 (cu_seqlens_q, cu_seqlens_k),603 (max_seqlen_in_batch_q, max_seqlen_in_batch_k),604 )605 606 607class LLaDASdpaAttention(LLaDAAttention):608 """609 LLaDA attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from610 `LLaDAAttention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to611 SDPA API.612 """613 614 # Adapted from LLaDAAttention.forward615 def forward(616 self,617 hidden_states: torch.Tensor,618 attention_mask: Optional[torch.Tensor] = None,619 position_ids: Optional[torch.LongTensor] = None,620 past_key_value: Optional[Cache] = None,621 output_attentions: bool = False,622 use_cache: bool = False,623 cache_position: Optional[torch.LongTensor] = None,624 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:625 if output_attentions:626 # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"` once this is implemented.627 logger.warning_once(628 "LLaDAModel is using LLaDASdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to the manual attention implementation, "629 '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.'630 )631 return super().forward(632 hidden_states=hidden_states,633 attention_mask=attention_mask,634 position_ids=position_ids,635 past_key_value=past_key_value,636 output_attentions=output_attentions,637 use_cache=use_cache,638 cache_position=cache_position,639 )640 641 bsz, q_len, _ = hidden_states.size()642 643 query_states = self.q_proj(hidden_states)644 key_states = self.k_proj(hidden_states)645 value_states = self.v_proj(hidden_states)646 647 query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)648 key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)649 value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)650 651 cos, sin = self.rotary_emb(value_states, position_ids)652 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)653 654 # In case static cache is used, it is an instance attribute.655 past_key_value = getattr(self, "past_key_value", past_key_value)656 657 if past_key_value is not None:658 # sin and cos are specific to RoPE models; cache_position needed for the static cache659 cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}660 key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)661 662 key_states = repeat_kv(key_states, self.num_key_value_groups)663 value_states = repeat_kv(value_states, self.num_key_value_groups)664 665 causal_mask = attention_mask666 # if attention_mask is not None and cache_position is not None:667 if attention_mask is not None:668 causal_mask = causal_mask[:, :, :, : key_states.shape[-2]]669 670 # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask,671 # Reference: https://github.com/pytorch/pytorch/issues/112577.672 if query_states.device.type == "cuda" and causal_mask is not None:673 query_states = query_states.contiguous()674 key_states = key_states.contiguous()675 value_states = value_states.contiguous()676 677 attn_output = torch.nn.functional.scaled_dot_product_attention(678 query_states,679 key_states,680 value_states,681 attn_mask=causal_mask,682 is_causal=False, # Modify: MDM683 dropout_p=self.attention_dropout if self.training else 0.0,684 )685 686 attn_output = attn_output.transpose(1, 2).contiguous()687 attn_output = attn_output.view(bsz, q_len, self.hidden_size)688 689 attn_output = self.o_proj(attn_output)690 691 return attn_output, None, past_key_value692 693 694LLaDA_ATTENTION_CLASSES = {695 "eager": LLaDAAttention,696 "flash_attention_2": LLaDAFlashAttention2,697 "sdpa": LLaDASdpaAttention,698}699 700 701class LLaDADecoderLayer(nn.Module):702 def __init__(self, config: LLaDAConfig, layer_idx: int):703 super().__init__()704 self.hidden_size = config.hidden_size705 706 self.self_attn = LLaDA_ATTENTION_CLASSES[config._attn_implementation](config=config, layer_idx=layer_idx)707 708 self.mlp = LLaDAMLP(config)709 self.input_layernorm = LLaDARMSNorm(config.hidden_size, eps=config.rms_norm_eps)710 self.post_attention_layernorm = LLaDARMSNorm(config.hidden_size, eps=config.rms_norm_eps)711 712 def forward(713 self,714 hidden_states: torch.Tensor,715 attention_mask: Optional[torch.Tensor] = None,716 position_ids: Optional[torch.LongTensor] = None,717 past_key_value: Optional[Tuple[torch.Tensor]] = None,718 output_attentions: Optional[bool] = False,719 use_cache: Optional[bool] = False,720 cache_position: Optional[torch.LongTensor] = None,721 **kwargs,722 ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:723 """724 Args:725 hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`726 attention_mask (`torch.FloatTensor`, *optional*):727 attention mask of size `(batch_size, sequence_length)` if flash attention is used or `(batch_size, 1,728 query_sequence_length, key_sequence_length)` if default attention is used.729 output_attentions (`bool`, *optional*):730 Whether or not to return the attentions tensors of all attention layers. See `attentions` under731 returned tensors for more detail.732 use_cache (`bool`, *optional*):733 If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding734 (see `past_key_values`).735 past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states736 """737 if "padding_mask" in kwargs:738 warnings.warn(739 "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"740 )741 742 residual = hidden_states743 744 hidden_states = self.input_layernorm(hidden_states)745 746 # Self Attention747 hidden_states, self_attn_weights, present_key_value = self.self_attn(748 hidden_states=hidden_states,749 attention_mask=attention_mask,750 position_ids=position_ids,751 past_key_value=past_key_value,752 output_attentions=output_attentions,753 use_cache=use_cache,754 cache_position=cache_position,755 **kwargs,756 )757 hidden_states = residual + hidden_states758 759 # Fully Connected760 residual = hidden_states761 hidden_states = self.post_attention_layernorm(hidden_states)762 hidden_states = self.mlp(hidden_states)763 hidden_states = residual + hidden_states764 765 outputs = (hidden_states,)766 767 if output_attentions:768 outputs += (self_attn_weights,)769 770 if use_cache:771 outputs += (present_key_value,)772 773 return outputs774 775 776LLaDA_START_DOCSTRING = r"""777 This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the778 library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads779 etc.)780 781 This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.782 Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage783 and behavior.784 785 Parameters:786 config ([`LLaDAConfig`]):787 Model configuration class with all the parameters of the model. Initializing with a config file does not788 load the weights associated with the model, only the configuration. Check out the789 [`~PreTrainedModel.from_pretrained`] method to load the model weights.790"""791 792 793@add_start_docstrings(794 "The bare LLaDA Model outputting raw hidden-states without any specific head on top.",795 LLaDA_START_DOCSTRING,796)797class LLaDAPreTrainedModel(PreTrainedModel):798 config_class = LLaDAConfig799 base_model_prefix = "model"800 supports_gradient_checkpointing = True801 _no_split_modules = ["LLaDADecoderLayer"]802 _skip_keys_device_placement = ["past_key_values"]803 _supports_flash_attn_2 = True804 _supports_sdpa = True805 _supports_cache_class = True806 807 def _init_weights(self, module):808 std = self.config.initializer_range809 if isinstance(module, nn.Linear):810 module.weight.data.normal_(mean=0.0, std=std)811 if module.bias is not None:812 module.bias.data.zero_()813 elif isinstance(module, nn.Embedding):814 module.weight.data.normal_(mean=0.0, std=std)815 if module.padding_idx is not None:816 module.weight.data[module.padding_idx].zero_()817 818 def _setup_cache(self, cache_cls, max_batch_size, max_cache_len: Optional[int] = None):819 if self.config._attn_implementation == "flash_attention_2" and cache_cls == StaticCache:820 raise ValueError(821 "`static` cache implementation is not compatible with `attn_implementation==flash_attention_2` "822 "make sure to use `sdpa` in the mean time, and open an issue at https://github.com/huggingface/transformers"823 )824 825 for layer in self.model.layers:826 device = layer.input_layernorm.weight.device827 if hasattr(self.config, "_pre_quantization_dtype"):828 dtype = self.config._pre_quantization_dtype829 else:830 dtype = layer.self_attn.o_proj.weight.dtype831 layer.self_attn.past_key_value = cache_cls(832 self.config, max_batch_size, max_cache_len, device=device, dtype=dtype833 )834 835 def _reset_cache(self):836 for layer in self.model.layers:837 layer.self_attn.past_key_value = None838 839 840LLaDA_INPUTS_DOCSTRING = r"""841 Args:842 input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):843 Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide844 it.845 846 Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and847 [`PreTrainedTokenizer.__call__`] for details.848 849 [What are input IDs?](../glossary#input-ids)850 attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):851 Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:852 853 - 1 for tokens that are **not masked**,854 - 0 for tokens that are **masked**.855 856 [What are attention masks?](../glossary#attention-mask)857 858 Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and859 [`PreTrainedTokenizer.__call__`] for details.860 861 If `past_key_values` is used, optionally only the last `input_ids` have to be input (see862 `past_key_values`).863 864 If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`]865 and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more866 information on the default strategy.867 868 - 1 indicates the head is **not masked**,869 - 0 indicates the head is **masked**.870 position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):871 Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,872 config.n_positions - 1]`.873 874 [What are position IDs?](../glossary#position-ids)875 past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*):876 Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention877 blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values`878 returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`.879 880 Two formats are allowed:881 - a [`~cache_utils.Cache`] instance;882 - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of883 shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy884 cache format.885 886 The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the887 legacy cache format will be returned.888 889 If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't890 have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids`891 of shape `(batch_size, sequence_length)`.892 inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):893 Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This894 is useful if you want more control over how to convert `input_ids` indices into associated vectors than the895 model's internal embedding lookup matrix.896 use_cache (`bool`, *optional*):897 If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see898 `past_key_values`).899 output_attentions (`bool`, *optional*):900 Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned901 tensors for more detail.902 output_hidden_states (`bool`, *optional*):903 Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for904 more detail.905 return_dict (`bool`, *optional*):906 Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.907 cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):908 Indices depicting the position of the input sequence tokens in the sequence. Contrarily to `position_ids`,909 this tensor is not affected by padding. It is used to update the cache in the correct position and to infer910 the complete sequence length.911"""912 913 914@add_start_docstrings(915 "The bare LLaDA Model outputting raw hidden-states without any specific head on top.",916 LLaDA_START_DOCSTRING,917)918class LLaDAModel(LLaDAPreTrainedModel):919 """920 Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`LLaDADecoderLayer`]921 922 Args:923 config: LLaDAConfig924 """925 926 def __init__(self, config: LLaDAConfig):927 super().__init__(config)928 self.padding_idx = config.pad_token_id929 self.vocab_size = config.vocab_size930 931 self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)932 self.layers = nn.ModuleList(933 [LLaDADecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]934 )935 self.norm = LLaDARMSNorm(config.hidden_size, eps=config.rms_norm_eps)936 self.gradient_checkpointing = False937 938 # Initialize weights and apply final processing939 self.post_init()940 941 def get_input_embeddings(self):942 return self.embed_tokens943 944 def set_input_embeddings(self, value):945 self.embed_tokens = value946 947 @add_start_docstrings_to_model_forward(LLaDA_INPUTS_DOCSTRING)948 def forward(949 self,950 input_ids: torch.LongTensor = None,951 attention_mask: Optional[torch.Tensor] = None,952 position_ids: Optional[torch.LongTensor] = None,953 past_key_values: Optional[List[torch.FloatTensor]] = None,954 inputs_embeds: Optional[torch.FloatTensor] = None,955 use_cache: Optional[bool] = None,956 output_attentions: Optional[bool] = None,957 output_hidden_states: Optional[bool] = None,958 return_dict: Optional[bool] = None,959 cache_position: Optional[torch.LongTensor] = None,960 **kwargs,961 ) -> Union[Tuple, BaseModelOutputWithPast]:962 # Add Basic MDM Model config check963 assert (past_key_values is None and not use_cache), "The kvcache is not suppotred for MDM."964 965 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions966 output_hidden_states = (967 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states968 )969 use_cache = use_cache if use_cache is not None else self.config.use_cache970 return_dict = return_dict if return_dict is not None else self.config.use_return_dict971 972 if (input_ids is None) ^ (inputs_embeds is not None):973 raise ValueError(974 "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one"975 )976 977 if self.gradient_checkpointing and self.training and use_cache:978 logger.warning_once(979 "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`."980 )981 use_cache = False982 983 if inputs_embeds is None:984 inputs_embeds = self.embed_tokens(input_ids)985 986 past_seen_tokens = 0987 if use_cache: # kept for BC (cache positions)988 if not isinstance(past_key_values, StaticCache):989 past_key_values = DynamicCache.from_legacy_cache(past_key_values)990 past_seen_tokens = past_key_values.get_seq_length()991 992 if cache_position is None:993 if isinstance(past_key_values, StaticCache):994 raise ValueError("cache_position is a required argument when using StaticCache.")995 cache_position = torch.arange(996 past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device997 )998 999 if position_ids is None:1000 position_ids = cache_position.unsqueeze(0)1001 1002 causal_mask = self._update_causal_mask(attention_mask, inputs_embeds, cache_position, is_causal=False) # Modify: MDM1003 1004 # embed positions1005 hidden_states = inputs_embeds1006 1007 # decoder layers1008 all_hidden_states = () if output_hidden_states else None1009 all_self_attns = () if output_attentions else None1010 next_decoder_cache = None1011 1012 for decoder_layer in self.layers:1013 if output_hidden_states:1014 all_hidden_states += (hidden_states,)1015 1016 if self.gradient_checkpointing and self.training:1017 layer_outputs = self._gradient_checkpointing_func(1018 decoder_layer.__call__,1019 hidden_states,1020 causal_mask,1021 position_ids,1022 past_key_values,1023 output_attentions,1024 use_cache,1025 cache_position,1026 **kwargs,1027 )1028 else:1029 layer_outputs = decoder_layer(1030 hidden_states,1031 attention_mask=causal_mask,1032 position_ids=position_ids,1033 past_key_value=past_key_values,1034 output_attentions=output_attentions,1035 use_cache=use_cache,1036 cache_position=cache_position,1037 **kwargs,1038 )1039 1040 hidden_states = layer_outputs[0]1041 1042 if use_cache:1043 next_decoder_cache = layer_outputs[2 if output_attentions else 1]1044 1045 if output_attentions:1046 all_self_attns += (layer_outputs[1],)1047 1048 hidden_states = self.norm(hidden_states)1049 1050 # add hidden states from the last decoder layer1051 if output_hidden_states:1052 all_hidden_states += (hidden_states,)1053 1054 next_cache = None1055 if use_cache:1056 next_cache = (1057 next_decoder_cache.to_legacy_cache() if isinstance(next_decoder_cache, Cache) else next_decoder_cache1058 )1059 if not return_dict:1060 return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)1061 return BaseModelOutputWithPast(1062 last_hidden_state=hidden_states,1063 past_key_values=next_cache,1064 hidden_states=all_hidden_states,1065 attentions=all_self_attns,1066 )1067 1068 # TODO: As of torch==2.2.0, the `attention_mask` passed to the model in `generate` is 2D and of dynamic length even when the static1069 # KV cache is used. This is an issue for torch.compile which then recaptures cudagraphs at each decode steps due to the dynamic shapes.1070 # (`recording cudagraph tree for symint key 13`, etc.), which is VERY slow. A workaround is `@torch.compiler.disable`, but this prevents using1071 # `fullgraph=True`. See more context in https://github.com/huggingface/transformers/pull/291141072 def _update_causal_mask(self, attention_mask, input_tensor, cache_position, is_causal=True):1073 if self.config._attn_implementation == "flash_attention_2":1074 if attention_mask is not None and 0.0 in attention_mask:1075 return attention_mask1076 return None1077 1078 dtype, device = input_tensor.dtype, input_tensor.device1079 min_dtype = torch.finfo(dtype).min1080 sequence_length = input_tensor.shape[1]1081 if hasattr(self.layers[0].self_attn, "past_key_value"): # static cache1082 target_length = self.config.max_position_embeddings1083 else: # dynamic cache1084 target_length = (1085 attention_mask.shape[-1] if isinstance(attention_mask, torch.Tensor) else cache_position[-1] + 11086 )1087 1088 causal_mask = torch.full((sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device)1089 if sequence_length != 1:1090 causal_mask = torch.triu(causal_mask, diagonal=1)1091 1092 if is_causal == False:1093 causal_mask = torch.zeros((sequence_length, target_length), dtype=dtype, device=device)1094 1095 causal_mask *= torch.arange(target_length, device=device) > cache_position.reshape(-1, 1)1096 causal_mask = causal_mask[None, None, :, :].expand(input_tensor.shape[0], 1, -1, -1)1097 if attention_mask is not None:1098 causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit1099 if attention_mask.dim() == 2:1100 # The position with 1 in attention_mask represents the place to be attended to, so here we need to mask the place where attention_mask is 01101 mask_length = attention_mask.shape[-1]1102 padding_mask = causal_mask[..., :mask_length].eq(0.0) * attention_mask[:, None, None, :].eq(0.0)1103 causal_mask[..., :mask_length] = causal_mask[..., :mask_length].masked_fill(padding_mask, min_dtype)1104 elif attention_mask.dim() == 4:1105 # The position with 1 in attention_mask represents the place to be attended to, so here we need to mask the place where attention_mask is 01106 # backwards compatibility: we allow passing a 4D attention mask shorter than the input length with1107 # cache. In that case, the 4D attention mask attends to the newest tokens only.1108 if attention_mask.shape[-2] < cache_position[0] + sequence_length:1109 offset = cache_position[0]1110 else:1111 offset = 01112 mask_shape = attention_mask.shape1113 mask_slice = (attention_mask.eq(0.0)).to(dtype=dtype) * min_dtype1114 causal_mask[1115 : mask_shape[0], : mask_shape[1], offset : mask_shape[2] + offset, : mask_shape[3]1116 ] = mask_slice1117 1118 if (1119 self.config._attn_implementation == "sdpa"1120 and attention_mask is not None1121 and attention_mask.device.type == "cuda"1122 ):1123 # TODO: For dynamo, rather use a check on fullgraph=True once this is possible (https://github.com/pytorch/pytorch/pull/120400).1124 is_tracing = (1125 torch.jit.is_tracing()1126 or isinstance(input_tensor, torch.fx.Proxy)1127 or (hasattr(torch, "_dynamo") and torch._dynamo.is_compiling())1128 )1129 if not is_tracing and torch.any(attention_mask != 1):1130 # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when1131 # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path.1132 # Details: https://github.com/pytorch/pytorch/issues/1102131133 causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype)1134 1135 return causal_mask1136 1137 1138class LLaDAModelLM(LLaDAPreTrainedModel):1139 _tied_weights_keys = ["lm_head.weight"]1140 1141 def __init__(self, config):1142 super().__init__(config)1143 self.model = LLaDAModel(config)1144 self.vocab_size = config.vocab_size1145 self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)1146 1147 # Initialize weights and apply final processing1148 self.post_init()1149 1150 def get_input_embeddings(self):1151 return self.model.embed_tokens1152 1153 def set_input_embeddings(self, value):1154 self.model.embed_tokens = value1155 1156 def get_output_embeddings(self):1157 return self.lm_head1158 1159 def set_output_embeddings(self, new_embeddings):1160 self.lm_head = new_embeddings1161 1162 def set_decoder(self, decoder):1163 self.model = decoder1164 1165 def get_decoder(self):1166 return self.model1167 1168 def _build_conversation_mask_optimized(self, conversation_ids):1169 # Reshape conversation_ids for broadcasting1170 ids_i = conversation_ids.unsqueeze(-1) # [batch_size, seq_len, 1]1171 ids_j = conversation_ids.unsqueeze(-2) # [batch_size, 1, seq_len]1172 1173 # Use broadcasting to compare all pairs of conversation IDs1174 conv_mask = (ids_j <= ids_i) # [batch_size, seq_len, seq_len]1175 1176 # Add the attention head dimension1177 return conv_mask.unsqueeze(1) # [batch_size, 1, seq_len, seq_len]1178 1179 @staticmethod1180 def add_gumbel_noise(logits, temperature):1181 '''1182 The Gumbel max is a method for sampling categorical distributions.1183 According to arXiv:2409.02908, for MDM, low-precision Gumbel Max improves perplexity score but reduces generation quality.1184 Thus, we use float64.1185 '''1186 if temperature == 0:1187 # When temperature=0, we can directly return the original logits. 1188 # without any noise or transformation1189 return logits1190 1191 # use float64 for more stable computation1192 logits = logits.to(torch.float64)1193 noise = torch.rand_like(logits, dtype=torch.float64)1194 gumbel_noise = (- torch.log(noise)) ** temperature1195 return logits.exp() / gumbel_noise1196 1197 @staticmethod1198 def get_num_transfer_tokens(mask_index, steps):1199 '''1200 Precompute the number of tokens to transition at each step.