Finjuster Whisper-modeller på Amazon SageMaker med LoRA | Amazon Web Services

Whisper er en Automatic Speech Recognition (ASR) model, der er blevet trænet ved hjælp af 680,000 timers overvågede data fra nettet, der omfatter en række sprog og opgaver. En af dens begrænsninger er den lave ydeevne på sprog med lav ressource, såsom marathisprog og dravidiske sprog, som kan afhjælpes med finjustering. Det er dog blevet en betydelig udfordring at finjustere en Whisper-model, både med hensyn til beregningsressourcer og lagerkrav. Fem til ti kørsler af fuld finjustering for Whisper-modeller kræver cirka 100 timers A100 GPU (40 GB SXM4) (varierer baseret på modelstørrelser og modelparametre), og hvert finjusteret kontrolpunkt kræver omkring 7 GB lagerplads. Denne kombination af høje beregnings- og lagerkrav kan udgøre betydelige forhindringer, især i miljøer med begrænsede ressourcer, hvilket ofte gør det usædvanligt vanskeligt at opnå meningsfulde resultater.

Low-Rank Adaptation, også kendt som LoRA, tager en unik tilgang til model finjustering. Det fastholder de forudtrænede modelvægte i en statisk tilstand og introducerer trenerbare rangnedbrydningsmatricer i hvert lag af transformatorstrukturen. Denne metode kan reducere antallet af oplærbare parametre, der er nødvendige for downstream-opgaver, med 10,000 gange og reducere GPU-hukommelseskravet med 3 gange. Med hensyn til modelkvalitet har LoRA vist sig at matche eller endda overgå ydeevnen af ​​traditionelle finjusteringsmetoder, på trods af at der arbejdes med færre parametre, der kan trænes (se resultaterne fra originalen LoRA papir). Det giver også fordelen ved øget træningsgennemstrømning. I modsætning til adapter metoder, introducerer LoRA ikke yderligere latens under inferens, og bibeholder derved modellens effektivitet under implementeringsfasen. Finjustering af Whisper ved hjælp af LoRA har vist lovende resultater. Tag Whisper-Large-v2, for eksempel: at køre 3-epoker med et 12-timers fælles stemmedatasæt på 8 GB hukommelse GPU tager 6-8 timer, hvilket er 5 gange hurtigere end fuld finjustering med sammenlignelig ydeevne.

Amazon SageMaker er en ideel platform til at implementere LoRA finjustering af Whisper. Amazon SageMaker giver dig mulighed for at bygge, træne og implementere maskinlæringsmodeller til enhver brug med fuldt administreret infrastruktur, værktøjer og arbejdsgange. Yderligere modeltræningsfordele kan omfatte lavere træningsomkostninger med Managed Spot Training, distribuerede træningsbiblioteker til opdeling af modeller og træningsdatasæt på tværs af AWS GPU-instanser og mere. De trænede SageMaker-modeller kan nemt implementeres til slutning direkte på SageMaker. I dette indlæg præsenterer vi en trin-for-trin guide til implementering af LoRA finjustering i SageMaker. Kildekoden forbundet med denne implementering kan findes på GitHub.

Forbered datasættet til finjustering

Vi bruger det ressourcesvage sprog marathi til finjusteringsopgaven. Bruger Knusende ansigt-datasæt bibliotek, kan du downloade og opdele Common Voice-datasættet i trænings- og testdatasæt. Se følgende kode:

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)

Whisper-talegenkendelsesmodellen kræver lydindgange 16kHz mono 16-bit signerede heltal WAV-filer. Da Common Voice-datasættet er 48K samplinghastighed, skal du først downsample lydfilerne. Derefter skal du anvende Whispers feature extractor til lyden for at udtrække log-mel spektrogram funktioner, og anvende Whispers tokenizer til de indrammede funktioner for at konvertere hver sætning i transskriptionen til et token ID. Se følgende kode:

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>

Når du har behandlet alle træningseksemplerne, skal du uploade de behandlede data til Amazon S3, så når du bruger de behandlede træningsdata i finjusteringsfasen, kan du bruge FastFile for at montere S3-filen direkte i stedet for at kopiere den til lokal disk:

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

Træn modellen

Til demonstration bruger vi whisper-large-v2 som den fortrænede model (whisper v3 er nu tilgængelig), som kan importeres gennem Hugging Face transformers bibliotek. Du kan bruge 8-bit kvantisering for yderligere at forbedre træningseffektiviteten. 8-bit kvantisering tilbyder hukommelsesoptimering ved at afrunde fra flydende komma til 8-bit heltal. Det er en almindeligt anvendt modelkomprimeringsteknik til at opnå besparelser ved reduceret hukommelse uden at ofre præcision under inferens for meget.

For at indlæse den præ-trænede model i 8-bit kvantiseret format, tilføjer vi blot argumentet load_in_8bit=True, når modellen instansieres, som vist i følgende kode. Dette vil indlæse modelvægtene kvantificeret til 8 bit, hvilket reducerer hukommelsesfodaftrykket.

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

Vi bruger LoRA-implementeringen fra Hugging Face's peft pakke. Der er fire trin til at finjustere en model ved hjælp af LoRA:

1. Instantiér en basismodel (som vi gjorde i sidste trin).
2. Opret en konfiguration (LoraConfig) hvor LoRA-specifikke parametre er defineret.
3. Pak basismodellen ind med get_peft_model() at få en træningsbar PeftModel.
4. Træn PeftModel som basismodel.

Se følgende kode:

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

At køre en SageMaker uddannelse job, vi medbringer vores egen Docker container. Du kan downloade Docker-billedet fra GitHub, hvor ffmpeg4 og git-lfs er pakket sammen med andre Python-krav. For at lære mere om, hvordan du tilpasser din egen Docker-container til at arbejde med SageMaker, se Tilpasning af din egen træningscontainer. Så kan du bruge Hugging Face Estimator og starte et SageMaker træningsjob:

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)

Implementeringen af ​​LoRA gjorde det muligt for os at køre Whisper store finjusteringsopgave på en enkelt GPU-instans (for eksempel ml.g5.2xlarge). Til sammenligning kræver Whisper store fuld finjusteringsopgave flere GPU'er (for eksempel ml.p4d.24xlarge) og en meget længere træningstid. Mere specifikt viste vores eksperiment, at den fulde finjusteringsopgave kræver 24 gange flere GPU-timer sammenlignet med LoRA-tilgangen.

Evaluer modellens ydeevne

For at evaluere ydeevnen af ​​den finjusterede Whisper-model beregner vi ordfejlfrekvensen (WER) på et udholdt testsæt. WER måler forskellen mellem den forudsagte transkription og grundsandhedens transkription. En lavere WER indikerer bedre ydeevne. Du kan køre følgende script mod den forudtrænede model og finjusterede model og sammenligne deres WER-forskel:

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

Konklusion

I dette indlæg demonstrerede vi finjustering af Whisper, en state-of-the-art talegenkendelsesmodel. Vi brugte især Hugging Faces PEFT LoRA og aktiverede 8-bit kvantisering for effektiv træning. Vi demonstrerede også, hvordan man kører træningsjobbet på SageMaker.

Selvom dette er et vigtigt første skridt, er der flere måder, du kan bygge videre på dette arbejde for yderligere at forbedre hviskemodellen. Fremover kan du overveje at bruge SageMaker distribueret træning til at skalere træning på et meget større datasæt. Dette vil give modellen mulighed for at træne på mere varierede og omfattende data, hvilket forbedrer nøjagtigheden. Du kan også optimere latens, når du betjener Whisper-modellen, for at aktivere talegenkendelse i realtid. Derudover kan du udvide arbejdet til at håndtere længere lydtransskriptioner, hvilket kræver ændringer af modelarkitektur og træningsskemaer.

Anerkendelse

Forfatterne takker Paras Mehra, John Sol og Evandro Franco for deres indsigtsfulde feedback og gennemgang af indlægget.

Om forfatterne

Jun Shi er Senior Solutions Architect hos Amazon Web Services (AWS). Hans nuværende fokusområder er AI/ML infrastruktur og applikationer. Han har over ti års erfaring i FinTech-industrien som softwareingeniør.

Dr. Changsha Ma er AI/ML Specialist hos AWS. Hun er teknolog med en ph.d. i datalogi, en kandidatgrad i pædagogisk psykologi og mange års erfaring inden for datavidenskab og selvstændig rådgivning inden for AI/ML. Hun brænder for at forske i metodiske tilgange til maskinel og menneskelig intelligens. Uden for arbejdet elsker hun at vandre, lave mad, jage mad og tilbringe tid med venner og familie.

• SEO Powered Content & PR Distribution. Bliv forstærket i dag.
• PlatoData.Network Vertical Generative Ai. Styrk dig selv. Adgang her.
• PlatoAiStream. Web3 intelligens. Viden forstærket. Adgang her.
• PlatoESG. Kulstof, CleanTech, Energi, Miljø, Solenergi, Affaldshåndtering. Adgang her.
• PlatoHealth. Bioteknologiske og kliniske forsøgs intelligens. Adgang her.
• Kilde: https://aws.amazon.com/blogs/machine-learning/fine-tune-whisper-models-on-amazon-sagemaker-with-lora/