Skip to content

Fix multi-GPU (DDP) training support#179

Open
u1vi wants to merge 1 commit into
OHF-Voice:mainfrom
u1vi:fix/multi-gpu-ddp-support
Open

Fix multi-GPU (DDP) training support#179
u1vi wants to merge 1 commit into
OHF-Voice:mainfrom
u1vi:fix/multi-gpu-ddp-support

Conversation

@u1vi

@u1vi u1vi commented Feb 17, 2026

Copy link
Copy Markdown

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:

  1. dataset.py — setup() crashes on non-rank-0 processes
    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.
  2. lightning.py — Missing toggle_optimizer/untoggle_optimizer
    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()).
  3. lightning.py — LR schedulers never stepped
    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.
  4. main.py — Auto-detect multi-GPU and set DDP strategy
    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

- __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)
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