Time each collate transform separately - #1274
Draft
EricBoittier wants to merge 7 commits into
Draft
EricBoittier wants to merge 7 commits into
EricBoittier wants to merge 7 commits into
Conversation
Before changing anything about the data pipeline, measure it: which fraction of a training step is spent waiting for the input pipeline rather than running the model. `metatrain.utils.timing` accumulates per-stage timings when the `METATRAIN_TIMING` environment variable is set, and does nothing at all otherwise. The PET training loop is split into `loader` (waiting for the next batch) and `step`, with `step` broken down into unpack, host to device transfer, forward, loss, backward and optimizer; `CollateFn` is split into group_and_join, transforms and serialize. The collate stages are recorded in whichever process runs the collate function, so they only show up with `num_workers=0`; with workers they are part of the `loader` wait instead. `benchmarks/benchmark_pipeline.py` runs a real `Trainer.train` with the timings on and prints the table. No public API changes.
`CollateFn` serializes systems and tensor maps into a byte blob so they survive the `DataLoader` worker boundary, and `unpack_batch()` immediately reverses that. Add a `Batch` dataclass holding the same content directly, and accept it in `unpack_batch()`. Nothing is faster yet, and the serialized tuple keeps working exactly as before, so every architecture and third-party caller is unaffected. This is the seam that lets a future transport skip the round trip.
`CollateFn` did two things: group samples into a batch and run the transformation callables, then serialize the result into a byte blob for transport. Extract those into `collate_batch()` (samples -> `Batch`) and `serialize_batch()` (`Batch` -> blob), and define `CollateFn` as their composition. Pure refactor: the code inside each is unchanged, and `CollateFn`'s signature and return value are identical, so plugins keep working. Callers that do not need the round trip can now collate without it, and the benchmark reports what that round trip costs (2.9 ms vs 0.4 ms for an 8-structure batch of qm9, no transformations).
5 tasks
Persistent workers hold their process state for as long as the loader lives, so the benchmark has to say whether that state is affordable, not only whether epochs got faster. Sample the proportional memory of the whole process tree (PSS, so pages the workers share with the parent are counted once) and report it as a delta over the baseline, which is what answers the question. psutil is optional: it is not declared anywhere, and `test_benchmark_script_runs` asserts this script exits 0. Also fix what the benchmark measures. The "realistic validation split" was not a split: training kept the whole dataset while validation took the first 20% of it, so a 100-structure run did 100 training plus 20 overlapping validation structures. Make it disjoint.
A `DataLoader` left on the default generator draws a worker seed from the global RNG every time an iterator is created. How often a loader is iterated therefore shifts every random number the training loop draws afterwards, and that count is not something the caller controls: the scaler fit iterates its loader twice when per-property scales are needed, and `get_num_workers()` returns `max(0, min(cpu_count - 4, 8))`, so whether workers exist at all depends on the host's core count. Give each loader its own generator, seeded from the global RNG at construction. The draws move off the global stream while `torch.manual_seed` stays in charge of the run, so a training trajectory no longer depends on how many times a loader is iterated, nor on how many cores the machine has. This moves two PET regression references once. The new values reproduce on 4 cores and on 24, where the old ones only held on the machine that generated them: 2.834107398987 at 24 cores against 2.805202722549 at 4, a 1% spread that was reading as reproducible only because CI and a workstation happen to sit on opposite sides of the worker-count threshold. The values now agree across core counts to float32 noise, ~1e-7 relative.
Worker processes were forked and torn down once per epoch, per loader. Set persistent_workers where the loader outlives a single pass, which cuts the steady-state epoch from 1.10 s to 0.74 s at four workers on a 100-structure carbon set, and lowers peak proportional memory at eight workers from 2.38 GB to 1.54 GB, since the repeated forking was itself the expensive part. The PET regression references do not move. They would have, had the previous commit not taken worker seeding off the global RNG: a persistent iterator is reset rather than recreated, so it used to draw one fewer seed and shift everything the training loop drew afterwards.
The `transforms` stage lumped together the atomic basis preparation, the rotational augmentation, the neighbor lists, the conditioning, the additive subtraction and the scaling, which is exactly the aggregate that has to be split before guessing at which one to move out of the epoch path. Time each callable in a stage named after it. On 216-atom periodic silicon that gives, per batch: neighbor lists 5.9 ms, additive subtraction 3.7 ms, augmentation 2.4 ms, scaling 0.5 ms and the atomic basis 0.01 ms, against 1.6 ms to group the samples and 11.1 ms to serialize them. The stage names come from the callables, so this reads best on top of the change that names them after what they do; without it the five closures are all called `transform` and share one stage.
EricBoittier
force-pushed
the
pr/collate-transform-timing
branch
from
September 18, 2026 08:53
ad2c99e to
5166057
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
The
transformsstage introduced in #1268 lumped together the atomic basispreparation, the rotational augmentation, the neighbor lists, the conditioning,
the additive subtraction and the scaling. That aggregate is exactly the thing
that has to be split before anyone can argue about which transform should move
out of the per-epoch path, so this times each callable in a stage named after
it.
Per batch on 216-atom periodic silicon:
The stage names come from the callables, so this reads best on top of #1272;
without it the five closures are all called
transformand share one stage.The first commit also fixes what the benchmark itself measures: a disjoint
train/validation split,
psutiloptional,nvidia-smisampling droppedbecause its own overhead was comparable to what it reported, and baseline /
peak / added memory reported separately.
This is instrumentation only, no pipeline behaviour changes. It is the evidence
step for deciding what preprocessing should be hoisted out of the epoch once
#1267 gives the trainer a place to declare it, rather than an argument for a
new preprocessing subsystem.
Stacked on #1271, #1270, #1269 and #1268; the new commits here are
Measure only what the benchmark can defendandTime each collate transform separately.Contributor (creator of pull-request) checklist
Maintainer/Reviewer checklist
Drafted with Cursor 2.4.26 in agent mode (Claude Opus 5), then read through and
edited by hand before pushing. Numbers from a Ryzen 9 5900X with an RTX 4060 Ti,
PyTorch 2.13.0+cu130; a Snakemake harness is in preparation to repeat the matrix
on cluster hardware.
📚 Documentation preview 📚: https://metatrain--1274.org.readthedocs.build/en/1274/