🚀 A simple way to launch, train, and use PyTorch models on almost any device and distributed configuration, automatic mixed precision (including fp8), and easy-to-configure FSDP and DeepSpeed support

GPUGrants - a list of GPU grants that I can think of

Course to get into Large Language Models (LLMs) with roadmaps and Colab notebooks.

Ongoing research training transformer models at scale

Development repository for the Triton language and compiler

Solve puzzles. Learn CUDA.

GRadient-INformed MoE

The simplest, fastest repository for training/finetuning medium-sized GPTs.

Building blocks for foundation models.

OLMoE: Open Mixture-of-Experts Language Models

Efficient Triton Kernels for LLM Training

⛷️ LLaMA-MoE: Building Mixture-of-Experts from LLaMA with Continual Pre-training (EMNLP 2024)

A simple, performant and scalable Jax LLM!

A curated reading list of research in Adaptive Computation, Inference-Time Computation & Mixture of Experts (MoE).

Optax is a gradient processing and optimization library for JAX.

PyTorch extensions for high performance and large scale training.

100 exercises to learn JAX

Flexible and powerful tensor operations for readable and reliable code (for pytorch, jax, TF and others)

Fast and memory-efficient exact attention

The AI Scientist: Towards Fully Automated Open-Ended Scientific Discovery 🧑‍🔬

🚴 Call stack profiler for Python. Shows you why your code is slow!

Facebook AI Research Sequence-to-Sequence Toolkit written in Python.

Agentic components of the Llama Stack APIs

Multi-Agent Reinforcement Learning with JAX