A numpy-only implementation of a Convolutional Neural Network, from the ground up.

The absolute most basic example of AlphaZero and Monte Carlo Tree Search I could come up with

Code for STaR: Bootstrapping Reasoning With Reasoning (NeurIPS 2022)

General Reasoner: Advancing LLM Reasoning Across All Domains [NeurIPS25]

LLM Unlearning

Implementation of FlashAttention in PyTorch

Variational Autoencoder implemented with PyTorch, Trained over CelebA Dataset

Recreating PyTorch from scratch (C/C++, CUDA, NCCL and Python, with multi-GPU support and automatic differentiation!)

Code for reproducing results in the sliced score matching paper (UAI 2019)

cnn

Flash Attention Triton kernel with support for second-order derivatives

Large scale 4D parallelism pre-training for 🤗 transformers in Mixture of Experts *(still work in progress)*

[ICML 2024 Oral] Official code repository for MLLM-as-a-Judge.

Pipeline Parallelism Emulation and Visualization

Utility scripts for PyTorch (e.g. Make Perfetto show some disappearing kernels, Memory profiler that understands more low-level allocations such as NCCL, ...)

[ICML 2024] On Discrete Prompt Optimization for Diffusion Models - Google

ICML 2024 - Official Repository for EXO: Towards Efficient Exact Optimization of Language Model Alignment

A re-implementation of the "Red Teaming Language Models with Language Models" paper by Perez et al., 2022

GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection

The official evaluation suite and dynamic data release for MixEval.

Safe-RLHF: Constrained Value Alignment via Safe Reinforcement Learning from Human Feedback

Example code demonstrating how LLMs could be Red Teamed.

A Port of Fréchet Inception Distance (FID score) to PyTorch