do not require comm in non-nccl environment

examples/transformer-xl/topk.py

-#!/usr/bin/env python
-# encoding: utf-8
-# File Name: topk.py
-# Author: Jiezhong Qiu
-# Create Time: 2020/11/24 20:23
-# TODO:
-
-import torch
-import time
-from mem_transformer import my_topk
-output = torch.rand(16, 512, 512).cuda()
-
-torch.cuda.synchronize()
-start = time.time()
-_, pred = output.topk(k=1, dim=-1, largest=True, sorted=False)
-torch.cuda.synchronize()
-print("torch.top1 Time :{}".format(time.time() - start))
-
-torch.cuda.synchronize()
-start = time.time()
-_, pred_ = my_topk(output, k=1, inplace=True)
-torch.cuda.synchronize()
-print("my top1 Time :{}".format(time.time() - start))
-
-torch.cuda.synchronize()
-start = time.time()
-_, pred = output.topk(k=2, dim=-1, largest=True, sorted=False)
-torch.cuda.synchronize()
-print("torch.top2 Time :{}".format(time.time() - start))
-
-torch.cuda.synchronize()
-start = time.time()
-_, pred_ = my_topk(output, k=2, inplace=True)
-torch.cuda.synchronize()
-print("my top2 Time :{}".format(time.time() - start))
-

fmoe/functions.py

@@ -21,9 +21,9 @@ def moe_prepare_forward(gate, num_expert, world_size, comm=None):
         world_size: number of workers that hold different experts.
         comm: the communicator of all workers in the expert-parallel group.
     """
-    if comm is None:
-        comm = get_torch_default_comm()
     if world_size > 1:
+        if comm is None:
+            comm = get_torch_default_comm()
         fmoe_cuda.ensure_nccl(comm, gate)
 
     with torch.no_grad():