tensor_parallel
v2.0.0: Sharding refactoring
尝试在Kaggle中的新40B LLMS演示
在多个GPU上以一条潜在的线性加速运行多个GPU上的大型Pytorch型号。
import transformers
import tensor_parallel as tp
tokenizer = transformers . AutoTokenizer . from_pretrained ( "facebook/opt-13b" )
model = transformers . AutoModelForCausalLM . from_pretrained ( "facebook/opt-13b" ) # use opt-125m for testing
model = tp . tensor_parallel ( model , [ "cuda:0" , "cuda:1" ]) # <- each GPU has half the weights
inputs = 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 usual 最新稳定版本（推荐）：
pip install tensor_parallel
出血边缘版本：
pip install https://github.com/BlackSamorez/tensor_parallel/archive/main.zip
只需将pytorch型号与tp.tensor_parallel包装并正常使用。为了获得最佳内存效率，请在模型仍在CPU上时致电tp.tensor_parallel 。
以下是一些用例：
examples/training_flan-t5-xl.ipynb在文本摘要中微调完整的flan-t5型号tensor_parallel int8 LLM使用llm.8bit + tensor_parallel的适配器调整大语言模型高级参数到tensor_parallel ：
device_ids: List[device] - 要使用的设备；默认为所有可用的GPUoutput_device: device - 模型输出将具有此设备tensor_parallel_config: tp.Config使用自定义并行策略，请参阅slicing_configs.pydistributed: bool如果为true，请使用Torch.Distribed Backend而不是线程（需要torchrun ）sharded: bool如果是的，请找到所有不被张量并行性拆分的可训练参数，并使用零3算法将它们分开。sharded_param_names: List[str] - 应该以这种方式碎片的参数名称，默认=自动找到为了保存模型，可以在非tensor_parallel上下文中使用，您应该使用save_tensor_parallel上下文包装器。
import torch
import transformers
import tensor_parallel as tp
model = tp . tensor_parallel (
transformers . AutoModelForCausalLM . from_pretrained ( "facebook/opt-13b" ),
)
# A whole lot of trainig...
with tp . save_tensor_parallel ( model ):
torch . save ( model . state_dict (), "/tmp/" )
# or
model . save_pretrained ( "/tmp/" )这样的代码保存了模型，就好像从未分开。它通过在state_dict创建期间收集模型零件来起作用。
通常，要通常创建并调度tensor_parallel模型，一个人需要内存中的整个模型。这可能很麻烦，但是还有另一种方法。
可以使用助手函数convert_state_dict将基本模型的state_dict转换为相应的tensor_parallel state_dict 。然后可以将状态DICE分配到模型中：
import accelerate
import transformers
import tensor_parallel as tp
# Initialize a weightless tensor_parallel model from MyModel
with accelerate . 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 MyModel
state_dict = torch . load ( "my_model_part_1_of_5.bin" )
# Convert it into a tensor_parallel state_dict
tensor_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 )
for param_name , param in state_dict . items ():
module_name = param_name
while len ( module_name ) > 0 and module_name not in device_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 )有了这一模型的一部分不超过一部分，就需要一次加载到内存中。
问：我没有多GPU服务器。我可以在Google Colab中使用Tensor_parally吗？
答： Colab有一个GPU，因此张量并行性没有意义。但是，Kaggle免费为所有电话验证的帐户免费提供两个T4。
问：什么是张量并行性？
答：您将每一层的权重分为一部分，将每个零件乘以单独的GPU，然后收集结果。在这里阅读更多
问：我应该使用TensorParallel还是DataParallel ？
答：对于大型型号的张量，较小的模型的数据平行
问：它如何与全沙dataparallear和零进行比较？
答：如果您可以安装大批量，则零是更好
为什么要使用tensor_parallel ...
alpaModel.parallelize()简而言之，请使用tensor_parallel在一台计算机上快速进行原型制作。使用DeepSpeed+Megatron或ALPA进行数百万美元的培训。
如果您遇到NCCL错误或随机悬挂，则可能会出现一些未正确显示的代码错误。要调试这些错误，我们建议使用export TENSOR_PARALLEL_USE_NATIVE=1或在单个设备上重新启动。
如果您发现错误或遇到问题，请向我们的问题跟踪器报告。我们将尽力提供帮助，但是我们可能需要一些时间才能实现。仅当您的问题专门使用tensor_parallel时，请创建问题。例如，如果您需要帮助安装transformers或优化代码，请在其他地方查找。
我们将黑色和ISORT用于所有拉请求。在进行代码之前，只需运行black . && isort .你会没事的。
