rtcwm/train/train.py

"""Causal (streaming) LoRA training of the CWM, initialised from the released CWM LoRA.

    torchrun --nnodes N --nproc-per-node 8 ... -m rtcwm.train.train --out DIR [--steps 150 | --epochs E]

* diffusion forcing on 6-chunk spans cut from the clips' latent streams, each with a drawn start frame
  (rtcwm.train.dataset, rtcwm.model.objectives)
* the frozen H3 blocks are FSDP-sharded over the 8 GPUs of each node; the LoRA weights are replicated
  and their gradients all-reduced asynchronously as each is produced in backward
* one window per GPU per step; steps = epochs * samples / world unless --steps is given
"""
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 .. import paths
from ..model import loading
from ..model.layout import SPAN_CHUNKS
from ..model.objectives import LOSSES, OBJECTIVES
from .dataset import StartFrames, list_samples, load_window, silence


def parse():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--out", type=Path, required=True)
    ap.add_argument("--objective", choices=OBJECTIVES, default="df")
    ap.add_argument("--init-lora", type=Path, default=paths.CWM_LORA)
    ap.add_argument("--steps", type=int, default=0, help="0 = epochs * samples / world")
    ap.add_argument("--epochs", type=float, default=2.0)
    ap.add_argument("--lr", type=float, default=2e-5)
    ap.add_argument("--lr-min", type=float, default=1e-6)
    ap.add_argument("--warmup", type=int, default=10)
    ap.add_argument("--save-every", type=int, default=10)
    ap.add_argument("--val-clips", type=int, default=12)
    ap.add_argument("--max-samples", type=int, default=0, help="smoke tests: cap the training set")
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--picture1", choices=["native", "2048"], default="native", help="Picture 1 resolution")
    ap.add_argument("--attn", choices=["fa3", "flex"], default="fa3", help="the same mask: FlashAttention-3 prefix calls or flex")
    ap.add_argument("--level-shift", type=float, default=3.0, help="DF: flow shift of the per-chunk noise levels")
    ap.add_argument("--act-offload", choices=["on", "off"], default="off", help="checkpointed block inputs on the CPU")
    return ap.parse_args()


def node_group(rank: int, world: int):
    """One FSDP shard group per node (NVLink); every rank must create every group."""
    lw = int(os.environ.get("LOCAL_WORLD_SIZE", torch.cuda.device_count()))
    groups = [dist.new_group(list(range(n * lw, (n + 1) * lw))) for n in range(world // lw)]
    return groups[rank // lw]


def resume_point(a):
    """(steps done, LoRA to start from): the latest checkpoint in --out (a requeued or resubmitted run --
    lowprio jobs are preempted and requeued -- continues from it), else --init-lora."""
    done = sorted(int(p.stem.split("step")[1]) for p in a.out.glob("lora_step*.safetensors"))
    return (done[-1], a.out / f"lora_step{done[-1]:05d}.safetensors") if done else (0, a.init_lora)


def cosine_with_warmup(opt, a, steps):
    return torch.optim.lr_scheduler.LambdaLR(opt, lambda i: min(1.0, (i + 1) / a.warmup) * (
        a.lr_min / a.lr + (1 - a.lr_min / a.lr) * 0.5 * (1 + math.cos(math.pi * min(i, steps) / steps))))


def main():
    a = parse()
    dist.init_process_group("nccl")
    rank, world = dist.get_rank(), dist.get_world_size()
    local = int(os.environ.get("LOCAL_RANK", 0))
    torch.cuda.set_device(local)
    dev = torch.device("cuda", local)
    log = (lambda *m: print(f"[{time.strftime('%H:%M:%S')}]", *m, flush=True)) if rank == 0 else (lambda *m: None)
    group = node_group(rank, world)

    train, val, dropped = list_samples(a.val_clips, a.seed, a.picture1)
    if a.max_samples:
        train = train[:a.max_samples]
    steps = a.steps or int(math.ceil(a.epochs * len(train) / world))
    log(f"objective={a.objective} span_chunks={SPAN_CHUNKS} level_shift={a.level_shift} attn={a.attn} act_offload={a.act_offload} world={world} train={len(train)} val={len(val)} dropped_flagged={dropped} steps={steps}")
    if rank == 0:
        a.out.mkdir(parents=True, exist_ok=True)
        (a.out / "val_samples.json").write_text(json.dumps(val))
        (a.out / "args.json").write_text(json.dumps({k: str(v) for k, v in vars(a).items()}, indent=1))

    tr = loading.load_base()
    start, init = resume_point(a)
    if start:
        log(f"resuming after step {start}: {init.name} (optimizer moments start afresh)")
    net = loading.attach_trainable_lora(tr, init)
    loading.shard_blocks(tr, group, dev)
    net.to(dev, torch.float32)
    tr.train()
    params = list(net.parameters())
    handles = []

    def reduce_hook(p):
        handles.append(dist.all_reduce(p.grad, op=dist.ReduceOp.AVG, async_op=True))
    for p in params:
        p.requires_grad_(True)
        p.register_post_accumulate_grad_hook(reduce_hook)
    log(f"LoRA trainable={sum(p.numel() for p in params) / 1e6:.1f}M mem_after_shard={torch.cuda.memory_allocated() / 2**30:.1f}G")
    opt = torch.optim.AdamW(params, lr=a.lr, betas=(0.9, 0.999), weight_decay=0.01)
    sched = cosine_with_warmup(opt, a, steps)
    for _ in range(start):                       # the schedule where the resumed run left it
        sched.step()
    loss_fn, audio = LOSSES[a.objective], silence()
    offload = a.act_offload == "on"
    gen = torch.Generator().manual_seed(a.seed * 1000 + rank)
    starts = StartFrames(train, a.seed * 1000 + rank, a.picture1)

    order, ptr, mask_cache, t0 = [], 0, {}, time.time()
    for it in range(steps):
        if ptr >= len(order):
            order = train[:]
            random.Random(a.seed + it).shuffle(order)
            order, ptr = order[rank::world], 0
        clip, j = order[ptr]
        ptr += 1
        j, m = starts.draw(clip, j)
        if it < start:                           # resumed: replay the sample order up to where it stopped
            continue
        loss = loss_fn(tr, load_window(clip, j, audio, start=m, res=a.picture1), dev, gen, mask_cache, offload=offload,
                       attn=a.attn, level_shift=a.level_shift)
        loss.backward()
        for h in handles:
            h.wait()
        handles.clear()
        gn = torch.nn.utils.clip_grad_norm_(params, 1.0)
        opt.step(); sched.step(); opt.zero_grad(set_to_none=True)
        stat = loss.detach().reshape(1)
        dist.all_reduce(stat, op=dist.ReduceOp.AVG)
        if rank == 0:
            el = time.time() - t0
            log(f"step {it + 1}/{steps} loss_v={stat[0]:.4f} gnorm={float(gn):.3f} lr={sched.get_last_lr()[0]:.2e} "
                f"mem={torch.cuda.max_memory_allocated() / 2**30:.0f}G {el / (it + 1 - start):.1f}s/step "
                f"eta={(steps - it - 1) * el / (it + 1 - start) / 3600:.2f}h")
            with open(a.out / "train_log.jsonl", "a") as fh:
                fh.write(json.dumps(dict(step=it + 1, loss_v=float(stat[0]), gnorm=float(gn), lr=sched.get_last_lr()[0],
                                         sec=el)) + "\n")
            if (it + 1) % a.save_every == 0 or it + 1 == steps:
                p = a.out / f"lora_step{it + 1:05d}.safetensors"
                net.save_weights(str(p), torch.bfloat16, {"base": "cwm", "step": str(it + 1), "objective": a.objective,
                                                    "span_chunks": str(SPAN_CHUNKS), "level_shift": str(a.level_shift)})
                log("saved", p)
    dist.barrier()
    dist.destroy_process_group()


if __name__ == "__main__":
    main()