flax
[NVIDIA] Use custom grad accumulation for FP8 params
#3623

Merged

[NVIDIA] Use custom grad accumulation for FP8 params #3623

copybara-service merged 3 commits into google:main from kaixih:fp8_meta_custom_grad_accumulation

Use custom grad accum for fp8 meta params

02913030

Add the full attr

54fa392b

Clean up

dd004c21

kaixih force pushed to dd004c21 2 years ago

cgarciae approved these changes on 2024-02-27

cgarciae added pull ready

copybara-service merged e4282ee1 into main 2 years ago

Reviewers

cgarciae

Assignees

No one assigned

Labels

pull ready

Milestone

No milestone