add example for fsdp

[JackCaoG]

@gkroiz FYI, for some reason I can't add you as a reviewer.

[alanwaketan]

Given FSDPv1 is more or less deprecated, I don't think we should add it as an example. We should just focus on SPMD instead.

[gkroiz]

Thanks for adding this @JackCaoG, I have a few questions.

1. Do we want to add a simple demonstration of checkpointing? I know there was some recent questions around this, namely Loading large sharded model weights #6989

2. How does this file's purpose differ to the FSDP example code in https://github.com/pytorch/xla/blob/master/test/test_train_mp_mnist_fsdp_with_ckpt.py

[JackCaoG]

Thanks for adding this @JackCaoG, I have a few questions.

1. Do we want to add a simple demonstration of checkpointing? I know there was some recent questions around this, namely Loading large sharded model weights #6989
2. How does this file's purpose differ to the FSDP example code in master/test/test_train_mp_mnist_fsdp_with_ckpt.py

Old example is kind of long, I want to make it very clear what's the minimal codes to run a feature. If user want to do a deeper dive we should directly them to the more completed example.

[JackCaoG]

Given FSDPv1 is more or less deprecated, I don't think we should add it as an example. We should just focus on SPMD instead.

I see your point, but there are users that want to use FSDP v1. It is similar that we don't really want user to use native DDP but I provide the example anyway. I think we can update our README to specify what do we recommened, in that section we can make it clear that we recommend xla_ddp and fsdp_v2.

[JackCaoG]

@gkroiz for checkponting and gradident checkpointing, I think we should provide a similar simple examples. Let me know if you are interested in contributing.

[gkroiz]

@gkroiz for checkponting and gradident checkpointing, I think we should provide a similar simple examples. Let me know if you are interested in contributing.

Happy to review, probably don't have time in the near future.

[Liyang90]

Why is the lambda functions needed? One case I can think of would be using it as a callable argument of auto_wrapper_callable. But since it is not the case, can we use FSDP() directly on the model for simplicity?

[JackCaoG]

good catch, I just copied it from the example. Let me clean it up

[gkroiz]

LGTM