diff --git a/README.md b/README.md
index 193fd94..7caa10a 100644
--- a/README.md
+++ b/README.md
@@ -16,6 +16,7 @@ This repository contains a variety of Determined examples that are not actively
| [LLM Finetuning 2](blog/llm-finetuning-2) | Finetuning Mistral-7B on Text-to-SQL using LoRA and DeepSpeed. |
| [LLM Finetuning 3](blog/llm-finetuning-3) | Finetuning Gemma-2B using DPO. |
| [Python SDK demo](blog/python_sdk_demo) | Example usage of the Determined Python SDK to run and administer experiments. |
+| [Tensor Parallelism](blog/tp) | Profiling tensor parallelism in PyTorch. |
## Computer Vision
diff --git a/blog/act-mem-2/README.md b/blog/act-mem-2/README.md
index 07f1547..9531f4c 100644
--- a/blog/act-mem-2/README.md
+++ b/blog/act-mem-2/README.md
@@ -9,3 +9,8 @@ memory.
- `attn_script.py` shows the cost of activation memory in the attention layer.
- Tests of the code are in `test.py`.
- See `requirements.txt` for versions the code was built against.
+
+
+## Contributors
+
+- [Garrett Goon](https://github.com/garrett361)
\ No newline at end of file
diff --git a/blog/tp/README.md b/blog/tp/README.md
new file mode 100644
index 0000000..79ea909
--- /dev/null
+++ b/blog/tp/README.md
@@ -0,0 +1,13 @@
+# Tensor Parallelism
+
+Code accompanying the deep-dive [blog post on Tensor Parallelism](https://determined.ai/blog/tp).
+
+- The MLP and TP MLP layers are in `layer.py`
+- Matmul profiling code in `matmul_profiling.py`
+- MLP TP profiling code in `tp_profiling.py`
+- Tests of the rearranging tensor sums are in `test_dot_product_{local,distributed}.py`
+
+
+## Contributors
+
+- [Garrett Goon](https://github.com/garrett361)
\ No newline at end of file
diff --git a/blog/tp/layers.py b/blog/tp/layers.py
new file mode 100644
index 0000000..f559742
--- /dev/null
+++ b/blog/tp/layers.py
@@ -0,0 +1,138 @@
+from typing import Any, Optional, Union
+
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+
+
+class MLP(nn.Module):
+ """
+ Basic MLP (multi-layer perceptron) layer. Dropout is neglected.
+ """
+
+ def __init__(
+ self,
+ d_model: int,
+ device: Optional[Union[str, torch.device]] = None,
+ dtype: Optional[torch.dtype] = None,
+ ) -> None:
+ super().__init__()
+
+ self.lin_0 = nn.Linear(d_model, 4 * d_model, device=device, dtype=dtype)
+ self.act_fn = nn.GELU()
+ self.lin_1 = nn.Linear(4 * d_model, d_model, device=device, dtype=dtype)
+
+ def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+ x = self.lin_0(inputs)
+ x = self.act_fn(x)
+ x = self.lin_1(x)
+ return x
+
+
+class AllReduceFwdIdentityBwd(torch.autograd.Function):
+ @staticmethod
+ def forward(
+ ctx: Any, inputs: torch.Tensor, group: Optional[dist.ProcessGroup] = None
+ ) -> torch.Tensor:
+ inputs = inputs.clone()
+ dist.all_reduce(inputs, group=group)
+ return inputs
+
+ @staticmethod
+ def backward(ctx: Any, grad_outputs: torch.Tensor) -> tuple[torch.Tensor, None]:
+ return grad_outputs, None
+
+
+class IdentityFwdAllReduceBwd(torch.autograd.Function):
+ @staticmethod
+ def forward(
+ ctx: Any, inputs: torch.Tensor, group: Optional[dist.ProcessGroup] = None
+ ) -> torch.Tensor:
+ ctx.group = group
+ return inputs
+
+ @staticmethod
+ def backward(ctx: Any, grad_outputs: torch.Tensor) -> tuple[torch.Tensor, None]:
+ grad_outputs = grad_outputs.clone()
+ dist.all_reduce(grad_outputs, group=ctx.group)
+ return grad_outputs, None
+
+
+class LinearShardedOutputs(nn.Linear):
+ def __init__(
+ self,
+ in_features: int,
+ out_features: int,
+ group: dist.ProcessGroup,
+ device: Optional[Union[str, torch.device]] = None,
+ dtype: Optional[torch.dtype] = None,
+ ) -> None:
+ sharded_out_features, remainder = divmod(out_features, group.size())
+ assert not remainder, "out_features must be divisible by the ProcessGroup size"
+ super().__init__(
+ in_features=in_features,
+ out_features=sharded_out_features,
+ device=device,
+ dtype=dtype,
+ )
+
+ self.group = group
+
+ def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+ # Wrap the unsharded inputs for backwards-pass correctness.
+ x = IdentityFwdAllReduceBwd.apply(inputs, self.group)
+ x = super().forward(x)
+ return x
+
+
+class LinearShardedInputs(nn.Linear):
+ def __init__(
+ self,
+ in_features: int,
+ out_features: int,
+ group: dist.ProcessGroup,
+ device: Optional[Union[str, torch.device]] = None,
+ dtype: Optional[torch.dtype] = None,
+ ) -> None:
+ sharded_in_features, remainder = divmod(in_features, group.size())
+ assert not remainder, "in_features must be divisible by the ProcessGroup size"
+ super().__init__(
+ in_features=sharded_in_features,
+ out_features=out_features,
+ device=device,
+ dtype=dtype,
+ )
+ self.group = group
+
+ def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+ x = inputs @ self.weight.T
+ # Wrap the mat-mul in an all-reduce forwards-pass correctness.
+ x = AllReduceFwdIdentityBwd.apply(x, self.group)
+ # Crucial: add the bias _after_ the all-reduce.
+ x = x + self.bias
+ return x
+
+
+class MLPTP(MLP):
+ """
+ Basic Tensor Parallel MLP (multi-layer perceptron) layer. Dropout is neglected.
+ """
+
+ def __init__(
+ self,
+ d_model: int,
+ group: Optional[dist.ProcessGroup] = None,
+ device: Optional[Union[str, torch.device]] = None,
+ dtype: Optional[torch.dtype] = None,
+ ) -> None:
+ nn.Module.__init__(self)
+ # Fallback to the WORLD process group, if None provided
+ group = group or dist.group.WORLD
+
+ self.lin_0 = LinearShardedOutputs(
+ d_model, 4 * d_model, group=group, device=device, dtype=dtype
+ )
+ self.act_fn = nn.GELU()
+ self.lin_1 = LinearShardedInputs(
+ 4 * d_model, d_model, group=group, device=device, dtype=dtype
+ )
diff --git a/blog/tp/matmul.png b/blog/tp/matmul.png
new file mode 100644
index 0000000..f23da74
Binary files /dev/null and b/blog/tp/matmul.png differ
diff --git a/blog/tp/matmul_profiling.py b/blog/tp/matmul_profiling.py
new file mode 100644
index 0000000..91ac4e6
--- /dev/null
+++ b/blog/tp/matmul_profiling.py
@@ -0,0 +1,91 @@
+import gc
+import logging
+
+import determined as det
+import torch
+
+import utils
+
+"""
+Script for profiling square matmuls on a single GPU.
+"""
+
+
+def profile_and_report(
+ core_context: det.core.Context,
+ d_model: int,
+ num_repeats: int,
+ num_warmups: int,
+ dtype: torch.dtype = torch.bfloat16,
+) -> None:
+ A = torch.randn(d_model, d_model, device="cuda", dtype=dtype)
+ B = torch.randn(d_model, d_model, device="cuda", dtype=dtype)
+
+ # Use CUDA events for accurate timing.
+ timer = utils.CUDAEventTimer()
+ torch.cuda.synchronize()
+
+ # Warmups
+ for _ in range(num_warmups):
+ A @ B
+
+ # Timed region.
+ for _ in range(num_repeats):
+ with timer:
+ A @ B
+
+ # Mean and std TFLOP computations
+ flops = 2 * d_model**3
+ time_s_t = torch.tensor(timer.time_s_list)
+ tflop_s_gpu_t = flops / time_s_t / 1e12
+ metrics = {
+ "d_model": d_model,
+ "time_s": timer.time_s_mean,
+ "time_s_std": timer.time_s_std,
+ "tflop_s_gpu": tflop_s_gpu_t.mean().item(),
+ "tflop_s_gpu_std": tflop_s_gpu_t.std().item(),
+ }
+
+ # Use d_model as the x-axis for plotting purposes.
+ core_context.train.report_metrics(group="matmul", steps_completed=d_model, metrics=metrics)
+
+ # Memory management
+ del A
+ del B
+ gc.collect()
+ torch.cuda.empty_cache()
+
+
+def main(
+ core_context: det.core.Context,
+ d_model_min: int,
+ d_model_max: int,
+ d_model_step: int,
+ num_repeats: int,
+ num_warmups: int,
+) -> None:
+ for d_model in range(d_model_min, d_model_max + 1, d_model_step):
+ profile_and_report(
+ core_context=core_context,
+ d_model=d_model,
+ num_repeats=num_repeats,
+ num_warmups=num_warmups,
+ )
+
+
+if __name__ == "__main__":
+ info = det.get_cluster_info()
+ assert info, "This script must run on a determined cluster."
+ hparams = info.trial.hparams
+
+ with det.core.init() as core_context:
+ logging.basicConfig(level=logging.INFO, format=det.LOG_FORMAT)
+
+ main(
+ core_context=core_context,
+ d_model_min=hparams["d_model_min"],
+ d_model_max=hparams["d_model_max"],
+ d_model_step=hparams["d_model_step"],
+ num_repeats=hparams["num_repeats"],
+ num_warmups=hparams["num_warmups"],
+ )
diff --git a/blog/tp/matmul_profiling.yaml b/blog/tp/matmul_profiling.yaml
new file mode 100644
index 0000000..97674dd
--- /dev/null
+++ b/blog/tp/matmul_profiling.yaml
@@ -0,0 +1,20 @@
+name: Matmul Profiling
+# Adjust the workspace and project names, as appropriate.
+workspace: TP Blog Post
+project: Matmul Profiling
+resources:
+ slots_per_trial: 1
+searcher:
+ name: single
+ metric: not_used
+ max_length: 1
+hyperparameters:
+ d_model_min: 256
+ d_model_max: 16384
+ d_model_step: 256
+ num_warmups: 5
+ num_repeats: 100
+entrypoint: >-
+ python3 -m determined.launch.torch_distributed
+ python3 matmul_profiling.py
+max_restarts: 0
diff --git a/blog/tp/mlp_tp.png b/blog/tp/mlp_tp.png
new file mode 100644
index 0000000..1aff737
Binary files /dev/null and b/blog/tp/mlp_tp.png differ
diff --git a/blog/tp/plots.ipynb b/blog/tp/plots.ipynb
new file mode 100644
index 0000000..a0663e6
--- /dev/null
+++ b/blog/tp/plots.ipynb
@@ -0,0 +1,470 @@
+{
+ "cells": [
+ {
+ "cell_type": "markdown",
+ "id": "6cf61f44-e76f-435a-a4c4-7882c1546d52",
+ "metadata": {},
+ "source": [
+ "# Plots"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 1,
+ "id": "1b461d37-b6de-42b7-9b03-063f3c46ef93",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import determined as det\n",
+ "from collections import defaultdict\n",
+ "from determined.experimental import client\n",
+ "import pandas as pd\n",
+ "import seaborn as sns\n",
+ "import matplotlib.pyplot as plt\n",
+ "\n",
+ "sns.set_theme(style=\"darkgrid\")\n",
+ "sns.set(rc={\"figure.figsize\": (7.5, 7.5)})\n"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "b1be1013-61b9-4290-95e7-891e9cbe0eb5",
+ "metadata": {},
+ "source": [
+ "Get all projects in the workspace"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 2,
+ "id": "27c9f4dc-f0b8-45c9-ba3e-5c1ff398b8d3",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "MLP TP Profiling\n",
+ "Matmul Profiling\n",
+ "obsolete\n"
+ ]
+ }
+ ],
+ "source": [
+ "workspace = client.get_workspace(\"TP Blog Post\")\n",
+ "projects_dict = {p.name: p for p in workspace.list_projects()}\n",
+ "for p in projects_dict:\n",
+ " print(p)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "5703a1f3-24b5-4d7c-acdd-3c14bf88f4af",
+ "metadata": {},
+ "source": [
+ "## Matmul Plots"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 3,
+ "id": "d94ce58b-497c-46de-ac9d-57c8a2c053e1",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "matmul_trial = projects_dict[\"Matmul Profiling\"].list_experiments()[0].list_trials()[0]"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 4,
+ "id": "996c7686-7c1d-4de2-8a8a-f46453272b41",
+ "metadata": {},
+ "outputs": [
+ {
+ "data": {
+ "text/html": [
+ "\n",
+ "\n",
+ "
\n",
+ " \n",
+ " \n",
+ " | \n",
+ " d_model | \n",
+ " TFLOP/sec/GPU | \n",
+ " TFLOP/sec/GPU_std | \n",
+ "
\n",
+ " \n",
+ " \n",
+ " \n",
+ " | 0 | \n",
+ " 256 | \n",
+ " 1.043684 | \n",
+ " 0.065087 | \n",
+ "
\n",
+ " \n",
+ " | 1 | \n",
+ " 512 | \n",
+ " 8.508206 | \n",
+ " 1.398408 | \n",
+ "
\n",
+ " \n",
+ " | 2 | \n",
+ " 768 | \n",
+ " 29.999840 | \n",
+ " 1.571498 | \n",
+ "
\n",
+ " \n",
+ " | 3 | \n",
+ " 1024 | \n",
+ " 49.371681 | \n",
+ " 3.053242 | \n",
+ "
\n",
+ " \n",
+ " | 4 | \n",
+ " 1280 | \n",
+ " 111.730354 | \n",
+ " 4.155525 | \n",
+ "
\n",
+ " \n",
+ " | ... | \n",
+ " ... | \n",
+ " ... | \n",
+ " ... | \n",
+ "
\n",
+ " \n",
+ " | 59 | \n",
+ " 15360 | \n",
+ " 267.599518 | \n",
+ " 2.319973 | \n",
+ "
\n",
+ " \n",
+ " | 60 | \n",
+ " 15616 | \n",
+ " 268.269440 | \n",
+ " 2.898224 | \n",
+ "
\n",
+ " \n",
+ " | 61 | \n",
+ " 15872 | \n",
+ " 267.012878 | \n",
+ " 4.683553 | \n",
+ "
\n",
+ " \n",
+ " | 62 | \n",
+ " 16128 | \n",
+ " 267.320831 | \n",
+ " 4.209327 | \n",
+ "
\n",
+ " \n",
+ " | 63 | \n",
+ " 16384 | \n",
+ " 268.287537 | \n",
+ " 3.332946 | \n",
+ "
\n",
+ " \n",
+ "
\n",
+ "
64 rows × 3 columns
\n",
+ "
"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ }
+ ],
+ "source": [
+ "matmul_plot = sns.scatterplot(matmul_results_df, x=\"d_model\", y=\"TFLOP/sec/GPU\")\n",
+ "plt.suptitle(\"Square Matmul TFLOP/sec\")\n",
+ "matmul_plot.axhline(y=312, color=\"r\", linestyle=\"--\")\n",
+ "matmul_plot.figure.savefig(\"matmul.png\", dpi=256, bbox_inches=\"tight\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "a309a6fa-8bd7-4559-b2e1-3774f0f530f8",
+ "metadata": {},
+ "source": [
+ "## MLP TP Plots"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 6,
+ "id": "5024162f-4d26-4d0e-a78d-6444e62248d7",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "mlp_tp_trial = projects_dict[\"MLP TP Profiling\"].list_experiments()[0].list_trials()[0]"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 7,
+ "id": "10445d10-70ba-4163-bfc4-a14cc3f7b64c",
+ "metadata": {},
+ "outputs": [
+ {
+ "data": {
+ "text/html": [
+ "\n",
+ "\n",
+ "
\n",
+ " \n",
+ " \n",
+ " | \n",
+ " tp_degree | \n",
+ " d_model | \n",
+ " time_s | \n",
+ " TFLOP/sec/GPU | \n",
+ " TFLOP/sec/GPU_std | \n",
+ "
\n",
+ " \n",
+ " \n",
+ " \n",
+ " | 0 | \n",
+ " 1 | \n",
+ " 1024 | \n",
+ " 0.000492 | \n",
+ " 139.760605 | \n",
+ " 1.070486 | \n",
+ "
\n",
+ " \n",
+ " | 1 | \n",
+ " 1 | \n",
+ " 1536 | \n",
+ " 0.000965 | \n",
+ " 160.297211 | \n",
+ " 0.191192 | \n",
+ "
\n",
+ " \n",
+ " | 2 | \n",
+ " 1 | \n",
+ " 2048 | \n",
+ " 0.001459 | \n",
+ " 188.433899 | \n",
+ " 3.754535 | \n",
+ "
\n",
+ " \n",
+ " | 3 | \n",
+ " 1 | \n",
+ " 2560 | \n",
+ " 0.002080 | \n",
+ " 206.486206 | \n",
+ " 2.993310 | \n",
+ "
\n",
+ " \n",
+ " | 4 | \n",
+ " 1 | \n",
+ " 3072 | \n",
+ " 0.002941 | \n",
+ " 210.291397 | \n",
+ " 2.333767 | \n",
+ "
\n",
+ " \n",
+ " | ... | \n",
+ " ... | \n",
+ " ... | \n",
+ " ... | \n",
+ " ... | \n",
+ " ... | \n",
+ "
\n",
+ " \n",
+ " | 151 | \n",
+ " 8 | \n",
+ " 18432 | \n",
+ " 0.012830 | \n",
+ " 216.637466 | \n",
+ " 2.120074 | \n",
+ "
\n",
+ " \n",
+ " | 152 | \n",
+ " 8 | \n",
+ " 18944 | \n",
+ " 0.013551 | \n",
+ " 216.700256 | \n",
+ " 3.428361 | \n",
+ "
\n",
+ " \n",
+ " | 153 | \n",
+ " 8 | \n",
+ " 19456 | \n",
+ " 0.014483 | \n",
+ " 213.789810 | \n",
+ " 2.577091 | \n",
+ "
\n",
+ " \n",
+ " | 154 | \n",
+ " 8 | \n",
+ " 19968 | \n",
+ " 0.015056 | \n",
+ " 216.667892 | \n",
+ " 3.947501 | \n",
+ "
\n",
+ " \n",
+ " | 155 | \n",
+ " 8 | \n",
+ " 20480 | \n",
+ " 0.015662 | \n",
+ " 219.187225 | \n",
+ " 4.835044 | \n",
+ "
\n",
+ " \n",
+ "
\n",
+ "
156 rows × 5 columns
\n",
+ "