added unit test for megatron (deepspeedai#102)

Author: mrwyattii

MANIFEST.in

@@ -1,2 +1,3 @@
 include megatron/data/Makefile
 include megatron/data/helpers.cpp
+recursive-include megatron/fused_kernels *.cpp *.h *.cu *.tr *.cuh *.cc

tests/run_megatron.py

@@ -0,0 +1,113 @@
+import torch
+import deepspeed
+import megatron
+from megatron import get_args
+from megatron import mpu
+from megatron.checkpointing import load_checkpoint
+from megatron.initialize import initialize_megatron
+from megatron.model import GPTModel
+from megatron.training import get_model
+from megatron.text_generation_utils import generate_samples_eval
+
+
+def model_provider(pre_process=True, post_process=True):
+    model = GPTModel(
+        num_tokentypes=0,
+        parallel_output=False,
+        pre_process=pre_process,
+        post_process=post_process,
+        return_moe_loss=False,
+    )
+    return model
+
+
+def add_text_generate_args(parser):
+    """Text generation arguments."""
+    group = parser.add_argument_group(title="text generation")
+
+    group.add_argument(
+        "--temperature", type=float, default=1.0, help="Sampling temperature."
+    )
+    group.add_argument(
+        "--greedy", action="store_true", default=False, help="Use greedy sampling."
+    )
+    group.add_argument("--top_p", type=float, default=0.0, help="Top p sampling.")
+    group.add_argument("--top_k", type=int, default=0, help="Top k sampling.")
+    group.add_argument(
+        "--out-seq-length",
+        type=int,
+        default=1024,
+        help="Size of the output generated text.",
+    )
+    group.add_argument(
+        "--sample-input-file",
+        type=str,
+        default=None,
+        help="Get input from file instead of interactive mode, "
+        "each line is an input.",
+    )
+    group.add_argument(
+        "--sample-output-file",
+        type=str,
+        default=None,
+        help="Output file got from --sample-input-file",
+    )
+    group.add_argument(
+        "--num-samples",
+        type=int,
+        default=0,
+        help="Number of samples to generate unconditionally, "
+        "defaults to 0 and interactive conditional sampling",
+    )
+    group.add_argument(
+        "--genfile", type=str, help="Output file when generating unconditionally"
+    )
+    group.add_argument(
+        "--recompute",
+        action="store_true",
+        help="During generation recompute all attention "
+        "instead of using previously computed keys/values.",
+    )
+    group.add_argument(
+        "--context-tokens", type=str, default="DeepSpeed is the greatest"
+    )
+    group.add_argument("--max-tokens", type=int, default=50)
+
+    return parser
+
+
+if __name__ == "__main__":
+    # initialize megatron
+    initialize_megatron(
+        extra_args_provider=add_text_generate_args,
+        args_defaults={
+            "tokenizer_type": "GPT2BPETokenizer",
+            "no_load_rng": True,
+            "no_load_optim": True,
+        },
+    )
+    args = get_args()
+
+    # setup model
+    model = get_model(model_provider)
+    _ = load_checkpoint(model, None, None)
+    model = model[0]
+    if args.ds_inference:
+        engine = deepspeed.init_inference(
+            model=model,
+            mp_size=args.tensor_model_parallel_size,
+            tensor_parallel={"mpu": mpu},
+            dtype=torch.half,
+            replace_with_kernel_inject=True,
+            moe_experts=args.num_experts,
+            moe_type=args.mlp_type,
+        )
+        model = engine.module
+
+    # generate output
+    generate_samples_eval(
+        model, args.context_tokens, 1, 0
+    )  # Just so we don't get log output from DeepSpeed (this should be removed once we improve logging in DeepSpeed)
+    print("===START OUTPUT===")
+    print(generate_samples_eval(model, args.context_tokens, args.max_tokens, 0))
+    print("===END OUTPUT===")

tests/test_megatron.py

@@ -0,0 +1,61 @@
+import pytest
+import os
+import re
+import subprocess
+
+
+@pytest.fixture(params=[1])
+def moe_num_experts(request):
+    return str(request.param)
+
+
+@pytest.fixture(params=[1])
+def mp_size(request):
+    return str(request.param)
+
+
+@pytest.fixture
+def params(moe_num_experts, mp_size):
+    base_dir = os.getenv("MEGATRON_CKPT_DIR")
+    assert base_dir, "Please set MEGATRON_CKPT_DIR in your environment"
+
+    vocab_file = os.path.join(base_dir, "gpt2-vocab.json")
+    merge_file = os.path.join(base_dir, "gpt2-merges.txt")
+    ckpt_path = os.path.join(base_dir, "checkpoints/gpt2_345m")
+
+    return [
+        "--micro-batch-size", "1",
+        "--num-layers", "24",
+        "--hidden-size", "1024",
+        "--num-attention-heads", "16",
+        "--max-position-embeddings", "1024",
+        "--vocab-file", vocab_file,
+        "--merge-file", merge_file,
+        "--load", ckpt_path,
+        "--seq-length", "1024",
+        "--out-seq-length", "1024",
+        "--tensor-model-parallel-size", mp_size,
+        "--tokenizer-type", "GPT2BPETokenizer",
+        "--num-experts", moe_num_experts,
+        "--mlp-type", "standard",
+        "--num-samples", "0",
+        "--fp16",
+    ]
+
+
+def test_moe_megatron(params, mp_size):
+    output_re = r"===START OUTPUT===([\S\s]*)===END OUTPUT==="
+
+    # Run the baseline
+    baseline_cmd = ["deepspeed", "--num_gpus", mp_size, "./run_megatron.py"] + params
+    result = subprocess.run(baseline_cmd, stdout=subprocess.PIPE)
+    baseline_output = re.search(output_re, result.stdout.decode("utf-8")).group(1)
+
+    # Run with DeepSpeed
+    deepspeed_cmd = baseline_cmd + ["--ds-inference"]
+    result = subprocess.run(deepspeed_cmd, stdout=subprocess.PIPE)
+    deepspeed_output = re.search(output_re, result.stdout.decode("utf-8")).group(1)
+
+    assert (
+        baseline_output == deepspeed_output
+    ), f"outputs do not match: {baseline_output}\n{deepspeed_output}"