Opt in flax and tf #17388

ArthurZucker · 2022-05-24T07:47:26Z

What does this PR do?

Adds support for OPT in both Flax and TF

Who can review?

@patrickvonplaten, @LysandreJik @younesbelkada @patil-suraj @sgugger

HuggingFaceDocBuilderDev · 2022-05-24T07:58:34Z

The documentation is not available anymore as the PR was closed or merged.

patrickvonplaten · 2022-05-24T12:41:40Z

Should we close the other PR? Let me know once it's ready for a review :-)

patrickvonplaten · 2022-05-24T14:58:53Z

Superseeds #17227 and #17226

tests/models/opt/test_modeling_opt.py

tests/models/opt/test_modeling_tf_opt.py

patrickvonplaten · 2022-05-25T20:48:17Z

Cool, very nice job @ArthurZucker !

Could you as a final safety guard also add TFOPT and FlaxOPT to the documentation test suite?

See: https://github.com/huggingface/transformers/tree/main/docs#docstring-testing

patrickvonplaten

Looks very nice from my side!

patrickvonplaten · 2022-05-25T20:46:28Z

tests/models/opt/test_modeling_opt.py


        EXPECTED_OUTPUTS = [
-            "Today is a beautiful day and I want to thank",
+            "Today is a beautiful day and I want everyone",


great thanks!

ArthurZucker · 2022-05-26T15:58:41Z

Can I merge @LysandreJik @sgugger ? (failing test are not related to OPT)

patrickvonplaten · 2022-05-26T16:30:46Z

utils/documentation_tests.txt

 src/transformers/models/mobilebert/modeling_mobilebert.py
 src/transformers/models/mobilebert/modeling_tf_mobilebert.py
 src/transformers/models/opt/modeling_opt.py
+src/transformers/models/opt/modeling_tf_opt.py


patrickvonplaten · 2022-05-26T16:31:34Z

@patil-suraj could you quickly check Flax and maybe @gante go over TF OPT?

gante

There was a lot of work put into this, and it is close to completion 💪 I've added a few questions, suggestions, and corrections on the TF side, but they should be straightforward. Great work!

P.S.: double-checking: have you run the slow tests locally? If you did, it implies that TF XLA generation is working for OPT 🎉

src/transformers/models/opt/modeling_tf_opt.py

gante · 2022-05-27T14:00:20Z

src/transformers/models/opt/modeling_tf_opt.py

+        for idx, decoder_layer in enumerate(self.layers):
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+


The LayerDrop, present in the PT version (L640), is missing

Yes it will be removed in PT as well (should probably need another PR WDYT @patrickvonplaten

src/transformers/models/opt/modeling_tf_opt.py

tests/models/opt/test_modeling_tf_opt.py

gante · 2022-05-27T14:19:11Z

tests/models/opt/test_modeling_tf_opt.py

+        )
+        self.assertTrue(np.allclose(output[:, :3, :3], expected_slice, atol=4e-3))
+
+        xla_generate = tf.function(model, jit_compile=True)


😱 does this run? If it does, ignore my comment above about the _update_model_kwargs_for_xla_generation function

I did not try it yet (xla not compatible with M1 chip, will try on brutasse soon

If it doesn't, remove the XLA compilation in the tests. It's still very brittle :)

I tested it and it seems to be working fine ! 🥳

Co-authored-by: Joao Gante <[email protected]>

…rmers into opt-flax-tf

…to opt-flax-tf

gante

Good to go, from the TF end 👍

sgugger

A few more nits, but LGTM otherwise. Thanks a lot!

sgugger · 2022-05-31T15:07:02Z

src/transformers/models/opt/modeling_flax_opt.py

+            cached_value.value = value
+            num_updated_cache_vectors = query.shape[1]
+            cache_index.value = cache_index.value + num_updated_cache_vectors
+            # causal mask for cached decoder self-attention: our single query position should only attend to those key positions that have already been generated and cached, not the remaining zero elements.


Can we respect the 119 char limit here and split that comment on several lines? ;-)

The long comment comes from BART, it has to be kept like that otherwise the repo-consistency fails.
Should leave like that

src/transformers/models/opt/modeling_tf_opt.py

tests/models/opt/test_modeling_flax_opt.py

tests/models/opt/test_modeling_tf_opt.py

Co-authored-by: Sylvain Gugger <[email protected]>

ArthurZucker · 2022-05-31T16:11:52Z

Thanks all for the reviews 😄 🥳

* initial commit * add init file * update globakl init * update index and dummy objects * style * update modelling auto * fix initi typo in src/transformers * fix typo in modeling tf auto, opt was in wrong mapping name * fixed a slow test : saved_model * style * fix positionnal embedding if no position id is provided * update tf test * update test flax requirements * fixed serialization * update * update tf name to allow smooth convertion * update flax tests * style * fix test typo * fix tf typo test * add xla for generate support in causal LM * fixed bug * cleaned tf tests * style * removed from PT for slow tests * fix typp * opt test as slow * trying to fix GPT2 undefined * correct documentation and add to test doc * update tf doc * fix doc * fake commit * Apply suggestions from code review Co-authored-by: Joao Gante <[email protected]> * update test based on review * merged main layer for functionning test * fixup + quality * Apply suggestions from code review Co-authored-by: Sylvain Gugger <[email protected]> * update long comment * make fix copies Co-authored-by: Arthur <[email protected]> Co-authored-by: Joao Gante <[email protected]> Co-authored-by: Sylvain Gugger <[email protected]>

ArthurZucker added 4 commits May 24, 2022 09:32

initial commit

3bd7a70

add init file

b4e03b7

update globakl init

5f4ad0a

update index and dummy objects

9bf9b95

ArthurZucker self-assigned this May 24, 2022

ArthurZucker added the Core: Modeling Internals of the library; Models. label May 24, 2022

ArthurZucker added 2 commits May 24, 2022 09:48

style

7e53484

update modelling auto

b3a74bd

ArthurZucker added 8 commits May 24, 2022 10:11

fix initi typo in src/transformers

a8e76bc

fix typo in modeling tf auto, opt was in wrong mapping name

9641709

fixed a slow test : saved_model

18bc191

style

1257768

fix positionnal embedding if no position id is provided

6355a10

update tf test

89d4dc9

update test flax requirements

da57713

fixed serialization

4fea708

update

d6b4fd5

patrickvonplaten reviewed May 24, 2022

View reviewed changes

tests/models/opt/test_modeling_opt.py Show resolved Hide resolved

patrickvonplaten reviewed May 24, 2022

View reviewed changes

tests/models/opt/test_modeling_tf_opt.py Outdated Show resolved Hide resolved

patrickvonplaten reviewed May 24, 2022

View reviewed changes

tests/models/opt/test_modeling_tf_opt.py Show resolved Hide resolved

patrickvonplaten reviewed May 24, 2022

View reviewed changes

tests/models/opt/test_modeling_tf_opt.py Outdated Show resolved Hide resolved

ArthurZucker added 6 commits May 25, 2022 10:13

update tf name to allow smooth convertion

2215d1c

update flax tests

fd18014

style

156b983

fix test typo

9597454

fix tf typo test

72510ea

add xla for generate support in causal LM

353f62f

patrickvonplaten approved these changes May 25, 2022

View reviewed changes

ArthurZucker added 3 commits May 26, 2022 12:22

correct documentation and add to test doc

79f005c

update tf doc

875a75c

fix doc

a1bf76c

ArthurZucker force-pushed the opt-flax-tf branch from f267718 to a1bf76c Compare May 26, 2022 15:23

fake commit

1a0389b

patrickvonplaten reviewed May 26, 2022

View reviewed changes

patrickvonplaten requested review from gante, patil-suraj and sgugger May 26, 2022 16:31

gante reviewed May 27, 2022

View reviewed changes

ArthurZucker and others added 6 commits May 27, 2022 21:09

Apply suggestions from code review

0ba0740

Co-authored-by: Joao Gante <[email protected]>

update test based on review

9352406

Merge branch 'opt-flax-tf' of https://github.com/ArthurZucker/transfo…

3ba2bbe

…rmers into opt-flax-tf

Merge branch 'main' of https://github.com/huggingface/transformers in…

c9ef765

…to opt-flax-tf

merged main layer for functionning test

8b3cbd3

fixup + quality

d054906

gante approved these changes May 30, 2022

View reviewed changes

sgugger approved these changes May 31, 2022

View reviewed changes

ArthurZucker and others added 3 commits May 31, 2022 17:26

Apply suggestions from code review

5f1bb05

Co-authored-by: Sylvain Gugger <[email protected]>

update long comment

f9a2202

make fix copies

a6f9537

ArthurZucker merged commit 7822a9b into huggingface:main May 31, 2022

Opt in flax and tf #17388

Opt in flax and tf #17388

Uh oh!

Conversation

ArthurZucker commented May 24, 2022

What does this PR do?

Who can review?

Uh oh!

HuggingFaceDocBuilderDev commented May 24, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

patrickvonplaten commented May 24, 2022

Uh oh!

patrickvonplaten commented May 24, 2022

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

patrickvonplaten commented May 25, 2022

Uh oh!

patrickvonplaten left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ArthurZucker commented May 26, 2022

Uh oh!

Choose a reason for hiding this comment

Uh oh!

patrickvonplaten commented May 26, 2022

Uh oh!

gante left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gante May 28, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gante left a comment

Choose a reason for hiding this comment

Uh oh!

sgugger left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ArthurZucker commented May 31, 2022

Uh oh!

Reviewers

Assignees

HuggingFaceDocBuilderDev commented May 24, 2022 •

edited

Loading

gante left a comment •

edited

Loading

gante May 28, 2022 •

edited

Loading