fix: fix DDP race during subtomo update (rank-0-only subtomo rewrite + barrier) - #41
Open
uermel wants to merge 1 commit into
Open
fix: fix DDP race during subtomo update (rank-0-only subtomo rewrite + barrier)#41uermel wants to merge 1 commit into
uermel wants to merge 1 commit into
Conversation
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>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Under multi-GPU DDP, the periodic missing-wedge update had every rank write the
same subtomo
.ptfiles concurrently, truncating them and later crashing the runwith:
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
DDPStrategyprocess-group timeout to give the now-serial rewrite headroom.Cause:
on_train_epoch_endruns on all ranks and callsupdate_subtomo_missing_wedges(), which builds a non-distributedDataLoaderover the full dataset (the training
DistributedSampleris stripped via.dataset)and
torch.saves every result:torch.savetruncates 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.