[bugfix] Prevent a DDP failure using copy

[tchaton]

What does this PR do?

Fixes #8821

This error is critical as it blocks DDP. Observations:

• needs num_workers > 0
• only for ddp (spawn works)
• returning a loss has impact
• epoch=True breaks with any reduce_fx
• step=True, sync_dist=True breaks with reduce_fx != "mean"

The failure can be reproduced with this minimal repro: #8821 (comment)

We believe this to be a bug upstream on PyTorch but haven't been able to easily reproduce it without Lightning.

A test hasn't been added because pytest seems to hang on teardown due to the num_workers>0 requirement. Not sure why.

This would break on master.

@RunIf(min_gpus=1)
def test_ddp_requires_a_deepcopy_on_training_step_output(tmpdir):

    class TestModel(BoringModel):

        def training_step(self, batch, batch_idx):
            loss = self(batch).sum()
            self.log('foo', torch.tensor(1), on_epoch=True)
            self.log('bar', torch.tensor(1), on_step=True, reduce_fx="sum", sync_dist=True)
            # a loss needs to be returned!
            return loss

        def train_dataloader(self):
            # only fails with `num_workers>0`
            return DataLoader(RandomDataset(32, 64), batch_size=2, num_workers=1)
 
    model = TestModel()
    trainer = Trainer(
        default_root_dir=tmpdir,
        gpus=1,
        accelerator='ddp',
        limit_train_batches=1,
        max_epochs=5,
        checkpoint_callback=False,
        logger=False,
    )
    trainer.fit(model)

Does your PR introduce any breaking changes? If yes, please list them.

None

Before submitting

• Was this discussed/approved via a GitHub issue? (not for typos and docs)
• Did you read the contributor guideline, Pull Request section?
• Did you make sure your PR does only one thing, instead of bundling different changes together?
• [n/a] 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 internal minor changes/refactorings)

PR review

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

• 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

[codecov]

Codecov Report

Merging #9239 (c43190a) into master (3e71046) will increase coverage by 4%.
The diff coverage is 100%.

@@           Coverage Diff           @@
##           master   #9239    +/-   ##
=======================================
+ Coverage      88%     92%    +4%     
=======================================
  Files         176     176            
  Lines       14807   14810     +3     
=======================================
+ Hits        13043   13663   +620     
+ Misses       1764    1147   -617     

[awaelchli]

I tried to add a clone() here but doesn't solve the issue. I'm just perplexed why the deepcopy is necessary. Do you have any intuition what is going on?

[tchaton]

Hey @awaelchli,

Mind sharing more details on what you are trying to do ?

Best,
T.C