IDM method for rapid_moss (data-space): DMD direction + L_pop moment matching, no adversarial

training/rapid_moss_train_idm.py

@@ -0,0 +1,577 @@
+"""Distillation of MOSS-Image-SFT to an R-step data-space generator (arXiv 2610.12102),
+with method='idm' added by akhaliq (2026-10-31): IDM-style training (arXiv 2610.39633).
+
+Based on darask0's rapid_moss/train.py (Apache-2.0,
+hf.co/darask0/MOSS-Image_4-step_adapter_lora). The 'dmd' and 'dmad' methods are unchanged.
+
+method='idm': no discriminator, no critic. Per iteration:
+  fake_per_gen fake-score updates (flow matching on generated x0, conditioned on r)  [IDM L_fake]
+  one generator update with the DMD direction sg[v_real - v_fake] on x0             [IDM L_score]
+  plus L_pop: EMA feature moment matching (mean/std of last-block image-token features,
+  current sample + FIFO queue vs a reference queue built from teachergen samples)   [IDM L_pop]
+Deviation from the paper: the base loss is the data-space DMD fixed-point iteration
+(arXiv 2610.12102) rather than RCGM; features come from the EMA generator path as in the paper.
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import os
+import random
+import time
+from collections import deque
+from pathlib import Path
+
+import torch
+import torch.distributed as dist
+from torch.nn import functional as F
+
+from .model import load_base
+from .nets import Critic, Discriminator, FakeScore, Generator
+from .textenc import Cache, read_prompts
+
+DEFAULTS = {
+    'model_dir': '/workspace/models/MOSS-Image-SFT',
+    'lodestones_dir': '/workspace/models/MOSS-Image-SFT-8-steps-dmd',
+    'cache_dir': '/workspace/cache',
+    'prompts': '/workspace/prompts/train.jsonl',
+    'eval_prompts': '/workspace/prompts/eval.jsonl',
+    'out_dir': '/workspace/runs/main',
+    'steps': 4,
+    'gen_rank': 128, 'gen_alpha': 128,
+    'fake_rank': 64, 'fake_alpha': 64,
+    'disc_taps': [5, 11, 17, 23, 29, 35], 'disc_hidden': 256,
+    'iterations': 3000,
+    'micro': 2,
+    'fake_per_gen': 2,
+    'fake_warmup': 200,
+    'lr_gen': 1e-5, 'lr_time': 1e-5, 'lr_q': 1e-5, 'lr_fake': 1e-4, 'lr_disc': 5e-5,
+    'betas': [0.0, 0.99],
+    'ema': 0.999,
+    'cfg': [3.0, 3.0],
+    'skip_neg': [22],
+    'tau_beta': [2.5, 1.0], 'tau_clamp': [0.03, 0.985],
+    'dt': 1 / 64,
+    'adv_weight': 0.5, 'adv_start': 50, 'adv_warmup': 200,
+    'disc_real': 'adm', 'teacher_dir': '/workspace/teacher', 'p_gan': 0.5, 'gan_t': [0.02, 0.98],
+    'clip_gen': 0.0, 'clip_critic': 0.0,
+    'buckets': [[1024, 1024, 0.35], [896, 1152, 0.15], [832, 1216, 0.15], [1152, 896, 0.1],
+                [1216, 832, 0.1], [768, 1344, 0.075], [1344, 768, 0.075]],
+    'method': 'dmd',
+    'critic_rank': 128, 'critic_alpha': 128, 'critic_start': 21, 'critic_hidden': 512,
+    'critic_shift': 5.0, 'critic_sigma': [0.02, 0.98], 'lr_critic': 4e-5, 'wd_critic': 0.01,
+    'critic_warmup': 50,
+    # IDM (L_pop) settings
+    'lpop_weight': 0.5, 'lpop_queue': 64, 'lpop_sigma': [0.5, 0.98],
+    'ckpt': True,
+    'seed': 0,
+    'log_every': 10, 'sample_every': 100, 'save_every': 250, 'export_every': 500,
+    'eval_count': 8, 'eval_seed': 1234,
+}
+
+
+def parse_config():
+    ap = argparse.ArgumentParser()
+    ap.add_argument('--config')
+    ap.add_argument('--set', nargs='*', default=[])
+    args = ap.parse_args()
+    cfg = dict(DEFAULTS)
+    if args.config:
+        cfg.update(json.loads(Path(args.config).read_text()))
+    for item in args.set:
+        key, value = item.split('=', 1)
+        if key not in DEFAULTS:
+            raise SystemExit(f'unknown key {key}')
+        try:
+            cfg[key] = json.loads(value)
+        except json.JSONDecodeError:
+            cfg[key] = value
+    return cfg
+
+
+class Dist:
+    def __init__(self):
+        self.on = 'RANK' in os.environ
+        if self.on:
+            dist.init_process_group('nccl')
+        self.rank = dist.get_rank() if self.on else 0
+        self.world = dist.get_world_size() if self.on else 1
+        self.device = torch.device('cuda', int(os.environ.get('LOCAL_RANK', 0)))
+        torch.cuda.set_device(self.device)
+
+    def mean_grads(self, params):
+        grads = [p.grad for p in params]
+        if self.world == 1:
+            return
+        flat = torch._utils._flatten_dense_tensors(grads)
+        dist.all_reduce(flat)
+        flat /= self.world
+        for g, f in zip(grads, torch._utils._unflatten_dense_tensors(flat, grads)):
+            g.copy_(f)
+
+    def mean(self, values):
+        t = torch.tensor(values, device=self.device, dtype=torch.float64)
+        if self.on:
+            dist.all_reduce(t)
+            t /= self.world
+        return t.tolist()
+
+    def barrier(self):
+        if self.on:
+            dist.barrier()
+
+
+def zero_grads(params):
+    for p in params:
+        if p.grad is None:
+            p.grad = torch.zeros_like(p)
+        else:
+            p.grad.zero_()
+
+
+def load_teacher(dirs, rows, dev):
+    """Scan teachergen shards; returns the prompt rows that have a teacher latent
+    (with w/h from the latent shape) and a getter id -> [1, C, h, w] fp32 on dev."""
+    from safetensors import safe_open
+    files, where = [], {}
+    for d in [dirs] if isinstance(dirs, str) else dirs:
+        for path in sorted(Path(d).glob('shard-*.safetensors')):
+            try:
+                f = safe_open(path, 'pt')
+            except Exception as e:
+                print(json.dumps({'skip_shard': str(path), 'error': str(e)[:120]}), flush=True)
+                continue
+            files.append(f)
+            for k in f.keys():
+                _, h, w = f.get_slice(k).get_shape()
+                where[k] = (len(files) - 1, w * 16, h * 16)
+    pool = [dict(r, w=where[r['id']][1], h=where[r['id']][2]) for r in rows if r['id'] in where]
+
+    def get(key):
+        return files[where[key][0]].get_tensor(key)[None].to(dev, torch.float32)
+    return pool, get
+
+
+def decode(vae, x0):
+    z = x0.to(vae.dtype)
+    bn, eps = vae.bn, vae.config.batch_norm_eps
+    z = z * (bn.running_var + eps).sqrt().view(1, -1, 1, 1) + bn.running_mean.view(1, -1, 1, 1)
+    b, c, h, w = z.shape
+    z = z.reshape(b, c // 4, 2, 2, h, w).permute(0, 1, 4, 2, 5, 3).reshape(b, c // 4, 2 * h, 2 * w)
+    img = vae.decode(z, return_dict=False)[0].clamp(-1, 1)
+    return ((img[0].float() + 1) * 127.5).round().byte().permute(1, 2, 0).cpu().numpy()
+
+
+def contact_sheet(images, path, cell=384, cols=4):
+    from PIL import Image
+    rows = math.ceil(len(images) / cols)
+    sheet = Image.new('RGB', (cols * cell, rows * cell), 'white')
+    for i, im in enumerate(images):
+        im = Image.fromarray(im)
+        im.thumbnail((cell, cell))
+        sheet.paste(im, ((i % cols) * cell + (cell - im.width) // 2, (i // cols) * cell + (cell - im.height) // 2))
+    sheet.save(path, quality=90)
+
+
+def main():
+    C = parse_config()
+    D = Dist()
+    dev = D.device
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+    out = Path(C['out_dir'])
+    if D.rank == 0:
+        (out / 'samples').mkdir(parents=True, exist_ok=True)
+        (out / 'ckpt').mkdir(exist_ok=True)
+        (out / 'config.json').write_text(json.dumps(C, indent=1))
+    R = C['steps']
+
+    model = load_base(C['model_dir'], dev)
+    gen = Generator(model, C['gen_rank'], C['gen_alpha'], R).to(dev)
+    gen.load_lodestones(C['lodestones_dir'])
+    ema = Generator(model, C['gen_rank'], C['gen_alpha'], R).to(dev).requires_grad_(False)
+    ema.load_state_dict(gen.state_dict())
+    betas = tuple(C['betas'])
+    gen_params = list(gen.parameters())
+    opt_gen = torch.optim.AdamW([{'params': list(gen.lora.parameters()), 'lr': C['lr_gen']},
+                                 {'params': list(gen.time.parameters()), 'lr': C['lr_time']},
+                                 {'params': [gen.Q], 'lr': C['lr_q']}], betas=betas, weight_decay=0.0)
+    parts = {'gen': gen, 'ema': ema, 'opt_gen': opt_gen}
+    dmad = C['method'] == 'dmad'
+    idm = C['method'] == 'idm'
+    if dmad:
+        critic = Critic(model, C['critic_rank'], C['critic_alpha'], C['critic_start'], C['critic_hidden']).to(dev)
+        opt_critic = torch.optim.AdamW(critic.parameters(), lr=C['lr_critic'], betas=betas,
+                                       weight_decay=C['wd_critic'])
+        parts.update(critic=critic, opt_critic=opt_critic)
+        trainables = (gen_params, list(critic.parameters()))
+    elif C['method'] == 'dmd':
+        fake = FakeScore(model, C['fake_rank'], C['fake_alpha'], R).to(dev)
+        disc = Discriminator(model.dim, C['disc_taps'], C['disc_hidden']).to(dev)
+        opt_fake = torch.optim.AdamW(fake.parameters(), lr=C['lr_fake'], betas=betas, weight_decay=0.0)
+        opt_disc = torch.optim.AdamW(disc.parameters(), lr=C['lr_disc'], betas=betas, weight_decay=0.0)
+        parts.update(fake=fake, disc=disc, opt_fake=opt_fake, opt_disc=opt_disc)
+        trainables = (gen_params, list(fake.parameters()), list(disc.parameters()))
+    elif idm:
+        fake = FakeScore(model, C['fake_rank'], C['fake_alpha'], R).to(dev)
+        opt_fake = torch.optim.AdamW(fake.parameters(), lr=C['lr_fake'], betas=betas, weight_decay=0.0)
+        parts.update(fake=fake, opt_fake=opt_fake)
+        trainables = (gen_params, list(fake.parameters()))
+    else:
+        raise SystemExit(f"method must be dmd, dmad or idm, got {C['method']}")
+    for params in trainables:
+        zero_grads(params)
+
+    cache = Cache(C['cache_dir'])
+    rows = [r for r in read_prompts([C['prompts']]) if r['id'] in cache]
+    evals = [r for r in read_prompts([C['eval_prompts']]) if r['id'] in cache][:C['eval_count']]
+    neg = model.context(cache.get('').to(dev))
+    buckets = C['buckets']
+    teacher = None
+    if dmad or idm or C['disc_real'] == 'teacher':
+        trows, teacher = load_teacher(C['teacher_dir'], rows, dev)
+        if D.rank == 0:
+            print(json.dumps({'teacher_samples': len(trows)}), flush=True)
+    elif C['disc_real'] != 'adm':
+        raise SystemExit(f"disc_real must be adm or teacher, got {C['disc_real']}")
+
+    start, gen_updates = 0, 0
+    latest = out / 'ckpt' / 'latest.pt'
+    if latest.exists():
+        state = torch.load(latest, map_location=dev, weights_only=False)
+        for key, mod in parts.items():
+            mod.load_state_dict(state[key])
+        start, gen_updates = state['it'], state['gen_updates']
+        del state
+        if D.rank == 0:
+            print(json.dumps({'resume': start, 'gen_updates': gen_updates}), flush=True)
+    rng = random.Random(C['seed'] * 1_000_003 + D.rank * 7919 + start)
+    torch.manual_seed(C['seed'] * 1_000_003 + D.rank * 7919 + start)
+    beta = torch.distributions.Beta(*[torch.tensor(float(v)) for v in C['tau_beta']])
+    lo, hi = C['tau_clamp']
+    dt = C['dt']
+    vae = None
+
+    def draw(pool=None):
+        pool = pool or rows
+        row = pool[rng.randrange(len(pool))]
+        if 'w' in row and 'h' in row:
+            w, h = row['w'], row['h']
+        else:
+            w, h, _ = rng.choices(buckets, weights=[b[2] for b in buckets])[0]
+        ctx = model.context(cache.get(row['id']).to(dev))
+        z = torch.randn(1, model.tr.config.latent_channels, h // 16, w // 16, device=dev)
+        return ctx, z, row
+
+    def tau():
+        return float(beta.sample().clamp(lo, hi))
+
+    def r_index(phase, m):
+        return ((phase * D.world + D.rank) * C['micro'] + m) % R
+
+    # ---- IDM L_pop reference: moment stats of EMA features over teachergen samples ----
+    ref_mean = ref_std = None
+    pop_q = deque(maxlen=C['lpop_queue'])
+    if idm:
+        last = model.depth - 1
+        ref_rows = rng.sample(trows, min(C['lpop_queue'], len(trows)))
+        feats = []
+        with torch.no_grad():
+            for row in ref_rows:
+                ctx = model.context(cache.get(row['id']).to(dev))
+                x_t = teacher(row['id'])
+                s = rng.uniform(*C['lpop_sigma'])
+                xt = (1 - s) * x_t + s * torch.randn_like(x_t)
+                f = model(xt, s * 1000, ctx, lora=ema.lora, tp=ema.time, taps=(last,))[last]
+                feats.append(f.float().mean(1)[0].cpu())
+        ref = torch.stack(feats)
+        ref_mean, ref_std = ref.mean(0).to(dev), ref.std(0).to(dev)
+        if D.rank == 0:
+            print(json.dumps({'lpop_reference': len(feats)}), flush=True)
+
+    def fake_step(phase):
+        losses = []
+        gen.lora.freeze_cache()
+        for m in range(C['micro']):
+            ctx, z, _ = draw()
+            r = r_index(phase, m)
+            x0 = gen.rollout(model, z, ctx, r + 1)
+            t = tau()
+            eps = torch.randn_like(z)
+            v = fake(model, (1 - t) * x0 + t * eps, t * 1000, ctx, r, ckpt=C['ckpt'])
+            loss = F.mse_loss(v, eps - x0)
+            (loss / C['micro']).backward()
+            losses.append(loss.item())
+        gen.lora.drop_cache()
+        D.mean_grads(list(fake.parameters()))
+        opt_fake.step()
+        opt_fake.zero_grad(set_to_none=False)
+        return sum(losses) / len(losses)
+
+    def gen_step_idm(phase):
+        """IDM generator step: DMD direction + L_pop moment matching. No adversarial losses."""
+        stats = {k: 0.0 for k in ('dmd_rms', 'lpop', 'lpop_std_gap')}
+        fake.lora.freeze_cache()
+        gen.lora.freeze_cache()
+        last = model.depth - 1
+        for m in range(C['micro']):
+            ctx, z, row = draw(trows)
+            r = r_index(phase, m)
+            x_prev = gen.rollout(model, z, ctx, r)
+            x0_hat = gen.step(model, z, x_prev, r, ctx, ckpt=C['ckpt'])
+            x0d = x0_hat.detach()
+            t = tau()
+            eps = torch.randn_like(z)
+            xt = (1 - t) * x0d + t * eps
+            w = rng.uniform(*C['cfg'])
+            with torch.no_grad():
+                v_real = model.teacher(xt, t * 1000, ctx, neg, w, tuple(C['skip_neg']))
+                v_fake = fake(model, xt, t * 1000, ctx, r)
+            g_dmd = (v_real - v_fake) / x0d.numel()
+            stats['dmd_rms'] += (v_real - v_fake).square().mean().sqrt().item()
+            # L_pop: EMA feature moment matching; gradients flow through the frozen
+            # EMA net into x0_hat only (IDM paper, feature-moment matching).
+            s = rng.uniform(*C['lpop_sigma'])
+            eps2 = torch.randn_like(z)
+            f = model((1 - s) * x0_hat + s * eps2, s * 1000, ctx,
+                      lora=ema.lora, tp=ema.time, taps=(last,), ckpt=C['ckpt'])[last]
+            v = f.float().mean(1)[0]  # [D]
+            pool = torch.stack([v] + list(pop_q)) if pop_q else v[None]
+            mean_g, std_g = pool.mean(0), pool.std(0, unbiased=False)  # n=1 on first step: unbiased std is NaN
+            loss_pop = (mean_g - ref_mean).square().mean() + (std_g - ref_std).square().mean()
+            g_pop, = torch.autograd.grad(C['lpop_weight'] * loss_pop, x0_hat)
+            pop_q.append(v.detach())
+            stats['lpop'] += loss_pop.item()
+            stats['lpop_std_gap'] += (std_g - ref_std).abs().mean().item()
+            x0_hat.backward((g_dmd + g_pop) / C['micro'])
+            del x0_hat
+        gen.lora.drop_cache()
+        fake.lora.drop_cache()
+        D.mean_grads(gen_params)
+        gnorm = torch.nn.utils.clip_grad_norm_(gen_params, C['clip_gen'] or 1e12)
+        if torch.isfinite(gnorm):
+            opt_gen.step()
+            with torch.no_grad():
+                for pe, p in zip(ema.parameters(), gen.parameters()):
+                    pe.lerp_(p, 1 - C['ema'])
+        opt_gen.zero_grad(set_to_none=False)
+        stats['gnorm'] = gnorm.item()
+        return {k: v / C['micro'] for k, v in stats.items()}
+
+    def gen_step(phase):
+        ramp = 0.0 if gen_updates < C['adv_start'] else min(1.0, (gen_updates - C['adv_start'] + 1) / max(1, C['adv_warmup']))
+        stats = {k: 0.0 for k in ('dmd_rms', 'adv_ratio', 'lam', 'd_loss', 'd_real', 'd_fake', 'd_acc')}
+        n_gan = 0
+        fake.lora.freeze_cache()
+        gen.lora.freeze_cache()
+        for m in range(C['micro']):
+            use_teacher = teacher is not None and rng.random() < C['p_gan']
+            ctx, z, row = draw(trows if use_teacher else None)
+            r = r_index(phase, m)
+            x_prev = gen.rollout(model, z, ctx, r)
+            x0_hat = gen.step(model, z, x_prev, r, ctx, ckpt=C['ckpt'])
+            x0d = x0_hat.detach()
+            t = tau()
+            eps = torch.randn_like(z)
+            xt = (1 - t) * x0d + t * eps
+            w = rng.uniform(*C['cfg'])
+            with torch.no_grad():
+                v_real = model.teacher(xt, t * 1000, ctx, neg, w, tuple(C['skip_neg']))
+                v_fake = fake(model, xt, t * 1000, ctx, r)
+            g_dmd = (v_real - v_fake) / x0d.numel()
+            stats['dmd_rms'] += (v_real - v_fake).square().mean().sqrt().item()
+            _, _, h, wd = z.shape
+            leaf = x0d.clone().requires_grad_(True)
+            if teacher is None:
+                s = t - dt
+                x_fake = (1 - t) * leaf + t * eps - dt * v_fake
+                x_real = xt - dt * v_real
+            elif use_teacher:
+                s = rng.uniform(*C['gan_t'])
+                x_fake = (1 - s) * leaf + s * torch.randn_like(z)
+                x_real = (1 - s) * teacher(row['id']) + s * torch.randn_like(z)
+            else:
+                x0_hat.backward(g_dmd / C['micro'])
+                del x0_hat
+                continue
+            n_gan += 1
+            feats = model(x_fake, s * 1000, ctx, taps=disc.taps, ckpt=C['ckpt'])
+            logits_f = disc(feats, h, wd)
+            g_adv, = torch.autograd.grad(-logits_f.mean(), leaf)
+            ratio = (g_adv.norm() / (g_dmd.norm() + 1e-30)).item()
+            lam = C['adv_weight'] * ramp / max(ratio, 1e-12)
+            stats['adv_ratio'] += ratio
+            stats['lam'] += lam
+            x0_hat.backward((g_dmd + lam * g_adv) / C['micro'])
+            f_fake = {k: v.detach() for k, v in feats.items()}
+            del x0_hat, feats, logits_f
+            with torch.no_grad():
+                f_real = model(x_real, s * 1000, ctx, taps=disc.taps)
+            d_real, d_fake = disc(f_real, h, wd), disc(f_fake, h, wd)
+            d_loss = F.relu(1 - d_real).mean() + F.relu(1 + d_fake).mean()
+            (d_loss / C['micro']).backward()
+            stats['d_loss'] += d_loss.item()
+            stats['d_real'] += d_real.mean().item()
+            stats['d_fake'] += d_fake.mean().item()
+            stats['d_acc'] += 0.5 * ((d_real > 0).float().mean() + (d_fake < 0).float().mean()).item()
+        gen.lora.drop_cache()
+        fake.lora.drop_cache()
+        D.mean_grads(gen_params)
+        D.mean_grads(list(disc.parameters()))
+        gnorm = torch.nn.utils.clip_grad_norm_(gen_params, C['clip_gen'] or 1e12)
+        if torch.isfinite(gnorm):
+            opt_gen.step()
+            opt_disc.step()
+        opt_gen.zero_grad(set_to_none=False)
+        opt_disc.zero_grad(set_to_none=False)
+        with torch.no_grad():
+            for pe, p in zip(ema.parameters(), gen.parameters()):
+                pe.lerp_(p, 1 - C['ema'])
+        stats = {k: v / (C['micro'] if k == 'dmd_rms' else max(n_gan, 1)) for k, v in stats.items()}
+        stats['gan_frac'] = n_gan / C['micro']
+        stats['gnorm'] = gnorm.item()
+        stats['ramp'] = ramp
+        return stats
+
+    def critic_sigma():
+        u, k = rng.uniform(*C['critic_sigma']), C['critic_shift']
+        return k * u / (1 + (k - 1) * u)
+
+    def dmad_step(phase, gen_on):
+        stats = {k: 0.0 for k in ('g_logit', 'c_loss', 'd_teacher', 'd_gen', 'c_acc')}
+        stash = []
+        gen.lora.freeze_cache()
+        critic.requires_grad_(False)
+        for m in range(C['micro']):
+            ctx, z, row = draw(trows)
+            r = r_index(phase, m)
+            if gen_on:
+                x0_hat = gen.step(model, z, gen.rollout(model, z, ctx, r), r, ctx, ckpt=C['ckpt'])
+                s = critic_sigma()
+                logit = critic(model, (1 - s) * x0_hat + s * torch.randn_like(z), s * 1000, ctx, ckpt=C['ckpt'])
+                (-logit.mean() / C['micro']).backward()
+                stats['g_logit'] += logit.item()
+                x0 = x0_hat.detach()
+                del x0_hat, logit
+            else:
+                x0 = gen.rollout(model, z, ctx, r + 1)
+            stash.append((ctx, row, x0))
+        gen.lora.drop_cache()
+        if gen_on:
+            D.mean_grads(gen_params)
+            gnorm = torch.nn.utils.clip_grad_norm_(gen_params, C['clip_gen'] or 1e12)
+            if torch.isfinite(gnorm):
+                opt_gen.step()
+                with torch.no_grad():
+                    for pe, p in zip(ema.parameters(), gen.parameters()):
+                        pe.lerp_(p, 1 - C['ema'])
+                stats['gnorm'] = gnorm.item() * C['micro']
+            else:
+                stats['skip_gen'] = C['micro']
+            opt_gen.zero_grad(set_to_none=False)
+        critic.requires_grad_(True)
+        for ctx, row, x0 in stash:
+            s = critic_sigma()
+            d_q = critic(model, (1 - s) * teacher(row['id']) + s * torch.randn_like(x0), s * 1000, ctx, ckpt=C['ckpt'])
+            d_g = critic(model, (1 - s) * x0 + s * torch.randn_like(x0), s * 1000, ctx, ckpt=C['ckpt'])
+            loss = F.softplus(-d_q).mean() + F.softplus(d_g).mean()
+            (loss / C['micro']).backward()
+            stats['c_loss'] += loss.item()
+            stats['d_teacher'] += d_q.item()
+            stats['d_gen'] += d_g.item()
+            stats['c_acc'] += 0.5 * ((d_q > 0).float() + (d_g < 0).float()).item()
+        critic_params = list(critic.parameters())
+        D.mean_grads(critic_params)
+        cnorm = torch.nn.utils.clip_grad_norm_(critic_params, C['clip_critic'] or 1e12)
+        if torch.isfinite(cnorm):
+            opt_critic.step()
+            stats['cnorm'] = cnorm.item() * C['micro']
+        else:
+            stats['skip_critic'] = C['micro']
+        opt_critic.zero_grad(set_to_none=False)
+        return {k: v / C['micro'] for k, v in stats.items()}
+
+    @torch.no_grad()
+    def sample(tag):
+        nonlocal vae
+        if D.rank == 0:
+            if vae is None:
+                from diffusers import AutoencoderKLFlux2
+                vae = AutoencoderKLFlux2.from_pretrained(Path(C['model_dir']) / 'vae', torch_dtype=torch.float32,
+                                                         local_files_only=True).to(dev).eval()
+            d = out / 'samples' / tag
+            d.mkdir(parents=True, exist_ok=True)
+            images = []
+            for row in evals:
+                g = torch.Generator('cpu').manual_seed(C['eval_seed'])
+                z = torch.randn(1, model.tr.config.latent_channels, row['h'] // 16, row['w'] // 16, generator=g).to(dev)
+                x0 = ema.rollout(model, z, model.context(cache.get(row['id']).to(dev)), R)
+                images.append(decode(vae, x0))
+                from PIL import Image
+                Image.fromarray(images[-1]).save(d / f"{row['id']}.png")
+            contact_sheet(images, out / 'samples' / f'{tag}.jpg')
+        D.barrier()
+
+    def save(it):
+        if D.rank == 0:
+            state = {key: mod.state_dict() for key, mod in parts.items()}
+            state.update(it=it, gen_updates=gen_updates, config=C)
+            tmp = out / 'ckpt' / 'latest.tmp'
+            torch.save(state, tmp)
+            if latest.exists():
+                latest.replace(out / 'ckpt' / 'prev.pt')
+            tmp.replace(latest)
+        D.barrier()
+
+    if start == 0 and C['sample_every']:
+        sample('init')
+    t_last, log = time.time(), []
+    for it in range(start, C['iterations']):
+        entry = {'it': it}
+        if dmad:
+            gen_on = it >= C['critic_warmup']
+            entry.update(dmad_step(it, gen_on))
+            gen_updates += gen_on
+        elif idm:
+            entry['fake'] = sum(fake_step(it * C['fake_per_gen'] + k) for k in range(C['fake_per_gen'])) / C['fake_per_gen']
+            if it >= C['fake_warmup']:
+                entry.update(gen_step_idm(it))
+                gen_updates += 1
+        else:
+            entry['fake'] = sum(fake_step(it * C['fake_per_gen'] + k) for k in range(C['fake_per_gen'])) / C['fake_per_gen']
+            if it >= C['fake_warmup']:
+                entry.update(gen_step(it))
+                gen_updates += 1
+        log.append(entry)
+        done = it + 1
+        if done % C['log_every'] == 0:
+            keys = sorted({k for e in log for k in e if k != 'it'})
+            means = D.mean([sum(e.get(k, 0.0) for e in log) / len(log) for k in keys])
+            now = time.time()
+            if D.rank == 0:
+                rec = {'it': done, 'gen_updates': gen_updates, 's_per_it': round((now - t_last) / len(log), 2),
+                       'mem_gib': round(torch.cuda.max_memory_allocated() / 2**30, 1),
+                       **{k: round(v, 6) for k, v in zip(keys, means)}}
+                print(json.dumps(rec), flush=True)
+                with open(out / 'log.jsonl', 'a') as f:
+                    f.write(json.dumps(rec) + '\n')
+            t_last, log = now, []
+        if C['sample_every'] and done % C['sample_every'] == 0:
+            sample(f'{done:06d}')
+        if C['save_every'] and done % C['save_every'] == 0:
+            save(done)
+        if C['export_every'] and done % C['export_every'] == 0 and D.rank == 0:
+            ema.export(out / 'export' / f'{done:06d}', {'iteration': str(done)})
+    save(C['iterations'])
+    if D.rank == 0:
+        ema.export(out / 'export' / 'final', {'iteration': str(C['iterations'])})
+        gen.export(out / 'export' / 'final-raw', {'iteration': str(C['iterations'])})
+    if C['sample_every']:
+        sample('final')
+    if D.on:
+        dist.destroy_process_group()
+
+
+if __name__ == '__main__':
+    main()