Lightning-AI / Lightning-AI/lightning-thunder
Partial function is not supported in `grad_transform`
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
## 🚀 Feature
Hitting this assert below vvv
```py
root@847841b8737c:/opt/pytorch/lightning-thunder# python /volume/pooling.py
Traceback (most recent call last):
File "/volume/pooling.py", line 36, in
o = jit_model(image)
File "/opt/pytorch/lightning-thunder/thunder/__init__.py", line 632, in fn_
cache_entry, inps, pro_to_epi = get_computation_and_inputs(*args, **kwargs)
File "/opt/pytorch/lightning-thunder/thunder/__init__.py", line 265, in cache_info_wrapper
res = fn(*args, **kwargs)
File "/opt/pytorch/lightning-thunder/thunder/__init__.py", line 574, in get_computation_and_inputs
computation_trc, backward_trc = split_forward_backward(computation_trc, cd, cs, *inps)
File "/opt/pytorch/lightning-thunder/thunder/executors/torch_autograd.py", line 216, in split_forward_backward
fw_trace, bw_trace = forward_and_backward_from_trace(primal_trace, torch_autograd=True)
File "/opt/pytorch/lightning-thunder/thunder/core/transforms.py", line 3879, in forward_and_backward_from_trace
forward_trace = construct_trace()(augmented_forward_fn, *trace.args, **trace.kwargs)
File "/opt/pytorch/lightning-thunder/thunder/core/interpreter.py", line 1292, in fn_
return fn(*args, **kwargs)
File "/opt/pytorch/lightning-thunder/thunder/common.py", line 528, in _trace
result = fn(*proxyargs, **proxykwargs)
File "/opt/pytorch/lightning-thunder/thunder/core/transforms.py", line 3850, in augmented_forward_fn
result, env = augmented_forward_pass(*args, trace=trace, **kwargs)
File "/opt/pytorch/lightning-thunder/thunder/core/transforms.py", line 3461, in augmented_forward_pass
result, env = eval_trace(
File "/opt/pytorch/lightning-thunder/thunder/core/transforms.py", line 1698, in eval_trace
prim_func = symbol_mapper(symbol)
File "/opt/pytorch/lightning-thunder/thunder/core/transforms.py", line 3385, in vjp_symbol_mapper
vjp_impl, backward_fn = make_aug_forward_and_backward(symbol)
File "/opt/pytorch/lightning-thunder/thunder/core/vjp_utils.py", line 63, in make_aug_forward_and_backward
joint_trace = thunder.trace(inline_trace=False, use_dce=False)(joint_forward_backward, *bsym.args, **bsym.kwargs)
File "/opt/pytorch/lightning-thunder/thunder/core/interpreter.py", line 1292, in fn_
return fn(*args, **kwargs)
File "/opt/pytorch/lightning-thunder/thunder/common.py", line 506, in _trace
proxyargs, proxykwargs = _unpack_inputs(fn, trace, args, kwargs, rename_proxies=rename_proxies)
File "/opt/pytorch/lightning-thunder/thunder/common.py", line 273, in _unpack_inputs
si = get_siginfo(fn, args, kwargs)
File "/opt/pytorch/lightning-thunder/thunder/core/codeutils.py", line 313, in get_siginfo
check(
File "/opt/pytorch/lightning-thunder/thunder/core/baseutils.py", line 103, in check
raise exception_type(s())
NotImplementedError: Support for partials with positional args (like ('test',)) is not implemented yet
```
I was trying to use something like
```py
foo = partial(bar, pos_arg0)
OperatorExecutor.register_operator(..., grad_transform=fn)
```
This isn't a high priority issue, since we can easily work around it for now. Filing the issue just to keep track of missing feature.
Contributor guide
No contributing guide indexed for this repository
Research direction
Start in thunder/core/codeutils.py at get_siginfo, where the traceback reports positional partial arguments as unsupported. Trace how this is reached from thunder/core/vjp_utils.py through make_aug_forward_and_backward, reproducing the case with the reported pooling.py script. Done means grad_transform accepts a partial with positional arguments without raising this NotImplementedError.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- compilers, machine-learning
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100