Open
Description
🐛 Describe the bug
Problem
When using torch.matmul with inputs with more than three dimensions, torch can unnecessarily physically broadcast the inputs to unmanageable sizes.
In the following MWE, torch tries to allocate 40GB memory. I believe this is because it tries to broadcast B to shape (4096, 256, 100, 100).
import torch
A = torch.rand(4096, 256, 1, 100)
B = torch.rand(256, 100, 100)
A.matmul(B) # Tries to allocate 40 GB memoryNot sure if this is a bug or a feature request.
Relevant other issues / discussions
- matmul uses too much memory in some batched cases #18862
- https://discuss.pytorch.org/t/unexpected-huge-memory-cost-of-matmul/41642
Potential Solution
A potential solution using torch.bmm does not have this problem. However, it requires some complicated jiggling of dimensions because torch.bmm only accepts 3D tensors. Specifically, the batch dimensions are either "matched dimensions" (where both input and other have the same size, greater than 1) or "broadcast dimensions" (dimensions where one of input and other has size 1 ). We can group:
- The matched dimensions of
inputandotherinto the first dimension - The broadcast dimensions of
inputwith the first matrix dimension - The broadcast dimensions of
otherwith the second matrix dimension
The result is a batched matrix multiplication between tensors of shape(m, b1*p, q)and(m, q, b2*r)wheremis the product of sizes of matched dimensions,b1is the product of sizes of broadcast dimensions forinputandb2is the product of sizes of broadcast dimensions forother.
def matmul(input, other):
"""
Performs a generalized batched matmul between input (..., p, q) and other (..., q, r),
respecting matched and broadcasted batch dimensions.
"""
# Check shapes are broadcastable - this raises if they are not
torch.broadcast_shapes(input.shape[:-2], other.shape[:-2])
if len(input.shape) < 2 or len(other.shape) < 2:
raise ValueError(
f"input and other must both have at least two dimensions. "
f"Got {input.shape=}, {other.shape=}"
)
if input.shape[-1] != other.shape[-2]:
raise ValueError(
f"Incompatible matrix dimensions. {input.shape[-1]=} != {other.shape[-2]}"
)
# Pad the shorter shape with 1s on the left for alignment
if len(input.shape) < len(other.shape):
k = len(other.shape) - len(input.shape)
input = input.view((1,) * k + input.shape)
elif len(other.shape) < len(input.shape):
k = len(input.shape) - len(other.shape)
other = other.view((1,) * k + other.shape)
input_batch, p, q = input.shape[:-2], input.shape[-2], input.shape[-1]
other_batch, r = other.shape[:-2], other.shape[-1]
# Split batch dimensions into those which need to be matched between inputs and
# those which will be broadcasted. The broadcasted ones will be later moved into
# the matrix dimensions.
matched_dims = []
broadcast_dims = []
for i, (a, b) in enumerate(zip(input_batch, other_batch)):
if a == b and a != 1:
matched_dims.append(i)
elif a == 1 or b == 1:
broadcast_dims.append(i)
else:
# Shouldn't get hit since we've already checked broadcastability
raise RuntimeError(
f"Incompatible broadcast dimensions {i=}, {a=}, {b=}"
)
matched_shape = torch.tensor([input_batch[i] for i in matched_dims])
input_broadcast_shape = torch.tensor([input_batch[i] for i in broadcast_dims])
other_broadcast_shape = torch.tensor([other_batch[i] for i in broadcast_dims])
# Permute dimensions and reshape to get 3D tensors compatible with bbm
# For input we combine broadcast dimensions with the first matrix dimension (of size
# p) and for other we combine broadcast dimensions with the second matrix dimension
# (of size r)
m = torch.prod(matched_shape)
b1 = torch.prod(input_broadcast_shape)
b2 = torch.prod(other_broadcast_shape)
input_perm = input.permute(*matched_dims, *broadcast_dims, -2, -1)
input_reshape = input_perm.reshape(m, b1 * p, q)
other_perm = other.permute(*matched_dims, -2, -1, *broadcast_dims)
other_reshape = other_perm.reshape(m, q, r * b2)
matrix_products = torch.bmm(input_reshape, other_reshape) # shape: (m, b1 * p, r * b2)
# Reshape and permute back to original dimensions
# To compute the inverse permutation, first choose whether the broadcast dimension
# came from input or other (at least one will be 1)
use_input = input_broadcast_shape != 1
input_effective_br_shape = input_broadcast_shape[use_input]
other_effective_br_shape = other_broadcast_shape[~use_input]
input_br_dims = [ind for i, ind in enumerate(broadcast_dims) if use_input[i]]
other_br_dims = [ind for i, ind in enumerate(broadcast_dims) if ~use_input[i]]
matrix_products = matrix_products.view(
*matched_shape, *input_effective_br_shape, p, r, *other_effective_br_shape
)
# Compute effective_inds, the permutation we want to inverse
ndim = len(matrix_products.shape)
effective_inds = matched_dims + input_br_dims + [ndim - 2, ndim - 1] + other_br_dims
inverse_perm = [
i for _, i in sorted((orig_ind, i) for i, orig_ind in enumerate(effective_inds))
]
matrix_products_perm = matrix_products.permute(*inverse_perm)
return matrix_products_permAnd a unit test on a problem of small dimensions to check we get the same as torch.matmul.
def test_matmul():
mat1 = torch.randn(5, 1, 5, 2, 3)
mat2 = torch.randn(4, 5, 3, 9)
out = matmul(mat1, mat2)
out_torch = torch.matmul(mat1, mat2)
torch.testing.assert_close(out, out_torch)Versions
PyTorch version: 2.6.0+cpu
Is debug build: False
CUDA used to build PyTorch: Could not collect
ROCM used to build PyTorch: N/A
OS: Ubuntu 20.04.6 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.2) 9.4.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.31
Python version: 3.13.2 | packaged by conda-forge | (main, Feb 17 2025, 14:10:22) [GCC 13.3.0] (64-bit runtime)
Python platform: Linux-5.15.0-139-generic-x86_64-with-glibc2.31
Is CUDA available: False
CUDA runtime version: 10.1.243
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: Could not collect
Nvidia driver version: Could not collect
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True
CPU:
Architecture: x86_64
CPU op-mode(s): 32-bit, 64-bit
Byte Order: Little Endian
Address sizes: 39 bits physical, 48 bits virtual
CPU(s): 12
On-line CPU(s) list: 0-11
Thread(s) per core: 2
Core(s) per socket: 6
Socket(s): 1
NUMA node(s): 1
Vendor ID: GenuineIntel
CPU family: 6
Model: 165
Model name: Intel(R) Core(TM) i7-10850H CPU @ 2.70GHz
Stepping: 2
CPU MHz: 2700.000
CPU max MHz: 5100.0000
CPU min MHz: 800.0000
BogoMIPS: 5399.81
Virtualisation: VT-x
L1d cache: 192 KiB
L1i cache: 192 KiB
L2 cache: 1.5 MiB
L3 cache: 12 MiB
NUMA node0 CPU(s): 0-11
Flags: fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush dts acpi mmx fxsr sse sse2 ss ht tm pbe syscall nx pdpe1gb rdtscp lm constant_tsc art arch_perfmon pebs bts rep_good nopl xtopology nonstop_tsc cpuid aperfmperf pni pclmulqdq dtes64 monitor ds_cpl vmx smx est tm2 ssse3 sdbg fma cx16 xtpr pdcm pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand lahf_lm abm 3dnowprefetch cpuid_fault epb invpcid_single ssbd ibrs ibpb stibp ibrs_enhanced tpr_shadow vnmi flexpriority ept vpid ept_ad fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid mpx rdseed adx smap clflushopt intel_pt xsaveopt xsavec xgetbv1 xsaves dtherm ida arat pln pts hwp hwp_notify hwp_act_window hwp_epp pku ospke md_clear flush_l1d arch_capabilities
Versions of relevant libraries:
[pip3] botorch==0.14.0
[pip3] gpytorch==1.14
[pip3] mypy-extensions==1.0.0
[pip3] numpy==2.2.4
[pip3] torch==2.6.0+cpu
[conda] botorch 0.14.0 pypi_0 pypi
[conda] gpytorch 1.14 pypi_0 pypi
[conda] numpy 2.2.4 pypi_0 pypi
[conda] torch 2.6.0+cpu pypi_0 pypi