tensor_parallel实战训练FLAN-T5-xl模型的完整案例【免费下载链接】tensor_parallelAutomatically split your PyTorch models on multiple GPUs for training inference项目地址: https://gitcode.com/gh_mirrors/te/tensor_parallelTensor Parallel是一个强大的Python库能够自动将PyTorch模型分割到多个GPU上进行训练和推理帮助开发者轻松实现大规模模型的分布式训练。本文将通过一个完整案例展示如何使用tensor_parallel库训练FLAN-T5-xl模型即使在普通硬件上也能高效运行。准备工作环境搭建与依赖安装首先我们需要安装必要的依赖库。打开终端执行以下命令pip install -q tensor_parallel py7zr transformers datasets这个命令会安装tensor_parallel库及其依赖的transformers、datasets等工具包。其中tensor_parallel是实现模型并行的核心库transformers提供了FLAN-T5-xl模型的实现datasets则用于加载训练数据。案例背景FLAN-T5-xl模型与任务介绍本案例将使用FLAN-T5-xl模型在samsum对话摘要数据集上进行微调。FLAN-T5-xl是一个拥有30亿参数的大型预训练Transformer模型比BERT-large大10倍左右通常需要大量GPU内存才能运行。通过tensor_parallel库我们可以在普通硬件上训练这样的大型模型。例如原始案例在四台老旧的1080Ti显卡上成功运行证明了该方法的实用性。如果你使用更高级的硬件可以适当调整批次大小以获得更好的性能。步骤一加载预训练模型与分词器首先我们需要加载FLAN-T5-xl模型和对应的分词器。使用transformers库可以轻松实现这一点import torch import tensor_parallel as tp import transformers tokenizer transformers.AutoTokenizer.from_pretrained(google/flan-t5-xl) model transformers.T5ForConditionalGeneration.from_pretrained( google/flan-t5-xl, torch_dtypetorch.bfloat16, low_cpu_mem_usageTrue, offload_state_dictTrue )这里我们使用了bfloat16数据类型来减少内存占用并启用了低CPU内存使用模式帮助在资源有限的环境中加载模型。步骤二应用Tensor Parallel实现模型并行接下来我们只需一行代码即可将模型转换为并行模式使其能够在多个GPU上运行model tp.tensor_parallel(model)执行这行代码后tensor_parallel会自动将模型分割到可用的GPU上。如果一切顺利你会看到类似以下的输出Using automatic config: no tensor parallel config provided and no predefined configs can be used Using ZeRO-3 sharding for 249856 non tensor-parallel parameters这表明tensor_parallel正在使用自动配置将模型参数分配到多个GPU并对非张量并行参数使用ZeRO-3分片技术。步骤三验证模型并行是否正常工作为了确保模型并行设置正确我们可以进行一次简单的前向和反向传播测试input_ids tokenizer(A cat sat on a mat, return_tensorspt).input_ids.to(cuda) output_ids tokenizer(A cat sat did not sit on a mat, return_tensorspt).input_ids.to(cuda) # 前向和反向传播 loss model(input_idsinput_ids, labelsoutput_ids).loss loss.backward() # 检查nvidia-smi查看GPU内存使用情况运行这段代码后你可以通过nvidia-smi命令查看各个GPU的内存使用情况。如果模型成功并行化你应该会看到多个GPU都有内存占用。步骤四数据准备与预处理我们使用samsum数据集进行对话摘要任务。首先加载数据集import datasets data datasets.load_dataset(samsum) print(Example:, data[train][25])这会输出数据集中的一个示例包含对话内容和对应的摘要。接下来我们需要对数据进行预处理将对话文本转换为模型可以接受的格式def preprocess_function(examples, prefixsummarize:): inputs, targets examples[dialogue], examples[summary] inputs [prefix inp for inp in inputs] model_inputs tokenizer(inputs, max_length256, truncationTrue) labels tokenizer(text_targettargets, max_length256, truncationTrue) model_inputs[labels] labels[input_ids] return model_inputs train_data data[train].filter(lambda row: row[dialogue] and row[summary]).map( preprocess_function, batchedTrue, remove_columns[id, dialogue, summary] )这个预处理函数会将对话文本加上summarize:前缀然后使用分词器将文本转换为模型输入的token ID。同时我们也对标签摘要进行同样的处理。步骤五配置训练参数并开始训练现在我们可以配置训练参数并开始训练了。使用transformers库的Seq2SeqTrainer可以简化训练过程trainer transformers.Seq2SeqTrainer( modelmodel, train_datasettrain_data, argstransformers.Seq2SeqTrainingArguments( do_trainTrue, remove_unused_columnsFalse, per_device_train_batch_size4, gradient_accumulation_steps2, optimadafactor, warmup_steps250, max_steps1000, learning_rate1e-5, logging_steps1, output_diroutputs ), tokenizertokenizer, data_collatortransformers.DataCollatorForSeq2Seq( tokenizer, paddingTrue, max_length512, pad_to_multiple_of8 ) )这里我们设置了每设备训练批次大小为4梯度累积步数为2使用adafactor优化器并将最大训练步数设为1000。这些参数可以根据你的硬件配置进行调整。最后执行以下命令开始训练trainer.train()训练过程中你会看到类似以下的输出显示训练进度和损失值***** Running training ***** Num examples 14731 Num Epochs 1 Instantaneous batch size per device 4 Total train batch size (w. parallel, distributed accumulation) 8 Gradient Accumulation steps 2 Total optimization steps 1000 Number of trainable parameters 2849757184总结与扩展通过本文的案例我们展示了如何使用tensor_parallel库在多个GPU上训练大型FLAN-T5-xl模型。整个过程只需添加一行代码model tp.tensor_parallel(model)即可实现模型的自动并行化大大降低了分布式训练的门槛。这个案例保持了简洁性重点展示了tensor_parallel的使用方法。如果你想训练更高级的摘要模型可以参考transformers库的官方示例transformers/examples/pytorch/summarization只需在加载模型后添加tp.tensor_parallel(model)这一行代码即可实现并行训练。tensor_parallel库的核心实现位于src/tensor_parallel/tensor_parallel.py其中的TensorParallel类负责将模型分割到多个设备上并协调它们之间的通信。通过这种方式即使是没有分布式训练经验的开发者也能轻松地在多个GPU上训练大型模型。希望这个案例能帮助你更好地理解和使用tensor_parallel库让大规模模型训练变得更加简单 【免费下载链接】tensor_parallelAutomatically split your PyTorch models on multiple GPUs for training inference项目地址: https://gitcode.com/gh_mirrors/te/tensor_parallel创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考