From 258b6c5ce82f5baa22a3a4e22cba52c84c857e2a Mon Sep 17 00:00:00 2001 From: Slava Shen Date: Thu, 5 Nov 2026 21:07:25 +0500 Subject: [PATCH] Fix MAISI concatenation memory handling Replace the fixed dimension cutoff with GPU-first concatenation and CPU fallback on CUDA out-of-memory errors. Add an optional free-memory margin, a fallback warning, and correctness regression tests. Refs: Project-MONAI/MONAI#9250 Assisted-by: OpenAI Codex Signed-off-by: Slava Shen --- docs/source/apps.rst | 31 +++ .../maisi/networks/autoencoderkl_maisi.py | 145 ++++++++--- .../networks/test_maisi_concatenation.py | 242 ++++++++++++++++++ 3 files changed, 380 insertions(+), 38 deletions(-) create mode 100644 tests/apps/maisi/networks/test_maisi_concatenation.py diff --git a/docs/source/apps.rst b/docs/source/apps.rst index 32c41d9ed2d..d189e444073 100644 --- a/docs/source/apps.rst +++ b/docs/source/apps.rst @@ -256,6 +256,37 @@ FastMRIReader ~~~~~~~~~~~~~ .. autofunction:: monai.apps.reconstruction.complex_utils.complex_conj +`MAISI` +------- + +.. autoclass:: monai.apps.generation.maisi.networks.autoencoderkl_maisi.AutoencoderKlMaisi + :members: + +By default, ``AutoencoderKlMaisi`` concatenates split convolution outputs and normalization +groups on their original device, falling back to CPU concatenation if CUDA concatenation +runs out of memory. The result is returned to the original device. Set ``concat_memory_margin`` +to a nonnegative number of bytes to also check available CUDA memory before concatenation: + +.. code-block:: python + + from monai.apps.generation.maisi.networks.autoencoderkl_maisi import AutoencoderKlMaisi + + model = AutoencoderKlMaisi( + spatial_dims=3, + in_channels=1, + out_channels=1, + num_res_blocks=(1, 1, 1), + num_channels=(32, 64, 64), + attention_levels=(False, False, False), + concat_memory_margin=512 * 1024**2, # 512 MiB + ) + +With a margin set, concatenation uses the CPU when the estimated output size plus the margin +exceeds the currently free CUDA memory. This check uses a snapshot of driver-reported free +memory, excludes reusable allocator cache, and does not reserve memory or guarantee that +other model operations or concurrent allocations fit. The default ``None`` skips this check; +CPU inputs do not query CUDA memory. + `Vista3d` --------- .. automodule:: monai.apps.vista3d.inferer diff --git a/monai/apps/generation/maisi/networks/autoencoderkl_maisi.py b/monai/apps/generation/maisi/networks/autoencoderkl_maisi.py index 90ffff6c340..5980fdf9be7 100644 --- a/monai/apps/generation/maisi/networks/autoencoderkl_maisi.py +++ b/monai/apps/generation/maisi/networks/autoencoderkl_maisi.py @@ -11,8 +11,8 @@ from __future__ import annotations -import gc import logging +import warnings from collections.abc import Sequence import torch @@ -34,6 +34,52 @@ def _empty_cuda_cache(save_mem: bool) -> None: return +def _validate_concat_memory_margin(concat_memory_margin: int | None) -> None: + if concat_memory_margin is not None and ( + isinstance(concat_memory_margin, bool) or not isinstance(concat_memory_margin, int) or concat_memory_margin < 0 + ): + raise ValueError("concat_memory_margin must be a nonnegative integer number of bytes or None.") + + +def _concatenate_with_cpu_fallback( + tensors: list[torch.Tensor], dim: int, save_mem: bool = True, concat_memory_margin: int | None = None +) -> torch.Tensor: + """Concatenate on the input device, offloading CUDA tensors when necessary. + + The optional margin compares the output allocation plus the margin with a snapshot of free CUDA memory. + CPU offloading replaces the input list entries to release their CUDA storage where possible. + """ + target_device = tensors[0].device + if target_device.type != "cuda": + return torch.cat(tensors, dim=dim) + + offload = False + if concat_memory_margin is not None: + output_dtype = tensors[0].dtype + for index in range(1, len(tensors)): + output_dtype = torch.promote_types(output_dtype, tensors[index].dtype) + element_size = torch.empty((), dtype=output_dtype, device="cpu").element_size() + output_bytes = sum(tensor.numel() for tensor in tensors) * element_size + free_bytes, _ = torch.cuda.mem_get_info(target_device) + offload = output_bytes + concat_memory_margin > free_bytes + + if not offload: + try: + return torch.cat(tensors, dim=dim) + except torch.cuda.OutOfMemoryError: + warnings.warn( + "CUDA concatenation ran out of memory; falling back to CPU. This may reduce performance.", + RuntimeWarning, + stacklevel=1, + ) + + # Leave the exception handler before releasing CUDA tensors so its traceback cannot keep them alive. + for index in range(len(tensors)): + tensors[index] = tensors[index].to("cpu") + _empty_cuda_cache(save_mem) + return torch.cat(tensors, dim=dim).to(target_device) + + class MaisiGroupNorm3D(nn.GroupNorm): """ Custom 3D Group Normalization with optional print_info output. @@ -46,6 +92,8 @@ class MaisiGroupNorm3D(nn.GroupNorm): norm_float16: If True, convert output of MaisiGroupNorm3D to float16 format, default to `False`. print_info: Whether to print information, default to `False`. save_mem: Whether to clean CUDA cache in order to save GPU memory, default to `True`. + concat_memory_margin: Optional free CUDA memory margin in bytes for concatenation. Defaults to `None`, + which attempts CUDA concatenation without a memory estimate. CUDA out-of-memory errors fall back to CPU. """ def __init__( @@ -57,11 +105,15 @@ def __init__( norm_float16: bool = False, print_info: bool = False, save_mem: bool = True, + *, + concat_memory_margin: int | None = None, ): + _validate_concat_memory_margin(concat_memory_margin) super().__init__(num_groups, num_channels, eps, affine) self.norm_float16 = norm_float16 self.print_info = print_info self.save_mem = save_mem + self.concat_memory_margin = concat_memory_margin def forward(self, input: torch.Tensor) -> torch.Tensor: if self.print_info: @@ -86,7 +138,7 @@ def forward(self, input: torch.Tensor) -> torch.Tensor: del input _empty_cuda_cache(self.save_mem) - input = torch.cat(inputs, dim=1) if max(inputs[0].size()) < 500 else self._cat_inputs(inputs) + input = self._cat_inputs(inputs) input = input.view(param_n, param_c, param_d, param_h, param_w) if self.affine: @@ -97,22 +149,8 @@ def forward(self, input: torch.Tensor) -> torch.Tensor: return input - def _cat_inputs(self, inputs): - input_type = inputs[0].device.type - input = inputs[0].clone().to("cpu", non_blocking=True) if input_type == "cuda" else inputs[0].clone() - inputs[0] = 0 - _empty_cuda_cache(self.save_mem) - - for k in range(len(inputs) - 1): - input = torch.cat((input, inputs[k + 1].cpu()), dim=1) - inputs[k + 1] = 0 - _empty_cuda_cache(self.save_mem) - gc.collect() - - if self.print_info: - logger.info(f"MaisiGroupNorm3D concat progress: {k + 1}/{len(inputs) - 1}.") - - return input.to("cuda", non_blocking=True) if input_type == "cuda" else input + def _cat_inputs(self, inputs: list[torch.Tensor]) -> torch.Tensor: + return _concatenate_with_cpu_fallback(inputs, 1, self.save_mem, self.concat_memory_margin) class MaisiConvolution(nn.Module): @@ -127,6 +165,8 @@ class MaisiConvolution(nn.Module): dim_split: Dimension of splitting for the input tensor. print_info: Whether to print information. save_mem: Whether to clean CUDA cache in order to save GPU memory, default to `True`. + concat_memory_margin: Optional free CUDA memory margin in bytes for concatenation. Defaults to `None`, + which attempts CUDA concatenation without a memory estimate. CUDA out-of-memory errors fall back to CPU. Additional arguments for the convolution operation. https://monai.readthedocs.io/en/stable/networks.html#convolution """ @@ -154,7 +194,10 @@ def __init__( is_transposed: bool = False, padding: Sequence[int] | int | None = None, output_padding: Sequence[int] | int | None = None, + *, + concat_memory_margin: int | None = None, ) -> None: + _validate_concat_memory_margin(concat_memory_margin) super().__init__() self.conv = Convolution( spatial_dims=spatial_dims, @@ -181,6 +224,7 @@ def __init__( self.num_splits = num_splits self.print_info = print_info self.save_mem = save_mem + self.concat_memory_margin = concat_memory_margin def _split_tensor(self, x: torch.Tensor, split_size: int, padding: int) -> list[torch.Tensor]: overlaps = [0] + [padding] * (self.num_splits - 1) @@ -211,26 +255,7 @@ def _concatenate_tensors(self, outputs: list[torch.Tensor], split_size: int, pad for i in range(self.num_splits): logger.info(f"Output {i + 1}/{len(outputs)} size after: {outputs[i].size()}") - if max(outputs[0].size()) < 500: - x = torch.cat(outputs, dim=self.dim_split + 2) - else: - target_device = outputs[0].device - - x = outputs[0].clone().to("cpu", non_blocking=True) - outputs[0] = torch.Tensor(0) - _empty_cuda_cache(self.save_mem) - for k in range(len(outputs) - 1): - x = torch.cat((x, outputs[k + 1].cpu()), dim=self.dim_split + 2) - outputs[k + 1] = torch.Tensor(0) - _empty_cuda_cache(self.save_mem) - gc.collect() - if self.print_info: - logger.info(f"MaisiConvolution concat progress: {k + 1}/{len(outputs) - 1}.") - - if target_device.type != "cpu": - x = x.to(target_device, non_blocking=True) - - return x + return _concatenate_with_cpu_fallback(outputs, self.dim_split + 2, self.save_mem, self.concat_memory_margin) def forward(self, x: torch.Tensor) -> torch.Tensor: if self.print_info: @@ -297,6 +322,8 @@ class MaisiUpsample(nn.Module): dim_split: Dimension of splitting for the input tensor. print_info: Whether to print information. save_mem: Whether to clean CUDA cache in order to save GPU memory, default to `True`. + concat_memory_margin: Optional free CUDA memory margin in bytes for concatenation. Defaults to `None`, + which attempts CUDA concatenation without a memory estimate. CUDA out-of-memory errors fall back to CPU. """ def __init__( @@ -308,6 +335,8 @@ def __init__( dim_split: int, print_info: bool, save_mem: bool = True, + *, + concat_memory_margin: int | None = None, ) -> None: super().__init__() self.conv = MaisiConvolution( @@ -323,6 +352,7 @@ def __init__( dim_split=dim_split, print_info=print_info, save_mem=save_mem, + concat_memory_margin=concat_memory_margin, ) self.use_convtranspose = use_convtranspose self.save_mem = save_mem @@ -353,6 +383,8 @@ class MaisiDownsample(nn.Module): dim_split: Dimension of splitting for the input tensor. print_info: Whether to print information. save_mem: Whether to clean CUDA cache in order to save GPU memory, default to `True`. + concat_memory_margin: Optional free CUDA memory margin in bytes for concatenation. Defaults to `None`, + which attempts CUDA concatenation without a memory estimate. CUDA out-of-memory errors fall back to CPU. """ def __init__( @@ -363,6 +395,8 @@ def __init__( dim_split: int, print_info: bool, save_mem: bool = True, + *, + concat_memory_margin: int | None = None, ) -> None: super().__init__() self.pad = (0, 1) * spatial_dims @@ -378,6 +412,7 @@ def __init__( dim_split=dim_split, print_info=print_info, save_mem=save_mem, + concat_memory_margin=concat_memory_margin, ) def forward(self, x: torch.Tensor) -> torch.Tensor: @@ -402,6 +437,8 @@ class MaisiResBlock(nn.Module): norm_float16: If True, convert output of MaisiGroupNorm3D to float16 format, default to `False`. print_info: Whether to print information, default to `False`. save_mem: Whether to clean CUDA cache in order to save GPU memory, default to `True`. + concat_memory_margin: Optional free CUDA memory margin in bytes for concatenation. Defaults to `None`, + which attempts CUDA concatenation without a memory estimate. CUDA out-of-memory errors fall back to CPU. """ def __init__( @@ -416,6 +453,8 @@ def __init__( norm_float16: bool = False, print_info: bool = False, save_mem: bool = True, + *, + concat_memory_margin: int | None = None, ) -> None: super().__init__() self.in_channels = in_channels @@ -430,6 +469,7 @@ def __init__( norm_float16=norm_float16, print_info=print_info, save_mem=save_mem, + concat_memory_margin=concat_memory_margin, ) self.conv1 = MaisiConvolution( spatial_dims=spatial_dims, @@ -443,6 +483,7 @@ def __init__( dim_split=dim_split, print_info=print_info, save_mem=save_mem, + concat_memory_margin=concat_memory_margin, ) self.norm2 = MaisiGroupNorm3D( num_groups=norm_num_groups, @@ -452,6 +493,7 @@ def __init__( norm_float16=norm_float16, print_info=print_info, save_mem=save_mem, + concat_memory_margin=concat_memory_margin, ) self.conv2 = MaisiConvolution( spatial_dims=spatial_dims, @@ -465,6 +507,7 @@ def __init__( dim_split=dim_split, print_info=print_info, save_mem=save_mem, + concat_memory_margin=concat_memory_margin, ) self.nin_shortcut = ( @@ -480,6 +523,7 @@ def __init__( dim_split=dim_split, print_info=print_info, save_mem=save_mem, + concat_memory_margin=concat_memory_margin, ) if self.in_channels != self.out_channels else nn.Identity() @@ -533,6 +577,8 @@ class MaisiEncoder(nn.Module): norm_float16: If True, convert output of MaisiGroupNorm3D to float16 format, default to `False`. print_info: Whether to print information, default to `False`. save_mem: Whether to clean CUDA cache in order to save GPU memory, default to `True`. + concat_memory_margin: Optional free CUDA memory margin in bytes for concatenation. Defaults to `None`, + which attempts CUDA concatenation without a memory estimate. CUDA out-of-memory errors fall back to CPU. """ def __init__( @@ -554,6 +600,8 @@ def __init__( include_fc: bool = False, use_combined_linear: bool = False, use_flash_attention: bool = False, + *, + concat_memory_margin: int | None = None, ) -> None: super().__init__() @@ -582,6 +630,7 @@ def __init__( dim_split=dim_split, print_info=print_info, save_mem=save_mem, + concat_memory_margin=concat_memory_margin, ) ) @@ -604,6 +653,7 @@ def __init__( norm_float16=norm_float16, print_info=print_info, save_mem=save_mem, + concat_memory_margin=concat_memory_margin, ) ) input_channel = output_channel @@ -629,6 +679,7 @@ def __init__( dim_split=dim_split, print_info=print_info, save_mem=save_mem, + concat_memory_margin=concat_memory_margin, ) ) @@ -673,6 +724,7 @@ def __init__( norm_float16=norm_float16, print_info=print_info, save_mem=save_mem, + concat_memory_margin=concat_memory_margin, ) ) blocks.append( @@ -688,6 +740,7 @@ def __init__( dim_split=dim_split, print_info=print_info, save_mem=save_mem, + concat_memory_margin=concat_memory_margin, ) ) @@ -723,6 +776,8 @@ class MaisiDecoder(nn.Module): norm_float16: If True, convert output of MaisiGroupNorm3D to float16 format, default to `False`. print_info: Whether to print information, default to `False`. save_mem: Whether to clean CUDA cache in order to save GPU memory, default to `True`. + concat_memory_margin: Optional free CUDA memory margin in bytes for concatenation. Defaults to `None`, + which attempts CUDA concatenation without a memory estimate. CUDA out-of-memory errors fall back to CPU. """ def __init__( @@ -745,6 +800,8 @@ def __init__( use_combined_linear: bool = False, use_flash_attention: bool = False, use_convtranspose: bool = False, + *, + concat_memory_margin: int | None = None, ) -> None: super().__init__() self.print_info = print_info @@ -767,6 +824,7 @@ def __init__( dim_split=dim_split, print_info=print_info, save_mem=save_mem, + concat_memory_margin=concat_memory_margin, ) ) @@ -822,6 +880,7 @@ def __init__( norm_float16=norm_float16, print_info=print_info, save_mem=save_mem, + concat_memory_margin=concat_memory_margin, ) ) block_in_ch = block_out_ch @@ -849,6 +908,7 @@ def __init__( dim_split=dim_split, print_info=print_info, save_mem=save_mem, + concat_memory_margin=concat_memory_margin, ) ) @@ -861,6 +921,7 @@ def __init__( norm_float16=norm_float16, print_info=print_info, save_mem=save_mem, + concat_memory_margin=concat_memory_margin, ) ) blocks.append( @@ -876,6 +937,7 @@ def __init__( dim_split=dim_split, print_info=print_info, save_mem=save_mem, + concat_memory_margin=concat_memory_margin, ) ) @@ -914,6 +976,8 @@ class AutoencoderKlMaisi(AutoencoderKL): norm_float16: If True, convert output of MaisiGroupNorm3D to float16 format, default to `False`. print_info: Whether to print information, default to `False`. save_mem: Whether to clean CUDA cache in order to save GPU memory, default to `True`. + concat_memory_margin: Optional free CUDA memory margin in bytes for concatenation. Defaults to `None`, + which attempts CUDA concatenation without a memory estimate. CUDA out-of-memory errors fall back to CPU. """ def __init__( @@ -939,7 +1003,10 @@ def __init__( norm_float16: bool = False, print_info: bool = False, save_mem: bool = True, + *, + concat_memory_margin: int | None = None, ) -> None: + _validate_concat_memory_margin(concat_memory_margin) super().__init__( spatial_dims, in_channels, @@ -977,6 +1044,7 @@ def __init__( norm_float16=norm_float16, print_info=print_info, save_mem=save_mem, + concat_memory_margin=concat_memory_margin, ) self.decoder: nn.Module = MaisiDecoder( @@ -998,4 +1066,5 @@ def __init__( norm_float16=norm_float16, print_info=print_info, save_mem=save_mem, + concat_memory_margin=concat_memory_margin, ) diff --git a/tests/apps/maisi/networks/test_maisi_concatenation.py b/tests/apps/maisi/networks/test_maisi_concatenation.py new file mode 100644 index 00000000000..f340906ad66 --- /dev/null +++ b/tests/apps/maisi/networks/test_maisi_concatenation.py @@ -0,0 +1,242 @@ +# Copyright (c) MONAI Consortium +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import unittest +import warnings +from unittest.mock import patch + +import torch +import torch.nn.functional as F +from parameterized import parameterized + +from monai.apps.generation.maisi.networks.autoencoderkl_maisi import ( + AutoencoderKlMaisi, + MaisiConvolution, + MaisiGroupNorm3D, + _concatenate_with_cpu_fallback, +) + + +class _BoundaryChecks: + device = torch.device("cpu") + + @parameterized.expand([(499, False), (500, False), (512, False), (500, True)]) + def test_group_norm(self, length, norm_float16): + image = torch.randn(1, 6, length, 2, 3, device=self.device) + norm = MaisiGroupNorm3D(3, 6, norm_float16=norm_float16, save_mem=False).to(self.device) + expected = F.group_norm(image, 3, norm.weight, norm.bias, norm.eps) + if norm_float16: + expected = expected.half() + + with patch("torch.cat", wraps=torch.cat) as concatenate: + result = norm(image) + + torch.testing.assert_close(result, expected) + self.assertEqual(result.device, self.device) + self.assertEqual(concatenate.call_count, 1) + chunks = concatenate.call_args.args[0] + self.assertEqual(len(chunks), 3) + self.assertTrue(all(chunk.device == self.device for chunk in chunks)) + + @parameterized.expand([(499, 0), (500, 0), (512, 0), (500, 1), (500, 2)]) + def test_convolution_overlap(self, length, split_dimension): + shape = [1, 2, 2, 2, 2] + axis = split_dimension + 2 + shape[axis] = 6 + shape[2 + (split_dimension + 1) % 3] = length + image = torch.randn(shape, device=self.device, requires_grad=True) + outputs = [image.narrow(axis, start, size) for start, size in [(0, 3), (1, 4), (3, 3)]] + convolution = MaisiConvolution( + 3, 2, 2, num_splits=3, dim_split=split_dimension, print_info=False, save_mem=False, conv_only=True + ) + + with patch("torch.cat", wraps=torch.cat) as concatenate: + result = convolution._concatenate_tensors(outputs, split_size=2, padding=1) + + torch.testing.assert_close(result, image, rtol=0, atol=0) + self.assertEqual(result.device, self.device) + self.assertEqual(concatenate.call_count, 1) + self.assertTrue(all(chunk.device == self.device for chunk in concatenate.call_args.args[0])) + result.square().sum().backward() + torch.testing.assert_close(image.grad, 2 * image.detach()) + + +class TestMaisiConcatenationCPU(_BoundaryChecks, unittest.TestCase): + def test_cpu_preserves_promoted_dtype_and_gradients(self): + inputs = [torch.randn(2, 3, dtype=dtype, requires_grad=True) for dtype in (torch.float32, torch.float64)] + expected = torch.cat(inputs, dim=1) + + with warnings.catch_warnings(record=True) as emitted, patch("torch.cuda.mem_get_info") as memory_info: + warnings.simplefilter("always") + result = _concatenate_with_cpu_fallback(inputs, 1, concat_memory_margin=100) + + memory_info.assert_not_called() + self.assertEqual(emitted, []) + torch.testing.assert_close(result, expected, rtol=0, atol=0) + self.assertEqual(result.dtype, torch.float64) + gradients = torch.autograd.grad(result.square().sum(), inputs) + for tensor, gradient in zip(inputs, gradients): + torch.testing.assert_close(gradient, 2 * tensor.detach()) + + def test_cpu_oom_is_not_retried(self): + inputs = [torch.ones(2, 3), torch.ones(2, 3)] + error = torch.cuda.OutOfMemoryError("simulated allocation failure") + with patch("torch.cat", side_effect=error) as concatenate: + with self.assertRaises(torch.cuda.OutOfMemoryError) as raised: + _concatenate_with_cpu_fallback(inputs, 1) + self.assertIs(raised.exception, error) + concatenate.assert_called_once() + + @parameterized.expand([(-1,), (True,), (1.5,), ("100",)]) + def test_invalid_memory_margin(self, margin): + with self.assertRaises(ValueError): + MaisiGroupNorm3D(1, 2, concat_memory_margin=margin) + with self.assertRaises(ValueError): + MaisiConvolution(3, 2, 2, 2, 0, False, concat_memory_margin=margin) + with self.assertRaises(ValueError): + self._make_autoencoder(margin) + + def test_autoencoder_propagates_memory_margin(self): + model = self._make_autoencoder(12345) + leaves = [module for module in model.modules() if isinstance(module, (MaisiGroupNorm3D, MaisiConvolution))] + self.assertTrue(leaves) + for module in leaves: + with self.subTest(module=type(module).__name__): + self.assertEqual(module.concat_memory_margin, 12345) + + @staticmethod + def _make_autoencoder(margin): + return AutoencoderKlMaisi( + spatial_dims=3, + in_channels=1, + out_channels=1, + num_channels=(2, 2), + latent_channels=2, + attention_levels=(False, False), + num_res_blocks=(1, 1), + norm_num_groups=1, + with_encoder_nonlocal_attn=False, + with_decoder_nonlocal_attn=False, + num_splits=2, + print_info=False, + save_mem=False, + concat_memory_margin=margin, + ) + + +@unittest.skipUnless(torch.cuda.is_available(), "CUDA is required") +class TestMaisiConcatenationCUDA(_BoundaryChecks, unittest.TestCase): + device = torch.device("cuda:0") + + @parameterized.expand([(torch.float32,), (torch.float64,)]) + def test_oom_fallback_preserves_values_device_dtype_and_gradients(self, dtype): + inputs = [torch.randn(2, 3, device=self.device, dtype=dtype, requires_grad=True) for _ in range(3)] + expected = torch.cat(inputs, dim=1) + original_inputs = inputs.copy() + concatenate = torch.cat + attempts = [] + + def simulate_cuda_oom(tensors, dim): + attempts.append(tensors[0].device.type) + if tensors[0].is_cuda: + raise torch.cuda.OutOfMemoryError("simulated CUDA allocation failure") + return concatenate(tensors, dim=dim) + + with ( + patch("torch.cat", side_effect=simulate_cuda_oom), + self.assertWarnsRegex( + RuntimeWarning, + r"^CUDA concatenation ran out of memory; falling back to CPU\. This may reduce performance\.$", + ) as warning, + ): + result = _concatenate_with_cpu_fallback(inputs, 1, save_mem=False) + + self.assertEqual(warning.filename, _concatenate_with_cpu_fallback.__code__.co_filename) + self.assertEqual(attempts, ["cuda", "cpu"]) + self.assertEqual(result.device, self.device) + self.assertEqual(result.dtype, dtype) + torch.testing.assert_close(result, expected, rtol=0, atol=0) + gradients = torch.autograd.grad(result.square().sum(), original_inputs) + for tensor, gradient in zip(original_inputs, gradients): + torch.testing.assert_close(gradient, 2 * tensor.detach()) + + def test_non_oom_error_is_not_retried(self): + inputs = [torch.ones(2, 3, device=self.device), torch.ones(2, 3, device=self.device)] + error = RuntimeError("simulated invalid concatenation") + with patch("torch.cat", side_effect=error) as concatenate: + with self.assertRaises(RuntimeError) as raised: + _concatenate_with_cpu_fallback(inputs, 1, save_mem=False) + self.assertIs(raised.exception, error) + concatenate.assert_called_once() + + @parameterized.expand([(499, False), (500, False), (512, False), (499, True), (500, True), (512, True)]) + def test_group_norm_cpu_fallback_is_bit_exact(self, length, norm_float16): + image = torch.randn(1, 6, length, 2, 3, device=self.device) + norm = MaisiGroupNorm3D(3, 6, norm_float16=norm_float16, save_mem=False).to(self.device) + expected = norm(image) + norm.concat_memory_margin = 0 + + with ( + patch("torch.cuda.mem_get_info", return_value=(0, image.numel() * image.element_size())) as memory_info, + patch("torch.cat", wraps=torch.cat) as concatenate, + ): + result = norm(image) + + memory_info.assert_called_once_with(self.device) + self.assertTrue(all(chunk.device.type == "cpu" for chunk in concatenate.call_args.args[0])) + torch.testing.assert_close(result, expected, rtol=0, atol=0) + + @parameterized.expand([(0, -1), (0, 0), (100, -1), (100, 0)]) + def test_memory_margin_uses_promoted_output_bytes(self, margin, spare_bytes): + inputs = [torch.ones(2, 3, device=self.device, dtype=dtype) for dtype in (torch.float32, torch.float64)] + expected = torch.cat(inputs, dim=1) + required_bytes = expected.numel() * expected.element_size() + margin + + with ( + warnings.catch_warnings(record=True) as emitted, + patch( + "torch.cuda.mem_get_info", return_value=(required_bytes + spare_bytes, required_bytes * 2) + ) as memory_info, + patch("torch.cat", wraps=torch.cat) as concatenate, + ): + warnings.simplefilter("always") + result = _concatenate_with_cpu_fallback(inputs, 1, save_mem=False, concat_memory_margin=margin) + + self.assertEqual(emitted, []) + memory_info.assert_called_once_with(self.device) + self.assertEqual(concatenate.call_count, 1) + concat_device = concatenate.call_args.args[0][0].device + self.assertEqual(concat_device.type, "cpu" if spare_bytes < 0 else "cuda") + self.assertEqual(result.device, self.device) + torch.testing.assert_close(result, expected, rtol=0, atol=0) + + def test_default_does_not_query_free_memory(self): + inputs = [torch.ones(2, 3, device=self.device), torch.ones(2, 3, device=self.device)] + with warnings.catch_warnings(record=True) as emitted, patch("torch.cuda.mem_get_info") as memory_info: + warnings.simplefilter("always") + result = _concatenate_with_cpu_fallback(inputs, 1, save_mem=False) + memory_info.assert_not_called() + self.assertEqual(emitted, []) + self.assertEqual(result.device, self.device) + + @unittest.skipUnless(torch.cuda.device_count() > 1, "Multiple CUDA devices are required") + def test_fallback_restores_noncurrent_device(self): + target = torch.device("cuda:1") + inputs = [torch.ones(2, 3, device=target), torch.ones(2, 3, device=target)] + with torch.cuda.device(0), patch("torch.cuda.mem_get_info", return_value=(0, 100)): + result = _concatenate_with_cpu_fallback(inputs, 1, save_mem=False, concat_memory_margin=0) + self.assertEqual(result.device, target) + + +if __name__ == "__main__": + unittest.main()