[Dev] Add schedule-aware cross-stream tensor lifetime management

[lhb8125]

What does this PR do?

This PR adds an opt-in, schedule-aware lifetime path for the two cross-stream tensors retired by fine-grained EP overlap:

• forward inputs selected by the existing free_input policy;
• backward incoming gradients consumed on a stream different from the stream that produced them.

The implementation is deliberately local to one TransformerModelChunkSchedulePlan / microbatch. It does not maintain a process-wide tensor provenance registry, event generations, replay epochs, or autograd hooks.

For a normal schedule-local edge, the layer topology already identifies the owner stream:

producer/owner stream A -> consumer stream B -> later acquire of stream A

After the consumer returns, the manager stores one strong-reference payload in pending[owner_stream]. A later real node on the owner stream executes the existing plan-event wait first, then drains that queue:

stream B: wait(plan_event) -> consume -> record(plan_event) -> enqueue deferred release
stream A: wait(plan_event) -> resize storage / drop manager reference -> execute next node

The wait must precede resize_(0): resizing is a host-side allocator operation, so making the block allocator-visible before enqueueing the owner-stream wait would create a reuse race.

Forward retirement uses untyped_storage().resize_(0). Backward retirement only holds and later drops an extra strong reference; it never resizes gradient storage. Inputs from an unknown external boundary and detached extra gradients retain the conservative record_stream() fallback.

The feature is disabled by default:

--ep-overlap-use-scheduled-tensor-lifetime

Why?

record_stream() correctly prevents use-after-free, but a logically freed block remains pending until the recorded consumer event completes and the allocator observes completion. Under fine-grained overlap this can raise the allocator high-water mark when CPU enqueue gets ahead of GPU consumers or completion has a long tail.

The schedule already has the required producer/consumer dependency. Reusing that dependency lets the owner stream reclaim the tensor without introducing a host synchronize or an additional cross-microbatch dependency.

Corner cases

• Consecutive same-stream nodes retire immediately.
• NoopScheduleNode carries the same object through to the next real consumer.
• Dense and MTP layouts bind owner streams from their actual node topology; no N-2 offset is encoded.
• First/last-layer boundaries are drained by a terminal hand-back on the plan event.
• Pipeline-received gradients and detached extra gradients use record_stream() fallback.
• Each microbatch owns a separate manager, so one plan cannot drain another plan's queue.
• Partial CUDA Graph keeps the Python manager path; full capture records the wait/release edges, while replay does not execute Python manager code.

Validation

All tests and benchmarks below ran directly from commit 611cb8355. Focused unit tests used
one 8x B200 node; the wide/long-sequence memory and performance A/B used 32 GB300 GPUs.

• Lifetime manager, cross-microbatch isolation, Noop, terminal hand-back, external/detached fallback, allocator poison, and full graph replay: 10 passed.
• Model-chunk parity across feature off/on, MTP off/on, dense/MoE, alltoall/HybridEP, BF16/FP8, and multiple layer layouts: 97 passed.
• FSDP activation recompute/offload/early-release matrix: 8 passed.
• FSDP + delay-wgrad across alltoall/HybridEP and optim_grads_params/optim_grads: 4 passed.
• TE partial CUDA Graph warmup/capture/replay: 1 passed, 3 deselected.
• One long, combined multi-file run encountered transient NCCL/resource-state failures after extensive FSDP setup. The feature-specific FSDP matrices above passed when isolated in fresh containers. Separately, the full schedule-file run reached 109 passed before a baseline-only TP8 padding case hit a transient NCCL error; that exact case also passed in a fresh container.
• Megatron tools/autoformat.sh check and copyright check pass. The local lint environment reports the repository's existing missing dependency/type-stub mypy diagnostics; black, isort, pylint, and ruff pass.

Memory A/B

The hero case keeps DeepSeek-V3's full hidden/expert width while increasing the sequence length:
14 layers (3 dense + 11 MoE), hidden size 7168, FFN 18432, 256 experts, MoE FFN 2048,
MTP1, TP1/PP1/EP16, MBS2/GBS1024, sequence length 6144, MXFP8, HybridEP combined-1F1B,
eager mode, 32 GB300 GPUs. Each iteration contains 16 microbatches. Only rank 31 records
allocator history; control and feature snapshots contain 193,551 and 193,620 allocator events.

Rank-31 metric	Control	Feature	Delta
Target dispatch + grouped-GEMM simultaneous pending peak	16,210.877 MiB (13 blocks)	1,358.875 MiB (one transient block)	-14,852.002 MiB
All pending-free peak	23,132.062 MiB	7,827.518 MiB	-15,304.544 MiB
Max reserved	224,614 MiB	214,092 MiB	-10,522 MiB (-10.28 GiB)
Dispatch retirement p95 / max	17,591 / 1,884,967 us	1 / 20 us	-
Grouped-GEMM retirement p95 / max	10,340 / 1,883,448 us	1 / 22 us	-

The dominant dispatch/FC2 blocks are about 1.32 GiB each. Control has six dispatch and six
grouped-GEMM outputs simultaneously pending; feature reduces the target peak to one trace-level
free_requested -> free_completed transient, with zero target pending at its reserved high-water.
Reserved savings are reported from the actual allocator high-water and are not inferred from the
pending-byte delta.

Performance A/B

The same hero configuration was run without memory history for 12 steps per arm. To eliminate
GB300 node-group variation, feature explicitly reused the exact same eight nodes immediately after
control. Results exclude iterations 1-4 and the periodic-GC iteration 11.

Same-node steady metric (7 samples/arm)	Control	Feature	Delta
Iteration median	10,204.5 ms	10,203.4 ms	-0.011%
Iteration mean	10,201.6 ms	10,205.0 ms	+0.034%
Iteration CV	0.063%	0.078%	-
Throughput median	1,456.1 TFLOP/s/GPU	1,456.2 TFLOP/s/GPU	+0.007%

Both arms completed normally with no skipped or NaN iterations. Median and mean are well inside
the predefined +0.5% regression threshold, so this is treated as performance parity, not a speedup.

Scope

This PR does not change the other record_stream() call sites used by residual paths, shared experts, or loss-storage helpers. The new behavior remains opt-in while broader production-scale and full-iteration CUDA Graph coverage is collected.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.