import copy from dataclasses import dataclass from typing import Callable, Optional, Tuple, Union import random import os import sys import json import numpy as np import torch import torch.nn.functional as F from torch import nn from transformers.modeling_outputs import CausalLMOutputWithPast, BaseModelOutput from transformers.utils import ModelOutput from torch.nn.attention.flex_attention import BlockMask, flex_attention, create_block_mask, or_masks from transformers.modeling_flash_attention_utils import FlashAttentionKwargs from transformers.processing_utils import Unpack from transformers.cache_utils import Cache, DynamicCache from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss from transformers.generation import GenerationMixin from transformers.loss.loss_utils import LOSS_MAPPING import math from .chat_utils import generate_with_prefix_cache_block_diff from .modeling_ministral import Ministral3Model, Ministral3PreTrainedModel, Ministral3Attention, apply_rotary_pos_emb, repeat_kv, _get_llama_4_attn_scale from .configuration_nemotron_labs_diffusion_vlm import NemotronLabsDiffusionVLMConfig @dataclass class NemotronLabsDiffusionVLMOutputWithPast(ModelOutput): loss: torch.FloatTensor | None = None logits: torch.FloatTensor | None = None causal_logits: torch.FloatTensor | None = None past_key_values: Cache | None = None hidden_states: tuple[torch.FloatTensor, ...] | None = None attentions: tuple[torch.FloatTensor, ...] | None = None # @torch.compile(dynamic=True, mode="reduce-overhead") # @torch.compile(mode="default") # @torch.compile(fullgraph=True, mode="reduce-overhead", dynamic=False) @torch.compile(fullgraph=True, mode="max-autotune-no-cudagraphs", dynamic=False) def fused_flex_attention(q, k, v, block_mask=None): return flex_attention(q, k, v, block_mask=block_mask) def _crop_dynamic_cache(past_key_values: DynamicCache, max_length: int): """Crop a DynamicCache to max_length, compatible with both old and new transformers.""" if hasattr(past_key_values, 'crop'): past_key_values.crop(max_length) else: for layer_idx in range(len(past_key_values)): past_key_values.key_cache[layer_idx] = past_key_values.key_cache[layer_idx][:, :, :max_length] past_key_values.value_cache[layer_idx] = past_key_values.value_cache[layer_idx][:, :, :max_length] past_key_values._seen_tokens = max_length def _extract_draft_kv_cache(past_key_values: DynamicCache, clean_len: int, block_length: int): """After quadratic decoding, extract only draft tokens (first of each block) from cache.""" for layer_idx in range(len(past_key_values)): if hasattr(past_key_values, 'layers'): layer_cache = past_key_values.layers[layer_idx] k, v = layer_cache.keys, layer_cache.values else: k = past_key_values.key_cache[layer_idx] v = past_key_values.value_cache[layer_idx] clean_k, draft_k = k[:, :, :clean_len], k[:, :, clean_len::block_length + 1] clean_v, draft_v = v[:, :, :clean_len], v[:, :, clean_len::block_length + 1] new_k = torch.cat([clean_k, draft_k], dim=2) new_v = torch.cat([clean_v, draft_v], dim=2) if hasattr(past_key_values, 'layers'): layer_cache.keys = new_k layer_cache.values = new_v else: past_key_values.key_cache[layer_idx] = new_k past_key_values.value_cache[layer_idx] = new_v past_key_values._seen_tokens = clean_len + block_length # with reference to https://github.com/pytorch-labs/attention-gym/blob/main/examples/flex_attn.ipynb class NemotronLabsDiffusionVLMFlexAttention(Ministral3Attention): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.max_seq_length = getattr(self.config, 'max_seq_length', 4096) self.block_size_orig = self.config.block_size if self.config.dlm_paradigm == 'bidirectional': self.bidirectional_mask = self.compute_block_mask(mode='bidirectional') elif self.config.dlm_paradigm == 'autoregressive': self.autoregressive_mask = self.compute_block_mask(mode='autoregressive') elif self.config.dlm_paradigm == 'block_diff': self.block_diff_mask = None elif self.config.dlm_paradigm == 'sbd_block_diff': self.sbd_block_diff_mask = None else: raise ValueError(f"Unknown attention mode: {self.config.dlm_paradigm}") self.block_size = self.block_size_orig self.mode = self.config.dlm_paradigm self._quadratic_block_mask = {} import torch._dynamo.config as dcfg dcfg.cache_size_limit = 512 def _get_sbd_inference_quadratic_decoding_block_mask(self, block_length: int): if block_length not in self._quadratic_block_mask: draft_len = block_length * (block_length + 1) def quadratic(b, h, q_idx, kv_idx): first_clean = torch.logical_and( kv_idx % (block_length + 1) == 0, kv_idx < draft_len, ) first_clean = torch.logical_and(first_clean, q_idx >= kv_idx) block_q = q_idx // (block_length + 1) block_kv = kv_idx // (block_length + 1) same_block = torch.logical_and(block_q == block_kv, q_idx < draft_len) same_block_except_first = torch.logical_and( same_block, q_idx % (block_length + 1) != 0, ) draft_part = torch.logical_or(first_clean, same_block_except_first) clean_part = kv_idx >= draft_len return torch.logical_or(draft_part, clean_part) block_mask = create_block_mask( quadratic, B=None, H=None, Q_LEN=draft_len, KV_LEN=draft_len + self.config.max_position_embeddings, device="cuda", ) self._quadratic_block_mask[block_length] = block_mask return self._quadratic_block_mask[block_length] def set_attention_mode(self, mode, block_size=None): self.mode = mode self.block_size = block_size def compute_block_mask(self, mode, q_len=None, block_size=None): def bidirectional_mask(b, h, q, kv): return (q >= kv) | (q < kv) def autoregressive_mask(b, h, q, kv): return (q >= kv) def block_diff_mask(block_size, b, h, q_idx, kv_idx, n): """ Constructs the specialized block diffusion attention mask for training composed of three masks: - **Block Diagonal Mask (M_BD)**: Self-attention within noised blocks - **Offset Block Causal Mask (M_OBC)**: Cross-attention for conditional context - **Block Causal Mask (M_BC)**: Attention to update x0 Args: b, h: Batch and head indices (ignored for mask logic). q_idx, kv_idx: Query and Key indices. seq_len: Total sequence length. block_size: Defines the block structure. Returns: A boolean attention mask. """ # Indicate whether token belongs to xt or x0 x0_flag_q = (q_idx >= n) x0_flag_kv = (kv_idx >= n) # Compute block indices block_q = torch.where(x0_flag_q == 1, (q_idx - n) // block_size, q_idx // block_size) block_kv = torch.where(x0_flag_kv == 1, (kv_idx - n) // block_size, kv_idx // block_size) # **1. Block Diagonal Mask (M_BD) ** block_diagonal = (block_q == block_kv) & (x0_flag_q == x0_flag_kv) # **2. Offset Block-Causal Mask (M_OBC) ** offset_block_causal = ( (block_q > block_kv) & (x0_flag_kv == 1) & (x0_flag_q == 0) ) # **3. Block-Causal Mask (M_BC) ** block_causal = (block_q >= block_kv) & (x0_flag_kv == 1) & (x0_flag_q == 1) # **4. Combine Masks ** return block_diagonal | offset_block_causal | block_causal def sbd_block_diff_mask(block_size, b, h, q_idx, kv_idx, n): x0_flag_q = (q_idx >= n) x0_flag_kv = (kv_idx >= n) # Compute block indices block_q = torch.where(x0_flag_q == 1, (q_idx - n) // block_size, q_idx // block_size) block_kv = torch.where(x0_flag_kv == 1, (kv_idx - n) // block_size, kv_idx // block_size) # **1. Block Diagonal Mask (M_BD) ** block_diagonal = (block_q == block_kv) & (x0_flag_kv == 0) & (x0_flag_q == 0) # **2. Offset Block-Causal Mask (M_OBC) ** offset_block_causal = ( (block_q > block_kv) & (x0_flag_kv == 1) & (x0_flag_q == 0) ) # **3. Fully Causal Mask (M_BC) ** fully_causal = (q_idx >= kv_idx) & (x0_flag_kv == 1) & (x0_flag_q == 1) # **4. Combine Masks ** return block_diagonal | offset_block_causal | fully_causal if mode == 'bidirectional': attn_mask = bidirectional_mask elif mode == 'autoregressive': attn_mask = autoregressive_mask elif mode == 'block_diff': assert block_size is not None n = (q_len // 2) if q_len is not None else self.max_seq_length attn_mask = lambda b, h, q, kv: block_diff_mask(block_size, b, h, q, kv, n) elif mode == 'sbd_block_diff': assert block_size is not None n = (q_len // 2) if q_len is not None else self.max_seq_length attn_mask = lambda b, h, q, kv: sbd_block_diff_mask(block_size, b, h, q, kv, n) else: raise ValueError(f"Unknown attention mode: {mode}") if q_len is not None: Q_LEN = q_len else: if mode in ['block_diff', 'sbd_block_diff']: Q_LEN = self.max_seq_length * 2 else: Q_LEN = self.max_seq_length block_mask = create_block_mask( attn_mask, B=None, H=None, Q_LEN=Q_LEN, KV_LEN=Q_LEN ) return block_mask def forward( self, hidden_states: torch.Tensor, position_embeddings: Tuple[torch.Tensor, torch.Tensor], attention_mask: Optional[torch.Tensor], past_key_values: Optional[Cache] = None, cache_position: Optional[torch.LongTensor] = None, is_training: bool = True, **kwargs: Unpack[FlashAttentionKwargs], ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: bsz, q_len, _ = hidden_states.size() input_shape = hidden_states.shape[:-1] hidden_shape = (*input_shape, -1, self.head_dim) query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2) value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) cos, sin = position_embeddings if self.mode in ['block_diff', 'sbd_block_diff'] and is_training: # Split query and key states in half along sequence length dimension q1, q2 = query_states.chunk(2, dim=2) k1, k2 = key_states.chunk(2, dim=2) # Apply RoPE independently to each half q1, k1 = apply_rotary_pos_emb(q1, k1, cos, sin) q2, k2 = apply_rotary_pos_emb(q2, k2, cos, sin) # Recombine the halves query_states = torch.cat([q1, q2], dim=2) key_states = torch.cat([k1, k2], dim=2) else: query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) query_states = query_states * _get_llama_4_attn_scale( cache_position, self.config.rope_parameters.get("llama_4_scaling_beta"), self.config.rope_parameters.get("original_max_position_embeddings"), ).to(query_states.dtype) if past_key_values is not None: # sin and cos are specific to RoPE models; cache_position needed for the static cache cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs) self_spec_inference_mode = getattr(self.config, "self_spec_inference_mode", None) if self_spec_inference_mode is not None: if self_spec_inference_mode == "quadratic": block_length = getattr(self.config, "block_length", None) or getattr(self.config, "block_size", None) if block_length is None: raise ValueError("SBD quadratic decoding requires block_length in config.") if past_key_values is not None: seq_len = key_states.shape[2] draft_len = block_length * (block_length + 1) clean_keys = key_states[:, :, :-draft_len] draft_keys = key_states[:, :, -draft_len:] clean_values = value_states[:, :, :-draft_len] draft_values = value_states[:, :, -draft_len:] key_states = torch.cat([draft_keys, clean_keys], dim=2) value_states = torch.cat([draft_values, clean_values], dim=2) block_mask = self._get_sbd_inference_quadratic_decoding_block_mask( block_length=block_length ) block_mask.seq_lengths = (draft_len, seq_len) else: seq_len = query_states.shape[2] draft_len = block_length * (block_length + 1) clean_len = seq_len - draft_len def _causal_mask(b, h, q_idx, kv_idx): return torch.logical_and(q_idx >= kv_idx, q_idx < clean_len) def _draft2clean_mask(b, h, q_idx, kv_idx): full_clean = torch.logical_and(q_idx >= clean_len, kv_idx <= clean_len) first_clean = torch.logical_and( q_idx >= clean_len, (kv_idx - clean_len) % (block_length + 1) == 0 ) first_clean = torch.logical_and(first_clean, q_idx >= kv_idx) return torch.logical_or(full_clean, first_clean) def _draft_mask(b, h, q_idx, kv_idx): block_q = (q_idx - clean_len) // (block_length + 1) block_kv = (kv_idx - clean_len) // (block_length + 1) quadrant = torch.logical_and(q_idx >= clean_len, kv_idx >= clean_len) same_block = torch.logical_and(block_q == block_kv, quadrant) same_block_except_first = torch.logical_and( same_block, (q_idx - clean_len) % (block_length + 1) != 0, ) return torch.logical_and(block_q == block_kv, same_block_except_first) mask = or_masks(_causal_mask, _draft2clean_mask) mask = or_masks(mask, _draft_mask) block_mask = create_block_mask( mask, B=None, H=None, Q_LEN=seq_len, KV_LEN=seq_len, ) key_states = repeat_kv(key_states, self.num_key_value_groups) value_states = repeat_kv(value_states, self.num_key_value_groups) attn_output = flex_attention(query_states, key_states, value_states, block_mask=block_mask) attn_output = attn_output.transpose(1, 2).reshape(*input_shape, -1).contiguous() attn_output = self.o_proj(attn_output) return attn_output, None elif self_spec_inference_mode == "default": block_length = getattr(self.config, "block_length", None) or getattr(self.config, "block_size", None) if block_length is None: raise ValueError("SBD default decoding requires block_length in config.") seq_len = query_states.shape[2] prefix_len = seq_len - block_length def _clean_q_mask(b, h, q_idx, kv_idx): return torch.logical_and(q_idx >= kv_idx, q_idx < prefix_len) def _noisy_q_mask(b, h, q_idx, kv_idx): return q_idx >= prefix_len block_mask = create_block_mask( or_masks(_clean_q_mask, _noisy_q_mask), B=None, H=None, Q_LEN=seq_len, KV_LEN=seq_len, ) key_states = repeat_kv(key_states, self.num_key_value_groups) value_states = repeat_kv(value_states, self.num_key_value_groups) attn_output = flex_attention(query_states, key_states, value_states, block_mask=block_mask) attn_output = attn_output.transpose(1, 2).reshape(*input_shape, -1).contiguous() attn_output = self.o_proj(attn_output) return attn_output, None key_states = repeat_kv(key_states, self.num_key_value_groups) value_states = repeat_kv(value_states, self.num_key_value_groups) if self.mode == 'bidirectional': if self.bidirectional_mask is None or q_len != self.bidirectional_mask.shape[-2]: block_mask = self.compute_block_mask(mode='bidirectional', q_len=q_len) else: block_mask = self.bidirectional_mask elif self.mode == 'autoregressive': if self.autoregressive_mask is None or q_len != self.autoregressive_mask.shape[-2]: block_mask = self.compute_block_mask(mode='autoregressive', q_len=q_len) else: block_mask = self.autoregressive_mask elif self.mode == 'block_diff': if self.block_diff_mask is None or self.block_size != self.block_size_orig or q_len != self.block_diff_mask.shape[-2]: block_mask = self.compute_block_mask(mode='block_diff', block_size=self.block_size, q_len=q_len) else: block_mask = self.block_diff_mask elif self.mode == 'sbd_block_diff': if self.sbd_block_diff_mask is None or self.block_size != self.block_size_orig or q_len != self.sbd_block_diff_mask.shape[-2]: block_mask = self.compute_block_mask(mode='sbd_block_diff', block_size=self.block_size, q_len=q_len) else: block_mask = self.sbd_block_diff_mask else: raise ValueError(f"Unknown attention mode: {self.mode}") attn_output = fused_flex_attention(query_states, key_states, value_states, block_mask=block_mask) attn_output = attn_output.transpose(1, 2).reshape(*input_shape, -1).contiguous() attn_output = self.o_proj(attn_output) return attn_output, None def gumbel_topk(log_w: torch.Tensor, k: int) -> torch.Tensor: """Return a Bool mask of length len(log_w) with exactly k True.""" g = -torch.log(-torch.log(torch.rand_like(log_w) + 1e-9) + 1e-9) topk = torch.topk(log_w + g, k).indices mask = torch.zeros_like(log_w, dtype=torch.bool) mask[topk] = True return mask class NemotronLabsDiffusionVLMModel(Ministral3PreTrainedModel, GenerationMixin): """ A single model with: - a bidirectional encoder + diffusion‐LM head over A - a causal decoder + LM head over B, conditioned on F_A """ def __init__(self, config: NemotronLabsDiffusionVLMConfig): super().__init__(config) self.mask_token_id = config.mask_token_id diffusion_config = copy.deepcopy(config) diffusion_config.diffusion_lm = True use_flex = getattr(config, 'enable_self_spec', False) if config.dlm_paradigm in ['block_diff', 'sbd_block_diff']: diffusion_config.attn_class = NemotronLabsDiffusionVLMFlexAttention elif config.dlm_paradigm in ['bidirectional', 'autoregressive']: diffusion_config.attn_class = NemotronLabsDiffusionVLMFlexAttention if use_flex else Ministral3Attention if config.dlm_paradigm == 'autoregressive': diffusion_config.diffusion_lm = False else: raise ValueError(f"Unsupported DLM paradigm: {config.dlm_paradigm}") self.encoder = Ministral3Model(diffusion_config) self.diffusion_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) self.vocab_size = config.vocab_size self.current_iter_ratio = None self.mdm_loss_function = LOSS_MAPPING['ForMaskedLM'] self.causal_loss_function = LOSS_MAPPING['ForCausalLM'] self.post_init() def get_input_embeddings(self): return self.encoder.embed_tokens def set_input_embeddings(self, value): self.encoder.embed_tokens = value def get_output_embeddings(self): return self.diffusion_head def set_output_embeddings(self, new_embeddings): self.diffusion_head = new_embeddings def forward_process_complementary(self, input_ids, eps=1e-3, block_size=None, loss_mask=None): device = input_ids.device if self.config.dp_varying_mask_ratio: import torch.distributed as dist dp_rank = 0 if dist.is_initialized(): try: dp_rank = dist.get_rank() except Exception: dp_rank = 0 generator = torch.Generator(device=device) generator.manual_seed(torch.seed() + dp_rank) else: generator = None noisy_input_ids = input_ids.clone() input_ids_flat = input_ids.reshape(input_ids.shape[0] * input_ids.shape[1] // block_size, block_size) b, l = input_ids_flat.shape t = torch.rand((b,), device=input_ids.device, generator=generator) p_mask = (1 - eps) * t + eps p_mask = p_mask[:, None].repeat(1, l) masked_indices = (torch.rand((b, l), device=input_ids.device, generator=generator) < p_mask).reshape(noisy_input_ids.shape) input_ids_flat = input_ids_flat.reshape(noisy_input_ids.shape) complementary_noisy_input_ids = input_ids.clone() complementary_masked_indices = ~masked_indices if getattr(self.config, 'always_mask_im_end', False): im_end_mask = (input_ids == self.config.im_end_token_id) masked_indices = masked_indices | im_end_mask complementary_masked_indices = complementary_masked_indices | im_end_mask if loss_mask is not None: masked_indices[loss_mask == 0] = 0 complementary_masked_indices[loss_mask == 0] = 0 noisy_input_ids[masked_indices] = self.mask_token_id complementary_noisy_input_ids[complementary_masked_indices] = self.mask_token_id noisy_input_ids = torch.cat([noisy_input_ids, complementary_noisy_input_ids], dim=0) masked_indices = torch.cat([masked_indices, complementary_masked_indices], dim=0) return noisy_input_ids, masked_indices, None # ── Vision / multimodal helpers (ported from Mistral3Model) ────────── IMAGE_TOKEN_ID = 19 def get_image_features( self, pixel_values: torch.FloatTensor, image_sizes: torch.Tensor, ) -> torch.FloatTensor: """ Run the vision tower + multimodal projector and return a flat tensor of image features ready to be scattered into the text embeddings. Mirrors ``Mistral3Model.get_image_features`` from transformers/models/mistral3/modeling_mistral3.py. Returns: Flat (total_image_tokens, hidden_size) tensor. """ vision_feature_layer = getattr(self.config, "vision_feature_layer", -1) image_outputs = self.encoder.vision_tower( pixel_values, image_sizes=image_sizes, output_hidden_states=True, return_dict=True, ) if isinstance(vision_feature_layer, int): selected_image_feature = image_outputs.hidden_states[vision_feature_layer] else: hs_pool = [image_outputs.hidden_states[idx] for idx in vision_feature_layer] selected_image_feature = torch.cat(hs_pool, dim=-1) image_features = self.encoder.multi_modal_projector( selected_image_feature.squeeze(0), image_sizes, ) # Split per image, then re-cat into one flat tensor downsample_ratio = ( self.encoder.vision_tower.patch_size * getattr(self.config, "spatial_merge_size", 2) ) split_sizes = ( (torch.as_tensor(image_sizes, device=image_features.device) // downsample_ratio) .prod(dim=-1) .tolist() ) # per_image = torch.split(image_features.squeeze(0), split_sizes) per_image = torch.split(image_features, split_sizes) return torch.cat(per_image, dim=0) # (total_tokens, hidden) def _is_vision_frozen(self) -> bool: """True if vision_tower and multi_modal_projector have no parameters requiring grad (e.g. --freeze_vision_encoder).""" vt = self.encoder.vision_tower proj = self.encoder.multi_modal_projector vt_has_grad = any(p.requires_grad for p in vt.parameters()) proj_has_grad = any(p.requires_grad for p in proj.parameters()) return not vt_has_grad and not proj_has_grad def _embed_with_vision( self, input_ids: torch.LongTensor, pixel_values: torch.FloatTensor, image_sizes: torch.Tensor, ) -> torch.FloatTensor: """ Embed *input_ids* and scatter vision features into [IMG] pad positions. Returns: inputs_embeds (batch, seq_len, hidden_size) """ inputs_embeds = self.encoder.embed_tokens(input_ids) image_features = self.get_image_features(pixel_values, image_sizes) image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) # Boolean mask: positions that are IMG pad tokens special_image_mask = (input_ids == self.IMAGE_TOKEN_ID) if self.training: if self.config.complementary_mask: image_features = image_features.repeat(2, 1) if self.config.dlm_paradigm in ['block_diff', 'sbd_block_diff']: image_features = image_features.repeat(2, 1) assert special_image_mask.sum() == image_features.shape[0], f"special_image_mask.sum() = {special_image_mask.sum()}, image_features.shape[0] = {image_features.shape[0]}" # Expand to hidden dim for masked_scatter special_image_mask = special_image_mask.unsqueeze(-1).expand_as(inputs_embeds) inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features) return inputs_embeds def forward_process(self, input_ids, eps=1e-3, block_size=None, loss_mask=None): b, l = input_ids.shape device = input_ids.device if self.config.dp_varying_mask_ratio: # Enable different random seeds for each DP rank during sampling import torch.distributed as dist dp_rank = 0 if dist.is_initialized(): try: dp_rank = dist.get_rank() except Exception: dp_rank = 0 # Use a local generator to avoid affecting global RNG state generator = torch.Generator(device=device) generator.manual_seed(torch.seed() + dp_rank) else: generator = None if self.config.adaptive_mask_rate: assert block_size is not None # --- simple linear window mapping --- bs_min = getattr(self.config, "t_bs_min", 16) bs_max = getattr(self.config, "t_bs_max", 128) w = getattr(self.config, "t_window_width", 0.6) # fixed width # fraction in [0,1] (unclamped first) frac = (float(block_size) - float(bs_min)) / max(1.0, float(bs_max - bs_min)) # upper bound decreases linearly from 1.0 -> 0.5 u_max = 1.0 - w * frac # clamp to [0.6, 1.0] to handle bs outside [bs_min, bs_max] u_max = max(0.6, min(1.0, u_max)) u_min = u_max - w # ensures width = w # sample t ~ Uniform(u_min, u_max) t = u_min + (u_max - u_min) * torch.rand(b, device=device, generator=generator) else: t = torch.rand(b, device=device, generator=generator) p_mask = (1 - eps) * t + eps # shape: (b,) p_mask = p_mask[:, None].expand(-1, l) # shape: (b, l) masked_indices = torch.rand((b, l), device=device) < p_mask if loss_mask is not None: masked_indices[loss_mask == 0] = 0 noisy_batch = torch.where(masked_indices, self.mask_token_id, input_ids) return noisy_batch, masked_indices, p_mask def forward_process_exp( self, input_ids: torch.Tensor, eps: float = 1e-3, block_size: int | None = None, half_life_ratio: float = 0.25, # λ = ln 2 / (half_life_ratio·L) loss_mask: Optional[torch.Tensor] = None, ): """ Two-stage corruption with optional per-block sampling. • Stage 1: m ~ U(eps, 1) → k = round(m · len) (exact budget). • Stage 2: sample exactly k positions with weights w_i(m) = exp[ λ · (1−m) · i ] (late-heavy when m→0, uniform when m→1). If `block_size` is given, the procedure is run *independently* inside each contiguous block of that length (last block may be shorter). When block_size is provided, m is sampled per-block and p_mask is per-block. Args ---- input_ids : (B, L) LongTensor eps : minimum corruption ratio block_size: if not None, operate block-wise with per-block m sampling half_life_ratio : controls steepness when m→0 """ B, L = input_ids.shape device = input_ids.device dtype = torch.float32 masked_indices = torch.zeros((B, L), dtype=torch.bool, device=device) p_mask = torch.zeros((B, L), dtype=dtype, device=device) # ---------- Stage 1 & 2: whole-sentence or block-wise ------------------- for b in range(B): if block_size is None: # ---------- Per-batch sampling (original behavior) ---------- m = eps + (1.0 - eps) * torch.rand(1, device=device).item() # scalar k_tot = int(round(m * L)) k_tot = max(1, min(k_tot, L)) # clamp to [1, L] # Fill p_mask for this batch p_mask[b, :] = m slope = 1.0 - m # ∈ [0,1]; 0 ⇒ uniform, 1 ⇒ late-heavy # ------- single pool over the whole sentence ------------- lam_base = math.log(2.0) / (half_life_ratio * L) # base decay rate (λ when slope=1) pos = torch.arange(L, device=device, dtype=dtype) log_w = (lam_base * slope * pos).clone() masked_indices[b] = gumbel_topk(log_w, k_tot) else: # ---------- Per-block sampling ---------- num_blocks = math.ceil(L / block_size) lam_base = math.log(2.0) / (half_life_ratio * block_size) # base decay rate (λ when slope=1) for blk in range(num_blocks): start = blk * block_size end = min((blk + 1) * block_size, L) blk_len = end - start # Sample m per block m_blk = eps + (1.0 - eps) * torch.rand(1, device=device).item() # Fill p_mask for this block p_mask[b, start:end] = m_blk # per-block budget k_blk = int(round(m_blk * blk_len)) k_blk = max(0, min(k_blk, blk_len)) if k_blk == 0: continue slope = 1.0 - m_blk # ∈ [0,1]; 0 ⇒ uniform, 1 ⇒ late-heavy pos = torch.arange(blk_len, device=device, dtype=dtype) log_w = lam_base * slope * pos blk_mask = gumbel_topk(log_w, k_blk) masked_indices[b, start:end] = blk_mask if loss_mask is not None: masked_indices[loss_mask == 0] = 0 noisy_batch = torch.where(masked_indices, self.mask_token_id, input_ids) return noisy_batch, masked_indices, p_mask def forward( self, input_ids: torch.LongTensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, labels: Optional[torch.LongTensor] = None, split_len: Optional[int] = None, past_key_values: Optional[Cache] = None, block_size: Optional[int] = None, block_diff_ppl: bool = False, eps: float = 1e-3, is_teacher: bool = False, masked_indices: Optional[torch.Tensor] = None, p_mask: Optional[torch.Tensor] = None, teacher_logits: Optional[torch.Tensor] = None, masked_indices_teacher: Optional[torch.Tensor] = None, loss_mask: Optional[torch.Tensor] = None, ce_loss_weight: float = 1.0, output_last_hidden_states_only: bool = False, skip_loss: bool = False, pixel_values: Optional[torch.FloatTensor] = None, image_sizes: Optional[torch.Tensor] = None, **kwargs, ) -> CausalLMOutputWithPast: batch_size, seq_len = input_ids.shape if self.config.dlm_paradigm == 'bidirectional' or self.config.dlm_paradigm == 'autoregressive': if labels is not None and torch.rand(1) < self.config.random_length_prob: random_length = torch.randint(2, input_ids.shape[1] + 1, (1,)) input_ids = input_ids[:, :random_length] labels = labels[:, :random_length] if attention_mask is not None: attention_mask = attention_mask[:, :random_length] if position_ids is not None: position_ids = position_ids[:, :random_length] if loss_mask is not None: loss_mask = loss_mask[:, :random_length] elif self.config.dlm_paradigm in ['block_diff', 'sbd_block_diff']: if labels is not None and block_size is None: if torch.rand(1) < self.config.random_length_prob: block_size = torch.randint(1, 8, (1,)).item() * 4 ## [4, 32] divisible by 4 else: block_size = self.config.block_size else: raise ValueError(f"Unknown dLM paradigm: {self.config.dlm_paradigm}") if labels is not None and self.config.dlm_paradigm != 'autoregressive': if masked_indices is not None: # assert p_mask is not None if loss_mask is not None: masked_indices[loss_mask == 0] = 0 noisy_inputs = torch.where(masked_indices, self.mask_token_id, input_ids) else: if self.config.complementary_mask: loss_mask = (labels != -100) noisy_inputs, masked_indices, p_mask = self.forward_process_complementary(input_ids, eps=eps, block_size=block_size, loss_mask=loss_mask) else: if self.config.tok_mask_half_life_ratio is not None: noisy_inputs, masked_indices, p_mask = self.forward_process_exp(input_ids, eps=eps, block_size=block_size, half_life_ratio=self.config.tok_mask_half_life_ratio, loss_mask=loss_mask) else: noisy_inputs, masked_indices, p_mask = self.forward_process(input_ids, eps=eps, block_size=block_size, loss_mask=loss_mask) else: noisy_inputs = input_ids masked_indices = None p_mask = None if self.config.dlm_paradigm in ['block_diff', 'sbd_block_diff']: for layer in self.encoder.layers: if hasattr(layer.self_attn, 'set_attention_mode'): layer.self_attn.set_attention_mode(self.config.dlm_paradigm, block_size=block_size) input_ids_len = noisy_inputs.shape[1] if labels is not None and self.config.dlm_paradigm in ['block_diff', 'sbd_block_diff']: if position_ids is None: position_ids = torch.arange(input_ids_len, device=noisy_inputs.device).unsqueeze(0) if self.config.complementary_mask: noisy_inputs = torch.cat([noisy_inputs, torch.cat([input_ids, input_ids], dim=0)], dim=1) else: noisy_inputs = torch.cat([noisy_inputs, input_ids], dim=1) if block_diff_ppl: if position_ids is None: position_ids = torch.arange(input_ids_len // 2, device=noisy_inputs.device).unsqueeze(0) # ── Vision: replace IMG pad embeddings with image features ──────── if pixel_values is not None and image_sizes is not None: inputs_embeds = self._embed_with_vision(noisy_inputs, pixel_values, image_sizes) enc_out = self.encoder( past_key_values=past_key_values, inputs_embeds=inputs_embeds, attention_mask=attention_mask, position_ids=position_ids, is_training=(labels is not None) or (block_diff_ppl), **kwargs, ) elif self.training and pixel_values is None and not self._is_vision_frozen(): vt = self.encoder.vision_tower _p = vt.patch_size _merge = getattr(self.config, "spatial_merge_size", 2) _side = _p * _merge _c = getattr(vt.config, "num_channels", 3) _dtype = next(vt.parameters()).dtype dummy_pixel = torch.zeros( 1, _c, _side, _side, dtype=_dtype, device=noisy_inputs.device, ) dummy_image_sizes = torch.tensor( [(int(_side), int(_side))], dtype=torch.long, device=noisy_inputs.device, ) dummy_features = self.get_image_features(dummy_pixel, dummy_image_sizes) inputs_embeds = self.encoder.embed_tokens(noisy_inputs) inputs_embeds = inputs_embeds + dummy_features.sum() * 0 enc_out = self.encoder( past_key_values=past_key_values, inputs_embeds=inputs_embeds, attention_mask=attention_mask, position_ids=position_ids, is_training=(labels is not None) or (block_diff_ppl), **kwargs, ) else: enc_out = self.encoder( past_key_values=past_key_values, input_ids=noisy_inputs, attention_mask=attention_mask, position_ids=position_ids, is_training=(labels is not None) or (block_diff_ppl), **kwargs, ) if output_last_hidden_states_only: return BaseModelOutput(last_hidden_state=enc_out.last_hidden_state) logits = self.diffusion_head(enc_out.last_hidden_state) # (batch, len_B, vocab) causal_logits = None if labels is not None and self.config.dlm_paradigm in ['block_diff', 'sbd_block_diff']: if self.config.dlm_paradigm == 'sbd_block_diff': causal_logits = logits[:, input_ids_len:] else: causal_logits = None logits = logits[:, :input_ids_len] loss = None if getattr(self.config, 'complementary_mask', False) and self.config.dlm_paradigm == 'sbd_block_diff': _raw_nib = kwargs.get('num_items_in_batch', None) kwargs = {**kwargs, 'num_items_in_batch': 2 * kwargs.get('num_items_in_batch', 1)} if self.training and (not hasattr(self, '_nib_logged') or not self._nib_logged): import torch.distributed as dist _rank = dist.get_rank() if dist.is_initialized() else 0 if _rank == 0: print(f"[DEBUG-NIB] raw num_items_in_batch from Trainer: {_raw_nib}, " f"after 2x: {kwargs['num_items_in_batch']}, " f"labels non-(-100): {(labels != -100).sum().item() if labels is not None else 'N/A'}, " f"batch_size={input_ids.shape[0]}, seq_len={input_ids.shape[1]}", flush=True) self._nib_logged = True if labels is not None and not skip_loss: if self.config.dlm_paradigm == 'autoregressive': shift_logits = logits[..., :-1, :].contiguous() shift_labels = labels[..., 1:].contiguous() if loss_mask is None: loss_fct = CrossEntropyLoss() shift_logits = shift_logits.view(-1, shift_logits.size(-1)) shift_labels = shift_labels.view(-1) loss = loss_fct(shift_logits, shift_labels) else: loss_mask = loss_mask[..., 1:].contiguous() loss_fct = CrossEntropyLoss(reduction='none') shift_logits = shift_logits.view(-1, shift_logits.size(-1)) shift_labels = shift_labels.view(-1) shift_labels = shift_labels.to(shift_logits.device) token_losses = loss_fct(shift_logits, shift_labels) flat_loss_mask = loss_mask.reshape(-1) loss = token_losses[flat_loss_mask == 1].sum() / flat_loss_mask.sum() else: # Handle DREAM vs LLADA style losses if hasattr(self.config, 'dlm_type') and self.config.dlm_type == 'dream': logits = logits[..., :-1, :].contiguous() labels = labels[..., 1:].contiguous() masked_indices = masked_indices[:, 1:] if p_mask is not None: p_mask = p_mask[:, 1:] if self.config.ada_perm_ratio_per_block is not None: # Only compute loss for the top ada_perm_ratio_per_block tokens by confidence within each block block_size = self.config.block_size batch_size, seq_len = masked_indices.shape num_blocks = seq_len // block_size # Get the max logit (confidence) for each position confidence = logits.max(dim=-1).values.detach() # (batch_size, seq_len) # Create a mask for tokens to include in loss selected_mask = torch.zeros_like(masked_indices, dtype=torch.bool) for blk in range(num_blocks): start = blk * block_size end = min((blk + 1) * block_size, seq_len) # Get masked indices within this block block_masked = masked_indices[:, start:end] # (batch_size, block_len) block_confidence = confidence[:, start:end] # (batch_size, block_len) for b in range(batch_size): # Get positions that are masked in this block for this batch masked_positions = torch.where(block_masked[b])[0] num_masked = len(masked_positions) if num_masked > 0: # Number of tokens to keep (top by confidence) k = min(max(1, int(block_size * self.config.ada_perm_ratio_per_block)), num_masked) # Get confidence values for masked positions masked_confidence = block_confidence[b, masked_positions] # Get indices of top-k confident tokens _, topk_indices = torch.topk(masked_confidence, k) selected_positions = masked_positions[topk_indices] # Mark these positions in the selected mask selected_mask[b, start + selected_positions] = True # Calculate loss only for selected positions token_loss = torch.nn.functional.cross_entropy( logits[selected_mask], labels[selected_mask], reduction='none' ) / p_mask[selected_mask] num_mask_tokens = selected_mask.sum() elif getattr(self.config, 'complementary_mask', False): token_loss = self.mdm_loss_function( logits=logits[masked_indices], labels=torch.cat([labels, labels], dim=0)[masked_indices], vocab_size=self.config.vocab_size, **kwargs ) num_mask_tokens = masked_indices.sum() else: # Calculate token-wise cross entropy loss for masked positions in B token_loss = torch.nn.functional.cross_entropy( logits[masked_indices], labels[masked_indices], reduction='none' ) / p_mask[masked_indices] num_mask_tokens = masked_indices.sum() if self.config.global_loss_avg: loss = token_loss.sum() else: loss = token_loss.sum() / num_mask_tokens if self.config.ada_dlm_loss_ratio is not None: assert self.current_iter_ratio is not None assert self.config.dlm_loss_weight is not None dlm_loss_weight = min(self.config.dlm_loss_weight, self.current_iter_ratio / self.config.ada_dlm_loss_ratio * self.config.dlm_loss_weight) loss = dlm_loss_weight * loss elif self.config.dlm_loss_weight is not None: loss = self.config.dlm_loss_weight * loss if self.config.dlm_paradigm == 'sbd_block_diff': if getattr(self.config, 'complementary_mask', False): ar_loss = self.causal_loss_function( logits=causal_logits[:logits.shape[0] // 2, :], labels=labels, vocab_size=self.config.vocab_size, **kwargs ) _diff_val = loss.detach().item() _ar_val = ar_loss.detach().item() if not hasattr(self, '_loss_accum_count'): self._loss_diff_accum = 0.0 self._loss_ar_accum = 0.0 self._loss_accum_count = 0 self._loss_diff_accum += _diff_val self._loss_ar_accum += _ar_val self._loss_accum_count += 1 self.loss_diffusion = self._loss_diff_accum self.loss_ar = self._loss_ar_accum loss = loss + ar_loss else: causal_logits = causal_logits[..., :-1, :].contiguous() causal_logits = causal_logits.view(-1, causal_logits.size(-1)) if hasattr(self.config, 'dlm_type') and self.config.dlm_type == 'dream': causal_labels = labels.view(-1) else: causal_labels = labels[..., 1:].contiguous().view(-1) if self.config.global_loss_avg: loss_fct = CrossEntropyLoss(reduction='sum') ar_loss = loss_fct(causal_logits, causal_labels) self.loss_diffusion = loss.detach().item() / num_mask_tokens self.loss_ar = ar_loss.detach().item() / seq_len loss = loss + self.config.ar_loss_weight * ar_loss else: loss_fct = CrossEntropyLoss() ar_loss = loss_fct(causal_logits, causal_labels) self.loss_diffusion = loss.detach().item() self.loss_ar = ar_loss.detach().item() loss = loss + self.config.ar_loss_weight * ar_loss # if self.config.global_loss_avg: # if self.config.dlm_paradigm == 'sbd_block_diff': # loss = (loss, num_mask_tokens + int(self.config.ar_loss_weight * seq_len)) # else: # loss = (loss, num_mask_tokens) return NemotronLabsDiffusionVLMOutputWithPast( loss=loss if not is_teacher else logits, logits=logits, causal_logits=causal_logits, past_key_values=enc_out.past_key_values, hidden_states=None, attentions=None, ) def generate(self, prompt_ids, max_new_tokens, steps, block_length, shift_logits, threshold, causal_context=True, temperature=0, pixel_values=None, image_sizes=None, eos_token_id=None): out_ids, nfe = generate_with_prefix_cache_block_diff( model=self, prompt=prompt_ids, gen_length=max_new_tokens, steps=steps, block_length=block_length, remasking="low_confidence", temperature=temperature, mask_id=self.mask_token_id, threshold=threshold, shift_logits=shift_logits, neg_entropy=False, causal_context=causal_context, pixel_values=pixel_values, image_sizes=image_sizes, eos_token_id=eos_token_id, ) return out_ids, nfe @torch.no_grad() def sbd_inference_diffusion_quadratic( self, clean_input_ids: Optional[torch.Tensor], draft_input_ids: torch.Tensor, block_length: int, draft_only: bool = False, past_key_values: Optional[Cache] = None, use_cache: bool = False, pixel_values: Optional[torch.FloatTensor] = None, image_sizes: Optional[torch.Tensor] = None, ): enc_config = self.encoder.config enc_config.use_sbd_objective = True enc_config.block_length = block_length if draft_only: assert clean_input_ids is not None if use_cache and past_key_values is None: past_key_values = DynamicCache() enc_config.self_spec_inference_mode = "default" input_ids = torch.cat([clean_input_ids, draft_input_ids], dim=-1) if pixel_values is not None and image_sizes is not None: inputs_embeds = self._embed_with_vision(input_ids, pixel_values, image_sizes) outputs = self.encoder( inputs_embeds=inputs_embeds, position_ids=None, past_key_values=past_key_values, use_cache=use_cache, is_training=False, ) else: outputs = self.encoder( input_ids=input_ids, position_ids=None, past_key_values=past_key_values, use_cache=use_cache, is_training=False, ) hidden_states = outputs.last_hidden_state logits = self.diffusion_head(hidden_states) past_key_values = getattr(outputs, "past_key_values", None) if use_cache and past_key_values is not None: _crop_dynamic_cache(past_key_values, clean_input_ids.shape[1]) return logits, past_key_values else: enc_config.self_spec_inference_mode = "quadratic" draft_len = block_length * (block_length + 1) draft_input_ids = torch.cat( [ draft_input_ids.view(-1, block_length, 1), torch.full( (draft_input_ids.shape[0], block_length, block_length), fill_value=self.config.mask_token_id, device=draft_input_ids.device, ), ], dim=-1, ).view(-1, draft_len) if use_cache: assert past_key_values is not None, ( "Past key values should be provided when using cache, e.g. run draft_only=True first." ) assert clean_input_ids is None, ( "Clean input ids should already be in cache, thus none should be provided." ) clean_len = past_key_values.get_seq_length() input_ids = draft_input_ids else: clean_len = clean_input_ids.shape[1] input_ids = torch.cat([clean_input_ids, draft_input_ids], dim=-1) per_block_position_ids = torch.arange( clean_len, clean_len + block_length + 1, device=draft_input_ids.device )[None,].repeat(block_length, 1) per_block_position_ids += torch.arange(block_length, device=draft_input_ids.device).view(-1, 1) if use_cache: position_ids = per_block_position_ids.view(-1)[None,] else: clean_position_ids = torch.arange(clean_len, device=draft_input_ids.device) position_ids = torch.cat([clean_position_ids, per_block_position_ids.view(-1)], dim=-1)[None,] if pixel_values is not None and image_sizes is not None and not use_cache: inputs_embeds = self._embed_with_vision(input_ids, pixel_values, image_sizes) outputs = self.encoder( inputs_embeds=inputs_embeds, position_ids=position_ids, past_key_values=past_key_values, use_cache=use_cache, is_training=False, ) else: outputs = self.encoder( input_ids=input_ids, position_ids=position_ids, past_key_values=past_key_values, use_cache=use_cache, is_training=False, ) hidden_states = outputs.last_hidden_state logits = self.diffusion_head(hidden_states) past_key_values = getattr(outputs, "past_key_values", None) if use_cache and past_key_values is not None: _extract_draft_kv_cache(past_key_values, clean_len, block_length) return logits, past_key_values @torch.no_grad() def self_spec_generate( self, prompt_ids: torch.Tensor, max_new_tokens: int = 128, steps: int = 128, block_length: int = 16, ar_mix_weight: Optional[float] = None, temperature: float = 0.0, mask_token_id: Optional[int] = None, eos_token_id: Optional[int] = None, pixel_values: Optional[torch.FloatTensor] = None, image_sizes: Optional[torch.Tensor] = None, ): self.config.use_sbd_objective = True self.config.dlm_paradigm = "sbd" if prompt_ids.shape[0] != 1: raise ValueError("Self speculation quadratic decoding currently requires batch_size == 1") token_mask_id = mask_token_id if mask_token_id is not None else self.config.mask_token_id if eos_token_id is None: eos_token_id = getattr(self.config, "eos_token_id", None) x = torch.full( (1, prompt_ids.shape[1] + max_new_tokens + block_length * 2), token_mask_id, dtype=torch.long, device=prompt_ids.device, ) x[:, : prompt_ids.shape[1]] = prompt_ids.clone() if max_new_tokens % block_length != 0: raise ValueError("max_new_tokens must be divisible by block_length") num_blocks = max_new_tokens // block_length if steps % num_blocks != 0: raise ValueError("steps must be divisible by (max_new_tokens // block_length)") prompt_len = prompt_ids.shape[1] nfe = 0 nfe += 1 logits, past_key_values = self.sbd_inference_diffusion_quadratic( clean_input_ids=x[:, :prompt_len], draft_input_ids=x[:, prompt_len : prompt_len + block_length], block_length=block_length, draft_only=True, use_cache=True, pixel_values=pixel_values, image_sizes=image_sizes, ) logits_proposal = logits[:, prompt_len - 1 : prompt_len + block_length] logits_proposal[:, 1] = logits_proposal[:, 0] logits_proposal = logits_proposal[:, 1:] x0_proposal = torch.argmax(logits_proposal, dim=-1) x[:, prompt_len : prompt_len + block_length] = x0_proposal total_accept_token = 0 while True: nfe += 1 block_start = prompt_len + total_accept_token block_end = block_start + block_length draft_input_ids = x[:, block_start:block_end] logits, past_key_values = self.sbd_inference_diffusion_quadratic( clean_input_ids=None, draft_input_ids=draft_input_ids, block_length=block_length, draft_only=False, past_key_values=past_key_values, use_cache=True, pixel_values=pixel_values, image_sizes=image_sizes, ) useful_token_logits = logits.view(1, block_length, block_length + 1, -1) if ar_mix_weight is None: useful_token_logits[:, :, 1] = useful_token_logits[:, :, 0] else: if not (0.0 <= ar_mix_weight <= 1.0): raise ValueError("ar_mix_weight must be between 0 and 1") mix_logits = useful_token_logits[:, :, 0] * ar_mix_weight + useful_token_logits[:, :, 1] * (1 - ar_mix_weight) useful_token_logits[:, :, 0] = mix_logits useful_token_logits[:, :, 1] = mix_logits if temperature > 0: useful_token_logits = useful_token_logits / temperature useful_token_pred = torch.argmax(useful_token_logits, dim=-1) new_draft_input_ids = useful_token_pred[:, 0, 1:] accept_cnt = 1 while accept_cnt < block_length: if useful_token_pred[:, accept_cnt - 1, 0].item() != draft_input_ids[:, accept_cnt].item(): break new_draft_input_ids = useful_token_pred[:, accept_cnt, 1:] accept_cnt += 1 x[:, block_start : block_start + accept_cnt] = draft_input_ids[:, :accept_cnt] # EoS early stopping if eos_token_id is not None: accepted = x[0, block_start : block_start + accept_cnt] eos_positions = (accepted == eos_token_id).nonzero(as_tuple=True)[0] if len(eos_positions) > 0: first_eos_rel = eos_positions[0].item() total_accept_token += first_eos_rel + 1 output_end = prompt_len + total_accept_token return x[:, :output_end], nfe x[:, block_start + accept_cnt : block_start + accept_cnt + block_length] = new_draft_input_ids _crop_dynamic_cache(past_key_values, block_start + accept_cnt) total_accept_token += accept_cnt if total_accept_token >= max_new_tokens: break return x[:, : -(block_length * 2)], nfe __all__ = ["NemotronLabsDiffusionVLMModel", "NemotronLabsDiffusionVLMFlexAttention"]