Fix wrong behavior of `DDPStrategy` option with simple GAN training using DDP #20936

samsara-ku · 2025-06-25T02:45:01Z

What does this PR do?

Fixes #20866 #20328 #18740 #17212

This PR adds MultiModelDDPStrategy class and its simple execution example, for the multi-gpu training with GAN training.

Simply speaking:

* right way using DDP in pytorch.
# each model is divided and wrapped by DDP class
generator, discriminator = DDP(generator), DDP(discrimiantor)

* current pytorch-lightning GAN training.
# trainer contains two nn.Modules; generator and discriminator
DDP(trainer) -> trainer contains two nn.Modules; generator and discriminator

* proposed `MultiModelDDP` class can use this way.
# trainer also has two nn.Modules, generator and discriminator;
# this way allocate DDP class to the each model in the trainer, not the trainer.
MultiModelDDP(trainer)

Currently, pytorch lightning simple GAN training has has problem with DistributedDataParallel strategy. It tries to wrap pl.trainer, not the nn.Module models in the pl.trainer
Although we can activate find_unused_parameters=True options to avoid this issue but it is not right way; I think it is just a trick.
So the key idea to solve this issue is that we assign DistributedDataParallel to the each model in the pl.trainer, different from the previous strategy DDPStrategy.
I already tested with my GPUs to visaulize the result and tracked the gradients of model with each epoch; it works and you can see the visulized result in thie google drive link

Before submitting

Was this discussed/agreed via a GitHub issue? (not for typos and docs)
[v] Did you read the contributor guideline, Pull Request section?
[v] Did you make sure your PR does only one thing, instead of bundling different changes together?
Did you make sure to update the documentation with your changes? (if necessary)
Did you write any new necessary tests? (not for typos and docs)
Did you verify new and existing tests pass locally with your changes?
Did you list all the breaking changes introduced by this pull request?
Did you update the CHANGELOG? (not for typos, docs, test updates, or minor internal changes/refactors)

PR review

Anyone in the community is welcome to review the PR.
Before you start reviewing, make sure you have read the review guidelines. In short, see the following bullet-list:

Reviewer checklist

Is this pull request ready for review? (if not, please submit in draft mode)
Check that all items from Before submitting are resolved
Make sure the title is self-explanatory and the description concisely explains the PR
Add labels and milestones (and optionally projects) to the PR so it can be classified

📚 Documentation preview 📚: https://pytorch-lightning--20936.org.readthedocs.build/en/20936/

for more information, see https://pre-commit.ci

Borda · 2025-06-25T09:20:32Z

examples/pytorch/domain_templates/generative_adversarial_net_ddp.py

+        def block(in_feat, out_feat, normalize=True):
+            layers = [nn.Linear(in_feat, out_feat)]
+            if normalize:
+                layers.append(nn.BatchNorm1d(out_feat, 0.8))
+            layers.append(nn.LeakyReLU(0.2, inplace=True))
+            return layers


lets move it out s a funtion

src/lightning/pytorch/strategies/ddp.py

stale · 2025-07-19T05:54:37Z

This pull request has been automatically marked as stale because it has not had recent activity. It will be closed in 7 days if no further activity occurs. If you need further help see our docs: https://lightning.ai/docs/pytorch/latest/generated/CONTRIBUTING.html#pull-request or ask the assistance of a core contributor here or on Discord. Thank you for your contributions.

…dp-implementation test: cover MultiModelDDPStrategy

examples/pytorch/domain_templates/generative_adversarial_net.py

Borda · 2025-08-08T17:36:20Z

examples/pytorch/domain_templates/generative_adversarial_net_ddp.py

+        opt_d = torch.optim.Adam(self.discriminator.parameters(), lr=lr, betas=(b1, b2))
+        return opt_g, opt_d
+
+    # ! TESTING


What does this comment mean?

examples/pytorch/domain_templates/generative_adversarial_net_ddp.py

Borda · 2025-08-08T17:38:11Z

@SkafteNicki could you pls check too :)

Borda · 2025-08-08T21:47:27Z

src/lightning/pytorch/strategies/ddp.py

@@ -419,6 +419,39 @@ def teardown(self) -> None:
        super().teardown()


+class MultiModelDDPStrategy(DDPStrategy):
+    @override
+    def _setup_model(self, model: Module) -> Module:


The typing is not very happy here as the parent class has the following footprint:

def _setup_model(self, model: Module) -> DistributedDataParallel:

SkafteNicki · 2025-08-09T11:57:34Z

examples/pytorch/domain_templates/generative_adversarial_net_ddp.py

+    dm = MNISTDataModule()
+    trainer = Trainer(
+        accelerator="auto",
+        devices=[0, 1, 2, 3],
+        strategy=MultiModelDDPStrategy(),
+        max_epochs=100,
+    )
+
+    trainer.fit(model, dm)


maybe i am reading it wrong, but is this the only core difference between this new _ddp script and the single gpu example script?
In that case could we please just incorporate the changes into the other script and add whatever arguments is needed to the ArgumentParser to distinguish between the two cases

SkafteNicki · 2025-08-09T11:58:47Z

src/lightning/pytorch/strategies/ddp.py

@@ -419,6 +419,39 @@ def teardown(self) -> None:
        super().teardown()


+class MultiModelDDPStrategy(DDPStrategy):


please add docstring to class. What is the purpose of the class and when to use it compared to the standard DDPStrategy

SkafteNicki · 2025-08-09T12:04:41Z

tests/tests_pytorch/strategies/test_multi_model_ddp.py

+    with mock.patch("lightning.pytorch.strategies.ddp._register_ddp_comm_hook") as register_hook:
+        strategy._register_ddp_hooks()
+
+    register_hook.assert_not_called()


could we add one or two tests similar to
https://github.com/samsara-ku/pytorch-lightning/blob/ece7d38abb908e96eb72af363606d2585e8843f2/tests/tests_pytorch/strategies/test_ddp_integration.py#L40-L66
that actually tests that the new strategy works with trainer.fit as expected

add: MultiModelDDPStrategy and its execution codes

00727e8

samsara-ku requested review from lantiga, Borda, tchaton, justusschock and ethanwharris as code owners June 25, 2025 02:45

github-actions bot added the pl Generic label for PyTorch Lightning package label Jun 25, 2025

[pre-commit.ci] auto fixes from pre-commit.com hooks

e6b061a

for more information, see https://pre-commit.ci

Borda reviewed Jun 25, 2025

View reviewed changes

src/lightning/pytorch/strategies/ddp.py Show resolved Hide resolved

stale bot added the won't fix This will not be worked on label Jul 19, 2025

samsara-ku added 2 commits August 5, 2025 17:10

refactor: extract block helper in GAN example

aa9b027

Merge pull request #1 from samsara-ku/codex/add-tests-for-multimodeld…

5503d3a

…dp-implementation test: cover MultiModelDDPStrategy

stale bot removed the won't fix This will not be worked on label Aug 5, 2025

Borda and others added 2 commits August 8, 2025 20:32

Merge branch 'master' into bugfix/gan-ddp-training

dc128b4

with

1fb4027

Borda reviewed Aug 8, 2025

View reviewed changes

examples/pytorch/domain_templates/generative_adversarial_net.py Outdated Show resolved Hide resolved

Borda reviewed Aug 8, 2025

View reviewed changes

examples/pytorch/domain_templates/generative_adversarial_net_ddp.py Outdated Show resolved Hide resolved

Apply suggestions from code review

ec62397

Borda approved these changes Aug 8, 2025

View reviewed changes

formating

ece7d38

Borda reviewed Aug 8, 2025

View reviewed changes

SkafteNicki reviewed Aug 9, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Fix wrong behavior of `DDPStrategy` option with simple GAN training using DDP #20936

Fix wrong behavior of `DDPStrategy` option with simple GAN training using DDP #20936

Uh oh!

samsara-ku commented Jun 25, 2025 •

edited by github-actions bot

Loading

Uh oh!

Borda Jun 25, 2025

Uh oh!

Uh oh!

stale bot commented Jul 19, 2025

Uh oh!

Uh oh!

Borda Aug 8, 2025

Uh oh!

Uh oh!

Borda commented Aug 8, 2025

Uh oh!

Borda Aug 8, 2025

Uh oh!

SkafteNicki Aug 9, 2025

Uh oh!

SkafteNicki Aug 9, 2025

Uh oh!

SkafteNicki Aug 9, 2025

Uh oh!

Uh oh!

		@@ -419,6 +419,39 @@ def teardown(self) -> None:
		super().teardown()


		class MultiModelDDPStrategy(DDPStrategy):

Fix wrong behavior of DDPStrategy option with simple GAN training using DDP #20936

Are you sure you want to change the base?

Fix wrong behavior of DDPStrategy option with simple GAN training using DDP #20936

Uh oh!

Conversation

samsara-ku commented Jun 25, 2025 • edited by github-actions bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

PR review

Uh oh!

Borda Jun 25, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

stale bot commented Jul 19, 2025

Uh oh!

Uh oh!

Borda Aug 8, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Borda commented Aug 8, 2025

Uh oh!

Borda Aug 8, 2025

Choose a reason for hiding this comment

Uh oh!

SkafteNicki Aug 9, 2025

Choose a reason for hiding this comment

Uh oh!

SkafteNicki Aug 9, 2025

Choose a reason for hiding this comment

Uh oh!

SkafteNicki Aug 9, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Fix wrong behavior of `DDPStrategy` option with simple GAN training using DDP #20936

Fix wrong behavior of `DDPStrategy` option with simple GAN training using DDP #20936

samsara-ku commented Jun 25, 2025 •

edited by github-actions bot

Loading