Tensor Parallel Inside notebook

[nivibilla]

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]

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]

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

[yanboliang]

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