Stream Capture of NCCL into a conditional node's graph after the current cuda graph has already done stream capture results in cudaErrorStreamCaptureMerge

### 🐛 Describe the bug

Oh HEAD today, if you want to use NCCL inside of a cuda graph, as well as inside of a conditional node (in this case, an IF node via torch.cond()) after that first NCCL call, that will result in a failure.

The reason why is that the current device's cuda stream in [ncclStreams_](https://github.com/pytorch/pytorch/blob/466bed4c5df845ae5403742378d0ece0f37f8b39/torch/csrc/distributed/c10d/ProcessGroupNCCL.hpp#L1402) dictionary is keyed based on the current device. Pytorch was written with the idea that at most one stream capture will happen at a time.

However, with conditional nodes, we can actually have two active stream captures at once on the same device. The reason why is that a conditional node's graph is built via stream capture while the parent cuda graph still has not exited stream capture.

Once a stream is capturing to a particular graph, the only way to "reset it" so that it can capture to another graph is to end the first stream capture with cudaStreamEndCapture(), but that hasn't happened yet in this case. Thus, ncclStreams_ entry for this device will already be capturing to the parent graph by the time that the conditional node's graph tries to join against it via a cudaEvent_t. This results in the error: cudaErrorStreamCaptureMerge. The exact failing code is right here:

https://github.com/pytorch/pytorch/blob/466bed4c5df845ae5403742378d0ece0f37f8b39/torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp#L3782-L3787

This problem can be reproduced with this script. Sorry it's so long.

```
#!/usr/bin/env python3
"""Reproduce NCCL stream reuse across a parent and conditional child capture.

Run with one process per GPU, for example:

torchrun --standalone --nproc-per-node=2 \
agent_space/repro_nccl_parent_then_cond.py

The important ordering is:

1. Start the parent CUDA graph capture.
2. Run NCCL in the parent, making the communicator stream join that capture.
3. Enter torch.cond(), whose branch bodies are captured into child graphs.
4. Run NCCL again from a branch body.

The graph is intentionally captured with ``keep_graph=True`` and is not
instantiated. Instantiation would hit the separate restriction on event nodes
inside conditional-node bodies and obscure this capture-time failure.

After detecting the expected capture-merge error, each worker exits immediately.
CUDA has invalidated both nested captures at that point, so normal context and
process-group teardown would emit secondary errors and can wait indefinitely.
"""

import argparse
import os

import torch
import torch.distributed as dist
from torch._higher_order_ops.cudagraph_conditional_nodes import (
CUDAGraphCaptureControlFlowOpDispatchMode,
)

def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser()
parser.add_argument(
"--skip-parent-nccl",
action="store_true",
help="Control case: do not put NCCL's stream in the parent capture first.",
)
return parser.parse_args()

def main() -> None:
args = parse_args()
rank = int(os.environ["RANK"])
local_rank = int(os.environ["LOCAL_RANK"])

torch.cuda.set_device(local_rank)
# Pin this reproducer to the ProcessGroupNCCL implementation with one
# persistent communication stream per communicator.
dist.init_process_group("nccl-legacy")
group_name = dist.group.WORLD.group_name

def all_reduce(tensor: torch.Tensor) -> torch.Tensor:
result = torch.ops._c10d_functional.all_reduce(
tensor, "sum", group_name
)
return torch.ops._c10d_functional.wait_tensor(result)

def true_branch(tensor: torch.Tensor) -> torch.Tensor:
return all_reduce(tensor + 1)

def false_branch(tensor: torch.Tensor) -> torch.Tensor:
return tensor - 1

def model(
tensor: torch.Tensor, predicate: torch.Tensor
) -> torch.Tensor:
if not args.skip_parent_nccl:
tensor = all_reduce(tensor)
return torch.cond(
predicate,
true_branch,
false_branch,
(tensor,),
)

tensor = torch.full(
(4,), rank + 1, dtype=torch.float32, device=local_rank
)
predicate = torch.tensor(True, device=local_rank)
capture_stream = torch.cuda.Stream(device=local_rank)

# Warm up tracing and both NCCL branch paths eagerly. Do not use
# ControlFlowOpWarmupDispatchMode here: it instantiates a temporary
# conditional graph and hits the separate unsupported-event-node issue
# before the parent-versus-child capture conflict can be exercised.
with torch.cuda.stream(capture_stream):
model(tensor, predicate)
predicate.fill_(False)
model(tensor, predicate)
predicate.fill_(True)
torch.cuda.synchronize()
dist.barrier()

if rank == 0:
parent = "without" if args.skip_parent_nccl else "with"
print(
f"capturing torch.cond {parent} parent NCCL first",
flush=True,
)

graph = torch.cuda.CUDAGraph(keep_graph=True)
with (
torch.cuda.graph(graph, stream=capture_stream),
CUDAGraphCaptureControlFlowOpDispatchMode(),
):
try:
model(tensor, predicate)
except torch.AcceleratorError as error:
merge_error: BaseException | None = error
while (
merge_error is not None
and "merge of separate capture sequences" not in str(merge_error)
):
merge_error = merge_error.__cause__ or merge_error.__context__
if merge_error is None:
raise
print(
"rank "
f"{rank}: reproduced cudaErrorStreamCaptureMerge: {merge_error}",
flush=True,
)
os._exit(1)

if rank == 0:
print("capture succeeded (graph intentionally not instantiated)", flush=True)

graph.reset()
dist.destroy_process_group()

if __name__ == "__main__":
main()

```

Some more context.

First: I originally thought that NCCL itself was the cause of this issue, but I tried to reproducer this issue with a pure C file that directly calls NCCL to stream capture a cuda graph with conditional nodes. It turns out that the NCCL strong stream concept will actually create a new cudaStream_t object of its own any time a new stream capture is encountered (a stream capture can be uniquely identified by the capture ID returned by cudaStreamGetCaptureInfo()).

https://github.com/NVIDIA/nccl/blob/5067397c2676d5aed50042fc39e5c8ee96eb0027/src/misc/strongstream.cc#L211-L216

Second: Even if we succeed in stream capturing to a cuda graph, instantiation will fail with NCCL unless you run with [NCCL_GRAPH_MIXING_SUPPORT](https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/env.html#nccl-graph-mixing-support)=0 *AND* your workload does not require a proxy CPU thread. This is because conditional nodes don't support cuda event nodes or host nodes. Graph mixing support requires cuda event nodes, and proxy CPU threads require host nodes. This could be improved at some point.

Third: This issue is probably pretty straightforward to fix. We've already encountered something similar with RNG state. There was only a single global RNG state for stream capture, which did not mesh correctly with conditional node.s There is some preliminary work on that here: https://github.com/pytorch/pytorch/pull/176753

### Versions

Reproducible on HEAD today.

cc @awgu @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @pragupta @msaroufim @dcci @aditvenk @weifengpy @kapilsh