From b239b67ae1a2f764412ed6f3090a5a0449df0964 Mon Sep 17 00:00:00 2001
From: Franz Louis Cesista <franzlouiscesista@gmail.com>
Date: Wed, 1 May 2024 00:15:08 +0800
Subject: [PATCH] add modal script

---
 dev/cuda/README.md             |  6 +++
 dev/cuda/benchmark_on_modal.py | 82 ++++++++++++++++++++++++++++++++++
 2 files changed, 88 insertions(+)
 create mode 100644 dev/cuda/benchmark_on_modal.py

diff --git a/dev/cuda/README.md b/dev/cuda/README.md
index 041ff00..d020cf6 100644
--- a/dev/cuda/README.md
+++ b/dev/cuda/README.md
@@ -31,3 +31,9 @@ You'll see that this first forwards the reference code on the CPU, then it runs
 You'll see that this matches all the CPU results but runs much much faster. The typical process from here on is we copy paste the kernel that ran fastest, adjust it manually (e.g. to hardcode the best block size) and drop it into the training code file, e.g. `train_gpt2.cu`.
 
 To add a new version of a kernel, add the kernel to the corresponding file and adjust the docs. To add a new kernel, add the new file and adjust the Makefile. Run `make clean` to clean up binaries from your directory.
+
+If you do not have a GPU or is having trouble with CUDA dependencies, you can run the benchmarks on the [Modal platform](http://modal.com). For example, to run the benchmark for the attention forward pass on an A100 GPU with 80GB of memory, you can run the following command:
+
+```bash
+GPU_MEM=80 modal run benchmark_on_modal.py --compile-command "nvcc -O3 --use_fast_math attention_forward.cu -o attention_forward -lcublas" --run-command "./attention_forward 1"
+```
diff --git a/dev/cuda/benchmark_on_modal.py b/dev/cuda/benchmark_on_modal.py
new file mode 100644
index 0000000..e597538
--- /dev/null
+++ b/dev/cuda/benchmark_on_modal.py
@@ -0,0 +1,82 @@
+"""
+Script for running benchmarks on the Modal platform.
+This is useful for folks who do not have access to expensive GPUs locally.
+
+Example usage:
+GPU_MEM=80 modal run benchmark_on_modal.py \
+    --compile-command "nvcc -O3 --use_fast_math attention_forward.cu -o attention_forward -lcublas" \
+    --run-command "./attention_forward 1"
+
+This will mount the contents of the current directory to the remote container on modal,
+compile the `attention_forward.cu` file with `nvcc`, and run the resulting binary on a A100 GPU with 80GB of memory.
+"""
+
+import subprocess
+import os
+import sys
+
+import modal
+from modal import Image, Stub
+
+GPU_NAME_TO_MODAL_CLASS_MAP = {
+    "H100": modal.gpu.H100,
+    "A100": modal.gpu.A100,
+    "A10G": modal.gpu.A10G,
+}
+
+N_GPUS = int(os.environ.get("N_GPUS", 1))
+GPU_MEM = int(os.environ.get("GPU_MEM", 40))
+GPU_NAME = os.environ.get("GPU_NAME", "A100")
+GPU_CONFIG = GPU_NAME_TO_MODAL_CLASS_MAP[GPU_NAME](count=N_GPUS, memory=GPU_MEM)
+
+APP_NAME = "llm.c benchmark run"
+
+# We don't actually need to use the Axolotl image here, but it's reliable
+AXOLOTL_REGISTRY_SHA = (
+    "d5b941ba2293534c01c23202c8fc459fd2a169871fa5e6c45cb00f363d474b6a"
+)
+axolotl_image = (
+    Image.from_registry(f"winglian/axolotl@sha256:{AXOLOTL_REGISTRY_SHA}")
+    .run_commands(
+        "git clone https://github.com/OpenAccess-AI-Collective/axolotl /root/axolotl",
+        "cd /root/axolotl && git checkout v0.4.0",
+    )
+    .pip_install("huggingface_hub==0.20.3", "hf-transfer==0.1.5")
+    .env(
+        dict(
+            HUGGINGFACE_HUB_CACHE="/pretrained",
+            HF_HUB_ENABLE_HF_TRANSFER="1",
+            TQDM_DISABLE="true",
+        )
+    )
+)
+
+stub = Stub(APP_NAME)
+
+
+def execute_command(command: str):
+    command_args = command.split(" ")
+    print(f"{command_args = }")
+    subprocess.run(command_args, stdout=sys.stdout, stderr=subprocess.STDOUT)
+
+
+@stub.function(
+    gpu=GPU_CONFIG,
+    image=axolotl_image,
+    allow_concurrent_inputs=4,
+    container_idle_timeout=900,
+    # This copies everything in this folder to the remote root folder
+    mounts=[modal.Mount.from_local_dir("./", remote_path="/root/")]
+)
+def run_benchmark(compile_command: str, run_command: str):
+    execute_command("pwd")
+    execute_command("ls")
+    execute_command(compile_command)
+    execute_command(run_command)
+    return None
+
+
+@stub.local_entrypoint()
+def inference_main(compile_command: str, run_command: str):
+    results = run_benchmark.remote(compile_command, run_command)
+    return results