Hacktoberfest 2026: le issue che i maintainer hanno segnato per ottobre, aperte e adatte ai principianti. Sfoglia le issue Hacktoberfest

[Performance] FusedAdam capturable kernel is ~1.9x slower than regular Adam on GB300

Aperta
  • Qualcosa non va?
#3,777 0 commenti 0 reazioni 1 assegnatario Vedi su GitHub

I maintainer di solito rispondono entro 2 giorni

@ksivaman ci sta già lavorando.

Dal 29/10/2026.

  • #3782 di @ksivaman — aperta

Valutazione

Questa issue non è ancora stata valutata.

Descrizione

Transformer Engine's FusedAdam optimizer exhibits approximately 1.9x higher kernel latency with capturable=True compared to capturable=False on GB300.

This was observed in Nemotron 3.5 30B training when enabling optimizer.optimizer_cuda_graph=true in Megatron. The optimizer kernel changes from AdamFunctor to AdamCapturableFunctor, increasing average kernel execution time from approximately 104 µs to 196 µs.

GPU: NVIDIA GB300
Container: nvidia+nemo+26.08.01
Training workload: Nemotron 3.5 30B (BF16, 2 nodes)

To reproduce the issue, run the standalone benchmark: python3 adam_bench.py --graph

"""
Benchmark Transformer Engine FusedAdam: regular vs capturable.
Compares optimizer kernel latency with capturable=False and
capturable=True using large FP32 tensors.

Usage:
    python3 adam_bench.py
    python3 adam_bench.py --n-tensors 10 --steps 5
    python3 adam_bench.py --graph
"""

import argparse
import statistics

import torch
from torch.profiler import profile, ProfilerActivity
from transformer_engine.pytorch.optimizers import FusedAdam


parser = argparse.ArgumentParser()
parser.add_argument("--n-tensors", type=int, default=201)
parser.add_argument("--numel", type=int, default=320 * 65536)
parser.add_argument("--steps", type=int, default=6)
parser.add_argument("--graph", action="store_true")
args = parser.parse_args()

print(f"GPU: {torch.cuda.get_device_name(0)}")
print(f"Tensors: {args.n_tensors}, Elements/tensor: {args.numel:,}")


def benchmark(capturable):
    # Create parameters and gradients
    params = [
        torch.randn(args.numel, device="cuda", requires_grad=True)
        for _ in range(args.n_tensors)
    ]
    for p in params:
        p.grad = torch.randn_like(p) * 1e-3

    optimizer = FusedAdam(
        params,
        lr=1.6e-3,
        betas=(0.9, 0.95),
        eps=1e-8,
        weight_decay=0.1,
        adam_w_mode=True,
        capturable=capturable,
        exp_avg_dtype=torch.float32,
        exp_avg_sq_dtype=torch.float32,
    )

    # Warmup and initialize optimizer states
    for _ in range(2):
        optimizer.step()
    torch.cuda.synchronize()

    # Measure individual Adam CUDA kernels
    with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
        for _ in range(args.steps):
            optimizer.step()
        torch.cuda.synchronize()

    kernel_times = [
        e.device_time
        for e in prof.events()
        if e.device_type == torch.autograd.DeviceType.CUDA
        and "multi_tensor_apply_kernel" in e.name
    ]

    # Measure total optimizer step time
    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)

    start.record()
    for _ in range(args.steps):
        optimizer.step()
    end.record()
    torch.cuda.synchronize()

    step_ms = start.elapsed_time(end) / args.steps
    kernel_us = statistics.mean(kernel_times)

    print(f"\nCapturable: {capturable}")
    print(f"  Kernels/step: {len(kernel_times) / args.steps:.0f}")
    print(f"  Mean kernel:  {kernel_us:.1f} us")
    print(f"  Step time:    {step_ms:.2f} ms")

    # Optional CUDA graph benchmark
    if capturable and args.graph:
        graph = torch.cuda.CUDAGraph()
        stream = torch.cuda.Stream()
        stream.wait_stream(torch.cuda.current_stream())

        with torch.cuda.stream(stream):
            optimizer.step()

        torch.cuda.current_stream().wait_stream(stream)

        with torch.cuda.graph(graph):
            optimizer.step()

        start.record()
        for _ in range(args.steps):
            graph.replay()
        end.record()
        torch.cuda.synchronize()

        print(f"  Graph replay: {start.elapsed_time(end) / args.steps:.2f} ms")

    del optimizer, params
    torch.cuda.empty_cache()

    return kernel_us, step_ms


regular_us, regular_ms = benchmark(capturable=False)
capturable_us, capturable_ms = benchmark(capturable=True)

print("\n=== Comparison ===")
print(f"Kernel slowdown: {capturable_us / regular_us:.2f}x")
print(f"Step slowdown:   {capturable_ms / regular_ms:.2f}x")
Lingua principale
Python
Stelle
3.6k
Fork
883
Merge medio
4g 17h
PR unite (30g)
70

Preparare l'ambiente

Come iniziare

  1. Leggi tutta la issue e poi la guida ai contributi del progetto.
  2. Commenta sulla issue per dire che te ne occupi tu — evita che due persone facciano lo stesso lavoro.
  3. Fai un fork del repository e lavora su un branch.
  4. Apri una pull request che faccia riferimento al numero della issue.

Altre issue di NVIDIA/TransformerEngine

Tutte le issue di NVIDIA/TransformerEngine

Issue simili

Altre issue su Python

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.