Add DS inference

Closes EleutherAI#845

megatron/model/transformer.py

@@ -840,7 +840,7 @@ def __init__(
         else:
             from torch import distributed as dist
 
-            if self.num_experts > dist.get_world_size():
+            if neox_args.ds_inference or self.num_experts > dist.get_world_size():
                 moe_mp_size = 1
             else:
                 moe_mp_size = dist.get_world_size() // self.num_experts

megatron/neox_arguments/neox_args.py

@@ -1155,6 +1155,16 @@ class NeoXArgsTextgen(NeoXArgsTemplate):
     prefix to which to save evaluation results - final fp will be {eval_results_prefix}_eval_results_yy-mm-dd-HH-MM.json
     """
 
+    ds_inference: bool = False
+    """
+    Use DeepSpeed inference.
+    """
+
+    moe_type: str = "standard"
+    """
+    Specify the type of MoE layer. We have two types of MoE layer: standard and residual.
+    """
+
     eval_tasks: list = None
     """
     Tasks to evaluate on using lm_eval_harness

megatron/utils.py

@@ -452,9 +452,30 @@ def setup_for_inference_or_eval(use_cache=True, overwrite_values=None, input_arg
     print_rank_0("Finished loading model")
 
     model.module.inference_mode(use_cache=use_cache)
+
+    if neox_args.ds_inference:
+        model = ds_inference(model, neox_args)
+        print("> DeepSpeed Inference engine initialized")
+
     return model, neox_args
 
 
+def ds_inference(model, neox_args):
+    import deepspeed
+
+    engine = deepspeed.init_inference(
+        model=model,
+        mp_size=neox_args.model_parallel_size,
+        tensor_parallel={"mpu": mpu},
+        dtype=torch.half,
+        replace_with_kernel_inject=True,
+        moe_experts=[neox_args.num_experts],
+        moe_type=neox_args.moe_type,
+    )
+
+    return engine.module
+
+
 class CharCounter:
     """
     Wraps the data_iterator to count the number of characters in a batch