Fix multi-GPU (DDP) training support#179
Open
u1vi wants to merge 1 commit into
Open
Conversation
- __main__.py: Auto-detect multiple GPUs and set ddp_find_unused_parameters_true strategy, needed because VITS has conditional code paths that leave some parameters unused - dataset.py: Load PiperConfig from JSON in setup() for non-rank-0 processes, since prepare_data() only runs on rank 0 in DDP mode - lightning.py: Add toggle_optimizer/untoggle_optimizer for proper DDP gradient synchronization with multiple optimizers - lightning.py: Remove unnecessary retain_graph=True - lightning.py: Add on_train_epoch_end to manually step LR schedulers (required when automatic_optimization=False)
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.
Multi-GPU training with DDP strategy fails with AssertionError in dataset.py because prepare_data() only runs on rank 0 in DDP mode, leaving piper_config as None on other ranks.
This PR fixes four issues that prevent or affect multi-GPU training:
prepare_data() only runs on rank 0 in DDP mode (by Lightning design), but it's where self.piper_config gets set. setup() runs on all ranks and asserts self.piper_config is not None, which fails on ranks 1+. Fix: load the config from the JSON file that prepare_data() already writes to disk.
With manual optimization + multiple optimizers in DDP, toggle_optimizer is needed so DDP correctly synchronizes only the relevant gradients during each backward pass. This also makes retain_graph=True unnecessary since the generator and discriminator loss graphs are already independent (via y_hat.detach()).
With automatic_optimization = False, Lightning does not automatically step LR schedulers. Added on_train_epoch_end to step them manually. Note: this was also a bug on single-GPU training — the learning rate decay was never being applied.
VITS has conditional code paths (e.g., stochastic duration predictor, speaker embeddings) that leave some parameters unused in certain steps. DDP requires find_unused_parameters=True to handle this. Instead of requiring users to pass an extra flag, multi-GPU is now auto-detected and the correct strategy is set automatically.
Usage:
No extra flags needed — multi-GPU training works out of the box
Tested with: 4x GPU setup, PyTorch Lightning, Python 3.13