https://github.com/felafax/felafax Skip to content Navigation Menu Toggle navigation Sign in * Product + Actions Automate any workflow + Packages Host and manage packages + Security Find and fix vulnerabilities + Codespaces Instant dev environments + GitHub Copilot Write better code with AI + Code review Manage code changes + Issues Plan and track work + Discussions Collaborate outside of code Explore + All features + Documentation + GitHub Skills + Blog * Solutions By size + Enterprise + Teams + Startups By industry + Healthcare + Financial services + Manufacturing By use case + CI/CD & Automation + DevOps + DevSecOps * Resources Topics + AI + DevOps + Security + Software Development + View all Explore + Learning Pathways + White papers, Ebooks, Webinars + Customer Stories + Partners * Open Source + GitHub Sponsors Fund open source developers + The ReadME Project GitHub community articles Repositories + Topics + Trending + Collections * Enterprise + Enterprise platform AI-powered developer platform Available add-ons + Advanced Security Enterprise-grade security features + GitHub Copilot Enterprise-grade AI features + Premium Support Enterprise-grade 24/7 support * Pricing Search or jump to... Search code, repositories, users, issues, pull requests... Search [ ] Clear Search syntax tips Provide feedback We read every piece of feedback, and take your input very seriously. [ ] [ ] Include my email address so I can be contacted Cancel Submit feedback Saved searches Use saved searches to filter your results more quickly Name [ ] Query [ ] To see all available qualifiers, see our documentation. Cancel Create saved search Sign in Sign up Reseting focus You signed in with another tab or window. Reload to refresh your session. You signed out in another tab or window. Reload to refresh your session. You switched accounts on another tab or window. Reload to refresh your session. Dismiss alert {{ message }} felafax / felafax Public * Notifications You must be signed in to change notification settings * Fork 13 * Star 250 Felafax is building AI infra for non-NVIDIA GPUs felafax.ai License Apache-2.0 license 250 stars 13 forks Branches Tags Activity Star Notifications You must be signed in to change notification settings * Code * Issues 0 * Pull requests 0 * Actions * Projects 0 * Security * Insights Additional navigation options * Code * Issues * Pull requests * Actions * Projects * Security * Insights felafax/felafax This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository. main BranchesTags Go to file Code Folders and files Name Name Last commit Last commit message date Latest commit History 320 Commits gemma2 gemma2 llama3_jax llama3_jax llama3_pytorch_xla llama3_pytorch_xla utils utils .gitignore .gitignore LICENSE LICENSE README.md README.md launch_tuner.sh launch_tuner.sh pyproject.toml pyproject.toml requirements.txt requirements.txt View all files Repository files navigation * README * Apache-2.0 license Felafax -- tune LLaMa3.1 on Google Cloud TPUs for 30% lower cost and scale seamlessly! image RoadRunnerX is a framework for continued-training and fine-tuning open source LLMs using XLA runtime. We take care of necessary runtime setup and provide a Jupyter notebook out-of-box to just get started. * Easy to use. * Easy to configure all aspects of training (designed for ML researchers and hackers). * Easy to scale training from a single TPU VM with 8 cores to entire TPU Pod containing 6000 TPU cores (1000X)! Goal Our goal at felafax is to build infra to make it easier to run AI workloads on non-NVIDIA hardware (TPU, AWS Trainium, AMD GPU, and Intel GPU). Currently supported models * LLaMa-3.1 JAX Implementation $${\color{red}New!}$$ + Converted from PyTorch to JAX for improved performance + By default, runs 2-way data parallel and 2-way model parallel training (2 data parallel model copies and each model copy is sharded across two TPU chips). + Compatible with NVIDIA GPUs and TPUs + Full-precision training support * LLaMa-3/3.1 PyTorch XLA + LoRA and full-precision training support * Gemma2 Models (2B, 9B, 27B) + Optimized for Cloud TPUs + Fast full-precision training Setup For a hosted version with a seamless workflow, please request access here. Here is a demo of our platform (demo) . If you prefer a self-hosted training version, follow the instructions below. These steps will guide you through launching a TPU VM on your Google Cloud account and starting a Jupyter notebook. With just 3 simple steps, you'll be up and running in under 10 minutes. 1. Install gcloud command-line tool and authenticate your account (SKIP this STEP if you already have gcloud installed and have used TPUs before! ) # Download gcloud CLI curl https://sdk.cloud.google.com | bash source ~/.bashrc # Authenticate gcloud CLI gcloud auth login # Create a new project for now gcloud projects create LLaMa3-tunerX --set-as-default # Config SSH and add gcloud compute config-ssh --quiet # Set up default credentials gcloud auth application-default login # Enable Cloud TPU API access gcloud services enable compute.googleapis.com tpu.googleapis.com storage-component.googleapis.com aiplatform.googleapis.com 2. Spin up a TPU v5-8 VM . sh ./launch_tuner.sh Keep an eye on the terminal -- you might be asked to input SSH key password and need to put in your HuggingFace token. 3. Clone the repo and install dependencies git clone https://github.com/felafax/felafax.git cd felafax pip install -r requirements.txt 4. Open the Jupyter notebook at https://localhost:888 and start fine-tuning! Credits: * Google Deepmind's Gemma repo. * EasyLM for great work on llama models in JAX * PyTorch XLA FSDP and SPMD testing done by HeegyuKim. * Examples from PyTorch-XLA repo. Contact If you have any questions, please contact us at founders@felafax.ai. About Felafax is building AI infra for non-NVIDIA GPUs felafax.ai Topics ai infra jax xla Resources Readme License Apache-2.0 license Activity Custom properties Stars 250 stars Watchers 2 watching Forks 13 forks Report repository Releases No releases published Packages 0 No packages published Contributors 3 * @felarof99 felarof99 Felarof * @shadowfax92 shadowfax92 Nikhil * @emrekenci emrekenci Emre Kenci Languages * Python 58.5% * Jupyter Notebook 37.6% * Shell 3.9% Footer (c) 2024 GitHub, Inc. Footer navigation * Terms * Privacy * Security * Status * Docs * Contact * Manage cookies * Do not share my personal information You can't perform that action at this time.