Metadata-Version: 2.4
Name: chorus-torch
Version: 0.1.0
Summary: Train many PyTorch models at once, on one GPU.
Author: Joaquin Arroyo
License-Expression: MIT
Project-URL: Source, https://github.com/joaquinarroyo/chorus
Project-URL: Thesis, https://github.com/joaquinarroyo/thesis
Project-URL: Benchmarks, https://github.com/joaquinarroyo/fugue
Keywords: pytorch,ensemble,cross-validation,hyperparameter-search,gpu
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.13
Classifier: Programming Language :: Python :: 3.14
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Operating System :: POSIX :: Linux
Classifier: Operating System :: MacOS
Requires-Python: >=3.13
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.5
Requires-Dist: numpy>=1.24
Provides-Extra: dev
Requires-Dist: pytest>=8; extra == "dev"
Requires-Dist: pytest-cov; extra == "dev"
Requires-Dist: ruff==0.16.3; extra == "dev"
Requires-Dist: mypy==1.19.1; extra == "dev"
Dynamic: license-file

# 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.

```python
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

```bash
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:

```python
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_dict`s — 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:

```python
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

```python
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:

| | why |
|---|---|
| **class** | `vmap` applies *one* forward to K parameter sets |
| **parameter shapes** | they have to stack |
| **batch size** | one dispatched shape per step |
| **objective, optimizer family, AMP, compile** | one loss, one update rule, one decision per step |

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**.

| Flag | Default | What it does |
|---|---|---|
| `vectorize` | on | one `vmap` kernel for a bucket of K interchangeable models |
| `streams` | on | one CUDA stream per bucket, so buckets overlap on the device |
| `fused_loss` | on | one loss kernel over the stacked batch instead of K |
| `compile` | **off** | `torch.compile` on the bucket's forward — `mode="reduce-overhead"` on CUDA, `"default"` on CPU |
| `verify` | on | the equivalence check above (a safety net, not an extension) |
| `verify_atol`, `verify_rtol` | `1e-4`, `1e-3` | the tolerances that check compares with |
| `verify_samples` | `2` | how many samples its probe batch carries |

```python
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:

| Regime | Mechanism | Measured |
|---|---|---|
| **Heterogeneous** (buckets of one) | one CUDA stream per bucket, one process, one dataset | **1.47×** faster and **33 % less energy** |
| **Homogeneous** (buckets of K) | `vmap` over stacked parameters | **1.58×**, and **2.25×** with `torch.compile` |

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](https://dcc.fceia.unr.edu.ar), UNR), and it lives in
[joaquinarroyo/thesis](https://github.com/joaquinarroyo/thesis). The numbers
quoted throughout this README were produced by
[fugue](https://github.com/joaquinarroyo/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.
