Lightning-AI / Lightning-AI/lightning-thunder

Dynamic shape needs to be modeled in trace

Open
#471 0 comments 0 reactions 1 assignee View on GitHub

@jjsjann123 is already working on this.

Since Jun 3, 2024.

bug dynamic constraints
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

  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.

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.