Memory planner ignores recompute_granularity: selective recompute admitted a 32k-token group at 5.9 GB predicted, then OOM'd at >107 GB (TP 2/4, Qwen3.8-27B)

[bradhilton]

Summary

TrainerRank.forward_micro_batches admits a micro-batch after a memory-planner check, but the planner's activation estimate does not depend on recompute_granularity. With ART_MEGATRON_RECOMPUTE_GRANULARITY=selective (method/num_layers unset) on a 64-layer dense model (Qwen/Qwen3.8-27B, hidden 5120) the planner predicted 5.9 GB for a ~32.7k-token packed group and admitted it; the forward then hit CUDA OOM against a 107–119 GB budget on both TP 2 and TP 4:

art.trainer_rank.TrainerRankMemoryError: forward_micro_batches: CUDA OOM occurred despite the planner estimate.
packed_tokens=32710 logical_tokens=38443 predicted_peak_gb=5.894 usable_limit_gb=107.192.
Use smaller top-level items, reduce output requests, or call dp_rank_forward with already-DP-local smaller inputs.

(TP 4 rank: predicted_peak_gb=5.894 usable_limit_gb=119.289, same failure.) The OOM surfaces inside _decoder_hidden → transformer_block.forward → _forward_mlp → art.megatron.lora.forward on the very first forward of the run, so no learned memory profile exists yet and the static estimate is the only guard.

Why the estimate is off

_estimate_required_memory_bytes_from_values (src/art/trainer_rank/_impl.py) uses

activation_factor = max(4, min(16, num_layers // 4 + 4))      # 16 for 64 layers
static_compute = packed_tokens * hidden_size * dtype_size * activation_factor
required = (output_bytes + static_compute) * 1.10

For 32,710 × 5,120 × 2 B × 16 × 1.10 that is ≈5.9 GB, which matches the observed footprint under full recompute (each layer retains roughly one hidden tensor). Under selective every layer retains its attention/MLP inputs, the MLP intermediate (ffn_hidden ≫ hidden) and norm outputs, i.e. tens of hidden-tensor equivalents per layer, so the true peak is one to two orders of magnitude larger and sharding by TP does not close the gap. The estimator has no term for the recompute mode (there is no reference to recompute_granularity anywhere in trainer_rank/_impl.py), and the provider defaults to full for every model (megatron/provider.py), so the constant is implicitly calibrated for full recompute only.

Reproduction

• Qwen/Qwen3.8-27B, rank-1 LoRA, Trainer(..., cluster="H200:2", num_processes_per_node=2)
• env: ART_MEGATRON_CONTEXT_PARALLEL_SIZE=1 ART_MEGATRON_TENSOR_MODEL_PARALLEL_SIZE=2 ART_MEGATRON_RECOMPUTE_GRANULARITY=selective ART_MEGATRON_RECOMPUTE_METHOD=none ART_MEGATRON_RECOMPUTE_NUM_LAYERS=none
• one forward_micro_batches call whose top-level item is a group of two ~16–20k-token sequences (38,443 logical tokens, 32,710 packed) with hidden_states requested → OOM on step 1. The same call with the default full recompute runs at 41 GiB peak per rank on TP 2 (22 GiB on TP 4).

Driver: caladan experiments/058-rank1-probe-bases.py probe ... --gpus 2 --tp 2 --no-recompute; ART checkout bcfe13519 (2026-09-08); the cluster ran caladan's pinned build.

Suggested fix

Any of these would have avoided the crash:

1. Scale activation_factor by the provider's effective recompute mode (full → current constant; selective → a per-layer retained-activation estimate that includes ffn_hidden_size and the attention inputs; none → larger still) so the check declines / splits instead of admitting.
2. Failing that, treat an unprofiled signature under non-full recompute as untrusted: run a probe forward on a small slice first, or cap the first admitted micro-batch until a profile exists.
3. At minimum, mention the recompute mode in the TrainerRankMemoryError hint, since "use smaller top-level items" is not actionable when the item is already a single (indivisible) group and the real remedy is switching back to full recompute.

Related: #911 (context-parallel forward returns rank-local shards; that is why these runs were on TP rather than the default CP topology).

[bradhilton]

Data point from a caladan 059 lane on ART a3a248a (pre-#915, full recompute, Qwen3.6-35B-A3B rank-1 + rank-8 LoRAs, TP 2 on 2×H200), in case it is useful for the calibration work: a no-grad forward_micro_batches call over 64 dynamics-view sequences was packed into a single micro-batch of packed_tokens=1196952 logical_tokens=1344312 with predicted_peak_gb=92.955 against usable_limit_gb=93.122 (admitted at 99.8% of the budget). It then failed allocating 10.62 GiB with 95.4 GiB allocated and 30.5 GiB reserved-but-unallocated, i.e. fragmentation, not the estimate itself, consumed the margin. Two sibling lanes failed identically (10.3 and 11.2 GiB allocations).

Two thoughts: (1) the 10% safety factor is smaller than the allocator fragmentation we see at this scale, so admitting up to the limit on a ~1.2M-token no-grad pack is risky unless expandable_segments is on (caladan main enables it since #441; this pin predates it); (2) a no-grad pass that large still materializes full-vocabulary logits per chunk, which is the allocation that actually failed here. We moved to TP 4 with expandable segments.