TrainerRank CP2/EP1 normal admission reaches CUDA OOM in compiled MoE LoRA forward
Nobody has claimed this yet.
Assessment
- Difficulty
- 5/5
- Estimated time
- Over a week
- Newbie friendliness
- 25/100
- Issue type
- Bug
- Clarity
- Mostly clear
- Activity status
- Active
- Tech stack
- python, pytorch
- Domain
- ai-infra-agents, backend, machine-learning, performance
Research direction
The issue is a deep CUDA OOM in a compiled MoE LoRA forward pass under CP2/EP1 topology. Examine the _moe_output_bytes_per_token function and the generic static compute estimator for CP2. The failing generated code is in megatron.core.transformer.moe.experts.forward, specifically an Inductor-generated add. Look at the retained planner reports and the synthetic diagnostic configuration to understand the memory admission logic and the discrepancy between predicted and actual peak memory.
Written by the indexing model from the issue text.
Description
A fixed, offline synthetic diagnostic on two H200s, CP2/EP1 reached a CUDA OOM during the first large grouped forward after the planner selected admitted. Its selected required estimate was below the available budget. Oversized admission was enabled, but this event did not select admitted_oversized.
This is distinct from #948, which reported an explicit oversized override on CP1. Both traverse MoE/LoRA forward code; neither report establishes an identical cause. This new diagnostic preserves the original nested callback grouping, uses synthetic inputs/adapters, and does not reproduce the historical production batch or establish a regression on current main.
Actual failure, September 24, 2026, 7:01:25 AM MDT (13:01:25 UTC): rank 1, event 51a626adb7fe493c9726a222d92d9cf2, forward phase. The inner exception is torch.cuda.OutOfMemoryError, wrapped by TrainerRankPartialExecutionError. The outer harness's shallow cuda_oom:false classification describes the wrapper and is superseded by the retained inner exception/planner report.
The allocator message reports an attempted 6.68 GiB allocation with 6.07 GiB free, device capacity 139.81 GiB, process use including non-PyTorch memory 133.61 GiB, PyTorch allocated 113.02 GiB, and reserved but unallocated 14.74 GiB. These rounded message values do not prove fragmentation. The separately retained allocator counters are explicitly sampled post-unwind.
| Selected admission / report field | Bytes |
|---|---|
| Predicted peak above observation baseline | 44,681,050,584 |
| Admission required, safety factor 1.1 | 49,149,155,642 |
| Available admission budget | 66,481,491,149 |
| Physical free at selected admission | 70,985,121,792 |
| PyTorch allocated at selected admission | 70,191,249,408 |
| PyTorch reserved at selected admission | 72,706,162,688 |
| Device capacity | 150,121,021,440 |
| OOM partial peak above observation baseline | 54,895,467,520 |
| Observation baseline | 70,191,232,512 |
The partial peak is a lower bound from the interrupted observation, not a completed peak. It already exceeds the predicted peak and the admission estimate. observed_peak_bytes and error_pct are null; this report does not manufacture a completed underestimate percentage. Selected ordinal 5 carries the required estimate from ordinal 4; allocator is native, scope world, reserve 3%, no test memory cap. Initial required bytes were 1,944,298,502,665, followed by native replanning to the selected estimate.
Failed execution plan: 10 requests, 779,067 logical and packed tokens, two subforwards, with packed-token counts 400,511 and 378,556 and five GDN segments each. Gradients enabled, logits/hidden-state outputs disabled. Planner topology (DP, TP, CP, PP)=(1,1,2,1), coefficients [2,"gdn-moe-h2048-h200-bf16"]. No full-size forward wave completed and no full-size optimizer step occurred.
| Subforward | Request lengths in retained plan-group order | Logical/packed tokens |
|---|---|---|
| 0, synthetic policy | 94017, 89049, 80763, 72480, 64202 | 400511 |
| 1, synthetic adversary | 93189, 84906, 76618, 68338, 55505 | 378556 |
The retained stack ends in megatron.core.transformer.moe.experts.forward, ART expert FC2 LoRA, and an Inductor-generated add. The actual generated file was preserved before cache retirement. Its failing line 98 is:
s38 = arg1_1.size()[0]
assert_size_stride(arg1_1, (s38, 2048), (2048, 1))
assert_size_stride(arg2_1, (s38, 2048), (2048, 1))
buf0 = empty_strided_cuda((s38, 2048), (2048, 1), torch.bfloat16)
The runtime value of s38 was not retained, so the symbolic shape plus rounded 6.68 GiB message is not an exact element-count measurement.
Reproduction configuration and ordering:
- Model
Qwen/Qwen3.6-35B-A3B, revision995ad96eacd98c81ed38be0c5b274b04031597b0, BF16, LoRA rank 1 / alpha 32. DP1/TP1/CP2/PP1/EP1/ETP1 verified on both ranks; full uniform recomputation, one layer, compiled transformer layers. Each CP rank receives the identical complete input tree, not half of the 19 histories. - Python 3.12.13; Torch 2.11.0+cu128; CUDA 12.8; Megatron Core 0.17.0; Megatron Bridge 0.5.0+e1a207ac; Transformer Engine 2.11.0 with the frozen empty-group patch; Transformers 5.12.1. Two local H200s, 16 CPU affinity, OMP/MKL/OpenBLAS thread counts 1. This is an isolated exact runtime, not shared-environment package substitution.
- ART base
8ff205e543473d96557fe4b0ffbeaea19d13b420with existing feature/instrumentation backports; Caladan base257be1824d2dd7495e6e3509fbcad9af9ee74e9a. Executedtrainer_rank/_impl.pySHA-2566c0e9e68d32eb0067fa5206f0ff419c2840fd60922f0e90941c07b4486b1f313. This is not a clean-base repro. - Both actual ranks assert
ART_TRAINER_RANK_ALLOW_OVERSIZED_BATCHES=1andART_TRAINER_RANK_PLANNER_MISS_THRESHOLD_PCT=10. The estimator/admission implementation was unchanged by this diagnostic; bounded report retention was configured separately. - Use native collective adapter export and checkpoint loading to put identical synthetic weights into policy/adversary slots. Nonzero LoRA initialization, native runtime seed 593015. Rank 0's adapter is authoritative; synchronize before export and agree file completion before either rank loads. Reset both slots to the same adapter with absent optimizer state before each phase.
- First execute the original grouped callback on lengths
[64,56,72,64,64,56,72,64], two groups of four. This tiny phase passed forward/backward, one optimizer step, counters and CP replica checks. Reload both initial slots, restore baseline RNG, then run the large phase below; keep the warmed runtime/planner state.
lengths = [94017,93189,89049,84906,80763,76618,72480,68338,64202,
55505,51373,47238,43111,38974,32251,28111,23973,16565,12430]
g = torch.Generator(device="cpu").manual_seed(593016)
tokens = [torch.randint(16, vocab_size, (n + 1,), generator=g,
dtype=torch.int64) for n in lengths]
# Alternate policy/adversary slots by history index. Use the unchanged nested
# callback on groups of 10 and 9 whole histories, not singleton backwards.
# Sample mask: arange(n+1) % 7 == 2; first position false; logprobs -10.
# Synthetic group rewards +1/-1. Token values are independent uniform IDs.
These histories total 1,073,093 tokens. They deliberately lack production prefix sharing, routing values, trajectories and actual batch membership; the 10+9 grouping is synthetic. Original callback source SHA-256 55a7eea915d3709820b9a3cbd90d98468701edbbfe785d1552c9d6625f37dfab; exact two-rank harness SHA-256 0df7393c12ef8fa67665cbb3652d2fb1eabb56c7010bed41c028404a311110af. The complete harness, synthetic fixture and all five planner reports are retained. The metadata report alone has replay_complete=false (device_or_modified_input, layout_inputs_unavailable); this issue does not claim the JSON alone is an executable GPU reproducer.
Source questions: the frozen _moe_output_bytes_per_token returns zero outside TP1/CP1, and this CP2 report records zero. Full recomputation uses the generic static compute estimate; the split-memory floor is zero and is only learned after successful completion. These are concrete places to examine CP2 compiled MoE working-set coverage and smaller native subforward plans while preserving the callback grouping. They are not proven fixes, and the allocator's fragmentation suggestion has not been tested.
Expected behavior for normal fit admission is either successful execution or a conservative refusal/replan before this allocation. Actual outcome is terminal CUDA OOM. This diagnostic provides no full19 memory-fit or qualification credit and no cross-topology numerical-equivalence claim. No retry or proposed fix was run. All owned ranks, processes and guardian were verified gone, assigned GPUs empty, and private runtime/compile bulk retired; evidence remains outside cleanup targets.
Retained artifact SHA-256 values: original planner report b24df6f4d604e309247cad106557ac893ac5e9441ac1ca8e283d423cd214694f; sanitized planner projection b70ac2e7cd761ed4d3a69d18974eb4d277d82ea6f05807f07cdf97ae549bbf2f; failed generated code fb9c32f5ab99e0d07b892b29a53b9273a141b1622935b630cc123f0549a593fe. No private text, credentials or checkpoint weights are included in this issue.
- Dominant language
- Python
- Stars
- 10.8k
- Forks
- 989
- Avg merge
- 11h 53m
- Merged PRs (30d)
- 97
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
More from OpenPipe/ART
-
Difficulty 2/5 1-3 hours Newbie friendliness 72/100
-
Difficulty 5/5 Over a week Newbie friendliness 10/100
-
Difficulty 5/5 Over a week Newbie friendliness 42/100
-
enhancement
Difficulty 4/5 3-5 days Newbie friendliness 48/100
-
Difficulty 5/5 Over a week Newbie friendliness 35/100
Similar issues
-
agent-ready documentation needs-triage
Difficulty 1/5 1-3 hours Newbie friendliness 88/100
-
documentation
Difficulty 1/5 Under an hour Newbie friendliness 91/100
-
workflow-status page template still says reusable workflows are "triggered only by workflow_call:" Open
Difficulty 1/5 Under an hour Newbie friendliness 92/100
-
instance instance add
Difficulty 1/5 Under an hour Newbie friendliness 72/100
searxng/searx-instances#939 · 1 comment ·
-
area-deployment area-integrations triage:bot-seen
Difficulty 2/5 Half a day Newbie friendliness 86/100