# coding: utf-8 # Standard library imports import json import os import warnings from copy import deepcopy from dataclasses import asdict from datetime import datetime from time import time from typing import List, Optional warnings.filterwarnings( "ignore", message=".*torch\\.cuda\\.amp\\.custom_(fwd|bwd).*deprecated.*", category=FutureWarning, module="deepspeed.runtime.zero.linear", ) # Third-party package imports import numpy as np import torch import torch.distributed as dist import torch.nn as nn import wandb from safetensors.torch import load_file from torch.utils.data import DataLoader from transformers.optimization import ( get_constant_schedule_with_warmup, get_cosine_with_min_lr_schedule_with_warmup, ) from modeling.qwen2_5_vl.configuration_qwen2_5_vl import Qwen2_5_VLVisionConfig # Local repository imports from common.utils.basic import get_global_rank from common.utils.fs import download, mkdir from common.utils.misc import AutoEncoderParams, tuple_mul from common.val.utils import ( decode_text_interleave, decode_video_tensor, make_padded_latent, map_splits_to_samples, ) from config.config_factory import ModelArguments, DataArguments, TrainingArguments from data.dataset_base_train import DataConfig, PackedDataset, simple_custom_collate from modeling.lance import Lance, Qwen2ForCausalLM from modeling.qwen2 import Qwen2Tokenizer from modeling.qwen2.modeling_qwen2 import Qwen2Config from modeling.vae.wan.model import WanVideoVAE from modeling.vit.qwen2_5_vl_vit import Qwen2_5_VisionTransformerPretrainedModel from train.fsdp_utils import FSDPCheckpoint, FSDPConfig, fsdp_ema_setup, fsdp_ema_update def setup_output_paths( training_args: TrainingArguments, logger, global_rank: int, ): run_output_dir = os.path.join(training_args.outputs_dir, training_args.wandb_name) training_args.run_output_dir = run_output_dir training_args.config_dir = os.path.join(run_output_dir, "configs") training_args.ckpt_dir = os.path.join(run_output_dir, "checkpoints") training_args.wandb_dir = os.path.join(run_output_dir, "wandb") if global_rank == 0: mkdir(training_args.config_dir) mkdir(training_args.ckpt_dir) if training_args.wandb_offline: mkdir(training_args.wandb_dir) logger.info(f"training_args.run_output_dir: {training_args.run_output_dir}") logger.info(f"training_args.config_dir: {training_args.config_dir}") logger.info(f"training_args.ckpt_dir: {training_args.ckpt_dir}") if training_args.wandb_offline: logger.info(f"training_args.wandb_dir: {training_args.wandb_dir}") if dist.is_available() and dist.is_initialized(): dist.barrier() def setup_rank0_logging_and_wandb( model_args: ModelArguments, data_args: DataArguments, training_args: TrainingArguments, logger, global_rank: int, ): if global_rank != 0: return timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") wandb_name = training_args.wandb_name[:64] wandb_id = f"{wandb_name}-{timestamp}"[:64] wandb_init_kwargs = { "project": training_args.wandb_project, "name": wandb_name, "id": wandb_id, "resume": training_args.wandb_resume, "mode": "offline" if training_args.wandb_offline else "online", } if training_args.wandb_offline: wandb_init_kwargs["dir"] = training_args.wandb_dir wandb.init(**wandb_init_kwargs) wandb.config.update({**vars(training_args), **vars(model_args), **vars(data_args)}) def prepare_model_paths(model_args: ModelArguments, training_args: TrainingArguments): if training_args.load_from_lance_checkpoint: model_args.model_path = download(model_args.model_path, add_hash_suffix=True) required_files = [ "llm_config.json", "generation_config.json", "merges.txt", "tokenizer.json", "tokenizer_config.json", "vocab.json", ] missing_files = [ filename for filename in required_files if not os.path.exists(os.path.join(model_args.model_path, filename)) ] if missing_files: raise FileNotFoundError( "MODEL_PATH must contain all Lance LLM/tokenizer files when " f"load_from_lance_checkpoint=True. Missing in {model_args.model_path}: " f"{', '.join(missing_files)}" ) else: model_args.llm_path = download(model_args.llm_path) model_args.vit_path = download(model_args.vit_path) def freeze_model_components(model: Lance, training_args: TrainingArguments, log_rank0): if training_args.freeze_llm: model.language_model.eval() for param in model.language_model.parameters(): param.requires_grad = False log_rank0("Freeze all LLM parameters") if training_args.freeze_llm_embed_tokens: model.language_model.freeze_embed_tokens() model.language_model.freeze_lm_head() log_rank0("Freeze LLM token embeddings and lm_head") if training_args.freeze_und_params: model.language_model.freeze_und_params() log_rank0("Freeze UND parameters") if training_args.freeze_vit and training_args.visual_und: model.vit_model.eval() for param in model.vit_model.parameters(): param.requires_grad = False log_rank0("Freeze VIT parameters") if training_args.freeze_vit_connector: model.connector.eval() for param in model.connector.parameters(): param.requires_grad = False log_rank0("Freeze VIT connector parameters") def save_trainable_parameters(model: Lance, fsdp_model: torch.nn.Module, training_args: TrainingArguments, logger, global_rank: int): if global_rank != 0: return report_path = os.path.join(training_args.config_dir, "trainable_parameters.txt") sep = "=" * 40 + " check requires_grad " + "=" * 40 lines = [ sep, "fsdp_model:", str(fsdp_model), sep, ] for name, param in model.named_parameters(): if param.requires_grad: lines.append(f"{name}: {param.requires_grad}") lines.append(sep) mkdir(os.path.dirname(report_path)) with open(report_path, "w", encoding="utf-8") as f: f.write("\n".join(lines)) f.write("\n") logger.info(f"Saved trainable parameter report to {report_path}") def save_checkpoint_load_report( report_dir: Optional[str], report_name: str, title: str, matched: int, not_matched: int, missing, unexpected, log_rank0, ): log_rank0(f"{title}: matched={matched}, not_matched={not_matched}, missing={len(missing)}, unexpected={len(unexpected)}") if get_global_rank() != 0 or report_dir is None: return mkdir(report_dir) report_path = os.path.join(report_dir, report_name) lines = [ title, f"matched: {matched}", f"not_matched: {not_matched}", f"missing_count: {len(missing)}", f"unexpected_count: {len(unexpected)}", "", "missing_keys:", *[str(item) for item in missing], "", "unexpected_keys:", *[str(item) for item in unexpected], ] with open(report_path, "w", encoding="utf-8") as f: f.write("\n".join(lines)) f.write("\n") log_rank0(f"Saved checkpoint load report to {report_path}") def setup_model_components(model_args: ModelArguments, training_args: TrainingArguments, log_rank0): if training_args.load_from_lance_checkpoint: llm_config: Qwen2Config = Qwen2Config.from_json_file(os.path.join(model_args.model_path, "llm_config.json")) else: llm_config: Qwen2Config = Qwen2Config.from_pretrained(model_args.llm_path) llm_config.layer_module = model_args.layer_module llm_config.qk_norm = model_args.llm_qk_norm llm_config.qk_norm_und = model_args.llm_qk_norm_und llm_config.qk_norm_gen = model_args.llm_qk_norm_gen log_rank0(f"llm_config.qk_norm: {llm_config.qk_norm}, llm_config.qk_norm_und: {llm_config.qk_norm_und}, llm_config.qk_norm_gen: {llm_config.qk_norm_gen}") llm_config.tie_word_embeddings = model_args.tie_word_embeddings llm_config.freeze_und = training_args.freeze_und llm_config.apply_qwen_2_5_vl_pos_emb = training_args.apply_qwen_2_5_vl_pos_emb if training_args.load_from_lance_checkpoint or training_args.init_from_vlm_checkpoint: language_model: Qwen2ForCausalLM = Qwen2ForCausalLM(llm_config) else: language_model: Qwen2ForCausalLM = Qwen2ForCausalLM.from_pretrained(model_args.llm_path, config=llm_config) vit_config, vit_model = None, None if training_args.visual_und: if model_args.vit_type in ("qwen2_5_vl", "qwen_2_5_vl_original"): vit_config = Qwen2_5_VLVisionConfig.from_pretrained(model_args.vit_path) vit_model = Qwen2_5_VisionTransformerPretrainedModel(vit_config) vit_weights = load_file(os.path.join(model_args.vit_path, "vit.safetensors")) msg = vit_model.load_state_dict(vit_weights, strict=True) log_rank0(f"Load vit model weights: {msg}, from {model_args.vit_path}") else: raise ValueError(f"Unsupported vit_type: {model_args.vit_type}") del vit_weights import gc; gc.collect(); torch.cuda.empty_cache() if training_args.visual_gen: if training_args.vae_model_type.lower() in ("wan", "wanvideo", "wan-video"): vae_model = WanVideoVAE() else: raise ValueError(f"Unsupported vae_model_type: {training_args.vae_model_type}") vae_config: AutoEncoderParams = deepcopy(vae_model.vae_config) else: vae_model = None vae_config = None return llm_config, language_model, vit_config, vit_model, vae_model, vae_config def build_fsdp_config(training_args: TrainingArguments): return FSDPConfig( sharding_strategy=training_args.sharding_strategy, backward_prefetch=training_args.backward_prefetch, cpu_offload=training_args.cpu_offload, num_replicate=training_args.num_replicate, num_shard=training_args.num_shard, use_orig_params=True, ) def build_lr_scheduler(optimizer: torch.optim.Optimizer, training_args: TrainingArguments): if training_args.lr_scheduler == "cosine": return get_cosine_with_min_lr_schedule_with_warmup( optimizer=optimizer, num_warmup_steps=training_args.warmup_steps, num_training_steps=training_args.total_steps, min_lr=training_args.min_lr, num_cycles=5, ) if training_args.lr_scheduler == "constant": return get_constant_schedule_with_warmup( optimizer=optimizer, num_warmup_steps=training_args.warmup_steps, ) raise ValueError(f"Unknown lr_scheduler: {training_args.lr_scheduler}") def load_training_state( optimizer: torch.optim.Optimizer, scheduler, model_args: ModelArguments, data_args: DataArguments, training_args: TrainingArguments, resume_from, resume_model_only: bool, fsdp_config: FSDPConfig, global_rank: int, world_size: int, ): if not resume_model_only: return FSDPCheckpoint.try_load_train_state( resume_from, optimizer, scheduler, fsdp_config, ) train_step = 0 data_status = None if not training_args.load_data_status: return optimizer, scheduler, train_step, data_status try: train_step = int(os.path.basename(os.path.normpath(model_args.model_path)).split(".")[0]) + 1 data_status_path = os.path.join(model_args.model_path, "data_status.pt") data_status_all = torch.load(data_status_path, weights_only=True, map_location="cpu") assert world_size == len(data_status_all), f"WORLD_SIZE ({world_size}) must be equal to the length of data_status_all ({len(data_status_all)})" dataset_names, data_worker_ids = [], [] for d_status in data_status_all: for d_name, d_worker_info in d_status.items(): dataset_names.append(d_name) data_worker_ids.extend(d_worker_info.keys()) dataset_names = list(set(dataset_names)) data_worker_ids = list(set(data_worker_ids)) assert data_args.num_workers == len(data_worker_ids), f"num_workers ({data_args.num_workers}) must be equal to the length of data_worker_ids ({len(data_worker_ids)})" data_status = data_status_all[global_rank] print( f"Successfully load train_step and data_status ***** \n" f"train_step: {train_step}, data_status_all: {data_status_all}, global_rank: {global_rank}, data_status: {data_status}\n" ) except Exception: train_step = 0 data_status = None print( f"Failed to load train_step and data_status ***** \n" f"train_step: {train_step}, data_status: {data_status}\n" ) return optimizer, scheduler, train_step, data_status def build_train_dataset_config( data_args: DataArguments, model_args: ModelArguments, training_args: TrainingArguments, vae_config: Optional[AutoEncoderParams], ): dataset_config = DataConfig.from_yaml(data_args.dataset_config_file) if training_args.visual_und: dataset_config.vit_patch_size = model_args.vit_patch_size dataset_config.vit_patch_size_temporal = model_args.vit_patch_size_temporal dataset_config.vit_max_num_patch_per_side = model_args.vit_max_num_patch_per_side if training_args.visual_gen: assert len(model_args.latent_patch_size) == 3, "len(latent_patch_size) must be 3" vae_downsample = tuple_mul( model_args.latent_patch_size, ( vae_config.downsample_temporal, vae_config.downsample_spatial, vae_config.downsample_spatial, ), ) dataset_config.latent_patch_size = model_args.latent_patch_size dataset_config.vae_downsample = vae_downsample dataset_config.max_latent_size = model_args.max_latent_size dataset_config.max_num_frames = model_args.max_num_frames dataset_config.text_cond_dropout_prob = model_args.text_cond_dropout_prob dataset_config.vae_cond_dropout_prob = model_args.vae_cond_dropout_prob dataset_config.vit_cond_dropout_prob = model_args.vit_cond_dropout_prob return dataset_config def compute_training_loss( loss_dict: dict, data: dict, ce_loss_weights, training_args: TrainingArguments, device: int, world_size: int, ): loss = torch.tensor(0.0, device=device) ce = loss_dict["ce"] if ce is not None: total_ce_tokens = torch.tensor(len(data["ce_loss_indexes"]), device=device) dist.all_reduce(total_ce_tokens, op=dist.ReduceOp.SUM) if training_args.ce_loss_reweighting: ce = ce * ce_loss_weights total_ce_loss_weights = ce_loss_weights.sum() dist.all_reduce(total_ce_loss_weights, op=dist.ReduceOp.SUM) ce = ce.sum() * world_size / total_ce_loss_weights else: ce = ce.sum() * world_size / total_ce_tokens loss_dict["ce"] = ce.detach() loss = loss + ce * training_args.ce_weight else: loss_dict["ce"] = torch.tensor(0, device=device) total_ce_tokens = torch.tensor(0, device=device) total_mse_tokens = loss_dict.pop("total_mse_tokens") frame_mse = loss_dict.pop("frame_mse") if training_args.visual_gen: mse = loss_dict["mse"] if mse is not None: total_mse_tokens = torch.tensor(total_mse_tokens, device=device) dist.all_reduce(total_mse_tokens, op=dist.ReduceOp.SUM) mse = mse.mean(dim=-1).sum() * world_size / total_mse_tokens loss_dict["mse"] = mse.detach() loss = loss + mse * training_args.mse_weight else: loss_dict["mse"] = torch.tensor(0, device=device) total_mse_tokens = torch.tensor(0, device=device) if frame_mse is not None: total_frame_mse_tokens = torch.tensor(sum(data["key_frame_mask"] == 1), device=device) dist.all_reduce(total_frame_mse_tokens, op=dist.ReduceOp.SUM) frame_mse = frame_mse.mean(dim=-1).sum() * world_size / total_frame_mse_tokens loss_dict["frame_mse"] = frame_mse.detach() loss = loss + frame_mse * training_args.mse_weight else: loss_dict["frame_mse"] = torch.tensor(0, device=device) else: loss_dict["mse"] = torch.tensor(0, device=device) total_mse_tokens = torch.tensor(0, device=device) return loss, loss_dict, total_ce_tokens, total_mse_tokens def optimizer_step_with_ema( loss, fsdp_model: torch.nn.Module, ema_model, optimizer: torch.optim.Optimizer, scheduler, training_args: TrainingArguments, curr_step: int, log_rank0, ): optimizer.zero_grad() loss.backward() total_norm = fsdp_model.clip_grad_norm_(training_args.max_grad_norm) jump_first_step = getattr(training_args, "jump_first_step", False) and (curr_step == 0) if not jump_first_step: optimizer.step() scheduler.step() if training_args.use_ema and ema_model is not None: if curr_step == training_args.ema_start_steps: fsdp_ema_update(ema_model, fsdp_model, decay=0.0) log_rank0(f"[EMA] initialized at step {curr_step}") elif curr_step > training_args.ema_start_steps: fsdp_ema_update(ema_model, fsdp_model, decay=training_args.ema) else: log_rank0(f"Jump step #{curr_step} without updating parameters.") return total_norm def log_training_metrics( loss_dict: dict, total_mse_tokens, total_ce_tokens, total_norm, data: dict, optimizer: torch.optim.Optimizer, progress_bar, training_args: TrainingArguments, curr_step: int, start_time: float, device: int, world_size: int, global_rank: int, ): if curr_step % training_args.log_every != 0: return start_time total_samples = torch.tensor(len(data["sample_lens"]), device=device) dist.all_reduce(total_samples, op=dist.ReduceOp.SUM) torch.cuda.synchronize() end_time = time() steps_per_sec = training_args.log_every / (end_time - start_time) message = f"(step={curr_step:07d}) " wandb_log = {} for key, value in loss_dict.items(): avg_loss = torch.tensor(value.item(), device=device) dist.all_reduce(avg_loss, op=dist.ReduceOp.SUM) avg_loss = avg_loss.item() / world_size message += f"Train Loss {key}: {avg_loss:.4f}, " wandb_log[key] = avg_loss message += f"Train Steps/Sec: {steps_per_sec:.2f}, " logs = {"loss": message, "lr": optimizer.param_groups[0]["lr"]} progress_bar.set_postfix(**logs, refresh=False) progress_bar.update(training_args.log_every) wandb_log["lr"] = optimizer.param_groups[0]["lr"] wandb_log["total_mse_tokens"] = total_mse_tokens.item() wandb_log["total_ce_tokens"] = total_ce_tokens.item() wandb_log["total_norm"] = total_norm.item() wandb_log["total_samples"] = total_samples.item() mem_allocated = torch.tensor(torch.cuda.max_memory_allocated() / 1024**2, device=device) dist.all_reduce(mem_allocated, op=dist.ReduceOp.MAX) wandb_log["mem_allocated"] = mem_allocated mem_cache = torch.tensor(torch.cuda.max_memory_reserved() / 1024**2, device=device) dist.all_reduce(mem_cache, op=dist.ReduceOp.MAX) wandb_log["mem_cache"] = mem_cache if global_rank == 0: wandb.log(wandb_log, step=curr_step) return time() def setup_ema_and_load_checkpoint( model: Lance, training_args: TrainingArguments, fsdp_config: FSDPConfig, load_ckpt, resume_from=None, resume_model_only: bool = True, logger=None, ): ema_model = None if training_args.use_ema: ema_model = deepcopy(model) load_ckpt(ema_model, ema=True) load_ckpt(model, ema=False) if resume_from is not None and not resume_model_only: model, ema_model = FSDPCheckpoint.try_load_ckpt( resume_from=resume_from, logger=logger, model=model, ema_model=ema_model, resume_from_ema=training_args.finetune_from_ema, report_dir=training_args.config_dir, ) if ema_model is not None: ema_model = fsdp_ema_setup(ema_model, fsdp_config) return ema_model def prepare_checkpoint_loader( model_args: ModelArguments, training_args: TrainingArguments, llm_config: Qwen2Config, language_model: Qwen2ForCausalLM, tokenizer: Qwen2Tokenizer, num_new_tokens: int, log_rank0, report_dir: Optional[str] = None, ): if training_args.copy_init_moe: language_model.init_moe() log_rank0("Copy init moe params: copy llm weight to gen.") should_untie_lm_head = model_args.tie_word_embeddings if should_untie_lm_head: model_args.tie_word_embeddings = False llm_config.tie_word_embeddings = False def load_ckpt(model: Lance, ema: bool = False): if training_args.load_from_lance_checkpoint: load_from_lance_checkpoint(model, model_args, log_rank0, ema=ema, report_dir=report_dir) elif training_args.init_from_vlm_checkpoint: init_from_vlm_checkpoint( model, model_args, log_rank0, report_dir=report_dir, report_name=f"vlm_checkpoint_load_report_{'ema' if ema else 'model'}.txt", ) if num_new_tokens > 0: model.language_model.resize_token_embeddings(len(tokenizer)) model.config.llm_config.vocab_size = len(tokenizer) model.language_model.config.vocab_size = len(tokenizer) log_rank0(f"Note: {num_new_tokens} new tokens are added!") else: log_rank0("Note: NO new tokens!") if model_args.vit_type.lower() == "qwen2_5_vl": from common.model.hacks import hack_qwen2_5_vl_config hack_qwen2_5_vl_config(model.language_model) model.update_tokenizer(tokenizer=tokenizer) if should_untie_lm_head: model.language_model.untie_lm_head() model.language_model.copy_new_token_rows_to_lm_head(num_new_tokens) log_rank0("Note: copy embed tokens to lm_head and untie") else: assert ( model.language_model.get_input_embeddings().weight.data.data_ptr() != model.language_model.get_output_embeddings().weight.data.data_ptr() ), "tie_word_embeddings conflict" freeze_model_components(model, training_args, log_rank0) return load_ckpt def get_image_token_id(language_model: Qwen2ForCausalLM): return language_model.config.video_token_id # ============================================================================= # Checkpoint and resume utilities # ============================================================================= def get_latest_ckpt(checkpoint_dir): step_dirs = [d for d in os.listdir(checkpoint_dir) if os.path.isdir(os.path.join(checkpoint_dir, d))] if len(step_dirs) == 0: return None step_dirs = sorted(step_dirs, key=lambda x: int(x)) latest_step_dir = os.path.join(checkpoint_dir, step_dirs[-1]) return latest_step_dir def prepare_resume_and_finetune_settings(training_args: TrainingArguments): if training_args.auto_resume: resume_from = get_latest_ckpt(training_args.ckpt_dir) if resume_from is None: resume_from = training_args.resume_from resume_model_only = training_args.resume_model_only else: resume_model_only = False else: resume_from = training_args.resume_from resume_model_only = training_args.resume_model_only return resume_from, resume_model_only # ============================================================================= # Distributed validation logging utilities # ============================================================================= def _gather_for_rank0(local_list): """Gather arbitrary Python objects from all ranks to rank 0 and flatten.""" world_size = dist.get_world_size() rank = get_global_rank() if rank == 0: gathered = [None] * world_size else: gathered = None dist.gather_object(local_list, gathered if rank == 0 else None, dst=0) if rank == 0: return [item for sub in gathered for item in sub] return None def _log_media_across_ranks(local_media_data, tag, step, fps, logger=None): """ local_media_data: List[Dict] with keys {"data", "caption", "is_image"} Gathers from all ranks and logs to wandb on rank 0. """ flat = _gather_for_rank0(local_media_data) if flat is None: return all_media = [] for i, item in enumerate(flat): # TODO: support sequence-level video display if item.get("validation_log_type", "") == "table": if i == 0: all_media = wandb.Table(columns=["id", "condition_image", "condition_text", "target_text", "target_image", "condition_video", "target_video"]) tag = f"{tag}/step_{step:06d}" if item.get("data_vit", None) is not None: data_vit = item["data_vit"] data_vit_video = None if data_vit.ndim == 3: # image data_vit = wandb.Image(data_vit) else: data_vit_video = wandb.Video(data_vit, fps=fps, format="mp4") data_vit = None if i==0: data_vit = wandb.Image(np.zeros((1,224, 224, 3), dtype=np.uint8)) else: data_vit, data_vit_video = None, None if item.get("data", None) is not None: data = item["data"] data_video = None if data.ndim == 3: # image data = wandb.Image(data) else: data_video = wandb.Video(data, fps=fps, format="mp4") data = None if i==0: data = wandb.Image(np.zeros((1,224, 224, 3), dtype=np.uint8)) else: data, data_video = None, None all_media.add_data( i, data_vit, # wandb.Image(item["data_vit"]) if item.get("data_vit", None) is not None else None, item.get("caption", ""), item.get("cap_target", ""), data, # wandb.Image(item["data"]) if item.get("data", None) is not None else None, data_vit_video, data_video, ) else: if item.get("is_image", False): data = item["data"] if item.get("data", None) is not None else item.get("data_vit", None) caption = f"cap_condition: {item.get('caption', '')}---cap_target: {item.get('cap_target', '')}" all_media.append(wandb.Image(data, caption=caption)) else: data = item["data"] if item.get("data", None) is not None else item["data_vit"] caption = f"cap_condition: {item.get('caption', '')}---cap_target: {item.get('cap_target', '')}" all_media.append(wandb.Video(data, fps=fps, format="mp4", caption=caption)) if all_media: wandb.log({tag: all_media}, step=step) if logger is not None: logger.info(f"Logged {len(flat)} items to {tag}.") # ============================================================================= # Model checkpoint loading utilities # ============================================================================= def init_from_vlm_checkpoint(model: Qwen2ForCausalLM, model_args: ModelArguments, log_rank0, report_dir: Optional[str] = None, report_name: str = "vlm_checkpoint_load_report.txt"): # NOTE: VLM checkpoint initialization goes through this path def load_safetensors_state_dict(folder_path): # Select only safetensors files and sort by filename for deterministic order safetensor_files = sorted( f for f in os.listdir(folder_path) if f.endswith(".safetensors") ) state_dict = {} for filename in safetensor_files: file_path = os.path.join(folder_path, filename) state_dict.update(load_file(file_path)) return state_dict state_dict = load_safetensors_state_dict(model_args.llm_path) # Rename parameters to match the current Lance model names for k in list(state_dict.keys()): if "visual" in k: # ViT and connector state_dict[k.replace("visual", "vit_model")] = state_dict.pop(k) else: # Add the language_model prefix state_dict["language_model." + k] = state_dict.pop(k) result = model.load_state_dict(state_dict, strict=False) missing = result.missing_keys unexpected = result.unexpected_keys # Number of matched parameters matched = len(state_dict) - len(unexpected) # Number of unmatched parameters not_matched = len(missing) + len(unexpected) save_checkpoint_load_report( report_dir=report_dir, report_name=report_name, title="Init from pretrained VLM checkpoint", matched=matched, not_matched=not_matched, missing=missing, unexpected=unexpected, log_rank0=log_rank0, ) log_rank0(f"Loading from pretrained VLM {model_args.llm_path} is finished.") del state_dict import gc; gc.collect(); torch.cuda.empty_cache() def load_from_lance_checkpoint(model: Qwen2ForCausalLM, model_args: ModelArguments, log_rank0, ema=False, report_dir: Optional[str] = None): # NOTE: Fine-tuning from a Lance checkpoint goes through this higher-priority path; prefer ema.safetensors, then model.safetensors path_dir = model_args.model_path ema_path = os.path.join(path_dir, "ema.safetensors") model_path = os.path.join(path_dir, "model.safetensors") model_path_ft = None if ema and os.path.exists(ema_path): model_path_ft = ema_path log_rank0("Found preferred EMA checkpoint for fine-tuning.") elif os.path.exists(model_path): model_path_ft = model_path log_rank0("EMA checkpoint not found. Using fallback: 'model.safetensors'.") if model_path_ft: log_rank0(f"Loading fine-tune model from: {model_path_ft}") model_state_dict = load_file(model_path_ft, device="cpu") else: raise FileNotFoundError( f"Fine-tuning failed: No valid checkpoint ('ema.safetensors' or 'model.safetensors') found in {path_dir}" ) # NOTE: position embeds are fixed sinusoidal embeddings, so we can just pop it off, # which makes it easier to adapt to different resolutions. if 'latent_pos_embed.pos_embed' in model_state_dict: model_state_dict.pop('latent_pos_embed.pos_embed') log_rank0(f"Pop `latent_pos_embed.pos_embed` from hf model") # model_state_dict.pop('vit_pos_embed.pos_embed') # TODO: check whether vit_pos_embed.pos_embed is present msg = model.load_state_dict(model_state_dict, strict=False) # strict = True | False missing = msg.missing_keys unexpected = msg.unexpected_keys matched = len(model_state_dict) - len(unexpected) # Number of matched parameters not_matched = len(missing) + len(unexpected) # Number of unmatched parameters report_name = f"lance_checkpoint_load_report_{'ema' if ema else 'model'}.txt" save_checkpoint_load_report( report_dir=report_dir, report_name=report_name, title=f"Init from Lance checkpoint ({'ema' if ema else 'model'}): {model_args.model_path}", matched=matched, not_matched=not_matched, missing=missing, unexpected=unexpected, log_rank0=log_rank0, ) del model_state_dict import gc; gc.collect(); torch.cuda.empty_cache() return msg # ============================================================================= # Training config saving # ============================================================================= def save_training_config(model_args: ModelArguments, data_args: DataArguments, training_args: TrainingArguments, logger): import fsspec """Save training, model, and data arguments to a JSON file on rank 0.""" if get_global_rank() == 0: logger.info(f"Training arguments {training_args}") logger.info(f"Model arguments {model_args}") logger.info(f"Data arguments {data_args}") all_args = { "model_args": model_args.__dict__, "data_args": data_args.__dict__, "training_args": training_args.__dict__, } mkdir(training_args.config_dir) config_path = os.path.join(training_args.config_dir, "training_config.json") try: with fsspec.open(config_path, "w") as f: json.dump(all_args, f, indent=4, default=str) logger.info(f"Saved training configuration to {config_path}") except Exception as e: logger.error(f"Failed to save training configuration: {e}") # ============================================================================= # Validation data and evaluation utilities # ============================================================================= def _get_data_mode(val_data: dict) -> str: # Prefer val_data; if missing, infer from keys without mutating the original data where possible if "data_mode" in val_data: return val_data["data_mode"] return "offline" if "padded_latent" in val_data else "online" def get_fixed_validation_data( data_args: DataArguments, model_args: ModelArguments, training_args: TrainingArguments, tokenizer: Qwen2Tokenizer, new_token_ids, image_token_id: int, GLOBAL_RANK: int, WORLD_SIZE: int, DEVICE: int, log_rank0, ): """Build and return a fixed validation batch equivalent to the existing inline implementation.""" assert data_args.val_dataset_config_file is not None and os.path.exists(data_args.val_dataset_config_file) # 1) Load the independent validation dataset config and override dropout val_dataset_config = DataConfig.from_yaml(data_args.val_dataset_config_file) val_dataset_config.text_cond_dropout_prob = model_args.val_text_cond_dropout_prob val_dataset_config.vae_cond_dropout_prob = model_args.val_vae_cond_dropout_prob val_dataset_config.vit_cond_dropout_prob = model_args.val_vit_cond_dropout_prob val_dataset_config.latent_patch_size = model_args.latent_patch_size log_rank0( f"val_dataset_config.text_cond_dropout_prob: {val_dataset_config.text_cond_dropout_prob}, " f"val_dataset_config.vae_cond_dropout_prob: {val_dataset_config.vae_cond_dropout_prob}, " f"val_dataset_config.vit_cond_dropout_prob: {val_dataset_config.vit_cond_dropout_prob}" ) val_loader = None val_dataset = None val_data_args = deepcopy(data_args) val_data_args.num_workers = min(val_data_args.num_workers, 1) try: log_rank0("Fetching a fixed batch for validation...") # 2) Dataset: keep arguments consistent with the original implementation val_dataset = PackedDataset( val_dataset_config, tokenizer=tokenizer, special_tokens=new_token_ids, local_rank=GLOBAL_RANK, # global rank, not local rank world_size=WORLD_SIZE, interpolate_pos=model_args.interpolate_pos, use_flex=training_args.use_flex, data_status=None, apply_chat_template=training_args.apply_chat_template, image_token_id=image_token_id, **asdict(val_data_args), ) # Fix order and seed val_dataset.set_epoch(training_args.validation_data_seed) # 3) DataLoader: keep arguments consistent with the original implementation val_num_workers = 0 ctx = torch.multiprocessing.get_context("spawn") if val_num_workers > 0 else None val_loader = DataLoader( val_dataset, batch_size=1, num_workers=val_num_workers, pin_memory=True, collate_fn=simple_custom_collate, # Top-level function drop_last=True, prefetch_factor=1 if val_num_workers > 0 else None, persistent_workers=True if val_num_workers > 0 else False, multiprocessing_context=ctx, ) # 4) Fetch one fixed batch and convert it to a dict val_data_cpu = next(iter(val_loader)) # val_data_cpu = val_data_cpu.cuda(DEVICE).to_dict() log_rank0("Fixed validation batch fetched, val_loader and val_dataset deleted.") return val_data_cpu finally: if val_loader is not None: del val_loader if val_dataset is not None: del val_dataset import gc; gc.collect() log_rank0("Temporary validation resources have been released.") def validate_on_fixed_batch( fsdp_model: Lance, vae_model: Optional[WanVideoVAE], tokenizer: Qwen2Tokenizer, val_data_cpu: dict, training_args: TrainingArguments, model_args: ModelArguments, data_args: DataArguments, curr_step: int, logger, new_token_ids, image_token_id: int, device: int, ): """ Extracted validation logic equivalent to the original validation block in the for-loop. """ log_rank0 = (lambda msg: logger.info(msg)) if get_global_rank() == 0 else (lambda *_: None) val_data = val_data_cpu.cuda(device).to_dict() fsdp_model.eval() try: with torch.no_grad(), torch.amp.autocast("cuda", enabled=True, dtype=torch.bfloat16): log_rank0(f"Running validation on fixed batch at step {curr_step}...") # Decode text, matching the original logic # If the validation dataset did not apply template/pos_emb during construction, decode_text_* handles it according to the switches here val_texts_gen, val_texts_und = decode_text_interleave( tokenizer, val_data, ) val_attn_modes = val_data["attn_modes"] sample_splits = map_splits_to_samples(val_data["sample_lens"], val_data["split_lens"]) val_sample_N_target = val_data["sample_N_target"] val_sample_type = val_data["sample_type"] # -------------------- GEN branch -------------------- if training_args.validation_type.lower() in ("gen", "und_gen", "gen_und") and len(val_texts_gen) > 0: data_mode = _get_data_mode(val_data) # fix: no longer define only at step == 0 # Compute video_sizes without changing the original inference config if data_mode == "offline": if curr_step == 0: # Decode all at step 0 because GT is logged val_padded_videos = vae_model.vae_decode(list(val_data["padded_latent"])) first_shape = val_padded_videos[0].shape[1:] # T,H,W else: # Later steps decode one sample only to get size and avoid extra full decodes first_shape = vae_model.vae_decode([val_data["padded_latent"][0]])[0].shape[1:] video_sizes = [first_shape for _ in (val_data["padded_latent"])] else: video_sizes = [v.shape[1:] for v in val_data["padded_videos"]] if curr_step == 0: val_padded_videos = val_data["padded_videos"] # Log GT at step 0 if curr_step == 0: if model_args.val_text_cond_dropout_prob > 0: val_texts_gen = ["NULL"] * len(video_sizes) local_gt_media_data = [] curr_sample, curr_video_tensor_index, curr_video_tensor_index_vit = 0, 0, 0 for i_gt, N_target in enumerate(val_sample_N_target[:-1]): # Remove the final padding sample left, right = sample_splits[i_gt][0], sample_splits[i_gt][-1] + 1 N_target_VIT = val_attn_modes[left:right].count("full") # Non-zero only in ti2i if val_sample_type[i_gt] != "gen": curr_video_tensor_index_vit += N_target_VIT continue curr_sample += 1 if curr_sample > training_args.validation_max_samples: break video_tensor = val_padded_videos[curr_video_tensor_index : curr_video_tensor_index + N_target] # [N], each item is [C T H W] curr_video_tensor_index += N_target v_thwc = decode_video_tensor(video_tensor) is_image = v_thwc.shape[0] == 1 # ===== Handle VIT condition features for the GEN branch ===== if N_target_VIT > 0: video_tensor_vit = val_data["padded_videos_vit"][curr_video_tensor_index_vit : curr_video_tensor_index_vit + N_target_VIT] curr_video_tensor_index_vit += N_target_VIT v_thwc_vit = decode_video_tensor(video_tensor_vit, video_type="vit") media_data_vit = v_thwc_vit[0] if is_image else np.ascontiguousarray(v_thwc_vit.transpose(0, 3, 1, 2)) else: media_data_vit = None media_data = v_thwc[0] if is_image else np.ascontiguousarray(v_thwc.transpose(0, 3, 1, 2)) cap = val_texts_gen[curr_sample-1] if (curr_sample-1) < len(val_texts_gen) else "GT_text_wrong" caption = cap local_gt_media_data.append( {"data": media_data, "caption": caption, "is_image": is_image, "data_vit": media_data_vit, "validation_log_type": training_args.validation_log_type} ) _log_media_across_ranks(local_gt_media_data, "validation_gt_samples_gen", curr_step, training_args.validation_video_saving_fps, logger) # Compute padded_latent val_data["padded_latent"] = make_padded_latent(val_data['padded_videos'], val_data['vae_data_mode'], vae_model) # Sample generation while preserving original arguments with fsdp_model.summon_full_params(fsdp_model, writeback=False, rank0_only=False): denoise_latent = fsdp_model.validation_gen( val_packed_text_ids=val_data["packed_text_ids"], val_packed_text_indexes=val_data["packed_text_indexes"], val_sample_lens=val_data["sample_lens"], val_packed_position_ids=val_data["packed_position_ids"], val_split_lens=val_data["split_lens"], val_attn_modes=val_data["attn_modes"], val_sample_N_target=val_data["sample_N_target"], val_packed_vae_token_indexes=val_data["packed_vae_token_indexes"], timestep_shift=training_args.validation_timestep_shift, num_timesteps=training_args.validation_num_timesteps, val_mse_loss_indexes=val_data.get("mse_loss_indexes", None), val_padded_latent=val_data["padded_latent"], video_sizes=video_sizes, cfg_text_scale=model_args.cfg_text_scale, cfg_interval=training_args.cfg_interval, cfg_renorm_min=training_args.cfg_renorm_min, cfg_renorm_type=training_args.cfg_renorm_type, device=device, dtype=torch.bfloat16, new_token_ids=new_token_ids, max_samples=training_args.validation_max_samples, validation_noise_seed=training_args.validation_noise_seed, apply_chat_template=training_args.apply_chat_template, apply_qwen_2_5_vl_pos_emb=training_args.apply_qwen_2_5_vl_pos_emb, image_token_id=image_token_id, val_packed_vit_token_indexes=val_data.get("packed_vit_token_indexes", None), val_packed_vit_tokens=val_data.get("packed_vit_tokens", None), vit_video_grid_thw=val_data.get("vit_video_grid_thw", None), vae_video_grid_thw=val_data["vae_video_grid_thw"], video_grid_thw=val_data.get("video_grid_thw", None), sample_task=val_data["sample_task"], sample_modality=val_data["sample_modality"], cfg_type=training_args.cfg_type, cfg_uncond_token_id=training_args.cfg_uncond_token_id, ) # Decode and log local_media_data = [] for i_val, latent in enumerate(denoise_latent): v_list = [] for latent_ in latent: v_list.append(vae_model.vae_decode([latent_])[0]) v_thwc = decode_video_tensor(v_list) is_image = v_thwc.shape[0] == 1 media_data = v_thwc[0] if is_image else np.ascontiguousarray(v_thwc.transpose(0, 3, 1, 2)) cap = val_texts_gen[i_val] if i_val < len(val_texts_gen) else "GT_text_wrong" caption = cap local_media_data.append({"data": media_data, "caption": caption, "is_image": is_image, "validation_log_type": training_args.validation_log_type}) if dist.is_available() and dist.is_initialized(): dist.barrier() _log_media_across_ranks(local_media_data, "validation_pre_samples_gen", curr_step, training_args.validation_video_saving_fps, logger) log_rank0(f"Validation(gen) at step {curr_step} finished.") log_rank0(f"cfg type is: {training_args.cfg_type}.") # -------------------- UND branch -------------------- visual_first = 0 if training_args.validation_type.lower() in ("und", "und_gen", "gen_und") and len(val_texts_und) > 0: # Prepare input video for visualization every step to avoid depending on step==0 locals # fix vis_list = [] curr_sample, curr_video_tensor_index = 0, 0 for i_gt, N_target in enumerate(val_sample_N_target[:-1]): # Remove the final padding sample left, right = sample_splits[i_gt][0], sample_splits[i_gt][-1] + 1 N_target_VIT = val_attn_modes[left:right].count("full") if val_sample_type[i_gt] != "und": curr_video_tensor_index += N_target_VIT continue curr_sample += 1 if curr_sample > training_args.validation_max_samples: break if N_target_VIT != 0: video_tensor = val_data["padded_videos_vit"][curr_video_tensor_index : curr_video_tensor_index + N_target_VIT] # [N], each item is [C 2 H W] or [C T H W] curr_video_tensor_index += N_target_VIT v_thwc = decode_video_tensor(video_tensor, video_type="vit") is_image = v_thwc.shape[0] == 2 # Match the original convention: two frames are treated as an image media_data = v_thwc[0] if is_image else np.ascontiguousarray(v_thwc.transpose(0, 3, 1, 2)) else: media_data, is_image = None, True if media_data is None and vis_list == [] and "full" in val_attn_modes: # Avoid display degradation when the first-row image is None visual_first += 1 continue vis_list.append((media_data, is_image)) # Log GT at step 0 if curr_step == 0: local_gt_media_data = [] for i_gt, (media_data, is_image) in enumerate(vis_list): cap = val_texts_und[i_gt] if i_gt < len(val_texts_und) else "GT_text_wrong" bos_token_id = tokenizer.decode(new_token_ids["bos_token_id"]) cap = cap.split(bos_token_id) cap_target = cap[-1] # Use it as the target cap = (bos_token_id).join(cap[:-1]) caption = cap local_gt_media_data.append( {"data_vit": media_data, "caption": caption, "is_image": is_image, "cap_target": cap_target, "validation_log_type": training_args.validation_log_type} ) _log_media_across_ranks(local_gt_media_data, "validation_gt_samples_und", curr_step, training_args.validation_video_saving_fps, logger) vocab_size = len(tokenizer) with fsdp_model.summon_full_params(fsdp_model, writeback=False, rank0_only=False): generated_sequence_all = fsdp_model.validation_video_to_text( val_packed_text_ids=val_data["packed_text_ids"], val_packed_text_indexes=val_data["packed_text_indexes"], val_packed_position_ids=val_data["packed_position_ids"], val_sample_N_target=val_data["sample_N_target"], val_split_lens=val_data["split_lens"], val_attn_modes=val_data["attn_modes"], val_sample_lens=val_data["sample_lens"], val_sample_type=val_data["sample_type"], val_packed_vit_tokens=val_data["packed_vit_tokens"], val_vit_video_grid_thw=val_data["vit_video_grid_thw"], val_ce_loss_indexes=val_data["ce_loss_indexes"], max_samples=training_args.validation_max_samples, max_length=256, device=device, dtype=torch.bfloat16, new_token_ids=new_token_ids, pad_token_id=tokenizer.pad_token_id, vocab_size=vocab_size, tokenizer=tokenizer, apply_chat_template=training_args.apply_chat_template, apply_qwen_2_5_vl_pos_emb=training_args.apply_qwen_2_5_vl_pos_emb, do_sample=False, image_token_id=image_token_id, ) print(f"generated_sequence_all end") local_media_data = [] for i_val, generated_sequence in enumerate(generated_sequence_all): if i_val < visual_first: continue try: cap = tokenizer.decode(generated_sequence[:, 0]) except Exception: log_rank0(f"Rank {get_global_rank()} failed to decode sequence {i_val}: {generated_sequence}") continue caption = cap media_data, is_image = vis_list[i_val - visual_first] local_media_data.append( { "data_vit": media_data, "caption": val_texts_und[i_val], "cap_target": caption, "is_image": is_image, "validation_log_type": training_args.validation_log_type, } ) print('local_media_data: ',local_media_data) if dist.is_available() and dist.is_initialized(): dist.barrier() _log_media_across_ranks(local_media_data, "validation_pre_samples_und", curr_step, training_args.validation_video_saving_fps, logger) log_rank0(f"Validation(und) at step {curr_step} finished.") finally: # del val_data import gc gc.collect() torch.cuda.empty_cache() # dist.barrier() fsdp_model.train()