Would like to contribute FSDP functionality

[calmitchell617]

Hi @karpathy,

I recently modified this repo to shard everything with PyTorch FSDP and would like to contribute this functionality to the project.

If I cleaned up the code and submitted a PR, is that something you would have a look at?

[karpathy]

It's def on my todo list to incorporate FSDP into nanoGPT but I haven't looked into it in detail just yet. I also know that FSDP internals are being actively iterated on and might move around soon, so I'm waiting that out a little bit.

[calmitchell617]

Ok, should I bother putting together a PR?