# coding=utf-8 # Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Iterative denoising sampler for the uniform-only GIDD diffusion language model. For uniform-only GIDD the noise distribution is constant ``π = 1/V`` and the forward marginal is:: q(z_t = v | x) = α_t · 1[v = x] + (1 - α_t) / V Given a model that predicts ``x̂(z_t)``, ``z_s`` (with ``s < t`` in the denoising direction) is sampled from the analytic ancestral posterior:: p(z_s = v | z_t, x̂) ∝ q(z_s = v | x̂) · q(z_t | z_s = v) This is the Hugging Face port of ``tools/run_uniform_diffusion_sampling.py``; the same two samplers (``"ancestral"`` and ``"adaptive"``) are exposed through :meth:`SumiGenerationMixin.generate`. """ import math from dataclasses import dataclass from typing import Callable, List, Optional, Tuple import torch import torch.nn.functional as F from transformers.cache_utils import Cache from transformers.generation import GenerationConfig, GenerationMixin from transformers.utils import ModelOutput @dataclass class SumiGenerationOutput(ModelOutput): """Output class for uniform diffusion generation.""" sequences: torch.LongTensor past_key_values: Cache | None = None logits: None = None hidden_states: None = None canvas: torch.LongTensor | None = None # full untrimmed denoised canvas class SumiGenerationConfig(GenerationConfig): """Generation configuration for uniform diffusion language models.""" def __init__( self, num_denoising_steps: int = 128, sampler: str = "ancestral", schedule: str = "linear", min_log_snr: float = -9.0, max_log_snr: float = 9.0, tokens_per_step: int = 1, **kwargs, ): super().__init__(**kwargs) if num_denoising_steps < 1: raise ValueError("num_denoising_steps must be a positive integer.") self.num_denoising_steps = num_denoising_steps self.sampler = sampler self.schedule = schedule self.min_log_snr = min_log_snr self.max_log_snr = max_log_snr self.tokens_per_step = tokens_per_step @classmethod def from_model_config(cls, model_config): # transformers>=5.12's GenerationConfig.from_model_config compares the generation # config against a base GenerationConfig() that lacks this subclass's custom fields # (num_denoising_steps, sampler, ...), raising AttributeError. Build via the base # class, then re-wrap so our extra defaults are preserved. Robust on 5.8-5.12+. base = GenerationConfig.from_model_config(model_config) return cls(**base.to_dict()) def _make_log_snr_schedule( num_steps: int, schedule: str, min_log_snr: float, max_log_snr: float, device: torch.device, ) -> torch.Tensor: """``[num_steps + 1]`` log-SNR values, ascending (noisy → clean).""" if schedule == "linear": alpha_ts = torch.linspace(1e-4, 1.0 - 1e-4, num_steps + 1, device=device) elif schedule == "cosine": ts = torch.linspace(1.0 - 1e-3, 1e-3, num_steps + 1, device=device) alpha_ts = 0.5 + 0.5 * torch.cos(ts * math.pi) else: raise ValueError(f"Unknown schedule: {schedule!r}") alpha_ts = alpha_ts.clamp(min=1e-6, max=1.0 - 1e-6) log_snrs = torch.log(alpha_ts) - torch.log1p(-alpha_ts) return log_snrs.clamp(min=min_log_snr, max=max_log_snr) @torch.no_grad() def _compute_logits( model, z: torch.Tensor, attention_mask: torch.Tensor, vocab_size: int, ) -> torch.Tensor: """Forward pass → raw ``[B, S, vocab_size]`` logits truncated to the real vocab. Entries beyond ``vocab_size`` are dropped (matching how training masks padded vocab); returned as fp32. Shared by both the ancestral path (which softmaxes with temperature, via :func:`_compute_x_hat`) and the adaptive path (which needs the raw logits to separate the un-tempered selection distribution from the tempered commit distribution). """ logits = model( input_ids=z, attention_mask=attention_mask, use_cache=False, ).logits return logits[..., :vocab_size].float() @torch.no_grad() def _compute_x_hat( model, z: torch.Tensor, attention_mask: torch.Tensor, vocab_size: int, temperature: float = 1.0, ) -> torch.Tensor: """Forward pass → ``softmax(logits / temperature)`` truncated to the real vocab. ``temperature`` sharpens (``< 1``) or flattens (``> 1``) the model's predicted clean-token distribution before it feeds the analytic ancestral posterior. ``1.0`` reproduces full posterior sampling; lower values make denoising more decisive. The ancestral step that samples ``z_s`` from the posterior keeps its own inherent stochasticity either way. """ logits = _compute_logits(model, z, attention_mask, vocab_size) if temperature != 1.0: logits = logits / max(float(temperature), 1e-6) return F.softmax(logits, dim=-1) def _ancestral_step( z_t: torch.Tensor, x_hat: torch.Tensor, log_snr_t: float, log_snr_s: float, vocab_size: int, generator: Optional[torch.Generator], eps: float = 1e-12, ) -> torch.Tensor: """One ancestral denoising step ``z_t → z_s`` (with ``s < t``). Shapes: ``z_t [B, S]``, ``x_hat [B, S, V]``. Returns ``[B, S]``. """ device = z_t.device dtype = x_hat.dtype alpha_t = torch.sigmoid(torch.tensor(log_snr_t, device=device, dtype=dtype)) alpha_s = torch.sigmoid(torch.tensor(log_snr_s, device=device, dtype=dtype)) alpha_t_s = alpha_t / alpha_s.clamp(min=eps) beta_t = 1.0 - alpha_t beta_s = 1.0 - alpha_s beta_t_s = (1.0 - alpha_t_s).clamp(min=0.0) inv_v = 1.0 / vocab_size u_t = beta_t * inv_v u_s = beta_s * inv_v u_t_s = beta_t_s * inv_v q_s = alpha_s * x_hat + u_s one_hot_zt = F.one_hot(z_t, num_classes=vocab_size).to(dtype) q_t_given_s = alpha_t_s * one_hot_zt + u_t_s x_hat_at_zt = x_hat.gather(-1, z_t.unsqueeze(-1)).squeeze(-1) q_t_at_zt = (alpha_t * x_hat_at_zt + u_t).clamp(min=eps) posterior = q_s * q_t_given_s / q_t_at_zt.unsqueeze(-1) posterior = posterior.clamp(min=0.0) flat = posterior.reshape(-1, vocab_size) sampled = torch.multinomial(flat, num_samples=1, generator=generator).squeeze(-1) return sampled.view_as(z_t) def _adaptive_step( z_t: torch.Tensor, logits: torch.Tensor, noise_mask: torch.Tensor, tokens_per_step: int, temperature: float, generator: Optional[torch.Generator], ) -> Tuple[torch.Tensor, torch.Tensor]: """One confidence-based ("adaptive") denoising step ``z_t → z_s``. Implements the generalized confidence heuristic of von Rütte et al. (ICLR 2026, "Scaling Behavior of Discrete Diffusion Language Models", App. A.1, Eq. 9): conf(z_t) = p_prior(z_t) · ( max_z' p_θ(x=z'|z_t) − p_θ(x=z_t|z_t) ) For the UNIFORM-only model the prior is constant ``π = 1/V`` over the whole vocab (there is no mask token), so it drops out of the ranking — the confidence reduces to ``p_max − p_curr`` ("how much the model wants to change this token"). Unlike :func:`_ancestral_step`, which resamples EVERY denoise position each step, this commits ONLY the top ``tokens_per_step`` highest-confidence denoise positions to their predicted token (argmax when ``temperature == 0``, else sampled from ``softmax(logits/temperature)``) and leaves all other positions unchanged. The SNR schedule is unused, so more steps than tokens are allowed — a position may be revisited and overwritten on a later step. Shapes: ``z_t [B, S]``, ``logits [B, S, V]`` (un-tempered, fp32), ``noise_mask [B, S]`` (True = denoise here, False = frozen prompt/anchor). Returns ``(z_s [B, S], selected_pos [B, k])``. Frozen positions are guaranteed unchanged. """ vocab_size = logits.shape[-1] # Selection distribution: un-tempered softmax (the confidence compares the model's # raw belief in the best token vs. the current token). x_hat = F.softmax(logits, dim=-1) p_max = x_hat.max(dim=-1).values # [B, S] p_curr = x_hat.gather(-1, z_t.unsqueeze(-1)).squeeze(-1) # [B, S] conf = p_max - p_curr # >= 0 (uniform prior const) # Never pick a frozen position: force its confidence below any real one. conf = conf.masked_fill(~noise_mask, float("-inf")) k = max(1, min(int(tokens_per_step), conf.shape[-1])) next_pos = torch.topk(conf, k, dim=-1).indices # [B, k] # Commit distribution: tempered (temperature == 0 -> greedy argmax). if temperature and float(temperature) > 0.0: probs = F.softmax(logits / max(float(temperature), 1e-6), dim=-1) pred = torch.multinomial( probs.reshape(-1, vocab_size), num_samples=1, generator=generator ).view_as(z_t) else: pred = logits.argmax(dim=-1) z_s = z_t.clone() batch_idx = torch.arange(z_t.shape[0], device=z_t.device).unsqueeze(-1) z_s[batch_idx, next_pos] = pred[batch_idx, next_pos] # Restore every frozen position (prompt + anchors) unconditionally. return torch.where(noise_mask, z_s, z_t), next_pos def _greedy_step( z_t: torch.Tensor, logits: torch.Tensor, noise_mask: torch.Tensor, ) -> torch.Tensor: """One greedy denoising step ``z_t → z_s`` (the original naive Sumi sampler). Every denoise position is overwritten each step with the model's argmax prediction; there is no SNR schedule, no temperature and no stochasticity. Frozen positions (prompt + anchors) are restored via ``noise_mask``. Shapes: ``z_t [B, S]``, ``logits [B, S, V]``. Returns ``[B, S]``. """ pred = logits.argmax(dim=-1) return torch.where(noise_mask, pred, z_t) def _build_noise_mask( generated_ids: torch.Tensor, prompt_length: int, frozen: Optional[List] = None, denoise_end: Optional[List[int]] = None, ) -> torch.Tensor: """``[B, S]`` mask (``True`` = denoise here, ``False`` = frozen prompt/anchor). Prompt positions ``[:prompt_length]`` are frozen. ``frozen`` pins ``(position, token_id)`` pairs at fixed canvas positions and writes the token id into ``generated_ids`` in place (a flat ``[(pos, tid), ...]`` shared by every row, or a PER-ROW ``[[(pos, tid), ...], ...]`` of length ``batch_size``). ``denoise_end`` (per-row exclusive upper bound) freezes the tail at its prior-random init so the step budget concentrates on the content window. """ batch_size, total_length = generated_ids.shape device = generated_ids.device noise_mask = torch.zeros((batch_size, total_length), dtype=torch.bool, device=device) noise_mask[:, prompt_length:] = True if frozen: per_row = len(frozen) > 0 and isinstance(frozen[0], list) if per_row and len(frozen) != batch_size: raise ValueError( f"per-row frozen must have batch_size ({batch_size}) entries, got {len(frozen)}" ) for i in range(batch_size): row_frozen = frozen[i] if per_row else frozen for pos, tid in row_frozen: if 0 <= pos < total_length and bool(noise_mask[i, pos]): # not frozen by a prompt generated_ids[i, pos] = tid noise_mask[i, pos] = False if denoise_end is not None: if len(denoise_end) != batch_size: raise ValueError( f"denoise_end must have batch_size ({batch_size}) entries, got {len(denoise_end)}" ) for i in range(batch_size): end = int(denoise_end[i]) if 0 <= end < total_length: noise_mask[i, end:] = False # tail stays at its prior-random init return noise_mask class SumiGenerationMixin(GenerationMixin): """Generation mixin for uniform diffusion language models.""" @staticmethod def _resolve(value, name, default, generation_config, self_generation_config): """Explicit kwarg → passed generation_config → ``self.generation_config`` → default.""" if value is not None: return value for config in (generation_config, self_generation_config): if config is not None: resolved = getattr(config, name, None) if resolved is not None: return resolved return default def _trim_at_eos(self, generated_ids, prompt_length): """Cut each row at the first EOS in its generated region (dropping the EOS, the anchored BOS delimiter, and the denoised tail), then right-pad the batch back into a rectangular tensor. Prompt tokens are kept. Mirrors the eval harness's ``_extract_text`` (`gen_ids[:gen_ids.index(eos)]`). The pad filler is a special token, so ``decode(..., skip_special_tokens=True)`` yields clean text.""" eos_id = self.config.eos_token_id if eos_id is None: return generated_ids pad_id = self.config.pad_token_id pad_id = eos_id if pad_id is None else pad_id batch_size = generated_ids.shape[0] rows, max_len = [], prompt_length for i in range(batch_size): row = generated_ids[i] hit = (row[prompt_length:] == eos_id).nonzero(as_tuple=True)[0] cut = prompt_length + int(hit[0].item()) if hit.numel() > 0 else row.shape[0] rows.append(row[:cut]) max_len = max(max_len, cut) out = torch.full((batch_size, max(max_len, 1)), pad_id, dtype=generated_ids.dtype, device=generated_ids.device) for i, row in enumerate(rows): if row.shape[0] > 0: out[i, : row.shape[0]] = row return out @torch.no_grad() def generate( self, input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, max_new_tokens: Optional[int] = None, max_length: Optional[int] = None, num_denoising_steps: Optional[int] = None, seed: Optional[int] = None, generator: Optional[torch.Generator] = None, sampler: Optional[str] = None, schedule: Optional[str] = None, min_log_snr: Optional[float] = None, max_log_snr: Optional[float] = None, temperature: Optional[float] = None, tokens_per_step: Optional[int] = None, frozen: Optional[List] = None, denoise_end: Optional[List[int]] = None, canvas_length: Optional[int] = None, anchor_eosbos: Optional[bool] = None, trim_at_eos: Optional[bool] = None, progress_callback: Optional[Callable] = None, **kwargs, ) -> SumiGenerationOutput: generation_config = kwargs.pop("generation_config", None) inputs = kwargs.pop("inputs", None) inputs_embeds = kwargs.pop("inputs_embeds", None) self_generation_config = getattr(self, "generation_config", None) if inputs_embeds is not None: raise ValueError("Diffusion generation currently requires input_ids; inputs_embeds is not supported.") if input_ids is None: input_ids = inputs kwargs.pop("return_dict_in_generate", None) num_denoising_steps = self._resolve( num_denoising_steps, "num_denoising_steps", 128, generation_config, self_generation_config ) if num_denoising_steps < 1: raise ValueError("num_denoising_steps must be a positive integer.") sampler = self._resolve(sampler, "sampler", "ancestral", generation_config, self_generation_config) if sampler not in ("ancestral", "adaptive", "greedy"): raise ValueError( f"sampler must be 'ancestral', 'adaptive' or 'greedy', got {sampler!r}" ) schedule = self._resolve(schedule, "schedule", "linear", generation_config, self_generation_config) min_log_snr = self._resolve(min_log_snr, "min_log_snr", -9.0, generation_config, self_generation_config) max_log_snr = self._resolve(max_log_snr, "max_log_snr", 9.0, generation_config, self_generation_config) temperature = self._resolve(temperature, "temperature", 1.0, generation_config, self_generation_config) tokens_per_step = self._resolve( tokens_per_step, "tokens_per_step", 1, generation_config, self_generation_config ) canvas_length = self._resolve(canvas_length, "canvas_length", 2048, generation_config, self_generation_config) anchor_eosbos = self._resolve(anchor_eosbos, "anchor_eosbos", True, generation_config, self_generation_config) trim_at_eos = self._resolve(trim_at_eos, "trim_at_eos", True, generation_config, self_generation_config) device = self.model.embed_tokens.weight.device if input_ids is None: bos_token_id = kwargs.pop("bos_token_id", None) if bos_token_id is None and generation_config is not None: bos_token_id = getattr(generation_config, "bos_token_id", None) if bos_token_id is None: bos_token_id = self.config.bos_token_id if bos_token_id is None: raise ValueError("input_ids is required when bos_token_id is not configured.") input_ids = torch.tensor([[bos_token_id]], dtype=torch.long, device=device) elif not isinstance(input_ids, torch.Tensor): input_ids = torch.tensor(input_ids, dtype=torch.long, device=device) else: input_ids = input_ids.to(device=device, dtype=torch.long) if input_ids.dim() == 1: input_ids = input_ids.unsqueeze(0) if input_ids.dim() != 2: raise ValueError("input_ids must have shape (batch_size, sequence_length).") batch_size, prompt_length = input_ids.shape if max_new_tokens is None and generation_config is not None: max_new_tokens = getattr(generation_config, "max_new_tokens", None) if max_new_tokens is None and self_generation_config is not None: max_new_tokens = getattr(self_generation_config, "max_new_tokens", None) if max_new_tokens is None: if max_length is None and generation_config is not None: max_length = getattr(generation_config, "max_length", None) if max_length is None and self_generation_config is not None: max_length = getattr(self_generation_config, "max_length", None) if max_length is None: raise ValueError("Either max_new_tokens or max_length must be provided for diffusion generation.") max_new_tokens = max_length - prompt_length if max_new_tokens < 0: raise ValueError("max_new_tokens must be non-negative.") if max_new_tokens == 0: return SumiGenerationOutput(sequences=input_ids, canvas=input_ids) # The model is trained on a packed, fixed-length canvas, so generation runs on a # full `canvas_length` canvas (default 2048) rather than just prompt+max_new_tokens. # `max_new_tokens` is the content budget: the EOS,BOS document delimiter is anchored # at prompt_length+max_new_tokens, the rest of the canvas is denoised as context, and # decoding is cut at the first EOS (trim_at_eos, default True). ceiling = self.config.max_position_embeddings canvas_length = min(int(canvas_length), ceiling) reserve = 2 if anchor_eosbos else 0 if prompt_length + reserve >= canvas_length: raise ValueError( f"prompt_length ({prompt_length}) leaves no room in canvas_length ({canvas_length})." ) budget = max(1, min(max_new_tokens, canvas_length - prompt_length - reserve)) total_length = canvas_length completion_length = total_length - prompt_length # Anchor the EOS,BOS delimiter at the end of the content budget. The caller can # override by passing an explicit `frozen` set (then it takes precedence). if anchor_eosbos and frozen is None: eos_pos = prompt_length + budget frozen = [(eos_pos, self.config.eos_token_id), (eos_pos + 1, self.config.bos_token_id)] if attention_mask is None: attention_mask = torch.ones_like(input_ids) elif not isinstance(attention_mask, torch.Tensor): attention_mask = torch.tensor(attention_mask, device=device) else: attention_mask = attention_mask.to(device=device) if attention_mask.dim() == 1: attention_mask = attention_mask.unsqueeze(0) if attention_mask.shape != input_ids.shape: raise ValueError("attention_mask must have the same shape as input_ids.") if generator is None and seed is not None: try: generator = torch.Generator(device=device) except RuntimeError: generator = torch.Generator() generator.manual_seed(seed) completion_shape = (batch_size, completion_length) try: completion_ids = torch.randint( low=0, high=self.vocab_size, size=completion_shape, dtype=torch.long, device=device, generator=generator, ) except RuntimeError: completion_ids = torch.randint( low=0, high=self.vocab_size, size=completion_shape, dtype=torch.long, generator=generator, ).to(device) generated_ids = torch.cat([input_ids, completion_ids], dim=-1) completion_attention_mask = torch.ones( completion_shape, dtype=attention_mask.dtype, device=device, ) generation_attention_mask = torch.cat([attention_mask, completion_attention_mask], dim=-1) # noise_mask follows the reference _init_z: prompt frozen at the front, plus # optional fixed anchors (frozen) and a per-row active-region bound (denoise_end). noise_mask = _build_noise_mask(generated_ids, prompt_length, frozen, denoise_end) log_snrs = _make_log_snr_schedule( num_denoising_steps, schedule, min_log_snr, max_log_snr, device ) was_training = self.training self.eval() try: for step in range(num_denoising_steps): selected_pos = None log_snr_t = log_snr_s = None if sampler == "ancestral": log_snr_t = float(log_snrs[step]) log_snr_s = float(log_snrs[step + 1]) x_hat = _compute_x_hat( self, generated_ids, generation_attention_mask, self.vocab_size, temperature ) if progress_callback is not None: top1_prob, top1_z = x_hat.max(dim=-1) else: top1_prob = top1_z = None z_new = _ancestral_step( generated_ids, x_hat, log_snr_t, log_snr_s, self.vocab_size, generator ) else: # adaptive / greedy both read the raw logits; the SNR schedule is unused. logits = _compute_logits(self, generated_ids, generation_attention_mask, self.vocab_size) if progress_callback is not None: top1_prob, top1_z = F.softmax(logits, dim=-1).max(dim=-1) else: top1_prob = top1_z = None if sampler == "adaptive": # Confidence-based commit (Eq. 9). z_new, selected_pos = _adaptive_step( generated_ids, logits, noise_mask, tokens_per_step, temperature, generator ) else: # greedy: overwrite every denoise position with the argmax. z_new = _greedy_step(generated_ids, logits, noise_mask) generated_ids = torch.where(noise_mask, z_new, generated_ids) if progress_callback is not None: progress_callback( { "step": step + 1, "num_steps": num_denoising_steps, "log_snr_t": log_snr_t, "log_snr_s": log_snr_s, "z": generated_ids, "top1_z": top1_z, "top1_prob": top1_prob, "noise_mask": noise_mask, # Adaptive only: the [B, k] confidence-selected positions # this step committed (None for ancestral). "selected_pos": selected_pos, } ) finally: if was_training: self.train() sequences = self._trim_at_eos(generated_ids, prompt_length) if trim_at_eos else generated_ids return SumiGenerationOutput(sequences=sequences, canvas=generated_ids) __all__ = [ "SumiGenerationOutput", "SumiGenerationConfig", "SumiGenerationMixin", ]