Lightning-AI / Lightning-AI/lightning-thunder
how to provide overrides needed for model compatibility
@tfogal is already working on this.
Since Jul 22, 2024.
- 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
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Assessment
This issue has not been assessed yet.