Option 1 for dynamic batching OOM safety#75
Open
TonyChen06 wants to merge 2 commits into
Open
Conversation
Pad each batch to its longest member in the collate rather than padding every sample to llm_input_len, saving compute on short batches. llm_input_len becomes the truncation ceiling only. - reps: trunc_pad_input no longer pads (returns unpadded when <= cap); drop the == llm_input_len term from the length asserts - Base.pad_to_batch: left-pad input_ids/mask/labels to the batch max and re-run find_signal_token_indices so injection positions track the shift - collate_fn: pad to batch max before the existing shape-assert + default_collate - torch.compile(dynamic=True): one graph for the now-varying seq length Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
With dynamic padding a max-length batch can first appear mid-epoch, whereas fixed padding stressed it on step ~1. Before training, stretch the first batch to llm_input_len and run one backward (params + grads + max activations), then add the optimizer's resident state and require headroom. No optimizer step / weight mutation, so training state stays pristine. - optimizer state is queried from the optimizer (Optimizer.estimated_state_bytes), not hardcoded: Adam/AdamW = 2 tensors/param, Muon = 1 momentum buffer/param - GPU memory for a fixed (B, L) shape varies a few % with attention-mask/padding (SDPA kernel choice), so we keep a 15% headroom margin rather than treat the probe as exact, and fail fast with an actionable message otherwise Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
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.
Stacked on #73 (dynamic padding). With dynamic padding a max-length batch may not
show up until mid-epoch, so a run can train for a while and then OOM.
This option:
Before the epoch loop, run one forward+backward on a worst-case
(batch_size x llm_input_len) batch and refuse to start if it won't fit.
We stretch the first batch to llm_input_len and run a backward -> params + grads + max-length activations.
tensors/param, Muon 1 momentum buffer/param).
Results:
model | optimizer | estimate | real full-length step | est/real
---------------+-----------+----------+-----------------------+---------
llama-3.2-1b | adamw | 20.33 GB | 20.34 GB | 1.000
llama-3.2-1b | muon | 18.41 GB | 18.41 GB | 1.000
qwen2.5-1.5b | adamw | 27.71 GB | 27.71 GB | 1.000
qwen2.5-1.5b | muon | 25.18 GB | 25.18 GB | 1.000
qwen3-1.7b | adamw | 29.69 GB | 29.69 GB | 1.000
qwen3-1.7b | muon | 26.95 GB | 26.95 GB | 1.000
Estimate matches a real full-length step to est/real = 1.000 across 3 model families × {adamw, muon}; other tiny memory spikes are covered by the 10% margin.
Downsides:
It's a good estimate but the real peak varies a bit for the same shape (SDPA kernel choice with
different padding). The 10% margin covers that, but it's a heuristic (and a magic number). Optimizer state is modeled from the optimizer type, not queried (so its correct for adam/adamw/muon, won't auto-adapt to a new one).
Upsides:
Cheap, self-contained, fails at startup if something needs to change. It works well on many different tested settings and hasn't failed once yet.
Other option does it the other way:
fixed-pad the first N real steps, then go dynamic. It's fully real but not as self-contained.
🤖 Generated with Claude Code