Undocumented sampling API pitfalls (p_sample_loop_known vs ddim_sample_loop_known), EMA resume filename mismatch, unconditional DDP wrap, and fp16 dtype bug in highway branch

[Armstrong66]

Summary

We used MedSegDiff-V2 for a few-shot stroke lesion segmentation study on a real clinical NCCT cohort. During integration, we hit several issues around the conditional sampling API, checkpoint resume, and DDP/AMP handling that cost significant debugging time and aren't currently documented anywhere in the repo or README. Filing this as one issue covering all of them since they're related (mostly inference-time), happy to split into separate issues/PRs per maintainer preference.

Confidence is marked per item: Confirmed = verified by direct source inspection against our cloned copy of this repo. Inferred = strongly supported by crash behavior/stack traces but the exact internal mechanism wasn't independently confirmed (we didn't have full visibility into gaussian_diffusion.py's ddim_sample_loop_known implementation).

---

1. p_sample_loop_known and ddim_sample_loop_known have incompatible signatures — undocumented (Confirmed)

These are treated interchangeably in places (e.g. toggled via a single --use_ddim flag in scripts/segmentation_sample.py), but they are not drop-in replacements:

# Non-DDIM path
sample, x_noisy, org, cal, cal_out = diffusion.p_sample_loop_known(
    model, shape, img, step=diffusion_steps, clip_denoised=True,
    model_kwargs=model_kwargs,
)
# 5-tuple return. `step=` is required.

# DDIM path
sample, x_noisy, org = diffusion.ddim_sample_loop_known(
    model, shape, img, clip_denoised=True, model_kwargs=model_kwargs,
)
# 3-tuple return. Passing `step=` raises TypeError: got an unexpected
# keyword argument 'step'. cal/cal_out are never computed on this path.

Anyone writing custom inference code (e.g. for K-pass uncertainty sampling, batch evaluation scripts, etc. — a very natural thing to do given the model's native stochasticity) who assumes a shared signature between the two, or assumes cal_out is always available, will get either an immediate TypeError or a silent shape/None issue depending on which branch they hit first.

Suggested fix: Document this explicitly in the README/docstrings, or better, give both functions a consistent signature (accept and no-op on step= in the DDIM path, or always return a 5-tuple with cal=None, cal_out=None on DDIM) so calling code doesn't need to branch on use_ddim to know how many values to unpack.

---

2. timestep_respacing + *_known sampling functions → CUDA device-side assert (index out of bounds) (Inferred root cause, Confirmed reproduction)

Repro: Build a SpacedDiffusion with any non-empty timestep_respacing (e.g. "ddim20"), then call ddim_sample_loop_known (or p_sample_loop_known) on it:

diffusion = create_gaussian_diffusion(..., timestep_respacing="ddim20")
sample, x_noisy, org = diffusion.ddim_sample_loop_known(model, shape, img, ...)

This reliably crashes with:
../aten/src/ATen/native/cuda/IndexKernel.cu:111: ... Assertion
-sizes[i] <= index && index < sizes[i] && "index out of bounds" failed.
...
torch.AcceleratorError: CUDA error: device-side assert triggered
(async CUDA error reporting means the visible stack trace point — for us, inside timestep_embedding/[unet.py](http://unet.py/) forward — is downstream of the actual fault, not the fault itself.)

Setting timestep_respacing="" (no respacing) avoids the crash entirely, all else equal. This strongly suggests the *_known functions reference an internal timestep value assuming the full, un-respaced schedule length—e.g., a hardcoded intermediate timestep—which then gets indexed against [respace.py](http://respace.py/)'s _WrappedModel's self.timestep_map (built with only as many entries as the respaced schedule specifies). If that internal timestep exceeds len(timestep_map), map_tensor[ts] in _WrappedModel.__call__ is exactly the kind of out-of-bounds index this assert describes.

We were not able to fully confirm this by reading ddim_sample_loop_known's implementation directly (couldn't get clean access to that specific source), so I'm flagging the mechanism as inferred, not confirmed—but the reproduction itself (crashes with respacing, works without it, same model/input otherwise) is solid.

Suggested fix: Either make the *_known functions respacing-aware (derive their internal timestep(s) relative to the actual schedule length in use, not a hardcoded full-schedule value), or explicitly document that timestep_respacing is incompatible with p_sample_loop_known/ddim_sample_loop_known and that speed should instead be controlled via the step= parameter on the non-DDIM path.

---

3. find_ema_checkpoint() filename mismatch—EMA weights silently never resumed (Confirmed)

In guided_diffusion/train_util.py:

def find_ema_checkpoint(main_checkpoint, step, rate):
    ...
    filename = f"ema_{rate}_{(step):06d}.pt"

But save() in the same file actually writes:

filename = f"emasavedmodel_{rate}_{(step+resume_step):06d}.pt"

These never match. find_ema_checkpoint() always returns None on an otherwise-valid checkpoint directory, so _load_ema_parameters() silently falls back to re-initializing the EMA trajectory from the main model's current weights on every resume, instead of continuing the existing EMA average. No error or warning is raised—this is a silent correctness bug affecting anyone who resumes training and relies on EMA weights for final evaluation (which the repo's own inference script prefers for final results in most use cases).

Suggested fix: one line—align the filename pattern in find_ema_checkpoint() to match save()'s actual output.

---

4. Model is unconditionally wrapped in DistributedDataParallel regardless of world_size (Confirmed)

In TrainLoop.__init__:

if th.cuda.is_available():
    self.use_ddp = True
    self.ddp_model = DDP(
        self.model, device_ids=[dist_util.dev()], output_device=dist_util.dev(),
        broadcast_buffers=False, bucket_cap_mb=128, find_unused_parameters=True,
    )

This is gated purely on CUDA availability, not on dist.get_world_size() > 1. For any single-GPU-per-process training setup (e.g. a grid of independent single-GPU jobs each launched via CUDA_VISIBLE_DEVICES, never through torchrun/multi-process spawn—a very common pattern for hyperparameter/data-ablation sweeps like ours), this pays full NCCL process-group init, gradient-bucket construction, and (with find_unused_parameters=True) a complete extra autograd graph traversal every step—all for a "distributed" group of exactly one member, with zero communication benefit. In our measurements, this was a non-trivial, avoidable per-step overhead.

Suggested fix:

if th.cuda.is_available() and dist.get_world_size() > 1:
    self.use_ddp = True
    self.ddp_model = DDP(...)
else:
    self.use_ddp = False
    self.ddp_model = self.model

dist.get_rank()/dist.get_world_size()/dist.barrier() calls elsewhere in TrainLoop continue to work fine without this—they depend on the process group from dist_util.setup_dist(), not on the DDP wrapper itself, so this change is fully self-contained.

---

5. highway_forward's ad-hoc Conv2d breaks under mixed precision (Confirmed)

Inside the highway/anchor branch ([unet.py](http://unet.py/), UNetModel_newpreview's forward path), a Conv2d layer is constructed inside the forward call itself, not in __init__:

emb = conv_nd(2, x.size(1), 512, 1).to(device=x.device)(x)

Because this is instantiated fresh every forward pass, it is never a registered submodule, so any model.half() / mixed-precision weight-casting pass that walks registered submodules never touches it. It's always constructed at default float32. When x is float16 (e.g. under use_fp16=True training), this raises:
RuntimeError: Input type (c10::Half) and bias type (float) should be the same

Suggested fix: .to(device=x.device, dtype=x.dtype)(x) — one-line, makes the ad-hoc layer always match the input's current precision regardless of fp16/fp32 mode.

---

Environment

• MedSegDiff-V2 (version="new" / UNetModel_newpreview), cloned from ImprintLab/MedSegDiff
• PyTorch 2.x, single-GPU and dual-GPU (RTX 8000) configurations
• Custom 2.5D (3-slice triplet) conditioning input, in_ch=4

Happy to submit PRs for items 3–5 (small, self-contained, we already have working patches). Items 1–2 likely need maintainer input on intended design (e.g. whether step= should be added to the DDIM path, or whether respacing support is out of scope for the *_known functions) before a PR makes sense.

Fix EMA checkpoint filename mismatch on resume

File: guided_diffusion/train_util.py

--- a/guided_diffusion/train_util.py
+++ b/guided_diffusion/train_util.py
@@ -XXX,7 +XXX,10 @@ def find_ema_checkpoint(main_checkpoint, step, rate):
     if main_checkpoint is None:
         return None
-    filename = f"ema_{rate}_{(step):06d}.pt"
+    # Must match the filename actually written by TrainLoop.save() below
+    # (`emasavedmodel_{rate}_{step:06d}.pt`), not `ema_{rate}_{step:06d}.pt`.
+    # The mismatched pattern previously meant this always returned None on
+    # an otherwise-valid checkpoint dir, silently reinitializing the EMA
+    # trajectory from the main model's current weights on every resume
+    # instead of continuing the saved EMA average.
+    filename = f"emasavedmodel_{rate}_{(step):06d}.pt"
     path = bf.join(bf.dirname(main_checkpoint), filename)
     if bf.exists(path):
         return path
     return None

description:

find_ema_checkpoint() searches for ema_{rate}{step}.pt, but TrainLoop.save() writes emasavedmodel{rate}_{step}.pt. The mismatch means EMA checkpoints are never found on resume — training silently falls back to reinitializing the EMA average from the main model's current weights rather than continuing the saved trajectory, with no warning or error. One-line fix aligning the search pattern to the actual save pattern. Verified by reproducing the bug (resume + confirm find_ema_checkpoint returns None despite a valid EMA file existing at that path) and confirming the fix resolves it.

Skip DDP wrap when world_size == 1

File: guided_diffusion/train_util.py

--- a/guided_diffusion/train_util.py
+++ b/guided_diffusion/train_util.py
@@ -XXX,17 +XXX,26 @@ class TrainLoop:
-        if th.cuda.is_available():
+        # Only wrap in DDP when there is actually more than one process to
+        # synchronize with. `dist.get_world_size()` is 1 for any single-
+        # process-per-GPU launch pattern (e.g. CUDA_VISIBLE_DEVICES-scoped
+        # subprocesses run independently, never spawned via torchrun).
+        # Wrapping unconditionally on CUDA availability alone pays full
+        # NCCL process-group init, gradient-bucket construction, and (with
+        # find_unused_parameters=True) a complete extra autograd graph
+        # traversal every step, for a "distributed" group of exactly one
+        # member with no communication benefit. dist.get_rank()/
+        # get_world_size()/barrier() calls elsewhere in this class continue
+        # to work correctly without the DDP wrapper -- they depend on the
+        # process group from dist_util.setup_dist(), not on DDP itself.
+        if th.cuda.is_available() and dist.get_world_size() > 1:
             self.use_ddp = True
             self.ddp_model = DDP(
                 self.model,
                 device_ids=[dist_util.dev()],
                 output_device=dist_util.dev(),
                 broadcast_buffers=False,
                 bucket_cap_mb=128,
                 find_unused_parameters=True,
             )
         else:
-            if dist.get_world_size() > 1:
+            if dist.get_world_size() > 1:
+                # unreachable under CUDA now that the condition above
+                # already routes world_size>1 into the DDP branch; kept
+                # for the (non-CUDA) distributed-without-CUDA warning path
                 logger.warn(
                     "Distributed training requires CUDA. "
                     "Gradients will not be synchronized properly!"
                 )
             self.use_ddp = False
             self.ddp_model = self.model

description:

The model is currently wrapped in DistributedDataParallel whenever CUDA is available, regardless of world_size. For single-GPU-per-process training setups (e.g. independent CUDA_VISIBLE_DEVICES-scoped jobs run for a hyperparameter/data sweep, never launched via torchrun), this pays full DDP overhead — NCCL init, bucket construction, and (with find_unused_parameters=True) an extra full autograd graph traversal every step — for a synchronization group of one, with zero benefit. Gating on dist.get_world_size() > 1 fixes this; all other dist.* usage elsewhere in TrainLoop is unaffected since it depends on the process group from dist_util.setup_dist(), not on the DDP wrapper. Verified via profiling: measurable per-step overhead reduction with no change in training behavior/loss trajectory in a single-GPU run.

Fix fp16 dtype mismatch in highway branch's ad-hoc Conv2d

File: guided_diffusion/unet.py

--- a/guided_diffusion/unet.py
+++ b/guided_diffusion/unet.py
@@ -XXX,7 +XXX,12 @@ class UNetModel_newpreview(nn.Module):
     def highway_forward(self, x, hs=None):
-        emb = conv_nd(2, x.size(1), 512, 1).to(device=x.device)(x)
+        # This Conv2d is constructed fresh on every forward call rather than
+        # in __init__, so it is never a registered submodule -- any
+        # model.half() / mixed-precision weight-casting pass that walks
+        # registered submodules never reaches it, and it is always built at
+        # default float32. Under autocast/fp16 training, x here can be
+        # float16, causing:
+        #   RuntimeError: Input type (c10::Half) and bias type (float)
+        #   should be the same
+        # Matching the layer's dtype to the input's current dtype fixes
+        # this regardless of fp16/fp32 mode.
+        emb = conv_nd(2, x.size(1), 512, 1).to(device=x.device, dtype=x.dtype)(x)
         ...

description:

highway_forward constructs a Conv2d layer inline on every forward call instead of registering it in init. Because it's never a registered submodule, mixed-precision casting passes never touch it, so it's always built at float32. Under fp16/autocast training, this raises RuntimeError: Input type (c10::Half) and bias type (float) should be the same as soon as x is float16. Adding dtype=x.dtype to the existing .to() call is sufficient — the layer now always matches whatever precision the rest of the forward pass is running at. Verified: reproduced the crash under use_fp16=True, confirmed the one-line fix resolves it with no change to fp32 training behavior.