mirror of
https://github.com/Nighthawk42/llm-tts-factory.git
synced 2026-08-30 07:22:27 +00:00
Minor fixes, code parity.
This commit is contained in:
Nighthawk
@@ -0,0 +1,40 @@
|
||||
# Training Strategy & Learnings for LLM-Backbone TTS
|
||||
|
||||
---
|
||||
|
||||
## 1. The Prompt Format & Token Specialization
|
||||
Early experiments showed that simply feeding text and discrete audio tokens without clear boundaries leads to confusion. The LLM struggles to differentiate between "understanding text" and "generating audio representations."
|
||||
|
||||
**Learning:** We heavily customized the prompt and tokenizer to include explicit, distinct special tokens.
|
||||
- The format must explicitly anchor the model's intent: `[TEXT]<text prompt>[START]<audio tokens>[STOP]`.
|
||||
- We added dedicated special tokens (`<|audio_start|>`, `<|audio_end|>`, `<|text|>`, etc.) to the tokenizer vocabulary and resized the embedding matrices. This provides absolute clarity to the causal transformer regarding which modality it is currently processing.
|
||||
|
||||
---
|
||||
|
||||
## 2. LLM Training Stability Strategies
|
||||
Teaching a causal language model (Qwen backbone) to map text directly to audio-latent distributions often results in the model outputting garbage audio if not carefully stabilized.
|
||||
|
||||
**Learning & Strategy:** The loss calculation must be incredibly intentional.
|
||||
* **Batching:** Changed the batching from packing to dynamic batching - one sample per row. Felt like this could be more stable. Shall experiment training with packing. I think this should also work.
|
||||
* **Text Weighting:** Initially trained the llm with text_weight as 0, but turns out that the model wasn't able to learn this way. It was overfitting and wasn't able to generalize. I was guessing even with 0 weights for text tokens, the model should be able to learn the relationship between text and audio tokens and generalize well.
|
||||
If the text tokens are zero-weighted in the loss calculation aggressively from the start, the model can lose its language understanding capabilities, learning only to spout random audio formats. A critical learning is to experiment with text weighting strategies (e.g., maintaining a small loss weight on the text prediction part for the first N steps) so the LLM retains language comprehension and can generalize on new texts.
|
||||
* **Masking for Loss:** Ensure that padding tokens (`[PAD]`) or masked sections are strictly ignored in the `CrossEntropyLoss` calculation. Any gradient penalty on padding will rapidly destabilize autoregressive predictions.
|
||||
* **Gradual Length Training (Curriculum Learning):** The model often crashes or learns poorly if immediately exposed to full 1000+ token sequences. It's recommended starting with shorter text-audio pairs and gradually increasing the sequence context length (`max_length`) as the model stabilizes. Yet to try this.
|
||||
|
||||
---
|
||||
|
||||
## 3. Two-Stage Decoder (Vocoder) Training
|
||||
The role of the Decoder (Vocos) is to take the *continuous* hidden states evaluated by the LLM and construct high-fidelity audio waveforms.
|
||||
|
||||
**Learning & Strategy:** The Decoder must be trained independently in two distinct stages, with the LLM rigidly completely frozen.
|
||||
|
||||
#### Stage 1: Pure Reconstruction (Global Structure)
|
||||
* **Objective:** Teach the Decoder how to map hidden state dimensions to spectral layout.
|
||||
* **Method:** Train using purely reconstruction losses: **Mel-Spectrogram L1 Loss** and **Multi-Resolution STFT Loss**.
|
||||
* **Data:** This phase is executed over the *almost-full* generated sequences to ensure the Decoder learns long-term structural coherence and timing alignment.
|
||||
|
||||
#### Stage 2: Adversarial Refinement (Local Fidelity)
|
||||
* **Objective:** Eliminate robotic artifacts and "muffled" qualities to achieve crystal-clear, natural acoustic texture.
|
||||
* **Method:** Introduce Discriminator networks (Multi-Period and Multi-Scale).
|
||||
* **Crucial Learning (Cropping):** Running Discriminators (GANs) on full-length sequences destroys memory and stability. We transitioned to **Segment-Based Training (Random Cropping)**. During this stage, we extract random 1-second chunks (e.g., ~32,000 samples) from both the ground truth and the generated audio, feeding only these small crops into the Discriminator.
|
||||
* **Combined Loss:** The final phase loss comprises Reconstruction (Full Audio) + Adversarial GAN Loss (Cropped Audio) + Discriminator Feature Matching.
|
||||
+31
-12
@@ -11,6 +11,7 @@ global:
|
||||
num_workers: 4
|
||||
use_wandb: true
|
||||
wandb_project: "soprano-tts"
|
||||
tokenizer_name: "ekwek/Soprano-80M"
|
||||
|
||||
# ------------------------------------------------------------------------------
|
||||
# File Paths
|
||||
@@ -20,23 +21,30 @@ global:
|
||||
# Examples:
|
||||
# Linux: "/home/ubuntu/data/lj_speech/LJSpeech-1.1"
|
||||
# Windows: "C:/Users/Name/Documents/datasets/LJSpeech-1.1"
|
||||
# Relative: "./datasets/LJSpeech-1.1"
|
||||
# Relative: "./data/LJSpeech-1.1"
|
||||
# ------------------------------------------------------------------------------
|
||||
paths:
|
||||
# Path to LJSpeech-1.1 directory
|
||||
dataset_root: "./data/LJSpeech-1.1"
|
||||
|
||||
# Base directory to save all checkpoints and logs
|
||||
# Base directory to save all checkpoints, logs, and generated datasets
|
||||
save_dir: "./checkpoints"
|
||||
|
||||
# Pretrained model paths (set to null if training from scratch)
|
||||
# Linux ex: "/home/ubuntu/soma/ckpt/suprano/codec/step_42000.pt"
|
||||
# Pretrained model paths (set to null if training from scratch or not applicable)
|
||||
pretrained_codec_path: null
|
||||
pretrained_llm_path: null
|
||||
pretrained_decoder_path: null
|
||||
pretrained_discriminator_path: null
|
||||
|
||||
# ------------------------------------------------------------------------------
|
||||
# Codec Training Configuration
|
||||
# Data Generation Configuration (generate_dataset*.py)
|
||||
# ------------------------------------------------------------------------------
|
||||
data_generation:
|
||||
val_prop: 0.1
|
||||
val_max: 512
|
||||
|
||||
# ------------------------------------------------------------------------------
|
||||
# Codec Training Configuration (codec_train.py)
|
||||
# ------------------------------------------------------------------------------
|
||||
codec:
|
||||
sample_rate: 32000
|
||||
@@ -44,9 +52,14 @@ codec:
|
||||
num_epochs: 100
|
||||
learning_rate: 1.0e-4
|
||||
freeze_encoder: false
|
||||
# Encoder/Decoder architecture params
|
||||
encoder_dim: 768
|
||||
encoder_num_layers: 8
|
||||
bottleneck_channels: 5
|
||||
decoder_num_layers: 8
|
||||
|
||||
# ------------------------------------------------------------------------------
|
||||
# LLM Training Configuration
|
||||
# LLM Training Configuration (train_llm.py)
|
||||
# ------------------------------------------------------------------------------
|
||||
llm:
|
||||
from_scratch: false
|
||||
@@ -65,11 +78,11 @@ llm:
|
||||
weight_decay: 0.1
|
||||
|
||||
# ------------------------------------------------------------------------------
|
||||
# Decoder (Vocos) Training Configuration
|
||||
# Decoder (Vocos) Training Configuration (train_decoder.py)
|
||||
# ------------------------------------------------------------------------------
|
||||
decoder:
|
||||
use_discriminator: true
|
||||
batch_size: 64
|
||||
batch_size: 8 # Reduced to avoid OOM, adjust based on your GPU
|
||||
max_steps: 200000
|
||||
max_lr: 2.0e-4
|
||||
min_lr_ratio: 0.1
|
||||
@@ -79,9 +92,12 @@ decoder:
|
||||
seq_len: 1024
|
||||
segment_size_samples: 32768 # ~1 sec (16 tokens)
|
||||
val_freq: 250
|
||||
val_steps: 10
|
||||
save_freq: 3000
|
||||
text_factor: 0.0
|
||||
betas: [0.8, 0.99]
|
||||
weight_decay: 0.1
|
||||
start_step: 0
|
||||
|
||||
# Loss Weights
|
||||
lambda_mel: 45.0
|
||||
@@ -90,8 +106,11 @@ decoder:
|
||||
lambda_stft: 1.0
|
||||
|
||||
# ------------------------------------------------------------------------------
|
||||
# Data Generation Configuration (generate_dataset*.py)
|
||||
# Inference Configuration (simple_inference.py)
|
||||
# ------------------------------------------------------------------------------
|
||||
data_generation:
|
||||
val_prop: 0.1
|
||||
val_max: 512
|
||||
inference:
|
||||
temperature: 0.8
|
||||
top_k: 50
|
||||
top_p: 0.95
|
||||
repetition_penalty: 1.2
|
||||
max_new_tokens: 512
|
||||
+384
-376
@@ -2,14 +2,16 @@
|
||||
Training script for Soprano Decoder (Vocos).
|
||||
Freezes LLM and trains Decoder with GAN loss.
|
||||
"""
|
||||
import os
|
||||
import random
|
||||
import time
|
||||
import os
|
||||
import io
|
||||
import wandb
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torchaudio
|
||||
from torch.utils.data import DataLoader
|
||||
from tqdm import tqdm
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
|
||||
@@ -22,35 +24,33 @@ from decoder.losses import MelSpectrogramWrapper, feature_matching_loss, discrim
|
||||
|
||||
from config_loader import load_config
|
||||
|
||||
# Global Tokenizer for collate function
|
||||
tokenizer = AutoTokenizer.from_pretrained('ekwek/Soprano-80M')
|
||||
tokenizer.padding_side = 'right' # Essential for training!
|
||||
|
||||
def worker_seed_init(_):
|
||||
worker_seed = torch.initial_seed() % (2**32-1)
|
||||
np.random.seed(worker_seed)
|
||||
random.seed(worker_seed)
|
||||
|
||||
def get_lr(it, max_lr, min_lr, warmup_steps, cooldown_steps, max_steps): # WSD schedule
|
||||
|
||||
def get_lr(it, max_lr, min_lr, warmup_steps, cooldown_steps, max_steps):
|
||||
if it < warmup_steps:
|
||||
return max_lr * (it + 1) / warmup_steps
|
||||
if it < max_steps - cooldown_steps:
|
||||
return max_lr
|
||||
return min_lr + (max_lr - min_lr) * ((max_steps - it) / cooldown_steps)
|
||||
|
||||
def collate_pack(batch_in):
|
||||
# batch_in is list of (text, wav)
|
||||
|
||||
def collate_pack(batch_in, tokenizer):
|
||||
texts = [x[0] for x in batch_in]
|
||||
wavs = [x[1] for x in batch_in]
|
||||
aud_token_lens = [x[2] for x in batch_in]
|
||||
|
||||
# We need to process each sample to align audio
|
||||
# Since lengths vary, we process list then pad
|
||||
tokens_batch = tokenizer(texts, padding=True, return_tensors='pt')
|
||||
input_ids = tokens_batch['input_ids']
|
||||
|
||||
batch_tokens_list = []
|
||||
batch_audio_list = []
|
||||
|
||||
for i in range(len(texts)):
|
||||
# Get raw tokens without padding for alignment logic
|
||||
raw_tokens = tokenizer(texts[i], padding=False, truncation=False)['input_ids']
|
||||
tokens = torch.tensor(raw_tokens, dtype=torch.long)
|
||||
|
||||
@@ -73,368 +73,25 @@ def collate_pack(batch_in):
|
||||
batch_tokens_list.append(tokens)
|
||||
batch_audio_list.append(aligned_audio)
|
||||
|
||||
# Pad Tokens
|
||||
batch_tokens = torch.nn.utils.rnn.pad_sequence(batch_tokens_list, batch_first=True, padding_value=0)
|
||||
|
||||
# Pad Audio
|
||||
batch_audio = torch.nn.utils.rnn.pad_sequence(batch_audio_list, batch_first=True, padding_value=0.0)
|
||||
|
||||
x = batch_tokens[:, :-1]
|
||||
y = batch_tokens[:, 1:]
|
||||
|
||||
# Calculate max seq len of x
|
||||
max_len_x = x.size(1)
|
||||
gt_audio = batch_audio[:, :max_len_x * SAMPLES_PER_TOKEN]
|
||||
|
||||
# Create Audio Mask (True where token is audio)
|
||||
audio_mask = (y > 3) & (y <= 8003)
|
||||
|
||||
return x, y, gt_audio, audio_mask
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
# ------------------
|
||||
# Load Configuration
|
||||
# ------------------
|
||||
config = load_config("config.yaml")
|
||||
cfg_global = config["global"]
|
||||
cfg_paths = config["paths"]
|
||||
cfg_decoder = config["decoder"]
|
||||
|
||||
device = cfg_global["device"]
|
||||
seed = cfg_global["seed"]
|
||||
device_type = "cuda" if device.startswith("cuda") else "cpu"
|
||||
|
||||
torch.manual_seed(seed)
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.manual_seed(seed)
|
||||
torch.set_float32_matmul_precision('high')
|
||||
|
||||
# Setup directories
|
||||
train_dataset_path = os.path.join(cfg_paths["dataset_root"], "train.json")
|
||||
val_dataset_path = os.path.join(cfg_paths["dataset_root"], "val.json")
|
||||
save_path = os.path.join(cfg_paths["save_dir"], "decoder")
|
||||
os.makedirs(save_path, exist_ok=True)
|
||||
print(f"Save Path: {save_path}")
|
||||
|
||||
if cfg_global["use_wandb"]:
|
||||
wandb.init(project=cfg_global["wandb_project"], config=config)
|
||||
|
||||
# Initialize Mel Spectrogram Wrapper dynamically
|
||||
mel_fn = MelSpectrogramWrapper().to(device)
|
||||
|
||||
# ------------------
|
||||
# Hyperparameters
|
||||
# ------------------
|
||||
max_steps = cfg_decoder["max_steps"]
|
||||
max_lr = float(cfg_decoder["max_lr"])
|
||||
min_lr = cfg_decoder["min_lr_ratio"] * max_lr
|
||||
warmup_steps = int(max_steps * cfg_decoder["warmup_ratio"])
|
||||
cooldown_steps = int(max_steps * cfg_decoder["cooldown_ratio"])
|
||||
|
||||
batch_size = cfg_decoder["batch_size"]
|
||||
segment_size_samples = cfg_decoder["segment_size_samples"]
|
||||
val_freq = cfg_decoder["val_freq"]
|
||||
save_freq = cfg_decoder["save_freq"]
|
||||
betas = tuple(cfg_decoder["betas"])
|
||||
weight_decay = cfg_decoder["weight_decay"]
|
||||
start_step = cfg_decoder.get("start_step", 0)
|
||||
|
||||
# Loss Weights
|
||||
lambda_mel = cfg_decoder["lambda_mel"]
|
||||
lambda_fm = cfg_decoder["lambda_fm"]
|
||||
lambda_gen = cfg_decoder["lambda_gen"]
|
||||
lambda_stft = cfg_decoder["lambda_stft"]
|
||||
|
||||
# ------------------
|
||||
# 1. Load LLM and Freeze
|
||||
# ------------------
|
||||
print("Loading LLM...")
|
||||
llm_config = AutoConfig.from_pretrained('ekwek/Soprano-80M')
|
||||
model = AutoModelForCausalLM.from_config(llm_config)
|
||||
|
||||
pretrained_llm_path = cfg_paths["pretrained_llm_path"]
|
||||
if pretrained_llm_path and os.path.exists(pretrained_llm_path):
|
||||
print(f"Loading custom LLM checkpoint from {pretrained_llm_path}")
|
||||
if pretrained_llm_path.endswith('.safetensors'):
|
||||
state_dict = load_file(pretrained_llm_path)
|
||||
model.load_state_dict(state_dict)
|
||||
else:
|
||||
model = AutoModelForCausalLM.from_pretrained(pretrained_llm_path)
|
||||
else:
|
||||
print("Warning: No pretrained LLM provided for Decoder training. Using random init (not recommended).")
|
||||
|
||||
model.to(torch.bfloat16).to(device)
|
||||
model.eval()
|
||||
for param in model.parameters():
|
||||
param.requires_grad = False
|
||||
print("LLM Frozen.")
|
||||
|
||||
# ------------------
|
||||
# 2. Load Decoder
|
||||
# ------------------
|
||||
print("Loading Decoder...")
|
||||
decoder = SopranoDecoder()
|
||||
pretrained_decoder_path = cfg_paths["pretrained_decoder_path"]
|
||||
|
||||
if pretrained_decoder_path and os.path.exists(pretrained_decoder_path):
|
||||
print(f"Loading custom Decoder checkpoint from {pretrained_decoder_path}")
|
||||
decoder.load_state_dict(torch.load(pretrained_decoder_path, map_location='cpu'))
|
||||
else:
|
||||
print("Training Decoder from scratch.")
|
||||
|
||||
decoder.to(device)
|
||||
decoder.train()
|
||||
print("Decoder loaded.")
|
||||
|
||||
# Initialize MR-STFT Loss
|
||||
mr_stft = MultiResolutionSTFTLoss().to(device)
|
||||
|
||||
# ------------------
|
||||
# 3. Load Discriminator
|
||||
# ------------------
|
||||
discriminator = None
|
||||
if cfg_decoder["use_discriminator"]:
|
||||
print("Initializing Discriminator...")
|
||||
discriminator = Discriminator()
|
||||
pretrained_disc_path = cfg_paths["pretrained_discriminator_path"]
|
||||
|
||||
if pretrained_disc_path and os.path.exists(pretrained_disc_path):
|
||||
print(f"Loading custom Discriminator checkpoint from {pretrained_disc_path}")
|
||||
discriminator.load_state_dict(torch.load(pretrained_disc_path, map_location='cpu'))
|
||||
|
||||
discriminator.to(device)
|
||||
discriminator.train()
|
||||
else:
|
||||
print("Training WITHOUT Discriminator (Reconstruction only).")
|
||||
|
||||
# ------------------
|
||||
# 4. Dataset Setup
|
||||
# ------------------
|
||||
dataset = AudioDataset(train_dataset_path)
|
||||
dataloader = DataLoader(
|
||||
dataset,
|
||||
batch_size=batch_size,
|
||||
shuffle=True,
|
||||
num_workers=cfg_global["num_workers"],
|
||||
pin_memory=True,
|
||||
persistent_workers=True,
|
||||
worker_init_fn=worker_seed_init,
|
||||
collate_fn=collate_pack,
|
||||
)
|
||||
dataloader_it = iter(dataloader)
|
||||
|
||||
val_dataset = AudioDataset(val_dataset_path)
|
||||
val_dataloader = DataLoader(
|
||||
val_dataset,
|
||||
batch_size=max(1, batch_size // 4), # Reduce val batch size to prevent OOM
|
||||
shuffle=False,
|
||||
num_workers=max(1, cfg_global["num_workers"] // 2),
|
||||
pin_memory=True,
|
||||
persistent_workers=True,
|
||||
worker_init_fn=worker_seed_init,
|
||||
collate_fn=collate_pack,
|
||||
)
|
||||
val_dataloader_it = iter(val_dataloader)
|
||||
|
||||
# ------------------
|
||||
# 5. Optimizers
|
||||
# ------------------
|
||||
opt_g = torch.optim.AdamW(decoder.parameters(), max_lr, betas=betas, weight_decay=weight_decay)
|
||||
opt_d = None
|
||||
if cfg_decoder["use_discriminator"]:
|
||||
opt_d = torch.optim.AdamW(discriminator.parameters(), max_lr, betas=betas, weight_decay=weight_decay)
|
||||
|
||||
# ------------------
|
||||
# Training Loop
|
||||
# ------------------
|
||||
pbar = tqdm(range(start_step + 1, max_steps + 1), ncols=200, dynamic_ncols=True)
|
||||
|
||||
for step in pbar:
|
||||
start = time.time()
|
||||
|
||||
# Get Data
|
||||
try:
|
||||
batch_data = next(dataloader_it)
|
||||
if batch_data[0] is None:
|
||||
dataloader_it = iter(dataloader)
|
||||
batch_data = next(dataloader_it)
|
||||
x, y, gt_audio, audio_mask = batch_data
|
||||
except StopIteration:
|
||||
dataloader_it = iter(dataloader)
|
||||
batch_data = next(dataloader_it)
|
||||
x, y, gt_audio, audio_mask = batch_data
|
||||
|
||||
x, y = x.to(device), y.to(device)
|
||||
gt_audio = gt_audio.to(device) # (B, T_audio_samples)
|
||||
audio_mask = audio_mask.to(device)
|
||||
|
||||
# Forward LLM (No Grad)
|
||||
with torch.no_grad():
|
||||
with torch.autocast(device_type=device_type, dtype=torch.bfloat16):
|
||||
outputs = model(x, output_hidden_states=True)
|
||||
hidden_states = outputs.hidden_states[-1] # (B, T_total, D)
|
||||
hidden_states = hidden_states.to(torch.float32)
|
||||
|
||||
# GATHER AUDIO LATENTS Logic
|
||||
gathered_states_list = []
|
||||
for b_idx in range(hidden_states.size(0)):
|
||||
mask = audio_mask[b_idx]
|
||||
valid_states = hidden_states[b_idx][mask]
|
||||
gathered_states_list.append(valid_states)
|
||||
|
||||
decoder_in_padded = torch.nn.utils.rnn.pad_sequence(gathered_states_list, batch_first=True, padding_value=0.0)
|
||||
|
||||
bsz = decoder_in_padded.size(0)
|
||||
max_aud_len = decoder_in_padded.size(1)
|
||||
audio_loss_mask = torch.zeros((bsz, max_aud_len), dtype=torch.bool, device=device)
|
||||
for b_idx in range(bsz):
|
||||
length = gathered_states_list[b_idx].size(0)
|
||||
audio_loss_mask[b_idx, :length] = True
|
||||
|
||||
# ---------------------
|
||||
# Generator Forward
|
||||
# ---------------------
|
||||
decoder_in = decoder_in_padded.transpose(1, 2) # (B, C, T)
|
||||
fake_audio = decoder(decoder_in) # (B, 1, T_audio_gen)
|
||||
if fake_audio.size(1) == 1: fake_audio = fake_audio.squeeze(1)
|
||||
|
||||
min_len = min(fake_audio.size(1), gt_audio.size(1))
|
||||
fake_audio = fake_audio[:, :min_len]
|
||||
real_audio = gt_audio[:, :min_len]
|
||||
|
||||
# ---------------------
|
||||
# Train Discriminator
|
||||
# ---------------------
|
||||
d_loss_item = 0.0
|
||||
if cfg_decoder["use_discriminator"]:
|
||||
opt_d.zero_grad()
|
||||
|
||||
# --- Random Cropping Logic ---
|
||||
real_crop_list = []
|
||||
fake_crop_list = []
|
||||
|
||||
for b_idx in range(bsz):
|
||||
valid_len = gathered_states_list[b_idx].size(0) * SAMPLES_PER_TOKEN
|
||||
valid_len = min(valid_len, min_len)
|
||||
|
||||
if valid_len <= segment_size_samples:
|
||||
pad_len = segment_size_samples - valid_len
|
||||
r_c = torch.nn.functional.pad(real_audio[b_idx, :valid_len], (0, pad_len))
|
||||
f_c = torch.nn.functional.pad(fake_audio[b_idx, :valid_len], (0, pad_len))
|
||||
else:
|
||||
start_idx = random.randint(0, valid_len - segment_size_samples)
|
||||
r_c = real_audio[b_idx, start_idx : start_idx + segment_size_samples]
|
||||
f_c = fake_audio[b_idx, start_idx : start_idx + segment_size_samples]
|
||||
|
||||
real_crop_list.append(r_c)
|
||||
fake_crop_list.append(f_c)
|
||||
|
||||
real_crops = torch.stack(real_crop_list).unsqueeze(1)
|
||||
fake_crops = torch.stack(fake_crop_list).unsqueeze(1).detach()
|
||||
|
||||
y_d_rs, y_d_gs, _, _ = discriminator(real_crops, fake_crops)
|
||||
d_loss, _, _ = discriminator_loss(y_d_rs, y_d_gs)
|
||||
|
||||
d_loss.backward()
|
||||
torch.nn.utils.clip_grad_norm_(discriminator.parameters(), 1.0)
|
||||
opt_d.step()
|
||||
d_loss_item = d_loss.item()
|
||||
|
||||
# ---------------------
|
||||
# Train Generator
|
||||
# ---------------------
|
||||
opt_g.zero_grad()
|
||||
|
||||
# We need "fake_crops_g" (with grad) for generator loss
|
||||
real_crop_list_g = []
|
||||
fake_crop_list_g = []
|
||||
if cfg_decoder["use_discriminator"]:
|
||||
for b_idx in range(bsz):
|
||||
valid_len = gathered_states_list[b_idx].size(0) * SAMPLES_PER_TOKEN
|
||||
valid_len = min(valid_len, min_len)
|
||||
|
||||
if valid_len <= segment_size_samples:
|
||||
pad_len = segment_size_samples - valid_len
|
||||
r_c = torch.nn.functional.pad(real_audio[b_idx, :valid_len], (0, pad_len))
|
||||
f_c = torch.nn.functional.pad(fake_audio[b_idx, :valid_len], (0, pad_len))
|
||||
else:
|
||||
start_idx = random.randint(0, valid_len - segment_size_samples)
|
||||
r_c = real_audio[b_idx, start_idx : start_idx + segment_size_samples]
|
||||
f_c = fake_audio[b_idx, start_idx : start_idx + segment_size_samples]
|
||||
|
||||
real_crop_list_g.append(r_c)
|
||||
fake_crop_list_g.append(f_c)
|
||||
|
||||
real_crops_g = torch.stack(real_crop_list_g).unsqueeze(1)
|
||||
fake_crops_g = torch.stack(fake_crop_list_g).unsqueeze(1)
|
||||
|
||||
# Mel Loss
|
||||
frames_per_token = SAMPLES_PER_TOKEN // 512
|
||||
mel_mask = audio_loss_mask.repeat_interleave(frames_per_token, dim=1)
|
||||
|
||||
pred_mel = mel_fn(fake_audio)
|
||||
gt_mel = mel_fn(real_audio)
|
||||
|
||||
min_mel_len = min(pred_mel.size(2), gt_mel.size(2), mel_mask.size(1))
|
||||
pred_mel = pred_mel[:, :, :min_mel_len]
|
||||
gt_mel = gt_mel[:, :, :min_mel_len]
|
||||
mel_mask = mel_mask[:, :min_mel_len]
|
||||
|
||||
loss_mel_raw = torch.nn.functional.l1_loss(pred_mel, gt_mel, reduction='none')
|
||||
loss_mel = (loss_mel_raw * mel_mask.unsqueeze(1)).sum() / (mel_mask.sum() * pred_mel.size(1) + 1e-6)
|
||||
|
||||
# Multi-Resolution STFT Loss
|
||||
sample_mask = audio_loss_mask.repeat_interleave(SAMPLES_PER_TOKEN, dim=1)
|
||||
sample_mask = sample_mask[:, :min_len]
|
||||
|
||||
sc_loss, mag_loss = mr_stft(fake_audio * sample_mask, real_audio * sample_mask)
|
||||
|
||||
loss_fm = torch.tensor(0.0, device=device)
|
||||
loss_gen = torch.tensor(0.0, device=device)
|
||||
|
||||
if cfg_decoder["use_discriminator"]:
|
||||
y_d_rs, y_d_gs, fmap_rs, fmap_gs = discriminator(real_crops_g, fake_crops_g)
|
||||
loss_fm = feature_matching_loss(fmap_rs, fmap_gs)
|
||||
loss_gen, _ = generator_loss(y_d_gs)
|
||||
|
||||
total_loss_g = (lambda_mel * loss_mel) + (lambda_gen * loss_gen) + (lambda_fm * loss_fm) + (lambda_stft * (sc_loss + mag_loss))
|
||||
|
||||
total_loss_g.backward()
|
||||
norm_g = torch.nn.utils.clip_grad_norm_(decoder.parameters(), 1.0)
|
||||
|
||||
# LR Update
|
||||
lr = get_lr(step, max_lr, min_lr, warmup_steps, cooldown_steps, max_steps)
|
||||
for param_group in opt_g.param_groups: param_group['lr'] = lr
|
||||
if cfg_decoder["use_discriminator"]:
|
||||
for param_group in opt_d.param_groups: param_group['lr'] = lr / 2
|
||||
|
||||
opt_g.step()
|
||||
|
||||
end = time.time()
|
||||
dt = (end-start)*1000
|
||||
|
||||
tqdm_log = f'mel: {loss_mel.item():.3f} | gen: {loss_gen.item():.3f} | sc: {sc_loss.item():.3f} | mag: {mag_loss.item():.3f} | fm: {loss_fm.item():.3f} | d: {d_loss_item:.3f} | lr: {lr:.2e} | time: {dt:.2f} ms'
|
||||
pbar.set_description(tqdm_log)
|
||||
|
||||
# WandB Logging
|
||||
log_dict = {
|
||||
"train/loss_mel": loss_mel.item(),
|
||||
"train/loss_gen": loss_gen.item(),
|
||||
"train/loss_fm": loss_fm.item(),
|
||||
"train/loss_d": d_loss_item,
|
||||
"train/lr": lr,
|
||||
"train/total_loss_g": total_loss_g.item(),
|
||||
"train/loss_sc": sc_loss.item(),
|
||||
"train/loss_mag": mag_loss.item()
|
||||
}
|
||||
|
||||
# ---------------------
|
||||
# Validation Loop
|
||||
# ---------------------
|
||||
if step % val_freq == 0:
|
||||
def evaluate(step, val_dataloader_it, val_dataloader, model, decoder, discriminator,
|
||||
mel_fn, mr_stft, use_disc, device, device_type, val_steps, segment_size, use_wandb):
|
||||
decoder.eval()
|
||||
if discriminator: discriminator.eval()
|
||||
if use_disc and discriminator is not None:
|
||||
discriminator.eval()
|
||||
|
||||
val_mel_loss_accum = 0.0
|
||||
val_gen_loss_accum = 0.0
|
||||
@@ -442,7 +99,8 @@ if __name__ == '__main__':
|
||||
val_d_loss_accum = 0.0
|
||||
val_sc_loss_accum = 0.0
|
||||
val_mag_loss_accum = 0.0
|
||||
val_steps = 10
|
||||
|
||||
log_dict = {}
|
||||
|
||||
with torch.no_grad():
|
||||
for _ in range(val_steps):
|
||||
@@ -497,6 +155,7 @@ if __name__ == '__main__':
|
||||
|
||||
v_mel_loss_raw = torch.nn.functional.l1_loss(v_pred_mel, v_gt_mel, reduction='none')
|
||||
v_mel_loss = (v_mel_loss_raw * v_mel_mask.unsqueeze(1)).sum() / (v_mel_mask.sum() * v_pred_mel.size(1) + 1e-6)
|
||||
|
||||
val_mel_loss_accum += v_mel_loss.item()
|
||||
|
||||
v_sample_mask = v_audio_loss_mask.repeat_interleave(SAMPLES_PER_TOKEN, dim=1)[:, :min_len_v]
|
||||
@@ -504,7 +163,7 @@ if __name__ == '__main__':
|
||||
val_sc_loss_accum += v_sc_loss.item()
|
||||
val_mag_loss_accum += v_mag_loss.item()
|
||||
|
||||
if cfg_decoder["use_discriminator"]:
|
||||
if use_disc and discriminator is not None:
|
||||
v_real_crop_list = []
|
||||
v_fake_crop_list = []
|
||||
v_min_len = min(v_fake_audio.size(1), v_real_audio.size(1))
|
||||
@@ -513,14 +172,14 @@ if __name__ == '__main__':
|
||||
v_valid_len = v_gathered_states_list[b_idx].size(0) * SAMPLES_PER_TOKEN
|
||||
v_valid_len = min(v_valid_len, v_min_len)
|
||||
|
||||
if v_valid_len <= segment_size_samples:
|
||||
v_pad_len = segment_size_samples - v_valid_len
|
||||
if v_valid_len <= segment_size:
|
||||
v_pad_len = segment_size - v_valid_len
|
||||
vr_c = torch.nn.functional.pad(v_real_audio[b_idx, :v_valid_len], (0, v_pad_len))
|
||||
vf_c = torch.nn.functional.pad(v_fake_audio[b_idx, :v_valid_len], (0, v_pad_len))
|
||||
else:
|
||||
v_start_idx = random.randint(0, v_valid_len - segment_size_samples)
|
||||
vr_c = v_real_audio[b_idx, v_start_idx : v_start_idx + segment_size_samples]
|
||||
vf_c = v_fake_audio[b_idx, v_start_idx : v_start_idx + segment_size_samples]
|
||||
v_start_idx = random.randint(0, v_valid_len - segment_size)
|
||||
vr_c = v_real_audio[b_idx, v_start_idx : v_start_idx + segment_size]
|
||||
vf_c = v_fake_audio[b_idx, v_start_idx : v_start_idx + segment_size]
|
||||
|
||||
v_real_crop_list.append(vr_c)
|
||||
v_fake_crop_list.append(vf_c)
|
||||
@@ -537,19 +196,21 @@ if __name__ == '__main__':
|
||||
val_fm_loss_accum += v_fm_loss.item()
|
||||
val_d_loss_accum += v_d_loss.item()
|
||||
|
||||
# Average metrics
|
||||
val_log = {
|
||||
log_dict.update({
|
||||
"val/loss_mel": val_mel_loss_accum / val_steps,
|
||||
"val/loss_sc": val_sc_loss_accum / val_steps,
|
||||
"val/loss_mag": val_mag_loss_accum / val_steps
|
||||
})
|
||||
|
||||
if use_disc and discriminator is not None:
|
||||
log_dict.update({
|
||||
"val/loss_gen": val_gen_loss_accum / val_steps,
|
||||
"val/loss_fm": val_fm_loss_accum / val_steps,
|
||||
"val/loss_d": val_d_loss_accum / val_steps,
|
||||
"val/loss_sc": val_sc_loss_accum / val_steps,
|
||||
"val/loss_mag": val_mag_loss_accum / val_steps
|
||||
}
|
||||
log_dict.update(val_log)
|
||||
})
|
||||
|
||||
if use_wandb:
|
||||
# Generate Mel Images (from last val batch)
|
||||
if cfg_global["use_wandb"]:
|
||||
gen_mel = mel_fn(v_fake_audio[0:1]).squeeze(0).cpu().numpy()
|
||||
gt_mel = mel_fn(v_real_audio[0:1]).squeeze(0).cpu().numpy()
|
||||
|
||||
@@ -564,9 +225,352 @@ if __name__ == '__main__':
|
||||
plt.close(fig)
|
||||
|
||||
decoder.train()
|
||||
if discriminator: discriminator.train()
|
||||
if use_disc and discriminator is not None:
|
||||
discriminator.train()
|
||||
|
||||
return log_dict, val_dataloader_it
|
||||
|
||||
|
||||
def main():
|
||||
config = load_config("config.yaml")
|
||||
cfg_global = config["global"]
|
||||
cfg_paths = config["paths"]
|
||||
cfg_decoder = config["decoder"]
|
||||
|
||||
device = cfg_global["device"]
|
||||
seed = cfg_global["seed"]
|
||||
device_type = "cuda" if device.startswith("cuda") else "cpu"
|
||||
tokenizer_name = cfg_global.get("tokenizer_name", "ekwek/Soprano-80M")
|
||||
use_wandb = cfg_global["use_wandb"]
|
||||
|
||||
torch.manual_seed(seed)
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.manual_seed(seed)
|
||||
torch.set_float32_matmul_precision('high')
|
||||
|
||||
train_dataset_path = os.path.join(cfg_paths["dataset_root"], "train.json")
|
||||
val_dataset_path = os.path.join(cfg_paths["dataset_root"], "val.json")
|
||||
save_path = os.path.join(cfg_paths["save_dir"], "decoder")
|
||||
os.makedirs(save_path, exist_ok=True)
|
||||
print(f"Save Path: {save_path}")
|
||||
|
||||
if use_wandb:
|
||||
wandb.init(project=cfg_global["wandb_project"], config=config)
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
|
||||
tokenizer.padding_side = 'right'
|
||||
|
||||
mel_fn = MelSpectrogramWrapper().to(device)
|
||||
mr_stft = MultiResolutionSTFTLoss().to(device)
|
||||
|
||||
# ------------------
|
||||
# Hyperparameters
|
||||
# ------------------
|
||||
max_steps = cfg_decoder["max_steps"]
|
||||
max_lr = float(cfg_decoder["max_lr"])
|
||||
min_lr = cfg_decoder["min_lr_ratio"] * max_lr
|
||||
warmup_steps = int(max_steps * cfg_decoder["warmup_ratio"])
|
||||
cooldown_steps = int(max_steps * cfg_decoder["cooldown_ratio"])
|
||||
|
||||
batch_size = cfg_decoder["batch_size"]
|
||||
segment_size_samples = cfg_decoder["segment_size_samples"]
|
||||
val_freq = cfg_decoder["val_freq"]
|
||||
val_steps = cfg_decoder.get("val_steps", 10)
|
||||
save_freq = cfg_decoder["save_freq"]
|
||||
betas = tuple(cfg_decoder["betas"])
|
||||
weight_decay = cfg_decoder["weight_decay"]
|
||||
start_step = cfg_decoder.get("start_step", 0)
|
||||
use_disc = cfg_decoder["use_discriminator"]
|
||||
|
||||
lambda_mel = cfg_decoder["lambda_mel"]
|
||||
lambda_fm = cfg_decoder["lambda_fm"]
|
||||
lambda_gen = cfg_decoder["lambda_gen"]
|
||||
lambda_stft = cfg_decoder["lambda_stft"]
|
||||
|
||||
# ------------------
|
||||
# 1. Load LLM and Freeze
|
||||
# ------------------
|
||||
print("Loading LLM...")
|
||||
llm_config = AutoConfig.from_pretrained(tokenizer_name)
|
||||
model = AutoModelForCausalLM.from_config(llm_config)
|
||||
|
||||
pretrained_llm_path = cfg_paths["pretrained_llm_path"]
|
||||
if pretrained_llm_path and os.path.exists(pretrained_llm_path):
|
||||
if pretrained_llm_path.endswith('.safetensors'):
|
||||
state_dict = load_file(pretrained_llm_path)
|
||||
model.load_state_dict(state_dict)
|
||||
else:
|
||||
model = AutoModelForCausalLM.from_pretrained(pretrained_llm_path)
|
||||
else:
|
||||
print("Warning: Training Decoder without a pre-trained LLM. Make sure this is intended.")
|
||||
|
||||
model.to(torch.bfloat16).to(device)
|
||||
model.eval()
|
||||
for param in model.parameters():
|
||||
param.requires_grad = False
|
||||
print("LLM Frozen.")
|
||||
|
||||
# ------------------
|
||||
# 2. Load Decoder
|
||||
# ------------------
|
||||
decoder = SopranoDecoder()
|
||||
pretrained_decoder_path = cfg_paths["pretrained_decoder_path"]
|
||||
|
||||
if pretrained_decoder_path and os.path.exists(pretrained_decoder_path):
|
||||
print(f"Loading custom Decoder checkpoint from {pretrained_decoder_path}")
|
||||
decoder.load_state_dict(torch.load(pretrained_decoder_path, map_location='cpu'))
|
||||
else:
|
||||
print("Training Decoder from scratch.")
|
||||
|
||||
decoder.to(device)
|
||||
decoder.train()
|
||||
|
||||
# ------------------
|
||||
# 3. Load Discriminator
|
||||
# ------------------
|
||||
discriminator = None
|
||||
if use_disc:
|
||||
print("Initializing Discriminator...")
|
||||
discriminator = Discriminator()
|
||||
pretrained_disc_path = cfg_paths["pretrained_discriminator_path"]
|
||||
|
||||
if pretrained_disc_path and os.path.exists(pretrained_disc_path):
|
||||
print(f"Loading custom Discriminator checkpoint from {pretrained_disc_path}")
|
||||
discriminator.load_state_dict(torch.load(pretrained_disc_path, map_location='cpu'))
|
||||
else:
|
||||
print("Training Discriminator from scratch.")
|
||||
|
||||
discriminator.to(device)
|
||||
discriminator.train()
|
||||
else:
|
||||
print("Training WITHOUT Discriminator (Reconstruction only).")
|
||||
|
||||
# ------------------
|
||||
# 4. Dataset Setup
|
||||
# ------------------
|
||||
dataset = AudioDataset(train_dataset_path)
|
||||
dataloader = DataLoader(
|
||||
dataset,
|
||||
batch_size=batch_size,
|
||||
shuffle=True,
|
||||
num_workers=cfg_global["num_workers"],
|
||||
pin_memory=True,
|
||||
persistent_workers=True,
|
||||
worker_init_fn=worker_seed_init,
|
||||
collate_fn=lambda batch_in: collate_pack(batch_in, tokenizer),
|
||||
)
|
||||
dataloader_it = iter(dataloader)
|
||||
|
||||
val_dataset = AudioDataset(val_dataset_path)
|
||||
val_dataloader = DataLoader(
|
||||
val_dataset,
|
||||
batch_size=max(1, batch_size // 4),
|
||||
shuffle=False,
|
||||
num_workers=max(1, cfg_global["num_workers"] // 2),
|
||||
pin_memory=True,
|
||||
persistent_workers=True,
|
||||
worker_init_fn=worker_seed_init,
|
||||
collate_fn=lambda batch_in: collate_pack(batch_in, tokenizer),
|
||||
)
|
||||
val_dataloader_it = iter(val_dataloader)
|
||||
|
||||
opt_g = torch.optim.AdamW(decoder.parameters(), max_lr, betas=betas, weight_decay=weight_decay)
|
||||
opt_d = None
|
||||
if use_disc:
|
||||
opt_d = torch.optim.AdamW(discriminator.parameters(), max_lr, betas=betas, weight_decay=weight_decay)
|
||||
|
||||
# ------------------
|
||||
# Training Loop
|
||||
# ------------------
|
||||
pbar = tqdm(range(start_step, max_steps), ncols=200, dynamic_ncols=True)
|
||||
|
||||
for step in pbar:
|
||||
start = time.time()
|
||||
|
||||
try:
|
||||
batch_data = next(dataloader_it)
|
||||
if batch_data[0] is None:
|
||||
dataloader_it = iter(dataloader)
|
||||
batch_data = next(dataloader_it)
|
||||
x, y, gt_audio, audio_mask = batch_data
|
||||
except StopIteration:
|
||||
dataloader_it = iter(dataloader)
|
||||
batch_data = next(dataloader_it)
|
||||
x, y, gt_audio, audio_mask = batch_data
|
||||
|
||||
x, y = x.to(device), y.to(device)
|
||||
gt_audio = gt_audio.to(device)
|
||||
audio_mask = audio_mask.to(device)
|
||||
|
||||
with torch.no_grad():
|
||||
with torch.autocast(device_type=device_type, dtype=torch.bfloat16):
|
||||
outputs = model(x, output_hidden_states=True)
|
||||
hidden_states = outputs.hidden_states[-1]
|
||||
hidden_states = hidden_states.to(torch.float32)
|
||||
|
||||
gathered_states_list = []
|
||||
for b_idx in range(hidden_states.size(0)):
|
||||
mask = audio_mask[b_idx]
|
||||
valid_states = hidden_states[b_idx][mask]
|
||||
gathered_states_list.append(valid_states)
|
||||
|
||||
decoder_in_padded = torch.nn.utils.rnn.pad_sequence(gathered_states_list, batch_first=True, padding_value=0.0)
|
||||
|
||||
bsz = decoder_in_padded.size(0)
|
||||
max_aud_len = decoder_in_padded.size(1)
|
||||
audio_loss_mask = torch.zeros((bsz, max_aud_len), dtype=torch.bool, device=device)
|
||||
for b_idx in range(bsz):
|
||||
length = gathered_states_list[b_idx].size(0)
|
||||
audio_loss_mask[b_idx, :length] = True
|
||||
|
||||
d_loss_item = 0.0
|
||||
|
||||
if use_disc:
|
||||
opt_d.zero_grad()
|
||||
|
||||
decoder_in = decoder_in_padded.transpose(1, 2)
|
||||
fake_audio = decoder(decoder_in)
|
||||
if fake_audio.size(1) == 1: fake_audio = fake_audio.squeeze(1)
|
||||
|
||||
min_len = min(fake_audio.size(1), gt_audio.size(1))
|
||||
fake_audio = fake_audio[:, :min_len]
|
||||
real_audio = gt_audio[:, :min_len]
|
||||
|
||||
real_crop_list = []
|
||||
fake_crop_list = []
|
||||
|
||||
for b_idx in range(bsz):
|
||||
valid_len = gathered_states_list[b_idx].size(0) * SAMPLES_PER_TOKEN
|
||||
valid_len = min(valid_len, min_len)
|
||||
|
||||
if valid_len <= segment_size_samples:
|
||||
pad_len = segment_size_samples - valid_len
|
||||
r_c = torch.nn.functional.pad(real_audio[b_idx, :valid_len], (0, pad_len))
|
||||
f_c = torch.nn.functional.pad(fake_audio[b_idx, :valid_len], (0, pad_len))
|
||||
else:
|
||||
start_idx = random.randint(0, valid_len - segment_size_samples)
|
||||
r_c = real_audio[b_idx, start_idx : start_idx + segment_size_samples]
|
||||
f_c = fake_audio[b_idx, start_idx : start_idx + segment_size_samples]
|
||||
|
||||
real_crop_list.append(r_c)
|
||||
fake_crop_list.append(f_c)
|
||||
|
||||
real_crops = torch.stack(real_crop_list).unsqueeze(1)
|
||||
fake_crops = torch.stack(fake_crop_list).unsqueeze(1).detach()
|
||||
|
||||
y_d_rs, y_d_gs, _, _ = discriminator(real_crops, fake_crops)
|
||||
d_loss, _, _ = discriminator_loss(y_d_rs, y_d_gs)
|
||||
|
||||
d_loss.backward()
|
||||
torch.nn.utils.clip_grad_norm_(discriminator.parameters(), 1.0)
|
||||
opt_d.step()
|
||||
d_loss_item = d_loss.item()
|
||||
|
||||
opt_g.zero_grad()
|
||||
|
||||
decoder_in = decoder_in_padded.transpose(1, 2)
|
||||
fake_audio = decoder(decoder_in)
|
||||
if fake_audio.size(1) == 1: fake_audio = fake_audio.squeeze(1)
|
||||
|
||||
min_len = min(fake_audio.size(1), gt_audio.size(1))
|
||||
fake_audio = fake_audio[:, :min_len]
|
||||
real_audio = gt_audio[:, :min_len]
|
||||
|
||||
real_crop_list_g = []
|
||||
fake_crop_list_g = []
|
||||
if use_disc:
|
||||
for b_idx in range(bsz):
|
||||
valid_len = gathered_states_list[b_idx].size(0) * SAMPLES_PER_TOKEN
|
||||
valid_len = min(valid_len, min_len)
|
||||
|
||||
if valid_len <= segment_size_samples:
|
||||
pad_len = segment_size_samples - valid_len
|
||||
r_c = torch.nn.functional.pad(real_audio[b_idx, :valid_len], (0, pad_len))
|
||||
f_c = torch.nn.functional.pad(fake_audio[b_idx, :valid_len], (0, pad_len))
|
||||
else:
|
||||
start_idx = random.randint(0, valid_len - segment_size_samples)
|
||||
r_c = real_audio[b_idx, start_idx : start_idx + segment_size_samples]
|
||||
f_c = fake_audio[b_idx, start_idx : start_idx + segment_size_samples]
|
||||
|
||||
real_crop_list_g.append(r_c)
|
||||
fake_crop_list_g.append(f_c)
|
||||
|
||||
real_crops_g = torch.stack(real_crop_list_g).unsqueeze(1)
|
||||
fake_crops_g = torch.stack(fake_crop_list_g).unsqueeze(1)
|
||||
|
||||
frames_per_token = SAMPLES_PER_TOKEN // 512
|
||||
mel_mask = audio_loss_mask.repeat_interleave(frames_per_token, dim=1)
|
||||
|
||||
pred_mel = mel_fn(fake_audio)
|
||||
gt_mel = mel_fn(real_audio)
|
||||
|
||||
min_mel_len = min(pred_mel.size(2), gt_mel.size(2), mel_mask.size(1))
|
||||
pred_mel = pred_mel[:, :, :min_mel_len]
|
||||
gt_mel = gt_mel[:, :, :min_mel_len]
|
||||
mel_mask = mel_mask[:, :min_mel_len]
|
||||
|
||||
loss_mel_raw = torch.nn.functional.l1_loss(pred_mel, gt_mel, reduction='none')
|
||||
loss_mel = (loss_mel_raw * mel_mask.unsqueeze(1)).sum() / (mel_mask.sum() * pred_mel.size(1) + 1e-6)
|
||||
|
||||
sample_mask = audio_loss_mask.repeat_interleave(SAMPLES_PER_TOKEN, dim=1)[:, :min_len]
|
||||
sc_loss, mag_loss = mr_stft(fake_audio * sample_mask, real_audio * sample_mask)
|
||||
|
||||
loss_fm = torch.tensor(0.0, device=device)
|
||||
loss_gen = torch.tensor(0.0, device=device)
|
||||
|
||||
if use_disc:
|
||||
y_d_rs, y_d_gs, fmap_rs, fmap_gs = discriminator(real_crops_g, fake_crops_g)
|
||||
loss_fm = feature_matching_loss(fmap_rs, fmap_gs)
|
||||
loss_gen, _ = generator_loss(y_d_gs)
|
||||
|
||||
total_loss_g = (lambda_mel * loss_mel) + (lambda_gen * loss_gen) + (lambda_fm * loss_fm) + (lambda_stft * (sc_loss + mag_loss))
|
||||
|
||||
total_loss_g.backward()
|
||||
norm_g = torch.nn.utils.clip_grad_norm_(decoder.parameters(), 1.0)
|
||||
|
||||
lr = get_lr(step, max_lr, min_lr, warmup_steps, cooldown_steps, max_steps)
|
||||
for param_group in opt_g.param_groups: param_group['lr'] = lr
|
||||
if use_disc:
|
||||
for param_group in opt_d.param_groups: param_group['lr'] = lr / 2
|
||||
|
||||
opt_g.step()
|
||||
|
||||
end = time.time()
|
||||
dt = (end-start)*1000
|
||||
|
||||
tqdm_log = f'mel: {loss_mel.item():.3f} | gen: {loss_gen.item():.3f} | sc: {sc_loss.item():.3f} | mag: {mag_loss.item():.3f} | fm: {loss_fm.item():.3f} | d: {d_loss_item:.3f} | lr: {lr:.2e} | time: {dt:.2f} ms'
|
||||
pbar.set_description(tqdm_log)
|
||||
|
||||
log_dict = {
|
||||
"train/loss_mel": loss_mel.item(),
|
||||
"train/loss_gen": loss_gen.item(),
|
||||
"train/loss_fm": loss_fm.item(),
|
||||
"train/loss_d": d_loss_item,
|
||||
"train/lr": lr,
|
||||
"train/total_loss_g": total_loss_g.item(),
|
||||
"train/loss_sc": sc_loss.item(),
|
||||
"train/loss_mag": mag_loss.item()
|
||||
}
|
||||
|
||||
if step > 0 and step % val_freq == 0:
|
||||
val_log_dict, val_dataloader_it = evaluate(
|
||||
step=step,
|
||||
val_dataloader_it=val_dataloader_it,
|
||||
val_dataloader=val_dataloader,
|
||||
model=model,
|
||||
decoder=decoder,
|
||||
discriminator=discriminator,
|
||||
mel_fn=mel_fn,
|
||||
mr_stft=mr_stft,
|
||||
use_disc=use_disc,
|
||||
device=device,
|
||||
device_type=device_type,
|
||||
val_steps=val_steps,
|
||||
segment_size=segment_size_samples,
|
||||
use_wandb=use_wandb
|
||||
)
|
||||
log_dict.update(val_log_dict)
|
||||
|
||||
# Save Checkpoint
|
||||
if step > 0 and step % save_freq == 0:
|
||||
print(f"\nSaving checkpoint at step {step} to {save_path}...")
|
||||
ckpt_name_dec = f"decoder_step_{step}.pth"
|
||||
@@ -575,13 +579,17 @@ if __name__ == '__main__':
|
||||
if discriminator:
|
||||
torch.save(discriminator.state_dict(), os.path.join(save_path, ckpt_name_disc))
|
||||
|
||||
if cfg_global["use_wandb"]:
|
||||
if use_wandb:
|
||||
wandb.log(log_dict, step=step)
|
||||
|
||||
print(f"Training complete. Saving model at {save_path}")
|
||||
print(f"\nTraining complete. Saving model at {save_path}")
|
||||
torch.save(decoder.state_dict(), os.path.join(save_path, "decoder_trained.pth"))
|
||||
if discriminator:
|
||||
torch.save(discriminator.state_dict(), os.path.join(save_path, "discriminator_trained.pth"))
|
||||
|
||||
if cfg_global["use_wandb"]:
|
||||
if use_wandb:
|
||||
wandb.finish()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
+59
-44
@@ -3,8 +3,6 @@ Training script for Soprano LLM backbone.
|
||||
|
||||
Usage:
|
||||
python train_llm.py
|
||||
|
||||
Adapted from https://github.com/karpathy/nanoGPT
|
||||
"""
|
||||
import os
|
||||
import random
|
||||
@@ -21,36 +19,72 @@ from safetensors.torch import load_file
|
||||
from dataset import AudioDataset
|
||||
from config_loader import load_config
|
||||
|
||||
# Initialize tokenizer globally so it can be used in the collate functions
|
||||
tokenizer = AutoTokenizer.from_pretrained('ekwek/Soprano-80M')
|
||||
tokenizer.padding_side = 'right' # Essential for training!
|
||||
|
||||
def worker_seed_init(_):
|
||||
worker_seed = torch.initial_seed() % (2**32-1)
|
||||
np.random.seed(worker_seed)
|
||||
random.seed(worker_seed)
|
||||
|
||||
def get_lr(it, max_lr, min_lr, warmup_steps, cooldown_steps, max_steps): # WSD schedule
|
||||
|
||||
def get_lr(it, max_lr, min_lr, warmup_steps, cooldown_steps, max_steps):
|
||||
if it < warmup_steps:
|
||||
return max_lr * (it + 1) / warmup_steps
|
||||
if it < max_steps - cooldown_steps:
|
||||
return max_lr
|
||||
return min_lr + (max_lr - min_lr) * ((max_steps - it) / cooldown_steps)
|
||||
|
||||
def collate_dynamic(texts):
|
||||
# Dynamic Batching: Pad to the longest in this batch (max 2048 safety)
|
||||
|
||||
def collate_pack(texts, tokenizer, seq_len, batch_size):
|
||||
tokens_batch = tokenizer(texts, padding=False, truncation=False)
|
||||
batch = []
|
||||
cur_sample, cur_size = [], 0
|
||||
for i in range(len(texts)):
|
||||
tokens = torch.tensor(tokens_batch['input_ids'][i][:-1], dtype=torch.long)
|
||||
cur_size += tokens.size(0)
|
||||
cur_sample.append(tokens)
|
||||
if cur_size >= seq_len + 1:
|
||||
batch.append(torch.cat(cur_sample)[: seq_len + 1])
|
||||
cur_sample, cur_size = [], 0
|
||||
if len(batch) == batch_size:
|
||||
break
|
||||
if cur_sample and not batch:
|
||||
batch_item = torch.cat(cur_sample + [torch.zeros(seq_len, dtype=torch.long)])[: seq_len + 1]
|
||||
batch.append(batch_item)
|
||||
if len(batch) < batch_size:
|
||||
pad = batch[-1]
|
||||
while len(batch) < batch_size:
|
||||
batch.append(pad)
|
||||
batch = torch.stack(batch)
|
||||
x = batch[:, :-1]
|
||||
y = batch[:, 1:]
|
||||
return x, y
|
||||
|
||||
|
||||
def collate_dynamic(texts, tokenizer):
|
||||
tokenized = tokenizer(texts, padding=True, truncation=True, max_length=2048, return_tensors='pt', add_special_tokens=False)
|
||||
batch = tokenized['input_ids']
|
||||
attn_mask = tokenized['attention_mask']
|
||||
|
||||
x = batch[:, :-1]
|
||||
y = batch[:, 1:]
|
||||
# Attention mask needs to align with x. Since we shift x by removing the last token,
|
||||
# we should also remove the last token from the mask.
|
||||
attn_mask = attn_mask[:, :-1]
|
||||
|
||||
return x, y, attn_mask
|
||||
|
||||
|
||||
def collate_pack_val(texts, tokenizer, seq_len):
|
||||
out = tokenizer(texts, padding=True, truncation=True, max_length=seq_len+1, return_tensors='pt', add_special_tokens=False)
|
||||
batch = out['input_ids']
|
||||
|
||||
if batch.size(1) < seq_len + 1:
|
||||
pad_len = seq_len + 1 - batch.size(1)
|
||||
batch = torch.nn.functional.pad(batch, (0, pad_len), value=tokenizer.pad_token_id)
|
||||
|
||||
x = batch[:, :-1]
|
||||
y = batch[:, 1:]
|
||||
return x, y
|
||||
|
||||
|
||||
def compute_loss(x, logits, y, num_steps, mask=None):
|
||||
pred = logits.view(-1, logits.size(-1))
|
||||
labels = y.reshape(-1)
|
||||
@@ -60,28 +94,20 @@ def compute_loss(x, logits, y, num_steps, mask=None):
|
||||
mask = mask.reshape(-1)
|
||||
loss = loss * mask
|
||||
|
||||
# Audio tokens: >=3 and <=8003.
|
||||
# NOTE: If [STOP] is 3, it counts as audio.
|
||||
# We apply the mask to filter out padding.
|
||||
audio_mask_cond = torch.logical_and(labels >= 3, labels <= 8003)
|
||||
if mask is not None:
|
||||
audio_mask = audio_mask_cond & (mask > 0)
|
||||
else:
|
||||
audio_mask = audio_mask_cond
|
||||
|
||||
# Text tokens: The rest, BUT excluding masked (padding) tokens
|
||||
if mask is not None:
|
||||
text_mask = (~audio_mask_cond) & (mask > 0)
|
||||
else:
|
||||
text_mask = ~audio_mask_cond
|
||||
|
||||
# Avoid division by zero
|
||||
audio_mean = loss[audio_mask].mean() if audio_mask.sum() > 0 else torch.tensor(0.0, device=loss.device)
|
||||
text_mean = loss[text_mask].mean() if text_mask.sum() > 0 else torch.tensor(0.0, device=loss.device)
|
||||
|
||||
# Acc: only on non-masked tokens.
|
||||
# Current logic: (logits.argmax(dim=-1) == y).view(-1)[audio_mask]
|
||||
# This correctly calculates accuracy only on valid audio tokens.
|
||||
acc = (logits.argmax(dim=-1).view(-1) == labels).view(-1)[audio_mask].to(torch.float32).mean()
|
||||
if torch.isnan(acc): acc = torch.tensor(0.0, device=loss.device)
|
||||
|
||||
@@ -90,6 +116,7 @@ def compute_loss(x, logits, y, num_steps, mask=None):
|
||||
acc = acc / num_steps
|
||||
return audio_loss, text_loss, acc
|
||||
|
||||
|
||||
def evaluate(model, val_dataloader, step, device, use_wandb):
|
||||
model.eval()
|
||||
val_dataloader_it = iter(val_dataloader)
|
||||
@@ -122,10 +149,7 @@ def evaluate(model, val_dataloader, step, device, use_wandb):
|
||||
model.train()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
# ------------------
|
||||
# Load Configuration
|
||||
# ------------------
|
||||
def main():
|
||||
config = load_config("config.yaml")
|
||||
cfg_global = config["global"]
|
||||
cfg_paths = config["paths"]
|
||||
@@ -134,25 +158,26 @@ if __name__ == '__main__':
|
||||
device = cfg_global["device"]
|
||||
seed = cfg_global["seed"]
|
||||
device_type = "cuda" if device.startswith("cuda") else "cpu"
|
||||
tokenizer_name = cfg_global.get("tokenizer_name", "ekwek/Soprano-80M")
|
||||
|
||||
torch.manual_seed(seed)
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.manual_seed(seed)
|
||||
torch.set_float32_matmul_precision('high')
|
||||
|
||||
# Setup directories
|
||||
train_dataset_path = os.path.join(cfg_paths["dataset_root"], "train.json")
|
||||
val_dataset_path = os.path.join(cfg_paths["dataset_root"], "val.json")
|
||||
save_path = os.path.join(cfg_paths["save_dir"], "llm")
|
||||
os.makedirs(save_path, exist_ok=True)
|
||||
|
||||
print(f"Save Path: {save_path}")
|
||||
|
||||
if cfg_global["use_wandb"]:
|
||||
wandb.init(project=cfg_global["wandb_project"], config=config)
|
||||
|
||||
# ------------------
|
||||
# Hyperparameters
|
||||
# ------------------
|
||||
tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
|
||||
tokenizer.padding_side = 'right'
|
||||
|
||||
max_steps = cfg_llm["max_steps"]
|
||||
max_lr = float(cfg_llm["max_lr"])
|
||||
min_lr = cfg_llm["min_lr_ratio"] * max_lr
|
||||
@@ -168,31 +193,25 @@ if __name__ == '__main__':
|
||||
betas = tuple(cfg_llm["betas"])
|
||||
weight_decay = cfg_llm["weight_decay"]
|
||||
|
||||
# ------------------
|
||||
# Model Setup
|
||||
# ------------------
|
||||
if cfg_llm["from_scratch"]:
|
||||
print("Initializing model from scratch (random weights)...")
|
||||
m_config = AutoConfig.from_pretrained('ekwek/Soprano-80M')
|
||||
m_config = AutoConfig.from_pretrained(tokenizer_name)
|
||||
model = AutoModelForCausalLM.from_config(m_config)
|
||||
else:
|
||||
pretrained_path = cfg_paths["pretrained_llm_path"]
|
||||
if pretrained_path and os.path.exists(pretrained_path):
|
||||
print(f"Loading pretrained model weights from {pretrained_path}...")
|
||||
m_config = AutoConfig.from_pretrained('ekwek/Soprano-80M')
|
||||
m_config = AutoConfig.from_pretrained(tokenizer_name)
|
||||
model = AutoModelForCausalLM.from_config(m_config)
|
||||
state_dict = load_file(pretrained_path)
|
||||
model.load_state_dict(state_dict)
|
||||
else:
|
||||
print("Loading default pretrained model weights from HF...")
|
||||
model = AutoModelForCausalLM.from_pretrained('ekwek/Soprano-80M')
|
||||
model = AutoModelForCausalLM.from_pretrained(tokenizer_name)
|
||||
|
||||
model.to(device)
|
||||
model.train()
|
||||
|
||||
# ------------------
|
||||
# Dataset Setup
|
||||
# ------------------
|
||||
dataset = AudioDataset(train_dataset_path)
|
||||
dataloader = DataLoader(
|
||||
dataset,
|
||||
@@ -202,7 +221,7 @@ if __name__ == '__main__':
|
||||
pin_memory=True,
|
||||
persistent_workers=True,
|
||||
worker_init_fn=worker_seed_init,
|
||||
collate_fn=collate_dynamic,
|
||||
collate_fn=lambda texts: collate_dynamic(texts, tokenizer),
|
||||
)
|
||||
dataloader_it = iter(dataloader)
|
||||
|
||||
@@ -215,18 +234,11 @@ if __name__ == '__main__':
|
||||
pin_memory=True,
|
||||
persistent_workers=True,
|
||||
worker_init_fn=worker_seed_init,
|
||||
collate_fn=collate_dynamic,
|
||||
collate_fn=lambda texts: collate_dynamic(texts, tokenizer),
|
||||
)
|
||||
|
||||
# ------------------
|
||||
# Optimizer
|
||||
# ------------------
|
||||
opt = torch.optim.AdamW(model.parameters(), max_lr, betas=betas, weight_decay=weight_decay, fused=True)
|
||||
|
||||
# ------------------
|
||||
# Training Loop
|
||||
# ------------------
|
||||
# Determine start step based on loaded checkpoint if needed, defaulting to 1 for new runs
|
||||
start_step = 1
|
||||
pbar = tqdm(range(start_step, max_steps + 1), ncols=200, dynamic_ncols=True)
|
||||
|
||||
@@ -301,3 +313,6 @@ if __name__ == '__main__':
|
||||
|
||||
if cfg_global["use_wandb"]:
|
||||
wandb.finish()
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
Reference in New Issue
Block a user