Hacktoberfest 2026：メンテナが10月に向けて印を付けた、オープンで初心者向けの issue。 Hacktoberfest の issue を見る

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

オープン
  • 何か問題がありますか？
#3,777 コメント 0 件 リアクション 0 件 担当者 1 名 GitHub で見る

メンテナーはふだん 2 日以内に返信

@ksivaman がすでに取り組んでいます。

2026年10月29日 から。

  • #3782 @ksivaman による — オープン

評価

この issue はまだ評価されていません。

説明

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")
主要言語
Python
スター
3.6k
フォーク
883
平均マージ
4日 17時間
マージ済み PR（30日）
70

環境構築

はじめの一歩

  1. issue を最後まで読み、次にプロジェクトのコントリビューションガイドを読みます。
  2. 着手することを issue にコメントします — 二人が同じ作業をするのを防げます。
  3. リポジトリをフォークし、ブランチを切って変更します。
  4. issue 番号を参照したプルリクエストを送ります。

NVIDIA/TransformerEngine のほかの issue

NVIDIA/TransformerEngine の issue をすべて見る

似ている issue

Python の issue をもっと見る

新しい issue をメールで受け取る

初心者向けの GitHub issue を短くまとめたダイジェスト。