PR #17843 [Flax] Add remat (gradient checkpointing)

[Flax] Add remat (gradient checkpointing) #17843

sanchit-gandhi merged 11 commits into huggingface:main from sanchit-gandhi:flax-remat

sanchit-gandhi added WIP

sanchit-gandhi commented on 2022-06-23

sanchit-gandhi requested a review from

patil-suraj 3 years ago

sanchit-gandhi requested a review from

patrickvonplaten 3 years ago

sanchit-gandhi removed WIP

patrickvonplaten approved these changes on 2022-06-27

[Flax] Add remat (gradient checkpointing)

7dd4c58a

fix variable naming in test

03606c12

flip: checkpoint using a method

9b6e1644

fix naming

395221f8

fix class naming

7b165d53

apply PVP's suggestions from code review

d6040e01

make fix-copies

9972f381

sanchit-gandhi force pushed from c7915740 to 9972f381 3 years ago

fix big-bird, electra, roberta

0a0b6bfa

cookie-cutter

70b7175e

fix flax big-bird

0665a336

move test to common

66c4b14d

sanchit-gandhi merged 485bbe79 into main 3 years ago

sanchit-gandhi deleted the flax-remat branch 3 years ago

Reviewers

patrickvonplaten

patil-suraj

Assignees

No one assigned

Labels

None yet

Milestone

No milestone