pytorch - 💡(How to fix) Fix `insert_overlap_deps=True` in `schedule_overlap_bucketing` breaks inductor when `graphsafe_run_with_rng_state` is present [1 participants]

🐛 Describe the bug

Description:

When schedule_overlap_bucketing is called with insert_overlap_deps=True on a graph that contains graphsafe_run_with_rng_state nodes (from activation checkpointing with nondeterministic ops like SDPA), inductor fails with an AssertionError in GraphLowering.placeholder.

The issue: insert_overlap_deps wraps nodes with control_deps (a HigherOrderOperator), including graphsafe_run_with_rng_state nodes. Inductor's placeholder handler for torch.Generator inputs asserts that the Generator placeholder has exactly one user whose target is either graphsafe_run_with_rng_state or invoke_subgraph:

  # torch/_inductor/graph.py, in GraphLowering.placeholder
  elif isinstance(example, torch.Generator):
      assert len(V.graph.current_node.users) == 1 and next(
          iter(V.graph.current_node.users)
      ).target in (
          torch._prims.rng_prims.graphsafe_run_with_rng_state,
          torch.ops.higher_order.invoke_subgraph,
      )

After insert_overlap_deps, the Generator placeholder's user is a control_deps node (whose target is a ControlDeps object), not graphsafe_run_with_rng_state directly. This fails the assertion.

Repro:

Any model that combines:

1. checkpoint_wrapper (or torch.utils.checkpoint.checkpoint) around a module containing a nondeterministic op (e.g., scaled_dot_product_attention)
2. torch.compile
3. schedule_overlap_bucketing(..., insert_overlap_deps=True) set as a post_grad_custom_post_pass

The nondeterministic op inside the checkpoint triggers functionalize_rng_ops in the partitioner, which introduces graphsafe_run_with_rng_state with Generator placeholders. Then insert_overlap_deps wraps these with control_deps, breaking the inductor assertion.

Possible fixes:

• schedule_overlap_bucketing should skip graphsafe_run_with_rng_state nodes when inserting overlap deps (they are RNG synchronization points, not compute/comm nodes that need reordering protection)
• Or inductor's assertion should recognize ControlDeps as a transparent wrapper and look through it to find the actual target

Versions

Torch 2.12.0.dev20260323+cu128

cc @soulitzer @pbelevich @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @kadeng @muchulee8 @amjames @aakhundov @coconutruben @jataylo @ydwu4 @bdhirsh @bobrenjc93 @aorenste