[PyTorch] High CPU overhead on Grace systems

Nobody has claimed this yet.

Hello everyone,

we have been investigating low compute throughput for LLM trainings on GH200 and GB200 systems in comparison to H100 DGX (Intel Xeon Platinum 8480C). This already affects 1 GPU trainings and we tracked this down to CPU overhead of Transformer Engine (TE) modules on Grace systems.

To showcase this, I appended a small benchmarking script, that evaluates GPU kernel runtime and CPU overhead of the Linear module from PyTorch and Transformer Engine on GB200, GH200 and H100 DGX.
Our benchmarking does multiple runs and we evaluate the average time per Linear forward pass both on CPU and GPU. In the following run we look at BF16 precision with a hidden size of 2048 and 8192 tokens, these values are typical for smaller LLMs:

	backend	CPU runtime [µs]	GPU runtime [µs]	estimated throughput [TFLOP/s]
GB200	torch	37.08	41.79	1644.22
GB200	TE	99.66	99.73	689.08
GH200	torch	33.17	91.64	749.86
GH200	TE	90.53	92.14	745.80
H100 DGX	torch	13.48	100.06	686.79
H100 DGX	TE	49.19	101.88	674.51

So what we see is that:

The torch Linear has significantly lower CPU overhead in comparison to TE. This is probably expected as TE offers a lot more features.
On Grace systems, the CPU overhead is almost a factor of two higher in comparison to H100 DGX

This two factors together can severely limit the throughput as execution becomes CPU bound, see for example the GB200 numbers with TE. Although the GPU has 2.5x the FLOP/s, the scenario does not make use of it. Even the GH200 run with TE is very close to be severly CPU bound, especially if you consider that actual training code adds additional CPU instructions.

Generally, increasing the number of tokens minimizes this issue as it increases GPU runtime but CPU overhead is independent of the problem size. But for many scenarios other constraints like memory, global batch size and pipelining considerations limit the number of tokens per forward/backward pass and as such are infeasible. Putting this all together the CPU overhead can severly limit the compute throughput during training.

We already did profiling with nsys and cProfile but were not able to identify any "easy wins". It would be great if you could try to reproduce this and give a first assesment. Hope that you have a few ideas how to improve the situation 😄

Please let me know if you have further questions or need more details.

Steps to reproduce

Benchmark script

from typing import List
import torch
import time
import argparse

import transformer_engine.pytorch as te


def speedometer(
    module: torch.nn.Module,
    args: List[torch.Tensor],
    timing_iters: int = 500,
    warmup_iters: int = 50,
    num_rounds: int = 5,
) -> float:
    """Measure average run time for a PyTorch module"""
    for _ in range(warmup_iters):
        module(*args)

    gpu_times = []
    cpu_times = []
    for round_idx in range(num_rounds):
        start = torch.cuda.Event(enable_timing=True)
        end = torch.cuda.Event(enable_timing=True)
        torch.cuda.synchronize()
        start.record()
        cpu_start = time.time()
        for _ in range(timing_iters):
            module(*args)
        cpu_end = time.time()
        end.record()
        torch.cuda.synchronize()
        gpu_elapsed = start.elapsed_time(end)
        cpu_elapsed = (cpu_end - cpu_start) * 1000
        gpu_times.append(gpu_elapsed)
        cpu_times.append(cpu_elapsed)
        print(
            f"Round {round_idx+1}/{num_rounds}: GPU {gpu_elapsed/timing_iters*1000:.2f} µs, CPU {cpu_elapsed/timing_iters*1000:.2f} µs"
        )
    print(f"Average GPU time over {num_rounds} rounds: {sum(gpu_times)/(num_rounds*timing_iters)*1000:.2f} µs")
    print(f"Average CPU time over {num_rounds} rounds: {sum(cpu_times)/(num_rounds*timing_iters)*1000:.2f} µs")

    return sum(gpu_times) / num_rounds


def main():
    parser = argparse.ArgumentParser(description="Benchmark torch.nn.Linear performance and CPU overhead.")
    parser.add_argument("--hidden_size", type=int, default=2048, help="Hidden size")
    parser.add_argument("--seq_length", type=int, default=8192, help="Sequence length")
    parser.add_argument("--warmup", type=int, default=500, help="Number of warmup iterations")
    parser.add_argument("--timing_iters", type=int, default=500, help="Number of timing iterations per round")
    parser.add_argument("--num_rounds", type=int, default=3, help="Number of timing rounds")
    parser.add_argument(
        "--backend", type=str, choices=["torch", "te"], default="te", help="Linear backend: torch or te"
    )
    args = parser.parse_args()

    x = torch.randn((args.seq_length, args.hidden_size), dtype=torch.bfloat16, device="cuda", requires_grad=True)
    if args.backend == "torch":
        model = torch.nn.Linear(args.hidden_size, args.hidden_size, bias=False).to(torch.bfloat16).cuda()
    else:
        model = te.Linear(args.hidden_size, args.hidden_size, bias=False, device="cuda").to(torch.bfloat16)
    avg_gpu_time_per_round = speedometer(
        model, [x], timing_iters=args.timing_iters, warmup_iters=args.warmup, num_rounds=args.num_rounds
    )

    total_ops = 2 * args.hidden_size * args.hidden_size * args.seq_length * args.timing_iters

    tflops = total_ops / avg_gpu_time_per_round / 1e9
    print(f"Estimated TFLOP/s: {tflops:.2f}")


if __name__ == "__main__":
    main()

The numbers above can be reproduced by executing the given benchmarking script with:

python benchmark_linear_cpu_overhead.py --backend torch
python benchmark_linear_cpu_overhead.py --backend te

Environment overview

Docker container: nvcr.io/nvidia/pytorch:25.06-py3
GB200: NVIDIA GB200 Superpod with NVL72
GH200: NVIDIA GH200 480GB
H100 DGX: NVIDIA H100 80GB HBM3

Start with the benchmark script included in the issue and run the two provided commands on the listed Grace and H100 environments to confirm the CPU-overhead difference. Compare the PyTorch and Transformer Engine Linear paths with profiling, then validate any identified improvement against the benchmark results; the issue does not define a specific target beyond understanding and reducing the overhead.

Written by the indexing model from the issue text.

Description

Contributor guide

First steps

Research direction

Assessment