lcm/train/trainer.py

# Copyright (c) Meta Platforms, Inc. and affiliates
# All rights reserved.
#
#

import gc
import logging
import os
import sys
from abc import abstractmethod
from contextlib import nullcontext
from dataclasses import asdict, dataclass, field
from functools import cached_property
from itertools import count
from pathlib import Path
from pprint import pformat
from typing import (
    Any,
    ContextManager,
    Dict,
    Iterator,
    List,
    Literal,
    Mapping,
    Optional,
    Tuple,
)

import torch
import yaml
from fairseq2.assets import AssetCard, AssetCardFieldNotFoundError
from fairseq2.checkpoint import FileCheckpointManager
from fairseq2.gang import FakeGang, Gang, ReduceOperation, all_sum
from fairseq2.logging import get_log_writer
from fairseq2.metrics import (
    LogMetricRecorder,
    MetricBag,
    MetricRecorder,
    TensorBoardRecorder,
    record_metrics,
)
from fairseq2.nn.ddp import to_ddp
from fairseq2.nn.fsdp import to_fsdp
from fairseq2.nn.utils.gradient import (
    check_gradient_norms,
    clip_gradient_norm,
    scale_gradients,
)
from fairseq2.nn.utils.module import (
    _get_named_modules,
    freeze_parameters,
    to_device,
)
from fairseq2.optim import AdamW, DynamicLossScaler
from fairseq2.optim.lr_scheduler import AbstractLRScheduler, get_effective_lr
from fairseq2.recipes.utils.log import log_model
from fairseq2.utils.profiler import Profiler, Stopwatch
from fairseq2.utils.rng import RngBag
from fairseq2.utils.state import StatefulObjectBag
from omegaconf import MISSING
from stopes.core import Requirements
from torch.distributed.fsdp.fully_sharded_data_parallel import (
    FullyShardedDataParallel as FSDP,
)
from torch.nn import Module
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.optim import Optimizer
from torch.profiler import record_function
from torcheval.metrics import Mean

from lcm.datasets.configs import DataLoadingConfig, ValidationDataLoadingConfig
from lcm.datasets.dataloading import ds_name
from lcm.train.metrics import (
    LCMWandBRecorder,
    flatten_dict,
)
from lcm.train.optim import build_lr_scheduler
from lcm.utils.data_utils import update_dataclass
from lcm.utils.distributed import (
    SUPPORTED_FSDP_MEMORY_POLICIES,
    SUPPORTED_FSDP_WRAP_POLICIES,
    get_fsdp_memory_policy,
    get_fsdp_wrap_policy,
    init_process_group,
)
from lcm.utils.logging import (
    log_env_variables,
    setup_additional_logging,
)

logger = get_log_writer(__name__)


@dataclass
class TrainingConfig:
    """Holds the configuration of a training job."""

    training_data: Any = MISSING
    """The datasets to train with."""

    validation_data: Any = MISSING
    """The datasets to validate on."""

    model_arch: Optional[str] = None
    """Starting architecture for the model to train"""

    model_arch_overrides: Optional[Dict] = None
    """Dict of parameters to overwrite in `model_arch`"""

    model_config_or_name: Optional[Any] = None
    """The model configuration or name to train.
        This option cannot be paired with model_arch + model_arch_overrides
        If provided, this option supersedes model_arch + model_arch_overrides
    """
    output_dir: Path = MISSING
    """The output directory to store checkpoints and logs."""

    log_folder: Optional[Path] = None
    """The executor's log directory where stdout/stderr will be redirected.
        We will use this directory to optionally enable ATEN and NCCL
        logging (if debug is True) """

    tb_dir: Optional[Path] = None
    """The output directory to store tensorbaord logs"""

    # defaults to "uncategorized"
    wandb_project: Optional[str] = None
    wandb_run_name: Optional[str] = None
    wandb_entity: Optional[str] = None

    requirements: Requirements = field(
        default_factory=lambda: Requirements(
            nodes=1,
            tasks_per_node=8,
            gpus_per_node=8,
            cpus_per_task=8,
            mem_gb=256,
            timeout_min=3 * 24 * 60,
            constraint="volta32gb",
        )
    )
    """The scheduling requirements for this trainer"""

    data_loading_config: DataLoadingConfig = MISSING

    validation_data_loading_config: ValidationDataLoadingConfig = field(
        default_factory=lambda: ValidationDataLoadingConfig()
    )

    criterion: Any = MISSING

    dtype: str = "torch.float32"
    """The data type of the model."""

    lr_schedule: str = "myle"
    """The learning rate schedule out of
        `noop`: no learning rate schedule, just use the initial learning rate,
        `myle`: inv-sqrt as implemented in Fairseq,
        `cosine` cosine annealing schedule,
        `wsd` for  Warmup-Stable-Decay (WSD) or tri-stage """

    lr: float = 0.004
    """The initial (post-warm-up) learning rate for AdamW."""

    start_lr: float = 1e-7
    """The initial warmup learning rate."""

    final_lr: float = 1e-5
    """The final learning rate."""

    lr_stage_ratios: List[float] = field(default_factory=lambda: [0.1, 0.4, 0.5])
    """The ratios of the wsd (tri-stage) learning rate scheduler."""

    num_lr_warmup_steps: int = 800
    """The number of warm-up steps for the learning rate."""

    weight_decay: float = 0.1
    """The weight decay coefficient of AdamW (PyTorch default: 1e-2, Fs2 default: 0.0)."""

    adam_betas: List[float] = field(default_factory=lambda: [0.9, 0.98])
    """The beta coefficients of AdamW used for computing running averages of gradient and its square."""

    adam_eps: float = 1e-6
    """The term added to the denominator in AdamW to improve numerical stability.
        Default in FS2 and PyTorch is 1e-8. Previous hard coded value in our trainer is 1e-6"""

    use_optimizer_in_fp32: bool = True
    """if True, the optimizer (AdamW) will be initialized with `use_fp32 = True`
        i.e. we will store the optimizer state in single precision and convert
        gradients on-the-fly to single precision for numerical stability"""

    max_steps: int = 10_000
    """The maximum number of training steps."""

    max_grad_norm: float = 1000
    """Maximal gradient norm, for gradient clipping.
       gradients are multiplied by `torch.clamp(max_norm / (total_norm + 1e-6), max=1.0)`
       if max_norm is arbitrarily large, then we'll only report gradients norm
    """
    turn_off_grad_normalization: bool = False
    """If ``True``, Turn off gradient normalization"""

    gradient_accumulation: int = 1
    """The number of steps to accumulate gradients before an optimizer update."""

    validate_every_n_steps: int = 5000
    """The number of steps after which to validate the model."""

    checkpoint_every_n_steps: int = 5000
    """The number of steps after which to checkpoint."""

    keep_last_n_checkpoints: int = -1
    """The number of checkpoints to keep on disk."""

    save_model_every_n_steps: int = 5000
    """The number of steps after which to save a consolidated version of the model."""

    preserve_consolidated_models: bool = False
    """If `True`, only pt files excluding ones starting with `mdoel` will be deleted from the step checkpoint directory."""

    publish_metrics_every_n_steps: int = 1
    """The number of steps after which to publish training metrics."""

    gc_every_n_steps: int = 1000
    """The frequency of steps at which we collect garbage with `gc.collect()`."""

    seed: int = 2
    """The RNG seed to use while starting the job."""

    debug: bool = False
    """If ``True``, runs the trainer in debug mode"""

    profile: bool = False
    """If ``True``, runs the PyTorch profiler at the beginning of the training."""

    profiler_skip_first: int = 200

    profiler_active: int = 3
    """If profiling (``profile = True``), The profiler will skip the first ``skip_first`` steps, then do the active recording for the next ``active`` steps
    If planning to visualize the trace with tensorbaord, then ``active`` should be small (less than 10 steps), otherwise tb won't load!
    """
    loss_scaler_init_scale: float = 2.0**15
    """The initial scale for the gradient scaler, fairseq2's default is 2.0**15"""

    loss_scaler_scale_window: Optional[int] = None
    """The number of consecutive optimizer steps without inf/NaN gradients that must occur for the scale to be updated"""

    use_fsdp: bool = True
    """If ``True``, uses FSDP instead of DDP."""

    use_autocast: bool = False
    """If ``True``, wrap the forward pass in AMP autocast context.
        autocast is only needed if training with mixed precision.
        If training fails without it, check if some module with its weights is not properly cast
    """

    fsdp_wrap_granularity: SUPPORTED_FSDP_WRAP_POLICIES = "model"
    """The granularity at which to wrap the model."""

    fsdp_memory_policy: SUPPORTED_FSDP_MEMORY_POLICIES = "standard"
    """The FSDP memory policy."""

    fsdp_fp32_reduce: bool = False
    """ If ``True``, the gradients will be reduced in full precision even when dtype is `torch.float16`"""

    use_submitit: bool = True
    """If ``True``, setup the environment ti use submitit."""

    fake_gang_device: Optional[str] = None
    """If non-empty, the trainer will be set locally on a device, instead of distributed training."""

    experiment_name: Optional[str] = None
    """experiment name for job trackin, if None default to StopesModule naming"""

    raise_oom: bool = False
    """If ``True``, raise OOM errors when they occur, if ``False`` give it another try."""

    raise_nan_or_inf: bool = False
    """If ``True``, raise FloatingPointError with Nan/Inf losses, if ``False`` give it another try."""

    max_ooms: int = 10
    """If ```raise_oom`` is False, how many OOMs we can tolerate per rank before raising an error."""

    max_nans_or_infs: int = 10
    """If ```raise_nan_or_inf`` is False, how many Nan/Infs we can tolerate per rank before raising an error."""

    freeze_modules: Optional[List[str]] = None
    """Name of modules in the model to be frozen when training/finetuning"""

    freezing_strategy: Literal["none", "modules", "ffn", "ffn-adaln", "adaln"] = "none"
    """
    Freezing strategy to follow. Options are:
        1. none: Nothing will be frozen (default)
        2. modules: A list of modules to freeze will be read from `freeze_modules`
        3. ffn: All ffn sub-modules will be frozen
        4. ffn-adaln: all FFN and Adaln sub-modules will be frozen.
    """


class Trainer(StatefulObjectBag):
    config: TrainingConfig
    model: Module
    training_data_loader: Any
    validation_data_loader: Optional[Any]
    gang: Gang
    optimizer: Optimizer
    loss_scaler: DynamicLossScaler
    lr_scheduler: AbstractLRScheduler
    rng_bag: RngBag
    step_nr: int
    train_metric_bag: MetricBag
    valid_metric_bag: Mapping[str, MetricBag]
    metric_recorders: List[MetricRecorder]
    profiler: Profiler
    stopwatch: Stopwatch
    criterion: Any
    card_metdata: Dict
    _train_step_time: float
    _valid_step_time: float

    def __init__(
        self,
        config: TrainingConfig,
        model: Module,
        training_data_loader: Any,
        validation_data_loader: Optional[Any],
        gang: Gang,
        checkpoint_manager: FileCheckpointManager,
        rng_bag: RngBag,
        stopwatch: Stopwatch,
        card_metadata: Dict,
    ) -> None:
        super().__init__()

        self.config = config

        if self.config.debug:
            logger._logger.setLevel(logging.DEBUG)
            os.environ["CUDA_LAUNCH_BLOCKING"] = "1"

        self.card_metadata = card_metadata

        self.dtype = eval(config.dtype)

        self.model = model

        self.training_data_loader = training_data_loader

        # Skip saving and loading the state of validation dataloader
        self.register_non_stateful("validation_data_loader", validation_data_loader)

        self.gang = gang

        self.rng_bag = rng_bag

        self.step_nr = 1

        self.current_run_steps = 0

        self.checkpoint_manager = checkpoint_manager

        tb_dir = config.tb_dir or config.output_dir.joinpath("tb")

        self.metric_recorders = [LogMetricRecorder(logger)]

        if gang.rank == 0:
            self.metric_recorders.append(TensorBoardRecorder(tb_dir))
            self.metric_recorders.append(
                LCMWandBRecorder(
                    name=config.wandb_run_name,
                    project=config.wandb_project or "uncategorized",
                    output_dir=config.output_dir / "wandb",
                    config=self._tb_flat_config,
                )
            )

        self.profiler = Profiler(
            skip_first=config.profiler_skip_first,
            active=config.profiler_active,
            log_dir=tb_dir,
            gang=gang,
            enabled=config.profile,
        )

        self.stopwatch = stopwatch
        self._train_step_time = 0.0
        self._valid_step_time = 0.0

        self.criterion = None  # type: ignore

        self.loss_scaler = None  # type: ignore

    @property
    def is_fsdp(self) -> bool:
        return isinstance(self.model, FSDP)

    @property
    def is_ddp(self) -> bool:
        return isinstance(self.model, DDP)

    def setup(self) -> None:
        self.criterion = self.setup_criterion()

        self.setup_metric_bags()

        # Add the grad_norm metric to the training metric bag
        self.train_metric_bag.register_metric(
            "grad_norm", Mean(device=self.gang.device), persistent=False
        )
        self.train_metric_bag.register_metric(
            "raw_grad_norm", Mean(device=self.gang.device), persistent=False
        )

        self.setup_optimizer_and_lr_schedule()

    def setup_optimizer_and_lr_schedule(self):
        optimizer = AdamW(
            self.model.parameters(),
            lr=self.config.lr,
            betas=tuple(self.config.adam_betas),  # type: ignore
            eps=self.config.adam_eps,
            use_fp32=self.config.use_optimizer_in_fp32,
            weight_decay=self.config.weight_decay,
        )
        logger.info(
            (
                f"Setting up AdamW optimizer with betas={self.config.adam_betas}, "
                f"base lr={self.config.lr} and weight decay={self.config.weight_decay} "
                f"and use_fp32={self.config.use_optimizer_in_fp32}"
            )
        )

        self.register_stateful("optimizer", optimizer)

        self.loss_scaler = DynamicLossScaler(
            optimizer,
            gang=self.gang,
            init_scale=self.config.loss_scaler_init_scale,
            min_scale=0.0001,
            scale_window=self.config.loss_scaler_scale_window,
            enabled=self.dtype == torch.float16,
        )

        if self.loss_scaler.is_enabled:
            logger.info(
                f"Initializing DynamicLossScaler with init_scale={self.config.loss_scaler_init_scale}"
            )

        lr_scheduler = build_lr_scheduler(
            optimizer=self.optimizer,
            schedule=self.config.lr_schedule,
            lr=self.config.lr,
            warmup_steps=self.config.num_lr_warmup_steps,
            start_lr=self.config.start_lr,
            final_lr=self.config.final_lr,
            max_steps=self.config.max_steps,
            stage_ratio=tuple(self.config.lr_stage_ratios),
        )

        # Saving the lr_scheduler as well to properly resume training
        self.register_stateful("lr_scheduler", lr_scheduler)

    @abstractmethod
    def setup_criterion(self):
        """Define a criterion (loss / objective function to optimize)"""

    def setup_metric_bags(self):
        """Setup metric bags for tracking"""

        self.train_metric_bag = MetricBag(self.gang)

        self.register_non_stateful(
            "valid_metric_bag",
            {
                ds_name(dataset): MetricBag(self.gang)
                for dataset in self.config.validation_data
            },
        )

    def checkpoint_and_raise(self, exc) -> None:
        # Checkpoint before exiting
        if torch.cuda.is_available():
            torch.cuda.synchronize()
        logger.warning(f"R{self.gang.rank} checkpoint_and_raise - error={exc}")
        if self.current_run_steps > 100:
            # avoid checkpoining for early failures
            self._checkpoint(crash=exc)
        raise exc

    @cached_property
    def _tb_flat_config(self):
        """
        Prepare the flat config that will be used as HParams
        to record training metadata, namely config and environment hashes.
        """

        dict_config = flatten_dict(asdict(self.config))

        # Merge the data lists:
        def get_data_signature(dataset):
            return ":".join(
                map(str, (dataset["name"], dataset["weight"], dataset["filters"]))
            )

        dict_config["training_data"] = "+".join(
            get_data_signature(dataset) for dataset in dict_config["training_data"]
        )
        dict_config["validation_data"] = "+".join(
            get_data_signature(dataset) for dataset in dict_config["validation_data"]
        )

        # value should be one of int, float, str, bool, or torch.Tensor
        allowed_types = (int, float, str, bool, torch.Tensor)
        config_keys = list(dict_config)
        for k in config_keys:
            if not isinstance(dict_config[k], allowed_types):
                del dict_config[k]

        return dict_config

    def run(self) -> None:
        """Run the trainer for up to `max_steps`"""

        logger.info(f"Running training on {self.gang.size} device(s).")

        data_iter = self.training_data_loader.iterate_batches()

        logger.info(
            f"R{self.gang.rank} - waiting for all ranks to prepare a data iterator!"
        )
        self.gang.barrier()

        # These counters are rank-specific
        ooms, nans_or_infs = 0, 0

        # TODO: validate before training
        # logger.info(f"Starting with validation at step={self.step_nr}")
        # self._validate()

        with self.profiler:
            while self.step_nr <= self.config.max_steps:
                with record_function(f"step_{self.step_nr}"):
                    try:
                        # Main training step: forward -> backward -> optimizer.step -> log
                        stepped = self._train_step(data_iter)

                    except RuntimeError as e:
                        if "out of memory" in str(e):
                            self._log_oom(e)
                            ooms += 1
                            if self.config.raise_oom or ooms > self.config.max_ooms:
                                # Previous behaviour, no retries but still checkpointing
                                self.checkpoint_and_raise(e)

                            logger.warning(
                                f"Attempting to recover from OOM on R{self.gang.rank} (OOMS={ooms})"
                            )
                            stepped = True
                            # reset optimizer
                            self.optimizer.zero_grad(set_to_none=True)

                            # rollback updates
                            self.train_metric_bag.rollback_updates()

                            # Empty CUDA cache before trying again
                            if torch.cuda.is_available():
                                torch.cuda.empty_cache()

                        else:
                            # Other RuntimeErrors
                            self.checkpoint_and_raise(e)

                    except FloatingPointError as e:
                        if "Losses are Nan/Inf" in str(e):
                            self._log_nan_loss(e)
                            nans_or_infs += 1
                            if (
                                self.config.raise_nan_or_inf
                                or nans_or_infs > self.config.max_nans_or_infs
                            ):
                                self.checkpoint_and_raise(e)

                            logger.warning(
                                f"Attempting to recover from NaN/Inf loss on R{self.gang.rank} (NaNs/Infs={nans_or_infs})"
                            )
                            stepped = True
                            # reset optimizer
                            self.optimizer.zero_grad(set_to_none=True)

                            # rollback updates
                            self.train_metric_bag.rollback_updates()

                        else:
                            # Other FloatingPointErrors
                            self.checkpoint_and_raise(e)

                    except Exception as e:
                        self.checkpoint_and_raise(e)

                if stepped:
                    if self._should_publish_train_metrics():
                        self._publish_train_metrics()

                    if self._should_checkpoint():
                        self._checkpoint()

                    if self._should_validate():
                        self._validate()

                    if self._should_collect_garbage():
                        self._collect_garbage()

                    self.profiler.step()

                    self.step_nr += 1
                    self.current_run_steps += 1

                else:
                    logger.info(f"R{self.gang.rank} - Resetting the datapipeline")
                    self.training_data_loader.pipeline.reset()

                    logger.info(f"R{self.gang.rank} - Done resetting the datapipeline")
                    data_iter = self.training_data_loader.iterate_batches()

        self._save_model_card_for_last_checkpoint(to_checkpoint_dir=False)
        logger.info(f"Finished training after {self.step_nr - 1} step(s).")

        self.gang.close()

    def restore(self) -> None:
        logger.info("Attempting to load last checkpoint.")

        step_nr, checkpoint = self.checkpoint_manager.load_last_checkpoint()

        logger.info(f"Checkpoint loaded, restoring training from step {step_nr}.")

        self.load_state_dict(checkpoint)

        self.gang.barrier()

        logger.info("Training restored, resuming.")

        self.step_nr = step_nr + 1

    def _maybe_with_autocast(self) -> ContextManager[None]:
        # autocast is only needed if training with mixed precision.
        # If training fails without it, check if some module with its weights
        # is not properly cast
        if self.config.use_autocast:
            return torch.autocast(device_type="cuda", dtype=self.dtype)
        else:
            return nullcontext()

    def _train_step(self, data_iter: Iterator) -> bool:
        step_nr = self.step_nr

        step_stopwatch = Stopwatch(start=True, device=self.gang.device)

        stepped = False

        # We have to retry the step in case of a gradient overflow.
        while not stepped:
            batches = []

            # Collect batches.
            with record_function(f"step_{step_nr}_data_load"):
                for _ in range(self.config.gradient_accumulation):
                    try:
                        batches.append(next(data_iter))
                    except StopIteration:
                        break

            if len(batches) != self.config.gradient_accumulation:
                logger.info(
                    f"R{self.gang.rank} -End of data reached at training step {step_nr}."
                )

                return False

            # create a copy of the current metrics
            # any update to the metrics from this point will either be committed with `commit_updates`
            # or ignored with `rollback_updates`
            self.train_metric_bag.begin_updates()

            num_targets = 0

            # Accumulate gradients.
            for batch_nr, batch in enumerate(batches):
                with self._maybe_no_sync(batch_nr, len(batches)):
                    with record_function(f"step_{step_nr}_{batch_nr}_forward"):
                        # autocast should wrap only the forward pass(es)
                        # of your network, including the loss computation(s).
                        # Backward passes under autocast are not recommended.
                        with self._maybe_with_autocast():
                            loss = self.criterion(batch)

                    if not (
                        torch.isfinite(loss.value).all() or self.loss_scaler.is_enabled
                    ):
                        raise FloatingPointError("Losses are Nan/Inf.")

                    # update metrics
                    self.train_metric_bag.update([loss])

                    with record_function(f"step_{step_nr}_{batch_nr}_backward"):
                        self.loss_scaler.backward(loss.value)

                    num_targets += loss.num_target_elements

            # Record and clip gradient norm
            grad_norm, raw_grad_norm = self.process_gradients(step_nr, num_targets)

            # Update parameters.
            with record_function(f"step_{step_nr}_optimizer"):
                # scale_result: LossScaleResult(old_scale: float, new_scale: float, overflow: bool, min_reached: bool)
                _, scale_result = self.loss_scaler.run_optimizer_step(step_nr)

            if scale_result.overflow:
                # Walk back the metrics update:
                self.train_metric_bag.rollback_updates()
                logger.debug(
                    f"R{self.gang.rank} rolled back update {self.train_metric_bag._original_metrics is None}"
                )

                if scale_result.min_reached:
                    logger.error(f"Loss has started exploding at step {step_nr}. Stopping training.")  # fmt: skip

                    raise FloatingPointError("The training loss has exploded.")

                logger.debug(f"Repeating training step {step_nr}.")

            else:
                self.lr_scheduler.step()

                stepped = True

            # Reset.
            self.optimizer.zero_grad(set_to_none=True)

        # Stepped = True:
        with record_function(f"step_{step_nr}_metrics"):
            # do something with losses and grad_norm

            self.train_metric_bag.commit_updates()

            # gradient norm is common to workers
            self.train_metric_bag.grad_norm.update(grad_norm)
            self.train_metric_bag.raw_grad_norm.update(raw_grad_norm)

            if self.gang.rank == 0:
                # update elapsed time once
                self._train_step_time += step_stopwatch.get_elapsed_time()

        del batches
        return stepped

    def _maybe_no_sync(self, batch_nr: int, num_batches: int) -> ContextManager[None]:
        if batch_nr < num_batches - 1 and self.gang.size > 1:
            return self.model.no_sync()
        return nullcontext()

    def normalize_gradients(self, num_targets: int) -> None:
        """
        :param num_target:
            The number of targets used in loss computation in this process.

        If reduction = sum:
            similar to fairseq2's `normalize_gradients`, will normalize the gradients of the model by ``world_size/num_targets``.
        If reduction = mean:
            will simply multiply by world size i.e undo DDP/FSDP's default normalization
        """
        reduction = self.criterion.reduction
        if reduction == "sum":
            total_num_targets = torch.tensor(
                num_targets, device=self.gang.device, dtype=torch.int64
            )

            self.gang.all_reduce(total_num_targets, ReduceOperation.SUM)

            # Both DDP and FSDP divide gradients by the world size which we also undo.
            if total_num_targets > 0:
                grad_scale = self.gang.size / total_num_targets
            else:
                # If total_num_targets == 0, gradients will be zeroes anyway
                grad_scale = self.gang.size

        else:
            grad_scale = self.gang.size

        scale_gradients(self.model, grad_scale)

    def process_gradients(
        self, step_nr: int, num_targets: int
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        with record_function(f"step_{self.step_nr}_process_grads"):
            # Normalize gradients
            """
            Normalize and clip the gradients
            """
            # this raw grad norm is only used for debugging
            raw_grad_norm = clip_gradient_norm(
                self.model,
                max_norm=None,
            )

            if not self.config.turn_off_grad_normalization:
                self.normalize_gradients(num_targets=num_targets)

            # undo the GradScaler's scaling before clipping
            self.loss_scaler.unscale_gradients_()

            # Clip gradients
            # If DDP, we use torch.nn.utils.clip_grad_norm_, if FSDP,
            # we use torch.distributed.fsdp.FullyShardedDataParallel.clip_grad_norm_
            # this method handles the fact that gradients might be sharded across ranks.
            grad_norm = clip_gradient_norm(
                self.model,
                max_norm=self.config.max_grad_norm,
            )

            # Check for gradient consistency across workers:
            if not check_gradient_norms(grad_norm, self.gang, step_nr):
                raise FloatingPointError(
                    f"The gradients are inconsistent between processes at step {step_nr}. Training cannot continue."
                )

        return grad_norm, raw_grad_norm

    def _should_validate(self) -> bool:
        return self._should_do(self.config.validate_every_n_steps)

    def _should_collect_garbage(self) -> bool:
        return self._should_do(self.config.gc_every_n_steps)

    def _collect_garbage(self):
        logger.info("Collecting garbage...")
        gc.collect()

    @torch.inference_mode()
    def _validate(self) -> None:
        gc.collect()
        torch.cuda.empty_cache()

        if self.validation_data_loader is None:
            logger.info("Skip validation as the data loader is empty")
            return

        self.model.eval()

        logger.info(f"Starting validation after step {self.step_nr}.")

        self.validation_data_loader.pipeline.reset()

        data_iter = self.validation_data_loader.iterate_batches()
        data_dummy_iter = self.validation_data_loader.iterate_dummy_batches()

        logger.info(f"R{self.gang.rank} done creating the validation data iterator")

        for step_nr in count(start=1):
            step_stopwatch = Stopwatch(start=True, device=self.gang.device)

            try:
                batch = next(data_iter)
                true_batch = 1
            except StopIteration:
                batch = next(data_dummy_iter)
                true_batch = 0

            total_nb_batches = all_sum(self.gang, true_batch)

            if bool(total_nb_batches == 0):
                break
            # we apply model for all workers to avoid process groups sync issues
            loss = self.criterion(batch)

            if true_batch:
                self._valid_step_time += step_stopwatch.get_elapsed_time()
                self.valid_metric_bag[batch.name].update([loss])

        self._publish_validation_metrics()

        logger.info(
            f"R{self.gang.rank} Validation complete in {step_nr} steps, resuming training."
        )

        self.model.train()

    def _should_publish_train_metrics(self) -> bool:
        return self._should_do(self.config.publish_metrics_every_n_steps)

    def _set_elements_per_second(
        self, metric_values: Dict[str, Any], elapsed_time: float
    ) -> None:
        try:
            num_elements = metric_values[self.criterion.throughput_metric_name]
        except KeyError:
            return

        if not isinstance(num_elements, (int, float, torch.Tensor)):
            return

        if elapsed_time == 0.0:
            metric_values["elements_per_second"] = 0.0
        else:
            metric_values["elements_per_second"] = num_elements / elapsed_time

    def _publish_train_metrics(self) -> None:
        values = self.train_metric_bag.sync_and_compute_metrics()

        self.train_metric_bag.reset_non_persistent_metrics()

        # Only rank-0 to record and publish
        # since sync_and_compute_metrics's recipient rank is 0
        if self.gang.rank != 0:
            return

        assert values is not None

        values["lr"] = get_effective_lr(self.lr_scheduler)

        self._set_elements_per_second(values, self._train_step_time)

        if self.loss_scaler.is_enabled:
            values["grad_scale"] = self.loss_scaler.get_scale()

        values["wall_time"] = self.stopwatch.get_elapsed_time()
        values["elapsed_time"] = self._train_step_time

        record_metrics(self.metric_recorders, "Train", values, self.step_nr)

        self._train_step_time = 0.0

    def _publish_validation_metrics(self) -> None:
        values = {}
        for name, metric_bag in self.valid_metric_bag.items():
            values[name] = metric_bag.sync_and_compute_metrics()
            metric_bag.reset_non_persistent_metrics()

        # Only rank-0 to record and publish
        if self.gang.rank != 0:
            return

        for name, val in values.items():
            assert val is not None
            self._set_elements_per_second(val, self._valid_step_time)
            val["elapsed_time"] = self._valid_step_time
            val["wall_time"] = self.stopwatch.get_elapsed_time()
            valid_name = f"Valid | {name}"
            record_metrics(self.metric_recorders, valid_name, val, self.step_nr)

        # reset timers
        self._valid_step_time = 0.0

    def _should_checkpoint(self) -> bool:
        return self._should_do(self.config.checkpoint_every_n_steps)

    def _should_save_consolidated_model(self) -> bool:
        return self.is_fsdp and self._should_do(self.config.save_model_every_n_steps)

    def _checkpoint(self, crash=None) -> None:
        logger.info(f"Saving checkpoint at step {self.step_nr}")
        checkpoint = self.state_dict()

        metadata = {
            "config": self.config,
            "crash": crash,
        }

        self.checkpoint_manager.begin_checkpoint(self.step_nr)

        if self.is_fsdp:
            replicated_keys = None
        elif self.is_ddp:
            # If we do not shard, save the model and the optimizer only on rank 0.
            replicated_keys = {"model", "optimizer"}
        else:
            replicated_keys = {"*"}

        self.checkpoint_manager.save_state(checkpoint, replicated_keys=replicated_keys)

        self.checkpoint_manager.save_metadata(metadata)

        if self._should_save_consolidated_model():
            self._save_consolidated_model()

        # Create a model card only after creating model.pt
        # i.e., regular checkpointing with DDP or after consolidation with FSDP
        if not self.is_fsdp:
            self._save_model_card_for_last_checkpoint(to_checkpoint_dir=True)

        self.checkpoint_manager.commit_checkpoint()

        # Note that this logic looks at saved directories regardless of
        # the nature of the checkpointing, consolidated or not
        if self.config.keep_last_n_checkpoints != -1:
            self.checkpoint_manager.keep_last_n_checkpoints(
                self.config.keep_last_n_checkpoints,
                preserve_model=self.config.preserve_consolidated_models,
            )

        logger.info(f"Checkpoint saved by worker @rank={self.gang.rank}")

    def _save_consolidated_model(self) -> None:
        logger.info(f"Saving consolidated model at step {self.step_nr}.")
        self.checkpoint_manager.save_consolidated_fsdp_model(self.model)
        self._save_model_card_for_last_checkpoint(to_checkpoint_dir=True)
        logger.info("Consolidated model saved.")

    def _should_do(self, n_step: int) -> bool:
        return self.step_nr % n_step == 0

    def create_model_card_for_last_checkpoint(
        self, is_final: bool = False, **card_kwargs
    ) -> Optional[AssetCard]:
        """Create a model card based on the last saved checkpoint and the model config."""
        logger.warning(
            "Could not create a model card with a generic trainer.  Please use a model-specific one."
        )
        return None

    def _save_model_card_for_last_checkpoint(
        self, to_checkpoint_dir: bool = False
    ) -> None:
        """Save the model card for the last checkpoint to the checkpoint directory or the core output directory."""
        if self.gang.rank != 0:
            return

        if to_checkpoint_dir:
            current_step_nr = self.checkpoint_manager._checkpoint_step_nr
            output_dir = self.checkpoint_manager._checkpoint_dir.joinpath(
                f"step_{current_step_nr}.tmp"
            )
        else:
            output_dir = self.config.output_dir

        card = self.create_model_card_for_last_checkpoint(
            is_final=not to_checkpoint_dir
        )

        if card is not None:
            card_data = card._metadata  # TODO: use the exposed attribute when available
            with open(output_dir / "model_card.yaml", "w", encoding="utf-8") as outfile:
                yaml.dump(card_data, outfile, default_flow_style=False)
            logger.info(f"Model card saved in {output_dir}")

    def _log_oom(self, exc):
        logger.warning(
            f"OOM: Ran out of memory on R{self.gang.rank} with exception: {exc}"
        )

        if torch.cuda.is_available():
            for device_idx in range(torch.cuda.device_count()):
                logger.warning(torch.cuda.memory_summary(device=device_idx))

        sys.stderr.flush()

    def _log_nan_loss(self, exc):
        logger.warning(f"We hit a Nan/Inf Loss: raised with exception: {exc}")


class TrainerBuilder:
    def __init__(self, config: TrainingConfig):
        assert config.save_model_every_n_steps % config.checkpoint_every_n_steps == 0, (
            f"save_model_every_n_steps={config.save_model_every_n_steps} for saving consolidated models should be a multiplier of checkpoint_every_n_steps={config.checkpoint_every_n_steps}"
        )

        self.config = config

        self.stopwatch = Stopwatch(start=True)

        # In case we train on Ampere or later, use TF32.
        torch.set_float32_matmul_precision("high")

        if self.config.fake_gang_device is None:
            # By default, we work with a process group
            self.gang = init_process_group(config, logger=logger._logger)
        else:
            # For testing purposes, we use a fake gang on the chosen device
            self.gang = FakeGang(device=torch.device(self.config.fake_gang_device))

        self.gang_rank = self.gang.rank if self.gang else 0

        if self.gang.device.type == "cuda":
            # Setup ATEN and NCCL logging if in debug mode
            self._setup_additional_logging()

            # Dump environment variables:
            log_env_variables(self.gang.device)

        # A variable to carry fields necessary to build concise model cards
        self.card_metdata: Dict = {}

        if self.gang_rank == 0:
            logger.info(f"Job Config\n{pformat(config)}")

        self.device = self.gang.device

        rng_bag = RngBag.from_device_defaults(self.device)

        # Ensure that each run has deterministic behavior.
        rng_bag.manual_seed(config.seed)

        self.rng_bag = rng_bag

        self.dtype = eval(config.dtype)

        self.finetune: bool = False

        self.has_checkpoint: bool = False

    @property
    @abstractmethod
    def model_loader(self):
        """A fairseq2 ModelLoader"""

    @property
    def model_config_loader(self):
        """A fairseq2 ConfigLoader"""
        return self.model_loader._config_loader

    @abstractmethod
    def load_data(self):
        """Load training and validation data
        Returns one loader for training data and one for validation data
        """

    def create_model_config(self, set_finetune_flag: bool = False):
        """
        Given `model_config_or_name`, `model_arch` and `model_arch_overrides`
        create the model config dict
        if `set_finetune_flag` is `True` then the trainer's finetune flag will be set
        here inferred from the use of `model_config_or_name`
        """
        if self.config.model_config_or_name is not None:
            assert self.config.model_arch is None, (
                "We cannot set both `model_config_or_name` and `model_arch`"
            )

            if isinstance(self.config.model_config_or_name, str):
                # The config of a registered model i.e. we're finetuning
                logger.info(
                    f"Loading pretrained model from {self.config.model_config_or_name}"
                )

                model_config = self.model_config_loader(
                    self.config.model_config_or_name
                )
                finetune = True

                # Metadata for card creation
                source_card = self.model_config_loader._asset_store.retrieve_card(
                    self.config.model_config_or_name
                )
                try:
                    arch = source_card.field("model_arch").as_(str)
                except AssetCardFieldNotFoundError:
                    arch = None

                self.card_metadata = {
                    "model_config": model_config if arch is None else None,
                    "model_type": model_config.model_type,
                    "model_arch": arch,
                }

            else:
                # model_config_or_name is a dataclass
                logger.info(
                    "Creating a model from the provided config in model_config_or_name"
                )
                model_config = self.config.model_config_or_name

                self.card_metadata = {
                    "model_config": model_config,
                    "model_type": model_config.model_type,
                    "model_arch": None,
                }

                finetune = False

        elif self.config.model_arch is not None:
            assert (
                self.config.model_arch in self.model_config_loader._arch_configs.names()
            ), (
                f"Could not recognise {self.config.model_arch} as a registered architecture "
            )

            logger.info(
                f"Creating a model from registered arch {self.config.model_arch}"
            )

            finetune = False
            model_config = self.model_config_loader._arch_configs.get(
                self.config.model_arch
            )
            self.card_metadata = {
                "model_config": None,
                "model_type": model_config.model_type,
                "model_arch": self.config.model_arch,
            }

        # In all setups we can override some config parameters
        if self.config.model_arch_overrides is not None:
            try:
                update_dataclass(model_config, self.config.model_arch_overrides)

            except (TypeError, ValueError) as ex:
                raise ValueError(
                    "The model_arch_overrides contain one or more invalid keys"
                ) from ex

            self.card_metadata["model_arch"] = None
            self.card_metadata["model_config"] = model_config

            logger.info(
                f"Overwriting model config parameters with {self.config.model_arch_overrides}"
            )

        if set_finetune_flag:
            self.finetune = finetune

        return model_config

    def create_model(self):
        """
        Load the model to be trained.
        In case other models are developed following a different paradigm, we can create
        corresponding trainers by overriding `create_model`
        """
        logger.info("Initializing model.")

        model_config = self.create_model_config(set_finetune_flag=True)

        if self.gang_rank == 0:
            logger.info(f"Final model config:\n{pformat(model_config)}")

        model = self.model_loader._factory(
            model_config,
            device=self.device,
            dtype=self.dtype,
        )
        # log model before any wrapping:
        log_model(model, logger)

        return model

    def wrap_model_with_ddp(self, model) -> DDP:
        """Wrap the model with DDP"""

        try:
            ddp_model = to_ddp(
                model,
                self.gang,
            )

        except ValueError:
            logger.warning(
                "Using pytorch DDP instead of fairseq's `to_ddp`\
                - please check fairseq2 after a3de79dcc6a4ea34cde644e15b4056f1a808a6a8"
            )

            ddp_model = DDP(model)

        if self.gang_rank == 0:
            log_model(ddp_model, logger)

        return ddp_model

    def wrap_model_with_fsdp(self, model) -> FSDP:
        """Wrap the model with FSDP."""

        wrap_policy, ignored_modules = get_fsdp_wrap_policy(
            model, wrap_granularity=self.config.fsdp_wrap_granularity
        )
        memory_policy = get_fsdp_memory_policy(policy=self.config.fsdp_memory_policy)

        if self.dtype == torch.float32:
            mixed_precision_dtype = None
        else:
            mixed_precision_dtype = self.dtype

        skip_init = False
        broadcast_state = self.finetune and not self.has_checkpoint
        fp32_reduce = self.config.fsdp_fp32_reduce

        if self.gang.rank == 0:
            logger.info(
                (
                    f"FSDP init with: \n--- ignored_modules={ignored_modules}"
                    f"\n--- wrap_policy={wrap_policy}"
                    f"\n--- mixed_precision_dtype={mixed_precision_dtype}"
                    f"\n--- skip_init={skip_init}"
                    f"\n--- broadcast_state (FSDP's sync_module_states)={broadcast_state}"
                    f"\n--- fp32_reduce={fp32_reduce}"
                    f"\n--- memory_policy={memory_policy}"
                )
            )

        fsdp_model = to_fsdp(
            model,
            self.gang,
            wrap_policy,
            mixed_precision_dtype=mixed_precision_dtype,
            ignored_modules=ignored_modules,
            fp32_reduce=fp32_reduce,
            skip_init=skip_init,
            broadcast_state=broadcast_state,
            memory_policy=memory_policy,
        )

        if self.gang_rank == 0:
            log_model(fsdp_model, logger)

        return fsdp_model

    def maybe_load_model(self, model):
        """
        If we are finetuning and we don't have a checkpoint,
        load the pre-trained model and broadcast it to
        all gang processes from rank 0.
        """
        if not self.has_checkpoint and self.finetune:
            logger.info(f"Loading for finetuning: {self.config.model_config_or_name}")

            if self.gang_rank == 0:
                pretrained_model = self.model_loader(
                    model_name_or_card=self.config.model_config_or_name,
                    device=self.gang.device,
                    dtype=self.dtype,
                )  # type: ignore[arg-type]

                try:
                    model.load_state_dict(
                        pretrained_model.state_dict(),
                        strict=True,
                        assign=False,
                    )
                except (KeyError, ValueError) as ex:
                    raise ValueError(
                        f"The model state form {self.config.model_config_or_name} "
                        "cannot be loaded. See nested exception for details."
                    ) from ex

            self.gang.barrier()

            to_device(model, self.gang.device)

            logger.info(
                f"Done loading model for finetuning: {self.config.model_config_or_name}"
            )

        return model

    def maybe_freeze_parameters(self, model):
        assert (self.config.freezing_strategy == "modules") == (
            self.config.freeze_modules is not None
        ), (
            "For the `modules` freezing_strategy, we need a list of `freeze_modules`. "
            "If `freeze_modules` is provided, make sure to use freezing_strategy=modules"
        )

        if self.config.freezing_strategy == "none":
            return model

        if self.config.freezing_strategy == "modules":
            # Optionally freeze the parameters of sub-modules:
            if self.config.freeze_modules is not None:
                for module in self.config.freeze_modules:
                    logger.info(f"... Freezing module={module}")
                    freeze_parameters(getattr(model, module))
            return model

        if self.config.freezing_strategy == "ffn":
            for name, m in _get_named_modules(model):
                if "ffn" in name:
                    logger.info(f"... Freezing module={name}")
                    freeze_parameters(m)
            return model

        if self.config.freezing_strategy == "adaln":
            for name, m in _get_named_modules(model):
                if "modulator" in name:
                    logger.info(f"... Freezing module={name}")
                    freeze_parameters(m)
            return model

        if self.config.freezing_strategy == "ffn-adaln":
            for name, m in _get_named_modules(model):
                if "modulator" in name or "ffn" in name:
                    logger.info(f"... Freezing module={name}")
                    freeze_parameters(m)
            return model

        raise ValueError(f"Unknown freezing stratgey {self.config.freezing_strategy}")

    def _setup_additional_logging(self):
        if self.config.debug:
            assert self.config.log_folder is not None, (
                "Missing log_folder, \
            make sure the log_folder is properly set in the training config"
            )
            setup_additional_logging(log_folder=self.config.log_folder)

    @property
    def use_fsdp(self) -> bool:
        return self.config.use_fsdp

    @property
    def use_ddp(self) -> bool:
        """
        Whether DDP should be used.
        if selg.gang.size == 1: single worker, no parallelism
        if use_fsdp:  use FSDP instead
        """
        return not (self.gang.size == 1 or self.use_fsdp)

    @abstractmethod
    def build_trainer(self):
        """Build the trainer by loading data and
        setting up the model for training

        Returns trainer
        """