Add dqwen loader code

dqwen/__init__.py

@@ -0,0 +1 @@
+from .core import *  # noqa

dqwen/bart_fusion.py

@@ -0,0 +1,201 @@
+"""BART (facebook/bart-large, frozen FP32) -> Qwen fusion.
+
+Pipeline
+--------
+  1. Qwen masked-diffusion generates a DRAFT
+  2. BART encodes the draft text -> `bart_states` [B, S, 1024]
+  3. Qwen remasks its lowest-confidence draft tokens and re-denoises them while `BartFusion`
+     injects the BART context into the response-region input embeddings
+  -> generate(prompt) returns (draft, final)
+
+dtype contract (this is the fix for the earlier BF16/FP32 mismatch, built in):
+  * BART stays FP32 and frozen.
+  * BartFusion casts `bart_states` and `qwen_embeds` to *its own weight dtype/device* inside forward(),
+    and casts the result back to the dtype of `qwen_embeds`. Works whether fusion is FP32 or BF16.
+"""
+from __future__ import annotations
+
+import math
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .core import DiffusionLM, diffusion_generate, masked_diffusion_loss
+
+
+def fourier_pos(frac: torch.Tensor, dim: int = 64) -> torch.Tensor:
+    """frac: [...,] in roughly [0, 2] -> [..., dim] sin/cos features."""
+    n = dim // 2
+    freqs = (2.0 ** torch.linspace(0, 6, n, device=frac.device)) * math.pi
+    ang = frac[..., None].float() * freqs
+    return torch.cat([ang.sin(), ang.cos()], dim=-1)
+
+
+class BartFusion(nn.Module):
+    def __init__(self, qwen_dim: int = 5120, bart_dim: int = 1024, d_inner: int = 1024, n_heads: int = 16):
+        super().__init__()
+        assert d_inner % n_heads == 0
+        self.n_heads, self.d_inner = n_heads, d_inner
+        # NOTE: proj[0] is a Linear -> `fusion.proj[0].weight` is valid (used for dtype/device lookups)
+        self.proj = nn.Sequential(nn.Linear(bart_dim, qwen_dim), nn.GELU(), nn.Linear(qwen_dim, d_inner))
+        self.ln_q = nn.LayerNorm(qwen_dim)
+        self.q = nn.Linear(qwen_dim, d_inner)
+        self.k = nn.Linear(d_inner, d_inner)
+        self.v = nn.Linear(d_inner, d_inner)
+        self.o = nn.Linear(d_inner, qwen_dim)
+        self.pos_q = nn.Linear(64, d_inner)
+        self.pos_k = nn.Linear(64, d_inner)
+        self.gate = nn.Parameter(torch.zeros(1))     # tanh-gated, 0 at init => identical to plain Qwen
+
+    def forward(self, qwen_embeds, bart_states, prefix_len, bart_mask=None, resp_len=None):
+        """qwen_embeds [B,L,H]; bart_states [B,S,bart_dim]; prefix_len: int or LongTensor[B] (prompt length);
+        resp_len: int or LongTensor[B] = real (pre-EOS) response length, used for position alignment."""
+        w = self.proj[0].weight
+        tgt_dtype, tgt_dev = w.dtype, w.device
+        out_dtype = qwen_embeds.dtype
+        bart_states = bart_states.to(device=tgt_dev, dtype=tgt_dtype)          # <- dtype fix
+        x = qwen_embeds.to(device=tgt_dev, dtype=tgt_dtype)
+
+        B, L, _ = x.shape
+        S = bart_states.shape[1]
+        prefix = torch.as_tensor(prefix_len, device=tgt_dev).reshape(-1).expand(B).long()
+        rel = torch.arange(L, device=tgt_dev)[None, :] - prefix[:, None]        # [B,L] index inside the response
+        is_resp = rel >= 0
+        n_real = torch.as_tensor(resp_len if resp_len is not None else (L - prefix),
+                                 device=tgt_dev).reshape(-1).expand(B).float().clamp(min=1)
+        fq = rel.clamp(min=0).float() / n_real[:, None]
+        if bart_mask is not None:
+            s_len = bart_mask.sum(1).clamp(min=1).float().to(tgt_dev)
+        else:
+            s_len = torch.full((B,), float(S), device=tgt_dev)
+        fk = torch.arange(S, device=tgt_dev, dtype=torch.float32)[None, :] / s_len[:, None]
+
+        ctx = self.proj(bart_states)                                           # [B,S,d_inner]
+        q = self.q(self.ln_q(x)) + self.pos_q(fourier_pos(fq).to(tgt_dtype))
+        k = self.k(ctx) + self.pos_k(fourier_pos(fk).to(tgt_dtype))
+        v = self.v(ctx)
+        h = self.n_heads
+        sh = lambda t: t.view(B, -1, h, self.d_inner // h).transpose(1, 2)
+        am = bart_mask.to(tgt_dev).bool()[:, None, None, :] if bart_mask is not None else None
+        a = F.scaled_dot_product_attention(sh(q), sh(k), sh(v), attn_mask=am)
+        delta = self.o(a.transpose(1, 2).reshape(B, L, self.d_inner))
+
+        # keep the injected signal on the same scale as the embeddings it is added to
+        d_rms = delta.float().pow(2).mean(-1, keepdim=True).add(1e-6).rsqrt()
+        e_rms = x.float().pow(2).mean(-1, keepdim=True).sqrt().detach()
+        delta = (delta.float() * d_rms * e_rms) * torch.tanh(self.gate.float())
+        delta = delta * is_resp[..., None].float()                              # prompt positions untouched
+        return qwen_embeds + delta.to(device=qwen_embeds.device, dtype=out_dtype)
+
+
+# ------------------------------------------------------------------------------------------
+@torch.no_grad()
+def bart_encode(bart, btok, texts, device, max_len: int = 512):
+    enc = btok(list(texts), return_tensors="pt", padding=True, truncation=True, max_length=max_len)
+    ids = enc["input_ids"].to(device)
+    mask = enc["attention_mask"].to(device)
+    enc_dtype = next(bart.parameters()).dtype
+    states = bart.get_encoder()(input_ids=ids, attention_mask=mask).last_hidden_state
+    return states.to(enc_dtype), mask
+
+
+def _cut_at_eos(ids_1d: torch.Tensor, eos_id: int) -> int:
+    hit = (ids_1d == eos_id).nonzero(as_tuple=True)[0]
+    return int(hit[0]) if hit.numel() else int(ids_1d.numel())
+
+
+class Pipeline:
+    """Holds everything needed for inference: draft (Qwen diffusion) -> BART -> refine (Qwen + fusion)."""
+
+    def __init__(self, dlm: DiffusionLM, tok, bart, btok, fusion: BartFusion, eos_id: int, device,
+                 chat_kwargs: Optional[dict] = None):
+        self.dlm, self.tok, self.bart, self.btok, self.fusion = dlm, tok, bart, btok, fusion
+        self.eos_id, self.device = eos_id, device
+        self.chat_kwargs = chat_kwargs or {"enable_thinking": False}
+
+    def encode_prompt(self, prompt: str) -> torch.Tensor:
+        text = self.tok.apply_chat_template(
+            [{"role": "user", "content": prompt}], tokenize=False, add_generation_prompt=True, **self.chat_kwargs
+        )
+        return torch.tensor([self.tok(text, add_special_tokens=False)["input_ids"]], device=self.device)
+
+    def text_of(self, ids_1d: torch.Tensor) -> str:
+        n = _cut_at_eos(ids_1d, self.eos_id)
+        return self.tok.decode(ids_1d[:n].tolist(), skip_special_tokens=True)
+
+    @torch.no_grad()
+    def generate(self, prompt: str, gen_len=256, steps=64, block_len=32, temperature=0.0,
+                 refine_ratio=0.3, refine_steps=16, use_bart=True):
+        self.dlm.eval(); self.fusion.eval()
+        p = self.encode_prompt(prompt)
+        draft_ids, conf = diffusion_generate(self.dlm, p, gen_len=gen_len, steps=steps, block_len=block_len,
+                                             temperature=temperature, eos_id=self.eos_id, return_conf=True)
+        draft = self.text_of(draft_ids[0])
+        if not use_bart:
+            return draft, draft
+        n_real = _cut_at_eos(draft_ids[0], self.eos_id)
+        if n_real == 0:
+            return draft, draft
+        states, bmask = bart_encode(self.bart, self.btok, [draft if draft.strip() else " "], self.device)
+        fusion_fn = lambda emb, plen: self.fusion(emb, states, plen, bart_mask=bmask, resp_len=n_real)
+        k = max(1, int(refine_ratio * n_real))
+        low = conf[0, :n_real].topk(k, largest=False).indices
+        init = draft_ids.clone()
+        init[0, low] = self.dlm.mask_id
+        # tail after EOS stays EOS (not re-opened)
+        final_ids = diffusion_generate(self.dlm, p, steps=refine_steps, block_len=draft_ids.shape[1],
+                                       temperature=temperature, eos_id=self.eos_id, fusion=fusion_fn,
+                                       init_response=init)
+        return draft, self.text_of(final_ids[0])
+
+
+# ------------------------------------------------------------------------------------------
+@torch.no_grad()
+def make_drafts(dlm: DiffusionLM, ids, resp, t_lo=0.3, t_hi=0.9):
+    """One-shot noisy drafts that look like real diffusion errors: mask at rate t, fill every masked
+    position with argmax (+ record the commit-time confidence). Returns (draft_ids, conf)."""
+    B, L = ids.shape
+    resp = resp.clone(); resp[:, 0] = False
+    t = torch.rand(B, device=ids.device) * (t_hi - t_lo) + t_lo
+    masked = (torch.rand(B, L, device=ids.device) < t[:, None]) & resp
+    noisy = torch.where(masked, torch.full_like(ids, dlm.mask_id), ids)
+    h = dlm.hidden(noisy)
+    b, j = masked.nonzero(as_tuple=True)
+    lg = dlm.head(h[b, j - 1]).float()
+    pr = lg.softmax(-1)
+    cand = pr.argmax(-1)
+    draft = ids.clone()
+    conf = torch.ones(B, L, device=ids.device)
+    draft[b, j] = cand
+    conf[b, j] = pr.gather(-1, cand[:, None]).squeeze(-1)
+    return draft, conf
+
+
+def refine_loss(dlm, fusion, bart, btok, tok, ids, resp, plen, eos_id, device, ratio_lo=0.15, ratio_hi=0.6):
+    """Stage-2 loss: remask the lowest-confidence draft tokens, predict ground truth with BART context."""
+    draft, conf = make_drafts(dlm, ids, resp)
+    B, L = ids.shape
+    texts, n_real = [], []
+    for i in range(B):
+        r = draft[i, plen[i]:]
+        n = _cut_at_eos(r, eos_id)
+        n_real.append(max(n, 1))
+        t = tok.decode(r[:n].tolist(), skip_special_tokens=True)
+        texts.append(t if t.strip() else " ")
+    states, bmask = bart_encode(bart, btok, texts, device)
+
+    rho = torch.rand(B, device=device) * (ratio_hi - ratio_lo) + ratio_lo
+    masked = torch.zeros(B, L, dtype=torch.bool, device=device)
+    for i in range(B):
+        lo, hi = int(plen[i]), L
+        k = max(1, int(rho[i] * n_real[i]))
+        c = conf[i, lo:lo + n_real[i]]
+        idx = c.topk(min(k, c.numel()), largest=False).indices + lo
+        masked[i, idx] = True
+    nr = torch.tensor(n_real, device=device)
+    pl = torch.as_tensor(plen, device=device)
+    fusion_fn = lambda emb, plen_: fusion(emb, states, plen_, bart_mask=bmask, resp_len=nr)
+    return masked_diffusion_loss(dlm, ids, resp, weighting="uniform", fusion=fusion_fn, prefix_len=pl,
+                                 forced_masked=masked, context_ids=draft)

dqwen/core.py

@@ -0,0 +1,314 @@
+"""Masked-diffusion adaptation of Qwen3.5 (hybrid Gated-DeltaNet + full attention).
+
+Pieces
+------
+* patch_bidirectional : full-attention layers -> bidirectional, DeltaNet layers -> forward + reverse scan
+                        (reverse path gated by a zero-initialised per-channel gate, so step 0 == original DeltaNet)
+* DiffusionLM         : wrapper with a trainable [MASK] embedding, 4D/2D mask dict builder, hidden-state API
+* masked_diffusion_loss : LLaDA-style masked-token CE with "shift" (position j is predicted from hidden j-1,
+                        which keeps the AR-pretrained lm_head alignment, as in Dream)
+* diffusion_generate  : block-wise (semi-AR) confidence-based parallel unmasking
+"""
+from __future__ import annotations
+
+import math
+from typing import Callable, Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+try:  # class lookup is done lazily so the module imports even on older transformers
+    from transformers.models.qwen3_5.modeling_qwen3_5 import (
+        Qwen3_5Attention,
+        Qwen3_5GatedDeltaNet,
+        Qwen3_5TextModel,
+    )
+except Exception as e:  # pragma: no cover
+    raise ImportError(
+        "Qwen3.5 classes not found. Install a recent transformers: pip install -U transformers"
+    ) from e
+
+
+# --------------------------------------------------------------------------------------
+# 1. architecture surgery
+# --------------------------------------------------------------------------------------
+def _patch_deltanet(mod: Qwen3_5GatedDeltaNet, hidden_size: int, bidirectional: bool = True):
+    """Make one GatedDeltaNet layer bidirectional.
+
+    out = orig(h) + g * flip(orig(flip(h)))      (out_proj is linear & bias-free, so this equals
+                                                   out_proj(core_f + g * core_r) exactly)
+    g is a per-channel fp32 parameter initialised to 0.
+    """
+    if getattr(mod, "_dqwen_patched", False):
+        return
+    orig_forward = mod.forward
+    # The 4-bit model has already been dispatched to CUDA before patching.
+    # New parameters otherwise default to CPU and cause device mismatch.
+    reference = next(mod.parameters())
+    mod.rev_gate = nn.Parameter(torch.zeros(hidden_size, device=reference.device, dtype=torch.float32))
+    mod._dqwen_patched = True
+
+    def forward(hidden_states, cache_params=None, attention_mask=None, **kwargs):
+        y = orig_forward(hidden_states, cache_params=None, attention_mask=attention_mask, **dict(kwargs))
+        if not bidirectional:
+            return y
+        h_r = hidden_states.flip(1)
+        m_r = attention_mask.flip(1) if attention_mask is not None else None
+        y_r = orig_forward(h_r, cache_params=None, attention_mask=m_r, **dict(kwargs)).flip(1)
+        # The first device-fix only moved rev_gate; y_r may still differ.
+        # Move the reverse output to the forward output's device with autograd intact.
+        if not isinstance(y, torch.Tensor) or not isinstance(y_r, torch.Tensor):
+            raise TypeError(f"DeltaNet outputs: forward={type(y)!r}, reverse={type(y_r)!r}")
+        if y.device != y_r.device:
+            print(f"[DeltaNet device] forward={y.device}, reverse={y_r.device}, "
+                  f"gate={mod.rev_gate.device}; aligning reverse output")
+            y_r = y_r.to(device=y.device, dtype=y.dtype)
+        gate = mod.rev_gate.to(device=y.device, dtype=y.dtype)
+        try:
+            return y + gate * y_r
+        except RuntimeError as exc:
+            raise RuntimeError(
+                f"DeltaNet fusion failed: forward={y.device}/{y.dtype}/{tuple(y.shape)}, "
+                f"reverse={y_r.device}/{y_r.dtype}/{tuple(y_r.shape)}, "
+                f"gate={gate.device}/{gate.dtype}/{tuple(gate.shape)}; "
+                f"source_gate={mod.rev_gate.device}. Original error: {exc}"
+            ) from exc
+
+    mod.forward = forward
+
+
+def patch_bidirectional(model: nn.Module, deltanet_bidirectional: bool = True) -> dict:
+    """In-place surgery. Returns simple stats."""
+    n_full = n_delta = 0
+    text = find_text_model(model)
+    hidden = text.config.hidden_size
+    for m in text.modules():
+        if isinstance(m, Qwen3_5Attention):
+            m.is_causal = False
+            n_full += 1
+        elif isinstance(m, Qwen3_5GatedDeltaNet):
+            _patch_deltanet(m, hidden, deltanet_bidirectional)
+            n_delta += 1
+    return {"full_attention_layers": n_full, "deltanet_layers": n_delta}
+
+
+def find_text_model(model: nn.Module) -> Qwen3_5TextModel:
+    for m in model.modules():
+        if isinstance(m, Qwen3_5TextModel):
+            return m
+    raise ValueError("Qwen3_5TextModel not found inside the given module")
+
+
+def find_lm_head(model: nn.Module) -> nn.Module:
+    for name, m in model.named_modules():
+        if name.endswith("lm_head"):
+            return m
+    raise ValueError("lm_head not found")
+
+
+def rev_gate_state(model: nn.Module) -> dict:
+    return {n: p.detach().cpu() for n, p in model.named_parameters() if n.endswith("rev_gate")}
+
+
+def load_rev_gates(model: nn.Module, state: dict):
+    params = dict(model.named_parameters())
+    missing = []
+    for n, v in state.items():
+        # tolerate different peft prefixes by matching on suffix after 'layers.'
+        key = n if n in params else next((k for k in params if k.endswith(n.split("layers.", 1)[-1]) and "rev_gate" in k), None)
+        if key is None:
+            missing.append(n)
+            continue
+        params[key].data.copy_(v.to(params[key].dtype))
+    return missing
+
+
+# --------------------------------------------------------------------------------------
+# 2. wrapper
+# --------------------------------------------------------------------------------------
+class DiffusionLM(nn.Module):
+    """Wraps a (possibly PEFT-wrapped) Qwen3_5ForCausalLM.
+
+    `fusion` (optional) is a callable (embeds, prefix_len) -> embeds that injects extra context
+    (e.g. BART states) into the input embeddings of the response region.
+    """
+
+    def __init__(self, lm: nn.Module, mask_id: int, init_mask_from_token: Optional[int] = None):
+        super().__init__()
+        self.lm = lm
+        self.mask_id = int(mask_id)
+        self.text = find_text_model(lm)
+        self.head = find_lm_head(lm)
+        hidden = self.text.config.hidden_size
+        emb = self.text.embed_tokens.weight
+        init = emb[init_mask_from_token if init_mask_from_token is not None else mask_id].detach().float().clone()
+        if init.abs().sum() == 0:  # untrained row -> fall back to the mean embedding
+            init = emb.detach().float().mean(0)
+        self.mask_embed = nn.Parameter(init.to(emb.device))
+        self.hidden_size = hidden
+
+    # ---- masks ----------------------------------------------------------------------
+    def build_masks(self, attn2d: Optional[torch.Tensor], dtype=torch.bfloat16):
+        if attn2d is None:
+            return {"full_attention": None, "linear_attention": None}
+        a = attn2d.bool()
+        full = (a[:, None, :, None] & a[:, None, None, :])  # [B,1,L,L] True = attend
+        # a padded query row would be all False -> NaN in sdpa; let pads attend to themselves
+        eye = torch.eye(a.shape[1], dtype=torch.bool, device=a.device)[None, None]
+        full = full | (~a[:, None, :, None] & eye)
+        if getattr(self.text.config, "_attn_implementation", "sdpa") == "eager":
+            neg = torch.finfo(dtype).min
+            full = torch.zeros(full.shape, dtype=dtype, device=a.device).masked_fill(~full, neg)
+        return {"full_attention": full, "linear_attention": a.long()}
+
+    # ---- embeddings -----------------------------------------------------------------
+    def embed(self, input_ids: torch.Tensor) -> torch.Tensor:
+        e = self.text.embed_tokens(input_ids)
+        is_mask = (input_ids == self.mask_id)[..., None]
+        return torch.where(is_mask, self.mask_embed.to(e.dtype)[None, None, :], e)
+
+    def hidden(self, input_ids, attn2d=None, fusion: Optional[Callable] = None, prefix_len=None):
+        emb = self.embed(input_ids)
+        if fusion is not None:
+            emb = fusion(emb, prefix_len)
+        out = self.text(
+            inputs_embeds=emb,
+            attention_mask=self.build_masks(attn2d, emb.dtype),
+            use_cache=False,
+        )
+        return out.last_hidden_state
+
+    def logits_at(self, h: torch.Tensor) -> torch.Tensor:
+        return self.head(h)
+
+
+# --------------------------------------------------------------------------------------
+# 3. loss
+# --------------------------------------------------------------------------------------
+def _chunked_ce(head, hs: torch.Tensor, targets: torch.Tensor, chunk: int = 1024):
+    """CE per token without materialising [N, V] logits for all N at once."""
+    outs = []
+    for i in range(0, hs.shape[0], chunk):
+        lg = head(hs[i : i + chunk]).float()
+        outs.append(F.cross_entropy(lg, targets[i : i + chunk], reduction="none"))
+    return torch.cat(outs) if outs else hs.new_zeros(0)
+
+
+def masked_diffusion_loss(
+    dlm: DiffusionLM,
+    input_ids: torch.Tensor,      # [B, L] clean tokens (prompt + response + EOS padding)
+    resp_mask: torch.Tensor,      # [B, L] bool, True where the token belongs to the response (may be masked)
+    attn2d: Optional[torch.Tensor] = None,
+    t_min: float = 0.05,
+    weighting: str = "elbo",      # "elbo" -> 1/t (LLaDA/Dream), "uniform" -> 1
+    fusion: Optional[Callable] = None,
+    prefix_len=None,
+    forced_masked: Optional[torch.Tensor] = None,  # if given, mask exactly this set (used for refine-stage training)
+    context_ids: Optional[torch.Tensor] = None,    # tokens shown at un-masked positions (draft); targets stay `input_ids`
+):
+    B, L = input_ids.shape
+    dev = input_ids.device
+    resp_mask = resp_mask.clone()
+    resp_mask[:, 0] = False  # position 0 has no left neighbour to predict it (shift)
+
+    if forced_masked is None:
+        t = torch.rand(B, device=dev) * (1 - t_min) + t_min
+        masked = (torch.rand(B, L, device=dev) < t[:, None]) & resp_mask
+        # guarantee >=1 masked token per row
+        none = masked.sum(1) == 0
+        if none.any():
+            idx = torch.multinomial(resp_mask[none].float() + 1e-9, 1).squeeze(1)
+            rows = none.nonzero(as_tuple=True)[0]
+            masked[rows, idx] = True
+    else:
+        masked = forced_masked & resp_mask
+        t = masked.sum(1).float() / resp_mask.sum(1).clamp(min=1)
+        t = t.clamp(min=t_min)
+
+    base = input_ids if context_ids is None else context_ids
+    noisy = torch.where(masked, torch.full_like(input_ids, dlm.mask_id), base)
+    h = dlm.hidden(noisy, attn2d, fusion=fusion, prefix_len=prefix_len)
+
+    b_idx, j_idx = masked.nonzero(as_tuple=True)
+    hs = h[b_idx, j_idx - 1]
+    tgt = input_ids[b_idx, j_idx]
+    ce = _chunked_ce(dlm.head, hs, tgt)
+
+    if weighting == "elbo":
+        w = 1.0 / t[b_idx]
+    else:
+        w = torch.ones_like(ce)
+    n_resp = resp_mask.sum(1).clamp(min=1).float()
+    loss = (w * ce / n_resp[b_idx]).sum() / B
+
+    return loss, {"ce": ce.mean().detach(), "n_masked": int(masked.sum()), "t": t.mean().detach()}
+
+
+# --------------------------------------------------------------------------------------
+# 4. sampler
+# --------------------------------------------------------------------------------------
+@torch.no_grad()
+def diffusion_generate(
+    dlm: DiffusionLM,
+    prompt_ids: torch.Tensor,           # [1, P]
+    gen_len: int = 256,
+    steps: int = 64,
+    block_len: int = 32,
+    temperature: float = 0.0,
+    eos_id: Optional[int] = None,
+    fusion: Optional[Callable] = None,
+    init_response: Optional[torch.Tensor] = None,   # [1, gen_len] tokens; MASK id => to be generated
+    return_conf: bool = False,
+):
+    """Block-wise low-confidence remasking (LLaDA). Returns [1, gen_len] response tokens."""
+    assert prompt_ids.shape[0] == 1
+    dev = prompt_ids.device
+    P = prompt_ids.shape[1]
+    if init_response is None:
+        resp = torch.full((1, gen_len), dlm.mask_id, dtype=torch.long, device=dev)
+    else:
+        resp = init_response.clone().to(dev)
+        gen_len = resp.shape[1]
+    x = torch.cat([prompt_ids, resp], dim=1)
+    conf_all = torch.zeros(1, gen_len, device=dev)
+
+    n_blocks = math.ceil(gen_len / block_len)
+    steps_per_block = max(1, steps // n_blocks)
+
+    for b in range(n_blocks):
+        lo, hi = P + b * block_len, min(P + (b + 1) * block_len, P + gen_len)
+        blk_masked0 = (x[0, lo:hi] == dlm.mask_id).sum().item()
+        if blk_masked0 == 0:
+            continue
+        # tokens to commit per step (spread evenly)
+        base, rem = divmod(blk_masked0, steps_per_block)
+        schedule = [base + (1 if i < rem else 0) for i in range(steps_per_block)]
+        for k in schedule:
+            if k == 0:
+                continue
+            h = dlm.hidden(x, None, fusion=fusion, prefix_len=P)
+            pos = torch.arange(lo, hi, device=dev)
+            still = x[0, lo:hi] == dlm.mask_id
+            pos = pos[still]
+            if pos.numel() == 0:
+                break
+            lg = dlm.head(h[0, pos - 1]).float()                      # shift: predict j from j-1
+            if temperature > 0:
+                g = -torch.log(-torch.log(torch.rand_like(lg).clamp_min(1e-10)).clamp_min(1e-10))
+                cand = (lg / temperature + g).argmax(-1)
+            else:
+                cand = lg.argmax(-1)
+            prob = F.softmax(lg, -1).gather(-1, cand[:, None]).squeeze(-1)
+            kk = min(k, pos.numel())
+            top = prob.topk(kk).indices
+            x[0, pos[top]] = cand[top]
+            conf_all[0, pos[top] - P] = prob[top]
+        if eos_id is not None and (x[0, P : hi] == eos_id).any() and not (x[0, P:hi] == dlm.mask_id).any():
+            break
+    out = x[:, P:].clone()
+    if eos_id is not None:
+        hit = (out[0] == eos_id).nonzero(as_tuple=True)[0]
+        if hit.numel() > 0:
+            out[0, hit[0]:] = eos_id          # everything after the first EOS is padding
+    return (out, conf_all) if return_conf else out

dqwen/data.py

@@ -0,0 +1,93 @@
+"""Chat-SFT data -> fixed-length (prompt | response | EOS filler) tensors for masked diffusion.
+
+Following LLaDA, the response region is padded with EOS up to the fixed length and those filler tokens are
+part of the response (so the model learns *when to stop* by predicting EOS at the tail).
+"""
+from __future__ import annotations
+
+import json
+import random
+from typing import Iterable, Optional
+
+import torch
+
+
+def to_messages(row: dict) -> Optional[list]:
+    """Normalise common SFT row layouts to [{'role','content'}, ...] ending with an assistant turn."""
+    if "messages" in row and row["messages"]:
+        msgs = [{"role": m["role"], "content": m["content"]} for m in row["messages"]]
+    elif "conversations" in row and row["conversations"]:
+        rm = {"human": "user", "user": "user", "gpt": "assistant", "assistant": "assistant", "system": "system"}
+        msgs = [{"role": rm.get(m.get("from", m.get("role")), "user"), "content": m.get("value", m.get("content"))}
+                for m in row["conversations"]]
+    elif "instruction" in row and ("output" in row or "response" in row):
+        q = row["instruction"] + (("\n" + row["input"]) if row.get("input") else "")
+        msgs = [{"role": "user", "content": q}, {"role": "assistant", "content": row.get("output", row.get("response"))}]
+    elif "question" in row and "answer" in row:
+        msgs = [{"role": "user", "content": row["question"]}, {"role": "assistant", "content": row["answer"]}]
+    else:
+        return None
+    if len(msgs) < 2 or msgs[-1]["role"] != "assistant" or not msgs[-1]["content"]:
+        return None
+    return msgs
+
+
+def iter_rows(source: str, config: Optional[str] = None, split: str = "train", n: int = 20000, seed: int = 0) -> Iterable[dict]:
+    """`source` = path to .jsonl, or a Hugging Face dataset id (streamed, so nothing huge is downloaded)."""
+    if source.endswith(".jsonl") or source.endswith(".json"):
+        with open(source, encoding="utf-8") as f:
+            rows = [json.loads(l) for l in f if l.strip()]
+        random.Random(seed).shuffle(rows)
+        yield from rows[:n]
+        return
+    from datasets import load_dataset
+    ds = load_dataset(source, config, split=split, streaming=True).shuffle(seed=seed, buffer_size=10_000)
+    for i, r in enumerate(ds):
+        if i >= n:
+            break
+        yield r
+
+
+def build_dataset(tok, rows: Iterable[dict], max_len: int = 512, max_prompt: int = 256,
+                  eos_token: str = "<|im_end|>", chat_kwargs: Optional[dict] = None, verbose: bool = True):
+    chat_kwargs = chat_kwargs or {"enable_thinking": False}
+    eos_id = tok.convert_tokens_to_ids(eos_token)
+    ids_l, resp_l, plen_l = [], [], []
+    seen = kept = 0
+    for row in rows:
+        seen += 1
+        msgs = to_messages(row)
+        if msgs is None:
+            continue
+        ptxt = tok.apply_chat_template(msgs[:-1], tokenize=False, add_generation_prompt=True, **chat_kwargs)
+        p = tok(ptxt, add_special_tokens=False)["input_ids"]
+        r = tok(msgs[-1]["content"], add_special_tokens=False)["input_ids"] + [eos_id]
+        if len(p) > max_prompt or len(p) + len(r) > max_len:
+            continue                                   # drop (never truncate: a cut response would lose its EOS)
+        pad = max_len - len(p) - len(r)
+        seq = p + r + [eos_id] * pad
+        rm = [False] * len(p) + [True] * (len(r) + pad)
+        ids_l.append(seq); resp_l.append(rm); plen_l.append(len(p)); kept += 1
+    if verbose:
+        print(f"[data] kept {kept}/{seen} examples (max_len={max_len}, max_prompt={max_prompt})")
+    if kept == 0:
+        raise ValueError(
+            f"No example survived filtering (seen {seen}). The chat template alone may exceed max_prompt="
+            f"{max_prompt}; raise max_prompt/max_len or check the dataset columns.")
+    return {
+        "ids": torch.tensor(ids_l, dtype=torch.long),
+        "resp": torch.tensor(resp_l, dtype=torch.bool),
+        "plen": torch.tensor(plen_l, dtype=torch.long),
+    }
+
+
+def batches(ds: dict, batch_size: int, device, seed: int = 0, epochs: int = 10**9):
+    g = torch.Generator().manual_seed(seed)
+    N = ds["ids"].shape[0]
+    if N < batch_size:
+        raise ValueError(f"Dataset has {N} examples but batch_size={batch_size}; lower batch_size or add data.")
+    for _ in range(epochs):
+        perm = torch.randperm(N, generator=g)
+        for i in range(0, N - batch_size + 1, batch_size):
+            ix = perm[i:i + batch_size]
+            yield ds["ids"][ix].to(device), ds["resp"][ix].to(device), ds["plen"][ix].to(device)

dqwen/evalutil.py

@@ -0,0 +1,57 @@
+"""Evaluation: held-out masked-token CE, GSM8K accuracy (draft vs BART-refined)."""
+from __future__ import annotations
+
+import re
+from typing import Optional
+
+import torch
+
+from .core import DiffusionLM, masked_diffusion_loss
+from .data import batches
+
+
+@torch.no_grad()
+def heldout_ce(dlm: DiffusionLM, ds: dict, device, n_batches: int = 10, batch_size: int = 2, rate: float = 0.5, seed: int = 123):
+    """Mean CE (nats/token) on masked response tokens at a fixed masking rate -> lower is better."""
+    was = dlm.training
+    dlm.eval()
+    g = torch.Generator(device="cpu").manual_seed(seed)
+    tot, n = 0.0, 0
+    for i, (ids, resp, _) in enumerate(batches(ds, batch_size, device, seed=seed, epochs=1)):
+        if i >= n_batches:
+            break
+        r = resp.clone(); r[:, 0] = False
+        forced = (torch.rand(ids.shape, generator=g).to(device) < rate) & r
+        _, info = masked_diffusion_loss(dlm, ids, resp, forced_masked=forced, weighting="uniform")
+        tot += float(info["ce"]); n += 1
+    dlm.train(was)
+    return tot / max(n, 1)
+
+
+_num = re.compile(r"-?\d[\d,]*\.?\d*")
+
+
+def _last_number(text: str) -> Optional[str]:
+    m = _num.findall(text.replace("$", ""))
+    return m[-1].replace(",", "").rstrip(".") if m else None
+
+
+def extract_answer(text: str) -> Optional[str]:
+    if "####" in text:
+        return _last_number(text.split("####")[-1])
+    return _last_number(text)
+
+
+def gsm8k_eval(pipe, n: int = 50, gen_len: int = 256, steps: int = 128, block_len: int = 32, use_bart: bool = True, seed: int = 0):
+    from datasets import load_dataset
+    ds = load_dataset("openai/gsm8k", "main", split="test").shuffle(seed=seed).select(range(n))
+    ok_draft = ok_final = 0
+    for i, row in enumerate(ds):
+        gold = row["answer"].split("####")[-1].strip().replace(",", "")
+        q = row["question"] + "\nSolve step by step, then write the final number after '####'."
+        draft, final = pipe.generate(q, gen_len=gen_len, steps=steps, block_len=block_len, use_bart=use_bart)
+        ok_draft += extract_answer(draft) == gold
+        ok_final += extract_answer(final) == gold
+        if (i + 1) % 10 == 0:
+            print(f"  {i+1}/{n} draft acc {ok_draft/(i+1):.3f} | refined acc {ok_final/(i+1):.3f}")
+    return {"n": n, "draft_acc": ok_draft / n, "refined_acc": ok_final / n}

dqwen/train.py

@@ -0,0 +1,242 @@
+"""Loading, LoRA setup, stage-1 (diffusion adaptation) and stage-2 (BART fusion) training loops."""
+from __future__ import annotations
+
+import math
+import os
+import time
+from dataclasses import dataclass, field
+from typing import Optional
+
+import torch
+
+from .core import (DiffusionLM, diffusion_generate, find_text_model, load_rev_gates, masked_diffusion_loss,
+                   patch_bidirectional, rev_gate_state)
+from .data import batches
+
+LORA_TARGETS = ["q_proj", "k_proj", "v_proj", "o_proj",            # full-attention layers
+                "in_proj_qkv", "in_proj_z", "out_proj",            # Gated-DeltaNet layers
+                "gate_proj", "up_proj", "down_proj"]               # MLP
+
+
+@dataclass
+class Stage1Cfg:
+    max_steps: int = 300
+    batch_size: int = 2            # micro batch (sequences of max_len tokens)
+    grad_accum: int = 8
+    lr: float = 2e-4               # LoRA
+    gate_lr: float = 2e-3          # DeltaNet reverse gates
+    mask_lr: float = 1e-3          # [MASK] embedding
+    warmup: int = 20
+    min_lr_ratio: float = 0.1
+    weight_decay: float = 0.0
+    clip: float = 1.0
+    t_min: float = 0.05
+    weighting: str = "elbo"        # or "uniform"
+    log_every: int = 5
+    sample_every: int = 100
+    save_every: int = 100
+    save_dir: str = "/content/diffusion_qwen_bart"
+    seed: int = 0
+
+
+def load_dlm(model_id: str = "Qwen/Qwen3.5-9B", device: str = "cuda", dtype=torch.bfloat16,
+             load_in_4bit: bool = False, lora_r: int = 32, lora_alpha: int = 64, lora_dropout: float = 0.0,
+             mask_token: str = "<|fim_pad|>", grad_ckpt: bool = True, deltanet_bidirectional: bool = True,
+             lm: Optional[torch.nn.Module] = None, tok=None):
+    """Returns (dlm, tok). `lm`/`tok` can be passed in for tests (skips from_pretrained)."""
+    from peft import LoraConfig, get_peft_model
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+
+    if tok is None:
+        tok = AutoTokenizer.from_pretrained(model_id)
+    if lm is None:
+        kw = dict(dtype=dtype, attn_implementation="sdpa")
+        if load_in_4bit:
+            from transformers import BitsAndBytesConfig
+            kw["quantization_config"] = BitsAndBytesConfig(
+                load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_use_double_quant=True,
+                bnb_4bit_compute_dtype=dtype)
+        kw["device_map"] = {"": 0} if device == "cuda" else None
+        lm = AutoModelForCausalLM.from_pretrained(model_id, **kw)
+    lm.config.use_cache = False
+
+    stats = patch_bidirectional(lm, deltanet_bidirectional)
+    print("[patch]", stats)
+    if grad_ckpt:
+        lm.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
+
+    cfg = LoraConfig(r=lora_r, lora_alpha=lora_alpha, lora_dropout=lora_dropout, target_modules=LORA_TARGETS,
+                     bias="none")
+    lm = get_peft_model(lm, cfg)
+    # PEFT freezes non-LoRA weights. Reactivate and verify reverse gates.
+    for n, p in lm.named_parameters():
+        if n.endswith("rev_gate"):
+            p.requires_grad_(True)
+            if p.device.type != torch.device(device).type:
+                raise RuntimeError(f"{n} is on {p.device}, expected {device}; rerun model loading with patched core.py")
+
+    mask_id = tok.convert_tokens_to_ids(mask_token)
+    dlm = DiffusionLM(lm, mask_id=mask_id)
+    dlm.mask_embed.requires_grad_(True)
+    if device == "cuda":
+        dlm.mask_embed.data = dlm.mask_embed.data.cuda()
+    n_tr = sum(p.numel() for p in dlm.parameters() if p.requires_grad)
+    print(f"[model] trainable params: {n_tr/1e6:.1f}M | mask token {mask_token!r} id={mask_id}")
+    return dlm, tok
+
+
+def _groups(dlm, cfg: Stage1Cfg):
+    lora, gates, other = [], [], []
+    for n, p in dlm.named_parameters():
+        if not p.requires_grad:
+            continue
+        (gates if n.endswith("rev_gate") else other if n.endswith("mask_embed") else lora).append(p)
+    return [
+        {"params": lora, "lr": cfg.lr, "base": cfg.lr},
+        {"params": gates, "lr": cfg.gate_lr, "base": cfg.gate_lr},
+        {"params": other, "lr": cfg.mask_lr, "base": cfg.mask_lr},
+    ]
+
+
+def _lr_at(step, cfg: Stage1Cfg):
+    if step < cfg.warmup:
+        return (step + 1) / cfg.warmup
+    pr = (step - cfg.warmup) / max(1, cfg.max_steps - cfg.warmup)
+    return cfg.min_lr_ratio + (1 - cfg.min_lr_ratio) * 0.5 * (1 + math.cos(math.pi * min(1.0, pr)))
+
+
+def save_checkpoint(dlm: DiffusionLM, tok, save_dir: str, step: int, fusion=None):
+    os.makedirs(save_dir, exist_ok=True)
+    dlm.lm.save_pretrained(os.path.join(save_dir, "qwen_lora"))
+    tok.save_pretrained(os.path.join(save_dir, "qwen_lora"))
+    torch.save({"mask_embed": dlm.mask_embed.detach().cpu(), "rev_gates": rev_gate_state(dlm.lm),
+                "mask_id": dlm.mask_id, "step": step}, os.path.join(save_dir, "extras.pt"))
+    if fusion is not None:
+        torch.save(fusion.state_dict(), os.path.join(save_dir, "fusion.pt"))
+    print(f"[ckpt] saved step {step} -> {save_dir}")
+
+
+def load_checkpoint(dlm: DiffusionLM, save_dir: str, fusion=None):
+    from peft import set_peft_model_state_dict
+    from safetensors.torch import load_file
+    sd = load_file(os.path.join(save_dir, "qwen_lora", "adapter_model.safetensors"))
+    set_peft_model_state_dict(dlm.lm, sd)
+    ex = torch.load(os.path.join(save_dir, "extras.pt"), map_location="cpu")
+    dlm.mask_embed.data.copy_(ex["mask_embed"].to(dlm.mask_embed.device))
+    missing = load_rev_gates(dlm.lm, ex["rev_gates"])
+    if missing:
+        print("[ckpt] WARNING rev_gate keys not matched:", missing[:3])
+    fp = os.path.join(save_dir, "fusion.pt")
+    if fusion is not None and os.path.exists(fp):
+        fusion.load_state_dict(torch.load(fp, map_location="cpu"))
+    return ex.get("step", 0)
+
+
+@torch.no_grad()
+def quick_sample(dlm, tok, prompt: str, eos_id: int, device, gen_len=128, steps=48, block_len=32, chat_kwargs=None):
+    was = dlm.training
+    dlm.eval()
+    text = tok.apply_chat_template([{"role": "user", "content": prompt}], tokenize=False,
+                                   add_generation_prompt=True, **(chat_kwargs or {"enable_thinking": False}))
+    p = torch.tensor([tok(text, add_special_tokens=False)["input_ids"]], device=device)
+    out = diffusion_generate(dlm, p, gen_len=gen_len, steps=steps, block_len=block_len, eos_id=eos_id)
+    ids = out[0].tolist()
+    if eos_id in ids:
+        ids = ids[:ids.index(eos_id)]
+    dlm.train(was)
+    return tok.decode(ids, skip_special_tokens=True)
+
+
+def train_stage1(dlm: DiffusionLM, tok, ds: dict, cfg: Stage1Cfg, device="cuda", sample_prompts=None,
+                 eos_token="<|im_end|>", resume: bool = False):
+    torch.manual_seed(cfg.seed)
+    eos_id = tok.convert_tokens_to_ids(eos_token)
+    opt = torch.optim.AdamW(_groups(dlm, cfg), weight_decay=cfg.weight_decay, betas=(0.9, 0.99))
+    start = 0
+    if resume and os.path.exists(os.path.join(cfg.save_dir, "extras.pt")):
+        start = load_checkpoint(dlm, cfg.save_dir)
+        print(f"[resume] from step {start}")
+    dlm.train()
+    it = batches(ds, cfg.batch_size, device, seed=cfg.seed + start)
+    t0, tok_count, ema = time.time(), 0, None
+    for step in range(start, cfg.max_steps):
+        for g in opt.param_groups:
+            g["lr"] = g["base"] * _lr_at(step, cfg)
+        tot = 0.0
+        for _ in range(cfg.grad_accum):
+            ids, resp, _plen = next(it)
+            loss, info = masked_diffusion_loss(dlm, ids, resp, t_min=cfg.t_min, weighting=cfg.weighting)
+            (loss / cfg.grad_accum).backward()
+            tot += float(loss.detach()) / cfg.grad_accum
+            tok_count += ids.numel()
+        gn = torch.nn.utils.clip_grad_norm_([p for g in opt.param_groups for p in g["params"]], cfg.clip)
+        opt.step(); opt.zero_grad(set_to_none=True)
+        ema = tot if ema is None else 0.9 * ema + 0.1 * tot
+        if (step + 1) % cfg.log_every == 0 or step == start:
+            dt = time.time() - t0
+            print(f"step {step+1}/{cfg.max_steps} loss {tot:.4f} (ema {ema:.4f}) gnorm {float(gn):.2f} "
+                  f"lr {opt.param_groups[0]['lr']:.2e} | {tok_count/dt:,.0f} tok/s")
+        if sample_prompts and (step + 1) % cfg.sample_every == 0:
+            for sp in sample_prompts[:2]:
+                print(f"  [sample] {sp!r} -> {quick_sample(dlm, tok, sp, eos_id, device)!r}")
+        if (step + 1) % cfg.save_every == 0 or step + 1 == cfg.max_steps:
+            save_checkpoint(dlm, tok, cfg.save_dir, step + 1)
+    return dlm
+
+
+# ------------------------------------------------------------------------------------------
+@dataclass
+class Stage2Cfg:
+    max_steps: int = 200
+    batch_size: int = 2
+    grad_accum: int = 4
+    lr: float = 5e-4
+    gate_lr: float = 3e-3
+    warmup: int = 10
+    clip: float = 1.0
+    log_every: int = 5
+    save_every: int = 100
+    save_dir: str = "/content/diffusion_qwen_bart"
+    seed: int = 0
+
+
+def train_stage2_fusion(dlm: DiffusionLM, tok, bart, btok, fusion, ds: dict, cfg: Stage2Cfg, device="cuda",
+                        eos_token="<|im_end|>"):
+    """Only the fusion module is trained; Qwen (LoRA) and BART are frozen."""
+    from .bart_fusion import refine_loss
+    torch.manual_seed(cfg.seed)
+    eos_id = tok.convert_tokens_to_ids(eos_token)
+    for p in dlm.parameters():
+        p.requires_grad_(False)
+    for p in bart.parameters():
+        p.requires_grad_(False)
+    bart.eval()
+    fusion.train()
+    for p in fusion.parameters():
+        p.requires_grad_(True)
+    gate = [fusion.gate]
+    rest = [p for n, p in fusion.named_parameters() if n != "gate"]
+    opt = torch.optim.AdamW([{"params": rest, "lr": cfg.lr, "base": cfg.lr},
+                             {"params": gate, "lr": cfg.gate_lr, "base": cfg.gate_lr}], weight_decay=0.0)
+    dlm.train()
+    it = batches(ds, cfg.batch_size, device, seed=cfg.seed)
+    t0 = time.time()
+    for step in range(cfg.max_steps):
+        sc = min(1.0, (step + 1) / cfg.warmup) * (0.5 * (1 + math.cos(math.pi * step / cfg.max_steps)) * 0.9 + 0.1)
+        for g in opt.param_groups:
+            g["lr"] = g["base"] * sc
+        tot = 0.0
+        for _ in range(cfg.grad_accum):
+            ids, resp, plen = next(it)
+            loss, info = refine_loss(dlm, fusion, bart, btok, tok, ids, resp, plen, eos_id, device)
+            (loss / cfg.grad_accum).backward()
+            tot += float(loss.detach()) / cfg.grad_accum
+        gn = torch.nn.utils.clip_grad_norm_(fusion.parameters(), cfg.clip)
+        opt.step(); opt.zero_grad(set_to_none=True)
+        if (step + 1) % cfg.log_every == 0 or step == 0:
+            print(f"[fusion] step {step+1}/{cfg.max_steps} loss {tot:.4f} gnorm {float(gn):.2f} "
+                  f"gate {float(torch.tanh(fusion.gate)):+.4f} | {time.time()-t0:.0f}s")
+        if (step + 1) % cfg.save_every == 0 or step + 1 == cfg.max_steps:
+            os.makedirs(cfg.save_dir, exist_ok=True)
+            torch.save(fusion.state_dict(), os.path.join(cfg.save_dir, "fusion.pt"))
+    return fusion