chorus-torch 0.1.0__tar.gz

chorus_torch-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Joaquin Arroyo
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

chorus_torch-0.1.0/PKG-INFO

@@ -0,0 +1,432 @@
+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.