storytold/FineTrainers-Conditioning
2
0
Fork 0
mirror of https://github.com/storytold/FineTrainers-Conditioning.git synced 2026-10-18 07:33:20 +00:00
Code Issues Packages Projects Releases Wiki Activity

Low-bit memory optimizers, CpuOffloadOptimizer, Memory Reports (#3)

Browse Source 
* update

* update

* update

* update

* update

* update

* update

* update

* update
This commit is contained in:
Aryan
2024-10-04 18:18:54 +05:30
committed by GitHub
parent 512a8a6854
commit d81209bf49
11 changed files with 434 additions and 314 deletions
Download Patch File Download Diff File Expand all files Collapse all files
training/args.py
+20 -2
View File
@@ -301,13 +301,21 @@ def _get_optimizer_args(parser: argparse.ArgumentParser) -> None:
"--optimizer",
type=lambda s: s.lower(),
default="adam",
choices=["adam", "adamw", "prodigy"],
choices=["adam", "adamw", "prodigy", "came"],
help=("The optimizer type to use."),
)
parser.add_argument(
"--use_8bit",
action="store_true",
help="Whether or not to use 8-bit optimizers from `bitsandbytes`. Ignored if incompatible optimzer selected.",
help="Whether or not to use 8-bit optimizers from `bitsandbytes` or `bitsandbytes`.",
)
parser.add_argument(
"--use_4bit",
action="store_true",
help="Whether or not to use 4-bit optimizers from `torchao`.",
)
parser.add_argument(
"--use_torchao", action="store_true", help="Whether or not to use the `torchao` backend for optimizers."
)
parser.add_argument(
"--beta1",
@@ -355,6 +363,16 @@ def _get_optimizer_args(parser: argparse.ArgumentParser) -> None:
action="store_true",
help="Remove lr from the denominator of D estimate to avoid issues during warm-up stage.",
)
parser.add_argument(
"--use_cpu_offload_optimizer",
action="store_true",
help="Whether or not to use the CPUOffloadOptimizer from TorchAO to perform optimization step and maintain parameters on the CPU.",
)
parser.add_argument(
"--offload_gradients",
action="store_true",
help="Whether or not to offload the gradients to CPU when using the CPUOffloadOptimizer from TorchAO.",
)
def _get_configuration_args(parser: argparse.ArgumentParser) -> None:
training/cogvideox_text_to_video_lora.py
+38 -31
View File
@@ -25,7 +25,6 @@ from typing import Any, Dict
import diffusers
import torch
import transformers
import wandb
from accelerate import Accelerator
from accelerate.logging import get_logger
from accelerate.utils import (
@@ -43,11 +42,7 @@ from diffusers import (
from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution
from diffusers.optimization import get_scheduler
from diffusers.training_utils import cast_training_params
from diffusers.utils import (
convert_unet_state_dict_to_peft,
export_to_video,
is_wandb_available,
)
from diffusers.utils import convert_unet_state_dict_to_peft, export_to_video
from diffusers.utils.hub_utils import load_or_create_model_card, populate_model_card
from diffusers.utils.torch_utils import is_compiled_module
from huggingface_hub import create_repo, upload_folder
@@ -56,6 +51,8 @@ from torch.utils.data import DataLoader
from tqdm.auto import tqdm
from transformers import AutoTokenizer, T5EncoderModel
import wandb
from args import get_args # isort:skip
from dataset import BucketSampler, VideoDatasetWithResizing # isort:skip
@@ -230,10 +227,6 @@ def main(args):
if torch.backends.mps.is_available():
accelerator.native_amp = False
if args.report_to == "wandb":
if not is_wandb_available():
raise ImportError("Make sure to install wandb if you want to use it for logging during training.")
# Make one log on every process with the configuration for debugging.
logging.basicConfig(
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
@@ -436,6 +429,7 @@ def main(args):
"lr": args.learning_rate,
}
params_to_optimize = [transformer_parameters_with_lr]
num_trainable_parameters = sum(param.numel() for model in params_to_optimize for param in model["params"])
use_deepspeed_optimizer = (
accelerator.state.deepspeed_plugin is not None
@@ -459,7 +453,11 @@ def main(args):
prodigy_use_bias_correction=args.prodigy_use_bias_correction,
prodigy_safeguard_warmup=args.prodigy_safeguard_warmup,
use_8bit=args.use_8bit,
use_4bit=args.use_4bit,
use_torchao=args.use_torchao,
use_deepspeed=use_deepspeed_optimizer,
use_cpu_offload_optimizer=args.use_cpu_offload_optimizer,
offload_gradients=args.offload_gradients,
)
# Dataset and DataLoader
@@ -524,24 +522,31 @@ def main(args):
args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
overrode_max_train_steps = True
if use_deepspeed_scheduler:
from accelerate.utils import DummyScheduler
lr_scheduler = DummyScheduler(
name=args.lr_scheduler,
optimizer=optimizer,
total_num_steps=args.max_train_steps * accelerator.num_processes,
num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes,
if args.use_cpu_offload_optimizer:
lr_scheduler = None
accelerator.print(
"CPU Offload Optimizer cannot be used with DeepSpeed or builtin PyTorch LR Schedulers. If "
"you are training with those settings, they will be ignored."
)
else:
lr_scheduler = get_scheduler(
args.lr_scheduler,
optimizer=optimizer,
num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes,
num_training_steps=args.max_train_steps * accelerator.num_processes,
num_cycles=args.lr_num_cycles,
power=args.lr_power,
)
if use_deepspeed_scheduler:
from accelerate.utils import DummyScheduler
lr_scheduler = DummyScheduler(
name=args.lr_scheduler,
optimizer=optimizer,
total_num_steps=args.max_train_steps * accelerator.num_processes,
num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes,
)
else:
lr_scheduler = get_scheduler(
args.lr_scheduler,
optimizer=optimizer,
num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes,
num_training_steps=args.max_train_steps * accelerator.num_processes,
num_cycles=args.lr_num_cycles,
power=args.lr_power,
)
# Prepare everything with our `accelerator`.
transformer, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
@@ -567,7 +572,6 @@ def main(args):
# Train!
total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
num_trainable_parameters = sum(param.numel() for model in params_to_optimize for param in model["params"])
accelerator.print("***** Running training *****")
accelerator.print(f" Num trainable parameters = {num_trainable_parameters}")
@@ -624,6 +628,8 @@ def main(args):
torch.cuda.empty_cache()
torch.cuda.synchronize(accelerator.device)
alphas_cumprod = scheduler.alphas_cumprod.to(accelerator.device, dtype=torch.float32)
for epoch in range(first_epoch, args.num_train_epochs):
transformer.train()
@@ -691,8 +697,7 @@ def main(args):
model_pred = scheduler.get_velocity(model_output, noisy_model_input, timesteps)
alphas_cumprod = scheduler.alphas_cumprod[timesteps]
weights = 1 / (1 - alphas_cumprod)
weights = 1 / (1 - alphas_cumprod[timesteps])
while len(weights.shape) < len(model_pred.shape):
weights = weights.unsqueeze(-1)
@@ -714,7 +719,8 @@ def main(args):
optimizer.step()
optimizer.zero_grad()
lr_scheduler.step()
if not args.use_cpu_offload_optimizer:
lr_scheduler.step()
# Checks if the accelerator has performed an optimization step behind the scenes
if accelerator.sync_gradients:
@@ -747,9 +753,10 @@ def main(args):
accelerator.save_state(save_path)
logger.info(f"Saved state to {save_path}")
last_lr = lr_scheduler.get_last_lr()[0] if lr_scheduler is not None else args.learning_rate
logs = {
"loss": loss.detach().item(),
"lr": lr_scheduler.get_last_lr()[0],
"lr": last_lr,
"gradient_norm_before_clip": gradient_norm_before_clip,
"gradient_norm_after_clip": gradient_norm_after_clip,
}
training/cogvideox_text_to_video_sft.py
+38 -27
View File
@@ -25,7 +25,6 @@ from typing import Any, Dict
import diffusers
import torch
import transformers
import wandb
from accelerate import Accelerator
from accelerate.logging import get_logger
from accelerate.utils import (
@@ -43,7 +42,7 @@ from diffusers import (
from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution
from diffusers.optimization import get_scheduler
from diffusers.training_utils import cast_training_params
from diffusers.utils import export_to_video, is_wandb_available
from diffusers.utils import export_to_video
from diffusers.utils.hub_utils import load_or_create_model_card, populate_model_card
from diffusers.utils.torch_utils import is_compiled_module
from huggingface_hub import create_repo, upload_folder
@@ -51,6 +50,8 @@ from torch.utils.data import DataLoader
from tqdm.auto import tqdm
from transformers import AutoTokenizer, T5EncoderModel
import wandb
from args import get_args # isort:skip
from dataset import BucketSampler, VideoDatasetWithResizing # isort:skip
@@ -183,10 +184,6 @@ def main(args):
if torch.backends.mps.is_available():
accelerator.native_amp = False
if args.report_to == "wandb":
if not is_wandb_available():
raise ImportError("Make sure to install wandb if you want to use it for logging during training.")
# Make one log on every process with the configuration for debugging.
logging.basicConfig(
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
@@ -363,6 +360,7 @@ def main(args):
"lr": args.learning_rate,
}
params_to_optimize = [transformer_parameters_with_lr]
num_trainable_parameters = sum(param.numel() for model in params_to_optimize for param in model["params"])
use_deepspeed_optimizer = (
accelerator.state.deepspeed_plugin is not None
@@ -386,7 +384,11 @@ def main(args):
prodigy_use_bias_correction=args.prodigy_use_bias_correction,
prodigy_safeguard_warmup=args.prodigy_safeguard_warmup,
use_8bit=args.use_8bit,
use_4bit=args.use_4bit,
use_torchao=args.use_torchao,
use_deepspeed=use_deepspeed_optimizer,
use_cpu_offload_optimizer=args.use_cpu_offload_optimizer,
offload_gradients=args.offload_gradients,
)
# Dataset and DataLoader
@@ -451,24 +453,31 @@ def main(args):
args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
overrode_max_train_steps = True
if use_deepspeed_scheduler:
from accelerate.utils import DummyScheduler
lr_scheduler = DummyScheduler(
name=args.lr_scheduler,
optimizer=optimizer,
total_num_steps=args.max_train_steps * accelerator.num_processes,
num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes,
if args.use_cpu_offload_optimizer:
lr_scheduler = None
accelerator.print(
"CPU Offload Optimizer cannot be used with DeepSpeed or builtin PyTorch LR Schedulers. If "
"you are training with those settings, they will be ignored."
)
else:
lr_scheduler = get_scheduler(
args.lr_scheduler,
optimizer=optimizer,
num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes,
num_training_steps=args.max_train_steps * accelerator.num_processes,
num_cycles=args.lr_num_cycles,
power=args.lr_power,
)
if use_deepspeed_scheduler:
from accelerate.utils import DummyScheduler
lr_scheduler = DummyScheduler(
name=args.lr_scheduler,
optimizer=optimizer,
total_num_steps=args.max_train_steps * accelerator.num_processes,
num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes,
)
else:
lr_scheduler = get_scheduler(
args.lr_scheduler,
optimizer=optimizer,
num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes,
num_training_steps=args.max_train_steps * accelerator.num_processes,
num_cycles=args.lr_num_cycles,
power=args.lr_power,
)
# Prepare everything with our `accelerator`.
transformer, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
@@ -494,7 +503,6 @@ def main(args):
# Train!
total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
num_trainable_parameters = sum(param.numel() for model in params_to_optimize for param in model["params"])
accelerator.print("***** Running training *****")
accelerator.print(f" Num trainable parameters = {num_trainable_parameters}")
@@ -551,6 +559,8 @@ def main(args):
torch.cuda.empty_cache()
torch.cuda.synchronize(accelerator.device)
alphas_cumprod = scheduler.alphas_cumprod.to(accelerator.device, dtype=torch.float32)
for epoch in range(first_epoch, args.num_train_epochs):
transformer.train()
@@ -618,8 +628,7 @@ def main(args):
model_pred = scheduler.get_velocity(model_output, noisy_model_input, timesteps)
alphas_cumprod = scheduler.alphas_cumprod[timesteps]
weights = 1 / (1 - alphas_cumprod)
weights = 1 / (1 - alphas_cumprod[timesteps])
while len(weights.shape) < len(model_pred.shape):
weights = weights.unsqueeze(-1)
@@ -641,7 +650,8 @@ def main(args):
optimizer.step()
optimizer.zero_grad()
lr_scheduler.step()
if not args.use_cpu_offload_optimizer:
lr_scheduler.step()
# Checks if the accelerator has performed an optimization step behind the scenes
if accelerator.sync_gradients:
@@ -674,9 +684,10 @@ def main(args):
accelerator.save_state(save_path)
logger.info(f"Saved state to {save_path}")
last_lr = lr_scheduler.get_last_lr()[0] if lr_scheduler is not None else args.learning_rate
logs = {
"loss": loss.detach().item(),
"lr": lr_scheduler.get_last_lr()[0],
"lr": last_lr,
"gradient_norm_before_clip": gradient_norm_before_clip,
"gradient_norm_after_clip": gradient_norm_after_clip,
}
training/utils.py
+82 -29
View File
@@ -1,4 +1,5 @@
import gc
import inspect
from typing import Optional, Tuple, Union
import torch
@@ -22,7 +23,11 @@ def get_optimizer(
prodigy_use_bias_correction: bool = False,
prodigy_safeguard_warmup: bool = False,
use_8bit: bool = False,
use_4bit: bool = False,
use_torchao: bool = False,
use_deepspeed: bool = False,
use_cpu_offload_optimizer: bool = False,
offload_gradients: bool = False,
) -> torch.optim.Optimizer:
optimizer_name = optimizer_name.lower()
@@ -38,18 +43,32 @@ def get_optimizer(
weight_decay=weight_decay,
)
if use_8bit and use_4bit:
raise ValueError("Cannot set both `use_8bit` and `use_4bit` to True.")
if (use_torchao and (use_8bit or use_4bit)) or use_cpu_offload_optimizer:
try:
import torchao
torchao.__version__
except ImportError:
raise ImportError(
"To use optimizers from torchao, please install the torchao library: `USE_CPP=0 pip install torchao`."
)
if not use_torchao and use_4bit:
raise ValueError("4-bit Optimizers are only supported with torchao.")
# Optimizer creation
supported_optimizers = ["adam", "adamw", "prodigy"]
supported_optimizers = ["adam", "adamw", "prodigy", "came"]
if optimizer_name not in supported_optimizers:
logger.warning(
f"Unsupported choice of optimizer: {optimizer_name}. Supported optimizers include {supported_optimizers}. Defaulting to `AdamW`."
)
optimizer_name = "adamw"
if use_8bit and optimizer_name not in ["adam", "adamw"]:
logger.warning(
f"use_8bit_adam is ignored when optimizer is not set to 'Adam' or 'AdamW'. Optimizer was set to {optimizer_name}."
)
if (use_8bit or use_4bit) and optimizer_name not in ["adam", "adamw"]:
raise ValueError("`use_8bit` and `use_4bit` can only be used with the Adam and AdamW optimizers.")
if use_8bit:
try:
@@ -60,24 +79,32 @@ def get_optimizer(
)
if optimizer_name == "adamw":
optimizer_class = bnb.optim.AdamW8bit if use_8bit else torch.optim.AdamW
if use_torchao:
from torchao.prototype.low_bit_optim import AdamW4bit, AdamW8bit
optimizer = optimizer_class(
params_to_optimize,
betas=(beta1, beta2),
eps=epsilon,
weight_decay=weight_decay,
)
optimizer_class = AdamW8bit if use_8bit else AdamW4bit if use_4bit else torch.optim.AdamW
else:
optimizer_class = bnb.optim.AdamW8bit if use_8bit else torch.optim.AdamW
init_kwargs = {
"betas": (beta1, beta2),
"eps": epsilon,
"weight_decay": weight_decay,
}
elif optimizer_name == "adam":
optimizer_class = bnb.optim.Adam8bit if use_8bit else torch.optim.Adam
if use_torchao:
from torchao.prototype.low_bit_optim import Adam4bit, Adam8bit
optimizer = optimizer_class(
params_to_optimize,
betas=(beta1, beta2),
eps=epsilon,
weight_decay=weight_decay,
)
optimizer_class = Adam8bit if use_8bit else Adam4bit if use_4bit else torch.optim.Adam
else:
optimizer_class = bnb.optim.Adam8bit if use_8bit else torch.optim.Adam
init_kwargs = {
"betas": (beta1, beta2),
"eps": epsilon,
"weight_decay": weight_decay,
}
elif optimizer_name == "prodigy":
try:
@@ -92,17 +119,43 @@ def get_optimizer(
"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=learning_rate,
betas=(beta1, beta2),
beta3=beta3,
weight_decay=weight_decay,
eps=epsilon,
decouple=prodigy_decouple,
use_bias_correction=prodigy_use_bias_correction,
safeguard_warmup=prodigy_safeguard_warmup,
init_kwargs = {
"lr": learning_rate,
"betas": (beta1, beta2),
"beta3": beta3,
"eps": epsilon,
"weight_decay": weight_decay,
"decouple": prodigy_decouple,
"use_bias_correction": prodigy_use_bias_correction,
"safeguard_warmup": prodigy_safeguard_warmup,
}
elif optimizer_name == "came":
try:
import came_pytorch
except ImportError:
raise ImportError("To use CAME, please install the came-pytorch library: `pip install came-pytorch`")
optimizer_class = came_pytorch.CAME
init_kwargs = {
"lr": learning_rate,
"eps": (1e-30, 1e-16),
"betas": (beta1, beta2, beta3),
"weight_decay": weight_decay,
}
if use_cpu_offload_optimizer:
from torchao.prototype.low_bit_optim import CPUOffloadOptimizer
if "fused" in inspect.signature(optimizer_class.__init__).parameters:
init_kwargs.update({"fused": True})
optimizer = CPUOffloadOptimizer(
params_to_optimize, optimizer_class=optimizer_class, offload_gradients=offload_gradients, **init_kwargs
)
else:
optimizer = optimizer_class(params_to_optimize, **init_kwargs)
return optimizer