Skip to content

optimizer: keep the probe's margin when the whole job fits - #43

Merged
niklashoelter merged 1 commit into
niklashoelter:masterfrom
avnoordenne:fix/whole-job-probe-margin
Oct 5, 2026
Merged

niklashoelter merged 1 commit into
niklashoelter:masterfrom
avnoordenne:fix/whole-job-probe-margin

Conversation

@avnoordenne

Copy link
Copy Markdown
Contributor

The whole-job shortcut ran a job as one batch whenever its forward-pass peak fit in max_memory_padding of free memory. Optimization needs more than one forward pass (optimizer state, fragmentation), so jobs at 0.78-0.80 of the budget passed the check and then ran out of memory mid-run (UMA-s, 878 structures, peak 18.1-18.8 of 23.4 GB).

Require peak * memory_scaling_factor to fit instead. The probe returns a batch one growth step below a size that fit, so this gives the shortcut the same margin. Jobs over the threshold fall back to the probe.

The whole-job shortcut ran a job as one batch whenever its forward-pass
peak fit in max_memory_padding of free memory. Optimization needs more
than one forward pass (optimizer state, fragmentation), so jobs at
0.78-0.80 of the budget passed the check and then ran out of memory
mid-run (UMA-s, 878 structures, peak 18.1-18.8 of 23.4 GB).

Require peak * memory_scaling_factor to fit instead. The probe returns
a batch one growth step below a size that fit, so this gives the
shortcut the same margin. Jobs over the threshold fall back to the probe.
@niklashoelter
niklashoelter merged commit c35ff48 into niklashoelter:master Oct 5, 2026
1 check passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants