Lightning-AI / Lightning-AI/lightning-thunder

NotImplementedError with no further explanation when trying to run Hugging Face implementation of Phi 3 model

Open
#2,038 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

huggingface interpreter
Dominant language
Python
Stars
1.5k
Forks
121
PR merge metrics
No merged PRs in 30d

Description

## 🐛 Bug

The following code snippet with transformers==4.50.3 raises a NotImplementedError in Thunder interpreter:
```py
import thunder, torch
from transformers import AutoConfig, AutoModelForCausalLM

model_id = "microsoft/Phi-3-mini-128k-instruct"
configuration = AutoConfig.from_pretrained(
model_id,
vocab_size=16,
pad_token_id=15,
max_position_embeddings=32,
num_hidden_layers=1,
)
configuration.hidden_size = configuration.num_attention_heads

with torch.device("cuda"):
model = AutoModelForCausalLM.from_config(configuration).to(torch.bfloat16)
compiled_model = thunder.jit(model)
input_ids = torch.randint(0, configuration.vocab_size, (1, configuration.max_position_embeddings), device="cuda")
compiled_output = compiled_model(input_ids=input_ids, labels=input_ids)
```
Traceback:
```py
NotImplementedError Traceback (most recent call last)
Cell In[1], line 19
17 compiled_model = thunder.jit(model)
18 input_ids = torch.randint(0, configuration.vocab_size, (1, configuration.max_position_embeddings), device="cuda")
---> 19 compiled_output = compiled_model(input_ids=input_ids, labels=input_ids)

File ~/dev/pytorch/main/torch/nn/modules/module.py:1749, in Module._wrapped_call_impl(self, *args, **kwargs)
1747 return self._compiled_call_impl(*args, **kwargs) # type: ignore[misc]
1748 else:
-> 1749 return self._call_impl(*args, **kwargs)

File ~/dev/pytorch/main/torch/nn/modules/module.py:1760, in Module._call_impl(self, *args, **kwargs)
1755 # If we don't have any hooks, we want to skip the rest of the logic in
1756 # this function, and just call forward.
1757 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
1758 or _global_backward_pre_hooks or _global_backward_hooks
1759 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1760 return forward_call(*args, **kwargs)
1762 result = None
1763 called_always_called_hooks = set()

File ~/dev/lightning-thunder/thunder/core/module.py:80, in ThunderModule.forward(self, *args, **kwargs)
79 def forward(self, *args, **kwargs):
---> 80 res = self._forward_fn(*args, **kwargs)
81 return res

File ~/dev/lightning-thunder/thunder/__init__.py:821, in jit..update_call_statistics..wrapped(*args, **kwargs)
819 cs.last_trace_host_start = time.perf_counter_ns()
820 try:
--> 821 return fn(*args, **kwargs)
822 finally:
823 cs.last_trace_host_stop = time.perf_counter_ns()

File ~/dev/lightning-thunder/thunder/__init__.py:861, in jit..fn_(*args, **kwargs)
858 _recursive_jit_call_warning()
859 return fn(*args, **kwargs)
--> 861 cache_entry, inps, pro_to_epi = get_computation_and_inputs(*args, **kwargs)
863 result = cache_entry.computation_fn(*inps)
864 result = maybe_connect_to_autograd(cache_entry, result)

File ~/dev/lightning-thunder/thunder/__init__.py:800, in jit..decorate_computation_function..wrapped(*args, **kwargs)
799 def wrapped(*args, **kwargs):
--> 800 cache_entry, inps, pro_to_epi = get_computation_and_inputs_fn(*args, **kwargs)
801 decorated_computation_fn = cache_entry.computation_fn
802 for decorator in decorators:

File ~/dev/lightning-thunder/thunder/core/langctxs.py:136, in langctx.__call__.._fn(*args, **kwargs)
134 try:
135 tok = set_langctx(self.langctx)
--> 136 result = fn(*args, **kwargs)
137 return result
138 finally:

File ~/dev/lightning-thunder/thunder/__init__.py:239, in _with_cache_info_ctx..cache_info_wrapper(*args, **kwargs)
237 tok = _cache_info_ctx.set({})
238 try:
--> 239 res = fn(*args, **kwargs)
240 finally:
241 _cache_info_ctx.reset(tok)

File ~/dev/lightning-thunder/thunder/__init__.py:765, in jit..get_computation_and_inputs(*args, **kwargs)
762 # Resets use of compile flags
763 cs.last_compile_reasons = defaultdict(list)
--> 765 prologue_trc, computation_trc, epilogue_trc = acquire_initial_trace(fn, args, kwargs, cd, cs, ad_hoc_executor)
766 cache_entry = apply_transforms_and_build_cache_entry(
767 cd, cs, cache_info, prologue_trc, computation_trc, epilogue_trc
768 )
770 if cd.cache_option is not CACHE_OPTIONS.NO_CACHING:

File ~/dev/lightning-thunder/thunder/__init__.py:444, in jit..acquire_initial_trace(fn, args, kwargs, cd, cs, ad_hoc_executor)
442 prologue_trc: TraceCtx
443 computation_trc: TraceCtx
--> 444 jit_results: TraceResults = thunder_general_jit(
445 fn,
446 args,
447 kwargs,
448 ad_hoc_executor=ad_hoc_executor,
449 sharp_edges=cd.sharp_edges,
450 )
451 prologue_trc = jit_results.prologue_trace
452 computation_trc = jit_results.computation_trace

File ~/dev/lightning-thunder/thunder/core/jit_ext.py:2059, in thunder_general_jit(fn, args, kwargs, sharp_edges, ad_hoc_executor)
2057 with jit_ctx(ctx):
2058 with tracectx(computation_trace):
-> 2059 result = jfn(*args, **kwargs)
2060 computation_trace.set_current_source_location(None, None)
2061 process_recorded_modifications(ctx, epilogue_trace)

File ~/dev/lightning-thunder/thunder/core/interpreter.py:7567, in interpret..fn_(*args, **kwargs)
7565 # The below is "raise e" but deleting e from the scope
7566 try:
-> 7567 raise e
7568 except Exception:
7569 del e

File ~/dev/lightning-thunder/thunder/core/interpreter.py:7526, in interpret..fn_..getfn..fn_2()
7525 def fn_2(args, kwargs):
-> 7526 return fn(*args, **kwargs)

File ~/dev/lightning-thunder/thunder/core/interpreter.py:6840, in _call_dispatch.._impl()
6839 def _impl(fn, *args, **kwargs):
-> 6840 return fn.__func__(fn.__self__, *args, **kwargs)

File ~/dev/pytorch/main/torch/nn/modules/module.py:1749, in Module._wrapped_call_impl()
1747 return self._compiled_call_impl(*args, **kwargs) # type: ignore[misc]
1748 else:
-> 1749 return self._call_impl(*args, **kwargs)

File ~/dev/lightning-thunder/thunder/core/interpreter.py:6840, in _call_dispatch.._impl()
6839 def _impl(fn, *args, **kwargs):
-> 6840 return fn.__func__(fn.__self__, *args, **kwargs)

File ~/dev/pytorch/main/torch/nn/modules/module.py:1760, in Module._call_impl()
1755 # If we don't have any hooks, we want to skip the rest of the logic in
1756 # this function, and just call forward.
1757 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
1758 or _global_backward_pre_hooks or _global_backward_hooks
1759 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1760 return forward_call(*args, **kwargs)
1762 result = None
1763 called_always_called_hooks = set()

File ~/dev/lightning-thunder/thunder/core/interpreter.py:6840, in _call_dispatch.._impl()
6839 def _impl(fn, *args, **kwargs):
-> 6840 return fn.__func__(fn.__self__, *args, **kwargs)

File ~/miniforge3/envs/pytorch-cuda-dev/lib/python3.10/site-packages/transformers/utils/deprecation.py:172, in deprecate_kwarg..wrapper..wrapped_func()
168 elif minimum_action in (Action.NOTIFY, Action.NOTIFY_ALWAYS) and not is_torchdynamo_compiling():
169 # DeprecationWarning is ignored by default, so we use FutureWarning instead
170 warnings.warn(message, FutureWarning, stacklevel=2)
--> 172 return func(*args, **kwargs)

File ~/miniforge3/envs/pytorch-cuda-dev/lib/python3.10/site-packages/transformers/models/phi3/modeling_phi3.py:917, in Phi3ForCausalLM.forward()
914 return_dict = return_dict if return_dict is not None else self.config.use_return_dict
916 # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
--> 917 outputs = self.model(
918 input_ids=input_ids,
919 attention_mask=attention_mask,
920 position_ids=position_ids,
921 past_key_values=past_key_values,
922 inputs_embeds=inputs_embeds,
923 use_cache=use_cache,
924 output_attentions=output_attentions,
925 output_hidden_states=output_hidden_states,
926 return_dict=return_dict,
927 cache_position=cache_position,
928 **kwargs,
929 )
931 hidden_states = outputs[0]
932 # Only compute necessary logits, and do not upcast them to float if we are not computing the loss

File ~/dev/lightning-thunder/thunder/core/interpreter.py:6840, in _call_dispatch.._impl()
6839 def _impl(fn, *args, **kwargs):
-> 6840 return fn.__func__(fn.__self__, *args, **kwargs)

File ~/dev/pytorch/main/torch/nn/modules/module.py:1749, in Module._wrapped_call_impl()
1747 return self._compiled_call_impl(*args, **kwargs) # type: ignore[misc]
1748 else:
-> 1749 return self._call_impl(*args, **kwargs)

File ~/dev/lightning-thunder/thunder/core/interpreter.py:6840, in _call_dispatch.._impl()
6839 def _impl(fn, *args, **kwargs):
-> 6840 return fn.__func__(fn.__self__, *args, **kwargs)

File ~/dev/pytorch/main/torch/nn/modules/module.py:1760, in Module._call_impl()
1755 # If we don't have any hooks, we want to skip the rest of the logic in
1756 # this function, and just call forward.
1757 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
1758 or _global_backward_pre_hooks or _global_backward_hooks
1759 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1760 return forward_call(*args, **kwargs)
1762 result = None
1763 called_always_called_hooks = set()

File ~/dev/lightning-thunder/thunder/core/interpreter.py:6840, in _call_dispatch.._impl()
6839 def _impl(fn, *args, **kwargs):
-> 6840 return fn.__func__(fn.__self__, *args, **kwargs)

File ~/miniforge3/envs/pytorch-cuda-dev/lib/python3.10/site-packages/transformers/models/phi3/modeling_phi3.py:618, in Phi3Model.forward()
615 hidden_states = inputs_embeds
617 # create position embeddings to be shared across the decoder layers
--> 618 position_embeddings = self.rotary_emb(hidden_states, position_ids)
620 # decoder layers
621 all_hidden_states = () if output_hidden_states else None

File ~/dev/lightning-thunder/thunder/core/interpreter.py:6840, in _call_dispatch.._impl()
6839 def _impl(fn, *args, **kwargs):
-> 6840 return fn.__func__(fn.__self__, *args, **kwargs)

File ~/dev/pytorch/main/torch/nn/modules/module.py:1749, in Module._wrapped_call_impl()
1747 return self._compiled_call_impl(*args, **kwargs) # type: ignore[misc]
1748 else:
-> 1749 return self._call_impl(*args, **kwargs)

File ~/dev/lightning-thunder/thunder/core/interpreter.py:6840, in _call_dispatch.._impl()
6839 def _impl(fn, *args, **kwargs):
-> 6840 return fn.__func__(fn.__self__, *args, **kwargs)

File ~/dev/pytorch/main/torch/nn/modules/module.py:1760, in Module._call_impl()
1755 # If we don't have any hooks, we want to skip the rest of the logic in
1756 # this function, and just call forward.
1757 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
1758 or _global_backward_pre_hooks or _global_backward_hooks
1759 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1760 return forward_call(*args, **kwargs)
1762 result = None
1763 called_always_called_hooks = set()

File ~/dev/lightning-thunder/thunder/core/interpreter.py:6840, in _call_dispatch.._impl()
6839 def _impl(fn, *args, **kwargs):
-> 6840 return fn.__func__(fn.__self__, *args, **kwargs)

File ~/dev/pytorch/main/torch/utils/_contextlib.py:116, in context_decorator..decorate_context()
113 @functools.wraps(func)
114 def decorate_context(*args, **kwargs):
115 with ctx_factory():
--> 116 return func(*args, **kwargs)

NotImplementedError:
```

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.

Research direction

Reproduce the snippet with transformers==4.50.3 and trace the failure from thunder/core/interpreter.py through thunder.jit into transformers/models/phi3/modeling_phi3.py, especially Phi3Model.forward and rotary_emb. Determine which operation produces the bare NotImplementedError and document the expected completed behavior or diagnostic, with a regression check for the Phi-3 example.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
compilers, machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.