Option 2 for dynamic batching OOM safety#76
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>
Alternative to the fit-check. With dynamic padding a max-length batch may not appear until mid-epoch, so an OOM can strike late. Instead of a synthetic probe, fixed-pad the first WARMUP_FULL_BATCHES (20) real training batches to llm_input_len so the worst-case shape is exercised immediately -- real data, real optimizer state, real kernels. If it won't fit, it OOMs in the first step or two rather than mid-run. Padding is masked, so these are ordinary training steps (identical loss); after the warmup it falls back to dynamic padding. - gpu_manager.pad_batch_to_len left-pads a collated batch to llm_input_len - trainer pads the first 20 batches of epoch 0 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:
Fixed-pad the first 20 real training batches to llm_input_len, then fall back to
dynamic. The worst-case shape is exercised immediately, so if it won't fit the run
OOMs in the first step or two instead of mid-epoch.
No synthetic probe, these are ordinary training steps. Padding is masked, so the
loss is identical to dynamic (the warmup batches just train at full length). Real
data, real optimizer state, real kernels. After 20 batches it's back to dynamic.
Results (tested):
Downsides:
The 20 is a magic number.
Honestly, I thought this would have taken many more lines than option 1 but I guess it didn't need to. I would choose this one probably.
Upsides:
Simple, no synthetic batch, no modeled optimizer state.
Other option does it the other way:
a startup fit-check -- one synthetic forward+backward before the epoch loop, refuses
to start with an actionable message. Self-contained and fails earlier, but it's an
estimate (models the optimizer state + a 10% margin).
🤖 Generated with Claude Code