Peiyuan Zhang
8ec82cbc37
Delete Open Sora Plan Modeling amend log validation Change name to fast video Remove OSP modeling clean up name changing remove files Deleted unnecessary files commit first commit first training ok Update generate_synthetic.sh and deepspeed_zero2_config.yaml 14 debug overfitting .... debugging .. Debug successful! Add zero3 OK update optimizer load small bug random seed args sp enable & still has bug rename update inference sp code / can run / still has bug / don't output normal mp4 switch to deepspeed dummyoptim fix some bugs; still output green SP inference done! fix typos in readme and latent dataset debug file
71 lines
2.6 KiB
Python
71 lines
2.6 KiB
Python
from accelerate.logging import get_logger
|
|
import torch
|
|
logger = get_logger(__name__)
|
|
|
|
def get_optimizer(args, params_to_optimize, use_deepspeed: bool = False):
|
|
# Optimizer creation
|
|
supported_optimizers = ["adam", "adamw", "prodigy"]
|
|
if args.optimizer not in supported_optimizers:
|
|
logger.warning(
|
|
f"Unsupported choice of optimizer: {args.optimizer}. Supported optimizers include {supported_optimizers}. Defaulting to AdamW"
|
|
)
|
|
args.optimizer = "adamw"
|
|
|
|
if args.use_8bit_adam and not (args.optimizer.lower() not in ["adam", "adamw"]):
|
|
logger.warning(
|
|
f"use_8bit_adam is ignored when optimizer is not set to 'Adam' or 'AdamW'. Optimizer was "
|
|
f"set to {args.optimizer.lower()}"
|
|
)
|
|
|
|
if args.use_8bit_adam:
|
|
try:
|
|
import bitsandbytes as bnb
|
|
except ImportError:
|
|
raise ImportError(
|
|
"To use 8-bit Adam, please install the bitsandbytes library: `pip install bitsandbytes`."
|
|
)
|
|
|
|
if args.optimizer.lower() == "adamw":
|
|
optimizer_class = bnb.optim.AdamW8bit if args.use_8bit_adam else torch.optim.AdamW
|
|
|
|
optimizer = optimizer_class(
|
|
params_to_optimize,
|
|
betas=(args.adam_beta1, args.adam_beta2),
|
|
eps=args.adam_epsilon,
|
|
weight_decay=args.adam_weight_decay,
|
|
)
|
|
elif args.optimizer.lower() == "adam":
|
|
optimizer_class = bnb.optim.Adam8bit if args.use_8bit_adam else torch.optim.Adam
|
|
|
|
optimizer = optimizer_class(
|
|
params_to_optimize,
|
|
betas=(args.adam_beta1, args.adam_beta2),
|
|
eps=args.adam_epsilon,
|
|
weight_decay=args.adam_weight_decay,
|
|
)
|
|
elif args.optimizer.lower() == "prodigy":
|
|
try:
|
|
import prodigyopt
|
|
except ImportError:
|
|
raise ImportError("To use Prodigy, please install the prodigyopt library: `pip install prodigyopt`")
|
|
|
|
optimizer_class = prodigyopt.Prodigy
|
|
|
|
if args.learning_rate <= 0.1:
|
|
logger.warning(
|
|
"Learning rate is too low. When using prodigy, it's generally better to set learning rate around 1.0"
|
|
)
|
|
|
|
optimizer = optimizer_class(
|
|
params_to_optimize,
|
|
lr=args.learning_rate,
|
|
betas=(args.adam_beta1, args.adam_beta2),
|
|
beta3=args.prodigy_beta3,
|
|
weight_decay=args.adam_weight_decay,
|
|
eps=args.adam_epsilon,
|
|
decouple=args.prodigy_decouple,
|
|
use_bias_correction=args.prodigy_use_bias_correction,
|
|
safeguard_warmup=args.prodigy_safeguard_warmup,
|
|
)
|
|
|
|
return optimizer |