blog: tp (#22)

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
 

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)

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)

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

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"],
+        )

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