[TorchTitan] Add destory process group at the end of training (#223)

train.py

@@ -15,6 +15,7 @@
 
 import torch
 import torch.nn.functional as F
+from torch.distributed import destroy_process_group
 from torch.distributed.elastic.multiprocessing.errors import record
 from torch.distributed.tensor.parallel import loss_parallel
 
@@ -387,6 +388,8 @@ def loss_fn(pred, labels):
                 )
 
     metric_logger.close()
+    logger.info("Training completed.")
+    destroy_process_group()
 
 
 if __name__ == "__main__":