Spaces:
Running on Zero
Running on Zero
File size: 7,399 Bytes
dcc4594 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 | from contextlib import nullcontext
import torch
import os
from models.video_tokenizer import VideoTokenizer
from datasets.data_utils import visualize_reconstruction, load_data_and_data_loaders
from utils.scheduler_utils import create_cosine_scheduler
from tqdm import tqdm
import wandb
from utils.utils import readable_timestamp, save_training_state, prepare_stage_dirs, prepare_pipeline_run_root
from utils.config import VideoTokenizerConfig, load_stage_config_merged
from utils.utils import save_training_state, load_videotokenizer_from_checkpoint
from utils.wandb_utils import init_wandb, log_training_metrics, log_system_metrics, log_learning_rate, finish_wandb
from dataclasses import asdict
from utils.distributed import init_distributed_from_env, prepare_model_for_distributed, unwrap_model, print_param_count_if_main, cleanup_distributed
from torch.distributed.fsdp import FSDPModule
def main():
# vidtokenizer config merged with training_config.yaml (training takes priority), plus CLI overrides
args: VideoTokenizerConfig = load_stage_config_merged(VideoTokenizerConfig, default_config_path=os.path.join(os.getcwd(), 'configs', 'video_tokenizer.yaml'))
# DDP setup
dist_setup = init_distributed_from_env()
# run save dir if it doesn't exist (running not from full train)
timestamp = readable_timestamp()
run_root = os.environ.get('NG_RUN_ROOT_DIR')
if not run_root:
run_root, _ = prepare_pipeline_run_root(base_cwd=os.getcwd())
is_main = dist_setup['is_main']
stage_dir, checkpoints_dir, visualizations_dir = prepare_stage_dirs(run_root, 'video_tokenizer')
if is_main:
print(f"Video Tokenizer Training")
print(f'Results will be saved in {stage_dir}')
# dataloader
data_overrides = {}
if hasattr(args, 'fps') and args.fps is not None:
data_overrides['fps'] = args.fps
if hasattr(args, 'preload_ratio') and args.preload_ratio is not None:
data_overrides['preload_ratio'] = args.preload_ratio
training_data, validation_data, training_loader, validation_loader, x_train_var = load_data_and_data_loaders(
dataset=args.dataset,
batch_size=args.batch_size_per_gpu,
num_frames=args.context_length,
distributed=dist_setup['is_distributed'],
rank=dist_setup['device_mesh'].get_rank() if dist_setup['device_mesh'] is not None else 0,
world_size=dist_setup['world_size'],
**data_overrides,
)
# print("Length of training data:", len(training_data))
# print("Length of validation data:", len(validation_data))
# init model and optional ckpt load
model = VideoTokenizer(
frame_size=(args.frame_size, args.frame_size),
patch_size=args.patch_size,
embed_dim=args.embed_dim,
num_heads=args.num_heads,
hidden_dim=args.hidden_dim,
num_blocks=args.num_blocks,
latent_dim=args.latent_dim,
num_bins=args.num_bins,
).to(args.device)
if args.checkpoint:
model, _ = load_videotokenizer_from_checkpoint(
args.checkpoint,
args.device,
model,
dist_setup['is_distributed'],
)
# optional DDP, compile, param count, tf32
print_param_count_if_main(model, "VideoTokenizer", is_main)
if args.compile:
model = torch.compile(model, mode="reduce-overhead", fullgraph=False, dynamic=True)
model = prepare_model_for_distributed(
model,
args.distributed,
model_type=model.model_type,
device_mesh=dist_setup['device_mesh'],
)
if args.tf32:
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
# create optimizer(s) — AdamW or Muon+AdamW split
from utils.optimizer_utils import create_optimizer
optimizers = create_optimizer(model, args)
# cosine scheduler for lr warmup
schedulers = [create_cosine_scheduler(opt, args.n_updates) for opt in optimizers]
train_ctx = torch.amp.autocast(args.device, enabled=True, dtype=torch.bfloat16) if args.amp and not args.distributed.use_fsdp else nullcontext()
results = {
'n_updates': 0,
'loss_vals': [],
}
# init wandb
if args.use_wandb and is_main:
cfg = asdict(args)
cfg.update({'timestamp': timestamp})
run_name = f"video_tokenizer_{timestamp}"
init_wandb(args.wandb_project, cfg, run_name)
unwrap_model(model).train()
train_iter = iter(training_loader)
for i in tqdm(range(args.n_updates), disable=not is_main):
for opt in optimizers:
opt.zero_grad(set_to_none=True)
if isinstance(model, FSDPModule):
model.set_requires_gradient_sync(False)
if args.compile:
torch.compiler.cudagraph_mark_step_begin()
for micro_batch in range(args.gradient_accumulation_steps):
try:
(x, _) = next(train_iter)
except StopIteration:
train_iter = iter(training_loader) # reset iterator when epoch ends
(x, _) = next(train_iter)
x = x.to(args.device, non_blocking=True)
with train_ctx:
loss, x_hat = model(x)
loss /= args.gradient_accumulation_steps
if isinstance(model, FSDPModule):
if (micro_batch + 1) % args.gradient_accumulation_steps == 0:
model.set_requires_gradient_sync(True)
loss.backward()
torch.nn.utils.clip_grad_norm_(unwrap_model(model).parameters(), max_norm=1.0)
for opt in optimizers:
opt.step()
for sched in schedulers:
sched.step()
results["loss_vals"].append(loss.cpu().detach())
results["n_updates"] = i
# wandb logging
if args.use_wandb and is_main:
metrics = {
'loss': loss.item(),
'learning_rate': schedulers[0].get_last_lr()[0],
}
log_training_metrics(i, metrics, prefix='train')
log_system_metrics(i)
log_learning_rate(optimizers[0], i)
# save model and visualize results
if i % args.log_interval == 0:
if args.use_wandb:
with torch.no_grad():
indices = unwrap_model(model).tokenize(x)
unique_codes = torch.unique(indices).numel()
codebook_usage = unique_codes / unwrap_model(model).codebook_size
if is_main:
wandb.log({'train/codebook_usage': codebook_usage}, step=i)
hyperparameters = args.__dict__
save_training_state(model, optimizers[0], schedulers[0], hyperparameters, checkpoints_dir, prefix='video_tokenizer', step=i)
if is_main:
x_hat_vis = x_hat.detach().cpu()
x_vis = x.detach().cpu()
save_path = os.path.join(visualizations_dir, f'video_tokenizer_recon_step_{i}.png')
visualize_reconstruction(x_vis[:16], x_hat_vis[:16], save_path)
print('\n Step', i, 'Loss:', torch.mean(torch.stack(results["loss_vals"][-args.log_interval:])).item())
# finish wandb
if args.use_wandb and is_main:
finish_wandb()
cleanup_distributed(dist_setup['is_distributed'])
if __name__ == "__main__":
main() |