feat(examples): add FSDP

[tpoisonooo]

Content

Add an FSDP example for Muon optimizer.

Here is the difference:

1. Add DistributedSampler

    sampler = DistributedSampler(dataset=train_dataset,
                                 rank=rank,
                                 num_replicas=world_size,
                                 shuffle=True)

2. Add all reduce

dist.all_reduce(loss, op=dist.ReduceOp.AVG)

Test

loss:

GPU usage:

I am also a FSDP begginer, for any question, I would try to fix. #25

[tpoisonooo]

cc @toothacher17

[tpoisonooo]

Furthermore, toy_train_fsdp.py is compatible with toy_train.py, for single GPU, just:

export CUDA_VISIBLE_DEVICES="0"
python3 toy_train_fsdp.py

They are same.