|
|
| """Block diffusion generation utilities for Dream models.""" |
| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
| from typing import List, Optional, Sequence, Union |
|
|
| import torch |
| from torch.nn import functional as F |
| from transformers.cache_utils import DynamicCache |
| from transformers.utils import ModelOutput |
|
|
|
|
| def top_k_logits(logits: torch.Tensor, k: int) -> torch.Tensor: |
| if k <= 0: |
| return logits |
| values, _ = torch.topk(logits, k) |
| min_values = values[..., -1, None] |
| return torch.where(logits < min_values, torch.full_like(logits, float('-inf')), logits) |
|
|
|
|
| def top_p_logits(logits: torch.Tensor, p: float) -> torch.Tensor: |
| sorted_logits, sorted_indices = torch.sort(logits, descending=True) |
| cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1) |
| sorted_mask = cumulative_probs > p |
| sorted_mask[..., 1:] = sorted_mask[..., :-1].clone() |
| sorted_mask[..., 0] = False |
| mask_indices = torch.scatter( |
| torch.full_like(logits, False, dtype=torch.bool), |
| -1, |
| sorted_indices, |
| sorted_mask, |
| ) |
| return logits.masked_fill(mask_indices, float('-inf')) |
|
|
|
|
| def sample_with_temperature_topk_topp( |
| logits: torch.Tensor, |
| temperature: float = 1.0, |
| top_k: int = 0, |
| top_p: float = 1.0, |
| ) -> tuple[torch.Tensor, torch.Tensor]: |
| orig_shape = logits.shape[:-1] |
| vocab_size = logits.shape[-1] |
| logits = logits.reshape(-1, vocab_size) |
|
|
| if temperature > 0: |
| logits = logits / temperature |
| if top_k > 0: |
| logits = top_k_logits(logits, top_k) |
| if top_p < 1.0: |
| logits = top_p_logits(logits, top_p) |
|
|
| probs = F.softmax(logits, dim=-1) |
| if temperature > 0: |
| token = torch.multinomial(probs, num_samples=1) |
| else: |
| token = probs.argmax(dim=-1, keepdim=True) |
| token_prob = torch.gather(probs, -1, token) |
| return token.view(*orig_shape), token_prob.view(*orig_shape) |
|
|
|
|
| def get_num_transfer_tokens(block_length: int, steps: int) -> torch.Tensor: |
| base = block_length // steps |
| remainder = block_length % steps |
| num_transfer_tokens = torch.zeros(steps, dtype=torch.int64) + base |
| num_transfer_tokens[:remainder] += 1 |
| return num_transfer_tokens |
|
|
|
|
| def build_block_diffusion_attention_mask( |
| num_blocks: int, |
| block_length: int, |
| device: torch.device, |
| batch_size: int = 1, |
| ) -> torch.Tensor: |
| block_mask = torch.tril(torch.ones(num_blocks, num_blocks, device=device)) |
| return block_mask.repeat_interleave(block_length, dim=0).repeat_interleave(block_length, dim=1).unsqueeze(0).expand( |
| batch_size, -1, -1 |
| ) |
|
|
|
|
| def _resolve_stopping_ids(stopping_criteria_idx: Optional[Union[int, Sequence[int]]]) -> Optional[List[int]]: |
| if stopping_criteria_idx is None: |
| return None |
| if isinstance(stopping_criteria_idx, int): |
| return [stopping_criteria_idx] |
| return list(stopping_criteria_idx) |
|
|
|
|
| def _should_stop( |
| generated_ids: torch.Tensor, |
| prompt_length: int, |
| stopping_criteria_idx: Optional[List[int]], |
| ) -> bool: |
| if not stopping_criteria_idx: |
| return False |
| gen_part = generated_ids[:, prompt_length:] |
| return any((gen_part == stop_id).any().item() for stop_id in stopping_criteria_idx) |
|
|
|
|
| def _default_use_kv_cache(model: torch.nn.Module) -> bool: |
| """Dream / Dream1 models use prefix KV cache during block diffusion decode.""" |
| model_type = getattr(model.config, 'model_type', None) |
| if model_type is None: |
| return False |
| return model_type.lower() in ('dream', 'dream1') |
|
|
|
|
| def _select_transfer_index( |
| remasking_strategy: str, |
| mask_index: torch.Tensor, |
| x0: torch.Tensor, |
| x0_p: torch.Tensor, |
| num_transfer_tokens: torch.Tensor, |
| step: int, |
| confidence_threshold: float, |
| eb_threshold: Optional[float], |
| *, |
| force_accept: bool = False, |
| ) -> torch.Tensor: |
| if force_accept: |
| return mask_index.clone() |
|
|
| if remasking_strategy == 'sequential': |
| transfer_index = torch.zeros_like(x0, dtype=torch.bool) |
| for j in range(x0.shape[0]): |
| if not mask_index[j].any(): |
| continue |
| first_mask_index = mask_index[j].nonzero(as_tuple=True)[0].min().item() |
| end = first_mask_index + int(num_transfer_tokens[step].item()) |
| transfer_index[j, first_mask_index:end] = True |
| return transfer_index |
|
|
| if remasking_strategy == 'low_confidence_static': |
| confidence = torch.where(mask_index, x0_p, -torch.inf) |
| transfer_index = torch.zeros_like(x0, dtype=torch.bool) |
| k = max(1, int(num_transfer_tokens[step].item())) |
| for j in range(confidence.shape[0]): |
| _, idx = torch.topk(confidence[j], k) |
| transfer_index[j, idx] = True |
| return transfer_index |
|
|
| if remasking_strategy == 'low_confidence_dynamic': |
| confidence = torch.where(mask_index, x0_p, -torch.inf) |
| transfer_index = torch.zeros_like(x0, dtype=torch.bool) |
| k = max(1, int(num_transfer_tokens[step].item())) |
| for j in range(confidence.shape[0]): |
| high_conf_mask = confidence[j] > confidence_threshold |
| if int(high_conf_mask.sum().item()) >= k: |
| transfer_index[j] = high_conf_mask |
| else: |
| _, idx = torch.topk(confidence[j], k) |
| transfer_index[j, idx] = True |
| return transfer_index |
|
|
| if remasking_strategy == 'entropy_bounded': |
| if eb_threshold is None: |
| raise ValueError('eb_threshold is required for entropy_bounded remasking.') |
| eps = 1e-12 |
| entropies = -(x0_p.clamp_min(eps) * x0_p.clamp_min(eps).log()) |
| entropies = torch.where(mask_index, entropies, torch.inf) |
| ent_sorted, order = torch.sort(entropies, dim=1, descending=False) |
| cumsum = torch.cumsum(ent_sorted, dim=1) |
| transfer_index = torch.zeros_like(x0, dtype=torch.bool) |
| for j in range(x0_p.shape[0]): |
| k = torch.searchsorted( |
| cumsum[j], torch.tensor(eb_threshold, device=x0_p.device), right=False |
| ).item() |
| k = max(1, min(k, int(mask_index[j].sum().item()))) |
| transfer_index[j, order[j, :k]] = True |
| return transfer_index |
|
|
| raise ValueError(f'Unknown remasking strategy: {remasking_strategy}') |
|
|
|
|
| def _denoise_current_block( |
| model: torch.nn.Module, |
| x: torch.Tensor, |
| num_block: int, |
| block_length: int, |
| mask_id: int, |
| block_diffusion_attention_mask: torch.Tensor, |
| position_ids: torch.Tensor, |
| denoising_steps: int, |
| num_transfer_tokens: torch.Tensor, |
| temperature: float, |
| top_k: int, |
| top_p: float, |
| remasking_strategy: str, |
| confidence_threshold: float, |
| eb_threshold: Optional[float], |
| *, |
| use_kv_cache: bool, |
| past_key_values: Optional[DynamicCache], |
| ) -> tuple[torch.Tensor, Optional[DynamicCache], int]: |
| block_start = num_block * block_length |
| block_end = block_start + block_length |
| cur_x = x[:, block_start:block_end].clone() |
| nfe = 0 |
|
|
| for step in range(denoising_steps + 1): |
| mask_index = cur_x == mask_id |
| if mask_index.sum() == 0: |
| if use_kv_cache: |
| cur_attn_mask = block_diffusion_attention_mask[:, block_start:block_end, :block_end] |
| cur_position_ids = position_ids[:, block_start:block_end] |
| model( |
| cur_x, |
| attention_mask=cur_attn_mask, |
| position_ids=cur_position_ids, |
| past_key_values=past_key_values, |
| use_cache=True, |
| store_kv=True, |
| ) |
| nfe += 1 |
| break |
|
|
| force_accept = step == denoising_steps - 1 |
| if use_kv_cache: |
| cur_attn_mask = block_diffusion_attention_mask[:, block_start:block_end, :block_end] |
| cur_position_ids = position_ids[:, block_start:block_end] |
| logits = model( |
| cur_x, |
| attention_mask=cur_attn_mask, |
| position_ids=cur_position_ids, |
| past_key_values=past_key_values, |
| use_cache=True, |
| store_kv=False, |
| ).logits |
| else: |
| seq_end = block_end |
| attn_mask = block_diffusion_attention_mask[:, :seq_end, :seq_end] |
| pos_ids = position_ids[:, :seq_end] |
| logits = model( |
| x[:, :seq_end], |
| attention_mask=attn_mask, |
| position_ids=pos_ids, |
| use_cache=False, |
| ).logits[:, block_start:block_end] |
|
|
| nfe += 1 |
| x0, x0_p = sample_with_temperature_topk_topp( |
| logits, |
| temperature=temperature, |
| top_k=top_k, |
| top_p=top_p, |
| ) |
| x0 = torch.where(mask_index, x0, cur_x) |
| transfer_index = _select_transfer_index( |
| remasking_strategy, |
| mask_index, |
| x0, |
| x0_p, |
| num_transfer_tokens, |
| step, |
| confidence_threshold, |
| eb_threshold, |
| force_accept=force_accept, |
| ) |
| cur_x[transfer_index] = x0[transfer_index] |
| if not use_kv_cache: |
| x[:, block_start:block_end] = cur_x |
|
|
| return cur_x, past_key_values, nfe |
|
|
|
|
| @dataclass |
| class BlockDiffusionOutput(ModelOutput): |
| sequences: torch.LongTensor = None |
| nfe: Optional[int] = None |
| logits: Optional[tuple] = None |
|
|
|
|
| @torch.no_grad() |
| def block_diffusion_generate( |
| model: torch.nn.Module, |
| input_ids: torch.LongTensor, |
| mask_id: int, |
| gen_length: int = 128, |
| block_length: Optional[int] = None, |
| denoising_steps: Optional[int] = None, |
| temperature: float = 0.0, |
| top_k: int = 0, |
| top_p: float = 1.0, |
| remasking_strategy: str = 'low_confidence_dynamic', |
| confidence_threshold: float = 0.9, |
| eb_threshold: Optional[float] = 0.35, |
| stopping_criteria_idx: Optional[Union[int, Sequence[int]]] = None, |
| use_kv_cache: Optional[bool] = None, |
| return_dict_in_generate: bool = False, |
| ) -> Union[torch.LongTensor, BlockDiffusionOutput]: |
| """Block-wise diffusion decoding with optional prefix KV cache.""" |
| model.eval() |
| if input_ids.dim() != 2: |
| raise ValueError(f'input_ids must be 2D, got shape {tuple(input_ids.shape)}') |
|
|
| device = input_ids.device |
| batch_size, prompt_length = input_ids.shape |
| block_length = block_length or getattr(model.config, 'block_size', 4) |
| if denoising_steps is None: |
| denoising_steps = 1 if remasking_strategy == 'low_confidence_static' else block_length |
| stopping_criteria_idx = _resolve_stopping_ids(stopping_criteria_idx) |
|
|
| if use_kv_cache is None: |
| use_kv_cache = _default_use_kv_cache(model) |
|
|
| num_blocks = (prompt_length + gen_length + block_length - 1) // block_length |
| total_length = num_blocks * block_length |
| block_diffusion_attention_mask = build_block_diffusion_attention_mask( |
| num_blocks, block_length, device, batch_size=batch_size |
| ) |
| position_ids = torch.arange(total_length, device=device, dtype=torch.long).unsqueeze(0).expand(batch_size, -1) |
|
|
| x = torch.full((batch_size, total_length), mask_id, dtype=input_ids.dtype, device=device) |
| x[:, :prompt_length] = input_ids |
|
|
| prefill_blocks = prompt_length // block_length |
| prefill_length = prefill_blocks * block_length |
| past_key_values = DynamicCache() if use_kv_cache else None |
| nfe = 0 |
|
|
| if use_kv_cache and prefill_length > 0: |
| cur_x = x[:, :prefill_length] |
| cur_attn_mask = block_diffusion_attention_mask[:, :prefill_length, :prefill_length] |
| cur_position_ids = position_ids[:, :prefill_length] |
| model( |
| cur_x, |
| attention_mask=cur_attn_mask, |
| position_ids=cur_position_ids, |
| past_key_values=past_key_values, |
| use_cache=True, |
| store_kv=True, |
| ) |
| nfe += 1 |
|
|
| num_transfer_tokens = get_num_transfer_tokens(block_length, denoising_steps) |
|
|
| for num_block in range(prefill_blocks, num_blocks): |
| cur_x, past_key_values, block_nfe = _denoise_current_block( |
| model, |
| x, |
| num_block, |
| block_length, |
| mask_id, |
| block_diffusion_attention_mask, |
| position_ids, |
| denoising_steps, |
| num_transfer_tokens, |
| temperature, |
| top_k, |
| top_p, |
| remasking_strategy, |
| confidence_threshold, |
| eb_threshold, |
| use_kv_cache=use_kv_cache, |
| past_key_values=past_key_values, |
| ) |
| nfe += block_nfe |
| x[:, num_block * block_length:(num_block + 1) * block_length] = cur_x |
|
|
| if _should_stop(x, prompt_length, stopping_criteria_idx): |
| break |
|
|
| output_length = min(total_length, prompt_length + gen_length) |
| x = x[:, :output_length] |
|
|
| if return_dict_in_generate: |
| return BlockDiffusionOutput(sequences=x, nfe=nfe) |
| return x |
|
|
|
|
| |
| _UNSUPPORTED_HF_KEYS = ( |
| 'stopping_criteria', 'num_return_sequences', 'num_beams', 'num_beam_groups', |
| 'penalty_alpha', 'use_cache', 'output_logits', 'output_scores', 'output_attentions', |
| 'output_hidden_states', 'return_legacy_cache', 'synced_gpus', 'streamer', |
| 'logits_processor', 'logits_warper', 'generation_config', 'tokenizer', |
| 'min_length', 'min_new_tokens', 'pad_token_id', 'bos_token_id', 'eos_token_id', |
| ) |
|
|
|
|
| class BlockDiffusionGenerationMixin: |
|
|
| def _resolve_generation_mode(self, generation_mode: Optional[str] = None) -> str: |
| if generation_mode is not None: |
| return generation_mode |
| return getattr(self.config, 'generation_mode', 'block_diffusion') |
|
|
| @torch.no_grad() |
| def generate( |
| self, |
| input_ids: torch.LongTensor, |
| attention_mask: Optional[torch.Tensor] = None, |
| generation_mode: Optional[str] = None, |
| **kwargs, |
| ) -> Union[torch.LongTensor, BlockDiffusionOutput, ModelOutput]: |
| """HF-compatible ``generate`` entry point with pluggable decoding modes. |
| |
| Supported modes (``generation_mode`` kwarg or ``config.generation_mode``): |
| - ``block_diffusion`` (default): calls :meth:`block_diffusion_generate`. |
| - ``autoregressive``: delegates to ``GenerationMixin.generate``. |
| """ |
| mode = self._resolve_generation_mode(kwargs.pop('generation_mode', generation_mode)) |
| return_dict_in_generate = kwargs.pop('return_dict_in_generate', False) |
|
|
| if mode == 'autoregressive': |
| return super().generate( |
| input_ids=input_ids, |
| attention_mask=attention_mask, |
| return_dict_in_generate=return_dict_in_generate, |
| **kwargs, |
| ) |
| if mode != 'block_diffusion': |
| raise ValueError(f'Unknown generation_mode: {mode!r}. Supported: block_diffusion, autoregressive.') |
|
|
| |
| if kwargs.pop('do_sample', None) is False: |
| kwargs['temperature'] = 0.0 |
| if 'max_new_tokens' not in kwargs and (max_length := kwargs.pop('max_length', None)) is not None: |
| kwargs['max_new_tokens'] = max(max_length - input_ids.shape[-1], 0) |
| for key in _UNSUPPORTED_HF_KEYS: |
| kwargs.pop(key, None) |
|
|
| return self.block_diffusion_generate( |
| input_ids=input_ids, |
| return_dict_in_generate=return_dict_in_generate, |
| **kwargs, |
| ) |
|
|
| @torch.no_grad() |
| def block_diffusion_generate( |
| self, |
| input_ids: torch.LongTensor, |
| max_new_tokens: int = 128, |
| temperature: float = 0.0, |
| top_k: int = 0, |
| top_p: float = 1.0, |
| return_dict_in_generate: bool = False, |
| |
| block_length: Optional[int] = None, |
| denoising_steps: Optional[int] = None, |
| remasking_strategy: str = 'low_confidence_dynamic', |
| confidence_threshold: float = 0.9, |
| eb_threshold: Optional[float] = 0.35, |
| use_kv_cache: Optional[bool] = None, |
| mask_token_id: Optional[int] = None, |
| ) -> Union[torch.LongTensor, BlockDiffusionOutput]: |
| mask_token_id = mask_token_id if mask_token_id is not None else self.config.mask_token_id |
| if mask_token_id is None: |
| raise ValueError('mask_token_id must be provided or set in model.config.mask_token_id') |
| stopping_criteria_idx = self.config.eos_token_id if getattr(self.config, 'eos_token_id', None) is not None else None |
| return block_diffusion_generate( |
| self, |
| input_ids=input_ids, |
| mask_id=mask_token_id, |
| gen_length=max_new_tokens, |
| block_length=block_length, |
| denoising_steps=denoising_steps, |
| temperature=temperature, |
| top_k=top_k, |
| top_p=top_p, |
| remasking_strategy=remasking_strategy, |
| confidence_threshold=confidence_threshold, |
| eb_threshold=eb_threshold, |
| stopping_criteria_idx=stopping_criteria_idx, |
| use_kv_cache=use_kv_cache, |
| return_dict_in_generate=return_dict_in_generate, |
| ) |
|
|
