Hacktoberfest 2026: những issue maintainer đã đánh dấu cho tháng Mười, đang mở và phù hợp người mới. Xem issue Hacktoberfest

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

Đang mở
  • Có gì không ổn?
#3,777 0 bình luận 0 reaction 1 người được giao Xem trên GitHub

Maintainer thường phản hồi trong vòng 2 ngày

@ksivaman đang làm issue này rồi.

Từ ngày 29/10/2026.

  • #3782 của @ksivaman — đang mở

Đánh giá

Issue này chưa được đánh giá.

Mô tả

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")
Ngôn ngữ chính
Python
Star
3.6k
Fork
883
Merge trung bình
4 ngày 20 giờ
Pull request đã merge (30 ngày)
71

Chuẩn bị môi trường

Bắt đầu từ đâu

  1. Đọc hết issue, rồi đọc hướng dẫn đóng góp của dự án.
  2. Bình luận trên issue rằng bạn sẽ nhận — tránh hai người làm cùng một việc.
  3. Fork repository và làm thay đổi trên một nhánh.
  4. Mở pull request có tham chiếu số hiệu của issue.

Issue khác của NVIDIA/TransformerEngine

Tất cả issue của NVIDIA/TransformerEngine

Issue tương tự

Thêm issue về Python

Nhận issue mới trong hộp thư của bạn

Bản tóm tắt ngắn những issue GitHub phù hợp với người mới.