[dtensor] add checkpointing example (#94743)
This PR adds some DTensor sharding example on a simple MLP model
for checkpointing reference purposes
Note that checkpointing itself is not implemented yet.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/94743
Approved by: https://github.com/wz337