sCM Training

A multi-GPU training pipeline for continuous-time consistency models (sCM, Lu & Song 2024 — Simplifying, Stabilizing and Scaling Continuous-Time Consistency Models) on images (pixels or VAE latents), with post-hoc EMA and W&B logging. The student is parameterized in TrigFlow (x_t = cos(t)·x_0 + sin(t)·sigma_data·eps, t ∈ [0, π/2]) and is trained either:

• from scratch with the analytic v-target (no teacher), or
• distilled from a frozen TrigFlow teacher (e.g. one trained with the matching flow-matching codebase using a TrigInterpolant + LogNormalDist schedule).

The infrastructure (dnnlib, torch_utils, the U-Net architectures in training/networks.py, post-hoc EMA, dataset_tool.py, FID/FD-DINOv2 metrics, persistence-based pickling) is borrowed from NVIDIA's EDM2. The sCM model, loss, sampler, and configuration layout sit on top.

Layout

.
├── train.py                   # Build a config and launch sCM training.
├── generate_images.py         # Sample from a saved snapshot pickle.
├── calculate_metrics.py       # FID and FD-DINOv2 against a reference dataset.
├── reconstruct_phema.py       # Post-hoc EMA reconstruction (EDM2).
├── dataset_tool.py            # Pack a folder of images into a zip dataset.
├── training/
│   ├── training_loop.py       # The actual training loop (loads optional teacher).
│   ├── model.py               # SCMModel (TrigFlow consistency parameterization) + sample().
│   ├── loss.py                # SCMLoss (JVP-based sCM loss with the paper's tricks).
│   ├── networks.py            # SongUNet / DhariwalUNet (from EDM).
│   ├── encoders.py            # StandardRGB and Stability VAE encoders.
│   ├── schedulers.py          # constant_lr (default for sCM) and cosine_lr.
│   ├── phema.py               # Power-function and traditional EMA.
│   ├── monitoring.py          # W&B logging helpers.
│   └── dataset.py             # Streaming image dataset (zip or folder).
├── torch_utils/               # Distributed, persistence, training stats (EDM).
├── dnnlib/                    # EasyDict, class/func construction by name (EDM).
├── scripts/                   # Shell scripts: env setup, training, metrics.
├── datasets/                  # Place your packed datasets here.
├── training-runs/             # Output runs (one timestamped subdir per launch).
├── fid-refs/                  # Reference statistics for FID/FD-DINOv2.
└── out/                       # Generated images.

Setup

# Install the Python environment (CUDA wheel; adjust as needed).
bash scripts/module.sh

A Dockerfile is also provided.

End-to-end workflow

1. Pack the dataset

python dataset_tool.py convert \
    --source=raw_cifar/ \
    --dest=datasets/cifar10.zip \
    --resolution=32x32

For VAE-latent training, use dataset_tool.py encode to pre-encode images.

2. Train

The simplest single-GPU launch:

torchrun --standalone --nproc_per_node=1 train.py \
    --outdir=training-runs/cifar10 \
    --data=datasets/cifar10.zip \
    --preset=scm-cifar10

The scm-cifar10 preset is sCM-from-scratch with the analytic v-target. There are two orthogonal pickle-loading flags you can combine freely:

Flag	Loaded once at startup	Used during training
--pretrained-pkl	Initializes the student weights	Just a warm start — no further use
--teacher-pkl	Frozen teacher network (fp32, eval)	Provides v_target = teacher(x_t, t) in the loss

Distillation from a TrigFlow teacher (the typical use case):

torchrun --standalone --nproc_per_node=2 train.py \
    --outdir=training-runs/cifar10-scm \
    --data=datasets/cifar10.zip \
    --preset=scm-cifar10 \
    --pretrained-pkl=path/to/teacher-snapshot.pkl \
    --teacher-pkl=path/to/teacher-snapshot.pkl \
    --max-batch-gpu=256 \
    --status=20Ki --snapshot=1Mi --checkpoint=2Mi \
    --metrics=6Mi --metric-ref=fid-refs/cifar10.pkl

--pretrained-pkl and --teacher-pkl are independent — passing only one is fine. Names that don't match between the teacher and the student (e.g. the teacher's interpolant.*) are silently skipped during weight transfer.

Each launch creates a timestamped subdirectory inside --outdir (e.g. training-runs/cifar10-scm/260504_141200_scm-cifar10). Pointing --outdir at an existing run that contains a training-state-*.pt resumes from the latest checkpoint; --pretrained-pkl is rejected on resume (the weights already live in the checkpoint), but --teacher-pkl can be re-passed since it isn't part of the checkpointed state.

3. Reconstruct post-hoc EMA snapshots (optional)

python reconstruct_phema.py \
    --indir=training-runs/cifar10-scm/<run-dir> \
    --outdir=training-runs/cifar10-scm/<run-dir> \
    --outstd=0.100

4. Compute reference statistics for the dataset

bash scripts/metrics/ref50k.sh

5. Generate images

bash scripts/metrics/gen50k.sh

6. Compute FID / FD-DINOv2

bash scripts/metrics/fid50k.sh

sCM training step

Per minibatch the loss does the following (see training/loss.py):

1. Sample t. log(sigma) ~ N(P_mean, P_std), then t = atan(sigma / sigma_data) (TrigFlow log-normal schedule).
2. Forward process. x_t = cos(t)·x_0 + sin(t)·sigma_data·eps.
3. Target velocity. v_target = teacher(x_t, t) if a teacher is provided, else the analytic cos(t)·x_1 − sin(t)·x_0.
4. JVP along the trajectory. torch.func.jvp on forward_scaled with tangents (cos·sin·v_target, cos·sin·sigma_data). Returns the primal F_theta, the directional derivative JVP(F), and the per-sample adaptive log-variance logvar(t).
5. Stop-gradient target. y_sg = cos²·(sigma_data·sg(F) − v_target) + r_warmup·(cos·sin·x_t + sg(JVP)).
6. Three sCM tricks (all on by default; toggle with --no-tangent-norm, --no-warmup, --no-adaptive-weight):
   • Tangent normalization. y_sg ← y_sg / (||y_sg|| + c_norm) per sample.
   • JVP warm-up. r_warmup = min(1, step / warmup_steps).
   • Adaptive weighting. loss = 0.5·exp(−logvar)·(F − sg(F) + y_sg)² + logvar (EDM2 sign convention, equivalent to the 0.5·exp(w)·res² − w form in the sCM paper with w = −logvar).

The only term that carries autograd is the F − sg(F) part, so the optimizer steps only update the consistency-network parameters.

The scm-cifar10 preset values (overridable on the CLI):

Knob	Value	Flag
P_mean	-1.0	--p-mean
P_std	1.4	--p-std
c_norm	0.1	--c-norm
warmup_steps	10000	--warmup-steps
dropout	0.20	--dropout
lr (constant)	1e-4	--lr
p_uncond_labels	0.13	--p-uncond-labels
max_clip_norm	1.0	--max_clip_norm

--fp16 is off by default for sCM because the JVP through the network is most stable in fp32; the loss internally forces fp32 inside the network even when use_fp16=True.

Adding a new training run

The configuration is split into two preset dictionaries at the top of train.py:

• dataset_presets holds everything intrinsic to the data: sigma_data, the network architecture (net_kwargs), the sampler used during monitoring (sampler_kwargs, defaults to 2-step consistency sampling), and the LR scheduler family (lr_scheduler_kwargs, defaults to constant_lr).
• config_presets describes a particular training run on top of a dataset: which dataset to use, conditional vs unconditional, total nimg, batch size, classifier-free-guidance dropout, channel width, dropout, learning rate, gradient clipping, and the sCM-loss knobs (P_mean, P_std, c_norm, warmup_steps, the three trick toggles).

Adding a new run is mostly a matter of editing those two dictionaries and pointing --preset / --data at the new entry. Per-run overrides are exposed as CLI flags on train.py. The two dictionaries are required to have disjoint keys (asserted at startup) so it's always clear which preset a knob lives in.

Monitoring

Loss, learning rate, gradient norm, gradient-clip coefficient, the JVP warm-up ramp r_warmup, the adaptive logvar, and timing counters are pushed to W&B at every --status interval, alongside a grid of samples generated from the EMA model with the dataset's 2-step consistency sampler. Scalar metrics are duplicated against three x-axes (training step, images seen, wall-clock time); plots are tied to the training-step axis. The W&B project name is scm.

Credits

Built on top of NVIDIA's EDM and EDM2. The sCM training recipe follows Lu & Song, Simplifying, Stabilizing and Scaling Continuous-Time Consistency Models, 2024.