kakaobrain/torchgpipe
A GPipe implementation in PyTorch
This tool helps machine learning engineers train extremely large neural networks that might otherwise exceed the memory capacity of a single GPU. It takes your existing PyTorch model and training data, then intelligently splits the model and data across multiple GPUs. The output is a successfully trained, massive model, enabling research and development of state-of-the-art AI.
862 stars. No commits in the last 6 months. Available on PyPI.
Use this if you are a machine learning engineer or researcher encountering 'out of memory' errors when trying to train very large PyTorch models on GPUs.
Not ideal if your models are small enough to train comfortably on a single GPU or if you are not using PyTorch and CUDA-enabled devices.
Stars
862
Forks
98
Language
Python
License
BSD-3-Clause
Category
Last pushed
Jul 25, 2024
Commits (30d)
0
Dependencies
1
Get this data via API
curl "https://pt-edge.onrender.com/api/v1/quality/ml-frameworks/kakaobrain/torchgpipe"
Open to everyone — 100 requests/day, no key needed. Get a free key for 1,000/day.
Related frameworks
deepspeedai/DeepSpeed
DeepSpeed is a deep learning optimization library that makes distributed training and inference...
helmholtz-analytics/heat
Distributed tensors and Machine Learning framework with GPU and MPI acceleration in Python
hpcaitech/ColossalAI
Making large AI models cheaper, faster and more accessible
horovod/horovod
Distributed training framework for TensorFlow, Keras, PyTorch, and Apache MXNet.
bsc-wdc/dislib
The Distributed Computing library for python implemented using PyCOMPSs programming model for HPC.