search

pytorch-labs / gpt-fast

Simple and efficient pytorch-native transformer text generation in <1000 LOC of python.
BSD 3-Clause "New" or "Revised" License
5.68k stars 514 forks source link

Tensor Parallel Inside notebook #167

Open nivibilla opened 7 months ago

nivibilla commented 7 months ago

Hi,

Im trying to get an example working with Ray on Databricks. Essentially having multiple replicas of the model. Is it possible to load a model with tensor parallelism inside a notebook?

Thanks

yanboliang commented 2 months ago

Yes, you can. gpt-fast uses torchrun to automatically dispatch tasks to different GPU nodes. You can use it just like regular python script.

kwen2501 commented 2 months ago

Is notebook compatible with torchrun? Last time I heard it has limited (or none) multi-processing support.

yanboliang commented 2 months ago

Right, I think you have to find a notebook which support multiprocessing.