Lightning-AI / Lightning-AI/lightning-thunder
NotImplementedError with no further explanation when trying to run Hugging Face implementation of Phi 3 model
Nobody has claimed this yet.
- 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
- 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.
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