Add LoRA fine-tuning scripts (#5)
This commit is contained in:
mrfakename
@@ -0,0 +1 @@
|
||||
# VibeVoice finetuning module
|
||||
@@ -0,0 +1,362 @@
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import warnings
|
||||
import random
|
||||
|
||||
try:
|
||||
import librosa # type: ignore
|
||||
except Exception: # pragma: no cover
|
||||
librosa = None # Fallback: user must install librosa when using local audio paths
|
||||
|
||||
try:
|
||||
import resampy # type: ignore
|
||||
except Exception: # pragma: no cover
|
||||
resampy = None
|
||||
|
||||
|
||||
def _resample_if_needed(wav: np.ndarray, orig_sr: int, target_sr: int) -> np.ndarray:
|
||||
if orig_sr == target_sr:
|
||||
return wav.astype(np.float32, copy=False)
|
||||
if resampy is not None:
|
||||
return resampy.resample(wav.astype(np.float32), orig_sr, target_sr)
|
||||
if librosa is not None:
|
||||
return librosa.resample(y=wav.astype(np.float32), orig_sr=orig_sr, target_sr=target_sr)
|
||||
warnings.warn(
|
||||
"No resampler available; treating audio as target_sr without resampling. Install resampy or librosa.",
|
||||
RuntimeWarning,
|
||||
)
|
||||
return wav.astype(np.float32, copy=False)
|
||||
|
||||
|
||||
# Lightweight HF-style dataset wrapper (optional). Trainer can also pass raw HF datasets directly.
|
||||
class VibeVoiceDataset:
|
||||
def __init__(
|
||||
self,
|
||||
dataset: Any,
|
||||
text_column: str = "text",
|
||||
audio_column: str = "audio",
|
||||
voice_prompts_column: Optional[str] = "voice_prompts",
|
||||
) -> None:
|
||||
self.dataset = dataset
|
||||
self.text_column = text_column
|
||||
self.audio_column = audio_column
|
||||
self.voice_prompts_column = voice_prompts_column
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.dataset)
|
||||
|
||||
def __getitem__(self, idx: int) -> Dict[str, Any]:
|
||||
item = self.dataset[idx]
|
||||
data: Dict[str, Any] = {}
|
||||
data["text"] = item[self.text_column]
|
||||
data["audio"] = item[self.audio_column]
|
||||
|
||||
user_provided_prompt = None
|
||||
if self.voice_prompts_column and self.voice_prompts_column in item:
|
||||
user_provided_prompt = item[self.voice_prompts_column]
|
||||
|
||||
if user_provided_prompt:
|
||||
# A prompt was provided in the dataset, so we use it.
|
||||
if not isinstance(user_provided_prompt, list):
|
||||
data["voice_prompts"] = [user_provided_prompt]
|
||||
else:
|
||||
data["voice_prompts"] = user_provided_prompt
|
||||
else:
|
||||
# FALLBACK: No prompt provided, so we auto-generate one from the target audio.
|
||||
try:
|
||||
target_sr = 24000
|
||||
wav_array = _load_audio_to_24k(item[self.audio_column], target_sr=target_sr)
|
||||
audio_len_seconds = len(wav_array) / target_sr
|
||||
|
||||
min_len_sec = min(5.0, audio_len_seconds / 4.0)
|
||||
max_len_sec = min(15.0, audio_len_seconds / 2.0)
|
||||
|
||||
if min_len_sec > max_len_sec:
|
||||
min_len_sec = max_len_sec
|
||||
max_len_sec = min(max_len_sec, audio_len_seconds)
|
||||
|
||||
if max_len_sec > 0.1:
|
||||
prompt_len_sec = random.uniform(min_len_sec, max_len_sec)
|
||||
prompt_len_samples = int(prompt_len_sec * target_sr)
|
||||
|
||||
max_start_sample = len(wav_array) - prompt_len_samples
|
||||
start_sample = random.randint(0, max_start_sample)
|
||||
|
||||
prompt_crop = wav_array[start_sample : start_sample + prompt_len_samples]
|
||||
|
||||
data["voice_prompts"] = [prompt_crop]
|
||||
else:
|
||||
data["voice_prompts"] = None
|
||||
|
||||
except Exception as e:
|
||||
warnings.warn(f"Could not create voice prompt for item {idx}: {e}")
|
||||
data["voice_prompts"] = None
|
||||
return data
|
||||
|
||||
|
||||
def _load_audio_to_24k(audio: Union[str, np.ndarray, torch.Tensor, Dict[str, Any]], *, target_sr: int = 24000) -> np.ndarray:
|
||||
if isinstance(audio, np.ndarray):
|
||||
return audio.astype(np.float32)
|
||||
if isinstance(audio, torch.Tensor):
|
||||
return audio.detach().cpu().float().numpy()
|
||||
if isinstance(audio, str):
|
||||
if librosa is None:
|
||||
raise RuntimeError("librosa is required to load audio file paths. Please pip install librosa.")
|
||||
wav, sr = librosa.load(audio, sr=None, mono=True)
|
||||
wav = _resample_if_needed(wav, int(sr), target_sr)
|
||||
return wav
|
||||
if isinstance(audio, dict) and "array" in audio and "sampling_rate" in audio:
|
||||
arr = np.asarray(audio["array"], dtype=np.float32)
|
||||
sr = int(audio["sampling_rate"])
|
||||
arr = _resample_if_needed(arr, sr, target_sr)
|
||||
return arr
|
||||
raise ValueError(f"Unsupported audio type: {type(audio)}")
|
||||
|
||||
|
||||
@dataclass
|
||||
class VibeVoiceCollator:
|
||||
processor: Any # VibeVoiceProcessor
|
||||
max_length: Optional[int] = None
|
||||
speech_compress_ratio: int = 3200
|
||||
semantic_vae_dim: int = 128
|
||||
compute_semantics: bool = False
|
||||
debug_checks: bool = False
|
||||
|
||||
text_field: str = "text"
|
||||
audio_field: str = "audio"
|
||||
voice_prompts_field: str = "voice_prompts"
|
||||
voice_prompt_drop_rate: float = 0.0
|
||||
|
||||
def __call__(self, features: Sequence[Dict[str, Any]]) -> Dict[str, Any]:
|
||||
batch_size = len(features)
|
||||
|
||||
sample_input_ids: List[List[int]] = []
|
||||
sample_attention_masks: List[List[int]] = []
|
||||
sample_acoustic_input_masks: List[List[bool]] = []
|
||||
sample_acoustic_loss_masks: List[List[bool]] = []
|
||||
|
||||
all_speech_waveforms: List[np.ndarray] = []
|
||||
all_speech_latent_lengths: List[int] = []
|
||||
per_segment_is_target: List[bool] = []
|
||||
|
||||
for ex in features:
|
||||
text: str = ex.get(self.text_field, "")
|
||||
voice_prompts: Optional[List[Union[str, np.ndarray, torch.Tensor]]] = ex.get(self.voice_prompts_field)
|
||||
target_audio: Union[str, np.ndarray, torch.Tensor, Dict[str, Any]] = ex.get(self.audio_field)
|
||||
|
||||
# Clamp drop rate for safety
|
||||
_drop_rate = self.voice_prompt_drop_rate
|
||||
if _drop_rate < 0.0:
|
||||
_drop_rate = 0.0
|
||||
elif _drop_rate > 1.0:
|
||||
_drop_rate = 1.0
|
||||
|
||||
proc = self.processor(
|
||||
text=[text],
|
||||
voice_samples=[voice_prompts] if voice_prompts is not None and random.random() >= _drop_rate else None,
|
||||
padding=False,
|
||||
truncation=False,
|
||||
max_length=self.max_length,
|
||||
return_tensors="pt",
|
||||
)
|
||||
|
||||
ids = proc["input_ids"][0].tolist()
|
||||
attn = proc.get("attention_mask", torch.ones_like(proc["input_ids"]))[0].tolist()
|
||||
speech_input_mask = proc.get("speech_input_mask")
|
||||
if speech_input_mask is None:
|
||||
speech_input_mask = torch.zeros_like(proc["input_ids"], dtype=torch.bool)
|
||||
speech_input_mask_list = speech_input_mask[0].tolist()
|
||||
|
||||
wav_target = _load_audio_to_24k(target_audio, target_sr=24000)
|
||||
# Prefer exact frame count from acoustic tokenizer if available; fallback to compress ratio
|
||||
target_latent_len = None
|
||||
try:
|
||||
acoustic_tok = getattr(self.processor, "acoustic_tokenizer", None)
|
||||
if acoustic_tok is not None and hasattr(acoustic_tok, "encode"):
|
||||
enc_out = acoustic_tok.encode(wav_target)
|
||||
# Normalize various possible return formats to get time dimension
|
||||
T = None
|
||||
try:
|
||||
# Direct array-like with shape (T, D) or (T,)
|
||||
if hasattr(enc_out, "shape") and len(getattr(enc_out, "shape", [])) >= 1:
|
||||
T = int(enc_out.shape[0])
|
||||
else:
|
||||
# Nested lists/tuples or ModelOutput-like
|
||||
cand = enc_out
|
||||
# Drill down a couple of levels safely
|
||||
for _ in range(2):
|
||||
if isinstance(cand, (list, tuple)) and len(cand) > 0:
|
||||
cand = cand[0]
|
||||
if hasattr(cand, "shape") and len(getattr(cand, "shape", [])) >= 1:
|
||||
T = int(cand.shape[0])
|
||||
except Exception:
|
||||
T = None
|
||||
if T is not None and T > 0:
|
||||
target_latent_len = T
|
||||
except Exception:
|
||||
target_latent_len = None
|
||||
if target_latent_len is None:
|
||||
target_latent_len = max(1, int(math.ceil(len(wav_target) / float(self.speech_compress_ratio))))
|
||||
|
||||
speech_diff_id = self.processor.tokenizer.speech_diffusion_id
|
||||
target_placeholders = [speech_diff_id] * target_latent_len
|
||||
|
||||
ids_extended = ids + target_placeholders
|
||||
attn_extended = attn + [1] * target_latent_len
|
||||
|
||||
acoustic_input_mask = speech_input_mask_list + [True] * target_latent_len
|
||||
acoustic_loss_mask = ([False] * len(speech_input_mask_list)) + [True] * target_latent_len
|
||||
|
||||
speech_end_id = self.processor.tokenizer.speech_end_id
|
||||
ids_extended.append(speech_end_id)
|
||||
attn_extended.append(1)
|
||||
acoustic_input_mask.append(False)
|
||||
acoustic_loss_mask.append(False)
|
||||
|
||||
if self.max_length is not None and len(ids_extended) > self.max_length:
|
||||
cut = len(ids_extended) - int(self.max_length)
|
||||
leading_non_acoustic = 0
|
||||
for v in acoustic_input_mask:
|
||||
if v:
|
||||
break
|
||||
leading_non_acoustic += 1
|
||||
if cut > leading_non_acoustic:
|
||||
raise ValueError(
|
||||
f"--max_length={self.max_length} would truncate into acoustic tokens. "
|
||||
f"Needed cut={cut}, but only {leading_non_acoustic} leading non-acoustic tokens available. "
|
||||
"Increase max_length or shorten text/voice-prompt preamble."
|
||||
)
|
||||
ids_extended = ids_extended[cut:]
|
||||
attn_extended = attn_extended[cut:]
|
||||
acoustic_input_mask = acoustic_input_mask[cut:]
|
||||
acoustic_loss_mask = acoustic_loss_mask[cut:]
|
||||
|
||||
sample_input_ids.append(ids_extended)
|
||||
sample_attention_masks.append(attn_extended)
|
||||
sample_acoustic_input_masks.append(acoustic_input_mask)
|
||||
sample_acoustic_loss_masks.append(acoustic_loss_mask)
|
||||
|
||||
voice_speeches = []
|
||||
voice_latent_lengths = []
|
||||
if proc.get("speech_tensors") is not None:
|
||||
voice_np = proc["speech_tensors"].cpu().numpy()
|
||||
voice_masks = proc["speech_masks"].cpu().numpy().astype(bool)
|
||||
for seg_idx in range(voice_np.shape[0]):
|
||||
voice_speeches.append(voice_np[seg_idx])
|
||||
voice_latent_lengths.append(int(voice_masks[seg_idx].sum()))
|
||||
|
||||
all_speech_waveforms.extend(voice_speeches)
|
||||
all_speech_latent_lengths.extend(voice_latent_lengths)
|
||||
per_segment_is_target.extend([False] * len(voice_speeches))
|
||||
|
||||
all_speech_waveforms.append(wav_target)
|
||||
all_speech_latent_lengths.append(target_latent_len)
|
||||
per_segment_is_target.append(True)
|
||||
|
||||
max_seq_len = max(len(x) for x in sample_input_ids)
|
||||
padded_input_ids = []
|
||||
padded_attention_masks = []
|
||||
padded_acoustic_input_masks = []
|
||||
padded_acoustic_loss_masks = []
|
||||
tok = self.processor.tokenizer
|
||||
pad_token_id = getattr(tok, "pad_token_id", None)
|
||||
if pad_token_id is None or pad_token_id < 0:
|
||||
pad_token_id = getattr(tok, "eos_token_id", None)
|
||||
if pad_token_id is None or pad_token_id < 0:
|
||||
raise ValueError(
|
||||
"Tokenizer has no pad_token_id or eos_token_id; please set one or pass a valid pad id."
|
||||
)
|
||||
for ids, attn, ain_mask, aloss_mask in zip(
|
||||
sample_input_ids, sample_attention_masks, sample_acoustic_input_masks, sample_acoustic_loss_masks
|
||||
):
|
||||
pad_len = max_seq_len - len(ids)
|
||||
padded_input_ids.append(ids + [pad_token_id] * pad_len)
|
||||
padded_attention_masks.append(attn + [0] * pad_len)
|
||||
padded_acoustic_input_masks.append(ain_mask + [False] * pad_len)
|
||||
padded_acoustic_loss_masks.append(aloss_mask + [False] * pad_len)
|
||||
|
||||
input_ids_tensor = torch.tensor(padded_input_ids, dtype=torch.long)
|
||||
attention_mask_tensor = torch.tensor(padded_attention_masks, dtype=torch.long)
|
||||
acoustic_input_mask_tensor = torch.tensor(padded_acoustic_input_masks, dtype=torch.bool)
|
||||
acoustic_loss_mask_tensor = torch.tensor(padded_acoustic_loss_masks, dtype=torch.bool)
|
||||
|
||||
if all_speech_waveforms:
|
||||
max_wave_len = max(w.shape[0] for w in all_speech_waveforms)
|
||||
padded_speeches = np.zeros((len(all_speech_waveforms), max_wave_len), dtype=np.float32)
|
||||
for i, w in enumerate(all_speech_waveforms):
|
||||
L = w.shape[0]
|
||||
padded_speeches[i, :L] = w
|
||||
|
||||
max_latent_len = max(all_speech_latent_lengths) if all_speech_latent_lengths else 1
|
||||
speech_masks_np = np.zeros((len(all_speech_waveforms), max_latent_len), dtype=np.bool_)
|
||||
for i, L_lat in enumerate(all_speech_latent_lengths):
|
||||
speech_masks_np[i, :L_lat] = True
|
||||
|
||||
speech_tensors_tensor = torch.tensor(padded_speeches, dtype=torch.float32)
|
||||
speech_masks_tensor = torch.tensor(speech_masks_np, dtype=torch.bool)
|
||||
|
||||
speeches_loss_input_np = np.zeros_like(speech_masks_np, dtype=np.bool_)
|
||||
for i, is_target in enumerate(per_segment_is_target):
|
||||
if is_target:
|
||||
speeches_loss_input_np[i] = speech_masks_np[i]
|
||||
speeches_loss_input_tensor = torch.tensor(speeches_loss_input_np, dtype=torch.bool)
|
||||
|
||||
# Semantic features
|
||||
if self.compute_semantics and hasattr(self.processor, "semantic_tokenizer") and self.processor.semantic_tokenizer is not None:
|
||||
sem_feats: List[np.ndarray] = []
|
||||
for w in all_speech_waveforms:
|
||||
try:
|
||||
# Expect [T, D] where T ≈ ceil(len(w)/compress_ratio)
|
||||
sem = self.processor.semantic_tokenizer.encode(w)
|
||||
sem = np.asarray(sem, dtype=np.float32)
|
||||
except Exception:
|
||||
sem = np.zeros((0, self.semantic_vae_dim), dtype=np.float32)
|
||||
if sem.ndim != 2:
|
||||
raise RuntimeError(f"Semantic tokenizer returned unexpected shape {sem.shape}. Expect [T, D].")
|
||||
L = sem.shape[0]
|
||||
D = sem.shape[1]
|
||||
if D != self.semantic_vae_dim:
|
||||
if D < self.semantic_vae_dim:
|
||||
pad_d = np.zeros((L, self.semantic_vae_dim - D), dtype=np.float32)
|
||||
sem = np.concatenate([sem, pad_d], axis=1)
|
||||
else:
|
||||
sem = sem[:, : self.semantic_vae_dim]
|
||||
if L < max_latent_len:
|
||||
pad = np.zeros((max_latent_len - L, self.semantic_vae_dim), dtype=np.float32)
|
||||
sem = np.concatenate([sem, pad], axis=0)
|
||||
elif L > max_latent_len:
|
||||
sem = sem[:max_latent_len]
|
||||
sem_feats.append(sem.astype(np.float32))
|
||||
speech_semantic_tensors = torch.tensor(np.stack(sem_feats, axis=0), dtype=torch.float32)
|
||||
else:
|
||||
# Semantic tokenizer unavailable while semantics are required for training.
|
||||
# Raise to avoid silently degrading alignment with zeroed features.
|
||||
raise RuntimeError(
|
||||
"Semantic features are required but could not be computed. "
|
||||
"Ensure processor.semantic_tokenizer is available or precompute and provide features."
|
||||
)
|
||||
else:
|
||||
speech_tensors_tensor = None
|
||||
speech_masks_tensor = None
|
||||
speeches_loss_input_tensor = None
|
||||
speech_semantic_tensors = None # No segments in batch
|
||||
|
||||
if self.debug_checks:
|
||||
assert (input_ids_tensor >= 0).all(), "input_ids contains negative indices"
|
||||
if speech_tensors_tensor is not None:
|
||||
assert speech_tensors_tensor.dim() == 2, "Expected speech_tensors 2D [segments, samples]"
|
||||
|
||||
return {
|
||||
"input_ids": input_ids_tensor,
|
||||
"attention_mask": attention_mask_tensor,
|
||||
"speech_tensors": speech_tensors_tensor,
|
||||
"speech_masks": speech_masks_tensor,
|
||||
"speech_semantic_tensors": speech_semantic_tensors,
|
||||
"acoustic_input_mask": acoustic_input_mask_tensor,
|
||||
"acoustic_loss_mask": acoustic_loss_mask_tensor,
|
||||
"speeches_loss_input": speeches_loss_input_tensor,
|
||||
}
|
||||
@@ -0,0 +1,902 @@
|
||||
# train_vibevoice_lora.py
|
||||
import logging
|
||||
import os
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from datasets import load_dataset, DatasetDict, VerificationMode
|
||||
|
||||
from transformers import (
|
||||
HfArgumentParser,
|
||||
Trainer,
|
||||
set_seed,
|
||||
TrainerCallback,
|
||||
)
|
||||
from transformers import TrainingArguments as HfTrainingArguments
|
||||
|
||||
from peft import LoraConfig, get_peft_model, TaskType
|
||||
|
||||
from vibevoice.modular.modeling_vibevoice import VibeVoiceForConditionalGeneration
|
||||
from vibevoice.modular.configuration_vibevoice import VibeVoiceConfig
|
||||
from vibevoice.processor.vibevoice_processor import VibeVoiceProcessor
|
||||
|
||||
from vibevoice.finetune.data_vibevoice import VibeVoiceDataset, VibeVoiceCollator
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# ================== SAMPLE CALLBACK UTILS ==================
|
||||
|
||||
import copy
|
||||
import torch
|
||||
from transformers import TrainerCallback
|
||||
|
||||
class EmaCallback(TrainerCallback):
|
||||
def __init__(self, attr_path="model.prediction_head", decay=0.999, device="cpu"):
|
||||
"""
|
||||
attr_path: where the head lives under self.model (Trainer wraps your VibeVoiceForConditionalGeneration)
|
||||
decay: EMA decay (0.999 ~ stable, 0.9999 ~ very smooth, slower to adapt)
|
||||
"""
|
||||
self.attr_path = attr_path
|
||||
self.decay = float(decay)
|
||||
self.device = torch.device(device)
|
||||
self.shadow = None
|
||||
self._orig = None # store non-EMA weights when we swap
|
||||
|
||||
def _get_module(self, model):
|
||||
# Resolve dotted path like "model.prediction_head"
|
||||
mod = model
|
||||
for name in self.attr_path.split('.'):
|
||||
mod = getattr(mod, name)
|
||||
return mod
|
||||
|
||||
def on_train_begin(self, args, state, control, model=None, **kwargs):
|
||||
head = self._get_module(model)
|
||||
self.shadow = {k: p.detach().to(self.device).clone()
|
||||
for k, p in head.state_dict().items()}
|
||||
|
||||
def on_step_end(self, args, state, control, model=None, **kwargs):
|
||||
if self.shadow is None: return
|
||||
head = self._get_module(model)
|
||||
with torch.no_grad():
|
||||
for k, v in head.state_dict().items():
|
||||
self.shadow[k].mul_(self.decay).add_(v.detach().to(self.device), alpha=(1.0 - self.decay))
|
||||
|
||||
# ---- Swap helpers ----
|
||||
def _swap_in_ema(self, model):
|
||||
head = self._get_module(model)
|
||||
self._orig = copy.deepcopy(head.state_dict())
|
||||
head.load_state_dict(self.shadow, strict=False)
|
||||
|
||||
def _swap_back(self, model):
|
||||
if self._orig is None: return
|
||||
head = self._get_module(model)
|
||||
head.load_state_dict(self._orig, strict=False)
|
||||
self._orig = None
|
||||
|
||||
def on_evaluate(self, args, state, control, model=None, **kwargs):
|
||||
# use EMA during eval
|
||||
self._swap_in_ema(model)
|
||||
|
||||
def on_evaluate_end(self, args, state, control, model=None, **kwargs):
|
||||
self._swap_back(model)
|
||||
|
||||
def on_save(self, args, state, control, model=None, **kwargs):
|
||||
# temporarily swap to EMA, let Trainer save, then swap back
|
||||
self._swap_in_ema(model)
|
||||
|
||||
def on_save_end(self, args, state, control, model=None, **kwargs):
|
||||
self._swap_back(model)
|
||||
|
||||
def on_train_end(self, args, state, control, model=None, **kwargs):
|
||||
# final checkpoint: persist EMA
|
||||
self._swap_in_ema(model)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelArguments:
|
||||
model_name_or_path: Optional[str] = field(
|
||||
default=None, metadata={"help": "Path to VibeVoice base model with config.json"}
|
||||
)
|
||||
processor_name_or_path: Optional[str] = field(
|
||||
default=None, metadata={"help": "Path to processor dir (preprocessor_config.json). Defaults to model path."}
|
||||
)
|
||||
cache_dir: Optional[str] = field(default=None)
|
||||
freeze_acoustic_tokenizer: bool = field(default=True)
|
||||
freeze_semantic_tokenizer: bool = field(default=True)
|
||||
lora_r: int = field(default=8)
|
||||
lora_alpha: int = field(default=32)
|
||||
lora_dropout: float = field(default=0.05)
|
||||
lora_target_modules: str = field(
|
||||
default="q_proj,k_proj,v_proj,o_proj,gate_proj,up_proj,down_proj",
|
||||
metadata={"help": "Comma-separated list of target module names in the LLM blocks"},
|
||||
)
|
||||
lora_wrap_diffusion_head: bool = field(default=False, metadata={"help": "Wrap diffusion head with PEFT LoRA"})
|
||||
train_diffusion_head: bool = field(default=False, metadata={"help": "Train diffusion prediction head (full fine-tune)"})
|
||||
train_connectors: bool = field(default=False, metadata={"help": "Train acoustic/semantic connectors (full fine-tune)"})
|
||||
layers_to_freeze: Optional[str] = field(
|
||||
default=None,
|
||||
metadata={"help": "Comma-separated indices of diffusion head layers to freeze (e.g., '0,1,5,7,8')."}
|
||||
)
|
||||
|
||||
@dataclass
|
||||
class DataArguments:
|
||||
dataset_name: Optional[str] = field(default=None, metadata={"help": "HF dataset name or 'json' with --train_jsonl for local files"})
|
||||
dataset_config_name: Optional[str] = field(default=None)
|
||||
train_split_name: str = field(default="train")
|
||||
eval_split_name: Optional[str] = field(default="validation")
|
||||
text_column_name: str = field(default="text")
|
||||
audio_column_name: str = field(default="audio")
|
||||
voice_prompts_column_name: Optional[str] = field(default="voice_prompts")
|
||||
eval_split_size: float = field(default=0.0)
|
||||
ignore_verifications: bool = field(default=False)
|
||||
max_length: Optional[int] = field(default=None)
|
||||
train_jsonl: Optional[str] = field(default=None, metadata={"help": "Path to local train JSONL with {text, audio, [voice_prompts]}"})
|
||||
validation_jsonl: Optional[str] = field(default=None, metadata={"help": "Optional path to local validation JSONL"})
|
||||
voice_prompt_drop_rate: float = field(
|
||||
default=0.0,
|
||||
metadata={"help": "Probability to drop conditioning voice prompt during training (0.0 keep always, 1.0 drop always)."},
|
||||
)
|
||||
|
||||
@dataclass
|
||||
class CustomTrainingArguments(HfTrainingArguments):
|
||||
ddpm_batch_mul: int = field(default=1)
|
||||
ce_loss_weight: float = field(default=1.0)
|
||||
diffusion_loss_weight: float = field(default=1.0)
|
||||
debug_ce_details: bool = field(default=False)
|
||||
debug_ce_topk: int = field(default=5)
|
||||
debug_ce_max_examples: int = field(default=1)
|
||||
debug_ce_every_n_steps: int = field(default=200)
|
||||
gradient_clipping: bool = field(
|
||||
default=False,
|
||||
metadata={"help": "Enable gradient clipping using max_grad_norm (set via --max_grad_norm, default 1.0). When False, disables clipping by forcing max_grad_norm=0.0."},
|
||||
)
|
||||
debug_save: bool = field(
|
||||
default=False,
|
||||
metadata={"help": "If set, saves model components BEFORE training starts, into output_dir/debug_initial."},
|
||||
)
|
||||
|
||||
def build_lora_config(args: ModelArguments) -> LoraConfig:
|
||||
target_modules = [s.strip() for s in args.lora_target_modules.split(",") if s.strip()]
|
||||
return LoraConfig(
|
||||
r=args.lora_r,
|
||||
lora_alpha=args.lora_alpha,
|
||||
lora_dropout=args.lora_dropout,
|
||||
bias="none",
|
||||
task_type=TaskType.CAUSAL_LM,
|
||||
target_modules=target_modules,
|
||||
)
|
||||
|
||||
def build_head_lora_config(args: ModelArguments) -> LoraConfig:
|
||||
target_modules = ["noisy_images_proj","cond_proj","gate_proj","up_proj","down_proj","linear"]
|
||||
return LoraConfig(
|
||||
r=args.lora_r,
|
||||
lora_alpha=args.lora_alpha,
|
||||
lora_dropout=args.lora_dropout,
|
||||
bias="none",
|
||||
task_type=TaskType.FEATURE_EXTRACTION,
|
||||
target_modules=target_modules,
|
||||
)
|
||||
|
||||
def mask_for_ce(labels: torch.Tensor, attention_mask: torch.Tensor, acoustic_input_mask: torch.Tensor, pad_id: int = -100) -> torch.Tensor:
|
||||
shifted = labels[:, 1:].contiguous()
|
||||
base_mask = attention_mask[:, 1:].contiguous().eq(1) if (attention_mask is not None and attention_mask.numel() > 0) else torch.ones_like(shifted, dtype=torch.bool)
|
||||
label_is_acoustic = acoustic_input_mask[:, 1:].contiguous()
|
||||
final_mask = base_mask & (~label_is_acoustic)
|
||||
out = shifted.clone()
|
||||
out[~final_mask] = pad_id
|
||||
return out
|
||||
|
||||
def _patch_acoustic_encode_for_legacy_indexing(model_obj, logger_):
|
||||
try:
|
||||
acoustic = getattr(getattr(model_obj, "model", model_obj), "acoustic_tokenizer", None)
|
||||
if acoustic is None or not hasattr(acoustic, "encode"):
|
||||
logger_.warning("No acoustic_tokenizer.encode() found to patch.")
|
||||
return
|
||||
base_encode = acoustic.encode
|
||||
def encode_wrapped(*args, **kwargs):
|
||||
out = base_encode(*args, **kwargs)
|
||||
try:
|
||||
_ = out[0][0]
|
||||
return out
|
||||
except Exception:
|
||||
pass
|
||||
if isinstance(out, dict):
|
||||
for k in ("frames", "codes", "tokens", "latents", "hidden_states"):
|
||||
if k in out:
|
||||
return [[out[k]]]
|
||||
if len(out) > 0:
|
||||
return [[next(iter(out.values()))]]
|
||||
for attr in ("frames", "codes", "tokens", "latents", "hidden_states"):
|
||||
if hasattr(out, attr):
|
||||
return [[getattr(out, attr)]]
|
||||
try:
|
||||
if isinstance(out, torch.Tensor):
|
||||
return [[out]]
|
||||
except Exception:
|
||||
pass
|
||||
return [[out]]
|
||||
acoustic.encode = encode_wrapped
|
||||
logger_.info("Patched acoustic_tokenizer.encode() to return [[...]] for legacy indexing.")
|
||||
except Exception as e:
|
||||
logger_.warning(f"Failed to patch acoustic_tokenizer.encode(): {e}")
|
||||
|
||||
def main() -> None:
|
||||
parser = HfArgumentParser((ModelArguments, DataArguments, CustomTrainingArguments))
|
||||
model_args, data_args, training_args = parser.parse_args_into_dataclasses()
|
||||
|
||||
logging.basicConfig(
|
||||
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
|
||||
datefmt="%m/%d/%Y %H:%M:%S",
|
||||
level=logging.INFO if training_args.local_rank in [-1, 0] else logging.WARN,
|
||||
)
|
||||
logger.info("Training/evaluation parameters %s", training_args)
|
||||
set_seed(training_args.seed)
|
||||
|
||||
# Configure gradient clipping
|
||||
if not getattr(training_args, "gradient_clipping", False):
|
||||
if hasattr(training_args, "max_grad_norm"):
|
||||
training_args.max_grad_norm = 0.0
|
||||
logger.info("Gradient clipping disabled (set max_grad_norm=0.0). Use --gradient_clipping to enable.")
|
||||
else:
|
||||
if (not hasattr(training_args, "max_grad_norm")) or training_args.max_grad_norm is None or training_args.max_grad_norm <= 0:
|
||||
training_args.max_grad_norm = 1.0
|
||||
logger.info(f"Gradient clipping enabled: max_grad_norm={training_args.max_grad_norm}")
|
||||
|
||||
# Load processor
|
||||
processor_path = model_args.processor_name_or_path or model_args.model_name_or_path
|
||||
if processor_path is None:
|
||||
raise ValueError("--model_name_or_path (or --processor_name_or_path) must be provided")
|
||||
processor: VibeVoiceProcessor = VibeVoiceProcessor.from_pretrained(processor_path)
|
||||
|
||||
# Required special tokens
|
||||
tok = processor.tokenizer
|
||||
for required in ["speech_start_id", "speech_diffusion_id", "speech_end_id"]:
|
||||
if not hasattr(tok, required) or getattr(tok, required) is None:
|
||||
raise RuntimeError(f"Tokenizer missing required special id: {required}")
|
||||
|
||||
# Load model
|
||||
if model_args.model_name_or_path is None:
|
||||
raise ValueError("--model_name_or_path is required to load VibeVoice base model")
|
||||
dtype = torch.float32
|
||||
if training_args.bf16:
|
||||
dtype = torch.bfloat16
|
||||
elif getattr(training_args, "fp16", False):
|
||||
dtype = torch.float16
|
||||
model = VibeVoiceForConditionalGeneration.from_pretrained(
|
||||
model_args.model_name_or_path,
|
||||
torch_dtype=dtype,
|
||||
)
|
||||
_patch_acoustic_encode_for_legacy_indexing(model, logger)
|
||||
processor.semantic_tokenizer = getattr(model.model, "semantic_tokenizer", None)
|
||||
|
||||
# Diagnostics: LM head tie
|
||||
try:
|
||||
in_emb_mod = model.get_input_embeddings()
|
||||
out_emb_mod = model.get_output_embeddings()
|
||||
in_w = getattr(in_emb_mod, "weight", None)
|
||||
out_w = getattr(out_emb_mod, "weight", None)
|
||||
shared_ptr = bool(in_w is not None and out_w is not None and in_w.data_ptr() == out_w.data_ptr())
|
||||
values_equal = False
|
||||
if in_w is not None and out_w is not None and in_w.shape == out_w.shape:
|
||||
try:
|
||||
values_equal = bool(torch.allclose(in_w, out_w))
|
||||
except Exception:
|
||||
values_equal = False
|
||||
try:
|
||||
tie_cfg = getattr(getattr(model.config, "decoder_config", model.config), "tie_word_embeddings", None)
|
||||
except Exception:
|
||||
tie_cfg = getattr(model.config, "tie_word_embeddings", None)
|
||||
logger.info(f"LM head diagnostics -> shared_params={shared_ptr}, values_equal={values_equal}, tie_word_embeddings={tie_cfg}")
|
||||
if out_w is not None:
|
||||
logger.info(f"LM head requires_grad before freeze: {bool(out_w.requires_grad)}")
|
||||
except Exception as e:
|
||||
logger.warning(f"LM head tie diagnostics failed: {e}")
|
||||
|
||||
# Hard-tie LM head
|
||||
try:
|
||||
emb_module = model.get_input_embeddings()
|
||||
head_module = model.get_output_embeddings()
|
||||
if hasattr(emb_module, "weight") and hasattr(head_module, "weight"):
|
||||
if emb_module.weight.shape == head_module.weight.shape and emb_module.weight.data_ptr() != head_module.weight.data_ptr():
|
||||
with torch.no_grad():
|
||||
head_module.weight = emb_module.weight
|
||||
logger.info("Force-tied LM head weight to input embeddings (pointer share).")
|
||||
except Exception as e:
|
||||
logger.warning(f"Force-tie of LM head failed: {e}")
|
||||
|
||||
# Validate special IDs (info logs only)
|
||||
try:
|
||||
special_names = ["speech_start_id", "speech_diffusion_id", "speech_end_id"]
|
||||
try:
|
||||
vocab_size = int(getattr(model.config.decoder_config, "vocab_size", 0))
|
||||
except Exception:
|
||||
vocab_size = 0
|
||||
in_emb_mod = model.get_input_embeddings()
|
||||
out_emb_mod = model.get_output_embeddings()
|
||||
in_w = getattr(in_emb_mod, "weight", None)
|
||||
out_w = getattr(out_emb_mod, "weight", None)
|
||||
for name in special_names:
|
||||
val = getattr(tok, name, None)
|
||||
exists = (val is not None)
|
||||
in_range = (exists and isinstance(val, int) and 0 <= val < vocab_size)
|
||||
equal_row = None
|
||||
if in_range and in_w is not None and out_w is not None and in_w.shape == out_w.shape and in_w.size(0) > val:
|
||||
try:
|
||||
equal_row = bool(torch.allclose(in_w[val], out_w[val]))
|
||||
except Exception:
|
||||
equal_row = False
|
||||
decoded_str = None
|
||||
if exists and isinstance(val, int):
|
||||
try:
|
||||
decoded_str = tok.decode([val])
|
||||
except Exception:
|
||||
try:
|
||||
decoded_str = tok.convert_ids_to_tokens(val)
|
||||
except Exception:
|
||||
decoded_str = "<decode_failed>"
|
||||
logger.info(f"Special token check -> {name}={val}, decoded='{decoded_str}', exists={exists}, in_vocab_range={in_range}, emb_vs_head_row_equal={equal_row}")
|
||||
except Exception as e:
|
||||
logger.warning(f"Special token ID/row validation failed: {e}")
|
||||
|
||||
# Quick tokenizer diagnostics (optional)
|
||||
try:
|
||||
logger.info("=== TOKENIZER DIAGNOSTICS ===")
|
||||
logger.info(f"Tokenizer class: {type(tok).__name__}")
|
||||
logger.info(f"Tokenizer vocab_size: {tok.vocab_size}")
|
||||
# tiny CE smoke test
|
||||
with torch.no_grad():
|
||||
simple_text = "The cat sat on the mat."
|
||||
simple_ids = torch.tensor([tok.encode(simple_text, add_special_tokens=True)], device=model.device)
|
||||
simple_mask = torch.ones_like(simple_ids)
|
||||
x = model.get_input_embeddings()(simple_ids)
|
||||
outputs = model.model(inputs_embeds=x, attention_mask=simple_mask, return_dict=True)
|
||||
logits = model.lm_head(outputs.last_hidden_state)
|
||||
shift_logits = logits[:, :-1, :].contiguous()
|
||||
shift_labels = simple_ids[:, 1:].contiguous()
|
||||
ce_loss = F.cross_entropy(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1), reduction='mean')
|
||||
logger.info(f"Simple text CE loss: {ce_loss.item():.4f}")
|
||||
except Exception as e:
|
||||
logger.warning(f"Tokenizer diagnostics failed: {e}")
|
||||
|
||||
# Disable cache during training
|
||||
if hasattr(model.config, "use_cache") and training_args.do_train:
|
||||
model.config.use_cache = False
|
||||
|
||||
# Freeze tokenizers
|
||||
if model_args.freeze_acoustic_tokenizer and hasattr(model.model, "acoustic_tokenizer"):
|
||||
for p in model.model.acoustic_tokenizer.parameters():
|
||||
p.requires_grad = False
|
||||
if model_args.freeze_semantic_tokenizer and hasattr(model.model, "semantic_tokenizer"):
|
||||
for p in model.model.semantic_tokenizer.parameters():
|
||||
p.requires_grad = False
|
||||
|
||||
# LoRA wrap LLM (optional)
|
||||
lora_cfg = build_lora_config(model_args)
|
||||
tm_lower = [s.strip().lower() for s in model_args.lora_target_modules.split(",") if s.strip()]
|
||||
skip_lm_lora = (len(tm_lower) == 0) or all(t in ("none", "off", "disable", "disabled") for t in tm_lower)
|
||||
if not skip_lm_lora:
|
||||
model.model.language_model = get_peft_model(model.model.language_model, lora_cfg)
|
||||
else:
|
||||
logger.info("Skipping LLM LoRA wrapping (lora_target_modules indicates none).")
|
||||
|
||||
try:
|
||||
model.tie_weights()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Freeze all then enable trainable subsets
|
||||
for _, p in model.named_parameters():
|
||||
p.requires_grad = False
|
||||
|
||||
try:
|
||||
for n, p in model.model.language_model.named_parameters():
|
||||
if "lora_A" in n or "lora_B" in n:
|
||||
p.requires_grad = True
|
||||
except Exception:
|
||||
logger.warning("Could not re-enable LoRA params on language_model.")
|
||||
|
||||
# Diffusion head LoRA wrapping (optional)
|
||||
if getattr(model_args, "lora_wrap_diffusion_head", False) and hasattr(model.model, "prediction_head"):
|
||||
class _HeadForwardShim(nn.Module):
|
||||
def __init__(self, base: nn.Module): super().__init__(); self.base = base
|
||||
def forward(self, *args, **kwargs):
|
||||
if len(args) >= 3:
|
||||
noisy_images, timesteps, condition = args[:3]
|
||||
else:
|
||||
noisy_images = kwargs.get("noisy_images")
|
||||
timesteps = kwargs.get("timesteps")
|
||||
condition = kwargs.get("condition")
|
||||
return self.base(noisy_images, timesteps, condition)
|
||||
try:
|
||||
shim = _HeadForwardShim(model.model.prediction_head)
|
||||
model.model.prediction_head = get_peft_model(shim, build_head_lora_config(model_args))
|
||||
for n, p in model.model.prediction_head.named_parameters():
|
||||
if "lora_A" in n or "lora_B" in n:
|
||||
p.requires_grad = True
|
||||
except Exception as e:
|
||||
logger.warning(f"Could not LoRA-wrap diffusion head: {e}")
|
||||
|
||||
# Train full diffusion head (optional)
|
||||
if getattr(model_args, "train_diffusion_head", False) and hasattr(model.model, "prediction_head"):
|
||||
for p in model.model.prediction_head.parameters():
|
||||
p.requires_grad = True
|
||||
|
||||
# Freeze diffusion head layers (optional)
|
||||
if model_args.layers_to_freeze is not None and hasattr(model.model, "prediction_head"):
|
||||
head_params = list(model.model.prediction_head.named_parameters())
|
||||
try:
|
||||
indices_to_freeze = {int(x.strip()) for x in model_args.layers_to_freeze.split(',') if x.strip()}
|
||||
frozen_count = 0
|
||||
for i, (name, param) in enumerate(head_params):
|
||||
if i in indices_to_freeze:
|
||||
param.requires_grad = False
|
||||
frozen_count += 1
|
||||
logger.info(f"Froze layer [{i}]: {name}")
|
||||
logger.info(f"Successfully froze {frozen_count} parameter groups in the diffusion head.")
|
||||
except Exception as e:
|
||||
logger.error(f"Could not parse --layers_to_freeze: {e}")
|
||||
raise
|
||||
|
||||
# Connectors
|
||||
if getattr(model_args, "train_connectors", False):
|
||||
if hasattr(model.model, "acoustic_connector"):
|
||||
for p in model.model.acoustic_connector.parameters():
|
||||
p.requires_grad = True
|
||||
if hasattr(model.model, "semantic_connector"):
|
||||
for p in model.model.semantic_connector.parameters():
|
||||
p.requires_grad = True
|
||||
else:
|
||||
if hasattr(model.model, "acoustic_connector"):
|
||||
for p in model.model.acoustic_connector.parameters():
|
||||
p.requires_grad = False
|
||||
if hasattr(model.model, "semantic_connector"):
|
||||
for p in model.model.semantic_connector.parameters():
|
||||
p.requires_grad = False
|
||||
|
||||
# Freeze embedding + head
|
||||
try:
|
||||
emb = model.get_input_embeddings()
|
||||
if hasattr(emb, "weight"):
|
||||
emb.weight.requires_grad_(False)
|
||||
head = model.get_output_embeddings()
|
||||
if head is not None and hasattr(head, "weight"):
|
||||
head.weight.requires_grad_(False)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Diagnostics
|
||||
def _sum_params(named_iter):
|
||||
return sum(p.numel() for _, p in named_iter if p.requires_grad)
|
||||
try:
|
||||
lm_lora = _sum_params(model.model.language_model.named_parameters()) if hasattr(model.model, "language_model") else 0
|
||||
pred_head_train = _sum_params(model.model.prediction_head.named_parameters()) if hasattr(model.model, "prediction_head") else 0
|
||||
ac_conn_train = _sum_params(model.model.acoustic_connector.named_parameters()) if hasattr(model.model, "acoustic_connector") else 0
|
||||
se_conn_train = _sum_params(model.model.semantic_connector.named_parameters()) if hasattr(model.model, "semantic_connector") else 0
|
||||
total_trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
|
||||
logger.info(f"Trainable by block -> LLM-LoRA: {lm_lora:,} | diff_head: {pred_head_train:,} | ac_conn: {ac_conn_train:,} | se_conn: {se_conn_train:,}")
|
||||
logger.info("TOTAL trainable: %s", f"{total_trainable:,}")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Datasets
|
||||
verification_mode = VerificationMode.NO_CHECKS if data_args.ignore_verifications else VerificationMode.BASIC_CHECKS
|
||||
if data_args.train_jsonl is not None:
|
||||
data_files: Dict[str, str] = {"train": data_args.train_jsonl}
|
||||
if data_args.validation_jsonl is not None:
|
||||
data_files["validation"] = data_args.validation_jsonl
|
||||
raw = load_dataset("json", data_files=data_files, verification_mode=verification_mode, cache_dir=model_args.cache_dir)
|
||||
else:
|
||||
if data_args.dataset_name is None:
|
||||
raise ValueError("Provide --dataset_name (HF datasets) or use --train_jsonl/--validation_jsonl for local files.")
|
||||
raw = load_dataset(
|
||||
data_args.dataset_name,
|
||||
data_args.dataset_config_name,
|
||||
verification_mode=verification_mode,
|
||||
cache_dir=model_args.cache_dir,
|
||||
)
|
||||
train_ds = raw[data_args.train_split_name]
|
||||
eval_ds = None
|
||||
if training_args.do_eval:
|
||||
if data_args.eval_split_name and data_args.eval_split_name in raw:
|
||||
eval_ds = raw[data_args.eval_split_name]
|
||||
elif data_args.eval_split_size and data_args.eval_split_size > 0 and len(train_ds) > 1:
|
||||
split = train_ds.train_test_split(test_size=data_args.eval_split_size, seed=training_args.seed)
|
||||
train_ds, eval_ds = split["train"], split["test"]
|
||||
|
||||
train_dataset = VibeVoiceDataset(
|
||||
train_ds,
|
||||
text_column=data_args.text_column_name,
|
||||
audio_column=data_args.audio_column_name,
|
||||
voice_prompts_column=data_args.voice_prompts_column_name,
|
||||
)
|
||||
eval_dataset = None
|
||||
if eval_ds is not None:
|
||||
eval_dataset = VibeVoiceDataset(
|
||||
eval_ds,
|
||||
text_column=data_args.text_column_name,
|
||||
audio_column=data_args.audio_column_name,
|
||||
voice_prompts_column=data_args.voice_prompts_column_name,
|
||||
)
|
||||
|
||||
# Ratios/dims from processor+model
|
||||
speech_compress_ratio = getattr(processor, "speech_tok_compress_ratio", 3200)
|
||||
semantic_dim = getattr(model.config, "semantic_vae_dim", None)
|
||||
if semantic_dim is None:
|
||||
try:
|
||||
semantic_dim = int(getattr(model.config.semantic_tokenizer_config, "vae_dim", 128))
|
||||
except Exception:
|
||||
semantic_dim = 128
|
||||
|
||||
compute_semantics_flag = hasattr(processor, "semantic_tokenizer") and processor.semantic_tokenizer is not None
|
||||
|
||||
data_collator = VibeVoiceCollator(
|
||||
processor=processor,
|
||||
max_length=data_args.max_length,
|
||||
speech_compress_ratio=speech_compress_ratio,
|
||||
semantic_vae_dim=semantic_dim,
|
||||
compute_semantics=compute_semantics_flag,
|
||||
debug_checks=False,
|
||||
voice_prompt_drop_rate=data_args.voice_prompt_drop_rate,
|
||||
)
|
||||
|
||||
class LoRADebugCallback(TrainerCallback):
|
||||
def __init__(self, log_every_n_steps: int = 50):
|
||||
self.log_every_n_steps = max(1, int(log_every_n_steps))
|
||||
self.prev_param_norms: Dict[str, float] = {}
|
||||
self.lora_param_names: List[str] = []
|
||||
|
||||
def on_train_begin(self, args, state, control, model=None, **kwargs):
|
||||
try:
|
||||
if model is None:
|
||||
return
|
||||
named: Dict[str, torch.nn.Parameter] = dict(model.named_parameters())
|
||||
self.lora_param_names = [n for n in named.keys() if ("lora_A" in n or "lora_B" in n)]
|
||||
for n in self.lora_param_names:
|
||||
p = named[n]
|
||||
self.prev_param_norms[n] = float(p.data.norm().item())
|
||||
total = len(self.lora_param_names)
|
||||
req_grad = sum(1 for n in self.lora_param_names if named[n].requires_grad)
|
||||
num_A = sum(1 for n in self.lora_param_names if "lora_A" in n)
|
||||
num_B = sum(1 for n in self.lora_param_names if "lora_B" in n)
|
||||
zero_B = sum(1 for n in self.lora_param_names if ("lora_B" in n and float(named[n].data.norm().item()) == 0.0))
|
||||
logger.info(f"LoRA debug: found {total} LoRA params (A={num_A}, B={num_B}); trainable={req_grad}. Initial lora_B_zero={zero_B}.")
|
||||
if total == 0:
|
||||
logger.warning("LoRA debug: No LoRA parameters found. Check lora_target_modules.")
|
||||
if req_grad != total:
|
||||
logger.warning("LoRA debug: Some LoRA params are frozen. They should be trainable.")
|
||||
except Exception as e:
|
||||
logger.warning(f"LoRA debug (on_train_begin) failed: {e}")
|
||||
|
||||
def on_step_end(self, args, state, control, model=None, **kwargs):
|
||||
try:
|
||||
if model is None or len(self.lora_param_names) == 0:
|
||||
return
|
||||
step = int(getattr(state, "global_step", 0) or 0)
|
||||
if step % self.log_every_n_steps != 0 and step != 1:
|
||||
return
|
||||
named: Dict[str, torch.nn.Parameter] = dict(model.named_parameters())
|
||||
changed_A = 0
|
||||
changed_B = 0
|
||||
zero_B = 0
|
||||
eps = 1e-12
|
||||
for n in self.lora_param_names:
|
||||
p = named.get(n, None)
|
||||
if p is None:
|
||||
continue
|
||||
prev = self.prev_param_norms.get(n, 0.0)
|
||||
curr = float(p.data.norm().item())
|
||||
if "lora_A" in n and abs(curr - prev) > eps:
|
||||
changed_A += 1
|
||||
if "lora_B" in n:
|
||||
if abs(curr - prev) > eps:
|
||||
changed_B += 1
|
||||
if curr == 0.0:
|
||||
zero_B += 1
|
||||
self.prev_param_norms[n] = curr
|
||||
total_A = sum(1 for n in self.lora_param_names if "lora_A" in n)
|
||||
total_B = sum(1 for n in self.lora_param_names if "lora_B" in n)
|
||||
logger.info(f"LoRA debug step {step}: changed A {changed_A}/{total_A}, changed B {changed_B}/{total_B}, lora_B_zero_now={zero_B}.")
|
||||
except Exception as e:
|
||||
logger.warning(f"LoRA debug (on_step_end) failed: {e}")
|
||||
|
||||
class VibeVoiceTrainer(Trainer):
|
||||
def compute_loss(self, model: VibeVoiceForConditionalGeneration, inputs: Dict[str, Any], return_outputs=False, num_items_in_batch: Optional[int] = None):
|
||||
labels = inputs.get("input_ids")
|
||||
attention_mask = inputs.get("attention_mask")
|
||||
acoustic_input_mask = inputs.get("acoustic_input_mask")
|
||||
|
||||
# Ensure semantic tensors exist and have correct dtype/device
|
||||
sem = inputs.get("speech_semantic_tensors", None)
|
||||
try:
|
||||
target_dtype = next(model.model.semantic_connector.parameters()).dtype
|
||||
except Exception:
|
||||
target_dtype = model.get_input_embeddings().weight.dtype
|
||||
|
||||
if sem is None:
|
||||
sm = inputs.get("speech_masks")
|
||||
if sm is not None:
|
||||
zeros = torch.zeros(
|
||||
sm.size(0), sm.size(1),
|
||||
getattr(model.config, "semantic_vae_dim", 128),
|
||||
dtype=target_dtype,
|
||||
device=sm.device,
|
||||
)
|
||||
inputs["speech_semantic_tensors"] = zeros
|
||||
else:
|
||||
if isinstance(sem, torch.Tensor):
|
||||
inputs["speech_semantic_tensors"] = sem.to(dtype=target_dtype)
|
||||
|
||||
outputs = model(
|
||||
input_ids=inputs.get("input_ids"),
|
||||
attention_mask=attention_mask,
|
||||
speech_tensors=inputs.get("speech_tensors"),
|
||||
speech_masks=inputs.get("speech_masks"),
|
||||
speech_semantic_tensors=inputs.get("speech_semantic_tensors"),
|
||||
acoustic_input_mask=acoustic_input_mask,
|
||||
acoustic_loss_mask=inputs.get("acoustic_loss_mask"),
|
||||
speeches_loss_input=inputs.get("speeches_loss_input"),
|
||||
ddpm_batch_mul=training_args.ddpm_batch_mul,
|
||||
)
|
||||
|
||||
# Invariants: token/latent selection equality across views (warn, don't assert)
|
||||
try:
|
||||
al_mask = inputs.get("acoustic_loss_mask")
|
||||
sp_masks = inputs.get("speech_masks")
|
||||
sp_loss_sel = inputs.get("speeches_loss_input")
|
||||
num_tok_total = int(acoustic_input_mask.sum().item()) if acoustic_input_mask is not None else 0
|
||||
num_tok_loss = int(al_mask.sum().item()) if al_mask is not None else 0
|
||||
num_lat_total = int(sp_masks.sum().item()) if sp_masks is not None else 0
|
||||
num_lat_loss = int(((sp_loss_sel & sp_masks).sum().item())) if (sp_loss_sel is not None and sp_masks is not None) else 0
|
||||
self.log({
|
||||
"debug/num_tok_total": float(num_tok_total),
|
||||
"debug/num_tok_loss": float(num_tok_loss),
|
||||
"debug/num_lat_total": float(num_lat_total),
|
||||
"debug/num_lat_loss": float(num_lat_loss),
|
||||
})
|
||||
if sp_loss_sel is not None and sp_masks is not None and al_mask is not None:
|
||||
if num_tok_loss != num_lat_loss:
|
||||
logger.warning(f"Loss selection mismatch: acoustic_loss_mask={num_tok_loss} vs speeches_loss_input={num_lat_loss}")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# CE Loss
|
||||
logits = outputs.logits
|
||||
ce_labels = mask_for_ce(labels, attention_mask, acoustic_input_mask, pad_id=-100)
|
||||
shift_logits = logits[:, :-1, :].contiguous()
|
||||
loss_fct = nn.CrossEntropyLoss(ignore_index=-100)
|
||||
ce_loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), ce_labels.view(-1))
|
||||
|
||||
# Optional CE diagnostics
|
||||
try:
|
||||
self._debug_ce(shift_logits, ce_labels, attention_mask, acoustic_input_mask)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed invoking CE debug: {e}")
|
||||
|
||||
# Diffusion loss
|
||||
diffusion_loss = outputs.diffusion_loss if outputs.diffusion_loss is not None else torch.tensor(0.0, device=ce_loss.device)
|
||||
total = training_args.ce_loss_weight * ce_loss + training_args.diffusion_loss_weight * diffusion_loss
|
||||
|
||||
# Logs
|
||||
try:
|
||||
prefix = "train" if model.training else "eval"
|
||||
self.log({
|
||||
f"{prefix}/ce_loss": ce_loss.detach().item(),
|
||||
f"{prefix}/diffusion_loss": diffusion_loss.detach().item() if isinstance(diffusion_loss, torch.Tensor) else float(diffusion_loss),
|
||||
})
|
||||
if hasattr(self, "optimizer") and self.optimizer is not None and len(self.optimizer.param_groups) > 0:
|
||||
lr_val = self.optimizer.param_groups[0].get("lr", None)
|
||||
if lr_val is not None:
|
||||
self.log({"train/learning_rate_real": float(lr_val)})
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return (total, outputs) if return_outputs else total
|
||||
|
||||
def _debug_ce(self, shift_logits: torch.Tensor, ce_labels: torch.Tensor, attention_mask: Optional[torch.Tensor], acoustic_input_mask: Optional[torch.Tensor]):
|
||||
try:
|
||||
if not getattr(training_args, "debug_ce_details", False):
|
||||
return
|
||||
step = int(getattr(self.state, "global_step", 0) or 0)
|
||||
every_n = max(1, int(getattr(training_args, "debug_ce_every_n_steps", 200) or 200))
|
||||
if not (step <= 1 or (step % every_n == 0)):
|
||||
return
|
||||
|
||||
with torch.no_grad():
|
||||
vocab = shift_logits.size(-1)
|
||||
per_token_loss = F.cross_entropy(
|
||||
shift_logits.view(-1, vocab),
|
||||
ce_labels.view(-1),
|
||||
reduction="none",
|
||||
ignore_index=-100,
|
||||
).view_as(ce_labels)
|
||||
|
||||
valid_mask = ce_labels.ne(-100)
|
||||
num_valid = int(valid_mask.sum().item())
|
||||
avg_loss = float((per_token_loss[valid_mask].mean().item())) if num_valid > 0 else float("nan")
|
||||
|
||||
per_ex_avgs = []
|
||||
max_examples = max(1, int(getattr(training_args, "debug_ce_max_examples", 1) or 1))
|
||||
B = ce_labels.size(0)
|
||||
for b in range(min(B, max_examples)):
|
||||
vb = valid_mask[b]
|
||||
if int(vb.sum().item()) > 0:
|
||||
per_ex_avgs.append(float(per_token_loss[b][vb].mean().item()))
|
||||
else:
|
||||
per_ex_avgs.append(float("nan"))
|
||||
logger.info(f"CE debug: tokens_in_loss={num_valid}, avg_loss={avg_loss:.4f}, per_example_avgs={[round(x,4) if x==x else None for x in per_ex_avgs]}")
|
||||
except Exception as e:
|
||||
logger.warning(f"CE detailed debug failed: {e}")
|
||||
|
||||
# --------- CRITICAL SAVE OVERRIDES: also dump FULL head/connectors for inference ---------
|
||||
|
||||
|
||||
def _save(self, output_dir: Optional[str] = None, state_dict=None) -> None:
|
||||
try:
|
||||
target_dir = output_dir or self.args.output_dir
|
||||
lora_out = os.path.join(target_dir, "lora")
|
||||
os.makedirs(lora_out, exist_ok=True)
|
||||
|
||||
# --- LLM PEFT adapters (if LoRA-wrapped) ---
|
||||
language_model = getattr(self.model.model, "language_model", None)
|
||||
if hasattr(language_model, "save_pretrained"):
|
||||
language_model.save_pretrained(lora_out)
|
||||
|
||||
# --- Diffusion head PEFT adapters (if LoRA-wrapped) ---
|
||||
pred_head = getattr(self.model.model, "prediction_head", None)
|
||||
if hasattr(pred_head, "save_pretrained"):
|
||||
ph_dir = os.path.join(lora_out, "diffusion_head")
|
||||
os.makedirs(ph_dir, exist_ok=True)
|
||||
pred_head.save_pretrained(ph_dir)
|
||||
|
||||
# --- ALWAYS save FULL diffusion head state_dict for fallback ---
|
||||
if pred_head is not None and hasattr(pred_head, "state_dict"):
|
||||
sd = pred_head.state_dict()
|
||||
torch.save(sd, os.path.join(lora_out, "diffusion_head_full.bin"))
|
||||
ph_dir = os.path.join(lora_out, "diffusion_head")
|
||||
os.makedirs(ph_dir, exist_ok=True)
|
||||
torch.save(sd, os.path.join(ph_dir, "diffusion_head_full.bin"))
|
||||
|
||||
# --- Connectors (plain state_dicts) ---
|
||||
ac = getattr(self.model.model, "acoustic_connector", None)
|
||||
if ac is not None:
|
||||
ac_dir = os.path.join(lora_out, "acoustic_connector")
|
||||
os.makedirs(ac_dir, exist_ok=True)
|
||||
torch.save(ac.state_dict(), os.path.join(ac_dir, "pytorch_model.bin"))
|
||||
|
||||
se = getattr(self.model.model, "semantic_connector", None)
|
||||
if se is not None:
|
||||
se_dir = os.path.join(lora_out, "semantic_connector")
|
||||
os.makedirs(se_dir, exist_ok=True)
|
||||
torch.save(se.state_dict(), os.path.join(se_dir, "pytorch_model.bin"))
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to save LoRA assets: {e}")
|
||||
|
||||
|
||||
# ------------- Build the Trainer -------------
|
||||
|
||||
# Resolve which adapters to apply in samples
|
||||
|
||||
ema_cb = EmaCallback(attr_path="model.prediction_head", decay=0.999, device="cpu")
|
||||
|
||||
trainer = VibeVoiceTrainer(
|
||||
model=model,
|
||||
args=training_args,
|
||||
train_dataset=train_dataset,
|
||||
eval_dataset=eval_dataset,
|
||||
data_collator=data_collator,
|
||||
callbacks=[ema_cb, LoRADebugCallback(log_every_n_steps=(int(getattr(training_args, "logging_steps", 50) or 50)))],
|
||||
)
|
||||
|
||||
# Optional debug pre-training save
|
||||
if getattr(training_args, "debug_save", False):
|
||||
try:
|
||||
debug_dir = os.path.join(training_args.output_dir, "debug_initial")
|
||||
lora_out = os.path.join(debug_dir, "lora")
|
||||
os.makedirs(lora_out, exist_ok=True)
|
||||
logger.info(f"[debug_save] Saving initial (pre-training) model components to: {debug_dir}")
|
||||
# language model adapters / base
|
||||
try:
|
||||
if hasattr(model.model.language_model, "save_pretrained"):
|
||||
model.model.language_model.save_pretrained(lora_out)
|
||||
except Exception as e_lm:
|
||||
logger.warning(f"[debug_save] Failed to save language_model: {e_lm}")
|
||||
# diffusion head
|
||||
try:
|
||||
if hasattr(model.model, "prediction_head") and hasattr(model.model.prediction_head, "save_pretrained"):
|
||||
model.model.prediction_head.save_pretrained(os.path.join(lora_out, "diffusion_head"))
|
||||
except Exception as e_head:
|
||||
logger.warning(f"[debug_save] Failed to save prediction_head: {e_head}")
|
||||
# NEW: full diffusion head state_dict as fallback
|
||||
try:
|
||||
ph = getattr(model.model, "prediction_head", None)
|
||||
if ph is not None and hasattr(ph, "state_dict"):
|
||||
sd = ph.state_dict()
|
||||
torch.save(sd, os.path.join(lora_out, "diffusion_head_full.bin"))
|
||||
os.makedirs(os.path.join(lora_out, "diffusion_head"), exist_ok=True)
|
||||
torch.save(sd, os.path.join(lora_out, "diffusion_head", "diffusion_head_full.bin"))
|
||||
except Exception as e:
|
||||
logger.warning(f"[debug_save] Failed to save FULL diffusion head: {e}")
|
||||
# connectors
|
||||
try:
|
||||
ac_conn = getattr(model.model, "acoustic_connector", None)
|
||||
if ac_conn is not None:
|
||||
ac_dir = os.path.join(lora_out, "acoustic_connector")
|
||||
os.makedirs(ac_dir, exist_ok=True)
|
||||
torch.save(ac_conn.state_dict(), os.path.join(ac_dir, "pytorch_model.bin"))
|
||||
except Exception as e_ac:
|
||||
logger.warning(f"[debug_save] Failed to save acoustic_connector: {e_ac}")
|
||||
try:
|
||||
se_conn = getattr(model.model, "semantic_connector", None)
|
||||
if se_conn is not None:
|
||||
se_dir = os.path.join(lora_out, "semantic_connector")
|
||||
os.makedirs(se_dir, exist_ok=True)
|
||||
torch.save(se_conn.state_dict(), os.path.join(se_dir, "pytorch_model.bin"))
|
||||
except Exception as e_se:
|
||||
logger.warning(f"[debug_save] Failed to save semantic_connector: {e_se}")
|
||||
except Exception as e:
|
||||
logger.warning(f"[debug_save] Unexpected failure saving initial components: {e}")
|
||||
|
||||
if getattr(training_args, "gradient_checkpointing", False):
|
||||
try:
|
||||
model.gradient_checkpointing_enable()
|
||||
except Exception:
|
||||
logger.warning("Failed to enable gradient checkpointing on the model.")
|
||||
|
||||
if training_args.do_train:
|
||||
trainer.train(resume_from_checkpoint=training_args.resume_from_checkpoint)
|
||||
|
||||
lora_out = os.path.join(training_args.output_dir, "lora")
|
||||
os.makedirs(lora_out, exist_ok=True)
|
||||
|
||||
# LLM PEFT (if any)
|
||||
lm = getattr(model.model, "language_model", None)
|
||||
if hasattr(lm, "save_pretrained"):
|
||||
lm.save_pretrained(lora_out)
|
||||
|
||||
# Diffusion head PEFT (if any)
|
||||
ph = getattr(model.model, "prediction_head", None)
|
||||
if hasattr(ph, "save_pretrained"):
|
||||
ph_dir = os.path.join(lora_out, "diffusion_head")
|
||||
os.makedirs(ph_dir, exist_ok=True)
|
||||
ph.save_pretrained(ph_dir)
|
||||
|
||||
# ALWAYS: full diffusion head state_dict fallback
|
||||
try:
|
||||
if ph is not None and hasattr(ph, "state_dict"):
|
||||
sd = ph.state_dict()
|
||||
torch.save(sd, os.path.join(lora_out, "diffusion_head_full.bin"))
|
||||
ph_dir = os.path.join(lora_out, "diffusion_head")
|
||||
os.makedirs(ph_dir, exist_ok=True)
|
||||
torch.save(sd, os.path.join(ph_dir, "diffusion_head_full.bin"))
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to save FULL diffusion head at end: {e}")
|
||||
|
||||
# Connectors (if trained)
|
||||
try:
|
||||
ac = getattr(model.model, "acoustic_connector", None)
|
||||
if ac is not None:
|
||||
ac_dir = os.path.join(lora_out, "acoustic_connector")
|
||||
os.makedirs(ac_dir, exist_ok=True)
|
||||
torch.save(ac.state_dict(), os.path.join(ac_dir, "pytorch_model.bin"))
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to save acoustic_connector: {e}")
|
||||
|
||||
try:
|
||||
se = getattr(model.model, "semantic_connector", None)
|
||||
if se is not None:
|
||||
se_dir = os.path.join(lora_out, "semantic_connector")
|
||||
os.makedirs(se_dir, exist_ok=True)
|
||||
torch.save(se.state_dict(), os.path.join(se_dir, "pytorch_model.bin"))
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to save semantic_connector: {e}")
|
||||
|
||||
if training_args.do_eval and eval_dataset is not None:
|
||||
trainer.evaluate()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -368,11 +368,18 @@ class VibeVoiceForConditionalGeneration(VibeVoicePreTrainedModel):
|
||||
)
|
||||
if speech_tensors is not None:
|
||||
if semantic_speech_all_connect_features is not None:
|
||||
x[acoustic_input_mask] = speech_all_connect_features[speech_masks] + semantic_speech_all_connect_features[speech_masks]
|
||||
x[acoustic_input_mask] = (
|
||||
speech_all_connect_features[speech_masks]
|
||||
+ semantic_speech_all_connect_features[speech_masks]
|
||||
)
|
||||
else:
|
||||
x[acoustic_input_mask] = speech_all_connect_features[speech_masks]
|
||||
speech_features = speech_all_features[speeches_loss_input.unsqueeze(-1) & speech_masks] # only part audio need diffuse
|
||||
speech_connect_features = speech_all_connect_features[speeches_loss_input.unsqueeze(-1) & speech_masks]
|
||||
|
||||
# Select only the target segments' latents for diffusion loss.
|
||||
# Both masks are [num_segments, max_latent_len]; using 2D mask on [B,T,D] selects [N_true, D].
|
||||
target_latent_mask = speeches_loss_input & speech_masks
|
||||
speech_features = speech_all_features[target_latent_mask]
|
||||
speech_connect_features = speech_all_connect_features[target_latent_mask]
|
||||
else:
|
||||
speech_features, speech_connect_features = self.forward_speech_features(
|
||||
speech_tensors=speech_tensors.type_as(x) if speech_tensors is not None else None,
|
||||
@@ -485,4 +492,4 @@ __all__ = [
|
||||
"VibeVoiceForConditionalGeneration",
|
||||
"VibeVoiceCausalLMOutputWithPast",
|
||||
"VibeVoiceGenerationOutput",
|
||||
]
|
||||
]
|
||||
|
||||
@@ -426,6 +426,14 @@ class VibeVoiceProcessor:
|
||||
if isinstance(speaker_audio, str):
|
||||
# Load audio from file
|
||||
wav = self.audio_processor._load_audio_from_path(speaker_audio)
|
||||
elif isinstance(speaker_audio, dict):
|
||||
# Handle dict format with 'array' or 'audio' key
|
||||
if 'array' in speaker_audio:
|
||||
wav = np.array(speaker_audio['array'], dtype=np.float32)
|
||||
elif 'audio' in speaker_audio:
|
||||
wav = np.array(speaker_audio['audio'], dtype=np.float32)
|
||||
else:
|
||||
raise ValueError(f"Dictionary audio input must have 'array' or 'audio' key, got: {speaker_audio.keys()}")
|
||||
else:
|
||||
wav = np.array(speaker_audio, dtype=np.float32)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user