Skip to content

fix(patching): keep overlap-blend reassembly in the patch dtype (halve peak memory) - #27

Merged
vboussot merged 1 commit into
mainfrom
fix/accumulator-blend-dtype
Jul 6, 2026
Merged

fix(patching): keep overlap-blend reassembly in the patch dtype (halve peak memory)#27
vboussot merged 1 commit into
mainfrom
fix/accumulator-blend-dtype

Conversation

@vboussot

@vboussot vboussot commented Jul 5, 2026

Copy link
Copy Markdown
Member

Summary

Accumulator.assemble (overlap-blend reassembly) allocated weight_sum as a default float32
tensor. The final normalization result = result / weight_sum then promoted the whole
(channels × volume) accumulator to float32 even when the patches were float16 (as they are for
autocast/fp16 inference, where ModelComposite casts outputs to fp16).

For large multi-class reassemblies this silently doubled the peak memory of the assembled volume.
Concretely, TotalSegmentator's 118-class whole-body output at inference has an assembled tensor of
shape (118, D, H, W) — for total-3mm that is ~3.1e9 elements = 12.5 GB in fp32 vs 6.3 GB in fp16;
for the full total at 1.5 mm it is ~8× larger and was OOM-ing.

The fix

Allocate weight_sum with dtype=result.dtype, so result / weight_sum stays in the patch dtype and
the accumulator is never promoted to float32.

Effect (measured on TotalSegmentator-KonfAI:total-3mm, autocast fp16)

  • Assembled volume dtype: float32 → float16; peak of that tensor halved (12.5 GB → 6.3 GB).
  • Reassembly time: 17.65 s → 14.36 s.
  • Output: 27 differing voxels / 15,068,484 (0.0000018%) vs the prior float32 path — fp16 rounding at
    class-boundary ties. The pipeline is already fp16 end-to-end (composite casts to fp16 before
    accumulation), so the float32 division was a spurious upcast, not a source of truth.

Test plan

  • New test_blended_reassembly_preserves_patch_dtype: a multi-channel float16 Cosinus-blended
    reassembly must stay float16.
  • Existing blend/partition-of-unity tests (test_patch_overlap_border, test_patching) and
    test_predictor_memory still pass (17 passed).

Independent of #26 (device placement); together they let TotalSegmentator run without crashing or OOM.

Accumulator.assemble allocated weight_sum as a default float32 tensor, so the final
result / weight_sum normalization promoted the whole (channels x volume) accumulator to
float32 even when the patches were float16. For large multi-class reassemblies -- e.g. a
118-class whole-body segmentation -- this silently doubled the peak memory of the assembled
volume and could OOM. Allocate weight_sum in the result dtype so the division stays in the
patch dtype. Adds a dtype-preservation regression test.
@vboussot
vboussot merged commit 0441032 into main Jul 6, 2026
29 checks passed
@vboussot
vboussot deleted the fix/accumulator-blend-dtype branch July 6, 2026 08:43
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant