Hacktoberfest 2026: the issues maintainers tagged for October, open and beginner-friendly. Browse Hacktoberfest issues

TrainerRank CP2/EP1 normal admission reaches CUDA OOM in compiled MoE LoRA forward

Open
#949 3 comments 0 reactions 0 assignees View on GitHub

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

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, revision 995ad96eacd98c81ed38be0c5b274b04031597b0, 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 8ff205e543473d96557fe4b0ffbeaea19d13b420 with existing feature/instrumentation backports; Caladan base 257be1824d2dd7495e6e3509fbcad9af9ee74e9a. Executed trainer_rank/_impl.py SHA-256 6c0e9e68d32eb0067fa5206f0ff419c2840fd60922f0e90941c07b4486b1f313. This is not a clean-base repro.
  • Both actual ranks assert ART_TRAINER_RANK_ALLOW_OVERSIZED_BATCHES=1 and ART_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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from OpenPipe/ART

All issues in OpenPipe/ART

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.