Traced models serialization and torchscripting fix #17206

michaelbenayoun · 2022-05-12T09:12:28Z

What does this PR do?

Fixes the issue that was preventing traced models to be TorchScripted
Fixes the issue that was preventing trace models serialization
Fixes get_attr issues

Fixes #15974

@jamesr66a Can you try on your end and validate that it solves your issues?

HuggingFaceDocBuilderDev · 2022-05-12T09:28:55Z

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

sgugger

Thanks for the fixes!

jamesr66a

Solves our issue, thanks!

LysandreJik · 2022-05-16T21:26:45Z

There seems to remain a few failures in torch.fx tests; the PR can be merged after those are solved!

jamesr66a · 2022-05-17T04:33:54Z

@michaelbenayoun I'd like to propose some additional fixes that we discovered were needed to properly trace T5ForConditionalGeneration:

jamesr66a@1a75148

Can these be added?

michaelbenayoun · 2022-05-17T08:28:51Z

Yes, for some reason the tests do not pass for torch 1.11 (I tested locally on torch 1.10).
I will add those changes too.

jamesr66a · 2022-05-17T17:22:46Z

@michaelbenayoun Can I propose one final change to switch the graph surgery workaround to only trigger on older PyTorch versions where it's relevant?

jamesr66a@5ac7bb7

Otherwise, when we're working on PyTorch nightly, this^ code breaks because it's trying to remove nodes that still have uses

michaelbenayoun · 2022-05-18T09:57:20Z

@jamesr66a I added the gating, but only from version 1.12 as it was failing otherwise.

michaelbenayoun · 2022-05-18T09:59:43Z

src/transformers/models/decision_transformer/modeling_decision_transformer.py

@@ -187,7 +187,7 @@ def _attn(self, query, key, value, attention_mask=None, head_mask=None):
        if not self.is_cross_attention:
            # if only "normal" attention layer implements causal mask
            query_length, key_length = query.size(-2), key.size(-2)
-            causal_mask = self.bias[:, :, key_length - query_length : key_length, :key_length].bool()
+            causal_mask = self.bias[:, :, key_length - query_length : key_length, :key_length].to(torch.bool)


Needed for copy consistency.

michaelbenayoun · 2022-05-18T10:00:00Z

src/transformers/models/distilbert/modeling_distilbert.py

@@ -211,7 +211,7 @@ def unshape(x: torch.Tensor) -> torch.Tensor:
        q = q / math.sqrt(dim_per_head)  # (bs, n_heads, q_length, dim_per_head)
        scores = torch.matmul(q, k.transpose(2, 3))  # (bs, n_heads, q_length, k_length)
        mask = (mask == 0).view(mask_reshp).expand_as(scores)  # (bs, n_heads, q_length, k_length)
-        scores = scores.masked_fill(mask, -float("inf"))  # (bs, n_heads, q_length, k_length)
+        scores = scores.masked_fill(mask, torch.tensor(-float("inf")))  # (bs, n_heads, q_length, k_length)


Needed to be able to TorchScript the traced model.

michaelbenayoun · 2022-05-18T10:00:40Z

src/transformers/models/gpt2/modeling_gpt2.py

@@ -198,7 +198,7 @@ def _attn(self, query, key, value, attention_mask=None, head_mask=None):
        if not self.is_cross_attention:
            # if only "normal" attention layer implements causal mask
            query_length, key_length = query.size(-2), key.size(-2)
-            causal_mask = self.bias[:, :, key_length - query_length : key_length, :key_length].bool()
+            causal_mask = self.bias[:, :, key_length - query_length : key_length, :key_length].to(torch.bool)


Needed to be able to TorchScript the traced model, this should not break things because it is equivalent, from the docs:

self.bool() is equivalent to self.to(torch.bool). See to().

michaelbenayoun · 2022-05-18T10:02:11Z

src/transformers/models/gpt2/modeling_gpt2.py

@@ -1410,7 +1410,7 @@ def forward(
                    "unexpected if using padding tokens in conjunction with `inputs_embeds.`"
                )

-        pooled_logits = logits[torch.arange(batch_size, device=self.device), sequence_lengths]
+        pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths]


Same reason.
This should not break things because the tensor should be on the same device as logits anyway, right?

Yes, and better not rely on self.device anyway for model parallelism (I've made a few PRs to hunt most of those).

michaelbenayoun · 2022-05-18T10:02:35Z

src/transformers/models/gpt_neo/modeling_gpt_neo.py

@@ -188,7 +188,7 @@ def _attn(self, query, key, value, attention_mask=None, head_mask=None):
        attn_weights = torch.matmul(query, key.transpose(-1, -2))

        query_length, key_length = query.size(-2), key.size(-2)
-        causal_mask = self.bias[:, :, key_length - query_length : key_length, :key_length].bool()
+        causal_mask = self.bias[:, :, key_length - query_length : key_length, :key_length].to(torch.bool)


Same as for GPT-2

michaelbenayoun · 2022-05-18T10:03:02Z

src/transformers/models/gpt_neo/modeling_gpt_neo.py

@@ -147,8 +147,8 @@ def __init__(self, config, attention_type):
        self.register_buffer("bias", bias)
        self.register_buffer("masked_bias", torch.tensor(-1e9))

-        self.attn_dropout = nn.Dropout(config.attention_dropout)
-        self.resid_dropout = nn.Dropout(config.resid_dropout)
+        self.attn_dropout = nn.Dropout(float(config.attention_dropout))


TorchScripting fails otherwise, this should not change anything.

michaelbenayoun · 2022-05-18T10:04:31Z

src/transformers/models/gptj/modeling_gptj.py

@@ -69,7 +69,7 @@ def fixed_pos_embedding(x, seq_dim=1, seq_len=None):
 def rotate_every_two(x):
    x1 = x[:, :, :, ::2]
    x2 = x[:, :, :, 1::2]
-    x = torch.stack((-x2, x1), axis=-1)
+    x = torch.stack((-x2, x1), dim=-1)


Needed for TorchScript. This should be ok since stack can take the dim argument since the very beginning.

michaelbenayoun · 2022-05-18T10:04:47Z

src/transformers/models/gptj/modeling_gptj.py

@@ -163,7 +163,7 @@ def _attn(

        # compute causal mask from causal mask buffer
        query_length, key_length = query.size(-2), key.size(-2)
-        causal_mask = self.bias[:, :, key_length - query_length : key_length, :key_length].bool()
+        causal_mask = self.bias[:, :, key_length - query_length : key_length, :key_length].to(torch.bool)


Same thing as GPT-2.

michaelbenayoun · 2022-05-18T10:04:59Z

src/transformers/models/gptj/modeling_gptj.py

@@ -971,7 +971,7 @@ def forward(
                    "unexpected if using padding tokens in conjunction with `inputs_embeds.`"
                )

-        pooled_logits = logits[torch.arange(batch_size, device=self.device), sequence_lengths]
+        pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths]


Same thing as GPT-2.

sgugger · 2022-05-18T14:41:52Z

src/transformers/models/gpt2/modeling_gpt2.py

@@ -1410,7 +1410,7 @@ def forward(
                    "unexpected if using padding tokens in conjunction with `inputs_embeds.`"
                )

-        pooled_logits = logits[torch.arange(batch_size, device=self.device), sequence_lengths]
+        pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths]


Yes, and better not rely on self.device anyway for model parallelism (I've made a few PRs to hunt most of those).

sgugger · 2022-05-18T14:42:53Z

src/transformers/utils/fx.py

@@ -126,45 +128,45 @@ def _generate_supported_model_classes(
 )


-def embedding_override(self, input):
+def torch_nn_embedding(self, input):


Why rename all of those? I liked the fact it was clear they were "fake" functions.

Yes, but it was the case for all of them, I think that is clear enough since they are mapped in the _MANUAL_META_OVERRIDES dictionary

sgugger · 2022-05-18T14:44:09Z

tests/test_modeling_common.py

@@ -133,6 +134,7 @@ class ModelTesterMixin:
    all_model_classes = ()
    all_generative_model_classes = ()
    fx_compatible = False
+    fx_trace_can_be_torchscripted = True


No to the new flag though. I'd rather reuse the fx_compatible flag.

It's another one. Basically all models but Swin can be torch scripted for now, so I use this flag to enable / disable the torchscript part in the torch_fx test.

Yes, I see it's a new flag. That's what I would like to avoid.

Then what do you suggest?

Not adding a new flag.

This is just for Swin, so you can override the new test and skip the part inside the swin testing file. My concern is that this gives several flags to learn for every contributor of new models, and they won't know the difference between all of them (right now we already have two with fx and torchscript, that's why I don't want a third).

I don't see how it's bad?

I am testing that traced models can be torchscripted here, without the flag I cannot really know when this should be tested.
Instead of using the flag, I could hardcode the classes to ignore but that seems worse IMO.

jamesr66a · 2022-05-19T03:10:58Z

@michaelbenayoun Unfortunately, bumping the version check up to 1.12 breaks us. Actually, that was indirectly working around a semantic issue with deleting the concrete arg node. Do you mind augmenting the patch with this:

pbelevich@e3fce52

michaelbenayoun · 2022-05-19T12:51:15Z

@sgugger Replaced the creation of a new flag by setting a special value for the fx_compatible flag for models that can be traced but not torchscipted (-1).
This flag should take a boolean value 99% of the time anyways.

sgugger · 2022-05-19T13:53:03Z

@michaelbenayoun This doesn't really work either. I'm not trying to be gratuitously painful here, but the common model tester is at the core of our test suite for the new model addition PRs. Those PRs are huge, and it only thanks to a robust CI that we can make sure the models added actually work with the whole API Transformers offers.

Adding a new flag, or a magic value for an existing flag, just because there is one model that needs different testing is not something we usually do or allow. In both cases, either the contributor or the reviewer will have no idea what your new flag/magic value does, especially since there is no documentation of it anywhere.

As I said before, in those instances where we need to adapt a common test to a specific model, we override it in the tester of said model. cc @LysandreJik and @patrickvonplaten

sgugger

Looks good on the test side now, thanks a lot!

* Fix torch.jit.script and pickling issues * Fix get_attr issues * Fix import in function * Fix GPT-J and T5 tracing for torch=1.11 * Gate graph surgery on torch version * Modeling minor changes to enable TorchScripting * Model serialization / deserialization test * Remove _assert_is_none users

michaelbenayoun requested a review from sgugger May 13, 2022 09:36

sgugger approved these changes May 13, 2022

View reviewed changes

jamesr66a approved these changes May 13, 2022

View reviewed changes

LysandreJik approved these changes May 16, 2022

View reviewed changes

michaelbenayoun added 4 commits May 17, 2022 13:49

Fix torch.jit.script and pickling issues

5661ed4

Fix get_attr issues

b8e5d9a

Fix import in function

9268f53

Fix GPT-J and T5 tracing for torch=1.11

9309f27

michaelbenayoun force-pushed the fx_issues branch from 373feb6 to 9309f27 Compare May 17, 2022 11:51

Trigger CI

37c5358

michaelbenayoun added 4 commits May 18, 2022 10:32

Gate graph surgery on torch version

ddef9c3

Modeling minor changes to enable TorchScripting

af092a1

Model serialization / deserialization test

119c8c5

Small fix

e107121

michaelbenayoun commented May 18, 2022

View reviewed changes

michaelbenayoun requested a review from sgugger May 18, 2022 10:06

sgugger reviewed May 18, 2022

View reviewed changes

michaelbenayoun added 4 commits May 19, 2022 10:00

Remove _assert_is_none users

80c5f96

Workaround to not add a flag in the tests

d6f4763

Remove flag

7d6ae49

Trigger CI

cb9f2b9

michaelbenayoun added 3 commits May 23, 2022 09:59

Fix test

175582a

Make style

72e6fe9

Trigger CI

772dada

michaelbenayoun requested a review from sgugger May 23, 2022 09:27

sgugger approved these changes May 23, 2022

View reviewed changes

michaelbenayoun merged commit 2e7e428 into huggingface:main May 23, 2022

michaelbenayoun deleted the fx_issues branch May 23, 2022 15:50

michaelbenayoun mentioned this pull request Aug 22, 2022

remvoe _create_and_check_torch_fx_tracing in specific test files #18667

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Traced models serialization and torchscripting fix #17206

Traced models serialization and torchscripting fix #17206

michaelbenayoun commented May 12, 2022 •

edited

Loading

HuggingFaceDocBuilderDev commented May 12, 2022 •

edited

Loading

sgugger left a comment

jamesr66a left a comment

LysandreJik commented May 16, 2022

jamesr66a commented May 17, 2022 •

edited

Loading

michaelbenayoun commented May 17, 2022

jamesr66a commented May 17, 2022

michaelbenayoun commented May 18, 2022

michaelbenayoun May 18, 2022

michaelbenayoun May 18, 2022

michaelbenayoun May 18, 2022

michaelbenayoun May 18, 2022

sgugger May 18, 2022

michaelbenayoun May 18, 2022

michaelbenayoun May 18, 2022

michaelbenayoun May 18, 2022

michaelbenayoun May 18, 2022

michaelbenayoun May 18, 2022

sgugger May 18, 2022

sgugger May 18, 2022

michaelbenayoun May 18, 2022

sgugger May 18, 2022

michaelbenayoun May 18, 2022

sgugger May 18, 2022

michaelbenayoun May 18, 2022

sgugger May 18, 2022

sgugger May 18, 2022 •

edited

Loading

michaelbenayoun May 18, 2022

jamesr66a commented May 19, 2022

michaelbenayoun commented May 19, 2022

sgugger commented May 19, 2022

sgugger left a comment

Traced models serialization and torchscripting fix #17206

Traced models serialization and torchscripting fix #17206

Conversation

michaelbenayoun commented May 12, 2022 • edited Loading

What does this PR do?

HuggingFaceDocBuilderDev commented May 12, 2022 • edited Loading

sgugger left a comment

Choose a reason for hiding this comment

jamesr66a left a comment

Choose a reason for hiding this comment

LysandreJik commented May 16, 2022

jamesr66a commented May 17, 2022 • edited Loading

michaelbenayoun commented May 17, 2022

jamesr66a commented May 17, 2022

michaelbenayoun commented May 18, 2022

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

sgugger May 18, 2022 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jamesr66a commented May 19, 2022

michaelbenayoun commented May 19, 2022

sgugger commented May 19, 2022

sgugger left a comment

Choose a reason for hiding this comment

michaelbenayoun commented May 12, 2022 •

edited

Loading

HuggingFaceDocBuilderDev commented May 12, 2022 •

edited

Loading

jamesr66a commented May 17, 2022 •

edited

Loading

sgugger May 18, 2022 •

edited

Loading