train.py 31.3 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
from tqdm import tqdm
Yoach Lacombe's avatar
Yoach Lacombe committed
28
from pathlib import Path
29
30
31

import datasets
import torch
32
33
from torch.utils.data import DataLoader

Dan Lyth's avatar
Dan Lyth committed
34
from datasets import IterableDataset
35

Yoach Lacombe's avatar
Yoach Lacombe committed
36
from huggingface_hub import Repository, create_repo
37
38
39
40
import transformers
from transformers import (
    AutoFeatureExtractor,
    AutoTokenizer,
Dan Lyth's avatar
Dan Lyth committed
41
    HfArgumentParser,
42
)
Yoach Lacombe's avatar
Yoach Lacombe committed
43
from transformers.trainer_pt_utils import LengthGroupedSampler
Yoach Lacombe's avatar
Yoach Lacombe committed
44
from transformers.optimization import get_scheduler
Yoach Lacombe's avatar
Yoach Lacombe committed
45
from transformers.utils import send_example_telemetry
Yoach Lacombe's avatar
add DAC  
Yoach Lacombe committed
46

47
48

from accelerate import Accelerator
49
from accelerate.utils import set_seed, AutocastKwargs, InitProcessGroupKwargs, TorchDynamoPlugin
Yoach Lacombe's avatar
Yoach Lacombe committed
50
from accelerate.utils.memory import release_memory
51

Yoach Lacombe's avatar
Yoach Lacombe committed
52
53
from parler_tts import (
    ParlerTTSForConditionalGeneration,
Dan Lyth's avatar
Dan Lyth committed
54
    ParlerTTSConfig
Yoach Lacombe's avatar
Yoach Lacombe committed
55
)
56

Dan Lyth's avatar
Dan Lyth committed
57
58
from parler_tts.utils import get_last_checkpoint, rotate_checkpoints, log_pred, log_metric
from parler_tts.arguments import ModelArguments, DataTrainingArguments, ParlerTTSTrainingArguments
Dan Lyth's avatar
Dan Lyth committed
59
60
from parler_tts.data import DataCollatorParlerTTSWithPadding
from parler_tts.eval import clap_similarity, wer
61
62
63
64
65
66
67
68
69
70


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
71
    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, ParlerTTSTrainingArguments))
72
73
74
75
76
77
78
79
80
    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
81
    send_example_telemetry("run_parler_tts", model_args, data_args)
Yoach Lacombe's avatar
Yoach Lacombe committed
82

Yoach Lacombe's avatar
Yoach Lacombe committed
83
84
85
86
87
88
    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
89
90
91
92
93
94
95
96
97

    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`"
        )
98
99

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

101
    # Accelerator preparation
102
103
104
    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
105
106
        kwargs_handlers.append(TorchDynamoPlugin(backend="inductor", mode="default"))  # reduce-overhead

Yoach Lacombe's avatar
Yoach Lacombe committed
107
108
109
110
111
    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,
112
        kwargs_handlers=kwargs_handlers,
Yoach Lacombe's avatar
Yoach Lacombe committed
113
    )
Yoach Lacombe's avatar
Yoach Lacombe committed
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135

    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
136
    # Detecting last checkpoint and eventually continue from last checkpoint
137
138
139
140
141
142
143
144
    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
145
        elif last_checkpoint is not None and training_args.resume_from_checkpoint is None:
146
147
148
149
150
151
152
153
154
155
156
            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)],
    )
157
    logger.setLevel(logging.INFO if accelerator.is_main_process else logging.WARN)
158

Yoach Lacombe's avatar
Yoach Lacombe committed
159
    # Log a small summary on each proces
160
161
162
163
    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
164
165
166
167

    # 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()
168
        transformers.utils.logging.set_verbosity_info()
Yoach Lacombe's avatar
Yoach Lacombe committed
169
170
171
172
    else:
        datasets.utils.logging.set_verbosity_error()
        transformers.utils.logging.set_verbosity_error()

173
174
175
176
    logger.info("Training/evaluation parameters %s", training_args)

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

179
    # 1. First, let's instantiate the feature extractor (DAC), tokenizers and model
180
181
    # 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
182

183
184
185
186
187
188
189
190
    # 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
191

192
193
194
195
196
197
198
    # 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
199
        padding_side="left",  # prompt has to be padded on the left bc it's preprend to codebooks hidden states
200
    )
Yoach Lacombe's avatar
Yoach Lacombe committed
201

202
203
204
205
206
207
208
209
    # 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
210

211
    if model_args.use_fast_tokenizer:
Yoach Lacombe's avatar
Yoach Lacombe committed
212
213
214
        logger.warning(
            "Disabling fast tokenizer warning: https://github.com/huggingface/transformers/blob/main/src/transformers/tokenization_utils_base.py#L3231-L3235"
        )
215
216
        prompt_tokenizer.deprecation_warnings["Asking-to-pad-a-fast-tokenizer"] = True
        description_tokenizer.deprecation_warnings["Asking-to-pad-a-fast-tokenizer"] = True
217

218
    # 2. Now, let's load the dataset
219
    # TODO add MDS dataset loading here
220

221
    # 3. Next, let's load the config.
Yoach Lacombe's avatar
Yoach Lacombe committed
222
    config = ParlerTTSConfig.from_pretrained(
223
224
225
226
227
        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
228

229
    # update pad token id and decoder_start_token_id
Yoach Lacombe's avatar
Yoach Lacombe committed
230
231
232
233
    config.update(
        {
            "pad_token_id": model_args.pad_token_id
            if model_args.pad_token_id is not None
234
            else config.pad_token_id,
Yoach Lacombe's avatar
Yoach Lacombe committed
235
236
            "decoder_start_token_id": model_args.decoder_start_token_id
            if model_args.decoder_start_token_id is not None
237
            else config.decoder_start_token_id,
Yoach Lacombe's avatar
Yoach Lacombe committed
238
239
240
        }
    )

Yoach Lacombe's avatar
Yoach Lacombe committed
241
    # create model
Yoach Lacombe's avatar
Yoach Lacombe committed
242
    model = ParlerTTSForConditionalGeneration.from_pretrained(
243
244
245
246
247
248
        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
249

250
251
252
    # enable gradient checkpointing if necessary
    if training_args.gradient_checkpointing:
        model.gradient_checkpointing_enable()
Yoach Lacombe's avatar
Yoach Lacombe committed
253

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

257
    # derive max & min input length for sample rate & max duration
258
    sampling_rate = feature_extractor.sampling_rate
Yoach Lacombe's avatar
Yoach Lacombe committed
259

260
    # Freeze Encoders
261
    model.freeze_encoders(model_args.freeze_text_encoder) # TODO check this implementation
Yoach Lacombe's avatar
Yoach Lacombe committed
262

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

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

Yoach Lacombe's avatar
Yoach Lacombe committed
267
    # Define evaluation metrics during training, *i.e.* CLAP similarity
Yoach Lacombe's avatar
Yoach Lacombe committed
268
    def compute_metrics(audios, descriptions, prompts, device="cpu"):
Dan Lyth's avatar
Dan Lyth committed
269
        results = {}
Yoach Lacombe's avatar
Yoach Lacombe committed
270
        input_ids = descriptions
271
        texts = description_tokenizer.batch_decode(input_ids, skip_special_tokens=True)
Yoach Lacombe's avatar
Yoach Lacombe committed
272
273
        prompts = prompt_tokenizer.batch_decode(prompts, skip_special_tokens=True)
        audios = [a.cpu().numpy() for a in audios]
Dan Lyth's avatar
Dan Lyth committed
274
275
276
277
278
279
280
281
282
283
        
        clap_score = clap_similarity(model_args.clap_model_name_or_path, texts, audios, device)
        results["clap"] = clap_score

        word_error, transcriptions = wer(model_args.asr_model_name_or_path,
                                        prompts,
                                        audios,
                                        device,
                                        training_args.per_device_eval_batch_size,
                                        sampling_rate)
Yoach Lacombe's avatar
Yoach Lacombe committed
284
        results["wer"] = word_error
285

Yoach Lacombe's avatar
Yoach Lacombe committed
286
        return results, texts, prompts, audios, transcriptions
Yoach Lacombe's avatar
Yoach Lacombe committed
287

Yoach Lacombe's avatar
Yoach Lacombe committed
288
289
290
291
292
293
    # 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
294

Yoach Lacombe's avatar
Yoach Lacombe committed
295
296
    if training_args.max_steps < 0:
        num_epochs = int(training_args.num_train_epochs)
297
        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
298
299
300
301
302
303
304
305
306
        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
307
        logger.info(f"eval_steps is not set, evaluating at the end of each epoch")
Yoach Lacombe's avatar
Yoach Lacombe committed
308
309
310
        eval_steps = steps_per_epoch
    else:
        eval_steps = training_args.eval_steps
Yoach Lacombe's avatar
Yoach Lacombe committed
311

312
    # T5 doesn't support fp16
Yoach Lacombe's avatar
Yoach Lacombe committed
313
314
    autocast_kwargs = AutocastKwargs(enabled=(mixed_precision != "fp16"))

Yoach Lacombe's avatar
Yoach Lacombe committed
315
316
317
318
319
320
    # 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,
321
        weight_decay=training_args.weight_decay,
Yoach Lacombe's avatar
Yoach Lacombe committed
322
    )
323

Yoach Lacombe's avatar
Yoach Lacombe committed
324
325
326
327
    # 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
328
        num_warmup_steps=training_args.get_warmup_steps(total_train_steps) * accelerator.num_processes,
Yoach Lacombe's avatar
Yoach Lacombe committed
329
330
        num_training_steps=total_train_steps * accelerator.num_processes,
    )
331
332

    # Instantiate custom data collator
Yoach Lacombe's avatar
Yoach Lacombe committed
333
    data_collator = DataCollatorParlerTTSWithPadding(
Yoach Lacombe's avatar
Yoach Lacombe committed
334
335
336
337
338
339
        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,
340
        audio_max_length=audio_max_length, # TODO add this variable
341
    )
Yoach Lacombe's avatar
Yoach Lacombe committed
342

Yoach Lacombe's avatar
Yoach Lacombe committed
343
344
    # Prepare everything with accelerate
    model, optimizer, lr_scheduler = accelerator.prepare(model, optimizer, lr_scheduler)
Yoach Lacombe's avatar
Yoach Lacombe committed
345

Yoach Lacombe's avatar
Yoach Lacombe committed
346
347
    logger.info("***** Running training *****")
    logger.info(f"  Num examples = {total_train_steps * train_batch_size * gradient_accumulation_steps}")
348
    logger.info("  Instantaneous batch size per device =" f" {per_device_train_batch_size}")
Yoach Lacombe's avatar
Yoach Lacombe committed
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
    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
370

Yoach Lacombe's avatar
Yoach Lacombe committed
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
    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
388

Yoach Lacombe's avatar
Yoach Lacombe committed
389
390
391
392
393
394
    # 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
395
396
397
398
399
            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
400
401
                prompt_tokenizer.save_pretrained(training_args.output_dir)
            else:
Yoach Lacombe's avatar
Yoach Lacombe committed
402
                logger.warning(
Dan Lyth's avatar
Dan Lyth committed
403
                    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
404
                )
Yoach Lacombe's avatar
Yoach Lacombe committed
405
                prompt_tokenizer.save_pretrained(training_args.output_dir)
Yoach Lacombe's avatar
Yoach Lacombe committed
406

Yoach Lacombe's avatar
Yoach Lacombe committed
407
408
            feature_extractor.save_pretrained(training_args.output_dir)
            config.save_pretrained(training_args.output_dir)
Yoach Lacombe's avatar
Yoach Lacombe committed
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425

    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
426

Yoach Lacombe's avatar
Yoach Lacombe committed
427
428
        if training_args.max_steps < 0:
            # we know exactly the number of steps per epoch, so can skip through the required number of batches
429
            resume_step = (cur_step - epochs_trained * steps_per_epoch) * gradient_accumulation_steps
Yoach Lacombe's avatar
Yoach Lacombe committed
430
431
432
433
434
435
436
437
        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
438

Yoach Lacombe's avatar
Yoach Lacombe committed
439
440
    gen_kwargs = {
        "do_sample": model_args.do_sample,
yoach@huggingface.co's avatar
yoach@huggingface.co committed
441
        "temperature": model_args.temperature,
Yoach Lacombe's avatar
Yoach Lacombe committed
442
443
        "max_length": model_args.max_length,
    }
Yoach Lacombe's avatar
Yoach Lacombe committed
444

Yoach Lacombe's avatar
Yoach Lacombe committed
445
446
447
    # Define gradient update step fn
    def train_step(
        batch,
448
449
        accelerator,
        autocast_kwargs,
Yoach Lacombe's avatar
Yoach Lacombe committed
450
451
    ):
        model.train()
Yoach Lacombe's avatar
Yoach Lacombe committed
452

453
        if mixed_precision == "fp16":
454
455
            # fp16 doesn't work with T5-like models
            with accelerator.autocast(autocast_handler=autocast_kwargs):
456
                if training_args.parallel_mode.value != "distributed":
Yoach Lacombe's avatar
Yoach Lacombe committed
457
458
459
                    encoder_outputs = model.text_encoder(
                        input_ids=batch.get("input_ids"), attention_mask=batch.get("attention_mask", None)
                    )
460
                else:
Yoach Lacombe's avatar
Yoach Lacombe committed
461
462
463
                    encoder_outputs = model.module.text_encoder(
                        input_ids=batch.get("input_ids"), attention_mask=batch.get("attention_mask", None)
                    )
464
                batch["encoder_outputs"] = encoder_outputs
Yoach Lacombe's avatar
Yoach Lacombe committed
465

Yoach Lacombe's avatar
Yoach Lacombe committed
466
467
468
469
470
471
        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
472

Yoach Lacombe's avatar
Yoach Lacombe committed
473
    # Define eval fn
Yoach Lacombe's avatar
Yoach Lacombe committed
474
475
476
477
478
    def eval_step(
        batch,
        accelerator,
        autocast_kwargs,
    ):
Yoach Lacombe's avatar
Yoach Lacombe committed
479
480
481
        eval_model = model if not training_args.torch_compile else model._orig_mod
        eval_model.eval()

482
        if mixed_precision == "fp16":
483
484
            # fp16 doesn't work with T5-like models
            with accelerator.autocast(autocast_handler=autocast_kwargs):
Yoach Lacombe's avatar
Yoach Lacombe committed
485
486
                with torch.no_grad():
                    if training_args.parallel_mode.value != "distributed" or training_args.torch_compile:
Yoach Lacombe's avatar
Yoach Lacombe committed
487
488
489
                        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
490
                    else:
Yoach Lacombe's avatar
Yoach Lacombe committed
491
492
493
                        encoder_outputs = eval_model.module.text_encoder(
                            input_ids=batch.get("input_ids"), attention_mask=batch.get("attention_mask", None)
                        )
494
                batch["encoder_outputs"] = encoder_outputs
Yoach Lacombe's avatar
Yoach Lacombe committed
495
496

        with torch.no_grad():
Yoach Lacombe's avatar
Yoach Lacombe committed
497
            outputs = eval_model(**batch)
Yoach Lacombe's avatar
Yoach Lacombe committed
498
499
500
501
502
503
        # CE (data) loss
        ce_loss = outputs.loss
        metrics = {"loss": ce_loss}
        return metrics

    def generate_step(batch):
504
        batch.pop("decoder_attention_mask", None)
Yoach Lacombe's avatar
Yoach Lacombe committed
505
        eval_model = accelerator.unwrap_model(model, keep_fp32_wrapper=mixed_precision != "fp16").eval()
Yoach Lacombe's avatar
Yoach Lacombe committed
506
507
508
509
        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
510
511
512
513
514
        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
515
516
517
        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
518
519
520
521
        train_dataloader = DataLoader(
            vectorized_datasets["train"],
            collate_fn=data_collator,
            batch_size=per_device_train_batch_size,
Yoach Lacombe's avatar
Yoach Lacombe committed
522
            sampler=sampler,
Yoach Lacombe's avatar
Yoach Lacombe committed
523
524
525
526
527
528
529
530
531
532
533
534
535
536
            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:
            with accelerator.accumulate(model):
537
                loss, train_metric = train_step(batch, accelerator, autocast_kwargs)
Yoach Lacombe's avatar
Yoach Lacombe committed
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
                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
569
                    # safe_serialization=False to avoid shared tensors saving issue (TODO(YL): it's a temporary fix)
570
571
                    # 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
572
573
                    accelerator.wait_for_everyone()
                    if accelerator.is_main_process:
Dan Lyth's avatar
Dan Lyth committed
574
                        rotate_checkpoints(training_args.save_total_limit, output_dir=training_args.output_dir, logger=logger)
Yoach Lacombe's avatar
Yoach Lacombe committed
575
576
577

                        if cur_step == total_train_steps:
                            # un-wrap student model for save
Yoach Lacombe's avatar
Yoach Lacombe committed
578
579
                            unwrapped_model = accelerator.unwrap_model(model)
                            unwrapped_model.save_pretrained(training_args.output_dir)
Yoach Lacombe's avatar
Yoach Lacombe committed
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594

                        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
595

Yoach Lacombe's avatar
Yoach Lacombe committed
596
597
                    # release training input batch
                    batch = release_memory(batch)
Yoach Lacombe's avatar
Yoach Lacombe committed
598
599
600
601
602

                    validation_dataloader = DataLoader(
                        vectorized_datasets["eval"],
                        collate_fn=data_collator,
                        batch_size=per_device_eval_batch_size,
603
                        drop_last=False,
Yoach Lacombe's avatar
Yoach Lacombe committed
604
605
606
607
608
609
610
                        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,
611
                        desc=f"Evaluating - Inference ...",
Yoach Lacombe's avatar
Yoach Lacombe committed
612
613
614
615
                        position=2,
                        disable=not accelerator.is_local_main_process,
                    ):
                        # Model forward
616
                        eval_metric = eval_step(batch, accelerator, autocast_kwargs)
Yoach Lacombe's avatar
Yoach Lacombe committed
617
618
619
                        eval_metric = accelerator.gather_for_metrics(eval_metric)
                        eval_metrics.append(eval_metric)

620
621
622
623
624
625
626
627
628
629
                    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
630
                        # generation
631
                        for batch in tqdm(
Yoach Lacombe's avatar
Yoach Lacombe committed
632
633
634
635
636
                            validation_dataloader,
                            desc=f"Evaluating - Generation ...",
                            position=2,
                            disable=not accelerator.is_local_main_process,
                        ):
Yoach Lacombe's avatar
Yoach Lacombe committed
637
638
                            generated_audios = generate_step(batch)
                            # Gather all predictions and targets
Yoach Lacombe's avatar
Yoach Lacombe committed
639
640
641
642
643
644
                            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)
                            )
645
646
647
                            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
648
649
650
651

                    eval_time = time.time() - eval_start
                    # normalize eval metrics
                    eval_metrics = {
Yoach Lacombe's avatar
Yoach Lacombe committed
652
653
                        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
654
655
656
657
658
659
660
661
662
663
                    }

                    # 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()])
664
665
666
667
668
669
670
671
672
673
674
                        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
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689

                    # 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
690

691
692
693
694
695
696
697
                    # 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
698
699
700
701
702
703
704
705
706
707
708
709
                    # 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()
710
711
712


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