[None][perf] enable NCCL user-buffer registration for VisualGen collectives #15013

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed

pkisfaludi-nv wants to merge 10 commits into NVIDIA:main from pkisfaludi-nv:feat/nccl-buffer-reg

docs/source/models/visual-generation.md

-Original file line number
+Diff line change
@@ Expand Up @@
     [^1]: FLUX models use embedded guidance and do not have a separate negative prompt path, so CFG parallelism is not applicable.
-    [^2]: Qwen-Image ships a native BF16 implementation with per-module numerical parity vs `diffusers.QwenImagePipeline` (cosine >= 0.999 on the full 20B transformer) and `trtllm-serve` / `/v1/images/generations` support. FP8 blockwise and NVFP4 use VisualGen dynamic quantization from BF16 checkpoints; no pre-quantized checkpoint is required.
+    [^2]: Qwen-Image ships a native BF16 implementation with per-module numerical parity vs `diffusers.QwenImagePipeline` (cosine >= 0.999 on the full 20B transformer) and `trtllm-serve` / `/v1/images/generations` support. FP8 blockwise and NVFP4 use VisualGen dynamic quantization from BF16 checkpoints; no pre-quantized checkpoint is required.
     ## Quick Start
@@ Expand Down Expand Up @@
         - **Attention2D** (`attn2d_size: [N, M]`): Shards the sequence axis across an `N × M` device mesh (CP degree = `N · M`; total SP degree = `N · M · ulysses_size`).
         - **Ring Attention** (`ring_size: N`): Shards the sequence axis across a 1D ring of `N` ranks, streaming K/V blocks (CP degree = `N`; total SP degree = `N · ulysses_size`; mutually exclusive with Attention2D).
     - **Tensor Parallelism** (`tp_size: N`): Splits attention heads and transformer MLPs across GPUs for faster compute and reduced memory usage.
+    ### NCCL User-Buffer Registration
+    Setting `nccl_buffer_reg: true` in `parallel_config` enables NCCL zero-copy collectives for VisualGen's inter-GPU communication (Ulysses all-to-all, ring all-gather, VAE all-reduce). Two mechanisms activate automatically:
+    | Mechanism | When active | NCCL requirement |
+    |---|---|---|
+    | **CUMEM mode** (`NCCL_CUMEM_ENABLE=1`) | Always when `nccl_buffer_reg: true` | ≥ 2.21 |
+    | **Explicit `ncclCommRegister`** | When the raw `ncclComm_t` is accessible from the PyTorch process group | ≥ 2.19 |
+    **CUMEM mode** instructs NCCL to allocate its internal scratch buffers via `cuMemCreate` (VMM). On NVLink-connected GPUs (e.g. H100, B200 NVL), this enables NCCL to fuse collectives with NVLS (NVLink SHARP), reducing SM usage and latency. On IB-connected multi-node setups, set `NCCL_PXN_DISABLE=1` alongside `NCCL_CUMEM_ENABLE=1` — PXN is incompatible with buffer registration.
+    **Effect on VisualGen parallelism modes:**
+    - **Ulysses** (all-to-all around self-attention): most impactful — the all-to-all is on the critical path of every denoising step.
+    - **Ring Attention** (K/V streaming all-gather): moderate — one all-gather per layer per step.
+    - **CFG parallel** (all-gather of noise predictions): minor — one collective per step.
+    **Example config** (`examples/visual_gen/serve/configs/flux1_nccl_ub.yml`):
+    ```yaml
+    attention_config:
+      backend: VANILLA
+    parallel_config:
+      ulysses_size: 2          # Ulysses across 2 GPUs
+      nccl_buffer_reg: true    # enables NCCL_CUMEM_ENABLE + ncclCommRegister
+    ```
+    **Benchmarking** — collect baseline vs. UB-registration numbers with:
+    ```bash
+    # Baseline (no UB registration)
+    torchrun --nproc-per-node=2 examples/visual_gen/bench_nccl_ub.py --tier micro
+    # With NCCL_CUMEM_ENABLE=1
+    torchrun --nproc-per-node=2 examples/visual_gen/bench_nccl_ub.py --tier micro --nccl-cumem
+    # Full pipeline (requires model weights)
+    torchrun --nproc-per-node=2 examples/visual_gen/bench_nccl_ub.py \
+        --tier pipeline \
+        --flux-path /path/to/FLUX.1-dev \
+        --wan-path  /path/to/Wan2.1-T2V-1.3B-Diffusers
+    ```
+    Results are saved to `bench_nccl_ub_results.json`.
     ## Developer Guide
     ### Architecture Overview
@@ Expand Down @@

examples/visual_gen/bench_cold_ring.py

-Original file line number
+Diff line change
@@ -0,0 +1,79 @@
+    #!/usr/bin/env python3
+    """Cold-start ring-4 test — each run is a fresh process with no prior ring-4 calls."""
+    import os, statistics, time, json, argparse
+    def main():
+        ap = argparse.ArgumentParser()
+        ap.add_argument("--nccl-cumem", action="store_true")
+        ap.add_argument("--warmup", type=int, default=0)
+        ap.add_argument("--iters", type=int, default=7)
+        ap.add_argument("--tag", default="")
+        a = ap.parse_args()
+        if a.nccl_cumem:
+            os.environ["NCCL_CUMEM_ENABLE"] = "1"
+        else:
+            os.environ["NCCL_CUMEM_ENABLE"] = "0"
+        os.environ.setdefault("TRTLLM_DISABLE_COSMOS3_GUARDRAILS", "1")
+        from tensorrt_llm import VisualGen, VisualGenArgs, VisualGenParams
+        PROMPTS = [
+            "A majestic lion resting on a sun-drenched savanna",
+            "Futuristic city skyline reflected in a calm river",
+            "A waterfall cascading into a turquoise pool",
+            "Northern lights above a snow-covered forest",
+            "Ancient temple ruins illuminated by morning mist",
+            "Desert dunes at sunset with long shadows",
+            "Underwater coral reef teeming with tropical fish",
+        ]
+        tag = a.tag or ("cumem" if a.nccl_cumem else "base")
+        print(f"\n{'='*60}")
+        print(f"COLD RING-4  NCCL_CUMEM_ENABLE={os.environ['NCCL_CUMEM_ENABLE']}  tag={tag}")
+        print(f"warmup={a.warmup}  iters={a.iters}")
+        print(f"{'='*60}\n")
+        args = VisualGenArgs(
+            attention_config={"backend": "CUTEDSL"},
+            parallel_config={"ring_size": 4},
+        )
+        print("Loading model...")
+        vg = VisualGen("/workspace/models/Cosmos3-Nano", args=args)
+        gen_kwargs = dict(height=720, width=1280, num_inference_steps=20,
+                          guidance_scale=6.0, seed=42, num_frames=57)
+        for i in range(a.warmup):
+            print(f"  warmup {i+1}/{a.warmup}...")
+            vg.generate(inputs=PROMPTS[i % len(PROMPTS)], params=VisualGenParams(**gen_kwargs))
+            print(f"  warmup {i+1} done")
+        times = []
+        for i in range(a.iters):
+            t0 = time.perf_counter()
+            vg.generate(inputs=PROMPTS[i % len(PROMPTS)], params=VisualGenParams(**gen_kwargs))
+            elapsed = time.perf_counter() - t0
+            times.append(elapsed)
+            print(f"  iter {i+1}/{a.iters}: {elapsed:.2f}s")
+        vg.shutdown()
+        result = {
+            "tag": tag,
+            "nccl_cumem": os.environ["NCCL_CUMEM_ENABLE"],
+            "warmup": a.warmup,
+            "iters": a.iters,
+            "all_s": [round(t, 3) for t in times],
+            "min_s": round(min(times), 3),
+            "median_s": round(statistics.median(times), 3),
+            "mean_s": round(statistics.mean(times), 3),
+        }
+        print(f"\nResult: {result}")
+        out = f"/workspace/results/cold_ring4_{tag}.json"
+        os.makedirs("/workspace/results", exist_ok=True)
+        with open(out, "w") as f:
+            json.dump(result, f, indent=2)
+        print(f"Saved → {out}")
+    if __name__ == "__main__":
+        main()

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[None][perf] enable NCCL user-buffer registration for VisualGen collectives #15013

Uh oh!

Diff view

Diff view

There are no files selected for viewing

Uh oh!