A native PyTorch Library for large model training

PyTorch per step fault tolerance (actively under development)

NVIDIA Resiliency Extension is a python package for framework developers and users to implement fault-tolerant features. It improves the effective training time by minimizing the downtime due to fa…

Ongoing research training transformer models at scale

A scalable generative AI framework built for researchers and developers working on Large Language Models, Multimodal, and Speech AI (Automatic Speech Recognition and Text-to-Speech)

20+ high-performance LLMs with recipes to pretrain, finetune and deploy at scale.

Pretrain, finetune and deploy AI models on multiple GPUs, TPUs with zero code changes.

A PyTorch Native LLM Training Framework

AI on GKE is a collection of examples, best-practices, and prebuilt solutions to help build, deploy, and scale AI Platforms on Google Kubernetes Engine

A simple, performant and scalable Jax LLM!

xpk (Accelerated Processing Kit, pronounced x-p-k,) is a software tool to help Cloud developers to orchestrate training jobs on accelerators such as TPUs and GPUs on GKE.

The TinyLlama project is an open endeavor to pretrain a 1.1B Llama model on 3 trillion tokens.

Reference models and tools for Cloud TPUs.

Orbax provides common utility libraries for JAX users.

Ray is a unified framework for scaling AI and Python applications. Ray consists of a core distributed runtime and a toolkit of libraries (Ray AIR) for accelerating ML workloads.

🤗 Diffusers: State-of-the-art diffusion models for image and audio generation in PyTorch

Lingvo

A flexible, high-performance serving system for machine learning models

Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more

Testing framework for Deep Learning models (Tensorflow and PyTorch) on Google Cloud hardware accelerators (TPU and GPU)

Enabling PyTorch on Google TPU

Minimal library to train LLMs on TPU in JAX with pjit().

A repository for scripts and notebooks for the UCSD big data course