#!/usr/bin/env python3 """ InstanceV 训练脚本 基于 DiffSynth-Studio 框架，训练 WanVideo 的 InstanceV 模块： - IMCA (Instance-aware Masked Cross-Attention) - STAPE (Shared Timestep Adaptive Prompt Enhancement) - SAUG (Spatially-Aware Unconditional Guidance) 训练时的 dropout 主要参考： - wan_video_dit_instancev.py: 模型定义 - wan_video_instanceV.py: Pipeline 定义 - train.py: 基础训练框架 """ import torch import os import sys import json import argparse import accelerate import warnings from pathlib import Path from PIL import Image from typing import Optional from datetime import datetime import numpy as np # Wandb (可选) try: import wandb WANDB_AVAILABLE = True except ImportError: WANDB_AVAILABLE = False print("[Warning] wandb not installed. Run 'pip install wandb' to enable wandb logging.") # 添加项目根目录到 path sys.path.insert(0, str(Path(__file__).resolve().parents[3])) from diffsynth.core import UnifiedDataset, load_state_dict from diffsynth.core.data.operators import LoadVideo, ToAbsolutePath from diffsynth.pipelines.wan_video_instanceV import WanVideoPipeline, ModelConfig from diffsynth.diffusion import ( DiffusionTrainingModule, add_general_config, add_video_size_config, launch_training_task, launch_data_process_task, ModelLogger, ) os.environ["TOKENIZERS_PARALLELISM"] = "false" class LoadInstanceMasks: """ 自定义数据算子：加载 instance mask 序列 mask 目录结构： {mask_dir}/{frame_id:06d}_No.{instance_id}.png """ def __init__( self, num_frames: int, time_division_factor: int = 4, time_division_remainder: int = 1, target_height: Optional[int] = None, target_width: Optional[int] = None, ): self.num_frames = num_frames self.time_division_factor = time_division_factor self.time_division_remainder = time_division_remainder self.target_height = target_height self.target_width = target_width def _adjust_num_frames(self, num_frames: int) -> int: num_frames = min(int(num_frames), self.num_frames) while num_frames > 1 and num_frames % self.time_division_factor != self.time_division_remainder: num_frames -= 1 return num_frames def _crop_and_resize_mask(self, mask: Image.Image) -> Image.Image: if self.target_height is None or self.target_width is None: return mask width, height = mask.size scale = max(self.target_width / width, self.target_height / height) new_w = max(1, int(round(width * scale))) new_h = max(1, int(round(height * scale))) if (new_w, new_h) != mask.size: mask = mask.resize((new_w, new_h), resample=Image.NEAREST) left = max(0, (new_w - self.target_width) // 2) top = max(0, (new_h - self.target_height) // 2) right = left + self.target_width bottom = top + self.target_height return mask.crop((left, top, right, bottom)) def __call__(self, instance_mask_dirs: list) -> list: """ Args: instance_mask_dirs: list of {"mask_dir": str, "instance_id": int, "num_frames": int} Returns: list of list of PIL.Image: [Nins][num_frames] """ all_masks = [] for mask_info in instance_mask_dirs: mask_dir = mask_info["mask_dir"] inst_id = mask_info["instance_id"] raw_num_frames = mask_info.get("num_frames", self.num_frames) num_frames = self._adjust_num_frames(raw_num_frames) masks = [] blank_size = ( self.target_width if self.target_width is not None else 64, self.target_height if self.target_height is not None else 64, ) for frame_idx in range(num_frames): mask_path = os.path.join(mask_dir, f"{frame_idx:06d}_No.{inst_id}.png") if os.path.exists(mask_path): mask = Image.open(mask_path).convert("L") else: # 如果 mask 不存在，创建全黑 mask mask = Image.new("L", blank_size, 0) mask = self._crop_and_resize_mask(mask) masks.append(mask) # 采样到目标帧数 target_frames = (num_frames - self.time_division_remainder) // self.time_division_factor * self.time_division_factor + self.time_division_remainder if len(masks) > target_frames: # 均匀采样 indices = np.linspace(0, len(masks) - 1, target_frames, dtype=int) masks = [masks[i] for i in indices] all_masks.append(masks) return all_masks class InstanceVTrainingModule(DiffusionTrainingModule): """ InstanceV 训练模块 继承自 DiffusionTrainingModule，添加 InstanceV 特定的处理逻辑 """ def __init__( self, model_paths=None, model_id_with_origin_paths=None, tokenizer_path=None, trainable_models=None, lora_base_model=None, lora_target_modules="", lora_rank=32, lora_checkpoint=None, preset_lora_path=None, preset_lora_model=None, use_gradient_checkpointing=True, use_gradient_checkpointing_offload=False, extra_inputs=None, fp8_models=None, offload_models=None, device="cpu", task="sft", max_timestep_boundary=1.0, min_timestep_boundary=0.0, # InstanceV 特有参数 saug_drop_prob=0.1, saug_scale=0.0, ): super().__init__() # Gradient checkpointing 检查 if not use_gradient_checkpointing: warnings.warn( "Gradient checkpointing is disabled. This may increase VRAM usage and risk OOM." ) # 加载模型 model_configs = self.parse_model_configs( model_paths, model_id_with_origin_paths, fp8_models=fp8_models, offload_models=offload_models, device=device ) tokenizer_config = ( ModelConfig(model_id="Wan-AI/Wan2.1-T2V-1.3B", origin_file_pattern="google/umt5-xxl/") if tokenizer_path is None else ModelConfig(tokenizer_path) ) # 使用 InstanceV Pipeline self.pipe = WanVideoPipeline.from_pretrained( torch_dtype=torch.bfloat16, device=device, model_configs=model_configs, tokenizer_config=tokenizer_config, audio_processor_config=None, # InstanceV 不需要音频 ) # 升级模型：动态添加 InstanceV 模块 (IMCA, STAPE, mv) self._upgrade_dit_with_instancev() self.pipe = self.split_pipeline_units(task, self.pipe, trainable_models, lora_base_model) # 训练模式 self.switch_pipe_to_training_mode( self.pipe, trainable_models, lora_base_model, lora_target_modules, lora_rank, lora_checkpoint, preset_lora_path, preset_lora_model, task=task, ) # InstanceV 特有：冻结 backbone，只训练新增模块 (IMCA, STAPE, mv) # 论文 Table 1: 只增加 20.65% 参数 self._freeze_backbone_keep_instancev() # 存储配置 self.use_gradient_checkpointing = use_gradient_checkpointing self.use_gradient_checkpointing_offload = use_gradient_checkpointing_offload self.extra_inputs = extra_inputs.split(",") if extra_inputs is not None else [] self.fp8_models = fp8_models self.task = task self.max_timestep_boundary = max_timestep_boundary self.min_timestep_boundary = min_timestep_boundary # InstanceV 特有参数 self.saug_drop_prob = saug_drop_prob self.saug_scale = saug_scale # Loss 函数映射 self.task_to_loss = { "sft:data_process": lambda pipe, *args: args, "sft": self._compute_instancev_loss, "sft:train": self._compute_instancev_loss, } def _upgrade_dit_with_instancev(self): """ 动态升级原始 Wan 模型，添加 InstanceV 模块： - STAPE (Shared Timestep-Adaptive Prompt Enhancement) - IMCA (Instance-aware Masked Cross-Attention) for each block - mv (gated residual) for each block - norm_imca (LayerNorm) for each block - 替换 block 的 forward 方法为 InstanceV 版本 """ from diffsynth.models.wan_video_dit_instancev import ( SharedTimestepAdaptivePromptEnhancement, InstanceAwareMaskedCrossAttention, DiTBlock as InstanceVDiTBlock, modulate, ) import torch.nn as nn from typing import Optional if self.pipe.dit is None: print("Warning: No dit model found in pipeline") return dit = self.pipe.dit # 获取模型维度 dim = dit.dim num_heads = dit.blocks[0].self_attn.num_heads if hasattr(dit.blocks[0], 'self_attn') else 12 eps = 1e-6 # 设置 enable_instancev 标志 dit.enable_instancev = True # 添加共享的 STAPE 模块 if not hasattr(dit, 'stape') or dit.stape is None: dit.stape = SharedTimestepAdaptivePromptEnhancement( dim=dim, num_heads=num_heads, eps=eps ).to(device=next(dit.parameters()).device, dtype=next(dit.parameters()).dtype) print(f"InstanceV: Added STAPE module (dim={dim}, num_heads={num_heads})") # 定义 InstanceV 版本的 forward 方法 def instancev_forward( self, x: torch.Tensor, context: torch.Tensor, t_mod: torch.Tensor, freqs: torch.Tensor, instance_tokens: Optional[torch.Tensor] = None, instance_attn_mask: Optional[torch.Tensor] = None, empty_instance_tokens: Optional[torch.Tensor] = None, saug_drop_prob: float = 0.0, ): """InstanceV 版本的 DiTBlock forward""" has_seq = len(t_mod.shape) == 4 chunk_dim = 2 if has_seq else 1 shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( self.modulation.to(dtype=t_mod.dtype, device=t_mod.device) + t_mod ).chunk(6, dim=chunk_dim) if has_seq: shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( shift_msa.squeeze(2), scale_msa.squeeze(2), gate_msa.squeeze(2), shift_mlp.squeeze(2), scale_mlp.squeeze(2), gate_mlp.squeeze(2), ) # 1) Self-attention input_x = modulate(self.norm1(x), shift_msa, scale_msa) x = self.gate(x, gate_msa, self.self_attn(input_x, freqs)) # 2) IMCA + STAPE if self.enable_instancev and (self.imca is not None) and (instance_tokens is not None) and (instance_attn_mask is not None): # SAUG training-time drop if isinstance(saug_drop_prob, torch.Tensor): saug_p = float(saug_drop_prob.detach().cpu().item()) else: saug_p = float(saug_drop_prob) if self.training and saug_p > 0.0 and empty_instance_tokens is not None: if torch.rand((), device=x.device) < saug_p: instance_tokens_use = empty_instance_tokens else: instance_tokens_use = instance_tokens else: instance_tokens_use = instance_tokens # STAPE if self.stape is not None: alpha1 = gate_msa instance_tokens_use = self.stape(instance_tokens_use, context, alpha1=alpha1) # IMCA imca_out = self.imca(self.norm_imca(x), instance_tokens_use, instance_attn_mask) x = x + self.mv.to(dtype=x.dtype, device=x.device) * imca_out # 3) Cross-attention x = x + self.cross_attn(self.norm3(x), context) # 4) FFN input_x = modulate(self.norm2(x), shift_mlp, scale_mlp) x = self.gate(x, gate_mlp, self.ffn(input_x)) return x # 为每个 DiT block 添加 IMCA、mv、norm_imca，并替换 forward 方法 for block_idx, block in enumerate(dit.blocks): block.enable_instancev = True block.stape = dit.stape # 共享 STAPE # 添加 IMCA if not hasattr(block, 'imca') or block.imca is None: block.imca = InstanceAwareMaskedCrossAttention( dim=dim, num_heads=num_heads, eps=eps ).to(device=next(block.parameters()).device, dtype=next(block.parameters()).dtype) # 从 cross_attn 初始化权重（论文建议） if hasattr(block, 'cross_attn'): try: block.imca.attn.q.load_state_dict(block.cross_attn.q.state_dict()) block.imca.attn.k.load_state_dict(block.cross_attn.k.state_dict()) block.imca.attn.v.load_state_dict(block.cross_attn.v.state_dict()) block.imca.attn.o.load_state_dict(block.cross_attn.o.state_dict()) block.imca.attn.norm_q.load_state_dict(block.cross_attn.norm_q.state_dict()) block.imca.attn.norm_k.load_state_dict(block.cross_attn.norm_k.state_dict()) except Exception as e: print(f"Warning: Failed to init IMCA from cross_attn in block {block_idx}: {e}") # 添加 mv (gated residual, 初始化为 0) if not hasattr(block, 'mv') or block.mv is None: block.mv = nn.Parameter(torch.zeros(1, device=next(block.parameters()).device, dtype=next(block.parameters()).dtype)) # 添加 norm_imca if not hasattr(block, 'norm_imca') or block.norm_imca is None: block.norm_imca = nn.LayerNorm(dim, eps=eps, elementwise_affine=False).to( device=next(block.parameters()).device, dtype=next(block.parameters()).dtype ) # 替换 forward 方法为 InstanceV 版本 import types block.forward = types.MethodType(instancev_forward, block) print(f"InstanceV: Added IMCA, mv, norm_imca to {len(dit.blocks)} blocks and replaced forward methods") def _freeze_backbone_keep_instancev(self): """ 冻结原始 backbone，只保留 InstanceV 新增模块可训练： - STAPE (Shared Timestep-Adaptive Prompt Enhancement) - IMCA (Instance-aware Masked Cross-Attention) - mv (gated residual parameter) - norm_imca (LayerNorm for IMCA) 论文 Table 1: InstanceV 只增加 20.65% 参数 """ if self.pipe.dit is None: return dit = self.pipe.dit # 首先冻结整个 dit dit.requires_grad_(False) # 解冻 STAPE（共享模块） if hasattr(dit, 'stape') and dit.stape is not None: dit.stape.requires_grad_(True) dit.stape.train() print("InstanceV: Enabled training for dit.stape") # 解冻每个 block 中的 IMCA、mv、norm_imca trainable_params = 0 total_params = 0 for block_idx, block in enumerate(dit.blocks): # IMCA 模块 if hasattr(block, 'imca') and block.imca is not None: block.imca.requires_grad_(True) block.imca.train() trainable_params += sum(p.numel() for p in block.imca.parameters()) # mv 参数 (gated residual) if hasattr(block, 'mv') and block.mv is not None: block.mv.requires_grad_(True) trainable_params += block.mv.numel() # norm_imca (LayerNorm) if hasattr(block, 'norm_imca') and block.norm_imca is not None: block.norm_imca.requires_grad_(True) block.norm_imca.train() trainable_params += sum(p.numel() for p in block.norm_imca.parameters()) # 统计参数 total_params = sum(p.numel() for p in dit.parameters()) if hasattr(dit, 'stape') and dit.stape is not None: trainable_params += sum(p.numel() for p in dit.stape.parameters()) print(f"InstanceV: Trainable params: {trainable_params:,} / {total_params:,} " f"({100.0 * trainable_params / total_params:.2f}%)") def _compute_instancev_loss(self, pipe, inputs_shared, inputs_posi, inputs_nega): """ InstanceV 专用 Loss 计算 包含： 1. 标准 Flow Matching SFT Loss 2. SAUG training-time dropout（在 DiT block 中处理） """ from diffsynth.diffusion.loss import FlowMatchSFTLoss return FlowMatchSFTLoss(pipe, **inputs_shared, **inputs_posi) def parse_extra_inputs(self, data, extra_inputs, inputs_shared): """处理额外输入""" for extra_input in extra_inputs: if extra_input == "input_image": inputs_shared["input_image"] = data["video"][0] elif extra_input == "end_image": inputs_shared["end_image"] = data["video"][-1] elif extra_input in data: inputs_shared[extra_input] = data[extra_input] return inputs_shared def get_pipeline_inputs(self, data): """ 构建 Pipeline 输入 InstanceV 特有字段： - instance_prompts: list[str] - instance_masks: list[list[PIL.Image]] """ inputs_posi = {"prompt": data["prompt"]} inputs_nega = {} # 基础输入 inputs_shared = { "input_video": data["video"], "height": data["video"][0].size[1], "width": data["video"][0].size[0], "num_frames": len(data["video"]), "cfg_scale": 1, "tiled": False, "rand_device": self.pipe.device, "use_gradient_checkpointing": self.use_gradient_checkpointing, "use_gradient_checkpointing_offload": self.use_gradient_checkpointing_offload, "cfg_merge": False, "vace_scale": 1, "max_timestep_boundary": self.max_timestep_boundary, "min_timestep_boundary": self.min_timestep_boundary, } # InstanceV 特有输入 if "instance_prompts" in data: inputs_shared["instance_prompts"] = data["instance_prompts"] # 注意: metadata 中使用 "instance_mask_dirs"，pipeline 需要 "instance_masks" if "instance_mask_dirs" in data: inputs_shared["instance_masks"] = data["instance_mask_dirs"] elif "instance_masks" in data: inputs_shared["instance_masks"] = data["instance_masks"] # SAUG 参数 inputs_shared["saug_drop_prob"] = self.saug_drop_prob inputs_shared["saug_scale"] = self.saug_scale # 处理额外输入 inputs_shared = self.parse_extra_inputs(data, self.extra_inputs, inputs_shared) return inputs_shared, inputs_posi, inputs_nega def forward(self, data, inputs=None): """前向传播""" if inputs is None: inputs = self.get_pipeline_inputs(data) inputs = self.transfer_data_to_device(inputs, self.pipe.device, self.pipe.torch_dtype) for unit in self.pipe.units: inputs = self.pipe.unit_runner(unit, self.pipe, *inputs) loss = self.task_to_loss[self.task](self.pipe, *inputs) return loss def instancev_parser(): """命令行参数解析""" parser = argparse.ArgumentParser(description="InstanceV Training Script") parser = add_general_config(parser) parser = add_video_size_config(parser) # Tokenizer parser.add_argument( "--tokenizer_path", type=str, default=None, help="Path to tokenizer.", ) # InstanceV 特有参数 parser.add_argument( "--saug_drop_prob", type=float, default=0.1, help="SAUG training-time dropout probability (paper recommends 0.1).", ) parser.add_argument( "--saug_scale", type=float, default=0.0, help="SAUG unconditional guidance scale (training time, usually 0).", ) # 时间步边界 parser.add_argument( "--max_timestep_boundary", type=float, default=1.0, help="Max timestep boundary.", ) parser.add_argument( "--min_timestep_boundary", type=float, default=0.0, help="Min timestep boundary.", ) # 模型初始化 parser.add_argument( "--initialize_model_on_cpu", default=False, action="store_true", help="Whether to initialize models on CPU.", ) # Wandb 参数 parser.add_argument( "--use_wandb", default=False, action="store_true", help="Enable wandb logging.", ) parser.add_argument( "--wandb_project", type=str, default="instancev-training", help="Wandb project name.", ) parser.add_argument( "--wandb_run_name", type=str, default=None, help="Wandb run name. Default: auto-generated with timestamp.", ) parser.add_argument( "--wandb_log_every", type=int, default=10, help="Log to wandb every N steps.", ) # 断点续跑参数 parser.add_argument( "--resume_from_checkpoint", type=str, default=None, help="Path to a trainable checkpoint (.safetensors) to resume from.", ) return parser def _resolve_resume_target(model, remove_prefix_in_ckpt): if not remove_prefix_in_ckpt: return model prefix = remove_prefix_in_ckpt.rstrip(".") if not prefix: return model target = model for part in prefix.split("."): if not hasattr(target, part): return None target = getattr(target, part) return target def _load_resume_checkpoint(model, checkpoint_path, remove_prefix_in_ckpt, accelerator): if checkpoint_path is None: return if not os.path.exists(checkpoint_path): raise FileNotFoundError(f"Resume checkpoint not found: {checkpoint_path}") state_dict = load_state_dict(checkpoint_path) target = _resolve_resume_target(model, remove_prefix_in_ckpt) if target is None: target = model if accelerator.is_main_process: print( f"[Warning] remove_prefix_in_ckpt='{remove_prefix_in_ckpt}' does not map to a model attribute. " "Fallback to loading into the full model." ) load_result = target.load_state_dict(state_dict, strict=False) if accelerator.is_main_process: missing_keys, unexpected_keys = load_result print(f"[Resume] Loaded checkpoint: {checkpoint_path}") print(f"[Resume] Loaded keys: {len(state_dict)}") if missing_keys: print(f"[Resume] Missing keys: {len(missing_keys)}") if unexpected_keys: print(f"[Resume] Unexpected keys: {len(unexpected_keys)}") def print_training_info(args, accelerator): """打印训练配置信息""" if not accelerator.is_main_process: return print("\n" + "="*60) print(" InstanceV Training Configuration") print("="*60) print(f"\n[时间] {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}") print(f"\n[数据配置]") print(f" - 数据路径: {args.dataset_base_path}") print(f" - Metadata: {args.dataset_metadata_path}") print(f" - 数据集重复: {args.dataset_repeat}x") print(f"\n[模型配置]") print(f" - 模型来源: {args.model_id_with_origin_paths or args.model_paths}") print(f" - 可训练模块: STAPE, IMCA, mv, norm_imca") print(f" - 输出路径: {args.output_path}") print(f"\n[训练参数]") print(f" - 分辨率: {args.width}x{args.height}") print(f" - 帧数: {args.num_frames}") print(f" - 学习率: {args.learning_rate}") print(f" - Epochs: {args.num_epochs}") print(f" - 梯度累积: {args.gradient_accumulation_steps}") print(f" - 保存间隔: {args.save_steps} steps") if args.resume_from_checkpoint: print(f" - 断点续跑: {args.resume_from_checkpoint}") print(f"\n[InstanceV 参数]") print(f" - SAUG Dropout: {args.saug_drop_prob}") print(f" - SAUG Scale: {args.saug_scale}") print(f" - Timestep 边界: [{args.min_timestep_boundary}, {args.max_timestep_boundary}]") if args.use_wandb: print(f"\n[Wandb]") print(f" - Project: {args.wandb_project}") print(f" - Run Name: {args.wandb_run_name}") print(f" - Log Every: {args.wandb_log_every} steps") print(f"\n[GPU/分布式]") print(f" - Accelerator: {accelerator.device}") print(f" - 进程数: {accelerator.num_processes}") print(f" - 混合精度: {accelerator.mixed_precision}") print("\n" + "="*60 + "\n") def main(): parser = instancev_parser() args = parser.parse_args() # 初始化 Accelerator accelerator = accelerate.Accelerator( gradient_accumulation_steps=args.gradient_accumulation_steps, kwargs_handlers=[ accelerate.DistributedDataParallelKwargs(find_unused_parameters=args.find_unused_parameters) ], ) # 设置 wandb_mode（runner.py 需要） if args.use_wandb and WANDB_AVAILABLE: args.wandb_mode = "online" # 生成默认 run name if args.wandb_run_name is None: args.wandb_run_name = f"instancev_{datetime.now().strftime('%Y%m%d_%H%M%S')}" # 只在主进程初始化 wandb if accelerator.is_main_process: wandb.init( project=args.wandb_project, name=args.wandb_run_name, config={ "learning_rate": args.learning_rate, "num_epochs": args.num_epochs, "gradient_accumulation_steps": args.gradient_accumulation_steps, "height": args.height, "width": args.width, "num_frames": args.num_frames, "saug_drop_prob": args.saug_drop_prob, "saug_scale": args.saug_scale, "trainable_models": args.trainable_models, "model": "Wan2.1-T2V-1.3B + InstanceV", }, ) print(f"[Wandb] 初始化完成: {args.wandb_project}/{args.wandb_run_name}") else: args.wandb_mode = "disabled" # 打印训练配置 print_training_info(args, accelerator) # 构建 Dataset # 注意：data_file_keys 需要包含所有需要处理的键（UnifiedDataset 只处理这些键） # - video: 使用 main_data_operator（视频加载器） # - instance_mask_dirs: 使用 special_operator_map 中的 LoadInstanceMasks # - instance_prompts: 使用 special_operator_map 中的 lambda（透传） dataset = UnifiedDataset( base_path=args.dataset_base_path, metadata_path=args.dataset_metadata_path, repeat=args.dataset_repeat, data_file_keys=["video", "instance_mask_dirs", "instance_prompts"], # 所有需要处理的键 main_data_operator=UnifiedDataset.default_video_operator( base_path=args.dataset_base_path, max_pixels=args.max_pixels, height=args.height, width=args.width, height_division_factor=16, width_division_factor=16, num_frames=args.num_frames, time_division_factor=4, time_division_remainder=1, ), special_operator_map={ # InstanceV 特有：加载 instance masks "instance_mask_dirs": LoadInstanceMasks( num_frames=args.num_frames, time_division_factor=4, time_division_remainder=1, target_height=args.height, target_width=args.width, ), # instance_prompts 直接透传（字符串列表） "instance_prompts": lambda x: x, }, ) # 构建训练模块 model = InstanceVTrainingModule( model_paths=args.model_paths, model_id_with_origin_paths=args.model_id_with_origin_paths, tokenizer_path=args.tokenizer_path, trainable_models=args.trainable_models, lora_base_model=args.lora_base_model, lora_target_modules=args.lora_target_modules, lora_rank=args.lora_rank, lora_checkpoint=args.lora_checkpoint, preset_lora_path=args.preset_lora_path, preset_lora_model=args.preset_lora_model, use_gradient_checkpointing=args.use_gradient_checkpointing, use_gradient_checkpointing_offload=args.use_gradient_checkpointing_offload, extra_inputs=args.extra_inputs, fp8_models=args.fp8_models, offload_models=args.offload_models, task=args.task, device="cpu" if args.initialize_model_on_cpu else accelerator.device, max_timestep_boundary=args.max_timestep_boundary, min_timestep_boundary=args.min_timestep_boundary, saug_drop_prob=args.saug_drop_prob, saug_scale=args.saug_scale, ) # Model Logger model_logger = ModelLogger( args.output_path, remove_prefix_in_ckpt=args.remove_prefix_in_ckpt, ) # 断点续跑（只加载权重） if args.resume_from_checkpoint and not args.task.endswith(":data_process"): _load_resume_checkpoint( model, args.resume_from_checkpoint, args.remove_prefix_in_ckpt, accelerator, ) # 启动训练 launcher_map = { "sft:data_process": launch_data_process_task, "sft": launch_training_task, "sft:train": launch_training_task, } if accelerator.is_main_process: print(f"\n[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] 开始训练...\n") launcher_map[args.task](accelerator, dataset, model, model_logger, args=args) # 训练完成 if accelerator.is_main_process: print(f"\n[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] 训练完成!") print(f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] Checkpoints 保存至: {args.output_path}") # 关闭 wandb if args.use_wandb and WANDB_AVAILABLE: wandb.finish() print(f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] Wandb 已关闭") if __name__ == "__main__": main()