Skip to content

fix: fix DDP race during subtomo update (rank-0-only subtomo rewrite + barrier) - #41

Open
uermel wants to merge 1 commit into
MLI-lab:masterfrom
uermel:ddp-subtomo-race-fix
Open

fix: fix DDP race during subtomo update (rank-0-only subtomo rewrite + barrier)#41
uermel wants to merge 1 commit into
MLI-lab:masterfrom
uermel:ddp-subtomo-race-fix

Conversation

@uermel

@uermel uermel commented Jun 9, 2026

Copy link
Copy Markdown

Under multi-GPU DDP, the periodic missing-wedge update had every rank write the
same subtomo .pt files concurrently, truncating them and later crashing the run
with:

PytorchStreamReader failed reading zip archive: failed finding central directory

This PR makes the rewrite rank-0-only and adds a barrier so the other ranks
wait for rank 0 to finish before reading the files again. It also bumps the
DDPStrategy process-group timeout to give the now-serial rewrite headroom.

Cause:

on_train_epoch_end runs on all ranks and calls
update_subtomo_missing_wedges(), which builds a non-distributed DataLoader
over the full dataset (the training DistributedSampler is stripped via .dataset)
and torch.saves every result:

train_set = train_loader.dataset                       # full dataset
dataset   = torch.utils.data.ConcatDataset(datasets)   # full set
loader    = torch.utils.data.DataLoader(dataset, ...)  # no sampler
...
for subtomo, file in zip(subtomo_batch, batch["subtomo0_file"]):
    torch.save(subtomo.cpu().clone(), file)            # N ranks -> same paths

torch.save truncates in place, so N concurrent writers (or a concurrent reader)
corrupt the zip. Training reads are sharded by PL's injected
DistributedSampler; this write path is not — that asymmetry is the bug.

update_subtomo_missing_wedges ran on every DDP rank and built a
non-distributed DataLoader over the full subtomo set, so all ranks
torch.save()'d the same .pt paths concurrently. That truncated the zip
archives ("PytorchStreamReader failed reading zip archive: failed finding
central directory") and crashed the run on a later epoch's load.

Split the method into a DDP-safe wrapper and a worker. The wrapper runs the
rewrite only on global-rank 0 (mirroring the update_hparam guard), then all
ranks synchronize at trainer.strategy.barrier() so non-zero ranks read the
freshly written files (SubtomoDataset.__getitem__ always reloads from disk).
This is correct because DDP keeps weights identical across ranks, so rank 0's
output equals what every rank would write.

Also bump the DDPStrategy process-group timeout to 2h: the now-serial rank-0
rewrite is ~Nx slower while other ranks wait at the barrier and could trip the
default ~30 min timeout on large subtomo sets.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
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