_test_refresh.py

"""Ad-hoc correctness check for the 'refresh' shadow/parallel-lookahead branch.

Not part of the shipped package; run manually with:
    python _test_refresh.py

This covers the decoupled real (flash_attn/sdpa/eager, unmodified) + shadow (FlexAttention,
cross-attends into the real branch's per-layer K/V) implementation: the real branch is a completely
standard causal forward, and the shadow branch is a *compacted* sequence (cycle heads are skipped
entirely, since their output is mathematically identical to the real branch's own thinking output).
"""
import sys
import types
import pathlib

import torch

HERE = pathlib.Path(__file__).resolve().parent

# `Causal-Parallel-Refresh` has a hyphen, so it cannot be imported as a normal package.
# Register a synthetic package pointing at this directory so the modules' relative imports work.
pkg = types.ModuleType("dmtd_pkg")
pkg.__path__ = [str(HERE)]
sys.modules["dmtd_pkg"] = pkg

from dmtd_pkg.configuration_dmtdqwen3 import DMTDQwen3Config  # noqa: E402
from dmtd_pkg.modeling_dmtdqwen3 import DMTDQwen3ForCausalLM  # noqa: E402

torch.manual_seed(0)

# FlexAttention needs a CUDA device with a Triton backend.
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
if DEVICE != "cuda":
    raise SystemExit("[FAIL] this test requires CUDA (FlexAttention has no usable CPU backend here).")

MTP_HORIZON = 4
NUM_ENC, NUM_THINK, NUM_DEC = 0, 4, 2
SEQ_LEN = 12  # 3 full cycles
VOCAB = 64
MASK_TOKEN_ID = VOCAB - 1

config = DMTDQwen3Config(
    vocab_size=VOCAB,
    hidden_size=32,
    intermediate_size=64,
    num_hidden_layers=NUM_ENC + NUM_THINK + NUM_DEC,
    num_attention_heads=4,
    num_key_value_heads=2,
    # FlexAttention's Triton kernel requires head_dim >= 16.
    head_dim=16,
    max_position_embeddings=64,
    num_encoding_layers=NUM_ENC,
    num_thinking_layers=NUM_THINK,
    num_decoding_layers=NUM_DEC,
    mtp_horizon=MTP_HORIZON,
    mask_token_id=MASK_TOKEN_ID,
    attn_implementation="sdpa",
    tie_word_embeddings=True,
)

model = DMTDQwen3ForCausalLM(config).to(DEVICE)
model.eval()

input_ids = torch.randint(0, VOCAB - 1, (1, SEQ_LEN), device=DEVICE)  # avoid mask token id among "real" tokens

# --- Test 1: plain forward runs and produces the expected shape. ---
with torch.no_grad():
    out = model(input_ids=input_ids)
assert out.logits.shape == (1, SEQ_LEN, VOCAB), out.logits.shape
print("[ok] forward shape:", out.logits.shape)

# --- Test 2: the real branch is now a completely standard forward pass, so ANY attention
#     implementation the model is configured with should work (not just eager/sdpa) -- this is the
#     whole point of decoupling the shadow branch's custom mask into FlexAttention. We can't assume
#     `flash_attn` is actually installed in this dev environment, so exercise `eager` and `sdpa`
#     (both previously supported) plus confirm nothing in the forward pass hard-codes either. ---
for impl in ("eager", "sdpa"):
    model.config._attn_implementation = impl
    with torch.no_grad():
        out_impl = model(input_ids=input_ids)
    assert out_impl.logits.shape == (1, SEQ_LEN, VOCAB)
    print(f"[ok] forward runs with attn_implementation={impl!r}:", out_impl.logits.shape)
model.config._attn_implementation = "sdpa"

try:
    import flash_attn  # noqa: F401

    model.config._attn_implementation = "flash_attention_2"
    with torch.no_grad():
        out_fa2 = model(input_ids=input_ids)
    assert out_fa2.logits.shape == (1, SEQ_LEN, VOCAB)
    print("[ok] forward runs with attn_implementation='flash_attention_2':", out_fa2.logits.shape)
    model.config._attn_implementation = "sdpa"
except ImportError:
    print("[skip] flash_attn package not installed in this environment; skipping flash_attention_2 check")

# --- Test 3: cycle-head positions never run a shadow computation at all -- the combination at those
#     positions is, by construction, `encoding_hidden_states + real_thinking_output`. Confirm this by
#     comparing the model's actual output against a manual vanilla (real-branch-only, no shadow at
#     all) forward through the encoding+thinking layers. ---
m = model.model
with torch.no_grad():
    inputs_embeds = m.embed_tokens(input_ids)
    position_ids = torch.arange(SEQ_LEN, device=DEVICE).unsqueeze(0)
    seq_len = SEQ_LEN
    thk_end = NUM_ENC + NUM_THINK

    from transformers.masking_utils import create_causal_mask

    ref_mask = create_causal_mask(
        config=m.config, inputs_embeds=inputs_embeds, attention_mask=None,
        past_key_values=None, position_ids=position_ids,
    )
    position_embeddings = m.rotary_emb(inputs_embeds, position_ids)

    vanilla_hidden_states = inputs_embeds
    encoding_snapshot = vanilla_hidden_states
    for i, layer in enumerate(m.layers[:thk_end]):
        vanilla_hidden_states = layer(
            vanilla_hidden_states,
            attention_mask=ref_mask,
            position_embeddings=position_embeddings,
            position_ids=position_ids,
            past_key_values=None,
            use_cache=False,
        )
        if i + 1 == NUM_ENC:
            encoding_snapshot = vanilla_hidden_states
    vanilla_combined = encoding_snapshot + vanilla_hidden_states

    # Run the model itself and grab its pre-decoding-layer hidden state via a forward hook on the
    # first decoding layer (its input IS `encoding_hidden_states + parallel_hidden_states`).
    captured = {}

    def _grab_input(_module, args, _kwargs):
        captured["hidden_states"] = args[0]

    hook = m.layers[thk_end].register_forward_pre_hook(_grab_input, with_kwargs=True)
    try:
        model(input_ids=input_ids)
    finally:
        hook.remove()

    model_combined = captured["hidden_states"]

    is_cycle_head = (torch.arange(seq_len, device=DEVICE) % MTP_HORIZON) == 0
    head_positions = is_cycle_head.nonzero(as_tuple=True)[0]
    nonhead_positions = (~is_cycle_head).nonzero(as_tuple=True)[0]

    head_diff = (model_combined[:, head_positions, :] - vanilla_combined[:, head_positions, :]).abs().max().item()
    print(f"[info] max |model - vanilla| at cycle-head positions: {head_diff:.3e}")
    assert head_diff < 1e-4, "cycle-head positions must exactly match the vanilla (no-shadow) computation"
    print("[ok] cycle-head positions never run a shadow computation and match vanilla exactly")

# --- Test 4: refresh confirmed -- a later cycle's shadow branch must depend on an EARLIER cycle's
#     non-head (masked) position's REAL token (via the refreshed real-branch history), even though
#     that position's own shadow/lookahead output never sees it (it's always MASK). ---
with torch.no_grad():
    input_ids_2 = input_ids.clone()
    input_ids_2[0, 1] = (input_ids_2[0, 1] + 1) % (VOCAB - 1)

    def run_and_capture_parallel(ids):
        captured_inner = {}

        def _grab(_module, args, _kwargs):
            captured_inner["hidden_states"] = args[0]

        h = m.layers[thk_end].register_forward_pre_hook(_grab, with_kwargs=True)
        try:
            model(input_ids=ids)
        finally:
            h.remove()
        return captured_inner["hidden_states"]

    combined_orig = run_and_capture_parallel(input_ids)
    combined_mut = run_and_capture_parallel(input_ids_2)

    # `combined` is `encoding_hidden_states + parallel_hidden_states`. With NUM_ENC == 0,
    # `encoding_hidden_states` is exactly `inputs_embeds`, so subtracting it isolates
    # `parallel_hidden_states` (the shadow branch's own contribution at non-head positions), which is
    # what should be invariant to x1's own mutation at its own (non-head) position.
    parallel_orig = combined_orig - m.embed_tokens(input_ids)
    parallel_mut = combined_mut - m.embed_tokens(input_ids_2)

    # Position 1 is a non-head position: its own shadow input is always MASK regardless of its real
    # token, so its own parallel-branch output must be unchanged by the mutation.
    diff_pos1 = (parallel_orig[:, 1, :] - parallel_mut[:, 1, :]).abs().max().item()
    # Positions 4.. (later cycles) must change: they read x1's REAL value via the refreshed history.
    diff_later = (parallel_orig[:, 4:, :] - parallel_mut[:, 4:, :]).abs().max().item()
    print(f"[info] parallel-branch diff at mutated position itself (1): {diff_pos1:.3e}")
    print(f"[info] parallel-branch diff at later cycles (4..11) after mutating x1: {diff_later:.3e}")
    assert diff_pos1 < 1e-6, "position 1's shadow input is always MASK, independent of its own real token"
    assert diff_later > 1e-4, "later cycles' shadow branch should see x1's REAL value via the refreshed history"
    print("[ok] refresh confirmed: later cycles read x1's REAL value (not its masked lookahead) as history")

# --- Test 5: packed sequences (multiple documents concatenated, position_ids reset at each document
#     boundary) must not leak attention across the document boundary, in either the real or the shadow
#     stream, with the actual training-time call pattern (`use_cache=False`). ---
with torch.no_grad():
    packed_ids = torch.randint(0, VOCAB - 1, (1, 12), device=DEVICE)
    packed_position_ids = torch.cat([torch.arange(6), torch.arange(6)]).unsqueeze(0).to(DEVICE)
    doc2 = packed_ids[:, 6:]

    out_full = model.model(input_ids=packed_ids, position_ids=packed_position_ids, use_cache=False)
    out_doc2 = model.model(input_ids=doc2, use_cache=False)
    diff = (out_full.last_hidden_state[:, 6:, :] - out_doc2.last_hidden_state).abs().max().item()
    print(f"[info] packed vs standalone doc2 hidden-state diff (use_cache=False): {diff:.3e}")
    assert diff < 1e-4, "cross-document attention leakage detected in packed training mode!"
    print("[ok] no cross-document leakage in packed sequences (use_cache=False)")

# --- Test 5c: the whole point of bucketing -- for a FIXED `seq_len`, the shadow branch's Q_LEN/KV_LEN
#     fed into `flex_attention`/`create_block_mask` must be IDENTICAL no matter how many documents are
#     packed into it (i.e. no matter how many "extra" cycle-heads packing introduces). Verify this
#     directly (not just "it doesn't crash / isn't slow") by hooking `_compile_friendly_flex_attention`
#     and recording every Q_LEN/KV_LEN it's actually called with across several very different packing
#     patterns of the same total length. ---
import dmtd_pkg.modeling_dmtdqwen3 as dmtd_modeling  # noqa: E402

seen_shapes = []
_orig_flex_attn = dmtd_modeling._compile_friendly_flex_attention


def _spy_flex_attention(query, key, value, **kw):
    seen_shapes.append((query.shape[-2], key.shape[-2]))
    return _orig_flex_attn(query, key, value, **kw)


dmtd_modeling._compile_friendly_flex_attention = _spy_flex_attention
try:
    with torch.no_grad():
        fixed_len = 16
        patterns = [
            torch.arange(fixed_len),  # one document, no resets (the theoretical max shadow case)
            torch.cat([torch.arange(6), torch.arange(fixed_len - 6)]),  # 2 documents
            torch.cat([torch.arange(3), torch.arange(3), torch.arange(fixed_len - 6)]),  # 3 documents
            torch.cat([torch.arange(2)] * 8),  # 8 tiny documents -> many extra cycle-heads
        ]
        for pattern in patterns:
            seen_shapes.clear()
            ids = torch.randint(0, VOCAB - 1, (1, fixed_len), device=DEVICE)
            model.model(input_ids=ids, position_ids=pattern.unsqueeze(0).to(DEVICE), use_cache=False)
            assert seen_shapes, "expected the shadow branch's flex_attention to actually be called"
            assert len(set(seen_shapes)) == 1, (
                f"Q_LEN/KV_LEN must be identical across every thinking layer within one forward, got {seen_shapes}"
            )
        # And the shape must ALSO be identical ACROSS the different packing patterns above (not just
        # within one forward) -- this is the actual bucketing guarantee.
        all_first_shapes = []
        for pattern in patterns:
            seen_shapes.clear()
            ids = torch.randint(0, VOCAB - 1, (1, fixed_len), device=DEVICE)
            model.model(input_ids=ids, position_ids=pattern.unsqueeze(0).to(DEVICE), use_cache=False)
            all_first_shapes.append(seen_shapes[0])
        assert len(set(all_first_shapes)) == 1, (
            f"Q_LEN/KV_LEN must be bucketed to a fixed value across different packing patterns of the "
            f"same seq_len, got {all_first_shapes}"
        )
        print(f"[info] shadow Q_LEN/KV_LEN across {len(patterns)} different packing patterns: {all_first_shapes[0]}")
        print("[ok] shadow branch shape is fixed/bucketed regardless of the packing pattern")
finally:
    dmtd_modeling._compile_friendly_flex_attention = _orig_flex_attn

# --- Test 5b: batch size > 1 must work too (both the shadow `mask_mod`'s `packed_sequence_mask[b, ...]`
#     lookups AND the model's own `position_ids` broadcasting need the real batch size, not the
#     batch-size-1 default `position_ids` shape that `DMTDQwen3Model.forward` falls back to). Also
#     exercise heterogeneous per-row packing (each row packs different document boundaries), which is
#     exactly what ms-swift's `packing` produces in practice. ---
with torch.no_grad():
    bsz2_ids = torch.randint(0, VOCAB - 1, (2, SEQ_LEN), device=DEVICE)
    out_bsz2 = model(input_ids=bsz2_ids)
    assert out_bsz2.logits.shape == (2, SEQ_LEN, VOCAB)
    print("[ok] forward works with batch_size=2 (default, batch-size-1 position_ids):", out_bsz2.logits.shape)

    hetero_position_ids = torch.stack(
        [
            torch.cat([torch.arange(5, device=DEVICE), torch.arange(SEQ_LEN - 5, device=DEVICE)]),
            torch.arange(SEQ_LEN, device=DEVICE),
        ]
    )
    out_hetero = model(input_ids=bsz2_ids, position_ids=hetero_position_ids, use_cache=False)
    assert out_hetero.logits.shape == (2, SEQ_LEN, VOCAB)
    print("[ok] forward works with batch_size=2 and heterogeneous per-row packing boundaries")

model.train()
out_bsz2_bwd = model(input_ids=bsz2_ids, labels=bsz2_ids, use_cache=False)
out_bsz2_bwd.loss.backward()
model.eval()
print("[ok] backward works with batch_size=2, loss:", out_bsz2_bwd.loss.item())

# --- Test 6: gradient checkpointing must work end-to-end, AND must actually be engaged on every
#     `DMTDQwen3DecoderLayer` (real branch) as well as wrapped manually around `_forward_shadow_layer`
#     (shadow branch) -- setting `model.model.gradient_checkpointing` directly (bypassing
#     `gradient_checkpointing_enable()`) does NOT propagate to the per-layer flag each
#     `GradientCheckpointingLayer` actually checks, so it would silently look like a no-op bug here. ---
model.train()
model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
assert model.model.gradient_checkpointing is True
assert all(layer.gradient_checkpointing for layer in model.model.layers[:thk_end]), (
    "gradient_checkpointing_enable() must propagate to every encoding/thinking decoder layer"
)
labels = input_ids.clone()
out = model(input_ids=input_ids, labels=labels, use_cache=False)
out.loss.backward()
model.gradient_checkpointing_disable()
model.eval()
print("[ok] gradient checkpointing forward+backward works (properly engaged), loss:", out.loss.item())

# --- Test 7: gradient checkpointing must materially reduce peak activation memory for the shadow
#     branch too, not just the real branch. This directly targets the bug where `_forward_shadow_layer`
#     was called as a bare method (bypassing `GradientCheckpointingLayer.__call__`), so its activations
#     across ALL thinking layers were kept alive for backward regardless of the checkpointing setting. ---
BIG_NUM_THINK = 8
big_config = DMTDQwen3Config(
    vocab_size=VOCAB,
    hidden_size=256,
    intermediate_size=1024,
    num_hidden_layers=BIG_NUM_THINK + NUM_DEC,
    num_attention_heads=8,
    num_key_value_heads=2,
    head_dim=32,
    max_position_embeddings=256,
    num_encoding_layers=0,
    num_thinking_layers=BIG_NUM_THINK,
    num_decoding_layers=NUM_DEC,
    mtp_horizon=MTP_HORIZON,
    mask_token_id=MASK_TOKEN_ID,
    attn_implementation="sdpa",
    tie_word_embeddings=True,
)
BIG_SEQ_LEN = 256
big_input_ids = torch.randint(0, VOCAB - 1, (2, BIG_SEQ_LEN), device=DEVICE)
big_labels = big_input_ids.clone()


def _peak_mem_for(gradient_checkpointing: bool) -> float:
    torch.manual_seed(0)
    big_model = DMTDQwen3ForCausalLM(big_config).to(DEVICE)
    big_model.train()
    if gradient_checkpointing:
        big_model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
    torch.cuda.synchronize()
    torch.cuda.reset_peak_memory_stats()
    out = big_model(input_ids=big_input_ids, labels=big_labels, use_cache=False)
    out.loss.backward()
    torch.cuda.synchronize()
    peak = torch.cuda.max_memory_allocated() / (1024**2)
    del big_model, out
    torch.cuda.empty_cache()
    return peak


mem_without_ckpt = _peak_mem_for(gradient_checkpointing=False)
mem_with_ckpt = _peak_mem_for(gradient_checkpointing=True)
print(f"[info] peak activation memory WITHOUT gradient checkpointing: {mem_without_ckpt:.1f} MiB")
print(f"[info] peak activation memory WITH    gradient checkpointing: {mem_with_ckpt:.1f} MiB")
assert mem_with_ckpt < mem_without_ckpt * 0.8, (
    "gradient checkpointing should meaningfully reduce peak memory; if it doesn't, the shadow "
    "branch's per-layer activations are likely being retained for backward again"
)
print("[ok] gradient checkpointing meaningfully reduces peak memory (shadow branch included)")

print("\nAll refresh tests passed.")