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>