From e32d137aecaa997c0dc08ef779e13e4661854424 Mon Sep 17 00:00:00 2001 From: rlsu9 Date: Fri, 29 Nov 2024 11:11:46 +0400 Subject: [PATCH] add ema_transform --- fastvideo/distill.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/fastvideo/distill.py b/fastvideo/distill.py index b766c365a..92cff4807 100644 --- a/fastvideo/distill.py +++ b/fastvideo/distill.py @@ -88,7 +88,10 @@ def save_checkpoint(transformer: MochiTransformer3DModel, rank, output_dir, step main_print(f"--> checkpoint saved at step {step}") - +def reshard_fsdp(model): + for m in FSDP.fsdp_modules(model): + if m._has_params and m.sharding_strategy is not ShardingStrategy.NO_SHARD: + torch.distributed.fsdp._runtime_utils._reshard(m, m._handle, True) def train_one_step_mochi(transformer, teacher_transformer, ema_transformer, optimizer, lr_scheduler,loader, noise_scheduler, solver,noise_random_generator, gradient_accumulation_steps, sp_size, precondition_outputs, max_grad_norm, uncond_prompt_embed, uncond_prompt_mask, num_euler_timesteps, multiphase, not_apply_cfg_solver, distill_cfg): total_loss = 0.0 optimizer.zero_grad() @@ -196,6 +199,12 @@ def train_one_step_mochi(transformer, teacher_transformer, ema_transformer, opti total_loss += avg_loss.item() + # update ema + reshard_fsdp(ema_transformer) + for p_averaged, p_model in zip(ema_transformer.parameters(), transformer.parameters()): + with torch.no_grad(): + p_averaged.copy_(torch.lerp(p_averaged.detach(), p_model.detach(), 1 - 0.95)) + grad_norm = transformer.clip_grad_norm_(max_grad_norm) optimizer.step() lr_scheduler.step() @@ -511,6 +520,8 @@ def main(args): if args.log_validation and step % args.validation_steps == 0: log_validation(args, transformer, device, torch.bfloat16, step, scheduler_type=args.scheduler_type, shift=args.shift, num_euler_timesteps=args.num_euler_timesteps, linear_quadratic_threshold=args.linear_quadratic_threshold, ema=False) + log_validation(args, ema_transformer, device, + torch.bfloat16, step, scheduler_type=args.scheduler_type, shift=args.shift, num_euler_timesteps=args.num_euler_timesteps, linear_quadratic_threshold=args.linear_quadratic_threshold, ema=True) if args.use_lora: save_lora_checkpoint(transformer, optimizer, rank, args.output_dir, args.max_train_steps)