dinov3/configs/config.py

# 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 logging
import math
import os
import pathlib
import sys
from dataclasses import dataclass, field
from datetime import timedelta
from typing import Any, List, Optional, Sequence, Tuple

from omegaconf import DictConfig, OmegaConf

import dinov3.distributed as distributed
from dinov3.logging import cleanup_logging, setup_logging
from dinov3.utils import fix_random_seeds, get_conda_env, get_sha

logger = logging.getLogger("dinov3")


@dataclass
class DinoV3SetupArgs:
    config_file: str
    pretrained_weights: str | None = None
    shard_unsharded_model: bool = False
    output_dir: str = ""
    opts: List[Any] = field(default_factory=lambda: [])

    def __post_init__(self):
        # When loaded from benchmark.yaml, self.opts is a frozen omegaconf.ListConfig,
        # which works everywhere except when we want to modify it or when
        # we try to json-serialize it. So we convert it to a regular list here.
        if OmegaConf.is_config(self.opts):
            self.opts = OmegaConf.to_object(self.opts)


def apply_scaling_rules_to_cfg(cfg):  # to fix
    assert distributed.is_enabled(), "Setup distributed to get global size !"
    if "schedules" in cfg:
        # For schedules v2, the scaling rules are applied when building the schedules, the config is not modified
        return cfg

    if cfg.optim.scaling_rule == "linear_wrt_256":
        old_lr = cfg.optim.lr
        cfg.optim.lr *= cfg.train.batch_size_per_gpu * distributed.get_world_size() / 256.0
        logger.info(f"linear scaling learning rate; old: {old_lr}, new: {cfg.optim.lr}")
    elif cfg.optim.scaling_rule == "sqrt_wrt_1024":
        old_lr = cfg.optim.lr
        cfg.optim.lr *= 4 * math.sqrt(cfg.train.batch_size_per_gpu * distributed.get_world_size() / 1024.0)
        logger.info(f"sqrt scaling learning rate; old: {old_lr}, new: {cfg.optim.lr}")
    return cfg


def write_config(cfg, output_dir, name="config.yaml"):
    logger.info(OmegaConf.to_yaml(cfg))
    output_dir = os.path.abspath(output_dir)
    saved_cfg_path = os.path.join(output_dir, name)
    with open(saved_cfg_path, "w") as f:
        OmegaConf.save(config=cfg, f=f)
    return saved_cfg_path


def get_default_config() -> DictConfig:
    p = pathlib.Path(__file__).parent / "ssl_default_config.yaml"
    return OmegaConf.load(p)


def get_cfg_from_args(args: DinoV3SetupArgs, multidistillation=False, strict=True):
    overrides = [*args.opts]
    if args.output_dir is not None:
        overrides.append(f"train.output_dir={os.path.realpath(args.output_dir)}")

    # Config file
    cfg = OmegaConf.load(args.config_file)

    # Command line overrides
    opts_cfg = OmegaConf.from_cli(overrides)

    if multidistillation:
        cfg = OmegaConf.merge(cfg, opts_cfg)
    else:
        # Default config
        default_cfg = get_default_config()
        if strict:
            OmegaConf.set_struct(default_cfg, True)
        cfg = OmegaConf.merge(default_cfg, cfg, opts_cfg)
    return cfg


def setup_config(args: DinoV3SetupArgs, strict_cfg=True):
    """
    Create configs and perform basic setups.
    """
    # Create the cfg with OmegaConf
    cfg = get_cfg_from_args(args, strict=strict_cfg)
    # setup distributed, logging, and random seeds
    logger.info("\n".join("%s: %s" % (k, str(v)) for k, v in sorted(dict(vars(args)).items())))
    # dump config before modifying so it can be reloaded
    if args.output_dir is not None:
        write_config(cfg, args.output_dir)
    # modify the config inplace by applying scaling rules
    apply_scaling_rules_to_cfg(cfg)
    return cfg


def _enumerate_all_subgroup_ranks(all_subgroup_rank_spans: Sequence[Tuple[int, int]]):
    """Expands a specification of process subgroups from spans to enumerated ranks.

    Args:
       all_group_rank_spans: a sequence of rank spans (first rank, last rank),
           one for each process group. Example: ((0, 1), (2, 3), (4, 7)).
    """
    for first, last in all_subgroup_rank_spans:
        assert first <= last
    return tuple(tuple(range(first, last + 1)) for first, last in all_subgroup_rank_spans)


def setup_multidistillation(args: DinoV3SetupArgs):
    base_output_dir = args.output_dir
    os.makedirs(args.output_dir, exist_ok=True)
    # get config file for this rank
    base_cfg = OmegaConf.load(args.config_file)
    assert base_cfg.multidistillation.enabled

    global_batch_size = base_cfg.multidistillation.global_batch_size

    distributed.enable(overwrite=True)
    seed = getattr(args, "seed", 0)
    rank = distributed.get_rank()

    # build process subgroups
    all_subgroup_rank_spans = tuple(
        (student.ranks_range[0], student.ranks_range[1] - 1) for student in base_cfg.multidistillation.students
    )
    all_subgroup_ranks = _enumerate_all_subgroup_ranks(all_subgroup_rank_spans)
    distributed.new_subgroups(all_subgroup_ranks)

    found = False
    for student in base_cfg.multidistillation.students:
        if rank in range(*student.ranks_range):
            found = True
            break
    assert found, "rank of worker not in defined range"

    name = student.name
    config_path = student.config_path
    n_gpus = student.ranks_range[1] - student.ranks_range[0]
    assert global_batch_size % n_gpus == 0
    total_n_gpus = distributed.get_world_size()

    args.output_dir = os.path.join(base_output_dir, name)
    args.opts += [f"train.output_dir={args.output_dir}"]
    args.opts += [f"train.batch_size_per_gpu={global_batch_size // total_n_gpus}"]
    args.config_file = os.path.abspath(config_path)
    default_cfg = get_default_config()
    cfg = OmegaConf.load(args.config_file)
    cfg = OmegaConf.merge(default_cfg, cfg, base_cfg, OmegaConf.from_cli(args.opts))

    global logger
    setup_logging(output=args.output_dir, level=logging.INFO)

    fix_random_seeds(seed + rank)

    write_config(cfg, args.output_dir)
    apply_scaling_rules_to_cfg(cfg)

    return cfg


def setup_job(
    output_dir: Optional[str] = None,
    distributed_enabled: bool = True,
    logging_enabled: bool = True,
    seed: Optional[int] = 0,
    restrict_print_to_main_process: bool = True,
    distributed_timeout: timedelta | None = None,
):
    """
    Setup methods that should be done in every fairvit job
    Initializes logging, distributed, random seeds and other utilities.
    """
    if output_dir is not None:
        output_dir = os.path.realpath(output_dir)
        os.makedirs(output_dir, exist_ok=True)

    if logging_enabled:
        setup_logging(
            output=output_dir,
            level=logging.INFO,
            log_to_stdout_only_in_main_process=restrict_print_to_main_process,
        )

    if distributed_enabled:
        distributed.enable(
            overwrite=True,
            nccl_async_error_handling=True,
            restrict_print_to_main_process=restrict_print_to_main_process,
            timeout=distributed_timeout,
        )

    if seed is not None:
        rank = distributed.get_rank()
        fix_random_seeds(seed + rank)

    logger = logging.getLogger("dinov3")
    logger.info("git:\n  {}\n".format(get_sha()))

    # Log some python info
    conda_env_name, conda_env_path = get_conda_env()
    logger.info(f"conda env name: {conda_env_name}")
    logger.info(f"conda env path: {conda_env_path}")
    logger.info(f"python path: {sys.path}")


def exit_job(distributed_enabled: bool = True, logging_enabled: bool = True):
    if distributed_enabled:
        distributed.disable()
    if logging_enabled:
        cleanup_logging()