#!/usr/bin/env python3 """ Dhara: Diffusion LLM with Canon Layers Combines: 1. Dhara's masked diffusion training (bidirectional attention, high throughput) 2. Canon layers (local context mixing via causal depthwise convolutions) Canon layers from "Physics of Language Models: Part 4.1" by Zeyuan Allen-Zhu: - Position A: After input LayerNorm, before attention - Position C: After post-attention LayerNorm, before MLP - kernel_size=4, residual=True, activation=False (default) Expected benefits: - ~280-290 tok/s throughput (Dhara's parallel generation) - +0.25-0.5% accuracy improvement (Canon's local context mixing) """ import math import warnings from typing import Optional, Tuple, Union, List import torch import torch.nn as nn import torch.nn.functional as F from torch.nn import CrossEntropyLoss from transformers import PreTrainedModel from transformers.generation import GenerationMixin from transformers.modeling_outputs import BaseModelOutputWithPast, MaskedLMOutput from transformers.utils import logging from transformers.cache_utils import Cache, DynamicCache from transformers import PretrainedConfig logger = logging.get_logger(__name__) # Optional performance imports try: from flash_attn import flash_attn_func FLASH_ATTN_AVAILABLE = True except ImportError: FLASH_ATTN_AVAILABLE = False try: import xformers.ops as xops XFORMERS_AVAILABLE = True except ImportError: XFORMERS_AVAILABLE = False class DharaConfig(PretrainedConfig): """ Configuration for Dhara model. Combines Dhara diffusion config with Canon layer parameters. """ model_type = "dhara" def __init__( self, # Core architecture vocab_size: int = 50304, hidden_size: int = 384, num_hidden_layers: int = 32, num_attention_heads: int = 8, num_key_value_heads: int = 4, intermediate_size: int = 1024, head_dim: int = None, max_position_embeddings: int = 2048, # Model specifics hidden_act: str = "silu", rms_norm_eps: float = 1e-6, rope_theta: float = 10000.0, initializer_range: float = 0.02, tie_word_embeddings: bool = True, attention_dropout: float = 0.0, # Canon layer parameters canon_set: str = "AC", # Positions: A (before attn), C (before MLP) canon_kernel: int = 4, # Kernel size (2-4) canon_residual: bool = True, # Highly recommended canon_activation: bool = False, # NOT recommended for transformers canon_bias: bool = False, # Diffusion specific mask_token_id: int = None, # Will be set from tokenizer mask_epsilon: float = 0.001, # Minimum mask probability num_diffusion_steps: int = 1000, # Special tokens bos_token_id: int = 1, eos_token_id: int = 2, pad_token_id: int = 0, # Performance flags use_cache: bool = False, use_flash_attention: bool = True, use_xformers: bool = False, **kwargs ): super().__init__( bos_token_id=bos_token_id, eos_token_id=eos_token_id, pad_token_id=pad_token_id, tie_word_embeddings=tie_word_embeddings, **kwargs ) # Core architecture self.vocab_size = vocab_size self.hidden_size = hidden_size self.num_hidden_layers = num_hidden_layers self.num_attention_heads = num_attention_heads self.num_key_value_heads = num_key_value_heads self.intermediate_size = intermediate_size self.head_dim = head_dim or (hidden_size // num_attention_heads) self.max_position_embeddings = max_position_embeddings # Model specifics self.hidden_act = hidden_act self.rms_norm_eps = rms_norm_eps self.rope_theta = rope_theta self.initializer_range = initializer_range self.tie_word_embeddings = tie_word_embeddings self.attention_dropout = attention_dropout # Canon parameters self.canon_set = canon_set self.canon_kernel = canon_kernel self.canon_residual = canon_residual self.canon_activation = canon_activation self.canon_bias = canon_bias # Diffusion specific self.mask_token_id = mask_token_id if mask_token_id is not None else (vocab_size - 1) self.mask_epsilon = mask_epsilon self.num_diffusion_steps = num_diffusion_steps # Special tokens self.bos_token_id = bos_token_id self.eos_token_id = eos_token_id self.pad_token_id = pad_token_id # Performance self.use_cache = use_cache self.use_flash_attention = use_flash_attention self.use_xformers = use_xformers class CanonLayer(nn.Module): """ Canon Layer: Causal 1D depthwise convolution for local context mixing. From "Physics of Language Models: Part 4.1" by Zeyuan Allen-Zhu. Captures local sequential dependencies with O(n) complexity. """ def __init__( self, hidden_size: int, kernel_size: int = 4, use_residual: bool = True, use_activation: bool = False, use_bias: bool = False, ): super().__init__() self.hidden_size = hidden_size self.kernel_size = kernel_size self.use_residual = use_residual self.use_activation = use_activation # Depthwise causal convolution self.conv = nn.Conv1d( in_channels=hidden_size, out_channels=hidden_size, kernel_size=kernel_size, padding=kernel_size - 1, # Causal (left-pad) groups=hidden_size, # Depthwise bias=use_bias, ) # Initialize for stability nn.init.normal_(self.conv.weight, mean=0.0, std=0.02) if use_bias: nn.init.zeros_(self.conv.bias) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: """ Args: hidden_states: [batch_size, seq_len, hidden_size] Returns: output: [batch_size, seq_len, hidden_size] """ batch_size, seq_len, hidden_size = hidden_states.shape # Transpose for Conv1d: [B, H, L] x = hidden_states.transpose(1, 2) # Apply conv with causal padding out = self.conv(x) # Remove right padding to make it causal out = out[:, :, :seq_len] # Optional activation if self.use_activation: out = F.silu(out) # Transpose back: [B, L, H] out = out.transpose(1, 2) # Residual connection if self.use_residual: out = hidden_states + out return out class RMSNorm(nn.Module): """Root Mean Square Layer Normalization""" 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) class RotaryEmbedding(nn.Module): """Rotary Position Embeddings (RoPE)""" def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None): super().__init__() self.dim = dim self.max_position_embeddings = max_position_embeddings self.base = base inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim)) self.register_buffer("inv_freq", inv_freq, persistent=False) self._set_cos_sin_cache( seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype() ) def _set_cos_sin_cache(self, seq_len, device, dtype): self.max_seq_len_cached = seq_len t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype) freqs = torch.outer(t, self.inv_freq) emb = torch.cat((freqs, freqs), dim=-1) self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False) self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False) def forward(self, x, seq_len=None): if seq_len > self.max_seq_len_cached: self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype) return ( self.cos_cached[:seq_len].to(dtype=x.dtype), self.sin_cached[:seq_len].to(dtype=x.dtype), ) 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) def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1): """Applies Rotary Position Embedding to query and key tensors.""" cos = cos[position_ids].unsqueeze(unsqueeze_dim) sin = sin[position_ids].unsqueeze(unsqueeze_dim) # Cast to input dtype for consistency cos = cos.to(q.dtype) sin = sin.to(q.dtype) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed class DharaMLP(nn.Module): """MLP with SwiGLU activation""" def __init__(self, config): 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 = nn.SiLU() def forward(self, x): return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: """Repeat KV heads for GQA.""" batch, num_key_value_heads, slen, head_dim = hidden_states.shape if n_rep == 1: return hidden_states hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) class DharaAttention(nn.Module): """Multi-Head Bidirectional Attention with GQA support (for diffusion)""" def __init__(self, config: DharaConfig, layer_idx: Optional[int] = None): super().__init__() self.config = config self.layer_idx = layer_idx 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 self.num_key_value_heads = config.num_key_value_heads self.num_key_value_groups = self.num_heads // self.num_key_value_heads self.max_position_embeddings = config.max_position_embeddings self.rope_theta = config.rope_theta self.is_causal = False # CRITICAL: Dhara uses bidirectional attention if (self.head_dim * self.num_heads) != self.hidden_size: raise ValueError( f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}" f" and `num_heads`: {self.num_heads})." ) self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False) self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False) self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False) self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False) self.rotary_emb = RotaryEmbedding( self.head_dim, max_position_embeddings=self.max_position_embeddings, base=self.rope_theta, ) def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_value=None, output_attentions: bool = False, use_cache: bool = False, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: 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) query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) kv_seq_len = key_states.shape[-2] if past_key_value is not None: if self.layer_idx is None: raise ValueError( f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} " "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class " "with a layer index." ) kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx) cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len) query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids) if past_key_value is not None: cache_kwargs = {"sin": sin, "cos": cos} key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) key_states = repeat_kv(key_states, self.num_key_value_groups) value_states = repeat_kv(value_states, self.num_key_value_groups) # Flash Attention for bidirectional if FLASH_ATTN_AVAILABLE and self.config.use_flash_attention: query_states = query_states.transpose(1, 2).contiguous() key_states = key_states.transpose(1, 2).contiguous() value_states = value_states.transpose(1, 2).contiguous() if query_states.dtype not in [torch.float16, torch.bfloat16]: query_states = query_states.to(torch.bfloat16) key_states = key_states.to(torch.bfloat16) value_states = value_states.to(torch.bfloat16) attn_output = flash_attn_func( query_states, key_states, value_states, dropout_p=0.0, causal=False, # Bidirectional for diffusion ) attn_output = attn_output.view(bsz, q_len, self.hidden_size) else: # Standard attention attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim) if attention_mask is not None: attn_weights = attn_weights + attention_mask attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training) attn_output = torch.matmul(attn_weights, value_states) 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) if not output_attentions: attn_weights = None return attn_output, attn_weights, past_key_value class DharaDecoderLayer(nn.Module): """ Dhara decoder layer with Canon layers at positions A and C. Flow: x -> LayerNorm -> [CanonA] -> Attention -> + residual x -> LayerNorm -> [CanonC] -> MLP -> + residual """ def __init__(self, config: DharaConfig, layer_idx: int): super().__init__() self.hidden_size = config.hidden_size self.config = config # Pre-attention norm self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) # Canon-A: before attention self.canon_a = None if "A" in config.canon_set: self.canon_a = CanonLayer( hidden_size=config.hidden_size, kernel_size=config.canon_kernel, use_residual=config.canon_residual, use_activation=config.canon_activation, use_bias=config.canon_bias, ) # Attention self.self_attn = DharaAttention(config=config, layer_idx=layer_idx) # Post-attention norm self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) # Canon-C: before MLP self.canon_c = None if "C" in config.canon_set: self.canon_c = CanonLayer( hidden_size=config.hidden_size, kernel_size=config.canon_kernel, use_residual=config.canon_residual, use_activation=config.canon_activation, use_bias=config.canon_bias, ) # MLP self.mlp = DharaMLP(config) def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_value=None, output_attentions: Optional[bool] = False, use_cache: Optional[bool] = False, ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: residual = hidden_states # Pre-attention layernorm hidden_states = self.input_layernorm(hidden_states) # Canon-A (before attention) if self.canon_a is not None: hidden_states = self.canon_a(hidden_states) # Self Attention (bidirectional) hidden_states, self_attn_weights, present_key_value = self.self_attn( hidden_states=hidden_states, attention_mask=attention_mask, position_ids=position_ids, past_key_value=past_key_value, output_attentions=output_attentions, use_cache=use_cache, ) hidden_states = residual + hidden_states # MLP block residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) # Canon-C (before MLP) if self.canon_c is not None: hidden_states = self.canon_c(hidden_states) hidden_states = self.mlp(hidden_states) hidden_states = residual + hidden_states outputs = (hidden_states,) if output_attentions: outputs += (self_attn_weights,) if use_cache: outputs += (present_key_value,) return outputs class DharaPreTrainedModel(PreTrainedModel): config_class = DharaConfig base_model_prefix = "model" supports_gradient_checkpointing = True _no_split_modules = ["DharaDecoderLayer"] _skip_keys_device_placement = "past_key_values" _supports_flash_attn_2 = True _supports_cache_class = True def _init_weights(self, module): std = self.config.initializer_range if isinstance(module, nn.Linear): module.weight.data.normal_(mean=0.0, std=std) if module.bias is not None: module.bias.data.zero_() elif isinstance(module, nn.Embedding): module.weight.data.normal_(mean=0.0, std=std) if module.padding_idx is not None: module.weight.data[module.padding_idx].zero_() class DharaModel(DharaPreTrainedModel): """ Dhara base model with bidirectional attention and Canon layers. """ def __init__(self, config: DharaConfig): super().__init__(config) self.padding_idx = config.pad_token_id self.vocab_size = config.vocab_size self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) self.layers = nn.ModuleList( [DharaDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] ) self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.gradient_checkpointing = False self.config = config self.mask_token_id = config.mask_token_id self._use_flash_attention_2 = config.use_flash_attention and FLASH_ATTN_AVAILABLE self.post_init() def get_input_embeddings(self): return self.embed_tokens def set_input_embeddings(self, value): self.embed_tokens = value def forward( self, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values=None, inputs_embeds: Optional[torch.FloatTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, ) -> Union[Tuple, BaseModelOutputWithPast]: 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 ) use_cache = use_cache if use_cache is not None else self.config.use_cache return_dict = return_dict if return_dict is not None else self.config.use_return_dict if input_ids is not None and inputs_embeds is not None: raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") elif input_ids is not None: batch_size, seq_length = input_ids.shape[:2] elif inputs_embeds is not None: batch_size, seq_length = inputs_embeds.shape[:2] else: raise ValueError("You have to specify either input_ids or inputs_embeds") if self.gradient_checkpointing and self.training: if use_cache: logger.warning_once( "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." ) use_cache = False past_key_values_length = 0 if use_cache: use_legacy_cache = not isinstance(past_key_values, Cache) if use_legacy_cache: past_key_values = DynamicCache.from_legacy_cache(past_key_values) past_key_values_length = past_key_values.get_usable_length(seq_length) if position_ids is None: device = input_ids.device if input_ids is not None else inputs_embeds.device position_ids = torch.arange( past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device ) position_ids = position_ids.unsqueeze(0) if inputs_embeds is None: inputs_embeds = self.embed_tokens(input_ids) if self._use_flash_attention_2: attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None else: # Bidirectional attention mask (not causal) if attention_mask is not None: if attention_mask.dim() == 2: batch_size, seq_length = attention_mask.shape attention_mask_4d = attention_mask[:, None, None, :].expand( batch_size, 1, seq_length, seq_length ).to(dtype=inputs_embeds.dtype) attention_mask = torch.where( attention_mask_4d == 0, torch.tensor(float('-inf'), dtype=inputs_embeds.dtype, device=attention_mask_4d.device), torch.tensor(0.0, dtype=inputs_embeds.dtype, device=attention_mask_4d.device) ) else: attention_mask = attention_mask else: attention_mask = None hidden_states = inputs_embeds all_hidden_states = () if output_hidden_states else None all_self_attns = () if output_attentions 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, use_cache, ) 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, use_cache=use_cache, ) 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],) hidden_states = self.norm(hidden_states) if output_hidden_states: all_hidden_states += (hidden_states,) next_cache = None if use_cache: next_cache = next_decoder_cache.to_legacy_cache() if use_legacy_cache else next_decoder_cache if not return_dict: return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None) return BaseModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=next_cache, hidden_states=all_hidden_states, attentions=all_self_attns, ) def add_noise_to_tokens(self, input_ids: torch.LongTensor, t: torch.FloatTensor, eps: float = None): """ MDM-style masking: Replace tokens with [MASK] based on noise level t. Args: input_ids: Input token IDs [batch_size, seq_len] t: Noise level in [0, 1] [batch_size] eps: Minimum mask probability Returns: Tuple of (noisy_input_ids, corruption_mask, p_mask) """ batch_size, seq_len = input_ids.shape device = input_ids.device if eps is None: eps = getattr(self.config, 'mask_epsilon', 0.001) p_mask = (1 - eps) * t + eps p_mask = p_mask.unsqueeze(-1).expand(batch_size, seq_len) corruption_mask = torch.rand(batch_size, seq_len, device=device) < p_mask mask_token_id = self.mask_token_id noisy_input_ids = torch.where(corruption_mask, mask_token_id, input_ids) return noisy_input_ids, corruption_mask, p_mask class DharaForMaskedDiffusion(DharaPreTrainedModel, GenerationMixin): """Dhara Model with Masked Diffusion head for training""" _tied_weights_keys = ["lm_head.weight"] def __init__(self, config): super().__init__(config) self.model = DharaModel(config) self.vocab_size = config.vocab_size self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) self.config = config self.mask_token_id = config.mask_token_id self.post_init() def get_input_embeddings(self): return self.model.embed_tokens def set_input_embeddings(self, value): self.model.embed_tokens = 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 def forward( self, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values=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, return_dict: Optional[bool] = None, corruption_mask: Optional[torch.BoolTensor] = None, p_mask: Optional[torch.Tensor] = None, ) -> Union[Tuple, MaskedLMOutput]: 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 ) return_dict = return_dict if return_dict is not None else self.config.use_return_dict 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, return_dict=return_dict, ) hidden_states = outputs[0] if self.config.tie_word_embeddings: logits = F.linear(hidden_states, self.model.embed_tokens.weight) else: logits = self.lm_head(hidden_states) logits = logits.float() loss = None if labels is not None: loss = self.compute_diffusion_loss(logits, labels, corruption_mask, p_mask) if not return_dict: output = (logits,) + outputs[1:] return (loss,) + output if loss is not None else output return MaskedLMOutput( loss=loss, logits=logits, hidden_states=outputs.hidden_states, attentions=outputs.attentions, ) def compute_diffusion_loss(self, logits, labels, corruption_mask=None, p_mask=None): """ MDM loss with p_mask importance weighting. """ if corruption_mask is None or p_mask is None: raise ValueError( "MDM requires both corruption_mask and p_mask for loss computation." ) loss = F.cross_entropy( logits.view(-1, self.config.vocab_size), labels.view(-1), reduction='none' ) loss = loss.view(labels.shape) masked_losses = loss[corruption_mask] masked_p_mask = p_mask[corruption_mask] weighted_losses = masked_losses / masked_p_mask total_positions = labels.shape[0] * labels.shape[1] return weighted_losses.sum() / total_positions def add_noise_to_tokens(self, input_ids: torch.LongTensor, t: torch.FloatTensor, eps: float = None): """Delegate to the base model""" return self.model.add_noise_to_tokens(input_ids, t, eps) def prepare_inputs_for_generation( self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=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() else: cache_length = past_length = past_key_values[0][0].shape[2] max_cache_length = None 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) :] elif past_length < input_ids.shape[1]: input_ids = input_ids[:, past_length:] 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: 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 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 @torch.no_grad() def generate( self, input_ids: Optional[torch.LongTensor] = None, max_length: Optional[int] = None, max_new_tokens: Optional[int] = None, num_diffusion_steps: int = 10, temperature: float = 1.0, top_p: float = 0.9, top_k: int = 50, do_sample: bool = True, pad_token_id: Optional[int] = None, eos_token_id: Optional[int] = None, repetition_penalty: float = 1.2, **kwargs ) -> torch.LongTensor: """ Generate text using autoregressive sampling with the diffusion model. Since this model was converted from AR to diffusion via WSD training, we generate tokens one at a time left-to-right, using the model's next-token predictions at each position. Args: input_ids: Input prompt token IDs [batch_size, prompt_len] max_length: Maximum total sequence length (prompt + generation) max_new_tokens: Number of new tokens to generate (alternative to max_length) num_diffusion_steps: Number of refinement iterations per token (higher = better quality) temperature: Sampling temperature (higher = more random) top_p: Nucleus sampling threshold top_k: Top-k sampling threshold do_sample: Whether to sample or take argmax pad_token_id: Token ID for padding eos_token_id: Token ID for end of sequence repetition_penalty: Penalty for repeating tokens (>1 = less repetition) Returns: Generated token IDs including the prompt """ # Handle device and dtype device = input_ids.device if input_ids is not None else next(self.parameters()).device # Determine generation length if input_ids is not None: batch_size, prompt_len = input_ids.shape else: batch_size, prompt_len = 1, 0 input_ids = torch.empty(batch_size, 0, dtype=torch.long, device=device) if max_new_tokens is not None: gen_len = max_new_tokens elif max_length is not None: gen_len = max_length - prompt_len else: gen_len = 50 # Default generation length if gen_len <= 0: return input_ids # Get special token IDs mask_token_id = self.config.mask_token_id if pad_token_id is None: pad_token_id = self.config.pad_token_id if hasattr(self.config, 'pad_token_id') else 0 if eos_token_id is None: eos_token_id = self.config.eos_token_id if hasattr(self.config, 'eos_token_id') else 2 # Start with the prompt generated = input_ids.clone() # Track generated tokens for repetition penalty generated_set = set() for i in range(prompt_len): for b in range(batch_size): generated_set.add(input_ids[b, i].item()) # Generate tokens one at a time (autoregressive style) for pos in range(gen_len): # Add a mask token at the next position current_seq = torch.cat([ generated, torch.full((batch_size, 1), mask_token_id, dtype=torch.long, device=device) ], dim=1) # Get model predictions outputs = self(input_ids=current_seq) logits = outputs.logits # [batch, seq_len, vocab] # Get logits for the last (masked) position next_token_logits = logits[:, -1, :] # [batch, vocab] # Apply repetition penalty if repetition_penalty != 1.0: for b in range(batch_size): for prev_token in generated_set: if prev_token < next_token_logits.shape[1]: next_token_logits[b, prev_token] /= repetition_penalty # Apply temperature if temperature != 1.0 and temperature > 0: next_token_logits = next_token_logits / temperature if do_sample and temperature > 0: # Apply top-k filtering if top_k > 0: indices_to_remove = next_token_logits < torch.topk(next_token_logits, top_k)[0][..., -1, None] next_token_logits[indices_to_remove] = float('-inf') # Apply top-p (nucleus) filtering if top_p < 1.0: sorted_logits, sorted_indices = torch.sort(next_token_logits, descending=True) cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1) # Remove tokens with cumulative probability above threshold sorted_indices_to_remove = cumulative_probs > top_p # Shift the indices to the right to keep the first token above threshold sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone() sorted_indices_to_remove[..., 0] = False # Scatter sorted indices to original indexing indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove) next_token_logits[indices_to_remove] = float('-inf') # Sample from the filtered distribution probs = F.softmax(next_token_logits, dim=-1) next_tokens = torch.multinomial(probs, num_samples=1).squeeze(-1) else: # Greedy decoding next_tokens = next_token_logits.argmax(dim=-1) # Add to generated sequence generated = torch.cat([generated, next_tokens.unsqueeze(-1)], dim=1) # Update generated set for repetition penalty for b in range(batch_size): generated_set.add(next_tokens[b].item()) # Check for EOS if eos_token_id is not None and (next_tokens == eos_token_id).all(): break return generated def save_pretrained(self, save_directory, **kwargs): """Override to save in SafeTensors format by default""" kwargs['safe_serialization'] = kwargs.get('safe_serialization', True) return super().save_pretrained(save_directory, **kwargs) def count_parameters(model): """Count total and Canon-specific parameters.""" total = sum(p.numel() for p in model.parameters()) canon = sum(p.numel() for n, p in model.named_parameters() if 'canon' in n.lower()) return total, canon if __name__ == "__main__": # Quick test print("Testing Dhara model creation...") config = DharaConfig( vocab_size=50304, hidden_size=384, num_hidden_layers=32, num_attention_heads=8, num_key_value_heads=4, intermediate_size=1024, canon_set="AC", canon_kernel=4, canon_residual=True, ) model = DharaForMaskedDiffusion(config) total, canon = count_parameters(model) print(f"Model created successfully!") print(f"Total params: {total:,} ({total/1e6:.2f}M)") print(f"Canon params: {canon:,} ({100*canon/total:.3f}%)") print(f"Base Dhara would be: {total - canon:,}") # Test forward pass batch_size, seq_len = 2, 64 input_ids = torch.randint(0, 50304, (batch_size, seq_len)) # Test with diffusion noise t = torch.rand(batch_size) noisy_ids, corruption_mask, p_mask = model.add_noise_to_tokens(input_ids, t) with torch.no_grad(): outputs = model( input_ids=noisy_ids, labels=input_ids, corruption_mask=corruption_mask, p_mask=p_mask, ) print(f"Forward pass: loss={outputs.loss.item():.4f}") print("Ready for training!")