URSA/scripts/train_onestep_ursa_dimo.py

#!/usr/bin/env python3
# Copyright (c) 2024-present, BAAI. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# -----------------------------------------------------------------------
"""URSA → URSA one-step distillation via Di[M]O-style on-policy training.



Verified native inference regime (from A/B testing — ground truth):

  height=320, width=512, num_frames=49, guidance_scale=7, teacher_steps=50.

  no_cfg (guidance_scale=1) does NOT produce valid output for this URSA checkpoint.

  All defaults below align to this verified regime.



Algorithm (9 stages per iteration)

------------------------------------

  teacher  : frozen URSA — provides supervision at pseudo-intermediate x_t.

  student  : trainable copy — 1-step target.

  aux      : trainable copy — approximates teacher at x_t; reduces REINFORCE variance.



  Stage 1  : tokenise prompts (cond + uncond when CFG enabled) → txt_ids [B,L]

  Stage 2  : sample x_init [B,T,H,W] ~ Uniform(K) (+ optional p_init mixing)

  Stage 3  : student 1-step forward on x_init (cond only) → x_hat, logp, H

  Stage 4  : pseudo-intermediate x_t = scheduler.add_noise(x_hat, t)

  Stage 5  : teacher forward on x_t (CFG=7 dual-branch is the default)

  Stage 6  : aux forward → Jeffrey KD

  Stage 7  : student forward on x_t → KL KD

  Stage 8  : reward = -KL(z_T_cond, z_S_cond)  [detached]

  Stage 9  : two-backward student update



Usage:

  # Smoke test (verified native regime):

  python scripts/train_onestep_ursa_dimo.py \\

      --teacher_ckpt /path/to/URSA --prompt_file prompts.txt \\

      --enable_teacher_cfg --teacher_cfg_scale 7.0 \\

      --num_frames 49 --height 320 --width 512 --dry_run



  # Full training:

  python scripts/train_onestep_ursa_dimo.py \\

      --teacher_ckpt /path/to/URSA --prompt_file prompts.txt \\

      --enable_teacher_cfg --teacher_cfg_scale 7.0 \\

      --num_frames 49 --height 320 --width 512 \\

      --batch_size 1 --num_steps 10000 --out_dir ./outputs/dimo_cfg

"""

import argparse
import copy
import json
import math
import os
import sys

import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader

_REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
if _REPO_ROOT not in sys.path:
    sys.path.insert(0, _REPO_ROOT)

from diffnext.pipelines import URSAPipeline
from src.distill.prompt_dataset import InfiniteDataLoader, PromptDataset, make_collate_fn, CSVSpec
from src.distill.utils_ursa_inputs import (
    build_ursa_inputs,
    compute_latents_shape,
    corrupt_tokens,
    extract_visual_logits,
    sample_t_curriculum,
)

def _get_logits(out):
    if isinstance(out, (tuple, list)):
        return out[0]
    if hasattr(out, "sample"):
        return out.sample
    if hasattr(out, "logits"):
        return out.logits
    return out

# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def parse_args():
    p = argparse.ArgumentParser(description="URSA DiMO one-step distillation")

    # Model / data
    p.add_argument("--teacher_ckpt", required=True)
    p.add_argument("--prompt_file", required=True)
    p.add_argument("--out_dir", default="./outputs/dimo")

    # Video geometry (verified native: 320×512×49)
    p.add_argument("--num_frames", type=int, default=49)
    p.add_argument("--height", type=int, default=320)
    p.add_argument("--width", type=int, default=512)
    p.add_argument("--max_prompt_length", type=int, default=320)

    # Training
    p.add_argument("--batch_size", type=int, default=1)
    p.add_argument("--num_steps", type=int, default=10_000)
    p.add_argument("--lr_student", type=float, default=1e-5)
    p.add_argument("--lr_aux", type=float, default=1e-5)
    p.add_argument("--weight_decay", type=float, default=0.01)
    p.add_argument("--grad_clip", type=float, default=1.0)
    p.add_argument("--mixed_precision", default="bf16", choices=["fp16", "bf16", "fp32"])
    p.add_argument("--seed", type=int, default=42)
    p.add_argument("--log_every", type=int, default=50)
    p.add_argument("--save_every", type=int, default=1000)

    # Loss weights
    p.add_argument("--lambda_pg", type=float, default=1.0)
    p.add_argument("--lambda_kd", type=float, default=0.5)
    p.add_argument("--lambda_ent", type=float, default=0.01)
    p.add_argument("--tau", type=float, default=1.0, help="Student sampling temperature")
    p.add_argument("--tau_kd", type=float, default=1.0, help="KD softmax temperature")

    # ---- Teacher CFG (DiMO true_cfg style) ----------------------------
    p.add_argument("--enable_teacher_cfg", action="store_true", default=False,
                   help="Enable teacher-side CFG for KD target. "
                        "False → prior single-branch behavior (fallback).")
    p.add_argument("--teacher_cfg_scale", type=float, default=7.0,
                   help="CFG scale s (verified working value=7)")
    p.add_argument("--teacher_cfg_prob", type=float, default=1.0,
                   help="Max prob of using guided target per sample (after warmup)")
    p.add_argument("--teacher_cfg_warmup_steps", type=int, default=2000,
                   help="Steps to ramp teacher_cfg_prob 0 → teacher_cfg_prob")
    p.add_argument("--teacher_cfg_trunc", type=float, default=0.9,
                   help="t threshold: when t >= trunc, s=1. Set >=1.0 to disable.")
    p.add_argument("--lambda_kd_uncond", type=float, default=0.3,
                   help="Weight for uncond-branch KD / aux loss")
    p.add_argument("--reward_use_guided", action="store_true", default=False,
                   help="[RISKY] Use guided teacher logits for REINFORCE reward.")
    # ---- Eval CFG (inference-time) -----------------------------------
    p.add_argument("--eval_cfg_scale", type=float, default=7.0)
    p.add_argument("--use_cfg_eval", action="store_true", default=True)

    # DiMO extensions
    p.add_argument("--use_surrogate_grad", action="store_true",
                   help="DiMO surrogate MSE trick applied to Stage-3 logits")
    p.add_argument("--lambda_surr", type=float, default=1.0)
    p.add_argument("--fake_rounds", type=int, default=1,
                   help="Aux updates per generator update (DiMO=2)")

    # Stability
    p.add_argument("--t_curriculum_steps", type=int, default=10_000)
    p.add_argument("--p_mix_corrupt_frac", type=float, default=0.2)
    p.add_argument("--p_init_mix_ratio", type=float, default=0.2)
    p.add_argument("--collapse_warn_frac", type=float, default=0.2)

    # Debug
    p.add_argument("--dry_run", action="store_true",
                   help="Run 1 step + grad-flow check, then exit")
    p.add_argument("--debug_dump", type=int, default=0,
                   help="Dump token histogram + x_hat every N steps (0=off)")

    p.add_argument("--device", type=int, default=0)
    return p.parse_args()


# ---------------------------------------------------------------------------
# Checkpoint
# ---------------------------------------------------------------------------

def save_checkpoint(model, path: str, name: str = "student"):
    os.makedirs(path, exist_ok=True)
    ckpt_path = os.path.join(path, f"{name}.pt")
    torch.save(model.state_dict(), ckpt_path)
    print(f"[save] {ckpt_path}")


# ---------------------------------------------------------------------------
# Stable KL / Jeffrey divergence helpers (float32 + log_softmax)
# ---------------------------------------------------------------------------

def _stable_kl(z_p: torch.Tensor, z_q: torch.Tensor, tau: float = 1.0) -> torch.Tensor:
    """KL(p||q) from raw logits, float32 + log_softmax.  → [B] (mean over N tokens).



    p = softmax(z_p/tau),  q = softmax(z_q/tau)

    KL(p||q) = sum_k  p_k * (log p_k - log q_k)



    Both log_p and log_q are computed via log_softmax to avoid

    log(softmax(...)) numerical issues.

    """
    lp = F.log_softmax(z_p.float() / tau, dim=-1)   # [B, N, K]
    lq = F.log_softmax(z_q.float() / tau, dim=-1)   # [B, N, K]
    return (lp.exp() * (lp - lq)).sum(-1).mean(-1)  # [B]


def _stable_jeffrey(z_p: torch.Tensor, z_q: torch.Tensor, tau: float = 1.0) -> torch.Tensor:
    """Symmetric KL (Jeffrey) from logits, float32 + log_softmax. → [B]."""
    return _stable_kl(z_p, z_q, tau) + _stable_kl(z_q, z_p, tau)


# ---------------------------------------------------------------------------
# Batch-concat input builder (ONE forward for cond + uncond)
# ---------------------------------------------------------------------------

def _build_dual_inputs(teacher_ref, txt_cond, txt_uncond, x_t, latents_shape, device):
    """Concatenate cond+uncond into a single [2B] forward-pass input.



    Returns (ids_dual [2B, L+N+1], rpos_dual [2B, L+N+1, 3], N).

    After the forward: chunk(2, dim=0) → (z_cond [B], z_uncond [B]).



    All three models (teacher/aux/student) share the SAME ids_dual / rpos_dual

    so the tokens are constructed only once per step.

    """
    txt_dual = torch.cat([txt_cond, txt_uncond], dim=0)  # [2B, L]
    x_t_dual = torch.cat([x_t, x_t], dim=0)             # [2B, T, H, W]
    return build_ursa_inputs(teacher_ref, txt_dual, x_t_dual, latents_shape, device)


# ---------------------------------------------------------------------------
# flex_attn probe / reset helpers
# ---------------------------------------------------------------------------

def _probe_flex_attn(model, label: str = "") -> object:
    """Return the FlexAttentionCausal2D object if present, else None."""
    return getattr(model, "flex_attn", None)


def _print_flex_attn_state(model, label: str):
    fa = _probe_flex_attn(model, label)
    if fa is None:
        print(f"  [flex_attn/{label}] not present on model")
        return
    print(
        f"  [flex_attn/{label}] offsets={fa.offsets!r}  "
        f"block_mask={'set' if fa.block_mask is not None else 'None'}  "
        f"cu_offsets={'set' if fa.cu_offsets is not None else 'None'}"
    )


def _reset_flex_attn(model, label: str = "", verbose: bool = False):
    """Reset flex_attn to None offsets so standard causal attention is used.



    Our distillation training processes each sample independently (batch dim)

    so block-packed attention (offsets != None) is not needed and must be cleared

    to avoid cross-sample mask contamination.

    """
    fa = _probe_flex_attn(model, label)
    if fa is None:
        return
    old_offsets = fa.offsets
    fa.offsets = None
    fa.block_mask = None
    fa.cu_offsets = None
    if verbose:
        print(f"  [flex_attn/{label}] reset: was={old_offsets!r} → None (standard causal)")


# ---------------------------------------------------------------------------
# Teacher CFG target construction
# ---------------------------------------------------------------------------

def _compute_cfg_scale(t: torch.Tensor, cfg_scale: float, trunc: float) -> torch.Tensor:
    """Per-sample CFG scale [B]: s=cfg_scale when t < trunc, else s=1."""
    s = torch.full_like(t, cfg_scale)
    if trunc < 1.0:
        s = torch.where(t >= trunc, torch.ones_like(t), s)
    return s


def _cfg_warmup_prob(step: int, cfg_prob: float, warmup_steps: int) -> float:
    """Linear warmup: 0 → cfg_prob over warmup_steps steps."""
    if warmup_steps <= 0:
        return cfg_prob
    return cfg_prob * min(1.0, step / warmup_steps)


def _build_guided_logits(

    z_T_cond: torch.Tensor,    # [B, N, K] float32

    z_T_uncond: torch.Tensor,  # [B, N, K] float32

    t: torch.Tensor,           # [B] ∈ (0,1)

    cfg_scale: float,

    trunc: float,

) -> torch.Tensor:
    """z_guided = z_uncond + s*(z_cond - z_uncond), per-sample s [B,1,1]."""
    s = _compute_cfg_scale(t, cfg_scale, trunc).view(-1, 1, 1)  # [B,1,1]
    return z_T_uncond + s * (z_T_cond - z_T_uncond)             # [B, N, K]


def _select_target(

    z_guided: torch.Tensor,   # [B, N, K]

    z_cond: torch.Tensor,     # [B, N, K]

    use_guided: torch.Tensor, # [B] bool — per-sample selection

) -> torch.Tensor:
    """Per-sample: z_guided where use_guided[b]=True, else z_cond."""
    mask = use_guided.view(-1, 1, 1).expand_as(z_cond)
    return torch.where(mask, z_guided, z_cond)


# ---------------------------------------------------------------------------
# Gradient-flow debug
# ---------------------------------------------------------------------------

def debug_grad_flow(

    teacher, student, aux,

    txt_cond, txt_uncond, x_t, latents_shape, device, K, N, tau, tau_kd,

    enable_teacher_cfg,

):
    """One fwd+bwd without optimizer.step().



    Asserts:

      - teacher: zero grads (frozen)

      - aux:     non-zero grads after loss_aux.backward()

      - student: non-zero grads after loss_student.backward()



    All cond/uncond forwards are batch-concatenated per requirement (1).

    """
    print("\n" + "=" * 64)
    print("[grad_flow] Starting gradient flow debug …")
    B = txt_cond.size(0)

    # -- Stage 3: student on x_init (cond only) ----------------------
    x_init_dbg = torch.randint(0, K, x_t.shape, device=device, dtype=torch.long)
    ids_init, rpos_init, _ = build_ursa_inputs(teacher, txt_cond, x_init_dbg, latents_shape, device)
    logits_s = student(ids_init, rope_pos=rpos_init).sample
    z_s = extract_visual_logits(logits_s.float(), N, K)
    p_s = F.softmax(z_s / tau, dim=-1)
    x_hat = torch.multinomial(p_s.view(-1, K), 1).view(B, N)
    logp = p_s.clamp(1e-8).log().gather(-1, x_hat.unsqueeze(-1)).squeeze(-1).sum(-1)
    H_mean = -(p_s * p_s.clamp(1e-8).log()).sum(-1).mean()

    # -- Stage 5: teacher forward — [2B] if CFG, else [B] ------------
    if enable_teacher_cfg and txt_uncond is not None:
        ids_dual, rpos_dual, _ = _build_dual_inputs(teacher, txt_cond, txt_uncond, x_t, latents_shape, device)
        with torch.no_grad():
            logits_T_dual = teacher(ids_dual, rope_pos=rpos_dual).sample.float()
        z_T_dual = extract_visual_logits(logits_T_dual, N, K)
        z_T_cond_dbg, z_T_uncond_dbg = z_T_dual.chunk(2, dim=0)
        t_dbg = torch.full((B,), 0.5, device=device, dtype=torch.float32)
        z_T_guided_dbg = _build_guided_logits(
            z_T_cond_dbg.float(), z_T_uncond_dbg.float(), t_dbg, 3.0, 0.9)
        z_T_target_dbg = z_T_guided_dbg.detach()
        print(f"  [grad_flow] z_T_cond   shape={z_T_cond_dbg.shape}  "
              f"min={z_T_cond_dbg.min():.3f}  max={z_T_cond_dbg.max():.3f}")
        print(f"  [grad_flow] z_T_uncond shape={z_T_uncond_dbg.shape}  "
              f"min={z_T_uncond_dbg.min():.3f}  max={z_T_uncond_dbg.max():.3f}")
        print(f"  [grad_flow] z_T_guided shape={z_T_guided_dbg.shape}  "
              f"min={z_T_guided_dbg.min():.3f}  max={z_T_guided_dbg.max():.3f}")
        ids_t_ref    = ids_dual[:B]
        rpos_t_ref   = rpos_dual[:B]
        ids_fwd      = ids_dual
        rpos_fwd     = rpos_dual
    else:
        ids_t_ref, rpos_t_ref, _ = build_ursa_inputs(teacher, txt_cond, x_t, latents_shape, device)
        with torch.no_grad():
            logits_T = teacher(ids_t_ref, rope_pos=rpos_t_ref).sample.float()
        z_T_target_dbg = extract_visual_logits(logits_T, N, K).detach()
        ids_fwd  = ids_t_ref
        rpos_fwd = rpos_t_ref

    # Dual-path shape check (teacher vs student, same input)
    with torch.no_grad():
        z_T_ref2 = extract_visual_logits(
            teacher(ids_t_ref, rope_pos=rpos_t_ref).sample.float(), N, K)
    z_S_ref2 = extract_visual_logits(
        student(ids_t_ref.detach(), rope_pos=rpos_t_ref.detach()).sample.float(), N, K)
    if z_T_ref2.shape != z_S_ref2.shape:
        raise RuntimeError(
            f"[FATAL] Dual-path shape mismatch: z_T={z_T_ref2.shape} z_S={z_S_ref2.shape}"
        )
    print(f"  [grad_flow] Dual-path check OK: shape={z_T_ref2.shape}")

    # -- Aux backward — [2B] if CFG, else [B] -------------------------
    logits_A = aux(ids_fwd.detach(), rope_pos=rpos_fwd.detach()).sample
    if enable_teacher_cfg and txt_uncond is not None:
        z_A_dual2 = extract_visual_logits(logits_A.float(), N, K)
        z_A_cond_dbg, _ = z_A_dual2.chunk(2, dim=0)
    else:
        z_A_cond_dbg = extract_visual_logits(logits_A.float(), N, K)
    loss_aux_sample = _stable_jeffrey(z_T_target_dbg, z_A_cond_dbg, tau_kd)
    loss_aux = loss_aux_sample.mean()
    loss_aux.backward()

    teacher_grads = [p.grad for p in teacher.parameters() if p.grad is not None]
    aux_grads     = [p.grad.norm().item() for p in aux.parameters() if p.grad is not None]
    print(f"  [grad_flow] teacher grads with non-None grad: {len(teacher_grads)} (must be 0)")
    if aux_grads:
        print(f"  [grad_flow] aux grad norm  min={min(aux_grads):.3e}  "
              f"mean={sum(aux_grads)/len(aux_grads):.3e}  max={max(aux_grads):.3e}")
    else:
        print("  [grad_flow] ⚠️ aux has NO grads")
    for param in aux.parameters():
        param.grad = None

    # -- Student backward — [B] (cond only for simplicity) ------------
    logits_S = student(ids_t_ref.detach(), rope_pos=rpos_t_ref.detach()).sample
    z_S_cond = extract_visual_logits(logits_S.float(), N, K)
    loss_kd = _stable_kl(z_T_target_dbg, z_S_cond, tau_kd).mean()
    adv = (loss_aux_sample.detach() * 0 + 1.0)   # dummy advantage (shape check)
    assert not adv.requires_grad, "[BUG] adv must be detached"
    loss_student = -(adv * logp).mean() + loss_kd - 0.01 * H_mean
    loss_student.backward()

    student_grads = [p.grad.norm().item() for p in student.parameters() if p.grad is not None]
    if student_grads:
        print(f"  [grad_flow] student grad norm  min={min(student_grads):.3e}  "
              f"mean={sum(student_grads)/len(student_grads):.3e}  "
              f"max={max(student_grads):.3e}")
    else:
        print("  [grad_flow] ⚠️ student has NO grads — diagnosing:")
        print(f"    logp.requires_grad={logp.requires_grad}")
        print(f"    z_s.requires_grad={z_s.requires_grad}")

    assert len(teacher_grads) == 0,   "teacher has grads — not frozen"
    assert len(aux_grads)     > 0,    "aux has no grads after loss_aux.backward()"
    assert len(student_grads) > 0,    "student has no grads — grad flow broken"

    for m in (student, aux):
        for param in m.parameters():
            param.grad = None

    print("  [grad_flow] All gradient assertions PASSED ✓")
    print("=" * 64 + "\n")


# ---------------------------------------------------------------------------
# Main training loop
# ---------------------------------------------------------------------------

def main():
    args = parse_args()

    device = torch.device("cuda", args.device) if torch.cuda.is_available() else torch.device("cpu")
    dtype_map = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}
    compute_dtype = dtype_map[args.mixed_precision]
    torch.manual_seed(args.seed)
    os.makedirs(args.out_dir, exist_ok=True)

    # -- Verified regime validation ----------------------------------------
    _NATIVE = dict(height=320, width=512, num_frames=49, guidance_scale=7.0)
    is_native = (
        args.height == _NATIVE["height"]
        and args.width == _NATIVE["width"]
        and args.num_frames == _NATIVE["num_frames"]
    )
    print(f"[init] verified_native_regime={is_native}  "
          f"geometry=({args.num_frames}×{args.height}×{args.width})  "
          f"teacher_cfg_scale={args.teacher_cfg_scale if args.enable_teacher_cfg else 'OFF'}")
    if not is_native:
        print(f"[WARN] Current geometry ({args.num_frames}×{args.height}×{args.width}) "
              f"is not the verified native URSA regime "
              f"({_NATIVE['num_frames']}×{_NATIVE['height']}×{_NATIVE['width']}). "
              "Distillation quality may degrade or become invalid.")
    if not args.enable_teacher_cfg:
        print("[WARN] Teacher CFG is DISABLED.  no_cfg is known to produce "
              "blank/blurry output for this URSA checkpoint.  "
              "Distillation without CFG is unlikely to produce useful results.")
    elif args.teacher_cfg_scale != _NATIVE["guidance_scale"]:
        print(f"[WARN] teacher_cfg_scale={args.teacher_cfg_scale} differs from "
              f"the verified working value ({_NATIVE['guidance_scale']}).")

    if args.enable_teacher_cfg and args.reward_use_guided:
        print("[WARN] --reward_use_guided is ON — can cause mode collapse, watch tok_entropy.")

    # -- Load pipeline ---------------------------------------------------
    print(f"[init] Loading from {args.teacher_ckpt} …")
    pipe = URSAPipeline.from_pretrained(
        args.teacher_ckpt, torch_dtype=compute_dtype, trust_remote_code=True
    ).to(device)

    tokenizer = pipe.tokenizer
    scheduler = pipe.scheduler
    scheduler.to(device=device)

    vae_t_stride = getattr(pipe.vae.config, "temporal_stride", 4)
    vae_s_stride = getattr(pipe.vae.config, "spatial_stride", 8)
    latents_shape = compute_latents_shape(
        args.num_frames, args.height, args.width, vae_t_stride, vae_s_stride
    )
    T, H, W = latents_shape
    N = T * H * W
    K = scheduler.codebook_size
    print(
        f"[init] latents_shape=({T},{H},{W})  N={N}  K={K}  "
        f"CFG={'ON' if args.enable_teacher_cfg else 'OFF'}"
    )

    # -- Pre-compute uncond token IDs (empty string, [1, L]) --------------
    txt_uncond_base = tokenizer(
        [""], max_length=args.max_prompt_length, padding="max_length",
        padding_side="left", truncation=True, return_tensors="pt",
    ).input_ids.to(device)   # [1, L]

    # -- Three models ----------------------------------------------------
    teacher = pipe.transformer.eval().requires_grad_(False)
    student = copy.deepcopy(teacher).train().requires_grad_(True)
    aux     = copy.deepcopy(teacher).train().requires_grad_(True)

    # -- flex_attn: reset offsets to None (standard causal attn) ---------
    # Our training processes B independent sequences in a batch, so block-packed
    # offsets are not needed and must be cleared before any forward call.
    if args.dry_run:
        print("[init] flex_attn state before reset:")
        for m, lbl in ((teacher, "teacher"), (student, "student"), (aux, "aux")):
            _print_flex_attn_state(m, lbl)
    for m, lbl in ((teacher, "teacher"), (student, "student"), (aux, "aux")):
        _reset_flex_attn(m, lbl, verbose=True)
    if args.dry_run:
        print("[init] flex_attn state after reset:")
        for m, lbl in ((teacher, "teacher"), (student, "student"), (aux, "aux")):
            _print_flex_attn_state(m, lbl)

    opt_student = torch.optim.AdamW(
        student.parameters(), lr=args.lr_student, weight_decay=args.weight_decay
    )
    opt_aux = torch.optim.AdamW(
        aux.parameters(), lr=args.lr_aux, weight_decay=args.weight_decay
    )

    # -- Dataset ----------------------------------------------------------
    # dataset = PromptDataset(args.prompt_file, shuffle=True, seed=args.seed)
    collate = make_collate_fn(tokenizer, args.max_prompt_length, device)
    # loader  = DataLoader(
    #     dataset, batch_size=args.batch_size, shuffle=True,
    #     drop_last=True, num_workers=0, collate_fn=collate,
    # )
    dataset = PromptDataset(
        args.prompt_file,
        shuffle_files=True,
        shuffle_buffer=50000,   # 例如 50k buffer，够用且不占太多内存
        seed=args.seed,
        infinite=True,
        csv=CSVSpec(caption_field="caption"),  # Koala 默认就是 caption
    )

    loader = DataLoader(
        dataset,
        batch_size=args.batch_size,
        shuffle=False,          # IMPORTANT for IterableDataset
        drop_last=True,
        num_workers=2,          # 视 IO 调大
        collate_fn=collate,
        pin_memory=True,
    )
    inf_loader = InfiniteDataLoader(loader)

    # -- Pre-training sanity check ---------------------------------------
    _sanity_check_forward(teacher, scheduler, latents_shape, device, K, args.dry_run)

    # -- Training state --------------------------------------------------
    baseline_ema: float = 0.0
    x_hat_prev = None
    initial_tok_entropy: float = None
    dump_dir = os.path.join(args.out_dir, "debug_dumps") if args.debug_dump > 0 else None

    num_steps = 1 if args.dry_run else args.num_steps
    print(f"[train] {'DRY RUN' if args.dry_run else f'{num_steps} steps'} "
          f"| CFG={args.enable_teacher_cfg}")

    for step in range(1, num_steps + 1):

        # ----------------------------------------------------------------
        # Stage 1: Tokenise → txt_cond [B, L], txt_uncond [B, L]
        # ----------------------------------------------------------------
        txt_cond = next(inf_loader)    # [B, L]
        txt_cond = txt_cond.to(device, non_blocking=True)
        B = txt_cond.size(0)

        txt_uncond = None
        if args.enable_teacher_cfg:
            txt_uncond = txt_uncond_base.expand(B, -1)  # [B, L]

        # ----------------------------------------------------------------
        # Stage 2: x_init ~ Uniform(K) (+ optional p_init mixing)
        # ----------------------------------------------------------------
        x_init = _sample_x_init(B, T, H, W, K, device, x_hat_prev, args)

        # ----------------------------------------------------------------
        # Stage 3: Student 1-step forward on x_init — COND only.
        #
        # Gradient needed: logp and H flow back through p_s → student.
        # ----------------------------------------------------------------
        with torch.no_grad():
            ids_init, rpos_init, _ = build_ursa_inputs(
                teacher, txt_cond, x_init, latents_shape, device)
            logits_s_init = student(ids_init, rope_pos=rpos_init).sample     # [B, L+N+1, D]
            z_s  = extract_visual_logits(logits_s_init.float(), N, K)        # [B, N, K]
            p_s  = F.softmax(z_s / args.tau, dim=-1)                         # [B, N, K]
            x_hat = torch.multinomial(p_s.view(-1, K), 1).view(B, N)         # [B, N]
            # logp  = p_s.clamp(1e-8).log().gather(
            #     -1, x_hat.unsqueeze(-1)).squeeze(-1).sum(-1)                  # [B]
            # H_mean = -(p_s * p_s.clamp(1e-8).log()).sum(-1).mean()
        x_hat_4d = x_hat.view(B, T, H, W)

        # ----------------------------------------------------------------
        # Stage 4: Pseudo-intermediate x_t
        # ----------------------------------------------------------------
        t = sample_t_curriculum(B, device, step, warmup_steps=args.t_curriculum_steps)
        with torch.no_grad():
            x_t = scheduler.add_noise(x_hat_4d, t)   # [B, T, H, W], long

        # ----------------------------------------------------------------
        # Stage 5: Teacher forward — single [2B] forward when CFG enabled.
        #
        # ids_dual / rpos_dual are SHARED by teacher, aux, and student to
        # avoid redundant input construction.
        # ----------------------------------------------------------------
        with torch.no_grad():
            if args.enable_teacher_cfg:
                # ONE [2B] forward = cond (first B) + uncond (last B)
                ids_dual, rpos_dual, _ = _build_dual_inputs(
                    teacher, txt_cond, txt_uncond, x_t, latents_shape, device)
                logits_T_dual = teacher(ids_dual, rope_pos=rpos_dual).sample.float()
                z_T_dual      = extract_visual_logits(logits_T_dual, N, K)  # [2B, N, K]
                z_T_cond, z_T_uncond = z_T_dual.chunk(2, dim=0)             # [B, N, K] each
                ids_t   = ids_dual[:B]    # cond half — alias (no copy)
                rpos_t  = rpos_dual[:B]
            else:
                ids_t, rpos_t, _ = build_ursa_inputs(
                    teacher, txt_cond, x_t, latents_shape, device)
                logits_T  = teacher(ids_t, rope_pos=rpos_t).sample.float()
                z_T_cond  = extract_visual_logits(logits_T, N, K)  # [B, N, K]
                z_T_uncond = None
                ids_dual  = ids_t
                rpos_dual = rpos_t

        # -- CFG guided target (float32, per-sample Bernoulli) ----------
        z_T_guided = None
        if args.enable_teacher_cfg:
            z_T_cond_f   = z_T_cond.float()
            z_T_uncond_f = z_T_uncond.float()
            z_T_guided   = _build_guided_logits(
                z_T_cond_f, z_T_uncond_f, t,
                args.teacher_cfg_scale, args.teacher_cfg_trunc)

            # per-sample Bernoulli: use_guided[b] ~ Bernoulli(p_guided)
            p_guided      = _cfg_warmup_prob(
                step, args.teacher_cfg_prob, args.teacher_cfg_warmup_steps)
            use_guided     = torch.rand(B, device=device) < p_guided   # [B] bool
            use_guided_ratio = use_guided.float().mean().item()
            z_T_target    = _select_target(z_T_guided, z_T_cond_f, use_guided)  # [B, N, K]
        else:
            use_guided      = torch.zeros(B, dtype=torch.bool, device=device)
            use_guided_ratio = 0.0
            z_T_target      = z_T_cond.float()

        # z_T_target is the KD target — must have no grad path to teacher
        z_T_target = z_T_target.detach()

        # ----------------------------------------------------------------
        # Stage 6: Aux forward (fake_rounds) — single [2B] forward when CFG.
        # ----------------------------------------------------------------
        loss_aux_cond_v_last   = None
        loss_aux_uncond_v_last = None
        loss_aux_cond_sample_last = None

        for _fr in range(args.fake_rounds):
            opt_aux.zero_grad()

            if args.enable_teacher_cfg:
                # ONE [2B] forward: cond+uncond in one shot
                logits_A_dual = aux(ids_dual.detach(), rope_pos=rpos_dual.detach()).sample
                z_A_dual      = extract_visual_logits(logits_A_dual.float(), N, K)  # [2B, N, K]
                z_A_cond, z_A_uncond = z_A_dual.chunk(2, dim=0)

                # Cond: Jeffrey(z_T_target, z_A_cond)
                loss_aux_cond_sample = _stable_jeffrey(z_T_target, z_A_cond, args.tau_kd)  # [B]
                loss_aux_cond_v      = loss_aux_cond_sample.mean()

                # Uncond: Jeffrey(z_T_uncond, z_A_uncond)
                z_T_uncond_det = z_T_uncond.float().detach()
                loss_aux_uncond_sample = _stable_jeffrey(z_T_uncond_det, z_A_uncond, args.tau_kd)
                loss_aux_uncond_v      = loss_aux_uncond_sample.mean()

                loss_aux_v = loss_aux_cond_v + args.lambda_kd_uncond * loss_aux_uncond_v
            else:
                logits_A  = aux(ids_t.detach(), rope_pos=rpos_t.detach()).sample
                z_A_cond  = extract_visual_logits(logits_A.float(), N, K)

                loss_aux_cond_sample = _stable_jeffrey(z_T_target, z_A_cond, args.tau_kd)  # [B]
                loss_aux_cond_v      = loss_aux_cond_sample.mean()
                loss_aux_uncond_v    = torch.tensor(0.0, device=device)
                loss_aux_v           = loss_aux_cond_v

            loss_aux_v.backward()
            if args.grad_clip > 0:
                torch.nn.utils.clip_grad_norm_(aux.parameters(), args.grad_clip)
            opt_aux.step()
            # make sure aux grads are cleared and no graph is retained
            for p in aux.parameters():
                p.grad = None

        loss_aux_cond_v_last      = loss_aux_cond_v.detach()
        loss_aux_uncond_v_last    = loss_aux_uncond_v.detach()
        loss_aux_cond_sample_last = loss_aux_cond_sample.detach()   # [B]

        # # ----------------------------------------------------------------
        # # Stage 7: Student KD forward on x_t — single [2B] when CFG.
        # # Dual-path consistency check included.
        # # ----------------------------------------------------------------
        # if args.enable_teacher_cfg:
        #     # ONE [2B] forward
        #     logits_S_dual = student(ids_dual.detach(), rope_pos=rpos_dual.detach()).sample
        #     z_S_dual      = extract_visual_logits(logits_S_dual.float(), N, K)  # [2B, N, K]
        #     z_S_cond, z_S_uncond = z_S_dual.chunk(2, dim=0)
        # else:
        #     logits_S = student(ids_t.detach(), rope_pos=rpos_t.detach()).sample
        #     z_S_cond = extract_visual_logits(logits_S.float(), N, K)   # [B, N, K]
        #     z_S_uncond = None

        # # Dual-path shape consistency check
        # if z_T_cond.shape != z_S_cond.shape:
        #     raise RuntimeError(
        #         f"[FATAL] Dual-path shape mismatch: "
        #         f"z_T_cond={z_T_cond.shape}  z_S_cond={z_S_cond.shape} — "
        #         "vocab slicing inconsistency."
        #     )

        # # KD losses (from raw logits, float32 + log_softmax)
        # loss_kd_cond   = _stable_kl(z_T_target, z_S_cond, args.tau_kd).mean()
        # loss_kd_uncond_v = torch.tensor(0.0, device=device)

        # if args.enable_teacher_cfg and z_S_uncond is not None:
        #     z_T_uncond_det2 = z_T_uncond.float().detach()
        #     loss_kd_uncond_v = _stable_kl(z_T_uncond_det2, z_S_uncond, args.tau_kd).mean()

        # loss_kd = loss_kd_cond + args.lambda_kd_uncond * loss_kd_uncond_v

        # # ----------------------------------------------------------------
        # # Stage 8: REINFORCE reward + advantage
        # #
        # # INVARIANT: reward and adv MUST NOT carry student gradients.
        # #   - z_S_cond is detached before entering reward computation.
        # #   - adv is explicitly detached.
        # #   - Runtime assertions enforce this.
        # # ----------------------------------------------------------------
        # if args.enable_teacher_cfg:
        #     if args.reward_use_guided:
        #         z_T_for_rew = z_T_target   # already detached (guided, see §5)
        #     else:
        #         z_T_for_rew = z_T_cond.float().detach()  # non-guided cond (stable default)
        #     # Both inputs are detached: no student gradient leaks into reward.
        #     reward = -_stable_kl(
        #         z_T_for_rew.detach(), z_S_cond.detach(), args.tau)   # [B]
        # else:
        #     reward = -loss_aux_cond_sample_last   # [B], already detached

        # # Mandatory detach assertions: catch reward/adv gradient leaks early.
        # assert not reward.requires_grad, (
        #     "[BUG] reward.requires_grad=True — student gradient leaked into reward. "
        #     "Ensure z_S_cond is detached in reward computation."
        # )
        # baseline_ema = 0.99 * baseline_ema + 0.01 * reward.mean().item()
        # adv = (reward - baseline_ema).detach()   # [B]
        # assert not adv.requires_grad, "[BUG] adv.requires_grad=True — explicit detach failed"

        # loss_pg = -(adv * logp).mean()

        # # ----------------------------------------------------------------
        # # Stage 9: Student loss + update
        # # ----------------------------------------------------------------
        # opt_student.zero_grad()

        # lambda_ent_eff = args.lambda_ent * (1.0 + 2.0 * use_guided_ratio)
        # loss_student = (
        #     args.lambda_pg * loss_pg
        #     + args.lambda_kd * loss_kd
        #     - lambda_ent_eff * H_mean
        # )

        # # Optional surrogate gradient (DiMO MSE trick — applied to Stage-3 logits z_s)
        # loss_surr = None
        # if args.use_surrogate_grad:
        #     with torch.no_grad():
        #         logits_A_ref = aux(ids_t.detach(), rope_pos=rpos_t.detach()).sample
        #         z_A_ref      = extract_visual_logits(logits_A_ref.float(), N, K)
        #     # grad_surr = (p_A - p_T): pushes z_s toward teacher distribution
        #     p_A_ref  = F.softmax(z_A_ref.float() / args.tau_kd, dim=-1).detach()
        #     p_T_surr = F.softmax(z_T_target / args.tau_kd, dim=-1).detach()
        #     grad_surr = (p_A_ref - p_T_surr).detach()
        #     loss_surr = 0.5 * F.mse_loss(z_s, (z_s - grad_surr).detach())
        #     loss_student = loss_student + args.lambda_surr * loss_surr

        # loss_student.backward()
        # if args.grad_clip > 0:
        #     torch.nn.utils.clip_grad_norm_(student.parameters(), args.grad_clip)
        # opt_student.step()

        # # p_init mixing: save x_hat_4d for next step
        # x_hat_prev = x_hat_4d.detach().clone()
        
        # ----------------------------------------------------------------
        # Stage 7: Student KD forward on x_t — single [2B] when CFG.
        # ----------------------------------------------------------------
        if args.enable_teacher_cfg:
            logits_S_dual = _get_logits(student(ids_dual.detach(), rope_pos=rpos_dual.detach())).float()
            z_S_dual = extract_visual_logits(logits_S_dual, N, K)  # [2B, N, K]
            z_S_cond, z_S_uncond = z_S_dual.chunk(2, dim=0)
        else:
            logits_S = _get_logits(student(ids_t.detach(), rope_pos=rpos_t.detach())).float()
            z_S_cond = extract_visual_logits(logits_S, N, K)
            z_S_uncond = None

        if z_T_cond.shape != z_S_cond.shape:
            raise RuntimeError(f"[FATAL] Dual-path shape mismatch: z_T_cond={z_T_cond.shape} z_S_cond={z_S_cond.shape}")

        loss_kd_cond = _stable_kl(z_T_target, z_S_cond, args.tau_kd).mean()
        loss_kd_uncond_v = torch.tensor(0.0, device=device)
        if args.enable_teacher_cfg and (z_S_uncond is not None):
            loss_kd_uncond_v = _stable_kl(z_T_uncond.float().detach(), z_S_uncond, args.tau_kd).mean()
        loss_kd = loss_kd_cond + args.lambda_kd_uncond * loss_kd_uncond_v

        # ----------------------------------------------------------------
        # Stage 8: reward + advantage (detached)
        # ----------------------------------------------------------------
        if args.enable_teacher_cfg and args.reward_use_guided:
            z_T_for_rew = z_T_target  # already detached
        else:
            z_T_for_rew = z_T_cond.float().detach()

        reward = -_stable_kl(z_T_for_rew.detach(), z_S_cond.detach(), args.tau)  # [B]
        assert not reward.requires_grad

        baseline_ema = 0.99 * baseline_ema + 0.01 * reward.mean().item()
        adv = (reward - baseline_ema).detach()
        assert not adv.requires_grad

        # ----------------------------------------------------------------
        # Stage 9: update student in two backward passes (KD then PG/Ent)
        # ----------------------------------------------------------------
        opt_student.zero_grad(set_to_none=True)

        # (9a) KD backward first (frees KD graph)
        (args.lambda_kd * loss_kd).backward()

        # (9b) Policy + entropy: need a fresh forward on x_init WITH grad
        ids_init, rpos_init, _ = build_ursa_inputs(teacher, txt_cond, x_init, latents_shape, device)
        logits_s_pol = _get_logits(student(ids_init, rope_pos=rpos_init)).float()
        z_s_pol = extract_visual_logits(logits_s_pol, N, K)

        logp_tok = F.log_softmax(z_s_pol / args.tau, dim=-1)   # [B,N,K]
        p_s_pol = logp_tok.exp()

        # fixed action: x_hat sampled in Stage 3 (no_grad)
        logp_sum = logp_tok.gather(-1, x_hat.unsqueeze(-1)).squeeze(-1).sum(-1)  # [B], sum over N tokens
        logp = logp_sum / N                                                      # [B], per-token average logp (RECOMMENDED)
        
        H_mean = -(p_s_pol * logp_tok).sum(-1).mean()

        loss_pg = -(adv * logp).mean()
        
        lambda_ent_eff = args.lambda_ent * (1.0 + 2.0 * use_guided_ratio)
        (loss_pg * args.lambda_pg - H_mean * lambda_ent_eff).backward()

        # (optional) surrogate grad — put it here; WARNING: extra forward makes it heavier
        loss_surr = None
        if args.use_surrogate_grad:
            with torch.no_grad():
                logits_A_ref = _get_logits(aux(ids_t.detach(), rope_pos=rpos_t.detach())).float()
                z_A_ref = extract_visual_logits(logits_A_ref, N, K)
            p_A_ref = F.softmax(z_A_ref / args.tau_kd, dim=-1).detach()
            p_T_ref = F.softmax(z_T_target / args.tau_kd, dim=-1).detach()
            grad_surr = (p_A_ref - p_T_ref).detach()
            loss_surr = 0.5 * F.mse_loss(z_s_pol, (z_s_pol - grad_surr).detach())
            (args.lambda_surr * loss_surr).backward()

        if args.grad_clip > 0:
            torch.nn.utils.clip_grad_norm_(student.parameters(), args.grad_clip)
        opt_student.step()
        
        # p_init mixing: save x_hat_4d for next step
        x_hat_prev = x_hat_4d.detach() #.clone()

        # ----------------------------------------------------------------
        # Post-step: assertions (step 1), collapse detection, logging
        # ----------------------------------------------------------------
        if step == 1:
            _run_assertions(
                x_init, ids_init, rpos_init,
                z_s, p_s, logp,
                z_T_cond, z_S_cond, x_t, K, N, B, T, H, W,
                teacher.config.lm_vocab_size,
                z_T_uncond=z_T_uncond,
                z_T_guided=z_T_guided,
                dry_run=args.dry_run,
            )

        tok_entropy = _token_histogram_entropy(x_hat, K)
        if initial_tok_entropy is None:
            initial_tok_entropy = tok_entropy

        if tok_entropy < args.collapse_warn_frac * initial_tok_entropy:
            print(
                f"[COLLAPSE WARNING] step={step}  tok_entropy={tok_entropy:.3f}  "
                f"initial={initial_tok_entropy:.3f}  "
                f"ratio={tok_entropy/max(initial_tok_entropy, 1e-8):.2f} < "
                f"{args.collapse_warn_frac}. "
                "Increase --lambda_ent (try 0.05) or --tau."
            )

        if step % args.log_every == 0 or args.dry_run:
            surr_str = f"  loss_surr={loss_surr.item():.4f}" if loss_surr is not None else ""
            print(
                f"[step {step:>6d}] "
                f"loss_aux_cond={loss_aux_cond_v_last.item():.3e}  "
                f"loss_aux_uncond={loss_aux_uncond_v_last.item():.3e}  "
                f"loss_kd_cond={loss_kd_cond.item():.4f}  "
                f"loss_kd_uncond={loss_kd_uncond_v.item():.4f}  "
                f"loss_pg={loss_pg.item():.4f}"
                f"{surr_str}  "
                f"H={H_mean.item():.3f}  tok_H={tok_entropy:.3f}  "
                f"guided_ratio={use_guided_ratio:.2f}  "
                f"baseline={baseline_ema:.4f}  "
                f"mean_logp_tok={logp.mean().item():.3f}"
            )

        if args.debug_dump > 0 and step % args.debug_dump == 0:
            _dump_debug(dump_dir, step, x_hat, K)

        if not args.dry_run and step % args.save_every == 0:
            ckpt_dir = os.path.join(args.out_dir, f"step_{step:06d}")
            save_checkpoint(student, ckpt_dir, "student")
            save_checkpoint(aux, ckpt_dir, "aux")

    # -- dry_run: full grad-flow check after the single training step ----
    if args.dry_run:
        print("\n[dry_run] Running gradient flow debug …")
        txt_dbg   = next(inf_loader)
        B_dbg     = txt_dbg.size(0)
        x_t_dbg   = torch.randint(0, K, (B_dbg, T, H, W), device=device, dtype=torch.long)
        txt_u_dbg = (txt_uncond_base.expand(B_dbg, -1)
                     if args.enable_teacher_cfg else None)
        debug_grad_flow(
            teacher, student, aux,
            txt_dbg, txt_u_dbg, x_t_dbg, latents_shape, device, K, N,
            args.tau, args.tau_kd, args.enable_teacher_cfg,
        )
        _dry_run_patches_789(teacher, latents_shape, K, N, device)
        print("[dry_run] Done. All checks (1-9) PASSED. Exiting.")
        return

    # Final save
    final_dir = os.path.join(args.out_dir, "final")
    save_checkpoint(student, final_dir, "student")
    save_checkpoint(aux, final_dir, "aux")
    print("[done] Training complete.")


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _sample_x_init(B, T, H, W, K, device, x_hat_prev, args):
    x_init = torch.randint(0, K, (B, T, H, W), device=device, dtype=torch.long)
    if x_hat_prev is not None and args.p_init_mix_ratio > 0:
        n_mix = max(1, int(B * args.p_init_mix_ratio))
        x_init[:n_mix] = corrupt_tokens(x_hat_prev[:n_mix], r=args.p_mix_corrupt_frac, K=K)
    return x_init


def _token_histogram_entropy(x_hat: torch.Tensor, K: int) -> float:
    counts = x_hat.flatten().bincount(minlength=K).float()
    p = counts / counts.sum()
    p = p[p > 0]
    return float(-(p * p.log()).sum().item())


def _dump_debug(dump_dir: str, step: int, x_hat: torch.Tensor, K: int):
    os.makedirs(dump_dir, exist_ok=True)
    counts = x_hat.flatten().bincount(minlength=K).tolist()
    with open(os.path.join(dump_dir, f"step_{step:06d}_hist.json"), "w") as fh:
        json.dump({"step": step, "counts": counts}, fh)
    torch.save(x_hat.cpu(), os.path.join(dump_dir, f"step_{step:06d}_xhat.pt"))
    print(f"[debug_dump] step={step} saved to {dump_dir}")


def _run_assertions(

    x_init, ids_init, rpos_init,

    z_s, p_s, logp,

    z_T_cond, z_S_cond, x_t,

    K, N, B, T, H, W, lm_vocab_size,

    z_T_uncond=None, z_T_guided=None,

    dry_run=False,

):
    """Full shape / value-domain / consistency assertions (run at step=1)."""
    print("[assert] Running shape/value assertions …")

    L_plus_N1 = ids_init.size(1)
    txt_len   = L_plus_N1 - (N + 1)

    # x_init
    assert x_init.dtype == torch.long, f"x_init dtype={x_init.dtype}"
    assert x_init.min() >= 0 and x_init.max() < K, \
        f"x_init out of [0,K): [{x_init.min()}, {x_init.max()}]"

    # input_ids shape & token value ranges
    assert ids_init.shape == (B, L_plus_N1), f"ids_init.shape={ids_init.shape}"
    txt_part = ids_init[:, :txt_len]
    vis_part = ids_init[:, -N:]
    assert (txt_part < lm_vocab_size).all(), \
        f"text tokens bleed into visual range (max={txt_part.max()})"
    assert (vis_part >= lm_vocab_size).all(), \
        f"visual tokens not shifted (min={vis_part.min()}, lm_vocab_size={lm_vocab_size})"
    assert (vis_part < lm_vocab_size + K).all(), \
        f"visual tokens exceed lm_vocab_size+K (max={vis_part.max()})"

    # rope_pos
    assert rpos_init.shape == (B, L_plus_N1, 3), \
        f"rope_pos shape={rpos_init.shape} expected ({B},{L_plus_N1},3)"

    # z_s
    assert z_s.shape == (B, N, K), f"z_s.shape={z_s.shape}"
    p_err = (p_s.sum(-1) - 1).abs().max().item()
    assert p_err < 1e-3, f"p_s not normalised: max deviation={p_err:.2e}"

    # logp
    assert not torch.isnan(logp).any(), "logp contains NaN"
    assert not torch.isinf(logp).any(), "logp contains Inf"

    # x_t
    assert x_t.min() >= 0 and x_t.max() < K, \
        f"x_t out of [0,K) after add_noise: [{x_t.min()}, {x_t.max()}]"

    # Dual-path shape check
    assert z_T_cond.shape == z_S_cond.shape, \
        f"Dual-path mismatch: z_T_cond={z_T_cond.shape} z_S_cond={z_S_cond.shape}"
    assert z_T_cond.shape == (B, N, K), f"z_T_cond.shape={z_T_cond.shape}"

    # z_T logits printout (always in dry_run; also when uncond is available)
    if dry_run or z_T_uncond is not None:
        print(
            f"[assert] z_T_cond   shape={z_T_cond.shape}  "
            f"min={z_T_cond.min():.3f}  max={z_T_cond.max():.3f}  "
            f"mean={z_T_cond.mean():.3f}"
        )
    if z_T_uncond is not None:
        assert z_T_uncond.shape == (B, N, K), f"z_T_uncond.shape={z_T_uncond.shape}"
        print(
            f"[assert] z_T_uncond shape={z_T_uncond.shape}  "
            f"min={z_T_uncond.min():.3f}  max={z_T_uncond.max():.3f}  "
            f"mean={z_T_uncond.mean():.3f}"
        )
    if z_T_guided is not None:
        assert z_T_guided.shape == (B, N, K), f"z_T_guided.shape={z_T_guided.shape}"
        g_min = z_T_guided.min().item()
        g_max = z_T_guided.max().item()
        g_mean = z_T_guided.mean().item()
        print(
            f"[assert] z_T_guided shape={z_T_guided.shape}  "
            f"min={g_min:.3f}  max={g_max:.3f}  mean={g_mean:.3f}"
        )
        # Explosion guard: guided logits must be finite and not excessively large.
        assert not torch.isnan(z_T_guided).any(), "z_T_guided contains NaN"
        assert not torch.isinf(z_T_guided).any(), "z_T_guided contains Inf"
        assert abs(g_min) < 1e4 and abs(g_max) < 1e4, (
            f"z_T_guided magnitude too large: min={g_min:.1e}  max={g_max:.1e}. "
            f"Reduce --teacher_cfg_scale (currently may amplify outlier logits)."
        )

    print("[assert] All assertions PASSED ✓")


def _sanity_check_forward(teacher, scheduler, latents_shape, device, K, verbose=False):
    print("[init] Checking logit dimensions …")
    T, H, W = latents_shape
    N, B, L = T * H * W, 1, 16
    dummy_txt = torch.zeros(B, L, dtype=torch.long, device=device)
    dummy_vis = torch.zeros(B, T, H, W, dtype=torch.long, device=device)
    with torch.no_grad():
        ids, rpos, _ = build_ursa_inputs(teacher, dummy_txt, dummy_vis, latents_shape, device)
        logits = teacher(ids, rope_pos=rpos).sample
    lm_head_size = teacher.config.lm_head_size
    lm_vocab     = teacher.config.lm_vocab_size
    print(
        f"[init] logits={logits.shape}  K={K}  "
        f"lm_head={lm_head_size}  lm_vocab={lm_vocab}"
    )
    assert ids.shape    == (B, L + N + 1),         f"ids shape {ids.shape}"
    assert rpos.shape   == (B, L + N + 1, 3),      f"rpos shape {rpos.shape}"
    z = extract_visual_logits(logits.float(), N, K)
    assert z.shape      == (B, N, K),               f"z shape {z.shape}"
    assert lm_head_size >= K,                       f"lm_head_size={lm_head_size} < K={K}"
    if verbose:
        print("[init] flex_attn state during sanity check:")
        _print_flex_attn_state(teacher, "teacher")
    print("[init] Forward check OK ✓")


# ---------------------------------------------------------------------------
# Dry-run patches 7 / 8 / 9
# ---------------------------------------------------------------------------

def _dry_run_patches_789(teacher, latents_shape, K, N, device):
    """Three deep self-checks executed only during --dry_run.



    Patch 7 — extract_visual_logits end-to-end alignment:

        Run a real teacher forward, manually reconstruct z_manual from raw logits

        using the latent_shift / codebook_size convention, and assert the result

        matches extract_visual_logits().  Handles the common URSA case where

        lm_head outputs K logits directly (latent_shift not applied to logit dim).



    Patch 8 — flex_attn semantics sanity:

        If the model exposes set_offsets_by_lens, compare visual-logit mean-delta

        between offsets=None (standard causal) and a single-block offset.  A large

        delta is expected and confirms that our training correctly uses offsets=None.

        Gracefully skips when flex_attention is unavailable at runtime.



    Patch 9 — logp / token reshape consistency:

        With a small (T=3, H=4, W=5) shape, verify x_hat reshape round-trips and

        spot-check 10 token positions against manually computed log-probability.

    """
    T, H, W = latents_shape
    L_test, B_test = 16, 1

    print("\n" + "=" * 64)
    print("[patch 7/8/9] Running additional dry_run self-checks …")

    # -------------------------------------------------------------------------
    # Build shared dummy inputs used by both patch 7 and patch 8
    # -------------------------------------------------------------------------
    dummy_txt = torch.zeros(B_test, L_test, dtype=torch.long, device=device)
    dummy_vis = torch.zeros(B_test, T, H, W, dtype=torch.long, device=device)
    with torch.no_grad():
        ids_test, rpos_test, _ = build_ursa_inputs(
            teacher, dummy_txt, dummy_vis, latents_shape, device)
        logits_full = teacher(ids_test, rope_pos=rpos_test).sample.float()  # [1, L+N+1, D]

    D            = logits_full.size(-1)          # actual logit last-dim (lm_head_size)
    latent_shift = teacher.config.lm_vocab_size  # text-vocab offset for input token IDs

    # =========================================================================
    # Patch 7 — extract_visual_logits end-to-end alignment
    # =========================================================================
    print("\n[7] extract_visual_logits end-to-end alignment …")
    z_vis = extract_visual_logits(logits_full, N, K)   # [1, N, K]
    assert z_vis.shape == (B_test, N, K), f"z_vis.shape={z_vis.shape}"

    if D >= latent_shift + K:
        # Full-vocab head: logit dim covers text (0..latent_shift) + visual tokens.
        z_seq    = logits_full[:, -(N + 1) : -1]            # [1, N, D]
        z_manual = z_seq[..., latent_shift : latent_shift + K]  # [1, N, K]
        delta    = (z_vis - z_manual).abs().max().item()
        print(f"  [7] path=full-vocab  D={D}  latent_shift+K={latent_shift + K}")
        print(f"  [7] z_vis.shape={z_vis.shape}  max|z_vis - z_manual|={delta:.2e}")
        assert delta < 1e-5, (
            f"extract_visual_logits mismatch (full-vocab path): delta={delta:.2e}. "
            "The function should return logits[..., latent_shift:latent_shift+K]."
        )
        print("[7] extract_visual_logits alignment PASSED ✓")

    else:
        # Common URSA case: lm_head outputs K logits directly (lm_head_size ≈ K).
        # latent_shift is the input token-ID offset, NOT a logit-dimension offset.
        # extract_visual_logits handles this as D==K (happy path) or D>K (offset=D-K).
        z_seq = logits_full[:, -(N + 1) : -1]   # [1, N, D]
        if D == K:
            delta = (z_vis - z_seq).abs().max().item()
            print(
                f"  [7] SKIP latent_shift formula: D={D} == K={K}  "
                f"latent_shift={latent_shift}.\n"
                f"  [7] Explanation: URSA lm_head outputs K visual logits directly.\n"
                f"  [7]   latent_shift={latent_shift} is the input token-ID shift "
                f"(raw_code + lm_vocab_size), NOT a logit-dim offset.\n"
                f"  [7]   extract_visual_logits happy-path: z = logits[:, -(N+1):-1] "
                f"(no vocab-dim slicing).\n"
                f"  [7]   Fallback check: z_vis == raw causal slice  "
                f"max_delta={delta:.2e}"
            )
            assert delta < 1e-5, (
                f"z_vis != raw causal slice when D==K: delta={delta:.2e}"
            )
        else:
            # D > K but D < latent_shift + K  →  extract uses offset = D - K
            offset   = D - K
            z_manual = z_seq[..., offset:]
            delta    = (z_vis - z_manual).abs().max().item()
            print(
                f"  [7] SKIP latent_shift formula: D={D} < latent_shift+K={latent_shift + K}.\n"
                f"  [7] extract_visual_logits uses offset={offset} (D-K). "
                f"max_delta={delta:.2e}"
            )
            assert delta < 1e-5, (
                f"z_vis != z_seq[..., D-K:]: delta={delta:.2e}"
            )
        print("[7] extract_visual_logits alignment PASSED (fallback path) ✓")

    # =========================================================================
    # Patch 8 — flex_attn semantics sanity
    # =========================================================================
    print("\n[8] flex_attn semantics sanity …")
    fa = _probe_flex_attn(teacher)
    if fa is None or not hasattr(fa, "set_offsets_by_lens"):
        print("  [8] flex_attn.set_offsets_by_lens not available — skip")
        print("[8] flex_attn semantics sanity PASSED (skipped — no flex_attn) ✓")
    else:
        L_total   = ids_test.size(1)          # L_test + N + 1
        txt_block = L_test + (N + 1)          # single-block: all tokens in one block
        block_lens = [txt_block]

        try:
            # Forward A: offsets=None — standard causal attention (our training config)
            _reset_flex_attn(teacher, "teacher")
            with torch.no_grad():
                logits_A = teacher(ids_test, rope_pos=rpos_test).sample.float()
            z_A = extract_visual_logits(logits_A, N, K)

            # Forward B: set_offsets_by_lens with a single block.
            # A single block causes the mask to allow full (bidirectional) attention
            # within the block, which differs from standard causal attention.
            fa.set_offsets_by_lens(block_lens)
            with torch.no_grad():
                logits_B = teacher(ids_test, rope_pos=rpos_test).sample.float()
            z_B = extract_visual_logits(logits_B, N, K)

            delta_mean = (z_A - z_B).abs().mean().item()
            delta_max  = (z_A - z_B).abs().max().item()
            print(
                f"  [8] offsets=None  vs  set_offsets_by_lens({block_lens}):\n"
                f"  [8]   mean_abs_delta={delta_mean:.4e}  max_abs_delta={delta_max:.4e}"
            )
            if delta_mean > 1e-3:
                print(
                    f"  [8] WARNING: mean_delta={delta_mean:.2e} > 1e-3.\n"
                    "  [8]   Single-block flex_attn uses FULL (bidirectional) attention\n"
                    "  [8]   inside the block, whereas offsets=None gives standard CAUSAL\n"
                    "  [8]   attention. This difference is EXPECTED — it confirms our\n"
                    "  [8]   training correctly uses offsets=None (no packed sequences)."
                )
            else:
                print(f"  [8] delta ≤ 1e-3: attention semantics equivalent for this input.")
            print("[8] flex_attn semantics sanity PASSED ✓")

        except (NotImplementedError, RuntimeError, Exception) as exc:
            print(f"  [8] flex_attn runtime not available ({type(exc).__name__}: {exc}) — skip")
            print("[8] flex_attn semantics sanity PASSED (runtime skip) ✓")
        finally:
            _reset_flex_attn(teacher, "teacher")   # always restore clean state

    # =========================================================================
    # Patch 9 — logp / token reshape consistency
    # =========================================================================
    print("\n[9] logp/token reshape consistency …")
    T9, H9, W9 = 3, 4, 5
    N9, B9     = T9 * H9 * W9, 1      # 60 tokens, batch=1

    torch.manual_seed(99)
    z9 = torch.randn(B9, N9, K)
    p9 = F.softmax(z9 / 1.0, dim=-1)  # [1, 60, K]; each row sums to 1

    # ----- token sampling ---------------------------------------------------
    x_hat_flat = torch.multinomial(p9.view(-1, K), 1)      # [N9, 1]  (1 sample per row)
    x_hat_1d   = x_hat_flat.view(B9, N9)                   # [1, 60]
    x_hat_4d   = x_hat_1d.view(B9, T9, H9, W9)            # [1, 3, 4, 5]

    # reshape round-trip: 1d → 4d → 1d must be lossless
    x_hat_back = x_hat_4d.view(B9, N9)
    assert torch.equal(x_hat_1d, x_hat_back), (
        f"reshape round-trip FAILED: x_hat_1d != x_hat_4d.view(B,N)\n"
        f"  x_hat_1d.shape={x_hat_1d.shape}  x_hat_back.shape={x_hat_back.shape}"
    )

    # ----- logp computation (mirrors training code) -------------------------
    # logp_all[b, n] = log p9[b, n, x_hat_1d[b, n]]
    logp_all = (
        p9.clamp(1e-8).log()
        .gather(-1, x_hat_1d.unsqueeze(-1))
        .squeeze(-1)
    )  # [B9, N9]
    logp_sum = logp_all.sum(-1)  # [B9]

    # ----- spot-check 10 random token positions -----------------------------
    torch.manual_seed(7)
    positions = torch.randperm(N9)[:10].tolist()
    for pos in positions:
        tok_id    = x_hat_1d[0, pos].item()
        logp_man  = math.log(max(p9[0, pos, tok_id].item(), 1e-8))
        logp_gat  = logp_all[0, pos].item()
        diff      = abs(logp_man - logp_gat)
        assert diff < 1e-6, (
            f"logp mismatch at pos={pos}, tok={tok_id}: "
            f"manual={logp_man:.8f}  gathered={logp_gat:.8f}  diff={diff:.2e}"
        )

    # check logp_sum matches sum of logp_all
    logp_sum_manual = logp_all[0].sum().item()
    assert abs(logp_sum.item() - logp_sum_manual) < 1e-5, \
        f"logp_sum mismatch: {logp_sum.item():.6f} vs {logp_sum_manual:.6f}"

    print(
        f"  [9] T={T9},H={H9},W={W9}  N={N9}  K={K}  "
        f"x_hat reshape round-trip ✓  "
        f"10 logp spot-checks (pos={positions}) ✓  "
        f"logp_sum={logp_sum.item():.3f}"
    )
    print("[9] logp/token reshape consistency PASSED ✓")

    print("\n" + "=" * 64)
    print("[patch 7/8/9] All 3 additional dry_run checks PASSED ✓")
    print("=" * 64)


if __name__ == "__main__":
    main()