From 4d4f32577f6cf92e82194f5010a7c7793c5765a1 Mon Sep 17 00:00:00 2001 From: cedriclim1 Date: Tue, 9 Jun 2026 23:00:20 -0700 Subject: [PATCH 1/2] Build INR training batches on the compute device for single-process runs The per-pixel DataLoader path costs ~43 ms/batch on CPU (8192 Python __getitem__ calls + collate + H2D copy) while the GPU step itself takes ~20 ms, so single-GPU reconstructions idle the GPU half the time. DeviceBatchSampler keeps the tilt stack resident on the device and builds each batch with index arithmetic and two tensor lookups, yielding the same batch dicts (and drop_last/val-split semantics) as the DataLoader path. DDP runs keep the DataLoader + DistributedSampler path. Also disable bf16 autocast in the validation loop to match the training pass: the so3 pose solve (lu_factor) has no BFloat16 kernel, and a bf16 val loss is not comparable to the fp32 train loss. --- src/quantem/tomography/dataset_models.py | 60 ++++++++++ src/quantem/tomography/tomography.py | 48 ++++++-- tests/tomography/test_device_batch_sampler.py | 104 ++++++++++++++++++ 3 files changed, 203 insertions(+), 9 deletions(-) create mode 100644 tests/tomography/test_device_batch_sampler.py diff --git a/src/quantem/tomography/dataset_models.py b/src/quantem/tomography/dataset_models.py index 2255636c..00419695 100644 --- a/src/quantem/tomography/dataset_models.py +++ b/src/quantem/tomography/dataset_models.py @@ -440,6 +440,66 @@ def to(self, device: str | torch.device): self.device = device +class DeviceBatchSampler: + """Epoch iterator that builds INR training batches directly on a device. + + Replaces the per-pixel DataLoader path for single-process runs: the tilt + stack and angles are made resident on ``device`` once, so producing a + batch is index arithmetic plus two tensor lookups instead of + ``batch_size`` Python ``__getitem__`` calls, a collate, and a + host-to-device copy per step. On a GPU this removes the CPU dataloader + bottleneck entirely. + + Yields the same batch dicts as ``TomographyINRDataset.__getitem__`` + under a DataLoader collate (``projection_idx``, ``pixel_i``, + ``pixel_j``, ``phi``, ``target_value``), with the train loader's + ``drop_last=True`` semantics. + """ + + def __init__( + self, + dset: "TomographyINRDataset", + batch_size: int, + device: torch.device | str, + indices: torch.Tensor | None = None, + shuffle: bool = True, + ): + self.batch_size = batch_size + self.device = torch.device(device) + self.shuffle = shuffle + self._stack = dset.tilt_stack.to(self.device) + self._angles = dset.tilt_angles.to(self.device) + # __getitem__ decodes flat indices with shape[1] for both rows and + # columns; replicate it exactly. + self._s1 = dset.tilt_stack.shape[1] + self._s2 = dset.tilt_stack.shape[2] + if indices is None: + indices = torch.arange(len(dset), dtype=torch.int64) + self._indices = indices.to(self.device) + + def __len__(self) -> int: + return len(self._indices) // self.batch_size # drop_last=True + + def __iter__(self): + idx = self._indices + if self.shuffle: + idx = idx[torch.randperm(len(idx), device=self.device)] + per_proj = self._s1 * self._s2 + for k in range(len(self)): + sel = idx[k * self.batch_size : (k + 1) * self.batch_size] + proj = sel // per_proj + rem = sel - proj * per_proj + pixel_i = rem // self._s1 + pixel_j = rem - pixel_i * self._s1 + yield { + "projection_idx": proj, + "pixel_i": pixel_i, + "pixel_j": pixel_j, + "phi": self._angles[proj], + "target_value": self._stack[proj, pixel_i, pixel_j], + } + + class TomographyINRDataset(TomographyDatasetConstraints, Dataset): """ Dataset class for INR-based tomography. diff --git a/src/quantem/tomography/tomography.py b/src/quantem/tomography/tomography.py index f86095ad..2fdbb7cf 100644 --- a/src/quantem/tomography/tomography.py +++ b/src/quantem/tomography/tomography.py @@ -17,6 +17,7 @@ DatasetConstraintParams, DatasetConstraintsType, DatasetModelType, + DeviceBatchSampler, TomographyINRDataset, TomographyPixDataset, ) @@ -145,14 +146,7 @@ def reconstruct( self.scheduler_params = scheduler_params self.set_schedulers(self.scheduler_params, num_iter=num_iter) - self.dataloader, self.sampler, self.val_dataloader, self.val_sampler = ( - self.setup_dataloader( - self.dset, - batch_size, - num_workers=num_workers, - val_fraction=val_fraction, - ) - ) + self._setup_recon_dataloaders(batch_size, num_workers, val_fraction) # Type check for INR-based reconstruction if not isinstance(self.dset, TomographyINRDataset): @@ -285,10 +279,14 @@ def reconstruct( val_loss = torch.tensor(0.0, device=self.device) for batch in self.val_dataloader: + # Match the training pass (enabled=False): bf16 autocast + # breaks the so3 pose solve (lu_factor has no BFloat16 + # kernel) and would make the val loss inconsistent with + # the fp32 training loss it is compared to. with torch.autocast( device_type=self.device.type, dtype=torch.bfloat16, - enabled=True, + enabled=False, ): all_coords = self.dset.get_coords(batch, N, curr_num_samples_per_ray) @@ -411,6 +409,38 @@ def _rebuild_dataloader(self, batch_size: int, num_workers: int, val_fraction: f """ Rebuilds the dataloader due to persistent workers error when reloading the object. """ + self._setup_recon_dataloaders(batch_size, num_workers, val_fraction) + + def _setup_recon_dataloaders(self, batch_size: int, num_workers: int, val_fraction: float): + """Build the train/val batch iterators. + + Single-process runs use ``DeviceBatchSampler`` — batches are built + with tensor ops on the compute device from a device-resident tilt + stack, removing the per-pixel ``__getitem__`` / collate / H2D-copy + dataloader bottleneck (``num_workers`` is ignored on this path). + Multi-process (DDP) runs keep the DataLoader + DistributedSampler + path. + """ + if self.world_size == 1 and isinstance(self.dset, TomographyINRDataset): + n = len(self.dset) + n_val = int(n * val_fraction) + perm = torch.randperm(n) + self.dataloader = DeviceBatchSampler( + self.dset, batch_size, self.device, indices=perm[n_val:] + ) + # The val sampler keeps its own device-resident copy of the tilt + # stack; acceptable, since val_fraction > 0 is the rare case. + self.val_dataloader = ( + DeviceBatchSampler( + self.dset, batch_size, self.device, indices=perm[:n_val], shuffle=False + ) + if n_val > 0 + else None + ) + self.sampler = None + self.val_sampler = None + return + self.dataloader, self.sampler, self.val_dataloader, self.val_sampler = ( self.setup_dataloader( self.dset, diff --git a/tests/tomography/test_device_batch_sampler.py b/tests/tomography/test_device_batch_sampler.py new file mode 100644 index 00000000..e40b1284 --- /dev/null +++ b/tests/tomography/test_device_batch_sampler.py @@ -0,0 +1,104 @@ +"""Tests for ``DeviceBatchSampler`` and its wiring in ``Tomography.reconstruct``. + +The sampler must yield batches identical in content to the per-pixel +DataLoader path (same keys, same index decode as +``TomographyINRDataset.__getitem__``), cover each pixel exactly once per +epoch (minus the dropped tail batch), and respect the train/val index +split. Device-independent, so everything here runs on CPU. +""" + +import numpy as np +import torch + +from quantem.tomography.dataset_models import DeviceBatchSampler, TomographyINRDataset + + +def _dset(n_proj=4, n=10, seed=0): + rng = np.random.default_rng(seed) + stack = rng.random((n_proj, n, n)).astype(np.float32) + angles = np.linspace(-60, 60, n_proj).astype(np.float32) + return TomographyINRDataset.from_data(tilt_stack=stack, tilt_angles=angles) + + +def test_batches_match_getitem(): + dset = _dset() + sampler = DeviceBatchSampler(dset, batch_size=37, device="cpu", shuffle=False) + seen = 0 + for batch in sampler: + for k in range(len(batch["target_value"])): + item = dset[seen + k] + for key in ("projection_idx", "pixel_i", "pixel_j", "phi", "target_value"): + torch.testing.assert_close( + batch[key][k], item[key].to(batch[key].dtype), rtol=0, atol=0 + ) + seen += len(batch["target_value"]) + + +def test_epoch_covers_indices_once_with_drop_last(): + dset = _dset() + n = len(dset) + batch_size = 64 + sampler = DeviceBatchSampler(dset, batch_size=batch_size, device="cpu", shuffle=True) + assert len(sampler) == n // batch_size + per_proj = dset.tilt_stack.shape[1] * dset.tilt_stack.shape[2] + flat = [] + for batch in sampler: + assert len(batch["target_value"]) == batch_size + flat.append( + batch["projection_idx"] * per_proj + + batch["pixel_i"] * dset.tilt_stack.shape[1] + + batch["pixel_j"] + ) + flat = torch.cat(flat) + assert flat.unique().numel() == flat.numel() # no repeats within an epoch + assert flat.numel() == len(sampler) * batch_size + + +def test_shuffle_changes_order_between_epochs(): + dset = _dset() + sampler = DeviceBatchSampler(dset, batch_size=50, device="cpu", shuffle=True) + first = next(iter(sampler))["target_value"] + second = next(iter(sampler))["target_value"] + assert not torch.equal(first, second) + + +def test_val_split_is_disjoint(): + dset = _dset() + n = len(dset) + perm = torch.randperm(n) + n_val = n // 10 + train = DeviceBatchSampler(dset, 32, "cpu", indices=perm[n_val:]) + val = DeviceBatchSampler(dset, 32, "cpu", indices=perm[:n_val], shuffle=False) + assert len(train._indices) + len(val._indices) == n + assert torch.cat([train._indices, val._indices]).unique().numel() == n + + +def test_reconstruct_uses_sampler_single_process(): + """Smoke: the single-process reconstruct path builds DeviceBatchSamplers. + + Object models that go through ``setup_distributed`` must be built on CUDA + when a CUDA device exists, so pick the device accordingly. + """ + from quantem.core.ml.models.kplanes import KPlanesTILTED + from quantem.core.ml.optimizer_mixin import OptimizerParams + from quantem.tomography.object_models import ObjectINR + from quantem.tomography.tomography import Tomography + + device = "cuda:0" if torch.cuda.is_available() else "cpu" + dset = _dset(n_proj=3, n=8) + model = KPlanesTILTED(M_features=2, resolution=(8, 8, 8), multiscale_res_multipliers=[1], T=1) + obj = ObjectINR.from_model(model, shape=(8, 8, 8), device=device) + tomo = Tomography.from_models(dset=dset, obj_model=obj, device=device, verbose=False) + tomo.reconstruct( + num_iter=1, + batch_size=32, + num_workers=0, + val_fraction=0.25, + optimizer_params={ + "object": {"default": OptimizerParams.Adam(lr=1e-3)}, + "pose": OptimizerParams.Adam(lr=1e-3), + }, + ) + assert isinstance(tomo.dataloader, DeviceBatchSampler) + assert isinstance(tomo.val_dataloader, DeviceBatchSampler) + assert tomo.sampler is None From 9e7e505a30f45850d6117b231ac580e444e3bef0 Mon Sep 17 00:00:00 2001 From: cedriclim1 Date: Tue, 9 Jun 2026 23:30:59 -0700 Subject: [PATCH 2/2] Shard DeviceBatchSampler across DDP ranks Every rank derives the same epoch permutation from seed + epoch (CPU generator, identical across ranks and reproducible) and takes an equal-size contiguous shard, with the ragged tail dropped so per-rank batch counts always match and gradient sync cannot hang. The training loop's existing sampler.set_epoch call drives reshuffling, exactly like DistributedSampler; single-process runs auto-advance the epoch instead. The train/val split now uses a fixed-seed generator: identical across ranks (no leakage between a rank's train shard and another's val shard) and stable across save/reload, so resumed runs keep validating on the same held-out pixels. Verified with a 2-GPU torchrun run: equal batch counts on both ranks, finite identical reduced losses, no deadlock. --- src/quantem/tomography/dataset_models.py | 48 +++++++++++++++---- src/quantem/tomography/tomography.py | 33 ++++++++----- tests/tomography/test_device_batch_sampler.py | 40 +++++++++++++++- 3 files changed, 100 insertions(+), 21 deletions(-) diff --git a/src/quantem/tomography/dataset_models.py b/src/quantem/tomography/dataset_models.py index 00419695..af69d205 100644 --- a/src/quantem/tomography/dataset_models.py +++ b/src/quantem/tomography/dataset_models.py @@ -443,17 +443,25 @@ def to(self, device: str | torch.device): class DeviceBatchSampler: """Epoch iterator that builds INR training batches directly on a device. - Replaces the per-pixel DataLoader path for single-process runs: the tilt - stack and angles are made resident on ``device`` once, so producing a - batch is index arithmetic plus two tensor lookups instead of - ``batch_size`` Python ``__getitem__`` calls, a collate, and a - host-to-device copy per step. On a GPU this removes the CPU dataloader - bottleneck entirely. + Replaces the per-pixel DataLoader path: the tilt stack and angles are + made resident on ``device`` once, so producing a batch is index + arithmetic plus two tensor lookups instead of ``batch_size`` Python + ``__getitem__`` calls, a collate, and a host-to-device copy per step. + On a GPU this removes the CPU dataloader bottleneck entirely. Yields the same batch dicts as ``TomographyINRDataset.__getitem__`` under a DataLoader collate (``projection_idx``, ``pixel_i``, ``pixel_j``, ``phi``, ``target_value``), with the train loader's ``drop_last=True`` semantics. + + Distributed runs: pass ``rank``/``world_size`` and every rank derives + the *same* epoch permutation from ``seed + epoch`` (CPU generator, so + it is identical across ranks and reproducible), then takes an + equal-size contiguous shard — equal so per-rank batch counts match and + DDP gradient sync cannot hang on a ragged tail. The training loop's + ``sampler.set_epoch(epoch)`` drives reshuffling, exactly like + ``DistributedSampler``; without ``set_epoch`` the epoch advances + automatically on each ``__iter__``. """ def __init__( @@ -463,10 +471,17 @@ def __init__( device: torch.device | str, indices: torch.Tensor | None = None, shuffle: bool = True, + rank: int = 0, + world_size: int = 1, + seed: int = 0, ): self.batch_size = batch_size self.device = torch.device(device) self.shuffle = shuffle + self.rank = rank + self.world_size = world_size + self.seed = seed + self._epoch = 0 self._stack = dset.tilt_stack.to(self.device) self._angles = dset.tilt_angles.to(self.device) # __getitem__ decodes flat indices with shape[1] for both rows and @@ -476,14 +491,29 @@ def __init__( if indices is None: indices = torch.arange(len(dset), dtype=torch.int64) self._indices = indices.to(self.device) + self._per_rank = len(self._indices) // world_size + + def set_epoch(self, epoch: int) -> None: + """Set the epoch used to seed this epoch's shared permutation.""" + self._epoch = epoch def __len__(self) -> int: - return len(self._indices) // self.batch_size # drop_last=True + return self._per_rank // self.batch_size # drop_last=True - def __iter__(self): + def _epoch_shard(self) -> torch.Tensor: idx = self._indices if self.shuffle: - idx = idx[torch.randperm(len(idx), device=self.device)] + g = torch.Generator() + g.manual_seed(self.seed + self._epoch) + perm = torch.randperm(len(idx), generator=g).to(self.device) + idx = idx[perm] + self._epoch += 1 # auto-advance; set_epoch overrides per epoch + if self.world_size > 1: + idx = idx[self.rank * self._per_rank : (self.rank + 1) * self._per_rank] + return idx + + def __iter__(self): + idx = self._epoch_shard() per_proj = self._s1 * self._s2 for k in range(len(self)): sel = idx[k * self.batch_size : (k + 1) * self.batch_size] diff --git a/src/quantem/tomography/tomography.py b/src/quantem/tomography/tomography.py index 2fdbb7cf..4a68edc5 100644 --- a/src/quantem/tomography/tomography.py +++ b/src/quantem/tomography/tomography.py @@ -414,30 +414,41 @@ def _rebuild_dataloader(self, batch_size: int, num_workers: int, val_fraction: f def _setup_recon_dataloaders(self, batch_size: int, num_workers: int, val_fraction: float): """Build the train/val batch iterators. - Single-process runs use ``DeviceBatchSampler`` — batches are built - with tensor ops on the compute device from a device-resident tilt - stack, removing the per-pixel ``__getitem__`` / collate / H2D-copy + INR datasets use ``DeviceBatchSampler`` — batches are built with + tensor ops on the compute device from a device-resident tilt stack, + removing the per-pixel ``__getitem__`` / collate / H2D-copy dataloader bottleneck (``num_workers`` is ignored on this path). - Multi-process (DDP) runs keep the DataLoader + DistributedSampler - path. + Distributed runs shard the same seeded epoch permutation across + ranks (DistributedSampler semantics; the loop's ``set_epoch`` drives + reshuffling). Non-INR datasets keep the DataLoader path. """ - if self.world_size == 1 and isinstance(self.dset, TomographyINRDataset): + if isinstance(self.dset, TomographyINRDataset): n = len(self.dset) n_val = int(n * val_fraction) - perm = torch.randperm(n) + # Fixed-seed split: identical across DDP ranks (no train/val + # leakage between ranks) and stable across save/reload, so a + # resumed run keeps validating on the same held-out pixels. + split_gen = torch.Generator() + split_gen.manual_seed(0) + perm = torch.randperm(n, generator=split_gen) + ddp = dict(rank=self.global_rank, world_size=self.world_size) self.dataloader = DeviceBatchSampler( - self.dset, batch_size, self.device, indices=perm[n_val:] + self.dset, batch_size, self.device, indices=perm[n_val:], **ddp ) # The val sampler keeps its own device-resident copy of the tilt # stack; acceptable, since val_fraction > 0 is the rare case. - self.val_dataloader = ( + val = ( DeviceBatchSampler( - self.dset, batch_size, self.device, indices=perm[:n_val], shuffle=False + self.dset, batch_size, self.device, indices=perm[:n_val], shuffle=False, **ddp ) if n_val > 0 else None ) - self.sampler = None + # A per-rank val shard smaller than one batch would divide by + # zero in the val-loss average. + self.val_dataloader = val if val is not None and len(val) > 0 else None + # The training loop calls set_epoch on self.sampler. + self.sampler = self.dataloader self.val_sampler = None return diff --git a/tests/tomography/test_device_batch_sampler.py b/tests/tomography/test_device_batch_sampler.py index e40b1284..47bce96f 100644 --- a/tests/tomography/test_device_batch_sampler.py +++ b/tests/tomography/test_device_batch_sampler.py @@ -101,4 +101,42 @@ def test_reconstruct_uses_sampler_single_process(): ) assert isinstance(tomo.dataloader, DeviceBatchSampler) assert isinstance(tomo.val_dataloader, DeviceBatchSampler) - assert tomo.sampler is None + # the loop drives epoch reshuffling through set_epoch on self.sampler + assert tomo.sampler is tomo.dataloader + + +def _flat(batch, s1, s2): + return batch["projection_idx"] * (s1 * s2) + batch["pixel_i"] * s1 + batch["pixel_j"] + + +def test_ddp_shards_are_disjoint_and_equal(): + dset = _dset() + s1, s2 = dset.tilt_stack.shape[1], dset.tilt_stack.shape[2] + world = 4 + shards = [] + for rank in range(world): + sampler = DeviceBatchSampler(dset, 32, "cpu", rank=rank, world_size=world) + sampler.set_epoch(3) + flats = torch.cat([_flat(b, s1, s2) for b in sampler]) + shards.append(flats) + assert len(sampler) == (len(dset) // world) // 32 # equal on every rank + allv = torch.cat(shards) + assert allv.unique().numel() == allv.numel() # no pixel on two ranks + + +def test_ddp_epoch_permutation_shared_and_reproducible(): + dset = _dset() + s1, s2 = dset.tilt_stack.shape[1], dset.tilt_stack.shape[2] + + def epoch_flats(rank, epoch): + sampler = DeviceBatchSampler(dset, 32, "cpu", rank=rank, world_size=2) + sampler.set_epoch(epoch) + return torch.cat([_flat(b, s1, s2) for b in sampler]) + + # same (rank, epoch) on a fresh instance -> identical batches + torch.testing.assert_close(epoch_flats(0, 7), epoch_flats(0, 7), rtol=0, atol=0) + # different epoch -> different order + assert not torch.equal(epoch_flats(0, 7), epoch_flats(0, 8)) + # ranks of the same epoch are disjoint + both = torch.cat([epoch_flats(0, 7), epoch_flats(1, 7)]) + assert both.unique().numel() == both.numel()