pytorch - 💡(How to fix) Fix native_batch_norm_backward extremely slow on MPS for 3D tensors (slower than CPU) [6 comments, 3 participants]

import time
import torch
import torch.nn as nn

assert torch.backends.mps.is_available(), "MPS not available"

# Minimal 3D model with BatchNorm3d
class Simple3DCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv3d(1, 32, 3, padding=1),
            nn.BatchNorm3d(32),
            nn.ReLU(),
            nn.Conv3d(32, 64, 3, padding=1),
            nn.BatchNorm3d(64),
            nn.ReLU(),
            nn.AdaptiveAvgPool3d(1),
            nn.Flatten(),
            nn.Linear(64, 4),
        )
    def forward(self, x):
        return self.net(x)

def bench(device_name, steps=20):
    device = torch.device(device_name)
    model = Simple3DCNN().to(device)
    optimizer = torch.optim.Adam(model.parameters())
    criterion = nn.CrossEntropyLoss()
    x = torch.randn(4, 1, 64, 64, 64, device=device)
    y = torch.randint(0, 4, (4,), device=device)

    # Warmup
    for _ in range(3):
        loss = criterion(model(x), y)
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
    if device_name == "mps":
        torch.mps.synchronize()

    start = time.perf_counter()
    for _ in range(steps):
        loss = criterion(model(x), y)
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
    if device_name == "mps":
        torch.mps.synchronize()
    elapsed = time.perf_counter() - start
    print(f"{device_name.upper():>4s}: {elapsed:.2f}s ({4*steps/elapsed:.1f} vol/s)")

bench("cpu")
bench("mps")

Expected behavior

MPS should be faster than (or at least comparable to) CPU for 3D batch norm backward, similar to how 2D batch norm achieves excellent MPS speedups.

Environment

- PyTorch: 2.10.0 (also tested 2.7.1 — same behavior)
- OS: macOS 26.1 / 26.2
- Hardware: MacBook Pro M1 Pro (32 GB) and Mac mini M4 Pro (64 GB) — same issue on both
- Python: 3.12

cc: #154052 (context from the "most requested MPS ops" discussion)

---## 🐛 Describe the bug

3D CNN training on MPS is **slower than CPU** due to `native_batch_norm_backward` consuming ~70% of training time for 5D tensors (batch, channels, D, H, W).
2D training on MPS works excellently (9–12× speedup), so this appears specific to the 3D batch norm backward kernel.

### Profiling result

The bottleneck is `aten::native_batch_norm_backward` for 3D (5D) tensors on MPS. Forward pass and all other 3D ops run at reasonable speed.

### Benchmark results

| Test | Device | Time (s) | Throughput | Speedup |
|------|--------|-----------|------------|---------|
| 2D ResNet18 | CPU | 45.7 | 21 img/s | — |
| 2D ResNet18 | MPS | 5.1 | 188 img/s | **8.9×** |
| 3D ResNet18 | CPU | 119.8 | 1.0 vol/s | — |
| 3D ResNet18 | MPS | 154.9 | 0.8 vol/s | **0.77×** (slower!) |

Tested on both PyTorch 2.7.1 and 2.10.0 — no improvement in 3D MPS performance.

## Reproducer

```python
import time
import torch
import torch.nn as nn

assert torch.backends.mps.is_available(), "MPS not available"

# Minimal 3D model with BatchNorm3d
class Simple3DCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv3d(1, 32, 3, padding=1),
            nn.BatchNorm3d(32),
            nn.ReLU(),
            nn.Conv3d(32, 64, 3, padding=1),
            nn.BatchNorm3d(64),
            nn.ReLU(),
            nn.AdaptiveAvgPool3d(1),
            nn.Flatten(),
            nn.Linear(64, 4),
        )
    def forward(self, x):
        return self.net(x)

def bench(device_name, steps=20):
    device = torch.device(device_name)
    model = Simple3DCNN().to(device)
    optimizer = torch.optim.Adam(model.parameters())
    criterion = nn.CrossEntropyLoss()
    x = torch.randn(4, 1, 64, 64, 64, device=device)
    y = torch.randint(0, 4, (4,), device=device)

    # Warmup
    for _ in range(3):
        loss = criterion(model(x), y)
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
    if device_name == "mps":
        torch.mps.synchronize()

    start = time.perf_counter()
    for _ in range(steps):
        loss = criterion(model(x), y)
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
    if device_name == "mps":
        torch.mps.synchronize()
    elapsed = time.perf_counter() - start
    print(f"{device_name.upper():>4s}: {elapsed:.2f}s ({4*steps/elapsed:.1f} vol/s)")

bench("cpu")
bench("mps")

Expected behavior

MPS should be faster than (or at least comparable to) CPU for 3D batch norm backward, similar to how 2D batch norm achieves excellent MPS speedups.

Environment

- PyTorch: 2.10.0 (also tested 2.7.1 — same behavior)
- OS: macOS 26.1 / 26.2
- Hardware: MacBook Pro M1 Pro (32 GB) and Mac mini M4 Pro (64 GB) — same issue on both
- Python: 3.12

cc: #154052 (context from the "most requested MPS ops" discussion)

---

### Versions

native_batch_norm_backward extremely slow on MPS for 3D tensors (slower than CPU)

cc @jerryzh168 @kulinseth @malfet @DenisVieriu97 @jhavukainen @aditvenk