Automated PR - 2026-01-05
This commit is contained in:
sync-bot
@@ -0,0 +1,955 @@
|
||||
import os
|
||||
import time
|
||||
import warnings
|
||||
from pathlib import Path
|
||||
from typing import Callable
|
||||
|
||||
import torch
|
||||
import wandb
|
||||
import yaml
|
||||
from accelerate import Accelerator, DistributedType
|
||||
from accelerate.utils import set_seed
|
||||
from peft import LoraConfig, get_peft_model, get_peft_model_state_dict, set_peft_model_state_dict
|
||||
from peft.tuners.tuners_utils import BaseTunerLayer
|
||||
from peft.utils import ModulesToSaveWrapper
|
||||
from pydantic import BaseModel
|
||||
from safetensors.torch import load_file, save_file
|
||||
from torch import Tensor
|
||||
from torch.optim import AdamW
|
||||
from torch.optim.lr_scheduler import (
|
||||
CosineAnnealingLR,
|
||||
CosineAnnealingWarmRestarts,
|
||||
LinearLR,
|
||||
LRScheduler,
|
||||
PolynomialLR,
|
||||
StepLR,
|
||||
)
|
||||
from torch.utils.data import DataLoader
|
||||
from torchvision.transforms import functional as F # noqa: N812
|
||||
|
||||
from ltx_trainer import logger
|
||||
from ltx_trainer.config import LtxTrainerConfig
|
||||
from ltx_trainer.config_display import print_config
|
||||
from ltx_trainer.datasets import PrecomputedDataset
|
||||
from ltx_trainer.hf_hub_utils import push_to_hub
|
||||
from ltx_trainer.model_loader import load_model as load_ltx_model
|
||||
from ltx_trainer.model_loader import load_text_encoder
|
||||
from ltx_trainer.progress import TrainingProgress
|
||||
from ltx_trainer.quantization import quantize_model
|
||||
from ltx_trainer.timestep_samplers import SAMPLERS
|
||||
from ltx_trainer.training_strategies import get_training_strategy
|
||||
from ltx_trainer.utils import get_gpu_memory_gb, open_image_as_srgb, save_image
|
||||
from ltx_trainer.validation_sampler import CachedPromptEmbeddings, GenerationConfig, ValidationSampler
|
||||
from ltx_trainer.video_utils import read_video, save_video
|
||||
|
||||
# Disable irrelevant warnings from transformers
|
||||
os.environ["TOKENIZERS_PARALLELISM"] = "true"
|
||||
|
||||
# Silence bitsandbytes warnings about casting
|
||||
warnings.filterwarnings(
|
||||
"ignore", message="MatMul8bitLt: inputs will be cast from torch.bfloat16 to float16 during quantization"
|
||||
)
|
||||
|
||||
# Disable progress bars if not main process
|
||||
IS_MAIN_PROCESS = os.environ.get("LOCAL_RANK", "0") == "0"
|
||||
if not IS_MAIN_PROCESS:
|
||||
from transformers.utils.logging import disable_progress_bar
|
||||
|
||||
disable_progress_bar()
|
||||
|
||||
StepCallback = Callable[[int, int, list[Path]], None] # (step, total, list[sampled_video_path]) -> None
|
||||
|
||||
MEMORY_CHECK_INTERVAL = 200
|
||||
|
||||
|
||||
class TrainingStats(BaseModel):
|
||||
"""Statistics collected during training"""
|
||||
|
||||
total_time_seconds: float
|
||||
steps_per_second: float
|
||||
samples_per_second: float
|
||||
peak_gpu_memory_gb: float
|
||||
global_batch_size: int
|
||||
num_processes: int
|
||||
|
||||
|
||||
class LtxvTrainer:
|
||||
def __init__(self, trainer_config: LtxTrainerConfig) -> None:
|
||||
self._config = trainer_config
|
||||
if IS_MAIN_PROCESS:
|
||||
print_config(trainer_config)
|
||||
self._training_strategy = get_training_strategy(self._config.training_strategy)
|
||||
self._cached_validation_embeddings = self._load_text_encoder_and_cache_embeddings()
|
||||
self._load_models()
|
||||
self._setup_accelerator()
|
||||
self._collect_trainable_params()
|
||||
self._load_checkpoint()
|
||||
self._prepare_models_for_training()
|
||||
self._dataset = None
|
||||
self._global_step = -1
|
||||
self._checkpoint_paths = []
|
||||
self._init_wandb()
|
||||
|
||||
def train( # noqa: PLR0912, PLR0915
|
||||
self,
|
||||
disable_progress_bars: bool = False,
|
||||
step_callback: StepCallback | None = None,
|
||||
) -> tuple[Path, TrainingStats]:
|
||||
"""
|
||||
Start the training process.
|
||||
Returns:
|
||||
Tuple of (saved_model_path, training_stats)
|
||||
"""
|
||||
device = self._accelerator.device
|
||||
cfg = self._config
|
||||
start_mem = get_gpu_memory_gb(device)
|
||||
|
||||
train_start_time = time.time()
|
||||
|
||||
# Use the same seed for all processes and ensure deterministic operations
|
||||
set_seed(cfg.seed)
|
||||
logger.debug(f"Process {self._accelerator.process_index} using seed: {cfg.seed}")
|
||||
|
||||
self._init_optimizer()
|
||||
self._init_dataloader()
|
||||
data_iter = iter(self._dataloader)
|
||||
self._init_timestep_sampler()
|
||||
|
||||
# Synchronize all processes after initialization
|
||||
self._accelerator.wait_for_everyone()
|
||||
|
||||
Path(cfg.output_dir).mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Save the training configuration as YAML
|
||||
self._save_config()
|
||||
|
||||
logger.info("🚀 Starting training...")
|
||||
|
||||
# Create progress tracking (disabled for non-main processes or when explicitly disabled)
|
||||
progress_enabled = IS_MAIN_PROCESS and not disable_progress_bars
|
||||
progress = TrainingProgress(
|
||||
enabled=progress_enabled,
|
||||
total_steps=cfg.optimization.steps,
|
||||
)
|
||||
|
||||
if IS_MAIN_PROCESS and disable_progress_bars:
|
||||
logger.warning("Progress bars disabled. Intermediate status messages will be logged instead.")
|
||||
|
||||
self._transformer.train()
|
||||
self._global_step = 0
|
||||
|
||||
peak_mem_during_training = start_mem
|
||||
|
||||
sampled_videos_paths = None
|
||||
|
||||
with progress:
|
||||
# Initial validation before training starts
|
||||
if cfg.validation.interval and not cfg.validation.skip_initial_validation:
|
||||
sampled_videos_paths = self._sample_videos(progress)
|
||||
if IS_MAIN_PROCESS and sampled_videos_paths and self._config.wandb.log_validation_videos:
|
||||
self._log_validation_samples(sampled_videos_paths, cfg.validation.prompts)
|
||||
|
||||
self._accelerator.wait_for_everyone()
|
||||
|
||||
for step in range(cfg.optimization.steps * cfg.optimization.gradient_accumulation_steps):
|
||||
# Get next batch, reset the dataloader if needed
|
||||
try:
|
||||
batch = next(data_iter)
|
||||
except StopIteration:
|
||||
data_iter = iter(self._dataloader)
|
||||
batch = next(data_iter)
|
||||
|
||||
step_start_time = time.time()
|
||||
with self._accelerator.accumulate(self._transformer):
|
||||
is_optimization_step = (step + 1) % cfg.optimization.gradient_accumulation_steps == 0
|
||||
if is_optimization_step:
|
||||
self._global_step += 1
|
||||
|
||||
loss = self._training_step(batch)
|
||||
self._accelerator.backward(loss)
|
||||
|
||||
if self._accelerator.sync_gradients and cfg.optimization.max_grad_norm > 0:
|
||||
self._accelerator.clip_grad_norm_(
|
||||
self._trainable_params,
|
||||
cfg.optimization.max_grad_norm,
|
||||
)
|
||||
|
||||
self._optimizer.step()
|
||||
self._optimizer.zero_grad()
|
||||
|
||||
if self._lr_scheduler is not None:
|
||||
self._lr_scheduler.step()
|
||||
|
||||
# Run validation if needed
|
||||
if (
|
||||
cfg.validation.interval
|
||||
and self._global_step > 0
|
||||
and self._global_step % cfg.validation.interval == 0
|
||||
and is_optimization_step
|
||||
):
|
||||
if self._accelerator.distributed_type == DistributedType.FSDP:
|
||||
# FSDP: All processes must participate in validation
|
||||
sampled_videos_paths = self._sample_videos(progress)
|
||||
if IS_MAIN_PROCESS and sampled_videos_paths and self._config.wandb.log_validation_videos:
|
||||
self._log_validation_samples(sampled_videos_paths, cfg.validation.prompts)
|
||||
# DDP: Only main process runs validation
|
||||
elif IS_MAIN_PROCESS:
|
||||
sampled_videos_paths = self._sample_videos(progress)
|
||||
if sampled_videos_paths and self._config.wandb.log_validation_videos:
|
||||
self._log_validation_samples(sampled_videos_paths, cfg.validation.prompts)
|
||||
|
||||
# Save checkpoint if needed
|
||||
if (
|
||||
cfg.checkpoints.interval
|
||||
and self._global_step > 0
|
||||
and self._global_step % cfg.checkpoints.interval == 0
|
||||
and is_optimization_step
|
||||
):
|
||||
self._save_checkpoint()
|
||||
|
||||
self._accelerator.wait_for_everyone()
|
||||
|
||||
# Call step callback if provided
|
||||
if step_callback and is_optimization_step:
|
||||
step_callback(self._global_step, cfg.optimization.steps, sampled_videos_paths)
|
||||
|
||||
self._accelerator.wait_for_everyone()
|
||||
|
||||
# Update progress and log metrics
|
||||
current_lr = self._optimizer.param_groups[0]["lr"]
|
||||
step_time = (time.time() - step_start_time) * cfg.optimization.gradient_accumulation_steps
|
||||
|
||||
progress.update_training(
|
||||
loss=loss.item(),
|
||||
lr=current_lr,
|
||||
step_time=step_time,
|
||||
advance=is_optimization_step,
|
||||
)
|
||||
|
||||
# Log metrics to W&B (only on main process and optimization steps)
|
||||
if IS_MAIN_PROCESS and is_optimization_step:
|
||||
self._log_metrics(
|
||||
{
|
||||
"train/loss": loss.item(),
|
||||
"train/learning_rate": current_lr,
|
||||
"train/step_time": step_time,
|
||||
"train/global_step": self._global_step,
|
||||
}
|
||||
)
|
||||
|
||||
# Fallback logging when progress bars are disabled
|
||||
if disable_progress_bars and IS_MAIN_PROCESS and self._global_step % 20 == 0:
|
||||
elapsed = time.time() - train_start_time
|
||||
progress_percentage = self._global_step / cfg.optimization.steps
|
||||
if progress_percentage > 0:
|
||||
total_estimated = elapsed / progress_percentage
|
||||
total_time = f"{total_estimated // 3600:.0f}h {(total_estimated % 3600) // 60:.0f}m"
|
||||
else:
|
||||
total_time = "calculating..."
|
||||
logger.info(
|
||||
f"Step {self._global_step}/{cfg.optimization.steps} - "
|
||||
f"Loss: {loss.item():.4f}, LR: {current_lr:.2e}, "
|
||||
f"Time/Step: {step_time:.2f}s, Total Time: {total_time}",
|
||||
)
|
||||
|
||||
# Sample GPU memory periodically
|
||||
if step % MEMORY_CHECK_INTERVAL == 0:
|
||||
current_mem = get_gpu_memory_gb(device)
|
||||
peak_mem_during_training = max(peak_mem_during_training, current_mem)
|
||||
|
||||
# Collect final stats
|
||||
train_end_time = time.time()
|
||||
end_mem = get_gpu_memory_gb(device)
|
||||
peak_mem = max(start_mem, end_mem, peak_mem_during_training)
|
||||
|
||||
# Calculate steps/second over entire training
|
||||
total_time_seconds = train_end_time - train_start_time
|
||||
steps_per_second = cfg.optimization.steps / total_time_seconds
|
||||
|
||||
samples_per_second = steps_per_second * self._accelerator.num_processes * cfg.optimization.batch_size
|
||||
|
||||
stats = TrainingStats(
|
||||
total_time_seconds=total_time_seconds,
|
||||
steps_per_second=steps_per_second,
|
||||
samples_per_second=samples_per_second,
|
||||
peak_gpu_memory_gb=peak_mem,
|
||||
num_processes=self._accelerator.num_processes,
|
||||
global_batch_size=cfg.optimization.batch_size * self._accelerator.num_processes,
|
||||
)
|
||||
|
||||
saved_path = self._save_checkpoint()
|
||||
|
||||
if IS_MAIN_PROCESS:
|
||||
# Log the training statistics
|
||||
self._log_training_stats(stats)
|
||||
|
||||
# Upload artifacts to hub if enabled
|
||||
if cfg.hub.push_to_hub:
|
||||
push_to_hub(saved_path, sampled_videos_paths, self._config)
|
||||
|
||||
# Log final stats to W&B
|
||||
if self._wandb_run is not None:
|
||||
self._log_metrics(
|
||||
{
|
||||
"stats/total_time_minutes": stats.total_time_seconds / 60,
|
||||
"stats/steps_per_second": stats.steps_per_second,
|
||||
"stats/samples_per_second": stats.samples_per_second,
|
||||
"stats/peak_gpu_memory_gb": stats.peak_gpu_memory_gb,
|
||||
}
|
||||
)
|
||||
self._wandb_run.finish()
|
||||
|
||||
self._accelerator.wait_for_everyone()
|
||||
self._accelerator.end_training()
|
||||
|
||||
return saved_path, stats
|
||||
|
||||
def _training_step(self, batch: dict[str, dict[str, Tensor]]) -> Tensor:
|
||||
"""Perform a single training step using the configured strategy."""
|
||||
# Apply embedding connectors to transform pre-computed text embeddings
|
||||
conditions = batch["conditions"]
|
||||
video_embeds, audio_embeds, attention_mask = self._text_encoder._run_connectors(
|
||||
conditions["prompt_embeds"], conditions["prompt_attention_mask"]
|
||||
)
|
||||
conditions["video_prompt_embeds"] = video_embeds
|
||||
conditions["audio_prompt_embeds"] = audio_embeds
|
||||
conditions["prompt_attention_mask"] = attention_mask
|
||||
|
||||
# Use strategy to prepare training inputs (returns ModelInputs with Modality objects)
|
||||
model_inputs = self._training_strategy.prepare_training_inputs(batch, self._timestep_sampler)
|
||||
|
||||
# Run transformer forward pass with Modality-based interface
|
||||
video_pred, audio_pred = self._transformer(
|
||||
video=model_inputs.video,
|
||||
audio=model_inputs.audio,
|
||||
perturbations=None,
|
||||
)
|
||||
|
||||
# Use strategy to compute loss
|
||||
loss = self._training_strategy.compute_loss(video_pred, audio_pred, model_inputs)
|
||||
|
||||
return loss
|
||||
|
||||
def _load_text_encoder_and_cache_embeddings(self) -> list[CachedPromptEmbeddings] | None:
|
||||
"""Load text encoder, computes and returns validation embeddings."""
|
||||
|
||||
# This method:
|
||||
# 1. Loads the text encoder on GPU
|
||||
# 2. If validation prompts are configured, computes and caches their embeddings
|
||||
# 3. Unloads the heavy Gemma model while keeping the lightweight embedding connectors
|
||||
# The text encoder is kept (as self._text_encoder) but with model/tokenizer/feature_extractor
|
||||
# set to None. Only the embedding connectors remain for use during training.
|
||||
|
||||
# Load text encoder on GPU
|
||||
logger.debug("Loading text encoder...")
|
||||
if self._config.acceleration.load_text_encoder_in_8bit:
|
||||
logger.warning(
|
||||
"⚠️ load_text_encoder_in_8bit is set to True but 8-bit text encoder loading "
|
||||
"is not currently implemented. The text encoder will be loaded in bfloat16 precision."
|
||||
)
|
||||
|
||||
self._text_encoder = load_text_encoder(
|
||||
checkpoint_path=self._config.model.model_path,
|
||||
gemma_model_path=self._config.model.text_encoder_path,
|
||||
device="cuda",
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
|
||||
# Cache validation embeddings if prompts are configured
|
||||
cached_embeddings = None
|
||||
if self._config.validation.prompts:
|
||||
logger.info(f"Pre-computing embeddings for {len(self._config.validation.prompts)} validation prompts...")
|
||||
cached_embeddings = []
|
||||
with torch.inference_mode():
|
||||
for prompt in self._config.validation.prompts:
|
||||
v_ctx_pos, a_ctx_pos, _ = self._text_encoder(prompt)
|
||||
v_ctx_neg, a_ctx_neg, _ = self._text_encoder(self._config.validation.negative_prompt)
|
||||
|
||||
cached_embeddings.append(
|
||||
CachedPromptEmbeddings(
|
||||
video_context_positive=v_ctx_pos.cpu(),
|
||||
audio_context_positive=a_ctx_pos.cpu(),
|
||||
video_context_negative=v_ctx_neg.cpu() if v_ctx_neg is not None else None,
|
||||
audio_context_negative=a_ctx_neg.cpu() if a_ctx_neg is not None else None,
|
||||
)
|
||||
)
|
||||
|
||||
# Unload heavy components to free VRAM, keeping only the embedding connectors
|
||||
self._text_encoder.model = None
|
||||
self._text_encoder.tokenizer = None
|
||||
self._text_encoder.feature_extractor_linear = None
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
logger.debug("Validation prompt embeddings cached. Gemma model unloaded")
|
||||
return cached_embeddings
|
||||
|
||||
def _load_models(self) -> None:
|
||||
"""Load the LTX-2 model components."""
|
||||
# Load audio components if:
|
||||
# 1. Training strategy requires audio (training the audio branch), OR
|
||||
# 2. Validation is configured to generate audio (even if not training audio)
|
||||
load_audio = self._training_strategy.requires_audio or self._config.validation.generate_audio
|
||||
|
||||
# Check if we need VAE encoder (for image or reference video conditioning)
|
||||
need_vae_encoder = (
|
||||
self._config.validation.images is not None or self._config.validation.reference_videos is not None
|
||||
)
|
||||
|
||||
# Load all model components (except text encoder - already handled)
|
||||
components = load_ltx_model(
|
||||
checkpoint_path=self._config.model.model_path,
|
||||
device="cpu",
|
||||
dtype=torch.bfloat16,
|
||||
with_video_vae_encoder=need_vae_encoder, # Needed for image conditioning
|
||||
with_video_vae_decoder=True, # Needed for validation sampling
|
||||
with_audio_vae_decoder=load_audio,
|
||||
with_vocoder=load_audio,
|
||||
with_text_encoder=False, # Text encoder handled separately
|
||||
)
|
||||
|
||||
# Extract components
|
||||
self._transformer = components.transformer
|
||||
self._vae_decoder = components.video_vae_decoder.to(dtype=torch.bfloat16)
|
||||
self._vae_encoder = components.video_vae_encoder
|
||||
if self._vae_encoder is not None:
|
||||
self._vae_encoder = self._vae_encoder.to(dtype=torch.bfloat16)
|
||||
self._scheduler = components.scheduler
|
||||
self._audio_vae = components.audio_vae_decoder
|
||||
self._vocoder = components.vocoder
|
||||
# Note: self._text_encoder was set in _load_text_encoder_and_cache_embeddings
|
||||
|
||||
# Determine initial dtype based on training mode.
|
||||
# Note: For FSDP + LoRA, we'll cast to FP32 later in _prepare_models_for_training()
|
||||
# after the accelerator is set up, and we can detect FSDP.
|
||||
transformer_dtype = torch.bfloat16 if self._config.model.training_mode == "lora" else torch.float32
|
||||
self._transformer = self._transformer.to(dtype=transformer_dtype)
|
||||
|
||||
if self._config.acceleration.quantization is not None:
|
||||
if self._config.model.training_mode == "full":
|
||||
raise ValueError("Quantization is not supported in full training mode.")
|
||||
|
||||
logger.warning(f"Quantizing model with precision: {self._config.acceleration.quantization}")
|
||||
self._transformer = quantize_model(
|
||||
self._transformer,
|
||||
precision=self._config.acceleration.quantization,
|
||||
)
|
||||
|
||||
# Freeze all models. We later unfreeze the transformer based on training mode.
|
||||
# Note: embedding_connectors are already frozen (they come from the frozen text encoder)
|
||||
self._vae_decoder.requires_grad_(False)
|
||||
if self._vae_encoder is not None:
|
||||
self._vae_encoder.requires_grad_(False)
|
||||
self._transformer.requires_grad_(False)
|
||||
if self._audio_vae is not None:
|
||||
self._audio_vae.requires_grad_(False)
|
||||
if self._vocoder is not None:
|
||||
self._vocoder.requires_grad_(False)
|
||||
|
||||
def _collect_trainable_params(self) -> None:
|
||||
"""Collect trainable parameters based on training mode."""
|
||||
if self._config.model.training_mode == "lora":
|
||||
# For LoRA training, first set up LoRA layers
|
||||
self._setup_lora()
|
||||
elif self._config.model.training_mode == "full":
|
||||
# For full training, unfreeze all transformer parameters
|
||||
self._transformer.requires_grad_(True)
|
||||
else:
|
||||
raise ValueError(f"Unknown training mode: {self._config.model.training_mode}")
|
||||
|
||||
self._trainable_params = [p for p in self._transformer.parameters() if p.requires_grad]
|
||||
logger.debug(f"Trainable params count: {sum(p.numel() for p in self._trainable_params):,}")
|
||||
|
||||
def _init_timestep_sampler(self) -> None:
|
||||
"""Initialize the timestep sampler based on the config."""
|
||||
sampler_cls = SAMPLERS[self._config.flow_matching.timestep_sampling_mode]
|
||||
self._timestep_sampler = sampler_cls(**self._config.flow_matching.timestep_sampling_params)
|
||||
|
||||
def _setup_lora(self) -> None:
|
||||
"""Configure LoRA adapters for the transformer. Only called in LoRA training mode."""
|
||||
logger.debug(f"Adding LoRA adapter with rank {self._config.lora.rank}")
|
||||
lora_config = LoraConfig(
|
||||
r=self._config.lora.rank,
|
||||
lora_alpha=self._config.lora.alpha,
|
||||
target_modules=self._config.lora.target_modules,
|
||||
lora_dropout=self._config.lora.dropout,
|
||||
init_lora_weights=True,
|
||||
)
|
||||
# Wrap the transformer with PEFT to add LoRA layers
|
||||
# noinspection PyTypeChecker
|
||||
self._transformer = get_peft_model(self._transformer, lora_config)
|
||||
|
||||
def _load_checkpoint(self) -> None:
|
||||
"""Load checkpoint if specified in config."""
|
||||
if not self._config.model.load_checkpoint:
|
||||
return
|
||||
|
||||
checkpoint_path = self._find_checkpoint(self._config.model.load_checkpoint)
|
||||
if not checkpoint_path:
|
||||
logger.warning(f"⚠️ Could not find checkpoint at {self._config.model.load_checkpoint}")
|
||||
return
|
||||
|
||||
logger.info(f"📥 Loading checkpoint from {checkpoint_path}")
|
||||
|
||||
if self._config.model.training_mode == "full":
|
||||
self._load_full_checkpoint(checkpoint_path)
|
||||
else: # LoRA mode
|
||||
self._load_lora_checkpoint(checkpoint_path)
|
||||
|
||||
def _load_full_checkpoint(self, checkpoint_path: Path) -> None:
|
||||
"""Load full model checkpoint."""
|
||||
state_dict = load_file(checkpoint_path)
|
||||
self._transformer.load_state_dict(state_dict, strict=True)
|
||||
|
||||
logger.info("✅ Full model checkpoint loaded successfully")
|
||||
|
||||
def _load_lora_checkpoint(self, checkpoint_path: Path) -> None:
|
||||
"""Load LoRA checkpoint with DDP/FSDP compatibility."""
|
||||
state_dict = load_file(checkpoint_path)
|
||||
|
||||
# Adjust layer names to match internal format.
|
||||
# (Weights are saved in ComfyUI-compatible format, with "diffusion_model." prefix)
|
||||
state_dict = {k.replace("diffusion_model.", "", 1): v for k, v in state_dict.items()}
|
||||
|
||||
# Load LoRA weights and verify all weights were loaded
|
||||
base_model = self._transformer.get_base_model()
|
||||
set_peft_model_state_dict(base_model, state_dict)
|
||||
|
||||
logger.info("✅ LoRA checkpoint loaded successfully")
|
||||
|
||||
def _prepare_models_for_training(self) -> None:
|
||||
"""Prepare models for training with Accelerate."""
|
||||
|
||||
# For FSDP + LoRA: Cast entire model to FP32.
|
||||
# FSDP requires uniform dtype across all parameters in wrapped modules.
|
||||
# In LoRA mode, PEFT creates LoRA params in FP32 while base model is BF16.
|
||||
# We cast the base model to FP32 to match the LoRA params.
|
||||
if self._accelerator.distributed_type == DistributedType.FSDP and self._config.model.training_mode == "lora":
|
||||
logger.debug("FSDP: casting transformer to FP32 for uniform dtype")
|
||||
self._transformer = self._transformer.to(dtype=torch.float32)
|
||||
|
||||
# Enable gradient checkpointing if requested
|
||||
# For PeftModel, we need to access the underlying base model
|
||||
transformer = (
|
||||
self._transformer.get_base_model() if hasattr(self._transformer, "get_base_model") else self._transformer
|
||||
)
|
||||
|
||||
transformer.set_gradient_checkpointing(self._config.optimization.enable_gradient_checkpointing)
|
||||
|
||||
# Keep frozen models on CPU for memory efficiency
|
||||
self._vae_decoder = self._vae_decoder.to("cpu")
|
||||
if self._vae_encoder is not None:
|
||||
self._vae_encoder = self._vae_encoder.to("cpu")
|
||||
|
||||
# Embedding connectors are already on GPU from _load_text_encoder_and_cache_embeddings
|
||||
|
||||
# noinspection PyTypeChecker
|
||||
self._transformer = self._accelerator.prepare(self._transformer)
|
||||
|
||||
# Log GPU memory usage after model preparation
|
||||
vram_usage_gb = torch.cuda.memory_allocated() / 1024**3
|
||||
logger.debug(f"GPU memory usage after models preparation: {vram_usage_gb:.2f} GB")
|
||||
|
||||
@staticmethod
|
||||
def _find_checkpoint(checkpoint_path: str | Path) -> Path | None:
|
||||
"""Find the checkpoint file to load, handling both file and directory paths."""
|
||||
checkpoint_path = Path(checkpoint_path)
|
||||
|
||||
if checkpoint_path.is_file():
|
||||
if not checkpoint_path.suffix == ".safetensors":
|
||||
raise ValueError(f"Checkpoint file must have a .safetensors extension: {checkpoint_path}")
|
||||
return checkpoint_path
|
||||
|
||||
if checkpoint_path.is_dir():
|
||||
# Look for checkpoint files in the directory
|
||||
checkpoints = list(checkpoint_path.rglob("*step_*.safetensors"))
|
||||
|
||||
if not checkpoints:
|
||||
return None
|
||||
|
||||
# Sort by step number and return the latest
|
||||
def _get_step_num(p: Path) -> int:
|
||||
try:
|
||||
return int(p.stem.split("step_")[1])
|
||||
except (IndexError, ValueError):
|
||||
return -1
|
||||
|
||||
latest = max(checkpoints, key=_get_step_num)
|
||||
return latest
|
||||
|
||||
else:
|
||||
raise ValueError(f"Invalid checkpoint path: {checkpoint_path}. Must be a file or directory.")
|
||||
|
||||
def _init_dataloader(self) -> None:
|
||||
"""Initialize the training data loader using the strategy's data sources."""
|
||||
if self._dataset is None:
|
||||
# Get data sources from the training strategy
|
||||
data_sources = self._training_strategy.get_data_sources()
|
||||
|
||||
self._dataset = PrecomputedDataset(self._config.data.preprocessed_data_root, data_sources=data_sources)
|
||||
logger.debug(f"Loaded dataset with {len(self._dataset):,} samples from sources: {list(data_sources)}")
|
||||
|
||||
num_workers = self._config.data.num_dataloader_workers
|
||||
dataloader = DataLoader(
|
||||
self._dataset,
|
||||
batch_size=self._config.optimization.batch_size,
|
||||
shuffle=True,
|
||||
drop_last=True,
|
||||
num_workers=num_workers,
|
||||
pin_memory=num_workers > 0,
|
||||
persistent_workers=num_workers > 0,
|
||||
)
|
||||
|
||||
self._dataloader = self._accelerator.prepare(dataloader)
|
||||
|
||||
def _init_lora_weights(self) -> None:
|
||||
"""Initialize LoRA weights for the transformer."""
|
||||
logger.debug("Initializing LoRA weights...")
|
||||
for _, module in self._transformer.named_modules():
|
||||
if isinstance(module, (BaseTunerLayer, ModulesToSaveWrapper)):
|
||||
module.reset_lora_parameters(adapter_name="default", init_lora_weights=True)
|
||||
|
||||
def _init_optimizer(self) -> None:
|
||||
"""Initialize the optimizer and learning rate scheduler."""
|
||||
opt_cfg = self._config.optimization
|
||||
|
||||
lr = opt_cfg.learning_rate
|
||||
if opt_cfg.optimizer_type == "adamw":
|
||||
optimizer = AdamW(self._trainable_params, lr=lr)
|
||||
elif opt_cfg.optimizer_type == "adamw8bit":
|
||||
# noinspection PyUnresolvedReferences
|
||||
from bitsandbytes.optim import AdamW8bit # noqa: PLC0415
|
||||
|
||||
optimizer = AdamW8bit(self._trainable_params, lr=lr)
|
||||
else:
|
||||
raise ValueError(f"Unknown optimizer type: {opt_cfg.optimizer_type}")
|
||||
|
||||
# Add scheduler initialization
|
||||
lr_scheduler = self._create_scheduler(optimizer)
|
||||
|
||||
# noinspection PyTypeChecker
|
||||
self._optimizer, self._lr_scheduler = self._accelerator.prepare(optimizer, lr_scheduler)
|
||||
|
||||
def _create_scheduler(self, optimizer: torch.optim.Optimizer) -> LRScheduler | None:
|
||||
"""Create learning rate scheduler based on config."""
|
||||
scheduler_type = self._config.optimization.scheduler_type
|
||||
steps = self._config.optimization.steps
|
||||
params = self._config.optimization.scheduler_params or {}
|
||||
|
||||
if scheduler_type is None:
|
||||
return None
|
||||
|
||||
if scheduler_type == "linear":
|
||||
scheduler = LinearLR(
|
||||
optimizer,
|
||||
start_factor=params.pop("start_factor", 1.0),
|
||||
end_factor=params.pop("end_factor", 0.1),
|
||||
total_iters=steps,
|
||||
**params,
|
||||
)
|
||||
elif scheduler_type == "cosine":
|
||||
scheduler = CosineAnnealingLR(
|
||||
optimizer,
|
||||
T_max=steps,
|
||||
eta_min=params.pop("eta_min", 0),
|
||||
**params,
|
||||
)
|
||||
elif scheduler_type == "cosine_with_restarts":
|
||||
scheduler = CosineAnnealingWarmRestarts(
|
||||
optimizer,
|
||||
T_0=params.pop("T_0", steps // 4), # First restart cycle length
|
||||
T_mult=params.pop("T_mult", 1), # Multiplicative factor for cycle lengths
|
||||
eta_min=params.pop("eta_min", 5e-5),
|
||||
**params,
|
||||
)
|
||||
elif scheduler_type == "polynomial":
|
||||
scheduler = PolynomialLR(
|
||||
optimizer,
|
||||
total_iters=steps,
|
||||
power=params.pop("power", 1.0),
|
||||
**params,
|
||||
)
|
||||
elif scheduler_type == "step":
|
||||
scheduler = StepLR(
|
||||
optimizer,
|
||||
step_size=params.pop("step_size", steps // 2),
|
||||
gamma=params.pop("gamma", 0.1),
|
||||
**params,
|
||||
)
|
||||
elif scheduler_type == "constant":
|
||||
scheduler = None
|
||||
else:
|
||||
raise ValueError(f"Unknown scheduler type: {scheduler_type}")
|
||||
|
||||
return scheduler
|
||||
|
||||
def _setup_accelerator(self) -> None:
|
||||
"""Initialize the Accelerator with the appropriate settings."""
|
||||
|
||||
# All distributed setup (DDP/FSDP, number of processes, etc.) is controlled by
|
||||
# the user's Accelerate configuration (accelerate config / accelerate launch).
|
||||
self._accelerator = Accelerator(
|
||||
mixed_precision=self._config.acceleration.mixed_precision_mode,
|
||||
gradient_accumulation_steps=self._config.optimization.gradient_accumulation_steps,
|
||||
)
|
||||
|
||||
if self._accelerator.num_processes > 1:
|
||||
logger.info(
|
||||
f"{self._accelerator.distributed_type.value} distributed training enabled "
|
||||
f"with {self._accelerator.num_processes} processes"
|
||||
)
|
||||
|
||||
local_batch = self._config.optimization.batch_size
|
||||
global_batch = self._config.optimization.batch_size * self._accelerator.num_processes
|
||||
logger.info(f"Local batch size: {local_batch}, global batch size: {global_batch}")
|
||||
|
||||
# Log torch.compile status from Accelerate's dynamo plugin
|
||||
is_compile_enabled = (
|
||||
hasattr(self._accelerator.state, "dynamo_plugin") and self._accelerator.state.dynamo_plugin.backend != "NO"
|
||||
)
|
||||
if is_compile_enabled:
|
||||
plugin = self._accelerator.state.dynamo_plugin
|
||||
logger.info(f"🔥 torch.compile enabled via Accelerate: backend={plugin.backend}, mode={plugin.mode}")
|
||||
|
||||
if self._accelerator.distributed_type == DistributedType.FSDP:
|
||||
logger.warning(
|
||||
"⚠️ FSDP + torch.compile is experimental and may hang on the first training iteration. "
|
||||
"If this occurs, disable torch.compile by removing dynamo_config from your Accelerate config."
|
||||
)
|
||||
|
||||
if self._accelerator.distributed_type == DistributedType.FSDP and self._config.acceleration.quantization:
|
||||
logger.warning(
|
||||
f"FSDP with quantization ({self._config.acceleration.quantization}) may have compatibility issues."
|
||||
"Monitor training stability and consider disabling quantization if issues arise."
|
||||
)
|
||||
|
||||
# Note: Use @torch.no_grad() instead of @torch.inference_mode() to avoid FSDP inplace update errors after validation
|
||||
@torch.no_grad()
|
||||
def _sample_videos(self, progress: TrainingProgress) -> list[Path] | None:
|
||||
"""Run validation by generating videos from validation prompts."""
|
||||
use_images = self._config.validation.images is not None
|
||||
use_reference_videos = self._config.validation.reference_videos is not None
|
||||
generate_audio = self._config.validation.generate_audio
|
||||
inference_steps = self._config.validation.inference_steps
|
||||
|
||||
# Free up GPU memory before validation sampling.
|
||||
# Zero gradients and empty the cache to reclaim memory.
|
||||
self._optimizer.zero_grad(set_to_none=True)
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
# Start sampling progress tracking
|
||||
sampling_ctx = progress.start_sampling(
|
||||
num_prompts=len(self._config.validation.prompts),
|
||||
num_steps=inference_steps,
|
||||
)
|
||||
|
||||
# Create validation sampler with loaded models and progress tracking
|
||||
sampler = ValidationSampler(
|
||||
transformer=self._transformer,
|
||||
vae_decoder=self._vae_decoder,
|
||||
vae_encoder=self._vae_encoder,
|
||||
text_encoder=None,
|
||||
audio_decoder=self._audio_vae if generate_audio else None,
|
||||
vocoder=self._vocoder if generate_audio else None,
|
||||
sampling_context=sampling_ctx,
|
||||
)
|
||||
|
||||
output_dir = Path(self._config.output_dir) / "samples"
|
||||
output_dir.mkdir(exist_ok=True, parents=True)
|
||||
|
||||
video_paths = []
|
||||
width, height, num_frames = self._config.validation.video_dims
|
||||
|
||||
for prompt_idx, prompt in enumerate(self._config.validation.prompts):
|
||||
# Update progress to show current video
|
||||
sampling_ctx.start_video(prompt_idx)
|
||||
|
||||
# Load conditioning image if provided
|
||||
condition_image = None
|
||||
if use_images:
|
||||
image_path = self._config.validation.images[prompt_idx]
|
||||
image = open_image_as_srgb(image_path)
|
||||
# Convert PIL image to tensor [C, H, W] in [0, 1]
|
||||
condition_image = F.to_tensor(image)
|
||||
|
||||
# Load reference video if provided (for IC-LoRA)
|
||||
reference_video = None
|
||||
if use_reference_videos:
|
||||
ref_video_path = self._config.validation.reference_videos[prompt_idx]
|
||||
# read_video returns [F, C, H, W] in [0, 1]
|
||||
reference_video, _ = read_video(ref_video_path, max_frames=num_frames)
|
||||
|
||||
# Get cached embeddings for this prompt if available
|
||||
cached_embeddings = (
|
||||
self._cached_validation_embeddings[prompt_idx]
|
||||
if self._cached_validation_embeddings is not None
|
||||
else None
|
||||
)
|
||||
|
||||
# Create generation config
|
||||
gen_config = GenerationConfig(
|
||||
prompt=prompt,
|
||||
negative_prompt=self._config.validation.negative_prompt,
|
||||
height=height,
|
||||
width=width,
|
||||
num_frames=num_frames,
|
||||
frame_rate=self._config.validation.frame_rate,
|
||||
num_inference_steps=inference_steps,
|
||||
guidance_scale=self._config.validation.guidance_scale,
|
||||
seed=self._config.validation.seed,
|
||||
condition_image=condition_image,
|
||||
reference_video=reference_video,
|
||||
generate_audio=generate_audio,
|
||||
include_reference_in_output=self._config.validation.include_reference_in_output,
|
||||
cached_embeddings=cached_embeddings,
|
||||
stg_scale=self._config.validation.stg_scale,
|
||||
stg_blocks=self._config.validation.stg_blocks,
|
||||
stg_mode=self._config.validation.stg_mode,
|
||||
)
|
||||
|
||||
# Generate sample
|
||||
video, audio = sampler.generate(
|
||||
config=gen_config,
|
||||
device=self._accelerator.device,
|
||||
)
|
||||
|
||||
# Save output (image for single frame, video otherwise)
|
||||
if IS_MAIN_PROCESS:
|
||||
ext = "png" if num_frames == 1 else "mp4"
|
||||
output_path = output_dir / f"step_{self._global_step:06d}_{prompt_idx + 1}.{ext}"
|
||||
if num_frames == 1:
|
||||
save_image(video, output_path)
|
||||
else:
|
||||
save_video(
|
||||
video_tensor=video,
|
||||
output_path=output_path,
|
||||
fps=self._config.validation.frame_rate,
|
||||
audio=audio,
|
||||
audio_sample_rate=self._vocoder.output_sample_rate if audio is not None else None,
|
||||
)
|
||||
video_paths.append(output_path)
|
||||
|
||||
# Clean up progress tasks
|
||||
sampling_ctx.cleanup()
|
||||
|
||||
# Clear GPU cache after validation
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
rel_outputs_path = output_dir.relative_to(self._config.output_dir)
|
||||
logger.info(f"🎥 Validation samples for step {self._global_step} saved in {rel_outputs_path}")
|
||||
return video_paths
|
||||
|
||||
@staticmethod
|
||||
def _log_training_stats(stats: TrainingStats) -> None:
|
||||
"""Log training statistics."""
|
||||
stats_str = (
|
||||
"📊 Training Statistics:\n"
|
||||
f" - Total time: {stats.total_time_seconds / 60:.1f} minutes\n"
|
||||
f" - Training speed: {stats.steps_per_second:.2f} steps/second\n"
|
||||
f" - Samples/second: {stats.samples_per_second:.2f}\n"
|
||||
f" - Peak GPU memory: {stats.peak_gpu_memory_gb:.2f} GB"
|
||||
)
|
||||
if stats.num_processes > 1:
|
||||
stats_str += f"\n - Number of processes: {stats.num_processes}\n"
|
||||
stats_str += f" - Global batch size: {stats.global_batch_size}"
|
||||
logger.info(stats_str)
|
||||
|
||||
def _save_checkpoint(self) -> Path | None:
|
||||
"""Save the model weights."""
|
||||
is_lora = self._config.model.training_mode == "lora"
|
||||
is_fsdp = self._accelerator.distributed_type == DistributedType.FSDP
|
||||
|
||||
# Prepare paths
|
||||
save_dir = Path(self._config.output_dir) / "checkpoints"
|
||||
prefix = "lora" if is_lora else "model"
|
||||
filename = f"{prefix}_weights_step_{self._global_step:05d}.safetensors"
|
||||
saved_weights_path = save_dir / filename
|
||||
|
||||
# Get state dict (collective operation - all processes must participate)
|
||||
self._accelerator.wait_for_everyone()
|
||||
full_state_dict = self._accelerator.get_state_dict(self._transformer)
|
||||
|
||||
if not IS_MAIN_PROCESS:
|
||||
return None
|
||||
|
||||
save_dir.mkdir(exist_ok=True, parents=True)
|
||||
|
||||
# For LoRA: extract only adapter weights; for full: use as-is
|
||||
if is_lora:
|
||||
unwrapped = self._accelerator.unwrap_model(self._transformer, keep_torch_compile=False)
|
||||
# For FSDP, pass full_state_dict since model params aren't directly accessible
|
||||
state_dict = get_peft_model_state_dict(unwrapped, state_dict=full_state_dict if is_fsdp else None)
|
||||
|
||||
# Remove "base_model.model." prefix added by PEFT
|
||||
state_dict = {k.replace("base_model.model.", "", 1): v for k, v in state_dict.items()}
|
||||
|
||||
# Convert to ComfyUI-compatible format (add "diffusion_model." prefix)
|
||||
state_dict = {f"diffusion_model.{k}": v for k, v in state_dict.items()}
|
||||
|
||||
# Save to disk
|
||||
save_file(state_dict, saved_weights_path)
|
||||
else:
|
||||
# Save to disk
|
||||
self._accelerator.save(full_state_dict, saved_weights_path)
|
||||
|
||||
rel_path = saved_weights_path.relative_to(self._config.output_dir)
|
||||
logger.info(f"💾 {prefix.capitalize()} weights for step {self._global_step} saved in {rel_path}")
|
||||
|
||||
# Keep track of checkpoint paths, and cleanup old checkpoints if needed
|
||||
self._checkpoint_paths.append(saved_weights_path)
|
||||
self._cleanup_checkpoints()
|
||||
return saved_weights_path
|
||||
|
||||
def _cleanup_checkpoints(self) -> None:
|
||||
"""Clean up old checkpoints."""
|
||||
if 0 < self._config.checkpoints.keep_last_n < len(self._checkpoint_paths):
|
||||
checkpoints_to_remove = self._checkpoint_paths[: -self._config.checkpoints.keep_last_n]
|
||||
for old_checkpoint in checkpoints_to_remove:
|
||||
if old_checkpoint.exists():
|
||||
old_checkpoint.unlink()
|
||||
logger.info(f"Removed old checkpoints: {old_checkpoint}")
|
||||
# Update the list to only contain kept checkpoints
|
||||
self._checkpoint_paths = self._checkpoint_paths[-self._config.checkpoints.keep_last_n :]
|
||||
|
||||
def _save_config(self) -> None:
|
||||
"""Save the training configuration as a YAML file in the output directory."""
|
||||
if not IS_MAIN_PROCESS:
|
||||
return
|
||||
|
||||
config_path = Path(self._config.output_dir) / "training_config.yaml"
|
||||
with open(config_path, "w") as f:
|
||||
yaml.dump(self._config.model_dump(), f, default_flow_style=False, indent=2)
|
||||
|
||||
logger.info(f"💾 Training configuration saved to: {config_path.relative_to(self._config.output_dir)}")
|
||||
|
||||
def _init_wandb(self) -> None:
|
||||
"""Initialize Weights & Biases run."""
|
||||
if not self._config.wandb.enabled or not IS_MAIN_PROCESS:
|
||||
self._wandb_run = None
|
||||
return
|
||||
|
||||
wandb_config = self._config.wandb
|
||||
run = wandb.init(
|
||||
project=wandb_config.project,
|
||||
entity=wandb_config.entity,
|
||||
name=Path(self._config.output_dir).name,
|
||||
tags=wandb_config.tags,
|
||||
config=self._config.model_dump(),
|
||||
)
|
||||
self._wandb_run = run
|
||||
|
||||
def _log_metrics(self, metrics: dict[str, float]) -> None:
|
||||
"""Log metrics to Weights & Biases."""
|
||||
if self._wandb_run is not None:
|
||||
self._wandb_run.log(metrics)
|
||||
|
||||
def _log_validation_samples(self, sample_paths: list[Path], prompts: list[str]) -> None:
|
||||
"""Log validation samples (videos or images) to Weights & Biases."""
|
||||
if not self._config.wandb.log_validation_videos or self._wandb_run is None:
|
||||
return
|
||||
|
||||
# Determine if outputs are images or videos based on file extension
|
||||
is_image = sample_paths and sample_paths[0].suffix.lower() in (".png", ".jpg", ".jpeg", ".heic", ".webp")
|
||||
media_cls = wandb.Image if is_image else wandb.Video
|
||||
|
||||
samples = [media_cls(str(path), caption=prompt) for path, prompt in zip(sample_paths, prompts, strict=True)]
|
||||
self._wandb_run.log({"validation_samples": samples}, step=self._global_step)
|
||||
Reference in New Issue
Block a user