PR 10: Add refactor CI and recovery benchmarks

.github/workflows/refactor.yml

@@ -0,0 +1,51 @@
+name: refactor
+
+on:
+  pull_request:
+    branches: [refactor]
+  push:
+    branches: [refactor]
+
+jobs:
+  quality:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+      - uses: dtolnay/rust-toolchain@stable
+      - run: python -m pip install --upgrade pip
+      - run: python -m pip install -e ".[dev]"
+      - run: ruff check .
+      - run: ruff format --check .
+      - run: pyright
+      - run: cargo test --locked
+      - name: CPU control, planning, data, and operator tests
+        run: >-
+          pytest -q
+          tests/refactor/test_control.py
+          tests/refactor/test_control_service.py
+          tests/refactor/test_planning.py
+          tests/refactor/test_reconfiguration_objectives.py
+          tests/refactor/test_data_runtime.py
+          tests/refactor/test_dataloader_workers.py
+          tests/refactor/test_examples.py
+          tests/refactor/test_hostfile.py
+          tests/refactor/test_state.py
+          tests/refactor/test_state_bundles.py
+          tests/refactor/test_topology.py
+          tests/refactor/test_optimization.py
+          tests/refactor/test_benchmark.py
+
+  cuda-gloo:
+    runs-on: [self-hosted, linux, arm64, gpu]
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+      - run: python -m pip install -e ".[dev]"
+      - run: nvidia-smi -L
+      - run: python -c "import torch; assert torch.cuda.is_available() and torch.cuda.device_count() >= 1"
+      - run: pytest -q

benchmarks/__init__.py

@@ -0,0 +1 @@
+"""Reproducible Oobleck research benchmarks."""

benchmarks/recovery_schedule.py

@@ -0,0 +1,99 @@
+"""Machine-readable balanced-recovery scheduling benchmark."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import time
+from pathlib import Path
+
+from oobleck.state import LogicalStateEntry, StateManifest, plan_state_redistribution
+
+
+def run(
+    *, replicas: int, tensor_bytes: int, chunk_bytes: int, round_bytes: int | None = None
+) -> dict[str, object]:
+    if replicas < 1 or tensor_bytes < 1 or chunk_bytes < 1:
+        raise ValueError("replicas and byte sizes must be positive")
+
+    def entry(rank: int) -> LogicalStateEntry:
+        return LogicalStateEntry(
+            "layers.0.weight",
+            (tensor_bytes,),
+            (tensor_bytes,),
+            "uint8",
+            "parameter",
+            ("replicate",),
+            0,
+            rank,
+            1,
+        )
+
+    old = tuple(StateManifest(rank, (entry(rank),), 1) for rank in range(replicas))
+    destination_rank = replicas
+    new = (StateManifest(destination_rank, (entry(destination_rank),), 1),)
+    started = time.perf_counter()
+    schedule = plan_state_redistribution(
+        old,
+        new,
+        chunk_bytes=chunk_bytes,
+        round_bytes=round_bytes,
+        alignment=min(256, chunk_bytes),
+    )
+    elapsed = time.perf_counter() - started
+    balanced = dict(schedule.per_source_bytes)
+    source_round_loads: dict[tuple[int, int], int] = {}
+    destination_round_loads: dict[tuple[int, int], int] = {}
+    for transfer in schedule.transfers:
+        source_key = (transfer.round, transfer.source_rank)
+        destination_key = (transfer.round, transfer.destination_rank)
+        source_round_loads[source_key] = source_round_loads.get(source_key, 0) + transfer.byte_count
+        destination_round_loads[destination_key] = (
+            destination_round_loads.get(destination_key, 0) + transfer.byte_count
+        )
+    return {
+        "schema_version": 1,
+        "scenario": {
+            "surviving_replicas": replicas,
+            "destinations": 1,
+            "tensor_bytes": tensor_bytes,
+            "chunk_bytes": chunk_bytes,
+            "round_bytes": round_bytes or chunk_bytes * replicas,
+        },
+        "schedule_hash": schedule.schedule_hash,
+        "transfer_count": len(schedule.transfers),
+        "balanced_source_bytes": balanced,
+        "balanced_max_source_bytes": max(balanced.values(), default=0),
+        "legacy_first_source_max_bytes": tensor_bytes,
+        "max_source_load_reduction": 1.0 - max(balanced.values(), default=0) / tensor_bytes,
+        "round_count": len({item.round for item in schedule.transfers}),
+        "maximum_round_source_bytes": max(source_round_loads.values(), default=0),
+        "maximum_round_destination_bytes": max(destination_round_loads.values(), default=0),
+        "planning_seconds": elapsed,
+    }
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--replicas", type=int, default=4)
+    parser.add_argument("--tensor-bytes", type=int, default=64 * 1024 * 1024)
+    parser.add_argument("--chunk-bytes", type=int, default=4 * 1024 * 1024)
+    parser.add_argument("--round-bytes", type=int)
+    parser.add_argument("--output", type=Path)
+    arguments = parser.parse_args()
+    result = run(
+        replicas=arguments.replicas,
+        tensor_bytes=arguments.tensor_bytes,
+        chunk_bytes=arguments.chunk_bytes,
+        round_bytes=arguments.round_bytes,
+    )
+    payload = json.dumps(result, indent=2, sort_keys=True) + "\n"
+    if arguments.output is None:
+        print(payload, end="")
+    else:
+        arguments.output.parent.mkdir(parents=True, exist_ok=True)
+        arguments.output.write_text(payload)
+
+
+if __name__ == "__main__":
+    main()

benchmarks/results/synthetic_recovery.json

@@ -0,0 +1,25 @@
+{
+  "balanced_max_source_bytes": 16777216,
+  "balanced_source_bytes": {
+    "0": 16777216,
+    "1": 16777216,
+    "2": 16777216,
+    "3": 16777216
+  },
+  "legacy_first_source_max_bytes": 67108864,
+  "max_source_load_reduction": 0.75,
+  "maximum_round_destination_bytes": 16777216,
+  "maximum_round_source_bytes": 4194304,
+  "planning_seconds": 0.00039029866456985474,
+  "round_count": 4,
+  "scenario": {
+    "chunk_bytes": 4194304,
+    "destinations": 1,
+    "round_bytes": 16777216,
+    "surviving_replicas": 4,
+    "tensor_bytes": 67108864
+  },
+  "schedule_hash": "475c9bd581858e75ea756c8a47535aee45174a3530e0bd79a7aec04f015c6c88",
+  "schema_version": 1,
+  "transfer_count": 16
+}

benchmarks/results/transaction_overhead.json

@@ -0,0 +1,23 @@
+{
+  "baseline_step_seconds": 0.00040008220821619034,
+  "committed_step": 5,
+  "numerically_equivalent": true,
+  "oobleck_steady_step_seconds": 0.0013069706037640572,
+  "recovery_attempts": 2,
+  "recovery_extra_seconds": 0.0015370016917586327,
+  "recovery_step_seconds": 0.00284397229552269,
+  "replayed_sample_indices": [
+    16,
+    17,
+    18,
+    19
+  ],
+  "scenario": {
+    "device": "cpu",
+    "global_batch_size": 4,
+    "microbatch_size": 2,
+    "steady_steps": 4
+  },
+  "schema_version": 1,
+  "steady_state_overhead_ratio": 3.266755124131429
+}

benchmarks/transaction_overhead.py

@@ -0,0 +1,155 @@
+"""Machine-readable steady-state and transactional recovery benchmark."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import time
+from dataclasses import dataclass, replace
+from pathlib import Path
+
+import torch
+from torch.utils.data import DataLoader, Dataset
+
+from oobleck import OobleckConfig, OobleckParallelizationPlan
+
+
+class LinearDataset(Dataset):
+    def __init__(self, samples: int) -> None:
+        self.x = torch.arange(samples, dtype=torch.float32).unsqueeze(1) / samples
+        self.y = 2 * self.x + 0.5
+
+    def __len__(self) -> int:
+        return len(self.x)
+
+    def __getitem__(self, index: int) -> dict[str, torch.Tensor]:
+        return {"x": self.x[index], "y": self.y[index]}
+
+
+@dataclass(frozen=True)
+class ParallelConfig:
+    tensor_parallel_size: int = 1
+    pipeline_parallel_size: None = None
+    data_parallel_size: int = 1
+    context_parallel_size: int = 1
+    expert_parallel_size: int = 1
+
+
+def run(*, steady_steps: int = 4) -> dict[str, object]:
+    if steady_steps < 1:
+        raise ValueError("steady_steps must be positive")
+    torch.manual_seed(11)
+    initial = torch.nn.Linear(1, 1).state_dict()
+    dataset = LinearDataset((steady_steps + 1) * 4)
+    warmup = torch.nn.Linear(1, 1)
+    warmup_optimizer = torch.optim.SGD(warmup.parameters(), lr=0.05)
+    warmup_optimizer.zero_grad(set_to_none=True)
+    warmup_loss = torch.nn.functional.mse_loss(warmup(dataset.x[:4]), dataset.y[:4])
+    warmup_loss.backward()
+    warmup_optimizer.step()
+
+    reference = torch.nn.Linear(1, 1)
+    reference.load_state_dict(initial)
+    reference_optimizer = torch.optim.SGD(reference.parameters(), lr=0.05)
+    baseline_durations = []
+    for step in range(steady_steps + 1):
+        batch = {key: value[step * 4 : (step + 1) * 4] for key, value in vars(dataset).items()}
+        started = time.perf_counter()
+        reference_optimizer.zero_grad(set_to_none=True)
+        loss = torch.nn.functional.mse_loss(reference(batch["x"]), batch["y"])
+        loss.backward()
+        reference_optimizer.step()
+        baseline_durations.append(time.perf_counter() - started)
+
+    model = torch.nn.Linear(1, 1)
+    model.load_state_dict(initial)
+    plan = OobleckParallelizationPlan(
+        OobleckConfig(global_batch_size=4, microbatch_size=2, max_nodes=1, seed=11)
+    )
+    plan.parallelize(model, ParallelConfig())
+    context = plan.materialize("cpu")
+    sampler = context.create_batch_sampler(dataset, shuffle=False)
+    loader = context.prepare_dataloader(DataLoader(dataset, batch_sampler=sampler))
+    context.configure_optimization(
+        optimizer_factory=lambda parameters: torch.optim.SGD(parameters, lr=0.05)
+    )
+    iterator = iter(loader)
+    steady_durations = []
+    for _ in range(steady_steps):
+        batch = next(iterator)
+        started = time.perf_counter()
+        context.step(
+            batch,
+            lambda microbatch: model(microbatch["x"]),
+            criterion=lambda output, microbatch: torch.nn.functional.mse_loss(
+                output, microbatch["y"]
+            ),
+        )
+        steady_durations.append(time.perf_counter() - started)
+
+    recovery_batch = next(iterator)
+    newer = replace(
+        context.execution_plan,
+        generation=context.generation + 1,
+        previous_generation=context.generation,
+        plan_checksum="",
+    )
+    announced = False
+
+    def recovery_criterion(output, microbatch):
+        nonlocal announced
+        if not announced:
+            announced = True
+            context.announce_generation(newer)
+        return torch.nn.functional.mse_loss(output, microbatch["y"])
+
+    recovery_started = time.perf_counter()
+    recovery_result = context.step(
+        recovery_batch,
+        lambda microbatch: model(microbatch["x"]),
+        criterion=recovery_criterion,
+    )
+    recovery_seconds = time.perf_counter() - recovery_started
+    equivalent = all(
+        torch.allclose(left, right, rtol=1e-5, atol=1e-6)
+        for left, right in zip(model.parameters(), reference.parameters())
+    )
+    baseline_mean = sum(baseline_durations[:steady_steps]) / steady_steps
+    steady_mean = sum(steady_durations) / steady_steps
+    result = {
+        "schema_version": 1,
+        "scenario": {
+            "device": "cpu",
+            "steady_steps": steady_steps,
+            "global_batch_size": 4,
+            "microbatch_size": 2,
+        },
+        "baseline_step_seconds": baseline_mean,
+        "oobleck_steady_step_seconds": steady_mean,
+        "steady_state_overhead_ratio": steady_mean / baseline_mean,
+        "recovery_step_seconds": recovery_seconds,
+        "recovery_extra_seconds": max(0.0, recovery_seconds - steady_mean),
+        "recovery_attempts": recovery_result.attempts,
+        "committed_step": recovery_result.committed_step,
+        "replayed_sample_indices": list(recovery_batch.sample_indices),
+        "numerically_equivalent": equivalent,
+    }
+    context.close()
+    return result
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--steady-steps", type=int, default=4)
+    parser.add_argument("--output", type=Path)
+    arguments = parser.parse_args()
+    payload = json.dumps(run(steady_steps=arguments.steady_steps), indent=2, sort_keys=True) + "\n"
+    if arguments.output is None:
+        print(payload, end="")
+    else:
+        arguments.output.parent.mkdir(parents=True, exist_ok=True)
+        arguments.output.write_text(payload)
+
+
+if __name__ == "__main__":
+    main()

tests/refactor/test_benchmark.py

@@ -0,0 +1,26 @@
+from benchmarks.recovery_schedule import run
+from benchmarks.transaction_overhead import run as run_transaction_overhead
+
+
+def test_recovery_schedule_benchmark_is_machine_readable_and_balanced():
+    result = run(replicas=4, tensor_bytes=64, chunk_bytes=16)
+    assert result["schema_version"] == 1
+    assert result["balanced_max_source_bytes"] == 16
+    assert result["legacy_first_source_max_bytes"] == 64
+    assert result["max_source_load_reduction"] == 0.75
+    assert result["transfer_count"] == 4
+    assert result["round_count"] == 1
+    assert result["maximum_round_source_bytes"] == 16
+    assert result["maximum_round_destination_bytes"] == 64
+
+
+def test_transaction_overhead_benchmark_replays_once_and_matches_reference():
+    result = run_transaction_overhead(steady_steps=2)
+    assert result["schema_version"] == 1
+    assert result["baseline_step_seconds"] > 0
+    assert result["oobleck_steady_step_seconds"] > 0
+    assert result["recovery_step_seconds"] > 0
+    assert result["recovery_attempts"] == 2
+    assert result["committed_step"] == 3
+    assert result["numerically_equivalent"] is True
+    assert len(result["replayed_sample_indices"]) == 4