diff --git a/training/mochi-1/args.py b/training/mochi-1/args.py index ee7c463..1d41f5a 100644 --- a/training/mochi-1/args.py +++ b/training/mochi-1/args.py @@ -369,7 +369,7 @@ def _get_optimizer_args(parser: argparse.ArgumentParser) -> None: parser.add_argument( "--beta2", type=float, - default=0.95, + default=0.999, help="The beta2 parameter for the Adam and Prodigy optimizers.", ) parser.add_argument( diff --git a/training/mochi-1/text_to_video_lora.py b/training/mochi-1/text_to_video_lora.py index 239d215..9fc4ac6 100644 --- a/training/mochi-1/text_to_video_lora.py +++ b/training/mochi-1/text_to_video_lora.py @@ -483,6 +483,7 @@ def main(args): use_cpu_offload_optimizer=args.use_cpu_offload_optimizer, offload_gradients=args.offload_gradients, ) + accelerator.print(f"Using {optimizer.__class__.__name__} optimizer.") # Dataset and DataLoader dataset_init_kwargs = { @@ -635,9 +636,12 @@ def main(args): schedule_timesteps = noise_scheduler_copy.timesteps.to(accelerator.device) timesteps = timesteps.to(accelerator.device) # notice the reverse. - step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps][::-1] + step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps] sigma = sigmas[step_indices].flatten() + if "invert_sigmas" in noise_scheduler_copy.config and noise_scheduler_copy.config.invert_sigmas: + # https://github.com/huggingface/diffusers/blob/99c0483b67427de467f11aa35d54678fd36a7ea2/src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py#L209 + sigma = 1.0 - sigma while len(sigma.shape) < n_dim: sigma = sigma.unsqueeze(-1) return sigma @@ -938,7 +942,7 @@ def main(args): repo_id=repo_id, folder_path=args.output_dir, commit_message="End of training", - ignore_patterns=["step_*", "epoch_*"], + ignore_patterns=["step_*", "epoch_*", "*.bin"], ) accelerator.end_training() diff --git a/training/mochi-1/train.sh b/training/mochi-1/train.sh index ad71ed2..e9851cb 100644 --- a/training/mochi-1/train.sh +++ b/training/mochi-1/train.sh @@ -34,13 +34,13 @@ cmd="accelerate launch --config_file deepspeed.yaml --gpu_ids $GPU_IDS text_to_v --checkpointing_steps 50 \ --gradient_accumulation_steps 4 \ --gradient_checkpointing \ - --learning_rate 0.0001 \ + --learning_rate 1e-5 \ --lr_scheduler constant \ --lr_warmup_steps 0 \ --lr_num_cycles 1 \ --enable_slicing \ --enable_tiling \ - --optimizer adamw \ + --optimizer adamw --use_8bit \ --beta1 0.9 \ --beta2 0.95 \ --beta3 0.99 \