chorus

Train many PyTorch models at once, on one GPU.

Cross-validation, ensembles and hyperparameter search all end up as the same loop: for model in models: train(model). One model at a time, each one leaving most of the GPU idle. chorus runs them together — in one process, one CUDA context, one copy of the dataset, and, where the models allow it, one batched kernel instead of K sequential ones.

import chorus

report = chorus.ensemble(lambda: ResNet18(), train_data, n=8,
                         val_data=val_data, epochs=20, batch_size=64)

print(report.mean("val_score"))
best = report.best().model            # a plain nn.Module, trained

It is a library, not a framework: nothing here asks you to restructure your code, subclass anything, or hand over your training loop unless you want to.

Install

pip install -e .          # from a clone; not on PyPI yet

Only torch (>= 2.4) and numpy. No CUDA required to use it — everything works on CPU, just without the parts that need a device.

Linux and macOS. Linux with an NVIDIA GPU is where the batched path pays off; macOS runs the whole test suite on CPU, which is where this is developed. Windows is not tested and not supported.

Three levels

1. The primitive: Stack is an nn.Module

K models, one module, one (K, B, …) output. Drop it into whatever loop you already have:

stack = chorus.Stack([make_model() for _ in range(8)]).cuda()
opt = torch.optim.AdamW(stack.parameters(), lr=1e-3, fused=True)

for x, y in loader:                                   # your loader
    loss = chorus.stacked_cross_entropy(stack(x), y)  # one kernel, not eight
    loss.backward(); opt.step(); opt.zero_grad()

stack.sync_()        # your eight model objects are now trained

.to(), .train(), .eval(), state_dict(), AMP and torch.compile behave exactly as they do for any other module. stack.unstack(i) hands back member i as an instance of your own class, and stack.state_dicts() hands back K ordinary state_dicts — your key names, no leading K — for saving. (The stack's own state_dict() is the stack's: right for reloading a stack, wrong for your_model.load_state_dict.)

Different learning rate per member? chorus.StackedAdam(stack, lr=[1e-3, 3e-4, …]) is a torch.optim.Optimizer.

2. The trainer: heterogeneous workloads

When the models are not all alike — different architectures, batch sizes, epoch budgets — Trainer works out what can be batched together and what merely shares the device:

t = chorus.Trainer()
t.add(resnet_a,   train, val_data=val, epochs=10, batch_size=64)
t.add(resnet_b,   train, val_data=val, epochs=10, batch_size=64)
t.add(transformer, other, epochs=5,    batch_size=16, lr=3e-4)

print(t.plan())      # before spending anything
report = t.run()

3 models → 2 bucket(s)
  bucket 0   2 models vectorized  2 models interchangeable by structure (shared loader)
  bucket 1   1 model  eager       single model, nothing to batch (shared loader)

3. The workloads

chorus.ensemble(model_fn, data, n=10, epochs=20)
chorus.cross_validate(model_fn, data, k=5, epochs=20)
chorus.hpo(model_fn, data, {"lr": [1e-2, 1e-3, 1e-4]}, val_data=val)

Each is a few lines over Trainer. They take a factory because chorus has to build the K members itself — vmap needs K parameter sets of identical structure, and one instance used K times would train one model K ways.

Mix them freely: add_folds and add_grid on the same trainer bucket together whatever can be bucketed together.

How it decides what to batch

Two models share a batched forward when they agree on:

Everything else may differ, and that is exactly what the three workloads need: data split, learning rate, weight decay, number of epochs.

Same class and shapes is not enough, and this is the failure the library is most careful about. Two instances of one class can still compute different things — use_residual=False, a different activation, a dropout probability read from self. Neither the class nor the shapes show it, and vmap would train both with the first one's logic and return weights that are wrong without raising anything. So before committing a bucket, chorus runs one batch through both routes and compares; if they disagree, that bucket falls back to eager. One forward at startup, and a silent wrong answer becomes a slower right one.

That batch is taken from the first samples of your Dataset, which works when it yields (input, target) and not otherwise. When it cannot be built the bucket still vectorises, and the plan says so — vectorized … unchecked (no probe batch), never the same line as one that passed. Pass your own with t.add(model, data, example_input=batch) to get the check back.

The rule worth remembering:

• hyperparameters that preserve parameter shapes (learning rate, weight decay, momentum, seed, fold) → one bucket, fully batched.
• hyperparameters that change the topology (width, depth) → one bucket each, no batching — still overlapped on the device, but that is a different and smaller win.

Trainer.plan() tells you which one you have before you spend anything, and report.plan tells you what actually happened — including how many samples the tail drop costs each member per epoch, since a stacked step needs a constant batch shape and the short last batch cannot be one.

All of that goes through logging.getLogger("chorus") when verbose=True, so it can be silenced, reformatted or routed like any other library's output.

Turning the extensions on and off

Every extension is a keyword argument, and every report says which were on. The defaults are all on except compilation.

chorus.Trainer(streams=False)                    # all eight live here
chorus.Stack(models, vectorize=False)            # the five a stack can honour
chorus.ensemble(model_fn, data, n=8, compile=True)

cfg = {"vectorize": True, "streams": False, "fused_loss": False}
chorus.ensemble(model_fn, data, n=8, **cfg)      # one config, several runs

Where each one is accepted, and why it is not everywhere:

• Trainer takes all eight.
• Stack takes the five it can act on: vectorize, streams, verify, verify_atol, verify_rtol. A stack computes no loss (fused_loss), does not compile itself (the runtime does that to the bucket's forward), and receives its probe batch already built (verify_samples).
• ensemble, cross_validate, hpo take all eight and hand them to the trainer they build. Everything else — epochs, batch_size, lr, optimizer, amp, objective — goes to Trainer.add. Pass a trainer= you built yourself and the switches are refused rather than silently discarded.

compile and amp are the two that can also be set per member — t.add(model, data, compile=True). That is not a convention: they are part of the bucket key, so members that disagree land in different buckets. The rest decide how a bucket dispatches its step, so one bucket cannot hold two answers, and they stay run-wide.

Three things worth knowing before you tune these:

• compile and streams are mutually exclusive, and that is PyTorch's constraint: CUDA Graph capture does not tolerate dynamic ambient streams. A compiling bucket drops its streams. If you asked for both, chorus warns; if streams was merely the default, it is silent about it.
• compile works off the device too. There are no CUDA Graphs to capture on a CPU, so the mode changes and the win is smaller — inductor's fusion, not the launch overhead it removes — but the bucket does compile. AMP is the one that does not: chorus's mixed precision is autocast + GradScaler, which is CUDA's, and asking for amp=True elsewhere warns and trains in full precision rather than quietly substituting CPU bfloat16 autocast, which is a different mechanism with a different numerical story.
• fused_loss was measured below the noise floor on its own, across all six ablation configurations. It stays on by default for a different reason: it keeps the dispatched shape constant as members finish, which is what CUDA Graph capture requires. Do not expect it to be the flag that makes your run fast, and do not go looking for the win it does not have.

What to expect

Two mechanisms, and they are worth telling apart:

Those numbers come from the thesis this code was extracted from (T4 and L4, CIFAR-scale CNNs and transformers). Yours will differ. Two things they do say:

• Stacking models that have nothing in common already pays, before any vectorisation. Most of that comes from not paying K times for the process, the CUDA context and the dataset.
• Overlap does not fix a bottleneck outside the GPU. In a workload that was data-loading-bound (30 % device utilisation), streams bought nothing. Check where your time goes before expecting a speed-up here.

What it does not do

• No trial pruning. Every point of an hpo grid runs to its last epoch. For real searches, a scheduler that stops bad trials early (ASHA, Hyperband, median stopping) usually beats fusion — in the thesis's own measurements, pruning was worth ~30 % on its own, more than the entire difference between execution strategies. The two compose (a pruned trial is just a member frozen early, which the runtime already supports) and that is the top item on the roadmap. Until then: if your search is long and your trials are separable, Ray Tune will beat this.
• One GPU. No multi-device, no distributed training.
• No memory estimation, and no promise about fitting. chorus trains everything you add, at once. It does not profile your models, does not guess how many fit, and will not quietly train fewer than you asked for. If the total exceeds your VRAM you get a CUDA OOM — set max_parallel=N and it runs in chunks of N. The ceiling is yours to set, deliberately: the alternative is a guess that refuses workloads which would have run.
• No CPU offload of finished members' weights, so peak VRAM is the sum of everything co-resident for as long as the chunk lasts.
• Round-robin only between buckets. The original paper describes four schedulers; this implements one.
• vmap-able models only for the fast path. Data-dependent control flow in the forward falls back to eager, correctly and automatically.
• The equivalence check compares outputs, not buffers. It catches two models that compute different things; it does not catch a layer whose state the batching rule updates through a copy — the BatchNorm failure that norms.py handles by hand for the layers it knows about. A model with an unusual stateful layer is still a manual analysis. Closing that is item 2 of the roadmap.

Roadmap

These are the open lines from the thesis's future-work chapter, restated for the library and ordered by what they are worth. Each one shares a condition worth saying once: implementing it is not the hard part, measuring what it buys is.

1. Trial pruning for hyperparameter search. The gap the evaluation left most exposed. Ray won all four HPO experiments without compilation, and not by running faster but by running less — within Ray itself, ASHA took 134 s against 190 s for FIFO, so pruning was worth more than the entire difference between concurrency strategies. It is orthogonal to fusion and the two compose: the epoch mask that already freezes finished members works unchanged to retire a trial at whatever step a scheduler decides, and the bucket keeps running vectorised with whoever is left. Expect the gain to be multiplicative, not additive — fusion over the trials that survive pruning. Already here: Bucket.finish_member, the mask, and per-member learning rates. Missing: a scheduler interface, and a per-epoch validation signal to feed it.

2. Verification that compares buffers, not outputs. vmap imposes a condition on the code it transforms, and this port resolves it in a bounded way: by hand, for the stateful layers of the models that were evaluated. A model with a different stateful layer would need that analysis repeated, and nothing warns you. The answer is not a static purity analyser — even a perfect one looks at the wrong level, since the layer's source is impure and the thing that breaks (the batching rule interposing a copy) happens below what static analysis sees. It has to be dynamic, and it is cheap: on building a bucket, run one training step down each route and compare which buffers each one modified, falling back to eager when they differ. A module with no buffers has no state to lose, which makes a cheap pre-filter. The open question underneath: specialising a layer by hand costs a fixed per-step overhead amortised across K members, while relying on the check costs that bucket's whole vectorisation. Since fusion pays little at small K, there should be a K below which falling back is nearly free — finding it says when a new layer is worth specialising. Already here: the output-level check in vmap_matches_eager.

3. Where the ceiling actually is. chorus deliberately does not estimate how many models fit — but nobody knows how many should run at once either, because co-resident work contends for GPU internals that no static model of VRAM and cores can see. A sweep that raises the number of concurrent models until the aggregate-throughput curve bends would measure how much of the expected margin contention eats, and turn max_parallel from a guess into a measurement. Measuring SM occupancy directly is not the way: profiler-level counters (ncu, CUPTI) serialise launches and destroy the wall-clock number that matters.

4. The shared loader, in both directions. Within a bucket whose members read the same data, one DataLoader feeds all K — a large part of why this wins on data-bound workloads. It has a cost and a symmetry, and neither is explored. The cost: the K members consume the same tensor, so they share the shuffle and the augmentations, which correlates their errors and can only subtract from a committee's quality. Recovering the augmentation diversity is cheap — apply it per member on the already-loaded batch, so each gets its own crop of the same images, with no second decode. Recovering the ordering diversity is expensive: which samples land together is a property of sampling, and undoing it needs one loader per member, which is what sharing avoids. The symmetry: cross-validation folds do build K loaders, and with five folds the training sets overlap by 80 %, so every sample crosses the transform pipeline four times per epoch. Sharing the pipeline and selecting per member would remove that. The tension: the direct way — one common batch from which each member drops its own validation fold — makes the effective batch size stochastic, and with it the equivalence to an independently run cross-validation.

5. The parts of the original design still missing. When a sub-model finishes, UnifiedNN copies its weights to CPU memory and releases its share of VRAM progressively; that matters in heterogeneous workloads, where members finish at very different times, and it is what would let freed memory admit incoming work. In the same direction, the original defines four scheduling policies between sub-models and this implements one — round-robin, which the paper itself calls suboptimal when per-epoch durations differ widely. The two belong together: they pay off in the same scenario.

6. The cloud scenario. The approach was designed for multiple users submitting heterogeneous models and datasets, with dynamic schedulers by priority or job size. chorus, like the thesis, evaluates the single-tenant case. Adapting it would need continuous job arrivals, which changes the shape of the runtime — and would allow comparing both approaches on the scenario UnifiedNN was actually conceived for.

7. Coordinated scheduling and memory between buckets. Salus and Zico show that switching jobs at iteration boundaries and sharing ephemeral memory sustains more concurrent training than a static allocation admits — but both need an intermediary runtime, because the processes they schedule cannot observe each other. chorus has no such barrier: its members share a process, a context and an allocator, so each one's phase (forward, backward, update) is directly observable. Staggering buckets so that one's backward, which frees activations, overlaps the other's forward, which reserves them, would cut peak VRAM without touching the CUDA stack or the framework. Complementary to item 5: that one attacks the persistent memory of finished members, this one the ephemeral memory of active ones.

Fidelity to the executor it was extracted from

chorus was pulled out of an execution strategy that had been measured over a 504-run campaign, and the plan is to run both side by side on the same benchmark: is the port faithful, and where can it now be made faster. That only works if every difference is written down, so they are — docs/fidelity.md lists what was aligned back to the original, what is identical by construction, and what differs on purpose (no memory estimation, the vectorisation check, no inner stream for a bucket of one, and this library not touching cudnn.benchmark or any other global of your process).

Where this comes from

chorus implements the approach of Taki et al., UnifiedNN: Efficient Neural Network Training on the Cloud (2024), whose authors published no code, plus the extensions developed and measured in the thesis it was extracted from: vectorised forward via torch.vmap, a fused loss over the stacked batch, and per-member hyperparameters on the stacked optimizer — without which a hyperparameter sweep silently collapses onto the first trial's values.

That thesis is Desarrollo de estrategias para el entrenamiento paralelo de modelos neuronales en GPUs individuales (Joaquín Arroyo, Licenciatura en Ciencias de la Computación, UNR), and it lives in joaquinarroyo/thesis. The numbers quoted throughout this README were produced by fugue, the benchmark harness written for it — chorus is the unified-model executor from there, pulled out and made usable on its own.

The comments in src/chorus/ are the interesting part of that history: every long one marks something that produced plausible, wrong numbers before it was understood. norms.py (BatchNorm statistics frozen at (0, 1) under vmap + AMP), optim.py (the sweep that stopped sweeping), runtime.py (a stream race that only hurt at scale, and momentum carrying a "finished" member for tens of steps).

MIT licensed.