moge/train/debug.py

"""Debug dumping for training runs that hit NaNs or unusually large gradients."""
from pathlib import Path
from typing import *

import torch

from .utils import detach_to_cpu


class DebugDumper:
    """Collects dump triggers during a step and writes the offending state to disk.

    Call `add_reason(tag)` from anywhere in the step to flag it. Tags starting with
    `nan_` are fatal and count toward the abort threshold; anything else (e.g.
    `large_grad_norm_12.34`) is informational and capped so a pathological run
    cannot fill the disk. Call `begin_step()` once per step and `flush(...)` once per
    accumulation step.

    Attributing a non-finite gradient
    ---------------------------------
    `clip_grad_norm_` only runs on the sync micro-step, and by then `.grad` holds the
    sum over every micro-batch of the step *and* the DDP all-reduce across every rank.
    A non-finite norm there therefore says nothing about which sample caused it: the
    batch in scope is one of `num_processes * gradient_accumulation_steps` candidates,
    and because the all-reduce spreads the NaN, every rank reports it and dumps its own
    innocent micro-batch.

    `check_grads` closes that gap. It runs after every micro-step's backward and flags
    the transition from all-finite to non-finite, so the dump carries the micro-batch
    that actually introduced the corruption. Before the sync micro-step nothing has
    been all-reduced yet, so that attribution is rank-local and exact; on the sync
    micro-step the gradient is already pooled and the tag says so.
    """

    def __init__(
        self,
        workspace: Path,
        accelerator,
        model,
        dump_grad_norm_above: Optional[float] = None,
        max_nan_dumps_before_abort: int = 10,
        max_extra_dumps: int = 25,
        save_model_on_first_dump: bool = False,
    ):
        self.workspace = workspace
        self.accelerator = accelerator
        # Unwrapped, so parameter names match the checkpoint and the optimizer
        # assignment log rather than carrying DDP's `module.` prefix. DDP shares the
        # underlying parameter objects, so the gradients are the same tensors.
        self.model = accelerator.unwrap_model(model)
        self.dump_grad_norm_above = dump_grad_norm_above
        self.max_nan_dumps_before_abort = max_nan_dumps_before_abort
        self.max_extra_dumps = max_extra_dumps
        self.save_model_on_first_dump = save_model_on_first_dump
        self.reasons: List[str] = []
        self.nan_encountered_times = 0
        self.extra_dump_count = 0
        self.model_saved = False
        # Gradient attribution state, reset per step by `begin_step`.
        self.grads_flagged = False
        self.grad_nonfinite: List[Dict[str, Any]] = []

    def add_reason(self, tag: str):
        self.reasons.append(tag)

    def begin_step(self):
        """Reset the per-step gradient attribution state.

        Gradients accumulate across a step's micro-steps and are only zeroed on the
        sync one, so "have they already gone bad?" is a per-step question.
        """
        self.grads_flagged = False
        self.grad_nonfinite = []

    def note_grad_norm(self, grad_norm, grad_norm_is_finite: bool):
        """Flag an unusually large but finite gradient norm, if a threshold was configured.

        The threshold and quota are checked before reading the value, so the
        `.cpu()` sync only happens when the feature is actually armed.
        """
        if (
            self.dump_grad_norm_above is not None
            and grad_norm_is_finite
            and self.extra_dump_count < self.max_extra_dumps
        ):
            grad_norm_value = float(grad_norm.detach().cpu().item())
            if grad_norm_value > self.dump_grad_norm_above:
                self.add_reason(f'large_grad_norm_{grad_norm_value:.2f}')

    def check_grads(self, i_accumulate: int, synced: bool) -> bool:
        """Screen the accumulated gradients after one micro-step's backward.

        Returns whether they still look finite. Only the micro-step that first trips
        the screen is flagged; gradients accumulate, so every later micro-step in the
        same step inherits the corruption and would otherwise each claim it.
        """
        named_grads = [(name, p.grad) for name, p in self.model.named_parameters() if p.grad is not None]
        if not named_grads:
            return True

        # Fused screen: one multi-tensor kernel rather than one per parameter, which is
        # ~16x cheaper on a 370M-parameter model (1.6 ms vs 25 ms). This is the same
        # quantity `clip_grad_norm_` computes, so it errs in the useful direction: it
        # never misses a non-finite gradient, and it also trips on a finite-but-huge one
        # whose squares overflow -- exactly what would give a non-finite clipped norm.
        finite = bool(torch.stack(torch._foreach_norm([g for _, g in named_grads])).isfinite().all().item())
        if finite or self.grads_flagged:
            return finite

        self.grads_flagged = True
        # Rare path only: name the offending parameters and separate a NaN (an invalid
        # op such as 0/0) from an Inf (an overflow). Enough to localise the module
        # without dumping any weights.
        self.grad_nonfinite = [
            {
                'name': name,
                'num_nan': int(torch.isnan(grad).sum()),
                'num_inf': int(torch.isinf(grad).sum()),
                'numel': grad.numel(),
            }
            for name, grad in named_grads if not torch.isfinite(grad).all()
        ]
        where = 'after_sync' if synced else f'local_accum{i_accumulate}'
        if self.grad_nonfinite:
            self.add_reason(f'nan_grad_{where}')
        else:
            # Every individual gradient is finite yet their norm is not, so the sum of
            # squares overflowed rather than a NaN propagating: a different failure.
            self.add_reason(f'nan_gradnorm_overflow_{where}')
        return False

    def _save_model_once(self):
        """Snapshot the weights the first time we dump, so a step > 0 can be replayed.

        Step 0 needs no snapshot (the weights are still the initial checkpoint), and a
        per-event snapshot would be ~1.5 GB times every rank and every event, so this
        is one file for the whole run and only on the main process.
        """
        if self.model_saved or not self.accelerator.is_main_process:
            return
        self.model_saved = True
        path = Path(self.workspace, 'debug', 'model_at_first_dump.pt')
        path.parent.mkdir(parents=True, exist_ok=True)
        torch.save({'model': self.model.state_dict()}, path)

    def _dump(self, step: int, accumulate_step: int, batch: Any, output: Any, meta: Optional[Dict[str, Any]]) -> Path:
        """Dump the flagged micro-batch so the failing forward pass can be replayed offline."""
        dump_path = Path(
            self.workspace,
            'debug',
            f'step_{step:08d}_accum_{accumulate_step}_proc_{self.accelerator.process_index}'
            f'_reasons_{self.reasons[0]}.pkl',
        )
        dump_path.parent.mkdir(parents=True, exist_ok=True)
        with dump_path.open('wb') as f:
            torch.save({
                'batch': detach_to_cpu(batch),
                'output': detach_to_cpu(output),
                'reasons': self.reasons,
                'grad_nonfinite': self.grad_nonfinite,
                'meta': {
                    'step': step,
                    'accumulate_step': accumulate_step,
                    'process_index': self.accelerator.process_index,
                    'num_processes': self.accelerator.num_processes,
                    **(meta or {}),
                },
            }, f)
        return dump_path

    def flush(self, i_step: int, i_accumulate: int, batch: Any, output: Any, meta: Optional[Dict[str, Any]] = None):
        """Write a dump if anything flagged this step, and abort after too many NaNs."""
        if not self.reasons:
            return
        if any(r.startswith('nan_') for r in self.reasons):
            self.nan_encountered_times += 1
        else:
            self.extra_dump_count += 1
        # A gradient that only went bad once it was all-reduced is shared by every
        # rank, so every rank would write an identical dump of its own unrelated
        # micro-batch. Keep a single witness instead of `num_processes` of them.
        bystander = all(r.endswith('after_sync') for r in self.reasons)
        if not bystander or self.accelerator.is_main_process:
            if self.save_model_on_first_dump:
                self._save_model_once()
            self._dump(i_step, i_accumulate, batch, output, meta)
        self.reasons = []
        if self.nan_encountered_times >= self.max_nan_dumps_before_abort:
            raise RuntimeError('NaN encountered too many times, abort training.')