Disable non-functioning torch.compile #17
Conversation
| # compilation fails as it does not support torch.complex64 for RoPE | ||
| # compile = False |
There was a problem hiding this comment.
Ah this sucks :((
We could explore this in the future: Add a flag to our nano model that allows us to switch between the complex implementation and the previous real one. We would use the flag for comparision and inference using the meta checkpoints, but we would use the real real implementation when training from scratch.
There was a problem hiding this comment.
We still want to be able to compile for training, as that's where the speedup becomes more valuable (less $$$ spent).
Maybe we should look into getting the non-complex implementation to work as expected.
There was a problem hiding this comment.
That's what I'm explaining above. We would compile the non-complex version for training, but we would use the other for inference where we load the meta checkpoint.
Maybe we should look into getting the non-complex implementation to work as expected.
I spent the whole weekend on this so yeah, it's not like we didn't try already xD
| @@ -5,7 +5,7 @@ | |||
| import lightning as L | |||
|
|
|||
|
|
|||
| @torch.inference_mode() | |||
| @torch.no_grad() | |||
There was a problem hiding this comment.
No description provided.