[perf]: route SP all-to-all NCCL on a dedicated comm stream - #1395
Open
Mister-Raggs wants to merge 2 commits into
Open
[perf]: route SP all-to-all NCCL on a dedicated comm stream#1395Mister-Raggs wants to merge 2 commits into
Mister-Raggs wants to merge 2 commits into
Conversation
Adds a lazy class-level `_comm_stream` on `DistributedAutograd` and routes the two `dist.all_to_all_single` calls inside the SP `AllToAll4D.forward` collective onto it via the standard `wait_stream` discipline. Extracts the forward shape-munging into a static `_all_to_all_4D_forward` helper so the inference path can bypass `autograd.Function.apply` dispatch overhead when grad is off. **Scope (intentionally narrow).** Only `AllToAll4D` is wrapped, not the other NCCL collectives (`AllReduce`, `AllGather`, `Slice.backward`). Per the profiling work in [`nccl-optimization-status.md`](https://github.com/rich7420/fastvideo/blob/inference-profile/nccl-optimization-status.md), SendRecv from `all_to_all_4D` is 97.7% of NCCL on the 4×L40S Wan T2V baseline; the other collectives are negligible. Wrapping them adds noise without measurable benefit. **No `record_stream`.** Modern PyTorch defaults `TORCH_NCCL_AVOID_RECORD_STREAMS=1` (deprecation warning visible in worker logs); the NCCL backend tracks cross-stream lifetime itself. Explicit `record_stream()` on collective inputs/outputs would force the deprecated slow path and inflate the caching allocator's per-collective bookkeeping ~2×. `wait_stream` ordering alone is sufficient — see the docstring on `_comm_stream` for detail. **No e2e perf claim on this PR.** The dedicated stream is necessary but not sufficient for comm/compute overlap on PCIe (per `nccl-optimization-status.md`: NCCL on L40S is bandwidth-bound, not stream-bound; software stream rearrangement alone doesn't move the needle). This change lands the infrastructure that future overlap-based work (e.g. Hybrid Ulysses-Ring rotation) needs in place. The `_functional_collectives` alternative was already tried and rolled back (+1% wall regression from AsyncCollectiveTensor wrapper overhead) — see status doc. **Inference fast-path.** `DeviceCommunicatorBase.all_to_all_4D` now calls `_all_to_all_4D_forward` directly when `torch.is_grad_enabled()` is False, skipping the `autograd.Function.apply` machinery. Small per-call dispatch saving (the SP all-to-all is on every attention call inside the denoising loop); cumulative effect to be measured in a follow-up if signal worth chasing. Co-Authored-By: KUAN-HAO HUANG <rich7420@gmail.com>
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request introduces a dedicated CUDA stream for 4D all-to-all collectives to prevent serialization with the default compute stream, and optimizes inference by bypassing the autograd dispatch overhead when gradients are disabled. The reviewer provided valuable feedback pointing out that reusing a single stream across multiple devices can cause runtime errors, suggesting a device-mapped dictionary of streams instead. Additionally, the reviewer recommended guarding CUDA stream operations to prevent failures when processing CPU-only tensors.
Comment on lines
+20
to
+36
| # Dedicated stream for SP all-to-all collectives. Lazily created on first | ||
| # use to avoid touching CUDA before the device is selected. NCCL kernels | ||
| # issued on this stream no longer serialize with compute on the default | ||
| # stream, freeing GPU launch queue pressure. | ||
| # | ||
| # NOTE: we intentionally do NOT call .record_stream() on the input/output | ||
| # tensors. PyTorch's NCCL backend now defaults TORCH_NCCL_AVOID_RECORD_STREAMS=1 | ||
| # and handles cross-stream lifetime via its own mechanism — calling | ||
| # record_stream() explicitly forces the deprecated slow path and inflates | ||
| # caching-allocator bookkeeping by 2× per collective. | ||
| _comm_stream: "torch.cuda.Stream | None" = None | ||
|
|
||
| @staticmethod | ||
| def _get_comm_stream() -> "torch.cuda.Stream": | ||
| if DistributedAutograd._comm_stream is None: | ||
| DistributedAutograd._comm_stream = torch.cuda.Stream() | ||
| return DistributedAutograd._comm_stream |
Contributor
There was a problem hiding this comment.
Reusing a single torch.cuda.Stream across different devices (e.g., in multi-GPU testing or single-process multi-device environments) will cause runtime errors or silent correctness issues because a stream is bound to a specific device.
To fix this, we should store a dictionary of streams mapped by device, and lazily create them on the correct device context.
Suggested change
| # Dedicated stream for SP all-to-all collectives. Lazily created on first | |
| # use to avoid touching CUDA before the device is selected. NCCL kernels | |
| # issued on this stream no longer serialize with compute on the default | |
| # stream, freeing GPU launch queue pressure. | |
| # | |
| # NOTE: we intentionally do NOT call .record_stream() on the input/output | |
| # tensors. PyTorch's NCCL backend now defaults TORCH_NCCL_AVOID_RECORD_STREAMS=1 | |
| # and handles cross-stream lifetime via its own mechanism — calling | |
| # record_stream() explicitly forces the deprecated slow path and inflates | |
| # caching-allocator bookkeeping by 2× per collective. | |
| _comm_stream: "torch.cuda.Stream | None" = None | |
| @staticmethod | |
| def _get_comm_stream() -> "torch.cuda.Stream": | |
| if DistributedAutograd._comm_stream is None: | |
| DistributedAutograd._comm_stream = torch.cuda.Stream() | |
| return DistributedAutograd._comm_stream | |
| # Dedicated streams for SP all-to-all collectives, mapped by device. | |
| # Lazily created on first use to avoid touching CUDA before the device is selected. | |
| # NCCL kernels issued on these streams no longer serialize with compute on the default | |
| # stream, freeing GPU launch queue pressure. | |
| # | |
| # NOTE: we intentionally do NOT call .record_stream() on the input/output | |
| # tensors. PyTorch's NCCL backend now defaults TORCH_NCCL_AVOID_RECORD_STREAMS=1 | |
| # and handles cross-stream lifetime via its own mechanism — calling | |
| # record_stream() explicitly forces the deprecated slow path and inflates | |
| # caching-allocator bookkeeping by 2× per collective. | |
| _comm_streams: dict[torch.device, "torch.cuda.Stream"] = {} | |
| @staticmethod | |
| def _get_comm_stream(device: torch.device) -> "torch.cuda.Stream": | |
| if device not in DistributedAutograd._comm_streams: | |
| with torch.cuda.device(device): | |
| DistributedAutograd._comm_streams[device] = torch.cuda.Stream(device=device) | |
| return DistributedAutograd._comm_streams[device] |
Comment on lines
+141
to
+198
| @staticmethod | ||
| def _all_to_all_4D_forward(group: ProcessGroup, input_: Tensor, world_size: int, scatter_dim: int, | ||
| gather_dim: int) -> Tensor: | ||
| """Pure forward logic for 4D all-to-all. Used by both the autograd | ||
| Function and the inference fast-path that bypasses autograd.""" | ||
| if world_size == 1: | ||
| return input_ | ||
|
|
||
| if world_size == 1: | ||
| return input_ | ||
| assert input_.dim() == 4, f"input must be 4D tensor, got {input_.dim()} and shape {input_.shape}" | ||
|
|
||
| assert input_.dim() == 4, f"input must be 4D tensor, got {input_.dim()} and shape {input_.shape}" | ||
| comm = DistributedAutograd._get_comm_stream() | ||
| curr = torch.cuda.current_stream() | ||
|
|
||
| if scatter_dim == 2 and gather_dim == 1: | ||
| bs, shard_seqlen, hn, hd = input_.shape | ||
| seqlen = shard_seqlen * world_size | ||
| shard_hn = hn // world_size | ||
| if scatter_dim == 2 and gather_dim == 1: | ||
| bs, shard_seqlen, hn, hd = input_.shape | ||
| shard_hn = hn // world_size | ||
|
|
||
| input_ = input_.transpose(0, 2).contiguous() # hn, shard_seqlen, bs, hd | ||
| output = torch.empty_like(input_) | ||
| input_ = input_.transpose(0, 2).contiguous() # hn, shard_seqlen, bs, hd | ||
| output = torch.empty_like(input_) | ||
|
|
||
| # Issue NCCL on the dedicated comm stream so it doesn't serialize | ||
| # with compute on the default stream. No record_stream — see the | ||
| # _comm_stream class docstring for why. | ||
| comm.wait_stream(curr) | ||
| with torch.cuda.stream(comm): | ||
| dist.all_to_all_single(output, input_, group=group) # hn, shard_seqlen, bs, hd | ||
| curr.wait_stream(comm) | ||
|
|
||
| output = torch.cat(output.split(shard_hn), dim=1) # sharded hn, seqlen, bs, hd | ||
| output = torch.cat(output.split(shard_hn), dim=1) # sharded hn, seqlen, bs, hd | ||
|
|
||
| output = output.transpose(0, 2).contiguous() # bs, seqlen, sharded_hn, hd | ||
| output = output.transpose(0, 2).contiguous() # bs, seqlen, sharded_hn, hd | ||
|
|
||
| return output | ||
| elif scatter_dim == 1 and gather_dim == 2: | ||
| bs, seqlen, shard_hn, hd = input_.shape | ||
| hn = shard_hn * world_size | ||
| shard_seqlen = seqlen // world_size | ||
| return output | ||
| elif scatter_dim == 1 and gather_dim == 2: | ||
| bs, seqlen, shard_hn, hd = input_.shape | ||
| shard_seqlen = seqlen // world_size | ||
|
|
||
| input_ = input_.transpose(0, 2).contiguous() # shard_hn, seqlen, bs, hd | ||
| input_ = input_.transpose(0, 2).contiguous() # shard_hn, seqlen, bs, hd | ||
|
|
||
| input_ = input_.reshape(shard_hn, world_size, shard_seqlen, bs, | ||
| hd).transpose(0, 1).reshape(shard_hn * world_size, shard_seqlen, bs, | ||
| hd).contiguous() | ||
| input_ = input_.reshape(shard_hn, world_size, shard_seqlen, bs, | ||
| hd).transpose(0, 1).reshape(shard_hn * world_size, shard_seqlen, bs, | ||
| hd).contiguous() | ||
|
|
||
| output = torch.empty_like(input_) | ||
| output = torch.empty_like(input_) | ||
|
|
||
| comm.wait_stream(curr) | ||
| with torch.cuda.stream(comm): | ||
| dist.all_to_all_single(output, input_, group=group) | ||
| curr.wait_stream(comm) | ||
|
|
||
| output = output.transpose(0, 2).contiguous() # bs, seqlen, sharded_hn, hd | ||
| output = output.transpose(0, 2).contiguous() # bs, seqlen, sharded_hn, hd | ||
|
|
||
| return output | ||
| else: | ||
| raise RuntimeError( | ||
| f"Invalid scatter_dim={scatter_dim}, gather_dim={gather_dim}. " | ||
| f"Only (scatter_dim=2, gather_dim=1) and (scatter_dim=1, gather_dim=2) are supported.") | ||
| return output | ||
| else: | ||
| raise RuntimeError( | ||
| f"Invalid scatter_dim={scatter_dim}, gather_dim={gather_dim}. " | ||
| f"Only (scatter_dim=2, gather_dim=1) and (scatter_dim=1, gather_dim=2) are supported.") | ||
|
|
Contributor
There was a problem hiding this comment.
If the input tensor is on CPU (e.g., during CPU-only testing or execution), calling torch.cuda.Stream() or torch.cuda.current_stream() will raise a RuntimeError or AssertionError because CUDA is not available or not used.
We should guard the stream creation and synchronization so that they are only executed when the input tensor is on a CUDA device. Additionally, we should pass the input tensor's device to _get_comm_stream and torch.cuda.current_stream to ensure the correct device context is used.
@staticmethod
def _all_to_all_4D_forward(group: ProcessGroup, input_: Tensor, world_size: int, scatter_dim: int,
gather_dim: int) -> Tensor:
"""Pure forward logic for 4D all-to-all. Used by both the autograd
Function and the inference fast-path that bypasses autograd."""
if world_size == 1:
return input_
assert input_.dim() == 4, f"input must be 4D tensor, got {input_.dim()} and shape {input_.shape}"
use_cuda = input_.is_cuda
if use_cuda:
comm = DistributedAutograd._get_comm_stream(input_.device)
curr = torch.cuda.current_stream(device=input_.device)
if scatter_dim == 2 and gather_dim == 1:
bs, shard_seqlen, hn, hd = input_.shape
shard_hn = hn // world_size
input_ = input_.transpose(0, 2).contiguous() # hn, shard_seqlen, bs, hd
output = torch.empty_like(input_)
# Issue NCCL on the dedicated comm stream so it doesn't serialize
# with compute on the default stream. No record_stream — see the
# _comm_stream class docstring for why.
if use_cuda:
comm.wait_stream(curr)
with torch.cuda.stream(comm):
dist.all_to_all_single(output, input_, group=group) # hn, shard_seqlen, bs, hd
curr.wait_stream(comm)
else:
dist.all_to_all_single(output, input_, group=group)
output = torch.cat(output.split(shard_hn), dim=1) # sharded hn, seqlen, bs, hd
output = output.transpose(0, 2).contiguous() # bs, seqlen, sharded_hn, hd
return output
elif scatter_dim == 1 and gather_dim == 2:
bs, seqlen, shard_hn, hd = input_.shape
shard_seqlen = seqlen // world_size
input_ = input_.transpose(0, 2).contiguous() # shard_hn, seqlen, bs, hd
input_ = input_.reshape(shard_hn, world_size, shard_seqlen, bs,
hd).transpose(0, 1).reshape(shard_hn * world_size, shard_seqlen, bs,
hd).contiguous()
output = torch.empty_like(input_)
if use_cuda:
comm.wait_stream(curr)
with torch.cuda.stream(comm):
dist.all_to_all_single(output, input_, group=group)
curr.wait_stream(comm)
else:
dist.all_to_all_single(output, input_, group=group)
output = output.transpose(0, 2).contiguous() # bs, seqlen, sharded_hn, hd
return output
else:
raise RuntimeError(
f"Invalid scatter_dim={scatter_dim}, gather_dim={gather_dim}. "
f"Only (scatter_dim=2, gather_dim=1) and (scatter_dim=1, gather_dim=2) are supported.")There was a problem hiding this comment.
Pull request overview
Note
Copilot was unable to run its full agentic suite in this review.
This PR refactors the 4D all-to-all implementation to reduce inference overhead and improve GPU scheduling by moving NCCL collectives onto a dedicated CUDA stream and adding an autograd-bypass fast path when gradients are disabled.
Changes:
- Added a lazily-initialized dedicated CUDA communication stream for SP all-to-all collectives.
- Extracted 4D all-to-all forward logic into a shared helper used by both autograd and inference paths.
- Updated
all_to_all_4Dto bypassautograd.Functiondispatch whentorch.is_grad_enabled()is false.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+30
to
+37
| _comm_stream: "torch.cuda.Stream | None" = None | ||
|
|
||
| @staticmethod | ||
| def _get_comm_stream() -> "torch.cuda.Stream": | ||
| if DistributedAutograd._comm_stream is None: | ||
| DistributedAutograd._comm_stream = torch.cuda.Stream() | ||
| return DistributedAutograd._comm_stream | ||
|
|
Comment on lines
+151
to
+152
| comm = DistributedAutograd._get_comm_stream() | ||
| curr = torch.cuda.current_stream() |
|
|
||
| if world_size == 1: | ||
| return input_ | ||
| assert input_.dim() == 4, f"input must be 4D tensor, got {input_.dim()} and shape {input_.shape}" |
Comment on lines
+32
to
+36
| @staticmethod | ||
| def _get_comm_stream() -> "torch.cuda.Stream": | ||
| if DistributedAutograd._comm_stream is None: | ||
| DistributedAutograd._comm_stream = torch.cuda.Stream() | ||
| return DistributedAutograd._comm_stream |
mergify
Bot
added
type: perf
Contributor
Merge Protections🔴 1 of 1 protections blocking · waiting on 👀 reviews and 🤖 CI
🔴 PR merge requirementsWaiting for
This rule is failing.
|
Four reviewer findings on bc5ea6b, all in base_device_communicator.py: 1. Per-device stream cache (gemini-high + Copilot-high). CUDA streams are device-bound; the single global _comm_stream would crash any multi-GPU / single-process multi-device context. Replaced with _comm_streams: dict[torch.device, Stream], lazily populated under `with torch.cuda.device(device):`. 2. CPU-path guard (gemini-high). use_cuda = input_.is_cuda; stream wait/wrap only when CUDA, plain dist.all_to_all_single fallback for CPU paths used by some unit tests. 3. assert -> ValueError (Copilot-medium). Asserts strip under python -O; ValueError is the correct runtime validation here. 4. Thread-safe lazy init (Copilot-medium). threading.Lock + double-checked locking around dict insertion to handle concurrent first-access. Scheduling/lifecycle only — no change to collective semantics or output. Kuan's A/B numbers on bc5ea6b (SSIM 1.000000, wall delta inside noise) remain valid in spirit; happy to re-run on the new tip if a maintainer asks.
Mister-Raggs
added a commit
to Mister-Raggs/FastVideo
that referenced
this pull request
May 30, 2026
Single-container two-pass A/B harness modeled on Kuan's nccl_stream_ab.py pattern (used to validate hao-ai-lab#1395). Runs the same seed-pinned prompts twice through the same VideoGenerator (once sequential, once batched), computes pairwise SSIM via pytorch_msssim (matches fastvideo/tests/utils.py:compute_video_ssim_torchvision), prints a per-prompt + total wall/SSIM table. Architecture: the harness ferries pass config + records over JSON files and subprocesses the inner script via /opt/venv/bin/python. Modal's main function process runs in its own add_python="3.12" layer where FastVideo isn't importable; the venv subprocess works around that. Same reason ssim_test.py runs pytest as a subprocess. Features: - Mode tag in output dirs ({eager,compile}[-{dit_precision}]) so multiple legs don't clobber each other in the hf-model-weights Volume. - Recovery function reads existing mp4s by mtime and recomputes SSIM, for post-run gap-fill if the in-run SSIM step missed. - Prebuilt FA3 wheel install on Hopper (autodetected from gpu kwarg) via mjun0812/flash-attention-prebuild-wheels v0.9.4 — the same release that supplies the FA2 wheel already baked into the fastvideo-dev image. ~30s install vs Kuan's hao-ai-lab#1389 ~90min cold source build. - Modal Secret integration for HF token (no token-on-the-wire). - --enable-compile, --dit-precision, --num-prompts flags for flexible legs (eager/compile, bf16/fp16/fp32, short diagnostic runs vs full 5-prompt validation). Reproducible by reviewers — see PR body for exact commands.
4 tasks
Mister-Raggs
added a commit
to Mister-Raggs/FastVideo
that referenced
this pull request
May 31, 2026
Single-container two-pass A/B harness modeled on Kuan's nccl_stream_ab.py pattern (used to validate hao-ai-lab#1395). Runs the same seed-pinned prompts twice through the same VideoGenerator (once sequential, once batched), computes pairwise SSIM via pytorch_msssim (matches fastvideo/tests/utils.py:compute_video_ssim_torchvision), prints a per-prompt + total wall/SSIM table. Architecture: the harness ferries pass config + records over JSON files and subprocesses the inner script via /opt/venv/bin/python. Modal's main function process runs in its own add_python="3.12" layer where FastVideo isn't importable; the venv subprocess works around that. Same reason ssim_test.py runs pytest as a subprocess. Features: - Mode tag in output dirs ({eager,compile}[-{dit_precision}]) so multiple legs don't clobber each other in the hf-model-weights Volume. - Recovery function reads existing mp4s by mtime and recomputes SSIM, for post-run gap-fill if the in-run SSIM step missed. - Prebuilt FA3 wheel install on Hopper (autodetected from gpu kwarg) via mjun0812/flash-attention-prebuild-wheels v0.9.4 — the same release that supplies the FA2 wheel already baked into the fastvideo-dev image. ~30s install vs Kuan's hao-ai-lab#1389 ~90min cold source build. - Modal Secret integration for HF token (no token-on-the-wire). - --enable-compile, --dit-precision, --num-prompts flags for flexible legs (eager/compile, bf16/fp16/fp32, short diagnostic runs vs full 5-prompt validation). Reproducible by reviewers — see PR body for exact commands.
Mister-Raggs
added a commit
to Mister-Raggs/FastVideo
that referenced
this pull request
May 31, 2026
Single-container two-pass A/B harness modeled on Kuan's nccl_stream_ab.py pattern (used to validate hao-ai-lab#1395). Runs the same seed-pinned prompts twice through the same VideoGenerator (once sequential, once batched), computes pairwise SSIM via pytorch_msssim (matches fastvideo/tests/utils.py:compute_video_ssim_torchvision), prints a per-prompt + total wall/SSIM table. Architecture: the harness ferries pass config + records over JSON files and subprocesses the inner script via /opt/venv/bin/python. Modal's main function process runs in its own add_python="3.12" layer where FastVideo isn't importable; the venv subprocess works around that. Same reason ssim_test.py runs pytest as a subprocess. Features: - Mode tag in output dirs ({eager,compile}[-{dit_precision}]) so multiple legs don't clobber each other in the hf-model-weights Volume. - Recovery function reads existing mp4s by mtime and recomputes SSIM, for post-run gap-fill if the in-run SSIM step missed. - Prebuilt FA3 wheel install on Hopper (autodetected from gpu kwarg) via mjun0812/flash-attention-prebuild-wheels v0.9.4 — the same release that supplies the FA2 wheel already baked into the fastvideo-dev image. ~30s install vs Kuan's hao-ai-lab#1389 ~90min cold source build. - Modal Secret integration for HF token (no token-on-the-wire). - --enable-compile, --dit-precision, --num-prompts flags for flexible legs (eager/compile, bf16/fp16/fp32, short diagnostic runs vs full 5-prompt validation). Reproducible by reviewers — see PR body for exact commands.
Mister-Raggs
added a commit
to Mister-Raggs/FastVideo
that referenced
this pull request
May 31, 2026
Single-container two-pass A/B harness modeled on Kuan's nccl_stream_ab.py pattern (used to validate hao-ai-lab#1395). Runs the same seed-pinned prompts twice through the same VideoGenerator (once sequential, once batched), computes pairwise SSIM via pytorch_msssim (matches fastvideo/tests/utils.py:compute_video_ssim_torchvision), prints a per-prompt + total wall/SSIM table. Architecture: the harness ferries pass config + records over JSON files and subprocesses the inner script via /opt/venv/bin/python. Modal's main function process runs in its own add_python="3.12" layer where FastVideo isn't importable; the venv subprocess works around that. Same reason ssim_test.py runs pytest as a subprocess. Features: - Mode tag in output dirs ({eager,compile}[-{dit_precision}]) so multiple legs don't clobber each other in the hf-model-weights Volume. - Recovery function reads existing mp4s by mtime and recomputes SSIM, for post-run gap-fill if the in-run SSIM step missed. - Prebuilt FA3 wheel install on Hopper (autodetected from gpu kwarg) via mjun0812/flash-attention-prebuild-wheels v0.9.4 — the same release that supplies the FA2 wheel already baked into the fastvideo-dev image. ~30s install vs Kuan's hao-ai-lab#1389 ~90min cold source build. - Modal Secret integration for HF token (no token-on-the-wire). - --enable-compile, --dit-precision, --num-prompts flags for flexible legs (eager/compile, bf16/fp16/fp32, short diagnostic runs vs full 5-prompt validation). Reproducible by reviewers — see PR body for exact commands.
Mister-Raggs
added a commit
to Mister-Raggs/FastVideo
that referenced
this pull request
Jun 1, 2026
Single-container two-pass A/B harness modeled on Kuan's nccl_stream_ab.py pattern (used to validate hao-ai-lab#1395). Runs the same seed-pinned prompts twice through the same VideoGenerator (once sequential, once batched), computes pairwise SSIM via pytorch_msssim (matches fastvideo/tests/utils.py:compute_video_ssim_torchvision), prints a per-prompt + total wall/SSIM table. Architecture: the harness ferries pass config + records over JSON files and subprocesses the inner script via /opt/venv/bin/python. Modal's main function process runs in its own add_python="3.12" layer where FastVideo isn't importable; the venv subprocess works around that. Same reason ssim_test.py runs pytest as a subprocess. Features: - Mode-tagged output dirs ({eager,compile}) so multiple legs don't clobber each other in the hf-model-weights Volume. - Recovery function reads existing mp4s by mtime and recomputes SSIM for post-run gap-fill if the in-run SSIM step missed. - Prebuilt FA3 wheel install on Hopper (autodetected from gpu kwarg) via mjun0812/flash-attention-prebuild-wheels v0.9.4 — the same release that supplies the FA2 wheel already baked into the fastvideo-dev image. ~30s install vs Kuan's hao-ai-lab#1389 ~90min cold source build. - Modal Secret integration for HF token (no token-on-the-wire). - --enable-compile + --num-prompts flags for flexible legs. Reproducible by reviewers — see PR body for exact commands.
Mister-Raggs
added a commit
to Mister-Raggs/FastVideo
that referenced
this pull request
Jul 12, 2026
Single-container two-pass A/B harness modeled on Kuan's nccl_stream_ab.py pattern (used to validate hao-ai-lab#1395). Runs the same seed-pinned prompts twice through the same VideoGenerator (once sequential, once batched), computes pairwise SSIM via pytorch_msssim (matches fastvideo/tests/utils.py:compute_video_ssim_torchvision), prints a per-prompt + total wall/SSIM table. Architecture: the harness ferries pass config + records over JSON files and subprocesses the inner script via /opt/venv/bin/python. Modal's main function process runs in its own add_python="3.12" layer where FastVideo isn't importable; the venv subprocess works around that. Same reason ssim_test.py runs pytest as a subprocess. Features: - Mode-tagged output dirs ({eager,compile}) so multiple legs don't clobber each other in the hf-model-weights Volume. - Recovery function reads existing mp4s by mtime and recomputes SSIM for post-run gap-fill if the in-run SSIM step missed. - Prebuilt FA3 wheel install on Hopper (autodetected from gpu kwarg) via mjun0812/flash-attention-prebuild-wheels v0.9.4 — the same release that supplies the FA2 wheel already baked into the fastvideo-dev image. ~30s install vs Kuan's hao-ai-lab#1389 ~90min cold source build. - Modal Secret integration for HF token (no token-on-the-wire). - --enable-compile + --num-prompts flags for flexible legs. Reproducible by reviewers — see PR body for exact commands.
Mister-Raggs
added a commit
to Mister-Raggs/FastVideo
that referenced
this pull request
Jul 17, 2026
Single-container two-pass A/B harness modeled on Kuan's nccl_stream_ab.py pattern (used to validate hao-ai-lab#1395). Runs the same seed-pinned prompts twice through the same VideoGenerator (once sequential, once batched), computes pairwise SSIM via pytorch_msssim (matches fastvideo/tests/utils.py:compute_video_ssim_torchvision), prints a per-prompt + total wall/SSIM table. Architecture: the harness ferries pass config + records over JSON files and subprocesses the inner script via /opt/venv/bin/python. Modal's main function process runs in its own add_python="3.12" layer where FastVideo isn't importable; the venv subprocess works around that. Same reason ssim_test.py runs pytest as a subprocess. Features: - Mode-tagged output dirs ({eager,compile}) so multiple legs don't clobber each other in the hf-model-weights Volume. - Recovery function reads existing mp4s by mtime and recomputes SSIM for post-run gap-fill if the in-run SSIM step missed. - Prebuilt FA3 wheel install on Hopper (autodetected from gpu kwarg) via mjun0812/flash-attention-prebuild-wheels v0.9.4 — the same release that supplies the FA2 wheel already baked into the fastvideo-dev image. ~30s install vs Kuan's hao-ai-lab#1389 ~90min cold source build. - Modal Secret integration for HF token (no token-on-the-wire). - --enable-compile + --num-prompts flags for flexible legs. Reproducible by reviewers — see PR body for exact commands.
|
This pull request has been automatically marked as stale because it has not had any activity within 60 days. It will be automatically closed if no further activity occurs within 14 days. Leave a comment if you feel this pull request should remain open. Thank you! |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Routes the SP all-to-all NCCL collective onto a dedicated comm stream so it no longer serializes with compute on the default stream. Lays the infrastructure for future overlap-based work (e.g. Hybrid Ulysses-Ring rotation) without itself claiming an e2e perf delta on PCIe-bound topologies.
Scope (intentionally narrow)
Wraps only
AllToAll4D— the twodist.all_to_all_singlecalls inside_all_to_all_4D_forward. Other NCCL collectives (AllReduce,AllGather,Slice.backward) are left untouched. Per@rich7420's NCCL optimization status doc, SendRecv fromall_to_all_4Dis 97.7% of NCCL time on the 4× L40S Wan T2V SP=4 baseline; the others are <3% combined. Wrapping them would add noise without measurable benefit.Mechanism
A per-device stream cache on
DistributedAutograd(lazily created on first use, avoids touching CUDA before device selection; thread-safe via a lock). The twodist.all_to_all_singlecall sites inside the SP forward path are bracketed with the standard cross-stream discipline:Also extracts the forward shape-munging into a static
_all_to_all_4D_forwardhelper so the inference path can bypassautograd.Function.applydispatch overhead when grad is off (DeviceCommunicatorBase.all_to_all_4Dbranches ontorch.is_grad_enabled()).Two design choices worth flagging
No
record_stream. Modern PyTorch defaultsTORCH_NCCL_AVOID_RECORD_STREAMS=1(deprecation warning visible in worker logs); the NCCL backend tracks cross-stream tensor lifetime itself. Explicitrecord_stream()on collective inputs/outputs would force the deprecated slow path and inflate caching-allocator per-collective bookkeeping ~2×.wait_streamordering alone is sufficient. (Earlier internal design notes prescribedrecord_stream; that was wrong, corrected in this PR.)AllToAll4D-only scope (vs all 5 NCCL sites). Per the SendRecv-97.7%-of-NCCL number above. The
_functional_collectivesalternative path was already tried in @rich7420's exploration and rolled back (+~1% wall regression fromAsyncCollectiveTensorwrapper overhead).A/B validation — 4× L40S sp=4 PCIe (Wan T2V 1.3B)
Thanks to @rich7420 for the Modal A/B run. Harness:
fastvideo/tests/modal/nccl_stream_ab.py(single-container overlay-swap pattern, generic for validating single-file PRs). Baseline =bc5ea6bf^(commit-parent) so the delta isolates this one-file change. 5 prompts × Wan T2V 1.3B 720×1280 / 77 frames / 30 steps, seed-pinned, same Modal container (no cross-host variance).Reading: bit-exact output (NCCL is deterministic across streams for the same nranks / topology / algorithm). Wall delta inside run-to-run noise — exactly matches the PR thesis that a dedicated stream is necessary-but-not-sufficient for overlap on bandwidth-bound PCIe. No regression on quality or wall; lands cleanly as infrastructure.
Peak-memory was deliberately not re-measured: this patch touches only stream scheduling (no
record_stream, tensor lifetimes unchanged) — a null hypothesis there. Easy to add if a reviewer asks.(A/B reflects the original
bc5ea6bftip. The review-feedback fixes at29c752f5are scheduling/lifecycle only — per-device stream cache, CPU-path guard,ValueError, thread-safe init — and do not change collective semantics or output. Happy to re-run Kuan's harness against the new tip if maintainers want a refreshed table.)Review-feedback fixes (
29c752f5on top ofbc5ea6bf)Addresses the four open review comments on
base_device_communicator.py:_comm_streamwould crash any multi-GPU / single-process multi-device context. Replaced with_comm_streams: dict[torch.device, Stream], lazily populated underwith torch.cuda.device(device):.use_cuda = input_.is_cuda; stream wait/wrap only on CUDA, plaindist.all_to_all_singlefallback for CPU paths used by some unit tests.assert→ValueError(Copilot-medium). Asserts strip underpython -O;ValueErroris the correct runtime validation.threading.Lock+ double-checked locking around dict insertion.Correctness
The diff is stream-scheduling-only.
dist.all_to_all_single(output, input_, group=group)is called with identical arguments to before; only the stream context differs.comm.wait_stream(curr)/curr.wait_stream(comm)ensures cross-stream ordering. No tensor math changes; no NCCL arguments change. SSIM 1.000000 / LPIPS 0.000000 in the A/B above confirms bit-exactness.Files
fastvideo/distributed/device_communicators/base_device_communicator.py— per-device_comm_streamsdict + thread-safe_get_comm_stream(device)onDistributedAutograd;_all_to_all_4D_forwardstatic helper extracted;wait_streamdiscipline at the twodist.all_to_all_singlesites, gated oninput_.is_cudawith a CPU fallback;DeviceCommunicatorBase.all_to_all_4Dadds the inference fast-path.Credit
Patch authored by @rich7420 on his
inference-profilebranch (the corrections re:record_streamdeprecation + AllToAll4D-only scope are entirely his). This PR is the clean-extraction-+-PR-shepherding of his draft. A/B validation on Modal 4× L40S also by @rich7420.