diff --git a/.github/workflows/refactor.yml b/.github/workflows/refactor.yml new file mode 100644 index 00000000..bd12d234 --- /dev/null +++ b/.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 diff --git a/benchmarks/__init__.py b/benchmarks/__init__.py new file mode 100644 index 00000000..0f9019fa --- /dev/null +++ b/benchmarks/__init__.py @@ -0,0 +1 @@ +"""Reproducible Oobleck research benchmarks.""" diff --git a/benchmarks/recovery_schedule.py b/benchmarks/recovery_schedule.py new file mode 100644 index 00000000..d24e53c2 --- /dev/null +++ b/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() diff --git a/benchmarks/results/synthetic_recovery.json b/benchmarks/results/synthetic_recovery.json new file mode 100644 index 00000000..026c76c3 --- /dev/null +++ b/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 +} diff --git a/benchmarks/results/transaction_overhead.json b/benchmarks/results/transaction_overhead.json new file mode 100644 index 00000000..3d9af13e --- /dev/null +++ b/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 +} diff --git a/benchmarks/transaction_overhead.py b/benchmarks/transaction_overhead.py new file mode 100644 index 00000000..a9494320 --- /dev/null +++ b/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() diff --git a/tests/refactor/test_benchmark.py b/tests/refactor/test_benchmark.py new file mode 100644 index 00000000..05903f26 --- /dev/null +++ b/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