[akwarii]

Summary

With autobatcher=True, a system that fits on the device only once or twice was refused with ValueError: ... is greater than max_metric: the probe steps back to 1 copy, and the padding then puts the limit below the system's own metric. This PR makes the estimated limit safe by introducing the following changes:

• Never pad a limit below a system that fit. The probe of the largest system fits at least one copy of it, so every estimated limit is now at least the largest metric, after the padding. This covers BinningAutoBatcher.load_states, both estimates of InFlightAutoBatcher._get_first_batch (the bootstrap one included), and the estimate perf(autobatcher): reduce the number of memory probes #642 adds to InFlightAutoBatcher.load_states. A max_memory_scaler set by the user is left as given.

• Keep the metrics in float64. estimate_max_memory_scaler converted them to float32 (torch default dtype), which could introduce rounding errors. This resulted in a system limit rounding below its own metric, and the system being refused by the autobatcher.

• Raise a clear error when nothing fits. When the first batch size already runs out of memory, determine_max_batch_size now says that it does not fit on the device, instead of returning 1 copy that then fails with the same max_metric error.

Closes #641. Builds on #642.

GPU validation

I reran the reproducer of #641 (MACE-MH-1+cuEq) on this branch, on a H200 GPU. Main refused these jobs with ValueError: ... is greater than max_metric for one of three reasons, each fixed here:

1. The probe returned 1 copy (it ran out of memory at 2 or 3 copies, or stopped at the cap of max_atoms_to_try either not working as intended or misleading for memory estimation #276, max_atoms_to_try = the job's atoms), and the padding of 0.9 made the limit 0.9 × the system's own metric. Now the limit is never padded below the largest system, which the probe showed fits at least once.
2. With torch's default dtype set to float32, the limit of 1 copy rounded just below the system's float64 metric. Now the metrics stay in float64.
3. ts.optimize took a first limit of 0.8 × what fits of the first system, which refused a larger second system. Now, for a job whose systems are all known, the limit is estimated over all of them (perf(autobatcher): reduce the number of memory probes #642) and kept above the largest one.

Systems	Setting	Runners	Cause	Before	This PR
17,576 atoms (fits twice on the device): alone, 2 copies, or next to 8 atoms	autobatcher=True	static, integrate, optimize	1	ValueError	runs
8 atoms, or 8 + 64 atoms	#276 cap, padding 0.9	static, optimize	1	ValueError	runs
8 + 64 atoms	#276 cap, padding 1.0	optimize	3	ValueError	runs
8 atoms, or 8 + 64 atoms	#276 cap, padding 1.0, float32 default dtype	static	2	ValueError	runs

A 54,872-atom cell, which runs out of memory even alone, now gets RuntimeError: A batch of 1 system(s) with 54872 atoms does not fit on the device. with autobatcher=True, instead of a 1-copy limit.

Checklist

Before a pull request can be merged, the following items must be checked:

• Doc strings have been added in the Google docstring format.
• Run ruff on your code.
• Tests have been added for any new functionality or bug fixes.

[lil-lon]

I'll review this PR once #642 is merged. The implementation looks good to me, so I'll mainly look at the tests then.

Also, as discussed in #641, the tests here don't run on a GPU. If you have results from running this branch on a GPU, could you add them to the PR description?

[lil-lon]

@akwarii Since #642 was merged, could you rebase this onto the latest main and resolve the conflicts?
The first five tests here can be dropped.

I also have a few suggestions for the tests:

• Could you mock determine_max_batch_size to return 1? That way the tests don't depend on the two-step back-off of the probe.
• Could you add a test case for the changes to _get_first_batch, with a lazy iterator?
• The three cases of the runner test all check the same line in BinningAutoBatcher.load_states. Could you replace the runner test with a test of BinningAutoBatcher.load_states itself?