addressed Jared and Patrick comments.

megatron/mpu/random.py

@@ -92,6 +92,11 @@ def gather_split_1d_tensor(tensor):
     gathered = torch.empty(numel_gathered, dtype=tensor.dtype,
                            device=torch.cuda.current_device(),
                            requires_grad=False)
+    # TODO: This API is experimental in pytorch (as of Feb 2022) and
+    # this might break in future pytorch releases. We chose this API
+    # as opposed to torch.distributed.all_gather for efficiency reasons.
+    # This API calls directly NCCL all-gather versus the former does
+    # internal copies and can potentially cause slow down.
     torch.distributed._all_gather_base(gathered, tensor,
                                        group=get_tensor_model_parallel_group())
     return gathered