transformers
21912efc - Reduce peak memory in examples_torch CI job (OOM fix) (#48983)

Commit
14 hours ago
Reduce peak memory in examples_torch CI job (OOM fix) (#48983) The examples_torch CI job intermittently hit the pod cgroup limit of 14,250 MB (exit code 137) when 4 xdist workers ran concurrently. Root cause: two tests dominated memory usage: - test_run_wav2vec2_pretraining: ~6 GB transient in backward() with 5 s audio (seq_len ≈ 1,250 frames) - test_run_vit_mae_pretraining: no --model_name_or_path → full ViT-Base (86 M params, 768 hidden / 12 enc + 8 dec layers) from scratch Fixes applied (peak SYS: 14,250 MB → ~4,200 MB, 5 consecutive passes): 1. MemoryCleanupMixin on ExamplesTests: gc.collect() + cache clear between tests, reducing inter-test PSS residual from ~1,700 → ~495 MB 2. All dataloader_num_workers / preprocessing_num_workers set to 0: eliminates fork → Copy-on-Write page duplication 3. wav2vec2 pretraining: max_duration 5 s → 1 s (seq_len 1,250 → 250) Peak PSS: ~6,081 MB → ~807 MB (-87%) 4. vit_mae pretraining: --config_overrides to tiny 32-hidden/2-layer encoder+decoder instead of full ViT-Base Peak PSS: ~5,500 MB → ~775 MB (-86%) 5. squad: bert-base-uncased (110 M) → tiny-random-bert (83 K), 30 steps Peak PSS: ~3,215 MB → ~774 MB (-76%) 6. swag: create BertForMultipleChoice from tiny config on the fly (run_swag.py lacks ignore_mismatched_sizes), torch.manual_seed(42) Peak PSS: ~2,447 MB → ~683 MB (-72%) 7. NER: tiny-random-bert + --ignore_mismatched_sizes, 10 epochs Peak PSS: ~2,297 MB → ~577 MB (-75%) 8. CLM: GPT2LMHeadModel from tiny config (n_embd=32, 2 layers), 25 epochs Peak PSS: ~2,916 MB → ~860 MB (-70%) 9. squad_seq2seq: t5-small (60 M) → t5-tinier-random (~1 M), 30 steps Peak PSS: ~2,436 MB → ~632 MB (-74%) 10. speech recognition (CTC adapter + seq2seq): --max_duration_in_seconds 3.0 Peak PSS: ~1,887 / ~1,687 MB → ~773 / ~771 MB (-54–59%) 11. glue: distilbert-base-uncased (67 M) → tiny-random-bert (83 K), 30 steps Peak PSS: ~1,652 MB → ~546 MB (-67%) 12. MLM: distilroberta-base (125 M) → tiny-random-bert (83 K), 20 epochs, seed=42; perplexity threshold relaxed 42 → 100 (calibrated for tiny random init; epoch-1 perplexity=934, epoch-20=87.5 — genuine 10× drop) Peak PSS: ~2,380 MB → ~883 MB (-63%) Co-authored-by: ydshieh <ydshieh@users.noreply.github.com> Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
Author
Parents
Loading