Saving train state of step 80000

checkpoint-80000-epoch-5/optimizer.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9d85a6ebcfaea697deb81a9f43cf52b763722cc63c52559be99475194f3b1740
+size 3652769047

checkpoint-80000-epoch-5/pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:64ed77445bae74c0a44de014644335672082db652c167e03307f5aa69f497029
+size 2588465818

checkpoint-80000-epoch-5/random_states_0.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9bee040cadf66c07cc28951d8b63f5309317c467c935239899967dedde236356
+size 16100

checkpoint-80000-epoch-5/random_states_1.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ac2984f6c11e0d6baba6e153c36e02cf7d8849f435a8e1d10db3db9422a5e14f
+size 16100

checkpoint-80000-epoch-5/random_states_2.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0328ba0db514b5427f0cf5728abf82349345bec8a02097c30900ec57f433da11
+size 16100

checkpoint-80000-epoch-5/random_states_3.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:668cc4578435da562a838c3860a81ef62a34fdfed1920ccebb7657eb586c14cb
+size 16100

checkpoint-80000-epoch-5/random_states_4.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c856097a2b0832cc85eba37f9cd1be76e0d395fe4b2c99ef3153199386dae5e5
+size 16100

checkpoint-80000-epoch-5/random_states_5.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:01d11b3bb1f3959f58be8c638fb21de1e89e54c13aa716a178687cb1fc51ac99
+size 16100

checkpoint-80000-epoch-5/random_states_6.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:18447190b31cbfd6ac8565a46dfbdc82123f510bef0eb0a72313d1150001b8f4
+size 16100

checkpoint-80000-epoch-5/random_states_7.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:025b2bf49a9a6fe2a6fd8cb7a98ba98f4eb439e04b5aa68fd3285c8c109bdc4f
+size 16100

checkpoint-80000-epoch-5/scheduler.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fd55300ad4b9da07287969193b88215d1af9b5d03f7bf3a833f1b260aa6434c4
+size 1000

starting_point_0.01_rope.json

@@ -10,6 +10,7 @@
     "prompt_tokenizer_name":"google/flan-t5-base",
 
     "report_to": ["wandb"],
+    "wandb_run_name": "parler-tts-600M-cross-attention-rope",
     "overwrite_output_dir": false,
     "output_dir": "./",
 

training/arguments.py

@@ -218,7 +218,7 @@ class DataTrainingArguments:
         metadata={
             "help": (
                 "If set, filter samples with descriptions that are longer than `max_description_token_length` tokens."
-                "Also, used to set maximum desription token length if `pad_to_max_length=True`."
+                "Also, used to set maximum description token length if `pad_to_max_length=True`."
             )
         },
     )
@@ -277,6 +277,12 @@ class DataTrainingArguments:
         default="parler-speech",
         metadata={"help": "The name of the wandb project."},
     )
+    wandb_run_name: str = field(
+        default=None,
+        metadata={
+            "help": "If specified, the name of the run. If not specified, wandb will give a random name to this run."
+        },
+    )
     save_to_disk: str = field(
         default=None,
         metadata={

training/data.py

@@ -31,7 +31,12 @@ class DataCollatorEncodecWithPadding:
         audios = [feature[self.audio_column_name]["array"] for feature in features]
         len_audio = [len(audio) for audio in audios]
 
-        batch = self.feature_extractor(audios, return_tensors="pt", padding=self.padding, max_length=self.max_length)
+        # since resampling has already been performed in the 'load_multiple_datasets' function,
+        # a fixed sampling_rate(44100hz) is passed to the feature_extractor.
+        sampling_rate = self.feature_extractor.sampling_rate
+        batch = self.feature_extractor(
+            audios, sampling_rate=sampling_rate, return_tensors="pt", padding=self.padding, max_length=self.max_length
+        )
         batch["len_audio"] = torch.tensor(len_audio).unsqueeze(1)
         return batch
 

training/eval.py

@@ -47,8 +47,7 @@ def wer(asr_model_name_or_path, prompts, audios, device, per_device_eval_batch_s
     normalized_references = []
 
     for pred, ref in zip(transcriptions, prompts):
-        normalizer = english_normalizer
-
+        normalizer = english_normalizer if return_language and pred["chunks"][0]["language"] == "english" else basic_normalizer
         norm_ref = normalizer(ref)
         if len(norm_ref) > 0:
             norm_pred = normalizer(pred["text"])

training/run_parler_tts_training.py

@@ -98,9 +98,6 @@ def main():
 
     ####### A. Preparation
     kwargs_handlers = [InitProcessGroupKwargs(timeout=timedelta(minutes=60))]
-    if training_args.torch_compile:
-        # TODO(YL): add more compile modes?
-        kwargs_handlers.append(TorchDynamoPlugin(backend="inductor", mode="default"))  # reduce-overhead
 
     accelerator = Accelerator(
         gradient_accumulation_steps=training_args.gradient_accumulation_steps,
@@ -129,6 +126,7 @@ def main():
             "adam_beta2": training_args.adam_beta2,
             "temperature": model_args.temperature,
         },
+        init_kwargs={"wandb": {"name": data_args.wandb_run_name}} if data_args.wandb_run_name else {},
     )
 
     # Detecting last checkpoint and eventually continue from last checkpoint
@@ -314,6 +312,7 @@ def main():
         token=data_args.token,
         trust_remote_code=data_args.trust_remote_code,
     )
+    generation_config = model.generation_config
 
     # enable gradient checkpointing if necessary
     if training_args.gradient_checkpointing:
@@ -334,8 +333,8 @@ def main():
     feature_extractor_input_name = feature_extractor.model_input_names[0]
     audio_encoder_pad_token_id = config.decoder.pad_token_id
     audio_encoder_eos_token_id = config.decoder.eos_token_id
-    audio_encoder_bos_token_id = model.generation_config.decoder_start_token_id
-    max_length = model.generation_config.max_length
+    audio_encoder_bos_token_id = generation_config.decoder_start_token_id
+    max_length = generation_config.max_length
     num_codebooks = model.decoder.config.num_codebooks
     bandwidth = model_args.bandwidth
 
@@ -538,7 +537,7 @@ def main():
         logger.info(f"Dataset saved at {data_args.save_to_disk}")
 
     audio_max_length = None
-    if training_args.torch_compile:
+    if padding == "max_length":
         audio_max_length = max(vectorized_datasets["train"]["target_length"])
         with accelerator.main_process_first():
             max_sample = vectorized_datasets["train"].filter(
@@ -548,6 +547,18 @@ def main():
             )
         audio_max_length = torch.tensor(max_sample[0]["labels"]).shape[1]
 
+    if training_args.group_by_length:
+        # apply a simple heuristic to take into account audio and text lengths
+        def add_target_lengths(target_length, prompt, description):
+            return {"target_length": target_length + len(prompt) + len(description)}
+
+        with accelerator.main_process_first():
+            vectorized_datasets = vectorized_datasets.map(
+                add_target_lengths,
+                num_proc=num_workers,
+                input_columns=["target_length", "prompt_input_ids", "input_ids"],
+            )
+
     # for large datasets it is advised to run the preprocessing on a
     # single machine first with ``args.preprocessing_only`` since there will mostly likely
     # be a timeout when running the script in distributed mode.
@@ -670,6 +681,8 @@ def main():
         checkpoint = last_checkpoint
 
     if accelerator.is_main_process:
+        if training_args.output_dir is not None:
+            os.makedirs(training_args.output_dir, exist_ok=True)
         if training_args.push_to_hub:
             api = HfApi(token=training_args.hub_token)
 
@@ -682,8 +695,6 @@ def main():
             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()
 
     # Now save everything to be able to create a single processor later
@@ -740,7 +751,13 @@ def main():
         "do_sample": model_args.do_sample,
         "temperature": model_args.temperature,
         "max_length": model_args.max_length,
+        # Because of the delayed pattern mask, generation might stop earlier because of unexpected behaviour
+        # on the first tokens of the codebooks that are delayed.
+        # This fix the issue.
+        "min_new_tokens": num_codebooks + 1,
     }
+    for key in gen_kwargs:
+        generation_config.key = gen_kwargs[key]
 
     # Define gradient update step fn
     def train_step(
@@ -869,9 +886,11 @@ def main():
                     # safe_serialization=False to avoid shared tensors saving issue (TODO(YL): it's a temporary fix)
                     # https://github.com/huggingface/transformers/issues/27293#issuecomment-1872560074
                     accelerator.save_state(output_dir=intermediate_dir, safe_serialization=False)
+                    config.save_pretrained(intermediate_dir)
+                    generation_config.save_pretrained(intermediate_dir)
                     accelerator.wait_for_everyone()
                     if accelerator.is_main_process:
-                        rotate_checkpoints(
+                        checkpoints_to_be_deleted = rotate_checkpoints(
                             training_args.save_total_limit, output_dir=training_args.output_dir, logger=logger
                         )
 
@@ -886,6 +905,7 @@ def main():
                                 folder_path=training_args.output_dir,
                                 commit_message=f"Saving train state of step {cur_step}",
                                 run_as_future=True,
+                                delete_patterns=checkpoints_to_be_deleted,
                             )
 
                 if training_args.do_eval and (cur_step % eval_steps == 0 or cur_step == total_train_steps):

training/utils.py

@@ -3,7 +3,7 @@ import re
 import shutil
 from pathlib import Path
 from dataclasses import field
-from typing import Dict, List
+from typing import Dict, List, Union
 
 import torch
 from wandb import Audio
@@ -44,7 +44,7 @@ def sorted_checkpoints(output_dir=None, checkpoint_prefix="checkpoint") -> List[
     return checkpoints_sorted
 
 
-def rotate_checkpoints(save_total_limit=None, output_dir=None, checkpoint_prefix="checkpoint", logger=None) -> None:
+def rotate_checkpoints(save_total_limit=None, output_dir=None, checkpoint_prefix="checkpoint", logger=None) -> Union[List, None]:
     """Helper function to delete old checkpoints."""
     if save_total_limit is None or save_total_limit <= 0:
         return
@@ -58,6 +58,8 @@ def rotate_checkpoints(save_total_limit=None, output_dir=None, checkpoint_prefix
     for checkpoint in checkpoints_to_be_deleted:
         logger.info(f"Deleting older checkpoint [{checkpoint}] due to args.save_total_limit")
         shutil.rmtree(checkpoint, ignore_errors=True)
+    checkpoints_to_be_deleted = [f"*{Path(checkpoint).absolute().name}*"  for checkpoint in checkpoints_to_be_deleted]
+    return checkpoints_to_be_deleted
 
 
 def log_metric(