Lightning-AI / Lightning-AI/lightning-thunder

how to provide overrides needed for model compatibility

Open
#816 2 comments 0 reactions 1 assignee View on GitHub

@tfogal is already working on this.

Since Jul 22, 2024.

design interpreter program-coverage
Dominant language
Python
Stars
1.5k
Forks
121
PR merge metrics
No merged PRs in 30d

Description

While transformers BERT is not yet fully working, we are getting closer.
However, there is the need to disable some data-dependent control flow (typically: checks) to get it to work.
Transformers itself hides some but not all behind a check for compiling, e.g.

https://github.com/huggingface/transformers/blob/0fdea8607d7e01eb0e38a1ebeb7feee30a22f0cf/src/transformers/modeling_attn_mask_utils.py#L256-L260
```
is_tracing = (
torch.jit.is_tracing()
or isinstance(inputs_embeds, torch.fx.Proxy)
or (hasattr(torch, "_dynamo") and torch._dynamo.is_compiling())
)
```

So here is a candidate for how we might currently be able to run BERT after fixing #805 (and subsequent bugs):
```
import transformers, thunder, torch
@thunder.core.jit_ext.register_general_jit_lookaside(
transformers.modeling_utils.PreTrainedModel.warn_if_padding_and_no_attention_mask
)
@thunder.core.jit_ext.interpreter_needs_wrap
def dummy(*args):
pass

@thunder.core.jit_ext.register_general_jit_lookaside(
torch._dynamo.is_compiling
)
@thunder.core.jit_ext.interpreter_needs_wrap
def is_compiling():
return True

m = transformers.BertForSequenceClassification(transformers.BertConfig())
inp = torch.randint(1, 20, (1, 32))
jm = thunder.jit(m)
jm(inp)
```

We might submit an issue to transformers to check for is_tracing in `transformers.modeling_utils.PreTrainedModel.warn_if_padding_and_no_attention_mask`
but in general I wonder how to provide the such compatibility lookasides to users.

possible variants:
- A compat executor,
- detecting the module we are tracing and adding them just for transformers (too much magic),
- always having the `torch._dynamo.is_compiling`-lookaside?
...

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.