spikes/008-streaming-diloco/tests/test_diloco_smoke.py

"""Spike 008 — DiLoCo outer-loop smoke.

Verifies the framework's DiLoCo wrapper integrates cleanly with
`torchft.local_sgd.DiLoCo`. Tests follow torchft's own test pattern
(`torchft/local_sgd_test.py::DiLoCoTest`) — single-process, mock Manager,
verify that the outer optimizer machinery actually fires, NOT that two
replicas converge in single-process (which they cannot due to the post-hook
sequencing — see below).

Cross-replica convergence test deferred to multi-process integration tests
once we have real torch.distributed in CI (post-replication phase).

Per `docs/adrs/ADR-003-diloco-impl.md`.
"""
from __future__ import annotations

import sys
from pathlib import Path
from unittest.mock import create_autospec

import pytest
import torch
import torch.nn as nn
import torch.optim as optim

HERE = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(HERE))

from composer_diloco import (  # noqa: E402
    _TORCHFT_AVAILABLE,
    DiLoCo,
    Manager,
    _DummyWork,
)


pytestmark = pytest.mark.skipif(
    not _TORCHFT_AVAILABLE,
    reason="torchft not installed (pip install torchft-nightly)",
)


class TinyMLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(4, 8), nn.ReLU(), nn.Linear(8, 2))

    def forward(self, x):
        return self.net(x)


def _make_passthrough_manager():
    """Manager whose allreduce is a no-op pass-through.

    Why no-op (and not real-averaging): in single-process, replica A's
    `inner_a.step()` post-hook runs prepare_sync + perform_sync to completion
    BEFORE replica B's `inner_b.step()` is called. By the time replica B
    arrives at allreduce, replica A's outer optimizer has already stepped
    using A's local pseudogradient. There is no way to inject a true
    cross-replica barrier in single-process without rewriting torchft's
    internals — and since we're using upstream code, we don't.

    This means single-process tests verify the *machinery* (sync fires,
    outer optimizer steps, Nesterov state populates), not cross-replica
    convergence. True cross-replica convergence is verified in production
    by NCCL.

    This is also exactly the pattern torchft uses in their own
    `torchft/local_sgd_test.py::DiLoCoTest` — they do not test convergence
    in single-process.
    """
    mgr = create_autospec(Manager)
    mgr._use_async_quorum = False
    mgr.errored.return_value = None
    mgr.should_commit.return_value = True
    mgr.current_step.return_value = 0

    def passthrough(tensor: torch.Tensor, should_quantize: bool = False):
        return _DummyWork(tensor)

    mgr.allreduce.side_effect = passthrough
    return mgr


# ---------------------------------------------------------------------
# Acceptance test 1 — outer loop machinery fires on single replica
# ---------------------------------------------------------------------

def test_diloco_single_replica_machinery_fires():
    """Acceptance: 1 replica × 4 inner steps × 2 outer rounds.

    After 2 outer rounds:
    - allreduce was called once per parameter per round
    - start_quorum was called once per round
    - outer optimizer's Nesterov state is populated for every parameter
    - parameters moved from the initial state
    """
    torch.manual_seed(0)
    model = TinyMLP()
    initial = {n: p.detach().clone() for n, p in model.named_parameters()}
    inner = optim.AdamW(model.parameters(), lr=1e-3)
    outer = optim.SGD(model.parameters(), lr=0.7, momentum=0.9, nesterov=True)
    mgr = _make_passthrough_manager()

    SYNC_EVERY = 4
    OUTER_ROUNDS = 2
    n_params = len(list(model.parameters()))

    with DiLoCo(mgr, [model], inner, outer, sync_every=SYNC_EVERY) as dl:
        for _outer_round in range(OUTER_ROUNDS):
            for _inner_step in range(SYNC_EVERY):
                inner.zero_grad()
                x = torch.randn(8, 4)
                y = torch.randn(8, 2)
                ((model(x) - y) ** 2).mean().backward()
                inner.step()  # outer sync fires automatically inside post-hook

    # 1. allreduce was called n_params × OUTER_ROUNDS times
    assert mgr.allreduce.call_count == n_params * OUTER_ROUNDS, (
        f"expected {n_params * OUTER_ROUNDS} allreduce calls, got {mgr.allreduce.call_count}"
    )

    # 2. start_quorum was called once per outer round
    assert mgr.start_quorum.call_count == OUTER_ROUNDS, (
        f"expected {OUTER_ROUNDS} start_quorum calls, got {mgr.start_quorum.call_count}"
    )

    # 3. should_commit was called once per outer round
    assert mgr.should_commit.call_count == OUTER_ROUNDS

    # 4. Outer optimizer holds Nesterov momentum state for every parameter
    assert len(outer.state_dict()["state"]) == n_params, (
        f"expected {n_params} momentum buffers, got {len(outer.state_dict()['state'])}"
    )

    # 5. Parameters moved from θ_initial (outer optimizer actually applied updates)
    any_change = any(
        not torch.equal(p, initial[n]) for n, p in model.named_parameters()
    )
    assert any_change, "outer optimizer did not move the parameters"


# ---------------------------------------------------------------------
# Acceptance test 2 — torchft sign convention is what we expect
# ---------------------------------------------------------------------

def test_diloco_pseudogradient_sign_convention():
    """Verify torchft computes pseudograd = θ_initial − θ_local + outer SGD math.

    Setup:
        - inner LR = 0 (so inner steps don't move params; only outer sync moves them)
        - manually nudge params so θ_local ≠ θ_initial
        - outer LR = 1, momentum = 0 (plain SGD, no Nesterov complications)
        - sync_every = 2

    Math:
        pseudograd = θ_initial − θ_local = -nudge
        restore: p.data ← θ_initial
        outer step: p.data ← θ_initial - lr * pseudograd
                          = θ_initial - 1 * (-nudge)
                          = θ_initial + nudge
                          = θ_local_at_sync
        merge(alpha=0): p.data unchanged

    Expected after 1 outer round: final = θ_local_at_sync

    A sign flip in pseudograd would land us at `θ_initial - nudge` (movement
    in the wrong direction by 2*nudge total), which this test catches.
    """
    torch.manual_seed(0)
    model = TinyMLP()
    inner = optim.SGD(model.parameters(), lr=0.0)  # zero inner LR
    outer = optim.SGD(model.parameters(), lr=1.0, momentum=0.0)  # plain SGD
    mgr = _make_passthrough_manager()

    SYNC_EVERY = 2
    NUDGE = 0.5
    initial_param = next(model.parameters()).detach().clone()

    with DiLoCo(mgr, [model], inner, outer, sync_every=SYNC_EVERY) as dl:
        # Manually nudge AFTER the DiLoCo wrapper saved θ_initial so
        # θ_local ≠ θ_initial when prepare_sync runs.
        with torch.no_grad():
            for p in model.parameters():
                p.add_(NUDGE)
        local_param_after_nudge = next(model.parameters()).detach().clone()

        # Run inner steps with zero LR — the post-hook fires the outer sync
        # at step `sync_every` but the inner step itself doesn't move params.
        for _ in range(SYNC_EVERY):
            inner.zero_grad()
            x = torch.randn(8, 4)
            ((model(x) - torch.randn(8, 2)) ** 2).mean().backward()
            inner.step()

        final_param = next(model.parameters()).detach().clone()

    # Per the math above: final should equal θ_local_at_sync = θ_initial + NUDGE.
    expected = local_param_after_nudge
    diff = (final_param - expected).abs().max().item()

    # And the wrong-sign result would have been θ_initial - NUDGE
    wrong_sign = initial_param - NUDGE * torch.ones_like(initial_param)
    wrong_sign_diff = (final_param - wrong_sign).abs().max().item()

    assert diff < 1e-5, (
        f"sign convention violated. \n"
        f"  initial[0,0]={initial_param.flatten()[0].item():.6f}\n"
        f"  local_at_sync[0,0]={local_param_after_nudge.flatten()[0].item():.6f}\n"
        f"  final[0,0]={final_param.flatten()[0].item():.6f}\n"
        f"  expected[0,0]={expected.flatten()[0].item():.6f}\n"
        f"  max-abs-diff={diff:.6e}\n"
        f"  wrong-sign-diff={wrong_sign_diff:.6e}  (≈0 means sign flipped)\n"
    )


# ---------------------------------------------------------------------
# Acceptance test 3 — Spike 005 imports still work alongside torchft
# ---------------------------------------------------------------------

def test_no_regression_in_spike_005_imports():
    """Verify importing torchft + composer_diloco coexists with Spike 005.

    This is a lightweight import-side-effects test. The 38-test Spike 005
    suite runs separately and passes there.
    """
    spike_005 = HERE.parent / "005-integrated-trainer-skeleton"
    sys.path.insert(0, str(spike_005))
    from opsd_loss import generalized_jsd_loss  # noqa: F401
    from teacher_replay import extract_dpo_pairs  # noqa: F401

    # Construct a fresh DiLoCo and verify it can be entered + exited
    model = TinyMLP()
    inner = optim.AdamW(model.parameters(), lr=1e-3)
    outer = optim.SGD(model.parameters(), lr=0.7, momentum=0.9, nesterov=True)
    mgr = _make_passthrough_manager()
    with DiLoCo(mgr, [model], inner, outer, sync_every=2) as dl:
        assert dl is not None


# ---------------------------------------------------------------------
# Acceptance test 4 — wrapper smoke (make_diloco_outer_loop)
# ---------------------------------------------------------------------

def test_make_diloco_outer_loop_factory():
    """The framework's `make_diloco_outer_loop()` constructs a working DiLoCo."""
    from composer_diloco import make_diloco_outer_loop

    model = TinyMLP()
    inner = optim.AdamW(model.parameters(), lr=1e-3)
    mgr = _make_passthrough_manager()

    dl = make_diloco_outer_loop(
        manager=mgr,
        model_fragments=[model],
        inner_optimizer=inner,
        outer_lr=0.7,
        outer_momentum=0.9,
        nesterov=True,
        sync_every=4,
    )
    # Outer optimizer was constructed with our hyperparams
    assert dl._sync_every == 4
    assert dl is not None


# ---------------------------------------------------------------------
# Acceptance test 5 — Streaming DiLoCo config path (deferred to v0.2 but
# importable today)
# ---------------------------------------------------------------------

def test_streaming_diloco_with_two_fragments_constructs():
    """Streaming DiLoCo accepts 2 fragments + nonzero sync delay (config path)."""
    torch.manual_seed(0)
    model = TinyMLP()
    # Two-fragment split (each linear is its own fragment)
    fragments = [model.net[0], model.net[2]]
    inner = optim.AdamW(model.parameters(), lr=1e-3)
    outer = optim.SGD(model.parameters(), lr=0.7, momentum=0.9, nesterov=True)
    mgr = _make_passthrough_manager()

    # sync_every=4, 2 fragments → effective per-fragment sync_every=2.
    # fragment_sync_delay=0 = no delay (still vanilla DiLoCo per-fragment).
    with DiLoCo(
        mgr, fragments, inner, outer,
        sync_every=4, fragment_sync_delay=0, fragment_update_alpha=0.0,
    ) as dl:
        assert len(dl._fragments) == 2