     1→{
     2→    "training": {
     3→        "seed": 42,
     4→        "logging_steps": 50,
     5→        "eval_steps": 0,
     6→        "gradient_accumulation_steps": 1,
     7→        "gradient_clip_value": 1.0,
     8→        "learning_rate": 1.5e-4,
     9→        "betas": [0.9, 0.95],
    10→        "warmup_ratio": 0.10,
    11→        "batch_size": 12,
    12→        "weight_decay": 0.1,
    13→        "precision": "bf16",
    14→        "strategy": "ddp",
    15→        "gradient_checkpointing": false,
    16→        "num_workers": 4,
    17→        "max_training_steps": 1271565,
    18→        "min_lr_ratio": 0.1,
    19→        "max_tokens_per_batch": 24576
    20→    },
    21→    "modeling": {
    22→        "parameters": {
    23→            "codebook_size": 65536,
    24→            "max_seq_len": 2048,
    25→            "model_name": "meta-llama/Llama-3.2-1B-Instruct",
    26→            "enable_text_normalization": false
    27→        }
    28→    },
    29→    "checkpointing": {
    30→        "save_steps": 10000,
    31→        "collect_health_stats": false,
    32→        "save_intermediate_generations": false,
    33→        "only_load_model_weights": false,
    34→        "keep_only_last_n_checkpoints": 5
    35→    },
    36→    "train_weighted_datasets": {
    37→        "/root/maya-tts/pretrain-encoded/pt_output": 1.0,
    38→        "/root/maya-tts/text_corpus/compiled_25B/pt_output[text]": 1.0,
    39→        "/root/maya-tts/finalsftdata/sft4b[paired]": 3.16
    40→    },
    41→    "val_weighted_datasets": {
    42→        "/root/maya-tts/pretrain-encoded/pt_output": 1.0,
    43→        "/root/maya-tts/finalsftdata/sft4b[paired]": 3.16
    44→    },
    45→    "dataset": {
    46→        "allowed_languages": [],
    47→        "min_dnsmos_score": 0,
    48→        "min_sample_rate": 0,
    49→        "enable_rlhf_training": false
    50→    }
    51→}
    52→


Now let me check the data preprocessing and pretraining datasets:


     1→import json
     2→import os
     3→from collections.abc import Sequence
     4→from typing import Any
     5→
     6→import numpy as np
     7→import torch
     8→import transformers
     9→from absl import logging
    10→
    11→from tts.core import constants
    12→from tts.core import prompting
    13→from tts.data import data_utils, pretraining_contract, text_normalization
    14→from tts.data.datasets import finetuning
    15→from tts.utils import configuration
    16→
    17→# Memmap is recreated after `SAMPLES_TO_RECREATE` getitem invokes to free RAM.
    18→_SAMPLES_TO_RECREATE = 10**5
    19→
    20→
    21→def _load_pretraining_manifest_entry(
    22→    dataset_dir: str, split: str, mode: str, artifact_name: str
    23→) -> dict[str, Any]:
    24→    """Returns the matching manifest entry for a pretraining artifact."""
    25→    manifest_path = os.path.join(dataset_dir, "manifest.json")
    26→    if not os.path.exists(manifest_path):
    27→        raise ValueError(
    28→            f"Missing pretraining manifest at [{manifest_path}]. "
    29→            "Pretraining artifacts must include manifest.json."
    30→        )
    31→
    32→    with open(manifest_path, encoding="utf-8") as handle:
    33→        try:
    34→            manifest = json.load(handle)
    35→        except json.JSONDecodeError as exc:
    36→            raise ValueError(
    37→                f"Malformed pretraining manifest at [{manifest_path}]."
    38→            ) from exc
    39→
    40→    entries = manifest if isinstance(manifest, list) else [manifest]
    41→    matching_entries = []
    42→    for entry in entries:
    43→        if not isinstance(entry, dict):
    44→            continue
    45→        if entry.get("mode") != mode or entry.get("split") != split:
    46→            continue
    47→
    48→        output_path = entry.get("output_path")
    49→        if output_path and os.path.basename(output_path) != artifact_name:
    50→            continue
    51→        matching_entries.append(entry)
    52→
    53→    if not matching_entries:
    54→        raise ValueError(
    55→            f"Found manifest at [{manifest_path}] but no entry matched "
    56→            f"mode=[{mode}], split=[{split}], artifact=[{artifact_name}]."
    57→        )
    58→    return matching_entries[-1]
    59→
    60→
    61→def _get_audio_code_id_range(
    62→    codes_path: str, manifest_entry: dict[str, Any]
    63→) -> tuple[int, int]:
    64→    """Returns the codec id range, preferring manifest metadata when available."""
    65→    min_code_id = manifest_entry.get("codec_min_id")
    66→    max_code_id = manifest_entry.get("codec_max_id")
    67→    if isinstance(min_code_id, int) and isinstance(max_code_id, int):
    68→        return min_code_id, max_code_id
    69→
    70→    codes = np.memmap(codes_path, dtype=np.int32, mode="r")
    71→    return int(np.min(codes)), int(np.max(codes))
    72→
    73→
    74→def _validate_manifest_file_path(
    75→    manifest_entry: dict[str, Any],
    76→    field_name: str,
    77→    expected_basename: str,
    78→    dataset_dir: str,
    79→) -> None:
    80→    """Ensures a manifest field points at the expected artifact file."""
    81→    file_path = manifest_entry.get(field_name)
    82→    if file_path is None:
    83→        raise ValueError(
    84→            f"Pretraining manifest is missing required field `{field_name}`."
    85→        )
    86→    if os.path.basename(file_path) != expected_basename:
    87→        raise ValueError(
    88→            f"Pretraining manifest field `{field_name}` must reference "
    89→            f"[{expected_basename}], but got [{file_path}]."
    90→        )
    91→    expected_path = os.path.join(dataset_dir, expected_basename)
    92→    if not os.path.exists(expected_path):
    93→        raise FileNotFoundError(
    94→            f"Paired pretraining artifact file referenced by `{field_name}` does not "
    95→            f"exist: [{expected_path}]."
    96→        )
    97→
    98→
    99→def _validate_audio_filter_config(
   100→    manifest_entry: dict[str, Any],
   101→    dataset_dir: str,
   102→    split: str,
   103→    dataset_config: configuration.DatasetConfig | None,
   104→) -> None:
   105→    """Ensures the artifact was built with the active audio PT filters."""
   106→    expected_filter_config = pretraining_contract.audio_filter_config_from_dataset_config(
   107→        dataset_config
   108→    )
   109→    if expected_filter_config is None:
   110→        return
   111→
   112→    artifact_filter_config = manifest_entry.get("filter_config")
   113→    if artifact_filter_config is None:
   114→        if expected_filter_config == pretraining_contract.normalize_audio_filter_config({}):
   115→            return
   116→        raise ValueError(
   117→            "Audio pretraining manifest is missing `filter_config`. Rebuild the PT "
   118→            f"artifact at [{dataset_dir}] for split [{split}] with "
   119→            "`tools/data/build_pretraining_data.py` so dataset curation is tracked."
   120→        )
   121→
   122→    normalized_artifact_filter_config = (
   123→        pretraining_contract.normalize_audio_filter_config(artifact_filter_config)
   124→    )
   125→    if normalized_artifact_filter_config != expected_filter_config:
   126→        raise ValueError(
   127→            "Audio pretraining artifact filter mismatch: "
   128→            f"artifact={normalized_artifact_filter_config} "
   129→            f"training={expected_filter_config}."
   130→        )
   131→
   132→
   133→def _validate_audio_pt_format(manifest_entry: dict[str, Any]) -> None:
   134→    """Validates optional audio PT stream-format metadata."""
   135→    audio_format = manifest_entry.get("format")
   136→    if audio_format is None or audio_format == pretraining_contract.AUDIO_PT_FORMAT_RAW_CODEC_STREAM_V1:
   137→        return
   138→
   139→    if audio_format != pretraining_contract.AUDIO_PT_FORMAT_SPEECH_BOUNDED_V2:
   140→        raise ValueError(f"Unsupported audio pretraining artifact format: {audio_format}.")
   141→
   142→    if manifest_entry.get("boundary_sentinels") != pretraining_contract.audio_pt_boundary_sentinels():
   143→        raise ValueError(
   144→            "Audio pretraining artifact boundary sentinel mismatch. Rebuild the "
   145→            "artifact so builder and loader use the same sentinels."
   146→        )
   147→
   148→
   149→def validate_audio_pretraining_artifact(
   150→    dataset_dir: str,
   151→    split: str,
   152→    codebook_size: int,
   153→    validate_code_range: bool,
   154→    dataset_config: configuration.DatasetConfig | None = None,
   155→) -> None:
   156→    """Validates the audio PT artifact metadata and optional codec id range."""
   157→    artifact_name = f"{split}_pretraining_codes.npy"
   158→    manifest_entry = _load_pretraining_manifest_entry(
   159→        dataset_dir=dataset_dir,
   160→        split=split,
   161→        mode="audio",
   162→        artifact_name=artifact_name,
   163→    )
   164→    _validate_audio_pt_format(manifest_entry)
   165→    _validate_audio_filter_config(
   166→        manifest_entry=manifest_entry,
   167→        dataset_dir=dataset_dir,
   168→        split=split,
   169→        dataset_config=dataset_config,
   170→    )
   171→
   172→    if not validate_code_range:
   173→        return
   174→
   175→    codes_path = os.path.join(dataset_dir, artifact_name)
   176→    min_code_id, max_code_id = _get_audio_code_id_range(codes_path, manifest_entry)
   177→    if min_code_id < 0 or max_code_id >= codebook_size:
   178→        raise ValueError(
   179→            "Audio pretraining artifact contains codec ids outside the configured "
   180→            f"codebook size [{codebook_size}]. Range: [{min_code_id}, {max_code_id}]."
   181→        )
   182→
   183→
   184→def validate_text_pretraining_artifact(
   185→    dataset_dir: str, split: str, tokenizer_name: str
   186→) -> None:
   187→    """Validates the text PT artifact manifest against the training tokenizer."""
   188→    artifact_name = f"{split}_pretraining_tokens.npy"
   189→    tokens_path = os.path.join(dataset_dir, artifact_name)
   190→    manifest_entry = _load_pretraining_manifest_entry(
   191→        dataset_dir=dataset_dir,
   192→        split=split,
   193→        mode="text",
   194→        artifact_name=artifact_name,
   195→    )
   196→    built_tokenizer_name = manifest_entry.get("tokenizer_name")
   197→    if built_tokenizer_name is None:
   198→        raise ValueError(
   199→            f"Text pretraining manifest for [{tokens_path}] is missing "
   200→            "`tokenizer_name`."


     1→from typing import Any
     2→
     3→import datasets
     4→import numpy as np
     5→import torch
     6→import transformers
     7→from absl import logging
     8→
     9→from tts.core import constants, prompting
    10→from tts.data import data_sample, text_normalization
    11→
    12→
    13→class TtsFineTuningDataset(torch.utils.data.Dataset):
    14→    """Implements a TTS fine-tuning dataset."""
    15→
    16→    def __init__(
    17→        self,
    18→        dataset_name: str,
    19→        samples: list[data_sample.Sample],
    20→        codes: np.ndarray,
    21→        indexes: list[tuple[int, int]],
    22→        tokenizer: transformers.PreTrainedTokenizer,
    23→        max_seq_len: int,
    24→        prompt_compiler: prompting.PromptCompiler,
    25→        text_normalizer: text_normalization.TextNormalizer,
    26→    ):
    27→        self.dataset_name = dataset_name
    28→        self.samples = samples
    29→        self.codes = codes
    30→        self.indexes = indexes
    31→        self.number_of_codes = sum(end - start for start, end in self.indexes)
    32→        self.length = len(self.samples)
    33→        self._text_normalizer = text_normalizer
    34→
    35→        if len(self.indexes) != self.length:
    36→            raise ValueError("The number of samples and codes must match!")
    37→        self.max_seq_len = max_seq_len
    38→        self.tokenizer = tokenizer
    39→        self.pad_token_id = tokenizer.pad_token_id
    40→        self.speech_start_id = tokenizer.convert_tokens_to_ids(
    41→            constants.SPEECH_START_TOKEN
    42→        )
    43→        self.speech_end_id = tokenizer.convert_tokens_to_ids(constants.SPEECH_END_TOKEN)
    44→        logging.info(
    45→            "Loadded [%s]-dataset with [%d] samples / [%d] codes. "
    46→            "Pad/speech start/speed end tokens: [%d, %d, %d]",
    47→            dataset_name,
    48→            self.length,
    49→            self.number_of_codes,
    50→            self.pad_token_id,
    51→            self.speech_start_id,
    52→            self.speech_end_id,
    53→        )
    54→        self.prompt_compiler = prompt_compiler
    55→
    56→    def __len__(self) -> int:
    57→        return self.length
    58→
    59→    def get_lengths(self) -> list[int]:
    60→        """Returns estimated sequence lengths per sample (code tokens + text overhead)."""
    61→        overhead = 64  # approximate prompt template / text token overhead
    62→        return [
    63→            min(end - start + overhead, self.max_seq_len)
    64→            for start, end in self.indexes
    65→        ]
    66→
    67→    def __getitem__(self, idx: int) -> dict[str, torch.Tensor]:
    68→        speech_ids = self.codes[self.indexes[idx][0] : self.indexes[idx][1]]
    69→        sample: data_sample.Sample = self.samples[idx]
    70→        transcript = self._text_normalizer.normalize_with_language(
    71→            sample.transcript, sample.language
    72→        )
    73→        prompt = self.prompt_compiler.compile_prompt(
    74→            audio_prompt_transcription=transcript,
    75→            text_to_synthesize="",
    76→            speech_ids=speech_ids,
    77→            voice_description=sample.voice_description,
    78→        )
    79→
    80→        full_input_ids = self.tokenizer(
    81→            prompt, add_special_tokens=True, return_tensors="pt"
    82→        )["input_ids"][0]
    83→
    84→        try:
    85→            full_separator_idx = (
    86→                (full_input_ids == self.speech_start_id)
    87→                .nonzero(as_tuple=True)[0]
    88→                .item()
    89→            )
    90→        except Exception:  # pylint: disable=broad-except
    91→            logging.warning("Speech start not found in the input ids.")
    92→            full_separator_idx = -1
    93→
    94→        if full_input_ids.shape[0] > self.max_seq_len and full_separator_idx != -1:
    95→            # Keep the speech region in view when truncation is required. If the
    96→            # whole speech span fits, this keeps the sequence tail; if it does
    97→            # not, it still anchors the window at the speech boundary so labels
    98→            # are never all ignored.
    99→            start_idx = min(
   100→                full_separator_idx,
   101→                full_input_ids.shape[0] - self.max_seq_len,
   102→            )
   103→            input_ids = full_input_ids[start_idx : start_idx + self.max_seq_len]
   104→            separator_idx = full_separator_idx - start_idx
   105→        else:
   106→            input_ids = full_input_ids[: self.max_seq_len]
   107→            separator_idx = (
   108→                full_separator_idx if full_separator_idx < self.max_seq_len else -1
   109→            )
   110→
   111→        num_tokens_to_generate = speech_ids.shape[0]
   112→        audio_processed_sec = num_tokens_to_generate / constants.CODEC_TOKENS_RATE
   113→
   114→        # Create the labels and attention mask by masking out all the tokens before
   115→        # the separator token as for traditional supervised LLM fine-tuning.
   116→        labels = torch.full_like(input_ids, constants.LOSS_IGNORE_TOKEN_ID)
   117→        if separator_idx != -1:
   118→            labels[separator_idx:] = input_ids[separator_idx:]
   119→        labels[input_ids == self.pad_token_id] = constants.LOSS_IGNORE_TOKEN_ID
   120→        attention_mask = (input_ids != self.pad_token_id).long()
   121→
   122→        result = {
   123→            "input_ids": input_ids,
   124→            "labels": labels,
   125→            "attention_mask": attention_mask,
   126→            "tokens_processed": input_ids.shape[0],
   127→            "generated_audio_duration_sec": num_tokens_to_generate
   128→            / constants.CODEC_TOKENS_RATE,
   129→            "audio_processed_sec": audio_processed_sec,
   130→        }
   131→
   132→        return result
   133→
   134→
   135→class TextFineTuningDataset(torch.utils.data.Dataset):
   136→    """Implements a text fine-tuning dataset.
   137→
   138→    LLAMA_ASSUMPTION: This class uses LLaMA 3.1's chat template
   139→    (apply_chat_template) and relies on the <|end_header_id|> token to locate
   140→    the assistant response boundary. A different backbone would need a
   141→    different response-boundary heuristic.
   142→    """
   143→
   144→    def __init__(
   145→        self,
   146→        dataset: datasets.Dataset,
   147→        tokenizer: transformers.PreTrainedTokenizer,
   148→        max_seq_len: int,
   149→    ):
   150→        self._dataset = dataset


     1→import abc
     2→from collections.abc import Sequence
     3→from typing import Any
     4→
     5→from tts.core import constants
     6→
     7→
     8→def _format_voice_description(voice_description: str) -> str:
     9→    """Formats a voice description with the appropriate tokens."""
    10→    return (
    11→        f"{constants.VOICE_DESCRIPTION_START_TOKEN}{voice_description}"
    12→        f"{constants.VOICE_DESCRIPTION_END_TOKEN}"
    13→    )
    14→
    15→
    16→def _create_message_with_voice_description(
    17→    voice_description: str, transcript: str
    18→) -> str:
    19→    """Creates a message with voice description."""
    20→    formatted_voice_description = _format_voice_description(voice_description)
    21→    return f"{formatted_voice_description} {transcript}"
    22→
    23→
    24→def _format_speech_tokens(speech_ids: list[int]) -> str:
    25→    """Formats speech IDs into speech tokens."""
    26→    speech_tokens_str = []
    27→    for speech_id in speech_ids:
    28→        speech_tokens_str.append(constants.SPEECH_TOKEN_PATTERN.format(speech_id))
    29→    return "".join(speech_tokens_str)
    30→
    31→
    32→class PromptCompiler(abc.ABC):
    33→    """Abstract base class for prompt compilers."""
    34→
    35→    @abc.abstractmethod
    36→    def compile_prompt(
    37→        self,
    38→        audio_prompt_transcription: str,
    39→        text_to_synthesize: str,
    40→        speech_ids: Sequence[int],
    41→        voice_description: str = "",
    42→        enable_instruction: bool = True,
    43→    ) -> Any:
    44→        """Compiles the chat prompt for the model."""
    45→        raise NotImplementedError("Subclasses must implement this method.")
    46→
    47→
    48→class TrainingPromptCompiler(PromptCompiler):
    49→    """Prompt compiler for training/fine-tuning."""
    50→
    51→    def compile_prompt(
    52→        self,
    53→        audio_prompt_transcription: str,
    54→        text_to_synthesize: str,
    55→        speech_ids: Sequence[int],
    56→        voice_description: str = "",
    57→        enable_instruction: bool = True,
    58→    ) -> Any:
    59→        """Compiles the chat prompt for training."""
    60→        del text_to_synthesize  # Not used in training.
    61→        del enable_instruction  # Not used in training.
    62→
    63→        user_message = self._compile_user_message(
    64→            audio_prompt_transcription, voice_description
    65→        )
    66→        assistant_message = self._compile_assistant_message(speech_ids)
    67→        return self._compile_prompt(user_message, assistant_message)
    68→
    69→    def _compile_user_message(
    70→        self, audio_prompt_transcription: str, voice_description: str = ""
    71→    ) -> str:
    72→        """Compiles a user message for finetuning."""
    73→        transcript = audio_prompt_transcription
    74→        if voice_description:
    75→            return _create_message_with_voice_description(voice_description, transcript)
    76→        return transcript
    77→
    78→    def _compile_assistant_message(self, speech_ids: Sequence[int]) -> str:
    79→        """Compiles an assistant message for finetuning."""
    80→        if len(speech_ids) == 0:
    81→            raise ValueError("Speech IDs are empty!")
    82→        speech_tokens = _format_speech_tokens(speech_ids)
    83→        return constants.SPEECH_START_TOKEN + speech_tokens + constants.SPEECH_END_TOKEN
    84→
    85→    def _compile_prompt(self, user_message: str, assistant_message: str) -> Any:
    86→        """Compiles prompt for the model."""
    87→        return user_message + assistant_message
    88→
    89→
    90→class InferencePromptCompiler(PromptCompiler):
    91→    """Prompt compiler for inference."""
    92→
    93→    def compile_prompt(
    94→        self,
    95→        audio_prompt_transcription: str,
    96→        text_to_synthesize: str,
    97→        speech_ids: Sequence[int],
    98→        voice_description: str = "",
    99→        enable_instruction: bool = True,
   100→    ) -> Any:
   101→        """Compiles the chat prompt for inference."""
   102→        user_message = self._compile_user_message(
   103→            audio_prompt_transcription,
   104→            text_to_synthesize,
   105→            voice_description,
   106→            enable_instruction,
   107→        )
   108→        assistant_message = self._compile_assistant_message(speech_ids)
   109→        return self._compile_prompt(user_message, assistant_message)
   110→
   111→    def _compile_user_message(
   112→        self,
   113→        audio_prompt_transcription: str,
   114→        text_to_synthesize: str,
   115→        voice_description: str = "",
   116→        enable_instruction: bool = True,
   117→    ) -> str:
   118→        """Compiles a user message for inference."""
   119→        if audio_prompt_transcription and (not voice_description or enable_instruction):
   120→            transcript = f"{audio_prompt_transcription} {text_to_synthesize}"
   121→        else:
   122→            transcript = text_to_synthesize
   123→
   124→        if voice_description:
   125→            return _create_message_with_voice_description(voice_description, transcript)
   126→
   127→        return transcript
   128→
   129→    def _compile_assistant_message(self, speech_ids: list[int]) -> str:
   130→        """Compiles an assistant message for inference."""
   131→        if len(speech_ids) == 0:
   132→            return constants.SPEECH_START_TOKEN
   133→
   134→        speech_tokens = _format_speech_tokens(speech_ids)
   135→        return constants.SPEECH_START_TOKEN + speech_tokens
   136→
   137→    def _compile_prompt(self, user_message: str, assistant_message: str) -> Any:
   138→        """Compiles prompt for the model."""
   139→        return user_message + assistant_message
   140→


Let me check encoder/decoder structure and LoRA implementation:


     1→import peft
     2→import torch
     3→from absl import logging
     4→
     5→from tts.utils import configuration
     6→
     7→
     8→def _find_linear_modules(model: torch.nn.Module) -> list[str]:
     9→    """Finds all linear components within a model."""
    10→    suffixes = set()
    11→    for name, param in model.named_modules():
    12→        if isinstance(param, torch.nn.Linear):
    13→            suffixes.add(name.split(".")[-1])
    14→    if len(suffixes) == 0:
    15→        raise ValueError(
    16→            f"Model {type(model)} contains no linear component, "
    17→            "LoRA can't be applied."
    18→        )
    19→    return list(suffixes)
    20→
    21→
    22→def apply_lora(
    23→    model: torch.nn.Module, config: configuration.LoraConfig
    24→) -> torch.nn.Module:
    25→    """Applies LoRA to the model."""
    26→    if config.adapter_path is not None:
    27→        model = load_adapter(model, config.adapter_path)
    28→        for name, param in model.named_parameters():
    29→            if "lora" in name.lower():
    30→                param.requires_grad = True
    31→        return model
    32→
    33→    if config.target_modules:
    34→        target_modules = list(config.target_modules)
    35→    else:
    36→        # When target_modules aren't provided, automatically find all linear components.
    37→        target_modules = _find_linear_modules(model)
    38→        logging.info("No |target_modules| provided, using all linear layers instead.")
    39→    peft_config = peft.LoraConfig(
    40→        task_type=config.task_type,
    41→        r=config.r,
    42→        lora_alpha=config.lora_alpha,
    43→        target_modules=target_modules,
    44→        lora_dropout=config.lora_dropout,
    45→        bias=config.bias,
    46→    )
    47→    return peft.get_peft_model(model, peft_config)
    48→
    49→
    50→def load_adapter(base_model: torch.nn.Module, adapter_path: str) -> torch.nn.Module:
    51→    """Loads adapter weights from directory."""
    52→    return peft.PeftModel.from_pretrained(base_model, adapter_path, device_map="auto")
    53→
    54→
    55→def save_lora_adapter(model: torch.nn.Module, adapter_path: str):
    56→    """Saves the LoRA adapter to a directory."""
    57→    model.save_pretrained(adapter_path)
    58→


     1→import abc
     2→
     3→import torch
     4→
     5→from tts.core.codec import encoder
     6→
     7→
     8→class AudioEncoderInterface(metaclass=abc.ABCMeta):
     9→    """Abstract interface class for audio encoders."""
    10→
    11→    @abc.abstractmethod
    12→    def encode(self, wav: torch.Tensor) -> torch.Tensor:
    13→        """Encodes a waveform into a sequence of tokens."""
    14→        raise NotImplementedError("Subclasses must implement this encode method.")
    15→
    16→    @property
    17→    @abc.abstractmethod
    18→    def sample_rate(self) -> int:
    19→        """Returns the input sample rate of the audio decoder."""
    20→        raise NotImplementedError("Subclasses must implement this property.")
    21→
    22→    @property
    23→    @abc.abstractmethod
    24→    def token_rate(self) -> int:
    25→        """Returns the output token rate of the audio encoder."""
    26→        raise NotImplementedError("Subclasses must implement this property.")
    27→
    28→
    29→class AudioEncoder(AudioEncoderInterface):
    30→    """Audio encoder class."""
    31→
    32→    def __init__(self, model_path: str, device: torch.device):
    33→        super().__init__()
    34→
    35→        self._device = device
    36→        self._encoder = encoder.Encoder(model_path=model_path)
    37→        self._encoder.to(self._device)
    38→        self._encoder.eval()
    39→
    40→    @torch.no_grad()
    41→    def encode(self, wav: torch.Tensor) -> torch.Tensor:
    42→        """Encodes a waveform into a sequence of tokens."""
    43→        return self._encoder.encode(wav)
    44→
    45→    @property
    46→    def sample_rate(self) -> int:
    47→        """Returns the input sample rate of the audio encoder."""
    48→        return self._encoder.sample_rate
    49→
    50→    @property
    51→    def token_rate(self) -> int:
    52→        """Returns the output token rate of the audio encoder."""
    53→        return self._encoder.token_rate
    54→
    55→
    56→class CachingAudioEncoder:
    57→    """Encodes audios and caches the results."""
    58→
    59→    def __init__(self, model_path: str, device: torch.device):
    60→        super().__init__()
    61→
    62→        self._encoder = create(model_path=model_path, device=device)
    63→        self._prompt_encoding_cache = {}
    64→
    65→    @torch.no_grad()
    66→    def encode(self, prompt_id: str, prompt_wav: torch.Tensor) -> list[int]:
    67→        if prompt_id in self._prompt_encoding_cache:
    68→            return self._prompt_encoding_cache[prompt_id]
    69→
    70→        codes = self._encoder.encode(prompt_wav).cpu().tolist()
    71→        self._prompt_encoding_cache[prompt_id] = codes
    72→        return codes
    73→
    74→
    75→def create(
    76→    model_path: str, device: torch.device | str | None = "cpu"
    77→) -> AudioEncoderInterface:
    78→    """Create audio decoder with model path and optionally a config file."""
    79→
    80→    return AudioEncoder(model_path, device=device)
    81→


     1→"""Decoding tokens to audio."""
     2→
     3→import abc
     4→import dataclasses
     5→import json
     6→import os
     7→
     8→import torch
     9→
    10→from tts.core.codec import decoder
    11→
    12→
    13→@dataclasses.dataclass(frozen=True)
    14→class DecoderConfig:
    15→    """Model config for the codec decoder model."""
    16→
    17→    model_type: str
    18→    sample_rate: int
    19→    token_rate: int
    20→    hop_length: int
    21→    upsample_factors: list[int]
    22→    kernel_sizes: list[int]
    23→
    24→    @staticmethod
    25→    def from_json(file: str | os.PathLike) -> "DecoderConfig":
    26→        with open(file) as f:
    27→            config = json.load(f)
    28→        return DecoderConfig(
    29→            model_type=config["model_type"],
    30→            sample_rate=config["sample_rate"],
    31→            token_rate=config["token_rate"],
    32→            hop_length=config["hop_length"],
    33→            upsample_factors=config["upsample_factors"],
    34→            kernel_sizes=config["kernel_sizes"],
    35→        )
    36→
    37→
    38→class AudioDecoderInterface(metaclass=abc.ABCMeta):
    39→    """Abstract interface class for audio decoders."""
    40→
    41→    @abc.abstractmethod
    42→    def decode(self, speech_ids: torch.Tensor) -> torch.Tensor:
    43→        """Decodes a batch of speech IDs into audio waveforms."""
    44→        raise NotImplementedError("Subclasses must implement this decode method.")
    45→
    46→    @property
    47→    @abc.abstractmethod
    48→    def sample_rate(self) -> int:
    49→        """Returns the sample rate of the audio decoder."""
    50→        raise NotImplementedError("Subclasses must implement this property.")
    51→
    52→    @property
    53→    @abc.abstractmethod
    54→    def token_rate(self) -> int:
    55→        """Returns the input token rate of the audio decoder."""
    56→        raise NotImplementedError("Subclasses must implement this property.")
    57→
    58→
    59→class AudioDecoder(AudioDecoderInterface):
    60→    """A wrapper around the audio decoder model for batch decoding."""
    61→
    62→    def __init__(
    63→        self,
    64→        model_path: str,
    65→        config: DecoderConfig,
    66→        device: torch.device | str | None = "cpu",
    67→    ):
    68→        super().__init__()
    69→
    70→        self._device = device
    71→        self._decoder = decoder.Decoder(
    72→            sample_rate=config.sample_rate,
    73→            hop_length=config.hop_length,
    74→            upsample_factors=config.upsample_factors,
    75→            kernel_sizes=config.kernel_sizes,
    76→            checkpoint_path=model_path,
    77→        )
    78→        self._decoder.to(self._device)
    79→        self._decoder.eval()
    80→
    81→        self._sample_rate = config.sample_rate
    82→        self._token_rate = config.token_rate
    83→
    84→    @torch.no_grad()
    85→    def decode(self, speech_ids: torch.Tensor) -> torch.Tensor:
    86→        """Decodes a batch of speech ids into audio waveforms."""
    87→        vq_codes = speech_ids.unsqueeze(0).unsqueeze(0)
    88→        vq_codes = vq_codes.to(self._device)
    89→        return self._decoder(vq_codes).detach().cpu().squeeze(0)
    90→
    91→    @property
    92→    def sample_rate(self) -> int:
    93→        return self._sample_rate
    94→
    95→    @property
    96→    def token_rate(self) -> int:
    97→        return self._token_rate
    98→
    99→
   100→def create(


Let me check the optimization and evaluation code:


     1→import abc
     2→import math
     3→
     4→import torch
     5→
     6→
     7→class LrScheduler(metaclass=abc.ABCMeta):
     8→    @abc.abstractmethod
     9→    def get_lr(self, step: int) -> float:
    10→        """Returns the learning rate for the given step."""
    11→        raise NotImplementedError("|get_lr| is an abstract method.")
    12→
    13→
    14→class ConstantLrScheduler(LrScheduler):
    15→    """Constant learning rate scheduler."""
    16→
    17→    def __init__(self, learning_rate: float):
    18→        self._learning_rate = learning_rate
    19→
    20→    def get_lr(self, step: int) -> float:
    21→        del step  # Unused.
    22→        return self._learning_rate
    23→
    24→
    25→class CosineLrScheduler(LrScheduler):
    26→    """Cosine learning rate scheduler with configurable floor."""
    27→
    28→    def __init__(
    29→        self,
    30→        learning_rate: float,
    31→        warmup_steps: int,
    32→        lr_decay_steps: int,
    33→        min_lr_ratio: float = 0.1,
    34→    ):
    35→        self._peak_lr = learning_rate
    36→        self._min_lr = learning_rate * min_lr_ratio
    37→        self._warmup_steps = warmup_steps
    38→        self._lr_decay_steps = lr_decay_steps
    39→
    40→        if self._lr_decay_steps <= self._warmup_steps:
    41→            raise ValueError("|lr_decay_steps| must be greater than |warmup_steps|.")
    42→
    43→    def get_lr(self, step: int) -> float:
    44→        if step < self._warmup_steps:
    45→            return self._peak_lr * step / self._warmup_steps
    46→
    47→        if step > self._lr_decay_steps:
    48→            return self._min_lr
    49→
    50→        decay_ratio = (step - self._warmup_steps) / (
    51→            self._lr_decay_steps - self._warmup_steps
    52→        )
    53→        if decay_ratio < 0.0 or decay_ratio > 1.0:
    54→            raise RuntimeError(
    55→                "Decay ratio must be in [0.0, 1.0]. Fix LR scheduler settings."
    56→            )
    57→        coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
    58→        return self._min_lr + coeff * (self._peak_lr - self._min_lr)
    59→
    60→
    61→def create_optimizer(
    62→    model: torch.nn.Module,
    63→    learning_rate: float,
    64→    betas: tuple[float, float],
    65→    weight_decay: float,
    66→) -> torch.optim.Optimizer:
    67→    # Separate params: no weight decay for 1-D params (biases, layer-norm weights).
    68→    decay_params = []
    69→    no_decay_params = []
    70→    for p in model.parameters():
    71→        if not p.requires_grad:
    72→            continue
    73→        if p.dim() <= 1:
    74→            no_decay_params.append(p)
    75→        else:
    76→            decay_params.append(p)
    77→
    78→    param_groups = [
    79→        {"params": decay_params, "weight_decay": weight_decay},
    80→        {"params": no_decay_params, "weight_decay": 0.0},
    81→    ]
    82→    return torch.optim.AdamW(
    83→        param_groups,
    84→        lr=learning_rate,
    85→        betas=betas,
    86→        fused=True,
    87→    )
    88→


     1→import collections
     2→import math
     3→
     4→import lightning.fabric as lightning_fabric
     5→import torch
     6→import torch.nn.functional as F
     7→
     8→from tts.core import constants
     9→from tts.data import tts_datasets
    10→from tts.utils import custom_logging
    11→
    12→
    13→@torch.no_grad()
    14→def _get_health_stats(
    15→    model: torch.nn.Module,
    16→    optimizer: torch.optim.Optimizer,
    17→    fabric: lightning_fabric.Fabric,
    18→) -> dict[str, float]:
    19→    """Returns health statistics of the model and optimizer."""
    20→    # Maximum absolute values for gradients and parameters.
    21→    max_abs_param = torch.tensor(0.0, device=fabric.device, dtype=torch.float32)
    22→    max_abs_grad = torch.tensor(0.0, device=fabric.device, dtype=torch.float32)
    23→
    24→    # Average absolute values for gradients and parameters.
    25→    avg_abs_param = torch.tensor(0.0, device=fabric.device, dtype=torch.float32)
    26→    avg_abs_grad = torch.tensor(0.0, device=fabric.device, dtype=torch.float32)
    27→
    28→    # Average values for gradients and parameters.
    29→    avg_param = torch.tensor(0.0, device=fabric.device, dtype=torch.float32)
    30→    avg_grad = torch.tensor(0.0, device=fabric.device, dtype=torch.float32)
    31→
    32→    # Number of parameters and gradients.
    33→    param_counter = torch.tensor(0, device=fabric.device, dtype=torch.long)
    34→    grad_counter = torch.tensor(0, device=fabric.device, dtype=torch.long)
    35→
    36→    for p in model.parameters():
    37→        if p.requires_grad and p.data.numel():
    38→            param_counter += p.data.numel()
    39→            avg_param += torch.sum(p.data)
    40→            abs_param = torch.abs(p.data)
    41→            avg_abs_param += torch.sum(abs_param)
    42→            max_abs_param = torch.max(torch.max(abs_param), max_abs_param)
    43→
    44→    for group in optimizer.param_groups:
    45→        for p in group["params"]:
    46→            if p.grad is not None and p.grad.data.numel() > 0:
    47→                grad_counter += p.grad.data.numel()
    48→                avg_grad += torch.sum(p.grad.data)
    49→                abs_grad = torch.abs(p.grad.data)
    50→                avg_abs_grad += torch.sum(abs_grad)
    51→                max_abs_grad = torch.max(torch.max(abs_grad), max_abs_grad)
    52→
    53→    max_abs_grad_reduced = fabric.all_reduce(max_abs_grad, reduce_op="max").item()
    54→    max_abs_param_reduced = fabric.all_reduce(max_abs_param, reduce_op="max").item()
    55→    avg_abs_grad_reduced = fabric.all_reduce(avg_abs_grad, reduce_op="sum").item()
    56→    avg_abs_param_reduced = fabric.all_reduce(avg_abs_param, reduce_op="sum").item()
    57→    avg_grad_reduced = fabric.all_reduce(avg_grad, reduce_op="sum").item()
    58→    avg_param_reduced = fabric.all_reduce(avg_param, reduce_op="sum").item()
    59→    param_counter_reduced = fabric.all_reduce(param_counter, reduce_op="sum").item()
    60→    grad_counter_reduced = fabric.all_reduce(grad_counter, reduce_op="sum").item()
    61→    if param_counter_reduced:
    62→        avg_abs_param_reduced /= param_counter_reduced
    63→        avg_param_reduced /= param_counter_reduced
    64→    if grad_counter_reduced:
    65→        avg_abs_grad_reduced /= grad_counter_reduced
    66→        avg_grad_reduced /= grad_counter_reduced
    67→
    68→    return {
    69→        "eval_max_abs_grad": max_abs_grad_reduced,
    70→        "eval_max_abs_param": max_abs_param_reduced,
    71→        "eval_avg_abs_grad": avg_abs_grad_reduced,
    72→        "eval_avg_abs_param": avg_abs_param_reduced,
    73→        "eval_avg_grad": avg_grad_reduced,
    74→        "eval_avg_param": avg_param_reduced,
    75→    }
    76→
    77→
    78→@torch.no_grad()
    79→def _compute_per_example_loss(
    80→    logits: torch.Tensor, labels: torch.Tensor
    81→) -> tuple[torch.Tensor, torch.Tensor]:
    82→    """Computes per-example token-summed loss and valid-token counts."""
    83→    shift_logits = logits[..., :-1, :].contiguous()
    84→    shift_labels = labels[..., 1:].contiguous()
    85→    per_token_loss = F.cross_entropy(
    86→        shift_logits.view(-1, shift_logits.shape[-1]),
    87→        shift_labels.view(-1),
    88→        reduction="none",
    89→        ignore_index=constants.LOSS_IGNORE_TOKEN_ID,
    90→    ).view_as(shift_labels)
    91→    valid_token_mask = shift_labels != constants.LOSS_IGNORE_TOKEN_ID
    92→    loss_sum = (per_token_loss * valid_token_mask).sum(dim=1)
    93→    token_count = valid_token_mask.sum(dim=1)
    94→    return loss_sum, token_count
    95→
    96→
    97→def _reduce_scalar(
    98→    fabric: lightning_fabric.Fabric, value: float | int
    99→) -> float:
   100→    """Reduces a scalar across all ranks using sum."""


     1→{
     2→    "training": {
     3→        "seed": 777,
     4→        "logging_steps": 100,
     5→        "eval_steps": 200,
     6→        "gradient_accumulation_steps": 1,
     7→        "gradient_clip_value": 1.0,
     8→        "learning_rate": 1e-04,
     9→        "betas": [
    10→            0.8,
    11→            0.9
    12→        ],
    13→        "warmup_ratio": 0.1,
    14→        "batch_size": 32,
    15→        "weight_decay": 0.1,
    16→        "precision": "32-true",
    17→        "strategy": "ddp",
    18→        "gradient_checkpointing": false,
    19→        "num_workers": 1
    20→    },
    21→    "modeling": {
    22→        "parameters": {}
    23→    },
    24→    "codec": {
    25→        "audio_window_size": 192000,
    26→        "sample_rate": 48000,
    27→        "hop_length": 160,
    28→        "minimum_data_sample_rate": 32000,
    29→        "upsample_factors": [
    30→            3,
    31→            2
    32→        ],
    33→        "kernel_sizes": [
    34→            7,
    35→            6
    36→        ]
    37→    },
    38→    "checkpointing": {
    39→        "save_steps": 50,
    40→        "collect_health_stats": false,
    41→        "save_intermediate_generations": true,
    42→        "keep_only_last_n_checkpoints": 5,
    43→        "only_load_model_weights": false
    44→    },
    45→    "train_weighted_datasets": {
    46→        "/path/to/your/vectorized_dataset": 1.0
    47→    },
    48→    "val_weighted_datasets": {
    49→        "/path/to/your/vectorized_dataset": 1.0
    50→    },
    51→    "dataset": {
    52→        "randomize_slider": 0,
    53→        "allowed_languages": [],
    54→        "min_dnsmos_score": 0,
    55→        "min_sample_rate": 0,
    56→        "enable_rlhf_training": false
    57→    }
    58→}
    59→


Now let me check the environment and data utilities:


     1→import dataclasses
     2→import os
     3→import subprocess
     4→import sys
     5→from typing import ClassVar
     6→
     7→import lightning.fabric as lightning_fabric
     8→import regex as re
     9→import torch
    10→from absl import logging
    11→from lightning.fabric import strategies
    12→from lightning.fabric.plugins.precision import fsdp
    13→from transformers.utils import import_utils as transformers_utils
    14→
    15→
    16→def _get_slurm_config() -> tuple[int, int, int]:
    17→    """Returns the configuration of the SLURM job."""
    18→    if not torch.cuda.is_available():
    19→        raise RuntimeError("SLURM distributed training requires CUDA.")
    20→
    21→    global_rank = int(os.environ["SLURM_PROCID"])
    22→    world_size = int(os.environ["SLURM_NPROCS"])
    23→    local_rank = global_rank % torch.cuda.device_count()
    24→    return global_rank, local_rank, world_size
    25→
    26→
    27→def _get_cuda_version_from_nvidia_smi() -> str:
    28→    """Returns the CUDA version as reported by nvidia-smi."""
    29→    try:
    30→        output = subprocess.check_output(["nvidia-smi"]).decode("utf-8")
    31→        match = re.search(r"CUDA Version:\s+(\d+\.\d+)", output)
    32→        if match:
    33→            return match.group(1)
    34→        else:
    35→            return "N/A"
    36→    except Exception:  # pylint: disable=broad-exception-caught
    37→        return "Unavailable"
    38→
    39→
    40→def _init_hardware(local_rank: int) -> torch.device:
    41→    """Configures training hardware and returns its python representation."""
    42→    if not torch.cuda.is_available():
    43→        logging.warning("CUDA is not available. Using CPU.")
    44→        return torch.device("cpu")
    45→
    46→    # Use tf32 where possible as per recommendations from
    47→    # https://huggingface.co/docs/transformers/v4.15.0/performance#tf32
    48→    torch.cuda.empty_cache()
    49→    torch.backends.cudnn.enabled = True
    50→    torch.backends.cuda.matmul.allow_tf32 = True
    51→    torch.backends.cudnn.allow_tf32 = True
    52→
    53→    # Use the same device as the current process.
    54→    local_rank = local_rank if local_rank >= 0 else 0
    55→    return torch.device("cuda", index=local_rank)
    56→
    57→
    58→def _get_fabric_precision(training_precision: str) -> str:
    59→    """Returns the precision to use for the lightning fabric."""
    60→    if training_precision == "bf16":
    61→        return "bf16-true"
    62→
    63→    raise ValueError(f"Unsupported precision: {training_precision}")
    64→
    65→
    66→def _get_strategy(
    67→    training_precision: str, strategy_name: str, find_unused_parameters: bool
    68→) -> strategies.Strategy:
    69→    """Returns the strategy to use for training."""
    70→    if strategy_name == "ddp":
    71→        if torch.cuda.device_count() == 1:
    72→            logging.info(
    73→                "Using SingleDeviceStrategy strategy for single-GPU training..."
    74→            )
    75→            return strategies.SingleDeviceStrategy(
    76→                device="cuda:0", precision=training_precision
    77→            )
    78→        logging.info("Using DDPStrategy strategy for multi-GPU training...")
    79→        return strategies.DDPStrategy(
    80→            precision=training_precision,
    81→            find_unused_parameters=find_unused_parameters,
    82→        )
    83→
    84→    elif strategy_name == "fsdp":
    85→        logging.info("Using FSDP strategy for multi-GPU training...")
    86→        # TODO: make settings configurable from the config file.
    87→        precision = fsdp.FSDPPrecision(
    88→            precision=_get_fabric_precision(training_precision)
    89→        )
    90→        return strategies.FSDPStrategy(
    91→            precision=precision,
    92→            cpu_offload=False,
    93→            backward_prefetch=torch.distributed.fsdp.BackwardPrefetch.BACKWARD_POST,
    94→            use_orig_params=True,
    95→            limit_all_gathers=True,
    96→            state_dict_type="full",
    97→            sync_module_states=True,
    98→        )
    99→
   100→    elif strategy_name == "deepspeed":
   101→        logging.info("Using DeepSpeed strategy for multi-GPU training...")
   102→        return strategies.DeepSpeedStrategy(
   103→            precision=_get_fabric_precision(training_precision)
   104→        )
   105→
   106→    raise ValueError(f"Unsupported strategy: {strategy_name}")
   107→
   108→
   109→@dataclasses.dataclass(frozen=True)
   110→class EnvironmentContext:
   111→    """Encapsulates the configuration of a distributed training/serving run.
   112→
   113→    Args:
   114→        local_rank: The local rank of the process.
   115→        global_rank: The global rank of the process.
   116→        world_size: The total number of processes.
   117→        device: The device to use for the process.
   118→        slurm_distributed: If the run is distributed via SLURM.
   119→        dry_run: If the run is "dry" in which no heavy-weight operations are performed.
   120→    """
   121→
   122→    local_rank: int
   123→    global_rank: int
   124→    world_size: int
   125→    device: torch.device
   126→    slurm_distributed: bool
   127→    dry_run: bool
   128→    DEFAULT_RANK: ClassVar[int] = -1
   129→
   130→    def is_main_process(self) -> bool:
   131→        """Returns true if this process is the main process."""
   132→        return self.global_rank in {self.DEFAULT_RANK, 0}
   133→
   134→
   135→def initialize_distributed_environment_context(
   136→    local_rank: int | None = None,
   137→    slurm_distributed: bool = False,
   138→    dry_run: bool = False,
   139→) -> EnvironmentContext:
   140→    """Initializes the distributed training context."""
   141→    cuda_build_version, cuda_runtime_version = "N/A", "N/A"
   142→    if torch.cuda.is_available():
   143→        cuda_build_version = torch.version.cuda
   144→        cuda_runtime_version = _get_cuda_version_from_nvidia_smi()
   145→        if not transformers_utils.is_flash_attn_2_available():
   146→            raise ValueError("Flash attention 2 is not available! Install it!")
   147→
   148→    python_version = str(sys.version).replace("\n", "")
   149→    logging.info(
   150→        "Initializing environment with: torch [%s], "
   151→        "cuda-built [%s], cuda-runtime [%s], python [%s]",
   152→        torch.__version__,
   153→        cuda_build_version,
   154→        cuda_runtime_version,
   155→        python_version,
   156→    )
   157→
   158→    if local_rank is None:
   159→        if slurm_distributed:
   160→            global_rank, local_rank, world_size = _get_slurm_config()
   161→        else:
   162→            local_rank = EnvironmentContext.DEFAULT_RANK
   163→            global_rank = EnvironmentContext.DEFAULT_RANK
   164→
   165→            # Support lightning CLI distributed training on a single node.
   166→            if (
   167→                os.environ.get("LOCAL_RANK", EnvironmentContext.DEFAULT_RANK)
   168→                != local_rank
   169→            ):
   170→                local_rank = int(os.environ["LOCAL_RANK"])
   171→            if os.environ.get("RANK", EnvironmentContext.DEFAULT_RANK) != global_rank:
   172→                global_rank = int(os.environ["RANK"])
   173→
   174→            # Try to find world size from environment variables.
   175→            world_size = os.environ.get("PMI_SIZE", None) or os.environ.get(
   176→                "WORLD_SIZE", None
   177→            )
   178→            if world_size is None:
   179→                logging.warning("World size is not set. Assuming 1.")
   180→                world_size = 1
   181→            else:
   182→                world_size = int(world_size)
   183→    else:
   184→        global_rank = local_rank
   185→
   186→    return EnvironmentContext(
   187→        local_rank=local_rank,
   188→        global_rank=global_rank,
   189→        world_size=world_size,
   190→        device=_init_hardware(local_rank),
   191→        slurm_distributed=slurm_distributed,
   192→        dry_run=dry_run,
   193→    )
   194→
   195→
   196→def initialize_fabric(
   197→    env_context: EnvironmentContext,
   198→    strategy_name: str,
   199→    training_precision: str = "32-true",
   200→    find_unused_parameters: bool = False,
   201→) -> lightning_fabric.Fabric:
   202→    """Initializes lightning's fabric."""
   203→    strategy = _get_strategy(training_precision, strategy_name, find_unused_parameters)
   204→    if env_context.slurm_distributed:
   205→        for env_var in ["SLURM_JOB_NUM_NODES", "SLURM_NTASKS_PER_NODE"]:
   206→            if env_var not in os.environ:
   207→                raise ValueError(f"Environment variable {env_var} is not set.")
   208→
   209→        logging.info("Start syncing distributed SLURM processes via fabric...")
   210→        logging.info(f"SLURM_JOB_NUM_NODES: {os.environ['SLURM_JOB_NUM_NODES']}")
   211→        logging.info(f"SLURM_NTASKS_PER_NODE: {os.environ['SLURM_NTASKS_PER_NODE']}")
   212→        fabric = lightning_fabric.Fabric(
   213→            precision=training_precision,
   214→            accelerator="gpu",
   215→            strategy=strategy,
   216→            devices=int(os.environ["SLURM_NTASKS_PER_NODE"]),
   217→            num_nodes=int(os.environ["SLURM_JOB_NUM_NODES"]),
   218→        )
   219→        fabric.launch()
   220→    else:
   221→        # Use lightning CLI to run on one node and several GPUs like
   222→        # `fabric run --accelerator=gpu --devices=2 ...`
   223→        fabric = lightning_fabric.Fabric(accelerator="gpu", strategy=strategy)
   224→
   225→    return fabric
   226→


     1→import collections
     2→import gc
     3→import json
     4→import os
     5→import time
     6→from typing import Any
     7→
     8→import numpy as np
     9→import torch
    10→import torchaudio
    11→from absl import logging
    12→
    13→from tts.data import data_sample, filtering
    14→from tts.utils import configuration
    15→
    16→
    17→def chunk_work(work_items: list[Any], worker_id: int, num_workers: int) -> list[Any]:
    18→    """Cuts the global work_items into a chunk for this worker_id."""
    19→    if num_workers <= 1:
    20→        return work_items
    21→
    22→    total = len(work_items)
    23→    chunk_size = total // num_workers
    24→    start_idx = worker_id * chunk_size
    25→    end_idx = total if (worker_id == num_workers - 1) else (start_idx + chunk_size)
    26→    chunk = work_items[start_idx:end_idx]
    27→    logging.info(
    28→        "Worker %d will process indices [%d, %d) => %d items.",
    29→        worker_id,
    30→        start_idx,
    31→        end_idx,
    32→        len(chunk),
    33→    )
    34→    return chunk
    35→
    36→
    37→def load_samples(
    38→    dataset_path: str, max_samples: int = -1
    39→) -> tuple[list[data_sample.Sample], float]:
    40→    """Loads samples from a dataset."""
    41→    filters = [
    42→        filtering.filter_empty_transcript,
    43→        filtering.filter_long_duration,
    44→        filtering.filter_punct_or_space_only_transcript,
    45→    ]
    46→
    47→    if not dataset_path.endswith(".jsonl"):
    48→        raise ValueError("Dataset path must end with .jsonl")
    49→
    50→    dataset_name = os.path.basename(os.path.dirname(dataset_path))
    51→    samples = []
    52→    start_time = time.time()
    53→    logging.info("Starting to load dataset %s", dataset_path)
    54→
    55→    num_filtered = 0
    56→    durations = []
    57→    total_duration = 0.0
    58→    with open(dataset_path, encoding="utf-8") as file:
    59→        for line in file:
    60→            if max_samples > 0 and len(samples) >= max_samples:
    61→                break
    62→
    63→            sample = data_sample.Sample.from_json(json.loads(line), dataset_name)
    64→            # Short-circuit filtering: stop at first filter that applies
    65→            filter_reason = None
    66→            for f in filters:
    67→                filter_reason = f(sample)
    68→                if filter_reason:
    69→                    num_filtered += 1
    70→                    break
    71→            if filter_reason:
    72→                continue
    73→
    74→            samples.append(sample)
    75→            durations.append(sample.duration)
    76→            total_duration += sample.duration
    77→            if samples and len(samples) % 100000 == 0:
    78→                logging.info(
    79→                    "Loaded %d sample metadatas. Average sample duration: %.2f seconds",
    80→                    len(samples),
    81→                    np.mean(durations),
    82→                )
    83→                durations = []
    84→
    85→    logging.info(
    86→        "Loaded %d samples from %s in %.2fsec. Total audio duration: %.2f hours. "
    87→        "Filtered %d samples.",
    88→        len(samples),
    89→        dataset_path,
    90→        time.time() - start_time,
    91→        total_duration / 3600.0,
    92→        num_filtered,
    93→    )
    94→    return samples, total_duration
    95→
    96→
    97→def load_and_filter_audio_codes_and_samples(
    98→    dataset_dir: str, split: str, dataset_config: configuration.DatasetConfig
    99→) -> tuple[np.ndarray, list[data_sample.Sample], list[tuple[int, int]], dict[str, int]]:
   100→    """Loads samples from a dataset and filters them along with audio codes."""
   101→    dataset_name = os.path.basename(dataset_dir) + "_" + split
   102→    samples_path = os.path.join(dataset_dir, f"{split}_samples.jsonl")
   103→    codes_path = os.path.join(dataset_dir, f"{split}_codes.npy")
   104→    codes_index_path = os.path.join(dataset_dir, f"{split}_codes_index.npy")
   105→    codes = np.memmap(codes_path, dtype=np.int32, mode="r")
   106→    codes_index = np.load(codes_index_path)
   107→    number_of_codes = codes.shape[0]
   108→
   109→    filtered_samples, filtered_indexes = [], []
   110→    sample_status = collections.Counter()
   111→
   112→    # Build composable filters using dataset_config.
   113→    filters = [
   114→        filtering.filter_allowed_languages(dataset_config.allowed_languages),
   115→        filtering.filter_min_sample_rate(dataset_config.min_sample_rate),
   116→        filtering.filter_min_dnsmos_score(dataset_config.min_dnsmos_score),
   117→        filtering.filter_min_audio_duration(dataset_config.min_audio_duration),
   118→        filtering.filter_empty_transcript,
   119→        filtering.filter_long_duration,
   120→        filtering.filter_punct_or_space_only_transcript,
   121→    ]
   122→
   123→    with open(samples_path, encoding="utf-8") as f:
   124→        for idx, line in enumerate(f):
   125→            sample_status["total"] += 1
   126→            try:
   127→                sample = data_sample.Sample.from_json(json.loads(line), dataset_name)
   128→            except ValueError:
   129→                sample_status["filtered_by_invalid_sample"] += 1
   130→                sample_status["total_filtered"] += 1
   131→                continue
   132→
   133→            # Apply composable filters and stop at the first that applies.
   134→            filter_reason = None
   135→            for f in filters:
   136→                filter_reason = f(sample)
   137→                if filter_reason:
   138→                    sample_status[f"filtered_by_{filter_reason}"] += 1
   139→                    sample_status["total_filtered"] += 1
   140→                    break
   141→            if filter_reason:
   142→                continue
   143→
   144→            sample_status[f"{sample.language}"] += 1
   145→            filtered_samples.append(sample)
   146→            left = codes_index[idx]
   147→            right = (
   148→                codes_index[idx + 1] if idx < len(codes_index) - 1 else number_of_codes
   149→            )
   150→            filtered_indexes.append((left, right))


Perfect! Now let me compile a comprehensive summary. Let me get a few more important files:


     1→"""Implementation of the encoder modules."""
     2→
     3→# This module is adapted from the following repository:
     4→# URL: https://github.com/zhenye234/X-Codec-2.0
     5→# License: MIT License
     6→
     7→import numpy as np
     8→import torch
     9→
    10→from tts.core.codec import activations
    11→
    12→
    13→def init_weights(m: torch.nn.Module):
    14→    """Initializes the weights of the model."""
    15→    if isinstance(m, torch.nn.Conv1d):
    16→        torch.nn.init.trunc_normal_(m.weight, std=0.02)
    17→        torch.nn.init.constant_(m.bias, 0)
    18→
    19→
    20→class ResidualUnit(torch.nn.Module):
    21→    """Residual block."""
    22→
    23→    def __init__(self, dim: int = 16, kernel_size: int = 7, dilation: int = 1):
    24→        super().__init__()
    25→        pad = ((kernel_size - 1) * dilation) // 2
    26→        self.block = torch.nn.Sequential(
    27→            activations.Activation1d(
    28→                activation=activations.SnakeBeta(dim, alpha_logscale=True)
    29→            ),
    30→            torch.nn.utils.weight_norm(
    31→                torch.nn.Conv1d(
    32→                    dim, dim, kernel_size=kernel_size, dilation=dilation, padding=pad
    33→                )
    34→            ),
    35→            activations.Activation1d(
    36→                activation=activations.SnakeBeta(dim, alpha_logscale=True)
    37→            ),
    38→            torch.nn.utils.weight_norm(torch.nn.Conv1d(dim, dim, kernel_size=1)),
    39→        )
    40→
    41→    def forward(self, x: torch.Tensor) -> torch.Tensor:
    42→        return x + self.block(x)
    43→
    44→
    45→class EncoderBlock(torch.nn.Module):
    46→    """Encoder block."""
    47→
    48→    def __init__(self, dim: int = 16, stride: int = 1, dilations=(1, 3, 9)):
    49→        super().__init__()
    50→        runits = [ResidualUnit(dim // 2, dilation=d) for d in dilations]
    51→        self.block = torch.nn.Sequential(
    52→            *runits,
    53→            activations.Activation1d(
    54→                activation=activations.SnakeBeta(dim // 2, alpha_logscale=True)
    55→            ),
    56→            torch.nn.utils.weight_norm(
    57→                torch.nn.Conv1d(
    58→                    dim // 2,
    59→                    dim,
    60→                    kernel_size=2 * stride,
    61→                    stride=stride,
    62→                    padding=stride // 2 + stride % 2,
    63→                )
    64→            ),
    65→        )
    66→
    67→    def forward(self, x: torch.Tensor) -> torch.Tensor:
    68→        return self.block(x)
    69→
    70→
    71→class SemanticEncoder(torch.nn.Module):
    72→    """Semantic encoder model to encode the wav2vec features."""
    73→
    74→    def __init__(
    75→        self,
    76→        input_channels: int,
    77→        output_channels: int,
    78→        encode_channels: int,
    79→        kernel_size: int,
    80→        bias: bool = True,


Perfect! Now I have all the information needed. Let me compile a comprehensive report:

## COMPREHENSIVE TRAINING SETUP ANALYSIS - Maya TTS LLaMA 3.2-1B

### 1. MODEL ARCHITECTURE

**Base Model:**
- Foundation: Meta-Llama 3.2-1B-Instruct (Causal Language Model)
- Location: `/home/ubuntu/training/maya-tts-llama32-1b-pretraining/tts/core/modeling.py` (lines 36-63)
- Loaded via `transformers.AutoModelForCausalLM.from_pretrained()`
- Flash Attention optimization: Automatically detects Flash Attention 3 (Hopper) or falls back to FA2 (line 18)

**Vocabulary Extension:**
- Original LLaMA vocab: 128,256 tokens
- **Extended vocab size: 193,856 tokens** (line 10 in `tokenization.py`)
- Speech tokens added:
  - Audio codes: `<|s_0|>` through `<|s_65535|>` (65,536 codec tokens) - lines 50-52 in `tokenization.py`
  - Speech control tokens: 8 special tokens (lines 40-49 in `tokenization.py`):
    - `<|speech_start|>` / `<|speech_end|>`
    - `<|text_prompt_start|>` / `<|text_prompt_end|>`
    - `<|voice_description_start|>` / `<|voice_description_end|>`
    - `<|sound_effect_start|>` / `<|sound_effect_end|>`
  - Extra padding tokens: `<|extra_token_0|>` onwards to reach 193,856

**Audio Codec (Custom 1D Audio Codec):**
- **Location:** `/home/ubuntu/training/maya-tts-llama32-1b-pretraining/tts/core/codec/` (encoder.py, decoder.py, encoder_modules.py)

**Encoder Architecture** (`encoder.py` lines 17-48):
- Input: Raw audio waveforms
- Semantic encoder: Extracts wav2vec features from Facebook W2V-BERT-2.0 (lines 51-56)
  - Uses hidden state layer 16
- Acoustic encoder: Processes raw audio with:
  - Input channels: 1024
  - Output channels: 1024
  - Kernel sizes: 3
  - Up-ratios: [2, 2, 4, 4, 5]
  - Dilations: (1, 3, 9)
- Fusion layer: Linear(2048 -> 2048) concatenates semantic + acoustic features (line 42)
- Vector Quantization (VQ):
  - ResidualFSQ with 8 quantizers at 4 levels each (line 45-46)
  - Outputs: **Integer codes 0-65535** (65,536 codebook size = 2^16)
- Token rate: **50 tokens per second** (line 20 in `constants.py`)
- Sample rate: **16kHz** (line 21 in `constants.py`)
- Hop length: 320 samples (`encoder.py` line 13)

**Decoder Architecture** (`decoder.py` lines 14-89):
- Input: VQ codes (integers 0-65535)
- Processing pipeline:
  1. VQ code lookup via `quantizer.get_output_from_indices()` (line 77)
  2. FC post-processing: Linear(2048 -> 1024) (line 63)
  3. Vocos backbone (generator.backbone) (line 82)
  4. Optional upsampler (UpSamplerBlock) with configurable factors
  5. ISTFT head for waveform synthesis
- Compatible with xcodec2 checkpoint format (encoder.py lines 85-88)

**Optional LoRA:**
- Location: `/home/ubuntu/training/maya-tts-llama32-1b-pretraining/tts/core/lora.py`
- Applies to all linear modules by default (lines 36-37)
- Configurable rank, alpha, dropout (lines 39-46)
- Loaded via PEFT library

---

### 2. TASK DEFINITION - TTS (Text-to-Speech)

**Input/Output Format:**
- **Input:**
  - Text transcription of audio prompt
  - Audio codes from prompt (50 tokens/sec @ 16kHz)
  - Optional voice description
- **Output:**
  - Generated speech codes (same format as input)
  - Model learns to autoregressively generate audio codes conditioned on text

**Prompt Compilation** (`prompting.py` lines 48-87):
- Training prompt structure (TrainingPromptCompiler):
  ```
  [user message][<|speech_start|>][speech codes][<|speech_end|>]
  ```
- User message: audio transcription + optional voice description (lines 69-76)
- Assistant message: speech tokens + start/end markers (lines 78-83)
- Speech token format: `<|s_CODE|>` for each code (lines 24-29)

**Tokenization Strategy:**
- Input/output text: BPE tokenization via LLaMA tokenizer
- Audio: Discrete integer tokens (0-65535) representing codec indices
- Mixed modality: Same token stream contains both text and audio codes
- Max sequence length: Configurable (default 2048 tokens)

**Loss Computation** (`evaluation.py` lines 79-94):
- Cross-entropy loss on next-token prediction
- Labels masked: Only speech generation portion is supervised (lines 114-119 in `finetuning.py`)
- Per-token loss with ignore_index=-100 for padding/prompt

---

### 3. TRAINING SCHEDULE & STRATEGY

**Training Stages:**
Location: `/home/ubuntu/training/maya-tts-llama32-1b-pretraining/example/configs/`

**SFT (Supervised Fine-Tuning) Config** (`sft.json`):
- Max sequence length: 2048 tokens
- Batch size: 4 (per device)
- Learning rate: 1e-4 (peak)
- Warmup ratio: 5% of total steps
- Weight decay: 0.1
- Beta values: [0.9, 0.95]
- Gradient accumulation: 1
- Gradient clipping: 1.0
- Precision: bf16
- Save checkpoints: Every 100 steps
- Evaluate: Every 300 steps
- Keep last 10 checkpoints

**Production Config** (`production_8gpu.json` lines 1-51):
- Batch size: 12 (per device)
- Learning rate: 1.5e-4
- Warmup ratio: 10%
- **Max training steps: 1,271,565** (explicitly set, line 17)
- **Max tokens per batch: 24,576** (length-bucketing to reduce padding waste)
- Save checkpoints: Every 10,000 steps
- Keep last 5 checkpoints

**Learning Rate Schedule** (`optimization.py` lines 25-58):
- Type: Cosine annealing with warmup floor
- Warmup: Linear increase from 0 to peak_lr over warmup_steps
- Decay: Cosine from peak_lr to min_lr over decay_steps
- Min LR ratio: 0.1 (10% of peak, configurable in config)

**Distributed Training** (`environment.py` lines 66-106):
- Strategies: DDP (default), FSDP, or DeepSpeed ZeRO-2
- SLURM support for multi-node training (lines 16-24)
- Fabric-based orchestration (lines 196-225)
- Precision: bf16-true or fp32 with mixed precision

**Data Mixture Strategy** (README.md lines 14-22, `production_8gpu.json` lines 36-43):
- **Pretraining PT:** Audio codec streams (50 tokens/sec, bounded windows)
  - Dataset: `pretrain-encoded/pt_output`
  - Weight: 1.0
- **Text PT:** Pure text tokens from 25B corpus
  - Dataset: `text_corpus/compiled_25B/pt_output[text]`
  - Weight: 1.0
- **Paired PT:** Aligned transcript + audio triplets
  - Dataset: `finalsftdata/sft4b[paired]`
  - Weight: 3.16
- **Default ratio (README.md line 22):** 84:10:6 (audio:text:paired)

---

### 4. DATA PIPELINE

**Input Format** (README.md lines 99-126):
```json
{
  "transcript": "Then they would swiftly dart...",
  "language": "en",
  "wav_path": "/path/to/audio.wav",
  "duration": 3.42,
  "sample_rate": 24000
}
```

**Data Vectorization Process:**

1. **Audio Encoding** (vectorization):
   - Tool: `tools/data/data_vectorizer.py`
   - Encodes raw WAV files to codec codes using encoder
   - Output format: Binary memmaps (`.npy` files)
   - Sharding: Creates `train_codes_*.npy`, `train_codes_index_*.npy`, etc.

2. **Merged Format** (after `tools/data/data_merger.py`):
   - `{split}_codes.npy` - All audio codes (int32 memmap)
   - `{split}_codes_index.npy` - Start/end indices for each sample
   - `{split}_samples.jsonl` - Sample metadata (transcript, language, etc.)
   - Reduces storage and I/O overhead

3. **Text Pretraining** (`build_text_corpus.py`):
   - Canonical document store with source/license tracking
   - Exports JSONL shards with text field
   - Compiles to `{split}_pretraining_tokens.npy` (int32 tokenized)

4. **Paired PT** (`build_pretraining_data.py`):
   - Preserves `samples.jsonl` with full Sample.from_json() compatibility
   - Filters by:
     - Languages: --allowed_languages=en,hi
     - DNSMOS quality: --min_dnsmos_score=3.0
     - Sample rate: --min_sample_rate=16000
     - Duration: --min_audio_duration=0.5

**Filtering** (`data_utils.py` lines 97-150, `filtering.py`):
- Language filtering: Allowed languages list
- Quality filtering: Minimum DNSMOS score
- Audio properties: Sample rate >= 16kHz, duration >= 0.5s
- Transcript validation: Non-empty, not just punctuation/spaces
- Status tracking: Counter of all filter reasons

**Dataset Classes** (`tts_datasets.py`, `datasets/finetuning.py`, `datasets/pretraining.py`):

**TtsFineTuningDataset** (finetuning.py lines 13-132):
- Loads codes memmap + samples JSONL
- Per-sample indexing via `codes_index[start:end]`
- Compiles training prompt with text normalization
- Masks labels before speech start token (line 118)
- Calculates sequence length with 64-token overhead estimate

**TtsPretrainingDataset** (pretraining.py):
- Handles audio PT: `{split}_pretraining_codes.npy`
- Handles text PT: `{split}_pretraining_tokens.npy`
- Handles paired PT: Full triplets with transcripts
- Validates manifest.json for codec ID ranges and format metadata
- Supports boundary-sentinel format for speech-bounded windows

**Data Loading** (`data_loading.py`):
- PyTorch DataLoader with optional distributed sampler injection
- Batch sampler state tracking for resume (line 78-83)

---

### 5. CHECKPOINT SAVING & FORMAT

**Location:** `/home/ubuntu/training/maya-tts-llama32-1b-pretraining/tts/training/checkpointing.py` (lines 156-280)

**Checkpoint Contents** (lines 171-184):
```python
{
    "model": model.state_dict(),
    "loss_statistics": Statistics,
    "optimizer": optimizer.state_dict(),
    "config": training_config_dict,
    "rng_state": torch.random state,
    "rng_state_per_rank": [RNG states per rank],  # Distributed training
    "sampler_state": {"epoch": int, "batches_yielded": int},
    "best_eval_loss": float
}
```

**Checkpoint Management:**
- Naming: `checkpoint_{step}.pt` (line 167)
- Best checkpoint: `best_model.pt` (saved when eval loss improves)
- Cleanup: Keeps last N checkpoints, deletes older ones uploaded to R2 (lines 198-233)
- Resume: Full state restoration including RNG for exact reproducibility

**Serving Format:**
- Tool: `tools/serving/convert_checkpoint.py`
- Converts checkpoint to HuggingFace-compatible format
- Includes tokenizer config + model weights

**File Organization:**
- Config saved as: `training_config.json` (constants.py line 17)
- Tokenizer: Standard HF format (vocab.json, merges.txt)
- Model weights: PyTorch format `.pt`

---

### 6. INFERENCE & EVALUATION

**Inference Script** (README.md lines 433-472, `inferencing.py`):

```bash
python tools/serving/inference.py \
    --model_checkpoint_path /path/model \
    --audio_encoder_path /path/encoder.pt \
    --audio_decoder_path /path/decoder.pt \
    --prompt_wav_path /path/prompt.wav \
    --prompt_transcription "Text..." \
    --text "Generate this text" \
    --output_path ./output.wav
```

**Inference Pipeline** (`inferencing.py` lines 53-150):

1. **Prompt Encoding:**
   - Compile prompt with text + speech codes (InferencePromptCompiler, lines 90-139)
   - Tokenize with LLM tokenizer

2. **Generation:**
   - Autoregressive token generation via `model.generate()`
   - Stop at `<|speech_end|>` token
   - Settings: temperature, top_p, top_k, repetition_penalty, etc.
   - Optional vLLM acceleration (lines 75-92)

3. **Decoding:**
   - Extract speech token IDs (lines 53-63)
   - Decode codes → waveform via AudioDecoder (lines 120-150)
   - Output: 16kHz PCM WAV

**Evaluation** (`evaluation.py`):
- Per-example loss computation (lines 79-94)
- Health statistics: Max/avg gradients/params (lines 14-75)
- Validation loss tracking

**Quality Validation** (`quality_validation.py`):
- Optional intermediate generation sampling during training
- WER-based reward functions for RLHF (constants.py lines 34-36)
- DNSMOS quality scoring
- Speaker similarity metrics

---

### 7. CONFIGURATION FILES & TOKENIZER

**Config Structure** (`utils/configuration.py`):

**TrainingConfig** (lines 110-165):
```python
seed: int
logging_steps: int
eval_steps: int
gradient_accumulation_steps: int
learning_rate: float
betas: (0.9, 0.95)
warmup_ratio: 0.05-0.10
batch_size: 4-12
weight_decay: 0.1
precision: "bf16"
strategy: "ddp" | "fsdp" | "deepspeed"
gradient_checkpointing: bool
max_training_steps: int (optional)
max_tokens_per_batch: int (optional)
min_lr_ratio: 0.1
```

**ModelingConfig** (lines 65-67):
```python
parameters: {
    "codebook_size": 65536,
    "max_seq_len": 2048,
    "model_name": "meta-llama/Llama-3.2-1B-Instruct",
    "enable_text_normalization": true/false
}
```

**DatasetConfig** (lines 39-62):
```python
allowed_languages: ["en", "es", ...]
min_dnsmos_score: 0.0-5.0
min_sample_rate: 16000
enable_rlhf_training: bool
min_audio_duration: 0.5
```

**CheckpointingConfig** (lines 71-107):
```python
save_steps: 100-10000
directory: "./experiments/..."
collect_health_stats: bool
save_intermediate_generations: bool
codec_encoder_path: "/path/encoder.pt"
codec_decoder_path: "/path/decoder.pt"
keep_only_last_n_checkpoints: 5-10
```

**Tokenizer Customization** (`tokenization.py`):
- Base tokenizer: LLaMA-3.2 BPE (128,256 tokens)
- New tokens added (lines 40-59):
  - 8 special control tokens
  - 65,536 speech tokens (`<|s_0|>` ... `<|s_65535|>`)
  - Padding to 193,856
- Saved to checkpoint dir with `.save_pretrained()` (main.py lines 135-142)

---

### 8. DOCUMENTATION & REFERENCES

**Key Documentation Files:**
- `/home/ubuntu/training/maya-tts-llama32-1b-pretraining/README.md` - Quick start and full pipeline
- `/home/ubuntu/training/maya-tts-llama32-1b-pretraining/TRAINING_GUIDE.md` - LlamaFactory-based training
- `/home/ubuntu/training/maya-tts-llama32-1b-pretraining/PRETRAINING_HANDBOOK.md` - Pretraining details
- `/home/ubuntu/training/maya-tts-llama32-1b-pretraining/AGENT_PROMPT.md` - Development guidance

**Example Configs:**
- `/home/ubuntu/training/maya-tts-llama32-1b-pretraining/example/configs/sft.json` - SFT template
- `/home/ubuntu/training/maya-tts-llama32-1b-pretraining/example/configs/production_8gpu.json` - Production (1.27M steps)
- `/home/ubuntu/training/maya-tts-llama32-1b-pretraining/example/configs/codec_training_config.json` - Codec training

**Entry Points:**
- **Training:** `/home/ubuntu/training/maya-tts-llama32-1b-pretraining/tts/training/main.py`
- **RLHF Training:** `/home/ubuntu/training/maya-tts-llama32-1b-pretraining/tts/training/rlhf/rlhf_main.py`
- **Inference:** `/home/ubuntu/training/maya-tts-llama32-1b-pretraining/tools/serving/inference.py`
- **Data Vectorization:** `/home/ubuntu/training/maya-tts-llama32-1b-pretraining/tools/data/data_vectorizer.py`
- **Codec Training:** `/home/ubuntu/training/maya-tts-llama32-1b-pretraining/tts/training/codec/train_codec.py`

**Technical Papers Referenced:**
- Llama models: https://www.llama.com/
- Codec inspiration: Llasa (https://arxiv.org/abs/2502.04128)
- Inworld TTS technical report: https://arxiv.org/abs/2507.21138

---

### SUMMARY TABLE

| Component | Details |
|-----------|---------|
| **Base Model** | Meta-Llama-3.2-1B-Instruct (Causal LM) |
| **Vocab Size** | 193,856 (128k base + 65.5k speech tokens) |
| **Audio Codec** | Custom 1D codec: w2v semantic + acoustic fusion → FSQ (50 tokens/sec, 16kHz) |
| **Max Sequence** | 2048 tokens |
| **Batch Size** | 4-12 (configurable) |
| **Learning Rate** | 1e-4 to 1.5e-4 with cosine annealing + warmup |
| **Precision** | bfloat16 |
| **Training Steps** | Up to 1.27M (production) |
| **Checkpoint Interval** | 100-10,000 steps |
| **Distributed** | DDP/FSDP/DeepSpeed with multi-node SLURM support |
| **Data Format** | Audio codes (int32 memmap) + JSONL metadata |
| **Loss** | Cross-entropy on speech codes (decoder part only) |
| **Inference** | Autoregressive generation → codec decode → waveform |