Cladtorch and GPT-2 implementation (#1463)

* Tensor implementation

* GPT-2 implmentation: Add optimized training implementation for clad

---------

Co-authored-by: Rohan-T144 <74475417+Rohan-T144@users.noreply.github.com>

Author: r-timmaraju

.github/workflows/ci.yml

@@ -115,6 +115,7 @@ jobs:
             compiler: gcc-13
             clang-runtime: '20'
             extra_cmake_options: '-DCLAD_ENABLE_BENCHMARKS=On -DCLAD_ENABLE_ENZYME_BACKEND=On'
+            extra_packages: "libomp-20-dev libopenblas-dev"
             benchmark: true
 
           - name: ubu22-clang13-runtime13-cuda

.github/workflows/clang-tidy-review.yml

@@ -34,7 +34,7 @@ jobs:
         id: review
         with:
           build_dir: build
-          apt_packages: cmake,libxml2,libxml2-dev,libtinfo-dev,zlib1g-dev,libzstd-dev,libthrust-dev,libbenchmark-dev
+          apt_packages: cmake,libxml2,libxml2-dev,libtinfo-dev,zlib1g-dev,libzstd-dev,libthrust-dev,libbenchmark-dev,libomp-20-dev,libopenblas-dev
           exclude: "test/*,unittests/*"
           split_workflow: true
           cmake_command: >
@@ -44,6 +44,7 @@ jobs:
             -DCMAKE_BUILD_TYPE="Release"
             -DLLVM_EXTERNAL_LIT="`which lit`"
             -DCMAKE_EXPORT_COMPILE_COMMANDS=On
+            -DCMAKE_CXX_FLAGS="-fexceptions" # Clad demos can have exceptions
 
       - name: Upload artifacts
         uses: ZedThree/clang-tidy-review/upload@v0.20.1

CMakeLists.txt

@@ -6,7 +6,7 @@ endif()
 
 enable_language(CXX)
 set(CMAKE_CXX_EXTENSIONS NO)
-
+set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
 include(GNUInstallDirs)
 
 # MUST be done before call to clad project

benchmark/CMakeLists.txt

@@ -32,10 +32,33 @@ target_compile_definitions(tapenade_support PRIVATE
 CB_ADD_GBENCHMARK(VectorModeComparison VectorModeComparison.cpp)
 CB_ADD_GBENCHMARK(MemoryComplexity_tapenade MemoryComplexity.cpp)
 CB_ADD_GBENCHMARK(Multithreading Multithreading.cpp)
+CB_ADD_GBENCHMARK(Hessians Hessians.cpp)
+CB_ADD_GBENCHMARK(GPT2Training GPT2Training.cpp)
+
+if(APPLE)
+  # On macOS, we want to explicitly use the high-performance Accelerate framework for BLAS.
+  # We also need to give CMake a hint to find the Homebrew OpenMP library.
+  list(APPEND CMAKE_PREFIX_PATH "/opt/homebrew/opt/libomp")
+  find_package(OpenMP REQUIRED)
+  # Link explicitly against Accelerate and the found OpenMP library.
+  target_link_libraries(GPT2Training PRIVATE "-framework Accelerate" OpenMP::OpenMP_CXX)
+else()
+  # On all other platforms (e.g., Linux), find the best available BLAS and OpenMP.
+  find_package(BLAS REQUIRED)
+  find_package(OpenMP REQUIRED)
+  # find_package(MKL CONFIG REQUIRED)
+  target_link_libraries(GPT2Training PRIVATE ${BLAS_LIBRARIES} "-fopenmp=libomp" OpenMP::OpenMP_CXX)
+  target_compile_options(GPT2Training PRIVATE -I/usr/include/mkl)
+endif()
+target_compile_options(GPT2Training PRIVATE -O3 -ffast-math)
+target_compile_definitions(GPT2Training PRIVATE OMP)
+if (BLAS_FOUND)
+    target_compile_definitions(GPT2Training PRIVATE HAVE_CBLAS)
+endif()
 
 target_link_libraries(MemoryComplexity_tapenade PRIVATE tapenade_support)
 
 add_custom_target(benchmark-clad COMMAND ${CMAKE_CTEST_COMMAND} -V
   DEPENDS ${CLAD_BENCHMARK_DEPS} WORKING_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR})
 
-set_target_properties(benchmark-clad PROPERTIES FOLDER "Clad benchmarks")
+set_target_properties(benchmark-clad PROPERTIES FOLDER "Clad benchmarks")

benchmark/GPT2Training.cpp

@@ -0,0 +1,179 @@
+#include "benchmark/benchmark.h"
+
+#include <clad/Differentiator/CladtorchBuiltins.h>
+#include <clad/Differentiator/Differentiator.h>
+#include <clad/Differentiator/STLBuiltins.h>
+#include <cstddef>
+#include <string>
+#include "../demos/cladtorch/llm.hpp"
+#include "../demos/cladtorch/llm_opt.hpp"
+
+// NOLINTBEGIN(cppcoreguidelines-*)
+class GPT2Optimized : public benchmark::Fixture {
+public:
+  GPT2* model;
+  GPT2* d_model;
+  int* inputs;
+  int* targets;
+
+  void SetUp(const ::benchmark::State& state) override {
+    GPT2Config config{};
+    config.max_seq_len = 1024;
+    config.vocab_size = 50257;
+    config.padded_vocab_size = 50304;
+    config.num_layers = 12;
+    config.num_heads = 12;
+    config.channels = 768;
+
+    model = new GPT2(config);
+    d_model = new GPT2(config);
+
+    // Get batch size (B) and sequence length (T) from the benchmark state
+    int B = (int)state.range(0);
+    int T = (int)state.range(1);
+
+    model->allocate(B, T);
+    d_model->allocate(B, T);
+
+    // Allocate and fill dummy input data
+    inputs = new int[B * T];
+    targets = new int[B * T];
+    for (int i = 0; i < B * T; ++i) {
+      inputs[i] = i % model->config.vocab_size;
+      targets[i] = (i + 1) % model->config.vocab_size;
+    }
+  }
+
+  void TearDown(const ::benchmark::State& state) override {
+    // This runs once after each benchmark test
+    delete model;
+    delete d_model;
+    delete[] inputs;
+    delete[] targets;
+  }
+};
+
+static float gpt2forw_opt(GPT2* model, const int* inputs, const int* targets) {
+  model->forward(inputs, targets);
+  return model->mean_loss;
+}
+
+// The benchmark itself
+BENCHMARK_DEFINE_F(GPT2Optimized, FullTrainingIteration)
+(benchmark::State& state) {
+  auto grad = clad::gradient(gpt2forw_opt, "0");
+  int B = state.range(0);
+  int T = state.range(1);
+
+  for (auto _ : state) {
+    state.PauseTiming();
+    d_model->zero_all();
+
+    state.ResumeTiming();
+    // The single training iteration:
+    // forward pass (calculated as part of gradient), backward pass, and
+    // update
+    grad.execute(model, inputs, targets, d_model);
+    model->update(d_model, /*lr=*/1e-3F);
+  }
+  state.SetLabel("B=" + std::to_string(B) + " T=" + std::to_string(T));
+}
+
+BENCHMARK_REGISTER_F(GPT2Optimized, FullTrainingIteration)
+    ->Args({1, 16}) // B=1, T=16
+    ->Args({1, 32}) // B=1, T=32
+    ->Args({2, 16}) // B=2, T=16
+    ->Args({1, 64}) // B=1, T=64
+    ->Args({2, 32})
+    ->Args({4, 32})
+    ->Args({4, 64}) // B=4, T=64
+    ->Unit(benchmark::kMillisecond);
+
+class GPT2Cladtorch : public benchmark::Fixture {
+public:
+  gpt2::GPT2* model;
+  gpt2::GPT2* d_model;
+  int* inputs;
+  int* targets;
+
+  void SetUp(const ::benchmark::State& state) override {
+    const gpt2::Config config = {
+        .max_seq_len = 1024,
+        .vocab_size = 50257,
+        .padded_vocab_size = 50304,
+        .num_layers = 12,
+        .num_heads = 12,
+        .channels = 768,
+    };
+    model = new gpt2::GPT2(config);
+    d_model = new gpt2::GPT2(config);
+
+    // Get batch size (B) and sequence length (T) from the benchmark state
+    int B = state.range(0);
+    int T = state.range(1);
+
+    // Allocate and fill dummy input data
+    inputs = new int[B * T];
+    targets = new int[B * T];
+    for (int i = 0; i < B * T; ++i) {
+      inputs[i] = i % model->config.vocab_size;
+      targets[i] = (i + 1) % model->config.vocab_size;
+    }
+  }
+
+  void TearDown(const ::benchmark::State& state) override {
+    // This runs once after each benchmark test
+    delete model;
+    delete d_model;
+    delete[] inputs;
+    delete[] targets;
+  }
+};
+
+static float gpt2_loss(const gpt2::GPT2& model, const gpt2::ITensor& input,
+                       const gpt2::ITensor& targets) {
+  auto probs = model.forward(input);
+  auto loss = cross_entropy_loss(probs, targets);
+  return loss.scalar();
+}
+
+// The benchmark itself
+BENCHMARK_DEFINE_F(GPT2Cladtorch, FullTrainingIteration)
+(benchmark::State& state) {
+  auto grad = clad::gradient(gpt2_loss, "0");
+  int B = (int)state.range(0);
+  int T = (int)state.range(1);
+  const gpt2::ITensor inp({B, T}, inputs);
+  const gpt2::ITensor tar({B, T}, targets);
+  for (auto _ : state) {
+    state.PauseTiming();
+    d_model->for_each_parameter([&](gpt2::FTensor* t) { t->fill(0); });
+    state.ResumeTiming();
+    // The single training iteration: forward pass, backward pass, and update
+    grad.execute(*model, inp, tar, d_model);
+    std::vector<gpt2::FTensor*> params = model->get_parameter_tensors();
+    std::vector<gpt2::FTensor*> grads = d_model->get_parameter_tensors();
+    for (size_t i = 0; i < params.size(); ++i) {
+      // Update parameters with a learning rate of 1e-4
+      *params[i] += (*grads[i]) * -1e-3F;
+    }
+  }
+
+  // You can set custom counters to report B and T
+  state.SetLabel("B=" + std::to_string(B) + " T=" + std::to_string(T));
+}
+
+// Register the benchmark with different arguments
+// This will run the benchmark for various combinations of batch size (B) and
+// sequence length (T)
+BENCHMARK_REGISTER_F(GPT2Cladtorch, FullTrainingIteration)
+    ->Args({1, 16}) // B=1, T=16
+    ->Args({1, 32}) // B=1, T=32
+    ->Args({2, 16}) // B=2, T=16
+    ->Args({1, 64}) // B=1, T=64
+    ->Args({2, 32})
+    ->Unit(benchmark::kMillisecond);
+
+// Define our main.
+BENCHMARK_MAIN();
+// NOLINTEND(cppcoreguidelines-*)