[BUG] multi_tensor_apply: int32 overflow in TensorListMetadata::sizes causes illegal memory access for tensors with numel > INT_MAX

[yezhengmao1]

multi_tensor_apply silently truncates per-tensor sizes from 64-bit to 32-bit, causing illegal memory access when any input tensor has numel() > INT_MAX (2,147,483,647).

In transformer_engine/common/multi_tensor/multi_tensor_apply.cuh, TensorListMetadataBase::sizes is declared as int sizes[...] (int32), but it is populated from Tensor::numel() (size_t / int64):

  // multi_tensor_apply.cuh:24
  int sizes[depth_to_max_tensors[n - 1]];
  ...
  // multi_tensor_apply.cuh:68
  tl.sizes[loc_tensor_info] = tensor_lists[0][t]->numel();   // int64 -> int32, silent truncation

For a tensor with numel = 2,476,250,368 (e.g. an embedding of shape [19345706, 128]), this field becomes 2476250368 - 2^32 = -1,818,716,928. The resulting negative / bogus size is then consumed by downstream kernels (e.g. multi_tensor_l2norm_kernel) which compute element offsets from it, producing out-of-bounds global-memory accesses and the following error at the next CUDA sync:

RuntimeError: .../multi_tensor_apply.cuh:92 in function multi_tensor_apply:
CUDA Error: an illegal memory access was encountered

This is hit by any real-world use that feeds a tensor with numel > 2^31 to TE's multi_tensor utilities. In particular, megatron.training.utils.calc_params_l2_norm →
multi_tensor_applier(multi_tensor_l2norm, ...) crashes for any model containing a single parameter with >2.14B elements (common for large-vocab embeddings, tied output layers, over-encoding tables, etc.).

[beccohov]

Hi @ksivaman @timmoon10 — I'd like to take this

[CodersAcademy006]

Confirming the root cause and adding one detail that matters for the fix: there are actually two truncation layers, so widening only the struct field is not sufficient.

1. The struct field (root cause, as reported)

transformer_engine/common/multi_tensor/multi_tensor_apply.cuh:24

int sizes[depth_to_max_tensors[n - 1]];
...
// line 68
tl.sizes[loc_tensor_info] = tensor_lists[0][t]->numel();  // int64 -> int32 truncation

2. The consumers re-truncate

Every kernel reads the size back into an int local before using it, so even after widening the struct, these would truncate again:

• l2norm.cu:126, l2norm.cu:196, l2norm.cu:266
• scale.cu:49
• compute_scale.cu:37, compute_scale.cu:64
• sgd.cu:48
• adam.cu:405, adam.cu:499

int n = tl.sizes[tensor_loc];   // truncates the full numel before n -= chunk_idx * chunk_size

Note that adam.cu already uses a templated index_t for three of its kernels (adam.cu:74, :189, :312), so the large-index intent exists there, but the int struct field undermines it. A consistent fix would make all of these 64-bit.

Suggested fix

• multi_tensor_apply.cuh: int sizes[...] -> int64_t sizes[...].
• Change the int n = tl.sizes[tensor_loc] locals listed above to int64_t (or index_t where the kernel is already templated), so the per-tensor size survives until n -= chunk_idx * chunk_size. The per-chunk remainder after that subtraction is bounded by chunk_size, so the inner loop indices do not need to change as long as the element base offset uses the existing chunk_idx * chunk_size (already 64-bit via the int64_t chunk_size).

Kernel argument size budget

The struct is passed by value as a kernel argument, and line 17 flags a possible sub-4KB limit on some cards, so it is worth checking. Widening sizes to int64_t adds at most depth_to_max_tensors[0] * 4 = 440 bytes (the n = 1 case). That brings the largest non-FP8 metadata struct to roughly 3.3KB, still under 4KB. The FP8 specialization already exceeds 4KB today because of fp8_meta_addresses, so it already relies on the larger parameter space on supported GPUs, and this change does not introduce a new platform constraint there.

@beccohov since you mentioned taking this, hopefully the consumer-side list and the budget check above save you some time. Happy to help review or split the work if useful.