"""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.")