Kakao Brain announces torchgpipe 33, an implementation of GPipe 9 in PyTorch as a handy library.
from torchgpipe import GPipe
model = nn.Sequential(a, b, c, d)
model = GPipe(model, balance=[1, 1, 1, 1], chunks=8)
output = model(input)

GPipe is a scalable pipeline parallelism library published by Google Brain. It leverages the training of a giant model which requires much memory. For instance, Google trained AmoebaNet-B with 557M parameters over GPipe.

  
      
      GitHub
  
  
    

kakaobrain/torchgpipe 33

A GPipe implementation in PyTorch. Contribute to kakaobrain/torchgpipe development by creating an account on GitHub.