tensor_parallel实战：训练FLAN-T5-xl模型的完整案例

tensor_parallel实战：训练FLAN-T5-xl模型的完整案例

【免费下载链接】tensor_parallelAutomatically split your PyTorch models on multiple GPUs for training & inference项目地址: https://gitcode.com/gh_mirrors/te/tensor_parallel

Tensor 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_dtype=torch.bfloat16, low_cpu_mem_usage=True, offload_state_dict=True )

这里，我们使用了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_tensors="pt").input_ids.to("cuda") output_ids = tokenizer("A cat sat did not sit on a mat", return_tensors="pt").input_ids.to("cuda") # 前向和反向传播 loss = model(input_ids=input_ids, labels=output_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, prefix="summarize:"): inputs, targets = examples['dialogue'], examples['summary'] inputs = [prefix + inp for inp in inputs] model_inputs = tokenizer(inputs, max_length=256, truncation=True) labels = tokenizer(text_target=targets, max_length=256, truncation=True) model_inputs["labels"] = labels["input_ids"] return model_inputs train_data = data['train'].filter(lambda row: row['dialogue'] and row['summary']).map( preprocess_function, batched=True, remove_columns=['id', 'dialogue', 'summary'] )

这个预处理函数会将对话文本加上"summarize:"前缀，然后使用分词器将文本转换为模型输入的token ID。同时，我们也对标签（摘要）进行同样的处理。

步骤五：配置训练参数并开始训练

现在，我们可以配置训练参数并开始训练了。使用transformers库的Seq2SeqTrainer可以简化训练过程：

trainer = transformers.Seq2SeqTrainer( model=model, train_dataset=train_data, args=transformers.Seq2SeqTrainingArguments( do_train=True, remove_unused_columns=False, per_device_train_batch_size=4, gradient_accumulation_steps=2, optim='adafactor', warmup_steps=250, max_steps=1000, learning_rate=1e-5, logging_steps=1, output_dir='outputs' ), tokenizer=tokenizer, data_collator=transformers.DataCollatorForSeq2Seq( tokenizer, padding=True, max_length=512, pad_to_multiple_of=8 ) )

这里，我们设置了每设备训练批次大小为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