BIGAUD Nathan authored 1 year ago and ANDREY Paul committed 1 year ago

Can be heavily simplified if implementation of model seeding
does not require anything fancier than a seed integer.

Add support for Torch 2.0

Closes #27

See merge request !49

- Torch 2.0 introduced `torch.compile`, a novel utilty to optimize
  computations via their JIT-compiling into optimized kernels. At
  the moment, compiled modules cannot be saved (but their states,
  which are shared with the underlying original module, can), and
  are not compatible with `torch.func` functional execution either.
- As declearn vows to support Torch 2.0, it seems crucial that end-
  users may use `torch.compile`. However, as long as Torch 1.10-13
  versions (and previous versions of DecLearn 2.X) are supported,
  this handling should not break backward compatibility.
- An initial approach was to enable compiling the handled module as
  part of `TorchModel.__init__`. However, this proves impractical,
  as it takes away the assumption that end-users should be able to
  use their customly-prepared module as-is - including pre-compiled
  ones. This is all the more true as their are many options to the
  `torch.compile` function, that DecLearn has no purpose handling.
- Therefore, the approach implemented here is to detect and handle
  models that were compiled prior to being input into `TorchModel`.

- Here are a few notes and caveats regarding the current implementation:
  - Some impractical steps were taken to ensure weights and gradients
    have the same name regardless of whether the module was compiled
    or not, _and regardless of the specific 2.x version of declearn_.
    When we move to declearn 3, it will be worth revising.
  - A positive consequence of the previous note is that the compilation
    of the module should not impair cases where some clients are using
    torch 1 and/or an older 2.x version of declearn.
  - The will to retain user-defined compilation option is mostly lost
    due to the current lack of recording of these info by torch. This
    is however expected to evolve in the future, which should enable
    sharing instructions with clients. See issue 101107 of pytorch:
    https://github.com/pytorch/pytorch/issues/101107
  - A clumsy bugfix was introduced to avoid an issue where the wrapped
    compiled model would not take weights updates into account when
    running in evaluation mode. The status of this hack should be kept
    under close look as the issue I opened to report the bug is treated:
    https://github.com/pytorch/pytorch/issues/104984

- Turn "all" and "torch" extras more loose, enabling (thus, in most cases,
  resulting in) installing torch >=2.0 rather than 1.1[0-3].
- Add "torch1" and "torch2" extras to specify explicitly whether the user
  would rather use Torch 1.13 or >=2.0. Opacus is made part of these two
  extras, as its version is also co-dependent with Torch/Functorch.
- Document this change as part of the package installation guide.

- Torch 2.0 was released in march 2023, introducing a number of
  new features to the torch ecosystem, while being non-breaking
  wrt past versions on most points.
- One of the salient changes is the introduction of `torch.func`,
  that integrates features previously introduced as part of the
  `functorch` package, which we have been relying upon in order
  to efficiently compute and clip sample-wise gradients.
- This commit therefore introduces a new backend code branch so
  as to make use of `torch.func` when `torch~=2.0` is used, and
  retain the existing `functorch`-based branch when older torch
  versions are used.
- As part of this effort, some of the `TorchModel` backend code
  was refactored, notably deferring the functional transform of
  the input `torch.nn.Module` to the first call for sample-wise
  gradients computation. This has the nice side effect of fixing
  cases when users want to train a non-functorch-compatible model
  without DP-SGD.

Refactor tests pipeline, introducing 'scripts/run_tests.sh'.

See merge request !52

- Add the 'scripts/run_tests.sh' bash script, which enables running
  various blocks of tests. This mostly refactors instructions that
  used to be part of 'tox.ini', in a more convenient way:
  - some tests are grouped (commands in a group _all_ run, but the
    group fails if any has failed once all are done);
  - instructions are (hopefully) easier to read through and should
    be easier to maintain in the future;
  - it should be easier to decline more groups in the future (e.g.
    to have a routine with only unit tests);
  - groups of tests can be run easily, including without using tox.
- Have the 'tox.ini' file trigger some groups of tests by calling
  the new bash script. This makes the tox file easier to read and
  parse through, and merely oriented at providing run isolation.
- Split test groups into various envs in tox, which may be slightly
  more costly to run, but enables writing simpler CI/CD pipelines
  that merely pass arguments to the tox command.
- Make it explicit in 'tox.ini' that tests can be run with python
  3.8 and onwards.
- Keep a "main" job that groups all others. This is useful in the
  Inria shared Gitlab Runner, as caching is disabled, so that the
  use of distinct jobs would result in redundant costful package
  installations.

Revise the MNIST example, cleaning up the code and documentation.

See merge request !51

Fix compatibility issues with scikit-learn 1.3.0 and tensorflow 2.13.0

Closes #28 and #29

See merge request !50

BIGAUD Nathan authored 1 year ago and ANDREY Paul committed 1 year ago

Co-authored-by: Paul ANDREY <paul.andrey@inria.fr>

The legacy 'test_main.py' integration tests are now mostly redundant
with other integration tests, adding little marginal value to the CI/CD
tests pipeline. From a Coverage POV, the only thing that is not already
tested elsewhere is the actual use of Scaffold.

Hence the test scenarios collected during PyTest runs are restricted to:
- FedAvg and Scaffold variants of Optimization strategies.
  (prior: FedAvg + FedAvgM + Scaffold)
- In non-fulltest, only the Scaffold strategy is used.
  (prior: FedAvg + Scaffold)

This results in testing 3 scenarios in the limited pipeline and 18 in the
"fulltest" one, when there previously would be 6 or 27 scenarios.

- Scikit-Learn 1.3.0 introduced the possibility to use any dtype for
  SGD-based models' weights.
- As a consequence, this commit introduces the optional 'dtype' argument
  to 'SklearnSGDModel.__init__' and '.from_parameters' methods, which is
  only used with `sickit-learn >=1.3` and an un-initialized model.
- Coherently (and for all scikit-learn versions), additional type/dtype
  verification has been implemented under the hood of `_unpack_batch`.

- Fix device-selection utils for Jax/Haiku.
- Ensure any GPU is left untouched when the device policy specifies
  that the CPU must be used.
- Fix instructions in both HaikuModel and its tests' setup code to
  achieve the former.
- Note: `jax.jit(..., device=<D>)` is used to prevent Haiku from
  using another device than the desired one; however this argument
  is being deprecated, making this solution short-lived. Since the
  supported Jax version is hard-coded however, we keep this patch
  for now, in wait of future framework evolutions.
- Add the use of the optional, custom `--cpu-only` flag in Haiku
  unit tests. This enabled testing the former changes in an env
  that has a detectable GPU, but an unproper driver version (so
  that any undesired use of that GPU would make the tests fail).

Minor gardening around the package

See merge request !44

- First run unit tests (which are the more problematic).
- Then run code quality verifications (pylint / mypy / black).
- Finally run integration tests.

The rationale to inverting the last two items relative to what was
done prior to this update is to keep the longest-running tests for
the end (rather than re-play a long pipeline due to late-detected
formatting issues).