diff --git a/README.md b/README.md index 196b48f62..7e3de0448 100644 --- a/README.md +++ b/README.md @@ -150,6 +150,49 @@ result = whisper.decode(model, mel, options) print(result.text) ``` +## Fine-tuning + +Whisper can be fine-tuned on your own data with the `whisper-train` command. Prepare a JSONL manifest with one sample per line (a CSV file with the same columns also works; relative paths are resolved against the manifest's directory): + +```json +{"audio": "clips/0001.wav", "text": "the transcript", "language": "en"} +``` + +Two modes are supported: + +- **LoRA** (`--finetune lora`, the default) freezes the model and trains small low-rank adapters, about 0.2–0.4% of the parameters. It needs much less GPU memory and saves a few-MB adapter file. +- **Full fine-tuning** (`--finetune full`) updates all weights (or only the decoder with `--freeze_encoder`). It needs much more memory and saves a complete model checkpoint. + +The number of trainable parameters in each mode: + +| Model | LoRA (default) | Full fine-tuning | Full, `--freeze_encoder` | +|:----------:|:--------------:|:----------------:|:------------------------:| +| tiny | 0.15 M (0.40%) | 37 M (100%) | 30 M (79%) | +| base | 0.29 M (0.41%) | 72 M (100%) | 52 M (72%) | +| small | 0.88 M (0.37%) | 241 M (100%) | 154 M (64%) | +| medium | 2.36 M (0.31%) | 762 M (100%) | 457 M (60%) | +| large-v3 | 3.93 M (0.26%) | 1542 M (100%) | 907 M (59%) | +| turbo | 1.64 M (0.20%) | 807 M (100%) | 172 M (21%) | + +LoRA adapts the query and value projections with rank 8 by default; `--lora_r` and `--lora_targets all` add capacity (e.g. 17.3 M parameters, 2.27%, for `medium` with rank 16 on all layers). + +```bash +whisper-train --model small --language en \ + --train_manifest train.jsonl --val_manifest val.jsonl --output_dir runs/small-lora + +whisper-train --finetune full --model small --language en \ + --train_manifest train.jsonl --val_manifest val.jsonl --output_dir runs/small-full +``` + +Training reports the word and character error rates on the validation set before and after fine-tuning, and `--precision` selects fp32, fp16 or bf16 (by default bf16 if the GPU supports it, otherwise fp16). To use the results: + +```bash +whisper audio.wav --model small --lora runs/small-lora/adapter_best.pt +whisper audio.wav --model runs/small-full/model_best.pt +``` + +In Python, use `whisper.load_model("small", lora="adapter_best.pt")` for an adapter. Run `whisper-train --help` to view all options. + ## More examples Please use the [🙌 Show and tell](https://github.com/openai/whisper/discussions/categories/show-and-tell) category in Discussions for sharing more example usages of Whisper and third-party extensions such as web demos, integrations with other tools, ports for different platforms, etc. diff --git a/pyproject.toml b/pyproject.toml index 21b90e737..72bc0d602 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,6 +33,7 @@ dependencies = [ optional-dependencies.dev = [ "black", "flake8", "isort", "pytest", "scipy" ] urls = { Homepage = "https://github.com/openai/whisper" } scripts.whisper = "whisper.transcribe:cli" +scripts.whisper-train = "whisper.train:cli" [tool.setuptools] py-modules = [ "whisper" ] diff --git a/tests/test_lora.py b/tests/test_lora.py new file mode 100644 index 000000000..cf5c665d5 --- /dev/null +++ b/tests/test_lora.py @@ -0,0 +1,171 @@ +import os + +import pytest +import torch + +import whisper +from whisper.lora import ( + LoRALinear, + apply_lora, + has_lora, + load_lora, + lora_state_dict, + merge_lora, + save_lora, +) +from whisper.model import Linear + + +@pytest.fixture(scope="module") +def checkpoint(): + model = whisper.load_model("tiny", device="cpu") + return model.dims, model.state_dict() + + +def make_model(checkpoint): + dims, state_dict = checkpoint + model = whisper.model.Whisper(dims) + model.load_state_dict(state_dict) + return model.eval() + + +def make_inputs(model): + torch.manual_seed(0) + mel = torch.randn(2, model.dims.n_mels, 3000) + tokens = torch.randint(0, 50000, (2, 12)) + return mel, tokens + + +def randomize_lora(model): + torch.manual_seed(1) + for module in model.modules(): + if isinstance(module, LoRALinear): + torch.nn.init.normal_(module.lora_B, std=0.02) + + +def test_output_unchanged_at_init(checkpoint): + model = make_model(checkpoint) + mel, tokens = make_inputs(model) + with torch.no_grad(): + expected = model(mel, tokens) + apply_lora(model, r=8, alpha=16, dropout=0.0) + actual = model(mel, tokens) + assert torch.equal(expected, actual) + + +def test_only_lora_is_trainable(checkpoint): + model = make_model(checkpoint) + apply_lora(model, r=4, target_modules=("query", "value")) + + trainable = {n for n, p in model.named_parameters() if p.requires_grad} + assert trainable and all("lora_" in n for n in trainable) + + dims = model.dims + per_layer_enc = 4 * (dims.n_audio_state + dims.n_audio_state) + per_layer_dec = 4 * (dims.n_text_state + dims.n_text_state) + # query and value in the encoder self-attention, and in both the decoder + # self-attention and cross-attention + expected = ( + 2 * dims.n_audio_layer * per_layer_enc + 4 * dims.n_text_layer * per_layer_dec + ) + assert sum(p.numel() for p in model.parameters() if p.requires_grad) == expected + + +def test_target_options(checkpoint): + model = make_model(checkpoint) + apply_lora(model, target_modules="all", apply_to="decoder") + block = model.decoder.blocks[0] + for attn in (block.attn, block.cross_attn): + for name in ("query", "key", "value", "out"): + assert isinstance(getattr(attn, name), LoRALinear) + assert isinstance(block.mlp[0], LoRALinear) and isinstance(block.mlp[2], LoRALinear) + assert not has_lora(model.encoder) + + with pytest.raises(ValueError): + apply_lora(make_model(checkpoint), target_modules=("qkv",)) + with pytest.raises(RuntimeError): + apply_lora(model) + + +def test_save_load_round_trip(checkpoint, tmp_path): + model = make_model(checkpoint) + apply_lora(model, r=4, alpha=8, target_modules=("query", "key", "value")) + randomize_lora(model) + path = os.path.join(tmp_path, "adapter.pt") + save_lora(model, path, base_model="tiny") + + adapter = torch.load(path, weights_only=True) + assert adapter["config"]["base_model"] == "tiny" + assert adapter["config"]["r"] == 4 + assert all(v.dtype == torch.float32 for v in adapter["state_dict"].values()) + assert set(adapter["state_dict"]) == set(lora_state_dict(model)) + + loaded = load_lora(make_model(checkpoint), path).eval() + mel, tokens = make_inputs(model) + with torch.no_grad(): + assert torch.equal(model(mel, tokens), loaded(mel, tokens)) + + +def test_load_model_with_lora(checkpoint, tmp_path): + model = make_model(checkpoint) + apply_lora(model) + randomize_lora(model) + path = os.path.join(tmp_path, "adapter.pt") + save_lora(model, path, base_model="tiny") + + loaded = whisper.load_model("tiny", device="cpu", lora=path).eval() + mel, tokens = make_inputs(model) + with torch.no_grad(): + assert torch.equal(model(mel, tokens), loaded(mel, tokens)) + + +def test_mismatched_base_model(checkpoint, tmp_path): + model = make_model(checkpoint) + apply_lora(model) + path = os.path.join(tmp_path, "adapter.pt") + save_lora(model, path, base_model="tiny") + + dims, state_dict = checkpoint + other = whisper.model.Whisper( + whisper.model.ModelDimensions(**{**dims.__dict__, "n_text_layer": 2}) + ) + with pytest.raises(ValueError, match="different base model"): + load_lora(other, path) + + +def test_merge_matches_unmerged(checkpoint, tmp_path): + model = make_model(checkpoint) + apply_lora(model, target_modules="all") + randomize_lora(model) + mel, tokens = make_inputs(model) + with torch.no_grad(): + expected = model(mel, tokens) + merge_lora(model) + actual = model(mel, tokens) + + assert not has_lora(model) + assert all( + type(m) is Linear for m in model.modules() if isinstance(m, torch.nn.Linear) + ) + assert torch.allclose(expected, actual, atol=1e-3, rtol=1e-3) + + # the merged model is a plain Whisper checkpoint + path = os.path.join(tmp_path, "merged.pt") + torch.save( + {"dims": model.dims.__dict__, "model_state_dict": model.state_dict()}, path + ) + reloaded = whisper.load_model(path, device="cpu").eval() + with torch.no_grad(): + assert torch.equal(actual, reloaded(mel, tokens)) + + +def test_transcribe_with_lora(checkpoint): + model = make_model(checkpoint) + expected = model.transcribe( + os.path.join(os.path.dirname(__file__), "jfk.flac"), fp16=False + )["text"] + apply_lora(model) + actual = model.transcribe( + os.path.join(os.path.dirname(__file__), "jfk.flac"), fp16=False + )["text"] + assert actual == expected diff --git a/tests/test_train.py b/tests/test_train.py new file mode 100644 index 000000000..8c8b8e2bb --- /dev/null +++ b/tests/test_train.py @@ -0,0 +1,282 @@ +import json +import os +import wave + +import numpy as np +import pytest +import torch + +import whisper +from whisper import precision as precision_module +from whisper.audio import SAMPLE_RATE +from whisper.lora import load_lora +from whisper.precision import bf16_supported, resolve_precision +from whisper.train import ( + TokenizerCache, + build_parser, + collate, + edit_distance, + error_rates, + read_manifest, + train, +) + +WORDS = ["one", "two", "three", "four"] + + +def write_wav(path: str, frequency: float, seconds: float = 1.0): + t = np.arange(int(seconds * SAMPLE_RATE)) / SAMPLE_RATE + audio = (0.3 * np.sin(2 * np.pi * frequency * t) * 32767).astype(np.int16) + with wave.open(path, "wb") as f: + f.setnchannels(1) + f.setsampwidth(2) + f.setframerate(SAMPLE_RATE) + f.writeframes(audio.tobytes()) + + +@pytest.fixture +def manifest(tmp_path): + os.makedirs(tmp_path / "clips") + with open(tmp_path / "train.jsonl", "w") as f: + for i, word in enumerate(WORDS): + write_wav(str(tmp_path / "clips" / f"{i}.wav"), 220 * (i + 1)) + f.write(json.dumps({"audio": f"clips/{i}.wav", "text": word}) + "\n") + return str(tmp_path / "train.jsonl") + + +def run_training(manifest, output_dir, *extra, language="en"): + args = build_parser().parse_args( + [ + "--model", "tiny", + "--train_manifest", manifest, + "--val_manifest", manifest, + "--output_dir", str(output_dir), + *(["--language", language] if language else []), + "--batch_size", "2", + "--lr", "1e-2", + "--warmup_steps", "0", + "--num_workers", "0", + "--log_every", "1", + *extra, + ] + ) # fmt: skip + return train(args) + + +def test_read_manifest(tmp_path, manifest): + samples = read_manifest(manifest) + assert len(samples) == len(WORDS) + assert samples[0]["audio"] == str(tmp_path / "clips" / "0.wav") + assert "language" not in samples[0] + + csv_path = tmp_path / "train.csv" + csv_path.write_text('audio,text,language\nclips/0.wav,"one, two",en\n') + samples = read_manifest(str(csv_path)) + assert samples == [ + { + "audio": str(tmp_path / "clips" / "0.wav"), + "text": "one, two", + "language": "en", + } + ] + + +@pytest.mark.parametrize("text", ["hello world", "আর কিছু", "नमस्ते", "你好"]) +def test_first_target_token_is_never_blank(text): + model = whisper.load_model("tiny", device="cpu") + tokenizers = TokenizerCache(model, "en", "transcribe") + tokens, prompt_length = tokenizers.encode(dict(text=text)) + tokenizer = tokenizers.get() + assert tokens[prompt_length] not in tokenizer.encode(" ") + assert tokenizer.decode(tokens[prompt_length:-1]).strip() == text + + +def test_collate_masks_prompt_and_padding(): + eot = 50257 + items = [ + dict(mel=torch.zeros(80, 3000), tokens=[1, 2, 3, 10, 11, eot], prompt_length=3, text="a", language="en"), + dict(mel=torch.zeros(80, 3000), tokens=[1, 2, 3, 12, eot], prompt_length=3, text="b", language="en"), + None, + ] # fmt: skip + batch = collate(items, pad_token=eot) + assert batch["inputs"].tolist() == [[1, 2, 3, 10, 11], [1, 2, 3, 12, eot]] + assert batch["labels"].tolist() == [ + [-100, -100, 10, 11, eot], + [-100, -100, 12, eot, -100], + ] + assert collate([None], pad_token=eot) is None + + +def test_error_rates(): + assert edit_distance("abc", "abd") == 1 + rates = error_rates(["Hello, world!"], ["hello world"], ["en"]) + assert rates == dict(wer=0.0, cer=0.0) + # combining vowel signs must not be stripped or turned into word breaks + rates = error_rates(["नमस्ते दुनिया"], ["नमस्ते दुनिय"], ["hi"]) + assert rates["wer"] == 0.5 + + +def test_training_reduces_loss(manifest, tmp_path): + results = run_training(manifest, tmp_path / "run", "--max_steps", "12", "--precision", "fp32", "--device", "cpu") # fmt: skip + assert results["precision"] == "fp32" + assert results["final"]["loss"] < results["baseline"]["loss"] + for name in ("adapter_best.pt", "adapter_final.pt", "run_config.json", "metrics.jsonl", "results.json"): # fmt: skip + assert os.path.isfile(tmp_path / "run" / name) + with open(tmp_path / "run" / "metrics.jsonl") as f: + losses = [r["loss"] for r in map(json.loads, f) if r["type"] == "train"] + assert all(np.isfinite(losses)) + + +def test_per_sample_language_without_language_flag(tmp_path): + os.makedirs(tmp_path / "clips") + with open(tmp_path / "train.jsonl", "w") as f: + for i, (word, language) in enumerate(zip(WORDS, ["en", "fr", "en", "fr"])): + write_wav(str(tmp_path / "clips" / f"{i}.wav"), 220 * (i + 1)) + sample = {"audio": f"clips/{i}.wav", "text": word, "language": language} + f.write(json.dumps(sample) + "\n") + manifest = str(tmp_path / "train.jsonl") + results = run_training(manifest, tmp_path / "run", "--max_steps", "2", "--precision", "fp32", "--device", "cpu", language=None) # fmt: skip + assert np.isfinite(results["final"]["loss"]) + + +def test_nonfinite_loss_is_skipped(manifest, tmp_path, monkeypatch): + from whisper import train as train_module + + compute_loss = train_module.compute_loss + calls = [] + + def flaky_loss(*args, **kwargs): + calls.append(None) + loss = compute_loss(*args, **kwargs) + return loss * float("nan") if len(calls) == 2 else loss + + monkeypatch.setattr(train_module, "compute_loss", flaky_loss) + run_training(manifest, tmp_path / "run", "--max_steps", "3", "--precision", "fp32", "--device", "cpu", "--skip_baseline") # fmt: skip + adapter = torch.load(tmp_path / "run" / "adapter_final.pt", weights_only=True) + assert all(torch.isfinite(v).all() for v in adapter["state_dict"].values()) + + +@pytest.mark.parametrize( + "mode,lr,filename,key", + [ + ("lora", "1e-2", "adapter_final.pt", "state_dict"), + ("full", "1e-4", "model_final.pt", "model_state_dict"), + ], +) +def test_resume_matches_uninterrupted_run(manifest, tmp_path, mode, lr, filename, key): + common = ["--finetune", mode, "--lr", lr, "--max_steps", "4", "--precision", "fp32", "--device", "cpu", "--skip_baseline"] # fmt: skip + run_training(manifest, tmp_path / "full", *common) + run_training(manifest, tmp_path / "part", *common, "--save_every", "1", "--keep_checkpoints", "2") # fmt: skip + assert sorted(os.listdir(tmp_path / "part" / "checkpoints")) == ["step_3.pt", "step_4.pt"] # fmt: skip + os.remove(tmp_path / "part" / "checkpoints" / "step_4.pt") + run_training(manifest, tmp_path / "part", *common, "--resume", "latest") + + full = torch.load(tmp_path / "full" / filename, weights_only=True) + part = torch.load(tmp_path / "part" / filename, weights_only=True) + for name, value in full[key].items(): + assert torch.allclose(value, part[key][name], atol=1e-6), name + + +def test_full_finetune(manifest, tmp_path): + results = run_training(manifest, tmp_path / "run", "--finetune", "full", "--lr", "1e-4", "--max_steps", "12", "--precision", "fp32", "--device", "cpu") # fmt: skip + assert results["finetune"] == "full" + assert results["trainable_percent"] == 100 + assert results["final"]["loss"] < results["baseline"]["loss"] + assert not os.path.exists(tmp_path / "run" / "adapter_final.pt") + + base = whisper.load_model("tiny", device="cpu") + tuned = whisper.load_model(str(tmp_path / "run" / "model_best.pt"), device="cpu") + assert tuned.dims == base.dims + assert not any(hasattr(m, "lora_A") for m in tuned.modules()) + changed = [ + not torch.equal(p, base.state_dict()[n]) for n, p in tuned.named_parameters() + ] + assert any(changed) + + +def test_full_finetune_freeze_encoder(manifest, tmp_path): + results = run_training(manifest, tmp_path / "run", "--finetune", "full", "--freeze_encoder", "--lr", "1e-4", "--max_steps", "4", "--precision", "fp32", "--device", "cpu", "--skip_baseline") # fmt: skip + base = whisper.load_model("tiny", device="cpu") + tuned = whisper.load_model(str(tmp_path / "run" / "model_final.pt"), device="cpu") + for name, p in tuned.encoder.state_dict().items(): + assert torch.equal(p, base.encoder.state_dict()[name]), name + assert not torch.equal( + tuned.decoder.token_embedding.weight, base.decoder.token_embedding.weight + ) + n_decoder = sum(p.numel() for p in base.decoder.parameters()) + assert results["trainable_parameters"] == n_decoder + + +@pytest.mark.parametrize( + "flags", [["--freeze_encoder"], ["--finetune", "full", "--export_merged"]] +) +def test_invalid_finetune_options(manifest, tmp_path, flags): + with pytest.raises(ValueError): + run_training(manifest, tmp_path / "run", "--max_steps", "1", "--device", "cpu", *flags) # fmt: skip + + +def test_precision_auto(monkeypatch): + assert resolve_precision("auto", "cpu").name == "fp32" + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "is_bf16_supported", lambda **kwargs: True) + assert resolve_precision("auto", "cuda").name == "bf16" + assert not resolve_precision("bf16", "cuda").use_grad_scaler + monkeypatch.setattr(torch.cuda, "is_bf16_supported", lambda **kwargs: False) + auto = resolve_precision("auto", "cuda") + assert auto.name == "fp16" and auto.use_grad_scaler + assert resolve_precision("fp32", "cuda").name == "fp32" + + +def test_precision_bf16_unsupported(monkeypatch): + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "get_device_name", lambda *args: "Tesla T4") + monkeypatch.setattr(precision_module, "bf16_supported", lambda device_type: False) + with pytest.raises(ValueError, match="--precision fp16"): + resolve_precision("bf16", "cuda") + with pytest.raises(ValueError): + resolve_precision("fp16", "cpu") + with pytest.raises(ValueError): + resolve_precision("fp8", "cpu") + + +@pytest.mark.requires_cuda +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_fp16_training_on_gpu(manifest, tmp_path): + results = run_training(manifest, tmp_path / "run", "--max_steps", "8", "--precision", "fp16", "--device", "cuda", "--gradient_checkpointing") # fmt: skip + assert results["precision"] == "fp16" + assert np.isfinite(results["final"]["loss"]) + assert results["final"]["loss"] < results["baseline"]["loss"] + assert results["peak_memory_gb"] > 0 + + model = whisper.load_model("tiny", device="cuda") + load_lora(model, str(tmp_path / "run" / "adapter_best.pt")) + assert all( + p.dtype == torch.float32 for n, p in model.named_parameters() if "lora_" in n + ) + + +@pytest.mark.requires_cuda +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_fp16_full_finetune_on_gpu(manifest, tmp_path): + results = run_training(manifest, tmp_path / "run", "--finetune", "full", "--lr", "1e-4", "--max_steps", "8", "--precision", "fp16", "--device", "cuda", "--gradient_checkpointing") # fmt: skip + assert np.isfinite(results["final"]["loss"]) + assert results["final"]["loss"] < results["baseline"]["loss"] + whisper.load_model(str(tmp_path / "run" / "model_best.pt"), device="cuda") + + +@pytest.mark.requires_cuda +@pytest.mark.skipif(not bf16_supported("cuda"), reason="requires a GPU with bf16") +def test_bf16_training_on_gpu(manifest, tmp_path): + results = run_training(manifest, tmp_path / "run", "--max_steps", "4", "--precision", "bf16", "--device", "cuda") # fmt: skip + assert np.isfinite(results["final"]["loss"]) + assert 0 <= results["final"]["wer"] + + +@pytest.mark.requires_cuda +@pytest.mark.skipif(not bf16_supported("cuda"), reason="requires a GPU with bf16") +def test_bf16_full_finetune_on_gpu(manifest, tmp_path): + results = run_training(manifest, tmp_path / "run", "--finetune", "full", "--lr", "1e-4", "--max_steps", "8", "--precision", "bf16", "--device", "cuda", "--gradient_checkpointing") # fmt: skip + assert np.isfinite(results["final"]["loss"]) + assert results["final"]["loss"] < results["baseline"]["loss"] + assert 0 <= results["final"]["wer"] + whisper.load_model(str(tmp_path / "run" / "model_best.pt"), device="cuda") diff --git a/whisper/__init__.py b/whisper/__init__.py index f284ec045..86f211f84 100644 --- a/whisper/__init__.py +++ b/whisper/__init__.py @@ -105,6 +105,7 @@ def load_model( device: Optional[Union[str, torch.device]] = None, download_root: str = None, in_memory: bool = False, + lora: Optional[str] = None, ) -> Whisper: """ Load a Whisper ASR model @@ -120,6 +121,8 @@ def load_model( path to download the model files; by default, it uses "~/.cache/whisper" in_memory: bool whether to preload the model weights into host memory + lora: str + path to a LoRA adapter file (see `whisper.lora`) to apply to the model Returns ------- @@ -158,4 +161,10 @@ def load_model( if alignment_heads is not None: model.set_alignment_heads(alignment_heads) + if lora is not None: + from .lora import load_lora + + load_lora(model, lora, base_model=name) + model.eval() # disable LoRA dropout for inference + return model.to(device) diff --git a/whisper/lora.py b/whisper/lora.py new file mode 100644 index 000000000..88773677b --- /dev/null +++ b/whisper/lora.py @@ -0,0 +1,288 @@ +""" +Low-Rank Adaptation (LoRA) for Whisper, implemented in plain PyTorch. + +LoRA freezes the pretrained weights and learns a low-rank update for selected +linear layers: ``W x + (alpha / r) * B A x``, where ``A`` has shape ``(r, in)`` +and ``B`` has shape ``(out, r)``. ``B`` is initialized to zero, so a freshly +wrapped model produces exactly the same outputs as the original model. + +See https://arxiv.org/abs/2106.09685 +""" + +import math +import warnings +from dataclasses import asdict +from typing import Dict, Iterable, List, Optional, Union + +import torch +import torch.nn.functional as F +from torch import Tensor, nn + +from .model import Linear, MultiHeadAttention, ResidualAttentionBlock, Whisper + +LORA_FORMAT = "whisper-lora" +LORA_FORMAT_VERSION = 1 + +ATTENTION_TARGETS = ("query", "key", "value", "out") +ALL_TARGETS = ATTENTION_TARGETS + ("mlp",) +ALL_PARTS = ("encoder", "decoder") + + +class LoRALinear(nn.Module): + """Wraps a linear layer with a trainable low-rank update; the base layer is frozen.""" + + def __init__(self, base: nn.Linear, r: int = 8, alpha: float = 16, dropout=0.0): + super().__init__() + if r <= 0: + raise ValueError(f"LoRA rank must be positive, got r={r}") + + self.base = base + self.r = r + self.alpha = alpha + self.scaling = alpha / r + self.in_features = base.in_features + self.out_features = base.out_features + + for param in self.base.parameters(): + param.requires_grad = False + + # LoRA weights are always kept in fp32 (master weights), regardless of the + # dtype of the base layer; they are cast to the input dtype in forward(). + device = base.weight.device + self.lora_A = nn.Parameter( + torch.empty(r, self.in_features, device=device, dtype=torch.float32) + ) + self.lora_B = nn.Parameter( + torch.zeros(self.out_features, r, device=device, dtype=torch.float32) + ) + nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5)) + self.dropout = nn.Dropout(dropout) if dropout > 0 else nn.Identity() + + def forward(self, x: Tensor) -> Tensor: + result = self.base(x) + update = F.linear(self.dropout(x), self.lora_A.to(x.dtype)) + update = F.linear(update, self.lora_B.to(x.dtype)) + return result + update * self.scaling + + @torch.no_grad() + def merged_weight(self) -> Tensor: + delta = (self.lora_B @ self.lora_A) * self.scaling + return (self.base.weight.float() + delta).to(self.base.weight.dtype) + + def extra_repr(self) -> str: + return ( + f"in_features={self.in_features}, out_features={self.out_features}, " + f"r={self.r}, alpha={self.alpha}" + ) + + +def _normalize( + values: Union[str, Iterable[str]], allowed: Iterable[str], name: str +) -> List[str]: + if isinstance(values, str): + values = values.split(",") + values = [v.strip() for v in values if v.strip()] + if values == ["all"]: + return list(allowed) + unknown = [v for v in values if v not in allowed] + if unknown or not values: + raise ValueError( + f"Invalid {name} {unknown or values}; choose from {list(allowed)} or 'all'" + ) + return list(dict.fromkeys(values)) + + +def _target_linears(block: ResidualAttentionBlock, targets: List[str]): + """Yields (parent module, attribute name) pairs for the targeted linear layers""" + for attn in (block.attn, block.cross_attn): + if isinstance(attn, MultiHeadAttention): + for name in ATTENTION_TARGETS: + if name in targets: + yield attn, name + if "mlp" in targets: + for index, layer in enumerate(block.mlp): + if isinstance(layer, (nn.Linear, LoRALinear)): + yield block.mlp, str(index) + + +def apply_lora( + model: Whisper, + r: int = 8, + alpha: float = 16, + dropout: float = 0.05, + target_modules: Union[str, Iterable[str]] = ("query", "value"), + apply_to: Union[str, Iterable[str]] = ("encoder", "decoder"), +) -> Whisper: + """ + Wrap the selected linear layers of a Whisper model with LoRALinear, in place. + + Parameters + ---------- + model: Whisper + The model to adapt + r: int + The rank of the low-rank update + alpha: float + The scaling numerator; the update is scaled by ``alpha / r`` + dropout: float + Dropout applied to the input of the LoRA branch + target_modules: Iterable[str] + Any of "query", "key", "value", "out" (applied to self- and cross-attention) + and "mlp" (both MLP layers), or "all" + apply_to: Iterable[str] + "encoder", "decoder", or both + + Returns + ------- + model : Whisper + The same model, with LoRA layers added and only the LoRA parameters trainable + """ + targets = _normalize(target_modules, ALL_TARGETS, "target_modules") + parts = _normalize(apply_to, ALL_PARTS, "apply_to") + + for part in parts: + for block in getattr(model, part).blocks: + for parent, name in _target_linears(block, targets): + layer = getattr(parent, name) + if isinstance(layer, LoRALinear): + raise RuntimeError("LoRA has already been applied to this model") + lora_layer = LoRALinear(layer, r, alpha, dropout) + setattr(parent, name, lora_layer.train(model.training)) + + model.lora_config = dict( + r=r, alpha=alpha, dropout=dropout, target_modules=targets, apply_to=parts + ) + mark_only_lora_as_trainable(model) + return model + + +def mark_only_lora_as_trainable(model: nn.Module) -> None: + for name, param in model.named_parameters(): + param.requires_grad = "lora_" in name + + +def has_lora(model: nn.Module) -> bool: + return any(isinstance(m, LoRALinear) for m in model.modules()) + + +def lora_state_dict(model: nn.Module) -> Dict[str, Tensor]: + """Returns only the LoRA weights, as fp32 tensors on the CPU""" + return { + name: param.detach().float().cpu() + for name, param in model.named_parameters() + if "lora_" in name + } + + +def save_lora(model: Whisper, path: str, base_model: Optional[str] = None) -> None: + """Save the LoRA adapter (weights and config) of a model to a file""" + if not has_lora(model): + raise ValueError("The model has no LoRA layers; call apply_lora() first") + config = dict(getattr(model, "lora_config", {})) + if base_model is not None: + config["base_model"] = base_model + torch.save( + { + "format": LORA_FORMAT, + "version": LORA_FORMAT_VERSION, + "config": config, + "dims": asdict(model.dims), + "state_dict": lora_state_dict(model), + }, + path, + ) + + +def load_lora( + model: Whisper, path: Union[str, dict], base_model: Optional[str] = None +) -> Whisper: + """ + Apply a saved LoRA adapter to a model, in place. + + If the model does not have LoRA layers yet, they are created using the + configuration stored in the adapter file. + + Parameters + ---------- + model: Whisper + The base model the adapter was trained on + path: Union[str, dict] + Path to an adapter file written by `save_lora()`, or its loaded contents + base_model: Optional[str] + The name of the base model, used to warn about mismatches + """ + if isinstance(path, dict): + adapter = path + else: + kwargs = {"weights_only": True} if torch.__version__ >= "1.13" else {} + adapter = torch.load(path, map_location="cpu", **kwargs) + + if not isinstance(adapter, dict) or adapter.get("format") != LORA_FORMAT: + raise ValueError(f"{path} is not a Whisper LoRA adapter file") + + config = adapter["config"] + trained_on = config.get("base_model") + if adapter["dims"] != asdict(model.dims): + raise ValueError( + f"LoRA adapter was trained on a different base model " + f"({trained_on or 'unknown'}); its dimensions {adapter['dims']} do not " + f"match the loaded model's {asdict(model.dims)}" + ) + if trained_on and base_model and trained_on != base_model: + warnings.warn( + f"LoRA adapter was trained on '{trained_on}' but is being loaded into " + f"'{base_model}'; the dimensions match but the results may be poor" + ) + + if not has_lora(model): + apply_lora( + model, + r=config["r"], + alpha=config["alpha"], + dropout=config.get("dropout", 0.0), + target_modules=config["target_modules"], + apply_to=config["apply_to"], + ) + + state_dict = adapter["state_dict"] + expected = {name for name, _ in model.named_parameters() if "lora_" in name} + missing = sorted(expected - state_dict.keys()) + unexpected = sorted(state_dict.keys() - expected) + if missing or unexpected: + raise ValueError( + f"LoRA adapter does not match the model's LoRA layers; " + f"missing keys: {missing[:5]}, unexpected keys: {unexpected[:5]}" + ) + model.load_state_dict(state_dict, strict=False) + model.lora_config = dict(config) + return model + + +@torch.no_grad() +def merge_lora(model: Whisper) -> Whisper: + """ + Fold the LoRA updates into the base weights and replace every LoRALinear with + a plain Linear layer, in place. The result is a standard Whisper model that + can be saved as a regular checkpoint. + """ + for module in list(model.modules()): + for name, child in list(module.named_children()): + if isinstance(child, LoRALinear): + base = child.base + merged = Linear( + base.in_features, + base.out_features, + bias=base.bias is not None, + device=base.weight.device, + dtype=base.weight.dtype, + ) + merged.weight.copy_(child.merged_weight()) + if base.bias is not None: + merged.bias.copy_(base.bias) + setattr(module, name, merged) + + if hasattr(model, "lora_config"): + del model.lora_config + for param in model.parameters(): + param.requires_grad = True + return model diff --git a/whisper/precision.py b/whisper/precision.py new file mode 100644 index 000000000..357b95ce2 --- /dev/null +++ b/whisper/precision.py @@ -0,0 +1,94 @@ +""" +Mixed-precision settings for training and evaluation. + +Model weights and optimizer state are always kept in fp32; reduced precision is +only used for the forward and backward passes, via torch.autocast. +""" + +import contextlib +from dataclasses import dataclass +from typing import Optional, Union + +import torch + +PRECISIONS = ("auto", "fp32", "fp16", "bf16") + + +@dataclass(frozen=True) +class Precision: + name: str # one of "fp32", "fp16", "bf16" + device_type: str # "cuda" or "cpu" + + @property + def dtype(self) -> torch.dtype: + return { + "fp32": torch.float32, + "fp16": torch.float16, + "bf16": torch.bfloat16, + }[self.name] + + @property + def use_grad_scaler(self) -> bool: + # fp16 has a narrow exponent range, so small gradients underflow to zero + # unless the loss is scaled up; bf16 has the same range as fp32. + return self.name == "fp16" + + def autocast(self): + if self.name == "fp32": + return contextlib.nullcontext() + return torch.autocast(device_type=self.device_type, dtype=self.dtype) + + def grad_scaler(self): + enabled = self.use_grad_scaler + try: + return torch.amp.GradScaler(self.device_type, enabled=enabled) + except (AttributeError, TypeError): # torch < 2.3 + return torch.cuda.amp.GradScaler(enabled=enabled) + + +def bf16_supported(device_type: str) -> bool: + if device_type != "cuda" or not torch.cuda.is_available(): + return False + try: + # recent PyTorch versions report emulated (slow) bf16 support on older GPUs + # such as the T4 unless including_emulation=False is passed + return torch.cuda.is_bf16_supported(including_emulation=False) + except TypeError: + return torch.cuda.is_bf16_supported() + + +def resolve_precision( + precision: str = "auto", device: Optional[Union[str, torch.device]] = None +) -> Precision: + """ + Resolve a precision name to the settings to use on the given device. + + "auto" selects bf16 on GPUs that support it, fp16 on other GPUs, and fp32 on + the CPU. Explicitly requesting a precision that the device cannot use raises + a ValueError instead of silently falling back to another one. + """ + if precision not in PRECISIONS: + raise ValueError(f"Unknown precision '{precision}'; choose from {PRECISIONS}") + + if device is None: + device = "cuda" if torch.cuda.is_available() else "cpu" + device_type = torch.device(device).type + + if precision == "auto": + if device_type != "cuda": + return Precision("fp32", device_type) + return Precision("bf16" if bf16_supported(device_type) else "fp16", device_type) + + if device_type != "cuda" and precision != "fp32": + raise ValueError( + f"--precision {precision} requires a CUDA GPU; use --precision fp32 " + f"on {device_type}" + ) + if precision == "bf16" and not bf16_supported(device_type): + name = torch.cuda.get_device_name(device) if torch.cuda.is_available() else "" + raise ValueError( + f"bf16 is not supported on this GPU ({name}); bf16 requires compute " + f"capability 8.0 or newer (e.g. A100, RTX 30xx/40xx, H100). " + f"Use --precision fp16 instead." + ) + return Precision(precision, device_type) diff --git a/whisper/train.py b/whisper/train.py new file mode 100644 index 000000000..0252b3fb7 --- /dev/null +++ b/whisper/train.py @@ -0,0 +1,905 @@ +""" +Fine-tune Whisper with LoRA adapters (default) or by updating all of its weights. + +Usage: + whisper-train --model small --train_manifest train.jsonl --val_manifest val.jsonl \ + --language en --output_dir runs/small-lora + + # full fine-tuning; saves regular Whisper checkpoints instead of adapters + whisper-train --finetune full --model small --train_manifest train.jsonl \ + --val_manifest val.jsonl --language en --output_dir runs/small-full + +Each line of a JSONL manifest describes one sample: + {"audio": "clips/0001.wav", "text": "the transcript", "language": "en"} +"language" is optional and defaults to --language. A CSV file with "audio" and +"text" columns (and optionally "language") is also accepted. Relative audio paths +are resolved against the directory containing the manifest. +""" + +import argparse +import csv +import functools +import json +import math +import os +import random +import re +import sys +import time +import unicodedata +from dataclasses import asdict +from typing import Dict, Iterator, List, Optional + +import numpy as np +import torch +import torch.nn.functional as F +from torch.utils.checkpoint import checkpoint +from torch.utils.data import DataLoader, Dataset, Sampler + +from .audio import N_SAMPLES, SAMPLE_RATE, load_audio, log_mel_spectrogram, pad_or_trim +from .decoding import DecodingOptions +from .lora import ALL_PARTS, ALL_TARGETS, apply_lora, load_lora, merge_lora, save_lora +from .model import Whisper +from .normalizers import EnglishTextNormalizer +from .precision import PRECISIONS, Precision, resolve_precision +from .tokenizer import LANGUAGES, TO_LANGUAGE_CODE, Tokenizer, get_tokenizer + +IGNORE_INDEX = -100 +FINETUNE_MODES = ("lora", "full") +DEFAULT_LR = {"lora": 5e-4, "full": 1e-5} + + +def log(message: str) -> None: + print(message, flush=True) + + +def warn(message: str) -> None: + print(f"WARNING: {message}", file=sys.stderr, flush=True) + + +# --------------------------------------------------------------------------- +# Data +# --------------------------------------------------------------------------- + + +def read_manifest(path: str) -> List[Dict[str, str]]: + """ + Read a JSONL or CSV manifest. Returns a list of dicts with the keys "audio", + "text" and, if present, "language", with audio paths resolved relative to the + directory of the manifest. + """ + root = os.path.dirname(os.path.abspath(path)) + + if path.lower().endswith(".csv"): + with open(path, newline="", encoding="utf-8") as f: + rows = list(csv.DictReader(f)) + if rows and not {"audio", "text"} <= set(rows[0].keys()): + raise ValueError(f"{path}: CSV manifests need 'audio' and 'text' columns") + else: + rows = [] + with open(path, encoding="utf-8") as f: + for line_number, line in enumerate(f, 1): + if not line.strip(): + continue + try: + row = json.loads(line) + except json.JSONDecodeError as e: + raise ValueError(f"{path}:{line_number}: invalid JSON: {e}") + if not isinstance(row, dict) or "audio" not in row or "text" not in row: + raise ValueError( + f"{path}:{line_number}: each line needs 'audio' and 'text'" + ) + rows.append(row) + + samples = [] + for row in rows: + sample = { + "audio": os.path.join(root, os.path.expanduser(str(row["audio"]))), + "text": "" if row["text"] is None else str(row["text"]), + } + if row.get("language"): + sample["language"] = str(row["language"]) + samples.append(sample) + return samples + + +class TokenizerCache: + """Builds the token sequences for each sample, with one tokenizer per language""" + + def __init__(self, model: Whisper, language: Optional[str], task: str): + self.multilingual = model.is_multilingual + self.num_languages = model.num_languages + self.language = language + self.task = task + self.tokenizers: Dict[Optional[str], Tokenizer] = {} + + def get(self, language: Optional[str] = None) -> Tokenizer: + if not self.multilingual: + language = None + else: + language = language or self.language + if language is None: + raise ValueError( + "A language is required for multilingual models: set --language " + "or add a 'language' field to every sample in the manifest" + ) + language = TO_LANGUAGE_CODE.get(language.lower(), language.lower()) + if language not in self.tokenizers: + self.tokenizers[language] = get_tokenizer( + self.multilingual, + num_languages=self.num_languages, + language=language, + task=self.task, + ) + return self.tokenizers[language] + + def encode(self, sample: Dict[str, str]): + """Returns (tokens, prompt_length) for a sample""" + tokenizer = self.get(sample.get("language")) + prompt = list(tokenizer.sot_sequence_including_notimestamps) + text = tokenizer.encode(" " + sample["text"].strip()) + if text[:1] == tokenizer.encode(" "): + # In some scripts (e.g. Bengali) the leading space is not merged into + # the first word and becomes a lone space token, which the decoder + # never samples first (see SuppressBlank in decoding.py); drop it. + text = text[1:] + return prompt + text + [tokenizer.eot], len(prompt) + + +class SpeechDataset(Dataset): + def __init__( + self, samples: List[Dict[str, str]], tokenizers: TokenizerCache, n_mels: int + ): + self.samples = samples + self.tokenizers = tokenizers + self.n_mels = n_mels + + def __len__(self): + return len(self.samples) + + def __getitem__(self, index: int): + sample = self.samples[index] + try: + audio = load_audio(sample["audio"]) + except Exception as e: + warn(f"skipping {sample['audio']}: failed to load audio ({e})") + return None + if len(audio) > N_SAMPLES: + warn( + f"skipping {sample['audio']}: {len(audio) / SAMPLE_RATE:.1f}s is longer " + f"than the 30s that Whisper can process in one window" + ) + return None + + mel = log_mel_spectrogram(pad_or_trim(audio), n_mels=self.n_mels) + tokens, prompt_length = self.tokenizers.encode(sample) + return dict( + mel=mel, + tokens=tokens, + prompt_length=prompt_length, + text=sample["text"], + language=self.tokenizers.get(sample.get("language")).language, + ) + + +def collate(items: List[Optional[dict]], pad_token: int) -> Optional[dict]: + items = [item for item in items if item is not None] + if not items: + return None + + max_length = max(len(item["tokens"]) for item in items) - 1 + inputs = torch.full((len(items), max_length), pad_token, dtype=torch.long) + labels = torch.full((len(items), max_length), IGNORE_INDEX, dtype=torch.long) + for i, item in enumerate(items): + tokens = torch.tensor(item["tokens"], dtype=torch.long) + inputs[i, : len(tokens) - 1] = tokens[:-1] + labels[i, : len(tokens) - 1] = tokens[1:] + # don't train on predicting the start-of-transcript prompt; the first + # supervised position is the one predicting the first text token + labels[i, : item["prompt_length"] - 1] = IGNORE_INDEX + + return dict( + mel=torch.stack([item["mel"] for item in items]), + inputs=inputs, + labels=labels, + texts=[item["text"] for item in items], + languages=[item["language"] for item in items], + ) + + +def filter_samples( + samples: List[Dict[str, str]], tokenizers: TokenizerCache, n_text_ctx: int +) -> List[Dict[str, str]]: + """Drop samples whose audio file is missing or whose transcript is too long""" + kept, missing, too_long = [], [], [] + for sample in samples: + if not os.path.isfile(sample["audio"]): + missing.append(sample["audio"]) + continue + tokens, _ = tokenizers.encode(sample) + if len(tokens) > n_text_ctx: + too_long.append(sample["audio"]) + continue + kept.append(sample) + + for paths, reason in [ + (missing, "audio file not found"), + (too_long, f"transcript longer than the decoder context ({n_text_ctx} tokens)"), + ]: + if paths: + examples = ", ".join(paths[:3]) + (", ..." if len(paths) > 3 else "") + warn(f"skipping {len(paths)} sample(s): {reason}: {examples}") + return kept + + +class EpochSampler(Sampler): + """A seeded shuffling sampler that can start from an offset, for exact resumption""" + + def __init__(self, length: int, seed: int, shuffle: bool = True): + self.length = length + self.seed = seed + self.shuffle = shuffle + self.epoch = 0 + self.start = 0 + + def set_epoch(self, epoch: int, start: int = 0): + self.epoch = epoch + self.start = start + + def __iter__(self) -> Iterator[int]: + if self.shuffle: + generator = torch.Generator().manual_seed(self.seed + self.epoch) + order = torch.randperm(self.length, generator=generator).tolist() + else: + order = list(range(self.length)) + return iter(order[self.start :]) + + def __len__(self): + return self.length - self.start + + +# --------------------------------------------------------------------------- +# Metrics +# --------------------------------------------------------------------------- + + +def edit_distance(reference: List[str], hypothesis: List[str]) -> int: + previous = list(range(len(hypothesis) + 1)) + for i, ref in enumerate(reference, 1): + current = [i] + [0] * len(hypothesis) + for j, hyp in enumerate(hypothesis, 1): + current[j] = min( + previous[j] + 1, current[j - 1] + 1, previous[j - 1] + (ref != hyp) + ) + previous = current + return previous[-1] + + +def normalize_text(text: str, language: Optional[str]) -> str: + """ + Normalize a transcript before computing WER/CER. + + English uses Whisper's EnglishTextNormalizer. Other languages only have + punctuation and symbols removed and are lowercased; combining marks are kept, + since in many scripts (e.g. Devanagari, Bengali, Tamil, Thai) they are vowel + signs that are part of the word rather than optional diacritics. + """ + if language in (None, "en"): + return EnglishTextNormalizer()(text) + text = unicodedata.normalize("NFKC", text).lower() + text = "".join(" " if unicodedata.category(c)[0] in "PS" else c for c in text) + return re.sub(r"\s+", " ", text).strip() + + +def error_rates( + references: List[str], hypotheses: List[str], languages: List[Optional[str]] +) -> Dict[str, float]: + word_errors = word_total = char_errors = char_total = 0 + for reference, hypothesis, language in zip(references, hypotheses, languages): + reference = normalize_text(reference, language) + hypothesis = normalize_text(hypothesis, language) + word_errors += edit_distance(reference.split(), hypothesis.split()) + word_total += len(reference.split()) + char_errors += edit_distance(list(reference), list(hypothesis)) + char_total += len(reference) + return dict( + wer=word_errors / max(word_total, 1), cer=char_errors / max(char_total, 1) + ) + + +# --------------------------------------------------------------------------- +# Training utilities +# --------------------------------------------------------------------------- + + +def enable_gradient_checkpointing(model: Whisper) -> None: + """Recompute each residual block's activations in the backward pass to save memory""" + + def wrap(forward): + @functools.wraps(forward) + def checkpointed(*args, **kwargs): + if torch.is_grad_enabled(): + return checkpoint(forward, *args, use_reentrant=False, **kwargs) + return forward(*args, **kwargs) + + return checkpointed + + for block in [*model.encoder.blocks, *model.decoder.blocks]: + block.forward = wrap(block.forward) + + +def lr_lambda(step: int, warmup: int, total: int, schedule: str) -> float: + if step < warmup: + return (step + 1) / (warmup + 1) + if schedule == "constant": + return 1.0 + progress = (step - warmup) / max(1, total - warmup) + progress = min(progress, 1.0) + if schedule == "cosine": + return 0.5 * (1.0 + math.cos(math.pi * progress)) + return 1.0 - progress + + +def compute_loss( + model: Whisper, batch: dict, precision: Precision, train_encoder: bool +) -> torch.Tensor: + device = model.device + mel = batch["mel"].to(device, non_blocking=True) + inputs = batch["inputs"].to(device, non_blocking=True) + labels = batch["labels"].to(device, non_blocking=True) + + with precision.autocast(): + if train_encoder: + audio_features = model.encoder(mel) + else: + with torch.no_grad(): + audio_features = model.encoder(mel) + logits = model.decoder(inputs, audio_features) + + return F.cross_entropy( + logits.float().reshape(-1, logits.shape[-1]), + labels.reshape(-1), + ignore_index=IGNORE_INDEX, + ) + + +@torch.no_grad() +def evaluate( + model: Whisper, + loader: DataLoader, + precision: Precision, + task: str, + beam_size: Optional[int] = None, +) -> Dict[str, float]: + model.eval() + total_loss, total_batches = 0.0, 0 + references, hypotheses, languages = [], [], [] + + for batch in loader: + if batch is None: + continue + total_loss += compute_loss(model, batch, precision, True).item() + total_batches += 1 + + mel = batch["mel"].to(model.device) + for language in dict.fromkeys(batch["languages"]): + index = [i for i, lang in enumerate(batch["languages"]) if lang == language] + options = DecodingOptions( + task=task, + language=language, + without_timestamps=True, + beam_size=beam_size, + fp16=precision.name == "fp16", + ) + # fp16 decoding uses Whisper's own fp16 inference path; bf16 decoding + # runs the fp32 path, since decode() only accepts fp16 or fp32 audio + # features and autocast would produce bf16 ones + results = model.decode(mel[index], options) + hypotheses.extend(result.text for result in results) + references.extend(batch["texts"][i] for i in index) + languages.extend([language] * len(index)) + + model.train() + metrics = dict(loss=total_loss / max(total_batches, 1)) + metrics.update(error_rates(references, hypotheses, languages)) + metrics["samples"] = len(references) + metrics["predictions"] = [ + dict(reference=r, hypothesis=h, language=lang) + for r, h, lang in zip(references, hypotheses, languages) + ] + return metrics + + +def set_seed(seed: int) -> None: + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + + +def trainable_parameters(model: Whisper): + return [p for p in model.parameters() if p.requires_grad] + + +def save_checkpoint(path: str, model, optimizer, scheduler, scaler, state: dict): + trainable = { + n: p.detach().cpu() for n, p in model.named_parameters() if p.requires_grad + } + torch.save( + dict( + model=trainable, + optimizer=optimizer.state_dict(), + scheduler=scheduler.state_dict(), + scaler=scaler.state_dict(), + state=state, + rng=dict( + python=random.getstate(), + numpy=np.random.get_state(), + torch=torch.get_rng_state(), + cuda=( + torch.cuda.get_rng_state_all() if torch.cuda.is_available() else [] + ), + ), + ), + path, + ) + + +def load_checkpoint(path: str, model, optimizer, scheduler, scaler) -> dict: + # resumable checkpoints contain the optimizer and RNG states, not just tensors + kwargs = {"weights_only": False} if torch.__version__ >= "1.13" else {} + checkpoint = torch.load(path, map_location="cpu", **kwargs) + missing, unexpected = model.load_state_dict(checkpoint["model"], strict=False) + if unexpected: + raise ValueError( + f"checkpoint {path} does not match the model: {unexpected[:5]}" + ) + optimizer.load_state_dict(checkpoint["optimizer"]) + scheduler.load_state_dict(checkpoint["scheduler"]) + scaler.load_state_dict(checkpoint["scaler"]) + rng = checkpoint["rng"] + random.setstate(rng["python"]) + np.random.set_state(rng["numpy"]) + torch.set_rng_state(rng["torch"]) + if rng["cuda"] and torch.cuda.is_available(): + torch.cuda.set_rng_state_all(rng["cuda"]) + return checkpoint["state"] + + +def save_model(model: Whisper, path: str) -> None: + """Save the whole model as a regular Whisper checkpoint, loadable by load_model()""" + # stored in fp16 like the official checkpoints; load_model() casts it back + state_dict = {k: v.detach().half().cpu() for k, v in model.state_dict().items()} + torch.save({"dims": asdict(model.dims), "model_state_dict": state_dict}, path) + + +def list_checkpoints(output_dir: str) -> List[str]: + """Paths of the resumable checkpoints in output_dir, oldest first""" + directory = os.path.join(output_dir, "checkpoints") + if not os.path.isdir(directory): + return [] + steps = sorted( + int(m.group(1)) + for name in os.listdir(directory) + if (m := re.fullmatch(r"step_(\d+)\.pt", name)) + ) + return [os.path.join(directory, f"step_{step}.pt") for step in steps] + + +def latest_checkpoint(output_dir: str) -> Optional[str]: + checkpoints = list_checkpoints(output_dir) + return checkpoints[-1] if checkpoints else None + + +class MetricsLogger: + def __init__(self, output_dir: str): + self.path = os.path.join(output_dir, "metrics.jsonl") + + def __call__(self, **record): + with open(self.path, "a", encoding="utf-8") as f: + f.write(json.dumps(record) + "\n") + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + + +def train(args: argparse.Namespace) -> Dict[str, object]: + from . import load_model + + full = args.finetune == "full" + if not full and args.freeze_encoder: + raise ValueError( + "--freeze_encoder only applies to --finetune full; " + "use --lora_apply_to decoder for LoRA" + ) + if full and args.export_merged: + raise ValueError("--export_merged only applies to --finetune lora") + lr = args.lr if args.lr is not None else DEFAULT_LR[args.finetune] + + device = torch.device(args.device) + if device.type == "cuda" and device.index is not None: + torch.cuda.set_device(device) + precision = resolve_precision(args.precision, device) + + os.makedirs(os.path.join(args.output_dir, "checkpoints"), exist_ok=True) + set_seed(args.seed) + log(f"Precision: {precision.name} (requested: {args.precision}) on {device}") + log(f"Fine-tuning: {args.finetune}, learning rate {lr:g}") + + model = load_model(args.model, device="cpu", download_root=args.model_dir) + if not model.is_multilingual and args.language not in (None, "en", "english"): + warn(f"{args.model} is English-only; ignoring --language {args.language}") + + tokenizers = TokenizerCache(model, args.language, args.task) + train_samples = read_manifest(args.train_manifest) + train_samples = filter_samples(train_samples, tokenizers, model.dims.n_text_ctx) + val_samples = [] + if args.val_manifest: + val_samples = read_manifest(args.val_manifest) + val_samples = filter_samples(val_samples, tokenizers, model.dims.n_text_ctx) + if args.max_eval_samples: + val_samples = val_samples[: args.max_eval_samples] + if not train_samples: + raise ValueError("no usable training samples") + log(f"Samples: {len(train_samples)} train, {len(val_samples)} validation") + + if full: + train_encoder = not args.freeze_encoder + for p in model.parameters(): + p.requires_grad = True + if not train_encoder: + model.encoder.requires_grad_(False) + else: + apply_lora( + model, + r=args.lora_r, + alpha=args.lora_alpha, + dropout=args.lora_dropout, + target_modules=args.lora_targets, + apply_to=args.lora_apply_to, + ) + train_encoder = "encoder" in model.lora_config["apply_to"] + # the parameters stay in fp32 and act as master weights under autocast + model.to(device) + if args.gradient_checkpointing: + enable_gradient_checkpointing(model) + + parameters = trainable_parameters(model) + n_trainable = sum(p.numel() for p in parameters) + n_total = sum(p.numel() for p in model.parameters()) + log( + f"Trainable parameters: {n_trainable:,} / {n_total:,} " + f"({100 * n_trainable / n_total:.2f}%)" + ) + + # the EOT token is the same for every language, so this works even when the + # languages only come from the manifest + pad_token = get_tokenizer( + model.is_multilingual, num_languages=model.num_languages + ).eot + loader_args = dict( + batch_size=args.batch_size, + num_workers=args.num_workers, + collate_fn=functools.partial(collate, pad_token=pad_token), + pin_memory=device.type == "cuda", + # a separate generator keeps data loading from consuming the global RNG, + # so that resumed runs reproduce the same dropout masks + generator=torch.Generator().manual_seed(args.seed), + ) + sampler = EpochSampler(len(train_samples), args.seed) + train_loader = DataLoader( + SpeechDataset(train_samples, tokenizers, model.dims.n_mels), + sampler=sampler, + drop_last=False, + **loader_args, + ) + val_loader = None + if val_samples: + val_loader = DataLoader( + SpeechDataset(val_samples, tokenizers, model.dims.n_mels), **loader_args + ) + + batches_per_epoch = math.ceil(len(train_samples) / args.batch_size) + steps_per_epoch = math.ceil(batches_per_epoch / args.grad_accum) + total_steps = args.max_steps or steps_per_epoch * args.epochs + warmup_steps = ( + args.warmup_steps if args.warmup_steps is not None else int(0.05 * total_steps) + ) + + optimizer = torch.optim.AdamW(parameters, lr=lr, weight_decay=args.weight_decay) + scheduler = torch.optim.lr_scheduler.LambdaLR( + optimizer, + functools.partial( + lr_lambda, + warmup=warmup_steps, + total=total_steps, + schedule=args.lr_scheduler, + ), + ) + scaler = precision.grad_scaler() + + state = dict(step=0, epoch=0, batch_in_epoch=0, best_wer=None, best_loss=None) + resume = args.resume + if resume == "latest": + resume = latest_checkpoint(args.output_dir) + if resume is None: + log("No checkpoint found to resume from; starting from scratch") + if resume: + state = load_checkpoint(resume, model, optimizer, scheduler, scaler) + log(f"Resumed from {resume} at step {state['step']}") + + run_config = dict( + vars(args), + lr=lr, + precision_resolved=precision.name, + device_name=( + torch.cuda.get_device_name(device) if device.type == "cuda" else "cpu" + ), + torch_version=torch.__version__, + lora_config=getattr(model, "lora_config", None), + dims=asdict(model.dims), + trainable_parameters=n_trainable, + total_parameters=n_total, + total_steps=total_steps, + warmup_steps=warmup_steps, + train_samples=len(train_samples), + val_samples=len(val_samples), + ) + with open(os.path.join(args.output_dir, "run_config.json"), "w") as f: + json.dump(run_config, f, indent=2) + + metrics_log = MetricsLogger(args.output_dir) + prefix = "model" if full else "adapter" + best_path = os.path.join(args.output_dir, f"{prefix}_best.pt") + final_path = os.path.join(args.output_dir, f"{prefix}_final.pt") + results: Dict[str, object] = dict(finetune=args.finetune, precision=precision.name) + + def save_weights(path: str) -> None: + if full: + save_model(model, path) + else: + save_lora(model, path, base_model=args.model) + + def run_eval(step: int, tag: str) -> Optional[dict]: + if val_loader is None: + return None + start = time.time() + metrics = evaluate(model, val_loader, precision, args.task, args.eval_beam_size) + log( + f"[eval {tag}] step {step}: loss {metrics['loss']:.4f} " + f"WER {100 * metrics['wer']:.2f}% CER {100 * metrics['cer']:.2f}% " + f"({metrics['samples']} samples, {time.time() - start:.1f}s)" + ) + predictions = metrics.pop("predictions") + path = os.path.join(args.output_dir, "predictions", f"step_{step}.jsonl") + os.makedirs(os.path.dirname(path), exist_ok=True) + with open(path, "w", encoding="utf-8") as f: + for prediction in predictions: + f.write(json.dumps(prediction, ensure_ascii=False) + "\n") + metrics_log(type="eval", tag=tag, step=step, **metrics) + return metrics + + def maybe_save_best(metrics: Optional[dict]) -> None: + if metrics is None: + return + if state["best_wer"] is None or metrics["wer"] < state["best_wer"]: + state["best_wer"] = metrics["wer"] + state["best_loss"] = metrics["loss"] + save_weights(best_path) + log(f"New best WER {100 * metrics['wer']:.2f}%; saved {best_path}") + + if state["step"] == 0 and not args.skip_baseline: + # nothing has been trained yet (LoRA B matrices start at zero), so this is + # the unmodified base model + results["baseline"] = run_eval(0, "baseline") + + if device.type == "cuda": + torch.cuda.reset_peak_memory_stats(device) + + model.train() + nonfinite_streak = 0 + window_loss, window_batches, window_steps = 0.0, 0, 0 + step_times: List[float] = [] + micro_batch = 0 + step_start = time.time() + + while state["step"] < total_steps: + sampler.set_epoch(state["epoch"], state["batch_in_epoch"] * args.batch_size) + for batch in train_loader: + state["batch_in_epoch"] += 1 + if batch is None: + continue + + if micro_batch % args.grad_accum == 0: + step_start = time.time() + loss = compute_loss(model, batch, precision, train_encoder) + if not torch.isfinite(loss): + nonfinite_streak += 1 + warn(f"non-finite loss {loss.item()} at step {state['step']}") + if nonfinite_streak >= 10: + raise RuntimeError( + "loss was NaN/inf for 10 consecutive batches; try a lower " + "learning rate or --precision fp32/bf16" + ) + continue # skip the batch rather than propagating NaN/inf + nonfinite_streak = 0 + window_loss += loss.item() + window_batches += 1 + scaler.scale(loss / args.grad_accum).backward() + micro_batch += 1 + + if micro_batch % args.grad_accum != 0: + continue + + scaler.unscale_(optimizer) + grad_norm = torch.nn.utils.clip_grad_norm_(parameters, args.max_grad_norm) + if not precision.use_grad_scaler and not torch.isfinite(grad_norm): + # GradScaler skips such steps in fp16; do the same for fp32/bf16 + warn(f"non-finite gradient norm at step {state['step']}; skipping") + optimizer.zero_grad(set_to_none=True) + continue + scaler.step(optimizer) + scaler.update() + optimizer.zero_grad(set_to_none=True) + scheduler.step() + state["step"] += 1 + step_times.append(time.time() - step_start) + window_steps += 1 + + step = state["step"] + if step % args.log_every == 0 or step == total_steps: + record = dict( + type="train", + step=step, + epoch=state["epoch"], + loss=window_loss / max(window_batches, 1), + lr=scheduler.get_last_lr()[0], + grad_norm=float(grad_norm), + seconds_per_step=float(np.mean(step_times[-window_steps:])), + ) + if precision.use_grad_scaler: + record["loss_scale"] = scaler.get_scale() + if device.type == "cuda": + record["peak_memory_gb"] = ( + torch.cuda.max_memory_allocated(device) / 1024**3 + ) + log( + " ".join( + f"{k} {v:.4g}" if isinstance(v, float) else f"{k} {v}" + for k, v in record.items() + ) + ) + metrics_log(**record) + window_loss, window_batches, window_steps = 0.0, 0, 0 + + if args.eval_every and step % args.eval_every == 0 and step < total_steps: + maybe_save_best(run_eval(step, "step")) + if args.save_every and step % args.save_every == 0: + path = os.path.join(args.output_dir, "checkpoints", f"step_{step}.pt") + save_checkpoint(path, model, optimizer, scheduler, scaler, state) + log(f"Saved checkpoint {path}") + if args.keep_checkpoints: + for old in list_checkpoints(args.output_dir)[ + : -args.keep_checkpoints + ]: + os.remove(old) + if step >= total_steps: + break + else: + state["epoch"] += 1 + state["batch_in_epoch"] = 0 + if not args.eval_every: + maybe_save_best(run_eval(state["step"], f"epoch {state['epoch']}")) + continue + break + + save_weights(final_path) + log(f"Saved final {prefix} to {final_path}") + final = run_eval(state["step"], "final") + maybe_save_best(final) + + results.update( + final=final, + best_wer=state["best_wer"], + steps=state["step"], + trainable_parameters=n_trainable, + trainable_percent=100 * n_trainable / n_total, + seconds_per_step=float(np.median(step_times)) if step_times else None, + ) + if device.type == "cuda": + results["peak_memory_gb"] = torch.cuda.max_memory_allocated(device) / 1024**3 + if results.get("baseline") and final: + log( + f"Baseline WER {100 * results['baseline']['wer']:.2f}% -> " + f"fine-tuned WER {100 * final['wer']:.2f}% " + f"(best {100 * state['best_wer']:.2f}%)" + ) + + if args.export_merged: + adapter = best_path if os.path.isfile(best_path) else final_path + merged = load_model(args.model, device="cpu", download_root=args.model_dir) + merge_lora(load_lora(merged, adapter, base_model=args.model)) + merged_path = os.path.join(args.output_dir, "merged_model.pt") + save_model(merged, merged_path) + log(f"Exported merged model ({adapter}) to {merged_path}") + results["merged_model"] = merged_path + + with open(os.path.join(args.output_dir, "results.json"), "w") as f: + json.dump(results, f, indent=2) + return results + + +def build_parser() -> argparse.ArgumentParser: + from . import available_models + + def valid_model_name(name): + if name in available_models() or os.path.exists(name): + return name + raise ValueError( + f"model should be one of {available_models()} or path to a model checkpoint" + ) + + def valid_language(name): + if name.lower() in LANGUAGES or name.lower() in TO_LANGUAGE_CODE: + return name + raise ValueError(f"unsupported language: {name}") + + # fmt: off + parser = argparse.ArgumentParser(description="Fine-tune Whisper with LoRA or full fine-tuning", formatter_class=argparse.ArgumentDefaultsHelpFormatter) + parser.add_argument("--model", default="small", type=valid_model_name, help="name of the Whisper model to fine-tune, or path to a checkpoint") + parser.add_argument("--model_dir", type=str, default=None, help="the path to save model files; uses ~/.cache/whisper by default") + parser.add_argument("--train_manifest", type=str, required=True, help="JSONL or CSV file with the training samples") + parser.add_argument("--val_manifest", type=str, default=None, help="JSONL or CSV file with the validation samples") + parser.add_argument("--output_dir", "-o", type=str, required=True, help="directory to save adapters or models, checkpoints and logs") + parser.add_argument("--language", type=valid_language, default=None, help="language of samples that don't specify one") + parser.add_argument("--task", type=str, default="transcribe", choices=["transcribe", "translate"], help="whether to train X->X transcription or X->English translation") + + parser.add_argument("--finetune", type=str, default="lora", choices=FINETUNE_MODES, help="lora: train small adapters on a frozen model; full: update all weights (needs much more memory)") + parser.add_argument("--freeze_encoder", action="store_true", help="with --finetune full, train only the decoder") + parser.add_argument("--lora_r", type=int, default=8, help="rank of the LoRA update") + parser.add_argument("--lora_alpha", type=float, default=16, help="LoRA scaling numerator; the update is scaled by alpha / r") + parser.add_argument("--lora_dropout", type=float, default=0.05, help="dropout applied to the input of the LoRA layers") + parser.add_argument("--lora_targets", type=str, default="query,value", help=f"comma-separated layers to adapt, from {','.join(ALL_TARGETS)}, or 'all'") + parser.add_argument("--lora_apply_to", type=str, default="encoder,decoder", help=f"comma-separated parts of the model to adapt, from {','.join(ALL_PARTS)}") + + parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu", help="device to train on, e.g. cuda, cuda:1 or cpu") + parser.add_argument("--precision", type=str, default="auto", choices=PRECISIONS, help="auto: bf16 if the GPU supports it, else fp16 on GPU and fp32 on CPU") + parser.add_argument("--epochs", type=int, default=3, help="number of passes over the training set") + parser.add_argument("--max_steps", type=int, default=None, help="number of optimizer steps; overrides --epochs") + parser.add_argument("--batch_size", type=int, default=8, help="number of samples per forward pass") + parser.add_argument("--grad_accum", type=int, default=1, help="number of batches to accumulate per optimizer step") + parser.add_argument("--lr", type=float, default=None, help=f"peak learning rate; defaults to {DEFAULT_LR['lora']:g} for LoRA and {DEFAULT_LR['full']:g} for full fine-tuning") + parser.add_argument("--weight_decay", type=float, default=0.0, help="AdamW weight decay") + parser.add_argument("--lr_scheduler", type=str, default="linear", choices=["linear", "cosine", "constant"], help="learning rate decay after warmup") + parser.add_argument("--warmup_steps", type=int, default=None, help="linear warmup steps; defaults to 5%% of the total steps") + parser.add_argument("--max_grad_norm", type=float, default=1.0, help="gradient clipping threshold") + parser.add_argument("--gradient_checkpointing", action="store_true", help="recompute activations in the backward pass to save memory") + parser.add_argument("--num_workers", type=int, default=4, help="number of data loading processes") + parser.add_argument("--seed", type=int, default=0, help="random seed") + + parser.add_argument("--eval_every", type=int, default=0, help="evaluate every N steps; 0 evaluates at the end of each epoch") + parser.add_argument("--eval_beam_size", type=int, default=None, help="beam size for evaluation decoding; greedy if not set") + parser.add_argument("--max_eval_samples", type=int, default=None, help="evaluate on at most this many validation samples") + parser.add_argument("--skip_baseline", action="store_true", help="don't evaluate the base model before training") + parser.add_argument("--log_every", type=int, default=10, help="log training loss every N steps") + parser.add_argument("--save_every", type=int, default=0, help="save a resumable checkpoint every N steps; 0 disables") + parser.add_argument("--keep_checkpoints", type=int, default=0, help="keep only the N most recent resumable checkpoints; 0 keeps all") + parser.add_argument("--resume", type=str, default=None, help="checkpoint to resume from, or 'latest' to use the newest one in --output_dir") + parser.add_argument("--export_merged", action="store_true", help="with LoRA, also save the best adapter merged into a full Whisper checkpoint") + # fmt: on + return parser + + +def cli(): + args = build_parser().parse_args() + try: + train(args) + except ValueError as e: + print(f"error: {e}", file=sys.stderr) + sys.exit(2) + + +if __name__ == "__main__": + cli() diff --git a/whisper/transcribe.py b/whisper/transcribe.py index 56b29fbc8..dd159866a 100644 --- a/whisper/transcribe.py +++ b/whisper/transcribe.py @@ -529,6 +529,7 @@ def valid_model_name(name): parser.add_argument("audio", nargs="+", type=str, help="audio file(s) to transcribe") parser.add_argument("--model", default="turbo", type=valid_model_name, help="name of the Whisper model to use") parser.add_argument("--model_dir", type=str, default=None, help="the path to save model files; uses ~/.cache/whisper by default") + parser.add_argument("--lora", type=str, default=None, help="path to a LoRA adapter file to apply to the model") parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu", help="device to use for PyTorch inference") parser.add_argument("--output_dir", "-o", type=str, default=".", help="directory to save the outputs") parser.add_argument("--output_format", "-f", type=str, default="all", choices=["txt", "vtt", "srt", "tsv", "json", "jsonl", "all"], help="format of the output file; if not specified, all available formats will be produced") @@ -569,6 +570,7 @@ def valid_model_name(name): args = parser.parse_args().__dict__ model_name: str = args.pop("model") model_dir: str = args.pop("model_dir") + lora: Optional[str] = args.pop("lora") output_dir: str = args.pop("output_dir") output_format: str = args.pop("output_format") device: str = args.pop("device") @@ -592,7 +594,7 @@ def valid_model_name(name): from . import load_model - model = load_model(model_name, device=device, download_root=model_dir) + model = load_model(model_name, device=device, download_root=model_dir, lora=lora) writer = get_writer(output_format, output_dir) word_options = [