Repository navigation
fix(ml): bound raw-model training memory for long windows - #92
Merged
Merged
Conversation
Attention memory grows with max_packets squared; a batch of 16 windows at 2048 packets needs about 12 GiB and the ONNX parity run about 17 GiB, so train jobs were OOM-killed in an 8 GiB pod. Cap the rows per forward pass from an activation estimate and a budget (memory_budget_mb, CORVUS_ML_MEMORY_BUDGET_MB, else half of the cgroup memory limit) for training, pretraining, the LR range test, validation, scoring, the kNN bank and the ONNX parity and benchmark runs. The configured batch stays the optimiser batch via gradient accumulation.
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
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.
Problem
A
trainjob withraw_hier_anomalyon a small snapshot was OOMKilled in an 8Gi pod afterthe 0.5.0 default of 2048 packets per window. The packet transformer attends over all packets,
so activation memory grows with
max_packetssquared.Measurements (P = 2048, default model, peak RSS growth)
score_batch16Change
corvus_ml/membudget.py: documented per-window activation estimate (attention + activations),budget (
memory_budget_mb,CORVUS_ML_MEMORY_BUDGET_MB, else half of the cgroupmemory.max,else RAM) and the micro-batch cap. A reduced micro-batch is logged.
fit,find_lrandpretrainaccumulate gradients over micro-batches, sobatch_sizestaysthe optimiser batch. The LR range test uses the candidate batch and the same micro-batch.
no_grad), the ONNX parity run and thebenchmark batch are capped too.
flatten_windowsandpretrainno longer copy the window tensor when not needed.ml/scripts/measure_raw_memory.pyreproduces the numbers.The masked-token loss is normalised per micro-batch, so accumulation is exact for mean losses
and approximate for it.