Share the batch-size probe and its wiring with every trainer

[klangenk]

Motivation

Every trainer that probes its batch size repeats the same code around the step: claim a safety margin, search, tell an out-of-memory failure from a bug, read the requested bound from the hyperparameters, and declare a --vram-limit-gb flag for itself and for the process it spawns. This PR moves all of it into the library, so a trainer only supplies the step it wants measured.

Implementation

Probing

• Extend probe_batch_size with on_out_of_memory (drop what a failed trial left on the card), minimum (for steps that cannot run on a single sample, e.g. BatchNorm over a 1×1 feature map) and candidate (an exact, non-power-of-two size that was asked for and then measured)
• Add measure_batch_size, the whole decision: honour the requested bound, bound the search by dataset_limit(sample_count), and return the settled size without writing it back anywhere

The requested bound

• Read it from the max_batch_size hyperparameter, named by REQUESTED_BATCH_SIZE and read through requested_batch_size (absent, None, '' and 0 all mean "the card decides")
• Keep batch_size free for reporting the result. The node saves the hyperparameters with a training (LastTrainingIO) and a training resumed after a restart reads them back, so reading the input from the same key would make it take its first run's measurement as its bound. The loop itself never hands reported values to a later training

VRAM limit

• Add node_parser(vram_limit=True), which declares --vram-limit-gb / VRAM_LIMIT_GB with one shared help text (VRAM_LIMIT_GB_FLAG, VRAM_LIMIT_GB_HELP)
• Add add_vram_limit_argument for the script a trainer spawns. The cap does not survive a spawn, so that script declares the same flag and calls limit_cuda_memory itself

Docs and tests

• Update AGENTS.md and list VRAM_LIMIT_GB in the README
• Add unit tests for every new parameter and for the VRAM flag: default, env var, flag precedence, legacy prefix

Compatibility

All new parameters are keyword-only with defaults that keep 0.23.0 behaviour. New public names: measure_batch_size, requested_batch_size, REQUESTED_BATCH_SIZE, add_vram_limit_argument, VRAM_LIMIT_GB_FLAG, VRAM_LIMIT_GB_HELP, and the vram_limit parameter of node_parser.

The node repositories pin learning_loop_node==0.24.0 on their companion branches, so this needs a v0.24.0 release (the version comes from the tag).

Checks

python -m pytest learning_loop_node/tests/unit: 276 passed. uvx ruff check is clean on every touched file. The unit suite fakes torch; what a real step costs can only be measured on a GPU in the nodes.

---

⬛ claude-opus-5-5

[jfrieli]

I've added some suggestions, please have a look. Most of them are cheap now but breaking after v0.24.0 is tagged.

Looking forward to move all nodes to the batch size probing. Especially yolo5, where your first tests doubled the estimated batch size!

[jfrieli]

This keyword carries the requested upper bound (the max_batch_size hyperparameter), but it is named after the result a trainer reports as batch_size. All three companion PRs end up writing

measure_batch_size(step, batch_size=requested_batch_size(hp))

which reads as "use this batch size" and is exactly the input/output mix-up this PR sets out to prevent. Suggest renaming it to max_batch_size:

measure_batch_size(step, max_batch_size=requested_batch_size(hp))

including the :param: line and the ValueError message. Since this is not released yet, it only costs one line in each of yolov5_node#65, classification_node#20 and dfine_node#52; after v0.24.0 it would be a breaking change.

[jfrieli]

Please log a warning when the trainer minimum overrides an explicitly requested maximum. For example, measure_batch_size(..., batch_size=1, minimum=2) currently selects 2 silently. Giving the technically required minimum precedence is fine, but the user should see why their setting was exceeded: 'Requested max_batch_size=1 is below the trainer minimum of 2; using 2.' Please also document this precedence in the parameter descriptions, which currently describe the requested value as the largest batch the training may use.

[jfrieli]

Suggestion: let measure_batch_size take a factory for the step instead of a step that is already built. It's cheap now, while the interface is unreleased and all three node PRs are adopting it anyway, and it protects every node added or reworked later.

What this finally can't release. It drops the margin, but the step (the throwaway model, its optimizer state, the EMA copy) is still referenced by the caller, so free_cuda_memory() can't reclaim it. That is why every node repeats the same lifecycle around the call:

# classification_node               # dfine_node              # yolov5_node
finally:                            finally:                  finally:
    del step                            del probe                 step.release()
    free_cuda_memory()                  free_cuda_memory()
    self.model.to('cuda')

classification and dfine probe in the process that then trains. A node that forgets this finally, or gets it slightly wrong, keeps the throwaway model on the card for the whole training. The training then has less memory than the probe measured and can run out of memory partway through. No test and no log line would show it. Right now avoiding this is up to each node's author; with a factory the library makes it impossible.

What a factory changes

# today (dfine)
probe = _Probe(config_file, dynamic_config)
try:
    return measure_batch_size(probe.run_trial, max_batch_size=..., sample_count=probe.sample_count,
                              on_out_of_memory=probe.drop_gradients, ...)
finally:
    del probe
    free_cuda_memory()

# with a factory
return measure_batch_size(lambda: _Probe(config_file, dynamic_config), max_batch_size=..., ...)

The library then owns the whole order: check for a GPU → reserve the margin → build the step → search → release the step → free memory. As a side effect:

• on_out_of_memory moves from the interface onto the step, because it is knowledge about the step.
• The no-GPU fallback becomes reachable. Today every caller has touched CUDA before calling (torch.cuda.init(), .to('cuda'), or dfine's own guard), so no node ever reaches it.
• One library test can assert that nothing references the step after the call, instead of each node having to get it right.

Cost. All three nodes already have a step class (TrainingStep, _TrainingStep, _Probe), so each needs roughly 10–20 lines, mostly moved:

• dfine: pass a lambda, and drop its own no-GPU guard and the finally.
• classification: the model.cpu() … model.to('cuda') bracket becomes a context manager.
• yolov5: pass a lambda. calc logs step.img_size after measuring, so that needs computing up front.

The measured sizes should not change, but a short run on n9 per node would confirm it.

To decide

1. How the library releases the step: a context manager (with build() as step, which fits classification's before/after) or an optional release() on the step.
2. Whether sample_count stays a parameter (the smaller change for yolov5 and classification) or moves onto the step.
3. probe_batch_size for the detection passes keeps taking a plain callable. Nothing is built there, so I think that asymmetry is fine.

While the signature changes anyway. The probe's vram_limit_gb repeats what limit_cuda_memory has already set in the same process: every node calls one with v and then passes the same v to the other. If limit_cuda_memory remembered the budget, the probe could read it from there, and the parameter could leave measure_batch_size and probe_batch_size in the same breaking change.

[jfrieli]

Could probe_batch_size keep only the parameters its own callers use? Its two external callers are the detection passes (classification's inference probe and dfine's _run_detection), and both pass only probe, limit and vram_limit_gb. candidate, minimum and on_out_of_memory are there only so measure_batch_size can forward them.

candidate matters most. Once it is public, a caller has to know from the docstring that an image count or a dataset bound must never be passed as a candidate. Kept inside, that rule is enforced by construction: only measure_batch_size can name a candidate, and it only ever names the requested size.

Suggestion: move the full-parameter version into a private core that both call, and keep the public probe_batch_size(run_batch, *, probe, limit, vram_limit_gb). The duplicated minimum tests in test_cuda.py could then collapse onto find_batch_size and measure_batch_size.

Same timing as the other interface comments: all three parameters are new in this PR, so narrowing it now breaks nothing. After v0.24.0 it would.

[jfrieli]

This rounds minimum down, so the search can start below the lower bound the caller named: minimum=3 tries 2 first, 6 tries 4, 12 tries 8.

minimum is documented as the size below which the step cannot run at all, so the trial below it doesn't run out of memory: it fails some other way, e.g. BatchNorm's Expected more than 1 value per channel. measured_fits correctly treats that as not-an-OOM and re-raises it, so the probe aborts with what looks like a bug in the node's step, although the library was the one that went below the minimum.

Nobody hits this today, because all three nodes use MIN_BATCH_SIZE = 2, which is already a power of two. It is a trap for the first node that needs 3 or 6, e.g. a training split across 3 GPUs that needs at least 2 samples per GPU.

Suggestion: round up to the next power of two, so the smallest trial is never below minimum. The same applies to the no-GPU fallback in probe_batch_size (cuda.py:101). test_a_minimum_is_rounded_down_to_a_power_of_two pins the current behaviour in both test_batch_size.py and test_cuda.py, and the :param minimum: text says "rounded down". Since this changes how a new public parameter behaves, now is the cheap moment, before v0.24.0.