search

Lightning-AI / lightning-thunder

Make PyTorch models up to 40% faster! Thunder is a source to source compiler for PyTorch. It enables using different hardware executors at once; across one or thousands of GPUs.
Apache License 2.0
1.08k stars 62 forks source link

Implement TensorBase.__setitem__ #341

Open athitten opened 2 months ago

athitten commented 2 months ago

🚀 Feature

Implement torch.Tensor.__setitem__

Motivation

NeMo's MegatronNeVaModel

cc @tfogal

riccardofelluga commented 1 week ago

Hi @athitten!

Would it be possible for you to share a small repro to run the NeVa model? I think this is the last operation that is needed to unlock GPT and NeVa from NeMo.

I see there are the examples in the NeMo repo but they don't seem to work even in the NeMo container :(

tfogal commented 1 week ago

Would it be possible for you to share a small repro to run the NeVa model? I think this is the last operation that is needed to unlock GPT and NeVa from NeMo.

Thanks for your interest! I pasted some ways to run this in #660.

@athitten going forward let's create a branch and include something similar to the instructions in #660 in the "root" issues (i.e. #343 in this case) so that other people can reproduce.