[PT-D][Checkpoint] Update import and update docstring for distributed checkpoint (#89256)
Update test import and docstring as we have moved distributed checkpointing from torch.distributed._shard.checkpoint to torch.distributed.checkpoint (https://github.com/pytorch/pytorch/pull/88698).
Test: CI
Pull Request resolved: https://github.com/pytorch/pytorch/pull/89256
Approved by: https://github.com/fduwjj
