在亚马逊SageMaker上使用LoRA对Whisper模型进行微调

在亚马逊SageMaker上使用LoRA微调Whisper模型

Whisper是通过使用网络上的680,000小时的监督数据进行训练得到的自动语音识别(ASR)模型,该数据涵盖了多种语言和任务。其中一个限制是在Marathi语言和德拉维语等低资源语言上的性能较低,可以通过微调来弥补。然而,对Whisper模型进行微调已经成为一个较大的挑战,无论是在计算资源还是存储需求方面。对Whisper模型进行完整微调的五到十次运行大约需要100小时的A100 GPU(40 GB SXM4)(根据模型大小和参数而有所变化),而每个微调后的检查点大约需要7 GB的存储空间。这种高计算和存储需求的组合在资源有限的环境中可能会造成重大障碍,往往使得实现有意义的结果异常困难。

低秩适应,也称为LoRA,采用了一种独特的模型微调方法。它将预训练模型权重保持在静态状态,并在Transformer结构的每个层中引入可训练的秩分解矩阵。这种方法可以将下游任务所需的可训练参数数量减少10,000倍,并将GPU内存要求降低3倍。在模型质量方面,尽管使用更少的可训练参数,LoRA已经显示出与传统微调方法相当甚至超过其性能的特点(请参阅原始的LoRA论文中的结果)。它还提供了增加的训练吞吐量的好处。与adapter方法不同,LoRA在推断过程中不会引入额外的延迟,从而在部署阶段保持模型的效率。使用LoRA对Whisper进行微调显示出了有希望的结果。例如,以Whisper-Large-v2为例:在8 GB内存GPU上运行3个epoch的共同语音数据集大约需要6-8小时,比具有可比性性能的完整微调快5倍。

Amazon SageMaker是实施Whisper的LoRA微调的理想平台。Amazon SageMaker使您能够构建、训练和部署任何用例的机器学习模型,拥有完全托管的基础设施、工具和工作流程。附加的模型训练优势包括通过托管的Spot培训降低培训成本,分布式训练库将模型和训练数据集分布在AWS GPU实例上,以及更多。训练的SageMaker模型可以直接在SageMaker上进行推断部署。在本文中,我们介绍了在SageMaker中实施LoRA微调的逐步指南。与此实现相关的源代码可以在GitHub上找到。

准备微调的数据集

我们使用低资源语言Marathi进行微调任务。使用Hugging Face datasets库,您可以下载和拆分Common Voice数据集为训练和测试数据集。请参阅以下代码:

from datasets import load_dataset, DatasetDictlanguage = "Marathi"language_abbr = "mr"task = "transcribe"dataset_name = "mozilla-foundation/common_voice_11_0"common_voice = DatasetDict()common_voice["train"] = load_dataset(dataset_name, language_abbr, split="train+validation", use_auth_token=True)common_voice["test"] = load_dataset(dataset_name, language_abbr, split="test", use_auth_token=True)

Whisper语音识别模型要求音频输入为16kHz单声道16位有符号整数WAV文件。由于Common Voice数据集的采样率为48K,您需要先将音频文件进行降采样。然后,您需要将Whisper的特征提取器应用于音频,提取出对数梅尔频谱特征,并将Whisper的标记器应用于分帧特征,将转录中的每个句子转换为令牌ID。请参阅以下代码:

from transformers import WhisperFeatureExtractorfrom transformers import WhisperTokenizer

feature_extractor = WhisperFeatureExtractor.from_pretrained(model_name_or_path)
tokenizer = WhisperTokenizer.from_pretrained(model_name_or_path, language=language, task=task)

def prepare_dataset(batch):
    # 从48KHz的音频数据中加载并重新采样到16KHz
    audio = batch["audio"]
    # 从输入音频数组中计算log-Mel输入特征
    batch["input_features"] = feature_extractor(audio["array"], sampling_rate=audio["sampling_rate"]).input_features[0]
    # 将目标文本编码为标签id
    batch["labels"] = tokenizer(batch["sentence"]).input_ids
    return batch

# 对我们的Fine-tuning数据集中的所有样本应用数据准备函数,使用dataset的.map方法
common_voice = common_voice.map(prepare_dataset, remove_columns=common_voice.column_names["train"], num_proc=2)
common_voice.save_to_disk("marathi-common-voice-processed")

# 将所有经过处理的训练数据上传到Amazon S3,以便在Fine-tuning阶段使用已处理的训练数据时,可以直接使用FastFile挂载S3文件而不是将其复制到本地磁盘
!aws s3 cp --recursive "marathi-common-voice-processed" s3://

在处理完所有训练样本后,将处理后的数据上传到Amazon S3。这样,在使用Fine-tuning阶段中的处理后的训练数据时,可以使用FastFile直接挂载S3文件,而不是将其复制到本地磁盘。

from sagemaker.inputs import TrainingInput

training_input_path = s3uri

training = TrainingInput(
    s3_data_type='S3Prefix', # Available Options: S3Prefix | ManifestFile | AugmentedManifestFile
    s3_data=training_input_path,
    distribution='FullyReplicated', # Available Options: FullyReplicated | ShardedByS3Key
    input_mode='FastFile'
)

训练模型

为了演示目的,我们使用whisper-large-v2作为预训练模型(现在已有whisper v3可用),可以通过Hugging Face transformers库导入。您可以使用8位量化来进一步提高训练效率。 8位量化通过将浮点数四舍五入为8位整数来实现内存优化。 这是一种常用的模型压缩技术,可在不太牺牲推理精度的情况下减少内存占用。

要以8位量化格式加载预训练模型,只需在实例化模型时添加load_in_8bit=True参数,如下所示。 这将加载量化为8位的模型权重,减小内存占用。

from transformers import WhisperForConditionalGeneration

model = WhisperForConditionalGeneration.from_pretrained(model_name_or_path, load_in_8bit=True, device_map="auto")

我们使用Hugging Face的peft包中的LoRA实现。 使用LoRA对模型进行Fine-tuning包括以下四个步骤:

  1. 实例化一个基础模型(就像上一步所做的那样)。
  2. 创建一个LoraConfig配置,其中定义了LoRA特定参数。
  3. 使用get_peft_model()将基础模型封装为可训练的PeftModel。
  4. 将PeftModel作为基础模型进行训练。

请参阅以下代码:

from peft import LoraConfig, get_peft_model

config = LoraConfig(
    r=32,
    lora_alpha=64,
    target_modules=["q_proj", "v_proj"],
    lora_dropout=0.05,
    bias="none"
)

model = get_peft_model(model, config)

training_args = Seq2SeqTrainingArguments(
    output_dir=args.model_dir,
    per_device_train_batch_size=int(args.train_batch_size),
    gradient_accumulation_steps=1,
    learning_rate=float(args.learning_rate),
    warmup_steps=args.warmup_steps,
    num_train_epochs=args.num_train_epochs,
    evaluation_strategy="epoch",
    fp16=True,
    per_device_eval_batch_size=args.eval_batch_size,
    generation_max_length=128,
    logging_steps=25,
    remove_unused_columns=False,
    label_names=["labels"],
)

trainer = Seq2SeqTrainer(
    args=training_args,
    model=model,
    train_dataset=train_dataset["train"],
    eval_dataset=train_dataset.get("test", train_dataset["test"]),
    data_collator=data_collator,
    tokenizer=processor.feature_extractor
)

要运行SageMaker训练作业,我们需要自带Docker容器。可以从GitHub下载Docker镜像,其中包含了ffmpeg4和git-lfs以及其他Python要求。有关如何使自己的Docker容器适用于SageMaker的更多信息,请参阅”Adapting your own training container”。然后,您可以使用Hugging Face Estimator开始一个SageMaker训练作业:

OUTPUT_PATH= f's3://{BUCKET}/{PREFIX}/{TRAINING_JOB_NAME}/output/'huggingface_estimator = HuggingFace(entry_point='train.sh',source_dir='./src',output_path= OUTPUT_PATH,instance_type=instance_type,instance_count=1,# transformers_version='4.17.0',# pytorch_version='1.10.2',py_version='py310',image_uri=<ECR-PATH>,role=ROLE,metric_definitions = metric_definitions,volume_size=200,distribution=distribution,keep_alive_period_in_seconds=1800,environment=environment,)huggingface_estimator.fit(job_name=TRAINING_JOB_NAME, wait=False)

实施 LoRA 让我们能够在单个 GPU 实例上运行 Whisper 大型微调任务(例如 ml.g5.2xlarge)。相比之下,Whisper 大型完全微调任务需要多个 GPU(例如 ml.p4d.24xlarge)和更长的训练时间。具体而言,我们的实验表明,相比于 LoRA 方法,完全微调任务需要 24 倍的 GPU 小时。

评估模型性能

为了评估微调的 Whisper 模型性能,我们在一个保留的测试集上计算了单词错误率(WER)。WER 测量预测转录和实际转录之间的差异。较低的 WER 表示更好的性能。您可以运行以下脚本对预训练模型和微调模型运行,并比较它们的 WER 差异:

metric = evaluate.load("wer")eval_dataloader = DataLoader(common_voice["test"], batch_size=8, collate_fn=data_collator)model.eval()for step, batch in enumerate(tqdm(eval_dataloader)):with torch.cuda.amp.autocast():with torch.no_grad():generated_tokens = (model.generate(input_features=batch["input_features"].to("cuda"),decoder_input_ids=batch["labels"][:, :4].to("cuda"),max_new_tokens=255,).cpu().numpy())labels = batch["labels"].cpu().numpy()labels = np.where(labels != -100, labels, tokenizer.pad_token_id)decoded_preds = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)metric.add_batch(predictions=decoded_preds,references=decoded_labels,)del generated_tokens, labels, batchgc.collect()wer = 100 * metric.compute()print(f"{wer=}")

结论

在本文中,我们演示了 Whisper 微调,这是一种最先进的语音识别模型。我们特别使用了 Hugging Face 的 PEFT LoRA,并启用了 8 位量化以进行高效训练。我们还演示了如何在 SageMaker 上运行训练作业。

虽然这是重要的第一步,但您还可以通过以下几种方式进一步改进 Whisper 模型。未来,考虑使用 SageMaker 分布式训练来扩展在更大数据集上的训练。这将允许模型在更多样化和全面的数据上进行训练,提高准确性。您还可以在提供 Whisper 模型时优化延迟,以实现实时语音识别。此外,您可以扩展工作以处理更长的音频转录,这需要对模型架构和训练方案进行更改。

致谢

作者对 Paras Mehra,John Sol 和 Evandro Franco 表示衷心的感谢,感谢他们对本文的富有见地的反馈和审查。