[isWangjunjie]

Summary

Opt-in speedups for MedNeXt affinity training. On 4×H200 (BF16, MedNeXt-L, global batch 8) step time goes from 0.600 to 0.228 s/step (~2.6×), so a 200k-step run takes ~13 h instead of ~33 h. Every new knob defaults to off. The only default change is optimization.benchmark: true in tutorials/banis.yaml.

Commit 1: training knobs + shift-sum depthwise conv

• optimization.benchmark / deterministic are now forwarded to the Lightning Trainer. Both fields existed but were never forwarded, so cudnn.benchmark was always off.
• optimizer.fused / merge_param_groups add fused Adam/AdamW and one param group per (lr, wd).
• Gradient clipping for fused optimizers is done directly when there is no GradScaler. Lightning refuses this case even under bf16. A preflight rejects fused + clipping under fp16 precision.
• model.compile / compile_mode compiles the model in place via configure_model. MetaTensor batches are unwrapped only when compiling.
• model.mednext_depthwise_impl: shift_sum (requires model.compile on CUDA) computes stride-1 depthwise Conv3d as shifted multiply-adds fused by inductor, replacing the slow conv_depthwise3d kernels. state_dict keys are unchanged.
• model.mednext_freeze_dummy_tensor lets MedNeXt run DDP with find_unused_parameters=False when it has no deep supervision and no heads. The new optimization.ddp_find_unused_parameters / ddp_static_graph fields replace an undeclared cfg.model read. Default DDP decisions are unchanged.
• Two always-on rewrites, both tested for equivalence against the old code:
  • PerChannelBCEWithLogitsLoss mean reduction is vectorized and makes no host sync. Weight shapes outside the vectorized path use the old loop. If every weight is invalid, the returned zero now stays attached to the graph.
  • EMA updates use foreach for fp32/fp64 groups, bit-identical to before. Low-precision groups keep the scalar path.

Commit 2: numba fast resample for RandAffined

• data.augmentation.affine.fast_resample replaces MONAI Resample with numba kernels for 3D reflection nearest/bilinear. Output equals Resample(dtype=float32) voxel for voxel.
• Unsupported inputs fall back to MONAI: non-CPU tensors, autograd, fp16/bf16 images, and other modes. numba is imported only when the option is on.
• RandAffined was 45% of per-sample loader CPU. It drops from 393 to 150 ms per call.

Benchmark (4×H200, 32 CPU, 30 warmup + 100 timed steps):

• 0.600 s/step: original code, 2 workers/GPU.
• 0.316 s/step: the knobs above with 6 workers/GPU. Of the savings, cudnn.benchmark gives −27%, compile −20%, and fused + merged groups −12%.
• 0.228 s/step: adding shift_sum, fast_resample and 8 workers/GPU.

With the faster GPU path, 2 workers/GPU leave ~20% of each step waiting on data, so use ≥6 workers/GPU.

A full 200k-step BANIS+ reproduction (BF16, 4×H200) with all of these knobs enabled scored test NERL (mt1) 60.04%. The published checkpoint scores 58.65%, and the recipe's reported range is 59.6 ± 1.1%.

Recommended for MedNeXt affinity training:

• model.compile: true
• model.mednext_depthwise_impl: shift_sum
• optimization.optimizer.fused: true
• optimization.optimizer.merge_param_groups: true
• data.augmentation.affine.fast_resample: true
• precision bf16-mixed
• model.mednext_freeze_dummy_tensor: true (without deep supervision or heads)

Test plan

• New tests:
  • test_train_speed_knobs.py
  • test_train_speed_loss.py
  • test_fast_resample.py
  • test_fast_resample_import.py: subprocesses with numba absent or broken.
  • test_ddp_mednext_freeze.py: 2-process gloo DDP.
• CUDA (1×H200): 179 passed, 0 skipped. Covers:
  • compiled shift_sum vs native conv, fp32 and bf16, including grads
  • the EMA dtype matrix
  • sync-free BCE under set_sync_debug_mode("error")
  • compiled fused-AdamW save/resume
• gloo DDP: 4/4 passed.
• Unit regression on loss/optim/ema/trainer/augment/mednext/config: 668 passed, 0 failed.
• Guard tests and validate_tutorial_configs.py: no new failures vs origin/master. Existing failures come from the neuron_j0126 tutorials and missing cloud-volume.

🤖 Generated with Claude Code