NVFP4 quantize on sm_120 reserves ~3.4 GiB of device memory via a TMA driver syscall

Offen
  • Stimmt etwas nicht?
#3,765 1 Kommentar 0 Reaktionen 0 zugewiesene Personen Auf GitHub ansehen

Maintainer antworten meist innerhalb von 2 Tagen

Dieses Issue hat noch niemand übernommen.

Bewertung

Schwierigkeit
4/5
Geschätzter Aufwand
3-5 Tage
Anfängerfreundlichkeit
55/100
Issue-Typ
Bug
Klarheit
Klar beschrieben
Aktivitätsstatus
Aktiv
Tech-Stack
python, pytorch
Bereich
machine-learning, performance

Rechercherichtung

Start in transformer_engine/common/util/ptx.cuh and compare the 2D TMA helpers, especially cp_async_bulk_tensor_2d_global_to_shared and the CTA variant. Check the callers’ cluster size and shared memory assumptions, then build and run the standalone tma_syscall_repro.cu or NVFP4 quantization on sm_120. Done means the change removes the extra device memory reservation while preserving quantized outputs and training behavior.

Vom Indexierungsmodell aus dem Issue-Text verfasst.

Beschreibung

Environment: Transformer Engine 2.20.2 (PyPI wheel transformer_engine_cu13), PyTorch 2.12.0+cu130, CUDA 13.2 toolkit,
driver 595.71.05, NVIDIA GeForce RTX 5090 (sm_120, 32 GB).

Symptom

The first NVFP4 quantization in a process permanently takes ~3.3 GiB of device memory outside the PyTorch caching
allocator (about 10% of a 32 GB card):

import torch
from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer
x = torch.randn(256, 1024, device="cuda", dtype=torch.bfloat16)
free0, total = torch.cuda.mem_get_info()
NVFP4Quantizer(rowwise=True, columnwise=False)(x); torch.cuda.synchronize()
free1, _ = torch.cuda.mem_get_info()
print((free0 - free1) / 2**30)   # ~3.3 GiB; cudaLimitStackSize goes from 1024 to 14608 bytes

FP8 current scaling and MXFP8 quantization do not show it; cuBLAS GEMMs (BF16/FP8/NVFP4) and the PyTorch SDPA
backends do not either. CUDA_MODULE_LOADING=EAGER does not change it (it is not module code).

Cause

  • The launch of quantize_transpose_nvfp4_tuned_1D_kernel raises cudaLimitStackSize from 1024 B to 14608 B, and the
    driver reserves that much local memory for every resident thread (170 SMs x 1536 threads), although the kernel
    itself reports 8 bytes of local memory.
  • The kernel's ELF lists an extern __cuda_syscall_cp_async_bulk_tensor_2d_tile_unicast: the 2D TMA load in
    ptx::cp_async_bulk_tensor_2d_global_to_shared (transformer_engine/common/util/ptx.cuh) uses a shared::cluster
    destination (cp.async.bulk.tensor.2d.shared::cluster.global.tile...). On sm_120 this compiles to
    UTMALDG plus CALL.ABS.NOINC into a driver syscall, and the syscall's stack requirement causes the reservation.
  • A standalone 20-line TMA kernel reproduces it (tma_syscall_repro.cu, below): with a shared::cluster destination
    the launch costs +3.41 GiB; with shared::cta the SASS is a bare UTMALDG, local memory is 0 and nothing is reserved.
    Same result for -arch=sm_120, sm_120a and sm_120f.

Proposed fix

TE already has cp_async_bulk_tensor_2d_global_to_shared_cta (used only by the FP8 block-scaling grouped kernels).
The callers of the 2D shared::cluster helper launch with cluster size 1 and pass CTA-local shared memory and mbarrier
addresses, so switching the helper (or the callers) to the shared::cta destination is semantically identical:

-      "cp.async.bulk.tensor.2d.shared::cluster.global.tile"
+      "cp.async.bulk.tensor.2d.shared::cta.global.tile"

Patch validated

Rebuilt libtransformer_engine.so from v2.20.2 with only this change (sm_120, CUDA 13.2), RTX 5090:

stock wheel patched
device memory outside the PyTorch allocator after NVFP4 quantize 3.92 GiB 0.61 GiB (CUDA context only)
quantized outputs (rowwise; rowwise+columnwise+RHT; stochastic rounding) bit-identical (data, scales, amax)
Qwen3-0.6B-shaped full training step, 4K tokens, NVFP4 recipe 155.4 ms 155.3 ms
same, 16K tokens, no activation recomputation OOM (3/3 runs) 964.3 ms

Impact

On 32 GB consumer Blackwell cards this is the difference between fitting and not fitting full-parameter training of
~1.7B models with fp32 master weights and AdamW (usable memory 27.4 GiB vs ~30.7 GiB).

Standalone reproducer (tma_syscall_repro.cu, no TE dependency)

// REPORT 26.2: on the RTX 5090 (sm_120, driver 595.71, CUDA 13.2) a 2D TMA load written with a shared::cluster destination
// compiles to CALL.ABS.NOINC (a driver syscall) next to UTMALDG; its first launch raises cudaLimitStackSize from 1 KiB to
// 14608 B and the driver reserves local memory for every resident thread of the GPU: +3.41 GiB, kept for the process lifetime.
// The shared::cta form compiles to a bare UTMALDG and reserves nothing. TE 2.20.2 uses the shared::cluster form in
// ptx::cp_async_bulk_tensor_2d_global_to_shared (common/util/ptx.cuh), which every NVFP4 quantize kernel calls.
//   nvcc -O2 -std=c++17 -arch=sm_120a tma_syscall_repro.cu -o r_cluster -lcuda && ./r_cluster
//   nvcc -O2 -std=c++17 -arch=sm_120a -DCTA tma_syscall_repro.cu -o r_cta -lcuda && ./r_cta
//   cuobjdump -sass r_cluster | grep -E "CALL|UTMALDG"
#ifdef CTA
#define TMA_DST "shared::cta"
#else
#define TMA_DST "shared::cluster"
#endif
#include <cuda.h>
#include <cuda_runtime.h>
#include <cudaTypedefs.h>
#include <cstdio>
#include <cstdint>
__global__ void k(const __grid_constant__ CUtensorMap map, float* out) {
  __shared__ alignas(128) float tile[32 * 32];
  __shared__ alignas(8) uint64_t bar;
  uint32_t sb = (uint32_t)__cvta_generic_to_shared(&bar), st = (uint32_t)__cvta_generic_to_shared(tile);
  if (threadIdx.x == 0) {
    asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" :: "r"(sb));
    asm volatile("fence.proxy.async.shared::cta;");
    asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" :: "r"(sb), "r"(32 * 32 * 4));
    asm volatile("cp.async.bulk.tensor.2d." TMA_DST ".global.tile.mbarrier::complete_tx::bytes [%0], [%1, {%2, %3}], [%4];"
                 :: "r"(st), "l"(&map), "r"(0), "r"(0), "r"(sb) : "memory");
    asm volatile("{ .reg .pred p; W: mbarrier.try_wait.parity.shared::cta.b64 p, [%0], 0; @!p bra W; }" :: "r"(sb));
  }
  __syncthreads();
  out[threadIdx.x] = tile[threadIdx.x];
}
int main() {
  float *src, *out; cudaMalloc(&src, 64 * 64 * 4); cudaMalloc(&out, 1024 * 4); cudaMemset(src, 0, 64 * 64 * 4);
  size_t lim0, lim1, f0, f1, t; cudaDeviceGetLimit(&lim0, cudaLimitStackSize); cudaMemGetInfo(&f0, &t);
  PFN_cuTensorMapEncodeTiled_v12000 enc; cudaDriverEntryPointQueryResult q;
  cudaGetDriverEntryPointByVersion("cuTensorMapEncodeTiled", (void**)&enc, 12000, cudaEnableDefault, &q);
  CUtensorMap m; cuuint64_t dims[2] = {64, 64}, strides[1] = {64 * 4}; cuuint32_t box[2] = {32, 32}, es[2] = {1, 1};
  CUresult r = enc(&m, CU_TENSOR_MAP_DATA_TYPE_FLOAT32, 2, src, dims, strides, box, es, CU_TENSOR_MAP_INTERLEAVE_NONE,
                   CU_TENSOR_MAP_SWIZZLE_NONE, CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE);
  cudaFuncAttributes a; cudaFuncGetAttributes(&a, k);
  k<<<1, 1024>>>(m, out); cudaError_t e = cudaDeviceSynchronize();
  cudaDeviceGetLimit(&lim1, cudaLimitStackSize); cudaMemGetInfo(&f1, &t);
  printf("%-15s encode %d launch %s  localSizeBytes %zu  stack limit %zu -> %zu  device memory used +%.2f GiB\n", TMA_DST, r, cudaGetErrorString(e),
         a.localSizeBytes, lim0, lim1, (double)(f0 - f1) / (1 << 30));
}

Output on the RTX 5090:

shared::cluster encode 0 launch no error  localSizeBytes 8  stack limit 1024 -> 14608  device memory used +3.41 GiB
shared::cta     encode 0 launch no error  localSizeBytes 0  stack limit 1024 -> 1024  device memory used +0.00 GiB
Vorherrschende Sprache
Python
Sterne
3.6k
Forks
883
Ø Merge
4 T. 17 Std.
Gemergte PRs (30 T.)
70

Entwicklungsumgebung

Erste Schritte

  1. Lesen Sie das ganze Issue und danach den Beitragsleitfaden des Projekts.
  2. Schreiben Sie ins Issue, dass Sie es übernehmen — das erspart doppelte Arbeit.
  3. Forken Sie das Repository und arbeiten Sie in einem Branch.
  4. Öffnen Sie einen Pull Request, der die Issue-Nummer nennt.

Mehr aus NVIDIA/TransformerEngine

Alle Issues in NVIDIA/TransformerEngine

Ähnliche Issues

Weitere Issues zu Python

Neue Issues direkt in Ihr Postfach

Eine kurze Übersicht über anfängerfreundliche GitHub-Issues.