# -*- coding: utf-8 -*-
from __future__ import annotations
from collections.abc import Mapping
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, GroupedRMSNorm
from Model.layers.vision_cross_attention import RaggedVisionCrossAttention
from Model.ocr.position_contract import (
BOUNDARY_V1,
LEGACY_SEQUENTIAL_V0,
validate_ocr_position_contract,
)
from Model.recurrent import RecurrentCore
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)
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: OCR alignment and 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
# OMVT-v2 detail memory is opt-in. Keeping this as ``None`` means a
# v1 model creates no bridge parameters and executes no bridge path.
self.vision_cross_attention: RaggedVisionCrossAttention | None = None
self.apply(self._init_weights)
self._scale_residual_projections()
for module in self.modules():
if isinstance(module, (RMSNorm, GroupedRMSNorm)):
module.token_chunk_size = cfg.norm_token_chunk_size
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 install_vision_cross_attention(
self,
*,
memory_dim: int,
n_heads: int | None = None,
dropout: float = 0.0,
) -> RaggedVisionCrossAttention:
"""Explicitly install the teacher-forced OMVT-v2 detail bridge.
Generation and incremental cache behavior are deliberately not
changed here. A second install is rejected so a trained bridge cannot
be silently reinitialized.
"""
if self.vision_cross_attention is not None:
raise RuntimeError("vision cross-attention is already installed")
bridge = RaggedVisionCrossAttention(
d_model=self.cfg.d_model,
memory_dim=memory_dim,
n_heads=self.cfg.n_heads if n_heads is None else n_heads,
dropout=dropout,
)
bridge.apply(self._init_weights)
bridge.zero_init_output_projection()
bridge.to(device=self.embed.weight.device, dtype=self.embed.weight.dtype)
bridge.train(self.training)
self.vision_cross_attention = bridge
return bridge
def _forward_decode(
self,
input_ids: torch.Tensor,
word_pos: torch.Tensor,
morph_depth: torch.Tensor,
cache,
pixel_values: torch.Tensor | Mapping[str, torch.Tensor] | None = None,
visual_features: torch.Tensor | None = None,
detail_memory: torch.Tensor | None = None,
detail_cu_seqlens: 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 != "two_stage":
raise NotImplementedError(
"incremental KV/state cache is implemented for core_type="
"'two_stage' only"
)
bsz, seq_len = input_ids.shape
pos_offset = cache.seq_len
h = self._embed_with_activation_dtype(input_ids)
if pixel_values is not None or visual_features is not None:
h = self.vision(
h,
input_ids,
pixel_values,
visual_features=visual_features,
)
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,
)
if (detail_memory is None) != (detail_cu_seqlens is None):
raise ValueError(
"detail_memory and detail_cu_seqlens must be supplied together"
)
if detail_memory is not None:
if self.vision_cross_attention is None:
raise RuntimeError(
"detail memory requires install_vision_cross_attention()"
)
h = self.vision_cross_attention(
h,
detail_memory,
detail_cu_seqlens,
)
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 | Mapping[str, torch.Tensor] | None = None,
visual_features: torch.Tensor | None = None,
word_pos: torch.Tensor | None = None,
morph_depth: torch.Tensor | None = None,
morphology_track_table: torch.Tensor | None = None,
steps: int | None = None,
bptt_window: int | None = None,
return_logits: bool = True,
logits_to_keep: int = 0,
loss_chunk_size: int | None = None,
pixel_repeats: int = 1,
position_contract: str | None = None,
detail_memory: torch.Tensor | None = None,
detail_cu_seqlens: torch.Tensor | None = None,
all_tokens_valid: bool = False,
) -> dict[str, torch.Tensor | dict | None]:
self._check_inputs(input_ids, attention_mask, labels)
if type(logits_to_keep) is not int or logits_to_keep < 0:
raise ValueError("logits_to_keep must be a non-negative integer")
if logits_to_keep and (labels is not None or not return_logits):
raise ValueError("logits_to_keep is only available for inference without 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 all_tokens_valid and (pixel_values is not None or visual_features is not None):
raise ValueError("the unpadded long-context path is text-only")
layer_mask = None if all_tokens_valid else attention_mask
if position_contract is not None:
position_contract = validate_ocr_position_contract(position_contract)
if morphology_track_table is not None:
raise ValueError(
"morphology_track_table cannot be combined with an explicit "
"OCR position contract"
)
if (word_pos is None) != (morph_depth is None):
raise ValueError("word_pos and morph_depth must be supplied together")
if word_pos is None:
if position_contract is not None:
word_pos, morph_depth = self._morph_info_for_position_contract(
input_ids,
attention_mask,
position_contract,
)
elif morphology_track_table is not None:
word_pos, morph_depth = self._morph_info_from_track_table(
input_ids,
morphology_track_table,
)
else:
word_pos, morph_depth = self._default_morph_info(
input_ids=input_ids,
attention_mask=attention_mask,
)
if (detail_memory is None) != (detail_cu_seqlens is None):
raise ValueError(
"detail_memory and detail_cu_seqlens must be supplied together"
)
if detail_memory is not None and self.vision_cross_attention is None:
raise RuntimeError(
"detail memory requires install_vision_cross_attention()"
)
h = self._embed_with_activation_dtype(input_ids)
if pixel_values is not None or visual_features is not None:
h = self.vision(
h,
input_ids,
pixel_values,
visual_features=visual_features,
pixel_repeats=pixel_repeats,
)
for block in self.prelude:
h = self._maybe_ckpt(
block,
h,
word_pos=word_pos,
morph_depth=morph_depth,
attn_mask=layer_mask,
causal=True,
)
if detail_memory is not None:
assert self.vision_cross_attention is not None
h = self.vision_cross_attention(
h,
detail_memory,
detail_cu_seqlens,
query_mask=attention_mask,
)
e0 = h
h, rec_info = self.recurrent(
e0,
word_pos=word_pos,
morph_depth=morph_depth,
attn_mask=layer_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=layer_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,
loss_chunk_size,
)
else:
logits = self.lm_head(h)
loss, loss_parts = self._losses(h, logits, labels)
elif return_logits:
logits = self.lm_head(h[:, -logits_to_keep:] if logits_to_keep else h)
return {
"loss": loss,
"logits": logits if return_logits else None,
"loss_parts": loss_parts,
"rec_info": rec_info,
}
def _embed_with_activation_dtype(self, input_ids):
hidden = self.embed(input_ids)
if self.cfg.low_precision_residual and torch.is_autocast_enabled(hidden.device.type):
hidden = hidden.to(torch.get_autocast_dtype(hidden.device.type))
return hidden
def _losses(
self,
h: torch.Tensor,
logits: torch.Tensor,
labels: torch.Tensor,
) -> 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())
return loss, parts
def _losses_chunked(
self,
h: torch.Tensor,
labels: torch.Tensor,
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())
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(())
def loss_piece(states, expected):
logits = F.linear(states, head.weight, head.bias)
return F.cross_entropy(logits, expected, reduction="sum")
for start in range(0, hidden.size(0), chunk_size):
end = min(start + chunk_size, hidden.size(0))
if self.cfg.grad_ckpt_loss and torch.is_grad_enabled():
part = checkpoint(loss_piece, hidden[start:end], targets[start:end], use_reentrant=False)
else:
part = loss_piece(hidden[start:end], targets[start:end])
loss_sum = loss_sum + part
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)
return word_pos.to(torch.long), depth.to(torch.long)
def _morph_info_for_position_contract(
self,
input_ids: torch.Tensor,
attention_mask: torch.Tensor,
position_contract: str,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Derive positions for one declared OCR train/inference contract."""
contract = validate_ocr_position_contract(position_contract)
if contract == BOUNDARY_V1:
word_pos, morph_depth = self._default_morph_info(
input_ids,
attention_mask,
)
if self.cfg.max_morph_depth > 0:
morph_depth = morph_depth.clamp(
max=self.cfg.max_morph_depth - 1
)
return word_pos, morph_depth
if contract == LEGACY_SEQUENTIAL_V0:
seq_len = input_ids.shape[1]
word_pos = torch.arange(
seq_len,
device=input_ids.device,
dtype=torch.long,
).unsqueeze(0).expand_as(input_ids)
word_pos = torch.where(
attention_mask.to(torch.bool),
word_pos,
torch.zeros_like(word_pos),
)
return word_pos, torch.zeros_like(word_pos)
raise AssertionError(f"unreachable position contract: {contract}")
def _morph_info_from_track_table(
self,
input_ids: torch.Tensor,
track_table: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Vectorized span-equivalent morphology for tokenizer-routed ids.
Track ``0`` resets the current word, ``1`` is a Mongolian word piece,
and ``2`` is a general-script word piece. Strict OCR data admits only
text for which this id-only representation exactly matches the
tokenizer's span-aware pretraining representation.
"""
if track_table.ndim != 1:
raise ValueError("morphology_track_table must be 1-D")
if track_table.device != input_ids.device:
raise ValueError(
"morphology_track_table must be on the same device as input_ids"
)
tracks = track_table[input_ids].to(torch.long)
is_word = tracks > 0
previous_word = torch.zeros_like(is_word)
previous_word[:, 1:] = is_word[:, :-1]
same_track = torch.zeros_like(is_word)
same_track[:, 1:] = tracks[:, 1:] == tracks[:, :-1]
word_start = is_word & ~(previous_word & same_track)
starts_seen = word_start.long().cumsum(dim=1)
word_pos = (starts_seen - 1).clamp(min=0)
positions = torch.arange(
input_ids.shape[1],
device=input_ids.device,
dtype=torch.long,
).unsqueeze(0).expand_as(input_ids)
start_positions = torch.where(
word_start,
positions,
torch.full_like(positions, -1),
)
latest_start = start_positions.cummax(dim=1).values
morph_depth = torch.where(
is_word,
positions - latest_start.clamp(min=0),
torch.zeros_like(positions),
)
return word_pos, morph_depth
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 its core-aware
``effective_depth`` (including one-pass encoder/local stages and the
repeated refinement stage), 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 == "two_stage"
)
@torch.no_grad()
def _adaptive_depth_logits(
self,
window: torch.Tensor,
pixel_values: torch.Tensor | Mapping[str, torch.Tensor] | None,
visual_features: torch.Tensor | None = None,
attention_mask: torch.Tensor | None = None,
morphology_track_table: torch.Tensor | None = None,
word_pos: torch.Tensor | None = None,
morph_depth: torch.Tensor | None = None,
detail_memory: torch.Tensor | None = None,
detail_cu_seqlens: torch.Tensor | None = 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
if (word_pos is None) != (morph_depth is None):
raise ValueError("word_pos and morph_depth must be supplied together")
morphology_kwargs = (
{"word_pos": word_pos, "morph_depth": morph_depth}
if word_pos is not None
else {}
if morphology_track_table is None
else {"morphology_track_table": morphology_track_table}
)
for d in range(1, max_steps + 1):
if pixel_values is not None or visual_features is not None:
out = self.forward(
window,
attention_mask=attention_mask,
steps=d,
return_logits=True,
pixel_values=pixel_values,
visual_features=visual_features,
detail_memory=detail_memory,
detail_cu_seqlens=detail_cu_seqlens,
**morphology_kwargs,
)
else:
out = self.forward(
window,
attention_mask=attention_mask,
steps=d,
return_logits=True,
detail_memory=detail_memory,
detail_cu_seqlens=detail_cu_seqlens,
**morphology_kwargs,
)
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,
on_sample: (
Callable[[int, torch.Tensor, torch.Tensor], None] | None
) = None,
recurrent_steps: int | None = None,
pixel_values: torch.Tensor | Mapping[str, torch.Tensor] | None = None,
morphology_track_table: torch.Tensor | None = None,
pixel_repeats: int = 1,
position_contract: str | None = None,
detail_memory: torch.Tensor | None = None,
detail_cu_seqlens: 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. ``</think>`` 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.
``on_sample`` additionally receives
``on_sample(step, token_ids, token_log_probs)``. The log-probabilities
are captured from the exact filtered distribution used by sampling,
before any later policy update. Finished rows receive zero because
their appended PAD is not a sampled action. OCR-GRPO uses this callback
to verify synchronous on-policy rollout/scoring consistency.
``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 pixels for prompts containing
``<image_patch>`` 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. OMVT
mapping inputs are encoded/projected once per generation call and their
visual features are reused across decode steps.
"""
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 type(pixel_repeats) is not int or pixel_repeats <= 0:
raise ValueError("pixel_repeats must be a positive integer")
if (detail_memory is None) != (detail_cu_seqlens is None):
raise ValueError(
"detail_memory and detail_cu_seqlens must be supplied together"
)
if detail_memory is not None:
if self.vision_cross_attention is None:
raise RuntimeError(
"detail memory requires install_vision_cross_attention()"
)
if pixel_repeats != 1:
raise ValueError(
"detail memory does not support pixel_repeats; repeat its "
"ragged offsets explicitly"
)
if position_contract is not None:
position_contract = validate_ocr_position_contract(position_contract)
if morphology_track_table is not None:
raise ValueError(
"morphology_track_table cannot be combined with an explicit "
"OCR position contract"
)
if use_cache and self.cfg.core_type != "two_stage":
raise NotImplementedError(
"use_cache=True is only supported for core_type="
"'two_stage'"
)
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."
)
if morphology_track_table is not None:
if morphology_track_table.ndim != 1:
raise ValueError("morphology_track_table must be 1-D")
if morphology_track_table.shape[0] < self.cfg.vocab_size:
raise ValueError(
"morphology_track_table is shorter than model vocab_size"
)
if morphology_track_table.device != input_ids.device:
raise ValueError(
"morphology_track_table must share input_ids device"
)
if bool(
(
(morphology_track_table < 0)
| (morphology_track_table > 2)
).any()
):
raise ValueError(
"morphology_track_table values must be reset/mn/general"
)
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)
seq_attention = (
(seq != pad_id).long()
if pad_id is not None
else torch.ones_like(seq, dtype=torch.long)
)
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:
# The cache-free official-Mamba path re-runs the text model for every
# generated token. Re-running the unchanged OMVT tower there would
# multiply visual compute by max_new_tokens. Cache only Mapping
# (production OMVT) inputs; retain the legacy Tensor path verbatim.
cached_visual_features = None
runtime_pixel_values = pixel_values
if isinstance(pixel_values, Mapping):
cached_visual_features = self.vision.encode_visual(
pixel_values,
pixel_repeats=pixel_repeats,
)
runtime_pixel_values = None
elif pixel_repeats != 1:
raise ValueError(
"pixel_repeats is only supported for OMVT mapping inputs"
)
generic_morphology_kwargs = (
{}
if morphology_track_table is None
else {"morphology_track_table": morphology_track_table}
)
for step in range(max_new_tokens):
if use_cache:
if decode_cache.seq_len == 0:
step_ids = seq
step_pixels = runtime_pixel_values
step_features = cached_visual_features
else:
step_ids = seq[:, -1:]
# The image was folded into the cache at prefill; later
# steps process only the new token and carry no
# <image_patch> slots, so pixel_values must be dropped.
step_pixels = None
step_features = None
if position_contract is not None:
word_pos, morph_depth = (
self._morph_info_for_position_contract(
seq,
seq_attention,
position_contract,
)
)
elif morphology_track_table is not None:
word_pos, morph_depth = (
self._morph_info_from_track_table(
seq,
morphology_track_table,
)
)
else:
word_pos, morph_depth = self._default_morph_info(
seq, seq_attention
)
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,
visual_features=step_features,
detail_memory=detail_memory,
detail_cu_seqlens=detail_cu_seqlens,
steps=recurrent_steps,
)[:, -1, :].float()
else:
window = seq
window_attention = seq_attention
if window.shape[1] > cfg.max_seq_len:
window = window[:, -cfg.max_seq_len:]
window_attention = window_attention[:, -cfg.max_seq_len:]
if pixel_values is not None:
# Left-truncation could drop <image_patch> 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 <image_patch> slots would "
"desync from pixel_values. Use use_cache=True or "
"keep L + max_new_tokens <= max_seq_len."
)
position_kwargs = {}
if position_contract is not None:
word_pos, morph_depth = (
self._morph_info_for_position_contract(
window,
window_attention,
position_contract,
)
)
position_kwargs = {
"word_pos": word_pos,
"morph_depth": morph_depth,
}
if self._kl_exit_active(recurrent_steps):
logits = self._adaptive_depth_logits(
window,
runtime_pixel_values,
visual_features=cached_visual_features,
attention_mask=window_attention,
morphology_track_table=morphology_track_table,
detail_memory=detail_memory,
detail_cu_seqlens=detail_cu_seqlens,
**position_kwargs,
)
elif (
runtime_pixel_values is not None
or cached_visual_features is not None
or detail_memory is not None
):
out = self.forward(
window,
attention_mask=window_attention,
steps=recurrent_steps,
return_logits=True,
pixel_values=runtime_pixel_values,
visual_features=cached_visual_features,
detail_memory=detail_memory,
detail_cu_seqlens=detail_cu_seqlens,
**position_kwargs,
**generic_morphology_kwargs,
)
logits = out["logits"][:, -1, :].float()
else:
out = self.forward(
window,
attention_mask=window_attention,
steps=recurrent_steps,
return_logits=True,
**position_kwargs,
**generic_morphology_kwargs,
)
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)
sampling_logp = (
F.log_softmax(logits.float(), dim=-1)
.gather(-1, next_token.unsqueeze(-1))
.squeeze(-1)
if on_sample is not None
else None
)
else:
logits = logits / temperature
logits = self._filter_logits(logits, top_k, top_p, min_p)
if on_sample is not None:
sampling_distribution = F.log_softmax(
logits.float(), dim=-1
)
probs = sampling_distribution.exp()
else:
sampling_distribution = None
probs = F.softmax(logits, dim=-1)
next_token = torch.multinomial(probs, num_samples=1).squeeze(-1)
sampling_logp = (
sampling_distribution.gather(
-1, next_token.unsqueeze(-1)
).squeeze(-1)
if sampling_distribution is not None
else None
)
was_finished = finished
next_token = torch.where(
was_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)
next_attention = (~was_finished).to(seq_attention.dtype).unsqueeze(1)
seq_attention = torch.cat([seq_attention, next_attention], dim=1)
if on_token is not None:
on_token(step, next_token)
if on_sample is not None:
assert sampling_logp is not None
on_sample(
step,
next_token,
torch.where(
was_finished,
torch.zeros_like(sampling_logp),
sampling_logp,
),
)
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()
# -*- coding: utf-8 -*-
from __future__ import annotations
from collections.abc import Mapping
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, GroupedRMSNorm
from Model.layers.vision_cross_attention import RaggedVisionCrossAttention
from Model.ocr.position_contract import (
BOUNDARY_V1,
LEGACY_SEQUENTIAL_V0,
validate_ocr_position_contract,
)
from Model.recurrent import RecurrentCore
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)
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: OCR alignment and 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
# OMVT-v2 detail memory is opt-in. Keeping this as ``None`` means a
# v1 model creates no bridge parameters and executes no bridge path.
self.vision_cross_attention: RaggedVisionCrossAttention | None = None
self.apply(self._init_weights)
self._scale_residual_projections()
for module in self.modules():
if isinstance(module, (RMSNorm, GroupedRMSNorm)):
module.token_chunk_size = cfg.norm_token_chunk_size
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 install_vision_cross_attention(
self,
*,
memory_dim: int,
n_heads: int | None = None,
dropout: float = 0.0,
) -> RaggedVisionCrossAttention:
"""Explicitly install the teacher-forced OMVT-v2 detail bridge.
Generation and incremental cache behavior are deliberately not
changed here. A second install is rejected so a trained bridge cannot
be silently reinitialized.
"""
if self.vision_cross_attention is not None:
raise RuntimeError("vision cross-attention is already installed")