NotImplementedError: VJP for PrimIDs.ITEM is not implemented for ThunderFX and litgGPT models. #1479
Comments
IvanYashchuk
added
nemo
|
Here's a simple reproducer for the problem: import torch
import thunder
@thunder.jit
def f(x): return x.item()
a = torch.randn(1, requires_grad=True)
f(a)File ~/dev/lightning-thunder/thunder/core/trace_interpreter.py:63, in interpret_trace(trace, symbol_mapper, with_env, *args, **kwargs)
61 args = tree_map(read, symbol.args)
62 kwargs = tree_map(read, symbol.kwargs)
---> 63 prim_func = symbol_mapper(symbol) if symbol_mapper is not None else symbol.sym
64 if prim_func is None:
65 continue
File ~/dev/lightning-thunder/thunder/core/transforms.py:2541, in vjp_symbol_mapper(symbol, *args, **kwargs)
2539 if symbol.sym.id == "torch.nn.functional.dropout":
2540 return None
-> 2541 raise NotImplementedError(f"VJP for {symbol.sym.id} is not implemented")
2543 def _vjp_impl(*args, **kwargs):
2544 primals, kwargs = tree_map(lambda x: x.primal if isinstance(x, VJPDual) else x, (args, kwargs))
NotImplementedError: VJP for PrimIDs.ITEM is not implementedA "grad rule" is missing for |
This was referenced Nov 27, 2024
|
Just to let you know - this issue was discovered in canary runs of LitGPT models, I'm not sure if it's related to nemo, as the label suggest. |
|
To chip in - this problem also occurs for SFT benchmarks, and shadows this issue #1482 in newer pjnl container release (tested on pjnl-20241126). Breaks for phi-3 also. We don't have NeMo deps in the SFT script, so it's probably unrelated. |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
🐛 Bug
For Llama-3-8B, Mistral-7B-v0.1 and Phi-3-mini-4k-instruct we get the following error:
To Reproduce
Please use:
1 GPU (H100)
Image "INTERNAL_IMAGE:pjnl-20241125"
Training script:
Expected behavior
Environment
As in the pjnl image
cc @tfogal
The text was updated successfully, but these errors were encountered: