Experiment: convert a Qwen3.5 hybrid (Gated DeltaNet + gated softmax attention) into a tri-mode language model in the style of NVIDIA Nemotron-Labs-Diffusion (arXiv 2607.05722): one set of weights that runs as autoregressive, block diffusion, or self-speculative (diffusion draft, AR verify, lossless w.r.t. AR greedy).
The question this repo tests: does block diffusion train when only the softmax-attention layers (1 in 4) see the block bidirectionally, with Gated DeltaNet kept causal? That is the constraint the larger Qwen3.5 / Qwen3.8 hybrids impose (e.g. Flash-Next: 36 GDN + 12 attention layers). So far this has been tried on Qwen3.5-0.8B only. The goal is to apply the same recipe to the larger MoE hybrids (Qwen3.5-35B-A3B, then Qwen3.8-Flash-Next) with experts frozen.
Small-scale results on Qwen3.5-0.8B, fineweb-edu tokens, one 16 GB GPU. Nothing here is tuned or scaled; treat the numbers as a feasibility signal, not a benchmark.
Run A: 3000 steps x 4 x 1024 tokens (~12M tokens), training attention + norms + mask embedding only
(44M params); GDN and MLP frozen. Loss = L_AR + 0.3 * L_diff.
| metric (val) | before | after |
|---|---|---|
| AR loss | 2.99 | 2.88 |
| block-diffusion loss (masked positions) | 8.60 | 4.61 |
| self-spec accepted tokens per round (draft 2 steps, block 32) | 1.16 | ~2.6 |
- AR mode is not hurt by the joint objective at this budget (it improves slightly).
- Self-speculation output equals AR greedy output (modulo bf16 tie flips).
- Best self-spec layout so far is the single-forward "fan" variant in
decode.py: 1.64-1.69 accepted tokens per forward at draft length 8-12, vs 1.24 for draft+verify as two forwards. - Pure block-diffusion decoding is not usable at this budget: few-step confidence unmasking collapses to repeated tokens; threshold unmasking stays coherent but commits ~1 token per forward. Published diffusion conversions use two to three orders of magnitude more tokens.
trimode/dual.py— dual-stream forward[clean | noisy blocks]. Clean stream is exactly the causal AR forward. Noisy block j: attention sees clean tokens before its start plus its own block bidirectionally (FlexAttention block-sparse mask); GDN runs causally inside the block seeded from the clean recurrent state at the block start. Block-start states come from oneflacall by inserting probe tokens (SCAN_MODE="probe"); a sequential reference scan is kept ("loop").trimode/data.py,trimode/prepare_data.py— flat uint32 token file + per-block masking.trimode/train.py— joint AR + diffusion training with freeze groups (attn,gdn,mlp,norm,embed,mask), resumable checkpoints. Mask token = a spare embedding row (248319 for Qwen3.5-0.8B).trimode/distill.py— stage 2: self-distil the one-step drafter onto the model's own AR greedy continuation (keeps the AR loss).trimode/decode.py— AR, block diffusion (fixed-k and threshold unmasking), self-spec (two-forward, merged, fan).trimode/bench_decode.py,trimode/bench_diffusion.py,trimode/demo.py,trimode/export_ckpt.py.tests/test_dual.py— leakage / equivalence checks for the dual forward.scripts/long_run.ps1— the two-stage run (attention-only, then full-param) used above.
python -m venv .venv && .venv\Scripts\activate
pip install -r requirements.txt
python -m trimode.prepare_data --tokens 100_000_000 # -> data/fineweb_edu.bin (or set TRIMODE_DATA)
python -m tests.test_dual
python -m trimode.train --out runs/smoke --steps 400 --train attn,norm,mask --grad-ckpt --save
python -m trimode.demo runs/smoke/model
python -m trimode.bench_decode runs/smoke/model
Notes: transformers>=5.16 is needed for qwen3_5; fla-core provides chunk_gated_delta_rule.
FlexAttention with head_dim 256 needs BLOCK_M/N=32 kernel options on consumer GPUs (set in dual.py).
torch.compile works on Windows with triton-windows; the compile caches default to .tc/ in the repo
to stay under MAX_PATH.
- Longer runs and full-parameter training on the 0.8B testbed.
- Self-distillation targeting the fan layout (the fan draft sees one token less than the two-forward draft).
- Same recipe on Qwen3.5-35B-A3B with experts frozen, then Qwen3.8-Flash-Next.
Apache-2.0.