Introduction

I am Rohan Timmaraju, a Computer Science student at Columbia University. During Google Summer of Code 2025, I will be working on the “Enhancing LLM Training Efficiency with Clad for Automatic Differentiation” project with the Compiler Research group.

Mentors: Vassil Vassilev, David Lange, Jonas Rembser, Christina Koutsou

About LLM Training

Large Language Models (LLMs) like ChatGPT have revolutionized AI, but their training is incredibly computationally intensive. Currently, Python-based frameworks such as PyTorch and TensorFlow are the go-to tools. While they offer excellent flexibility and a rich ecosystem, their reliance on interpreted execution and dynamic computation graphs can lead to performance bottlenecks and high memory consumption. This is particularly noticeable when we consider deploying or training these models in resource-constrained environments or within C++-centric high-performance computing (HPC) setups, which are common in scientific research.

While C++ provides the tools for fine-grained control over system resources and has proven its capabilities in efficient LLM inference (as seen with projects like llama.cpp), the critical component for training – flexible and efficient Automatic Differentiation (AD) – presents an ongoing challenge for C++ solutions.

Why Use Clad?

This project proposes to tackle this challenge by integrating Clad, an Automatic Differentiation plugin for the Clang compiler. Unlike traditional AD libraries that often operate at runtime, Clad performs source-to-source transformation. It analyzes the C++ Abstract Syntax Tree (AST) at compile time and generates optimized C++ code for computing derivatives. This compiler-level approach has the potential to reduce runtime overhead and improve memory efficiency compared to dynamic methods.

To facilitate this integration, I am developing a custom C++ tensor library to be used in neural network training. Inspired by the powerful approaches of libraries such as llm.c and pytorch, this library is being designed from the ground up with Clad compatibility in mind. The core idea is to replace manual or internally managed gradient computations with Clad’s reverse-mode AD (as in clad::gradient) for key LLM operations like matrix multiplications, activation functions, normalization layers, and the final loss function.

Implementation Plan

  1. Foundation & Baseline: The implementation will start by implementing a complete GPT-2 training loop in C++ without Clad. This will serve as our performance baseline. GPT-2 is chosen here as a relatively simple open-source LLM architecture capable of being trained on local devices. This could be extended to other architectures like Llama or Mistral.
  2. Core Clad Integration Strategy: We will investigate and evaluate different strategies for applying Clad to tensor network gradient calculations, potentially also identifying potential areas where Clad itself could be enhanced for deep learning workloads.
  3. Expanding Integration: Once a promising strategy is identified and validated on simpler operations, we’ll systematically integrate Clad into more complex components of the GPT-2 architecture.
  4. Benchmarking & Optimization: Benchmarking against our baseline will be crucial to quantify the performance gains (speed, memory). We’ll also use profiling tools to identify bottlenecks and optimize the tensor library with Clad. OpenMP may be employed for parallelization to further boost performance.
  5. Documentation & Potential Extensions: Thorough documentation of the tensor library, the Clad integration process, and our findings will also be a primary focus. Time permitting, we’ll explore extending this work to other LLM architectures like Llama.

Conclusion

By successfully integrating Clad into a C++ LLM training pipeline, we aim to:

  • Demonstrate Performance Gains: Show tangible improvements in training speed and memory efficiency.
  • Clad for ML: Provide a significant real-world use case, potentially identifying areas for Clad’s improvement in supporting ML tasks.
  • Offer a C++ Alternative: Provide a foundation for more efficient, compiler-driven LLM training within the C++ ecosystems.
  • Learn and Share: Gain insights into the practicalities of applying compiler-based AD to complex ML problems and share these learnings with the community.

I believe this project has the potential to make a valuable contribution to both the compiler research field and the ongoing efforts to make powerful AI models more accessible and efficient to train.