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