| from __future__ import annotations
|
|
|
| import torch
|
|
|
| from conveyer.config import ConveyerConfig, ModelConfig, load_config
|
| from conveyer.data.synthetic import SyntheticLookaheadDataset
|
| from conveyer.infer import generate_conveyer, stream_conveyer_generation
|
| from conveyer.losses import build_horizon_weights, compute_unrolled_conveyer_loss, seed_corrupted_belt
|
| from conveyer.model import BeltState, ConveyerDiffusionLM
|
|
|
|
|
| def _small_model_cfg(
|
| belt_size: int = 64,
|
| use_thought_highway: bool = True,
|
| causal_belt_only: bool = False,
|
| ) -> ModelConfig:
|
| return ModelConfig(
|
| vocab_size=64,
|
| d_model=64,
|
| n_heads=4,
|
| n_layers=2,
|
| d_ff=128,
|
| max_seq_len=256,
|
| dropout=0.0,
|
| belt_size=belt_size,
|
| tau_min=0.5,
|
| tau_max=4.0,
|
| tau_gamma=1.5,
|
| gate_bias_init=1.0,
|
| use_thought_highway=use_thought_highway,
|
| causal_belt_only=causal_belt_only,
|
| )
|
|
|
|
|
| def test_prefix_kv_cache_chunked_equivalence() -> None:
|
| """
|
| Verifies that encoding a prefix in one shot (`prefill_prefix(x_{0:P})`) produces the exact same
|
| KV cache and downstream belt logits as incremental token-by-token `append_committed_token`.
|
| """
|
| torch.manual_seed(123)
|
| cfg = _small_model_cfg(belt_size=16)
|
| model = ConveyerDiffusionLM(cfg).eval()
|
|
|
| prefix = torch.randint(1, 60, (2, 10), dtype=torch.long)
|
| belt = model.init_empty_belt(batch_size=2, tick=10, device=prefix.device, deterministic=True)
|
|
|
|
|
| kv_full = model.prefill_prefix(prefix)
|
| out_full = model.forward_belt(belt, kv_cache=kv_full)
|
|
|
|
|
| kv_inc = model.prefill_prefix(prefix[:, :4])
|
| for t in range(4, 10):
|
| kv_inc = model.append_committed_token(prefix[:, t], kv_inc)
|
| out_inc = model.forward_belt(belt, kv_cache=kv_inc)
|
|
|
| assert kv_full.seq_len == kv_inc.seq_len == 10
|
| for l_idx in range(cfg.n_layers):
|
| assert torch.allclose(kv_full.k_cache[l_idx], kv_inc.k_cache[l_idx], atol=1e-5)
|
| assert torch.allclose(kv_full.v_cache[l_idx], kv_inc.v_cache[l_idx], atol=1e-5)
|
| assert torch.allclose(out_full.logits, out_inc.logits, atol=1e-5)
|
|
|
|
|
| def test_slot0_lexical_anchoring_alpha_zero() -> None:
|
| """
|
| Verifies that alpha_0 == 0.0 at Position 0:
|
| Perturbing `m_thought[:, 0, :]` at Slot 0 must have ZERO effect on the composed dough input at Slot 0,
|
| forcing Slot 0 to rely 100% on its crisp lexical projection `e_soft[:, 0, :]` before popping out.
|
| Meanwhile, perturbing `m_thought[:, k, :]` for k > 0 MUST change `dough[:, k, :]`.
|
| """
|
| torch.manual_seed(42)
|
| cfg = _small_model_cfg(belt_size=64, use_thought_highway=True)
|
| model = ConveyerDiffusionLM(cfg).eval()
|
|
|
| belt1 = model.init_empty_belt(batch_size=2, tick=5, device=torch.device("cpu"), deterministic=True)
|
| m_mod = belt1.m_thought.clone()
|
| m_mod[:, 0, :] += 100.0
|
| belt2 = BeltState(e_soft=belt1.e_soft.clone(), m_thought=m_mod, tick=5)
|
|
|
| dough1 = model.compose_dough_inputs(belt1)
|
| dough2 = model.compose_dough_inputs(belt2)
|
|
|
| assert float(model.alpha_schedule[0].item()) == 0.0
|
| assert float(model.alpha_schedule[-1].item()) == 1.0
|
| assert torch.allclose(dough1[:, 0, :], dough2[:, 0, :], atol=1e-6)
|
|
|
|
|
| m_mod63 = belt1.m_thought.clone()
|
| m_mod63[:, 63, :] += 10.0
|
| belt3 = BeltState(e_soft=belt1.e_soft.clone(), m_thought=m_mod63, tick=5)
|
| dough3 = model.compose_dough_inputs(belt3)
|
| assert not torch.allclose(dough1[:, 63, :], dough3[:, 63, :], atol=1e-4)
|
|
|
|
|
| def test_bidirectional_belt_attention_vs_causal_ar() -> None:
|
| """
|
| Verifies that in Conveyor mode (`causal_belt_only=False`), Slot 0's output logits actively depend on
|
| downstream dough balls at the back of the belt (e.g., Slot 63), whereas in AR baseline mode
|
| (`causal_belt_only=True`), Slot 0 is completely blind to Slots 1..63.
|
| """
|
| torch.manual_seed(7)
|
|
|
| cfg_bi = _small_model_cfg(belt_size=64, causal_belt_only=False)
|
| model_bi = ConveyerDiffusionLM(cfg_bi).eval()
|
|
|
| belt_a = model_bi.init_empty_belt(batch_size=1, tick=4, device=torch.device("cpu"), deterministic=True)
|
| belt_b = BeltState(e_soft=belt_a.e_soft.clone(), m_thought=belt_a.m_thought.clone(), tick=4)
|
| belt_b.e_soft[:, 63, :] += 5.0
|
|
|
| out_bi_a = model_bi.forward_belt(belt_a)
|
| out_bi_b = model_bi.forward_belt(belt_b)
|
|
|
| assert not torch.allclose(out_bi_a.slot0_logits, out_bi_b.slot0_logits, atol=1e-5)
|
|
|
|
|
| cfg_ar = _small_model_cfg(belt_size=64, causal_belt_only=True)
|
| model_ar = ConveyerDiffusionLM(cfg_ar).eval()
|
| out_ar_a = model_ar.forward_belt(belt_a)
|
| out_ar_b = model_ar.forward_belt(belt_b)
|
|
|
| assert torch.allclose(out_ar_a.slot0_logits, out_ar_b.slot0_logits, atol=1e-6)
|
|
|
|
|
| def test_left_shift_mechanics_and_elastic_planning_zone() -> None:
|
| """
|
| Verifies both:
|
| 1. Strict 1-by-1 conveyor shift (`use_elastic_planning=False`), where slots 1..63 shift into 0..62.
|
| 2. Elastic Two-Zone Planning Belt (`use_elastic_planning=True`), where the front 25% (`0..14`) shifts
|
| strictly from `1..15`, Slot 15 blends mid-belt sentence expansion (`expand_cand`), and the back 75%
|
| (`16..63`) pauses elastically via `shift_prob < 1.0`.
|
| """
|
| torch.manual_seed(99)
|
|
|
| cfg_strict = _small_model_cfg(belt_size=64)
|
| cfg_strict.use_elastic_planning = False
|
| model_strict = ConveyerDiffusionLM(cfg_strict).eval()
|
|
|
| belt = model_strict.init_empty_belt(batch_size=1, tick=0, device=torch.device("cpu"), deterministic=True)
|
| step_out = model_strict.forward_belt(belt)
|
| shifted = model_strict.shift_belt_left(step_out, current_tick=0, deterministic_spawn=True)
|
|
|
| assert shifted.tick == 1
|
| assert shifted.e_soft.shape == (1, 64, cfg_strict.d_model)
|
| assert torch.allclose(shifted.e_soft[:, :63, :], step_out.e_soft[:, 1:, :], atol=1e-6)
|
| assert torch.allclose(shifted.m_thought[:, :63, :], step_out.m_updated[:, 1:, :], atol=1e-6)
|
|
|
|
|
| cfg_elastic = _small_model_cfg(belt_size=64)
|
| cfg_elastic.use_elastic_planning = True
|
| cfg_elastic.front_conveyer_slots = 16
|
| model_elastic = ConveyerDiffusionLM(cfg_elastic).eval()
|
|
|
| step_el = model_elastic.forward_belt(belt)
|
| shifted_el = model_elastic.shift_belt_left(step_el, current_tick=0, deterministic_spawn=True)
|
|
|
|
|
| assert torch.allclose(shifted_el.e_soft[:, :15, :], step_el.e_soft[:, 1:16, :], atol=1e-6)
|
| assert torch.allclose(shifted_el.m_thought[:, :15, :], step_el.m_updated[:, 1:16, :], atol=1e-6)
|
|
|
| assert float(step_el.shift_prob[:, 16:].mean().item()) < 0.50
|
|
|
|
|
| def test_prompt_overlap_prebaking_clamps_input_and_emits() -> None:
|
| """
|
| Verifies the `n - W_pre` cold-start solution:
|
| When `prompt_overlap = 6` on a prompt of length `n = 10`:
|
| - Warm-up runs for ticks `t = 4, 5, 6, 7, 8, 9` (6 warm-up ticks), clamping emitted tokens to `prompt[t]`.
|
| - Generation runs for `max_new_tokens = 8` ticks (`t = 10..17`), returning a sequence of length `10 + 8 = 18`.
|
| """
|
| torch.manual_seed(11)
|
| cfg = _small_model_cfg(belt_size=32)
|
| model = ConveyerDiffusionLM(cfg).eval()
|
|
|
| prompt = torch.arange(5, 15, dtype=torch.long).unsqueeze(0)
|
| full_seq, snapshots = generate_conveyer(
|
| model=model,
|
| prompt_tokens=prompt,
|
| max_new_tokens=8,
|
| prompt_overlap=6,
|
| bake_steps=32,
|
| temperature=0.0,
|
| deterministic_spawn=True,
|
| )
|
|
|
| assert full_seq.shape == (1, 18)
|
| assert torch.equal(full_seq[:, :10], prompt)
|
| assert len(snapshots) == 6 + 8
|
|
|
|
|
| for i in range(6):
|
| assert snapshots[i].is_warmup is True
|
| assert snapshots[i].tick == 4 + i
|
| assert snapshots[i].emitted_token == int(prompt[0, 4 + i].item())
|
|
|
|
|
| for j in range(8):
|
| snap = snapshots[6 + j]
|
| assert snap.is_warmup is False
|
| assert snap.tick == 10 + j
|
| assert snap.emitted_token == int(full_seq[0, 10 + j].item())
|
|
|
|
|
| def test_unrolled_bptt_gradient_flow_and_slot63_zero_leakage() -> None:
|
| """
|
| Verifies that:
|
| 1. `seed_corrupted_belt` puts zero target signal at Slot 63 (pure noise).
|
| 2. `compute_unrolled_conveyer_loss` backpropagates non-zero gradients into both the lexical embeddings
|
| and the continuous thought highway (`thought_cand`, `thought_gate`, `slot_emb`, `m_init`).
|
| """
|
| torch.manual_seed(21)
|
| full_cfg = ConveyerConfig(model=_small_model_cfg(belt_size=16, use_thought_highway=True))
|
| full_cfg.train.unroll_steps = 4
|
| model = ConveyerDiffusionLM(full_cfg.model).train()
|
|
|
|
|
| torch.manual_seed(55)
|
| t_a = torch.ones((1, 16), dtype=torch.long) * 5
|
| b_a = seed_corrupted_belt(model, t_a, tick_start=2, cold_drop_prob=0.0)
|
| torch.manual_seed(55)
|
| t_b = torch.ones((1, 16), dtype=torch.long) * 5
|
| t_b[:, -1] = 42
|
| b_b = seed_corrupted_belt(model, t_b, tick_start=2, cold_drop_prob=0.0)
|
| assert torch.allclose(b_a.e_soft[:, -1, :], b_b.e_soft[:, -1, :], atol=1e-6)
|
|
|
|
|
| seq = torch.randint(4, 50, (3, 40), dtype=torch.long)
|
| plen = torch.tensor([8, 8, 8], dtype=torch.long)
|
| metrics = compute_unrolled_conveyer_loss(
|
| model=model,
|
| seq_tokens=seq,
|
| cfg=full_cfg,
|
| prompt_lengths=plen,
|
| tick_start=6,
|
| )
|
| metrics.loss.backward()
|
|
|
| assert model.tok_emb.weight.grad is not None and model.tok_emb.weight.grad.abs().sum() > 0
|
| assert model.slot_emb.weight.grad is not None and model.slot_emb.weight.grad.abs().sum() > 0
|
| assert model.thought_cand.weight.grad is not None and model.thought_cand.weight.grad.abs().sum() > 0
|
| assert model.thought_gate.weight.grad is not None and model.thought_gate.weight.grad.abs().sum() > 0
|
|
|
|
|
| def test_synthetic_lookahead_dataset_and_config_loader() -> None:
|
| cfg = load_config("configs/synthetic.yaml", ["train.max_steps=10", "data.num_train=60"])
|
| assert cfg.train.max_steps == 10
|
| assert cfg.data.num_train == 60
|
|
|
| ds = SyntheticLookaheadDataset(num_samples=12, seq_len=96, task="all", vocab_size=128)
|
| assert len(ds) == 12
|
| for i in range(6):
|
| item = ds[i]
|
| assert item["input_ids"].shape == (96,)
|
| assert 4 < item["prompt_len"] < 60
|
| assert item["task_id"] in (0, 1, 2)
|
|
|
