diff --git a/dist_ir/backend/torch.py b/dist_ir/backend/torch.py
index 42be0bd9..09c9bf94 100644
--- a/dist_ir/backend/torch.py
+++ b/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:
diff --git a/dist_ir/executor/calibrate_simulator.py b/dist_ir/executor/calibrate_simulator.py
index 1ef42778..32240115 100644
--- a/dist_ir/executor/calibrate_simulator.py
+++ b/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,7 +148,7 @@ 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
@@ -152,12 +156,12 @@ def calibrate_network_bandwidth():
                 _, 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)
diff --git a/dist_ir/executor/numpy_register.py b/dist_ir/executor/numpy_register.py
index dffe868b..10282d2f 100644
--- a/dist_ir/executor/numpy_register.py
+++ b/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):
diff --git a/dist_ir/executor/simulator.py b/dist_ir/executor/simulator.py
index f80a17e8..10f74b8a 100644
--- a/dist_ir/executor/simulator.py
+++ b/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])
 
diff --git a/dist_ir/ir/type.py b/dist_ir/ir/type.py
index 2c68238f..ca52cbc2 100644
--- a/dist_ir/ir/type.py
+++ b/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)
diff --git a/dist_ir/transforms/gpt2_dhp_transform.py b/dist_ir/transforms/gpt2_dhp_transform.py
index 404ed5a4..52c38546 100644
--- a/dist_ir/transforms/gpt2_dhp_transform.py
+++ b/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"]}
diff --git a/examples/gpt2.py b/examples/gpt2.py
index 9cffee2f..e434dedb 100644
--- a/examples/gpt2.py
+++ b/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, Float32, Tensor, Type, abstract_values
 from dist_ir.transforms import (
     gpt2_dhp_transform,
     sanitize_unhashable_attributes,
@@ -31,6 +31,41 @@ 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),
+        )
+        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 +384,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 +398,22 @@ 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:
-                input_data_map[inp] = inp.type
+    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:
+                if dtype == "fp16" and isinstance(inp.type.dtype, Float32):
+                    input_data_map[inp] = Tensor(
+                        shape=inp.type.shape,
+                        dtype=Float16(inp.type.dtype.device),
+                        device=inp.type.device,
+                    )
+                else:
+                    input_data_map[inp] = inp.type
+            elif dtype == "fp16" and input_data_map[inp].dtype == np.float32:
+                input_data_map[inp] = input_data_map[inp].astype(np.float16)
     input_data = list(input_data_map.values())
 
     return function, input_data
@@ -488,6 +538,7 @@ def transform(
 
 def get_transformed_function_and_input_data(
     model_path,
+    dtype,
     device_throughput,
     dram_bandwidth,
     kernel_launch_overhead,
@@ -515,6 +566,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,13 +606,17 @@ 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:
             return torch.int32
         elif dtype == np.int64:
             return torch.int64
+        elif dtype == np.float16:
+            return torch.float16
         elif dtype == np.float32:
             return torch.float32
         else:
@@ -616,6 +672,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,
diff --git a/examples/gpt2_grid_search.py b/examples/gpt2_grid_search.py
index 8f085d12..2d4c2265 100644
--- a/examples/gpt2_grid_search.py
+++ b/examples/gpt2_grid_search.py
@@ -10,6 +10,7 @@ class GPTGridSearch(GridSearch):
     def __init__(
         self,
         backend,
+        dtype,
         use_gpu,
         output_file,
         device_throughput,
@@ -38,6 +39,7 @@ def __init__(
         super().__init__(
             model_params,
             backend,
+            dtype,
             use_gpu,
             output_file,
             device_throughput,
@@ -51,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 = {}
diff --git a/examples/grid_search.py b/examples/grid_search.py
index 47b0dd75..1f12d3a7 100644
--- a/examples/grid_search.py
+++ b/examples/grid_search.py
@@ -43,6 +43,7 @@ def __init__(
         self,
         model_params,
         backend,
+        dtype,
         use_gpu,
         output_file,
         device_throughput,
@@ -55,6 +56,7 @@ def __init__(
     ):
         self.model_params = model_params
         self.backend = backend
+        self.dtype = dtype
         self.use_gpu = use_gpu
         self.output_file = output_file
         self.device_throughput = device_throughput
@@ -263,15 +265,16 @@ def run_grid_search(args, grid_search_cls):
     if args.simulation_parameters_file is not None:
         with open(args.simulation_parameters_file, "r") as f:
             simulation_parameters = json.load(f)
-        args.device_throughput = simulation_parameters["device_throughput"]
-        args.dram_bandwidth = simulation_parameters["dram_bandwidth"]
-        args.kernel_launch_overhead = simulation_parameters["kernel_launch_overhead"]
+        args.device_throughput = 1.0 / simulation_parameters["device_parameters"][0]
+        args.dram_bandwidth = 1.0 / simulation_parameters["device_parameters"][1]
+        args.kernel_launch_overhead = simulation_parameters["device_parameters"][2]
         args.network_bandwidth = simulation_parameters["network_bandwidth"]
         args.allreduce_parameters = {
             int(k): v for k, v in simulation_parameters["allreduce_parameters"].items()
         }
     grid_search = grid_search_cls(
         args.backend,
+        args.dtype,
         args.use_gpu,
         args.output_file,
         args.device_throughput,
diff --git a/examples/mlp.py b/examples/mlp.py
index 949c2b9a..0346d59e 100644
--- a/examples/mlp.py
+++ b/examples/mlp.py
@@ -5,8 +5,14 @@
 import torch
 
 from dist_ir.ir import FunctionMaker, Topology, get_uniform_topology, Value
-from dist_ir.ir.type import Int32, Float32, Tensor, abstract_values
-from dist_ir.executor import CostModel, Simulator, infer_types
+from dist_ir.ir.type import Int32, Float16, Float32, Tensor, abstract_values
+from dist_ir.executor import (
+    CostModel,
+    Simulator,
+    ConcreteValue,
+    infer_types,
+    sequentially_execute,
+)
 from dist_ir.transforms import mlp_dhp_transform
 from .parser import Parser
 import dist_ir.backend.torch as torch_backend
@@ -39,42 +45,40 @@ def get_typed_input_values(inputs, batch_size, input_dim, output_dim):
     return tuple(typed_inputs)
 
 
-def get_input_data(batch_size, dim, num_layers):
-    x = np.random.normal(size=(batch_size, dim))
-    z = np.random.normal(size=(batch_size, dim))
-    n = batch_size
-    weights = [np.random.normal(size=(dim, dim))]
-    for i in range(1, num_layers - 1):
-        weights.append(np.random.normal(size=(dim, dim)))
-    weights.append(np.random.normal(size=(dim, dim)))
+def get_input_data(inputs, batch_size, input_dim, output_dim, device, dtype):
+    input_data = []
+    x = np.random.normal(0, 0.02, size=(batch_size, input_dim))
+    z = np.random.normal(0, 0.02, size=(batch_size, output_dim))
+    n = np.int64(batch_size)
+    weights = [np.random.normal(0, 0.02, size=inp.type.shape) for inp in inputs[3:]]
     input_data = [x, z, n] + weights
-    input_data = [
-        v.astype(np.float32) if i != 2 else v for i, v in enumerate(input_data)
-    ]
+    input_data = [v.astype(dtype) if i != 2 else v for i, v in enumerate(input_data)]
+    input_data = [ConcreteValue(v, device) for v in input_data]
+    assert len(input_data) == len(inputs)
     return input_data
 
 
-def mlp(input_dim, hidden_dim, output_dim, num_hidden_layers, device):
+def mlp(input_dim, hidden_dim, output_dim, num_hidden_layers, device, dtype):
     function = FunctionMaker(name="mlp")
     x = function.add_input_value(
         "x",
-        Tensor(dtype=Float32(), shape=None, device=device),
+        Tensor(dtype=dtype(), shape=None, device=device),
     )
     z = function.add_input_value(
         "z",
-        Tensor(dtype=Float32(), shape=None, device=device),
+        Tensor(dtype=dtype(), shape=None, device=device),
     )
     n = function.add_input_value("n", Int32(device=device))
     weights = []
     for i in range(num_hidden_layers - 1):
         w = function.add_input_value(
             f"w{chr(ord('A')+i)}",
-            Tensor(dtype=Float32(), shape=(input_dim, hidden_dim), device=device),
+            Tensor(dtype=dtype(), shape=(input_dim, hidden_dim), device=device),
         )
         weights.append(w)
     w = function.add_input_value(
         f"w{chr(ord('A')+i+1)}",
-        Tensor(dtype=Float32(), shape=(hidden_dim, output_dim), device=device),
+        Tensor(dtype=dtype(), shape=(hidden_dim, output_dim), device=device),
     )
     weights.append(w)
 
@@ -107,24 +111,24 @@ def mlp(input_dim, hidden_dim, output_dim, num_hidden_layers, device):
 
 
 def mlp_inference(
-    batch_size, input_dim, hidden_dim, output_dim, num_hidden_layers, device
+    batch_size, input_dim, hidden_dim, output_dim, num_hidden_layers, device, dtype
 ):
     function = FunctionMaker(name="mlp")
     weights = []
     for i in range(num_hidden_layers - 1):
         w = function.add_input_value(
             f"w{chr(ord('A')+i)}",
-            Tensor(dtype=Float32(), shape=(input_dim, hidden_dim), device=device),
+            Tensor(dtype=dtype(), shape=(input_dim, hidden_dim), device=device),
         )
         weights.append(w)
     w = function.add_input_value(
         f"w{chr(ord('A')+i+1)}",
-        Tensor(dtype=Float32(), shape=(hidden_dim, output_dim), device=device),
+        Tensor(dtype=dtype(), shape=(hidden_dim, output_dim), device=device),
     )
     weights.append(w)
     x = function.add_input_value(
         "x",
-        Tensor(dtype=Float32(), shape=(batch_size, input_dim), device=device),
+        Tensor(dtype=dtype(), shape=(batch_size, input_dim), device=device),
     )
 
     a = x
@@ -136,7 +140,7 @@ def mlp_inference(
 
 
 def mlp_inference_dp(
-    batch_size, input_dim, hidden_dim, output_dim, num_hidden_layers, devices
+    batch_size, input_dim, hidden_dim, output_dim, num_hidden_layers, devices, dtype
 ):
     num_devices = len(devices)
     assert batch_size % num_devices == 0
@@ -147,16 +151,16 @@ def mlp_inference_dp(
         for i in range(num_hidden_layers - 1):
             weights[i, d] = function.add_input_value(
                 f"w{chr(ord('A')+i)}_{d.device_id}",
-                Tensor(dtype=Float32(), shape=(input_dim, hidden_dim), device=d),
+                Tensor(dtype=dtype(), shape=(input_dim, hidden_dim), device=d),
             )
         weights[num_hidden_layers - 1, d] = function.add_input_value(
             f"w{chr(ord('A')+i+1)}_{d.device_id}",
-            Tensor(dtype=Float32(), shape=(hidden_dim, output_dim), device=d),
+            Tensor(dtype=dtype(), shape=(hidden_dim, output_dim), device=d),
         )
         x[d] = function.add_input_value(
             f"x_{d.device_id}",
             Tensor(
-                dtype=Float32(), shape=(batch_size // num_devices, input_dim), device=d
+                dtype=dtype(), shape=(batch_size // num_devices, input_dim), device=d
             ),
         )
 
@@ -274,12 +278,31 @@ def simulate(function, input_types, topology, allreduce_parameters=None):
     return simulation
 
 
-def run_pytorch(function, input_data, world_size, use_gpu=True):
+def run_pytorch(function, input_data, world_size, use_gpu=torch.cuda.is_available()):
+    # TODO: Move this to a utils file
+    def _resolve_dtype(dtype):
+        if dtype == np.int32:
+            return torch.int32
+        elif dtype == np.int64:
+            return torch.int64
+        elif dtype == np.float16:
+            return torch.float16
+        elif dtype == np.float32:
+            return torch.float32
+        else:
+            raise NotImplementedError(dtype)
+
     if use_gpu and world_size > torch.cuda.device_count():
         raise ValueError(
             f"Specified world size is {world_size}, but only "
             f"{torch.cuda.device_count()} GPUs available"
         )
+    pytorch_input_data = [
+        torch.tensor(x.val, dtype=_resolve_dtype(x.val.dtype))
+        if isinstance(x.val, np.ndarray)
+        else torch.tensor(x.val, dtype=torch.int32)
+        for x in input_data
+    ]
     input_types = abstract_values(
         input_data,
         tuple(
@@ -287,7 +310,6 @@ def run_pytorch(function, input_data, world_size, use_gpu=True):
             for i in range(len(input_data))
         ),
     )
-    pytorch_input_data = [torch.tensor(x.val, dtype=torch.float32) for x in input_data]
     per_rank_outputs, runtimes = torch_backend.run_pytorch(
         function,
         pytorch_input_data,
@@ -302,6 +324,7 @@ def run_pytorch(function, input_data, world_size, use_gpu=True):
 def run_mlp(
     phase,
     backend,
+    dtype,
     use_gpu,
     batch_size,
     input_dim,
@@ -319,6 +342,8 @@ def run_mlp(
     trace_file,
     verbose=False,
 ):
+    dist_ir_dtype = Float32 if dtype == "fp32" else Float16
+    numpy_dtype = np.float32 if dtype == "fp32" else np.float16
     world_size = dp_degree * hp_degree * pp_degree
     topology = get_uniform_topology(
         world_size,
@@ -335,6 +360,7 @@ def run_mlp(
             output_dim,
             num_hidden_layers,
             topology.devices[0],
+            dist_ir_dtype,
         )
     elif phase == "inference":
         fn = mlp_inference(
@@ -343,6 +369,7 @@ def run_mlp(
             output_dim,
             num_hidden_layers,
             topology.devices[0],
+            dist_ir_dtype,
         )
 
     if verbose:
@@ -350,6 +377,16 @@ def run_mlp(
         print("Parameter count:", parameter_count_str)
         print("Model size:", model_size_str)
 
+    if backend == "pytorch":
+        input_data = get_input_data(
+            fn.inputs,
+            batch_size,
+            input_dim,
+            output_dim,
+            topology.devices[0],
+            numpy_dtype,
+        )
+
     if world_size > 1:
         init_fn, transformed_fn = mlp_dhp_transform(
             fn,
@@ -365,6 +402,8 @@ def run_mlp(
         init_fn = infer_types(init_fn, typed_inputs)
         transformed_fn = infer_types(transformed_fn, init_fn.outputs)
         input_types = tuple(output.type for output in init_fn.outputs)
+        if backend == "pytorch":
+            transformed_input_data = sequentially_execute(init_fn, input_data)
     else:
         typed_inputs = get_typed_input_values(
             fn.inputs, batch_size, input_dim, output_dim
@@ -372,6 +411,8 @@ def run_mlp(
         fn = infer_types(fn, typed_inputs)
         transformed_fn = fn
         input_types = tuple(inp.type for inp in fn.inputs)
+        if backend == "pytorch":
+            transformed_input_data = input_data
     transformed_fn = add_optimizer_ops(transformed_fn)
     if backend == "simulate":
         simulation = simulate(transformed_fn, input_types, topology)
@@ -381,19 +422,14 @@ def run_mlp(
             simulation.dump_chrome_trace(trace_file)
         return simulation
     elif backend == "pytorch":
-        input_data = [
-            ConcreteValue(
-                np.random.normal(size=typ.size).astype(np.float32), device=typ.device
-            )
-            for typ in input_types
-        ]
-        return run_pytorch(fn, input_data, world_size, use_gpu)
+        return run_pytorch(transformed_fn, transformed_input_data, world_size, use_gpu)
 
 
 def main(args):
     run_mlp(
         args.phase,
         args.backend,
+        args.dtype,
         args.use_gpu,
         args.batch_size,
         args.input_dim,
diff --git a/examples/mlp_grid_search.py b/examples/mlp_grid_search.py
index c138f188..9668c2d5 100644
--- a/examples/mlp_grid_search.py
+++ b/examples/mlp_grid_search.py
@@ -1,5 +1,7 @@
+import numpy as np
+
 from dist_ir.ir import Value
-from dist_ir.ir.type import Tensor
+from dist_ir.ir.type import Tensor, Float32, Float16
 from dist_ir.executor import infer_types, sequentially_execute, ConcreteValue
 from dist_ir.transforms import mlp_dhp_transform
 from . import mlp
@@ -11,6 +13,7 @@ class MLPGridSearch(GridSearch):
     def __init__(
         self,
         backend,
+        dtype,
         use_gpu,
         output_file,
         device_throughput,
@@ -30,6 +33,7 @@ def __init__(
         super().__init__(
             model_params,
             backend,
+            dtype,
             use_gpu,
             output_file,
             device_throughput,
@@ -46,16 +50,25 @@ def get_model_and_input_data(self, batch_size, model_size):
         if model_size not in self.models:
             num_layers, dim = self.model_params[model_size]
             self.models[model_size] = mlp.mlp(
-                dim, dim, dim, num_layers, self.topology.devices[0]
+                dim,
+                dim,
+                dim,
+                num_layers,
+                self.topology.devices[0],
+                Float32 if self.dtype == "fp32" else Float16,
             )
 
         fn = self.models[model_size]
         num_layers, dim = self.model_params[model_size]
         if self.backend == "pytorch":
-            input_data = mlp.get_input_data(batch_size, dim, num_layers)
-            input_data = tuple(
-                ConcreteValue(t, inp.type.device)
-                for t, inp in zip(input_data, fn.inputs)
+            dtype = np.float32 if self.dtype == "fp32" else np.float16
+            input_data = mlp.get_input_data(
+                fn.inputs,
+                batch_size,
+                dim,
+                dim,
+                self.topology.devices[0],
+                dtype,
             )
         else:
             input_data = mlp.get_typed_input_values(fn.inputs, batch_size, dim, dim)
diff --git a/examples/mlsys_experiments.py b/examples/mlsys_experiments.py
index e4b49486..04fc4485 100644
--- a/examples/mlsys_experiments.py
+++ b/examples/mlsys_experiments.py
@@ -11,6 +11,10 @@
 
 
 def calibrate_parameters(args):
+    if args.calibrate_all:
+        args.calibrate_device_parameters = True
+        args.calibrate_network_bandwidth = True
+        args.calibrate_allreduce_parameters = True
     if args.output_file is None:
         raise ValueError(
             "Output file must be specified to calibrate simulation parameters"
@@ -19,16 +23,8 @@ def calibrate_parameters(args):
         print(f"Reading simulation parameters from {args.output_file}...")
         with open(args.output_file, "r") as f:
             simulation_parameters = json.load(f)
-        if "device_throughput" in simulation_parameters:
-            device_throughput = simulation_parameters["device_throughput"]
-        else:
-            assert args.calibrate_device_parameters
-        if "dram_bandwidth" in simulation_parameters:
-            dram_bandwidth = simulation_parameters["dram_bandwidth"]
-        else:
-            assert args.calibrate_device_parameters
-        if "kernel_launch_overhead" in simulation_parameters:
-            kernel_launch_overhead = simulation_parameters["kernel_launch_overhead"]
+        if "device_parameters" in simulation_parameters:
+            device_parameters = simulation_parameters["device_parameters"]
         else:
             assert args.calibrate_device_parameters
         if "network_bandwidth" in simulation_parameters:
@@ -44,27 +40,21 @@ def calibrate_parameters(args):
     update_simulation_parameters = False
     if args.calibrate_device_parameters:
         print("Calibrating device parameters...")
-        (
-            dram_bandwidth,
-            device_throughput,
-            kernel_launch_overhead,
-        ) = calibrate_device_parameters()
+        device_parameters = calibrate_device_parameters(args.dtype)
         update_simulation_parameters = True
-        print(f"DRAM bandwidth: {dram_bandwidth:.2e}")
-        print(f"Device throughput: {device_throughput:.2e}")
-        print(f"Kernel launch overhead: {kernel_launch_overhead:.2e}")
+        print(f"DRAM bandwidth: {1.0 / device_parameters[0]:.2e}")
+        print(f"Device throughput: {1.0 / device_parameters[1]:.2e}")
+        print(f"Kernel launch overhead: {device_parameters[2]:.2e}")
     if args.calibrate_network_bandwidth:
-        network_bandwidth = calibrate_network_bandwidth()
+        network_bandwidth = calibrate_network_bandwidth(args.dtype)
         update_simulation_parameters = True
         print(f"Network bandwidth: {network_bandwidth}")
     if args.calibrate_allreduce_parameters:
-        allreduce_parameters = calibrate_allreduce_parameters()
+        allreduce_parameters = calibrate_allreduce_parameters(args.dtype)
         update_simulation_parameters = True
         print(f"Allreduce parameters: {allreduce_parameters}")
     if update_simulation_parameters:
-        simulation_parameters["dram_bandwidth"] = dram_bandwidth
-        simulation_parameters["device_throughput"] = device_throughput
-        simulation_parameters["kernel_launch_overhead"] = kernel_launch_overhead
+        simulation_parameters["device_parameters"] = device_parameters
         simulation_parameters["network_bandwidth"] = network_bandwidth
         simulation_parameters["allreduce_parameters"] = allreduce_parameters
         with open(args.output_file, "w") as f:
@@ -139,6 +129,15 @@ def prepare_accuracy_sample_configs(args):
         default=False,
         help="Calibrate allreduce parameters",
     )
+    parser.add_argument(
+        "--calibrate_all",
+        action="store_true",
+        default=False,
+        help="Calibrate all parameters",
+    )
+    parser.add_argument(
+        "--dtype", choices=["fp32", "fp16"], required=True, help="Dtype"
+    )
 
     args = parser.parse_args()
     assert args.mode is not None
diff --git a/examples/parser.py b/examples/parser.py
index f0567d90..c8eb5595 100644
--- a/examples/parser.py
+++ b/examples/parser.py
@@ -75,6 +75,7 @@ def add_simulation_config_arguments(self):
 
     def add_execution_mode_config_arguments(self):
         self.add_argument("--backend", choices=["simulate", "pytorch"], required=True)
+        self.add_argument("--dtype", choices=["fp32", "fp16"], default="fp32")
 
     def add_simulation_output_config_arguments(self):
         self.add_argument("--trace_file", type=str, default=None, help="Trace file")
@@ -89,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",
         )
 
@@ -186,7 +187,3 @@ def add_gpt2_model_path_config_arguments(self):
                 "text/machine_comprehension/gpt-2/model/gpt2-10.onnx?raw=True)"
             ),
         )
-
-    def add_calibration_arguments(self):
-        # TODO: Add for simulator accuracy
-        pass
diff --git a/test/test_gpt2_dhp_transform.py b/test/test_gpt2_dhp_transform.py
index fc9e8a6d..8ff72e32 100644
--- a/test/test_gpt2_dhp_transform.py
+++ b/test/test_gpt2_dhp_transform.py
@@ -6,8 +6,9 @@
 
 from dist_ir.executor import sequentially_execute, ConcreteValue
 from dist_ir.ir import cpprint
-from examples.gpt2 import get_transformed_function_and_input_data, run_pytorch
+from dist_ir.ir.type import Float16, Float32
 from dist_ir.utils import constants
+from examples.gpt2 import get_transformed_function_and_input_data, simulate, run_pytorch
 
 # Assume the onnx file is stored in the repository root
 MODEL_PATH = (Path(__file__).parent.parent / "gpt2-10.onnx").absolute()
@@ -16,6 +17,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 +40,7 @@ def _run_gpt(
         topology,
     ) = get_transformed_function_and_input_data(
         MODEL_PATH,
+        dtype,
         device_throughput,
         dram_bandwidth,
         kernel_launch_overhead,
@@ -73,10 +76,13 @@ def _run_gpt(
         else:
             outputs = sequentially_execute(transformed_function, initialized_input_data)
         return outputs
+    else:
+        return simulate(transformed_function, initialized_input_data, topology)
 
 
 def _test(
     original_outputs,
+    dtype,
     dp_degree=1,
     hp_degree=1,
     pp_degree=1,
@@ -86,6 +92,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,
@@ -94,14 +101,28 @@ def _test(
     )
     assert len(transformed_outputs) == dp_degree * hp_degree
     for i in range(len(transformed_outputs)):
+        if dtype == "fp32":
+            assert transformed_outputs[i].val.dtype == np.float32
+        else:
+            assert transformed_outputs[i].val.dtype == np.float16
         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()
+    if torch.cuda.is_available():
+        return {
+            "fp16": _run_gpt(dtype="fp16", use_pytorch_backend=True),
+            "fp32": _run_gpt(dtype="fp32", use_pytorch_backend=True),
+        }
+    else:
+        return {
+            "fp32": _run_gpt(dtype="fp32", use_pytorch_backend=True),
+        }
 
 
 @pytest.mark.parametrize(
@@ -110,7 +131,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 +141,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 +167,16 @@ 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):
-    _run_gpt(
+def test_mixed_simulation(dtype, dp_degree, hp_degree, pp_degree):
+    simulation = _run_gpt(
+        dtype=dtype,
         dp_degree=dp_degree,
         hp_degree=hp_degree,
         pp_degree=pp_degree,
         num_microbatches=pp_degree,
         use_real_weights=False,
     )
+    # TODO: Verify that output dtypes are correct
diff --git a/test/test_grid_search.py b/test/test_grid_search.py
index 0ae951c8..185a6950 100644
--- a/test/test_grid_search.py
+++ b/test/test_grid_search.py
@@ -17,11 +17,22 @@
 
 
 @pytest.mark.parametrize(
-    ("backend"),
-    ["simulate", "pytorch"],
+    "backend, dtype",
+    [
+        ("simulate", "fp32"),
+        ("simulate", "fp16"),
+        ("pytorch", "fp32"),
+        pytest.param(
+            "pytorch",
+            "fp16",
+            marks=pytest.mark.skipif(
+                not torch.cuda.is_available(), reason="fp16 only available on GPU"
+            ),
+        ),
+    ],
 )
-def test_mlp_grid_search(backend):
-    all_world_sizes = [1, 2, 4]
+def test_mlp_grid_search(backend, dtype):
+    all_world_sizes = [1, 2]
     all_batch_sizes = [256]
     all_model_sizes = ["mlp-xs"]
     with tempfile.NamedTemporaryFile() as tf:
@@ -30,6 +41,7 @@ def test_mlp_grid_search(backend):
             writer.writeheader()
         grid_search = MLPGridSearch(
             backend=backend,
+            dtype=dtype,
             use_gpu=torch.cuda.is_available(),
             output_file=tf.name,
             device_throughput=constants.DEFAULT_DEVICE_THROUGHPUT,
@@ -56,6 +68,7 @@ def test_mlp_grid_search(backend):
                 simulation = mlp.run_mlp(
                     phase="training",
                     backend="simulate",
+                    dtype=dtype,
                     use_gpu=False,
                     batch_size=all_batch_sizes[0],
                     input_dim=dim,
@@ -88,11 +101,22 @@ def test_mlp_grid_search(backend):
 
 
 @pytest.mark.parametrize(
-    ("backend"),
-    ["simulate", "pytorch"],
+    "backend, dtype",
+    [
+        ("simulate", "fp32"),
+        ("simulate", "fp16"),
+        ("pytorch", "fp32"),
+        pytest.param(
+            "pytorch",
+            "fp16",
+            marks=pytest.mark.skipif(
+                not torch.cuda.is_available(), reason="fp16 only available on GPU"
+            ),
+        ),
+    ],
 )
-def test_gpt_grid_search(backend):
-    all_world_sizes = [1, 2, 4]
+def test_gpt_grid_search(backend, dtype):
+    all_world_sizes = [1, 2]
     all_batch_sizes = [64]
     all_model_sizes = ["gpt2-xs"]
     with tempfile.NamedTemporaryFile() as tf:
@@ -101,6 +125,7 @@ def test_gpt_grid_search(backend):
             writer.writeheader()
         grid_search = GPTGridSearch(
             backend=backend,
+            dtype=dtype,
             use_gpu=torch.cuda.is_available(),
             output_file=tf.name,
             device_throughput=constants.DEFAULT_DEVICE_THROUGHPUT,
@@ -131,6 +156,7 @@ def test_gpt_grid_search(backend):
                     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,
@@ -159,3 +185,10 @@ def test_gpt_grid_search(backend):
                     & (df["num_microbatches"] == p)
                 ]["latency"].values[0]
                 assert math.isclose(latency, grid_search_latency, abs_tol=10 ** -8)
+
+
+if __name__ == "__main__":
+    print(f"MLP fp32")
+    test_mlp_grid_search("pytorch", "fp32")
+    print(f"MLP fp16")
+    test_mlp_grid_search("pytorch", "fp16")
diff --git a/test/test_mlp_dhp_transform.py b/test/test_mlp_dhp_transform.py
index bc9ebfc1..1a969a74 100644
--- a/test/test_mlp_dhp_transform.py
+++ b/test/test_mlp_dhp_transform.py
@@ -6,8 +6,8 @@
 
 from examples import mlp
 from dist_ir.ir import FunctionMaker, get_uniform_topology
+from dist_ir.ir.type import Float32, Float16
 from dist_ir.executor import infer_types, sequentially_execute, ConcreteValue
-from dist_ir.ir.type import Float32, Tensor
 from dist_ir.transforms import mlp_dhp_transform
 
 BATCH_SIZE = 64
@@ -51,19 +51,22 @@ def _verify_hp(function, transformed_function, outputs, transformed_outputs, dp=
 
 
 @pytest.mark.parametrize(
-    ("dp_degree", "hp_degree", "pp_degree"),
-    list(itertools.product([1, 2], [1, 2], [1, 2])),
+    ("dp_degree", "hp_degree", "pp_degree", "dtype"),
+    list(itertools.product([1, 2], [1, 2], [1, 2], ["fp32", "fp16"])),
 )
 def test_mlp_dhp_transform(
     dp_degree,
     hp_degree,
     pp_degree,
+    dtype,
     batch_size=BATCH_SIZE,
     num_hidden_layers=8,
     input_dim=INPUT_DIM,
 ):
     num_microbatches = pp_degree
     world_size = dp_degree * hp_degree * pp_degree
+    dist_ir_dtype = Float32 if dtype == "fp32" else Float16
+    numpy_dtype = np.float32 if dtype == "fp32" else np.float16
     topology = get_uniform_topology(world_size)
     function = mlp.mlp(
         input_dim,
@@ -71,6 +74,7 @@ def test_mlp_dhp_transform(
         input_dim,
         num_hidden_layers,
         topology.devices[0],
+        dist_ir_dtype,
     )
     typed_inputs = mlp.get_typed_input_values(
         function.inputs, batch_size, input_dim, input_dim
@@ -90,13 +94,14 @@ def test_mlp_dhp_transform(
     transformed_function = infer_types(transformed_function, init_function.outputs)
     transformed_function = mlp.add_optimizer_ops(transformed_function)
 
-    input_data = [
-        ConcreteValue(
-            np.random.normal(size=inp.type.shape) if i != 2 else batch_size,
-            topology.devices[0],
-        )
-        for i, inp in enumerate(typed_inputs)
-    ]
+    input_data = mlp.get_input_data(
+        init_function.inputs,
+        batch_size,
+        input_dim,
+        input_dim,
+        topology.devices[0],
+        numpy_dtype,
+    )
     outputs = sequentially_execute(function, input_data)
     dist_input_data = sequentially_execute(init_function, input_data)
     transformed_outputs = sequentially_execute(transformed_function, dist_input_data)
diff --git a/test/test_pytorch_backend.py b/test/test_pytorch_backend.py
index 9e01b98c..cc8578e7 100644
--- a/test/test_pytorch_backend.py
+++ b/test/test_pytorch_backend.py
@@ -9,7 +9,7 @@
 from dist_ir.executor.simulator import Simulator
 from dist_ir.executor.type_inference import infer_types
 from dist_ir.ir import Device, FunctionMaker, cpprint, Value
-from dist_ir.ir.type import Float32, Tensor
+from dist_ir.ir.type import Float16, Float32, Tensor
 from dist_ir.ir.topology import Topology, get_uniform_topology
 
 # TODO make examples submodule of dist_ir?
@@ -79,20 +79,34 @@ def create_owt_model(num_devices, num_layers):
 
 
 @pytest.mark.parametrize(
-    "num_devices, num_layers, use_gpu",
+    "num_devices, num_layers, use_gpu, dtype",
     [
-        (2, 4, False),
+        (2, 4, False, "fp32"),
         pytest.param(
             2,
             4,
             True,
+            "fp32",
+            marks=pytest.mark.skipif(
+                torch.cuda.device_count() < 2, reason="Not enough available GPUs"
+            ),
+        ),
+        pytest.param(
+            2,
+            4,
+            True,
+            "fp16",
             marks=pytest.mark.skipif(
                 torch.cuda.device_count() < 2, reason="Not enough available GPUs"
             ),
         ),
     ],
 )
-def test_owt(num_devices, num_layers, use_gpu):
+def test_owt(num_devices, num_layers, use_gpu, dtype):
+    dist_ir_dtype = Float32 if dtype == "fp32" else Float16
+    numpy_dtype = np.float32 if dtype == "fp32" else np.float16
+    torch_dtype = torch.float32 if dtype == "fp32" else torch.float16
+
     fn = create_owt_model(num_devices, num_layers)
 
     devices = [Device(0, "cpu")]
@@ -110,11 +124,11 @@ def test_owt(num_devices, num_layers, use_gpu):
             else:
                 shape = (hidden_dim, hidden_dim // num_devices)
             # w{l}_{d}:
-            input_vals.append(Value("", Tensor(Float32(), shape, devices[d])))
+            input_vals.append(Value("", Tensor(dist_ir_dtype(), shape, devices[d])))
     for d in range(1, num_devices + 1):
         # x_{d}:
         shape = (batch_size // num_devices, hidden_dim)
-        input_vals.append(Value("", Tensor(Float32(), shape, devices[d])))
+        input_vals.append(Value("", Tensor(dist_ir_dtype(), shape, devices[d])))
 
     # Test type inference:
     fn = infer_types(fn, input_vals)
@@ -125,8 +139,11 @@ def test_owt(num_devices, num_layers, use_gpu):
 
     # Test with sequential executor:
     np.random.seed(0)
-    weights = [np.random.randn(hidden_dim, hidden_dim) for l in range(num_layers)]
-    x = np.random.randn(batch_size, hidden_dim)
+    weights = [
+        np.random.normal(0, 0.02, size=(hidden_dim, hidden_dim)).astype(numpy_dtype)
+        for l in range(num_layers)
+    ]
+    x = np.random.normal(0, 0.02, size=(batch_size, hidden_dim)).astype(numpy_dtype)
 
     # Split inputs for distributed function
     input_arrays = []
@@ -150,7 +167,7 @@ def test_owt(num_devices, num_layers, use_gpu):
 
     # Run per-rank modules using PyTorch backend:
     per_rank_outputs, _ = run_pytorch(
-        fn, [torch.tensor(a) for a in input_arrays], use_gpu=use_gpu
+        fn, [torch.tensor(a).to(torch_dtype) for a in input_arrays], use_gpu=use_gpu
     )
 
     # Check outputs:
@@ -249,18 +266,27 @@ def test_send_recv(use_gpu):
 
 
 @pytest.mark.parametrize(
-    "use_gpu",
+    "use_gpu, dtype",
     [
-        False,
+        (False, "fp32"),
+        (False, "fp16"),
+        pytest.param(
+            True,
+            "fp16",
+            marks=pytest.mark.skipif(
+                torch.cuda.device_count() < 2, reason="Not enough available GPUs"
+            ),
+        ),
         pytest.param(
             True,
+            "fp32",
             marks=pytest.mark.skipif(
                 torch.cuda.device_count() < 2, reason="Not enough available GPUs"
             ),
         ),
     ],
 )
-def test_dp_mlp(use_gpu):
+def test_dp_mlp(use_gpu, dtype):
     num_devices = 2
     num_layers = 4
     batch_size = 4
@@ -268,7 +294,13 @@ def test_dp_mlp(use_gpu):
     devices = [Device(d, "gpu") for d in range(num_devices + 1)]
 
     fn = mlp_inference_dp(
-        batch_size, hidden_dim, hidden_dim, hidden_dim, num_layers, devices[1:]
+        batch_size,
+        hidden_dim,
+        hidden_dim,
+        hidden_dim,
+        num_layers,
+        devices[1:],
+        Float32 if dtype == "fp32" else Float16,
     )
     fn = infer_types(fn, fn.inputs)
     cpprint(fn)
@@ -293,7 +325,9 @@ def new_inputs():
 
     # Project and run on backend:
     per_rank_outputs, runtimes = run_pytorch(
-        fn, convert_inputs_dp(weights, x), use_gpu=use_gpu
+        fn,
+        convert_inputs_dp(weights, x),
+        use_gpu=use_gpu,
     )
 
     # Check outputs: