fix megatron adapter for swipe

fmoe/gates/__init__.py

@@ -7,3 +7,5 @@ from .noisy_gate import NoisyGate
 
 from .gshard_gate import GShardGate
 from .switch_gate import SwitchGate
+
+from .swipe_gate import SwipeGate

fmoe/gates/base_gate.py

@@ -23,3 +23,7 @@ class BaseGate(nn.Module):
         if clear:
             self.loss = None
         return loss
+
+    @property
+    def has_loss(self):
+        return self.loss is not None

fmoe/megatron/balance.py

@@ -51,9 +51,12 @@ def add_balance_log(model, writer, iteration):
     while hasattr(model, 'module'):
         model = model.module
 
-    balance_dict_tensor = torch.vstack(
-        [l.mlp.gate.get_loss(clear=True) for l in model.language_model.transformer.layers]
-    ).detach()
+    losses = [l.mlp.gate.get_loss(clear=True)
+            for l in model.language_model.transformer.layers
+            if l.mlp.gate.has_loss]
+    if len(losses) == 0:
+        return
+    balance_dict_tensor = torch.vstack(losses).detach()
     world_group = get_torch_default_comm()
     world_size = torch.distributed.get_world_size(group=world_group)
     torch.distributed.all_reduce(balance_dict_tensor, group=world_group)

fmoe/megatron/layers.py

@@ -95,6 +95,9 @@ class MegatronMLP(FMoETransformerMLP):
         elif args.balance_strategy == "switch":
             from fmoe.gates import SwitchGate
             gate = SwitchGate
+        elif args.balance_strategy == "swipe":
+            from fmoe.gates import SwipeGate
+            gate = SwipeGate
         elif gate is None:
             assert False, "Undefined balance strategy {}" % (args.balance_strategy)
 

fmoe/megatron/patch.py

@@ -20,15 +20,19 @@ def patch_forward_step(forward_step_func):
         args = get_args()
         output = forward_step_func(data_iterator, model, input_tensor)
 
-        if not is_pipeline_last_stage() or not args.balance_strategy or args.balance_strategy == 'naive':
+        if not is_pipeline_last_stage() or not args.balance_strategy:
             return output
-        loss_name = args.balance_strategy + "_loss"
 
         while hasattr(model, 'module'):
             model = model.module
 
         loss_list = [l.mlp.gate.get_loss(clear=False).view(1)
-                for l in model.language_model.transformer.layers]
+                for l in model.language_model.transformer.layers
+                if l.mlp.gate.has_loss]
+        if len(loss_list) == 0:
+            return output
+
+        loss_name = args.balance_strategy + "_loss"
         (loss, state_dict), bal_loss = (
             output,
             torch.cat(loss_list).mean() * args.balance_loss_weight