Clean-up unnecessary warnings (including update to PyTorch 2.0) (#581)

Summary:
This PR is a collection of smaller fixes that will save us some deprecation issues in the future

## 1. Updating to PyTorch 2.0

**Key files: grad_sample/functorch.py, requirements.txt**

`functorch` has been a part of core PyTorch since 1.13.
Now they're going a step further and changing the API, while deprecating the old one.

There's a [guide](https://pytorch.org/docs/master/func.migrating.html) on how to migrate. TL;DR - `make_functional` will no longer be part of the API, with `torch.func.functional_call()` being (non drop-in) replacement.

They key difference for us is `make_functional()` creates a fresh copy of the module, while `functional_call()` uses existing module. As a matter of fact, we need the fresh copy (otherwise all the hooks start firing and you enter nested madness), so I've copy-pasted a [gist](https://gist.github.com/zou3519/7769506acc899d83ef1464e28f22e6cf) from the official guide on how to get a full replacement for `make_functional`.

## 2. New mechanism for gradient accumulation detection

**Key file: privacy_engine.py, grad_sample_module.py**

As [reported](https://discuss.pytorch.org/t/gan-raises-userwarning-using-a-non-full-backward-hook-when-the-forward-contains-multiple/175638/2) on the forum, clients are still getting "non-full backward hook" warning even when using `grad_sample_mode="ew"`. Naturally, `functorch` and `hooks` modes rely on backward hooks and can't be migrated to full hooks because [reasons](#328 (comment)). However, `ew` doesn't rely on hooks and it's unclear why the message should appear.

The reason, however, is simple. If the client is using poisson sampling we add an extra check to prohibit gradient accumulation (two poisson batches combined is not a poisson batch), and we do that by the means of backward hooks.

~In this case, backward hook serves a simple purpose and there shouldn't be any problems with migrating to the new method, however that involved changing the checking method. That's because `register_backward_hook` is called *after* hooks on submodule, but `register_full_backward_hook` is called before.~

Strikethrough solution didn't work, because hook order execution is weird for complex graphs, e.g. for GANs. For example, if your forward call looks like this:
```
Discriminator(Generator(x))
```
then top-level module hook will precede submodule's hooks for `Generator`, but not for `Discriminator`

As such, I've realised that gradient accumulation is not even supported in `ExpandedWeights`, so we don't have to worry about that. And the other two modes are both hooks-based, so we can just check the accumulation in the existing backward hook, no need for an extra hook. Deleted some code, profit.

## 3. Refactoring `wrap_collate_with_empty` to please pickle

Now here're two facts I didn't know before

1) You can't pickle a nested function, e.g. you can't do the following
```python
def foo():
    def bar():
        <...>

    return bar

pickle.dump(foo(), ...)
```

2) Whether or not `multiprocessing` uses pickle is python- and platform- dependant.

This affects our tests when we test `DataLoader` with multiple workers. As such, our data loaders tests:
* Pass on CircleCI with python3.9
* Fail on my local machine with python3.9
* Pass on my local machine with python3.7

I'm not sure how cow common the issue is, but it's safer to just refactor `wrap_collate_with_empty` to avoid nested functions.

## 4. Fix benchmark tests

We don't really run `benchmarks/tests` on a regular basis, and some of them were broken since we've upgraded to PyTorch 1.13 (`API_CUTOFF_VERSION` doesn't exist anymore)

## 4. Fix flake8 config

Flake8 config no [longer support](https://flake8.pycqa.org/en/latest/user/configuration.html) inline comments, fix is due

Pull Request resolved: #581

Reviewed By: alexandresablayrolles

Differential Revision: D44749760

Pulled By: ffuuugor

fbshipit-source-id: cf225f4134c049da4ee2eef53e1af3ef54d090bf

Author: Igor Shilov

.circleci/config.yml

@@ -6,16 +6,16 @@ version: 2.1
 
 commands:
 
-  py_3_7_setup:
-    description: "Install and switch to Python 3.7.5; also install pip and pytest."
+  py_3_9_setup:
+    description: "Install and switch to Python 3.9; also install pip and pytest."
     steps:
       - run:
-          name: "Setup Python v3.7.5 environment"
+          name: "Setup Python v3.9 environment"
           command: |
             cd /opt/circleci/.pyenv && git pull && cd -
-            pyenv install -s 3.7.5
-            pyenv global 3.7.5
-            pyenv local 3.7.5
+            pyenv install -s 3.9.4
+            pyenv global 3.9.4
+            pyenv local 3.9.4
             pyenv versions
             echo "In venv: $(pyenv local) - $(python -V), $(pip -V)"
             sudo "$(which python)" -m pip install --upgrade pip
@@ -283,24 +283,16 @@ commands:
 
 jobs:
 
-  lint_py37_torch_release:
+  lint_py39_torch_release:
     docker:
-      - image: cimg/python:3.7.5
+      - image: cimg/python:3.9
     steps:
       - checkout
       - pip_dev_install
       - lint_flake8
       - lint_black
       - isort
 
-  unittest_py37_torch_release:
-    docker:
-      - image: cimg/python:3.7.5
-    steps:
-      - checkout
-      - pip_dev_install
-      - unit_tests
-
   unittest_py38_torch_release:
     docker:
       - image: cimg/python:3.8
@@ -328,34 +320,34 @@ jobs:
 
   prv_accountant_values:
     docker:
-      - image: cimg/python:3.7.5
+      - image: cimg/python:3.9
     steps:
       - checkout
-      - py_3_7_setup
+      - py_3_9_setup
       - pip_dev_install
       - run:
           name: "Unit test prv accountant"
           no_output_timeout: 1h
           command: |
             python -m unittest opacus.tests.prv_accountant
 
-  integrationtest_py37_torch_release_cpu:
+  integrationtest_py39_torch_release_cpu:
     docker:
-      - image: cimg/python:3.7.5
+      - image: cimg/python:3.9
     steps:
       - checkout
-      - py_3_7_setup
+      - py_3_9_setup
       - pip_dev_install
       - mnist_integration_test:
           device: "cpu"
 
-  integrationtest_py37_torch_release_cuda:
+  integrationtest_py39_torch_release_cuda:
     machine:
       resource_class: gpu.nvidia.small.multi
       image: ubuntu-2004-cuda-11.4:202110-01
     steps:
       - checkout
-      - py_3_7_setup
+      - py_3_9_setup
       - pip_dev_install
       - run_nvidia_smi
       - mnist_integration_test:
@@ -369,13 +361,13 @@ jobs:
       - dcgan_integration_test:
           device: "cuda"
 
-  micro_benchmarks_py37_torch_release_cuda:
+  micro_benchmarks_py39_torch_release_cuda:
     machine:
       resource_class: gpu.nvidia.small.multi
       image: ubuntu-2004-cuda-11.4:202110-01
     steps:
       - checkout
-      - py_3_7_setup
+      - py_3_9_setup
       - pip_dev_install
       - run_nvidia_smi
       - benchmark_layers_integration_test:
@@ -459,7 +451,7 @@ jobs:
       image: ubuntu-2004-cuda-11.4:202110-01
     steps:
       - checkout
-      - py_3_7_setup
+      - py_3_9_setup
       - pip_dev_install
       - run_nvidia_smi
       - command_unit_tests_multi_gpu
@@ -493,9 +485,7 @@ workflows:
       not:
         equal: [ scheduled_pipeline, << pipeline.trigger_source >> ]
     jobs:
-      - lint_py37_torch_release:
-          filters: *exclude_ghpages
-      - unittest_py37_torch_release:
+      - lint_py39_torch_release:
           filters: *exclude_ghpages
       - unittest_py38_torch_release:
           filters: *exclude_ghpages
@@ -505,9 +495,9 @@ workflows:
           filters: *exclude_ghpages
       - unittest_multi_gpu:
           filters: *exclude_ghpages
-      - integrationtest_py37_torch_release_cpu:
+      - integrationtest_py39_torch_release_cpu:
           filters: *exclude_ghpages
-      - integrationtest_py37_torch_release_cuda:
+      - integrationtest_py39_torch_release_cuda:
           filters: *exclude_ghpages
       - prv_accountant_values:
           filters: *exclude_ghpages
@@ -518,12 +508,12 @@ workflows:
     jobs:
       - unittest_py39_torch_nightly:
           filters: *exclude_ghpages
-      - integrationtest_py37_torch_release_cpu:
+      - integrationtest_py39_torch_release_cpu:
           filters: *exclude_ghpages
-      - integrationtest_py37_torch_release_cuda:
+      - integrationtest_py39_torch_release_cuda:
           filters: *exclude_ghpages
-      - lint_py37_torch_release:
+      - lint_py39_torch_release:
           filters: *exclude_ghpages
-      - micro_benchmarks_py37_torch_release_cuda:
+      - micro_benchmarks_py39_torch_release_cuda:
           filters: *exclude_ghpages
 

.circleci/flake8_config.ini

@@ -5,60 +5,113 @@ max-line-length = 80
 ignore =
   # Black conflicts and overlaps.
   # Found in https://github.com/psf/black/issues/429
-  B950, # Line too long. (Use `arc lint`'s LINEWRAP instead)
-  E111, # Indentation is not a multiple of four.
-  E115, # Expected an indented block (comment).
-  E117, # Over-indented.
-  E121, # Continuation line under-indented for hanging indent.
-  E122, # Continuation line missing indentation or outdented.
-  E123, # Closing bracket does not match indentation of opening bracket's line.
-  E124, # Closing bracket does not match visual indentation.
-  E125, # Continuation line with same indent as next logical line.
-  E126, # Continuation line over-indented for hanging indent.
-  E127, # Continuation line over-indented for visual indent.
-  E128, # Continuation line under-indented for visual indent.
-  E129, # Visually indented line with same indent as next logical line.
-  E201, # Whitespace after '('.
-  E202, # Whitespace before ')'.
-  E203, # Whitespace before ':'.
-  E221, # Multiple spaces before operator.
-  E222, # Multiple spaces after operator.
-  E225, # Missing whitespace around operator.
-  E226, # Missing whitespace around arithmetic operator.
-  E227, # Missing whitespace around bitwise or shift operator.
-  E231, # Missing whitespace after ',', ';', or ':'.
-  E241, # Multiple spaces after ','.
-  E251, # Unexpected spaces around keyword / parameter equals.
-  E261, # At least two spaces before inline comment.
-  E262, # Inline comment should start with '# '.
-  E265, # Block comment should start with '# '.
-  E271, # Multiple spaces after keyword.
-  E272, # Multiple spaces before keyword.
-  E301, # Expected 1 blank line, found 0.
-  E302, # Expected 2 blank lines, found 0.
-  E303, # Too many blank lines (3).
-  E305, # Expected 2 blank lines after end of function or class.
-  E306, # Expected 1 blank line before a nested definition.
-  E501, # Line too long (82 > 79 characters).
-  E502, # The backslash is redundant between brackets.
-  E701, # Multiple statements on one line (colon).
-  E702, # Multiple statements on one line (semicolon).
-  E703, # Statement ends with a semicolon.
-  E704, # Multiple statements on one line (def).
-  W291, # Trailing whitespace.
-  W292, # No newline at end of file.
-  W293, # Blank line contains whitespace.
-  W391, # Blank line at end of file.
+  # B950: Line too long. (Use `arc lint`'s LINEWRAP instead)
+  B950,
+  # E111: Indentation is not a multiple of four.
+  E111,
+  # E115: Expected an indented block (comment).
+  E115,
+  # E117: Over-indented.
+  E117,
+  # E121: Continuation line under-indented for hanging indent.
+  E121,
+  # E122: Continuation line missing indentation or outdented.
+  E122,
+  # E123: Closing bracket does not match indentation of opening bracket's line.
+  E123,
+  # E124: Closing bracket does not match visual indentation.
+  E124,
+  # E125: Continuation line with same indent as next logical line.
+  E125,
+  # E126: Continuation line over-indented for hanging indent.
+  E126,
+  # E127: Continuation line over-indented for visual indent.
+  E127,
+  # E128: Continuation line under-indented for visual indent.
+  E128,
+  # E129: Visually indented line with same indent as next logical line.
+  E129,
+  # E201: Whitespace after '('.
+  E201,
+  # E202: Whitespace before ')'.
+  E202,
+  # E203: Whitespace before ':'.
+  E203,
+  # E221: Multiple spaces before operator.
+  E221,
+  # E222: Multiple spaces after operator.
+  E222,
+  # E225: Missing whitespace around operator.
+  E225,
+  # E226: Missing whitespace around arithmetic operator.
+  E226,
+  # E227: Missing whitespace around bitwise or shift operator.
+  E227,
+  # E231: Missing whitespace after ',', ';', or ':'.
+  E231,
+  # E241: Multiple spaces after ','.
+  E241,
+  # E251: Unexpected spaces around keyword / parameter equals.
+  E251,
+  # E261: At least two spaces before inline comment.
+  E261,
+  # E262: Inline comment should start with '# '.
+  E262,
+  # E265: Block comment should start with '# '.
+  E265,
+  # E271: Multiple spaces after keyword.
+  E271,
+  # E272: Multiple spaces before keyword.
+  E272,
+  # E301: Expected 1 blank line, found 0.
+  E301,
+  # E302: Expected 2 blank lines, found 0.
+  E302,
+  # E303: Too many blank lines (3).
+  E303,
+  # E305: Expected 2 blank lines after end of function or class.
+  E305,
+  # E306: Expected 1 blank line before a nested definition.
+  E306,
+  # E501: Line too long (82 > 79 characters).
+  E501,
+  # E502: The backslash is redundant between brackets.
+  E502,
+  # E701: Multiple statements on one line (colon).
+  E701,
+  # E702: Multiple statements on one line (semicolon).
+  E702,
+  # E703: Statement ends with a semicolon.
+  E703,
+  # E704: Multiple statements on one line (def).
+  E704,
+  # W291: Trailing whitespace.
+  W291,
+  # W292: No newline at end of file.
+  W292,
+  # W293: Blank line contains whitespace.
+  W293,
+  # W391: Blank line at end of file.
+  W391,
 
   # Too opinionated.
-  E265, # Block comment should start with '# '.
-  E266, # Too many leading '#' for block comment.
-  E402, # Module level import not at top of file.
-  E722, # Do not use bare except, specify exception instead. (Duplicate of B001)
-  F811, # Redefinition of unused name from line n.
-  P207, # (Duplicate of B003)
-  P208, # (Duplicate of C403)
-  W503  # Line break occurred before a binary operator.
+  # E265: Block comment should start with '# '.
+  E265,
+  # E266: Too many leading '#' for block comment.
+  E266,
+  # E402: Module level import not at top of file.
+  E402,
+  # E722: Do not use bare except, specify exception instead. (Duplicate of B001)
+  E722,
+  # F811: Redefinition of unused name from line n.
+  F811,
+  # P207: (Duplicate of B003)
+  P207,
+  # P208: (Duplicate of C403)
+  P208,
+  # W503: Line break occurred before a binary operator.
+  W503
+
 exclude =
   .hg,
   __pycache__,

benchmarks/tests/test_layers.py

@@ -15,36 +15,21 @@
 import copy
 import logging
 import random
-from typing import Any, Dict, List, Set, Tuple, Type
+from typing import Any, Dict, List, Tuple, Type
 
 import pytest
 import torch
 import torch.nn as nn
 from helpers import skipifnocuda
 from opacus.grad_sample import GradSampleModule
-from opacus.grad_sample.gsm_exp_weights import (
-    API_CUTOFF_VERSION,
-    GradSampleModuleExpandedWeights,
-)
+from opacus.grad_sample.gsm_exp_weights import GradSampleModuleExpandedWeights
 from opacus.layers import DPGRU, DPLSTM, DPRNN, DPMultiheadAttention
 
 from benchmarks.layers import LayerFactory
 from benchmarks.utils import reset_peak_memory_stats
 
 
-def _gsm_modes() -> Set[str]:
-    ret = ["baseline", "hooks"]
-    try:
-        import functorch  # noqa: F401, Checking for import errors
-
-        ret += ["functorch"]
-    except ImportError:
-        pass
-
-    if torch.__version__ >= API_CUTOFF_VERSION:
-        ret += ["ew"]
-    return set(ret)
-
+GSM_MODES = {"baseline", "hooks", "ew", "functorch"}
 
 PARAMETERS = [
     (
@@ -121,7 +106,7 @@ def test_layer_modules(
     """
 
     for layer_name, module, gsm_mode_blocklist in layer_list:
-        for gsm_mode in _gsm_modes() - set(gsm_mode_blocklist):
+        for gsm_mode in GSM_MODES - set(gsm_mode_blocklist):
             if gsm_mode in gsm_mode_blocklist:
                 continue
 
@@ -161,7 +146,7 @@ def test_to_device(
     assert reset_peak_memory_stats(cuda).cur_mem == 0
 
     for layer_name, module, gsm_mode_blocklist in layer_list:
-        for gsm_mode in _gsm_modes() - set(gsm_mode_blocklist):
+        for gsm_mode in GSM_MODES - set(gsm_mode_blocklist):
             layer = LayerFactory.create(
                 layer_name=layer_name,
                 batch_size=64,
@@ -207,7 +192,7 @@ def test_layer_outputs(
     }
 
     for layer_name, module, gsm_mode_blocklist in layer_list:
-        for gsm_mode in _gsm_modes() - set(gsm_mode_blocklist):
+        for gsm_mode in GSM_MODES - set(gsm_mode_blocklist):
             for random_seed in (random_seed_a, random_seed_b):
                 logging.error(f"{gsm_mode}, {layer_name}")
                 layer = LayerFactory.create(
@@ -246,7 +231,7 @@ def test_forward_backward(
         layer_config: config for instantiating the layers in layer_list
     """
     for layer_name, module, gsm_mode_blocklist in layer_list:
-        for gsm_mode in _gsm_modes() - set(gsm_mode_blocklist):
+        for gsm_mode in GSM_MODES - set(gsm_mode_blocklist):
             layer = LayerFactory.create(
                 layer_name=layer_name,
                 batch_size=64,