diff --git "a/diffretriever_trainable.py" "b/diffretriever_trainable.py" new file mode 100644--- /dev/null +++ "b/diffretriever_trainable.py" @@ -0,0 +1,2465 @@ +""" +DiffRetriever — fine-tunable variant (paper §3.4 supervised fine-tuning). + +Contrastively fine-tunable wrapper around Dream / LLaDA backbones, +sharing the masked-position prediction encoder with the zero-shot +DreamRetriever / LLaDA2Retriever classes. Supports Dream, LLaDA1/1.5, +LLaDA2 backbones via backbone_adapters.py. + +Encoding: + steps=1 (fast): Single forward pass over [prefix][text][suffix][MASK×K]. + repr_hidden[:, s, :] from MASK position s; quotation_emb from the + closing " token. + steps>1 (rich): Iterative denoising loop with mixed representations. + At each step, gen positions already decoded in a prior step contribute + their frozen hidden state (no gradient); positions still MASK contribute + the current step's hidden (with gradient). + Uniform LLaDA-style unmasking schedule: n_per_step = K // n_steps + tokens are decoded per step. + +Loss (training): + K > 1: ColBERT MaxSim InfoNCE over all K mixed vectors. + steps>1: loss computed only at the final denoising step so gradient + flows from the fully-contextualized representation without dilution. + If progressive_step_weight > 0, retrieval loss is also applied at + each intermediate step with linearly increasing weight (t/T), + giving direct supervision to early denoising steps. + Optional: sparse InfoNCE + FLOPS L1 regularization. + +At eval time the same encode() output (repr_hidden, quotation_emb, +sparse_acts) supports all zero-shot retrieval modes: single_dense, +multi_dense_k*, multi-denoise-step variants, and sparse versions of each. +""" + +import re +import math +import torch +import torch.nn as nn +import torch.nn.functional as F +from typing import Dict, List, Optional, Tuple +from pathlib import Path +import logging + +from .sparse_utils import filter_sparse +from .backbone_adapters import get_adapter, BackboneAdapter + +logger = logging.getLogger(__name__) + +_NUM_WORDS = ['one', 'two', 'three', 'four', 'five', + 'six', 'seven', 'eight', 'nine', 'ten'] + + +class TrainableDiffusionRetriever(nn.Module): + """ + Trainable Diffusion Retriever. + + Build via from_backbone(); use tokenize() in your data collator, then pass + pre-tokenized tensors to forward() during training. + """ + + def __init__( + self, + backbone: nn.Module, + tokenizer, + mask_token_id: int, + hidden_size: int, + query_prefix_ids: List[int], + query_suffix_ids: List[int], + passage_prefix_ids: List[int], + passage_suffix_ids: List[int], + max_length: int = 512, + n_gen_tokens: int = 4, + n_gen_q_tokens: Optional[int] = None, # asymmetric K_q (defaults to n_gen_tokens) + n_gen_p_tokens: Optional[int] = None, # asymmetric K_p (defaults to n_gen_tokens) + temperature: float = 0.01, + num_denoise_steps: int = 4, + sparse_weight: float = 1.0, + normalize: bool = True, + flash_attn: bool = False, + use_eos: bool = False, + # K-adapter (joint training of K-router + retriever). When enabled, + # the model encodes at K_max and at training time computes per-cell + # InfoNCE losses; a tiny MLP head maps the query's quotation_emb to a + # softmax over K-cell choices and is supervised by KL(teacher || head) + # where teacher = softmax(-per_cell_loss / τ_T). See compute_loss(). + use_k_adapter: bool = False, + adapter_weight: float = 1.0, + teacher_temperature: float = 1.0, + k_adapter_options: Optional[Tuple[int, ...]] = None, + # Two-stage K-pre-encoder (true encoding savings via per-item K). + # When enabled, a tiny MLP runs over the embedding-layer output and + # predicts K_q (for queries) and K_p (for passages) per item. Each + # item's input is sliced to its predicted K (removing K_max-K MASK + # tokens), the batch is padded to max-K-in-batch, and main encoder + # runs at variable length. See _forward_with_k_pre_encoder(). + use_k_pre_encoder: bool = False, + gumbel_temperature: float = 1.0, + k_cost_lambda: float = 0.001, + k_pre_encoder_options: Optional[Tuple[int, ...]] = None, + ): + super().__init__() + self.backbone = backbone + self.tokenizer = tokenizer + self.mask_token_id = mask_token_id + self.hidden_size = hidden_size + + # Prompt IDs (cached lists, not parameters) + self.register_buffer('_dummy', torch.zeros(1)) # for .device + self._query_prefix_ids = list(query_prefix_ids) + self._query_suffix_ids = list(query_suffix_ids) + self._passage_prefix_ids = list(passage_prefix_ids) + self._passage_suffix_ids = list(passage_suffix_ids) + + self.max_length = max_length + # Per-side K. When the user passes only `n_gen_tokens`, both sides + # use it (symmetric — original behaviour). Pass `n_gen_q_tokens` / + # `n_gen_p_tokens` to train asymmetric (e.g. K_q=4, K_p=16). + self.n_gen_tokens = n_gen_tokens # legacy, kept for back-compat + self.n_gen_q_tokens = n_gen_q_tokens if n_gen_q_tokens is not None else n_gen_tokens + self.n_gen_p_tokens = n_gen_p_tokens if n_gen_p_tokens is not None else n_gen_tokens + self.use_eos = use_eos # kept for config compat; _n_tail is authoritative + # _n_tail is the structural tail (quote + chat_end + EOS); it's the + # same regardless of K_q or K_p as long as at least one side has K>0. + max_k = max(self.n_gen_q_tokens, self.n_gen_p_tokens) + self._n_tail = 3 if max_k > 0 else 0 + self._tail_ids: List[int] = self._build_tail_ids() # cached at init + self.temperature = temperature + # Denoise step count is capped by the larger K (the side with more + # mask tokens has more positions to iteratively decode). + self.num_denoise_steps = min(num_denoise_steps, max_k) + self.sparse_weight = sparse_weight + self.dense_weight = 1.0 + self.normalize = normalize + self.flash_attn = flash_attn + self.model_type = 'unknown' # set by from_backbone() / load() + self.adapter: Optional[BackboneAdapter] = None # set by factory methods + self._hook_registered = False # set by _setup_hook() + + # Auxiliary losses (activated by setting weight > 0) + self.denoising_weight = 0.0 # diffusion-native masked text denoising + self.diversity_weight = 0.0 # explicit multi-vector diversity + self.denoise_mask_ratio = 0.15 # fraction of text tokens to mask for denoising + self.progressive_step_weight = 0.0 # progressive step supervision (multi-step only) + self.use_fresh_final = False # "fresh" ablation: use current hidden for ALL K at final step + self.soft_denoising = False # soft-token multi-step: differentiable embeddings instead of hard tokens + self.soft_temperature = 1.0 # temperature for softmax in soft-token denoising + self.corruption_rate = 0.0 # max text corruption rate for denoising-conditioned training (0 = off) + self.debug_dense_metrics = False + self.debug_compare_hidden_once = False + self._debug_hidden_compared = False + + # ── K-adapter (joint K-router + retriever) ───────────────────────── + self.use_k_adapter = bool(use_k_adapter) + self.adapter_weight = float(adapter_weight) + self.teacher_temperature = float(teacher_temperature) + # Default K options: factors of K_max that fall within the active K + # range. E.g., K_max=16 → {1,2,4,8,16}; K_max=4 → {1,2,4}. + if k_adapter_options is None: + k_max_active = max(self.n_gen_q_tokens, self.n_gen_p_tokens) + base = (1, 2, 4, 8, 16) + self.k_adapter_options: Tuple[int, ...] = tuple( + k for k in base if k <= k_max_active + ) + else: + self.k_adapter_options = tuple(int(k) for k in k_adapter_options) + self.n_K = len(self.k_adapter_options) + + if self.use_k_adapter: + # Tiny MLP: hidden_size → hidden_size//4 → n_K * n_K. + # Trained jointly with the backbone via the per-cell loss teacher. + # Kept in fp32 to match the .float() conversion of representations + # in encode(); HF Trainer's bf16 autocast handles mixed precision + # at the call site without needing a manual dtype cast here. + head_dim = max(64, hidden_size // 4) + self.k_adapter = nn.Sequential( + nn.Linear(hidden_size, head_dim), + nn.GELU(), + nn.Linear(head_dim, self.n_K * self.n_K), + ) + logger.info( + f"KAdapter enabled: K options={self.k_adapter_options}, " + f"adapter_weight={self.adapter_weight}, " + f"teacher_temperature={self.teacher_temperature}" + ) + else: + self.k_adapter = None + + # ── K-pre-encoder (two-stage encoding) ───────────────────────────── + self.use_k_pre_encoder = bool(use_k_pre_encoder) + self.gumbel_temperature = float(gumbel_temperature) + self.k_cost_lambda = float(k_cost_lambda) + if k_pre_encoder_options is None: + k_max_active_pe = max(self.n_gen_q_tokens, self.n_gen_p_tokens) + base_pe = (1, 2, 4, 8, 16) + self.k_pre_encoder_options: Tuple[int, ...] = tuple( + k for k in base_pe if k <= k_max_active_pe + ) + else: + self.k_pre_encoder_options = tuple(int(k) for k in k_pre_encoder_options) + self.n_K_pe = len(self.k_pre_encoder_options) + # Buffer of K values for differentiable expected-K (cost regularizer). + self.register_buffer( + 'k_pe_options_tensor', + torch.tensor(self.k_pre_encoder_options, dtype=torch.float32), + persistent=False, + ) + if self.use_k_pre_encoder: + head_dim_pe = max(64, hidden_size // 4) + # Separate q + p heads — query and passage have different optimal K + # distributions and texts differ enough to justify dedicated heads. + # Both heads operate on the (shared) embedding layer's output — + # mean-pooled over real tokens. Cheap (~1% of full forward). + self.k_pre_encoder_q = nn.Sequential( + nn.Linear(hidden_size, head_dim_pe), + nn.GELU(), + nn.Linear(head_dim_pe, self.n_K_pe), + ) + self.k_pre_encoder_p = nn.Sequential( + nn.Linear(hidden_size, head_dim_pe), + nn.GELU(), + nn.Linear(head_dim_pe, self.n_K_pe), + ) + logger.info( + f"K pre-encoder enabled: K options={self.k_pre_encoder_options}, " + f"gumbel_τ={self.gumbel_temperature}, " + f"cost_λ={self.k_cost_lambda}" + ) + else: + self.k_pre_encoder_q = None + self.k_pre_encoder_p = None + + def _k(self, is_query: bool) -> int: + """Return the appropriate per-side K. Lets the rest of the code + say `K = self._k(is_query)` and stay correct under both symmetric + (K_q == K_p) and asymmetric configurations.""" + return self.n_gen_q_tokens if is_query else self.n_gen_p_tokens + + # ---------------------------------------------------------------- + # Build from pretrained backbone + # ---------------------------------------------------------------- + + @classmethod + def from_backbone( + cls, + model_name: str, + model_type: str, + query_prompt: str, + passage_prompt: str, + max_length: int = 512, + n_gen_tokens: int = 4, + n_gen_q_tokens: Optional[int] = None, + n_gen_p_tokens: Optional[int] = None, + temperature: float = 0.01, + num_denoise_steps: int = 4, + sparse_weight: float = 1.0, + normalize: bool = True, + gradient_checkpointing: bool = True, + lora_rank: int = 0, + lora_alpha: int = 64, + lora_dropout: float = 0.05, + device_map=None, + use_eos: bool = False, + disable_hidden_hook: bool = False, + # K-adapter passthrough + use_k_adapter: bool = False, + adapter_weight: float = 1.0, + teacher_temperature: float = 1.0, + k_adapter_options: Optional[Tuple[int, ...]] = None, + # K pre-encoder passthrough + use_k_pre_encoder: bool = False, + gumbel_temperature: float = 1.0, + k_cost_lambda: float = 0.001, + k_pre_encoder_options: Optional[Tuple[int, ...]] = None, + ) -> 'TrainableDiffusionRetriever': + from transformers import AutoTokenizer + + adapter = get_adapter(model_type) + + tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True) + tokenizer.padding_side = 'left' + + backbone = adapter.load_backbone(model_name, device_map=device_map) + + if lora_rank > 0: + from peft import get_peft_model + lora_cfg = adapter.get_lora_config(lora_rank, lora_alpha, lora_dropout) + backbone = get_peft_model(backbone, lora_cfg) + backbone.print_trainable_parameters() + + if gradient_checkpointing: + backbone.enable_input_require_grads() + adapter.enable_gradient_checkpointing(backbone) + + mask_token_id = adapter.mask_token_id + logger.info(f"Mask token: {model_type} → {mask_token_id}") + + # Resolve per-side K (default to symmetric n_gen_tokens). The + # YAML adaptation ('the relevant document is one word' → 'four + # words') uses the per-side K so the natural-language prompt + # matches the actual MASK count the model will see. + k_q = n_gen_q_tokens if n_gen_q_tokens is not None else n_gen_tokens + k_p = n_gen_p_tokens if n_gen_p_tokens is not None else n_gen_tokens + + # Build prompt token IDs (per-side adaptation) + q_yaml = cls._load_yaml(query_prompt) + p_yaml = cls._load_yaml(passage_prompt) + q_yaml = dict(q_yaml) + q_yaml['user_suffix'] = cls._adapt_for_k(q_yaml.get('user_suffix', ''), k_q) + q_yaml['assistant_prefix'] = cls._adapt_for_k(q_yaml.get('assistant_prefix', ''), k_q) + p_yaml = dict(p_yaml) + p_yaml['user_suffix'] = cls._adapt_for_k(p_yaml.get('user_suffix', ''), k_p) + p_yaml['assistant_prefix'] = cls._adapt_for_k(p_yaml.get('assistant_prefix', ''), k_p) + + q_prefix_ids, q_suffix_ids = cls._build_prompt_ids(tokenizer, q_yaml) + p_prefix_ids, p_suffix_ids = cls._build_prompt_ids(tokenizer, p_yaml) + + logger.info(f"Query prompt: {len(q_prefix_ids)} prefix + {len(q_suffix_ids)} suffix tokens") + logger.info(f"Passage prompt: {len(p_prefix_ids)} prefix + {len(p_suffix_ids)} suffix tokens") + if k_q == k_p: + logger.info(f"n_gen_tokens={k_q} (symmetric), num_denoise_steps={num_denoise_steps}") + else: + logger.info(f"n_gen_q_tokens={k_q}, n_gen_p_tokens={k_p} (asymmetric), " + f"num_denoise_steps={num_denoise_steps}") + if sparse_weight > 0: + logger.info(f"Sparse loss: sparse_weight={sparse_weight}") + + if num_denoise_steps is None: + num_denoise_steps = max(k_q, k_p) + + hidden_size = backbone.config.hidden_size + + model = cls( + backbone=backbone, + tokenizer=tokenizer, + mask_token_id=mask_token_id, + hidden_size=hidden_size, + query_prefix_ids=q_prefix_ids, + query_suffix_ids=q_suffix_ids, + passage_prefix_ids=p_prefix_ids, + passage_suffix_ids=p_suffix_ids, + max_length=max_length, + n_gen_tokens=n_gen_tokens, + n_gen_q_tokens=k_q, + n_gen_p_tokens=k_p, + temperature=temperature, + num_denoise_steps=num_denoise_steps, + sparse_weight=sparse_weight, + normalize=normalize, + flash_attn=adapter.flash_attn, + use_eos=use_eos, + use_k_adapter=use_k_adapter, + adapter_weight=adapter_weight, + teacher_temperature=teacher_temperature, + k_adapter_options=k_adapter_options, + use_k_pre_encoder=use_k_pre_encoder, + gumbel_temperature=gumbel_temperature, + k_cost_lambda=k_cost_lambda, + k_pre_encoder_options=k_pre_encoder_options, + ) + model.model_type = model_type + model.adapter = adapter + model.lora_rank = lora_rank + model.lora_alpha = lora_alpha + + # Hook on output projection for efficient hidden state extraction + model._last_hidden: Dict[str, torch.Tensor] = {} + if disable_hidden_hook: + logger.info(f"{model_type}: hidden hook disabled by flag") + model._hook_registered = False + else: + model._hook_registered = adapter.register_hidden_hook( + backbone, model._last_hidden) + if not model._hook_registered: + logger.info(f"{model_type}: no hook — will use output_hidden_states=True") + + return model + + # ---------------------------------------------------------------- + # Prompt helpers (mirrors DreamRetriever) + # ---------------------------------------------------------------- + + @staticmethod + def _load_yaml(path: str) -> dict: + import yaml + p = Path(path) + if not p.exists(): + raise FileNotFoundError(f"Prompt YAML not found: {path}") + return yaml.safe_load(p.read_text()) + + @staticmethod + def _adapt_for_k(text: str, k: int) -> str: + if k <= 1 or not text: + return text + count = _NUM_WORDS[k - 1] if k <= len(_NUM_WORDS) else str(k) + result = re.sub( + r'\b(?:' + '|'.join(_NUM_WORDS) + r')\b(\s+words?)', + lambda m: f'{count} words', text, + ) + return re.sub(r'\bword is\b', 'words are', result) + + @staticmethod + def _build_prompt_ids(tokenizer, yaml_dict: dict) -> Tuple[List[int], List[int]]: + system = yaml_dict.get('system', '') + user_prefix = yaml_dict.get('user_prefix', '') + user_suffix = yaml_dict.get('user_suffix', '') + assistant_prefix = yaml_dict.get('assistant_prefix', '') + + SENTINEL = "XSENTINELX" + messages = [] + if system: + messages.append({"role": "system", "content": system}) + messages.append({"role": "user", "content": user_prefix + SENTINEL + user_suffix}) + + full_str = tokenizer.apply_chat_template( + messages, tokenize=False, add_generation_prompt=True + ) + pre_str, post_str = full_str.split(SENTINEL, 1) + prefix_ids = tokenizer.encode(pre_str, add_special_tokens=False) + suffix_ids = tokenizer.encode(post_str + assistant_prefix, add_special_tokens=False) + return prefix_ids, suffix_ids + + # ---------------------------------------------------------------- + # Structural tail tokens (matches zero-shot dream/llada retrievers) + # ---------------------------------------------------------------- + + def _build_tail_ids(self) -> List[int]: + """Build structural tail tokens: [quote, chat_end, EOS]. + + Matches the zero-shot DreamRetriever._build_tail_ids and + LLaDA2Retriever._build_tail_ids so that training sees the same + sequence layout as zero-shot inference. + """ + if self._n_tail <= 0: + return [] + + eos_id = self.tokenizer.eos_token_id + tail: List[int] = [] + + # Closing quote " + quote_ids = self.tokenizer.encode('"', add_special_tokens=False) + if len(quote_ids) == 1: + tail.append(quote_ids[0]) + else: + logger.warning('quote \'"\' tokenises to %d tokens; using EOS in slot 0', len(quote_ids)) + tail.append(eos_id) + + # Chat-template end token (im_end for Dream/Qwen, eot_id for LLaDA) + _KNOWN_EOT = {'<|im_end|>': None, '<|eot_id|>': 126348} # fallback for LLaDA + for tok_str, fallback_id in _KNOWN_EOT.items(): + tid = self.tokenizer.convert_tokens_to_ids(tok_str) + unk = getattr(self.tokenizer, 'unk_token_id', None) + if tid is not None and tid >= 0 and tid != unk and tid != eos_id: + tail.append(tid) + break + elif fallback_id is not None and fallback_id != eos_id: + tail.append(fallback_id) + break + + # Fill remaining slots with EOS + while len(tail) < self._n_tail: + tail.append(eos_id) + return tail[:self._n_tail] + + # ---------------------------------------------------------------- + # Tokenization (called from the data collator, runs on CPU) + # ---------------------------------------------------------------- + + def tokenize( + self, texts: List[str], is_query: bool + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Returns (input_ids, attention_mask) tensors on CPU.""" + prefix_ids = self._query_prefix_ids if is_query else self._passage_prefix_ids + suffix_ids = self._query_suffix_ids if is_query else self._passage_suffix_ids + gen_ids = [self.mask_token_id] * self._k(is_query) + mask_block = gen_ids + self._tail_ids # [MASK×K | " | chat_end | EOS] + max_text_len = self.max_length - len(prefix_ids) - len(suffix_ids) + + enc = self.tokenizer( + texts, + padding=False, + truncation=True, + max_length=max_text_len, + return_attention_mask=False, + return_token_type_ids=False, + add_special_tokens=False, + ) + enc['input_ids'] = [ + prefix_ids + ids + suffix_ids + mask_block + for ids in enc['input_ids'] + ] + collated = self.tokenizer.pad( + enc, + padding=True, + return_attention_mask=True, + return_tensors='pt', + ) + return collated['input_ids'], collated['attention_mask'] + + # ---------------------------------------------------------------- + # Attention mask: bidirectional (padding-only masking) + # ---------------------------------------------------------------- + + def _build_4d_mask( + self, seq_len: int, attention_mask: torch.Tensor + ) -> torch.Tensor: + """4D [B, 1, S, S] bidirectional attention mask.""" + dtype = next(self.backbone.parameters()).dtype + min_val = torch.finfo(dtype).min + B = attention_mask.size(0) + mask_4d = torch.zeros(B, 1, seq_len, seq_len, + device=attention_mask.device, dtype=dtype) + pad = ~attention_mask.bool() + mask_4d = mask_4d.masked_fill(pad.unsqueeze(1).unsqueeze(1), min_val) + mask_4d = mask_4d.masked_fill(pad.unsqueeze(1).unsqueeze(3), min_val) + return mask_4d + + # ---------------------------------------------------------------- + # Backbone forward + # ---------------------------------------------------------------- + + def _fwd( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor, + need_logits: bool = False, + mask_4d: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + """Single backbone pass. Returns (hidden [B,L,H], logits [B,L,V] or None). + + mask_4d: pre-computed 4D attention mask — pass it in the multi-step loop + to avoid recomputing the same mask at every denoising step. + inputs_embeds: if provided, use these embeddings instead of input_ids. + """ + # Adapter determines mask format: 4D [B,1,L,L] or 2D [B,L] + use_4d = self.adapter.needs_4d_mask() if self.adapter else True + if use_4d: + if mask_4d is None: + seq_len = inputs_embeds.size(1) if inputs_embeds is not None else input_ids.size(1) + mask_4d = self._build_4d_mask(seq_len, attention_mask) + fwd_mask = mask_4d + else: + fwd_mask = attention_mask + + # If hook is registered, skip output_hidden_states (saves ~1.2GB). + need_hidden_states = not self._hook_registered + fwd_kwargs = dict( + attention_mask=fwd_mask, + output_hidden_states=need_hidden_states, + return_dict=True, + ) + if inputs_embeds is not None: + fwd_kwargs['inputs_embeds'] = inputs_embeds + else: + fwd_kwargs['input_ids'] = input_ids + outputs = self.backbone(**fwd_kwargs) + # Keep hidden in native dtype (bf16). Callers convert only small slices. + if self._hook_registered and 'h' in self._last_hidden: + hidden = self._last_hidden.pop('h') + elif getattr(outputs, 'last_hidden_state', None) is not None: + hidden = outputs.last_hidden_state + elif hasattr(outputs, 'hidden_states') and outputs.hidden_states is not None: + hidden = outputs.hidden_states[-1] + else: + hidden = outputs[0] + + # Keep logits in native dtype (bf16) to avoid doubling memory. + # Downstream consumers convert small slices to float32 as needed. + logits = outputs.logits if need_logits and hasattr(outputs, 'logits') else None + + if (self.debug_compare_hidden_once + and self._hook_registered + and not self._debug_hidden_compared): + should_log = (not torch.distributed.is_initialized() + or torch.distributed.get_rank() == 0) + if should_log: + with torch.no_grad(): + ref_outputs = self.backbone( + input_ids=input_ids, + attention_mask=fwd_mask, + output_hidden_states=True, + return_dict=True, + ) + if getattr(ref_outputs, 'last_hidden_state', None) is not None: + ref_hidden = ref_outputs.last_hidden_state + elif hasattr(ref_outputs, 'hidden_states') and ref_outputs.hidden_states is not None: + ref_hidden = ref_outputs.hidden_states[-1] + else: + ref_hidden = ref_outputs[0] + diff = (hidden.float() - ref_hidden.float()).abs() + logger.warning( + "Hidden-hook check: model=%s shape=%s mean_abs_diff=%.6g max_abs_diff=%.6g", + self.model_type, tuple(hidden.shape), + diff.mean().item(), diff.max().item(), + ) + self._debug_hidden_compared = True + return hidden, logits + + # ---------------------------------------------------------------- + # Confidence-based unmasking + # ---------------------------------------------------------------- + + @staticmethod + def _sample_with_confidence(logits: torch.Tensor, alg: str = 'entropy'): + """Greedy decode + confidence score for discrete diffusion unmasking. + + Args: + logits: [N, V] logits at the N currently-masked positions + alg: 'entropy' (neg entropy; higher = more certain) or 'max_prob' + Returns: + (confidence [N], x0 [N]) — confidence scores and predicted token IDs + """ + probs = torch.softmax(logits.float(), dim=-1) + x0 = probs.argmax(dim=-1) + if alg == 'entropy': + log_probs = torch.log(probs + 1e-10) + confidence = (probs * log_probs).sum(dim=-1) # neg entropy + else: + confidence = probs.max(dim=-1).values + return confidence, x0 + + @torch.no_grad() + def _unmask_step( + self, + curr_ids: torch.Tensor, # [B, L] + logits: torch.Tensor, # [B, L, V] + K: int, + n_per_step: int, + ) -> Tuple[torch.Tensor, List[List[int]]]: + """Unmask n_per_step most-confident gen tokens (uniform LLaDA-style schedule). + + Returns: + new_ids: updated token IDs (same shape as curr_ids) + newly_decoded: list-of-lists; newly_decoded[i] = gen-block positions + [0..K-1] that were just decoded for example i + """ + new_ids = curr_ids.clone() + B, L = curr_ids.shape + g_start = L - K - self._n_tail + newly_decoded: List[List[int]] = [[] for _ in range(B)] + + for i in range(B): + gen_ids = new_ids[i, g_start:g_start + K] + mask_bool = (gen_ids == self.mask_token_id) + n_masked = mask_bool.sum().item() + if n_masked == 0: + continue + # Convert only the small slice [n_masked, V] to float32, not full logits + ml = logits[i, g_start:g_start + K][mask_bool].detach().float() + if not torch.isfinite(ml).all(): + continue + conf, x0 = self._sample_with_confidence(ml) + n_tr = min(n_per_step, n_masked) + _, xfer = torch.topk(conf, n_tr) + masked_abs = torch.where(mask_bool)[0] + selected = masked_abs[xfer] + new_ids[i, g_start + selected] = x0[xfer] + newly_decoded[i] = selected.tolist() + + return new_ids, newly_decoded + + # ---------------------------------------------------------------- + # Soft-token multi-step helpers + # ---------------------------------------------------------------- + + def _get_embed_layer(self) -> nn.Module: + """Return the input embedding layer (works through PEFT wrapper).""" + return self.backbone.get_input_embeddings() + + def _soft_unmask_step( + self, + curr_embeds: torch.Tensor, # [B, L, H_emb] — current embeddings (differentiable) + curr_ids: torch.Tensor, # [B, L] �� tracking which positions are still MASK + logits: torch.Tensor, # [B, L, V] — logits from current forward pass + K: int, + n_per_step: int, + temperature: float = 1.0, + ) -> Tuple[torch.Tensor, torch.Tensor, List[List[int]]]: + """Soft-unmask n_per_step most confident gen tokens. + + Instead of hard token replacement (argmax → embed lookup), computes + differentiable soft embeddings: softmax(logits / τ) @ embedding_matrix. + Gradient flows through logits → softmax → matmul → next-step embeddings. + + Returns: + new_embeds: [B, L, H_emb] with soft tokens at decoded positions + new_ids: [B, L] updated token IDs (hard argmax, for tracking only) + newly_decoded: list of gen-block positions [0..K-1] decoded this step + """ + embed_weight = self._get_embed_layer().weight # [V, H_emb] + B, L = curr_ids.shape + g_start = L - K - self._n_tail + new_ids = curr_ids.clone() + newly_decoded: List[List[int]] = [[] for _ in range(B)] + + # Compute soft embeddings for all K gen positions (differentiable) + gen_logits = logits[:, g_start:g_start + K, :] # [B, K, V] + soft_probs = F.softmax(gen_logits / temperature, dim=-1) # [B, K, V] + soft_embs = soft_probs @ embed_weight # [B, K, H_emb] + + # Determine which positions to unmask (most confident, same schedule as hard) + replace_mask = torch.zeros(B, K, dtype=torch.bool, device=curr_embeds.device) + for i in range(B): + gen_ids = curr_ids[i, g_start:g_start + K] + mask_bool = (gen_ids == self.mask_token_id) + n_masked = mask_bool.sum().item() + if n_masked == 0: + continue + ml = gen_logits[i][mask_bool].detach().float() + if not torch.isfinite(ml).all(): + continue + conf, x0 = self._sample_with_confidence(ml) + n_tr = min(n_per_step, n_masked) + _, xfer = torch.topk(conf, n_tr) + masked_abs = torch.where(mask_bool)[0] + selected = masked_abs[xfer] + replace_mask[i, selected] = True + new_ids[i, g_start + selected] = x0[xfer] + newly_decoded[i] = selected.tolist() + + # Mix: replace selected gen positions with soft embeddings, keep rest + curr_gen_embs = curr_embeds[:, g_start:g_start + K, :] # [B, K, H_emb] + new_gen_embs = torch.where( + replace_mask.unsqueeze(-1), # [B, K, 1] + soft_embs, # [B, K, H_emb] — differentiable + curr_gen_embs, # [B, K, H_emb] + ) + + # Rebuild full sequence (avoids in-place ops on grad-carrying tensor) + new_embeds = torch.cat([ + curr_embeds[:, :g_start, :], + new_gen_embs, + curr_embeds[:, g_start + K:, :], + ], dim=1) + + return new_embeds, new_ids, newly_decoded + + def _soft_multistep_forward( + self, + query_input_ids: torch.Tensor, + query_attention_mask: torch.Tensor, + passage_input_ids: torch.Tensor, + passage_attention_mask: torch.Tensor, + query_content_ids: Optional[List], + passage_content_ids: Optional[List], + n_steps: int, + ) -> Dict[str, torch.Tensor]: + """Multi-step denoising training with soft (differentiable) token replacement. + + Unlike hard multi-step where decoded positions are frozen (.detach()), + soft-token mode keeps gradient flowing through ALL K positions at every + step via softmax(logits/τ) @ embedding_matrix. + + Per-step contrastive loss (GIRCSE-style) provides direct supervision at + each denoising step. A monotonicity regularizer penalizes regressions. + """ + K_q, K_p = self.n_gen_q_tokens, self.n_gen_p_tokens + n_tail = self._n_tail + n_per_step_q = max(1, K_q // n_steps) + n_per_step_p = max(1, K_p // n_steps) + device = query_input_ids.device + B_q = query_input_ids.size(0) + B_p = passage_input_ids.size(0) + n_paq = B_p // B_q + L_q = query_input_ids.size(1) + L_p = passage_input_ids.size(1) + need_sparse = self.sparse_weight > 0 + + # Pre-compute 4D masks (sequence structure doesn't change across steps) + use_4d = self.adapter.needs_4d_mask() if self.adapter else True + q_mask_4d = self._build_4d_mask(L_q, query_attention_mask) if use_4d else None + p_mask_4d = self._build_4d_mask(L_p, passage_attention_mask) if use_4d else None + + # Initial embeddings from the backbone's embedding layer + embed_layer = self._get_embed_layer() + q_embeds = embed_layer(query_input_ids) + p_embeds = embed_layer(passage_input_ids) + q_curr_ids = query_input_ids.clone() + p_curr_ids = passage_input_ids.clone() + + q_g = L_q - K_q - n_tail + p_g = L_p - K_p - n_tail + + step_losses: List[torch.Tensor] = [] + final_loss: Dict[str, torch.Tensor] = {} + + for step in range(n_steps): + is_last = (step == n_steps - 1) + + # Forward pass using embeddings (step 0: original, step 1+: soft tokens) + q_h, q_logits = self._fwd( + q_curr_ids, query_attention_mask, + need_logits=True, mask_4d=q_mask_4d, + inputs_embeds=q_embeds if step > 0 else None, + ) + p_h, p_logits = self._fwd( + p_curr_ids, passage_attention_mask, + need_logits=True, mask_4d=p_mask_4d, + inputs_embeds=p_embeds if step > 0 else None, + ) + + # Extract representations — ALL K positions have gradient + q_repr_hidden = q_h[:, q_g:q_g + K_q, :].float() + q_quotation_emb = q_h[:, q_g - 1, :].float() + p_repr_hidden = p_h[:, p_g:p_g + K_p, :].float() + p_quotation_emb = p_h[:, p_g - 1, :].float() + + q_sparse_max = None + p_sparse_max = None + # Monotonic trick: log1p(relu(.)) is non-decreasing, so + # max_k log1p(relu(x_k)) == log1p(relu(max_k x_k)) + # Maxing in bf16 first avoids materialising [B, K, V] fp32 intermediate. + if need_sparse and q_logits is not None: + q_sparse_max = torch.log1p(torch.relu( + q_logits[:, q_g:q_g + K_q, :].max(dim=1).values)) + if need_sparse and p_logits is not None: + p_sparse_max = torch.log1p(torch.relu( + p_logits[:, p_g:p_g + K_p, :].max(dim=1).values)) + + if self.normalize: + q_quotation_emb = F.normalize(q_quotation_emb, p=2, dim=-1) + q_repr_hidden = F.normalize(q_repr_hidden, p=2, dim=-1) + p_quotation_emb = F.normalize(p_quotation_emb, p=2, dim=-1) + p_repr_hidden = F.normalize(p_repr_hidden, p=2, dim=-1) + + q_repr = {'repr_hidden': q_repr_hidden, 'quotation_emb': q_quotation_emb} + p_repr = {'repr_hidden': p_repr_hidden, 'quotation_emb': p_quotation_emb} + if q_sparse_max is not None: + if query_content_ids is not None: + q_sparse_max = filter_sparse(q_sparse_max, query_content_ids) + q_repr['sparse_acts'] = q_sparse_max + if p_sparse_max is not None: + if passage_content_ids is not None: + p_sparse_max = filter_sparse(p_sparse_max, passage_content_ids) + p_repr['sparse_acts'] = p_sparse_max + + # Cross-GPU negative sharing + q_repr = self._gather_repr(q_repr) + p_repr = self._gather_repr(p_repr) + B_q_g = q_repr['repr_hidden'].size(0) + B_p_g = p_repr['repr_hidden'].size(0) + labels = torch.arange(B_q_g, device=device) * (B_p_g // B_q_g) + + # Per-step contrastive loss + step_loss_dict = self.compute_loss(q_repr, p_repr, labels) + step_losses.append(step_loss_dict['loss']) + + if is_last: + final_loss = step_loss_dict + + # Soft unmask for next step (per-side K) + if not is_last: + q_embeds, q_curr_ids, _ = self._soft_unmask_step( + q_embeds, q_curr_ids, q_logits, K_q, n_per_step_q, + self.soft_temperature) + p_embeds, p_curr_ids, _ = self._soft_unmask_step( + p_embeds, p_curr_ids, p_logits, K_p, n_per_step_p, + self.soft_temperature) + + # Accumulate per-step losses (linearly increasing weight) + if len(step_losses) > 1: + progressive_loss = sum( + (s + 1) / n_steps * loss + for s, loss in enumerate(step_losses[:-1]) + ) + final_loss['loss'] = final_loss['loss'] + progressive_loss + final_loss['loss_progressive'] = progressive_loss.detach() + + # Monotonicity regularizer: penalize if later step is worse + mono_penalty = torch.stack([ + torch.relu(step_losses[k + 1] - step_losses[k].detach()) + for k in range(len(step_losses) - 1) + ]).mean() + final_loss['loss'] = final_loss['loss'] + 0.1 * mono_penalty + final_loss['loss_monotonicity'] = mono_penalty.detach() + + # DDP scaling + if torch.distributed.is_initialized() and torch.distributed.get_world_size() > 1: + final_loss['loss'] = final_loss['loss'] * torch.distributed.get_world_size() + + # Denoising auxiliary (same as other paths) + if self.denoising_weight > 0 and final_loss: + p_corrupted, p_denoise_targets, mask_ratio = self._apply_text_masking( + passage_input_ids, passage_attention_mask) + _, p_logits_dn = self._fwd(p_corrupted, passage_attention_mask, + need_logits=True, mask_4d=p_mask_4d) + denoise_loss = self.compute_denoising_loss( + p_logits_dn, p_denoise_targets, mask_ratio) + final_loss['loss'] = final_loss['loss'] + self.denoising_weight * denoise_loss + final_loss['loss_denoising'] = denoise_loss.detach() + + return final_loss + + # ---------------------------------------------------------------- + # K pre-encoder (two-stage variable-length training) + # ---------------------------------------------------------------- + + def _pre_encoder_logits( + self, + input_ids: torch.Tensor, # [B, L] + attention_mask: torch.Tensor, # [B, L] + is_query: bool, + ) -> torch.Tensor: + """Tiny embedding-only K-router. Mean-pools the embedding-layer + output over real (non-pad) tokens, then runs a small MLP head. + Returns [B, n_K_pe] logits over self.k_pre_encoder_options. + + Cost ≈ embedding lookup + 2-layer MLP — under 1% of full forward. + """ + head = self.k_pre_encoder_q if is_query else self.k_pre_encoder_p + # Use main encoder's embedding layer (parameter-shared). + embed_layer = self._get_embed_layer() + emb = embed_layer(input_ids) # [B, L, H] + # Mean-pool over real tokens (attention_mask = 0 at pads). + m = attention_mask.float().unsqueeze(-1) + pooled = (emb.float() * m).sum(dim=1) / m.sum(dim=1).clamp(min=1.0) # [B, H] + # Match head dtype to avoid AMP/DeepSpeed dtype mismatch. + head_dtype = next(head.parameters()).dtype + logits = head(pooled.to(dtype=head_dtype)).float() # [B, n_K_pe] + return logits + + def _slice_to_per_item_K( + self, + input_ids: torch.Tensor, # [B, L_max] left-padded, K_max MASKs at end + attention_mask: torch.Tensor, + K_per_item: torch.Tensor, # [B] integers in self.k_pre_encoder_options + K_max_side: int, # n_gen_q_tokens or n_gen_p_tokens for this side + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Slice each item's input from K_max MASK tokens down to K MASK tokens + (per-item K), preserving the structural tail [", chat_end, EOS]. + + Layout (all left-padded — content right-aligned): + [PAD ... PAD | prefix | text | suffix | K_max MASKs | tail] + After slicing item with K_i: + [PAD ... PAD | prefix | text | suffix | K_i MASKs | tail] + Batch is then padded to max-K-in-batch length. + """ + B = input_ids.size(0) + L_in = input_ids.size(1) + n_tail = self._n_tail + K_per_item = K_per_item.tolist() if torch.is_tensor(K_per_item) else list(K_per_item) + # Per-item sliced sequences. + sliced_ids: List[torch.Tensor] = [] + sliced_mask: List[torch.Tensor] = [] + for i in range(B): + K_i = int(K_per_item[i]) + n_remove = K_max_side - K_i + # Cut tokens [-(n_tail + n_remove) : -n_tail). + # Equivalent: keep input_ids[: L_in - n_tail - n_remove] + input_ids[-n_tail:] + keep_until = L_in - n_tail - n_remove + new_ids = torch.cat([input_ids[i, :keep_until], input_ids[i, -n_tail:] + if n_tail > 0 else input_ids.new_empty(0)]) + new_mask = torch.cat([attention_mask[i, :keep_until], attention_mask[i, -n_tail:] + if n_tail > 0 else attention_mask.new_empty(0)]) + sliced_ids.append(new_ids) + sliced_mask.append(new_mask) + + # Pad LEFT to max length so the structural tail stays at position -1. + L_out = max(s.size(0) for s in sliced_ids) + pad_id = (self.tokenizer.pad_token_id or self.tokenizer.eos_token_id or 0) + out_ids = input_ids.new_full((B, L_out), pad_id) + out_mask = attention_mask.new_zeros((B, L_out)) + for i, (s_ids, s_mask) in enumerate(zip(sliced_ids, sliced_mask)): + L_i = s_ids.size(0) + out_ids[i, L_out - L_i:] = s_ids + out_mask[i, L_out - L_i:] = s_mask + return out_ids, out_mask + + def _extract_repr_per_K( + self, + hidden: torch.Tensor, # [B, L_out, H] + K_per_item: torch.Tensor, # [B] int K per item + n_K_max: int, # max K across self.k_pre_encoder_options + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Extract per-item repr_hidden using each item's K. + Returns: + repr_hidden: [B, n_K_max, H] — first K_i positions filled, rest zero + quotation_emb: [B, H] — position before first MASK per item + With left-padding, the rightmost positions per item are: + [..., gen_K_i MASKs, tail tokens (n_tail)] + So mask block = positions [L_out - n_tail - K_i, L_out - n_tail). + Quotation = position L_out - n_tail - K_i - 1. + """ + B, L_out, H = hidden.shape + n_tail = self._n_tail + K_per_item = K_per_item.tolist() if torch.is_tensor(K_per_item) else list(K_per_item) + repr_hidden = hidden.new_zeros(B, n_K_max, H) + quotation_emb = hidden.new_zeros(B, H) + for i in range(B): + K_i = int(K_per_item[i]) + mask_end = L_out - n_tail + mask_start = mask_end - K_i + repr_hidden[i, :K_i] = hidden[i, mask_start:mask_end].float() + if mask_start - 1 >= 0: + quotation_emb[i] = hidden[i, mask_start - 1].float() + return repr_hidden, quotation_emb + + def _gumbel_st_sample_K( + self, + logits: torch.Tensor, # [B, n_K_pe] + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Gumbel-softmax sample with straight-through. + + Returns: + K_per_item: [B] int — actual K integer values to use for slicing + soft_one_hot: [B, n_K_pe] — hard one-hot in fwd, soft in backward + (used for cost regularizer + ST credit) + chosen_prob: [B] — one-hot's value at chosen index (always 1.0 in fwd) + used to inject gradient back through ST + """ + soft_one_hot = F.gumbel_softmax( + logits, tau=max(self.gumbel_temperature, 1e-3), hard=True + ) # [B, n_K_pe] + # Integer K from one-hot + K_idx = soft_one_hot.argmax(dim=-1) # [B] + K_per_item = self.k_pe_options_tensor[K_idx].long() # [B] + # ST credit: chosen_prob is always 1.0 in forward, has gradient via ST + chosen_prob = soft_one_hot.gather(-1, K_idx.unsqueeze(-1)).squeeze(-1) + return K_per_item, soft_one_hot, chosen_prob + + def _forward_with_k_pre_encoder( + self, + query_input_ids: torch.Tensor, + query_attention_mask: torch.Tensor, + passage_input_ids: torch.Tensor, + passage_attention_mask: torch.Tensor, + query_content_ids: Optional[List], + passage_content_ids: Optional[List], + labels: torch.Tensor, + ) -> Dict[str, torch.Tensor]: + """Two-stage forward: pre-encoder predicts K_q, K_p per item; main + encoder runs at variable length (per-item K). Real encoding savings + when batch's max-K is below K_max. + + Backprop: + - Main encoder: gradient flows through standard InfoNCE on per-item + sliced representations. + - Pre-encoder: gradient flows via Gumbel straight-through on + chosen_prob (always 1.0 in fwd; gradient injected through ST). + Plus a cost regularizer λ * E[K] discourages collapsing to K=K_max. + """ + device = labels.device + K_max_q = self.n_gen_q_tokens + K_max_p = self.n_gen_p_tokens + + # ── 1. Pre-encoder forward (cheap — embedding lookup + small MLP). + q_logits = self._pre_encoder_logits( + query_input_ids, query_attention_mask, is_query=True) + p_logits = self._pre_encoder_logits( + passage_input_ids, passage_attention_mask, is_query=False) + + # ── 2. Gumbel-ST sample K per item (queries + passages). + K_q_per_item, q_soft, q_chosen = self._gumbel_st_sample_K(q_logits) + K_p_per_item, p_soft, p_chosen = self._gumbel_st_sample_K(p_logits) + + # Clamp at the per-side K_max in case the pre-encoder options exceed it. + # (Shouldn't happen by construction but a safety net.) + K_q_per_item = K_q_per_item.clamp(min=1, max=K_max_q) + K_p_per_item = K_p_per_item.clamp(min=1, max=K_max_p) + + # ── 3. Slice each item's input to its K (variable-length batch). + q_ids_sliced, q_mask_sliced = self._slice_to_per_item_K( + query_input_ids, query_attention_mask, K_q_per_item, K_max_q) + p_ids_sliced, p_mask_sliced = self._slice_to_per_item_K( + passage_input_ids, passage_attention_mask, K_p_per_item, K_max_p) + + # ── 4. Main encoder forward (one pass each at variable length). + need_sparse = self.sparse_weight > 0 + q_hidden, q_logits_main = self._fwd( + q_ids_sliced, q_mask_sliced, need_logits=need_sparse) + p_hidden, p_logits_main = self._fwd( + p_ids_sliced, p_mask_sliced, need_logits=need_sparse) + + # ── 5. Per-item repr extraction. + n_K_max = max(self.k_pre_encoder_options) + q_repr_hidden, q_quotation_emb = self._extract_repr_per_K( + q_hidden, K_q_per_item, n_K_max) + p_repr_hidden, p_quotation_emb = self._extract_repr_per_K( + p_hidden, K_p_per_item, n_K_max) + + # ── 6. Sparse: max-pool over each item's K positions. + q_sparse_max = None + p_sparse_max = None + if need_sparse and q_logits_main is not None: + # For each item, compute log(1+relu(...)) over its K MASK positions, + # then max-pool across positions. Same as legacy sparse computation. + B_q = q_ids_sliced.size(0) + L_q_out = q_ids_sliced.size(1) + q_sparse_max = q_logits_main.new_zeros(B_q, q_logits_main.size(-1)) + for i in range(B_q): + K_i = int(K_q_per_item[i].item()) + mask_end = L_q_out - self._n_tail + mask_start = mask_end - K_i + slogits = q_logits_main[i, mask_start:mask_end] + # Monotonic trick: max over K positions first → log1p(relu()) on [V]. + q_sparse_max[i] = torch.log1p(torch.relu(slogits.max(dim=0).values)) + if need_sparse and p_logits_main is not None: + B_p = p_ids_sliced.size(0) + L_p_out = p_ids_sliced.size(1) + p_sparse_max = p_logits_main.new_zeros(B_p, p_logits_main.size(-1)) + for i in range(B_p): + K_i = int(K_p_per_item[i].item()) + mask_end = L_p_out - self._n_tail + mask_start = mask_end - K_i + slogits = p_logits_main[i, mask_start:mask_end] + # Monotonic trick: max over K positions first → log1p(relu()) on [V]. + p_sparse_max[i] = torch.log1p(torch.relu(slogits.max(dim=0).values)) + + # ── 7. Normalize. + if self.normalize: + q_repr_hidden = F.normalize(q_repr_hidden, p=2, dim=-1) + q_quotation_emb = F.normalize(q_quotation_emb, p=2, dim=-1) + p_repr_hidden = F.normalize(p_repr_hidden, p=2, dim=-1) + p_quotation_emb = F.normalize(p_quotation_emb, p=2, dim=-1) + + # ── 8. ST credit injection — multiply repr by chosen_prob (always 1.0 + # in fwd, has gradient through ST in backward). This is the standard + # Gumbel-ST trick for credit assignment when downstream operations + # (slicing, reshape) are non-differentiable. + q_repr_hidden = q_repr_hidden * q_chosen.float().unsqueeze(-1).unsqueeze(-1) + p_repr_hidden = p_repr_hidden * p_chosen.float().unsqueeze(-1).unsqueeze(-1) + q_quotation_emb = q_quotation_emb * q_chosen.float().unsqueeze(-1) + p_quotation_emb = p_quotation_emb * p_chosen.float().unsqueeze(-1) + + # Build repr dicts and apply content filter to sparse. + q_repr = {'repr_hidden': q_repr_hidden, 'quotation_emb': q_quotation_emb} + p_repr = {'repr_hidden': p_repr_hidden, 'quotation_emb': p_quotation_emb} + if q_sparse_max is not None: + if query_content_ids is not None: + q_sparse_max = filter_sparse(q_sparse_max, query_content_ids) + q_repr['sparse_acts'] = q_sparse_max + if p_sparse_max is not None: + if passage_content_ids is not None: + p_sparse_max = filter_sparse(p_sparse_max, passage_content_ids) + p_repr['sparse_acts'] = p_sparse_max + + # ── 9. Cross-GPU gather + recompute labels for gathered batch. + q_repr = self._gather_repr(q_repr) + p_repr = self._gather_repr(p_repr) + B_q_all = q_repr['repr_hidden'].size(0) + B_p_all = p_repr['repr_hidden'].size(0) + n_paq_g = B_p_all // B_q_all + labels_g = torch.arange(B_q_all, device=device) * n_paq_g + + # ── 10. Standard InfoNCE on per-item-K representations (zero-padded + # positions contribute nothing thanks to MaxSim's max-over-positions). + loss_dict = self.compute_loss(q_repr, p_repr, labels_g) + + # ── 11. Cost regularizer — λ * E[K_q + K_p] using soft probs (so it's + # differentiable w.r.t. logits). Discourages collapsing to K_max. + if self.k_cost_lambda > 0: + K_options_t = self.k_pe_options_tensor # [n_K_pe] + expected_K_q = (q_soft * K_options_t).sum(dim=-1).mean() + expected_K_p = (p_soft * K_options_t).sum(dim=-1).mean() + cost_loss = (expected_K_q + expected_K_p) / float(max(K_options_t.max().item(), 1.0)) + loss_dict['loss'] = loss_dict['loss'] + self.k_cost_lambda * cost_loss + loss_dict['k_cost'] = cost_loss.detach() + loss_dict['expected_K_q'] = expected_K_q.detach() + loss_dict['expected_K_p'] = expected_K_p.detach() + + # Diagnostics: marginal K distributions. + with torch.no_grad(): + for i, K in enumerate(self.k_pre_encoder_options): + loss_dict[f'pe_kq_p{K}'] = q_soft[:, i].mean().detach() + loss_dict[f'pe_kp_p{K}'] = p_soft[:, i].mean().detach() + + # DDP loss scaling (counters HF Trainer's gradient averaging). + if torch.distributed.is_initialized() and torch.distributed.get_world_size() > 1: + loss_dict['loss'] = loss_dict['loss'] * torch.distributed.get_world_size() + + return loss_dict + + # ---------------------------------------------------------------- + # Encode — single-pass or multi-step denoising (inference) + # ---------------------------------------------------------------- + + def encode( + self, + input_ids: torch.Tensor, # [B, L] + attention_mask: torch.Tensor, # [B, L] + is_query: bool = False, + compute_sparse: Optional[bool] = None, + content_token_ids: Optional[List] = None, + ) -> Dict[str, torch.Tensor]: + """ + PromptReps encoding. + + num_denoise_steps == 1: single forward pass (fast path, differentiable, + used in both training and inference). + num_denoise_steps > 1: iterative denoising loop (inference). + repr_hidden[i, k] is captured at the step when gen token k transitions + MASK → decoded, benefiting from previously-decoded token context. + Uniform schedule: n_per_step = K // n_steps tokens decoded per step. + quotation_emb is captured at the final step. + Call under torch.no_grad() for inference; use forward() for training. + + Returns: + repr_hidden [B, K, H] ColBERT repr (one vector per gen token) + quotation_emb [B, H] closing " token (single-dense repr) + sparse_acts [B, V] sum of log(1+relu(logit)) across K positions + """ + K = self._k(is_query) + B = input_ids.size(0) + L = input_ids.size(1) + n_steps = self.num_denoise_steps + device = input_ids.device + g_start = L - K - self._n_tail # left-padded: MASK block then structural tail + + if compute_sparse is None: + need_sparse = self.sparse_weight > 0 + else: + need_sparse = compute_sparse + + # Initialised here so both the single- and multi-step branches always leave + # these bound — the multi-step path only ever populates sparse_max. + sparse_max: Optional[torch.Tensor] = None + sparse_per_pos: Optional[torch.Tensor] = None + + # ── Single-pass fast path (n_steps ≤ 1) ────────────────────────────── + if n_steps <= 1: + hidden, logits = self._fwd(input_ids, attention_mask, need_logits=need_sparse) + repr_hidden = hidden[:, g_start:g_start + K, :].float() # [B, K, H] + quotation_emb = hidden[:, g_start - 1, :].float() # [B, H] + if need_sparse and logits is not None: + # Per-position sparse acts for ColBERT-style multi-vector sparse + # (caller may pop result['sparse_acts_per_pos'] to do per-K topk). + # Keep in bf16 — the downstream topk → ×100 → int round path + # doesn't need fp32 precision, and the fp32 cast was the + # dominant memory bandwidth cost at K=16 (312 MB → 156 MB/batch + # for V=152k, B=32; 2× reduction). + gen_logits = logits[:, g_start:g_start + K, :] # [B, K, V] bf16 + sparse_per_pos = torch.log1p(torch.relu(gen_logits)) # [B, K, V] bf16 + # sparse_max stays fp32 (small, [B, V]) for downstream consumers. + sparse_max = sparse_per_pos.max(dim=1).values.float() # [B, V] fp32 + + # ── Multi-step denoising loop (n_steps > 1, inference) ─────────────── + else: + curr_ids = input_ids.clone() + n_per_step = max(1, K // n_steps) + vocab_size = self.backbone.config.vocab_size + + repr_buf = torch.zeros(B, K, self.hidden_size, device=device) + repr_saved = torch.zeros(B, K, dtype=torch.bool, device=device) + sparse_max = (torch.zeros(B, vocab_size, device=device) + if need_sparse else None) + quotation_emb = None + mask_4d = self._build_4d_mask(L, attention_mask) if (self.adapter.needs_4d_mask() if self.adapter else True) else None + + for step in range(n_steps): + hidden, logits = self._fwd(curr_ids, attention_mask, need_logits=True, mask_4d=mask_4d) + is_last = (step == n_steps - 1) + + # ── Vectorized over batch (replaces `for i in range(B): ... .item() ...`). + # Behavior preserved: greedy argmax, neg-entropy confidence, top-n_per_step + # transitions per item; on last step all remaining masks transition. + gen_ids_all = curr_ids[:, g_start:g_start + K] # [B, K] int + mask_bool_all = (gen_ids_all == self.mask_token_id) # [B, K] bool + ml_all = logits[:, g_start:g_start + K, :] # [B, K, V] bf16 + + # Confidence + argmax (matches _sample_with_confidence(alg='entropy')) + probs = F.softmax(ml_all.float(), dim=-1) # [B, K, V] fp32 + log_probs = torch.log(probs + 1e-10) + conf = (probs * log_probs).sum(dim=-1) # [B, K] (neg entropy) + x0_pred = probs.argmax(dim=-1) # [B, K] + # Restrict topk to currently-masked positions + conf = torch.where(mask_bool_all, conf, torch.full_like(conf, -float('inf'))) + + if is_last: + xfer_mask = mask_bool_all # all remaining masks + else: + n_tr = min(n_per_step, K) + _, top_idx = conf.topk(n_tr, dim=-1) # [B, n_tr] + xfer_mask = torch.zeros_like(mask_bool_all) + xfer_mask.scatter_(1, top_idx, True) + xfer_mask = xfer_mask & mask_bool_all # only actually-masked positions + + # ── Vectorized state updates ──────────────────────────────────── + # 1. curr_ids: replace MASK with predicted at xfer positions + if xfer_mask.any(): + new_K = torch.where(xfer_mask, x0_pred, gen_ids_all) + # In-place via clone (curr_ids may be a view of input_ids) + curr_ids = curr_ids.clone() + curr_ids[:, g_start:g_start + K] = new_K + + # 2. repr_buf: save hidden at to_save positions (xfer & not yet saved) + to_save = xfer_mask & ~repr_saved + if to_save.any(): + hidden_chunk = hidden[:, g_start:g_start + K, :] # [B, K, H] + repr_buf = torch.where(to_save.unsqueeze(-1), hidden_chunk, repr_buf) + # 3. sparse_max: max over to_save positions of log(1+relu(logits)) + if sparse_max is not None: + sparse_chunk = torch.log1p(torch.relu(ml_all)) # [B, K, V] bf16 + masked_sparse = torch.where( + to_save.unsqueeze(-1).expand_as(sparse_chunk), + sparse_chunk, + torch.full_like(sparse_chunk, -float('inf'))) + new_max = masked_sparse.max(dim=1).values.float() # [B, V] fp32 + sparse_max = torch.max(sparse_max, new_max) + repr_saved = repr_saved | to_save + + if is_last: + quotation_emb = hidden[:, g_start - 1, :].float() + # Catch positions that were never masked (rare but possible + # if input already had a non-MASK token in the K block). + unsaved = ~repr_saved + if unsaved.any(): + hidden_chunk = hidden[:, g_start:g_start + K, :] + repr_buf = torch.where(unsaved.unsqueeze(-1), hidden_chunk, repr_buf) + if sparse_max is not None: + sparse_chunk = torch.log1p(torch.relu(ml_all)) + masked_sparse = torch.where( + unsaved.unsqueeze(-1).expand_as(sparse_chunk), + sparse_chunk, + torch.full_like(sparse_chunk, -float('inf'))) + new_max = masked_sparse.max(dim=1).values.float() + sparse_max = torch.max(sparse_max, new_max) + repr_saved = repr_saved | unsaved + + repr_buf = torch.nan_to_num(repr_buf) + repr_hidden = repr_buf + quotation_emb = torch.nan_to_num(quotation_emb) + + # ── Normalize ───────────────────────────────────────────────────────── + if self.normalize: + quotation_emb = F.normalize(quotation_emb, p=2, dim=-1) + repr_hidden = F.normalize(repr_hidden, p=2, dim=-1) + + result: Dict[str, torch.Tensor] = { + 'repr_hidden': repr_hidden, + 'quotation_emb': quotation_emb, + } + if sparse_max is not None: + if content_token_ids is not None: + sparse_max = filter_sparse(sparse_max, content_token_ids) + result['sparse_acts'] = sparse_max + # Per-position sparse [B, K, V] — single-pass only, K>1 + if sparse_per_pos is not None and K > 1: + result['sparse_acts_per_pos'] = sparse_per_pos # unfiltered; caller applies content filter + return result + + # ---------------------------------------------------------------- + # Cross-GPU negative sharing + # ---------------------------------------------------------------- + + @staticmethod + def _dist_gather(t: torch.Tensor) -> torch.Tensor: + """All-gather tensors across GPUs with gradient passthrough.""" + if not (torch.distributed.is_initialized() and torch.distributed.get_world_size() > 1): + return t + gathered = [torch.zeros_like(t) for _ in range(torch.distributed.get_world_size())] + torch.distributed.all_gather(gathered, t.contiguous()) + # Replace own shard with original (keeps gradient) + gathered[torch.distributed.get_rank()] = t + return torch.cat(gathered, dim=0) + + def _gather_repr(self, repr_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: + """Gather all representation tensors across GPUs.""" + out = {} + for k, v in repr_dict.items(): + out[k] = self._dist_gather(v) + return out + + # ---------------------------------------------------------------- + # MaxSim (ColBERT-style) + # ---------------------------------------------------------------- + + @staticmethod + def maxsim(q: torch.Tensor, p: torch.Tensor) -> torch.Tensor: + """ + q: [B_q, k, H], p: [B_p, k, H] → [B_q, B_p] MaxSim scores. + """ + sims = torch.einsum('ikh,jlh->ijkl', q, p) # [B_q, B_p, k_q, k_p] + return sims.max(dim=-1).values.sum(dim=-1) # [B_q, B_p] + + # ---------------------------------------------------------------- + # Auxiliary losses (diffusion-native) + # ---------------------------------------------------------------- + + @torch.no_grad() + def _corrupt_text( + self, + input_ids: torch.Tensor, # [B, L] + attention_mask: torch.Tensor, # [B, L] + rate: float, + ) -> torch.Tensor: + """Randomly replace rate% of text tokens with MASK for corruption augmentation. + + Only corrupts real text tokens (not padding, not gen MASKs, not EOS). + Returns a new tensor (does not modify input in-place). + """ + B, L = input_ids.shape + K = self.n_gen_p_tokens # only called on passages + g_start = L - K - self._n_tail + + corrupted = input_ids.clone() + is_real = attention_mask.bool() + is_gen = torch.zeros(L, dtype=torch.bool, device=input_ids.device) + is_gen[g_start:] = True + is_candidate = is_real & ~is_gen.unsqueeze(0) & (input_ids != self.mask_token_id) + + for i in range(B): + cand_idx = torch.where(is_candidate[i])[0] + if cand_idx.numel() == 0: + continue + n_mask = max(1, int(cand_idx.numel() * rate)) + perm = torch.randperm(cand_idx.numel(), device=input_ids.device)[:n_mask] + corrupted[i, cand_idx[perm]] = self.mask_token_id + + return corrupted + + def _apply_text_masking( + self, + input_ids: torch.Tensor, # [B, L] + attention_mask: torch.Tensor, # [B, L] + ) -> Tuple[torch.Tensor, torch.Tensor, float]: + """Randomly mask text tokens (NOT generation MASKs) for denoising auxiliary loss. + + Returns: + corrupted_ids: input_ids with some text tokens replaced by MASK + denoise_targets: original token IDs at masked positions (-100 elsewhere) + mask_ratio: fraction of tokens masked (t for 1/t weighting) + """ + B, L = input_ids.shape + K = self.n_gen_p_tokens # called on passages + g_start = L - K - self._n_tail # left-padded: gen block + structural tail + + corrupted = input_ids.clone() + targets = torch.full_like(input_ids, -100) # -100 = ignore in CE + + # Candidate text positions: not padding, not in gen/EOS block, not already MASK + is_real = attention_mask.bool() # [B, L] + is_gen = torch.zeros(L, dtype=torch.bool, device=input_ids.device) + is_gen[g_start:] = True # marks MASK positions + EOS as non-candidate + is_candidate = is_real & ~is_gen.unsqueeze(0) & (input_ids != self.mask_token_id) + + for i in range(B): + cand_idx = torch.where(is_candidate[i])[0] + if cand_idx.numel() == 0: + continue + n_mask = max(1, int(cand_idx.numel() * self.denoise_mask_ratio)) + perm = torch.randperm(cand_idx.numel(), device=input_ids.device)[:n_mask] + mask_positions = cand_idx[perm] + targets[i, mask_positions] = input_ids[i, mask_positions] + corrupted[i, mask_positions] = self.mask_token_id + + mask_ratio = self.denoise_mask_ratio + return corrupted, targets, mask_ratio + + def compute_denoising_loss( + self, + logits: torch.Tensor, # [B, L, V] from forward pass on corrupted input + targets: torch.Tensor, # [B, L] with -100 for non-masked positions + mask_ratio: float, # t for 1/t weighting + ) -> torch.Tensor: + """Compute Dream/LLaDA-style denoising loss: weighted CE on masked text tokens.""" + # Flatten and compute CE (ignores -100 positions automatically) + loss = F.cross_entropy( + logits.view(-1, logits.size(-1)), + targets.view(-1), + ignore_index=-100, + ) + # Weight by 1/t (matching Dream/LLaDA pre-training ELBO) + return loss / max(mask_ratio, 1e-3) + + def compute_diversity_loss( + self, + repr_hidden: torch.Tensor, # [B, K, H] + ) -> torch.Tensor: + """Push K representation vectors to be diverse (low pairwise cosine similarity).""" + K = repr_hidden.size(1) + if K <= 1: + return torch.tensor(0.0, device=repr_hidden.device) + + # Pairwise cosine similarity [B, K, K] — repr_hidden is already L2-normalised + sim_matrix = torch.bmm(repr_hidden, repr_hidden.transpose(1, 2)) # [B, K, K] + + # Mean of upper triangle (exclude diagonal) + mask = torch.triu(torch.ones(K, K, device=repr_hidden.device), diagonal=1).bool() + pairwise_sims = sim_matrix[:, mask] # [B, K*(K-1)/2] + + # Hinge: penalize similarities above 0 (push toward orthogonal) + diversity_loss = torch.relu(pairwise_sims).mean() + return diversity_loss + + @staticmethod + def _mean_offdiag_cos(repr_hidden: torch.Tensor) -> torch.Tensor: + """Mean off-diagonal cosine similarity over K vectors.""" + K = repr_hidden.size(1) + if K <= 1: + return torch.tensor(0.0, device=repr_hidden.device) + normed = F.normalize(repr_hidden.float(), p=2, dim=-1) + sim = torch.bmm(normed, normed.transpose(1, 2)) + mask = torch.triu(torch.ones(K, K, device=repr_hidden.device), diagonal=1).bool() + return sim[:, mask].mean() + + def _dense_debug_stats( + self, + q_repr: Dict[str, torch.Tensor], + p_repr: Dict[str, torch.Tensor], + labels: torch.Tensor, + q_raw: Optional[torch.Tensor] = None, + p_raw: Optional[torch.Tensor] = None, + ) -> Dict[str, torch.Tensor]: + """Summaries that explain whether dense loss comes from poor score gaps or collapsed vectors.""" + scores = self.maxsim(q_repr['repr_hidden'], p_repr['repr_hidden']) + row_idx = torch.arange(scores.size(0), device=scores.device) + pos = scores[row_idx, labels] + neg = scores.clone() + neg[row_idx, labels] = float('-inf') + hard_neg = neg.max(dim=1).values + + stats = { + 'debug_pos_score': pos.mean().detach(), + 'debug_hardneg_score': hard_neg.mean().detach(), + 'debug_score_gap': (pos - hard_neg).mean().detach(), + 'debug_pos_beats_hardneg': (pos > hard_neg).float().mean().detach(), + 'debug_q_mask_cos': self._mean_offdiag_cos(q_repr['repr_hidden']).detach(), + 'debug_p_mask_cos': self._mean_offdiag_cos(p_repr['repr_hidden']).detach(), + } + + if q_raw is not None: + stats['debug_q_mask_norm'] = q_raw.float().norm(dim=-1).mean().detach() + stats['debug_q_mask_raw_cos'] = self._mean_offdiag_cos(q_raw).detach() + if p_raw is not None: + stats['debug_p_mask_norm'] = p_raw.float().norm(dim=-1).mean().detach() + stats['debug_p_mask_raw_cos'] = self._mean_offdiag_cos(p_raw).detach() + + return stats + + # ---------------------------------------------------------------- + # Loss + # ---------------------------------------------------------------- + + def compute_loss( + self, + q_repr: Dict[str, torch.Tensor], + p_repr: Dict[str, torch.Tensor], + labels: torch.Tensor, # [B_q] index of positive in B_p + ) -> Dict[str, torch.Tensor]: + """ + Dense loss: ColBERT MaxSim on repr_hidden (MASK positions). + For K=1 this reduces to dot product of the single MASK vector. + Sparse InfoNCE (if sparse_weight > 0): raw dot product + temperature. + + K-adapter: when self.use_k_adapter is True, dispatch to the per-cell + loss path that trains both the retriever and the K-routing head. + """ + if self.use_k_adapter: + return self._compute_loss_with_k_adapter(q_repr, p_repr, labels) + + device = labels.device + K = q_repr['repr_hidden'].size(1) + + total_loss = torch.tensor(0.0, device=device) + + # Primary dense: ColBERT MaxSim on repr_hidden (MASK positions). + colbert_scores = self.maxsim(q_repr['repr_hidden'], p_repr['repr_hidden']) + colbert_loss = F.cross_entropy(colbert_scores / K / self.temperature, labels) + total_loss = self.dense_weight * colbert_loss + + # Sparse InfoNCE — raw dot product, same as inference. + # Content-token filtering (PromptReps-style) applied in forward() keeps + # scores small (~10-100), so no clamping or normalization needed. + sparse_loss = torch.tensor(0.0, device=device) + if (self.sparse_weight > 0 + and 'sparse_acts' in q_repr + and 'sparse_acts' in p_repr): + sparse_scores = q_repr['sparse_acts'] @ p_repr['sparse_acts'].T + sparse_loss = F.cross_entropy(sparse_scores, labels) + total_loss = total_loss + self.sparse_weight * sparse_loss + + # Auxiliary: diversity loss on repr_hidden (if enabled) + diversity_loss = torch.tensor(0.0, device=device) + if self.diversity_weight > 0 and K > 1: + # Average over query and passage batches (both need diverse vectors) + diversity_loss = 0.5 * ( + self.compute_diversity_loss(q_repr['repr_hidden']) + + self.compute_diversity_loss(p_repr['repr_hidden']) + ) + total_loss = total_loss + self.diversity_weight * diversity_loss + + result = { + 'loss': total_loss, + 'loss_dense': (self.dense_weight * colbert_loss).detach(), + 'loss_sparse': sparse_loss.detach(), + } + if self.diversity_weight > 0: + result['loss_diversity'] = diversity_loss.detach() + # loss_denoising is added in forward() since it needs logits + return result + + # ---------------------------------------------------------------- + # K-adapter loss + # ---------------------------------------------------------------- + + def _compute_loss_with_k_adapter( + self, + q_repr: Dict[str, torch.Tensor], + p_repr: Dict[str, torch.Tensor], + labels: torch.Tensor, + ) -> Dict[str, torch.Tensor]: + """Joint training of retriever + K-routing head. + + For each (K_q, K_p) cell in `self.k_adapter_options × self.k_adapter_options`: + 1. Slice prefix of q_repr / p_repr at that K. + 2. Compute MaxSim → fuse with sparse → CE loss vs in-batch negatives. + 3. Per-query loss tensor of shape [B_q, n_K_q, n_K_p]. + + Teacher distribution per query: softmax(-per_cell_loss / τ_T). + Adapter (small MLP on quotation_emb) predicts a softmax over cells. + + Adapter loss = KL(teacher || adapter_probs) — the adapter learns to + mimic which cell the model itself prefers (lowest CE) without any + oracle/MRR involvement. + + Retriever loss = sum_c adapter_probs[i, c] * per_cell_loss[i, c]: + - early in training (uniform adapter): ~ matryoshka-style avg over + all cells, every prefix gets gradient → all K positions trained. + - late in training (peaky adapter): gradient concentrates on + adapter-preferred cell → cell specialization. + """ + device = labels.device + K_opts = self.k_adapter_options + n_K = self.n_K + K_max_q = q_repr['repr_hidden'].size(1) + K_max_p = p_repr['repr_hidden'].size(1) + # If model was instantiated at smaller K than max(k_adapter_options), + # filter the option list down to what we actually have positions for. + valid_q = [k for k in K_opts if k <= K_max_q] + valid_p = [k for k in K_opts if k <= K_max_p] + # Index lookup back into the n_K-wide cell grid (so options outside + # the model's K-range get teacher prob 0 / adapter logits ignored). + idx_q = [K_opts.index(k) for k in valid_q] + idx_p = [K_opts.index(k) for k in valid_p] + + # ── 1. Sparse score (K-invariant — added to every cell uniformly). + # Stays max-pooled so it doesn't depend on K choice. + sparse_scores = None + if (self.sparse_weight > 0 + and 'sparse_acts' in q_repr + and 'sparse_acts' in p_repr): + sparse_scores = q_repr['sparse_acts'] @ p_repr['sparse_acts'].T + # Note: sparse uses raw dot product (no temperature), matching the + # legacy compute_loss path. Adding it BEFORE softmax in CE won't + # change ranking since it shifts every doc by the same per-query + # amount — but it DOES affect per-cell CE values (different cells + # have different dense magnitudes; sparse adds a constant). This + # is intentional: the per-cell loss reflects the actual fused-score + # contrast the model is trained on. + # + # We deliberately apply temperature scaling to the dense part only, + # mirroring how inference fuses dense (MaxSim/K) + sparse_weight*sparse. + B_q = q_repr['repr_hidden'].size(0) + + # ── 2. Per-cell dense scores → per-cell CE losses. + # [B_q, B_p, K_max_q, K_max_p] is the full per-position dot product. + # We compute MaxSim per (K_q, K_p) by slicing prefixes. + per_cell_losses = q_repr['repr_hidden'].new_zeros(B_q, n_K, n_K) + # Precompute the "all positions" inner-product tensor once: + # sims_full[b_q, b_p, i, j] = q_repr[b_q, i] · p_repr[b_q_p_pair, j] + # then maxsim(K_q, K_p) = sims_full[:, :, :K_q, :K_p].max(-1).sum(-1) + sims_full = torch.einsum('bih,djh->bdij', + q_repr['repr_hidden'].float(), + p_repr['repr_hidden'].float()) # [B_q, B_p, K_max_q, K_max_p] + + for ii, K_q in enumerate(valid_q): + ki = idx_q[ii] + for jj, K_p in enumerate(valid_p): + kj = idx_p[jj] + # Slice → max over passage positions → sum over query positions. + cell_dense = sims_full[:, :, :K_q, :K_p].max(dim=-1).values.sum(dim=-1) + cell_dense = cell_dense / K_q # MaxSim normalization, matches legacy + cell_score = self.dense_weight * cell_dense / self.temperature + if sparse_scores is not None: + cell_score = cell_score + self.sparse_weight * sparse_scores + # Per-query CE — reduction='none' for [B_q] tensor, no batch mean. + cell_loss = F.cross_entropy(cell_score, labels, reduction='none') # [B_q] + per_cell_losses[:, ki, kj] = cell_loss + + # ── 3. Teacher distribution: softmax(-loss/τ_T) over cells. + # Lower loss (better separation) → higher teacher weight. + flat_losses = per_cell_losses.view(B_q, n_K * n_K) # [B_q, n_K^2] + teacher_probs = F.softmax( + -flat_losses.detach() / max(self.teacher_temperature, 1e-6), dim=-1 + ) + + # ── 4. Adapter prediction. Input: query's quotation_emb (K-invariant + # content-aware summary). Output: distribution over n_K^2 cells. + if 'quotation_emb' not in q_repr: + raise RuntimeError( + "K-adapter requires 'quotation_emb' in q_repr. " + "encode() must populate it before compute_loss()." + ) + # Match adapter weight dtype (DeepSpeed/AMP may cast adapter to bf16). + # quotation_emb is fp32 from encode(); cast it to whatever the linear + # weight is, run the MLP, then cast logits back to fp32 for stable + # log_softmax + KL arithmetic. + adapter_dtype = next(self.k_adapter.parameters()).dtype + adapter_input = q_repr['quotation_emb'].to(dtype=adapter_dtype) # [B_q, H] + adapter_logits = self.k_adapter(adapter_input).float() # [B_q, n_K^2] + adapter_log_probs = F.log_softmax(adapter_logits, dim=-1) + adapter_probs = adapter_log_probs.exp() + + # ── 5. Adapter loss: KL( teacher || adapter ). + L_adapter = F.kl_div(adapter_log_probs, teacher_probs, reduction='batchmean') + + # ── 6. Retriever loss: soft mixture over cells weighted by adapter. + # Early training: adapter ~uniform → every cell gets gradient + # (matryoshka-style — all K prefixes become useful). + # Late training: adapter peaky → mostly the chosen cell gets + # gradient → specialization at that K. + L_retriever = (adapter_probs * flat_losses).sum(dim=-1).mean() + + total_loss = L_retriever + self.adapter_weight * L_adapter + + # Diagnostics: marginal K_q and K_p distributions averaged over batch. + with torch.no_grad(): + cell_grid = adapter_probs.view(B_q, n_K, n_K) # [B_q, n_K_q, n_K_p] + kq_dist = cell_grid.sum(dim=-1).mean(dim=0) # [n_K] + kp_dist = cell_grid.sum(dim=-2).mean(dim=0) # [n_K] + entropy = -(adapter_probs * adapter_log_probs).sum(dim=-1).mean() + # Standalone sparse-only CE for logging (mirrors the legacy + # `loss_sparse` field). This is the InfoNCE on raw sparse scores + # — it's NOT used for the gradient (sparse already contributes + # via the per-cell dense+sparse fused score). Reported just so + # users can monitor sparse-only retrieval quality during training. + if sparse_scores is not None: + sparse_only_loss = F.cross_entropy(sparse_scores, labels) + else: + sparse_only_loss = torch.tensor(0.0, device=device) + + result = { + 'loss': total_loss, + 'loss_dense': L_retriever.detach(), + 'loss_sparse': sparse_only_loss.detach(), # diagnostic only — not used for gradient + 'loss_adapter': L_adapter.detach(), + 'adapter_entropy': entropy.detach(), + } + # log marginals: kq_p1, kq_p2, kq_p4, kq_p8, kq_p16, kp_p1, ... + for i, K in enumerate(K_opts): + if i < kq_dist.shape[0]: + result[f'kq_p{K}'] = kq_dist[i].detach() + result[f'kp_p{K}'] = kp_dist[i].detach() + return result + + # ---------------------------------------------------------------- + # HF Trainer-compatible forward + # ---------------------------------------------------------------- + + def forward( + self, + query_input_ids: torch.Tensor, + query_attention_mask: torch.Tensor, + passage_input_ids: torch.Tensor, + passage_attention_mask: torch.Tensor, + query_content_ids: Optional[List] = None, + passage_content_ids: Optional[List] = None, + ) -> Dict[str, torch.Tensor]: + """ + passages layout: [pos_0, neg_0_0, …, neg_0_M, pos_1, neg_1_0, …] + positive for query i is at index i * (1 + n_neg). + + num_denoise_steps == 1: single-pass fast path. + num_denoise_steps > 1: multi-step denoising training. + Runs the full denoising loop (each step decodes n_per_step tokens, + storing frozen hidden states in repr_buf), then computes loss once + at the final step. Gradient flows from the final mixed representation + only — no dilution across steps. Uniform unmasking: n_per_step = K // n_steps. + """ + B_q = query_input_ids.size(0) + B_p = passage_input_ids.size(0) + assert B_p % B_q == 0, f"B_p ({B_p}) must be divisible by B_q ({B_q})" + n_paq = B_p // B_q + device = query_input_ids.device + labels = torch.arange(B_q, device=device) * n_paq + n_steps = self.num_denoise_steps + + # ── Corruption augmentation (denoising-conditioned training) ───────── + # Randomly mask text tokens (NOT gen MASKs) in passages to force robust, + # diverse multi-vector representations. Unique to diffusion: the model + # was pretrained on this exact corruption — zero distribution shift. + if self.training and self.corruption_rate > 0: + import random as _rng + t = _rng.uniform(0.0, self.corruption_rate) + if t > 0.01: # skip very small rates + passage_input_ids = self._corrupt_text( + passage_input_ids, passage_attention_mask, t) + + # ── Soft-token multi-step (differentiable denoising) ───────────────── + if self.soft_denoising and n_steps > 1: + return self._soft_multistep_forward( + query_input_ids, query_attention_mask, + passage_input_ids, passage_attention_mask, + query_content_ids, passage_content_ids, + n_steps, + ) + + # ── K pre-encoder (two-stage encoding) ─────────────────────────────── + # Predict K per item BEFORE main encoder, slice each input to its K, + # then run main encoder at variable-length batch. True encoding + # savings (vs the K-adapter which always encodes at K_max). + if self.use_k_pre_encoder: + return self._forward_with_k_pre_encoder( + query_input_ids, query_attention_mask, + passage_input_ids, passage_attention_mask, + query_content_ids, passage_content_ids, + labels, + ) + + # ── Single-step fast path ───────────────────────────────────────────── + if n_steps <= 1: + K_q, K_p = self.n_gen_q_tokens, self.n_gen_p_tokens + n_tail = self._n_tail + need_sparse = self.sparse_weight > 0 + + # Optionally corrupt passages for denoising auxiliary loss + if self.denoising_weight > 0: + p_corrupted, p_denoise_targets, mask_ratio = self._apply_text_masking( + passage_input_ids, passage_attention_mask) + p_ids_for_fwd = p_corrupted + else: + p_ids_for_fwd = passage_input_ids + + need_logits = need_sparse or (self.denoising_weight > 0) + L_q, L_p = query_input_ids.size(1), p_ids_for_fwd.size(1) + + if K_q == K_p: + # ── Symmetric path: concat Q+P into one _fwd (memory-efficient). + if L_q < L_p: + pad = L_p - L_q + q_ids_cat = F.pad(query_input_ids, (pad, 0)) # left-pad + q_mask_cat = F.pad(query_attention_mask, (pad, 0)) + p_ids_cat, p_mask_cat = p_ids_for_fwd, passage_attention_mask + elif L_p < L_q: + pad = L_q - L_p + q_ids_cat, q_mask_cat = query_input_ids, query_attention_mask + p_ids_cat = F.pad(p_ids_for_fwd, (pad, 0)) # left-pad + p_mask_cat = F.pad(passage_attention_mask, (pad, 0)) + else: + q_ids_cat, q_mask_cat = query_input_ids, query_attention_mask + p_ids_cat, p_mask_cat = p_ids_for_fwd, passage_attention_mask + + all_ids = torch.cat([q_ids_cat, p_ids_cat], dim=0) + all_mask = torch.cat([q_mask_cat, p_mask_cat], dim=0) + all_hidden, all_logits = self._fwd(all_ids, all_mask, + need_logits=need_logits) + L_all = all_ids.size(1) + g_q_in_all = L_all - K_q - n_tail + g_p_in_all = L_all - K_p - n_tail + q_hidden, p_hidden = all_hidden[:B_q], all_hidden[B_q:] + q_logits = all_logits[:B_q] if all_logits is not None else None + p_logits = all_logits[B_q:] if all_logits is not None else None + else: + # ── Asymmetric path: K_q != K_p means different gen-block sizes, + # so a concatenated forward would mis-align the slices. Run + # two separate _fwd calls (~2× compute but correct). + q_hidden, q_logits = self._fwd( + query_input_ids, query_attention_mask, + need_logits=need_logits) + p_hidden, p_logits = self._fwd( + p_ids_for_fwd, passage_attention_mask, + need_logits=need_logits) + g_q_in_all = L_q - K_q - n_tail + g_p_in_all = L_p - K_p - n_tail + # No padding/concat: each side keeps its own length. + L_all = max(L_q, L_p) # only used for denoising slice + + # ── Extract Q repr ───────────────────────────────────────────── + q_repr_hidden = q_hidden[:, g_q_in_all:g_q_in_all + K_q, :].float() + q_quotation_emb = q_hidden[:, g_q_in_all - 1, :].float() + q_sparse_max = None + if need_sparse and q_logits is not None: + # Monotonic trick: max in bf16 first → log1p(relu()) on [B, V]. + q_sparse_max = torch.log1p(torch.relu( + q_logits[:, g_q_in_all:g_q_in_all + K_q, :].max(dim=1).values)) + if self.normalize: + q_quotation_emb = F.normalize(q_quotation_emb, p=2, dim=-1) + q_repr_hidden = F.normalize(q_repr_hidden, p=2, dim=-1) + q_repr = {'repr_hidden': q_repr_hidden, 'quotation_emb': q_quotation_emb} + if q_sparse_max is not None: + if query_content_ids is not None: + q_sparse_max = filter_sparse(q_sparse_max, query_content_ids) + q_repr['sparse_acts'] = q_sparse_max + + # ── Extract P repr ───────────────────────────────────────────── + p_repr_hidden = p_hidden[:, g_p_in_all:g_p_in_all + K_p, :].float() + p_quotation_emb = p_hidden[:, g_p_in_all - 1, :].float() + p_sparse_max = None + if need_sparse and p_logits is not None: + # Monotonic trick: max in bf16 first → log1p(relu()) on [B, V]. + p_sparse_max = torch.log1p(torch.relu( + p_logits[:, g_p_in_all:g_p_in_all + K_p, :].max(dim=1).values)) + if self.normalize: + p_quotation_emb = F.normalize(p_quotation_emb, p=2, dim=-1) + p_repr_hidden = F.normalize(p_repr_hidden, p=2, dim=-1) + p_repr = {'repr_hidden': p_repr_hidden, 'quotation_emb': p_quotation_emb} + if p_sparse_max is not None: + if passage_content_ids is not None: + p_sparse_max = filter_sparse(p_sparse_max, passage_content_ids) + p_repr['sparse_acts'] = p_sparse_max + + # ── Cross-GPU negative sharing ──────────────────────────────── + q_repr = self._gather_repr(q_repr) + p_repr = self._gather_repr(p_repr) + # Recompute labels for gathered batch + B_q_all = q_repr['repr_hidden'].size(0) + B_p_all = p_repr['repr_hidden'].size(0) + n_paq_g = B_p_all // B_q_all + labels_g = torch.arange(B_q_all, device=device) * n_paq_g + + # ── Loss ────────────────────────────────────────────────────── + loss_dict = self.compute_loss(q_repr, p_repr, labels_g) + if self.debug_dense_metrics: + loss_dict.update(self._dense_debug_stats( + q_repr, p_repr, labels_g, + q_raw=q_repr_hidden, p_raw=p_repr_hidden, + )) + + # Scale loss to counter DDP gradient averaging + if torch.distributed.is_initialized() and torch.distributed.get_world_size() > 1: + loss_dict['loss'] = loss_dict['loss'] * torch.distributed.get_world_size() + + if self.denoising_weight > 0: + # Denoising targets were created from original passage length; + # if passages were left-padded to match queries, slice + # accordingly. In asymmetric mode (separate forward) no + # padding was applied, so the slice is a no-op (dn_start=0). + p_seq_len = p_logits.size(1) if p_logits is not None else L_p + dn_start = p_seq_len - L_p + p_logits_dn = p_logits[:, dn_start:, :] + denoise_loss = self.compute_denoising_loss( + p_logits_dn, p_denoise_targets, mask_ratio) + loss_dict['loss'] = loss_dict['loss'] + self.denoising_weight * denoise_loss + loss_dict['loss_denoising'] = denoise_loss.detach() + + return loss_dict + + # ── Multi-step: denoising training ─────────────────────────────────── + K_q, K_p = self.n_gen_q_tokens, self.n_gen_p_tokens + n_tail = self._n_tail + n_per_step_q = max(1, K_q // n_steps) + n_per_step_p = max(1, K_p // n_steps) + device = query_input_ids.device + H = self.hidden_size + V = self.backbone.config.vocab_size + L_q = query_input_ids.size(1) + L_p = passage_input_ids.size(1) + + # Pre-compute 4D masks once (only for models that need them) + q_mask_4d = self._build_4d_mask(query_input_ids.size(1), query_attention_mask) if (self.adapter.needs_4d_mask() if self.adapter else True) else None + p_mask_4d = self._build_4d_mask(passage_input_ids.size(1), passage_attention_mask) if (self.adapter.needs_4d_mask() if self.adapter else True) else None + + q_curr = query_input_ids.clone() + p_curr = passage_input_ids.clone() + need_sparse = self.sparse_weight > 0 + + # Frozen buffers for positions already decoded in prior steps + q_repr_buf = torch.zeros(B_q, K_q, H, device=device) + p_repr_buf = torch.zeros(B_p, K_p, H, device=device) + q_decoded_mask = torch.zeros(B_q, K_q, dtype=torch.bool, device=device) + p_decoded_mask = torch.zeros(B_p, K_p, dtype=torch.bool, device=device) + q_sparse_decoded = torch.zeros(B_q, V, device=device) if need_sparse else None + p_sparse_decoded = torch.zeros(B_p, V, device=device) if need_sparse else None + + progressive = self.progressive_step_weight > 0 + + final_loss: Dict[str, torch.Tensor] = {} + progressive_loss_sum = torch.tensor(0.0, device=device) + + def _build_repr_at_step(h, logits, repr_buf, decoded_mask, sparse_decoded, + K_side: int, + content_ids=None, use_fresh=False): + """Build retrieval representation from current hidden states. + Used for both intermediate (progressive) and final steps. + Snapshot buffers with .clone() to isolate from later in-place updates. + + K_side: per-side K (n_gen_q_tokens for queries, n_gen_p_tokens + for passages — needed under asymmetric configs). + + use_fresh: if True, ALL K positions use current hidden states + (ignoring frozen repr_buf). Lets all tokens benefit from + fully-decoded context at the final step. + """ + Ls = h.size(1) + gs = Ls - K_side - n_tail # left-padded: gen block then structural tail + curr_gen_h = h[:, gs:gs + K_side, :].float() # [B_loc, K_side, H] + + decoded_snap = decoded_mask.clone() + repr_snap = repr_buf.detach().clone() + + mask_sparse = None + if need_sparse and logits is not None: + sp_all = torch.log(1.0 + torch.relu(logits[:, gs:gs + K_side, :].float())) + if use_fresh: + mask_sparse = sp_all.max(dim=1).values + else: + is_mask = (~decoded_snap).float().unsqueeze(-1) + mask_sparse = (sp_all * is_mask).max(dim=1).values + + if use_fresh: + mixed = curr_gen_h + else: + mixed = torch.where(decoded_snap.unsqueeze(-1), repr_snap, curr_gen_h) + quotation_emb = h[:, gs - 1, :].float() + + if self.normalize: + mixed = F.normalize(mixed, p=2, dim=-1) + quotation_emb = F.normalize(quotation_emb, p=2, dim=-1) + + result = {'repr_hidden': mixed, 'quotation_emb': quotation_emb} + if need_sparse and sparse_decoded is not None: + sparse_snap = sparse_decoded.detach().clone() + combined_sparse = torch.max(sparse_snap, mask_sparse) + if content_ids is not None: + combined_sparse = filter_sparse(combined_sparse, content_ids) + result['sparse_acts'] = combined_sparse + return result + + q_g = L_q - K_q - n_tail # left-padded: gen block starts here + p_g = L_p - K_p - n_tail + + for step in range(n_steps): + is_last = (step == n_steps - 1) + + q_h, q_logits = self._fwd(q_curr, query_attention_mask, + need_logits=True, mask_4d=q_mask_4d) + p_h, p_logits = self._fwd(p_curr, passage_attention_mask, + need_logits=True, mask_4d=p_mask_4d) + + if progressive or is_last: + _fresh = self.use_fresh_final and is_last + q_repr_s = _build_repr_at_step( + q_h, q_logits, + q_repr_buf, q_decoded_mask, q_sparse_decoded, K_q, + content_ids=query_content_ids, use_fresh=_fresh) + p_repr_s = _build_repr_at_step( + p_h, p_logits, + p_repr_buf, p_decoded_mask, p_sparse_decoded, K_p, + content_ids=passage_content_ids, use_fresh=_fresh) + + # Cross-GPU gather for multi-step + q_repr_s = self._gather_repr(q_repr_s) + p_repr_s = self._gather_repr(p_repr_s) + B_q_g = q_repr_s['repr_hidden'].size(0) + B_p_g = p_repr_s['repr_hidden'].size(0) + labels_ms = torch.arange(B_q_g, device=device) * (B_p_g // B_q_g) + step_loss_dict = self.compute_loss(q_repr_s, p_repr_s, labels_ms) + + if is_last: + final_loss = step_loss_dict + if self.debug_dense_metrics: + final_loss.update(self._dense_debug_stats(q_repr_s, p_repr_s, labels_ms)) + elif progressive: + step_weight = (step + 1) / n_steps + progressive_loss_sum = progressive_loss_sum + step_weight * step_loss_dict['loss'] + + if not is_last: + q_curr, q_newly = self._unmask_step( + q_curr, q_logits, K_q, n_per_step_q) + p_curr, p_newly = self._unmask_step( + p_curr, p_logits, K_p, n_per_step_p) + + with torch.no_grad(): + for i, pos_list in enumerate(q_newly): + for tok_pos in pos_list: + q_repr_buf[i, tok_pos] = q_h[i, q_g + tok_pos].detach() + q_decoded_mask[i, tok_pos] = True + if q_sparse_decoded is not None and q_logits is not None: + q_sparse_decoded[i] = torch.max(q_sparse_decoded[i], + torch.log(1.0 + torch.relu(q_logits[i, q_g + tok_pos])).detach()) + for i, pos_list in enumerate(p_newly): + for tok_pos in pos_list: + p_repr_buf[i, tok_pos] = p_h[i, p_g + tok_pos].detach() + p_decoded_mask[i, tok_pos] = True + if p_sparse_decoded is not None and p_logits is not None: + p_sparse_decoded[i] = torch.max(p_sparse_decoded[i], + torch.log(1.0 + torch.relu(p_logits[i, p_g + tok_pos])).detach()) + + # Add progressive step loss to final loss + if progressive and final_loss: + final_loss['loss'] = final_loss['loss'] + self.progressive_step_weight * progressive_loss_sum + final_loss['loss_progressive'] = progressive_loss_sum.detach() + + # Scale loss for DDP (multi-step path) + if final_loss and torch.distributed.is_initialized() and torch.distributed.get_world_size() > 1: + final_loss['loss'] = final_loss['loss'] * torch.distributed.get_world_size() + + # Denoising auxiliary for multi-step (same as single-step: separate corrupted forward) + if self.denoising_weight > 0 and final_loss: + p_corrupted, p_denoise_targets, mask_ratio = self._apply_text_masking( + passage_input_ids, passage_attention_mask) + _, p_logits_dn = self._fwd(p_corrupted, passage_attention_mask, + need_logits=True, mask_4d=p_mask_4d) + denoise_loss = self.compute_denoising_loss(p_logits_dn, p_denoise_targets, mask_ratio) + final_loss['loss'] = final_loss['loss'] + self.denoising_weight * denoise_loss + final_loss['loss_denoising'] = denoise_loss.detach() + + return final_loss + + # ---------------------------------------------------------------- + # Save / load helpers + # ---------------------------------------------------------------- + + def _save_retriever_config(self, output_dir: str): + import json, os + config = { + 'model_type': self.model_type, + 'mask_token_id': self.mask_token_id, + 'hidden_size': self.hidden_size, + 'max_length': self.max_length, + 'n_gen_tokens': self.n_gen_tokens, + 'n_gen_q_tokens': self.n_gen_q_tokens, + 'n_gen_p_tokens': self.n_gen_p_tokens, + 'temperature': self.temperature, + 'flops_weight': 0.0, # deprecated, kept for backwards compat + 'num_denoise_steps': self.num_denoise_steps, + 'sparse_weight': self.sparse_weight, + 'normalize': self.normalize, + 'query_prefix_ids': self._query_prefix_ids, + 'query_suffix_ids': self._query_suffix_ids, + 'passage_prefix_ids': self._passage_prefix_ids, + 'passage_suffix_ids': self._passage_suffix_ids, + 'lora_rank': getattr(self, 'lora_rank', 0), + 'lora_alpha': getattr(self, 'lora_alpha', 64), + 'use_eos': self.use_eos, + 'n_tail': self._n_tail, + # K-adapter persistence + 'use_k_adapter': self.use_k_adapter, + 'adapter_weight': self.adapter_weight, + 'teacher_temperature': self.teacher_temperature, + 'k_adapter_options': list(self.k_adapter_options), + # K pre-encoder persistence + 'use_k_pre_encoder': self.use_k_pre_encoder, + 'gumbel_temperature': self.gumbel_temperature, + 'k_cost_lambda': self.k_cost_lambda, + 'k_pre_encoder_options': list(self.k_pre_encoder_options), + } + with open(os.path.join(output_dir, 'retriever_config.json'), 'w') as f: + json.dump(config, f, indent=2) + # Save adapter state separately (PEFT save_pretrained only saves LoRA; + # the small KAdapter MLP isn't a LoRA module so we persist it ourselves). + if self.use_k_adapter and self.k_adapter is not None: + torch.save(self.k_adapter.state_dict(), + os.path.join(output_dir, 'k_adapter.bin')) + if self.use_k_pre_encoder and self.k_pre_encoder_q is not None: + torch.save({'q': self.k_pre_encoder_q.state_dict(), + 'p': self.k_pre_encoder_p.state_dict()}, + os.path.join(output_dir, 'k_pre_encoder.bin')) + + @classmethod + def load(cls, model_dir: str, **fallback_kwargs) -> 'TrainableDiffusionRetriever': + """Load a fine-tuned TrainableDiffusionRetriever from a saved directory. + + If retriever_config.json is missing (e.g. mid-training checkpoint or old checkpoint + with matryoshka_config.json), falls back to from_backbone() using fallback_kwargs: + model_type, query_prompt, passage_prompt, n_gen_tokens, + sparse_weight, max_length. + """ + import json + from transformers import AutoTokenizer + + # Support both new name and old name (matryoshka_config.json) for backwards compat + config_path = Path(model_dir) / 'retriever_config.json' + if not config_path.exists(): + config_path = Path(model_dir) / 'matryoshka_config.json' + if not config_path.exists(): + model_type = fallback_kwargs.get('model_type') + query_prompt = fallback_kwargs.get('query_prompt') + passage_prompt = fallback_kwargs.get('passage_prompt') + if not model_type or not query_prompt or not passage_prompt: + raise FileNotFoundError( + f"No retriever_config.json in {model_dir}. " + "Pass model_type, query_prompt, passage_prompt as fallback kwargs.") + + _fallback_adapter = get_adapter(model_type) + source_model = (fallback_kwargs.get('original_model') + or _fallback_adapter.hub_model_name) + logger.info(f"No retriever_config.json — loading architecture from {source_model}, " + f"weights from {model_dir}") + + # Load state dict first so we can detect LoRA rank before building model + checkpoint_dir = Path(model_dir) + weight_file = checkpoint_dir / 'model.safetensors' + bin_files = sorted(checkpoint_dir.glob('model-*.safetensors')) + bin_pt = sorted(checkpoint_dir.glob('pytorch_model*.bin')) + + state_dict = None + if weight_file.exists() or bin_files: + from safetensors.torch import load_file + if weight_file.exists(): + state_dict = load_file(str(weight_file)) + else: + state_dict = {} + for f in bin_files: + state_dict.update(load_file(str(f))) + elif bin_pt: + state_dict = torch.load(str(bin_pt[0]), map_location='cpu') + + # Detect LoRA rank from key names (keys look like backbone.base_model.model.*) + detected_lora_rank = fallback_kwargs.get('lora_rank', 0) + if state_dict is not None and detected_lora_rank == 0: + for k, v in state_dict.items(): + if 'lora_A.default.weight' in k: + detected_lora_rank = v.shape[0] + logger.info(f"Detected LoRA rank {detected_lora_rank} from checkpoint keys") + break + + model = cls.from_backbone( + model_name=source_model, + model_type=model_type, + query_prompt=query_prompt, + passage_prompt=passage_prompt, + n_gen_tokens=fallback_kwargs.get('n_gen_tokens', 4), + n_gen_q_tokens=fallback_kwargs.get('n_gen_q_tokens', None), + n_gen_p_tokens=fallback_kwargs.get('n_gen_p_tokens', None), + temperature=fallback_kwargs.get('temperature', 0.02), + num_denoise_steps=fallback_kwargs.get('num_denoise_steps', None), + sparse_weight=fallback_kwargs.get('sparse_weight', 1.0), + normalize=fallback_kwargs.get('normalize', True), + max_length=fallback_kwargs.get('max_length', 256), + lora_rank=detected_lora_rank, + device_map='auto', + ) + + if state_dict is not None: + # Load full state dict directly into the model (keys have backbone. prefix) + missing, unexpected = model.load_state_dict(state_dict, strict=False) + backbone_keys_loaded = sum(1 for k in state_dict if k.startswith('backbone.')) + logger.info(f"Loaded checkpoint weights from {model_dir} " + f"(lora_rank={detected_lora_rank}, " + f"{backbone_keys_loaded} backbone keys, " + f"{len(missing)} missing, {len(unexpected)} unexpected)") + else: + logger.warning(f"No weight file found in {model_dir} — using original model weights") + + # Merge LoRA weights into base model for faster inference + if detected_lora_rank > 0 and hasattr(model.backbone, 'merge_and_unload'): + model.backbone = model.backbone.merge_and_unload() + logger.info("Merged LoRA adapters for inference") + + return model + + with open(config_path) as f: + cfg = json.load(f) + + model_type = cfg.get('model_type', 'dream') + adapter = get_adapter(model_type) + + tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True) + tokenizer.padding_side = 'left' + + load_weights_separately = False + + # ── Released PEFT adapter repo (the layout published on the Hub) ────── + # A released checkpoint ships only the LoRA adapter + # (adapter_config.json + adapter_model.safetensors/.bin) alongside + # retriever_config.json and the tokenizer — there is no base config.json + # or full-weight file. Detect that and load the base backbone from the + # Hub, then attach the adapter in place. PeftModel.from_pretrained + # restores the trained LoRA weights, so no separate weight load is + # needed and the merge step below folds them in for fast inference. + _adapter_cfg = Path(model_dir) / 'adapter_config.json' + _base_cfg = Path(model_dir) / 'config.json' + _adapter_weights = ((Path(model_dir) / 'adapter_model.safetensors').exists() + or (Path(model_dir) / 'adapter_model.bin').exists()) + if _adapter_cfg.exists() and _adapter_weights and not _base_cfg.exists(): + from peft import PeftModel + logger.info(f"Detected released PEFT adapter in {model_dir} — " + f"loading base {adapter.hub_model_name} from Hub and " + f"attaching the LoRA adapter") + base = adapter.load_backbone(adapter.hub_model_name, device_map='auto') + backbone = PeftModel.from_pretrained(base, str(model_dir)) + else: + try: + backbone = adapter.load_backbone(model_dir, device_map='auto') + except ValueError: + # model_dir has no config.json — load architecture from Hub + logger.info(f"No valid config.json in {model_dir} — " + f"loading architecture from {adapter.hub_model_name}") + backbone = adapter.load_backbone(adapter.hub_model_name, device_map='auto') + load_weights_separately = True + + # Detect if model.safetensors was saved as a full PEFT state dict by DeepSpeed + _ckpt = Path(model_dir) / 'model.safetensors' + if _ckpt.exists(): + from safetensors import safe_open as _so + with _so(str(_ckpt), framework='pt') as _f: + _sample = list(_f.keys())[:8] + if any('base_model.model' in k for k in _sample): + logger.info("Detected PEFT/LoRA state dict — " + "applying LoRA to backbone before loading weights") + _lora_rank = cfg.get('lora_rank', 0) + _lora_alpha = cfg.get('lora_alpha', 64) + if _lora_rank == 0: + from safetensors.torch import load_file as _lf + _peek = _lf(str(_ckpt)) + _la_keys = [k for k in _peek if 'lora_A.default.weight' in k] + if _la_keys: + _lora_rank = _peek[_la_keys[0]].shape[0] + if _lora_rank > 0 and not getattr(backbone, 'peft_config', None): + from peft import get_peft_model + lora_cfg = adapter.get_lora_config(_lora_rank, _lora_alpha) + backbone = get_peft_model(backbone, lora_cfg) + logger.info(f"Applied LoRA rank={_lora_rank} alpha={_lora_alpha}") + elif getattr(backbone, 'peft_config', None): + logger.info("Backbone already has PEFT — skipping duplicate LoRA") + load_weights_separately = True + + model = cls( + backbone=backbone, + tokenizer=tokenizer, + mask_token_id=cfg['mask_token_id'], + hidden_size=cfg['hidden_size'], + query_prefix_ids=cfg['query_prefix_ids'], + query_suffix_ids=cfg['query_suffix_ids'], + passage_prefix_ids=cfg['passage_prefix_ids'], + passage_suffix_ids=cfg['passage_suffix_ids'], + max_length=cfg['max_length'], + n_gen_tokens=cfg['n_gen_tokens'], + n_gen_q_tokens=cfg.get('n_gen_q_tokens'), # back-compat: missing → fall through to n_gen_tokens + n_gen_p_tokens=cfg.get('n_gen_p_tokens'), + temperature=cfg['temperature'], + num_denoise_steps=cfg['num_denoise_steps'], + sparse_weight=cfg['sparse_weight'], + normalize=cfg['normalize'], + flash_attn=adapter.flash_attn, + use_eos=cfg.get('use_eos', False), + # K-adapter restore (back-compat: keys missing → defaults disable) + use_k_adapter=cfg.get('use_k_adapter', False), + adapter_weight=cfg.get('adapter_weight', 1.0), + teacher_temperature=cfg.get('teacher_temperature', 1.0), + k_adapter_options=tuple(cfg['k_adapter_options']) + if cfg.get('k_adapter_options') is not None else None, + # K pre-encoder restore + use_k_pre_encoder=cfg.get('use_k_pre_encoder', False), + gumbel_temperature=cfg.get('gumbel_temperature', 1.0), + k_cost_lambda=cfg.get('k_cost_lambda', 0.001), + k_pre_encoder_options=tuple(cfg['k_pre_encoder_options']) + if cfg.get('k_pre_encoder_options') is not None else None, + ) + # Load adapter weights if a saved file exists. + if model.use_k_adapter and model.k_adapter is not None: + adapter_path = Path(model_dir) / 'k_adapter.bin' + if adapter_path.exists(): + model.k_adapter.load_state_dict( + torch.load(str(adapter_path), map_location='cpu')) + logger.info(f"Loaded KAdapter state from {adapter_path}") + else: + logger.warning(f"use_k_adapter=True but {adapter_path} missing; " + f"adapter starts from random init.") + # Load pre-encoder weights if a saved file exists. + if model.use_k_pre_encoder and model.k_pre_encoder_q is not None: + pe_path = Path(model_dir) / 'k_pre_encoder.bin' + if pe_path.exists(): + pe_state = torch.load(str(pe_path), map_location='cpu') + model.k_pre_encoder_q.load_state_dict(pe_state['q']) + model.k_pre_encoder_p.load_state_dict(pe_state['p']) + logger.info(f"Loaded KPreEncoder state from {pe_path}") + else: + logger.warning(f"use_k_pre_encoder=True but {pe_path} missing; " + f"pre-encoder starts from random init.") + # Restore _n_tail from config. Old checkpoints (before structural tail) + # don't have 'n_tail' — infer from use_eos for backward compat. + if 'n_tail' in cfg: + model._n_tail = cfg['n_tail'] + model._tail_ids = model._build_tail_ids() + else: + # Legacy: use_eos=True → 1 tail token (EOS only), use_eos=False → 0 + model._n_tail = 1 if cfg.get('use_eos', False) else 0 + eos_id = model.tokenizer.eos_token_id + model._tail_ids = [eos_id] if model._n_tail == 1 else [] + logger.info(f"Legacy checkpoint: n_tail={model._n_tail} (from use_eos={cfg.get('use_eos', False)})") + model.model_type = model_type + model.adapter = adapter + + # Hook for efficient hidden state extraction + model._last_hidden: Dict[str, torch.Tensor] = {} + model._hook_registered = adapter.register_hidden_hook( + backbone, model._last_hidden) + + if load_weights_separately: + checkpoint_dir = Path(model_dir) + weight_file = checkpoint_dir / 'model.safetensors' + bin_files = sorted(checkpoint_dir.glob('model-*.safetensors')) + bin_pt = sorted(checkpoint_dir.glob('pytorch_model*.bin')) + if weight_file.exists() or bin_files: + from safetensors.torch import load_file + state_dict = load_file(str(weight_file)) if weight_file.exists() else {} + for f in bin_files: + state_dict.update(load_file(str(f))) + backbone_dict = {k[len('backbone.'):]: v for k, v in state_dict.items() + if k.startswith('backbone.')} + model.backbone.load_state_dict(backbone_dict or state_dict, strict=False) + logger.info(f"Loaded fine-tuned weights from {model_dir}") + elif bin_pt: + sd = torch.load(str(bin_pt[0]), map_location='cpu') + backbone_dict = {k[len('backbone.'):]: v for k, v in sd.items() + if k.startswith('backbone.')} + model.backbone.load_state_dict(backbone_dict or sd, strict=False) + logger.info(f"Loaded fine-tuned weights from {model_dir}") + else: + logger.warning(f"No weight file found in {model_dir} — using Hub weights") + + # Merge LoRA adapters into the base weights for inference. Training + # keeps them separate for gradient flow; at inference each LoRA adds + # per-layer overhead otherwise. + # + # Two shapes seen in saved checkpoints: + # (a) `model.backbone` is a PeftModel wrapper → call + # `merge_and_unload()` directly. + # (b) `model.backbone` is the underlying DreamModel / LLaDAModel + # with `LoraLinear` modules still in place (the top-level + # PeftModel wrapper was dropped during save). `merge_and_unload` + # is missing but each inner `LoraLinear` still exposes `.merge()` + # which folds its `lora_A @ lora_B` into `base_layer.weight` in + # place and flips `merged=True`, disabling the LoRA forward path. + # We walk the module tree and merge case-by-case. + if hasattr(model.backbone, 'merge_and_unload'): + try: + model.backbone = model.backbone.merge_and_unload() + logger.info("Merged LoRA adapters into base backbone for inference") + except Exception as exc: + logger.warning(f"merge_and_unload failed ({exc}); " + f"continuing with un-merged LoRA (slower inference)") + else: + merged_modules = 0 + for _mod in model.backbone.modules(): + if (hasattr(_mod, 'lora_A') and hasattr(_mod, 'merge') + and callable(getattr(_mod, 'merge', None)) + and not getattr(_mod, 'merged', False)): + try: + _mod.merge() + merged_modules += 1 + except Exception as exc: + logger.warning(f"per-module LoRA merge failed ({exc}) " + f"— continuing with un-merged layer") + if merged_modules > 0: + logger.info(f"Merged {merged_modules} LoraLinear modules into " + f"base weights for inference") + + logger.info(f"Loaded TrainableDiffusionRetriever from {model_dir}") + return model + + def save(self, output_dir: str): + import os + os.makedirs(output_dir, exist_ok=True) + + backbone = self.backbone + if hasattr(backbone, 'save_pretrained'): + backbone.save_pretrained(output_dir) + else: + backbone.base_model.save_pretrained(output_dir) + self.tokenizer.save_pretrained(output_dir) + self._save_retriever_config(output_dir) + logger.info(f"Saved to {output_dir}") + + @property + def config(self): + # HF Trainer / DeepSpeed expect model.config to exist. + return self.backbone.config + + def gradient_checkpointing_enable(self, **kwargs): + if self.adapter: + self.adapter.enable_gradient_checkpointing(self.backbone, **kwargs) + else: + self.backbone.gradient_checkpointing_enable(**kwargs)