Flaky test `test_fine_grained_activation_offloading.py::test_fine_grained_activation_offload_with_ep_a2a_overlap_compatibility`

[ko3n1g]

Describe the bug

This test is occasionally hanging on:

9-task-1-0/0 [default0]:PASSED
9-task-1-0/0 [default0]:tests/unit_tests/pipeline_parallel/test_fine_grained_activation_offloading.py::test_fine_grained_activation_offload_with_ep_a2a_overlap_compatibility[alltoall-False-offload_modules8] [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:WARNING:megatron.core.tensor_parallel.random:CPU RNG state changed within GPU RNG context
9-task-1-0/0 [default0]:
9-task-1-0/0 [default0]:================================================================================
9-task-1-0/0 [default0]:                        Activation Offload Summary (MB)                         
9-task-1-0/0 [default0]:================================================================================
9-task-1-0/0 [default0]:Rank       attn_proj   core_attn  expert_fc1    mlp_norm     moe_act       Total
9-task-1-0/0 [default0]:--------------------------------------------------------------------------------
9-task-1-0/0 [default0]:Rank 0         28.00     1204.00       51.90       28.00      207.80     1519.69
9-task-1-0/0 [default0]:Rank 1         28.00     1204.00       52.66       28.00      210.83     1523.49
9-task-1-0/0 [default0]:Rank 2         28.00     1204.00       57.17       28.00      228.91     1546.08
9-task-1-0/0 [default0]:Rank 3         28.00     1204.00       62.27       28.00      249.34     1571.61
9-task-1-0/0 [default0]:Rank 4         28.00     1204.00       51.90       28.00      207.80     1519.69
9-task-1-0/0 [default0]:Rank 5         28.00     1204.00       52.66       28.00      210.83     1523.49
9-task-1-0/0 [default0]:Rank 6         28.00     1204.00       57.17       28.00      228.91     1546.08
9-task-1-0/0 [default0]:Rank 7         28.00     1204.00       62.27       28.00      249.34     1571.61
9-task-1-0/0 [default0]:--------------------------------------------------------------------------------
9-task-1-0/0 [default0]:Total         224.00     9632.00      448.00      224.00     1793.75    12321.75
9-task-1-0/0 [default0]:================================================================================
9-task-1-0/0 [default0]:

Steps/Code to reproduce bug

Please list minimal steps or code snippet for us to be able to reproduce the bug.

A helpful guide on on how to craft a minimal bug report http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports.

Expected behavior

A clear and concise description of what you expected to happen.

Additional context

Add any other context about the problem here.

[ko3n1g]

Investigation summary (2026-05-26)

Current status

The test test_fine_grained_activation_offload_with_ep_a2a_overlap_compatibility is currently marked with @pytest.mark.flaky_in_dev on main (see tests/unit_tests/pipeline_parallel/test_fine_grained_activation_offloading.py:307). The marker was introduced when the test landed on main via #4500 (commit 3b1521e, 2026-05-01) — tests/unit_tests/run_ci_test.sh:108 excludes flaky_in_dev in the dev environment, so this case is not currently blocking CI.

No new CI failure issues referencing this test name have been filed since this issue was opened on 2026-03-19. The two sister issues (#4350, #4351) target different cases in the same file (test_fine_grained_activation_offloading_with_cuda_graph and test_mhc_fine_grained_activation_offloading_with_cuda_graph) with hard assertion failures, not hangs.

Observed symptom

Hang occurs after the test prints the "Activation Offload Summary" table — i.e. after the offload-enabled second 1F1B schedule completes its compute phase. No assertion failure, no stack trace; pure deadlock. Stage of execution suggests a NCCL collective or stream-sync mismatch during the post-iteration grad/output capture or teardown of the second model, when overlap_moe_expert_parallel_comm=True is combined with fine-grained activation offload and the largest module set (offload_modules8 = ["attn_norm", "core_attn", "attn_proj", "mlp_norm", "expert_fc1", "moe_act"]).

Likely root cause (not yet verified)

The test interleaves three asynchronous subsystems on different streams:

1. Fine-grained activation offload (D2H copies tied to a background stream via PipelineOffloadManager)
2. EP all-to-all overlap (overlap_moe_expert_parallel_comm, delay_wgrad_compute)
3. The TransformerModelChunkSchedulePlan.run() 1F1B scheduler

Plus singleton state — off_interface / PipelineOffloadManager.get_instance() — that persists across the baseline-and-offload model rebuild. The warmup→off_interface.reset()→steady-state pattern at lines 515–518 is the most suspicious window: if a prefetch / D2H copy from the warmup iteration is still in flight on the offload stream when reset() clears bookkeeping, a subsequent collective on the EP overlap stream can wait on a stale event.

Recommendation

The flaky_in_dev marker is doing its job; this issue is no longer P0 because CI is unblocked. A real fix needs (a) a reliable reproducer on Blackwell/H100 with the full 9-case parametrization run in a loop, and (b) an nsys trace showing the stream the hang lives on. Both require GPU time that I cannot allocate from this autonomous-triage path.

Suggested next steps for the assignee (@lhb8125):

• Audit PipelineOffloadManager.reset() / FineGrainedActivationOffloadingInterface.reset() for cross-stream event handling vs. in-flight D2H copies.
• Confirm set_streams() in megatron/core/pipeline_parallel/utils.py produces deterministic stream assignments across the baseline→offload model rebuild within a single test.
• If the test cannot be fixed in the next CI cycle, downgrade the parametrization to drop offload_modules8 (the all-modules case) and offload_modules7 (its MLA twin) — these are the cases that have actually been observed hanging in this issue and on an internal branch (84191ee7c9, c3ff98aba6 on helenn-make-hybrid-seqpack-graphable).

Stopping autonomous triage here per "if uncertain, comment investigation and stop" rule.

[lhb8125]

Hi @ko3n1g , can you share the repro steps? I cannot repro it on eos.

[ko3n1g]

Maybe this issue is healed nowadays? Did you remove the flaky_in_dev marker before trying to reproduce? Feel free to close the issue if the test passes without the marker

[lhb8125]

No, but I tried to repro it without -m 'not flaky_in_dev', so the test was executed without errors. Do you mean we can remove the marker by a PR? The description shows "This test is occasionally hanging on ...", but I am not sure when it fails.

[ko3n1g]

For repo, make sure to mount MLM into a CI container at /opt/megatron-lm. Remove the flaky_in_dev marker, then

bash /opt/megatron-lm/tests/unit_tests/run_ci_test.sh \
    --tag latest \
    --environment dev \
    --bucket 'tests/unit_tests/pipeline_parallel/**/*.py'
    --unit-test-repeat 5

this should reproduce CI behavior.. if it passes, we can open a PR to main that removes the marker