New loss function for training (#9)
This commit is contained in:
mrfakename
@@ -603,6 +603,164 @@ def main() -> None:
|
||||
logger.warning(f"LoRA debug (on_step_end) failed: {e}")
|
||||
|
||||
class VibeVoiceTrainer(Trainer):
|
||||
def training_forward(self, model: VibeVoiceForConditionalGeneration, inputs: Dict[str, Any]):
|
||||
"""Custom forward pass for training with new diffusion loss calculation."""
|
||||
# Extract inputs
|
||||
input_ids = inputs.get("input_ids")
|
||||
attention_mask = inputs.get("attention_mask")
|
||||
position_ids = inputs.get("position_ids")
|
||||
past_key_values = inputs.get("past_key_values")
|
||||
inputs_embeds = inputs.get("inputs_embeds")
|
||||
use_cache = inputs.get("use_cache", False)
|
||||
output_attentions = inputs.get("output_attentions")
|
||||
output_hidden_states = inputs.get("output_hidden_states")
|
||||
return_dict = inputs.get("return_dict", True)
|
||||
cache_position = inputs.get("cache_position")
|
||||
|
||||
# Speech-related inputs
|
||||
speech_tensors = inputs.get("speech_tensors")
|
||||
speech_masks = inputs.get("speech_masks")
|
||||
speeches_loss_input = inputs.get("speeches_loss_input")
|
||||
speech_semantic_tensors = inputs.get("speech_semantic_tensors")
|
||||
acoustic_input_mask = inputs.get("acoustic_input_mask")
|
||||
acoustic_loss_mask = inputs.get("acoustic_loss_mask")
|
||||
ddmp_batch_mul = training_args.ddpm_batch_mul
|
||||
kwargs = {}
|
||||
|
||||
# --- START: Copy of model forward logic with new diffusion loss ---
|
||||
x = model.get_input_embeddings()(input_ids)
|
||||
|
||||
semantic_speech_all_connect_features = model.model.semantic_connector(speech_semantic_tensors)
|
||||
if speeches_loss_input is not None:
|
||||
# only part audio need diffuse
|
||||
speech_all_features, speech_all_connect_features = model.forward_speech_features(
|
||||
speech_tensors=speech_tensors.type_as(x) if speech_tensors is not None else None,
|
||||
speech_masks=speech_masks,
|
||||
speech_type=kwargs.get("speech_type", "audio"),
|
||||
return_unmask=True
|
||||
)
|
||||
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]
|
||||
else:
|
||||
x[acoustic_input_mask] = speech_all_connect_features[speech_masks]
|
||||
speech_features = speech_all_features[speeches_loss_input & speech_masks] # only part audio need diffuse
|
||||
speech_connect_features = speech_all_connect_features[speeches_loss_input & speech_masks]
|
||||
# Forward-time consistency check: selected latent count should match number of acoustic placeholders
|
||||
try:
|
||||
if acoustic_input_mask is not None:
|
||||
assert speech_connect_features.shape[0] == int(acoustic_input_mask.sum().item()), (
|
||||
f"Mismatch between selected speech connectors ({speech_connect_features.shape[0]}) and acoustic_input_mask sum ({int(acoustic_input_mask.sum().item())})"
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
else:
|
||||
speech_features, speech_connect_features = model.forward_speech_features(
|
||||
speech_tensors=speech_tensors.type_as(x) if speech_tensors is not None else None,
|
||||
speech_masks=speech_masks,
|
||||
speech_type=kwargs.get("speech_type", "audio"),
|
||||
)
|
||||
if speech_tensors is not None:
|
||||
x[acoustic_input_mask] = speech_connect_features
|
||||
|
||||
outputs = model.model(
|
||||
input_ids=None,
|
||||
attention_mask=attention_mask,
|
||||
position_ids=position_ids,
|
||||
past_key_values=past_key_values,
|
||||
inputs_embeds=x,
|
||||
use_cache=use_cache,
|
||||
output_attentions=output_attentions,
|
||||
output_hidden_states=False,
|
||||
return_dict=return_dict,
|
||||
cache_position=cache_position,
|
||||
)
|
||||
|
||||
hidden_states = outputs.last_hidden_state
|
||||
logits = model.lm_head(hidden_states)
|
||||
|
||||
loss = None
|
||||
|
||||
# --- NEW Diffusion Loss Calculation ---
|
||||
diffusion_loss = None
|
||||
# This block is executed only if we are in a context that involves speech.
|
||||
if speech_tensors is not None and acoustic_loss_mask.sum().item() > 0:
|
||||
# Build conditioning mask from positions whose NEXT token is a speech latent (shift left by 1)
|
||||
cond_mask = torch.zeros_like(acoustic_loss_mask, dtype=torch.bool)
|
||||
cond_mask[:, :-1] = acoustic_loss_mask[:, 1:]
|
||||
cond_mask[:, 0] = False
|
||||
condition_features = hidden_states[cond_mask]
|
||||
|
||||
speech_len, latent_size = speech_features.shape
|
||||
# Sanity check: ensure 1:1 alignment between selected conditions and latents
|
||||
try:
|
||||
assert condition_features.shape[0] == speech_len, (
|
||||
f"Mismatch: condition_features={condition_features.shape[0]} vs speech_features={speech_len}"
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
noise = torch.randn(
|
||||
(speech_len * ddmp_batch_mul, latent_size),
|
||||
device=hidden_states.device,
|
||||
dtype=hidden_states.dtype
|
||||
)
|
||||
|
||||
timesteps = torch.multinomial(
|
||||
torch.ones(model.config.diffusion_head_config.ddpm_num_steps),
|
||||
speech_len * ddmp_batch_mul,
|
||||
replacement=True,
|
||||
).to(hidden_states.device)
|
||||
|
||||
speech_features_repeated = speech_features.repeat_interleave(ddmp_batch_mul, dim=0)
|
||||
condition_features_repeated = condition_features.repeat_interleave(ddmp_batch_mul, dim=0)
|
||||
|
||||
noisy_speech_features = model.model.noise_scheduler.add_noise(
|
||||
speech_features_repeated, noise, timesteps
|
||||
)
|
||||
|
||||
model_output = model.model.prediction_head(
|
||||
noisy_speech_features,
|
||||
timesteps.type_as(x),
|
||||
condition_features_repeated
|
||||
)
|
||||
|
||||
prediction_type = model.config.diffusion_head_config.prediction_type
|
||||
if prediction_type == "epsilon":
|
||||
target_for_loss = noise
|
||||
elif prediction_type == "v_prediction":
|
||||
target_for_loss = model.model.noise_scheduler.get_velocity(
|
||||
speech_features_repeated, noise, timesteps
|
||||
)
|
||||
else:
|
||||
raise NotImplementedError(f"Prediction type {prediction_type} not implemented")
|
||||
|
||||
diffusion_loss = F.mse_loss(model_output.float(), target_for_loss.float(), reduction='sum')
|
||||
if latent_size > 0 and ddmp_batch_mul > 0:
|
||||
# Normalize by latent dim, number of sampled diffusion steps per latent, and number of speech tokens
|
||||
diffusion_loss = diffusion_loss / latent_size / ddmp_batch_mul / max(speech_len, 1)
|
||||
else:
|
||||
diffusion_loss = torch.tensor(0.0, device=diffusion_loss.device)
|
||||
|
||||
else:
|
||||
# Dummy loss for DDP to work when there are no speech samples in a batch,
|
||||
# but we are in a speech context.
|
||||
diffusion_loss = sum(p.sum() for p in model.model.prediction_head.parameters()) * 0.0
|
||||
diffusion_loss += sum(p.sum() for p in model.model.acoustic_connector.parameters()) * 0.0
|
||||
diffusion_loss += sum(p.sum() for p in model.model.semantic_connector.parameters()) * 0.0
|
||||
# --- End NEW Diffusion Loss Calculation ---
|
||||
|
||||
from vibevoice.modular.modeling_vibevoice import VibeVoiceCausalLMOutputWithPast
|
||||
return VibeVoiceCausalLMOutputWithPast(
|
||||
loss=loss,
|
||||
diffusion_loss=diffusion_loss,
|
||||
speech_token_num=speech_len if speech_tensors is not None else 0,
|
||||
logits=logits,
|
||||
past_key_values=outputs.past_key_values,
|
||||
hidden_states=outputs.hidden_states,
|
||||
attentions=outputs.attentions,
|
||||
)
|
||||
|
||||
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")
|
||||
@@ -629,17 +787,8 @@ def main() -> None:
|
||||
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,
|
||||
)
|
||||
# Use custom training forward pass with new diffusion loss
|
||||
outputs = self.training_forward(model, inputs)
|
||||
|
||||
# Invariants: token/latent selection equality across views (warn, don't assert)
|
||||
try:
|
||||
|
||||
Reference in New Issue
Block a user