github-actions[bot]
1232 lines
55 KiB
Python
1232 lines
55 KiB
Python
import os
|
|
import re
|
|
import time
|
|
import warnings
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any, 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_core.text_encoders.gemma import convert_to_additive_mask
|
|
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.gpu_utils import free_gpu_memory, free_gpu_memory_context, get_gpu_memory_gb
|
|
from ltx_trainer.hf_hub_utils import push_to_hub
|
|
from ltx_trainer.model_loader import load_embeddings_processor, load_text_encoder
|
|
from ltx_trainer.model_loader import load_model as load_ltx_model
|
|
from ltx_trainer.progress import TrainingProgress
|
|
from ltx_trainer.quantization import quantize_model
|
|
from ltx_trainer.sigma_tracker import SigmaBucketTracker
|
|
from ltx_trainer.timestep_samplers import SAMPLERS
|
|
from ltx_trainer.training_state import ConfigFingerprint, RngStates, TrainingState
|
|
from ltx_trainer.training_strategies import get_training_strategy
|
|
from ltx_trainer.utils import 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
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class TrainingStepOutput:
|
|
"""Output from a single training step."""
|
|
|
|
loss: Tensor # [B,] per-element loss (unreduced)
|
|
sigma: Tensor # [B,] sampled sigma, detached from computational graph
|
|
|
|
|
|
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._loaded_checkpoint_path: Path | None = None
|
|
self._load_checkpoint()
|
|
self._prepare_models_for_training()
|
|
self._dataset = None
|
|
self._global_step = -1
|
|
self._checkpoint_paths: list[Path] = []
|
|
self._training_state_paths: list[Path] = []
|
|
self._training_state_size_warned = False
|
|
self._wandb_run = None
|
|
self._sigma_tracker = SigmaBucketTracker()
|
|
|
|
def train( # noqa: PLR0912, PLR0915
|
|
self,
|
|
disable_progress_bars: bool = False,
|
|
step_callback: StepCallback | None = None,
|
|
) -> tuple[Path, TrainingStats]:
|
|
"""
|
|
Start the training process.
|
|
Args:
|
|
disable_progress_bars: Disable Rich progress bars (useful for multi-process runs).
|
|
step_callback: Optional callback invoked after each optimization step.
|
|
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()
|
|
|
|
initial_step, training_state = self._resume_state
|
|
resuming = training_state is not None
|
|
|
|
set_seed(cfg.seed)
|
|
logger.debug(f"Process {self._accelerator.process_index} using seed: {cfg.seed}")
|
|
|
|
self._init_optimizer()
|
|
|
|
if training_state is not None and not self._restore_training_state(training_state):
|
|
initial_step = 0
|
|
resuming = False
|
|
|
|
# Initialize W&B after restore so we only resume the run when state restore succeeds.
|
|
resume_run_id = training_state.wandb_run_id if resuming and training_state is not None else None
|
|
self._init_wandb(resume_run_id=resume_run_id)
|
|
|
|
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()
|
|
|
|
remaining_steps = cfg.optimization.steps - initial_step
|
|
if remaining_steps <= 0:
|
|
raise ValueError(
|
|
f"No remaining training steps: initial_step={initial_step} >= "
|
|
f"target_steps={cfg.optimization.steps}. Nothing to train."
|
|
)
|
|
|
|
if resuming:
|
|
logger.info(f"🚀 Resuming training from step {initial_step} → {cfg.optimization.steps}")
|
|
else:
|
|
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=remaining_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 = initial_step
|
|
|
|
peak_mem_during_training = start_mem
|
|
|
|
sampled_videos_paths = None
|
|
|
|
with progress:
|
|
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(remaining_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
|
|
|
|
output = self._training_step(batch)
|
|
self._accelerator.backward(output.loss.mean())
|
|
|
|
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
|
|
step_loss = output.loss.detach().mean().item()
|
|
|
|
progress.update_training(
|
|
loss=step_loss,
|
|
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:
|
|
# Track per-element loss by sigma bucket
|
|
self._sigma_tracker.update(output.sigma.cpu().tolist(), output.loss.detach().cpu().tolist())
|
|
metrics = {
|
|
"train/loss": step_loss,
|
|
"train/learning_rate": current_lr,
|
|
"train/step_time": step_time,
|
|
"train/global_step": self._global_step,
|
|
}
|
|
metrics.update(self._sigma_tracker.get_metrics())
|
|
self._log_metrics(metrics)
|
|
|
|
# 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
|
|
steps_done = self._global_step - initial_step
|
|
if steps_done > 0:
|
|
total_estimated = elapsed / steps_done * remaining_steps
|
|
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: {step_loss:.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 = remaining_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]]) -> TrainingStepOutput:
|
|
"""Perform a single training step using the configured strategy."""
|
|
# Apply embedding connectors to transform pre-computed text embeddings
|
|
conditions = batch["conditions"]
|
|
|
|
if "video_prompt_embeds" in conditions:
|
|
# New format: separate video/audio features from precompute()
|
|
video_features = conditions["video_prompt_embeds"]
|
|
audio_features = conditions.get("audio_prompt_embeds")
|
|
else:
|
|
# Legacy format: single prompt_embeds tensor — duplicate for both modalities
|
|
video_features = conditions["prompt_embeds"]
|
|
audio_features = conditions["prompt_embeds"]
|
|
|
|
mask = conditions["prompt_attention_mask"]
|
|
additive_mask = convert_to_additive_mask(mask, video_features.dtype)
|
|
video_embeds, audio_embeds, attention_mask = self._embeddings_processor.create_embeddings(
|
|
video_features, audio_features, additive_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)
|
|
sigma = model_inputs.video.sigma.detach() if model_inputs.video.enabled else model_inputs.audio.sigma.detach()
|
|
|
|
return TrainingStepOutput(loss=loss, sigma=sigma)
|
|
|
|
@free_gpu_memory_context(after=True)
|
|
def _load_text_encoder_and_cache_embeddings(self) -> list[CachedPromptEmbeddings] | None:
|
|
"""Load text encoder + embeddings processor, compute and cache validation embeddings."""
|
|
|
|
# This method:
|
|
# 1. Loads the pure Gemma text encoder on GPU
|
|
# 2. Loads the embeddings processor (feature extractor + connectors)
|
|
# 3. If validation prompts are configured, computes and caches their embeddings
|
|
# 4. Unloads the Gemma model entirely, keeps the embeddings processor for training
|
|
|
|
# Load text encoder (pure Gemma LLM) on GPU
|
|
logger.debug("Loading text encoder...")
|
|
text_encoder = load_text_encoder(
|
|
gemma_model_path=self._config.model.text_encoder_path,
|
|
device="cuda",
|
|
dtype=torch.bfloat16,
|
|
load_in_8bit=self._config.acceleration.load_text_encoder_in_8bit,
|
|
)
|
|
|
|
# Load embeddings processor (feature extractor + connectors)
|
|
logger.debug("Loading embeddings processor...")
|
|
self._embeddings_processor = load_embeddings_processor(
|
|
checkpoint_path=self._config.model.model_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:
|
|
pos_hs, pos_mask = text_encoder.encode(prompt)
|
|
pos_out = self._embeddings_processor.process_hidden_states(pos_hs, pos_mask)
|
|
|
|
neg_hs, neg_mask = text_encoder.encode(self._config.validation.negative_prompt)
|
|
neg_out = self._embeddings_processor.process_hidden_states(neg_hs, neg_mask)
|
|
|
|
cached_embeddings.append(
|
|
CachedPromptEmbeddings(
|
|
video_context_positive=pos_out.video_encoding.cpu(),
|
|
audio_context_positive=pos_out.audio_encoding.cpu(),
|
|
video_context_negative=neg_out.video_encoding.cpu(),
|
|
audio_context_negative=(
|
|
neg_out.audio_encoding.cpu() if neg_out.audio_encoding is not None else None
|
|
),
|
|
)
|
|
)
|
|
|
|
# Unload Gemma model and feature extractor, keep only connectors for training
|
|
del text_encoder
|
|
self._embeddings_processor.feature_extractor = None
|
|
|
|
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._embeddings_processor 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.info(f'Quantizing model with "{self._config.acceleration.quantization}". This may take a while...')
|
|
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, then resolve resume state."""
|
|
if not self._config.model.load_checkpoint:
|
|
self._resume_state: tuple[int, TrainingState | None] = (0, None)
|
|
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}")
|
|
self._resume_state = (0, None)
|
|
return
|
|
|
|
self._loaded_checkpoint_path = checkpoint_path
|
|
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)
|
|
|
|
self._resume_state = self._resolve_resume_state()
|
|
|
|
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 _resolve_resume_state(self) -> tuple[int, TrainingState | None]:
|
|
"""Determine resume state by looking for a training state file next to the loaded checkpoint.
|
|
Returns (initial_step, TrainingState or None).
|
|
If no_resume config is set, no checkpoint loaded, or no state file found: returns (0, None).
|
|
"""
|
|
if self._config.checkpoints.no_resume or self._loaded_checkpoint_path is None:
|
|
return 0, None
|
|
|
|
state = self._load_training_state(self._loaded_checkpoint_path)
|
|
if state is None:
|
|
return 0, None
|
|
|
|
fp = state.config_fingerprint
|
|
cfg = self._config
|
|
mismatches: list[str] = []
|
|
if fp.optimizer_type != cfg.optimization.optimizer_type:
|
|
mismatches.append(f"optimizer_type: {fp.optimizer_type} → {cfg.optimization.optimizer_type}")
|
|
if fp.scheduler_type != cfg.optimization.scheduler_type:
|
|
mismatches.append(f"scheduler_type: {fp.scheduler_type} → {cfg.optimization.scheduler_type}")
|
|
if fp.training_mode != cfg.model.training_mode:
|
|
mismatches.append(f"training_mode: {fp.training_mode} → {cfg.model.training_mode}")
|
|
if (
|
|
cfg.model.training_mode == "lora"
|
|
and cfg.lora is not None
|
|
and fp.lora_rank is not None
|
|
and fp.lora_rank != cfg.lora.rank
|
|
):
|
|
mismatches.append(f"lora_rank: {fp.lora_rank} → {cfg.lora.rank}")
|
|
if mismatches:
|
|
logger.warning(
|
|
f"⚠️ Training state config mismatch ({', '.join(mismatches)}). "
|
|
"Starting from step 0. Set checkpoints.no_resume=true to silence this warning."
|
|
)
|
|
return 0, None
|
|
|
|
if state.global_step < 0:
|
|
logger.warning(f"⚠️ Training state has invalid global_step={state.global_step!r}. Starting from step 0.")
|
|
return 0, None
|
|
logger.info(f"📌 Resuming from step {state.global_step}")
|
|
return state.global_step, state
|
|
|
|
@staticmethod
|
|
def _load_training_state(checkpoint_path: Path) -> TrainingState | None:
|
|
"""Load training state file that corresponds to a checkpoint weights file."""
|
|
match = re.search(r"step_(\d+)", checkpoint_path.name)
|
|
if not match:
|
|
return None
|
|
|
|
step_str = match.group(1)
|
|
state_path = checkpoint_path.parent / f"training_state_step_{step_str}.pt"
|
|
|
|
if not state_path.exists():
|
|
return None
|
|
|
|
try:
|
|
raw: dict = torch.load(state_path, map_location="cpu", weights_only=False)
|
|
state = TrainingState.from_save_dict(raw)
|
|
logger.info(f"📥 Loaded training state from {state_path}")
|
|
return state
|
|
except Exception as e:
|
|
logger.warning(f"⚠️ Failed to load training state from {state_path}: {e}. Starting from step 0.")
|
|
return None
|
|
|
|
def _restore_training_state(self, training_state: TrainingState) -> bool:
|
|
"""Restore optimizer, scheduler, and RNG states from a loaded TrainingState.
|
|
Must be called after _init_optimizer() (which calls accelerator.prepare).
|
|
Returns True if restore succeeded, False if it failed (caller should fall back to step 0).
|
|
"""
|
|
try:
|
|
if training_state.optimizer_state_dict is not None:
|
|
self._optimizer.load_state_dict(training_state.optimizer_state_dict)
|
|
logger.debug("Restored optimizer state (full mode)")
|
|
|
|
if training_state.lr_scheduler_state_dict is not None and self._lr_scheduler is not None:
|
|
self._lr_scheduler.load_state_dict(training_state.lr_scheduler_state_dict)
|
|
logger.debug("Restored LR scheduler state")
|
|
except Exception as e:
|
|
logger.warning(f"⚠️ Failed to restore training state: {e}. Starting from step 0.")
|
|
return False
|
|
|
|
rng = training_state.rng_states
|
|
if self._accelerator.num_processes > 1:
|
|
logger.debug("Skipping RNG restore in multi-process mode (only main process state was saved)")
|
|
else:
|
|
if rng.torch_state is not None:
|
|
torch.random.set_rng_state(rng.torch_state)
|
|
if rng.cuda_state is not None and torch.cuda.is_available():
|
|
torch.cuda.set_rng_state(rng.cuda_state)
|
|
logger.debug("Restored RNG states")
|
|
|
|
return True
|
|
|
|
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}")
|
|
|
|
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),
|
|
T_mult=params.pop("T_mult", 1),
|
|
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()
|
|
@free_gpu_memory_context(after=True)
|
|
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
|
|
|
|
# Zero gradients and free GPU memory to reclaim memory before validation sampling
|
|
self._optimizer.zero_grad(set_to_none=True)
|
|
free_gpu_memory()
|
|
|
|
# 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,
|
|
reference_downscale_factor=self._config.validation.reference_downscale_factor,
|
|
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_sampling_rate if audio is not None else None,
|
|
)
|
|
video_paths.append(output_path)
|
|
|
|
# Clean up progress tasks
|
|
sampling_ctx.cleanup()
|
|
|
|
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)
|
|
|
|
# Determine save precision
|
|
save_dtype = torch.bfloat16 if self._config.checkpoints.precision == "bfloat16" else torch.float32
|
|
|
|
# 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()}
|
|
|
|
# Cast to configured precision
|
|
state_dict = {k: v.to(save_dtype) if isinstance(v, Tensor) else v for k, v in state_dict.items()}
|
|
|
|
# Build metadata for safetensors file
|
|
metadata = self._build_checkpoint_metadata()
|
|
|
|
# Save to disk with metadata
|
|
save_file(state_dict, saved_weights_path, metadata=metadata)
|
|
else:
|
|
# Cast to configured precision
|
|
full_state_dict = {k: v.to(save_dtype) if isinstance(v, Tensor) else v for k, v in full_state_dict.items()}
|
|
|
|
# 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}")
|
|
|
|
self._checkpoint_paths.append(saved_weights_path)
|
|
self._cleanup_checkpoints()
|
|
|
|
self._save_training_state(save_dir)
|
|
|
|
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 checkpoint: {old_checkpoint}")
|
|
self._checkpoint_paths = self._checkpoint_paths[-self._config.checkpoints.keep_last_n :]
|
|
|
|
def _save_training_state(self, save_dir: Path) -> None:
|
|
"""Save training state alongside checkpoint for resume.
|
|
Respects checkpoints.save_training_state config:
|
|
- "full": optimizer + scheduler + RNG + step + wandb_run_id
|
|
- "minimal": scheduler + RNG + step + wandb_run_id
|
|
- "off": skip entirely
|
|
"""
|
|
if not IS_MAIN_PROCESS:
|
|
return
|
|
|
|
mode = self._config.checkpoints.save_training_state
|
|
if mode == "off":
|
|
return
|
|
|
|
is_fsdp = self._accelerator.distributed_type == DistributedType.FSDP
|
|
|
|
optimizer_state = None
|
|
if mode == "full":
|
|
if is_fsdp:
|
|
logger.warning(
|
|
"⚠️ save_training_state='full' is not supported with FSDP. "
|
|
"Saving 'minimal' state (scheduler + RNG only)."
|
|
)
|
|
else:
|
|
optimizer_state = self._optimizer.state_dict()
|
|
|
|
state = TrainingState(
|
|
global_step=self._global_step,
|
|
config_fingerprint=ConfigFingerprint(
|
|
optimizer_type=self._config.optimization.optimizer_type,
|
|
scheduler_type=self._config.optimization.scheduler_type,
|
|
training_mode=self._config.model.training_mode,
|
|
lora_rank=self._config.lora.rank if self._config.lora is not None else None,
|
|
),
|
|
rng_states=RngStates(
|
|
torch_state=torch.random.get_rng_state(),
|
|
cuda_state=torch.cuda.get_rng_state() if torch.cuda.is_available() else None,
|
|
),
|
|
lr_scheduler_state_dict=self._lr_scheduler.state_dict() if self._lr_scheduler is not None else None,
|
|
optimizer_state_dict=optimizer_state,
|
|
wandb_run_id=self._wandb_run.id if self._wandb_run is not None else None,
|
|
)
|
|
|
|
state_path = save_dir / f"training_state_step_{self._global_step:05d}.pt"
|
|
tmp_path = state_path.with_suffix(".pt.tmp")
|
|
try:
|
|
torch.save(state.to_save_dict(), tmp_path)
|
|
except Exception:
|
|
if tmp_path.exists():
|
|
tmp_path.unlink()
|
|
raise
|
|
tmp_path.rename(state_path)
|
|
|
|
file_size_gb = state_path.stat().st_size / (1024**3)
|
|
if file_size_gb > 1.0 and not self._training_state_size_warned:
|
|
self._training_state_size_warned = True
|
|
logger.warning(
|
|
f"⚠️ Training state file is {file_size_gb:.1f} GB (full mode includes optimizer state). "
|
|
f'Set checkpoints.save_training_state="minimal" to save only scheduler/RNG/step (~few KB), '
|
|
f'or "off" to disable entirely.'
|
|
)
|
|
|
|
if not self._training_state_paths or self._training_state_paths[-1] != state_path:
|
|
self._training_state_paths.append(state_path)
|
|
self._cleanup_training_states()
|
|
|
|
rel_path = state_path.relative_to(self._config.output_dir)
|
|
logger.debug(f"Training state saved to {rel_path}")
|
|
|
|
def _cleanup_training_states(self) -> None:
|
|
"""Clean up old training state files, using the same keep_last_n as checkpoints."""
|
|
keep_n = self._config.checkpoints.keep_last_n
|
|
if 0 < keep_n < len(self._training_state_paths):
|
|
to_remove = self._training_state_paths[:-keep_n]
|
|
for old_state in to_remove:
|
|
if old_state.exists():
|
|
old_state.unlink()
|
|
logger.debug(f"Removed old training state: {old_state}")
|
|
self._training_state_paths = self._training_state_paths[-keep_n:]
|
|
|
|
def _build_checkpoint_metadata(self) -> dict[str, str]:
|
|
"""Build metadata dictionary for safetensors checkpoint.
|
|
Delegates to the training strategy to get strategy-specific metadata
|
|
that downstream inference pipelines may need.
|
|
Returns:
|
|
Dictionary of string key-value pairs for safetensors metadata.
|
|
Values are converted to strings for safetensors compatibility.
|
|
"""
|
|
raw_metadata = self._training_strategy.get_checkpoint_metadata()
|
|
# Convert all values to strings for safetensors compatibility
|
|
metadata = {k: str(v) for k, v in raw_metadata.items()}
|
|
if metadata:
|
|
logger.info(f"Saving checkpoint metadata: {metadata}")
|
|
return metadata
|
|
|
|
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, resume_run_id: str | None = None) -> 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
|
|
init_kwargs: dict[str, Any] = {
|
|
"project": wandb_config.project,
|
|
"entity": wandb_config.entity,
|
|
"name": Path(self._config.output_dir).name,
|
|
"tags": wandb_config.tags,
|
|
"config": self._config.model_dump(),
|
|
}
|
|
if resume_run_id is not None:
|
|
init_kwargs["id"] = resume_run_id
|
|
init_kwargs["resume"] = "allow"
|
|
run = wandb.init(**init_kwargs)
|
|
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")
|
|
|
|
if is_image:
|
|
samples = [
|
|
wandb.Image(str(path), caption=prompt) for path, prompt in zip(sample_paths, prompts, strict=True)
|
|
]
|
|
else:
|
|
samples = [
|
|
wandb.Video(str(path), caption=prompt, format=path.suffix.lower().lstrip("."))
|
|
for path, prompt in zip(sample_paths, prompts, strict=True)
|
|
]
|
|
self._wandb_run.log({"validation_samples": samples}, step=self._global_step)
|