Package: tensor-parallel
tensor_parallel🚀Try new 40B LLMs demo in KaggleRun large PyTorch models on multiple GPUs in one line of code with potentially linear speedup.importtransformersimporttensor_parallelastptokenizer=transformers.AutoTokenizer.from_pretrained("facebook/opt-13b")model=transformers.AutoModelForCausalLM.from_pretrained("facebook/opt-13b")# use opt-125m for testingmodel=tp.tensor_parallel(model,["cuda:0","cuda:1"])# <- each GPU has half the weightsinputs=tokenizer("A cat sat",return_tensors="pt")["input_ids"].to("cuda:0")outputs=model.generate(inputs,num_beams=5)print(tokenizer.decode(outputs[0]))# A cat sat on my lap for a few minutes ...model(input_ids=inputs,labels=inputs).loss.backward()# training works as usualInstallationLatest stable version (recommended):pip install tensor_parallelBleeding edge version:pip install https://github.com/BlackSamorez/tensor_parallel/archive/main.zipUsageSimply wrap your PyTorch model withtp.tensor_paralleland use it normally.
For best memory efficiency, calltp.tensor_parallelwhile the model is still on CPU.Here are a few use cases:examples/training_flan-t5-xl.ipynb- fine-tune full FLAN-T5 model on text summarizationtensor_parallel int8 LLM- adapter-tuning a large language model with LLM.8bit + tensor_parallelTBA- defining custom parallelism strategyAdvanced parameters totensor_parallel:device_ids: List[device]- which devices to use; defaults to all available GPUsoutput_device: device- model outputs will have this devicetensor_parallel_config: tp.Config- use custom parallelism strategy, seeslicing_configs.pydistributed: bool- if True, use torch.distributed backend instead of threading (requirestorchrun)sharded: bool- if True, find all trainable parameters that weren't split by Tensor Parallelism and split them usingZeRO-3 algorithm.weights will be split between GPUs and re-assembled before each forward passTL;DR use this when training to avoid duplicate parameters (enabled by default!)sharded_param_names: List[str]- parameter names that should be sharded this way, default = found automaticallySaving the modelTo save a model such that it could be used in a nontensor_parallelcontext, you should use asave_tensor_parallelcontext wrapper.importtorchimporttransformersimporttensor_parallelastpmodel=tp.tensor_parallel(transformers.AutoModelForCausalLM.from_pretrained("facebook/opt-13b"),)# A whole lot of trainig...withtp.save_tensor_parallel(model):torch.save(model.state_dict(),"/tmp/")# ormodel.save_pretrained("/tmp/")Such code saves a model as if it was never split. It works by gathering model parts duringstate_dictcreation.Memory efficient dispatchNormally, to normally create and dispatch atensor_parallelmodel, one needs the whole model in memory. This can be troublesome, but there is another way.It's possible to convert astate_dictof a basic model into the correspondingtensor_parallelstate_dictusing a helper functionconvert_state_dict. The state dict can then be dispatched and loaded into the model:importaccelerateimporttransformersimporttensor_parallelastp# Initialize a weightless tensor_parallel model from MyModelwithaccelerate.init_empty_weights():model=tp.TensorParallel(MyModel(),device_ids=[0,1]# and prepare it to be put on GPUs 0 and 1)# Load partial state_dict for MyModelstate_dict=torch.load("my_model_part_1_of_5.bin")# Convert it into a tensor_parallel state_dicttensor_parallel_state_dict=tp.convert_state_dict(state_dict,tensor_parallel_config=model.tensor_parallel_config,world_size=len(model.devices),)# Dispatch the partial state_dict (load_state_dict doesn't work with meta so here I use accelerate)device_map=tp.infer_sharded_device_map(model)forparam_name,paraminstate_dict.items():module_name=param_namewhilelen(module_name)>0andmodule_namenotindevice_map:module_name=".".join(module_name.split(".")[:-1])param_device=device_map[module_name]accelerate.utils.set_module_tensor_to_device(model,param_name,param_device,value=param)With this no more than one part of the model needs to be loaded into memory at once.FAQQ:I don't have a multi-GPU server. Can I use tensor_parallel in Google Colab?A:Colab has a single GPU, so there's no point in tensor parallelism. However,Kaggle offers two T4 for freeto all phone-verified accounts.Q:What is tensor parallelism?A:You split each layer's weights into parts, multiply each part on a separate GPU, then gather results. Read morehereQ:Should I useTensorParallelorDataParallel?A:TensorParallel for large models, DataParallel for smaller onesQ:How does it compare against FullyShardedDataParallel and ZeRO?A:ZeRO is better if you can fit a large batch, TensorParallel is better for small batchesWhy usetensor_parallel...v.s.DeepSpeedandFairScaleDeepSpeed has many parallelization strategies, but requires careful configurationtensor_parallel has one strategy that works with 1 line of codetensor_parallel works in a jupyter notebookv.s.MegatronLMMegatronLM hasgreattensor parallelism for one model architecturetensor_parallel hasgoodparallelism for any architecturetensor_parallel is way easier to installv.s.parallelformersparallelformers is inference-only, tensor_parallel supports trainingv.s.alpaalpa is a powerful tool for automatic distributed training / inference in JAXtensor_parallel works with PyTorchv.s.Model.parallelize()both are easy to use, both fit large modelsin parallelize, one GPU works at a timein tensor_parallel, GPUs work in parallelIn short, usetensor_parallelfor quick prototyping on a single machine.
Use DeepSpeed+Megatron or alpa for million-dollar training runs.TroubleshootingIf you experience NCCL errors, or random hanging, you may have some code errors that are not displayed properly.
To debug these errors, we recommend restarting withexport TENSOR_PARALLEL_USE_NATIVE=1or on a single device.If you found a bug or encountered a problem, please report it toour issue tracker.
We will do our best to help, but it may take some time before we get to it.
Please create issues only if your problem is specifically withtensor_parallel.
For example, if you need help installingtransformersor optimizing your code, please seek it elsewhere.Code styleWe useblackandisortfor all pull requests.
Before committing your code, simply runblack . && isort .and you will be fine.