code/rapid_moss/train.py

"""Distillation of MOSS-Image-SFT to an R-step data-space generator (arXiv 2610.12102).

    torchrun --nproc-per-node N -m rapid_moss.train --config configs/main.json [--set key=value ...]

method=dmad (arXiv 2610.02188): no fake score and no teacher evaluation during training. A critic
(base + LoRA on the later blocks + a head on the last block) separates noised teacher samples
(rapid_moss.teachergen) from noised generator samples of the same prompts with a logistic loss;
the generator minimises -critic(noised x0). Prompts are drawn from the teacher-sample pool.

method=dmd, per iteration: fake_per_gen fake-score updates, then (after fake_warmup) one generator update
with the DMD direction sg[v_real - v_fake] on x0 (arXiv 2610.12102) plus the ADM generator loss
-D(x_fake at tau - dt) (arXiv 2507.18569), and one discriminator hinge update on the same pair
(x_real = x_tau - dt * v_teacher_cfg, x_fake = x_tau - dt * v_fake). Each GPU runs `micro`
single-sample micro-steps per update (samples have different resolutions and text lengths).
"""
from __future__ import annotations

import argparse
import json
import math
import os
import random
import time
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,
    # discriminator real side: 'adm' = teacher ODE step from the same noised generator sample
    # (ADM, x_real/x_fake at tau - dt); 'teacher' = noised teacher samples of the same prompt
    # (ADP-style endpoints from rapid_moss.teachergen), drawn on a p_gan share of generator steps
    '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,  # 0 = no clipping; a non-finite gradient norm always skips the step
    '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,
    '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 (any number of dirs); 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:  # a shard cut off when its teachergen run was stopped
                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'
    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()))
    else:
        raise SystemExit(f"method must be dmd or dmad, 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 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

    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(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:
                # ADM: generator loss -D(x_tau - dt * v_fake) at tau - dt; gradient taken w.r.t. x0 only
                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'])
            # discriminator hinge on the same pair (fake features reused from the pass above)
            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')  # one older state survives a bad save
            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
        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()