Optimice los modelos Whisper en Amazon SageMaker con LoRA

Whisper es un modelo de reconocimiento automático de voz (ASR) entrenado con 680.000 horas de datos monitoreados desde Internet, que cubre una variedad de idiomas y tareas. Una de sus limitaciones es el bajo rendimiento en idiomas de bajos recursos, como el marathi y el dravidiano, que se puede solucionar mediante ajustes. Sin embargo, ajustar un modelo de susurro se ha convertido en un desafío importante en términos tanto de recursos computacionales como de requisitos de almacenamiento. De cinco a diez ejecuciones de ajuste completo para los modelos Whisper requieren aproximadamente 100 horas de GPU A100 (40 GB SXM4) (varía según el tamaño y los parámetros del modelo), y cada punto de control ajustado requiere aproximadamente 7 GB de almacenamiento. Esta combinación de altos requisitos computacionales y de almacenamiento puede presentar obstáculos importantes, particularmente en entornos con recursos limitados, y a menudo puede hacer que sea extremadamente difícil lograr resultados significativos.

La adaptación de bajo rango, también conocida como LoRA, adopta un enfoque único para el ajuste del modelo. Mantiene los pesos del modelo previamente entrenado en un estado estático e introduce matrices de descomposición de rangos entrenables en cada capa de la estructura del Transformador. Este método puede reducir 10.000 veces la cantidad de parámetros entrenables necesarios para tareas posteriores y reducir los requisitos de memoria de GPU 3 veces. En términos de calidad del modelo, se ha demostrado que LoRA iguala o incluso supera el rendimiento de los métodos de ajuste tradicionales a pesar de trabajar con menos parámetros entrenables (consulte los resultados del artículo original de LoRA). También ofrece la ventaja de un mayor rendimiento de la capacitación. A diferencia de los métodos de adaptador, LoRA no genera latencia adicional durante la inferencia, lo que mantiene la eficiencia del modelo durante la fase de implementación. El ajuste fino de Whisper utilizando LoRA ha mostrado resultados prometedores. Tomemos como ejemplo Whisper-Large-v2: ejecutar 3 épocas en un conjunto de datos de voz compartido de 12 horas en una GPU de memoria de 8 GB tarda entre 6 y 8 horas, lo que es 5 veces más rápido que el ajuste completo con un rendimiento comparable.

Amazon SageMaker es una plataforma ideal para implementar el ajuste de LoRA de Whisper. Amazon SageMaker le permite crear, entrenar e implementar modelos de aprendizaje automático para cualquier caso de uso con infraestructura, herramientas y flujos de trabajo totalmente administrados. Otros beneficios del entrenamiento de modelos pueden incluir menores costos de entrenamiento con Managed Spot Training, bibliotecas de entrenamiento distribuidas para dividir modelos y conjuntos de datos de entrenamiento en instancias de GPU de AWS, y más. Los modelos entrenados de SageMaker se pueden utilizar fácilmente para realizar inferencias directamente en SageMaker. En esta publicación, presentamos una guía paso a paso para implementar el ajuste de LoRA en SageMaker. El código fuente asociado con esta implementación se puede encontrar en GitHub.

Prepare el conjunto de datos para realizar ajustes

Para la tarea de ajuste, utilizamos el idioma marathi, de bajos recursos. Con la biblioteca de conjuntos de datos de Hugging Face, puede descargar el conjunto de datos de Common Voice y dividirlo en conjuntos de datos de entrenamiento y prueba. Vea el siguiente código:

from datasets import load_dataset, DatasetDict

language = "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)

El modelo de reconocimiento de voz Whisper requiere que las entradas de audio sean archivos WAV enteros y con signo mono de 16 kHz y 16 bits. Dado que el conjunto de datos de Common Voice tiene una frecuencia de muestreo de 48 KB, primero debe reducir la escala de los archivos de audio. Luego debe aplicar el extractor de funciones de Whisper al audio para extraer las funciones del espectrograma log-mel y aplicar el tokenizador de Whisper a las funciones del cuadro para convertir cada oración de la transcripción en una identificación simbólica. Vea el siguiente código:

from transformers import WhisperFeatureExtractor
from 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):
# load and resample audio data from 48 to 16kHz
audio = batch["audio"]

# compute log-Mel input features from input audio array
batch["input_features"] = feature_extractor(audio["array"], sampling_rate=audio["sampling_rate"]).input_features[0]

# encode target text to label ids
batch["labels"] = tokenizer(batch["sentence"]).input_ids
return batch

#apply the data preparation function to all of our fine-tuning dataset samples using dataset's .map method.
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")
!aws s3 cp --recursive "marathi-common-voice-processed" s3://<Your-S3-Bucket>

Después de procesar todas las muestras de entrenamiento, cargue los datos procesados ​​en Amazon S3, de modo que cuando use los datos de entrenamiento procesados ​​en la fase de ajuste, pueda usar FastFile para entregar directamente el archivo S3 en lugar de copiarlo al disco local:

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"
)

Entrenar el modelo

Para demostrarlo, usaremos Whisper-large-v2 como modelo previamente entrenado (whisper v3 ya está disponible), que se puede importar a través de la biblioteca Hugging Face Transformer. Puede utilizar la cuantificación de 8 bits para mejorar aún más la eficiencia del entrenamiento. La cuantificación de 8 bits proporciona optimización de la memoria al redondear números de punto flotante a enteros de 8 bits. Es una técnica de compresión de modelos comúnmente utilizada para lograr ahorros de memoria reducida sin sacrificar demasiada precisión de inferencia.

Para cargar el modelo previamente entrenado en formato cuantificado de 8 bits, simplemente agregamos el argumento «load_in_8bit=True» al crear una instancia del modelo, como se muestra en el siguiente código. Esto carga los pesos del modelo cuantificados a 8 bits, lo que reduce el requisito de memoria.

from transformers import WhisperForConditionalGeneration

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

Usamos la implementación LoRA del paquete Peft de Hugging Face. El ajuste de un modelo utilizando LoRA se realiza en cuatro pasos:

1. Cree una instancia de un modelo base (como hicimos en el último paso).
2. Crear una configuración (LoraConfig), donde se definen los parámetros específicos de LoRA.
3. Envuelva el modelo base con get_peft_model() para conseguir uno entrenable PeftModel.
4. Entrenalos PeftModel como modelo básico.

Vea el siguiente código:

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,
)

Para ejecutar un trabajo de capacitación de SageMaker, traemos nuestro propio contenedor Docker. Puede descargar la imagen de Docker desde GitHub, donde se empaquetan ffmpeg4 y git-lfs junto con otros requisitos de Python. Para obtener más información sobre cómo personalizar su propio contenedor Docker para que funcione con SageMaker, consulte Personalizar su propio contenedor de entrenamiento. Luego puede utilizar Hugging Face Estimator y comenzar un trabajo de capacitación en 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)

La implementación de LoRA nos permitió ejecutar la tarea de ajuste fino de Whisper a gran escala en una sola instancia de GPU (por ejemplo, ml.g5.2xlarge). En comparación, la tarea de ajuste completo de Whisper Large requiere múltiples GPU (por ejemplo, ml.p4d.24xlarge) y un tiempo de entrenamiento mucho más largo. Más específicamente, nuestro experimento mostró que la tarea de ajuste completo requiere 24 veces más horas de GPU en comparación con el enfoque LoRA.

Evaluar el rendimiento del modelo

Para evaluar el rendimiento del modelo Whisper ajustado, calculamos la tasa de error de palabras (WER) utilizando un conjunto de pruebas sostenidas. La OMS mide la diferencia entre la transcripción prevista y la transcripción real. Un WER más bajo indica un mejor rendimiento. Puede ejecutar el siguiente script en el modelo previamente entrenado y el modelo ajustado y comparar sus diferencias 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, batch
gc.collect()
wer = 100 * metric.compute()
print(f"{wer=}")

Diploma

En esta publicación, demostramos el ajuste de Whisper, un modelo de reconocimiento de voz de última generación. Específicamente, utilizamos PEFT LoRA de Hugging Face y habilitamos la cuantificación de 8 bits para un entrenamiento eficiente. También mostramos cómo ejecutar el trabajo de capacitación en SageMaker.

Aunque este es un primer paso importante, hay varias maneras de aprovechar este trabajo para mejorar aún más el modelo Whisper. En el futuro, considere utilizar la capacitación distribuida de SageMaker para escalar la capacitación a un conjunto de datos mucho más grande. Esto permite entrenar el modelo con datos más diversos y completos, lo que mejora la precisión. También puede optimizar la latencia al implementar el modelo Whisper para permitir el reconocimiento de voz en tiempo real. Además, podría ampliar el trabajo para manejar transcripciones de audio más largas, lo que requiere cambios en la arquitectura del modelo y los esquemas de capacitación.

Saber

Los autores desean agradecer a Paras Mehra, John Sol y Evandro Franco por sus valiosos comentarios y revisión del artículo.

Sobre los autores

Jun Shi es arquitecto senior de soluciones en Amazon Web Services (AWS). Su enfoque actual es la infraestructura y aplicaciones de AI/ML. Tiene más de una década de experiencia en la industria FinTech como desarrollador de software.

DR. Changsha Ma es un especialista en IA/ML en AWS. Es tecnóloga con un doctorado en informática, una maestría en psicología educativa y años de experiencia en ciencia de datos y consultoría independiente en el campo de AI/ML. Le apasiona investigar enfoques metodológicos sobre la inteligencia humana y de máquinas. Fuera del trabajo, le gusta hacer senderismo, cocinar, buscar comida y pasar tiempo con amigos y familiares.