Add fp16 support

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
@@ -14,7 +15,7 @@
 from ..executor.rank_projector import project
 from ..ir import Function, cpprint
 from ..ir.device import Device
-from ..ir.type import Int32, Int64, Float32, Type
+from ..ir.type import Int32, Int64, Float16, Float32, Type
 
 # NOTE: The code currently suffers from this issue, more investigation needed:
 # https://github.com/pytorch/pytorch/issues/11201
@@ -166,6 +167,8 @@ def _recv(shape=None, from_d=None, group=None, dtype=None, ctx=None):
         x = torch.zeros(shape).int()
     elif isinstance(dtype, Int64):
         x = torch.zeros(shape).long()
+    elif isinstance(dtype, Float16):
+        x = torch.zeros(shape).half()
     elif isinstance(dtype, Float32):
         x = torch.zeros(shape).float()
     else:
@@ -236,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):
@@ -415,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/calibrate_simulator.py

@@ -6,44 +6,46 @@
 from tqdm import tqdm
 import pandas as pd
 
-from dist_ir.ir import FunctionMaker, Topology, cpprint
-from dist_ir.ir.type import Device, Float32, Tensor
-from dist_ir.backend.torch import run_pytorch
+from ..ir import FunctionMaker, Topology, cpprint
+from ..ir.type import Device, Float16, Float32, Tensor
+from ..backend.torch import run_pytorch
 from .type_inference import infer_types
 from .cost_model import CostModel
 from .simulator import Simulator
 
 BYTES_IN_Gb = 1.25e8
+NUM_WARMUP = 5
+NUM_REPETITIONS = 30
 
 
-def _matmul(batch_size, input_dim, output_dim, device):
+def _matmul(batch_size, input_dim, output_dim, device, dtype):
     fn = FunctionMaker(name="matmul")
     x = fn.add_input_value(
         "x",
-        Tensor(shape=(batch_size, input_dim), dtype=Float32(), device=device),
+        Tensor(shape=(batch_size, input_dim), dtype=dtype(), device=device),
     )
     w = fn.add_input_value(
         "w",
-        Tensor(shape=(input_dim, output_dim), dtype=Float32(), device=device),
+        Tensor(shape=(input_dim, output_dim), dtype=dtype(), device=device),
     )
     y = fn.add_op(op_type="MatMul", inputs=[x, w], output_names=["y"])
     return fn.finalize()
 
 
-def _send(src, dst, m=1024, n=1024):
+def _send(src, dst, dtype, m=1024, n=1024):
     fn = FunctionMaker(name=f"send_{src.device_id}_to_{dst.device_id}")
-    x = fn.add_input_value("x", Tensor(shape=(m, n), dtype=Float32(), device=src))
+    x = fn.add_input_value("x", Tensor(shape=(m, n), dtype=dtype(), device=src))
     y = fn.add_op(
         op_type="Send", inputs=[x], attributes={"device": dst}, output_names=["y"]
     )
     return fn.finalize()
 
 
-def _allreduce(devices, m=1024, n=1024):
+def _allreduce(devices, dtype, m=1024, n=1024):
     fn = FunctionMaker(name=f"allreduce")
     xs = [
         fn.add_input_value(
-            f"x{i}", Tensor(shape=(m, n), dtype=Float32(), device=devices[i])
+            f"x{i}", Tensor(shape=(m, n), dtype=dtype(), device=devices[i])
         )
         for i in range(len(devices))
     ]
@@ -93,8 +95,8 @@ def network_bandwidth_debug():
                     for i in range(len(fn.inputs))
                 ],
                 use_gpu=True,
-                num_repetitions=10,
-                num_warmup=5,
+                num_repetitions=NUM_REPETITIONS,
+                num_warmup=NUM_WARMUP,
             )
             real_latency = np.median(runtimes[0])
             ex = Simulator(CostModel(topology))
@@ -129,7 +131,9 @@ def network_bandwidth_debug():
     print(df)
 
 
-def calibrate_network_bandwidth():
+def calibrate_network_bandwidth(dtype):
+    dist_ir_dtype = Float32 if dtype == "fp32" else Float16
+    pytorch_dtype = torch.float32 if dtype == "fp32" else torch.float16
     bandwidths = []
     all_sizes = [1024, 2048, 4096, 8192]
     n = len(all_sizes)
@@ -144,20 +148,20 @@ def calibrate_network_bandwidth():
             if src == dst:
                 continue
             for i, size in enumerate(tqdm(all_sizes)):
-                fn = _send(src, dst, m=size, n=size)
+                fn = _send(src, dst, dist_ir_dtype, m=size, n=size)
                 fn = infer_types(fn, fn.inputs)
                 X[i][0] = fn.inputs[0].type.size() / BYTES_IN_Gb
                 X[i][1] = 1
 
                 _, runtimes = run_pytorch(
                     fn=fn,
                     inputs=[
-                        torch.randn(size=fn.inputs[i].type.shape, dtype=torch.float32)
+                        torch.randn(size=fn.inputs[i].type.shape, dtype=pytorch_dtype)
                         for i in range(len(fn.inputs))
                     ],
                     use_gpu=True,
-                    num_repetitions=10,
-                    num_warmup=5,
+                    num_repetitions=NUM_REPETITIONS,
+                    num_warmup=NUM_WARMUP,
                 )
                 pytorch_latency = np.median(runtimes[0])
                 Y[i] = pytorch_latency
@@ -170,18 +174,25 @@ def calibrate_network_bandwidth():
     return bandwidths
 
 
-def calibrate_device_parameters():
-    all_batch_sizes = [1024, 2048, 4096]
-    all_input_dims = [1024, 2048, 4096]
-    all_output_dims = [1024, 2048, 4096]
+def calibrate_device_parameters(dtype):
+    dist_ir_dtype = Float32 if dtype == "fp32" else Float16
+    pytorch_dtype = torch.float32 if dtype == "fp32" else torch.float16
+    all_batch_sizes = [2 ** i for i in range(14, 16)]
+    all_input_dims = [2 ** i for i in range(14, 16)]
+    all_output_dims = [2 ** i for i in range(14, 16)]
+    if dtype == "fp16":
+        all_batch_sizes = [2 * v for v in all_batch_sizes]
+        all_input_dims = [2 * v for v in all_input_dims]
+        all_output_dims = [2 * v for v in all_output_dims]
     n = len(all_batch_sizes) * len(all_input_dims) * len(all_output_dims)
     X = np.zeros(shape=(n, 3))
     Y = np.zeros(shape=(n,))
+    data = []
     device = Device(0, "gpu")
     for i, (batch_size, input_dim, output_dim) in enumerate(
         tqdm(list(itertools.product(all_batch_sizes, all_input_dims, all_output_dims)))
     ):
-        fn = _matmul(batch_size, input_dim, output_dim, device)
+        fn = _matmul(batch_size, input_dim, output_dim, device, dist_ir_dtype)
         x = fn.inputs[0].type
         y = fn.inputs[1].type
         data_size = x.dtype.size() * (x.shape[0] * x.shape[1] + y.shape[0] * y.shape[1])
@@ -193,21 +204,37 @@ def calibrate_device_parameters():
         _, runtimes = run_pytorch(
             fn=fn,
             inputs=[
-                torch.randn(size=fn.inputs[0].type.shape, dtype=torch.float32),
-                torch.randn(size=fn.inputs[1].type.shape, dtype=torch.float32),
+                torch.randn(size=fn.inputs[0].type.shape, dtype=pytorch_dtype),
+                torch.randn(size=fn.inputs[1].type.shape, dtype=pytorch_dtype),
             ],
             use_gpu=True,
-            num_repetitions=10,
-            num_warmup=5,
+            num_repetitions=NUM_REPETITIONS,
+            num_warmup=NUM_WARMUP,
         )
         pytorch_latency = np.median(runtimes[0])
         Y[i] = pytorch_latency
+        data.append(
+            {
+                "m": batch_size,
+                "n": input_dim,
+                "k": output_dim,
+                "data_size": data_size,
+                "flops": flops,
+                "latency": pytorch_latency,
+            }
+        )
+
+    df = pd.DataFrame(data)
+    df.to_csv("matmul_benchmark.csv")
 
     reg = LinearRegression(positive=True, fit_intercept=False).fit(X, Y)
-    return 1.0 / reg.coef_[0], 1.0 / reg.coef_[1], reg.coef_[2]
+
+    return (reg.coef_[0], reg.coef_[1], reg.coef_[2])
 
 
-def calibrate_allreduce_parameters():
+def calibrate_allreduce_parameters(dtype):
+    dist_ir_dtype = Float32 if dtype == "fp32" else Float16
+    pytorch_dtype = torch.float32 if dtype == "fp32" else torch.float16
     all_input_dims = [2048, 4096, 8192]
     all_output_dims = [2048, 4096, 8192]
     n = len(all_input_dims) * len(all_output_dims)
@@ -222,7 +249,9 @@ def calibrate_allreduce_parameters():
         for i, (input_dim, output_dim) in enumerate(
             tqdm(list(itertools.product(all_input_dims, all_output_dims)))
         ):
-            fn = _allreduce(devices[1 : num_devices + 1], input_dim, output_dim)
+            fn = _allreduce(
+                devices[1 : num_devices + 1], dist_ir_dtype, input_dim, output_dim
+            )
             fn = infer_types(fn, fn.inputs)
             X[i][0] = fn.inputs[0].type.size() / BYTES_IN_Gb
             X[i][1] = num_devices
@@ -231,12 +260,12 @@ def calibrate_allreduce_parameters():
             _, runtimes = run_pytorch(
                 fn=fn,
                 inputs=[
-                    torch.randn(size=fn.inputs[i].type.shape, dtype=torch.float32)
+                    torch.randn(size=fn.inputs[i].type.shape, dtype=pytorch_dtype)
                     for i in range(len(fn.inputs))
                 ],
                 use_gpu=True,
-                num_repetitions=10,
-                num_warmup=5,
+                num_repetitions=NUM_REPETITIONS,
+                num_warmup=NUM_WARMUP,
             )
             pytorch_latency = np.median(runtimes[0])
             Y[i] = pytorch_latency
@@ -246,7 +275,7 @@ def calibrate_allreduce_parameters():
     return params
 
 
-def calibrate_simulator():
-    device_parameters = calibrate_device_parameters()
-    network_bandwidth = calibrate_network_bandwidth()
+def calibrate_simulator(dtype):
+    device_parameters = calibrate_device_parameters(dtype)
+    network_bandwidth = calibrate_network_bandwidth(dtype)
     return (*device_parameters, network_bandwidth)

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):

dist_ir/executor/simulator.py

@@ -5,7 +5,7 @@
 from warnings import warn
 
 from ..ir import Function, Device, Op
-from ..ir.type import Type, abstract_values
+from ..ir.type import Tensor, Type, abstract_values
 from .absint import (
     AbstractState,
     interpreter,
@@ -44,10 +44,28 @@ def __init__(self, function: Function, inputs: Sequence[Any]):
         self.trace = []
         self._function_inputs_set = set(function.inputs)
 
-        for inp in function.inputs:
+        for i, inp in enumerate(function.inputs):
             if inp.type is None or inp.type.device is None:
-                continue
-            self.peak_memory[inp.type.device] += inp.type.size()
+                if (
+                    isinstance(inputs[i], ConcreteValue)
+                    and inputs[i].device is not None
+                ):
+                    self.peak_memory[inputs[i].device] += inputs[i].val.nbytes
+                elif (
+                    isinstance(inputs[i], Tensor)
+                    and inputs[i].shape is not None
+                    and inputs[i].dtype is not None
+                    and inputs[i].device is not None
+                ):
+                    self.peak_memory[inputs[i].device] += inputs[i].size()
+                else:
+                    warn_msg = (
+                        f"No input type or device for input {inp} ({type(inputs[i])})"
+                    )
+                    warn(warn_msg)
+                    continue
+            else:
+                self.peak_memory[inp.type.device] += inp.type.size()
         for device in self.peak_memory:
             self.live_memory[device][0] = (0, self.peak_memory[device])
 

dist_ir/ir/type.py

@@ -148,7 +148,7 @@ def from_concrete(concrete_value):
             np.float16: Float16,
             np.float32: Float32,
             np.float64: Float64,
-            np.bool: Bool,
+            bool: Bool,
         }  # TODO does this map exist/belong somewhere else?
         dtype = dtype_to_type[concrete_value.val.dtype.type](concrete_value.device)
         return Tensor(dtype, concrete_value.val.shape, concrete_value.device)

dist_ir/transforms/gpt2_dhp_transform.py

@@ -450,7 +450,7 @@ def update_attributes(
             and value.shape == (1,)
             and value[0] == old_n_head
         ):
-            value = np.array([new_n_head])
+            value = np.array([new_n_head], dtype=value.dtype)
             sanitized_value = value.tobytes()
             attributes = frozendict(
                 {"value": sanitized_value, "device": attributes["device"]}