Add fp16 support for GPT2

Author: santhnm2

dist_ir/backend/torch.py

@@ -6,6 +6,7 @@
 from time import perf_counter
 from traceback import print_exc
 from typing import Any, Dict, Iterable, List, NamedTuple, Sequence, Tuple
+from warnings import warn
 
 import torch
 import torch.distributed as dist
@@ -238,8 +239,7 @@ def _slice(x, starts, ends, axes, steps=None, ctx=None):
 
 
 def _softmax(x, axis, ctx=None):
-    exp = torch.exp(x)
-    return exp / torch.sum(exp, dim=axis, keepdim=True)
+    return torch.nn.functional.softmax(x, dim=axis)
 
 
 def _split(x, axis, split, ctx=None):
@@ -417,8 +417,12 @@ def print_memory_usage():
             assert isinstance(output, tuple)
             for i, v in enumerate(op.outputs):
                 value_map[v] = output[i]
+                if torch.any(torch.isnan(output[i])):
+                    warn(f"NaNs in op {op} output {i}")
         elif len(op.outputs) == 1:
             value_map[op.outputs[0]] = output
+            if torch.any(torch.isnan(output)):
+                warn(f"NaNs in op {op.name} output {0}")
 
         # Free tensors that are not used again
         for v in op.inputs:

dist_ir/executor/numpy_register.py

@@ -1,4 +1,5 @@
 import numpy as np
+import scipy
 
 
 def _handle_negative_axis(axis, tensor_rank):
@@ -372,8 +373,7 @@ def slice_conc(op, x, starts, ends, axes, steps=None):
 
 def softmax(op, x):
     axis = op.attributes["axis"]
-    exp = np.exp(x)
-    return exp / np.sum(exp, axis=axis, keepdims=True)
+    return scipy.special.softmax(x, axis=axis)
 
 
 def softmax_grad(op, dY, Y):

examples/gpt2.py

@@ -13,8 +13,8 @@
     ConcreteValue,
 )
 from dist_ir.importer import import_from_onnx
-from dist_ir.ir import FunctionMaker, Op, get_uniform_topology
-from dist_ir.ir.type import Tensor, Type, abstract_values
+from dist_ir.ir import FunctionMaker, Op, Value, get_uniform_topology
+from dist_ir.ir.type import Float16, Tensor, Type, abstract_values
 from dist_ir.transforms import (
     gpt2_dhp_transform,
     sanitize_unhashable_attributes,
@@ -31,6 +31,42 @@ def _to_numpy(x):
     return x
 
 
+def _cast_to_fp16(function):
+    is_weight = lambda x: "weight" in x or "bias" in x
+
+    fp16_function = FunctionMaker(function.name)
+    value_map = {}
+    for i, inp in enumerate(function.inputs):
+        if is_weight(inp.name):
+            fp16_inp = fp16_function.add_input_value(
+                inp.name,
+                Tensor(
+                    shape=inp.type.shape,
+                    device=inp.type.device,
+                    dtype=Float16(device=inp.type.dtype.device),
+                ),
+            )
+        else:
+            fp16_inp = fp16_function.add_input_value(inp.name, inp.type)
+        value_map[inp] = fp16_inp
+    for op in function.ops:
+        inputs = [value_map[inp] for inp in op.inputs]
+        fp16_op = Op(
+            op_type=op.op_type,
+            name=op.name,
+            inputs=tuple(value_map[inp] for inp in op.inputs),
+            attributes=op.attributes,
+            subfunctions=op.subfunctions,
+            output_names=tuple(output.name for output in op.outputs),
+            output_types=tuple(None for output in op.outputs),
+            # output_types=tuple(output.type for output in op.outputs),
+        )
+        fp16_function.ops.append(fp16_op)
+        for output, fp16_output in zip(op.outputs, fp16_op.outputs):
+            value_map[output] = fp16_output
+    return fp16_function.finalize()
+
+
 def _filter_extra_outputs(function):
     function, attribute_map = sanitize_unhashable_attributes(function)
 
@@ -349,8 +385,11 @@ def _get_stats(function):
 def import_function_and_get_input_data(
     model_path,
     default_device,
+    dtype,
     use_real_weights=False,
 ):
+    is_input_or_weight = lambda x: "input" in x or "weight" in x or "bias" in x
+
     function, input_data_map = import_from_onnx(
         model_path,
         name="GPT-2",
@@ -360,10 +399,15 @@ def import_function_and_get_input_data(
 
     function = _filter_extra_outputs(function)
 
-    if not use_real_weights:
-        for inp in input_data_map:
-            if "input" in inp.name or "weight" in inp.name or "bias" in inp.name:
+    if dtype == "fp16":
+        function = _cast_to_fp16(function)
+
+    for inp in input_data_map:
+        if is_input_or_weight(inp.name):
+            if not use_real_weights:
                 input_data_map[inp] = inp.type
+            elif dtype == "fp16" and "input" not in inp.name:
+                input_data_map[inp] = input_data_map[inp].astype(np.float16)
     input_data = list(input_data_map.values())
 
     return function, input_data
@@ -488,6 +532,7 @@ def transform(
 
 def get_transformed_function_and_input_data(
     model_path,
+    dtype,
     device_throughput,
     dram_bandwidth,
     kernel_launch_overhead,
@@ -515,6 +560,7 @@ def get_transformed_function_and_input_data(
     function, input_data = import_function_and_get_input_data(
         model_path,
         default_device=topology.devices[0],
+        dtype=dtype,
         use_real_weights=use_real_weights,
     )
 
@@ -554,7 +600,9 @@ def simulate(function, input_data, topology, allreduce_parameters=None):
     return simulation
 
 
-def run_pytorch(function, input_data, world_size, use_gpu=True, debug_stacktrace=False):
+def run_pytorch(
+    function, input_data, world_size, use_gpu=False, debug_stacktrace=False
+):
     # TODO: Move this to a utils file
     def _resolve_dtype(dtype):
         if dtype == np.int32:
@@ -618,6 +666,7 @@ def main(args):
         topology,
     ) = get_transformed_function_and_input_data(
         args.model_path,
+        args.dtype,
         args.device_throughput,
         args.dram_bandwidth,
         args.kernel_launch_overhead,

examples/gpt2_grid_search.py

@@ -53,6 +53,7 @@ def __init__(
         self.base_model, self.base_input_data = gpt2.import_function_and_get_input_data(
             self.model_path,
             self.topology.devices[0],
+            self.dtype,
             use_real_weights=(self.backend == "pytorch"),
         )
         self.models_and_input_data = {}

examples/parser.py

@@ -90,7 +90,7 @@ def add_backend_config_arguments(self):
         self.add_argument(
             "--use_gpu",
             action="store_true",
-            default=torch.cuda.is_available(),
+            default=False,
             help="Use GPU with PyTorch backend",
         )
 

test/test_gpt2_dhp_transform.py

@@ -16,6 +16,7 @@
 
 
 def _run_gpt(
+    dtype="fp32",
     device_throughput=constants.DEFAULT_DEVICE_THROUGHPUT,
     dram_bandwidth=constants.DEFAULT_DRAM_BANDWIDTH,
     kernel_launch_overhead=constants.DEFAULT_KERNEL_LAUNCH_OVERHEAD,
@@ -38,6 +39,7 @@ def _run_gpt(
         topology,
     ) = get_transformed_function_and_input_data(
         MODEL_PATH,
+        dtype,
         device_throughput,
         dram_bandwidth,
         kernel_launch_overhead,
@@ -77,6 +79,7 @@ def _run_gpt(
 
 def _test(
     original_outputs,
+    dtype,
     dp_degree=1,
     hp_degree=1,
     pp_degree=1,
@@ -86,6 +89,7 @@ def _test(
 
     # Test with real weights
     transformed_outputs = _run_gpt(
+        dtype=dtype,
         dp_degree=dp_degree,
         hp_degree=hp_degree,
         pp_degree=pp_degree,
@@ -95,13 +99,18 @@ def _test(
     assert len(transformed_outputs) == dp_degree * hp_degree
     for i in range(len(transformed_outputs)):
         np.testing.assert_array_almost_equal(
-            original_outputs[0].val, transformed_outputs[i].val, decimal=2
+            original_outputs[0].val,
+            transformed_outputs[i].val,
+            decimal=(2 if dtype == "fp32" else 1),
         )
 
 
 @pytest.fixture(scope="session")
 def original_outputs():
-    return _run_gpt()
+    return {
+        "fp16": _run_gpt(dtype="fp16", use_pytorch_backend=True),
+        "fp32": _run_gpt(dtype="fp32", use_pytorch_backend=True),
+    }
 
 
 @pytest.mark.parametrize(
@@ -110,7 +119,8 @@ def original_outputs():
 )
 def test_reference_execution(original_outputs, dp_degree, hp_degree, pp_degree):
     _test(
-        original_outputs,
+        original_outputs["fp32"],
+        dtype="fp32",
         dp_degree=dp_degree,
         hp_degree=hp_degree,
         pp_degree=pp_degree,
@@ -119,12 +129,23 @@ def test_reference_execution(original_outputs, dp_degree, hp_degree, pp_degree):
 
 
 @pytest.mark.parametrize(
-    ("dp_degree", "hp_degree", "pp_degree"),
-    list(itertools.product([1, 2], [1, 2], [1, 2])),
+    ("dtype", "dp_degree", "hp_degree", "pp_degree"),
+    list(
+        itertools.product(
+            ["fp16", "fp32"] if torch.cuda.is_available() else ["fp32"],
+            [1, 2],
+            [1, 2],
+            [1, 2],
+        )
+    ),
 )
-def test_pytorch_backend(original_outputs, dp_degree, hp_degree, pp_degree):
+def test_pytorch_backend(original_outputs, dtype, dp_degree, hp_degree, pp_degree):
+    world_size = dp_degree * hp_degree * pp_degree
+    if dtype == "fp16" and world_size > torch.cuda.device_count():
+        pytest.skip("Not enough GPUs available")
     _test(
-        original_outputs,
+        original_outputs[dtype],
+        dtype,
         dp_degree=dp_degree,
         hp_degree=hp_degree,
         pp_degree=pp_degree,
@@ -134,14 +155,24 @@ def test_pytorch_backend(original_outputs, dp_degree, hp_degree, pp_degree):
 
 
 @pytest.mark.parametrize(
-    ("dp_degree", "hp_degree", "pp_degree"),
-    list(itertools.product([1, 2], [1, 2], [1, 2])),
+    ("dtype", "dp_degree", "hp_degree", "pp_degree"),
+    list(itertools.product(["fp16", "fp32"], [1, 2], [1, 2], [1, 2])),
 )
-def test_mixed_simulation(dp_degree, hp_degree, pp_degree):
+def test_mixed_simulation(dtype, dp_degree, hp_degree, pp_degree):
     _run_gpt(
+        dtype=dtype,
         dp_degree=dp_degree,
         hp_degree=hp_degree,
         pp_degree=pp_degree,
         num_microbatches=pp_degree,
         use_real_weights=False,
     )
+
+if __name__=="__main__":
+    original_outputs = {
+        "fp16": _run_gpt(dtype="fp16", use_pytorch_backend=True),
+        "fp32": _run_gpt(dtype="fp32", use_pytorch_backend=True),
+    }
+    for dtype, dp_degree, hp_degree, pp_degree in list(itertools.product(["fp16", "fp32"], [1, 2], [1, 2], [1, 2])):
+        print(f"dtype={dtype}, dp_degree={dp_degree}, hp_degree={hp_degree}, pp_degree={pp_degree}")
+        test_pytorch_backend(original_outputs, dtype, dp_degree, hp_degree, pp_degree)

test/test_grid_search.py

@@ -156,6 +156,7 @@ def test_gpt_grid_search(backend, dtype):
                     topology,
                 ) = gpt2.get_transformed_function_and_input_data(
                     model_path=GPT2_MODEL_PATH,
+                    dtype=dtype,
                     device_throughput=constants.DEFAULT_DEVICE_THROUGHPUT,
                     dram_bandwidth=constants.DEFAULT_DRAM_BANDWIDTH,
                     kernel_launch_overhead=constants.DEFAULT_KERNEL_LAUNCH_OVERHEAD,