from __future__ import annotations import math from typing import Any import torch import torch.nn.functional as F from torch import Tensor from transformers import LlamaForCausalLM from transformers.modeling_outputs import CausalLMOutputWithPast class CIDDiffusionForMaskedLM(LlamaForCausalLM): """Bidirectional masked-diffusion language model backed by Llama weights. The parameter layout is intentionally identical to LlamaForCausalLM so AR-to-diffusion checkpoints can be published without rewriting weights. """ def __init__(self, config) -> None: super().__init__(config) mask_token_id = getattr(config, "mask_token_id", None) if mask_token_id is None: raise ValueError("CID diffusion checkpoints require config.mask_token_id") self.mask_token_id = int(mask_token_id) self.config.use_cache = False self._set_bidirectional_attention() def _set_bidirectional_attention(self) -> None: for layer in self.model.layers: layer.self_attn.is_causal = False def _bidirectional_hidden_states( self, *, input_ids: Tensor | None = None, inputs_embeds: Tensor | None = None, attention_mask: Tensor | None = None, position_ids: Tensor | None = None, output_hidden_states: bool = False, ) -> tuple[Tensor, tuple[Tensor, ...] | None]: if (input_ids is None) == (inputs_embeds is None): raise ValueError("provide exactly one of input_ids or inputs_embeds") if inputs_embeds is None: inputs_embeds = self.model.embed_tokens(input_ids) batch_size, sequence_length = inputs_embeds.shape[:2] if position_ids is None: position_ids = torch.arange( sequence_length, device=inputs_embeds.device, dtype=torch.long, ).unsqueeze(0).expand(batch_size, -1) if attention_mask is None: attention_mask = torch.ones( (batch_size, sequence_length), device=inputs_embeds.device, dtype=torch.bool, ) elif attention_mask.ndim != 2: raise ValueError("attention_mask must have shape [batch, tokens]") self._set_bidirectional_attention() key_mask = None if not bool(attention_mask.bool().all()): minimum = torch.finfo(inputs_embeds.dtype).min valid_keys = attention_mask.to( device=inputs_embeds.device, dtype=torch.bool, ) key_mask = torch.zeros( (batch_size, 1, 1, sequence_length), device=inputs_embeds.device, dtype=inputs_embeds.dtype, ) key_mask = key_mask.masked_fill( ~valid_keys[:, None, None, :], minimum, ) hidden_states = inputs_embeds position_embeddings = self.model.rotary_emb(hidden_states, position_ids) all_hidden_states: list[Tensor] | None = [] if output_hidden_states else None if all_hidden_states is not None: all_hidden_states.append(hidden_states) for layer in self.model.layers[: self.config.num_hidden_layers]: layer_output = layer( hidden_states, attention_mask=key_mask, position_ids=position_ids, past_key_values=None, use_cache=False, cache_position=None, position_embeddings=position_embeddings, ) hidden_states = ( layer_output[0] if isinstance(layer_output, tuple) else layer_output ) if all_hidden_states is not None: all_hidden_states.append(hidden_states) hidden_states = self.model.norm(hidden_states) if all_hidden_states is not None: all_hidden_states[-1] = hidden_states return hidden_states, tuple(all_hidden_states) return hidden_states, None def forward( self, input_ids: Tensor | None = None, attention_mask: Tensor | None = None, position_ids: Tensor | None = None, labels: Tensor | None = None, inputs_embeds: Tensor | None = None, output_attentions: bool | None = None, output_hidden_states: bool | None = None, return_dict: bool | None = None, **_: Any, ) -> CausalLMOutputWithPast | tuple[Tensor, ...]: if output_attentions: raise ValueError( "CIDDiffusionForMaskedLM does not expose per-layer attention weights" ) output_hidden_states = bool(output_hidden_states) return_dict = self.config.use_return_dict if return_dict is None else return_dict hidden_states, all_hidden_states = self._bidirectional_hidden_states( input_ids=input_ids, inputs_embeds=inputs_embeds, attention_mask=attention_mask, position_ids=position_ids, output_hidden_states=output_hidden_states, ) logits = self.lm_head(hidden_states) loss = None if labels is not None: loss = F.cross_entropy( logits.float().reshape(-1, logits.shape[-1]), labels.reshape(-1), ignore_index=-100, ) if not return_dict: output = (logits, None, all_hidden_states, None) return ((loss,) + output) if loss is not None else output return CausalLMOutputWithPast( loss=loss, logits=logits, past_key_values=None, hidden_states=all_hidden_states, attentions=None, ) @staticmethod def _filter_top_p(logits: Tensor, top_p: float) -> Tensor: if top_p >= 1.0: return logits if not 0.0 < top_p <= 1.0: raise ValueError("top_p must be in (0, 1]") sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1) cumulative = torch.cumsum( torch.softmax(sorted_logits, dim=-1), dim=-1, ) remove = cumulative > top_p remove[..., 1:] = remove[..., :-1].clone() remove[..., 0] = False removal_mask = torch.zeros_like(remove).scatter( -1, sorted_indices, remove, ) return logits.masked_fill(removal_mask, -torch.inf) def _predict_tokens( self, logits: Tensor, *, temperature: float, top_p: float, generator: torch.Generator | None, ) -> tuple[Tensor, Tensor]: logits = logits.float() if temperature <= 0.0: probabilities = torch.softmax(logits, dim=-1) confidence, predicted = probabilities.max(dim=-1) return predicted, confidence filtered = self._filter_top_p(logits / temperature, top_p) probabilities = torch.softmax(filtered, dim=-1) flat = probabilities.reshape(-1, probabilities.shape[-1]) predicted = torch.multinomial( flat, num_samples=1, generator=generator, ).reshape(probabilities.shape[:-1]) confidence = probabilities.gather( -1, predicted.unsqueeze(-1), ).squeeze(-1) return predicted, confidence @torch.inference_mode() def denoise( self, input_ids: Tensor, *, editable_mask: Tensor | None = None, steps: int | None = None, temperature: float = 0.0, top_p: float = 1.0, generator: torch.Generator | None = None, ) -> Tensor: """Iteratively resolve MASK tokens already present in input_ids.""" if input_ids.ndim != 2: raise ValueError("input_ids must have shape [batch, tokens]") result = input_ids.clone() unresolved = result.eq(self.mask_token_id) if editable_mask is not None: if editable_mask.shape != result.shape: raise ValueError("editable_mask must match input_ids") unresolved &= editable_mask.bool() else: editable_mask = unresolved.clone() initial = unresolved.sum(dim=1) max_masks = int(initial.max().item()) if initial.numel() else 0 if max_masks == 0: return result if steps is None: steps = max_masks if steps <= 0: raise ValueError("steps must be positive") for step_index in range(steps): unresolved = result.eq(self.mask_token_id) & editable_mask if not bool(unresolved.any()): break logits = self(result).logits predicted, confidence = self._predict_tokens( logits, temperature=temperature, top_p=top_p, generator=generator, ) remaining_steps = steps - step_index for batch_index in range(result.shape[0]): positions = torch.nonzero( unresolved[batch_index], as_tuple=False, ).flatten() if positions.numel() == 0: continue reveal_count = math.ceil(positions.numel() / remaining_steps) ranked = positions[ confidence[batch_index, positions].argsort(descending=True) ] selected = ranked[:reveal_count] result[batch_index, selected] = predicted[batch_index, selected] unresolved = result.eq(self.mask_token_id) & editable_mask if bool(unresolved.any()): logits = self(result).logits result[unresolved] = logits.argmax(dim=-1)[unresolved] return result @torch.inference_mode() def diffusion_generate( self, input_ids: Tensor, *, max_new_tokens: int = 64, steps: int | None = None, block_length: int | None = 16, temperature: float = 0.0, top_p: float = 1.0, generator: torch.Generator | None = None, ) -> Tensor: """Append a masked canvas and fill it with block-wise diffusion.""" if input_ids.ndim != 2: raise ValueError("input_ids must have shape [batch, tokens]") if max_new_tokens <= 0: raise ValueError("max_new_tokens must be positive") if steps is None: steps = max_new_tokens if steps <= 0: raise ValueError("steps must be positive") if block_length is None: block_length = max_new_tokens if block_length <= 0: raise ValueError("block_length must be positive") batch_size, prompt_length = input_ids.shape result = torch.full( (batch_size, prompt_length + max_new_tokens), self.mask_token_id, dtype=input_ids.dtype, device=input_ids.device, ) result[:, :prompt_length] = input_ids block_count = math.ceil(max_new_tokens / block_length) base_steps = steps // block_count extra_steps = steps % block_count for block_index in range(block_count): start = prompt_length + block_index * block_length stop = min(start + block_length, result.shape[1]) block_steps = base_steps + (1 if block_index < extra_steps else 0) block_steps = max(1, block_steps) editable = torch.zeros_like(result, dtype=torch.bool) editable[:, start:stop] = True result = self.denoise( result, editable_mask=editable, steps=block_steps, temperature=temperature, top_p=top_p, generator=generator, ) return result def generate(self, input_ids: Tensor | None = None, **kwargs: Any) -> Tensor: """Diffusion-aware replacement for autoregressive GenerationMixin.generate.""" if input_ids is None: input_ids = kwargs.pop("inputs", None) if input_ids is None: raise ValueError("input_ids are required for diffusion generation") attention_mask = kwargs.pop("attention_mask", None) if attention_mask is not None and not bool(attention_mask.bool().all()): raise ValueError( "batched padded prompts are not supported by diffusion generate; " "generate each prompt separately" ) generation_config = kwargs.pop("generation_config", None) max_new_tokens = kwargs.pop("max_new_tokens", None) if max_new_tokens is None: max_length = kwargs.pop("max_length", None) if max_length is not None: max_new_tokens = int(max_length) - input_ids.shape[1] if max_new_tokens is None: max_new_tokens = 64 do_sample = kwargs.pop("do_sample", None) temperature = kwargs.pop("temperature", None) top_p = kwargs.pop("top_p", None) if generation_config is not None: if do_sample is None: do_sample = bool(getattr(generation_config, "do_sample", False)) if temperature is None: temperature = float(getattr(generation_config, "temperature", 1.0)) if top_p is None: top_p = float(getattr(generation_config, "top_p", 1.0)) do_sample = bool(do_sample) if do_sample is not None else False temperature = float(temperature) if temperature is not None else 1.0 top_p = float(top_p) if top_p is not None else 1.0 if not do_sample: temperature = 0.0 num_beams = int(kwargs.pop("num_beams", 1)) num_return_sequences = int(kwargs.pop("num_return_sequences", 1)) if num_beams != 1 or num_return_sequences != 1: raise ValueError( "beam search and multiple return sequences are not defined for this " "diffusion sampler" ) steps = int(kwargs.pop("diffusion_steps", kwargs.pop("steps", max_new_tokens))) block_length = kwargs.pop("block_length", 16) generator = kwargs.pop("generator", None) for ignored in ( "pad_token_id", "eos_token_id", "bos_token_id", "use_cache", "return_dict_in_generate", "output_scores", ): kwargs.pop(ignored, None) if kwargs: unsupported = ", ".join(sorted(kwargs)) raise ValueError(f"unsupported diffusion generation arguments: {unsupported}") return self.diffusion_generate( input_ids, max_new_tokens=int(max_new_tokens), steps=steps, block_length=None if block_length is None else int(block_length), temperature=temperature, top_p=top_p, generator=generator, )