Skip to content

fix(ml): bound raw-model training memory for long windows - #92

Merged
TheMeinerLP merged 1 commit into
mainfrom
fix/ml-raw-memory
Oct 16, 2026
Merged

TheMeinerLP merged 1 commit into
mainfrom
fix/ml-raw-memory

Conversation

@TheMeinerLP

Copy link
Copy Markdown
Contributor

Problem

A train job with raw_hier_anomaly on a small snapshot was OOMKilled in an 8Gi pod after
the 0.5.0 default of 2048 packets per window. The packet transformer attends over all packets,
so activation memory grows with max_packets squared.

Measurements (P = 2048, default model, peak RSS growth)

Stage Before After
training step / LR range test, batch 16 ~12.4 GiB (6.5 GiB at batch 8, measured) 3.4 GiB (micro-batch 4)
scoring, score_batch 16 ~4.2 GiB (1.05 GiB at 4) 3.4 GiB (micro-batch 13)
validation pass over a split ~0.26 GiB per window, all at once capped like scoring
ONNX parity run, 64 windows 16.8 GiB 3.8 GiB (13 rows per run)

Change

  • corvus_ml/membudget.py: documented per-window activation estimate (attention + activations),
    budget (memory_budget_mb, CORVUS_ML_MEMORY_BUDGET_MB, else half of the cgroup memory.max,
    else RAM) and the micro-batch cap. A reduced micro-batch is logged.
  • fit, find_lr and pretrain accumulate gradients over micro-batches, so batch_size stays
    the optimiser batch. The LR range test uses the candidate batch and the same micro-batch.
  • Validation, scoring, the kNN memory bank (now under no_grad), the ONNX parity run and the
    benchmark batch are capped too.
  • flatten_windows and pretrain no longer copy the window tensor when not needed.
  • ml/scripts/measure_raw_memory.py reproduces the numbers.
  • Docs: memory behaviour and rule of thumb in raw-byte-tokenizer.md, parameter in ml-jobs.md.

The masked-token loss is normalised per micro-batch, so accumulation is exact for mean losses
and approximate for it.

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.
@TheMeinerLP
TheMeinerLP merged commit f2fdfa4 into main Oct 16, 2026
13 checks passed
@TheMeinerLP
TheMeinerLP deleted the fix/ml-raw-memory branch October 16, 2026 14:02
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.

1 participant