# Copyright (c) Meta Platforms, Inc. and affiliates. # # This software may be used and distributed in accordance with # the terms of the DINOv3 License Agreement. import gc import logging from functools import partial import math import os import sys from pathlib import Path from typing import Callable import dinov3.distributed as distributed import torch from dinov3.checkpointer import ( find_latest_checkpoint, keep_last_n_checkpoints, load_checkpoint, save_checkpoint, ) from dinov3.configs import setup_job from dinov3.data import SamplerType, make_data_loader, make_dataset from dinov3.eval.text.build_dinotxt import build_model_and_tokenizer from dinov3.eval.text.clip_loss import memory_efficient_clip_loss from dinov3.eval.text.dinotxt_model import DINOTxt, DINOTxtConfig from dinov3.eval.text.gram_loss import gram_loss_fn from dinov3.logging import MetricLogger, setup_logging from dinov3.train.cosine_lr_scheduler import linear_warmup_cosine_decay from omegaconf import OmegaConf from torch import optim logger = logging.getLogger("dinov3") def unwrap_model(model): return getattr(model, "module", model) def test( model: DINOTxt, iteration: str, output_dir: str, ): eval_dir = Path(output_dir) / "eval" / str(iteration) if distributed.is_subgroup_main_process(): eval_dir.mkdir(parents=True, exist_ok=True) ckpt_dir = eval_dir / str("sharded_model_checkpoint") save_checkpoint(ckpt_dir, iteration=iteration, model=model) def apply_learning_rate(optimizer, lr: float): for param_group in optimizer.param_groups: param_group["lr"] = lr def exclude(n: str, p: torch.Tensor) -> bool: return p.ndim < 2 or "bn" in n or "ln" in n or "bias" in n or "logit_scale" in n def include(n: str, p: torch.Tensor) -> bool: return not exclude(n, p) def train( train_dataset, model: DINOTxt, tokenizer: Callable, max_iteration: int, warmup_length: int, checkpointing_period: int, output_dir: str, dtype_str: str, sampler_type: SamplerType, lr_scheduler_type: str, lr: float, weight_decay: float, batch_size: int = 256, beta1: float = 0.9, beta2: float = 0.99, eps: float = 1e-8, num_workers: int = 10, eval_freq: int = 1000, gc_freq: int = 100, use_gram_loss: bool = False, patch_sampling_rate_for_gram_loss: float = 0.5, normalize_patch_tokens_for_gram_loss: bool = False, gram_loss_weight: float = 1.0, max_checkpoints_to_keep: int = None, resume: bool = False, seed: int = 11, ): named_parameters = list(model.named_parameters()) gain_or_bias_params = [ p for n, p in named_parameters if exclude(n, p) and p.requires_grad ] gain_or_bias_params_names = [ n for n, p in named_parameters if exclude(n, p) and p.requires_grad ] rest_params = [p for n, p in named_parameters if include(n, p) and p.requires_grad] rest_params_names = [ n for n, p in named_parameters if include(n, p) and p.requires_grad ] logger.info(f"Gain or bias params: {gain_or_bias_params_names}") logger.info(f"Rest params: {rest_params_names}") logger.info( f"Learning rate: {lr}, batch_size_per_gpu: {batch_size}, weight_decay: {weight_decay}" ) optimizer = optim.AdamW( [ {"params": gain_or_bias_params, "weight_decay": 0.0}, {"params": rest_params, "weight_decay": weight_decay}, ], lr=lr, betas=(beta1, beta2), eps=eps, ) learning_rates = linear_warmup_cosine_decay( 0, lr, 0, warmup_iterations=warmup_length, total_iterations=max_iteration ) logger.info( f"Init lr scheduler: {lr_scheduler_type}, warmup length: {warmup_length}, base_lr: {lr}, max iter: {max_iteration}" ) if ( resume and (ckpt_dir := find_latest_checkpoint(os.path.join(output_dir, "ckpt"))) is not None ): iteration = load_checkpoint(ckpt_dir, model=model, optimizer=optimizer) start_iteration = iteration + 1 del iteration, ckpt_dir else: logger.info("Initializing from scratch") start_iteration = 0 def collate_fn(batch): images, captions = list(zip(*batch))[:2] return torch.stack(images), tokenizer.tokenize(captions) train_data_loader = make_data_loader( dataset=train_dataset, batch_size=batch_size, num_workers=num_workers, shuffle=False, seed=seed, sampler_type=sampler_type, sampler_advance=start_iteration, drop_last=True, persistent_workers=True, collate_fn=collate_fn, ) rank = distributed.get_rank() world_size = distributed.get_world_size() logger.info( f"Init loss function: rank: {distributed.get_rank()}, world size: {world_size}" ) clip_loss = partial(memory_efficient_clip_loss, group=torch.distributed.group.WORLD) cur_iteration = start_iteration logger.info(f"Starting training from iteration {start_iteration}...") header = "Training" metric_logger = MetricLogger(delimiter=" ") gc.disable() device_id = rank % torch.cuda.device_count() for batch in metric_logger.log_every( train_data_loader, 10, header, max_iteration, start_iteration, ): images, text_tokens = batch images = images.to(device=device_id, non_blocking=True) text_tokens = text_tokens.to(device=device_id, non_blocking=True) ( image_embeddings, text_embeddings, logit_scale, patch_tokens, backbone_patch_tokens, ) = model(images, text_tokens) contrastive_loss = clip_loss(image_embeddings, text_embeddings, logit_scale) total_loss = contrastive_loss if use_gram_loss: gram_loss = gram_loss_fn( patch_tokens, backbone_patch_tokens, patch_sampling_rate_for_gram_loss, normalize_patch_tokens_for_gram_loss, ) total_loss = contrastive_loss + gram_loss_weight * gram_loss if total_loss.isnan(): msg = f"Loss is NaN at iteration {cur_iteration}, aborting..." logger.error(msg) raise RuntimeError(msg) apply_learning_rate(optimizer=optimizer, lr=learning_rates[cur_iteration]) optimizer.zero_grad() total_loss.backward() optimizer.step() # This clamping trick is from OpenCLIP reposistory which MetaCLIP follows. Orginally used in CLIP training. # NOTE: we clamp to 4.6052 = ln(100), as in the original paper. with torch.no_grad(): unwrap_model(model).logit_scale.clamp_(0, math.log(100)) metric_logger.update(contrastive_loss=contrastive_loss.item()) if use_gram_loss: metric_logger.update(gram_loss=gram_loss.item()) metric_logger.update(lr=optimizer.param_groups[0]["lr"]) metric_logger.update(logit_scale=logit_scale.item()) is_last_iteration = (cur_iteration + 1) == max_iteration is_ckpt_iteration = ( (cur_iteration + 1) % checkpointing_period == 0 ) or is_last_iteration if is_ckpt_iteration: ckpt_dir = Path(output_dir, "ckpt").expanduser() save_checkpoint( ckpt_dir / str(cur_iteration), iteration=cur_iteration, model=model, optimizer=optimizer, ) if distributed.is_main_process(): keep_last_n_checkpoints(ckpt_dir, max_checkpoints_to_keep) if eval_freq > 0 and (cur_iteration + 1) % eval_freq == 0: test( model, iteration=f"training_{cur_iteration}", batch_size=batch_size, num_workers=num_workers, output_dir=output_dir, dtype_str=dtype_str, ) torch.cuda.synchronize() if (cur_iteration + 1) % gc_freq == 0: logger.info("Garbage collection...") gc.collect() cur_iteration += 1 def write_config( model_config: DINOTxtConfig, output_dir, name="clip_model_config.yaml" ): logger.info(OmegaConf.to_yaml(model_config)) saved_cfg_path = os.path.join(output_dir, name) with open(saved_cfg_path, "w") as f: OmegaConf.save(config=model_config, f=f) return saved_cfg_path def main(argv=None): if argv is None: argv = sys.argv[1:] args_dict = OmegaConf.to_container(OmegaConf.from_cli(argv)) logger.info(args_dict) config = OmegaConf.load(args_dict["trainer_config_file"]) logger.info(config) if "output_dir" in args_dict: config.output_dir = args_dict["--output-dir"] setup_job(output_dir=config.output_dir, seed=config.seed) setup_logging(output=os.path.join(config.output_dir, "nan_logs"), name="nan_logger") logger.info("Trainer config:") logger.info(config) model_config = DINOTxtConfig( embed_dim=config.embed_dim, text_backbone_config=config.text_backbone_config, vision_backbone_config=config.vision_backbone_config, text_backbone_pretrained_weights=config.text_backbone_pretrained_weights, vision_backbone_pretrained_weights=config.vision_backbone_pretrained_weights, vision_model_train_img_size=config.vision_model_train_img_size, vision_model_use_class_token=config.vision_model_use_class_token, vision_model_use_patch_tokens=config.vision_model_use_patch_tokens, vision_model_num_head_blocks=config.vision_model_num_head_blocks, vision_model_head_blocks_drop_path=config.vision_model_head_blocks_drop_path, vision_model_use_linear_projection=config.vision_model_use_linear_projection, vision_model_patch_tokens_pooler_type=config.vision_model_patch_tokens_pooler_type, vision_model_patch_token_layer=config.vision_model_patch_token_layer, text_model_freeze_backbone=config.text_model_freeze_backbone, text_model_num_head_blocks=config.text_model_num_head_blocks, text_model_head_blocks_is_causal=config.text_model_head_blocks_is_causal, text_model_head_blocks_drop_prob=config.text_model_head_blocks_drop_prob, text_model_tokens_pooler_type=config.text_model_tokens_pooler_type, text_model_use_linear_projection=config.text_model_use_linear_projection, text_vocab_path_or_url=config.text_vocab_path_or_url, init_logit_scale=config.init_logit_scale, freeze_logit_scale=config.freeze_logit_scale, init_logit_bias=config.init_logit_bias, ) write_config(model_config=model_config, output_dir=config.output_dir) model, transform, tokenizer = build_model_and_tokenizer( model_config, use_fsdp=config.use_fsdp, do_compile=config.do_compile, use_ac=config.use_ac, use_cuda_graphs=config.use_cuda_graphs, ) train_dataset = make_dataset( dataset_str=config.train_dataset_str, transform=transform, ) sampler_type = ( SamplerType.SHARDED_INFINITE if config.dataset_use_cache else SamplerType.INFINITE ) train( train_dataset=train_dataset, model=model, tokenizer=tokenizer, max_iteration=config.max_iteration, warmup_length=config.warmup_length, checkpointing_period=config.checkpointing_period, output_dir=config.output_dir, dtype_str=config.dtype_str, lr_scheduler_type=config.lr_scheduler_type, lr=config.lr, weight_decay=config.weight_decay, batch_size=config.batch_size, beta1=config.beta1, beta2=config.beta2, eps=config.eps, sampler_type=sampler_type, eval_freq=config.eval_freq, gc_freq=config.gc_freq, max_checkpoints_to_keep=config.max_checkpoints_to_keep, use_gram_loss=config.vision_model_use_gram_loss, patch_sampling_rate_for_gram_loss=config.vision_model_patch_sampling_rate_for_gram_loss, normalize_patch_tokens_for_gram_loss=config.vision_model_normalize_patch_tokens_for_gram_loss, gram_loss_weight=config.vision_model_gram_loss_weight, resume=not config.no_resume, ) if __name__ == "__main__": main()