Add support for Torch 2.0

Torch 2.0 has been out for a couple of months now, and brings the promise of faster model training thanks to their new compilation engine - in addition to being the new basis to which future features will be added. Overall, Torch 2.0 preserves compatibility with the previous versions, meaning that most of the declearn codebase is natively compatible with 2.0 simply by being compatible with 1.10-1.13 versions. This is, however, not the case of the differential-privacy-oriented code, which has been relying on functorch, which is now deprecated in favor of torch.func, with the latter building on the previous logic and preserving most of the function names and signatures, but still introducing some breaking changes.

Hence, this issue is about adding support for Torch 2.0, while retaining support for versions 1.10-1.13 (or is it really 1.11-1.13, due to functorch?), since 2.0 is very young, and these versions are quite recent as well (1.10 was released in October 2021, 1.11 in March 2022).

Current advancement:

• Write an alternative 'TorchModel' for Torch-2.0.
• Refactor backend code to shrink and isolate version-dependent parts.
• Enable dynamically (and quietly) loading whichever version should be.
• Specify two alternative dependency specifications (torch 1 & functorch | torch 2) in pyproject.toml.
• Handle torch-version-based linter issues, with as little brutal silencing as possible.
• Enable testing both cases in the CI/CD, for proper maintenance.
• Optimize models using torch.compile in Torch-2.0.

added backend feature-request refactoring labels

assigned to @paandrey

See branch torch-2.0 for the on-going implementation effort.

For now, I was able to write an alternative TorchModel for torch 2.0, and incrementally refactor the code to end up with as little repetition a possible. Here are the steps I took:

• 0: have two very-redundant classes;
• 1: use composition with two partly-redundant classes handling sample-wise gradients computations;
• 2: refactor torch-1.1X code to shrink version-dependent part;
• 3: refactor into a version-dependent (private) function, that is imported and used in the backend of TorchModel.

I could test locally that the code works with both torch 1.13 and 2.0. Note that pip warns about functorch and torch 2.0 being installed in the same environment; hence specifying the alternative dependencies in pyproject.toml might be a burden. I will deal with CI/CD considerations later (one possibility being to run tests with the latest torch version, and then forcefully retrograde to torch 1.1X+functorch to run test/models/test_torch.py only).

Regarding performance: on my CPU-only machine, using the toy "MLP-tune" test case, I achieve better runtime performance on clipped-gradients-computation with 1.13 than 2.0, thanks to the use of functorch.compile. If I remove the latter, performance become similar for the two versions of the package and code. I have not yet succeeded in using torch.compile in 2.0, which will hopefully yield similar acceleration as functorch.compile. I am still to run more thorough performance tests on CPU and GPU using our various test cases (or a real-life example such as a small-ish CNN on MNIST with DP-SGD training).

marked the checklist item Specify two alternative dependency specifications (torch 1 & functorch | torch 2) in pyproject.toml. as completed

The pyproject.toml and installation guide have now been updated to enable installing either torch 1 or torch 2 dependencies. There is also a 2.0 version of functorch, which we will not use but the installation of which should please the linters when parsing the torch 1.10-1.13 backend code branch for sample-wise clipped gradients computation.

I have also updated both backend code branches to add support for model buffers, i.e. non-differentiable tensors that are to be used in forward passes (and may be updated inplace as part of it).

marked the checklist item Handle torch-version-based linter issues, with as little brutal silencing as possible. as completed

marked the checklist item Optimize models using torch.compile in Torch-2.0. as completed

• When pip install declearn[all] is used, torch 2.X and functorch 2.0 get installed, resulting in both backend code branches being acceptable to linters. However, if one develops with torch 1.13 installed, the (unused) torch-2.0-only backend code (logically) raises some warnings and errors.
• If we are somewhat lazy, we may rely on this to have CI/CD pass thanks to the installation instructions. It may still be worth it to install [torch1] extra dependencies and re-run pytest test/model/test_torch.py for the sake of controlling that this (supposedly frozen) code is not broken (and count it in coverage stats), as a late step in the tests pipeline.
• The (optional) use of torch.compile has been implemented, in a minimal way that I think should be reviewed based on practical use cases prior to merging.

created branch 27-add-support-for-torch-2-0 to address this issue

mentioned in merge request !49 (merged)

marked the checklist item Enable testing both cases in the CI/CD, for proper maintenance. as completed

mentioned in commit 861d559b

closed with merge request !49 (merged)