Upload dblock training helper

dblocks_train.py

@@ -0,0 +1,432 @@
+"""DiffusionBlocks training mode folded into AGILLM-4 (gated by --dblock).
+
+Block-wise EDM denoising on the real Encoder blocks, supervising AR + SAT(fixed+var)
++ NAT each step on ONE block, with grad-checkpointed layers and fused vocab-streaming
+CE. Reuses the live data stream / optimizer / checkpointing of nB300_agillm4.
+Lazy-imports nB300 inside functions to avoid a circular import.
+"""
+import math
+import random
+import time
+from collections import defaultdict
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.checkpoint as _ck
+from fused_ce import fused_ce
+
+SD = 0.5
+
+
+
+
+def _profile_active(state, args):
+    limit = int(getattr(args, "profile_steps", 0) or 0)
+    return limit > 0 and int(state.get("profile_n", 0)) < limit
+
+
+def _profile_add(state, name, seconds):
+    if seconds is None:
+        return
+    prof = state.setdefault("profile_times", defaultdict(float))
+    prof[name] += float(seconds)
+
+
+def _profile_tic(enabled):
+    if not enabled:
+        return None
+    if torch.cuda.is_available():
+        torch.cuda.synchronize()
+    return time.perf_counter()
+
+
+def _profile_toc(state, name, start):
+    if start is None:
+        return
+    if torch.cuda.is_available():
+        torch.cuda.synchronize()
+    _profile_add(state, name, time.perf_counter() - start)
+
+
+def _profile_step_done(state, args):
+    limit = int(getattr(args, "profile_steps", 0) or 0)
+    if limit <= 0:
+        return
+    n_prev = int(state.get("profile_n", 0))
+    if n_prev >= limit:
+        return
+    state["profile_n"] = n_prev + 1
+    n = int(state["profile_n"])
+    log_every = max(1, int(getattr(args, "profile_log_every", 25) or 25))
+    if n % log_every != 0 and n != limit:
+        return
+    times = state.get("profile_times", {})
+    keys = [
+        "data_stream", "tensor", "setup",
+        "ar_forward", "ar_ce", "ar_backward",
+        "sat_forward", "sat_ce", "sat_backward",
+        "nat_forward", "nat_ce", "nat_backward",
+        "opt_step", "step_total",
+    ]
+    parts = []
+    for key in keys:
+        val = float(times.get(key, 0.0)) * 1000.0 / max(1, n)
+        if val > 0.01:
+            parts.append(f"{key}={val:.2f}ms")
+    print(f"[profile] n={n}/{limit} avg " + " ".join(parts), flush=True)
+
+def _cdf(x):
+    return 0.5 * (1 + math.erf(x / math.sqrt(2)))
+
+
+def _ppf(p):
+    return float(torch.erfinv(torch.tensor(2 * p - 1.0)) * math.sqrt(2))
+
+
+def _block_sigmas(B, smin=0.002, smax=80.0, pm=-1.2, ps=1.2):
+    a, b = _cdf((math.log(smin) - pm) / ps), _cdf((math.log(smax) - pm) / ps)
+    return [float(np.exp(pm + ps * _ppf(a + (b - a) * (i / B)))) for i in range(B + 1)]
+
+
+def _edm_pre(s):
+    s = s[:, None, None]
+    return SD**2 / (s**2 + SD**2), s * SD / (s**2 + SD**2) ** 0.5, 1 / (s**2 + SD**2) ** 0.5
+
+
+def _edm_w(s, wmax=5.0):
+    return float(((s**2 + SD**2) / (s * SD) ** 2).clamp(max=wmax).mean())
+
+
+def _dblock_init(core, args):
+    B = int(getattr(args, "dblock_blocks", 4))
+    L = len(core.blocks)
+    sp = max(1, L // B)
+    asg = [list(range(i * sp, (i + 1) * sp)) for i in range(B)]
+    asg[-1] = list(range((B - 1) * sp, L))
+    bsig = _block_sigmas(B)
+    schedule = getattr(args, "dblock_schedule", "loss_balanced")
+    print(f"[dblock] DiffusionBlocks mode: {L} layers -> {B} blocks {asg}")
+    print(f"[dblock] schedule={schedule} sigma boundaries: {[round(x, 3) for x in bsig]}")
+    return {
+        "B": B,
+        "assign": asg,
+        "bsig": bsig,
+        "step": 0,
+        "counts": [0 for _ in range(B)],
+        "loss_ema": [None for _ in range(B)],
+    }
+
+
+def _choose_block(state, args):
+    B = state["B"]
+    schedule = str(getattr(args, "dblock_schedule", "loss_balanced") or "loss_balanced").lower()
+    step = int(state.get("step", 0))
+    counts = state.setdefault("counts", [0 for _ in range(B)])
+    emas = state.setdefault("loss_ema", [None for _ in range(B)])
+    if schedule == "random":
+        return random.randrange(B)
+    if schedule == "roundrobin":
+        return step % B
+    explore = float(getattr(args, "dblock_explore", 0.05))
+    warmup = int(getattr(args, "dblock_warmup_steps", max(8, B * 2)))
+    if step < warmup or any(c == 0 for c in counts):
+        return min(range(B), key=lambda i: (counts[i], i))
+    if explore > 0.0 and random.random() < explore:
+        return min(range(B), key=lambda i: (counts[i], i))
+    return max(range(B), key=lambda i: (-1.0 if emas[i] is None else emas[i], -counts[i]))
+
+
+def _sample_sigma(ids, lo, hi, args, state):
+    cur_step = int(state.get("step", 0))
+    curriculum = int(getattr(args, "dblock_sigma_curriculum_steps", 0))
+    if curriculum > 0:
+        frac = min(1.0, max(0.05, (cur_step + 1) / float(curriculum)))
+        hi = lo * ((hi / max(lo, 1e-8)) ** frac)
+    sig_np = np.exp(
+        np.random.uniform(
+            math.log(max(lo, 1e-4)),
+            math.log(max(hi, lo + 1e-4)),
+            ids.size(0),
+        ).astype("float32")
+    )
+    return torch.from_numpy(sig_np).to(ids.device)
+
+
+def _maybe_log(state, args, bi, layers, ar_val, sat_val, nat_val, total_val, peak_alloc, peak_reserved, objective=None):
+    log_every = int(getattr(args, "dblock_log_every", 50))
+    step = int(state.get("step", 0))
+    if log_every <= 0 or step % log_every != 0:
+        return
+    counts = ",".join(str(x) for x in state.get("counts", []))
+    emas = ",".join("nan" if x is None else f"{x:.2f}" for x in state.get("loss_ema", []))
+    mem = ""
+    if peak_alloc is not None:
+        mem = f" peak_alloc={peak_alloc:.2f}GB peak_reserved={peak_reserved:.2f}GB"
+    print(
+        f"[dblock] step={step} block={bi} obj={objective or 'mixed'} layers={layers} "
+        f"loss={total_val:.3f} ar={ar_val:.3f} sat={sat_val:.3f} nat={nat_val:.3f} "
+        f"counts=[{counts}] ema=[{emas}]{mem}",
+        flush=True,
+    )
+
+
+def _update_stats(state, bi, loss_value):
+    B = state["B"]
+    counts = state.setdefault("counts", [0 for _ in range(B)])
+    emas = state.setdefault("loss_ema", [None for _ in range(B)])
+    counts[bi] += 1
+    prev = emas[bi]
+    beta = 0.96
+    emas[bi] = float(loss_value) if prev is None else beta * float(prev) + (1.0 - beta) * float(loss_value)
+    state["step"] = int(state.get("step", 0)) + 1
+
+
+def _activation_offload_enabled(args):
+    return bool(getattr(args, "dblock_activation_offload", False)) and torch.cuda.is_available()
+
+
+def _activation_offload_hooks(args):
+    min_bytes = int(float(getattr(args, "dblock_activation_offload_min_mb", 1.0) or 1.0) * 1024 * 1024)
+
+    def pack(t):
+        if not torch.is_tensor(t) or not t.is_cuda or not t.is_floating_point() or t.numel() * t.element_size() < min_bytes:
+            return t
+        return ("cpu_offload", t.device, t.detach().to("cpu", non_blocking=True))
+
+    def unpack(x):
+        if isinstance(x, tuple) and len(x) == 3 and x[0] == "cpu_offload":
+            _, dev, cpu_t = x
+            return cpu_t.to(dev, non_blocking=True)
+        return x
+
+    return torch.autograd.graph.saved_tensors_hooks(pack, unpack)
+
+
+def _run_block(block, x, mask, use_checkpoint, args=None):
+    if use_checkpoint:
+        return _ck.checkpoint(lambda y, block=block: block(y, mask), x, use_reentrant=False)
+    if args is not None and _activation_offload_enabled(args):
+        with _activation_offload_hooks(args):
+            return block(x, mask)
+    return block(x, mask)
+
+
+def _dblock_checkpoint_this_layer(args, base_enabled, layer_pos, layer_count=None):
+    if not base_enabled:
+        return False
+    pos = int(layer_pos)
+    count = int(layer_count or 0)
+    skip_tail = max(0, int(getattr(args, "dblock_checkpoint_skip_tail", 0) or 0))
+    if skip_tail > 0 and count > 0 and pos >= max(0, count - skip_tail):
+        return False
+    stride = int(getattr(args, "dblock_checkpoint_stride", 1) or 1)
+    if stride <= 0:
+        return False
+    if stride == 1:
+        return True
+    return (pos % stride) == 0
+
+
+def _sample_token_loss_inputs(hidden, targets, max_tokens):
+    max_tokens = int(max_tokens or 0)
+    if max_tokens <= 0:
+        return hidden.contiguous(), targets.contiguous(), int(targets.numel()), int(targets.numel())
+    flat_targets = targets.reshape(-1)
+    total = int(flat_targets.numel())
+    if total <= max_tokens:
+        return hidden.contiguous(), targets.contiguous(), total, total
+    # With-replacement sampling avoids building a full randperm each step; the sampled
+    # mean remains an unbiased estimator of the dense token CE mean.
+    idx = torch.randint(total, (max_tokens,), device=targets.device)
+    flat_hidden = hidden.reshape(total, hidden.size(-1))
+    return flat_hidden.index_select(0, idx).contiguous(), flat_targets.index_select(0, idx).contiguous(), int(max_tokens), total
+
+
+def _choose_objectives(state, args, ar_weight, sat_weight, nat_weight, do_sat_periodic, do_nat_periodic):
+    mode = str(getattr(args, "dblock_objective_mode", "periodic") or "periodic").lower()
+    if mode != "stochastic":
+        return ar_weight > 0.0, sat_weight > 0.0 and do_sat_periodic, nat_weight > 0.0 and do_nat_periodic, "periodic"
+    choices = []
+    probs = []
+    if ar_weight > 0.0:
+        choices.append("ar")
+        probs.append(max(0.0, float(getattr(args, "dblock_ar_prob", 0.80))))
+    if sat_weight > 0.0 and not getattr(args, "ar_only", False):
+        choices.append("sat")
+        probs.append(max(0.0, float(getattr(args, "dblock_sat_prob", 0.10))))
+    if nat_weight > 0.0 and not getattr(args, "ar_only", False):
+        choices.append("nat")
+        probs.append(max(0.0, float(getattr(args, "dblock_nat_prob", 0.10))))
+    if not choices:
+        return False, False, False, "none"
+    total = sum(probs)
+    if total <= 0.0:
+        probs = [1.0 / len(choices) for _ in choices]
+    else:
+        probs = [p / total for p in probs]
+    picked = random.choices(choices, weights=probs, k=1)[0]
+    return picked == "ar", picked == "sat", picked == "nat", picked
+
+
+def _dblock_step(core, ar_h, sat_h, nat_h, opt, scaler, args, ids, state):
+    import nB300_agillm4 as M
+
+    prof = _profile_active(state, args)
+    _step_t = _profile_tic(prof)
+    if torch.cuda.is_available():
+        torch.cuda.reset_peak_memory_stats()
+
+    _setup_t = _profile_tic(prof)
+    B = state["B"]
+    asg = state["assign"]
+    bs = state["bsig"]
+    T = ids.size(1)
+    use_layer_checkpoint = bool(getattr(args, "grad_checkpoint", False))
+    bi = _choose_block(state, args)
+    lo, hi = sorted([bs[bi], bs[bi + 1]])
+    layers = asg[bi]
+    sig = _sample_sigma(ids, lo, hi, args, state)
+    cs, co, ci = _edm_pre(sig)
+    w = _edm_w(sig, float(getattr(args, "dblock_edm_wmax", 5.0)))
+    SATB = M.SAT_BLOCK
+    ar_weight = float(getattr(args, "dblock_ar_weight", 1.0))
+    sat_weight = float(getattr(args, "dblock_sat_weight", 1.0))
+    nat_weight = float(getattr(args, "dblock_nat_weight", 1.0)) * float(getattr(args, "nat_loss_weight", 1.0))
+    do_sat_periodic = (not getattr(args, "ar_only", False)) and (
+        int(getattr(args, "sat_every", 1)) <= 1 or ((int(state.get("step", 0)) + 1) % int(getattr(args, "sat_every", 1)) == 0)
+    )
+    do_nat_periodic = (
+        nat_h is not None
+        and (not getattr(args, "ar_only", False))
+        and int(getattr(args, "nat_every", 1)) > 0
+        and (
+            int(getattr(args, "nat_every", 1)) <= 1
+            or ((int(state.get("step", 0)) + 1) % int(getattr(args, "nat_every", 1)) == 0)
+        )
+    )
+    run_ar, run_sat, run_nat, objective = _choose_objectives(
+        state, args, ar_weight, sat_weight, nat_weight, do_sat_periodic, do_nat_periodic
+    )
+    _profile_toc(state, "setup", _setup_t)
+
+    ar_val = 0.0
+    sat_val = 0.0
+    nat_val = 0.0
+
+    if run_ar:
+        causal = M.causal_mask(T, structured=M.use_structured_masks(args))
+        _t = _profile_tic(prof)
+        with M.amp(args.amp):
+            emb = core.emb(ids)
+            zt = emb + sig[:, None, None] * torch.randn_like(emb)
+            h = ci * zt
+            for lpos, li in enumerate(layers):
+                h = _run_block(core.blocks[li], h, causal, _dblock_checkpoint_this_layer(args, use_layer_checkpoint, lpos, len(layers)), args)
+            Dn = core.ln(cs * zt + co * h)
+        _profile_toc(state, "ar_forward", _t)
+        _t = _profile_tic(prof)
+        ar_hidden, ar_targets, ar_used, ar_total = _sample_token_loss_inputs(
+            Dn[:, :-1], ids[:, 1:], int(getattr(args, "dblock_ar_loss_tokens", 0))
+        )
+        ar = ar_weight * w * fused_ce(ar_hidden, ar_h.proj.weight, ar_targets)
+        ar_val = float(ar.detach())
+        _profile_toc(state, "ar_ce", _t)
+        _t = _profile_tic(prof)
+        scaler.scale(ar).backward()
+        _profile_toc(state, "ar_backward", _t)
+        del causal, emb, zt, h, Dn, ar_hidden, ar_targets, ar, ar_used, ar_total
+
+    if run_sat:
+        smask = M.sat_mask(T, structured=M.use_structured_masks(args))
+        _t = _profile_tic(prof)
+        with M.amp(args.amp):
+            emb2 = core.emb(ids)
+            zt2 = emb2 + sig[:, None, None] * torch.randn_like(emb2)
+            h2 = ci * zt2
+            for lpos, li in enumerate(layers):
+                h2 = _run_block(core.blocks[li], h2, smask, _dblock_checkpoint_this_layer(args, use_layer_checkpoint, lpos, len(layers)), args)
+            Ds = core.ln(cs * zt2 + co * h2)
+            last = Ds[:, -SATB:]
+        _profile_toc(state, "sat_forward", _t)
+        _t = _profile_tic(prof)
+        sat_hidden, sat_targets, sat_used, sat_total = _sample_token_loss_inputs(
+            last, ids[:, 1 : SATB + 1], int(getattr(args, "dblock_sat_loss_tokens", 0))
+        )
+        with M.amp(args.amp):
+            satf = fused_ce(sat_hidden, sat_h.proj.weight, sat_targets)
+            satv = (
+                M.EMIT_LAMBDA
+                * F.cross_entropy(
+                    sat_h.gate(Ds[:, 0].float()),
+                    torch.ones(ids.size(0), dtype=torch.long, device=ids.device),
+                )
+                if sat_h.gate is not None
+                else 0.0
+            )
+            sat = sat_weight * w * (satf + satv)
+        _profile_toc(state, "sat_ce", _t)
+        sat_val = float(sat.detach())
+        _t = _profile_tic(prof)
+        scaler.scale(sat).backward()
+        _profile_toc(state, "sat_backward", _t)
+        del smask, emb2, zt2, h2, Ds, last, sat_hidden, sat_targets, satf, satv, sat
+
+    if run_nat:
+        ratio = min(max(float(getattr(args, "nat_mask_ratio", 0.5)), 0.05), 0.95)
+        nat_ids = M._nat_ids_for_training(ids, int(getattr(args, "nat_max_tokens", 0)))
+        _t = _profile_tic(prof)
+        with M.amp(args.amp):
+            nat_in = nat_ids.clone()
+            m = torch.rand(nat_ids.shape, device=nat_ids.device) < ratio
+            if not bool(m.any()):
+                m[..., -1] = True
+            nat_in[m] = M.BLANK
+            hn = core.emb(nat_in)
+            for lpos, li in enumerate(layers):
+                hn = _run_block(core.blocks[li], hn, None, _dblock_checkpoint_this_layer(args, use_layer_checkpoint, lpos, len(layers)), args)
+            Dnat = core.ln(hn)
+        _profile_toc(state, "nat_forward", _t)
+        _t = _profile_tic(prof)
+        nat_hidden = Dnat[m]
+        nat_targets = nat_ids[m]
+        nat_hidden, nat_targets, nat_used, nat_total = _sample_token_loss_inputs(
+            nat_hidden.unsqueeze(0), nat_targets.unsqueeze(0), int(getattr(args, "dblock_nat_loss_tokens", 0))
+        )
+        nat = nat_weight * fused_ce(nat_hidden, nat_h.proj.weight, nat_targets)
+        nat_val = float(nat.detach())
+        _profile_toc(state, "nat_ce", _t)
+        _t = _profile_tic(prof)
+        scaler.scale(nat).backward()
+        _profile_toc(state, "nat_backward", _t)
+        del nat_ids, nat_in, m, hn, Dnat, nat_hidden, nat_targets, nat, nat_used, nat_total
+
+    total_val = ar_val + sat_val + nat_val
+    if not math.isfinite(total_val):
+        opt.zero_grad(set_to_none=True)
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        print(f"[dblock] non-finite loss {total_val}; skipped optimizer step", flush=True)
+        _profile_toc(state, "step_total", _step_t)
+        _profile_step_done(state, args)
+        _update_stats(state, bi, total_val)
+        return total_val
+
+    _t = _profile_tic(prof)
+    scaler.unscale_(opt)
+    nn.utils.clip_grad_norm_([p for g in opt.param_groups for p in g["params"]], 1.0)
+    scaler.step(opt)
+    scaler.update()
+    opt.zero_grad(set_to_none=True)
+    _profile_toc(state, "opt_step", _t)
+
+    peak_alloc = None
+    peak_reserved = None
+    if torch.cuda.is_available():
+        peak_alloc = torch.cuda.max_memory_allocated() / (1024**3)
+        peak_reserved = torch.cuda.max_memory_reserved() / (1024**3)
+    _profile_toc(state, "step_total", _step_t)
+    _profile_step_done(state, args)
+    _update_stats(state, bi, total_val)
+    _maybe_log(state, args, bi, layers, ar_val, sat_val, nat_val, total_val, peak_alloc, peak_reserved, objective=objective)
+    return total_val