Hugging Face's logo

lafa-hf
/
lafa-7B

Text Generation
Transformers
Safetensors
lafa
feature-extraction
qwen2
flow-matching
speculative-decoding
custom_code
conversational
Model card Files
xet
 Community
lafa-7B / flow_model.py
david3684's picture
david3684
Sync model code with the release repository
a1202ae verified
Raw History Blame Contribute Delete
105 kB
"""Embedding-space flow model with DiT time conditioning."""
import contextlib
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor
from transformers.modeling_outputs import CausalLMOutputWithPast
from .base_wrapper import BaseModelWrapper
from .attention import create_packed_attention_mask
from .cache import trim_kv_cache, gather_kv_cache
from .layers import TimestepEmbedder, DiTAdaLN
from .loss import cross_entropy_loss
class FlowModelWrapper(BaseModelWrapper):
"""Wrap a causal LM for block flow training and decoding."""
def __init__(
self,
backbone: nn.Module,
bd_size: int = 8,
noise_scale_mult: float = 4.0,
adaln_init_scale: float = 1e-3,
adaln_no_shift: bool = True,
adaln_gate_init: float = 0.0,
adaln_share_layers: bool = False,
two_time: bool = False,
distill_two_time_mode: str = "shared_embedding",
distill_second_time_input: str = "target",
distill_second_time_bias: bool = True,
enable_self_cond: bool = True,
sc_mode: str = "residual",
sc_drop_rate: float = 0.5,
freeze_embed: bool = True,
freeze_lm_head: bool = True,
embed_norm_project: bool = False,
block_causal: bool = False,
noisy_causal: bool = False,
logit_shift: bool = False,
ar_loss: bool = False,
ar_loss_weight: float = 1.0,
flow_ce_weight: float = 1.0,
denoise_mse_weight: float = 0.0,
ema_decay: float = 0.0,
time_warp_num_bins: int = 100,
time_warp_per_position: bool = True,
time_warp_init_path: str = None,
flow_ce_enable: bool = True,
time_warp_start_step: int = 0,
warp_uniform_est: bool = False,
warp_est_every: int = 4,
warp_est_samples: int = 4,
noise_sigma: float | None = None,
):
super().__init__(backbone)
if bd_size < 1:
raise ValueError("bd_size must be positive")
self.bd_size = bd_size
self.block_causal = block_causal
self.noisy_causal = noisy_causal
self.noise_scale_mult = noise_scale_mult
# The warper is saved separately from model.state_dict().
self.time_warp_num_bins = time_warp_num_bins
self.time_warp_per_position = time_warp_per_position
self.time_warp_init_path = time_warp_init_path
self.flow_ce_enable = flow_ce_enable
self.time_warp_start_step = time_warp_start_step
self.warp_uniform_est = warp_uniform_est
self.warp_est_every = warp_est_every
self.warp_est_samples = warp_est_samples
self.time_warper = None
self._global_step = 0 # set by TimeWarpCallback for warm-up
self._last_sampled_t = None # for the sampled-time histogram
self.enable_self_cond = enable_self_cond
self.sc_mode = sc_mode
self.sc_drop_rate = sc_drop_rate
self.logit_shift = logit_shift
self.ar_loss = ar_loss
self.ar_loss_weight = ar_loss_weight
self.flow_ce_weight = flow_ce_weight
self.denoise_mse_weight = denoise_mse_weight
# Weight EMA is saved separately from model.state_dict().
self.ema_decay = ema_decay
self._ema_shadow = None
# Detached losses for FlowTrainer logging.
self._loss_components = {}
self.__dict__["_mask_cache"] = {}
# ZeRO-3 may partition embedding weights during construction.
_bcfg = self.backbone.config
self._d_model = _bcfg.hidden_size
# Noise scale uses the full embedding table, including under ZeRO-3.
_E = self.embed_tokens.weight
if noise_sigma is None:
with self._maybe_gather([_E]):
self._noise_sigma = _E.data.std().item()
else:
self._noise_sigma = float(noise_sigma)
# The persistent buffer preserves embedding row norms across resumes.
self._embed_norm_project = embed_norm_project
V = _bcfg.vocab_size
self.register_buffer("_embed_init_norm", torch.zeros(V, 1), persistent=True)
self._adaln_init_scale = adaln_init_scale
self._adaln_no_shift = adaln_no_shift
self._adaln_gate_init = adaln_gate_init
self._adaln_share_layers = bool(adaln_share_layers)
self._two_time = bool(two_time)
self._distill_two_time_mode = str(distill_two_time_mode).strip().lower()
if self._distill_two_time_mode not in {
"shared_embedding",
"source_only",
}:
raise ValueError(
"Two-time inference supports "
"distill_two_time_mode='shared_embedding' or 'source_only'"
)
self._distill_second_time_input = (
str(distill_second_time_input or "target").strip().lower()
)
if self._distill_second_time_input != "target":
raise ValueError(
"Two-time inference supports absolute target-time conditioning only"
)
self._distill_second_time_bias = bool(distill_second_time_bias)
if self._two_time and not self._distill_second_time_bias:
raise ValueError(
"Bias-free two-time embeddings are not supported by this inference"
)
self._freeze_embed = freeze_embed
self._freeze_lm_head = freeze_lm_head
self._setup_backbone()
def _setup_backbone(self):
if self._freeze_embed:
self.embed_tokens.weight.requires_grad_(False)
if self._freeze_lm_head:
self.lm_head.weight.requires_grad_(False)
self._setup_dit_adaln()
if self.enable_self_cond:
self._setup_self_cond()
@staticmethod
def _maybe_gather(params):
"""Gather ZeRO-3 parameters when needed; otherwise use a null context."""
try:
from transformers.integrations.deepspeed import is_deepspeed_zero3_enabled
if is_deepspeed_zero3_enabled():
import deepspeed
return deepspeed.zero.GatheredParameters(params, modifier_rank=None)
except Exception:
pass
return contextlib.nullcontext()
def _setup_dit_adaln(self):
d = self._d_model
self.time_embed = TimestepEmbedder(d)
if self._two_time and self._distill_two_time_mode != "source_only":
self.time_embed_tgt = TimestepEmbedder(d)
if self._distill_two_time_mode == "shared_embedding":
nn.init.zeros_(self.time_embed_tgt.mlp[2].weight)
nn.init.zeros_(self.time_embed_tgt.mlp[2].bias)
layers = self.backbone.model.layers
def _mk():
return DiTAdaLN(
d,
init_scale=self._adaln_init_scale,
no_shift=self._adaln_no_shift,
gate_init=self._adaln_gate_init,
)
if self._adaln_share_layers:
# Keep the state_dict prefix and AdaLN optimizer group unchanged.
shared = _mk()
self.adaln_layers = nn.ModuleList([shared])
for layer in layers:
shared.patch_layer(layer)
else:
self.adaln_layers = nn.ModuleList([_mk() for _ in layers])
for layer, adaln in zip(layers, self.adaln_layers):
adaln.patch_layer(layer)
def _setup_self_cond(self):
d = self._d_model
if self.sc_mode == "concat":
self.sc_proj = nn.Linear(d * 2, d, bias=False)
nn.init.zeros_(self.sc_proj.weight)
self.sc_proj.weight.data[:d, :d] = torch.eye(d)
elif self.sc_mode == "residual":
self.sc_lin_a = nn.Linear(d, d, bias=False)
self.sc_lin_b = nn.Linear(d, d, bias=False)
nn.init.zeros_(self.sc_lin_a.weight)
nn.init.zeros_(self.sc_lin_b.weight)
elif self.sc_mode == "simple":
self.sc_lin_b = nn.Linear(d, d, bias=False)
nn.init.zeros_(self.sc_lin_b.weight)
elif self.sc_mode == "mlp":
self.sc_mlp = nn.Sequential(
nn.Linear(d, d, bias=False),
nn.SiLU(),
nn.Linear(d, d, bias=False),
)
nn.init.zeros_(self.sc_mlp[2].weight)
elif self.sc_mode == "scale":
self.sc_scale = nn.Parameter(torch.zeros(1))
@property
def embed_tokens(self) -> nn.Embedding:
return self.backbone.model.embed_tokens
@property
def lm_head(self) -> nn.Linear:
return self.backbone.lm_head
def state_dict(self, *args, **kwargs):
return nn.Module.state_dict(self, *args, **kwargs)
def load_state_dict(self, state_dict, *args, **kwargs):
# Older ZeRO-3 checkpoints may have a malformed embedding norm buffer.
# Keep the current buffer so those checkpoints still load.
k = "_embed_init_norm"
cur = getattr(self, k, None)
if k in state_dict and cur is not None and state_dict[k].shape != cur.shape:
print(
f"[flow] load_state_dict: checkpoint '{k}' {tuple(state_dict[k].shape)} "
f"!= model {tuple(cur.shape)} (ZeRO-3 capture artifact) -> keeping fresh "
f"buffer",
flush=True,
)
state_dict = dict(state_dict)
state_dict[k] = cur.detach().clone()
return nn.Module.load_state_dict(self, state_dict, *args, **kwargs)
def _sample_noise(self, like: Tensor) -> Tensor:
return torch.randn_like(like) * self._noise_sigma * self.noise_scale_mult
def _interpolate(self, e_0: Tensor, e_1: Tensor, t: Tensor) -> Tensor:
"""e_t = (1-t)*e_0 + t*e_1. t: [B,1,1]-broadcastable."""
return (1 - t) * e_0 + t * e_1
def _compute_t_bn(
self,
t: Tensor,
tgt: bool = False,
) -> Tensor | None:
"""Embed a source time or the PSD checkpoint's absolute target time."""
if tgt and self._distill_two_time_mode == "source_only":
return None
embedder = self.time_embed_tgt if tgt else self.time_embed
return embedder(t)
def _compose_flowmap_time_condition(
self,
source_t: Tensor,
target_t: Tensor,
target_mask: Tensor | None = None,
) -> Tensor:
"""Build shared-AdaLN conditioning for mixed clean and PSD positions.
Clean verifier/cache positions use only ``E_source``. Noisy flow-map
draft positions use ``E_source + E_target``. Masking the target term is
what lets one fused forward retain the checkpoint's ordinary AR path
while also evaluating its two-time PSD map.
"""
source = self._compute_t_bn(source_t)
if not self._two_time or self._distill_two_time_mode == "source_only":
return source
if self._distill_two_time_mode != "shared_embedding":
raise RuntimeError(
"Fused PSD self-speculation requires shared_embedding "
"two-time conditioning"
)
target = self._compute_t_bn(
target_t,
tgt=True,
).to(source.dtype)
if target_mask is not None:
target = target * target_mask.to(target.dtype).unsqueeze(-1)
return source + target
def _precompute_shared_adaln(self, t_bn: Tensor):
"""Return reusable AdaLN modulation for a static inference layout.
Shared-AdaLN checkpoints patch the same module into every decoder layer,
so repeatedly evaluating ``adaln(t_bn)`` produces identical tensors.
Non-shared checkpoints compute each layer's modulation separately.
"""
if not getattr(self, "_adaln_share_layers", False):
return None
return self.adaln_layers[0](t_bn)
def _ensure_time_warper(self, device):
"""Lazily build the TimeWarping module on the right device (t in [0,1])."""
if self.time_warper is None:
from .time_warping import TimeWarping
self.time_warper = TimeWarping(
num_bins=self.time_warp_num_bins,
t_min=0.0,
t_max=1.0,
device=device,
per_position=self.time_warp_per_position,
bd_size=self.bd_size,
)
_twip = getattr(self, "time_warp_init_path", None)
if _twip:
import os
p = (
os.path.join(_twip, "time_warper.pt")
if os.path.isdir(_twip)
else _twip
)
self.time_warper.load_state_dict(torch.load(p, map_location=device))
elif self.time_warper.loss_ema.device != torch.device(device):
self.time_warper.to(device)
return self.time_warper
def load_time_warper(self, ckpt_path, device=None):
"""Load a trained TimeWarping for generation: prefer
ckpt_path/time_warper.pt, else fall back to the latest
checkpoint-*/time_warper.pt in the parent run dir. Returns the source
path, or None (leaving a fresh ~uniform warper) if nothing is found."""
import glob
import os
import re
if device is None:
device = self.embed_tokens.weight.device
src = None
direct = os.path.join(ckpt_path, "time_warper.pt")
if os.path.exists(direct):
src = direct
else:
parent = os.path.dirname(os.path.normpath(ckpt_path))
cks = glob.glob(os.path.join(parent, "checkpoint-*", "time_warper.pt"))
if cks:
src = max(
cks, key=lambda p: int(re.search(r"checkpoint-(\d+)", p).group(1))
)
if src is None:
return None
self._ensure_time_warper(device)
self.time_warper.load_state_dict(torch.load(src, map_location=device))
return src
def _warp_active(self):
"""True once warm-up is over (uniform-t before time_warp_start_step)."""
return self._global_step >= self.time_warp_start_step
def _sample_t_units(self, n, device):
"""Sample learned times after the uniform warm-up."""
wp = self._ensure_time_warper(device)
if self._warp_active():
return wp.sample(n, device=device).detach()
return torch.rand(n, device=device)
def _embedding_shard_slices(self, embedding):
"""Yield contiguous portions of each row in this ZeRO-3 shard."""
import torch.distributed as dist
shard = embedding.ds_tensor
group = getattr(embedding, "ds_process_group", None)
rank = dist.get_rank(group=group) if dist.is_initialized() else 0
width = self._d_model
start = rank * shard.numel()
remaining = max(0, min(shard.numel(), embedding.ds_numel - start))
offset = 0
row = start // width
if start % width and remaining:
length = min(remaining, width - start % width)
yield row, shard[:length]
offset += length
row += 1
while remaining - offset >= width:
n_rows = min((remaining - offset) // width, 4096)
length = n_rows * width
yield (
slice(row, row + n_rows),
shard[offset : offset + length].view(n_rows, width),
)
offset += length
row += n_rows
if offset < remaining:
yield row, shard[offset:remaining]
def _embedding_row_norms(self, embedding):
"""Compute full row norms without gathering a ZeRO-3 embedding table."""
shard = getattr(embedding, "ds_tensor", None)
if shard is None:
return embedding.detach().float().norm(dim=-1)
import torch.distributed as dist
squared = torch.zeros(
self._embed_init_norm.shape[0], device=shard.device, dtype=torch.float32
)
for rows, values in self._embedding_shard_slices(embedding):
squared[rows] = values.float().square().sum(dim=-1)
if dist.is_initialized():
dist.all_reduce(squared, group=getattr(embedding, "ds_process_group", None))
return squared.sqrt()
def _capture_embed_norms(self):
"""Capture pretrained row norms once, including from ZeRO-3 shards."""
if not self._embed_norm_project or self._freeze_embed:
return
if float(self._embed_init_norm.abs().sum()) > 0:
return
with torch.no_grad():
norms = self._embedding_row_norms(self.embed_tokens.weight)
self._embed_init_norm.copy_(
norms.unsqueeze(-1).to(self._embed_init_norm.device)
)
print(
f"[EmbedNormProject] captured per-row norms "
f"(mean={self._embed_init_norm.mean().item():.4f}).",
flush=True,
)
def _post_step_project(self):
"""Restore captured row norms after an optimizer step."""
if not self._embed_norm_project or self._freeze_embed:
return
if float(self._embed_init_norm.abs().sum()) == 0:
return
with torch.no_grad():
embedding = self.embed_tokens.weight
current = self._embedding_row_norms(embedding).clamp_min(1e-8)
target = self._embed_init_norm.flatten().to(current.device)
scale = target.float() / current
shard = getattr(embedding, "ds_tensor", None)
if shard is None:
embedding.mul_(scale[:, None].to(embedding.dtype))
else:
for rows, values in self._embedding_shard_slices(embedding):
factor = scale[rows]
if isinstance(rows, slice):
factor = factor[:, None]
values.mul_(factor.to(values.dtype))
# ZeRO-3 updates local EMA shards and gathers them only for saving.
@staticmethod
def _ema_param_view(p):
"""The tensor the EMA shadow tracks for param p: its ZeRO-3 local shard
(p.ds_tensor) if partitioned, else the full param. Always materialized."""
ds = getattr(p, "ds_tensor", None)
return ds if ds is not None else p
def ema_initialize(self):
"""Initialize fp32 EMA shadows for trainable parameters."""
if self.ema_decay <= 0 or self._ema_shadow is not None:
return
self._ema_shadow = {
n: self._ema_param_view(p).detach().float().clone()
for n, p in self.named_parameters()
if p.requires_grad
}
@torch.no_grad()
def ema_update(self):
"""EMA step: shadow = decay*shadow + (1-decay)*param. Call AFTER optimizer.step().
Shard-local under ZeRO-3 (no gather); shape-guarded so a transient mismatch skips
rather than crashes."""
if not self._ema_shadow:
return
d = self.ema_decay
for n, p in self.named_parameters():
s = self._ema_shadow.get(n)
if s is None:
continue
v = self._ema_param_view(p)
if v.shape != s.shape:
# A ZeRO-3 shard may not be materialized on this step.
continue
s.mul_(d).add_(v.detach().float(), alpha=1.0 - d)
@property
def has_ema(self):
return bool(self._ema_shadow)
def ema_state_dict(self):
"""Return full EMA weights for saving.
ZeRO-3 gathers shards collectively on all ranks. Only rank 0 receives
the complete state.
"""
if not self._ema_shadow:
return None
z3 = any(
getattr(p, "ds_tensor", None) is not None
for _, p in self.named_parameters()
)
if not z3:
return self._ema_shadow
import torch.distributed as dist
rank = dist.get_rank() if dist.is_initialized() else 0
full = {}
name_to_p = dict(self.named_parameters())
for n, s in self._ema_shadow.items():
p = name_to_p.get(n)
if p is None:
continue
# NCCL gathers on GPU; assemble on CPU to limit peak GPU memory.
shard = s.detach().float().contiguous()
if dist.is_initialized():
shards = [torch.empty_like(shard) for _ in range(dist.get_world_size())]
dist.all_gather(shards, shard)
if rank == 0:
flat = torch.cat([sh.cpu() for sh in shards])
else:
del shards
continue
else:
flat = shard.cpu()
# Remove ZeRO shard padding.
true_numel = getattr(p, "ds_numel", None) or p.numel()
full[n] = (
flat[:true_numel].view(getattr(p, "ds_shape", None) or p.shape).clone()
)
return full if rank == 0 else None
def load_ema_state_dict(self, sd, device=None):
"""Load a FULL-tensor EMA state (as written by ema_state_dict). Under ZeRO-3,
re-shard each full tensor into this rank's local shard to match the live layout."""
dev = device or self.embed_tokens.weight.device
z3 = any(
getattr(p, "ds_tensor", None) is not None
for _, p in self.named_parameters()
)
if not z3:
self._ema_shadow = {k: v.float().to(dev) for k, v in sd.items()}
return
import torch.distributed as dist
rank = dist.get_rank() if dist.is_initialized() else 0
world = dist.get_world_size() if dist.is_initialized() else 1
shadow = {}
name_to_p = dict(self.named_parameters())
for n, full in sd.items():
p = name_to_p.get(n)
if p is None:
continue
flat = full.detach().float().flatten()
# Match ZeRO's contiguous shard layout.
shard_numel = getattr(p, "ds_tensor", p).numel()
pad = shard_numel * world - flat.numel()
if pad > 0:
flat = torch.cat([flat, flat.new_zeros(pad)])
shadow[n] = (
flat[rank * shard_numel : (rank + 1) * shard_numel].to(dev).clone()
)
self._ema_shadow = shadow
@contextlib.contextmanager
def ema_swap(self):
"""Temporarily load EMA weights into the live params; restore on exit.
No-op if no EMA shadow exists. Under ZeRO-3 this swaps each rank's local SHARD
(via _ema_param_view -> p.ds_tensor), so the subsequently-gathered params in the
eval forward reflect EMA — no all-gather needed here."""
if not self._ema_shadow:
yield
return
backup = {}
with torch.no_grad():
for n, p in self.named_parameters():
s = self._ema_shadow.get(n)
if s is None:
continue
v = self._ema_param_view(p)
if v.shape != s.shape:
continue
backup[n] = v.detach().clone()
v.data.copy_(s.to(v.dtype))
try:
yield
finally:
with torch.no_grad():
for n, p in self.named_parameters():
if n in backup:
self._ema_param_view(p).data.copy_(backup[n])
def _prepare_flow_input(self, input_ids: Tensor, cond_mask: Tensor):
"""input_ids [B, L] → e_1, e_t, t_bn, position_ids, attn_mask, t.
`t` is the sampled time, per position for a coupled warper or per block
during warm-up and for a shared warper. It is returned for warper fitting.
cond_mask [B, L]: 1 = answer (corrupt), 0 =
prompt/pad (kept CLEAN, conditioned at t=1)."""
B, L = input_ids.shape
device = input_ids.device
e_1 = self.embed_tokens(input_ids) # [B, L, d]
# Couple position-specific times with one quantile per block.
num_blocks = L // self.bd_size
if self.time_warp_per_position and self._warp_active():
wp = self._ensure_time_warper(device)
tau = torch.rand(B * num_blocks, device=device) # one tau per block
tc = wp.sample_tau_coupled( # [B*nb, bd_size]
tau, device=device
).detach()
# Positions follow absolute token columns; boundary blocks have no flow loss.
t_pos = tc.view(B, num_blocks, self.bd_size).reshape(B, L)
t = t_pos
else:
t = self._sample_t_units(B * num_blocks, device).view(B, num_blocks)
t_pos = t.repeat_interleave(self.bd_size, dim=1) # [B, L]
e_0 = self._sample_noise(e_1) # [B, L, d]
e_t = self._interpolate(
e_0, e_1, t_pos.unsqueeze(-1).to(e_1.dtype)
) # [B, L, d]
# Keep prompt/pad positions clean and condition them at t=1.
m_e = cond_mask.unsqueeze(-1).to(e_t.dtype) # [B, L, 1]
e_t = m_e * e_t + (1.0 - m_e) * e_1
m_t = cond_mask.to(t_pos.dtype) # [B, L]
t_pos = m_t * t_pos + (1.0 - m_t) * 1.0
ones_clean = torch.ones(B, L, device=device, dtype=t_pos.dtype)
per_pos_t = torch.cat([ones_clean, t_pos], dim=1) # [B, 2L]
t_bn = self._compute_t_bn(per_pos_t) # [B, 2L, d]
# Position ids: [0..L-1, 0..L-1] — noisy half shares positions with clean
pos = torch.arange(L, device=device)
position_ids = torch.cat([pos, pos]).unsqueeze(0).expand(B, -1) # [B, 2L]
cache_key = (
L,
self.bd_size,
self.block_causal,
self.noisy_causal,
device,
e_1.dtype,
)
if cache_key not in self._mask_cache:
self._mask_cache[cache_key] = create_packed_attention_mask(
L,
self.bd_size,
device=device,
dtype=e_1.dtype,
block_causal=self.block_causal,
noisy_causal=self.noisy_causal,
)
attn_mask = self._mask_cache[cache_key]
return e_1, e_t, t_bn, position_ids, attn_mask, t
def _full_block_loss_mask(self, valid: Tensor) -> Tensor:
"""Keep positions only when their entire block is supervised."""
B, L = valid.shape
block_full = valid.view(B, L // self.bd_size, self.bd_size).all(dim=-1)
return valid & block_full.repeat_interleave(self.bd_size, dim=1)
@staticmethod
def _answer_corruption_mask(labels: Tensor):
"""Map supervised tokens to 1 and prompt/pad tokens to clean 0."""
return (labels != -100).float()
def _denoising_squared_error(self, noisy_logits, e_1):
probs = F.softmax(noisy_logits, dim=-1)
e_denoised = probs @ self.embed_tokens.weight
target = F.pad(e_1, (0, 0, 0, 1))[:, 1:, :] if self.logit_shift else e_1
return ((e_denoised - target.to(e_denoised.dtype)) ** 2).sum(dim=-1)
def _warp_loss_per_token(self, ce_pp, sq_err):
loss = (self.flow_ce_weight if self.flow_ce_enable else 0.0) * ce_pp.float()
if sq_err is not None:
loss = loss + self.denoise_mse_weight * sq_err.float()
return loss
@torch.no_grad()
def _warp_fit_inputs(self, ce_pp, sq_err, valid, t_raw):
"""Reduce supervised token losses to position or block histogram samples."""
B, L = valid.shape
nb = L // self.bd_size
loss_pp = self._warp_loss_per_token(ce_pp, sq_err)
if self.time_warp_per_position:
if t_raw.numel() == B * L:
t_pos = t_raw.reshape(B, L)
elif t_raw.numel() == B * nb:
t_pos = t_raw.reshape(B, nb).repeat_interleave(self.bd_size, dim=1)
else:
raise ValueError("Expected per-position or per-block sampled times")
positions = (
torch.arange(L, device=loss_pp.device)
.remainder(self.bd_size)
.unsqueeze(0)
.expand(B, L)
)
t_by_pos, loss_by_pos = [], []
for pos in range(self.bd_size):
selected = valid & (positions == pos)
t_by_pos.append(t_pos[selected])
loss_by_pos.append(loss_pp[selected])
return t_by_pos, loss_by_pos
complete_blocks = valid.view(B, nb, self.bd_size).all(dim=-1)
return (
t_raw.reshape(B, nb)[complete_blocks],
loss_pp.view(B, nb, self.bd_size).mean(dim=-1)[complete_blocks],
)
@torch.no_grad()
def _warp_block_fit_inputs(
self,
e_1: Tensor,
labels: Tensor,
position_ids: Tensor,
per_position: bool = False,
):
"""Concatenate independent warp estimates without increasing peak batch size."""
n_samples = int(getattr(self, "warp_est_samples", 1) or 1)
if n_samples <= 1:
return self._warp_block_fit_once(e_1, labels, position_ids, per_position)
if per_position:
t_acc = [[] for _ in range(self.bd_size)]
l_acc = [[] for _ in range(self.bd_size)]
for _ in range(n_samples):
tbp, lbp = self._warp_block_fit_once(e_1, labels, position_ids, True)
for i in range(self.bd_size):
t_acc[i].append(tbp[i])
l_acc[i].append(lbp[i])
t_by_pos = [
torch.cat(t_acc[i]) if t_acc[i] else torch.empty(0, device=e_1.device)
for i in range(self.bd_size)
]
loss_by_pos = [
torch.cat(l_acc[i]) if l_acc[i] else torch.empty(0, device=e_1.device)
for i in range(self.bd_size)
]
return t_by_pos, loss_by_pos
t_all, l_all = [], []
for _ in range(n_samples):
tu, lu = self._warp_block_fit_once(e_1, labels, position_ids, False)
t_all.append(tu)
l_all.append(lu)
return torch.cat(t_all), torch.cat(l_all)
def _warp_block_fit_once(
self,
e_1: Tensor,
labels: Tensor,
position_ids: Tensor,
per_position: bool = False,
):
"""Estimate loss at a uniform anchor time for each block."""
B, L = e_1.shape[0], e_1.shape[1]
device = e_1.device
nb = L // self.bd_size
t = torch.rand(B * nb, device=device).view(B, nb) # [B, nb]
if per_position:
# A random position anchors each block at uniform t; the others share its tau.
anchor = torch.randint(0, self.bd_size, (B * nb,), device=device)
tc = self._ensure_time_warper(device).coupled_fit_times(
t.reshape(B * nb), anchor
) # [B*nb, bd_size]
t_pos = tc.view(B, nb, self.bd_size).reshape(B, L) # [B, L] per-position
else:
t_pos = t.repeat_interleave(self.bd_size, dim=1) # [B, L] flat shared-t
e_0 = self._sample_noise(e_1)
e_t = self._interpolate(e_0, e_1, t_pos.unsqueeze(-1).to(e_1.dtype))
response = self._answer_corruption_mask(labels).bool()
e_t = torch.where(response.unsqueeze(-1), e_t, e_1)
t_condition = torch.where(response, t_pos, torch.ones_like(t_pos))
ones_clean = torch.ones(B, L, device=device, dtype=t_pos.dtype)
t_bn = self._compute_t_bn(
torch.cat([ones_clean, t_condition], dim=1)
) # [B, 2L, d]
cache_key = (
L,
self.bd_size,
self.block_causal,
self.noisy_causal,
device,
e_1.dtype,
)
attn_mask = self._mask_cache.get(cache_key)
if attn_mask is None: # mirror _prepare_flow_input
attn_mask = create_packed_attention_mask(
L,
self.bd_size,
device=device,
dtype=e_1.dtype,
block_causal=self.block_causal,
noisy_causal=self.noisy_causal,
)
self._mask_cache[cache_key] = attn_mask
packed = torch.cat([e_1, e_t], dim=1) # [B, 2L, d]
logits, _, _ = self._forward_backbone(
packed, t_bn, attn_mask, use_cache=False, position_ids=position_ids
)
noisy_logits = logits[:, L:, :].contiguous() # [B, L, V]
_ce_tgt = (
F.pad(labels, (0, 1), value=-100)[:, 1:] if self.logit_shift else labels
)
valid = self._full_block_loss_mask(_ce_tgt != -100)
# Gather supports non-contiguous shifted targets.
ce_pp = (
-F.log_softmax(noisy_logits.float(), dim=-1)
.gather(-1, _ce_tgt.clamp(min=0).long().unsqueeze(-1))
.squeeze(-1)
) # [B, L]
sq_err = (
self._denoising_squared_error(noisy_logits, e_1)
if self.denoise_mse_weight > 0
else None
)
loss_pp = self._warp_loss_per_token(ce_pp, sq_err)
if per_position:
col = torch.arange(L, device=device)
pos = col % self.bd_size # [L] block-local position
t_by_pos, loss_by_pos = [], []
for i in range(self.bd_size):
m = valid & (pos == i).unsqueeze(0) # [B, L]
t_by_pos.append(t_pos[m])
loss_by_pos.append(loss_pp[m])
return t_by_pos, loss_by_pos
complete_blocks = valid.view(B, nb, self.bd_size).all(dim=-1)
return (
t[complete_blocks],
loss_pp.view(B, nb, self.bd_size).mean(dim=-1)[complete_blocks],
)
def _use_self_cond_this_step(self) -> bool:
"""Choose whether this step pays for the self-conditioning bootstrap.
The decision is step-seeded so every distributed rank takes the same
backbone-forward path. Divergent choices can deadlock sharded training
because only some ranks would enter the bootstrap forward's collectives.
"""
if not (self.enable_self_cond and self.training):
return False
step = int(getattr(self, "_global_step", 0) or 0)
generator = torch.Generator(device="cpu")
generator.manual_seed(0x5C0FF + step)
return torch.rand(1, generator=generator).item() > self.sc_drop_rate
def _apply_training_self_condition(
self,
e_1: Tensor,
e_t: Tensor,
t_bn: Tensor,
attn_mask,
cond_mask: Tensor,
position_ids=None,
) -> Tensor:
"""Apply predicted self-conditioning, bypassing it on dropped steps."""
if not (
self.enable_self_cond and self.training and self._use_self_cond_this_step()
):
return e_t
with torch.no_grad():
packed = torch.cat([e_1, e_t], dim=1)
logits_first, _, _ = self._forward_backbone(
packed,
t_bn,
attn_mask,
use_cache=False,
position_ids=position_ids,
)
length = e_t.shape[1]
noisy_logits_sc = logits_first[:, length:, :]
probs_sc = F.softmax(noisy_logits_sc.float(), dim=-1).to(e_1.dtype)
e_prev = probs_sc @ self.embed_tokens.weight
e_t = self._apply_self_cond(e_t, e_prev.detach())
# Self-cond transforms (notably residual's A(e_t) branch) perturb every
# position. Keep prompt/padding positions exactly on their clean embeds
# when self-conditioning is active.
mask = cond_mask.unsqueeze(-1).to(e_t.dtype)
e_t = mask * e_t + (1.0 - mask) * e_1
return e_t
def _apply_self_cond(self, e_t: Tensor, e_prev: Tensor) -> Tensor:
if self.sc_mode == "concat":
return self.sc_proj(torch.cat([e_t, e_prev], dim=-1))
elif self.sc_mode == "residual":
return e_t + self.sc_lin_a(e_t) + self.sc_lin_b(e_prev)
elif self.sc_mode == "simple":
return e_t + self.sc_lin_b(e_prev)
elif self.sc_mode == "mlp":
return e_t + self.sc_mlp(e_prev)
elif self.sc_mode == "scale":
return e_t + self.sc_scale * e_prev
return e_t
def _forward_backbone(
self,
packed_embeds: Tensor,
t_bn: Tensor,
attn_mask,
past_key_values=None,
use_cache: bool = False,
position_ids=None,
adaln_modulation=None,
t_bn_tgt=None,
compute_logits=True,
):
"""Run backbone with time conditioning.
Returns (logits, past_key_values, hidden_states).
past_key_values is None when use_cache=False.
"""
if t_bn_tgt is not None:
t_bn = t_bn + t_bn_tgt
extra_kwargs = {}
if adaln_modulation is not None:
extra_kwargs["dit_adaln_modulation"] = adaln_modulation
outputs = self.backbone.model(
inputs_embeds=packed_embeds,
attention_mask=attn_mask,
position_ids=position_ids,
past_key_values=past_key_values,
use_cache=use_cache,
dit_t_bn=t_bn,
return_dict=True,
**extra_kwargs,
)
hidden_states = outputs.last_hidden_state
logits = self.lm_head(hidden_states) if compute_logits else None
past_kv = outputs.past_key_values if use_cache else None
return logits, past_kv, hidden_states
def _get_generation_time_grid(
self,
num_steps: int,
device,
grid_mode: str = "uniform_range",
t_max: float = 0.2,
) -> Tensor:
"""Build Euler step schedule for generation."""
if num_steps < 1:
raise ValueError("num_steps must be positive")
if grid_mode == "uniform_range":
return torch.linspace(0, t_max, num_steps + 1).to(device)
elif grid_mode == "uniform":
return torch.linspace(0, 1, num_steps + 1).to(device)
elif grid_mode == "warp":
self._ensure_time_warper(device)
if getattr(self, "time_warp_per_position", False):
return self.time_warper.grid_per_position(num_steps, device=device).to(
device
)
return self.time_warper.grid(num_steps, device=device).to(device)
elif grid_mode == "warp_tau":
# Tau=1 maps to noisy t=0. Use the top quantiles, sort for Euler,
# and leave endpoints unpinned to preserve the restricted span.
self._ensure_time_warper(device)
w = self.time_warper
tau = torch.linspace(1.0 - float(t_max), 1.0, num_steps + 1, device=device)
if getattr(self, "time_warp_per_position", False):
rows = []
for i in range(w.bd_size):
edges = w._pos_edges(i)
t, _ = torch.sort(w._invert_u_row(tau, *edges))
rows.append(t)
return torch.stack(rows, dim=0).to(device) # [bd_size, num_steps+1]
edges = w._get_bin_widths_and_edges()
t, _ = torch.sort(w._invert_u_row(tau, *edges))
return t.to(device)
else:
raise ValueError(f"Unknown grid_mode: {grid_mode}")
def _euler_denoise_block(
self,
past_key_values,
block_pos_ids: Tensor,
time_grid: Tensor,
block_attn_mask=None,
prev_logit: Tensor = None,
ode_argmax: bool = False,
sde_gamma: float = 0.0,
sde_alpha: float = 0.0,
sde_vp: bool = False,
sde_noiseless: bool = False,
euler_alpha: float = 1.0,
cfg_scale: float = 1.0,
) -> tuple[Tensor, int]:
"""Denoise one block via Euler ODE with optional self-conditioning.
Assumes context is already encoded in past_key_values.
Trims stale KV entries after each step (does NOT extend cache with result).
Args:
prev_logit: [1, V] — last logit from prefix/previous block. Required
when logit_shift=True to correct ODE trajectory alignment.
Returns:
(logits [1, bd_size, V], num_forward_passes)
"""
device = block_pos_ids.device
dtype = self.embed_tokens.weight.dtype
cache_len = past_key_values.get_seq_length()
num_fwd = 0
E = self.embed_tokens.weight
use_cfg = abs(cfg_scale - 1.0) > 1e-6
per_pos = time_grid.dim() == 2
n_steps = (time_grid.shape[-1] - 1) if per_pos else (len(time_grid) - 1)
# The unconditional mask hides prefix keys from the block.
uncond_mask = None
if use_cfg:
Kc = cache_len + self.bd_size
uncond_mask = (
block_attn_mask.clone()
if block_attn_mask is not None
else torch.zeros(1, 1, self.bd_size, Kc, device=device, dtype=dtype)
)
uncond_mask[..., :cache_len] = float("-inf") # block prompt keys
def to_denoised(probs_raw):
# position-aligned predicted-clean embedding (logit_shift handling)
if self.logit_shift and prev_logit is not None:
prev_probs = F.softmax(prev_logit.float(), dim=-1).to(dtype)
pa = torch.cat([prev_probs, probs_raw[:-1]], dim=0)
return (
E[pa.argmax(-1)].unsqueeze(0)
if ode_argmax
else (pa @ E).unsqueeze(0)
)
return (
E[probs_raw.argmax(-1)].unsqueeze(0)
if ode_argmax
else (probs_raw @ E).unsqueeze(0)
)
e_t = (
torch.randn(1, self.bd_size, self._d_model, device=device, dtype=dtype)
* self._noise_sigma
* self.noise_scale_mult
)
# Conditional and unconditional passes carry separate self-conditioning.
e_self_cond = e_self_cond_u = None
if self.enable_self_cond:
init_t = (
time_grid[:, 0].to(dtype)[None, :]
if per_pos
else torch.full(
(1, self.bd_size), time_grid[0].item(), device=device, dtype=dtype
)
)
init_t_bn = self._compute_t_bn(init_t)
il_c, _, _ = self._forward_backbone(
e_t,
init_t_bn,
block_attn_mask,
past_key_values=past_key_values,
use_cache=False,
position_ids=block_pos_ids,
)
trim_kv_cache(past_key_values, past_key_values.get_seq_length() - cache_len)
e_self_cond = (
(F.softmax(il_c[0].float(), -1).to(dtype) @ E).unsqueeze(0).detach()
)
num_fwd += 1
if use_cfg:
il_u, _, _ = self._forward_backbone(
e_t,
init_t_bn,
uncond_mask,
past_key_values=past_key_values,
use_cache=False,
position_ids=block_pos_ids,
)
trim_kv_cache(
past_key_values, past_key_values.get_seq_length() - cache_len
)
e_self_cond_u = (
(F.softmax(il_u[0].float(), -1).to(dtype) @ E).unsqueeze(0).detach()
)
num_fwd += 1
for step in range(n_steps):
if per_pos:
t_val_vec = time_grid[:, step] # [bd_size]
t_out_vec = time_grid[:, step + 1] # [bd_size]
if sde_alpha > 0.0 or sde_gamma > 0.0:
dt_vec = t_out_vec - t_val_vec
if sde_alpha > 0.0:
alpha_v = torch.full_like(t_val_vec, sde_alpha)
else:
alpha_v = (1.0 - sde_gamma * dt_vec).clamp(0.0, 1.0)
# Noise at t=0 cannot be backtracked.
alpha_v = torch.where(
t_val_vec > 0.0, alpha_v, torch.ones_like(alpha_v)
)
if sde_noiseless:
# Draw noise below to keep the seeded RNG stream aligned.
coef_v = torch.zeros_like(alpha_v)
elif sde_vp:
coef_v = torch.sqrt(
(
(1.0 - alpha_v)
* (1.0 + alpha_v - 2.0 * alpha_v * t_val_vec)
).clamp(min=0.0)
)
else:
coef_v = 1.0 - alpha_v
a_b = alpha_v.to(dtype)[None, :, None] # [1, bd_size, 1]
c_b = coef_v.to(dtype)[None, :, None]
noise = (
torch.randn_like(e_t)
* self._noise_sigma
* self.noise_scale_mult
)
z_fwd = a_b * e_t + c_b * noise
t_fwd_vec = alpha_v * t_val_vec # conditioning time = alpha*t
else:
t_fwd_vec = t_val_vec
z_fwd = e_t
t_bn_block = self._compute_t_bn(t_fwd_vec.to(dtype)[None, :])
else:
t_val = time_grid[step].item()
t_out = time_grid[step + 1].item()
dt = t_out - t_val
if (sde_alpha > 0.0 or sde_gamma > 0.0) and t_val > 0.0:
alpha = (
sde_alpha
if sde_alpha > 0.0
else max(0.0, min(1.0, 1.0 - sde_gamma * dt))
)
t_fwd = alpha * t_val
coef = (
0.0
if sde_noiseless
else (
math.sqrt(
max(
0.0,
(1.0 - alpha) * (1.0 + alpha - 2.0 * alpha * t_val),
)
)
if sde_vp
else (1.0 - alpha)
)
)
noise = (
torch.randn_like(e_t)
* self._noise_sigma
* self.noise_scale_mult
)
z_fwd = alpha * e_t + coef * noise
else:
t_fwd, z_fwd = t_val, e_t
t_bn_block = self._compute_t_bn(
torch.full((1, self.bd_size), t_fwd, device=device, dtype=dtype)
)
e_in = (
self._apply_self_cond(z_fwd, e_self_cond)
if self.enable_self_cond
else z_fwd
)
logits, _, _ = self._forward_backbone(
e_in,
t_bn_block,
block_attn_mask,
past_key_values=past_key_values,
use_cache=False,
position_ids=block_pos_ids,
)
trim_kv_cache(past_key_values, past_key_values.get_seq_length() - cache_len)
num_fwd += 1
probs_c = F.softmax(logits[0].float(), dim=-1).to(dtype)
e_denoised = to_denoised(probs_c)
if use_cfg:
e_in_u = (
self._apply_self_cond(z_fwd, e_self_cond_u)
if self.enable_self_cond
else z_fwd
)
logits_u, _, _ = self._forward_backbone(
e_in_u,
t_bn_block,
uncond_mask,
past_key_values=past_key_values,
use_cache=False,
position_ids=block_pos_ids,
)
trim_kv_cache(
past_key_values, past_key_values.get_seq_length() - cache_len
)
num_fwd += 1
probs_u = F.softmax(logits_u[0].float(), dim=-1).to(dtype)
e_uncond = to_denoised(probs_u)
e_denoised = e_uncond + cfg_scale * (e_denoised - e_uncond)
if self.enable_self_cond:
e_self_cond = (probs_c @ E).unsqueeze(0).detach()
if use_cfg:
e_self_cond_u = (probs_u @ E).unsqueeze(0).detach()
if step == n_steps - 1:
break
# Broadcast position times over embeddings and compute in fp32.
if per_pos:
tf = t_fwd_vec[None, :, None].float() # [1, bd_size, 1]
to = t_out_vec[None, :, None].float()
e_t = (
z_fwd.float()
+ (to - tf)
* (e_denoised.float() - z_fwd.float())
/ (1.0 - tf).clamp(min=1e-6)
).to(dtype)
if euler_alpha != 1.0:
c_v = (to - tf) / (1.0 - tf).clamp(min=1e-6)
e_t = (
e_t.float() - (1.0 - c_v) * (1.0 - euler_alpha) * z_fwd.float()
).to(dtype)
else:
e_t = z_fwd + (t_out - t_fwd) * (e_denoised - z_fwd) / max(
1.0 - t_fwd, 1e-6
)
if euler_alpha != 1.0:
c = (t_out - t_fwd) / max(1.0 - t_fwd, 1e-6)
e_t = e_t - (1.0 - c) * (1.0 - euler_alpha) * z_fwd
return logits, num_fwd
def forward(self, input_ids=None, labels=None):
if input_ids is None or labels is None:
raise ValueError(
"Flow training requires input_ids and response-masked labels"
)
if input_ids.shape != labels.shape:
raise ValueError("Flow input_ids and labels must have the same shape")
if input_ids.shape[1] < self.bd_size or input_ids.shape[1] % self.bd_size != 0:
raise ValueError("Flow training sequences must align to bd_size")
B, L = input_ids.shape
cond_mask = self._answer_corruption_mask(labels)
e_1, e_t, t_bn, position_ids, attn_mask, t_raw = self._prepare_flow_input(
input_ids, cond_mask=cond_mask
)
e_t = self._apply_training_self_condition(
e_1=e_1,
e_t=e_t,
t_bn=t_bn,
attn_mask=attn_mask,
position_ids=position_ids,
cond_mask=cond_mask,
)
packed_embeds = torch.cat([e_1, e_t], dim=1) # [B, 2L, d]
logits, _, _ = self._forward_backbone(
packed_embeds, t_bn, attn_mask, use_cache=False, position_ids=position_ids
)
noisy_logits = logits[:, L:, :].contiguous() # [B, L, V]
_ce_tgt = (
F.pad(labels, (0, 1), value=-100)[:, 1:] if self.logit_shift else labels
)
# Only complete response blocks contribute to flow loss.
flow_valid = self._full_block_loss_mask(_ce_tgt != -100)
ce_valid = flow_valid.float() # [B, L]
ce_pp = (
-F.log_softmax(noisy_logits.float(), dim=-1)
.gather(-1, _ce_tgt.clamp(min=0).long().unsqueeze(-1))
.squeeze(-1)
) # [B, L]
flow_loss = (ce_pp * ce_valid).sum() / ce_valid.sum().clamp(min=1.0)
comps = {"flow_ce": float(flow_loss.detach())}
sq_err = None
flow_term = (
self.flow_ce_weight * flow_loss
if getattr(self, "flow_ce_enable", True)
else 0.0 * flow_loss
)
if self.ar_loss:
clean_logits = logits[:, :L, :].contiguous()
ar_loss_val = cross_entropy_loss(clean_logits, labels, shift=True)
comps["ar"] = float(ar_loss_val.detach())
loss = flow_term + self.ar_loss_weight * ar_loss_val
else:
loss = flow_term
if self.denoise_mse_weight > 0:
sq_err = self._denoising_squared_error(noisy_logits, e_1)
sq_err_masked = torch.where(flow_valid, sq_err, torch.zeros_like(sq_err))
n_valid = flow_valid.sum().clamp(min=1).to(sq_err_masked.dtype)
denoise_mse_val = sq_err_masked.sum().float() / n_valid.float()
comps["denoise_mse"] = float(denoise_mse_val.detach())
loss = loss + self.denoise_mse_weight * denoise_mse_val
# Warm-up samples cover the full time range for the EMA fit.
if self.training and self.time_warper is not None:
if self.warp_uniform_est:
# Every rank uses the same step gate before the extra forward.
if self._global_step % max(1, self.warp_est_every) == 0:
with torch.no_grad():
if self.time_warp_per_position:
t_by_pos, loss_by_pos = self._warp_block_fit_inputs(
e_1, labels, position_ids, per_position=True
)
for i in range(self.bd_size):
self.time_warper.update_warping(
t_by_pos[i], loss_by_pos[i], pos=i
)
else:
t_u, loss_u = self._warp_block_fit_inputs(
e_1, labels, position_ids
)
self.time_warper.update_warping(t_u, loss_u)
t_units = None
elif self.time_warp_per_position:
# Every rank updates position buckets in the same order.
t_by_pos, loss_by_pos = self._warp_fit_inputs(
ce_pp, sq_err, flow_valid, t_raw
)
for i in range(self.bd_size):
self.time_warper.update_warping(
t_by_pos[i].detach(), loss_by_pos[i].detach(), pos=i
)
t_units = None
else:
t_units, loss_units = self._warp_fit_inputs(
ce_pp, sq_err, flow_valid, t_raw
)
# Call on every rank, including ranks with no valid samples, so the
# per-bin count/loss all-reduce remains in the same order.
if t_units is not None:
self.time_warper.update_warping(t_units.detach(), loss_units.detach())
if self.training:
with torch.no_grad():
self._last_sampled_t = t_raw.flatten().detach().float().cpu()
comps["total"] = float(loss.detach())
comps.setdefault("flow_loss", comps.get("flow_ce"))
self._loss_components = comps
return CausalLMOutputWithPast(loss=loss, logits=noisy_logits)
@torch.no_grad()
def generate_ar(
self,
prompt_ids: Tensor,
max_new_tokens: int = 128,
temperature: float = 0.0,
seed: int = 42,
):
"""Decode from the clean causal stream after next-token training."""
device = next(self.backbone.parameters()).device
dtype = self.embed_tokens.weight.dtype
torch.manual_seed(seed)
if prompt_ids.dim() == 1:
prompt_ids = prompt_ids.unsqueeze(0)
prompt_ids = prompt_ids.to(device)
eos_id = self.backbone.config.eos_token_id
orig_attn = getattr(self.backbone.config, "_attn_implementation", None)
self.backbone.config._attn_implementation = "eager"
gen_ids = prompt_ids[0].tolist()
n = 0
try:
embeds = self.embed_tokens(prompt_ids)
t_bn = self._compute_t_bn(
torch.ones(1, prompt_ids.shape[1], device=device, dtype=dtype)
)
logits, past, _ = self._forward_backbone(
embeds, t_bn, attn_mask=None, use_cache=True
)
prev_logit = logits[0, -1, :].float() # P(next token | prompt)
while n < max_new_tokens:
if temperature > 0:
probs = F.softmax(prev_logit / temperature, dim=-1)
tok = int(torch.multinomial(probs, 1).item())
else:
tok = int(prev_logit.argmax().item())
gen_ids.append(tok)
n += 1
if tok == eos_id or n >= max_new_tokens:
break
emb = self.embed_tokens(torch.tensor([[tok]], device=device))
t1 = self._compute_t_bn(torch.ones(1, 1, device=device, dtype=dtype))
pos = torch.tensor([[past.get_seq_length()]], device=device)
logits, past, _ = self._forward_backbone(
emb,
t1,
attn_mask=None,
past_key_values=past,
use_cache=True,
position_ids=pos,
)
prev_logit = logits[0, -1, :].float()
finally:
self.backbone.config._attn_implementation = orig_attn
out = torch.tensor(gen_ids, device=device).unsqueeze(0)
return out, {
"mode": "ar",
"generated": n,
"total_steps": n,
"logical_steps": n,
"tokens_per_step": 1.0,
}
@torch.no_grad()
def generate_ar_batched(
self,
prompt_ids: Tensor,
attention_mask: Tensor,
max_new_tokens: int = 128,
temperature: float = 0.0,
seed: int = 42,
):
"""Decode left-padded prompts in a batch with masked cache padding."""
device = next(self.backbone.parameters()).device
dtype = self.embed_tokens.weight.dtype
torch.manual_seed(seed)
gen = torch.Generator(device=device).manual_seed(seed)
prompt_ids = prompt_ids.to(device)
attention_mask = attention_mask.to(device)
B, L = prompt_ids.shape
eos_id = self.backbone.config.eos_token_id
NEG = torch.finfo(dtype).min
orig_attn = getattr(self.backbone.config, "_attn_implementation", None)
self.backbone.config._attn_implementation = "eager"
cache_key_bias = (1.0 - attention_mask.to(dtype)) * NEG # [B, L]
gen_tokens = [[] for _ in range(B)]
finished = torch.zeros(B, dtype=torch.bool, device=device)
try:
embeds = self.embed_tokens(prompt_ids)
t_bn = self._compute_t_bn(torch.ones(B, L, device=device, dtype=dtype))
causal = torch.triu(
torch.full((L, L), NEG, device=device, dtype=dtype), diagonal=1
)
prefill_mask = causal[None, None] + cache_key_bias[:, None, None, :]
logits, past, _ = self._forward_backbone(
embeds, t_bn, prefill_mask, use_cache=True
)
prev = logits[:, -1, :].float() # [B, V]
gen_pos = L
n = 0
while n < max_new_tokens and not bool(finished.all()):
if temperature > 0:
probs = F.softmax(prev / temperature, dim=-1)
tok = torch.multinomial(probs, 1, generator=gen).squeeze(1)
else:
tok = prev.argmax(dim=-1) # [B]
tok_l = tok.tolist()
for b in range(B):
if finished[b]:
continue
gen_tokens[b].append(tok_l[b])
if eos_id is not None and tok_l[b] == eos_id:
finished[b] = True
n += 1
if n >= max_new_tokens or bool(finished.all()):
break
emb = self.embed_tokens(tok.unsqueeze(1)) # [B, 1, d]
t1 = self._compute_t_bn(torch.ones(B, 1, device=device, dtype=dtype))
cache_len = past.get_seq_length()
step_mask = torch.zeros(
B, 1, 1, cache_len + 1, dtype=dtype, device=device
)
step_mask[:, :, :, :cache_len] = cache_key_bias[:, None, None, :]
pos = torch.arange(gen_pos, gen_pos + 1, device=device)[None].expand(
B, -1
)
logits, past, _ = self._forward_backbone(
emb,
t1,
step_mask,
past_key_values=past,
use_cache=True,
position_ids=pos,
)
prev = logits[:, -1, :].float()
cache_key_bias = torch.cat(
[cache_key_bias, torch.zeros(B, 1, device=device, dtype=dtype)],
dim=1,
)
gen_pos += 1
finally:
if orig_attn is not None:
self.backbone.config._attn_implementation = orig_attn
out_list = []
for b in range(B):
real_prompt = prompt_ids[b][attention_mask[b].bool()].tolist()
out_list.append(real_prompt + gen_tokens[b])
return out_list, {
"mode": "ar-batched", "generated": n, "batch": B,
"forwards": n,
"per_seq_forwards": [len(tokens) for tokens in gen_tokens],
}
@staticmethod
def _debug_block_row(committed, blk_tok, tokenizer):
"""One line showing the current block state: committed positions show
their token, uncommitted positions show [noise]."""
cells = []
for i, tok in enumerate(blk_tok):
if committed[i]:
cells.append(
tokenizer.decode([tok]) if tokenizer is not None else f"<{tok}>"
)
else:
cells.append("[noise]")
return "".join(cells)
@staticmethod
def _make_frame(gen_ids, prompt_len, committed, blk_tok):
"""Snapshot of the full sequence at one decoding step, for visualization.
Returns a list of (token_id_or_None, state) where state is one of
'prompt' (context), 'clean' (already-finalized blocks), 'active'
(committed in the current block) or 'noise' (uncommitted in the
current block). Future, not-yet-started blocks are omitted.
"""
frame = [
(tid, "prompt" if j < prompt_len else "clean")
for j, tid in enumerate(gen_ids)
]
for i, tid in enumerate(blk_tok):
frame.append((tid, "active") if committed[i] else (None, "noise"))
return frame
def _denoise_flowmap_block(self, past, positions, grid, mask, **kwargs):
"""Two-time transitions; terminal prediction uses the vocabulary head."""
dtype, device = self.embed_tokens.weight.dtype, positions.device
length = positions.shape[1]
state = torch.randn(1, length, self._d_model, device=device, dtype=dtype)
state = state * self._noise_sigma * self.noise_scale_mult
path = (grid.to(dtype).view(-1, 1, 1).expand(-1, 1, length)
if grid.ndim == 1 else grid[:length].to(dtype).T.unsqueeze(1))
cache_len = past.get_seq_length()
for step, (source, target) in enumerate(zip(path[:-1], path[1:])):
condition = self._compose_flowmap_time_condition(source, target)
logits, _, _ = self._forward_backbone(
state, condition, mask, past_key_values=past,
use_cache=False, position_ids=positions)
trim_kv_cache(past, past.get_seq_length() - cache_len)
if step == len(path) - 2:
return logits.float(), len(path) - 1
delta = F.softmax(logits.float(), dim=-1).to(dtype) @ self.embed_tokens.weight
if grid.ndim == 1:
s, t = grid[step].item(), grid[step + 1].item()
updated = state + (t - s) * (delta - state) / max(1 - s, 1e-6)
else:
s, t = source.float().unsqueeze(-1), target.float().unsqueeze(-1)
updated = (state.float() + (t-s) * (delta.float()-state.float())
/ (1-s).clamp_min(1e-6)).to(dtype)
state = torch.where((target > source + 1e-8).unsqueeze(-1), updated, state)
@torch.no_grad()
def generate(
self,
prompt_ids: Tensor,
max_new_tokens: int = 128,
num_steps: int = 64,
threshold: float = 0.0,
temperature: float = 0.0,
seed: int = 42,
grid_mode: str = "uniform_range",
t_max: float = 0.2,
ode_argmax: bool = False,
sde_gamma: float = 0.0,
sde_alpha: float = 0.0,
sde_vp: bool = False,
sde_noiseless: bool = False,
euler_alpha: float = 1.0,
cfg_scale: float = 1.0,
):
"""Decode blocks with Euler steps and a KV-cached clean prefix."""
if self._two_time and (cfg_scale != 1 or ode_argmax or sde_gamma or sde_alpha or sde_vp or sde_noiseless or euler_alpha != 1):
raise ValueError("Student flow-map decoding requires unguided deterministic updates")
device = next(self.backbone.parameters()).device
dtype = self.embed_tokens.weight.dtype
torch.manual_seed(seed)
if prompt_ids.dim() == 1:
prompt_ids = prompt_ids.unsqueeze(0)
prompt_ids = prompt_ids.to(device)
prompt_len = prompt_ids.shape[1]
eos_id = self.backbone.config.eos_token_id
time_grid = self._get_generation_time_grid(
num_steps, device, grid_mode=grid_mode, t_max=t_max
)
# Eager attention supports the dense block mask.
orig_attn = getattr(self.backbone.config, "_attn_implementation", None)
self.backbone.config._attn_implementation = "eager"
gen_ids = prompt_ids[0].tolist()
tokens_generated = 0
total_forward = 0
# Logical steps exclude CFG and self-conditioning bootstrap forwards.
total_logical = 0
try:
prompt_embeds = self.embed_tokens(prompt_ids) # [1, P, d]
t_bn_prompt = self._compute_t_bn(
torch.ones(1, prompt_len, device=device, dtype=dtype)
)
prompt_logits, past_key_values, _ = self._forward_backbone(
prompt_embeds,
t_bn_prompt,
attn_mask=None,
use_cache=True,
)
if prompt_logits.dim() == 2: # some envs return [L, V]
prompt_logits = prompt_logits.unsqueeze(0)
prev_logit = prompt_logits[0, -1:, :] # [1, V]
# Only shifted logits make the prefix logit a next-token prediction.
ar_mode = self.bd_size == 1 and num_steps == 1 and self.logit_shift
while tokens_generated < max_new_tokens:
block_start_pos = len(gen_ids)
if ar_mode:
total_forward += 1
total_logical += 1
final_logits = prev_logit[0].float() # [V]
if temperature > 0:
final_probs = F.softmax(final_logits / temperature, dim=-1)
tok = torch.multinomial(final_probs.unsqueeze(0), 1).item()
else:
tok = final_logits.argmax().item()
new_tokens = [tok]
else:
block_pos_ids = torch.arange(
block_start_pos,
block_start_pos + self.bd_size,
device=device,
).unsqueeze(0)
cache_len = past_key_values.get_seq_length()
if self.noisy_causal:
block_attn_mask = None
else:
block_attn_mask = torch.zeros(
1,
1,
self.bd_size,
cache_len + self.bd_size,
dtype=dtype,
device=device,
)
denoise = self._denoise_flowmap_block if self._two_time else self._euler_denoise_block
logits, num_fwd = denoise(
past_key_values,
block_pos_ids,
time_grid,
block_attn_mask,
prev_logit=prev_logit if self.logit_shift else None,
ode_argmax=ode_argmax,
sde_gamma=sde_gamma,
sde_alpha=sde_alpha,
sde_vp=sde_vp,
sde_noiseless=sde_noiseless,
euler_alpha=euler_alpha,
cfg_scale=cfg_scale,
)
total_forward += num_fwd
total_logical += num_steps
if self.logit_shift:
logits = torch.cat([prev_logit.unsqueeze(0), logits], dim=1)
remaining = max_new_tokens - tokens_generated
final_logits = logits[0, : min(self.bd_size, remaining)].float()
if temperature > 0:
final_probs = F.softmax(final_logits / temperature, dim=-1)
preds = torch.multinomial(final_probs, 1).squeeze(-1).tolist()
else:
final_probs = F.softmax(final_logits, dim=-1)
preds = final_logits.argmax(dim=-1).tolist()
# Keep the longest prefix above the confidence threshold.
if threshold > 0.0:
confidences = final_probs.max(dim=-1).values
new_tokens = [preds[0]]
for j in range(1, len(preds)):
if confidences[j].item() >= threshold:
new_tokens.append(preds[j])
else:
break
else:
new_tokens = preds
if eos_id is not None and eos_id in new_tokens:
new_tokens = new_tokens[: new_tokens.index(eos_id) + 1]
gen_ids.extend(new_tokens)
tokens_generated += len(new_tokens)
if tokens_generated >= max_new_tokens or (
eos_id is not None and eos_id in new_tokens
):
break
block_ids = torch.tensor([new_tokens], device=device)
block_embeds = self.embed_tokens(block_ids)
accepted_len = len(new_tokens)
accepted_pos_ids = torch.arange(
block_start_pos,
block_start_pos + accepted_len,
device=device,
).unsqueeze(0)
t_bn_final = self._compute_t_bn(
torch.ones(1, accepted_len, device=device, dtype=dtype)
)
re_logits, past_key_values, _ = self._forward_backbone(
block_embeds,
t_bn_final,
attn_mask=None,
past_key_values=past_key_values,
use_cache=True,
position_ids=accepted_pos_ids,
)
prev_logit = re_logits[0, -1:, :]
finally:
if orig_attn is not None:
self.backbone.config._attn_implementation = orig_attn
generated = gen_ids[prompt_len : prompt_len + max_new_tokens]
result = torch.tensor(
[gen_ids[:prompt_len] + generated], device=device, dtype=torch.long
)
return result, {
"threshold": threshold,
"total_steps": total_forward,
"logical_steps": total_logical,
"gen_len": len(generated),
"hit_eos": eos_id is not None and eos_id in generated,
"tokens_per_step": tokens_generated / max(total_forward, 1),
"logical_tokens_per_step": tokens_generated / max(total_logical, 1),
"block_size": self.bd_size,
"num_euler_steps": num_steps,
}
@torch.no_grad()
def generate_self_speculative(
self,
prompt_ids: Tensor,
max_new_tokens: int = 128,
seed: int = 42,
grid_mode: str = "uniform_range",
t_max: float = 0.2,
self_cond: bool = False,
):
"""Self-speculative single-step decoding (batch size 1).
Each round drafts N = bd_size tokens with a bidirectional block. The
first proposal is checked against the prefix's next-token prediction.
On a miss, a compact correction pass commits that token and drafts the
next block. On a hit, a fused pass verifies the remaining proposals
and drafts a successor block for every possible accepted prefix. Useful
K/V is retained by index selection.
Layout of the fused forward (after the committed prefix, which lives in
the KV cache):
d_0 [noise a_0] d_1 [noise a_1] ... d_{N-1} [noise a_{N-1}]
Mask rules (each row also sees the whole prefix):
- draft token d_i: sees d_0..d_i (causal over drafts, incl. itself), no
noise. Its clean next-token logit verifies d_{i+1}.
- noise block a_i (N slots): sees d_0..d_i and its own slots
(bidirectional). It is the next-block draft assuming d_0..d_i accepted.
- recovery correction c sees only prefix + itself; its noise block
sees prefix + c + all of its own bidirectional slots.
Verification is greedy. Decoding uses one flow step and no SDE.
Optional self-conditioning adds a bootstrap pass before each draft.
Returns (output_ids [1, P+gen_len], stats_dict).
"""
device = next(self.backbone.parameters()).device
dtype = self.embed_tokens.weight.dtype
torch.manual_seed(seed)
if prompt_ids.dim() == 1:
prompt_ids = prompt_ids.unsqueeze(0)
prompt_ids = prompt_ids.to(device)
prompt_len = prompt_ids.shape[1]
eos_id = self.backbone.config.eos_token_id
N = self.bd_size
d = self._d_model
E = self.embed_tokens.weight
# Draft self-conditioning adds a bootstrap forward; verification stays causal.
sc_on = bool(self_cond and getattr(self, "enable_self_cond", False))
tg = self._get_generation_time_grid(1, device, grid_mode=grid_mode, t_max=t_max)
if tg.dim() == 2: # per-position warp grid [N, 2]
t_start = tg[:, 0].to(dtype) # [N] block-local start times
t_target = tg[:, 1].to(dtype) # [N] flow-map target times
else:
t_start = tg[0].to(dtype).expand(N).clone()
t_target = tg[1].to(dtype).expand(N).clone()
orig_attn = getattr(self.backbone.config, "_attn_implementation", None)
self.backbone.config._attn_implementation = "eager"
gen_ids = prompt_ids[0].tolist()
committed = 0
fused_fwd = boot_fwd = corr_fwd = 0
first_token_rejects = 0
def _noise(n):
return (
torch.randn(1, n, d, device=device, dtype=dtype)
* self._noise_sigma
* self.noise_scale_mult
)
def _draft_from_cache(past, base_pos):
"""Bidirectional N-slot draft over the whole cache. Does NOT extend the
cache. Returns (draft token id list, num_forwards). With self-cond a
bootstrap forward first estimates the clean embedding (2 forwards)."""
cache_len = past.get_seq_length()
e = _noise(N)
pos = torch.arange(base_pos, base_pos + N, device=device).unsqueeze(0)
mask = torch.zeros(1, 1, N, cache_len + N, dtype=dtype, device=device)
source_t = t_start.unsqueeze(0)
t_bn = self._compose_flowmap_time_condition(
source_t,
t_target.unsqueeze(0),
)
adaln_modulation = self._precompute_shared_adaln(t_bn)
nf = 0
if sc_on:
lb, _, _ = self._forward_backbone(
e,
t_bn,
mask,
past_key_values=past,
use_cache=False,
position_ids=pos,
adaln_modulation=adaln_modulation,
)
trim_kv_cache(past, past.get_seq_length() - cache_len)
nf += 1
e_sc = (F.softmax(lb[0].float(), -1).to(dtype) @ E).unsqueeze(0)
e = self._apply_self_cond(e, e_sc)
logits, _, _ = self._forward_backbone(
e,
t_bn,
mask,
past_key_values=past,
use_cache=False,
position_ids=pos,
adaln_modulation=adaln_modulation,
)
trim_kv_cache(past, past.get_seq_length() - cache_len)
nf += 1
return logits[0].argmax(-1).tolist(), nf
# ---- static fused-layout index maps (depend only on N) ----
S = N * (N + 1)
draft_idx = [i * (N + 1) for i in range(N)] # d_i fused index
noise_pos = [draft_idx[i] + 1 + k for i in range(N) for k in range(N)]
blk = [0] * S
isd = [False] * S
for i in range(N):
blk[draft_idx[i]] = i
isd[draft_idx[i]] = True
for k in range(N):
blk[draft_idx[i] + 1 + k] = i
blk_t = torch.tensor(blk, device=device)
isd_t = torch.tensor(isd, device=device)
# allowed_fused[q,k]: draft-key k visible iff its block <= q's block;
# noise-key k visible iff q is a noise slot in the SAME block.
kd = isd_t[None, :]
allowed_fused = (kd & (blk_t[None, :] <= blk_t[:, None])) | (
(~kd) & (~isd_t)[:, None] & (blk_t[None, :] == blk_t[:, None])
)
# per-position t: 1.0 at drafts, t_start[k] at noise slot k of each block
ppt = torch.ones(1, S, device=device, dtype=dtype)
target_ppt = torch.ones(1, S, device=device, dtype=dtype)
target_mask = torch.zeros(1, S, device=device, dtype=torch.bool)
for i in range(N):
block_noise = [draft_idx[i] + 1 + k for k in range(N)]
ppt[0, block_noise] = t_start
target_ppt[0, block_noise] = t_target
target_mask[0, block_noise] = True
# relative position ids (absolute = rel + committed length P)
rel_pos = torch.zeros(S, dtype=torch.long, device=device)
for i in range(N):
rel_pos[draft_idx[i]] = i
for k in range(N):
rel_pos[draft_idx[i] + 1 + k] = i + 1 + k
neg = torch.full((), float("-inf"), device=device, dtype=dtype)
zero = torch.zeros((), device=device, dtype=dtype)
fused_allow = torch.where(allowed_fused, zero, neg) # [S, S] additive
R = N + 1
recovery_ppt = torch.cat(
(torch.ones(1, 1, device=device, dtype=dtype), t_start.unsqueeze(0)),
dim=1,
)
recovery_target_ppt = torch.cat(
(torch.ones(1, 1, device=device, dtype=dtype), t_target.unsqueeze(0)),
dim=1,
)
recovery_target_mask = torch.zeros(1, R, device=device, dtype=torch.bool)
recovery_target_mask[:, 1:] = True
recovery_t_bn = self._compose_flowmap_time_condition(
recovery_ppt,
recovery_target_ppt,
target_mask=recovery_target_mask,
)
recovery_adaln = self._precompute_shared_adaln(recovery_t_bn)
recovery_allow = torch.zeros(R, R, device=device, dtype=dtype)
recovery_allow[0, 1:] = float("-inf")
recovery_rel_pos = torch.arange(R, device=device)
def _recover_and_draft(past, base_pos, correction):
"""Commit a rejected first draft token and propose its successor."""
cache_len = past.get_seq_length()
assert cache_len == base_pos
emb = torch.empty(1, R, d, device=device, dtype=dtype)
emb[:, :1] = self.embed_tokens(
torch.tensor([[correction]], device=device)
).to(dtype)
emb[:, 1:] = _noise(N)
mask = torch.zeros(R, cache_len + R, device=device, dtype=dtype)
mask[:, cache_len:] = recovery_allow
pos = (recovery_rel_pos + base_pos).unsqueeze(0)
nf = 0
if sc_on:
lb, _, _ = self._forward_backbone(
emb,
recovery_t_bn,
mask[None, None],
past_key_values=past,
use_cache=False,
position_ids=pos,
adaln_modulation=recovery_adaln,
)
trim_kv_cache(past, past.get_seq_length() - cache_len)
nf += 1
e_sc = (F.softmax(lb[0, 1:].float(), -1).to(dtype) @ E).unsqueeze(0)
emb[:, 1:] = self._apply_self_cond(emb[:, 1:], e_sc)
logits, past, _ = self._forward_backbone(
emb,
recovery_t_bn,
mask[None, None],
past_key_values=past,
use_cache=True,
position_ids=pos,
adaln_modulation=recovery_adaln,
)
nf += 1
past = gather_kv_cache(past, torch.arange(cache_len + 1, device=device))
return past, logits[0, 0], logits[0, 1:].argmax(-1).tolist(), nf
try:
# Prefill prompt (clean, causal) → cache + prev_logit predicting d_0.
pe = self.embed_tokens(prompt_ids)
t_bn_p = self._compute_t_bn(
torch.ones(1, prompt_len, device=device, dtype=dtype)
)
prompt_adaln = self._precompute_shared_adaln(t_bn_p)
pl, past, _ = self._forward_backbone(
pe,
t_bn_p,
attn_mask=None,
use_cache=True,
adaln_modulation=prompt_adaln,
)
prev_logit = pl[0, -1, :] # predicts next token = d_0
# The prompt modulation can be much larger than the static decode
# layouts and is never reused after prefill.
del prompt_adaln, t_bn_p, pe
D, nf = _draft_from_cache(past, past.get_seq_length()) # bootstrap draft
boot_fwd += nf
t_bn_fused = self._compose_flowmap_time_condition(
ppt,
target_ppt,
target_mask=target_mask,
)
fused_adaln = self._precompute_shared_adaln(t_bn_fused)
while committed < max_new_tokens:
P = past.get_seq_length()
correction = int(prev_logit.argmax().item())
if correction != D[0]:
first_token_rejects += 1
new_tokens = [correction]
past, prev_logit, D, nf = _recover_and_draft(
past, P, correction
)
corr_fwd += nf
else:
emb = torch.empty(1, S, d, device=device, dtype=dtype)
emb[0, draft_idx] = self.embed_tokens(
torch.tensor(D, device=device)
).to(dtype)
emb[0, noise_pos] = _noise(N * N)[0]
m = torch.empty(S, P + S, device=device, dtype=dtype)
m[:, :P] = 0.0
m[:, P:] = fused_allow
pos = (rel_pos + P).unsqueeze(0)
if sc_on:
lb, _, _ = self._forward_backbone(
emb,
t_bn_fused,
m[None, None],
past_key_values=past,
use_cache=False,
position_ids=pos,
adaln_modulation=fused_adaln,
)
trim_kv_cache(past, past.get_seq_length() - P)
fused_fwd += 1
e_sc = F.softmax(lb[0, noise_pos].float(), -1).to(dtype) @ E
emb[0, noise_pos] = self._apply_self_cond(
emb[:, noise_pos, :], e_sc.unsqueeze(0)
)[0]
logits, past, _ = self._forward_backbone(
emb,
t_bn_fused,
m[None, None],
past_key_values=past,
use_cache=True,
position_ids=pos,
adaln_modulation=fused_adaln,
)
fused_fwd += 1
logf = logits[0]
j = 1
for i in range(1, N):
if logf[draft_idx[i - 1]].argmax().item() == D[i]:
j += 1
else:
break
new_tokens = D[:j]
keep = list(range(P)) + [P + draft_idx[i] for i in range(j)]
past = gather_kv_cache(past, torch.tensor(keep, device=device))
prev_logit = logf[draft_idx[j - 1]]
win = [draft_idx[j - 1] + 1 + k for k in range(N)]
D = logf[win].argmax(-1).tolist()
if eos_id is not None and eos_id in new_tokens:
new_tokens = new_tokens[: new_tokens.index(eos_id) + 1]
gen_ids.extend(new_tokens)
committed += len(new_tokens)
break
gen_ids.extend(new_tokens)
committed += len(new_tokens)
finally:
if orig_attn is not None:
self.backbone.config._attn_implementation = orig_attn
generated = gen_ids[prompt_len : prompt_len + max_new_tokens]
total_fwd = fused_fwd + boot_fwd + corr_fwd
result = torch.tensor(
[gen_ids[:prompt_len] + generated], device=device, dtype=torch.long
)
apf = len(generated) / max(total_fwd, 1)
return result, {
"total_steps": total_fwd,
"logical_steps": total_fwd,
"gen_len": len(generated),
"hit_eos": eos_id is not None and eos_id in generated,
"tokens_per_step": apf,
"logical_tokens_per_step": apf,
"accepted_per_forward": apf,
"fused_forwards": fused_fwd,
"bootstrap_forwards": boot_fwd,
"correction_forwards": corr_fwd,
"first_token_rejects": first_token_rejects,
"self_cond": sc_on,
"block_size": N,
"num_euler_steps": 1,
}
@torch.no_grad()
def generate_batched(
self,
prompt_ids: Tensor,
attention_mask: Tensor,
max_new_tokens: int = 128,
num_steps: int = 64,
temperature: float = 0.0,
seed: int = 42,
grid_mode: str = "uniform_range",
t_max: float = 0.2,
ode_argmax: bool = False,
sde_gamma: float = 0.0,
sde_alpha: float = 0.0,
sde_vp: bool = False,
sde_noiseless: bool = False,
euler_alpha: float = 1.0,
cfg_scale: float = 1.0,
):
"""Decode full blocks for left-padded prompts with a masked KV cache.
Shared initial noise makes deterministic decoding nearly batch invariant.
SDE churn uses independent noise for each batch row.
"""
device = next(self.backbone.parameters()).device
dtype = self.embed_tokens.weight.dtype
torch.manual_seed(seed)
noise_gen = torch.Generator(device=device).manual_seed(seed)
prompt_ids = prompt_ids.to(device)
attention_mask = attention_mask.to(device)
B, L = prompt_ids.shape
bd = self.bd_size
eos_id = self.backbone.config.eos_token_id
NEG = torch.finfo(dtype).min
time_grid = self._get_generation_time_grid(
num_steps, device, grid_mode=grid_mode, t_max=t_max
)
per_pos = time_grid.dim() == 2
n_steps = (time_grid.shape[-1] - 1) if per_pos else (len(time_grid) - 1)
orig_attn = getattr(self.backbone.config, "_attn_implementation", None)
self.backbone.config._attn_implementation = "eager"
cache_key_bias = (1.0 - attention_mask.to(dtype)) * NEG # [B, L]
gen_tokens = [[] for _ in range(B)]
finished = torch.zeros(B, dtype=torch.bool, device=device)
total_forward = 0
tokens_generated = 0
E = self.embed_tokens.weight
def shared_noise():
"""One standalone-equivalent RNG draw, broadcast across batch rows."""
return (
(
torch.randn(
1,
bd,
self._d_model,
generator=noise_gen,
device=device,
dtype=dtype,
)
* self._noise_sigma
* self.noise_scale_mult
)
.expand(B, bd, self._d_model)
.contiguous()
)
try:
prompt_embeds = self.embed_tokens(prompt_ids)
t_bn_p = self._compute_t_bn(torch.ones(B, L, device=device, dtype=dtype))
causal = torch.triu(
torch.full((L, L), NEG, device=device, dtype=dtype), diagonal=1
)
prefill_mask = causal[None, None] + cache_key_bias[:, None, None, :]
prompt_logits, past, _ = self._forward_backbone(
prompt_embeds, t_bn_p, prefill_mask, use_cache=True
)
if prompt_logits.dim() == 2: # some envs return [B*L, V]
prompt_logits = prompt_logits.reshape(B, -1, prompt_logits.shape[-1])
prev_logit = prompt_logits[:, -1:, :] # [B, 1, V]
gen_pos = L
while tokens_generated < max_new_tokens and not bool(finished.all()):
block_pos_ids = torch.arange(gen_pos, gen_pos + bd, device=device)[
None
].expand(B, -1)
cache_len = past.get_seq_length()
block_mask = torch.zeros(
B, 1, bd, cache_len + bd, dtype=dtype, device=device
)
block_mask[:, :, :, :cache_len] = cache_key_bias[:, None, None, :]
if self.noisy_causal:
cw = torch.triu(
torch.full((bd, bd), NEG, device=device, dtype=dtype),
diagonal=1,
)
block_mask[:, :, :, cache_len:] = cw[None, None]
e_t = shared_noise()
# Unconditional guidance hides prefix keys.
use_cfg = abs(cfg_scale - 1.0) > 1e-6
uncond_mask = None
if use_cfg:
uncond_mask = block_mask.clone()
uncond_mask[:, :, :, :cache_len] = NEG
use_shift = self.logit_shift
def to_den(pr): # probs [B, bd, V] -> predicted-clean embedding
if use_shift and prev_logit is not None:
pp = F.softmax(prev_logit.float(), dim=-1).to(dtype)
pa = torch.cat([pp, pr[:, :-1]], dim=1)
return E[pa.argmax(-1)] if ode_argmax else pa @ E
return E[pr.argmax(-1)] if ode_argmax else pr @ E
e_self_cond = e_self_cond_u = None
if self.enable_self_cond:
init_t = (
time_grid[:, 0].to(dtype)[None, :].expand(B, bd)
if per_pos
else torch.full(
(B, bd), time_grid[0].item(), device=device, dtype=dtype
)
)
init_t_bn = self._compute_t_bn(init_t)
init_logits, _, _ = self._forward_backbone(
e_t,
init_t_bn,
block_mask,
past_key_values=past,
use_cache=False,
position_ids=block_pos_ids,
)
trim_kv_cache(past, past.get_seq_length() - cache_len)
total_forward += 1
e_self_cond = (
F.softmax(init_logits.float(), -1).to(dtype) @ E
).detach()
if use_cfg:
il_u, _, _ = self._forward_backbone(
e_t,
init_t_bn,
uncond_mask,
past_key_values=past,
use_cache=False,
position_ids=block_pos_ids,
)
trim_kv_cache(past, past.get_seq_length() - cache_len)
total_forward += 1
e_self_cond_u = (
F.softmax(il_u.float(), -1).to(dtype) @ E
).detach()
logits = None
for step in range(n_steps):
if per_pos:
t_val_vec = time_grid[:, step] # [bd]
t_out_vec = time_grid[:, step + 1] # [bd]
if sde_alpha > 0.0 or sde_gamma > 0.0:
dt_vec = t_out_vec - t_val_vec
if sde_alpha > 0.0:
alpha_v = torch.full_like(t_val_vec, sde_alpha)
else:
alpha_v = (1.0 - sde_gamma * dt_vec).clamp(0.0, 1.0)
alpha_v = torch.where(
t_val_vec > 0.0, alpha_v, torch.ones_like(alpha_v)
)
if sde_noiseless:
coef_v = torch.zeros_like(
alpha_v
) # deterministic α-rewind, no injection
elif sde_vp:
coef_v = torch.sqrt(
(
(1.0 - alpha_v)
* (1.0 + alpha_v - 2.0 * alpha_v * t_val_vec)
).clamp(min=0.0)
)
else:
coef_v = 1.0 - alpha_v
a_b = alpha_v.to(dtype)[
None, :, None
] # [1, bd, 1] → bcast over B,d
c_b = coef_v.to(dtype)[None, :, None]
noise = (
torch.randn_like(e_t)
* self._noise_sigma
* self.noise_scale_mult
)
z_fwd = a_b * e_t + c_b * noise
t_fwd_vec = (
alpha_v * t_val_vec
) # conditioning time = alpha*t
else:
t_fwd_vec = t_val_vec
z_fwd = e_t
t_bn_block = self._compute_t_bn(
t_fwd_vec.to(dtype)[None, :].expand(B, bd)
)
else:
t_val = time_grid[step].item()
t_out = time_grid[step + 1].item()
dt = t_out - t_val
# sde_vp uses variance-preserving noise at the rewound time.
if (sde_alpha > 0.0 or sde_gamma > 0.0) and t_val > 0.0:
alpha = (
sde_alpha
if sde_alpha > 0.0
else max(0.0, min(1.0, 1.0 - sde_gamma * dt))
)
t_fwd = alpha * t_val
coef = (
0.0
if sde_noiseless
else (
math.sqrt(
max(
0.0,
(1.0 - alpha)
* (1.0 + alpha - 2.0 * alpha * t_val),
)
)
if sde_vp
else (1.0 - alpha)
)
)
# Draw even in noiseless mode to keep seeded RNG aligned.
bk = (
torch.randn(
B,
bd,
self._d_model,
generator=noise_gen,
device=device,
dtype=dtype,
)
* self._noise_sigma
* self.noise_scale_mult
)
z_fwd = alpha * e_t + coef * bk
else:
t_fwd, z_fwd = t_val, e_t
t_bn_block = self._compute_t_bn(
torch.full((B, bd), t_fwd, device=device, dtype=dtype)
)
e_t_input = (
self._apply_self_cond(z_fwd, e_self_cond)
if self.enable_self_cond
else z_fwd
)
logits, _, _ = self._forward_backbone(
e_t_input,
t_bn_block,
block_mask,
past_key_values=past,
use_cache=False,
position_ids=block_pos_ids,
)
trim_kv_cache(past, past.get_seq_length() - cache_len)
total_forward += 1
probs_c = F.softmax(logits.float(), dim=-1).to(dtype)
e_denoised = to_den(probs_c)
if use_cfg:
e_in_u = (
self._apply_self_cond(z_fwd, e_self_cond_u)
if self.enable_self_cond
else z_fwd
)
logits_u, _, _ = self._forward_backbone(
e_in_u,
t_bn_block,
uncond_mask,
past_key_values=past,
use_cache=False,
position_ids=block_pos_ids,
)
trim_kv_cache(past, past.get_seq_length() - cache_len)
total_forward += 1
probs_u = F.softmax(logits_u.float(), dim=-1).to(dtype)
e_uncond = to_den(probs_u)
e_denoised = e_uncond + cfg_scale * (e_denoised - e_uncond)
if self.enable_self_cond:
e_self_cond = (probs_c @ E).detach()
if use_cfg:
e_self_cond_u = (probs_u @ E).detach()
if step == n_steps - 1:
break
# Broadcast position times over embeddings in fp32.
if per_pos:
tf = t_fwd_vec[None, :, None].float() # [1, bd, 1]
to = t_out_vec[None, :, None].float()
e_t = (
z_fwd.float()
+ (to - tf)
* (e_denoised.float() - z_fwd.float())
/ (1.0 - tf).clamp(min=1e-6)
).to(dtype)
if euler_alpha != 1.0:
c_v = (to - tf) / (1.0 - tf).clamp(min=1e-6)
e_t = (
e_t.float()
- (1.0 - c_v) * (1.0 - euler_alpha) * z_fwd.float()
).to(dtype)
else:
e_t = z_fwd + (t_out - t_fwd) * (e_denoised - z_fwd) / max(
1.0 - t_fwd, 1e-6
)
if euler_alpha != 1.0:
c = (t_out - t_fwd) / max(1.0 - t_fwd, 1e-6)
e_t = e_t - (1.0 - c) * (1.0 - euler_alpha) * z_fwd
if use_shift:
final_logits = torch.cat([prev_logit, logits], dim=1).float()
else:
final_logits = logits.float() # [B, bd, V]
if temperature > 0:
fl = final_logits / temperature
probs = F.softmax(fl, dim=-1)
preds = torch.multinomial(
probs.reshape(-1, probs.shape[-1]), 1, generator=noise_gen
).reshape(B, -1)
else:
preds = final_logits.argmax(dim=-1) # [B, bd(+1)]
preds = preds[:, :bd] # keep bd tokens
pred_l = preds.tolist()
for b in range(B):
if finished[b]:
continue
toks = pred_l[b]
if eos_id is not None and eos_id in toks:
toks = toks[: toks.index(eos_id) + 1]
finished[b] = True
gen_tokens[b].extend(toks)
tokens_generated += bd
# Re-encode every row to keep the cache rectangular.
block_embeds = self.embed_tokens(preds)
t_bn_final = self._compute_t_bn(
torch.ones(B, bd, device=device, dtype=dtype)
)
re_cache_len = past.get_seq_length()
recode_mask = torch.zeros(
B, 1, bd, re_cache_len + bd, dtype=dtype, device=device
)
recode_mask[:, :, :, :re_cache_len] = cache_key_bias[:, None, None, :]
cw = torch.triu(
torch.full((bd, bd), NEG, device=device, dtype=dtype), diagonal=1
)
recode_mask[:, :, :, re_cache_len:] = cw[None, None]
re_logits, past, _ = self._forward_backbone(
block_embeds,
t_bn_final,
recode_mask,
past_key_values=past,
use_cache=True,
position_ids=block_pos_ids,
)
prev_logit = re_logits[:, -1:, :]
cache_key_bias = torch.cat(
[cache_key_bias, torch.zeros(B, bd, device=device, dtype=dtype)],
dim=1,
)
gen_pos += bd
finally:
if orig_attn is not None:
self.backbone.config._attn_implementation = orig_attn
out_list = []
for b in range(B):
real_prompt = prompt_ids[b][attention_mask[b].bool()].tolist()
out_list.append(real_prompt + gen_tokens[b][:max_new_tokens])
return out_list, {
"mode": "flow-batched",
"total_steps": total_forward,
"block_size": bd,
"num_euler_steps": num_steps,
"batch": B,
}