Thomas Viehmann

improve shape accuracy in transform output by providing better tooling

5

Some transforms, notably FSDP and TensorParallel ones, change shapes, but currently do not completely update them (it does for the linear that follows, but not for the activation etc.). We...

Support PyTorch custom ops when they have meta functions (also autograd)

2

PyTorch allows to register meta functions with custom ops With that information, we should be able to apply the fallback mechanism that @kiya00 developed for PyTorch operations in thunder. https://pytorch.org/docs/stable/library.html...

provide trace checker and debug mode with it enabled

2

We are currently plagued a bit with inconsistency issues in traces and might have good use for a checker and then having a mode where we do this a lot:...

trace: add cursor to for bsym insertion

1

currently bsyms are appended to the trace. for transformations, we usually have a place in mind where these should go instead. the current way is to ``` trace.push_scope([]) out =...

quantization: process tensors on meta device directly, maybe implement CPU quantization (if it is easy)

4

Currently the BitsAndBytesLinearQuant4bit for submodule always calls `bitsandbytes.functional.quantize_4bit`. This is somewhat touchy for CPU tensors because `quantize_4bit` only works on GPU tensors but it is outright not so nice for...

enable using python_callable without mapping symbols to their impls

1

Currently, we can't use python_callable if we want to retrace because it does half of transform for execution. We should absolutely get a flag on python_callable for that because we...

Currently, the torch compile executors have their own way of creating callables operator by operator. It would be more principled to use transform for execution to do so, however when...

Make autocast a Transform

2

Currently autocast (`transforms/autocast.py`) is carried over from the old-style transforms on callable and invoked like ``` jit(autocast(fn)) ``` it would be good to make this a `AutocastTransform` to be invoked...

### Bug description there have been backward-incompatible changes to - Phi models (lm_head.bias) - LlamaRotaryEmbedding signature ### What operating system are you using? Unknown ### LitGPT Version ``` ```

### Bug description https://github.com/Lightning-AI/litgpt/pull/1796 pins the version to