Lightning-AI / Lightning-AI/lightning-thunder
Dynamic shape needs to be modeled in trace
@jjsjann123 is already working on this.
Since Jun 3, 2024.
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
## 🐛 Bug
dynamic shape computation inside prim_meta is opaque in trace. (see [comment](https://github.com/Lightning-AI/lightning-thunder/issues/262#issuecomment-2127729530))
This would give us a `TensorProxy` with a shape which contains magic `NumberProxy` without its operations.
e.g. the output TensorProxy below has shape `[8, 16, i42]`. However, the trace does not model how i42 is produced.
`# t43 = prims.slice_prim(t_0, [0, 0, i9], [8, 16, 32], [1, 1, 1]) # t43: "cuda:0 f32[8, 16, [IntegerProxy name=i42, value=22]]"`
Currently this triggers an assert in unpack_inputs during prologue trace construction. commit 775aa7405b8f0e9539326a069522d247c6363cd8
```py
import thunder
def foo(a, dim0, dim1):
dim2 = a.size(2) - dim0 - dim1
x = torch.split(a, (dim0, dim1, dim2), dim=2)
dim3 = x[2].size(2) + dim2
return x, dim3
import torch
dtype = torch.float32
size0 = 8
size1 = 16
size2 = 32
a = torch.randn(size0, size1, size2, device="cuda").to(dtype=dtype)
dim0=2
dim1=8
a_ref = a.detach()
out_ref = foo(a_ref, dim0, dim1)
jfoo = thunder.jit(foo)
out = jfoo(a, dim0, dim1)
```
hits an error:
```
Traceback (most recent call last):
File "/volume/thunder_scalar.py", line 35, in
out = jfoo(a, dim0, dim1)
File "/opt/pytorch/lightning-thunder/thunder/__init__.py", line 617, in fn_
cache_entry, inps, pro_to_epi = get_computation_and_inputs(*args, **kwargs)
File "/opt/pytorch/lightning-thunder/thunder/__init__.py", line 202, in cache_info_wrapper
res = fn(*args, **kwargs)
File "/opt/pytorch/lightning-thunder/thunder/__init__.py", line 473, in get_computation_and_inputs
jit_results: TraceResults = interpreter(
File "/opt/pytorch/lightning-thunder/thunder/__init__.py", line 190, in _general_frontend
return thunder_general_jit(fn, args, kwargs, sharp_edges=sharp_edges, record_history=record_history)
File "/opt/pytorch/lightning-thunder/thunder/core/jit_ext.py", line 1598, in thunder_general_jit
pro_to_comp_proxies, pro_to_epi_proxies = unpack_inputs(
File "/opt/pytorch/lightning-thunder/thunder/core/jit_ext.py", line 1401, in unpack_inputs
pro_to_comp = tuple(sorted((unpack(v) for v in pro_to_comp_inps), key=lambda x: param_ordering[id(x)][1]))
File "/opt/pytorch/lightning-thunder/thunder/core/jit_ext.py", line 1401, in
pro_to_comp = tuple(sorted((unpack(v) for v in pro_to_comp_inps), key=lambda x: param_ordering[id(x)][1]))
File "/opt/pytorch/lightning-thunder/thunder/core/jit_ext.py", line 1189, in unpack
assert p.history is not None, f"{p} has history None"
AssertionError: [IntegerProxy name=i42, value=22] has history None
```
### Expected behavior
To handle the problem, we can have shape inference for outputs explicit in traces via one of the two approaches:
- lift all arithmetic operations on NumerProxy related to shape from `prims` to higher level like `clang`, or
- trace inside primitive?! which I feel will be messy, but I'm open to suggestions.
### Additional context
Meta operations that produces `TensorProxy(..., shape=xxx, ...)`
Operations that could produce new `NumberProxy` in shape in meta ops:
- [ ] cat_meta
- [ ] pad_meta
- [ ] slice_meta
- [ ] transpose_meta
- [ ] unfold_meta
- [ ] _reduction_meta
- [ ] _argmin_argmax_meta
- [ ] convolution_meta (double check this)
no new NumberProxy created:
- broadcast_in_dim_meta
- reshape_meta
- squeeze_meta
- take_meta
- take_along_axis_meta
- gather_meta
- topk_meta
- matmul_meta (conditional on ndim is considered static)
- embedding_meta
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.