# -*- coding: utf-8 -*- from __future__ import annotations from typing import Callable, Sequence import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.checkpoint import checkpoint from Model.blocks import StandardBlock from Model.config import RDTConfig from Model.layers.rmsnorm import RMSNorm from Model.recurrent import RecurrentCore from Model.segmented import SegmentedCore from Model.two_stage import TwoStageCore from Model.vision import VisionInjector class RDTForCausalLM(nn.Module): def __init__(self, cfg: RDTConfig, patch_pixels: int = 14 * 14 * 3): super().__init__() self.cfg = cfg self.patch_pixels = patch_pixels self.embed = nn.Embedding( cfg.vocab_size, cfg.d_model, padding_idx=cfg.pad_id, ) self.vision = VisionInjector(cfg, patch_pixels) self.prelude = nn.ModuleList( StandardBlock(cfg, layer_idx=i) for i in range(cfg.n_prelude) ) if cfg.core_type == "two_stage": self.recurrent = TwoStageCore(cfg) elif cfg.core_type == "segmented": self.recurrent = SegmentedCore(cfg) else: self.recurrent = RecurrentCore(cfg) self.coda = nn.ModuleList( StandardBlock(cfg, layer_idx=cfg.n_prelude + i) for i in range(cfg.n_coda) ) self.final_norm = RMSNorm(cfg.d_model, eps=cfg.rmsnorm_eps) self.lm_head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False) self.bidirectional = cfg.bidirectional if self.bidirectional: self.reverse_head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False) else: self.reverse_head = None # Runtime toggle: alignment phases (SFT/DPO/GRPO) disable the # bidirectional reverse-LM auxiliary loss so it does not pollute the # supervised/preference gradient. Pretraining keeps it enabled. self.reverse_loss_enabled = True self.apply(self._init_weights) self._scale_residual_projections() if cfg.tie_word_embeddings: self.lm_head.weight = self.embed.weight if self.reverse_head is not None: self.reverse_head.weight = self.embed.weight def _forward_decode( self, input_ids: torch.Tensor, word_pos: torch.Tensor, morph_depth: torch.Tensor, cache, pixel_values: torch.Tensor | None = None, steps: int | None = None, ) -> torch.Tensor: """Incremental forward over the new tokens ``input_ids`` (``[B, m]``). ``word_pos`` / ``morph_depth`` are the *absolute* per-position values of the new tokens (computed by the caller from the full running sequence). Reuses ``cache`` so the result is bit-exact with a full :meth:`forward` over the whole prefix. Returns logits ``[B, m, vocab]``; callers typically read the last position. Generation is pad-free, so ``attn_mask`` is omitted throughout. """ if self.cfg.core_type not in {"two_stage", "segmented"}: raise NotImplementedError( "incremental KV/state cache is implemented for core_type=" "'two_stage' / 'segmented' only" ) bsz, seq_len = input_ids.shape pos_offset = cache.seq_len h = self.embed(input_ids) if pixel_values is not None: h = self.vision(h, input_ids, pixel_values) for i, block in enumerate(self.prelude): h = block( h, word_pos=word_pos, morph_depth=morph_depth, attn_mask=None, causal=True, cache=cache.mla_cache(f"prelude.{i}"), pos_offset=pos_offset, ) h, _rec_info = self.recurrent( h, word_pos=word_pos, morph_depth=morph_depth, attn_mask=None, causal=True, cache=cache, pos_offset=pos_offset, steps=steps, ) for i, block in enumerate(self.coda): h = block( h, word_pos=word_pos, morph_depth=morph_depth, attn_mask=None, causal=True, cache=cache.mla_cache(f"coda.{i}"), pos_offset=pos_offset, ) h = self.final_norm(h) logits = self.lm_head(h) cache.seq_len = pos_offset + seq_len return logits def forward( self, input_ids: torch.Tensor, attention_mask: torch.Tensor | None = None, labels: torch.Tensor | None = None, pixel_values: torch.Tensor | None = None, word_pos: torch.Tensor | None = None, morph_depth: torch.Tensor | None = None, steps: int | None = None, bptt_window: int | None = None, return_logits: bool = True, loss_chunk_size: int | None = None, ) -> dict[str, torch.Tensor | dict | None]: self._check_inputs(input_ids, attention_mask, labels) if loss_chunk_size is None: loss_chunk_size = self.cfg.loss_chunk_size elif loss_chunk_size <= 0: raise ValueError("loss_chunk_size must be positive") bsz, seq_len = input_ids.shape if attention_mask is None: attention_mask = (input_ids != self.cfg.pad_id).long() if word_pos is None or morph_depth is None: word_pos, morph_depth = self._default_morph_info( input_ids=input_ids, attention_mask=attention_mask, ) h = self.embed(input_ids) if pixel_values is not None: h = self.vision(h, input_ids, pixel_values) for block in self.prelude: h = self._maybe_ckpt( block, h, word_pos=word_pos, morph_depth=morph_depth, attn_mask=attention_mask, causal=True, ) e0 = h h, rec_info = self.recurrent( e0, word_pos=word_pos, morph_depth=morph_depth, attn_mask=attention_mask, causal=True, steps=steps, bptt_window=bptt_window, ) for block in self.coda: h = self._maybe_ckpt( block, h, word_pos=word_pos, morph_depth=morph_depth, attn_mask=attention_mask, causal=True, ) h = self.final_norm(h) logits = None loss = None loss_parts: dict[str, float] = {} if labels is not None: if not return_logits and loss_chunk_size is not None: loss, loss_parts = self._losses_chunked( h, labels, rec_info, loss_chunk_size, ) else: logits = self.lm_head(h) loss, loss_parts = self._losses(h, logits, labels, rec_info) elif return_logits: logits = self.lm_head(h) return { "loss": loss, "logits": logits if return_logits else None, "loss_parts": loss_parts, "rec_info": rec_info, } def _losses( self, h: torch.Tensor, logits: torch.Tensor, labels: torch.Tensor, rec_info: dict, ) -> tuple[torch.Tensor, dict[str, float]]: forward = self._causal_loss(logits, labels) loss = forward parts = {"forward": float(forward.detach())} if self.reverse_head is not None and self.reverse_loss_enabled: rev_logits = self.reverse_head(h) reverse = self._reverse_loss(rev_logits, labels) loss = loss + self.cfg.reverse_loss_weight * reverse parts["reverse"] = float(reverse.detach()) ponder = rec_info.get("ponder_cost") if self.cfg.use_act and isinstance(ponder, torch.Tensor): loss = loss + self.cfg.act_ponder_cost * ponder parts["ponder"] = float(ponder.detach()) return loss, parts def _losses_chunked( self, h: torch.Tensor, labels: torch.Tensor, rec_info: dict, chunk_size: int, ) -> tuple[torch.Tensor, dict[str, float]]: forward = self._chunked_causal_loss(h, labels, self.lm_head, chunk_size) loss = forward parts = {"forward": float(forward.detach())} if self.reverse_head is not None and self.reverse_loss_enabled: reverse = self._chunked_reverse_loss( h, labels, self.reverse_head, chunk_size, ) loss = loss + self.cfg.reverse_loss_weight * reverse parts["reverse"] = float(reverse.detach()) ponder = rec_info.get("ponder_cost") if self.cfg.use_act and isinstance(ponder, torch.Tensor): loss = loss + self.cfg.act_ponder_cost * ponder parts["ponder"] = float(ponder.detach()) return loss, parts def _causal_loss(self, logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor: return self._token_loss( logits[:, :-1].reshape(-1, logits.size(-1)), labels[:, 1:].reshape(-1), ) def _reverse_loss( self, rev_logits: torch.Tensor, labels: torch.Tensor, ) -> torch.Tensor: return self._token_loss( rev_logits[:, 1:].reshape(-1, rev_logits.size(-1)), labels[:, :-1].reshape(-1), ) def _token_loss(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor: valid = targets != self.cfg.ignore_index if logits.numel() == 0 or not bool(valid.any()): return logits.sum() * 0.0 return F.cross_entropy( logits, targets, ignore_index=self.cfg.ignore_index, ) def _chunked_causal_loss( self, h: torch.Tensor, labels: torch.Tensor, head: nn.Linear, chunk_size: int, ) -> torch.Tensor: return self._chunked_token_loss( h[:, :-1].reshape(-1, h.size(-1)), labels[:, 1:].reshape(-1), head, chunk_size, ) def _chunked_reverse_loss( self, h: torch.Tensor, labels: torch.Tensor, head: nn.Linear, chunk_size: int, ) -> torch.Tensor: return self._chunked_token_loss( h[:, 1:].reshape(-1, h.size(-1)), labels[:, :-1].reshape(-1), head, chunk_size, ) def _chunked_token_loss( self, hidden: torch.Tensor, targets: torch.Tensor, head: nn.Linear, chunk_size: int, ) -> torch.Tensor: valid = targets != self.cfg.ignore_index if hidden.numel() == 0 or not bool(valid.any()): return hidden.sum() * 0.0 hidden = hidden[valid] targets = targets[valid] loss_sum = hidden.new_zeros(()) for start in range(0, hidden.size(0), chunk_size): end = min(start + chunk_size, hidden.size(0)) logits = F.linear(hidden[start:end], head.weight, head.bias) loss_sum = loss_sum + F.cross_entropy( logits, targets[start:end], reduction="sum", ) return loss_sum / targets.numel() def _maybe_ckpt( self, block: nn.Module, h: torch.Tensor, word_pos: torch.Tensor | None, morph_depth: torch.Tensor | None, attn_mask: torch.Tensor | None, causal: bool, ) -> torch.Tensor: if ( getattr(self.cfg, "grad_ckpt_prelude_coda", False) and self.training and h.requires_grad ): def _fn(x): return block( x, word_pos=word_pos, morph_depth=morph_depth, attn_mask=attn_mask, causal=causal, ) return checkpoint(_fn, h, use_reentrant=False) return block( h, word_pos=word_pos, morph_depth=morph_depth, attn_mask=attn_mask, causal=causal, ) def _default_morph_info( self, input_ids: torch.Tensor, attention_mask: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: """Derive ``(word_pos, morph_depth)`` from boundary token IDs. This is a vectorized counterpart of :func:`Tokenizer.pretraining.derive_morph_info_from_boundary_ids`. Callers that already supply ``word_pos`` / ``morph_depth`` skip this. Semantics (see tokenizer-side reference for details): * ``word_boundary_id`` opens a new word at depth 0. * ``morpheme_boundary_id`` keeps the current word and bumps depth. * Other special tokens (ids in ``[0, 256)``) reset depth to 0 and stay on the current ``word_pos``; the next non-special content token will inherit ``word_pos`` until a new ``word_boundary``. * Padding positions (``attention_mask == 0``) keep their derived ``word_pos`` so downstream RoPE never sees ``-1``; the mask itself is what zeros out padded contributions. """ device = input_ids.device bsz, seq_len = input_ids.shape cfg = self.cfg wb = int(cfg.word_boundary_id) mb = int(cfg.morpheme_boundary_id) special_hi = 256 is_wb = input_ids == wb is_mb = input_ids == mb is_special = (input_ids >= 0) & (input_ids < special_hi) is_content = ~is_special is_other_special = is_special & ~is_wb & ~is_mb # word_pos: cumulative count of word_boundary occurrences with a # per-row shift. If the very first word_boundary appears before # any content token, the first wb anchors word 0 (shift = -1). # Otherwise content tokens implicitly open word 0 and the first # wb opens word 1 (shift = 0). wb_cum = is_wb.long().cumsum(dim=1) inf = seq_len + 1 any_wb = is_wb.any(dim=1) any_content = is_content.any(dim=1) first_wb = torch.where( any_wb, is_wb.long().argmax(dim=1), torch.full((bsz,), inf, device=device, dtype=torch.long), ) first_content = torch.where( any_content, is_content.long().argmax(dim=1), torch.full((bsz,), inf, device=device, dtype=torch.long), ) shift = torch.where( first_wb < first_content, torch.full((bsz,), -1, device=device, dtype=torch.long), torch.zeros(bsz, device=device, dtype=torch.long), ) word_pos = (wb_cum + shift.unsqueeze(-1)).clamp(min=0) # morph_depth: cumulative morpheme_boundary count since the most # recent reset (word_boundary or other special). The reset # position itself reads depth 0. cum_inc = is_mb.long().cumsum(dim=1) reset_positions = is_wb | is_other_special reset_value = torch.where( reset_positions, cum_inc, torch.full_like(cum_inc, -1) ) last_reset, _ = reset_value.cummax(dim=1) depth = cum_inc - last_reset.clamp(min=0) depth = torch.where(reset_positions, torch.zeros_like(depth), depth) if cfg.max_morph_depth > 0: depth = depth.clamp(max=cfg.max_morph_depth - 1) return word_pos.to(torch.long), depth.to(torch.long) def _check_inputs( self, input_ids: torch.Tensor, attention_mask: torch.Tensor | None, labels: torch.Tensor | None, ) -> None: if input_ids.ndim != 2: raise ValueError("input_ids must have shape [B, L]") if input_ids.numel() == 0: raise ValueError("input_ids cannot be empty") # Dtype + shape checks only; value-range checks are skipped here # to avoid host syncs on the hot path. The embedding lookup will # raise an out-of-range error if vocab bounds are violated, and # we rely on the data pipeline to keep ids well-formed. if input_ids.dtype not in (torch.int32, torch.int64): raise TypeError("input_ids must be int32 or int64") if attention_mask is not None and attention_mask.shape != input_ids.shape: raise ValueError("attention_mask must have shape [B, L]") if labels is not None and labels.shape != input_ids.shape: raise ValueError("labels must have shape [B, L]") if input_ids.shape[1] > self.cfg.max_seq_len: raise ValueError("sequence length exceeds max_seq_len") def _init_weights(self, module: nn.Module) -> None: std = self.cfg.init_std if isinstance(module, nn.Linear): nn.init.normal_(module.weight, mean=0.0, std=std) if module.bias is not None: nn.init.zeros_(module.bias) elif isinstance(module, nn.Embedding): nn.init.normal_(module.weight, mean=0.0, std=std) if module.padding_idx is not None: with torch.no_grad(): module.weight[module.padding_idx].zero_() # Last linear of every residual write: attention output, SwiGLU down # projection, and the SSM output projection (both NaiveSSM and the # official Mamba3 expose it as ``out_proj``). _RESIDUAL_PROJ_SUFFIXES = ("attn.o_proj", "ffn.w_down", "mamba.out_proj") @torch.no_grad() def _scale_residual_projections(self) -> None: """GPT-2-style depth-scaled init, extended to the unrolled loop. Every sublayer writes ``x + f(x)`` into the residual stream. With a weight-shared recurrent core the stream sees ``effective_depth`` layers (prelude + block_layers x recurrent_steps + coda), so leaving the output projections at ``init_std`` makes the hidden-state norm grow linearly across recurrent steps (and costs bf16 mantissa precision late in the loop). Scaling them by ``1/sqrt(2 * effective_depth)`` keeps the residual variance roughly depth-invariant at init. """ scale = (2.0 * self.cfg.effective_depth) ** -0.5 for container in (self.prelude, self.recurrent, self.coda): for name, module in container.named_modules(): if isinstance(module, nn.Linear) and name.endswith( self._RESIDUAL_PROJ_SUFFIXES ): module.weight.mul_(scale) def _kl_exit_active(self, recurrent_steps: int | None) -> bool: """Whether zero-shot KL early-exit governs this decode step. Gated on ``cfg.kl_exit_threshold > 0`` and only on the cache-free path with no explicit ``recurrent_steps`` override; when disabled the decode is bit-exact with a single fixed-depth forward. """ return ( self.cfg.kl_exit_threshold > 0.0 and recurrent_steps is None and self.cfg.core_type in {"two_stage", "segmented"} ) @torch.no_grad() def _adaptive_depth_logits( self, window: torch.Tensor, pixel_values: torch.Tensor | None, ) -> torch.Tensor: """Last-position logits with per-token adaptive recurrent depth. Runs the refinement at increasing depth ``d = 1 .. recurrent_steps`` and stops as soon as the symmetric-free KL between the depth-``d`` and depth-``d-1`` last-token distributions drops below ``cfg.kl_exit_threshold`` for every row (Huginn arXiv:2502.05171 §6.1, zero-shot convergence of the latent recurrence). Saves test-time compute on "easy" tokens while spending full depth on hard / rare ones. Always returns the deepest distribution actually computed. """ max_steps = self.cfg.recurrent_steps threshold = self.cfg.kl_exit_threshold prev_logp = None logits = None for d in range(1, max_steps + 1): if pixel_values is not None: out = self.forward( window, steps=d, return_logits=True, pixel_values=pixel_values ) else: out = self.forward(window, steps=d, return_logits=True) logits = out["logits"][:, -1, :].float() logp = F.log_softmax(logits, dim=-1) if prev_logp is not None: # KL(p_d || p_{d-1}) per row; exit when all rows have converged. kl = (logp.exp() * (logp - prev_logp)).sum(dim=-1) if bool((kl < threshold).all()): break prev_logp = logp return logits @torch.no_grad() def generate( self, input_ids: torch.Tensor, max_new_tokens: int = 64, temperature: float = 1.0, top_k: int | None = None, top_p: float | None = None, min_p: float | None = None, greedy: bool = False, eos_id: int | None = None, pad_id: int | None = None, repetition_penalty: float = 1.0, use_cache: bool = False, stop_ids: Sequence[int] | None = None, on_token: Callable[[int, torch.Tensor], None] | None = None, recurrent_steps: int | None = None, pixel_values: torch.Tensor | None = None, ) -> torch.Tensor: """Autoregressively continue ``input_ids`` (``[B, L]``) with sampling. Sampling controls: ``temperature`` then ``min_p`` (ICLR 2025), ``top_k`` and nucleus ``top_p`` truncation, plus CTRL-style ``repetition_penalty``; ``greedy=True`` takes the argmax. With ``use_cache=True`` (``core_type='two_stage'`` only) decoding runs an incremental KV/state cache that is numerically identical to the cache- free path but processes one token per step instead of re-running the whole prefix: * **Stage 1 (Mamba)** keeps a constant-size ``(conv_state, ssm_state)`` per layer and steps the selective SSM once per new token -- O(1) time and memory per step (Gu & Dao, "Mamba", 2023/2024). * **Stage 2 (MLA)** appends the new token's ``K``/``V`` to a per-layer cache and attends over the frozen prefix; because each refinement pass is causal and earlier positions are frozen once computed, one cache is kept per ``(refinement step, layer)`` pair. The cached path keeps the full context (no sliding-window eviction), so ``L + max_new_tokens`` must stay within ``max_seq_len``; the cache-free path instead slides a ``max_seq_len`` window. ``use_cache`` is only supported for ``core_type='two_stage'``. Returns the full sequence ``[B, L + n]`` where ``n <= max_new_tokens``. Generation stops early for a row once it emits ``eos_id`` (subsequent positions are filled with ``pad_id``). ``stop_ids`` adds extra stop tokens: a row also finishes once it emits any id in ``stop_ids`` (the stop token itself is kept in the output, so a harness can see *which* delimiter — e.g. ```` or a tool-call close token — ended the segment, then resume by calling ``generate`` again on the returned sequence). ``eos_id`` is always a stop token. ``on_token`` is an optional callback invoked as ``on_token(step, tok)`` after each decoding step with the step index and the ``[B]`` tensor of ids that were just appended to the sequence; use it for streaming. Note rows that have already finished receive ``pad_id`` (not a freshly sampled id), matching exactly what is written to the output. Both are backward compatible: when unset, behaviour is identical to before. ``recurrent_steps`` overrides the recurrent-depth refinement count for this decode call (default ``None`` -> ``cfg.recurrent_steps``). Because RDT reasons in latent depth, raising this spends more test-time compute ("think harder") per token without emitting any extra tokens; lowering it trades quality for speed. The override is constant for the whole call so the incremental KV/state cache stays consistent across positions. ``pixel_values`` optionally supplies image features for prompts containing ```` slots. Cached decoding consumes them only during the prefill step; cache-free decoding refuses sliding-window truncation with images because dropping patch slots would desync the visual payload. """ if input_ids.dim() != 2: raise ValueError("input_ids must have shape [B, L]") if max_new_tokens < 0: raise ValueError("max_new_tokens must be non-negative") if temperature <= 0: raise ValueError("temperature must be positive") if top_k is not None and top_k <= 0: raise ValueError("top_k must be positive when set") if top_p is not None and not (0.0 < top_p <= 1.0): raise ValueError("top_p must be in (0, 1] when set") if min_p is not None and not (0.0 < min_p <= 1.0): raise ValueError("min_p must be in (0, 1] when set") if repetition_penalty <= 0: raise ValueError("repetition_penalty must be positive") if recurrent_steps is not None and recurrent_steps <= 0: raise ValueError("recurrent_steps must be positive when set") if use_cache and self.cfg.core_type not in {"two_stage", "segmented"}: raise NotImplementedError( "use_cache=True is only supported for core_type=" "'two_stage' / 'segmented'" ) if use_cache and self.cfg.use_official_mamba: raise NotImplementedError( "use_cache=True requires the NaiveSSM fallback backend; " "official Mamba kernels are not steppable by the decode cache. " "Pass --mamba=naive for cached decoding or use_cache=False." ) cfg = self.cfg eos_id = cfg.eos_id if eos_id is None else eos_id pad_id = cfg.pad_id if pad_id is None else pad_id was_training = self.training self.eval() seq = input_ids device = seq.device finished = torch.zeros(seq.shape[0], dtype=torch.bool, device=device) stop_token_ids: list[int] = [] if eos_id is not None: stop_token_ids.append(int(eos_id)) if stop_ids is not None: stop_token_ids.extend(int(s) for s in stop_ids) stop_tensor = ( torch.tensor(sorted(set(stop_token_ids)), device=device, dtype=seq.dtype) if stop_token_ids else None ) decode_cache = None if use_cache: from Model.inference.cache import DecodeCache # The incremental cache prefills every prompt position into the # MLA/Mamba state with an all-ones mask, so a padded prompt would # fold pad embeddings into the cache and change logits for shorter # rows. Reject padded prompts (use the cache-free path for those). if pad_id is not None and bool((input_ids == pad_id).any()): raise ValueError( "use_cache=True requires pad-free prompts; the incremental " "cache cannot mask padding. Use use_cache=False for padded " "batches." ) # The cache has no eviction, so the whole run must fit in context. # Guard up front (counting the tokens about to be appended) instead # of generating one token past the limit before raising. if input_ids.shape[1] + max_new_tokens > cfg.max_seq_len: raise ValueError( "cached generation requires L + max_new_tokens <= " f"max_seq_len ({cfg.max_seq_len}); got L={input_ids.shape[1]}" f" + max_new_tokens={max_new_tokens}. Reduce max_new_tokens " "or use use_cache=False." ) decode_cache = DecodeCache() try: for step in range(max_new_tokens): if use_cache: if decode_cache.seq_len == 0: step_ids = seq step_pixels = pixel_values else: step_ids = seq[:, -1:] # The image was folded into the cache at prefill; later # steps process only the new token and carry no # slots, so pixel_values must be dropped. step_pixels = None mask = (seq != pad_id).long() word_pos, morph_depth = self._default_morph_info(seq, mask) m = step_ids.shape[1] logits = self._forward_decode( step_ids, word_pos=word_pos[:, -m:], morph_depth=morph_depth[:, -m:], cache=decode_cache, pixel_values=step_pixels, steps=recurrent_steps, )[:, -1, :].float() else: window = seq if window.shape[1] > cfg.max_seq_len: window = window[:, -cfg.max_seq_len:] if pixel_values is not None: # Left-truncation could drop slots while # pixel_values still holds the full visual payload, # desyncing the injector. Refuse loudly. raise ValueError( "cache-free image generation cannot truncate the " "context (L + generated tokens exceeded " "max_seq_len); the slots would " "desync from pixel_values. Use use_cache=True or " "keep L + max_new_tokens <= max_seq_len." ) if self._kl_exit_active(recurrent_steps): logits = self._adaptive_depth_logits( window, pixel_values ) elif pixel_values is not None: out = self.forward( window, steps=recurrent_steps, return_logits=True, pixel_values=pixel_values, ) logits = out["logits"][:, -1, :].float() else: out = self.forward( window, steps=recurrent_steps, return_logits=True ) logits = out["logits"][:, -1, :].float() if repetition_penalty != 1.0: logits = self._apply_repetition_penalty( logits, seq, repetition_penalty ) if greedy: next_token = torch.argmax(logits, dim=-1) else: logits = logits / temperature logits = self._filter_logits(logits, top_k, top_p, min_p) probs = F.softmax(logits, dim=-1) next_token = torch.multinomial(probs, num_samples=1).squeeze(-1) next_token = torch.where( finished, torch.full_like(next_token, pad_id), next_token ) next_token = next_token.to(seq.dtype) seq = torch.cat([seq, next_token.unsqueeze(1)], dim=1) if on_token is not None: on_token(step, next_token) if stop_tensor is not None: finished = finished | torch.isin(next_token, stop_tensor) else: finished = finished | (next_token == eos_id) if bool(finished.all()): break finally: if was_training: self.train() return seq @staticmethod def _apply_repetition_penalty( logits: torch.Tensor, seq: torch.Tensor, penalty: float ) -> torch.Tensor: """Divide logits of already-seen tokens by ``penalty`` (CTRL-style).""" for row in range(seq.shape[0]): seen = torch.unique(seq[row]) row_logits = logits[row, seen] logits[row, seen] = torch.where( row_logits > 0, row_logits / penalty, row_logits * penalty ) return logits @staticmethod def _filter_logits( logits: torch.Tensor, top_k: int | None, top_p: float | None, min_p: float | None = None, ) -> torch.Tensor: """Apply top-k, nucleus (top-p) and min-p masking to ``[B, V]`` logits. min-p (Nguyen et al., "Turning Up the Heat: Min-p Sampling for Creative and Coherent LLM Outputs", ICLR 2025) keeps only tokens whose probability is at least ``min_p * p_max`` where ``p_max`` is the top token's probability. The candidate pool scales with the model's own confidence: sharp distributions prune hard, flat ones stay permissive. It is applied before top-k/top-p so it can act as the primary truncation. """ if min_p is not None: probs = F.softmax(logits, dim=-1) p_max = probs.max(dim=-1, keepdim=True).values logits = logits.masked_fill(probs < min_p * p_max, float("-inf")) if top_k is not None: k = min(top_k, logits.shape[-1]) kth = torch.topk(logits, k, dim=-1).values[:, -1, None] logits = logits.masked_fill(logits < kth, float("-inf")) if top_p is not None: sorted_logits, sorted_idx = torch.sort(logits, descending=True, dim=-1) cum_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1) remove = cum_probs > top_p remove[:, 1:] = remove[:, :-1].clone() remove[:, 0] = False remove_idx = remove.scatter(1, sorted_idx, remove) logits = logits.masked_fill(remove_idx, float("-inf")) return logits def count_params(self, trainable_only: bool = False) -> int: seen: set[int] = set() total = 0 for param in self.parameters(): if trainable_only and not param.requires_grad: continue ptr = param.data_ptr() if ptr in seen: continue seen.add(ptr) total += param.numel() return total def _check() -> None: from Model.config import tiny_config torch.manual_seed(0) cfg = tiny_config() model = RDTForCausalLM(cfg) print("RDTForCausalLM") print(f" params: {model.count_params():,}") print(f" actual_layers: {cfg.actual_layers}") print(f" effective_depth: {cfg.effective_depth}") bsz, seq_len = 2, 32 input_ids = torch.randint(256, 24576, (bsz, seq_len)) input_ids[:, 0] = cfg.bos_id attention_mask = torch.ones(bsz, seq_len, dtype=torch.long) labels = input_ids.clone() out = model( input_ids=input_ids, attention_mask=attention_mask, labels=labels, ) print("Text") print(f" logits: {tuple(out['logits'].shape)}") print(f" loss: {out['loss'].item():.6f}") print(f" loss_parts: {out['loss_parts']}") print(f" steps_used: {out['rec_info'].get('steps_used')}") out["loss"].backward() grad_sq = 0.0 for p in model.parameters(): if p.grad is not None: grad_sq += p.grad.norm().item() ** 2 print(f" grad_norm: {grad_sq**0.5:.6f}") ids2 = torch.tensor( [ [ cfg.bos_id, 300, cfg.image_start_id, cfg.image_patch_id, cfg.image_patch_id, cfg.image_end_id, 301, cfg.eos_id, ] ] ) labels2 = ids2.clone() labels2[ids2 == cfg.image_start_id] = cfg.ignore_index labels2[ids2 == cfg.image_patch_id] = cfg.ignore_index labels2[ids2 == cfg.image_end_id] = cfg.ignore_index pixel_values = torch.randn(2, model.patch_pixels) out2 = model( input_ids=ids2, labels=labels2, pixel_values=pixel_values, ) print("ImageText") print(f" loss: {out2['loss'].item():.6f}") model.eval() with torch.no_grad(): for steps in [2, 8]: out_step = model(input_ids=input_ids, steps=steps) print( f" steps={steps}, logits_norm={out_step['logits'].norm().item():.6f}" ) if __name__ == "__main__": _check()