train/whisper-large-lora-DDP.ipynb

{
	"cells": [
		{
			"cell_type": "markdown",
			"metadata": {},
			"source": [
				"# fine-tine whisper large with PEFT-LoRA + int4 + DDP\n",
				"\n",
				"better use of multi GPU with Distributed Data Parallelism (the other notebook use naive model parallelism)\n",
				"\n",
				"**attention**: in this case batch size value is total on all GPU and is evenly splited across (instead of batch size × GPU count)"
			]
		},
		{
			"cell_type": "code",
			"execution_count": null,
			"metadata": {
				"collapsed": false,
				"jupyter": {
					"outputs_hidden": false
				},
				"trusted": true
			},
			"outputs": [],
			"source": [
				"from huggingface_hub import notebook_login\n",
				"notebook_login()\n",
				"# !huggingface-cli login --token=███"
			]
		},
		{
			"cell_type": "code",
			"execution_count": null,
			"metadata": {
				"collapsed": false,
				"jupyter": {
					"outputs_hidden": false
				},
				"scrolled": true,
				"trusted": true
			},
			"outputs": [],
			"source": [
				"# workaround for a bug in `datasets` package\n",
				"%pip uninstall -y cudf dask-cuda dask-cudf\n",
				"%pip install -q cudf-cu12 --extra-index-url=https://pypi.nvidia.com\n",
				"%pip install -qU \"datasets[audio]\" accelerate transformers bitsandbytes peft\n",
				"# no compute metrics so no `jiwer`"
			]
		},
		{
			"cell_type": "code",
			"execution_count": null,
			"metadata": {
				"collapsed": false,
				"jupyter": {
					"outputs_hidden": false
				},
				"trusted": true
			},
			"outputs": [],
			"source": [
				"# everything must be inside this function\n",
				"def train_ddp(pretrained_model, batch_size, total_steps, save_path, resume_training):\n",
				"\tfrom dataclasses import dataclass\n",
				"\timport torch\n",
				"\timport datasets as hugDS\n",
				"\tfrom transformers import WhisperForConditionalGeneration, WhisperFeatureExtractor, WhisperTokenizer, BitsAndBytesConfig, Seq2SeqTrainingArguments, Seq2SeqTrainer\n",
				"\tfrom accelerate import Accelerator\n",
				"\timport peft\n",
				"\n",
				"\thas_bf16 = torch.cuda.is_bf16_supported()  # GPU Ampere or later\n",
				"\taccelerator = Accelerator(project_dir=save_path, log_with=\"tensorboard\", mixed_precision=\"bf16\" if has_bf16 else \"fp16\")\n",
				"\n",
				"\tFEATURE_EXTRACTOR = WhisperFeatureExtractor.from_pretrained(pretrained_model)\n",
				"\tTOKENIZER = WhisperTokenizer.from_pretrained(pretrained_model, language=\"vi\", task=\"transcribe\")\n",
				"\tMODEL = WhisperForConditionalGeneration.from_pretrained(\n",
				"\t\tpretrained_model, use_cache=False, device_map={\"\": accelerator.device},\n",
				"\t\tquantization_config=BitsAndBytesConfig(\n",
				"\t\t\tload_in_4bit=True, bnb_4bit_use_double_quant=True, bnb_4bit_quant_type=\"nf4\",\n",
				"\t\t\tbnb_4bit_compute_dtype=torch.bfloat16 if has_bf16 else torch.float16\n",
				"\t\t)\n",
				"\t)\n",
				"\tMODEL.config.forced_decoder_ids = None\n",
				"\tMODEL.config.suppress_tokens = []\n",
				"\n",
				"\tDUMMY_TOKEN = -100\n",
				"\n",
				"\tMODEL_BIS = peft.get_peft_model(\n",
				"\t\tpeft.prepare_model_for_kbit_training(MODEL, use_gradient_checkpointing=True, gradient_checkpointing_kwargs={\"use_reentrant\": False}),\n",
				"\t\tpeft.LoraConfig(r=32, lora_alpha=64, target_modules=[\"q_proj\", \"v_proj\"], lora_dropout=.05, bias=\"none\")\n",
				"\t\t# peft.AdaLoraConfig(…)  # higher number of trainable parameters\n",
				"\t)\n",
				"\tMODEL_BIS.model.model.encoder.conv1.register_forward_hook(lambda module, input, output: output.requires_grad_(True))  # re-enable grad computation for conv layer\n",
				"\n",
				"\tif accelerator.is_main_process:\n",
				"\t\tMODEL_BIS.print_trainable_parameters()  # 16 millions = 1% of 1.6 billions params of whisper large\n",
				"\n",
				"\t###############################################################################\n",
				"\t# prepare data\n",
				"\n",
				"\tSAMPLING_RATE = 16_000\n",
				"\tdef load_my_data(**kwargs):\n",
				"\t\treturn hugDS.load_dataset(**kwargs, split=\"train\", trust_remote_code=True, streaming=True).cast_column(\"audio\", hugDS.Audio(sampling_rate=SAMPLING_RATE))\n",
				"\n",
				"\n",
				"\twith accelerator.main_process_first():\n",
				"\t\tMY_DATA = hugDS.concatenate_datasets([  # total: 86k samples\n",
				"\t\t\tload_my_data(path=\"doof-ferb/fpt_fosd\"),  # 25.9k\n",
				"\t\t\tload_my_data(path=\"doof-ferb/infore1_25hours\"),  # 14.9k\n",
				"\t\t\tload_my_data(path=\"doof-ferb/LSVSC\").select_columns([\"audio\", \"transcription\"]),  # 45k\n",
				"\t\t])\n",
				"\n",
				"\n",
				"\tdef prepare_dataset(batch):\n",
				"\t\taudio = batch[\"audio\"]\n",
				"\t\tbatch[\"input_length\"] = len(audio[\"array\"])  # compute input length\n",
				"\t\tbatch[\"input_features\"] = FEATURE_EXTRACTOR(audio[\"array\"], sampling_rate=SAMPLING_RATE).input_features[0]  # compute log-Mel input features\n",
				"\t\tbatch[\"labels\"] = TOKENIZER(batch[\"transcription\"]).input_ids  # encode target text to label ids\n",
				"\t\tbatch[\"labels_length\"] = len(batch[\"labels\"])  # compute labels length\n",
				"\t\treturn batch\n",
				"\n",
				"\tdef filter_inputs(input_length):\n",
				"\t\t\"\"\"Filter inputs with zero input length or longer than 30s\"\"\"\n",
				"\t\treturn 0 < input_length < 48e4  # 30s × 16kHz\n",
				"\n",
				"\tdef filter_labels(labels_length):\n",
				"\t\t\"\"\"Filter label sequences longer than max length 448 tokens\"\"\"\n",
				"\t\treturn labels_length < 448  # MODEL.config.max_length\n",
				"\n",
				"\tMY_DATA = (MY_DATA\n",
				"\t\t.map(prepare_dataset)  # no `num_proc` coz streaming\n",
				"\t\t.filter(filter_inputs, input_columns= [\"input_length\"])\n",
				"\t\t.filter(filter_labels, input_columns=[\"labels_length\"])\n",
				"\t)\n",
				"\n",
				"\t@dataclass\n",
				"\tclass DataCollatorSpeechSeq2SeqWithPadding:\n",
				"\t\tdef __call__(self, features):\n",
				"\t\t\t# split inputs and labels since they have to be of different lengths and need different padding methods\n",
				"\t\t\tinput_features = [{\"input_features\": feature[\"input_features\"]} for feature in features]\n",
				"\t\t\tlabel_features = [{\"input_ids\"     : feature[\"labels\"]        } for feature in features]  # get the tokenized label sequences\n",
				"\n",
				"\t\t\tbatch = FEATURE_EXTRACTOR.pad(input_features, return_tensors=\"pt\")  # treat the audio inputs by simply returning torch tensors\n",
				"\t\t\tlabels_batch =  TOKENIZER.pad(label_features, return_tensors=\"pt\")  # pad the labels to max length\n",
				"\t\t\tlabels = labels_batch[\"input_ids\"].masked_fill(labels_batch.attention_mask.ne(1), DUMMY_TOKEN)  # replace padding with -100 to ignore loss correctly\n",
				"\n",
				"\t\t\tif (labels[:, 0] == TOKENIZER.bos_token_id).all().cpu().item():  # if bos token is appended in previous tokenization step,\n",
				"\t\t\t\tlabels = labels[:, 1:]  # cut bos token here as it’s append later anyways\n",
				"\n",
				"\t\t\tbatch[\"labels\"] = labels\n",
				"\t\t\treturn batch\n",
				"\n",
				"\tDATA_COLLATOR = DataCollatorSpeechSeq2SeqWithPadding()\n",
				"\n",
				"\t###############################################################################\n",
				"\t# training setup\n",
				"\n",
				"\t# a practical learning rate while fine-tuning is a value 40× smaller than original used for pre-training\n",
				"\tif \"tiny\" in pretrained_model:\n",
				"\t\tLEARNING_RATE = 3.75e-5\n",
				"\telif \"base\" in pretrained_model:\n",
				"\t\tLEARNING_RATE = 2.5e-5\n",
				"\telif \"small\" in pretrained_model:\n",
				"\t\tLEARNING_RATE = 1.25e-5\n",
				"\telif \"medium\" in pretrained_model:\n",
				"\t\tLEARNING_RATE = 6.25e-6\n",
				"\telif \"large\" in pretrained_model:\n",
				"\t\tLEARNING_RATE = 5e-6\n",
				"\telse:\n",
				"\t\tLEARNING_RATE = 5e-5\n",
				"\n",
				"\n",
				"\tTRAINING_ARGS = Seq2SeqTrainingArguments(\n",
				"\t\toutput_dir=save_path,\n",
				"\t\tper_device_train_batch_size=batch_size,\n",
				"\t\t# per_device_eval_batch_size=batch_size,\n",
				"\t\tfp16=not has_bf16,\n",
				"\t\tbf16=has_bf16, tf32=has_bf16,\n",
				"\t\t# torch_compile=True,  # SDPA not support whisper yet\n",
				"\t\treport_to=[\"tensorboard\"],\n",
				"\n",
				"\t\tmax_steps=total_steps,\n",
				"\t\tlogging_steps=25,\n",
				"\t\tsave_steps=50,\n",
				"\t\tevaluation_strategy=\"no\",\n",
				"\t\tsave_total_limit=5,\n",
				"\t\taccelerator_config={\"split_batches\": True},  # mandatory for streaming datasets\n",
				"\n",
				"\t\toptim=\"adamw_bnb_8bit\",  # 8-bit AdamW optimizer: lower vram usage than default AdamW\n",
				"\t\tlearning_rate=LEARNING_RATE,\n",
				"\t\twarmup_steps=.05,  # keep between 5-15%\n",
				"\t\tgradient_accumulation_steps=1 if batch_size >= 8 else 8 // batch_size,  # keep effective batch size as min 8 per device\n",
				"\t\tremove_unused_columns=False, label_names=[\"labels\"],  # required by PEFT\n",
				"\t\t# predict_with_generate=True,  # must disable coz PEFT\n",
				"\t)\n",
				"\n",
				"\tTRAINER = Seq2SeqTrainer(\n",
				"\t\targs=TRAINING_ARGS,\n",
				"\t\tmodel=MODEL_BIS,\n",
				"\t\ttrain_dataset=MY_DATA,\n",
				"\t\tdata_collator=DATA_COLLATOR,\n",
				"\t\t# compute_metrics=compute_metrics,  # must disable coz PEFT\n",
				"\t\ttokenizer=FEATURE_EXTRACTOR,  # not TOKENIZER\n",
				"\t)\n",
				"\n",
				"\tTRAINER.train(resume_from_checkpoint=resume_training)\n",
				"\n",
				"\taccelerator.wait_for_everyone()\n",
				"\tif accelerator.is_main_process:\n",
				"\t\tTRAINER.save_model()"
			]
		},
		{
			"cell_type": "code",
			"execution_count": null,
			"metadata": {
				"collapsed": false,
				"jupyter": {
					"outputs_hidden": false
				},
				"trusted": true
			},
			"outputs": [],
			"source": [
				"from accelerate import notebook_launcher"
			]
		},
		{
			"cell_type": "code",
			"execution_count": null,
			"metadata": {
				"collapsed": false,
				"jupyter": {
					"outputs_hidden": false
				},
				"trusted": true
			},
			"outputs": [],
			"source": [
				"%cd /kaggle/working\n",
				"!rm -rf ./my-whisper-lora"
			]
		},
		{
			"cell_type": "code",
			"execution_count": null,
			"metadata": {
				"collapsed": false,
				"jupyter": {
					"outputs_hidden": false
				},
				"scrolled": true,
				"trusted": true
			},
			"outputs": [],
			"source": [
				"notebook_launcher(train_ddp, args=(\"vinai/PhoWhisper-large\", 10, 8000, \"./my-whisper-lora\", False), mixed_precision=\"fp16\", num_nodes=1, num_processes=2)"
			]
		},
		{
			"cell_type": "code",
			"execution_count": null,
			"metadata": {
				"collapsed": false,
				"jupyter": {
					"outputs_hidden": false
				},
				"trusted": true
			},
			"outputs": [],
			"source": [
				"!zip -FSr res.zip ./my-whisper-lora"
			]
		}
	],
	"metadata": {
		"accelerator": "GPU",
		"colab": {
			"gpuType": "T4",
			"private_outputs": true,
			"provenance": []
		},
		"kaggle": {
			"accelerator": "nvidiaTeslaT4",
			"dataSources": [],
			"isGpuEnabled": true,
			"isInternetEnabled": true,
			"language": "python",
			"sourceType": "notebook"
		},
		"kernelspec": {
			"display_name": "Python 3",
			"language": "python",
			"name": "python3"
		},
		"language_info": {
			"name": "python",
			"version": "3.11.7"
		}
	},
	"nbformat": 4,
	"nbformat_minor": 0
}