train.py 32.6 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
#!/usr/bin/env python
# coding=utf-8
# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

Yoach Lacombe's avatar
Yoach Lacombe committed
17
""" Train Parler-TTS using 🤗 Accelerate"""
18
19
20
21
22

import logging
import os
import re
import sys
Yoach Lacombe's avatar
Yoach Lacombe committed
23
import time
24
from multiprocess import set_start_method
25
from datetime import timedelta
26

27

Yoach Lacombe's avatar
Yoach Lacombe committed
28
import evaluate
29
from tqdm import tqdm
Yoach Lacombe's avatar
Yoach Lacombe committed
30
from pathlib import Path
31
32
33

import datasets
import torch
34
35
from torch.utils.data import DataLoader

Dan Lyth's avatar
Dan Lyth committed
36
from datasets import DatasetDict, Dataset, IterableDataset, concatenate_datasets
37

Yoach Lacombe's avatar
Yoach Lacombe committed
38
from huggingface_hub import Repository, create_repo
39
40
41
42
43
44
import transformers
from transformers import (
    AutoFeatureExtractor,
    AutoModel,
    AutoProcessor,
    AutoTokenizer,
Dan Lyth's avatar
Dan Lyth committed
45
    HfArgumentParser
46
)
Yoach Lacombe's avatar
Yoach Lacombe committed
47
from transformers.trainer_pt_utils import LengthGroupedSampler
Yoach Lacombe's avatar
Yoach Lacombe committed
48
49
from transformers import pipeline
from transformers.optimization import get_scheduler
Yoach Lacombe's avatar
Yoach Lacombe committed
50
from transformers.utils import send_example_telemetry
Yoach Lacombe's avatar
Yoach Lacombe committed
51
from transformers import AutoModel
Yoach Lacombe's avatar
add DAC  
Yoach Lacombe committed
52

53
54

from accelerate import Accelerator
55
from accelerate.utils import set_seed, AutocastKwargs, InitProcessGroupKwargs, TorchDynamoPlugin
Yoach Lacombe's avatar
Yoach Lacombe committed
56
from accelerate.utils.memory import release_memory
57

Yoach Lacombe's avatar
Yoach Lacombe committed
58
59
60
61
62
from parler_tts import (
    ParlerTTSForConditionalGeneration,
    ParlerTTSConfig,
    build_delay_pattern_mask,
)
63

Dan Lyth's avatar
Dan Lyth committed
64
65
from parler_tts.utils import get_last_checkpoint, rotate_checkpoints, log_pred, log_metric
from parler_tts.arguments import ModelArguments, DataTrainingArguments, ParlerTTSTrainingArguments
66
from parler_tts.data import DataCollatorParlerTTSWithPadding, DataCollatorEncodecWithPadding
67
68
69
70
71
72
73
74
75
76


logger = logging.getLogger(__name__)


def main():
    # See all possible arguments in src/transformers/training_args.py
    # or by passing the --help flag to this script.
    # We now keep distinct sets of args, for a cleaner separation of concerns.

Yoach Lacombe's avatar
Yoach Lacombe committed
77
    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, ParlerTTSTrainingArguments))
78
79
80
81
82
83
84
85
86
    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
        # If we pass only one argument to the script and it's the path to a json file,
        # let's parse it to get our arguments.
        model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
    else:
        model_args, data_args, training_args = parser.parse_args_into_dataclasses()

    # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
    # information sent is the one passed as arguments along with your Python/PyTorch versions.
Yoach Lacombe's avatar
Yoach Lacombe committed
87
    send_example_telemetry("run_parler_tts", model_args, data_args)
Yoach Lacombe's avatar
Yoach Lacombe committed
88

Yoach Lacombe's avatar
Yoach Lacombe committed
89
90
91
92
93
94
    if training_args.dtype == "float16":
        mixed_precision = "fp16"
    elif training_args.dtype == "bfloat16":
        mixed_precision = "bf16"
    else:
        mixed_precision = "no"
Yoach Lacombe's avatar
Yoach Lacombe committed
95
96
97
98
99
100
101
102
103

    if data_args.pad_to_max_length and (
        data_args.max_duration_in_seconds is None
        or data_args.max_prompt_token_length is None
        or data_args.max_description_token_length is None
    ):
        raise ValueError(
            "`pad_to_max_length` is `True` but one of the following parameters has not been set: `max_duration_in_seconds`, `max_prompt_token_length`, `max_description_token_length`"
        )
104
105

    padding = "max_length" if data_args.pad_to_max_length else "longest"
106

107
    # Accelerator preparation
108
109
110
    kwargs_handlers = [InitProcessGroupKwargs(timeout=timedelta(minutes=60))]
    if training_args.torch_compile:
        # TODO(YL): add more compile modes?
Yoach Lacombe's avatar
Yoach Lacombe committed
111
112
        kwargs_handlers.append(TorchDynamoPlugin(backend="inductor", mode="default"))  # reduce-overhead

Yoach Lacombe's avatar
Yoach Lacombe committed
113
114
115
116
117
    accelerator = Accelerator(
        gradient_accumulation_steps=training_args.gradient_accumulation_steps,
        mixed_precision=mixed_precision,
        log_with=training_args.report_to,
        project_dir=training_args.output_dir,
118
        kwargs_handlers=kwargs_handlers,
Yoach Lacombe's avatar
Yoach Lacombe committed
119
    )
Yoach Lacombe's avatar
Yoach Lacombe committed
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141

    accelerator.init_trackers(
        project_name=data_args.wandb_project,
        config={
            "learning_rate": training_args.learning_rate,
            "model_name_or_path": model_args.model_name_or_path,
            "num_train_epochs": training_args.num_train_epochs,
            "gradient_accumulation_steps": training_args.gradient_accumulation_steps,
            "per_device_train_batch_size": training_args.per_device_train_batch_size,
            "global_batch_size": training_args.per_device_train_batch_size * accelerator.num_processes,
            "mixed_precision": mixed_precision,
            "lr_scheduler_type": training_args.lr_scheduler_type,
            "warmup_steps": training_args.warmup_steps,
            "freeze_text_encoder": model_args.freeze_text_encoder,
            "max_duration_in_seconds": data_args.max_duration_in_seconds,
            "weight_decay": training_args.weight_decay,
            "adam_beta1": training_args.adam_beta1,
            "adam_beta2": training_args.adam_beta2,
            "temperature": model_args.temperature,
        },
    )

Yoach Lacombe's avatar
Yoach Lacombe committed
142
    # Detecting last checkpoint and eventually continue from last checkpoint
143
144
145
146
147
148
149
150
    last_checkpoint = None
    if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir:
        last_checkpoint = get_last_checkpoint(training_args.output_dir)
        if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 0:
            raise ValueError(
                f"Output directory ({training_args.output_dir}) already exists and is not empty. "
                "Use --overwrite_output_dir to overcome."
            )
Yoach Lacombe's avatar
Yoach Lacombe committed
151
        elif last_checkpoint is not None and training_args.resume_from_checkpoint is None:
152
153
154
155
156
157
158
159
160
161
162
            logger.info(
                f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change "
                "the `--output_dir` or add `--overwrite_output_dir` to train from scratch."
            )

    # Setup logging
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        handlers=[logging.StreamHandler(sys.stdout)],
    )
163
    logger.setLevel(logging.INFO if accelerator.is_main_process else logging.WARN)
164

Yoach Lacombe's avatar
Yoach Lacombe committed
165
    # Log a small summary on each proces
166
167
168
169
    logger.warning(
        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}, "
        f"distributed training: {training_args.parallel_mode.value == 'distributed'}, 16-bits training: {training_args.fp16}"
    )
Yoach Lacombe's avatar
Yoach Lacombe committed
170
171
172
173

    # Set the verbosity to info of the Transformers logger (on main process only)
    if accelerator.is_local_main_process:
        datasets.utils.logging.set_verbosity_warning()
174
        transformers.utils.logging.set_verbosity_info()
Yoach Lacombe's avatar
Yoach Lacombe committed
175
176
177
178
    else:
        datasets.utils.logging.set_verbosity_error()
        transformers.utils.logging.set_verbosity_error()

179
180
181
182
    logger.info("Training/evaluation parameters %s", training_args)

    # Set seed before initializing model.
    set_seed(training_args.seed)
183
    num_workers = data_args.preprocessing_num_workers
Yoach Lacombe's avatar
Yoach Lacombe committed
184

185
    # 1. First, let's instantiate the feature extractor (DAC), tokenizers and model
186
187
    # Note for distributed training, the .from_pretrained methods guarantee that only
    # one local process can concurrently download model & vocab.
Yoach Lacombe's avatar
Yoach Lacombe committed
188

189
190
191
192
193
194
195
196
    # load feature extractor
    feature_extractor = AutoFeatureExtractor.from_pretrained(
        model_args.feature_extractor_name or model_args.model_name_or_path,
        cache_dir=model_args.cache_dir,
        token=data_args.token,
        trust_remote_code=data_args.trust_remote_code,
    )
    sampling_rate = feature_extractor.sampling_rate
Yoach Lacombe's avatar
Yoach Lacombe committed
197

198
199
200
201
202
203
204
    # load prompt tokenizer
    prompt_tokenizer = AutoTokenizer.from_pretrained(
        model_args.prompt_tokenizer_name or model_args.description_tokenizer_name or model_args.model_name_or_path,
        cache_dir=model_args.cache_dir,
        token=data_args.token,
        trust_remote_code=data_args.trust_remote_code,
        use_fast=model_args.use_fast_tokenizer,
Yoach Lacombe's avatar
Yoach Lacombe committed
205
        padding_side="left",  # prompt has to be padded on the left bc it's preprend to codebooks hidden states
206
    )
Yoach Lacombe's avatar
Yoach Lacombe committed
207

208
209
210
211
212
213
214
215
    # load description tokenizer
    description_tokenizer = AutoTokenizer.from_pretrained(
        model_args.description_tokenizer_name or model_args.model_name_or_path,
        cache_dir=model_args.cache_dir,
        token=data_args.token,
        trust_remote_code=data_args.trust_remote_code,
        use_fast=model_args.use_fast_tokenizer,
    )
Yoach Lacombe's avatar
Yoach Lacombe committed
216

217
    if model_args.use_fast_tokenizer:
Yoach Lacombe's avatar
Yoach Lacombe committed
218
219
220
        logger.warning(
            "Disabling fast tokenizer warning: https://github.com/huggingface/transformers/blob/main/src/transformers/tokenization_utils_base.py#L3231-L3235"
        )
221
222
        prompt_tokenizer.deprecation_warnings["Asking-to-pad-a-fast-tokenizer"] = True
        description_tokenizer.deprecation_warnings["Asking-to-pad-a-fast-tokenizer"] = True
223

224
    # 2. Now, let's load the dataset
225
    # TODO add MDS dataset loading here
226

227
    # 3. Next, let's load the config.
Yoach Lacombe's avatar
Yoach Lacombe committed
228
    config = ParlerTTSConfig.from_pretrained(
229
230
231
232
233
        model_args.model_name_or_path,
        cache_dir=model_args.cache_dir,
        token=data_args.token,
        trust_remote_code=data_args.trust_remote_code,
    )
Yoach Lacombe's avatar
Yoach Lacombe committed
234

235
    # update pad token id and decoder_start_token_id
Yoach Lacombe's avatar
Yoach Lacombe committed
236
237
238
239
    config.update(
        {
            "pad_token_id": model_args.pad_token_id
            if model_args.pad_token_id is not None
240
            else config.pad_token_id,
Yoach Lacombe's avatar
Yoach Lacombe committed
241
242
            "decoder_start_token_id": model_args.decoder_start_token_id
            if model_args.decoder_start_token_id is not None
243
            else config.decoder_start_token_id,
Yoach Lacombe's avatar
Yoach Lacombe committed
244
245
246
        }
    )

Yoach Lacombe's avatar
Yoach Lacombe committed
247
    # create model
Yoach Lacombe's avatar
Yoach Lacombe committed
248
    model = ParlerTTSForConditionalGeneration.from_pretrained(
249
250
251
252
253
254
        model_args.model_name_or_path,
        cache_dir=model_args.cache_dir,
        config=config,
        token=data_args.token,
        trust_remote_code=data_args.trust_remote_code,
    )
Yoach Lacombe's avatar
Yoach Lacombe committed
255

256
257
258
    # enable gradient checkpointing if necessary
    if training_args.gradient_checkpointing:
        model.gradient_checkpointing_enable()
Yoach Lacombe's avatar
Yoach Lacombe committed
259

260
    # 4. Now we preprocess the datasets including loading the audio, resampling and normalization
261
    # TODO add MDS dataset preprocessing here (only thing we'll need is the delay pattern)
Yoach Lacombe's avatar
Yoach Lacombe committed
262

263
    # derive max & min input length for sample rate & max duration
264
    sampling_rate = feature_extractor.sampling_rate
Yoach Lacombe's avatar
Yoach Lacombe committed
265

266
    # Freeze Encoders
267
    model.freeze_encoders(model_args.freeze_text_encoder) # TODO check this implementation
Yoach Lacombe's avatar
Yoach Lacombe committed
268

269
    # 6. Next, we can prepare the training.
Yoach Lacombe's avatar
Yoach Lacombe committed
270

271
    # Let's use word CLAP similary and WER metrics as our evaluation metrics # TODO move this to seperate file
272

Yoach Lacombe's avatar
Yoach Lacombe committed
273
274
275
    # Define evaluation metrics during training, *i.e.* CLAP similarity
    clap = AutoModel.from_pretrained(model_args.clap_model_name_or_path)
    clap_processor = AutoProcessor.from_pretrained(model_args.clap_model_name_or_path)
Yoach Lacombe's avatar
Yoach Lacombe committed
276
    metric = evaluate.load("wer")
Yoach Lacombe's avatar
Yoach Lacombe committed
277

Yoach Lacombe's avatar
Yoach Lacombe committed
278
279
280
    def clap_similarity(texts, audios, device):
        clap_inputs = clap_processor(text=texts, audios=audios, padding=True, return_tensors="pt").to(device)
        clap.to(device)
281
        with torch.no_grad():
Yoach Lacombe's avatar
Yoach Lacombe committed
282
283
284
            text_features = clap.get_text_features(
                clap_inputs["input_ids"], attention_mask=clap_inputs.get("attention_mask", None)
            )
285
            audio_features = clap.get_audio_features(clap_inputs["input_features"])
Yoach Lacombe's avatar
Yoach Lacombe committed
286

287
            cosine_sim = torch.nn.functional.cosine_similarity(audio_features, text_features, dim=1, eps=1e-8)
Yoach Lacombe's avatar
Yoach Lacombe committed
288

Yoach Lacombe's avatar
Yoach Lacombe committed
289
290
        clap.to("cpu")
        clap_inputs.to("cpu")
291
        return cosine_sim.mean().to("cpu")
Yoach Lacombe's avatar
Yoach Lacombe committed
292

Yoach Lacombe's avatar
Yoach Lacombe committed
293
    def wer(prompts, audios, device):
Yoach Lacombe's avatar
Yoach Lacombe committed
294
        asr_pipeline = pipeline(model=model_args.asr_model_name_or_path, device=device)
Yoach Lacombe's avatar
Yoach Lacombe committed
295
296
297
298
299
300
301
302
303
        transcriptions = asr_pipeline(
            [{"raw": audio, "sampling_rate": sampling_rate} for audio in audios],
            batch_size=int(training_args.per_device_eval_batch_size),
        )

        word_error = 100 * metric.compute(
            predictions=[t["text"].lower() for t in transcriptions], references=[t.lower() for t in prompts]
        )

Yoach Lacombe's avatar
Yoach Lacombe committed
304
        return word_error, [t["text"] for t in transcriptions]
Yoach Lacombe's avatar
Yoach Lacombe committed
305

Yoach Lacombe's avatar
Yoach Lacombe committed
306
    eval_methods = {"clap": clap_similarity, "wer": wer}
307

Yoach Lacombe's avatar
Yoach Lacombe committed
308
309
    def compute_metrics(audios, descriptions, prompts, device="cpu"):
        input_ids = descriptions
310
        texts = description_tokenizer.batch_decode(input_ids, skip_special_tokens=True)
Yoach Lacombe's avatar
Yoach Lacombe committed
311
312
        prompts = prompt_tokenizer.batch_decode(prompts, skip_special_tokens=True)
        audios = [a.cpu().numpy() for a in audios]
Yoach Lacombe's avatar
Yoach Lacombe committed
313
        results = {"clap": eval_methods["clap"](texts, audios, device)}
Yoach Lacombe's avatar
Yoach Lacombe committed
314
315
        word_error, transcriptions = eval_methods["wer"](prompts, audios, device)
        results["wer"] = word_error
316

Yoach Lacombe's avatar
Yoach Lacombe committed
317
        return results, texts, prompts, audios, transcriptions
Yoach Lacombe's avatar
Yoach Lacombe committed
318

Yoach Lacombe's avatar
Yoach Lacombe committed
319
320
321
322
323
324
    # Define Training Schedule
    # Store some constants
    per_device_train_batch_size = int(training_args.per_device_train_batch_size)
    train_batch_size = per_device_train_batch_size * accelerator.num_processes
    gradient_accumulation_steps = int(training_args.gradient_accumulation_steps)
    per_device_eval_batch_size = int(training_args.per_device_eval_batch_size)
Yoach Lacombe's avatar
Yoach Lacombe committed
325

Yoach Lacombe's avatar
Yoach Lacombe committed
326
327
    if training_args.max_steps < 0:
        num_epochs = int(training_args.num_train_epochs)
328
        steps_per_epoch = len(vectorized_datasets["train"]) // (train_batch_size * gradient_accumulation_steps) # TODO fix this missing variable
Yoach Lacombe's avatar
Yoach Lacombe committed
329
330
331
332
333
334
335
336
337
        total_train_steps = steps_per_epoch * num_epochs
    elif training_args.max_steps > 0:
        logger.info("max_steps is given, it will override any value given in num_train_epochs")
        total_train_steps = int(training_args.max_steps)
        # Setting a very large number of epochs so we go as many times as necessary over the iterator.
        num_epochs = sys.maxsize
        steps_per_epoch = total_train_steps

    if training_args.eval_steps is None:
Yoach Lacombe's avatar
Yoach Lacombe committed
338
        logger.info(f"eval_steps is not set, evaluating at the end of each epoch")
Yoach Lacombe's avatar
Yoach Lacombe committed
339
340
341
        eval_steps = steps_per_epoch
    else:
        eval_steps = training_args.eval_steps
Yoach Lacombe's avatar
Yoach Lacombe committed
342

343
    # T5 doesn't support fp16
Yoach Lacombe's avatar
Yoach Lacombe committed
344
345
    autocast_kwargs = AutocastKwargs(enabled=(mixed_precision != "fp16"))

Yoach Lacombe's avatar
Yoach Lacombe committed
346
347
348
349
350
351
    # Define optimizer, LR scheduler, collator
    optimizer = torch.optim.AdamW(
        params=model.parameters(),
        lr=training_args.learning_rate,
        betas=(training_args.adam_beta1, training_args.adam_beta2),
        eps=training_args.adam_epsilon,
352
        weight_decay=training_args.weight_decay,
Yoach Lacombe's avatar
Yoach Lacombe committed
353
    )
354

Yoach Lacombe's avatar
Yoach Lacombe committed
355
356
357
358
    # LR scheduler gets stepped by `num_processes` each time -> account for this in warmup / total steps
    lr_scheduler = get_scheduler(
        name=training_args.lr_scheduler_type,
        optimizer=optimizer,
Yoach Lacombe's avatar
Yoach Lacombe committed
359
        num_warmup_steps=training_args.get_warmup_steps(total_train_steps) * accelerator.num_processes,
Yoach Lacombe's avatar
Yoach Lacombe committed
360
361
        num_training_steps=total_train_steps * accelerator.num_processes,
    )
362
363

    # Instantiate custom data collator
Yoach Lacombe's avatar
Yoach Lacombe committed
364
    data_collator = DataCollatorParlerTTSWithPadding(
Yoach Lacombe's avatar
Yoach Lacombe committed
365
366
367
368
369
370
        prompt_tokenizer=prompt_tokenizer,
        description_tokenizer=description_tokenizer,
        pad_to_multiple_of=data_args.pad_to_multiple_of,
        padding=padding,
        prompt_max_length=data_args.max_prompt_token_length,
        description_max_length=data_args.max_description_token_length,
371
        audio_max_length=audio_max_length, # TODO add this variable
372
    )
Yoach Lacombe's avatar
Yoach Lacombe committed
373

Yoach Lacombe's avatar
Yoach Lacombe committed
374
375
    # Prepare everything with accelerate
    model, optimizer, lr_scheduler = accelerator.prepare(model, optimizer, lr_scheduler)
Yoach Lacombe's avatar
Yoach Lacombe committed
376

Yoach Lacombe's avatar
Yoach Lacombe committed
377
378
    logger.info("***** Running training *****")
    logger.info(f"  Num examples = {total_train_steps * train_batch_size * gradient_accumulation_steps}")
379
    logger.info("  Instantaneous batch size per device =" f" {per_device_train_batch_size}")
Yoach Lacombe's avatar
Yoach Lacombe committed
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
    logger.info("  Gradient accumulation steps =" f" {gradient_accumulation_steps}")
    logger.info(
        f"  Total train batch size (w. parallel & distributed) = {train_batch_size * gradient_accumulation_steps}"
    )
    logger.info(f"  Total optimization steps = {total_train_steps}")

    # ======================== Training ================================
    train_time = 0
    train_start = time.time()
    steps_trained_progress_bar = tqdm(
        range(total_train_steps), desc="Train steps ... ", position=0, disable=not accelerator.is_local_main_process
    )
    continue_training = True
    epochs_trained = 0
    cur_step = 0

    checkpoint = None
    if training_args.resume_from_checkpoint is not None:
        checkpoint = training_args.resume_from_checkpoint
    elif last_checkpoint is not None:
        checkpoint = last_checkpoint
Yoach Lacombe's avatar
Yoach Lacombe committed
401

Yoach Lacombe's avatar
Yoach Lacombe committed
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
    if accelerator.is_main_process:
        if training_args.push_to_hub:
            # Retrieve of infer repo_name
            repo_name = training_args.hub_model_id
            if repo_name is None:
                repo_name = Path(training_args.output_dir).absolute().name
            # Create repo and retrieve repo_id
            repo_id = create_repo(repo_name, exist_ok=True, token=training_args.hub_token).repo_id
            # Clone repo locally
            repo = Repository(training_args.output_dir, clone_from=repo_id, token=training_args.hub_token)

            with open(os.path.join(training_args.output_dir, ".gitignore"), "w+") as gitignore:
                if "wandb" not in gitignore:
                    gitignore.write("wandb\n")
        elif training_args.output_dir is not None:
            os.makedirs(training_args.output_dir, exist_ok=True)
    accelerator.wait_for_everyone()
Yoach Lacombe's avatar
Yoach Lacombe committed
419

Yoach Lacombe's avatar
Yoach Lacombe committed
420
421
422
423
424
425
    # Now save everything to be able to create a single processor later
    # make sure all processes wait until data is saved
    with accelerator.main_process_first():
        # only the main process saves them
        if accelerator.is_main_process:
            # save feature extractor, tokenizer and config
Yoach Lacombe's avatar
Yoach Lacombe committed
426
427
428
429
430
            if (
                model_args.prompt_tokenizer_name is None
                and model_args.description_tokenizer_name
                or (model_args.prompt_tokenizer_name == model_args.description_tokenizer_name)
            ):
Yoach Lacombe's avatar
Yoach Lacombe committed
431
432
                prompt_tokenizer.save_pretrained(training_args.output_dir)
            else:
Yoach Lacombe's avatar
Yoach Lacombe committed
433
                logger.warning(
Dan Lyth's avatar
Dan Lyth committed
434
                    f"Prompt tokenizer ('{model_args.prompt_tokenizer_name}') and description tokenizer ('{model_args.description_tokenizer_name}') are not the same. Saving only the prompt tokenizer."
Yoach Lacombe's avatar
Yoach Lacombe committed
435
                )
Yoach Lacombe's avatar
Yoach Lacombe committed
436
                prompt_tokenizer.save_pretrained(training_args.output_dir)
Yoach Lacombe's avatar
Yoach Lacombe committed
437

Yoach Lacombe's avatar
Yoach Lacombe committed
438
439
            feature_extractor.save_pretrained(training_args.output_dir)
            config.save_pretrained(training_args.output_dir)
Yoach Lacombe's avatar
Yoach Lacombe committed
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456

    if checkpoint is not None:
        accelerator.load_state(checkpoint)
        # Find num steps and epoch from saved state string pattern
        pattern = r"checkpoint-(\d+)-epoch-(\d+)"
        match = re.search(pattern, checkpoint)
        cur_step = int(match.group(1))
        epochs_trained = int(match.group(2))

        logger.info("  Continuing training from checkpoint, will skip to saved global_step")
        logger.info(f"  Continuing training from epoch {epochs_trained}")
        logger.info(f"  Continuing training from global step {cur_step}")

        steps_trained_progress_bar.update(cur_step)

        for epoch in range(0, epochs_trained):
            vectorized_datasets["train"] = vectorized_datasets["train"].shuffle(training_args.seed)
Yoach Lacombe's avatar
Yoach Lacombe committed
457

Yoach Lacombe's avatar
Yoach Lacombe committed
458
459
        if training_args.max_steps < 0:
            # we know exactly the number of steps per epoch, so can skip through the required number of batches
460
            resume_step = (cur_step - epochs_trained * steps_per_epoch) * gradient_accumulation_steps
Yoach Lacombe's avatar
Yoach Lacombe committed
461
462
463
464
465
466
467
468
        else:
            # Currently we don't know how many steps we've taken in the current epoch
            # So we just shuffle the dataset one extra time and start from a fresh epoch
            # This is "good enough" for our purposes but not fully correct
            resume_step = None
            vectorized_datasets["train"] = vectorized_datasets["train"].shuffle(training_args.seed)
    else:
        resume_step = None
Yoach Lacombe's avatar
Yoach Lacombe committed
469

Yoach Lacombe's avatar
Yoach Lacombe committed
470
471
    gen_kwargs = {
        "do_sample": model_args.do_sample,
yoach@huggingface.co's avatar
yoach@huggingface.co committed
472
        "temperature": model_args.temperature,
Yoach Lacombe's avatar
Yoach Lacombe committed
473
474
        "max_length": model_args.max_length,
    }
Yoach Lacombe's avatar
Yoach Lacombe committed
475

Yoach Lacombe's avatar
Yoach Lacombe committed
476
477
478
    # Define gradient update step fn
    def train_step(
        batch,
479
480
        accelerator,
        autocast_kwargs,
Yoach Lacombe's avatar
Yoach Lacombe committed
481
482
    ):
        model.train()
Yoach Lacombe's avatar
Yoach Lacombe committed
483

484
        if mixed_precision == "fp16":
485
486
            # fp16 doesn't work with T5-like models
            with accelerator.autocast(autocast_handler=autocast_kwargs):
487
                if training_args.parallel_mode.value != "distributed":
Yoach Lacombe's avatar
Yoach Lacombe committed
488
489
490
                    encoder_outputs = model.text_encoder(
                        input_ids=batch.get("input_ids"), attention_mask=batch.get("attention_mask", None)
                    )
491
                else:
Yoach Lacombe's avatar
Yoach Lacombe committed
492
493
494
                    encoder_outputs = model.module.text_encoder(
                        input_ids=batch.get("input_ids"), attention_mask=batch.get("attention_mask", None)
                    )
495
                batch["encoder_outputs"] = encoder_outputs
Yoach Lacombe's avatar
Yoach Lacombe committed
496

Yoach Lacombe's avatar
Yoach Lacombe committed
497
498
499
500
501
502
        outputs = model(**batch)
        # CE (data) loss
        ce_loss = outputs.loss

        metrics = {"loss": ce_loss}
        return ce_loss, metrics
Yoach Lacombe's avatar
Yoach Lacombe committed
503

Yoach Lacombe's avatar
Yoach Lacombe committed
504
    # Define eval fn
Yoach Lacombe's avatar
Yoach Lacombe committed
505
506
507
508
509
    def eval_step(
        batch,
        accelerator,
        autocast_kwargs,
    ):
Yoach Lacombe's avatar
Yoach Lacombe committed
510
511
512
        eval_model = model if not training_args.torch_compile else model._orig_mod
        eval_model.eval()

513
        if mixed_precision == "fp16":
514
515
            # fp16 doesn't work with T5-like models
            with accelerator.autocast(autocast_handler=autocast_kwargs):
Yoach Lacombe's avatar
Yoach Lacombe committed
516
517
                with torch.no_grad():
                    if training_args.parallel_mode.value != "distributed" or training_args.torch_compile:
Yoach Lacombe's avatar
Yoach Lacombe committed
518
519
520
                        encoder_outputs = eval_model.text_encoder(
                            input_ids=batch.get("input_ids"), attention_mask=batch.get("attention_mask", None)
                        )
Yoach Lacombe's avatar
Yoach Lacombe committed
521
                    else:
Yoach Lacombe's avatar
Yoach Lacombe committed
522
523
524
                        encoder_outputs = eval_model.module.text_encoder(
                            input_ids=batch.get("input_ids"), attention_mask=batch.get("attention_mask", None)
                        )
525
                batch["encoder_outputs"] = encoder_outputs
Yoach Lacombe's avatar
Yoach Lacombe committed
526
527

        with torch.no_grad():
Yoach Lacombe's avatar
Yoach Lacombe committed
528
            outputs = eval_model(**batch)
Yoach Lacombe's avatar
Yoach Lacombe committed
529
530
531
532
533
534
        # CE (data) loss
        ce_loss = outputs.loss
        metrics = {"loss": ce_loss}
        return metrics

    def generate_step(batch):
535
        batch.pop("decoder_attention_mask", None)
Yoach Lacombe's avatar
Yoach Lacombe committed
536
        eval_model = accelerator.unwrap_model(model, keep_fp32_wrapper=mixed_precision != "fp16").eval()
Yoach Lacombe's avatar
Yoach Lacombe committed
537
538
539
540
        if training_args.torch_compile:
            eval_model = model._orig_mod

        output_audios = eval_model.generate(**batch, **gen_kwargs)
Yoach Lacombe's avatar
Yoach Lacombe committed
541
542
543
544
545
        output_audios = accelerator.pad_across_processes(output_audios, dim=1, pad_index=0)
        return output_audios

    for epoch in range(epochs_trained, num_epochs):
        vectorized_datasets["train"] = vectorized_datasets["train"].shuffle(training_args.seed)
Yoach Lacombe's avatar
Yoach Lacombe committed
546
547
548
        sampler = None
        if training_args.group_by_length:
            sampler = LengthGroupedSampler(train_batch_size, lengths=vectorized_datasets["train"]["target_length"])
Yoach Lacombe's avatar
Yoach Lacombe committed
549
550
551
552
        train_dataloader = DataLoader(
            vectorized_datasets["train"],
            collate_fn=data_collator,
            batch_size=per_device_train_batch_size,
Yoach Lacombe's avatar
Yoach Lacombe committed
553
            sampler=sampler,
Yoach Lacombe's avatar
Yoach Lacombe committed
554
555
556
557
558
559
560
561
562
563
564
565
566
            num_workers=training_args.dataloader_num_workers,
            pin_memory=training_args.dataloader_pin_memory,
        )
        train_dataloader = accelerator.prepare(train_dataloader)
        if hasattr(train_dataloader, "dataset") and isinstance(train_dataloader.dataset, IterableDataset):
            train_dataloader.dataset.set_epoch(epoch)

        if resume_step is not None:
            # Skip the first N batches in the dataloader when resuming from a checkpoint
            train_dataloader = accelerator.skip_first_batches(train_dataloader, resume_step)
            resume_step = None

        for batch in train_dataloader:
567
            breakpoint()
Yoach Lacombe's avatar
Yoach Lacombe committed
568
            with accelerator.accumulate(model):
569
                loss, train_metric = train_step(batch, accelerator, autocast_kwargs)
Yoach Lacombe's avatar
Yoach Lacombe committed
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
                accelerator.backward(loss)
                if accelerator.sync_gradients:
                    accelerator.clip_grad_norm_(model.parameters(), training_args.max_grad_norm)
                optimizer.step()
                lr_scheduler.step()
                optimizer.zero_grad()

            # Check if the accelerator has performed an optimization step behind the scenes
            if accelerator.sync_gradients:
                steps_trained_progress_bar.update(1)
                cur_step += 1

                if cur_step % training_args.logging_steps == 0:
                    steps_trained_progress_bar.write(
                        f"Step... ({cur_step} / {total_train_steps} | Loss:"
                        f" {train_metric['loss']}, Learning Rate:"
                        f" {lr_scheduler.get_last_lr()[0]})"
                    )
                    log_metric(
                        accelerator,
                        metrics=train_metric,
                        learning_rate=lr_scheduler.get_last_lr()[0],
                        train_time=train_time + time.time() - train_start,
                        step=cur_step,
                        epoch=epoch,
                        prefix="train",
                    )

                # save checkpoint and weights after each save_steps and at the end of training
                if (cur_step % training_args.save_steps == 0) or cur_step == total_train_steps:
                    intermediate_dir = os.path.join(training_args.output_dir, f"checkpoint-{cur_step}-epoch-{epoch}")
Yoach Lacombe's avatar
Yoach Lacombe committed
601
                    # safe_serialization=False to avoid shared tensors saving issue (TODO(YL): it's a temporary fix)
602
603
                    # https://github.com/huggingface/transformers/issues/27293#issuecomment-1872560074
                    accelerator.save_state(output_dir=intermediate_dir, safe_serialization=False)
Yoach Lacombe's avatar
Yoach Lacombe committed
604
605
                    accelerator.wait_for_everyone()
                    if accelerator.is_main_process:
Dan Lyth's avatar
Dan Lyth committed
606
                        rotate_checkpoints(training_args.save_total_limit, output_dir=training_args.output_dir, logger=logger)
Yoach Lacombe's avatar
Yoach Lacombe committed
607
608
609

                        if cur_step == total_train_steps:
                            # un-wrap student model for save
Yoach Lacombe's avatar
Yoach Lacombe committed
610
611
                            unwrapped_model = accelerator.unwrap_model(model)
                            unwrapped_model.save_pretrained(training_args.output_dir)
Yoach Lacombe's avatar
Yoach Lacombe committed
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626

                        if training_args.push_to_hub:
                            repo.push_to_hub(
                                commit_message=f"Saving train state of step {cur_step}",
                                blocking=False,
                            )

                if training_args.do_eval and (cur_step % eval_steps == 0 or cur_step == total_train_steps):
                    train_time += time.time() - train_start
                    # ======================== Evaluating ==============================
                    eval_metrics = []
                    eval_preds = []
                    eval_descriptions = []
                    eval_prompts = []
                    eval_start = time.time()
Yoach Lacombe's avatar
Yoach Lacombe committed
627

Yoach Lacombe's avatar
Yoach Lacombe committed
628
629
                    # release training input batch
                    batch = release_memory(batch)
Yoach Lacombe's avatar
Yoach Lacombe committed
630
631
632
633
634

                    validation_dataloader = DataLoader(
                        vectorized_datasets["eval"],
                        collate_fn=data_collator,
                        batch_size=per_device_eval_batch_size,
635
                        drop_last=False,
Yoach Lacombe's avatar
Yoach Lacombe committed
636
637
638
639
640
641
642
                        num_workers=training_args.dataloader_pin_memory,
                        pin_memory=training_args.dataloader_pin_memory,
                    )
                    validation_dataloader = accelerator.prepare(validation_dataloader)

                    for batch in tqdm(
                        validation_dataloader,
643
                        desc=f"Evaluating - Inference ...",
Yoach Lacombe's avatar
Yoach Lacombe committed
644
645
646
647
                        position=2,
                        disable=not accelerator.is_local_main_process,
                    ):
                        # Model forward
648
                        eval_metric = eval_step(batch, accelerator, autocast_kwargs)
Yoach Lacombe's avatar
Yoach Lacombe committed
649
650
651
                        eval_metric = accelerator.gather_for_metrics(eval_metric)
                        eval_metrics.append(eval_metric)

652
653
654
655
656
657
658
659
660
661
                    if training_args.predict_with_generate:
                        validation_dataloader = DataLoader(
                            vectorized_datasets["eval"],
                            collate_fn=data_collator,
                            batch_size=per_device_eval_batch_size,
                            drop_last=False,
                            num_workers=training_args.dataloader_pin_memory,
                            pin_memory=training_args.dataloader_pin_memory,
                        )
                        validation_dataloader = accelerator.prepare(validation_dataloader)
Yoach Lacombe's avatar
Yoach Lacombe committed
662
                        # generation
663
                        for batch in tqdm(
Yoach Lacombe's avatar
Yoach Lacombe committed
664
665
666
667
668
                            validation_dataloader,
                            desc=f"Evaluating - Generation ...",
                            position=2,
                            disable=not accelerator.is_local_main_process,
                        ):
Yoach Lacombe's avatar
Yoach Lacombe committed
669
670
                            generated_audios = generate_step(batch)
                            # Gather all predictions and targets
Yoach Lacombe's avatar
Yoach Lacombe committed
671
672
673
674
675
676
                            generated_audios, input_ids, prompts = accelerator.pad_across_processes(
                                (generated_audios, batch["input_ids"], batch["prompt_input_ids"]), dim=1, pad_index=0
                            )
                            generated_audios, input_ids, prompts = accelerator.gather_for_metrics(
                                (generated_audios, input_ids, prompts)
                            )
677
678
679
                            eval_preds.extend(generated_audios.to("cpu"))
                            eval_descriptions.extend(input_ids.to("cpu"))
                            eval_prompts.extend(prompts.to("cpu"))
Yoach Lacombe's avatar
Yoach Lacombe committed
680
681
682
683

                    eval_time = time.time() - eval_start
                    # normalize eval metrics
                    eval_metrics = {
Yoach Lacombe's avatar
Yoach Lacombe committed
684
685
                        key: torch.mean(torch.cat([d[key].unsqueeze(0) for d in eval_metrics]))
                        for key in eval_metrics[0]
Yoach Lacombe's avatar
Yoach Lacombe committed
686
687
688
689
690
691
692
693
694
695
                    }

                    # compute metrics
                    metrics_desc = ""
                    if training_args.predict_with_generate:
                        metric_values, pred_descriptions, pred_prompts, audios, transcriptions = compute_metrics(
                            eval_preds, eval_descriptions, eval_prompts, accelerator.device
                        )
                        eval_metrics.update(metric_values)
                        metrics_desc = " ".join([f"Eval {key}: {value} |" for key, value in metric_values.items()])
696
697
698
699
700
701
702
703
704
705
706
                        if "wandb" in training_args.report_to:
                            log_pred(
                                accelerator,
                                pred_descriptions,
                                pred_prompts,
                                transcriptions,
                                audios,
                                sampling_rate=sampling_rate,
                                step=cur_step,
                                prefix="eval",
                            )
Yoach Lacombe's avatar
Yoach Lacombe committed
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721

                    # Print metrics and update progress bar
                    steps_trained_progress_bar.write(
                        f"Eval results for step ({cur_step} / {total_train_steps} | Eval Loss: {eval_metrics['loss']} |"
                        f" {metrics_desc})"
                    )

                    log_metric(
                        accelerator,
                        metrics=eval_metrics,
                        train_time=eval_time,
                        step=cur_step,
                        epoch=epoch,
                        prefix="eval",
                    )
Yoach Lacombe's avatar
Yoach Lacombe committed
722

723
724
725
726
727
728
729
                    # release eval batch and relax metrics
                    eval_metrics = []
                    eval_preds = []
                    eval_descriptions = []
                    eval_prompts = []
                    batch = release_memory(batch)

Yoach Lacombe's avatar
Yoach Lacombe committed
730
731
732
733
734
735
736
737
738
739
740
741
                    # flush the train metrics
                    train_start = time.time()

                # break condition
                if cur_step == total_train_steps:
                    continue_training = False
                    break

        if not continue_training:
            break

    accelerator.end_training()
742
743
744


if __name__ == "__main__":
745
    set_start_method("spawn")
Yoach Lacombe's avatar
Yoach Lacombe committed
746
    main()