Avoid mixing allocators in mGPU fused losses

src/fvdb/detail/ops/gsplat/FusedL1.cu

@@ -166,14 +166,8 @@ forwardL1Tiles(const torch::Tensor &img1,
                const torch::Tensor &img2,
                int blockOffset,
                int blockCount,
-               float *loss,
+               float *partial_sums,
                cudaStream_t stream) {
-    if (!blockCount) {
-        C10_CUDA_CHECK(cudaMemsetAsync(loss, 0, sizeof(float), stream));
-        return;
-    }
-    auto partial_sums = torch::empty(
-        {blockCount}, img1.options().device(torch::kCUDA, c10::cuda::current_device()));
     dim3 grid(blockCount);
     dim3 block(BLOCK_X, BLOCK_Y);
     fusedL1Kernel<<<grid, block, 0, stream>>>(blockOffset,
@@ -183,10 +177,8 @@ forwardL1Tiles(const torch::Tensor &img1,
                                               img1.size(1),
                                               img1.const_data_ptr<float>(),
                                               img2.const_data_ptr<float>(),
-                                              partial_sums.data_ptr<float>());
+                                              partial_sums);
     C10_CUDA_KERNEL_LAUNCH_CHECK();
-    FVDB_CUB_WRAPPER_ASYNC(
-        stream, cub::DeviceReduce::Sum, partial_sums.const_data_ptr<float>(), loss, blockCount);
 }
 
 } // namespace
@@ -209,7 +201,13 @@ fusedL1CUDA(const torch::Tensor &img1, const torch::Tensor &img2) {
     }
     const int blockCount = ((W + BLOCK_X - 1) / BLOCK_X) * ((H + BLOCK_Y - 1) / BLOCK_Y) * B;
     auto loss            = torch::empty({}, img1.options());
-    forwardL1Tiles(img1, img2, 0, blockCount, loss.data_ptr<float>(), stream);
+    auto partial_sums    = torch::empty({blockCount}, img1.options());
+    forwardL1Tiles(img1, img2, 0, blockCount, partial_sums.data_ptr<float>(), stream);
+    FVDB_CUB_WRAPPER(cub::DeviceReduce::Sum,
+                     partial_sums.data_ptr<float>(),
+                     loss.data_ptr<float>(),
+                     blockCount,
+                     stream);
     return loss;
 }
 
@@ -277,18 +275,32 @@ fusedL1PrivateUse1(const torch::Tensor &img1, const torch::Tensor &img2) {
         C10_CUDA_CHECK(cudaSetDevice(deviceId));
         auto stream = c10::cuda::getCurrentCUDAStream(deviceId);
 
-        const auto chunk = imageBlockChunk(B, H, W, deviceId);
-        partial_losses.push_back(torch::empty({1}, img1.options().device(torch::kCUDA, deviceId)));
-        forwardL1Tiles(img1,
-                       img2,
-                       chunk.blockOffset,
-                       chunk.blockCount,
-                       partial_losses.back().data_ptr<float>(),
-                       stream);
+        const auto chunk    = imageBlockChunk(B, H, W, deviceId);
+        float *partial_loss = nullptr;
+        C10_CUDA_CHECK(cudaMallocAsync(&partial_loss, sizeof(float), stream));
+        // The CUDA device is only metadata for NCCL; from_blob does not allocate storage.
+        partial_losses.push_back(
+            torch::from_blob(partial_loss, {1}, img1.options().device(torch::kCUDA, deviceId)));
+        if (!chunk.blockCount) {
+            C10_CUDA_CHECK(cudaMemsetAsync(partial_loss, 0, sizeof(float), stream));
+            continue;
+        }
+        float *partial_sums = nullptr;
+        C10_CUDA_CHECK(cudaMallocAsync(&partial_sums, chunk.blockCount * sizeof(float), stream));
+        forwardL1Tiles(img1, img2, chunk.blockOffset, chunk.blockCount, partial_sums, stream);
+        FVDB_CUB_WRAPPER_ASYNC(
+            stream, cub::DeviceReduce::Sum, partial_sums, partial_loss, chunk.blockCount);
+        C10_CUDA_CHECK(cudaFreeAsync(partial_sums, stream));
     }
     auto nccl_loss =
         torch::from_blob(loss.data_ptr<float>(), {1}, partial_losses.front().options());
     torch::cuda::nccl::reduce(partial_losses, nccl_loss);
+    // Release each scalar on its producing stream after NCCL has consumed it.
+    for (const auto deviceId: c10::irange(c10::cuda::device_count())) {
+        C10_CUDA_CHECK(cudaSetDevice(deviceId));
+        auto stream = c10::cuda::getCurrentCUDAStream(deviceId);
+        C10_CUDA_CHECK(cudaFreeAsync(partial_losses[deviceId].data_ptr<float>(), stream));
+    }
     mergeStreams();
     return loss;
 }

src/fvdb/detail/ops/gsplat/FusedL1SSIM.cu

@@ -565,21 +565,13 @@ forwardTiles(const torch::Tensor &img1,
              float C2,
              int blockOffset,
              int blockCount,
-             float *loss,
+             float *loss_tiles,
              cudaStream_t stream,
              bool normalize_target) {
-    if (!blockCount) {
-        C10_CUDA_CHECK(cudaMemsetAsync(loss, 0, sizeof(float), stream));
-        return;
-    }
     const int B  = img1.size(0);
     const int CH = img1.size(1);
     const int H  = img1.size(2);
     const int W  = img1.size(3);
-    // Keep tile contributions on the GPU that produced them.
-    auto tile_losses = torch::empty(
-        {blockCount}, img1.options().device(torch::kCUDA, c10::cuda::current_device()));
-    auto *loss_tiles = tile_losses.data_ptr<float>();
 
     dim3 grid(blockCount);
     dim3 block(BLOCK_X, BLOCK_Y);
@@ -606,7 +598,6 @@ forwardTiles(const torch::Tensor &img1,
                    at::kFloat,
                    at::kByte);
     C10_CUDA_KERNEL_LAUNCH_CHECK();
-    FVDB_CUB_WRAPPER_ASYNC(stream, cub::DeviceReduce::Sum, loss_tiles, loss, blockCount);
 }
 
 void
@@ -622,9 +613,6 @@ backwardTiles(const torch::Tensor &img1,
               int blockCount,
               cudaStream_t stream,
               bool normalize_target) {
-    if (!blockCount) {
-        return;
-    }
     dim3 grid(blockCount);
     dim3 block(BLOCK_X, BLOCK_Y);
     AT_DISPATCH_V2(img2.scalar_type(),
@@ -676,6 +664,7 @@ fusedL1SSIMCUDA(const torch::Tensor &img1,
 
     auto stream          = c10::cuda::getCurrentCUDAStream(img1.device().index());
     const int blockCount = B * ((H + BLOCK_Y - 1) / BLOCK_Y) * ((W + BLOCK_X - 1) / BLOCK_X);
+    auto tile_losses     = torch::empty({blockCount}, img1.options());
     forwardTiles(img1,
                  img2,
                  dm_dmu1,
@@ -686,9 +675,14 @@ fusedL1SSIMCUDA(const torch::Tensor &img1,
                  static_cast<float>(C2),
                  0,
                  blockCount,
-                 loss.data_ptr<float>(),
+                 tile_losses.data_ptr<float>(),
                  stream,
                  normalize_target);
+    FVDB_CUB_WRAPPER(cub::DeviceReduce::Sum,
+                     tile_losses.data_ptr<float>(),
+                     loss.data_ptr<float>(),
+                     blockCount,
+                     stream);
     return std::make_tuple(loss, dm_dmu1, dm_dsigma1_sq, dm_dsigma12);
 }
 
@@ -719,9 +713,20 @@ fusedL1SSIMPrivateUse1(const torch::Tensor &img1,
     partial_losses.reserve(c10::cuda::device_count());
     for (const auto deviceId: c10::irange(c10::cuda::device_count())) {
         C10_CUDA_CHECK(cudaSetDevice(deviceId));
-        auto stream      = c10::cuda::getCurrentCUDAStream(deviceId);
-        const auto chunk = imageBlockChunk(B, H, W, deviceId);
-        partial_losses.push_back(torch::empty({1}, img1.options().device(torch::kCUDA, deviceId)));
+        auto stream         = c10::cuda::getCurrentCUDAStream(deviceId);
+        const auto chunk    = imageBlockChunk(B, H, W, deviceId);
+        float *partial_loss = nullptr;
+        C10_CUDA_CHECK(cudaMallocAsync(&partial_loss, sizeof(float), stream));
+        // The CUDA device is only metadata for NCCL; from_blob does not allocate storage.
+        partial_losses.push_back(
+            torch::from_blob(partial_loss, {1}, img1.options().device(torch::kCUDA, deviceId)));
+        if (!chunk.blockCount) {
+            C10_CUDA_CHECK(cudaMemsetAsync(partial_loss, 0, sizeof(float), stream));
+            continue;
+        }
+        // Keep tile contributions on the GPU that produced them.
+        float *loss_tiles = nullptr;
+        C10_CUDA_CHECK(cudaMallocAsync(&loss_tiles, chunk.blockCount * sizeof(float), stream));
         forwardTiles(img1,
                      img2,
                      dm_dmu1,
@@ -732,13 +737,22 @@ fusedL1SSIMPrivateUse1(const torch::Tensor &img1,
                      static_cast<float>(C2),
                      chunk.blockOffset,
                      chunk.blockCount,
-                     partial_losses.back().data_ptr<float>(),
+                     loss_tiles,
                      stream,
                      normalize_target);
+        FVDB_CUB_WRAPPER_ASYNC(
+            stream, cub::DeviceReduce::Sum, loss_tiles, partial_loss, chunk.blockCount);
+        C10_CUDA_CHECK(cudaFreeAsync(loss_tiles, stream));
     }
     auto nccl_loss =
         torch::from_blob(loss.data_ptr<float>(), {1}, partial_losses.front().options());
     torch::cuda::nccl::reduce(partial_losses, nccl_loss);
+    // Release each scalar on its producing stream after NCCL has consumed it.
+    for (const auto deviceId: c10::irange(c10::cuda::device_count())) {
+        C10_CUDA_CHECK(cudaSetDevice(deviceId));
+        auto stream = c10::cuda::getCurrentCUDAStream(deviceId);
+        C10_CUDA_CHECK(cudaFreeAsync(partial_losses[deviceId].data_ptr<float>(), stream));
+    }
     mergeStreams();
     return std::make_tuple(loss, dm_dmu1, dm_dsigma1_sq, dm_dsigma12);
 }
@@ -802,6 +816,9 @@ fusedL1SSIMBackwardPrivateUse1(const torch::Tensor &img1,
         C10_CUDA_CHECK(cudaSetDevice(deviceId));
         auto stream      = c10::cuda::getCurrentCUDAStream(deviceId);
         const auto chunk = imageBlockChunk(B, H, W, deviceId);
+        if (!chunk.blockCount) {
+            continue;
+        }
         backwardTiles(img1,
                       img2,
                       grad_loss,

tests/unit/test_fused_l1.py

@@ -1,48 +1,39 @@
 # Copyright Contributors to the OpenVDB Project
 # SPDX-License-Identifier: Apache-2.0
 
-from contextlib import ExitStack
-
 import pytest
 import torch
 import torch.nn.functional as nnf
 
 from fvdb.utils.metrics import fused_l1_loss, ssim
 
-
-@pytest.fixture(params=["cuda", "dgx"])
-def image_device(request):
-    if not torch.cuda.is_available():
-        pytest.skip("requires CUDA")
-    if request.param == "dgx":
-        pytest.importorskip("torch_dgx")
-    return request.param
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
 
 
 @pytest.mark.parametrize(
     "shape",
     [
-        (1, 3, 1, 1),  # More devices than image tile rows.
+        (1, 3, 1, 1),  # A single pixel in a partial tile.
         (1, 3, 16, 19),  # One complete tile row, with a partial width tile.
-        (1, 3, 33, 19),  # Unequal tile-row shards and a partial final row.
-        (3, 4, 17, 23),  # A device's shard can cross an image boundary.
+        (1, 3, 33, 19),  # Multiple tile rows with a partial final row.
+        (3, 4, 17, 23),  # Multiple images and channels.
         (2, 1, 32, 16),
     ],
 )
 @pytest.mark.parametrize("needs_grad", [(True, False), (False, True), (True, True)])
-def test_fused_l1_forward_backward(image_device, shape, needs_grad):
+def test_fused_l1_forward_backward(shape, needs_grad):
     generator = torch.Generator().manual_seed(17)
     img1 = torch.randn(shape, generator=generator)
     img2 = torch.randn(shape, generator=generator)
     img2.flatten()[::7] = img1.flatten()[::7]  # Include the zero subgradient.
     reference = [value.clone().requires_grad_(need) for value, need in zip((img1, img2), needs_grad)]
-    actual = [value.to(image_device).requires_grad_(need) for value, need in zip((img1, img2), needs_grad)]
+    actual = [value.to("cuda").requires_grad_(need) for value, need in zip((img1, img2), needs_grad)]
 
     expected_loss = nnf.l1_loss(*reference)
     actual_loss = fused_l1_loss(*actual)
     torch.testing.assert_close(actual_loss.cpu(), expected_loss.detach(), rtol=2e-6, atol=1e-6)
     expected_loss.backward(torch.tensor(-2.75))
-    actual_loss.backward(torch.tensor(-2.75, device=image_device))
+    actual_loss.backward(torch.tensor(-2.75, device="cuda"))
     for expected, result, need in zip(reference, actual, needs_grad):
         if need:
             torch.testing.assert_close(result.grad.cpu(), expected.grad)
@@ -51,11 +42,11 @@ def test_fused_l1_forward_backward(image_device, shape, needs_grad):
 
 
 @pytest.mark.parametrize("layout", ["strided", "offset"])
-def test_fused_l1_input_views(image_device, layout):
+def test_fused_l1_input_views(layout):
     generator = torch.Generator().manual_seed(23)
     values = [torch.randn(1, 3, 33, 38, generator=generator) for _ in range(2)]
     reference = [value.clone().requires_grad_() for value in values]
-    actual = [value.to(image_device).requires_grad_() for value in values]
+    actual = [value.to("cuda").requires_grad_() for value in values]
 
     def view(value):
         if layout == "strided":
@@ -73,8 +64,8 @@ def view(value):
 
 
 @pytest.mark.parametrize("shape", [(0, 3, 17, 19), (1, 0, 17, 19), (1, 3, 0, 19), (1, 3, 17, 0)])
-def test_fused_l1_empty(image_device, shape):
-    inputs = [torch.empty(shape, device=image_device, requires_grad=True) for _ in range(2)]
+def test_fused_l1_empty(shape):
+    inputs = [torch.empty(shape, device="cuda", requires_grad=True) for _ in range(2)]
     result = fused_l1_loss(*inputs)
     assert result.ndim == 0
     assert torch.isnan(result.cpu())
@@ -84,11 +75,11 @@ def test_fused_l1_empty(image_device, shape):
         assert value.grad.shape == value.shape
 
 
-def test_fused_l1_nonfinite_residuals(image_device):
+def test_fused_l1_nonfinite_residuals():
     img1 = torch.tensor([0.0, 2.0, -2.0, float("inf"), -float("inf"), float("nan"), float("inf")])
     img2 = torch.tensor([0.0, 1.0, -1.0, 0.0, 0.0, 1.0, float("inf")])
     reference = [value.view(1, 1, 1, -1).requires_grad_() for value in (img1, img2)]
-    actual = [value.detach().to(image_device).requires_grad_() for value in reference]
+    actual = [value.detach().to("cuda").requires_grad_() for value in reference]
     expected = nnf.l1_loss(*reference)
     result = fused_l1_loss(*actual)
     torch.testing.assert_close(result.cpu(), expected.detach(), equal_nan=True)
@@ -102,8 +93,8 @@ def test_fused_l1_nonfinite_residuals(image_device):
     "case,message",
     [("dimension", "NCHW"), ("shape", "identical image shapes"), ("dtype", "float32")],
 )
-def test_fused_l1_invalid_inputs(image_device, case, message):
-    img1 = torch.empty((1, 3, 17, 19), device=image_device)
+def test_fused_l1_invalid_inputs(case, message):
+    img1 = torch.empty((1, 3, 17, 19), device="cuda")
     img2 = torch.empty_like(img1)
     if case == "dimension":
         img1, img2 = img1[0], img2[0]
@@ -115,39 +106,34 @@ def test_fused_l1_invalid_inputs(image_device, case, message):
         fused_l1_loss(img1, img2)
 
 
-def test_fused_l1_nondefault_streams(image_device):
-    device_ids = range(torch.cuda.device_count()) if image_device == "dgx" else [torch.cuda.current_device()]
-    streams = [torch.cuda.Stream(device=device) for device in device_ids]
+def test_fused_l1_nondefault_streams():
+    stream = torch.cuda.Stream()
     values = [torch.randn(3, 3, 17, 19) for _ in range(2)]
     reference = [value.clone().requires_grad_() for value in values]
     expected = nnf.l1_loss(*reference)
     expected.backward()
 
-    with ExitStack() as stack:
-        for stream in streams:
-            stack.enter_context(torch.cuda.stream(stream))
-        actual = [value.to(image_device).requires_grad_() for value in values]
+    with torch.cuda.stream(stream):
+        actual = [value.to("cuda").requires_grad_() for value in values]
         current_device = torch.cuda.current_device()
         result = fused_l1_loss(*actual)
-        if image_device == "cuda":
-            assert torch.cuda.current_device() == current_device
+        assert torch.cuda.current_device() == current_device
         result.backward()
 
-    for stream in streams:
-        torch.cuda.current_stream(stream.device).wait_stream(stream)
+    torch.cuda.current_stream(stream.device).wait_stream(stream)
     torch.testing.assert_close(result.cpu(), expected.detach(), rtol=2e-6, atol=1e-6)
     for expected_input, actual_input in zip(reference, actual):
         torch.testing.assert_close(actual_input.grad.cpu(), expected_input.grad)
 
 
-def test_fused_l1_with_ssim(image_device):
+def test_fused_l1_with_ssim():
     generator = torch.Generator().manual_seed(41)
     img1 = torch.rand(1, 3, 33, 47, generator=generator)
     img2 = torch.rand(1, 3, 33, 47, generator=generator)
     reference = img1.to("cuda").requires_grad_()
     target = img2.to("cuda")
-    actual = img1.to(image_device).requires_grad_()
-    actual_target = img2.to(image_device)
+    actual = img1.to("cuda").requires_grad_()
+    actual_target = img2.to("cuda")
 
     expected = torch.lerp(nnf.l1_loss(reference, target), 1.0 - ssim(reference, target), 0.2)
     result = torch.lerp(fused_l1_loss(actual, actual_target), 1.0 - ssim(actual, actual_target), 0.2)