make last step save sharded ckpt as well

Author: lchu6

fms_fsdp/utils/train_utils.py

@@ -166,7 +166,7 @@ def train(
             ddp_stats.zero_()
         torch.cuda.reset_peak_memory_stats(device=torch.cuda.current_device())
 
-        if batch_idx % cfg.checkpoint_interval == 0:
+        if batch_idx % cfg.checkpoint_interval == 0 or batch_idx == cfg.num_steps:
             checkpointer.save(
                 batch_idx,
                 model,