"""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()