Repository: pytorch/opacus (language: python)

Issue:
Add opacus microbenchmarks code
Merge the code used for opacus microbenchmarks into opacus repo


Discussion / hints:

Fix patch:
diff --git a/benchmarks/README.md b/benchmarks/README.md
new file mode 100644
index 00000000..7910835f
--- /dev/null
+++ b/benchmarks/README.md
@@ -0,0 +1,114 @@
+# Benchmarks
+
+A set of micro-benchmarks that measure elapsed runtime and allocated CUDA memory for both [basic modules](https://github.com/pytorch/opacus/tree/main/opacus/grad_sample) and [more complex layers](https://github.com/pytorch/opacus/tree/main/opacus/layers), as well as their respective [torch.nn](https://pytorch.org/docs/stable/nn.html) counterparts.
+
+Requires PyTorch version >= 1.10.0.
+
+## Contents
+
+- [run_benchmarks.py](run_benchmarks.py) runs a list of benchmarks based on the given config file and writes out the results as pickle files to `results/raw` (by default)
+
+- [benchmark_layer.py](benchmark_layer.py) benchmarks a single layer at the given batch size for both runtime and memory
+
+- [config.json](config.json) contains an example JSON config for running benchmarks
+
+- [layers.py](layers.py) implements each layer with its corresponding inputs/outputs, moving each component to device, and collecting respective memory statistics
+
+- [utils.py](utils.py) implements helper functions, e.g. for saving results as pickle files
+
+
+## Benchmarks
+
+For each layer and batch size in [config.json](config.json), [run_benchmarks.py](run_benchmarks.py) will do the following:
+```
+Do this num_runs times:
+    Init layer, one batch of random input, one batch of random "labels"
+    Move each component to GPU and collect its allocated CUDA memory (if applicable)
+
+    Start timer
+
+    Do this num_repeats times:
+        preds = layer(input)
+        loss = criterion(preds, labels)
+        loss.backward()
+
+    Stop timer
+    
+    Return elapsed time / num_repeats and memory statistics
+```
+
+## Layers
+
+All layers follow the corresponding `torch.nn` module's interface with the same default values (if not specified in [config.json](config.json)).
+
+A note on `input_shape` in [config.json](config.json): parameters that are shared between the model and the input are listed separately. Therefore, the actual input shape will vary:
+
+- Linear: `(batch_size, *input_shape, in_features)`
+
+- Convolutional: `(batch_size, in_channels, *input_shape)`
+
+- LayerNorm:
+    - Input: `(batch_size, *input_shape)`
+    - Normalized shape: `(input_shape[-D:])`
+
+- InstanceNorm: `(batch_size, num_features, *input_shape)`
+
+- GroupNorm: `(batch_size, num_channels, *input_shape)`
+
+- Embedding: `(batch_size, *input_shape)`
+
+- MultiheadAttention:
+    - If not batch_first: `(targ_seq_len, batch_size, embed_dim)`
+    - Else: `(batch_size, targ_seq_len, embed_dim)`
+
+- RNN, GRU, LSTM:
+    - If not batch_first: `(seq_len, batch_size, input_size)`
+    - Else: `(batch_size, seq_len, input_size)`
+
+
+## Usage
+
+`run_benchmarks.py` without additional arguments will replicate the results in the [technical report introducing Opacus](https://arxiv.org/abs/2109.12298), which you can cite as follows:
+```
+@article{opacus,
+  title={Opacus: {U}ser-Friendly Differential Privacy Library in {PyTorch}},
+  author={Ashkan Yousefpour and Igor Shilov and Alexandre Sablayrolles and Davide Testuggine and Karthik Prasad and Mani Malek and John Nguyen and Sayan Ghosh and Akash Bharadwaj and Jessica Zhao and Graham Cormode and Ilya Mironov},
+  journal={arXiv preprint arXiv:2109.12298},
+  year={2021}
+}
+```
+
+If saving results, ensure that the `results/raw` (or otherwise specified) root directory exists.
+
+```
+usage: run_benchmarks.py [-h]
+                         [--layers {linear,gsm_linear,conv,gsm_conv,layernorm,gsm_layernorm,instancenorm,gsm_instancenorm,groupnorm,gsm_groupnorm,embedding,gsm_embedding,mha,dpmha,gsm_dpmha,rnn,dprnn,gsm_dprnn,gru,dpgru,gsm_dpgru,lstm,dplstm,gsm_dplstm} [{linear,gsm_linear,conv,gsm_conv,layernorm,gsm_layernorm,instancenorm,gsm_instancenorm,groupnorm,gsm_groupnorm,embedding,gsm_embedding,mha,dpmha,gsm_dpmha,rnn,dprnn,gsm_dprnn,gru,dpgru,gsm_dpgru,lstm,dplstm,gsm_dplstm} ...]]
+                         [--batch_sizes BATCH_SIZES [BATCH_SIZES ...]] [--num_runs NUM_RUNS]
+                         [--num_repeats NUM_REPEATS] [--forward_only] [--random_seed RANDOM_SEED]
+                         [-c CONFIG_FILE] [--cont] [--root ROOT] [--suffix SUFFIX] [--no_save]
+                         [-v]
+
+optional arguments:
+  -h, --help            show this help message and exit
+  --layers {linear,gsm_linear,conv,gsm_conv,layernorm,gsm_layernorm,instancenorm,gsm_instancenorm,groupnorm,gsm_groupnorm,embedding,gsm_embedding,mha,dpmha,gsm_dpmha,rnn,dprnn,gsm_dprnn,gru,dpgru,gsm_dpgru,lstm,dplstm,gsm_dplstm} [{linear,gsm_linear,conv,gsm_conv,layernorm,gsm_layernorm,instancenorm,gsm_instancenorm,groupnorm,gsm_groupnorm,embedding,gsm_embedding,mha,dpmha,gsm_dpmha,rnn,dprnn,gsm_dprnn,gru,dpgru,gsm_dpgru,lstm,dplstm,gsm_dplstm} ...]
+  --batch_sizes BATCH_SIZES [BATCH_SIZES ...]
+  --num_runs NUM_RUNS   number of benchmarking runs
+  --num_repeats NUM_REPEATS
+                        how many forward/backward passes per run
+  --forward_only        only run forward passes
+  --random_seed RANDOM_SEED
+                        random seed for the first run of each layer, subsequent runs increase the
+                        random seed by 1
+  -c CONFIG_FILE, --config_file CONFIG_FILE
+  --cont                only run missing experiments
+  --root ROOT           path to directory that benchmark results should be saved in
+  --suffix SUFFIX       suffix to append to each result file's name
+  --no_save
+  -v, --verbose
+```
+
+## Tests
+
+```python -m pytest tests/```
+
+Tests for CUDA memory benchmarking will be skipped on CPU.
diff --git a/benchmarks/benchmark_layer.py b/benchmarks/benchmark_layer.py
new file mode 100644
index 00000000..41466f59
--- /dev/null
+++ b/benchmarks/benchmark_layer.py
@@ -0,0 +1,103 @@
+import argparse
+import json
+from typing import Callable, Dict, Tuple
+
+import torch
+import torch.utils.benchmark as benchmark
+from layers import LayerFactory, LayerType
+from utils import get_layer_set, reset_peak_memory_stats
+
+
+def run_layer_benchmark(
+    num_repeats: int,
+    forward_only: bool = False,
+    create_layer: Callable = LayerFactory.create,
+    **kwargs,
+) -> Tuple[float, Dict[str, int]]:
+    """Benchmarks a single layer for runtime and CUDA memory (if applicable).
+
+    Args:
+        num_repeats: how many times to repeat the forward(/backward) pass
+        forward_only: whether to skip the backward pass
+        create_layer: function for creating the layer, takes **kwargs
+
+    Returns: a tuple consisting of
+        Runtime as a float, and
+        Memory statistics as a dict
+    """
+
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+    if torch.cuda.is_available():
+        assert reset_peak_memory_stats(device).cur_mem == 0
+
+    # setup layer
+    layer_fun = create_layer(**kwargs)
+
+    if forward_only:
+        layer_fun.module.eval()
+        benchmark_fun = layer_fun.forward_only
+    else:
+        layer_fun.module.train()
+        benchmark_fun = layer_fun.forward_backward
+
+    # move layer to device and get memory statistics
+    memory_stats = layer_fun.to(device=device)
+    assert sum(v for _, v in memory_stats.items()) == torch.cuda.memory_allocated(
+        device
+    )
+
+    # benchmark.Timer performs its own warmups
+    try:
+        timer = benchmark.Timer(
+            stmt="benchmark_fun()",
+            globals={"benchmark_fun": benchmark_fun},
+            num_threads=1,
+        )
+        runtime = timer.timeit(num_repeats).mean
+    except RuntimeError:
+        runtime = float("nan")
+
+    # get max memory allocated and reset memory statistics
+    memory_stats["max_memory"] = reset_peak_memory_stats(device).prev_max_mem
+
+    return runtime, memory_stats
+
+
+def main(args) -> None:
+
+    with open(args.config_file) as config_file:
+        config = json.load(config_file)
+
+    runtime, memory_stats = run_layer_benchmark(
+        num_repeats=args.num_repeats,
+        forward_only=args.forward_only,
+        layer_name=args.layer,
+        batch_size=args.batch_size,
+        random_seed=args.random_seed,
+        **config[get_layer_set(args.layer)],
+    )
+    print(f"Runtime (seconds): {runtime}")
+    print(f"Memory statistics (bytes): {memory_stats}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "layer",
+        choices=[v for k, v in LayerType.__dict__.items() if not k.startswith("__")],
+    )
+    parser.add_argument("--batch_size", default=64, type=int)
+    parser.add_argument(
+        "--num_repeats",
+        default=20,
+        type=int,
+        help="how many forward/backward passes to run",
+    )
+    parser.add_argument(
+        "--forward_only", action="store_true", help="only run forward passes"
+    )
+    parser.add_argument("--random_seed", default=0, type=int)
+    parser.add_argument("-c", "--config_file", type=str, default="config.json")
+    args = parser.parse_args()
+    main(args)
diff --git a/benchmarks/config.json b/benchmarks/config.json
new file mode 100644
index 00000000..ce4cd254
--- /dev/null
+++ b/benchmarks/config.json
@@ -0,0 +1,43 @@
+{
+	"linear": {
+		"input_shape": [],
+		"in_features": 512,
+		"out_features": 512
+	},
+	"conv": {
+		"in_channels": 64,
+		"input_shape": [50, 100],
+		"out_channels": 64,
+		"kernel_size": 8
+	},
+	"layernorm": {
+		"input_shape": [64],
+		"D": 1
+	},
+	"instancenorm": {
+		"num_features": 256,
+		"input_shape": [64],
+		"affine": true
+	},
+	"groupnorm": {
+		"input_shape": [],
+		"num_groups": 16,
+		"num_channels": 256
+	},
+	"embedding": {
+		"input_shape": [],
+		"num_embeddings": 20000,
+		"embedding_dim": 100
+	},
+	"mha": {
+		"source_seq_len": 128,
+		"targ_seq_len": 64,
+		"embed_dim": 100,
+		"num_heads": 4
+	},
+	"rnn_base": {
+		"seq_len": 128,
+		"input_size": 100,
+		"hidden_size": 100
+	}
+}
diff --git a/benchmarks/layers.py b/benchmarks/layers.py
new file mode 100644
index 00000000..d1220c31
--- /dev/null
+++ b/benchmarks/layers.py
@@ -0,0 +1,593 @@
+from typing import Callable, Dict, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from opacus.grad_sample import GradSampleModule
+from opacus.layers import DPGRU, DPLSTM, DPRNN, DPMultiheadAttention
+
+
+class LayerType:
+    LINEAR: str = "linear"
+    GSM_LINEAR: str = "gsm_linear"
+    CONV: str = "conv"
+    GSM_CONV: str = "gsm_conv"
+    LAYERNORM: str = "layernorm"
+    GSM_LAYERNORM: str = "gsm_layernorm"
+    INSTANCENORM: str = "instancenorm"
+    GSM_INSTANCENORM: str = "gsm_instancenorm"
+    GROUPNORM: str = "groupnorm"
+    GSM_GROUPNORM: str = "gsm_groupnorm"
+    EMBEDDING: str = "embedding"
+    GSM_EMBEDDING: str = "gsm_embedding"
+    MHA: str = "mha"
+    DPMHA: str = "dpmha"
+    GSM_DPMHA: str = "gsm_dpmha"
+    RNN: str = "rnn"
+    DPRNN: str = "dprnn"
+    GSM_DPRNN: str = "gsm_dprnn"
+    GRU: str = "gru"
+    DPGRU: str = "dpgru"
+    GSM_DPGRU: str = "gsm_dpgru"
+    LSTM: str = "lstm"
+    DPLSTM: str = "dplstm"
+    GSM_DPLSTM: str = "gsm_dplstm"
+
+
+class Layer:
+    _input_tensor: torch.Tensor
+    _module: nn.Module
+    _labels: torch.Tensor
+
+    def __init__(
+        self,
+        *,
+        random_seed: Optional[int] = None,
+        criterion: Callable = F.cross_entropy,
+    ) -> None:
+        """Sets random seed and criterion."""
+        if random_seed is not None:
+            torch.manual_seed(random_seed)
+        self._criterion = criterion
+
+    @staticmethod
+    def _get_memory_difference(device: torch.device, stats: Dict[str, int]) -> int:
+        """If applicable, computes the device's CUDA memory difference between the
+        sum of the values in the stats dict and the current allocated CUDA memory.
+
+        Args:
+            device: torch.device
+            stats: dictionary from item names (e.g. input_tensor, module, labels)
+            to their respective sizes
+
+        Returns:
+            device's CUDA memory difference between the sum of the values in the
+            stats dict and the current allocated CUDA memory if given a CUDA
+            device. 0 if given a CPU device.
+        """
+        if device.type == "cuda":
+            return torch.cuda.memory_allocated(device) - sum(stats.values())
+        return 0
+
+    def _inputs_to(self, device: torch.device, stats: Dict[str, int]) -> Dict[str, int]:
+        """Some modules (e.g. RNNs) take additional layer inputs such as initial
+        hidden state. These modules should override this function accordingly.
+
+        Args:
+            device: torch.device
+            stats: dictionary from item names (e.g. input_tensor, module, labels)
+            to their respective sizes
+
+        Returns: updated stats dictionary with input tensor sizes
+        """
+        self._input_tensor = self._input_tensor.to(device)
+        stats["input_tensor"] = self._get_memory_difference(device=device, stats=stats)
+        return stats
+
+    def to(self, device: torch.device) -> Dict[str, int]:
+        """Moves input_tensor, additional inputs, module, and labels to device.
+
+        Args:
+            device: torch.device
+
+        Returns:
+            Dictionary from item names (e.g. input_tensor, module, labels) to
+            their respective sizes if CUDA device, else respective sizes are all
+            set to 0.
+        """
+        stats: Dict[str, int] = {}
+        stats["offset"] = self._get_memory_difference(device=device, stats=stats)
+
+        # some modules take additional inputs such as hidden state
+        stats = self._inputs_to(device=device, stats=stats)
+
+        self._module = self._module.to(device)
+        stats["layer"] = self._get_memory_difference(device=device, stats=stats)
+
+        self._labels = self._labels.to(device)
+        stats["labels"] = self._get_memory_difference(device=device, stats=stats)
+
+        # check that all memory is accounted for
+        if device.type == "cuda":
+            assert torch.cuda.memory_allocated(device) == sum(stats.values())
+
+        return stats
+
+    def forward_only(self) -> torch.Tensor:
+        return self._module(self._input_tensor)
+
+    def forward_backward(self) -> None:
+        preds = self.forward_only()
+        loss = self._criterion(preds, self._labels)
+        loss.backward()
+        self._module.zero_grad()
+
+    def make_private(self) -> None:
+        self._module = GradSampleModule(self._module)
+
+    @property
+    def module(self):
+        return self._module
+
+    def __del__(self):
+        self.to(torch.device("cpu"))
+
+
+class LinearBase(Layer):
+    def __init__(
+        self,
+        *,
+        batch_size: int,
+        input_shape: Tuple[int, ...],
+        in_features: int,
+        out_features: int,
+        bias: bool = True,
+        random_seed: Optional[int] = None,
+        criterion: Callable = F.cross_entropy,
+    ) -> None:
+        super().__init__(random_seed=random_seed, criterion=criterion)
+        self._input_tensor = torch.randn(batch_size, *input_shape, in_features)
+        self._module = nn.Linear(
+            in_features=in_features,
+            out_features=out_features,
+            bias=bias,
+        )
+        self._labels = torch.randn(batch_size, *input_shape, out_features)
+
+
+class ConvBase(Layer):
+    def __init__(
+        self,
+        *,
+        batch_size: int,
+        in_channels: int,
+        input_shape: Tuple[int, ...],
+        out_channels: int,
+        kernel_size: Union[int, Tuple[int, ...]],
+        stride: Union[int, Tuple[int, ...]] = 1,
+        padding: Union[int, Tuple[int, ...]] = 0,
+        dilation: Union[int, Tuple[int, ...]] = 1,
+        groups: int = 1,
+        bias: bool = True,
+        padding_mode: str = "zeros",
+        random_seed: Optional[int] = None,
+        criterion: Callable = F.cross_entropy,
+    ) -> None:
+        super().__init__(random_seed=random_seed, criterion=criterion)
+
+        D = len(input_shape)
+        if D == 1:
+            self._module_name = nn.Conv1d
+        elif D == 2:
+            self._module_name = nn.Conv2d
+        elif D == 3:
+            self._module_name = nn.Conv3d
+        else:
+            raise Exception("Input shape must be between 1 and 3 long")
+
+        self._input_tensor = torch.randn(batch_size, in_channels, *input_shape)
+        self._module = self._module_name(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            groups=groups,
+            bias=bias,
+            padding_mode=padding_mode,
+        )
+        outputs = self._module(self._input_tensor)
+        self._labels = torch.randn(outputs.shape)
+        del outputs
+
+
+class LayerNormBase(Layer):
+    def __init__(
+        self,
+        *,
+        batch_size: int,
+        input_shape: Tuple[int, ...],
+        D: int,
+        eps: float = 1e-05,
+        elementwise_affine: bool = True,
+        random_seed: Optional[int] = None,
+        criterion: Callable = F.cross_entropy,
+    ) -> None:
+        super().__init__(random_seed=random_seed, criterion=criterion)
+
+        self._input_tensor = torch.randn(batch_size, *input_shape)
+        self._module = nn.LayerNorm(
+            normalized_shape=self._input_tensor.shape[-D:],
+            eps=eps,
+            elementwise_affine=elementwise_affine,
+        )
+        self._labels = torch.randn(self._input_tensor.shape)
+
+
+class InstanceNormBase(Layer):
+    def __init__(
+        self,
+        *,
+        batch_size: int,
+        num_features: int,
+        input_shape: Tuple[int, ...],
+        eps: float = 1e-05,
+        affine: bool = False,
+        track_running_stats: bool = False,
+        random_seed: Optional[int] = None,
+        criterion: Callable = F.cross_entropy,
+    ) -> None:
+        super().__init__(random_seed=random_seed, criterion=criterion)
+
+        D = len(input_shape)
+        if D == 1:
+            self._module_name = nn.InstanceNorm1d
+        elif D == 2:
+            self._module_name = nn.InstanceNorm2d
+        elif D == 3:
+            self._module_name = nn.InstanceNorm3d
+        else:
+            raise Exception("Input shape must be between 1 and 3 long")
+
+        self._input_tensor = torch.randn(batch_size, num_features, *input_shape)
+        self._module = self._module_name(
+            num_features=num_features,
+            eps=eps,
+            affine=affine,
+            track_running_stats=track_running_stats,
+        )
+        self._labels = torch.randn(self._input_tensor.shape)
+
+
+class GroupNormBase(Layer):
+    def __init__(
+        self,
+        *,
+        batch_size: int,
+        input_shape: Tuple[int, ...],
+        num_groups: int,
+        num_channels: int,
+        eps: float = 1e-05,
+        affine: bool = True,
+        random_seed: Optional[int] = None,
+        criterion: Callable = F.cross_entropy,
+    ) -> None:
+        super().__init__(random_seed=random_seed, criterion=criterion)
+
+        self._input_tensor = torch.randn(batch_size, num_channels, *input_shape)
+        self._module = nn.GroupNorm(
+            num_groups=num_groups, num_channels=num_channels, eps=eps, affine=affine
+        )
+        self._labels = torch.randn(self._input_tensor.shape)
+
+
+class EmbeddingBase(Layer):
+    def __init__(
+        self,
+        *,
+        batch_size: int,
+        input_shape: Tuple[int, ...],
+        num_embeddings: int,
+        embedding_dim: int,
+        padding_idx: Optional[int] = None,
+        max_norm: Optional[float] = None,
+        norm_type: float = 2.0,
+        scale_grad_by_freq: bool = False,
+        sparse: bool = False,
+        random_seed: Optional[int] = None,
+        criterion: Callable = F.cross_entropy,
+    ) -> None:
+        super().__init__(random_seed=random_seed, criterion=criterion)
+
+        self._input_tensor = torch.randint(
+            high=num_embeddings,
+            size=(batch_size, *input_shape),
+            dtype=torch.long,
+        )
+        self._module = nn.Embedding(
+            num_embeddings=num_embeddings,
+            embedding_dim=embedding_dim,
+            padding_idx=padding_idx,
+            max_norm=max_norm,
+            norm_type=norm_type,
+            scale_grad_by_freq=scale_grad_by_freq,
+            sparse=sparse,
+        )
+        self._labels = torch.randn(batch_size, *input_shape, embedding_dim)
+
+
+class MHABase(Layer):
+    def __init__(
+        self,
+        *,
+        layer: nn.Module,
+        batch_size: int,
+        source_seq_len: int,
+        targ_seq_len: int,
+        embed_dim: int,
+        num_heads: int,
+        dropout: float = 0.0,
+        bias: bool = True,
+        add_bias_kv: bool = False,
+        add_zero_attn: bool = False,
+        kdim: Optional[int] = None,
+        vdim: Optional[int] = None,
+        batch_first: bool = False,
+        random_seed: Optional[int] = None,
+        criterion: Callable = F.cross_entropy,
+    ) -> None:
+        super().__init__(random_seed=random_seed, criterion=criterion)
+
+        kdim = kdim if kdim else embed_dim
+        vdim = vdim if vdim else embed_dim
+
+        self._input_tensor = (
+            torch.randn(targ_seq_len, batch_size, embed_dim)
+            if not batch_first
+            else torch.randn(batch_size, targ_seq_len, embed_dim)
+        )
+        self._key = (
+            torch.randn(source_seq_len, batch_size, kdim)
+            if not batch_first
+            else torch.randn(batch_size, source_seq_len, kdim)
+        )
+        self._value = (
+            torch.randn(source_seq_len, batch_size, vdim)
+            if not batch_first
+            else torch.randn(batch_size, source_seq_len, vdim)
+        )
+        self._module = layer(
+            embed_dim,
+            num_heads,
+            dropout=dropout,
+            bias=bias,
+            add_bias_kv=add_bias_kv,
+            add_zero_attn=add_zero_attn,
+            kdim=kdim,
+            vdim=vdim,
+        )
+        self._labels = (
+            torch.randn(targ_seq_len, batch_size, embed_dim)
+            if not batch_first
+            else torch.randn(batch_size, targ_seq_len, embed_dim)
+        )
+
+    def _inputs_to(self, device: torch.device, stats: Dict[str, int]) -> Dict[str, int]:
+        """MultiheadAttention takes additional layer inputs key and value.
+
+        Args:
+            device: torch.device
+            stats: dictionary from item names (e.g. input_tensor, module, labels)
+            to their respective sizes
+
+        Returns: updated stats dictionary with input tensor, key, and value size
+        """
+        stats = super()._inputs_to(device=device, stats=stats)
+        self._key = self._key.to(device)
+        stats["key"] = self._get_memory_difference(device=device, stats=stats)
+        self._value = self._value.to(device)
+        stats["value"] = self._get_memory_difference(device=device, stats=stats)
+        return stats
+
+    def forward_only(self) -> torch.Tensor:
+        return self._module(self._input_tensor, self._key, self._value)[0]
+
+
+class RNNBase(Layer):
+    def __init__(
+        self,
+        *,
+        layer: nn.Module,
+        batch_size: int,
+        seq_len: int,
+        input_size: int,
+        hidden_size: int,
+        num_layers: int = 1,
+        bias: bool = False,
+        batch_first: bool = False,
+        dropout: float = 0,
+        bidirectional: bool = False,
+        random_seed: Optional[int] = None,
+        criterion: Callable = F.cross_entropy,
+        **kwargs,
+    ) -> None:
+        super().__init__(random_seed=random_seed, criterion=criterion)
+
+        self._input_tensor = (
+            torch.randn(
+                seq_len,
+                batch_size,
+                input_size,
+            )
+            if not batch_first
+            else torch.randn(batch_size, seq_len, input_size)
+        )
+
+        D = 2 if bidirectional else 1
+        self._h_0 = torch.randn(D * num_layers, batch_size, hidden_size)
+
+        self._module = layer(
+            input_size=input_size,
+            hidden_size=hidden_size,
+            num_layers=num_layers,
+            bias=bias,
+            batch_first=batch_first,
+            dropout=dropout,
+            bidirectional=bidirectional,
+            **kwargs,
+        )
+
+        self._labels = (
+            torch.randn(seq_len, batch_size, D * hidden_size)
+            if not batch_first
+            else torch.randn(batch_size, seq_len, D * hidden_size)
+        )
+
+    def _inputs_to(self, device: torch.device, stats: Dict[str, int]) -> Dict[str, int]:
+        """RNNs take additional layer inputs h_0.
+
+        Args:
+            device: torch.device
+            stats: dictionary from item names (e.g. input_tensor, module, labels)
+            to their respective sizes
+
+        Returns: updated stats dictionary with input tensor and h_0 size
+        """
+        stats = super()._inputs_to(device=device, stats=stats)
+        self._h_0 = self._h_0.to(device)
+        stats["h_0"] = self._get_memory_difference(device=device, stats=stats)
+        return stats
+
+    def forward_only(self) -> torch.Tensor:
+        return self._module(self._input_tensor, self._h_0)[0]
+
+
+class LSTMBase(RNNBase):
+    def __init__(
+        self,
+        *,
+        layer: nn.Module,
+        batch_size: int,
+        seq_len: int,
+        input_size: int,
+        hidden_size: int,
+        num_layers: int = 1,
+        bias: bool = False,
+        batch_first: bool = False,
+        dropout: float = 0,
+        bidirectional: bool = False,
+        proj_size: int = 0,
+        random_seed: Optional[int] = None,
+        criterion: Callable = F.cross_entropy,
+    ) -> None:
+        super().__init__(
+            layer=layer,
+            batch_size=batch_size,
+            seq_len=seq_len,
+            input_size=input_size,
+            hidden_size=hidden_size,
+            num_layers=num_layers,
+            bias=bias,
+            batch_first=batch_first,
+            dropout=dropout,
+            bidirectional=bidirectional,
+            proj_size=proj_size,
+            random_seed=random_seed,
+            criterion=criterion,
+        )
+        h_out = proj_size if proj_size > 0 else hidden_size
+        D = 2 if bidirectional else 1
+        self._h_0 = torch.randn(D * num_layers, batch_size, h_out)
+        self._c_0 = torch.randn(D * num_layers, batch_size, hidden_size)
+
+        self._labels = (
+            torch.randn(seq_len, batch_size, D * h_out)
+            if not batch_first
+            else torch.randn(batch_size, seq_len, D * h_out)
+        )
+
+    def _inputs_to(self, device: torch.device, stats: Dict[str, int]) -> Dict[str, int]:
+        """LSTMs take additional layer inputs h_0, c_0.
+
+        Args:
+            device: torch.device
+            stats: dictionary from item names (e.g. input_tensor, module, labels)
+            to their respective sizes
+
+        Returns: updated stats dictionary with input tensor, h_0, and c_0 size
+        """
+        stats = super()._inputs_to(device=device, stats=stats)
+        self._h_0 = self._h_0.to(device)
+        stats["h_0"] = self._get_memory_difference(device=device, stats=stats)
+
+        self._c_0 = self._c_0.to(device)
+        stats["c_0"] = self._get_memory_difference(device=device, stats=stats)
+
+        return stats
+
+    def forward_only(self) -> torch.Tensor:
+        return self._module(self._input_tensor, (self._h_0, self._c_0))[0]
+
+
+class LayerFactory:
+    @staticmethod
+    def make_private(layer: Layer) -> Layer:
+        layer.make_private()
+        return layer
+
+    @staticmethod
+    # flake8: noqa C901
+    def create(layer_name: str, **kwargs) -> Layer:
+        if layer_name == LayerType.LINEAR:
+            return LinearBase(**kwargs)
+        elif layer_name == LayerType.GSM_LINEAR:
+            return LayerFactory.make_private(LinearBase(**kwargs))
+        elif layer_name == LayerType.CONV:
+            return ConvBase(**kwargs)
+        elif layer_name == LayerType.GSM_CONV:
+            return LayerFactory.make_private(ConvBase(**kwargs))
+        elif layer_name == LayerType.LAYERNORM:
+            return LayerNormBase(**kwargs)
+        elif layer_name == LayerType.GSM_LAYERNORM:
+            return LayerFactory.make_private(LayerNormBase(**kwargs))
+        elif layer_name == LayerType.INSTANCENORM:
+            return InstanceNormBase(**kwargs)
+        elif layer_name == LayerType.GSM_INSTANCENORM:
+            return LayerFactory.make_private(InstanceNormBase(**kwargs))
+        elif layer_name == LayerType.GROUPNORM:
+            return GroupNormBase(**kwargs)
+        elif layer_name == LayerType.GSM_GROUPNORM:
+            return LayerFactory.make_private(GroupNormBase(**kwargs))
+        elif layer_name == LayerType.EMBEDDING:
+            return EmbeddingBase(**kwargs)
+        elif layer_name == LayerType.GSM_EMBEDDING:
+            return LayerFactory.make_private(EmbeddingBase(**kwargs))
+        elif layer_name == LayerType.RNN:
+            return RNNBase(layer=nn.RNN, **kwargs)
+        elif layer_name == LayerType.DPRNN:
+            return RNNBase(layer=DPRNN, **kwargs)
+        elif layer_name == LayerType.GSM_DPRNN:
+            return LayerFactory.make_private(RNNBase(layer=DPRNN, **kwargs))
+        elif layer_name == LayerType.GRU:
+            return RNNBase(layer=nn.GRU, **kwargs)
+        elif layer_name == LayerType.DPGRU:
+            return RNNBase(layer=DPGRU, **kwargs)
+        elif layer_name == LayerType.GSM_DPGRU:
+            return LayerFactory.make_private(RNNBase(layer=DPGRU, **kwargs))
+        elif layer_name == LayerType.LSTM:
+            return LSTMBase(layer=nn.LSTM, **kwargs)
+        elif layer_name == LayerType.DPLSTM:
+            return LSTMBase(layer=DPLSTM, **kwargs)
+        elif layer_name == LayerType.GSM_DPLSTM:
+            return LayerFactory.make_private(LSTMBase(layer=DPLSTM, **kwargs))
+        elif layer_name == LayerType.MHA:
+            return MHABase(layer=nn.MultiheadAttention, **kwargs)
+        elif layer_name == LayerType.DPMHA:
+            return MHABase(layer=DPMultiheadAttention, **kwargs)
+        elif layer_name == LayerType.GSM_DPMHA:
+            return LayerFactory.make_private(
+                MHABase(layer=DPMultiheadAttention, **kwargs)
+            )
+        else:
+            raise Exception(f"Invalid layer type: {layer_name}.")
diff --git a/benchmarks/run_benchmarks.py b/benchmarks/run_benchmarks.py
new file mode 100644
index 00000000..87300a9e
--- /dev/null
+++ b/benchmarks/run_benchmarks.py
@@ -0,0 +1,150 @@
+import argparse
+import json
+import logging
+from os.path import exists
+from typing import Any, Dict
+
+from benchmark_layer import run_layer_benchmark
+from layers import LayerType
+from utils import get_layer_set, get_path, save_results
+
+
+logger = logging.getLogger(__name__)
+
+
+def run_and_save_benchmark(
+    layer: LayerType,
+    batch_size: int,
+    args,
+    layer_config: Dict[str, Any],
+    root: str = "./results/raw/",
+    suffix: str = "",
+) -> None:
+    """Runs and saves (if desired) the benchmark for the given layer and batch size.
+
+    Args:
+        layer: the layer to run benchmarks for
+        batch_size: the batch size to run benchmarks for
+        args: additional arguments
+        layer_config: the settings for this layer
+        root: directory to write results to
+        suffix: optional string to append to result file name
+    """
+
+    logger.info(f"Benchmarking {layer} layer with batch size {batch_size}.")
+    results = []
+
+    for i in range(args.num_runs):
+        runtime, memory_stats = run_layer_benchmark(
+            num_repeats=args.num_repeats,
+            forward_only=args.forward_only,
+            layer_name=layer,
+            batch_size=batch_size,
+            random_seed=args.random_seed + i if args.random_seed else args.random_seed,
+            **layer_config,
+        )
+        res = {"runtime": runtime, "memory_stats": memory_stats}
+        results.append(res)
+        logger.info(res)
+
+    # save the benchmark results if desired
+    if not args.no_save:
+        save_results(
+            layer=layer,
+            batch_size=batch_size,
+            num_runs=args.num_runs,
+            num_repeats=args.num_repeats,
+            results=results,
+            config=layer_config,
+            random_seed=args.random_seed,
+            forward_only=args.forward_only,
+            root=root,
+            suffix=suffix,
+        )
+
+
+def main(args) -> None:
+
+    if args.verbose:
+        logger.setLevel(logging.DEBUG)
+
+    # get config for each layer
+    logger.info(f"Using {args.config_file} as config file.")
+    with open(args.config_file) as config_file:
+        config = json.load(config_file)
+
+    for layer in args.layers:
+        for batch_size in args.batch_sizes:
+
+            # skip benchmark for this layer and batch size if applicable
+            if args.cont and exists(
+                get_path(
+                    layer=layer,
+                    batch_size=batch_size,
+                    num_runs=args.num_runs,
+                    num_repeats=args.num_repeats,
+                    random_seed=args.random_seed,
+                    forward_only=args.forward_only,
+                    root=args.root,
+                    suffix=args.suffix,
+                )
+            ):
+                logger.info(f"Skipping {layer} at {batch_size} - already exists.")
+                continue
+
+            # run and save (if applicable) the benchmark for this layer and batch size
+            run_and_save_benchmark(
+                layer=layer,
+                batch_size=batch_size,
+                args=args,
+                layer_config=config[get_layer_set(layer)],
+                root=args.root,
+                suffix=args.suffix,
+            )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    layers = [v for k, v in LayerType.__dict__.items() if not k.startswith("__")]
+    parser.add_argument("--layers", choices=layers, default=layers, nargs="+")
+    parser.add_argument(
+        "--batch_sizes",
+        default=[16, 32, 64, 128, 256, 512],
+        nargs="+",
+        type=int,
+    )
+    parser.add_argument(
+        "--num_runs", default=100, type=int, help="number of benchmarking runs"
+    )
+    parser.add_argument(
+        "--num_repeats",
+        default=20,
+        type=int,
+        help="how many forward/backward passes per run",
+    )
+    parser.add_argument(
+        "--forward_only", action="store_true", help="only run forward passes"
+    )
+    parser.add_argument(
+        "--random_seed",
+        default=0,
+        type=int,
+        help="random seed for the first run of each layer, subsequent runs increase the random seed by 1",
+    )
+    parser.add_argument("-c", "--config_file", type=str, default="config.json")
+    parser.add_argument(
+        "--cont", action="store_true", help="only run missing experiments"
+    )
+    parser.add_argument(
+        "--root",
+        type=str,
+        default="./results/raw/",
+        help="path to directory where benchmark results should be saved in",
+    )
+    parser.add_argument(
+        "--suffix", type=str, help="suffix to append to each result file's name"
+    )
+    parser.add_argument("--no_save", action="store_true")
+    parser.add_argument("-v", "--verbose", action="store_true")
+    args = parser.parse_args()
+    main(args)
diff --git a/benchmarks/utils.py b/benchmarks/utils.py
new file mode 100644
index 00000000..dff34e9d
--- /dev/null
+++ b/benchmarks/utils.py
@@ -0,0 +1,151 @@
+import pickle
+from collections import namedtuple
+from typing import Any, Dict, List, Optional
+
+import torch
+from layers import LayerType
+
+
+Memory = namedtuple("Memory", "prev_max_mem, cur_mem")
+
+
+def reset_peak_memory_stats(device: torch.device) -> Memory:
+    """Safely resets CUDA peak memory statistics of device if it is
+    a CUDA device.
+
+    Notes: ``torch.cuda.reset_peak_memory_stats(device)`` will error
+    if no CUDA memory has been allocated to the device.
+
+    Args:
+        device: A torch.device
+
+    Returns:
+        max_memory_allocated before resetting the statistics and
+        memory_allocated, both in bytes
+    """
+    prev_max_memory = torch.cuda.max_memory_allocated(device)
+    memory_allocated = torch.cuda.memory_allocated(device)
+
+    if prev_max_memory != memory_allocated and prev_max_memory > 0:
+        # raises RuntimeError if no previous allocation occurred
+        torch.cuda.reset_peak_memory_stats(device)
+        assert torch.cuda.max_memory_allocated(device) == memory_allocated
+
+    return Memory(prev_max_memory, memory_allocated)
+
+
+def get_layer_set(layer: str) -> str:
+    """Layers in the same layer set share a config.
+
+    Args:
+        layer: Full name of the layer. This will be the PyTorch or Opacus
+        name of the layer in lower case (e.g. linear, rnn, dprnn), prefixed with
+        gsm_ (e.g. gsm_linear, gsm_dprnn) if DP is enabled. MultiheadAttention
+        is abbreviated to mha.
+
+    Returns:
+        The name of the layer set, where a set of layers are defined as layers
+        that share the same __init__ signature.
+
+    Notes:
+        All RNN-based models share a config.
+
+    """
+    layer_set = layer.replace("gsm_dp", "").replace("gsm_", "").replace("dp", "")
+
+    # all RNN-based model use the same config
+    if layer_set in ["rnn", "gru", "lstm"]:
+        layer_set = "rnn_base"
+
+    return layer_set
+
+
+def get_path(
+    layer: LayerType,
+    batch_size: int,
+    num_runs: int,
+    num_repeats: int,
+    random_seed: Optional[int] = None,
+    forward_only: bool = False,
+    root: str = "./results/raw/",
+    suffix: str = "",
+) -> str:
+    """Gets the path to the file where the corresponding results are located.
+    File is presumed to be a pickle file.
+
+    Args:
+        layer: full layer name
+        batch_size: batch size
+        num_runs: number of runs per benchmark
+        num_repeats: how many benchmarks were run
+        random_seed: the initial random seed
+        forward_only: whether backward passes were skipped
+        root: directory to write results to
+        suffix: optional string to append to file name
+
+    Returns:
+        Path to results pickle file
+    """
+    pickle_name = f"{layer}_bs_{batch_size}_runs_{num_runs}_repeats_{num_repeats}_seed_{random_seed}"
+    if forward_only:
+        pickle_name += "_forward_only"
+
+    if len(suffix) and not suffix.startswith("_"):
+        suffix = f"_{suffix}"
+
+    return f"{root}{pickle_name}{suffix}.pkl"
+
+
+def save_results(
+    layer: LayerType,
+    batch_size: int,
+    num_runs: int,
+    num_repeats: int,
+    results: List[Dict[str, Any]],
+    config: Dict[str, Any],
+    random_seed: Optional[int] = None,
+    forward_only: bool = False,
+    root: str = "./results/raw/",
+    suffix: str = "",
+) -> None:
+    """Saves the corresponding results as a pickle file.
+
+    Args:
+        layer: full layer name
+        batch_size: batch size
+        num_runs: number of runs per benchmark
+        num_repeats: how many benchmarks were run
+        runtimes: list of runtimes of length num_repeats
+        memory: list of memory stats of length num_repeats
+        config: layer config
+        random_seed: the initial random seed
+        forward_only: whether backward passes were skipped
+        root: directory to write results to
+        suffix: optional string to append to file name
+    """
+    path = get_path(
+        layer=layer,
+        batch_size=batch_size,
+        num_runs=num_runs,
+        num_repeats=num_repeats,
+        random_seed=random_seed,
+        forward_only=forward_only,
+        root=root,
+        suffix=suffix,
+    )
+
+    with open(path, "wb") as handle:
+        pickle.dump(
+            {
+                "layer": layer,
+                "batch_size": batch_size,
+                "num_runs": num_runs,
+                "num_repeats": num_repeats,
+                "random_seed": random_seed,
+                "forward_only": forward_only,
+                "results": results,
+                "config": config,
+            },
+            handle,
+            protocol=pickle.HIGHEST_PROTOCOL,
+        )


Test patch:
diff --git a/benchmarks/tests/helpers.py b/benchmarks/tests/helpers.py
new file mode 100644
index 00000000..9e113e8f
--- /dev/null
+++ b/benchmarks/tests/helpers.py
@@ -0,0 +1,48 @@
+from typing import Optional
+
+import pytest
+import torch
+
+
+skipifnocuda = pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="Requires CUDA."
+)
+
+
+def get_n_byte_tensor(n: int, device: Optional[torch.device] = None) -> torch.Tensor:
+    """Returns a torch.int8 tensor of size n.
+
+    Args:
+        n: size of the tensor to allocate
+        device: torch.device to allocate tensor on
+
+    Returns:
+        torch.int8 tensor of size n on device
+
+    Notes: Though 1 int8 = 1 byte, memory is allocated in blocks, such that the size
+    of the tensor in bytes >= n.
+    """
+    return torch.zeros(n, dtype=torch.int8, device=device)
+
+
+def get_actual_memory_allocated(n: int, device: torch.device) -> int:
+    """
+    Returns the CUDA memory allocated for a torch.int8 tensor of size n.
+
+    Args:
+        n: size of the tensor to get allocated memory for
+        device: torch.device to allocate tensor on
+
+    Returns:
+        Number of bytes of CUDA memory allocated for a torch.int8 tensor of size n.
+
+    Notes: Should only be called on CUDA devices. Resets CUDA memory statistics.
+    """
+    assert device.type == "cuda"
+    prev_memory_allocated = torch.cuda.memory_allocated(device)
+    tensor = get_n_byte_tensor(n, device=device)
+    memory_allocated = torch.cuda.memory_allocated(device)
+    del tensor
+    torch.cuda.reset_peak_memory_stats(device)
+    assert prev_memory_allocated == torch.cuda.memory_allocated(device)
+    return memory_allocated - prev_memory_allocated
diff --git a/benchmarks/tests/test_benchmark_layer.py b/benchmarks/tests/test_benchmark_layer.py
new file mode 100644
index 00000000..c40a2edd
--- /dev/null
+++ b/benchmarks/tests/test_benchmark_layer.py
@@ -0,0 +1,188 @@
+import math
+import time
+from itertools import product
+
+import pytest
+import torch
+import torch.nn as nn
+from benchmark_layer import run_layer_benchmark
+from helpers import get_actual_memory_allocated, get_n_byte_tensor, skipifnocuda
+from layers import Layer
+from utils import reset_peak_memory_stats
+
+
+class FakeModule(nn.Module):
+    def __init__(self, **kwargs):
+        super().__init__()
+        self._runtime = kwargs.get("runtime", 0)
+        self._pass_memory = kwargs.get("pass_memory", 0)
+        self._parameters = {
+            "fake_param": get_n_byte_tensor(kwargs.get("layer_memory", 0))
+        }
+
+
+class FakeLayer(Layer):
+    """Fake layer to test runtime and memory benchmarking.
+
+    Kwargs:
+        runtime: duration for one forward or backward pass
+        pass_memory: memory used during a forward or backward pass
+        layer_memory: memory taken up by the module
+    """
+
+    def __init__(self, **kwargs) -> None:
+        self._runtime = kwargs.get("runtime", 0)
+        self._pass_memory = kwargs.get("pass_memory", 0)
+        self._input_tensor = get_n_byte_tensor(0)
+        self._module = FakeModule(**kwargs)
+        self._labels = get_n_byte_tensor(0)
+
+    def forward_only(self) -> torch.Tensor:
+        """Wait for self.duration and allocate self.max_memory bytes"""
+        time.sleep(self._runtime)
+        tensor = get_n_byte_tensor(
+            self._pass_memory,
+            device=torch.device("cuda:0" if torch.cuda.is_available() else "cpu"),
+        )
+        return tensor
+
+    def forward_backward(self) -> None:
+        """2x runtime and memory of forward_only"""
+        _ = self.forward_only()
+        _ = self.forward_only()
+
+
+@pytest.mark.parametrize(
+    "duration, num_repeats, forward_only",
+    product(
+        [0, 0.005, 0.01, 0.05],
+        [10, 20, 50],
+        [False, True],
+    ),
+)
+def test_runtime_benchmark(
+    duration: float, num_repeats: int, forward_only: bool
+) -> None:
+    """Tests runtime benchmarks on a dummy layer.
+
+    Args:
+        duration: duration (seconds) to test runtime benchmarks for
+        num_repeats: number of repeats for each benchmark
+        forward_only: whether to benchmark forward passes or forward/backward passes
+    """
+    runtime, memory_stats = run_layer_benchmark(
+        num_repeats=num_repeats,
+        forward_only=forward_only,
+        create_layer=FakeLayer,
+        runtime=duration,
+    )
+    # account for small variations
+    assert abs(runtime - ((2 - forward_only) * duration)) < 0.01
+
+    # check that no memory allocation took place
+    assert memory_stats["layer"] == 0 and memory_stats["max_memory"] == 0
+
+
+@skipifnocuda
+@pytest.mark.parametrize(
+    "pass_bytes, layer_bytes, num_repeats, forward_only",
+    product(
+        [0, 1, 100, 500, 4096],
+        [0, 1, 256, 1024, 2000],
+        [10, 20, 50, 100],
+        [False, True],
+    ),
+)
+def test_memory_benchmark(
+    pass_bytes: int, layer_bytes: int, num_repeats: int, forward_only: bool
+) -> None:
+    """Tests CUDA memory benchmarks on a dummy layer.
+
+    Args:
+        pass_bytes: number of bytes allocated during a forward or forward/backward pass
+        layer_bytes: number of bytes allocated for the layer
+        num_repeats: number of repeats for each benchmark
+        forward_only: whether to benchmark forward passes or forward/backward passes
+    """
+    device = torch.device("cuda:0")
+
+    true_pass_memory = get_actual_memory_allocated(pass_bytes, device=device)
+    true_layer_memory = get_actual_memory_allocated(layer_bytes, device=device)
+
+    # reset memory stats and ensure there is no memory leakage
+    assert reset_peak_memory_stats(device).cur_mem == 0
+
+    runtime, memory_stats = run_layer_benchmark(
+        num_repeats=num_repeats,
+        forward_only=forward_only,
+        create_layer=FakeLayer,
+        layer_memory=layer_bytes,
+        pass_memory=pass_bytes,
+    )
+
+    assert memory_stats["layer"] == true_layer_memory
+    assert (
+        memory_stats["max_memory"]
+        == true_layer_memory + (2 - forward_only) * true_pass_memory
+    )
+
+    # reset memory stats and ensure there is no memory leakage
+    assert reset_peak_memory_stats(device).cur_mem == 0
+
+
+@skipifnocuda
+@pytest.mark.parametrize(
+    "pass_bytes, layer_bytes, num_repeats, forward_only",
+    product(
+        [0, 1, 100, 500, 4096],
+        [0, 1, 256, 1024, 2000],
+        [10, 20, 50, 100],
+        [False, True],
+    ),
+)
+def test_memory_benchmark_strict(
+    pass_bytes: int, layer_bytes: int, num_repeats: int, forward_only: bool
+) -> None:
+    """Tests CUDA memory benchmarks on a dummy layer by predicting each measurement
+    using the CUDA memory block size.
+
+    Notes:
+        During the experiments included in the paper, CUDA memory is allocated in
+        blocks, where block sizes vary across kernels. New CUDA memory is allocated
+        for each new tensor. This test will fail under a different allocation
+        scheme.
+
+    Args:
+        pass_bytes: number of bytes allocated during a forward or forward/backward pass
+        layer_bytes: number of bytes allocated for the layer
+        num_repeats: number of repeats for each benchmark
+        forward_only: whether to benchmark forward passes or forward/backward passes
+    """
+    device = torch.device("cuda:0")
+
+    # find the block size by creating a tensor of size 1 byte
+    tiny_tensor = get_n_byte_tensor(1, device=device)
+    BLOCK_SIZE = torch.cuda.max_memory_allocated(device)
+    del tiny_tensor
+
+    num_blocks_pass = math.ceil(pass_bytes / BLOCK_SIZE)
+    num_blocks_layer = math.ceil(layer_bytes / BLOCK_SIZE)
+
+    # reset memory stats and ensure there is no memory leakage
+    assert reset_peak_memory_stats(device).cur_mem == 0
+
+    runtime, memory_stats = run_layer_benchmark(
+        num_repeats=num_repeats,
+        forward_only=forward_only,
+        create_layer=FakeLayer,
+        layer_memory=layer_bytes,
+        pass_memory=pass_bytes,
+    )
+    assert memory_stats["layer"] == num_blocks_layer * BLOCK_SIZE
+    assert (
+        memory_stats["max_memory"]
+        == (num_blocks_layer + (2 - forward_only) * num_blocks_pass) * BLOCK_SIZE
+    )
+
+    # reset memory stats and ensure there is no memory leakage
+    assert reset_peak_memory_stats(device).cur_mem == 0
diff --git a/benchmarks/tests/test_layers.py b/benchmarks/tests/test_layers.py
new file mode 100644
index 00000000..1b6e7970
--- /dev/null
+++ b/benchmarks/tests/test_layers.py
@@ -0,0 +1,202 @@
+import copy
+import random
+from typing import Any, Dict, List, Tuple
+
+import pytest
+import torch
+import torch.nn as nn
+from helpers import skipifnocuda
+from layers import LayerFactory
+from opacus.grad_sample import GradSampleModule
+from opacus.layers import DPGRU, DPLSTM, DPRNN, DPMultiheadAttention
+from utils import reset_peak_memory_stats
+
+
+PARAMETERS = [
+    (
+        [("linear", nn.Linear), ("gsm_linear", nn.Linear)],
+        {"input_shape": [], "in_features": 512, "out_features": 512},
+    ),
+    (
+        [("conv", nn.Conv2d), ("gsm_conv", nn.Conv2d)],
+        {
+            "in_channels": 64,
+            "input_shape": [50, 100],
+            "out_channels": 64,
+            "kernel_size": 8,
+        },
+    ),
+    (
+        [("layernorm", nn.LayerNorm), ("gsm_layernorm", nn.LayerNorm)],
+        {"input_shape": [64], "D": 1},
+    ),
+    (
+        [
+            ("instancenorm", nn.InstanceNorm1d),
+            ("gsm_instancenorm", nn.InstanceNorm1d),
+        ],
+        {"num_features": 256, "input_shape": [64], "affine": True},
+    ),
+    (
+        [("groupnorm", nn.GroupNorm), ("gsm_groupnorm", nn.GroupNorm)],
+        {"input_shape": [], "num_groups": 16, "num_channels": 256},
+    ),
+    (
+        [("embedding", nn.Embedding), ("gsm_embedding", nn.Embedding)],
+        {"input_shape": [], "num_embeddings": 20000, "embedding_dim": 100},
+    ),
+    (
+        [
+            ("mha", nn.MultiheadAttention),
+            ("dpmha", DPMultiheadAttention),
+            ("gsm_dpmha", DPMultiheadAttention),
+        ],
+        {
+            "source_seq_len": 128,
+            "targ_seq_len": 64,
+            "embed_dim": 100,
+            "num_heads": 4,
+        },
+    ),
+    (
+        [
+            ("rnn", nn.RNN),
+            ("dprnn", DPRNN),
+            ("gsm_dprnn", DPRNN),
+            ("gru", nn.GRU),
+            ("dpgru", DPGRU),
+            ("gsm_dpgru", DPGRU),
+            ("lstm", nn.LSTM),
+            ("dplstm", DPLSTM),
+            ("gsm_dplstm", DPLSTM),
+        ],
+        {"seq_len": 128, "input_size": 100, "hidden_size": 100},
+    ),
+]
+
+
+@pytest.mark.parametrize("layer_list, layer_config", PARAMETERS)
+def test_layer_modules(
+    layer_list: List[Tuple[str, nn.Module]], layer_config: Dict[str, Any]
+) -> None:
+    """For each supported layer, tests that it is instantiated with the correct module
+    and DP support.
+
+    Args:
+        layer_list: list of tuples of form (layer_name, module)
+        layer_config: config for instantiating the layers in layer_list
+    """
+    for layer_name, module in layer_list:
+        layer = LayerFactory.create(
+            layer_name=layer_name,
+            batch_size=64,
+            **layer_config,
+        )
+
+        if "gsm" in layer_name:
+            assert isinstance(layer.module, GradSampleModule)
+            assert isinstance(layer.module.to_standard_module(), module)
+        else:
+            assert isinstance(layer.module, module)
+
+
+@skipifnocuda
+@pytest.mark.parametrize("layer_list, layer_config", PARAMETERS)
+def test_to_device(
+    layer_list: List[Tuple[str, nn.Module]], layer_config: Dict[str, Any]
+) -> None:
+    """Tests that inputs, labels, and module are initialized on CPU, and that moving
+    them to GPU and CPU works correctly.
+
+    Args:
+        layer_list: list of tuples of form (layer_name, module)
+        layer_config: config for instantiating the layers in layer_list
+    """
+    cuda = torch.device("cuda:0")
+    cpu = torch.device("cpu")
+    assert reset_peak_memory_stats(cuda).cur_mem == 0
+
+    for layer_name, module in layer_list:
+        layer = LayerFactory.create(
+            layer_name=layer_name,
+            batch_size=64,
+            **layer_config,
+        )
+        # layer should be initialized on CPU
+        assert torch.cuda.memory_allocated(cuda) == 0
+
+        mem_stats = layer.to(cuda)
+        allocated = torch.cuda.memory_allocated(cuda)
+        assert allocated > 0
+        # all allocated memory should be accounted for in the memory statistics
+        assert allocated == sum(v for _, v in mem_stats.items())
+
+        mem_stats = layer.to(cpu)
+        allocated = torch.cuda.memory_allocated(cuda)
+        assert allocated == 0
+        assert allocated == sum(v for _, v in mem_stats.items())
+
+    assert reset_peak_memory_stats(cuda).cur_mem == 0
+
+
+@pytest.mark.parametrize("layer_list, layer_config", PARAMETERS)
+def test_layer_outputs(
+    layer_list: List[Tuple[str, nn.Module]], layer_config: Dict[str, Any]
+) -> None:
+    """Layers in layer_list that share the same underlying module (either a
+    torch.nn.Module or opacus.layers.DPModule) should produce the same output
+    given the same random seed and different outputs given different random seeds.
+
+    Args:
+        layer_list: list of tuples of form (layer_name, module)
+        layer_config: config for instantiating the layers in layer_list
+    """
+    random_seed_a = random.randint(0, 100000)
+    random_seed_b = random.randint(100000, 200000)
+    outputs: Dict[int, Dict[str, torch.Tensor]] = {
+        random_seed_a: {},
+        random_seed_b: {},
+    }
+
+    for random_seed in (random_seed_a, random_seed_b):
+        for layer_name, module in layer_list:
+            layer = LayerFactory.create(
+                layer_name=layer_name,
+                batch_size=64,
+                random_seed=random_seed,
+                **layer_config,
+            )
+            if str(module) not in outputs[random_seed]:
+                outputs[random_seed][str(module)] = layer.forward_only()
+
+            # same module with same seed should result in same output
+            assert torch.equal(outputs[random_seed][str(module)], layer.forward_only())
+
+    # same module with different seed should result in different output
+    for module_name in outputs[random_seed_a]:
+        assert not torch.equal(
+            outputs[random_seed_a][module_name], outputs[random_seed_b][module_name]
+        )
+
+
+@pytest.mark.parametrize("layer_list, layer_config", PARAMETERS)
+def test_forward_backward(
+    layer_list: List[Tuple[str, nn.Module]], layer_config: Dict[str, Any]
+) -> None:
+    """Tests that Layer.forward_backward() runs for each layer in layer_list and that
+    the Layer is not modified.
+
+    Args:
+        layer_list: list of tuples of form (layer_name, module)
+        layer_config: config for instantiating the layers in layer_list
+    """
+    for layer_name, module in layer_list:
+        layer = LayerFactory.create(
+            layer_name=layer_name,
+            batch_size=64,
+            **layer_config,
+        )
+        layer_copy = copy.deepcopy(layer)
+        layer.forward_backward()
+        for p1, p2 in zip(layer.module.parameters(), layer_copy.module.parameters()):
+            assert torch.equal(p1.data, p2.data)
diff --git a/benchmarks/tests/test_utils.py b/benchmarks/tests/test_utils.py
new file mode 100644
index 00000000..6c1f5f76
--- /dev/null
+++ b/benchmarks/tests/test_utils.py
@@ -0,0 +1,177 @@
+import os
+import pickle
+import shutil
+from typing import Any, Dict, List, Tuple
+
+import pytest
+import torch
+from helpers import get_n_byte_tensor, skipifnocuda
+from utils import get_layer_set, get_path, reset_peak_memory_stats, save_results
+
+
+@pytest.mark.parametrize(
+    "layer_set, layers",
+    [
+        # Pytorch layers are named layer (no DP) or gsm_layer (DP)
+        ("linear", ["linear", "gsm_linear"]),
+        ("conv", ["conv", "gsm_conv"]),
+        # Opacus layers are named dplayer (no DP) or gsm_dplayer (DP)
+        ("mha", ["mha", "dpmha", "gsm_dpmha"]),
+        # RNN-based models share the same interface
+        ("rnn_base", ["rnn", "dprnn", "gsm_dprnn"]),
+        ("rnn_base", ["lstm", "dplstm", "gsm_dplstm"]),
+    ],
+)
+def test_get_layer_set(layer_set: str, layers: List[str]) -> None:
+    """Tests assignment of individual layers to the layer set.
+
+    Args:
+        layer_set: layer set (e.g. linear, rnn_base)
+        layers: non-exhaustive list of layers that belong to the layer_set
+    """
+    assert all(get_layer_set(layer) == layer_set for layer in layers)
+
+
+@skipifnocuda
+@pytest.mark.parametrize(
+    "prev_max_memory, allocated_memory",
+    [
+        # prev_max_memory = allocated_memory = 0 --> (0, 0)
+        (0, 0),
+        # prev_max_memory = allocated_memory > 0 --> (prev_max_memory, prev_max_memory)
+        (1, 1),
+        # prev_max_memory > allocated_memory = 0 --> (prev_max_memory, 0)
+        (1, 0),
+        # prev_max_memory > allocated_memory > 0 --> (prev_max_memory, allocated_memory)
+        (2, 1),
+    ],
+)
+def test_reset_peak_memory_stats(prev_max_memory: int, allocated_memory: int) -> None:
+    """Tests resetting of peak memory stats.
+
+    Notes: Only the relative and not the absolute sizes of prev_max_memory and
+    allocated_memory are relevant.
+
+    Args:
+        prev_max_memory: current maximum memory stat to simulate
+        allocated_memory: current allocated memory to simulate
+    """
+    device = torch.device("cuda:0")
+
+    # keep x, delete y
+    x = get_n_byte_tensor(allocated_memory, device=device)
+    y = get_n_byte_tensor(prev_max_memory - allocated_memory, device=device)
+    del y
+
+    # get the true allocated memory (CUDA memory is allocated in blocks)
+    prev_max_memory = torch.cuda.max_memory_allocated(device)
+    allocated_memory = torch.cuda.memory_allocated(device)
+    assert prev_max_memory >= allocated_memory
+    assert reset_peak_memory_stats(device) == (prev_max_memory, allocated_memory)
+
+    # clean up
+    del x
+    torch.cuda.reset_peak_memory_stats(device)
+    assert torch.cuda.max_memory_allocated(device) == 0
+    assert torch.cuda.memory_allocated(device) == 0
+
+
+@pytest.mark.parametrize(
+    "config, path",
+    [
+        (
+            {"layer": "linear", "batch_size": 64, "num_runs": 10, "num_repeats": 100},
+            "./results/raw/linear_bs_64_runs_10_repeats_100_seed_None.pkl",
+        ),
+        (
+            {
+                "layer": "gsm_rnn",
+                "batch_size": 128,
+                "num_runs": 5,
+                "num_repeats": 20,
+                "random_seed": 13964,
+                "forward_only": True,
+            },
+            "./results/raw/gsm_rnn_bs_128_runs_5_repeats_20_seed_13964_forward_only.pkl",
+        ),
+        (
+            {
+                "layer": "dpmha",
+                "batch_size": 16,
+                "num_runs": 20,
+                "num_repeats": 50,
+                "random_seed": 88362,
+                "root": "./results/tmp/",
+                "suffix": "no_3",
+            },
+            "./results/tmp/dpmha_bs_16_runs_20_repeats_50_seed_88362_no_3.pkl",
+        ),
+    ],
+)
+def test_get_path(config: Dict[str, Any], path: str) -> None:
+    """Tests result pickle path generation.
+
+    Args:
+        config: arguments to pass to get_path
+        path: corresponding path
+    """
+    assert path == get_path(**config)
+
+
+@pytest.fixture(scope="function")
+def pickle_data_and_config(
+    config: Dict[str, Any], root: str, suffix: str
+) -> Tuple[Dict[str, Any], Dict[str, Any]]:
+
+    # setup test directory and save results to pickle file
+    os.mkdir(root)
+    save_results(**config, results=[], config={}, root=root, suffix=suffix)
+
+    # return pickle data and the config
+    with open(get_path(**config, root=root, suffix=suffix), "rb") as f:
+        yield pickle.load(f), config
+
+    # remove directory
+    shutil.rmtree(root)
+
+
+@pytest.mark.parametrize(
+    "config, root, suffix",
+    [
+        (
+            {
+                "layer": "linear",
+                "batch_size": 64,
+                "num_runs": 10,
+                "num_repeats": 100,
+                "random_seed": 13964,
+            },
+            "tests/tmp/",
+            "",
+        ),
+        (
+            {
+                "layer": "dpmha",
+                "batch_size": 16,
+                "num_runs": 20,
+                "num_repeats": 50,
+                "random_seed": 88362,
+                "forward_only": True,
+            },
+            "tests/tmp1/",
+            "no_3",
+        ),
+    ],
+)
+def test_save_results(
+    pickle_data_and_config: Tuple[Dict[str, Any], Dict[str, Any]]
+) -> None:
+    """Tests saving benchmark results.
+
+    Args:
+        pickle_data_and_config: tuple consisting of the pickle data as a dict and the
+            original data as a dict, as given in the config
+    """
+    pickle_data, config = pickle_data_and_config
+    for key, value in config.items():
+        assert pickle_data[key] == value