Supporting Adding DDP Communication Hooks

[shuyingsunshine21]

What does this PR do?

Fixes #6727, #643

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?
• 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 update the CHANGELOG? (not for typos, docs, test updates, or internal minor changes/refactorings)

PR review

Anyone in the community is free to review the PR once the tests have passed.
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

Did you have fun?

Make sure you had fun coding 🙃

[awaelchli]

LGTM on a high level
don't have access to multi-gpu at the moment so can't test it :(

please see my comments for a few small suggestions for improvements!

[shuyingsunshine21]

tested locally :)

[carmocca]

This is epic! Thanks!

Minor comment but we use a 120 line length for both code and docs. Can you re-configure your formatter? You can also use our pre-commit setup

[carmocca]

Technically it's also available in 1.7.0 right? But protected with an underscore. Do we want to include it or were important improvements done from 1.7 to 1.8?

[shuyingsunshine21]

encountered import issue when I tried to import torch.distributed.algorithms for 1.7.0.
complaining ModuleNotFoundError: No module named torch.distributed.algorithms for conda tests (3.7, 1.7)

also, power SGD is introduced later

[shuyingsunshine21]

pre-commit and pull rebase

[SeanNaren]

Great work! having the all reduce in fp16 is a nice perf gain that @tchaton and @blefaudeux made me aware of. We should definitely run a few experiments to see what memory/speed/convergence looks like with these comm hooks :)

[blefaudeux]

Great work! having the all reduce in fp16 is a nice perf gain that @tchaton and @blefaudeux made me aware of. We should definitely run a few experiments to see what memory/speed/convergence looks like with these comm hooks :)

just to add some context, it's really interesting for multi-node, for a single node it may not bring much (or slow things down a tiny bit). When used with AMP there's no tradeoff really, the grads are really computed in fp16 anyway so this folds them back to what they were originally (they're upgraded to fp32 when leaving autocast)

[shuyingsunshine21]

yeah, from our experimental results, for multi-nodes, fp16 compress hook could give 1.5 speed up for XLM-R model not sacrificing the model accuracy, (even for using native AMP). Power SGD one give even bigger gain :)

[ananthsub]

@shuyingsunshine21 as a follow up we should add this to https://pytorch-lightning.readthedocs.io/en/latest/benchmarking/performance.html

[awaelchli]

Locally when running the special tests (bash tests/special_tests.sh) I get the following error:

====================================================== short test summary info =======================================================
FAILED tests/plugins/test_ddp_plugin_with_comm_hook.py::test_ddp_fp16_compress_comm_hook - AttributeError: 'torch._C._distributed_c...
=================================================== 1 failed, 2 warnings in 5.60s ====================================================
FAILED

============================================================== FAILURES ==============================================================
__________________________________________________ test_ddp_fp16_compress_comm_hook __________________________________________________

tmpdir = local('/tmp/pytest-of-adrian/pytest-33/test_ddp_fp16_compress_comm_ho0')

    @RunIf(skip_windows=True, min_torch="1.8.0", min_gpus=2, special=True)
    def test_ddp_fp16_compress_comm_hook(tmpdir):
        """Test for DDP FP16 compress hook."""
        model = BoringModel()
        training_type_plugin = DDPPlugin(
            ddp_comm_hook=default.fp16_compress_hook,
            sync_batchnorm=True,
        )
        trainer = Trainer(
            max_epochs=1,
            gpus=2,
            plugins=[training_type_plugin],
            default_root_dir=tmpdir,
            sync_batchnorm=True,
            fast_dev_run=True,
        )
        trainer.fit(model)
        trainer_comm_hook = (
>           trainer.accelerator.training_type_plugin._model.get_ddp_logging_data().comm_hook
        )
E       AttributeError: 'torch._C._distributed_c10d.DDPLoggingData' object has no attribute 'comm_hook'

I'm on torch 1.8.1.
It will be in 1.9, right? so we should update the min requirement in the test.

[shuyingsunshine21]

Ah, good catch, I am on 1.9, looks like this attribute is introduced later ... Let me update.