[NVIDIA] Use custom grad accumulation for FP8 params

[kaixih]

This pull request introduces a custom data type rule for the FP8 parameters to implement custom gradient accumulation. Specifically, when reusing the FP8 parameters, the autograd will accumulate their gradients. In this case, we aim for the accumulation to be a maximum operation instead of the default addition operation.

[kaixih]

cc. @nluehr @reedwm @zhangqiaorjc

[kaixih]

Also, cc. @mingxu1067

[kaixih]

Note, the failed tests on ValueError: Cannot convert_element_type to dtype=fp8_meta32 is probably because we might have utilized some feature only available from the jax nightly (0.4.24.devXXXX). Currently, Jax is on its 0.4.23. Maybe @mattjj knows better.

[cgarciae]

@kaixih it looks good but we will have to wait for JAX to push a new release to pypi for tests to pass (according to your comment).

[cgarciae]

@kaixih thanks! Seems that pytest is still broken.

[zhangqiaorjc]

@kaixih do you mind fixing the CI errors?

[kaixih]

It seems the CI is still on jax 0.4.23 (failed test)

Requirement already satisfied: jax>=0.4.19 in ./venv/lib/python3.10/site-packages (from flax==0.8.1) (0.4.23)
Requirement already satisfied: jaxlib in ./venv/lib/python3.10/site-packages (from flax==0.8.1) (0.4.23)

I re-tested on my machine but using 0.4.24 and then the tests can pass.

pip install --no-deps jax==0.4.24
pip install --no-deps jaxlib==0.4.24

@zhangqiaorjc @cgarciae Can you help check if the jax has already been updated to 0.4.24 or later?

[cgarciae]

Based on the output of the test run it seems not:

====== test config =======
PYTEST_OPTS: --cov=flax --cov-report=xml --cov-report=term --cov-config=pyproject.toml
RUN_DOCTEST: false
RUN_PYTEST: true
RUN_MYPY: false
RUN_PYTYPE: false
GH_VENV: true
WHICH PYTHON: /home/runner/work/flax/flax/venv/bin/python
jax: 0.4.23
flax: 0.8.1
==========================

[cgarciae]

As a quick fix maybe add

venv/bin/python3 -m pip install -U jax jaxlib

after this line:

flax/.github/workflows/build.yml

Line 120 in daf06ea

venv/bin/python3 -m pip install -e .[all,testing]

[kaixih]

@cgarciae Do you mean I should add this line in this PR?

[cgarciae]

Created a PR so you can rebase when merged.

[kaixih]

@zhangqiaorjc It seems all tests pass now. Can you take another look or reassign? Thx.

[zhangqiaorjc]

i'll let @cgarciae merge this and review the follow up in praxis

[zhangqiaorjc]

@cgarciae is this PR merge blocked on any internal error?