Optimize convolution and groupnorm

ScaFFold/unet/unet_parts.py

@@ -20,6 +20,7 @@
 
 from ScaFFold.utils.perf_measure import annotate
 
+from .conv3d import FastConv3d, FastConvTranspose3d
 from .group_norm import FastGroupNorm
 
 _doubleconv_annotate = annotate(fmt="DoubleConv.{}")
@@ -28,30 +29,126 @@
 _outconv_annotate = annotate(fmt="OutConv.{}")
 
 
-def _group_norm(num_groups, num_channels):
+def _group_norm(num_groups, num_channels, activation=None):
     if num_channels % num_groups != 0:
         raise ValueError(
             f"group_norm_groups={num_groups} must evenly divide num_channels={num_channels}"
         )
-    # FastGroupNorm is nn.GroupNorm plus a compiled GPU kernel; it holds the
-    # same parameters under the same names, so checkpoints are unaffected.
-    return FastGroupNorm(num_groups, num_channels)
+    # FastGroupNorm is nn.GroupNorm plus a Triton/compiled GPU kernel; it holds
+    # the same parameters under the same names, and `activation` is a plain
+    # attribute rather than a submodule, so checkpoints are unaffected.
+    return FastGroupNorm(num_groups, num_channels, activation=activation)
+
+
+def _conv3d(in_channels, out_channels, **kwargs):
+    """The model's non-transposed convolutions, in one place.
+
+    ``FastConv3d`` is ``nn.Conv3d`` plus a Triton GPU kernel; it holds the same
+    parameters under the same names and adds no buffers, so checkpoints are
+    unaffected in either direction.  It falls back to MIOpen for anything the
+    kernel does not serve, and it serves the sharded configurations too: it
+    performs the halo exchange itself, above autograd, rather than leaving it to
+    the one DistConv does below.  Only the split axis is halo'd, so ``padding=1``
+    survives on the other two and these convolutions are padded at every
+    configuration; see :mod:`ScaFFold.unet.conv3d`.
+
+    The ``nn.ConvTranspose3d`` in ``Up`` do not come through here: they are a
+    different operator, with the weight's channel axes the other way round and a
+    different set of kernels behind them, so they have a factory of their own
+    (:func:`_conv_transpose3d`).
+    """
+    return FastConv3d(in_channels, out_channels, **kwargs)
+
+
+def _conv_transpose3d(in_channels, out_channels, **kwargs):
+    """The model's transposed convolutions -- the decoder's upsamplers.
+
+    ``FastConvTranspose3d`` is ``nn.ConvTranspose3d`` plus a Triton GPU kernel;
+    it holds the same parameters (``weight`` *and* ``bias``, which these sites
+    have and the ordinary convolutions mostly do not) under the same names and
+    adds no buffers, so checkpoints are unaffected in either direction.  The
+    kernel serves only the ``kernel == stride``, no-padding upsample built
+    below; anything else falls back to MIOpen.
+    """
+    return FastConvTranspose3d(in_channels, out_channels, **kwargs)
+
+
+def _consumer_dtype(*tensors):
+    """The dtype the convolution consuming a concatenation will actually see.
+
+    Inside an enabled autocast region the answer is autocast's dtype, because
+    ``aten::convolution`` carries the ``lower_precision_fp`` cast policy and
+    casts whatever it is handed.  Producing that dtype from the concatenation is
+    *bitwise identical* to producing ATen's promoted dtype and letting the
+    convolution narrow it -- the promoted tensor holds exact widenings of both
+    sources, so narrowing before or after the copy rounds the same values once
+    -- while moving half the bytes.
+
+    Outside autocast the answer is ``torch.cat``'s ordinary promotion, so eval,
+    ``inference_mode`` and pure-fp32 runs are unchanged.
+    """
+    dtype = tensors[0].dtype
+    for tensor in tensors[1:]:
+        dtype = torch.promote_types(dtype, tensor.dtype)
+    device_type = tensors[0].device.type
+    try:
+        if not torch.is_autocast_enabled(device_type):
+            return dtype
+        autocast_dtype = torch.get_autocast_dtype(device_type)
+    except (RuntimeError, TypeError):  # a device type autocast does not know
+        return dtype
+    # Only ever narrow: if autocast's dtype is the wider of the two, keep the
+    # promotion ATen would have done.
+    if torch.promote_types(autocast_dtype, dtype) is autocast_dtype:
+        return dtype
+    return autocast_dtype
+
+
+def _skip_concat(skip, upsampled):
+    """``torch.cat([skip, upsampled], dim=1)`` at the consumer's dtype.
+
+    Under ``torch.autocast`` the two halves do not share a dtype: the skip comes
+    from a GroupNorm, an fp32-policy op, while the upsampled half comes from a
+    ``ConvTranspose3d`` and is bf16.  ``torch.cat`` carries the ``promote``
+    policy, so it widens the bf16 half to fp32, concatenates at fp32, and the
+    following convolution narrows the whole double-width result straight back
+    down -- three full-resolution passes to deliver one.  Casting the inputs
+    first collapses that to one, and the convolution reads the same bits either
+    way (see :func:`_consumer_dtype`).
+    """
+    dtype = _consumer_dtype(skip, upsampled)
+    return torch.cat([skip.to(dtype), upsampled.to(dtype)], dim=1)
 
 
 class DoubleConv(nn.Module):
-    """(convolution => GroupNorm => ReLU) * 2"""
+    """(convolution => GroupNorm => ReLU) * 2
+
+    The ReLU lives *inside* the GroupNorm (``activation="relu"``), because the
+    Triton GroupNorm kernel folds it into its forward store for free and thereby
+    removes an entire streaming pass.  ``FastGroupNorm`` applies the ReLU on
+    every path, including eager, so the network's function is unchanged; only
+    the number of memory passes differs.
+
+    The ``nn.ReLU`` slots are held open by ``nn.Identity`` rather than removed:
+    ``nn.Sequential`` names its children by position, so deleting them would
+    renumber the two convolutions and the second GroupNorm and invalidate every
+    existing checkpoint.  Neither ``nn.ReLU`` nor ``nn.Identity`` has parameters
+    or buffers, so with the placeholders in place the state dict is byte
+    identical to the pre-fusion model's (pinned by
+    ``tests/test_groupnorm.py::test_state_dict_matches_plain_groupnorm_model``).
+    """
 
     def __init__(self, in_channels, out_channels, group_norm_groups, mid_channels=None):
         super().__init__()
         if not mid_channels:
             mid_channels = out_channels
         self.double_conv = nn.Sequential(
-            nn.Conv3d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False),
-            _group_norm(group_norm_groups, mid_channels),
-            nn.ReLU(inplace=True),
-            nn.Conv3d(mid_channels, out_channels, kernel_size=3, padding=1, bias=False),
-            _group_norm(group_norm_groups, out_channels),
-            nn.ReLU(inplace=True),
+            _conv3d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False),
+            _group_norm(group_norm_groups, mid_channels, activation="relu"),
+            nn.Identity(),
+            _conv3d(mid_channels, out_channels, kernel_size=3, padding=1, bias=False),
+            _group_norm(group_norm_groups, out_channels, activation="relu"),
+            nn.Identity(),
         )
 
     @_doubleconv_annotate
@@ -75,7 +172,16 @@ def forward(self, x):
 
 
 class Up(nn.Module):
-    """Upscaling then double conv"""
+    """Upscaling then double conv
+
+    The skip concatenation goes through :func:`_skip_concat` rather than
+    ``torch.cat`` directly, so that it emits the dtype the following
+    convolution will use instead of ``torch.cat``'s promoted one.  The tensor
+    that convolution reads is bitwise unchanged either way; it is written and
+    read back at half the width, saving both time and peak memory.  This rests
+    on ``self.conv`` beginning with a convolution, which the constructor below
+    guarantees on either branch.
+    """
 
     def __init__(self, in_channels, out_channels, group_norm_groups, trilinear=True):
         super().__init__()
@@ -90,7 +196,7 @@ def __init__(self, in_channels, out_channels, group_norm_groups, trilinear=True)
                 in_channels // 2,
             )
         else:
-            self.up = nn.ConvTranspose3d(
+            self.up = _conv_transpose3d(
                 in_channels, in_channels // 2, kernel_size=2, stride=2
             )
             self.conv = DoubleConv(in_channels, out_channels, group_norm_groups)
@@ -118,14 +224,16 @@ def forward(self, x1, x2):
         # if you have padding issues, see
         # https://github.com/HaiyongJiang/U-Net-Pytorch-Unstructured-Buggy/commit/0e854509c2cea854e247a9c615f175f76fbb2e3a
         # https://github.com/xiaopeng-liao/Pytorch-UNet/commit/8ebac70e633bac59fc22bb5195e513d5832fb3bd
-        x = torch.cat([x2, x1], dim=1)
+        # torch.cat([x2, x1], dim=1) at the dtype the convolution below
+        # actually wants; see the class docstring.
+        x = _skip_concat(x2, x1)
         return self.conv(x)
 
 
 class OutConv(nn.Module):
     def __init__(self, in_channels, out_channels):
         super(OutConv, self).__init__()
-        self.conv = nn.Conv3d(in_channels, out_channels, kernel_size=1)
+        self.conv = _conv3d(in_channels, out_channels, kernel_size=1)
 
     @_outconv_annotate
     def forward(self, x):

ScaFFold/utils/trainer.py

@@ -29,6 +29,7 @@
 from torch.utils.data import DataLoader
 from tqdm import tqdm
 
+from ScaFFold.unet._rungs import format_kernel_selection, kernel_selection
 from ScaFFold.utils.checkpointing import CheckpointManager
 from ScaFFold.utils.data_loading import FractalDataset, SpatialShardSpec
 from ScaFFold.utils.data_types import AMP_DTYPE, VOLUME_TORCH_DTYPE
@@ -90,6 +91,9 @@ class BaseTrainer:
 
     def __init__(self, model, config, device, log):
         self.model = model
+        # One-shot guard for the startup kernel-selection line; see
+        # _log_kernel_selection for why it has two call sites.
+        self._kernel_selection_logged = False
         self.config = config
         self.device = device
         self.log = log
@@ -786,6 +790,30 @@ def _warmup_ragged_batches(self, batch):
                 log_prefix=f"warmup ragged ({ragged}): ",
             )
 
+    def _log_kernel_selection(self):
+        """Log which kernel each accelerated module is using, once, on rank 0.
+
+        Called from two places -- the end of :meth:`warmup` and after the first
+        training batch -- because the answer only exists once a forward has run:
+        ``_triton_ok`` is a latch set when a rung first answers a call, so
+        reporting at construction time would say "Native" about modules that
+        have simply not run yet.  ``warmup_batches <= 0`` makes :meth:`warmup`
+        return without running anything, so neither call site alone covers every
+        run; the flag below makes the pair idempotent.
+
+        Rank 0 only, and one rank's answer: each rank latches independently, so
+        under DDP this is representative rather than global.  Deliberately not a
+        collective -- gathering it would put a barrier on a path that has no
+        other reason for one.
+        """
+        if self._kernel_selection_logged or self.world_rank != 0:
+            return
+        self._kernel_selection_logged = True
+        model = getattr(self.model, "module", self.model)
+        self.log.info("Kernel selection (rank 0):")
+        for line in format_kernel_selection(kernel_selection(model)):
+            self.log.info(line)
+
     def warmup(self):
         """Run warmup iterations before the main training loop."""
         warmup_batches = self.config.warmup_batches
@@ -847,6 +875,7 @@ def warmup(self):
 
         torch.distributed.barrier()
         self.log.info(f"Done warmup. Took {int(time.time() - start_warmup)}s")
+        self._log_kernel_selection()
 
     def train(self, profiler=None):
         """
@@ -953,6 +982,8 @@ def train(self, profiler=None):
                         # not skew the epoch mean.
                         train_dice_total += batch_dice_score * batch_size
                         end_code_region("run_training_batch")
+                        if first_batch:
+                            self._log_kernel_selection()
 
                         # Update the loss
                         begin_code_region("update_loss")

ScaFFold/viz/standard_viz.py

@@ -49,7 +49,9 @@ def main(config: RunConfig):
                 if "val_loss_avg" in row:
                     val_loss.append(float(row["val_loss_avg"]))
 
-        plot_title = f"v={config.vol_size}, c={config.n_categories}, u={config.unet_layers}"
+        plot_title = (
+            f"v={config.vol_size}, c={config.n_categories}, u={config.unet_layers}"
+        )
         line_thickness = 2
         fontsize = 20
         tick_fontsize = 14
@@ -64,7 +66,9 @@ def main(config: RunConfig):
         plt.tick_params(axis="both", which="major", labelsize=tick_fontsize)
         plt.yscale("log")
         plt.title(plot_title, fontsize=12)
-        plt.legend(loc="upper left", bbox_to_anchor=legend_loc, fontsize=legend_fontsize)
+        plt.legend(
+            loc="upper left", bbox_to_anchor=legend_loc, fontsize=legend_fontsize
+        )
         plt.grid(True, axis="y")
         plt.savefig(figures_path / "train_loss.png", dpi=300, bbox_inches="tight")
         plt.close(figures[-1])
@@ -76,7 +80,9 @@ def main(config: RunConfig):
         plt.ylabel("Val dice score", fontsize=fontsize)
         plt.tick_params(axis="both", which="major", labelsize=tick_fontsize)
         plt.title(plot_title, fontsize=12)
-        plt.legend(loc="upper left", bbox_to_anchor=legend_loc, fontsize=legend_fontsize)
+        plt.legend(
+            loc="upper left", bbox_to_anchor=legend_loc, fontsize=legend_fontsize
+        )
         plt.grid(True, axis="y")
         plt.savefig(figures_path / "val_dice.png", dpi=300, bbox_inches="tight")
         plt.close(figures[-1])
@@ -89,7 +95,9 @@ def main(config: RunConfig):
             plt.ylabel("Val loss", fontsize=fontsize)
             plt.tick_params(axis="both", which="major", labelsize=tick_fontsize)
             plt.title(plot_title, fontsize=12)
-            plt.legend(loc="upper left", bbox_to_anchor=legend_loc, fontsize=legend_fontsize)
+            plt.legend(
+                loc="upper left", bbox_to_anchor=legend_loc, fontsize=legend_fontsize
+            )
             plt.grid(True, axis="y")
             plt.savefig(figures_path / "val_loss.png", dpi=300, bbox_inches="tight")
             plt.close(figures[-1])