nicolabs/Open-Sora
1
0
Fork 0
mirror of https://github.com/hpcaitech/Open-Sora.git synced 2026-05-21 11:59:01 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

[feature] support grad accumulation and tensorboard (#11)

Browse source 
* [feature] support grad accumulation

* [feature] support tensorboard

* [feature] support lr config
This commit is contained in:
Hongxin Liu 2024-02-26 17:06:06 +08:00 committed by GitHub
parent adba00f151
commit 9644774813
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 34 additions and 11 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

45
train.py
View file

@ -25,6 +25,7 @@ from colossalai.booster.plugin import LowLevelZeroPlugin
from colossalai.cluster import DistCoordinator
from colossalai.logging import get_dist_logger
from colossalai.utils import get_current_device
from torch.utils.tensorboard import SummaryWriter
from tqdm import tqdm
from transformers import AutoModel
@ -78,6 +79,8 @@ def main(args):
if coordinator.is_master():
os.makedirs(args.checkpoint_dir, exist_ok=True)
os.makedirs(args.tensorboard_dir, exist_ok=True)
writer = SummaryWriter(args.tensorboard_dir)
# Setup model
if len(args.vqvae) > 0:
@ -105,7 +108,7 @@ def main(args):
) # default: 1000 steps, linear noise schedule
# Setup optimizer (we used default Adam betas=(0.9, 0.999) and a constant learning rate of 1e-4 in our paper):
opt = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=0)
opt = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=0)
# Setup dataloader
dataset = load_datasets(args.dataset)
@ -125,13 +128,16 @@ def main(args):
ranks=[0],
)
num_steps_per_epoch = len(dataloader) // args.accumulation_steps
for epoch in range(args.epochs):
dataloader.sampler.set_epoch(epoch)
with tqdm(
desc=f"Epoch {epoch}",
disable=not coordinator.is_master(),
total=len(dataloader),
total=num_steps_per_epoch,
) as pbar:
total_loss = torch.tensor(0.0, device=get_current_device())
for step, batch in enumerate(dataloader):
batch = preprocess_batch(batch, patch_size, vqvae)
video_inputs = batch.pop("video_latent_states")
@ -145,20 +151,34 @@ def main(args):
loss_dict = diffusion.training_losses(
model, video_inputs, t, batch, mask=mask
)
loss = loss_dict["loss"].mean()
loss = loss_dict["loss"].mean() / args.accumulation_steps
total_loss.add_(loss.data)
booster.backward(loss, opt)
opt.step()
opt.zero_grad()
update_ema(ema, model)
loss_no_grad = loss.data
all_reduce_mean(loss_no_grad)
pbar.set_postfix({"Loss": f"{loss_no_grad.item():.4f}"})
pbar.update()
if (step + 1) % args.accumulation_steps == 0:
opt.step()
opt.zero_grad()
update_ema(ema, model)
all_reduce_mean(total_loss)
pbar.set_postfix({"Loss": f"{total_loss.item():.4f}"})
if coordinator.is_master():
global_step = (epoch * num_steps_per_epoch) + (
step + 1
) // args.accumulation_steps
writer.add_scalar(
tag="Loss",
scalar_value=total_loss.item(),
global_step=global_step,
)
pbar.update()
total_loss.zero_()
# Save DiT checkpoint:
if (
args.save_interval > 0 and (step + 1) % args.save_interval == 0
args.save_interval > 0
and (step + 1) % (args.save_interval * args.accumulation_steps) == 0
) or (step + 1) == len(dataloader):
save_path = os.path.join(
args.checkpoint_dir, f"epoch-{epoch}-step-{step}"
@ -196,8 +216,11 @@ if __name__ == "__main__":
parser.add_argument("-e", "--epochs", type=int, default=10)
parser.add_argument("-b", "--batch_size", type=int, default=4)
parser.add_argument("-g", "--grad_checkpoint", action="store_true", default=False)
parser.add_argument("-a", "--accumulation_steps", default=1, type=int)
parser.add_argument("--lr", type=float, default=1e-4)
parser.add_argument("--save_interval", type=int, default=20)
parser.add_argument("--checkpoint_dir", type=str, default="checkpoints")
parser.add_argument("--tensorboard_dir", type=str, default="runs")
parser.add_argument("--vqvae", default="hpcai-tech/vqvae")
args = parser.parse_args()
main(args)