Blame for src/utils/distributed.py

"""FSDP readers/writers and one scalar EMA; no implicit fabric overrides."""

import os
import tempfile
from datetime import timedelta
from functools import partial

import torch
import torch.distributed as dist
from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import MixedPrecision, ShardingStrategy
from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy

from src.utils.checkpoint_state import (
    normalize_generator_fsdp_key,
    normalize_generator_fsdp_state_dict,
)


def launch_distributed_job():
    """Direct Python uses a local process group; torchrun uses its own rendezvous."""
    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
    torch.cuda.set_device(local_rank)
    if dist.is_initialized():
        return
    kwargs = {
        "backend": "nccl",
        "device_id": torch.device("cuda", local_rank),
        "timeout": timedelta(minutes=30),
    }
    if "RANK" in os.environ:
        dist.init_process_group(**kwargs)
    else:
        descriptor, filename = tempfile.mkstemp(prefix="looking-ahead-pg-")
        os.close(descriptor)
        os.unlink(filename)
        dist.init_process_group(init_method=f"file://{filename}", rank=0, world_size=1, **kwargs)


def barrier():
    if dist.is_initialized():
        dist.barrier()


def is_main_process():
    return not dist.is_initialized() or dist.get_rank() == 0


def fsdp_state_dict(model):
    """All ranks collect; only global rank zero retains ordinary weight tensors."""
    checkpoint = get_model_state_dict(
        model, options=StateDictOptions(full_state_dict=True, cpu_offload=True)
    )
    return checkpoint if is_main_process() else {}


def fsdp_wrap(module, *, sharding_strategy, mixed_precision, sync_module_states=False):
    precision = (
        MixedPrecision(
            param_dtype=torch.bfloat16,
            reduce_dtype=torch.float32,
            buffer_dtype=torch.float32,
            cast_forward_inputs=False,
        )
        if mixed_precision
        else None
    )
    strategy = {
        "full": ShardingStrategy.FULL_SHARD,
        "hybrid_full": ShardingStrategy.HYBRID_SHARD,
        "hybrid_zero2": ShardingStrategy._HYBRID_SHARD_ZERO2,
        "no_shard": ShardingStrategy.NO_SHARD,
    }[sharding_strategy]
    return FSDP(
        module,
        auto_wrap_policy=partial(size_based_auto_wrap_policy, min_num_params=int(5e7)),
        sharding_strategy=strategy,
        mixed_precision=precision,
        device_id=torch.cuda.current_device(),
        limit_all_gathers=True,
        use_orig_params=True,
        sync_module_states=sync_module_states,
    )


def _fsdp_sharding_strategy(module: torch.nn.Module):
    strategy = getattr(module, "sharding_strategy", None)
    if strategy is not None:
        return strategy
    fsdp_state = getattr(module, "_fsdp_state", None)
    if fsdp_state is not None:
        return getattr(fsdp_state, "sharding_strategy", None)
    return None


def _effective_no_shard(module: torch.nn.Module) -> bool:
    """Return true when all FSDP wrappers use effective ``NO_SHARD``.

    A plain module is treated like ``NO_SHARD`` for EMA serialization helpers;
    this keeps the direct save/load path testable without requiring CUDA/FSDP.
    """
    fsdp_modules = [submodule for submodule in module.modules() if isinstance(submodule, FSDP)]
    if not fsdp_modules:
        return True
    return all(
        _fsdp_sharding_strategy(submodule) == ShardingStrategy.NO_SHARD
        for submodule in fsdp_modules
    )


class EMA_FSDP:
    """
    Exponential Moving Average for an FSDP-wrapped model.

    Design notes
    ------------
    * **Sharded, on-GPU shadow.** Each rank stores an EMA shadow of *only its
      own local shard* of every ``flat_param`` (FSDP1) / ``DTensor`` (FSDP2),
      keeping the buffer in fp32 on the same device as the live parameter.
      This means ``update`` is a pure local elementwise op -- no
      ``summon_full_params``, no all-gather, no host-device traffic, no
      cross-rank synchronization -- so it is essentially free compared to
      the optimizer step.
    * **Collective only on save.** ``full_state_dict`` is the only path that
      materializes full EMA tensors for checkpointing. Sharded FSDP uses
      ``summon_full_params``; effective ``NO_SHARD`` returns the already-full
      local EMA shadows directly. It should only run at checkpoint time, never
      per-step.
    * Non-floating-point buffers (e.g. ``num_batches_tracked``) are never
      averaged.
    """

    def __init__(
        self,
        fsdp_module: torch.nn.Module,
        decay: float = 0.999,
    ):
        self.decay = decay
        self.num_updates = 0

        # name -> fp32 tensor on the *same device & with the same local
        # shape* as the corresponding live (sharded) parameter. Each rank
        # only holds its own shard.
        self.shadow: dict = {}
        self._is_main = is_main_process()

        self._init_shadow(fsdp_module)

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------
    def _iter_named_tensors(self, module):
        yield from module.named_parameters()

    # ------------------------------------------------------------------
    # Initialization
    # ------------------------------------------------------------------
    @torch.no_grad()
    def _init_shadow(self, fsdp_module):
        # No collective: each rank just clones its own shard in fp32. For
        # FSDP1 the local shard is a 1-D ``flat_param`` slice; for FSDP2 it
        # is a ``DTensor`` whose ``_local_tensor`` we read transparently
        # (``detach().clone()`` on a DTensor preserves the DTensor wrapping,
        # which is exactly what we want for arithmetic later).
        for n, t in self._iter_named_tensors(fsdp_module):
            self.shadow[n] = t.detach().to(dtype=torch.float32).clone()

    # ------------------------------------------------------------------
    # Update -- per-step hot path
    # ------------------------------------------------------------------
    @torch.no_grad()
    def update(self, fsdp_module):
        """Per-rank local EMA update. **No** collective communication.

        Cost is one fp32 mul_+add_ over the local shard, which is dominated
        by memory bandwidth and negligible next to the optimizer step.
        """
        self.num_updates += 1
        d = self.decay
        one_minus_d = 1.0 - d

        for n, t in self._iter_named_tensors(fsdp_module):
            shadow = self.shadow.get(n)
            if shadow is None:
                # New / unseen parameter -- initialize on the fly.
                self.shadow[n] = t.detach().to(dtype=torch.float32).clone()
                continue
            # Cast live param to fp32 *on device* (cheap if already fp32,
            # otherwise a fused kernel). No host transfer, no collective.
            src = t.detach()
            if src.dtype != torch.float32:
                src = src.to(dtype=torch.float32)
            shadow.mul_(d).add_(src, alpha=one_minus_d)

    # ------------------------------------------------------------------
    # (De)serialization helpers
    # ------------------------------------------------------------------
    def state_dict(self):
        """Return the *local-shard* shadow on every rank.

        Note: this is **not** a full state dict. Use :meth:`full_state_dict`
        if you want the rank-0-only fully-gathered version (e.g. for
        checkpointing). Each rank's tensors live on its own device.
        """
        return self.shadow

    @torch.no_grad()
    def load_full_state_dict(
        self, full_sd, fsdp_module, *, allow_missing_aux=False, expected_shapes=None
    ):
        """Restore ordinary EMA from each rank's validated full source state.

        All-rank writeback avoids PyTorch's unsupported rank0-only writeback.
        The temporary swap restores the original independent live shards.
        """
        full_sd = normalize_generator_fsdp_state_dict(full_sd)
        expected = {
            normalize_generator_fsdp_key(name): tensor
            for name, tensor in self._iter_named_tensors(fsdp_module)
        }
        missing = set(expected) - set(full_sd)
        unexpected = set(full_sd) - set(expected)
        if unexpected or (missing and not allow_missing_aux):
            raise RuntimeError("EMA checkpoint has missing or unexpected parameter keys.")
        if allow_missing_aux:
            from src.utils.predictive_training_checkpoint import is_predictive_aux_key

            if any(not is_predictive_aux_key(k) for k in missing):
                raise RuntimeError("EMA checkpoint is missing backbone parameters")
        if expected_shapes is not None:
            for name, value in full_sd.items():
                if tuple(value.shape) != expected_shapes[name]:
                    raise RuntimeError(f"EMA checkpoint shape mismatch: {name}")
        if _effective_no_shard(fsdp_module):
            for name, tensor in expected.items():
                if name in full_sd and full_sd[name].shape != tensor.shape:
                    raise RuntimeError(f"EMA checkpoint shape mismatch: {name}")
            new_shadow = {}
            for n, t in self._iter_named_tensors(fsdp_module):
                canonical_name = normalize_generator_fsdp_key(n)
                src = full_sd.get(canonical_name, t)
                new_shadow[n] = (
                    src.detach()
                    .to(
                        device=t.device,
                        dtype=torch.float32,
                    )
                    .clone()
                )
            self.shadow = new_shadow
            return

        # Backup live shards (each rank saves its own).
        backup = {n: t.data.clone() for n, t in self._iter_named_tensors(fsdp_module)}

        try:
            with FSDP.summon_full_params(
                fsdp_module, writeback=True, rank0_only=False, offload_to_cpu=False
            ):
                for name, tensor in self._iter_named_tensors(fsdp_module):
                    canonical_name = normalize_generator_fsdp_key(name)
                    if canonical_name in full_sd:
                        tensor.copy_(
                            full_sd[canonical_name].to(dtype=tensor.dtype, device=tensor.device)
                        )
            new_shadow = {
                name: tensor.detach().to(dtype=torch.float32).clone()
                for name, tensor in self._iter_named_tensors(fsdp_module)
            }
            self.shadow = new_shadow
        finally:
            for name, tensor in self._iter_named_tensors(fsdp_module):
                tensor.data.copy_(backup[name])

    @torch.no_grad()
    def full_state_dict(self, fsdp_module):
        """Return the *full* EMA state dict on rank 0 (empty dict elsewhere).

        This should only be used at checkpoint time. Effective ``NO_SHARD``
        returns the already-full EMA shadows directly. Sharded FSDP temporarily
        swaps the live FSDP parameters for the EMA shadows, summons them in
        full on rank 0, copies them out, then restores the live weights -- all
        without persisting any change.
        """
        if _effective_no_shard(fsdp_module):
            if not self._is_main:
                return {}
            return {
                normalize_generator_fsdp_key(n): shadow.detach()
                .to("cpu", dtype=torch.float32)
                .clone()
                for n, shadow in self.shadow.items()
            }

        # Clone the live local shards before copying EMA into their storage.
        # ``use_orig_params=True`` exposes views into FSDP's FlatParameter, so
        # retaining ``t.data`` here would only retain an alias. FSDP may reuse
        # and overwrite that storage while summoning full parameters, turning
        # an apparent restore into an accidental live <- EMA switch.
        backup = {}
        for n, t in self._iter_named_tensors(fsdp_module):
            shadow = self.shadow.get(n)
            if shadow is None:
                continue
            backup[n] = t.detach().clone()
            t.copy_(shadow.to(device=t.device, dtype=t.dtype), non_blocking=True)

        # CPU FSDP parameters are already host-resident. Asking PyTorch to
        # offload a rank0-only summon back to the same device can hang under
        # Gloo, while CUDA training still benefits from the original offload.
        offload_to_cpu = any(t.device.type != "cpu" for t in backup.values())

        try:
            with FSDP.summon_full_params(
                fsdp_module,
                writeback=False,
                rank0_only=True,
                offload_to_cpu=offload_to_cpu,
            ):
                if not self._is_main:
                    out = {}
                else:
                    out = {}
                    for n, t in self._iter_named_tensors(fsdp_module):
                        out[normalize_generator_fsdp_key(n)] = (
                            t.detach().to(device="cpu", dtype=torch.float32).clone()
                        )
        finally:
            # Restore the independent live-shard copies regardless of what
            # happened above. Copying into the post-context parameter views
            # preserves FSDP's storage/view invariants.
            for n, t in self._iter_named_tensors(fsdp_module):
                if n in backup:
                    t.copy_(backup[n], non_blocking=True)

        return out