""" Adapted from https://github.com/huggingface/transformers/blob/main/src/transformers/models/qwen3/modeling_qwen3.py """ from typing import Callable, Optional, Tuple, Union import torch from torch import nn from transformers import PreTrainedModel from transformers.activations import ACT2FN from transformers.utils import logging from .model_config import CoDAConfig from .attention import AttentionModule from .modeling_utils import ( HomogeneousSequential, RopeScaling, default_rope_frequencies, apply_rotary_pos_emb, transition, prefix_input_ids, truncate_input_ids, ) from .generation_utils import DLMGenerationMixin, DLMGenerationConfig logger = logging.get_logger(__name__) class CoDARMSNorm(nn.Module): def __init__(self, hidden_size, eps=1e-6): super().__init__() self.weight = nn.Parameter(torch.ones(hidden_size)) self.variance_epsilon = eps def forward(self, hidden_states): input_dtype = hidden_states.dtype hidden_states = hidden_states.to(torch.float32) variance = hidden_states.pow(2).mean(-1, keepdim=True) hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) return self.weight * hidden_states.to(input_dtype) def extra_repr(self): return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" class CoDAMLP(nn.Module): def __init__(self, config: CoDAConfig): super().__init__() self.config = config self.hidden_size = config.hidden_size self.intermediate_size = config.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): down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) return down_proj class CoDAAttention(nn.Module): """Multi-headed attention from 'Attention Is All You Need' paper""" def __init__(self, config: CoDAConfig, layer_idx: int | None = None): super().__init__() self.config = config self.attention_block = AttentionModule(config) self.layer_idx = layer_idx if layer_idx is None: logger.warning_once( f"Instantiating {self.__class__.__name__} without passing a `layer_idx` is not recommended and will " "lead to errors during the forward call if caching is used. Please make sure to provide a `layer_idx` " "when creating this class." ) self.hidden_size = config.hidden_size self.num_heads = config.num_attention_heads self.head_dim = getattr(config, "head_dim", self.hidden_size // self.num_heads) self.num_key_value_heads = config.num_key_value_heads self.num_key_value_groups = self.num_heads // self.num_key_value_heads self.scaling = self.head_dim**-0.5 self.attention_dropout = getattr(config, "attention_dropout", 0.0) # weiran: diffullama self.is_causal = False self.q_proj = nn.Linear( self.hidden_size, self.num_heads * self.head_dim, bias=getattr(config, "attention_bias", False), ) self.k_proj = nn.Linear( self.hidden_size, self.num_key_value_heads * self.head_dim, bias=getattr(config, "attention_bias", False), ) self.v_proj = nn.Linear( self.hidden_size, self.num_key_value_heads * self.head_dim, bias=getattr(config, "attention_bias", False), ) self.o_proj = nn.Linear( self.num_heads * self.head_dim, self.hidden_size, bias=getattr(config, "attention_bias", False), ) self.q_norm = CoDARMSNorm( self.head_dim, eps=getattr(config, "rms_norm_eps", 1e-6) ) self.k_norm = CoDARMSNorm( self.head_dim, eps=getattr(config, "rms_norm_eps", 1e-6) ) def forward( self, hidden_states: torch.Tensor, position_embeddings: Tuple[torch.Tensor, torch.Tensor], attention_mask: torch.Tensor | None = None, position_ids: torch.LongTensor | None = None, ) -> torch.FloatTensor: bsz, q_len, _ = hidden_states.size() query_states = self.q_proj(hidden_states) key_states = self.k_proj(hidden_states) value_states = self.v_proj(hidden_states) # Apply q_norm and k_norm to the head dimension query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim) key_states = key_states.view( bsz, q_len, self.num_key_value_heads, self.head_dim ) value_states = value_states.view( bsz, q_len, self.num_key_value_heads, self.head_dim ) # Apply normalization query_states = self.q_norm(query_states) key_states = self.k_norm(key_states) # Transpose to get the right shape for attention query_states = query_states.transpose(1, 2) key_states = key_states.transpose(1, 2) value_states = value_states.transpose(1, 2) cos, sin = position_embeddings query_states, key_states = apply_rotary_pos_emb( query_states, key_states, cos, sin ) attn_output = self.attention_block( query_states, key_states, value_states, attention_mask ) attn_output = attn_output.transpose(1, 2).contiguous() attn_output = attn_output.reshape(bsz, q_len, self.hidden_size) attn_output = self.o_proj(attn_output) return attn_output class CoDARotaryEmbedding(nn.Module): inv_freq: nn.Buffer def __init__( self, head_dim, rope_theta, scaling: RopeScaling | None = None, ): super().__init__() if scaling is None: inv_freq = default_rope_frequencies(head_dim, theta=rope_theta) else: raise NotImplementedError("Scaling is not implemented") self.register_buffer("inv_freq", inv_freq, persistent=False) @torch.no_grad() def forward(self, x, position_ids): # x: [bs, num_attention_heads, seq_len, head_size] inv_freq_expanded = ( self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1) ) position_ids_expanded = position_ids[:, None, :].float() # Force float32 since bfloat16 loses precision on long contexts # See https://github.com/huggingface/transformers/pull/29285 device_type = x.device.type device_type = ( device_type if isinstance(device_type, str) and device_type != "mps" else "cpu" ) with torch.autocast(device_type=device_type, enabled=False): freqs = ( inv_freq_expanded.float() @ position_ids_expanded.float() ).transpose(1, 2) emb = torch.cat((freqs, freqs), dim=-1) cos = emb.cos() sin = emb.sin() return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) class CoDADecoderLayer(nn.Module): def __init__(self, config: CoDAConfig, layer_idx: int): super().__init__() self.hidden_size = config.hidden_size self.layer_idx = layer_idx self.self_attn = CoDAAttention(config=config, layer_idx=layer_idx) self.mlp = CoDAMLP(config) self.input_layernorm = CoDARMSNorm( config.hidden_size, eps=getattr(config, "rms_norm_eps", 1e-6) ) self.post_attention_layernorm = CoDARMSNorm( config.hidden_size, eps=getattr(config, "rms_norm_eps", 1e-6) ) def forward( self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None, position_ids: torch.Tensor | None = None, position_embeddings: ( tuple[torch.Tensor, torch.Tensor] | None ) = None, # necessary, but kept here for BC ) -> torch.Tensor: """ 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. """ # This gives the `hidden_states` tensor a name so that we can layer specify # to offload this tensor to host RAM to save memory. This is not a standard # torch API because there is no such feature in PyTorch. Instead, the name # becomes node metadata during FX graph capture. residual = hidden_states hidden_states = self.input_layernorm(hidden_states) # Self Attention hidden_states = self.self_attn( hidden_states=hidden_states, attention_mask=attention_mask, position_ids=position_ids, position_embeddings=position_embeddings, ) hidden_states = residual + hidden_states # Fully Connected residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) hidden_states = self.mlp(hidden_states) hidden_states = residual + hidden_states return hidden_states class CoDAModel(PreTrainedModel): """ Transformer decoder consisting of *config.num_hidden_layers* layers. Args: config: FlexConfig """ config_class = CoDAConfig def __init__(self, config: CoDAConfig): super().__init__(config=config) self.vocab_size = config.vocab_size if "pad_token_id" not in config: self.padding_idx = None else: self.padding_idx = config.pad_token_id self.embed_tokens = nn.Embedding( config.vocab_size, config.hidden_size, padding_idx=self.padding_idx ) # `HomogeneousSequential` is similar to `nn.Sequential` but can be compiled with # `scan` described in https://pytorch.org/xla/release/r2.6/features/scan.html. self.layers = HomogeneousSequential( *[ CoDADecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers) ] ) self.norm = CoDARMSNorm( config.hidden_size, eps=getattr(config, "rms_norm_eps", 1e-6) ) rope_scaling = getattr(config, "rope_scaling", None) head_dim = getattr( config, "head_dim", config.hidden_size // config.num_attention_heads ) self.rope_theta = getattr(config, "rope_theta", 10000.0) if rope_scaling is not None: rope_scaling = RopeScaling(**rope_scaling) self.rotary_emb = CoDARotaryEmbedding( head_dim=head_dim, rope_theta=self.rope_theta, scaling=rope_scaling ) self.post_init() def _init_weights(self, module): std = getattr(self.config, "initializer_range", 0.02) if isinstance(module, nn.Linear): module.weight.data.normal_(mean=0.0, std=std) if module.bias is not None: module.bias.data.zero_() elif isinstance(module, nn.Embedding): module.weight.data.normal_(mean=0.0, std=std) if module.padding_idx is not None: module.weight.data[module.padding_idx].zero_() def forward( self, input_ids: torch.LongTensor, attention_mask: torch.FloatTensor | None = None, ) -> torch.Tensor: # convert input ids to embeddings inputs_embeds = self.embed_tokens(input_ids) seq_length = inputs_embeds.size(1) position_ids = ( torch.arange(seq_length, device=inputs_embeds.device).unsqueeze(0).float() ) # Create a causal attention mask causal_mask = torch.triu( torch.full( (seq_length, seq_length), float("-inf"), device=inputs_embeds.device ), diagonal=1, ) causal_mask = causal_mask.unsqueeze(0).unsqueeze( 0 ) # Add batch and head dimension if attention_mask is not None: causal_mask = causal_mask * attention_mask[:, None, None, :] 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 hidden_states = self.layers( hidden_states, attention_mask=causal_mask, position_ids=position_ids, position_embeddings=position_embeddings, ) hidden_states = self.norm(hidden_states) return hidden_states class CoDALanguageModel(DLMGenerationMixin, PreTrainedModel): config_class = CoDAConfig base_model_prefix = "model" is_parallelizable = True supports_gradient_checkpointing = False _no_split_modules = ["FlexDecoderLayer"] _skip_keys_device_placement = ["past_key_values"] _supports_flash_attn_2 = True _supports_sdpa = True _supports_cache_class = True def __init__(self, config: CoDAConfig): super().__init__(config) self.config = config self.model = CoDAModel(config) self.vocab_size = config.vocab_size self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) self.mask_token_id = config.mask_token_id self.generation_config = DLMGenerationConfig(mask_token_id=config.mask_token_id) self.apply(self._init_weights) def _init_weights(self, module): std = getattr(self.config, "initializer_range", 0.02) if isinstance(module, nn.Linear): module.weight.data.normal_(mean=0.0, std=std) if module.bias is not None: module.bias.data.zero_() elif isinstance(module, nn.Embedding): module.weight.data.normal_(mean=0.0, std=std) if module.padding_idx is not None: module.weight.data[module.padding_idx].zero_() def get_embeds(self, input_ids): """ Get input embeddings from the model. This method is used by the diffusion trainer to access embeddings. """ return self.model.embed_tokens(input_ids) def forward( self, input_ids: torch.LongTensor, labels: torch.LongTensor | None = None, attention_mask: torch.FloatTensor | None = None, src_mask: torch.BoolTensor | None = None, training_mode: str = "pretrain", **kwargs, ) -> tuple[torch.FloatTensor, torch.FloatTensor | None]: if not self.training: model_output = self.model( input_ids=input_ids, attention_mask=None ) hidden_states = model_output logits = self.lm_head(hidden_states) # NOTE: we shift logits at inference time return logits, None if training_mode == "sft" and src_mask is None: raise ValueError("SFT mode requires a non-null src_mask") epoch = kwargs.get("epoch", None) sampling_eps = getattr( self.config, "sampling_eps", 1e-3 ) # NOTE: use sampling_eps to control the noise level # If sampling_eps is a list, choose based on epoch if isinstance(sampling_eps, list): if epoch is None: # If epoch is not provided, use the first value sampling_eps = sampling_eps[0] else: # Use modulo to cycle through the list if epoch exceeds list length sampling_eps = sampling_eps[epoch % len(sampling_eps)] mask_token_id = self.mask_token_id loss_func = nn.CrossEntropyLoss(reduction="none") batch_size, seq_len = input_ids.shape # input_ids: [batch_size, seq_len] masking_schedule = kwargs.get("masking_schedule", None) # Create maskable_mask based on training mode and src_mask # For SFT: src_mask is provided, maskable_mask = ~src_mask # For pretrain: src_mask is None, maskable_mask = all True if src_mask is not None: maskable_mask = ~src_mask else: # pretrain or midtrain maskable_mask = torch.ones_like( input_ids, dtype=torch.bool, device=input_ids.device ) if masking_schedule is not None: prefix_probability = masking_schedule.get("prefix_probability", 0) truncate_probability = masking_schedule.get("truncate_probability", 0) else: prefix_probability = getattr(self.config, "prefix_probability", 0) truncate_probability = getattr(self.config, "truncate_probability", 0) if training_mode == "sft": prefix_probability = 0 truncate_probability = 0 # Generate random decisions for all batch items apply_prefix = ( torch.rand(batch_size, device=input_ids.device) < prefix_probability ) # Only apply truncation to rows that are NOT prefixed apply_truncate = ( torch.rand(batch_size, device=input_ids.device) < truncate_probability ) apply_truncate = apply_truncate & ~apply_prefix if prefix_probability > 0: maskable_mask = prefix_input_ids(input_ids, maskable_mask, apply_prefix) if truncate_probability > 0: input_ids = truncate_input_ids( input_ids, apply_truncate, self.config.pad_token_id ) maskable_mask = maskable_mask & (input_ids != self.config.pad_token_id) # add noise to input_ids sigma = (1 - sampling_eps) * torch.rand( input_ids.shape[0], device=input_ids.device ) + sampling_eps dsigma = torch.reciprocal(sigma) # Sample mask block size # Use mask_block_sizes from masking_probs if provided, otherwise fall back to config if masking_schedule is not None and "mask_block_sizes" in masking_schedule: mask_block_sizes = masking_schedule["mask_block_sizes"] else: mask_block_sizes = getattr(self.config, "mask_block_sizes", None) # Use masking_config if provided, otherwise fall back to config values if masking_schedule is not None: block_masking_probability = masking_schedule.get( "block_masking_probability", 0 ) else: block_masking_probability = getattr( self.config, "block_masking_probability", 0 ) if isinstance(block_masking_probability, list): if epoch is None: block_masking_probability = block_masking_probability[0] else: block_masking_probability = block_masking_probability[ epoch % len(block_masking_probability) ] if block_masking_probability > 0 and mask_block_sizes is not None and len(mask_block_sizes) > 0: mask_block_size = mask_block_sizes[ torch.randint(0, len(mask_block_sizes), (1,)).item() ] else: mask_block_size = 1 noisy_input_ids = transition( input_ids, sigma[:, None], maskable_mask=maskable_mask, mask_token_id=mask_token_id, mask_block_size=mask_block_size, ) loss_mask = noisy_input_ids == mask_token_id # Use gradient checkpointing if enabled if ( hasattr(self, "gradient_checkpointing") and self.gradient_checkpointing and self.training ): # Define a function for gradient checkpointing def custom_forward(*inputs): return self.model(*inputs) # Apply gradient checkpointing to the model forward pass hidden_states = self._gradient_checkpointing_func( custom_forward, noisy_input_ids, attention_mask ) else: hidden_states = self.model( input_ids=noisy_input_ids, attention_mask=attention_mask ) logits = self.lm_head(hidden_states) logits = logits.float() # logits: [bs, seq_len, vocab_size] # Shifted logits and labels # logits: [bs, seq_len-1, vocab_size] logits = logits[..., :-1, :].contiguous() # weiran: if the shifted token is not masked in the original input, the loss is 0 # loss_mask: [bs, seq_len-1] loss_mask = loss_mask[..., 1:].contiguous() target_ids = input_ids[..., 1:].contiguous() # loss: [bs, seq_len-1] loss = loss_func( logits.reshape(-1, logits.shape[-1]), target_ids.reshape(-1) ).reshape(target_ids.shape[0], -1) loss = loss.masked_fill(~loss_mask, 0) # weiran: divide by the number of tokens in the sequence instead of the number of masked tokens # justification is dsigma already accounts for the number of masked tokens # this is a hack to get something like per token loss # https://github.com/ML-GSAI/SMDM/blob/main/pretrain/train_mdm_rl.py#L281-L283 loss = (dsigma[:, None] * loss).sum() / ( input_ids.shape[0] * input_ids.shape[1] ) return logits, loss