Hacktoberfest 2026：维护者为十月标记出来的 issue，仍然开放、适合新手。 浏览 Hacktoberfest issue

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

未关闭
  • 哪里不对？
#3,777 0 条评论 0 个 reaction 已指派 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 天 20 小时
30 天内合并 PR
71

环境准备

从这里开始

  1. 先读完整个 Issue，再读项目的贡献指南。
  2. 在 Issue 下留言说明你要接手 —— 这能避免两个人做同样的事。
  3. Fork 仓库，在一个分支上完成修改。
  4. 提交 Pull Request，并在描述里引用这个 Issue 编号。

NVIDIA/TransformerEngine 的其他 Issue

查看 NVIDIA/TransformerEngine 的全部 Issue

相似的 Issue

更多 Python Issue

把新 issue 发到你的邮箱

精选适合新手参与的 GitHub issue 摘要。