torchft-nightly 2026.1.3__cp310-cp310-manylinux_2_24_x86_64.whl

torchft/collectives_test.py

@@ -0,0 +1,212 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import unittest
+from typing import Callable
+from unittest import skipUnless, TestCase
+
+import torch
+import torch.distributed as dist
+from parameterized import parameterized
+from torch import cuda
+from torch.distributed import ReduceOp, ReduceScatterOptions
+
+from torchft import _test_utils
+from torchft.process_group import ProcessGroup
+from torchft.process_group_test import MultiPgBaseTest
+
+try:
+    # pyre-ignore[21]: Could not find a module corresponding to import `triton`
+    import triton
+except ImportError:
+    pass
+else:
+    from torchft.collectives import (
+        allocate_reduce_scatter_output,
+        allreduce_quantized,
+        get_padded_sizes,
+        reduce_scatter_quantized,
+    )
+
+    def _check_result_tolerance(
+        actual: torch.Tensor, expected: torch.Tensor, tolerance: float
+    ) -> None:
+        diff = torch.abs(
+            (expected - actual).div(expected.to(torch.float32) + 0.0000001)
+        )
+        mean_diff = diff.mean().item()
+
+        if mean_diff > tolerance:
+            print(f"Diff: {diff=}\n{expected=}\n{actual=}")
+            raise AssertionError(f"Results not within tolerance {tolerance}")
+
+    @skipUnless(
+        torch.cuda.is_available() and torch.cuda.device_count() >= 2,
+        "2 CUDA devices are required for this test",
+    )
+    class QuantizedAllReduceTest(MultiPgBaseTest):
+        BACKEND = "nccl"
+        WORLD_SIZE = 2
+
+        def _run_parallel_collectives(
+            self, collective: Callable[[ProcessGroup, int, str], None]
+        ) -> None:
+            futures = []
+            for rank in range(self.WORLD_SIZE):
+                pg = self.pg_pool[rank]
+                device = f"cuda:{rank}"
+                fut = self.executor.submit(collective, pg, rank, device)
+                futures.append(fut)
+
+            self._collect(futures)
+
+        def _run_all_reduce_collective(
+            self,
+            pg: ProcessGroup,
+            device: str,
+            tensors_num: int,
+            tensor_size: int,
+            multiplier: float,
+            tolerance: float,
+            reduce_op: ReduceOp,
+            dtype: torch.dtype,
+        ) -> None:
+            cuda.set_device(device)
+            inp = (
+                torch.rand(
+                    tensors_num * tensor_size,
+                    dtype=dtype,
+                    device=device,
+                )
+                * multiplier
+            )
+            for split in _test_utils.gen_splits(inp, tensor_size):
+                actual = inp.clone()
+                expected = inp.clone()
+                tensors = [
+                    i.view(*s)
+                    for s, i in zip(
+                        split,
+                        torch.split(actual, tensor_size),
+                    )
+                ]
+
+                work = allreduce_quantized(tensors, reduce_op, pg)
+                work.wait()
+
+                work = pg.allreduce([expected], reduce_op)
+                work.get_future().wait()
+
+                _check_result_tolerance(actual, expected, tolerance)
+
+        def _run_reduce_scatter_collective(
+            self,
+            pg: ProcessGroup,
+            device: str,
+            tensors_num: int,
+            tensor_size: int,
+            multiplier: float,
+            tolerance: float,
+            reduce_op: ReduceOp,
+            dtype: torch.dtype,
+        ) -> None:
+            cuda.set_device(device)
+            inp = (
+                torch.rand(
+                    tensors_num * tensor_size,
+                    dtype=dtype,
+                    device=device,
+                )
+                * multiplier
+            )
+            world_size = pg.size()
+            for split in _test_utils.gen_splits(inp, tensor_size):
+                actual = inp.clone()
+                tensors = [
+                    i.view(*s)
+                    for s, i in zip(
+                        split,
+                        torch.split(actual, tensor_size),
+                    )
+                ]
+
+                actual_output, _ = allocate_reduce_scatter_output(
+                    tensors,
+                    world_size,
+                )
+
+                opts = ReduceScatterOptions()
+                opts.reduceOp = reduce_op
+
+                work = reduce_scatter_quantized(actual_output, tensors, opts, pg)
+                work.get_future().wait()
+
+                padded_sizes = get_padded_sizes(tensors, world_size)
+                padded_numel = sum(s.numel() for s in padded_sizes)
+
+                padded_input = torch.empty(padded_numel, dtype=dtype, device=device)
+                torch._chunk_cat(
+                    tensors, dim=0, num_chunks=world_size, out=padded_input
+                )
+
+                expected_output = torch.empty(
+                    padded_numel // world_size, dtype=dtype, device=device
+                )
+
+                work = pg.reduce_scatter([expected_output], [[padded_input]], opts)
+                work.get_future().wait()
+
+                _check_result_tolerance(actual_output, expected_output, tolerance)
+
+        END_TO_END_CONFIGS: list[tuple[int, float, ReduceOp, torch.dtype]] = [
+            (ts, m, o, t)
+            for ts in [256, 1024, 2048]
+            for m in [1.0, 100.0, 1000.0]
+            for o in [ReduceOp.AVG, ReduceOp.SUM]
+            for t in [torch.float32, torch.float16, torch.bfloat16]
+        ]
+
+        @parameterized.expand(END_TO_END_CONFIGS)
+        def test_all_reduce_collective(
+            self,
+            tensor_size: int,
+            multiplier: float,
+            reduce_op: ReduceOp,
+            dtype: torch.dtype,
+        ) -> None:
+            self._run_parallel_collectives(
+                lambda pg, _, device: self._run_all_reduce_collective(
+                    pg,
+                    device,
+                    2,
+                    tensor_size,
+                    multiplier,
+                    0.04,
+                    reduce_op,
+                    dtype,
+                )
+            )
+
+        @parameterized.expand(END_TO_END_CONFIGS)
+        def test_reduce_scatter_collective(
+            self,
+            tensor_size: int,
+            multiplier: float,
+            reduce_op: ReduceOp,
+            dtype: torch.dtype,
+        ) -> None:
+            self._run_parallel_collectives(
+                lambda pg, _, device: self._run_reduce_scatter_collective(
+                    pg,
+                    device,
+                    2,
+                    tensor_size,
+                    multiplier,
+                    0.05,
+                    reduce_op,
+                    dtype,
+                )
+            )

torchft/coordination.py

@@ -0,0 +1,39 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+Coordination (Low Level API)
+============================
+
+.. warning::
+    As torchft is still in development, the APIs in this module are subject to change.
+
+This module exposes low level coordination APIs to allow you to build your own
+custom fault tolerance algorithms on top of torchft.
+
+If you're looking for a more complete solution, please use the other modules in
+torchft.
+
+This provides direct access to the Lighthouse and Manager servers and clients.
+"""
+
+from torchft._torchft import (
+    LighthouseClient,
+    LighthouseServer,
+    ManagerClient,
+    ManagerServer,
+    Quorum,
+    QuorumMember,
+)
+
+__all__ = [
+    "LighthouseClient",
+    "LighthouseServer",
+    "ManagerServer",
+    "ManagerClient",
+    "Quorum",
+    "QuorumMember",
+]

torchft/coordination_test.py

@@ -0,0 +1,29 @@
+import inspect
+from unittest import TestCase
+
+from torchft.coordination import (
+    LighthouseClient,
+    LighthouseServer,
+    ManagerClient,
+    ManagerServer,
+    Quorum,
+    QuorumMember,
+)
+
+
+class TestCoordination(TestCase):
+    def test_coordination_docs(self) -> None:
+        classes = [
+            ManagerClient,
+            ManagerServer,
+            LighthouseServer,
+            LighthouseClient,
+            Quorum,
+            QuorumMember,
+        ]
+        for cls in classes:
+            self.assertIn("Args:", str(cls.__doc__), cls)
+            for name, method in inspect.getmembers(cls, predicate=inspect.ismethod):
+                if name.startswith("_"):
+                    continue
+                self.assertIn("Args:", str(cls.__doc__), cls)

torchft/data.py

@@ -0,0 +1,77 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+Data
+====
+
+This module provides helper classes to implement fault tolerant data loaders.
+
+We recommend using torchdata's StatefulDataLoader to checkpoint each replica's
+dataloader frequently to avoid duplicate batches.
+"""
+
+from typing import Optional
+
+import torch.distributed as dist
+from torch.utils import data
+
+
+# pyre-fixme[24]: expected generic parameter
+class DistributedSampler(data.distributed.DistributedSampler):
+    """
+    DistributedSampler extends the standard PyTorch DistributedSampler with a
+    `num_replica_groups` that is used to shard the data across the fault
+    tolerance replica groups.
+
+    torchft doesn't know how many replica groups ahead of time so we need to set
+    this to be the max number.
+
+    This sampler is inherently lossy when used with torchft. torchft
+    occasionally drops batches on rejoining and if a replica group is down that
+    group examples will never be used. This can lead to imbalances if using a
+    small dataset.
+
+    This will shard the input dataset into ``num_replicas*num_replica_group``
+    number of shards.
+
+    Each shard rank is calculated via: ``rank + num_replicas*replica_rank``
+
+    num_replicas and replica_rank must be the same on all workers.
+    """
+
+    def __init__(
+        self,
+        dataset: data.Dataset,
+        replica_rank: int,
+        num_replica_groups: int,
+        group_rank: Optional[int] = None,
+        num_replicas: Optional[int] = None,
+        **kwargs: object,
+    ) -> None:
+        """
+        Args:
+            data: the dataset to use
+            replica_rank: the group ID (0-num_replica_groups) to use for this shard of data.
+            num_replica_groups: the max number of global replica groups
+            rank: the local group rank
+            num_replicas: the local group world size
+        """
+        if group_rank is None:
+            group_rank = dist.get_rank()
+        if num_replicas is None:
+            num_replicas = dist.get_world_size()
+
+        self.global_rank: int = group_rank + num_replicas * replica_rank
+        self.global_world_size: int = num_replicas * num_replica_groups
+
+        super().__init__(
+            dataset,
+            rank=self.global_rank,
+            num_replicas=self.global_world_size,
+            # pyre-fixme[6]: got object
+            **kwargs,
+        )

torchft/data_test.py

@@ -0,0 +1,39 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from unittest import TestCase
+
+from torch.utils.data import Dataset
+
+from torchft.data import DistributedSampler
+
+
+class DummyDataset(Dataset):
+    def __init__(self, length: int) -> None:
+        self.length = length
+
+    def __len__(self) -> int:
+        return self.length
+
+    def __getitem__(self, idx: int) -> int:
+        return idx
+
+
+class TestData(TestCase):
+    def test_distributed_sampler(self) -> None:
+        dataset = DummyDataset(1000)
+        sampler = DistributedSampler(
+            dataset,
+            replica_rank=1,
+            num_replica_groups=2,
+            group_rank=3,
+            num_replicas=4,
+        )
+        self.assertEqual(sampler.global_rank, 3 + 1 * 4)
+        self.assertEqual(sampler.global_world_size, 2 * 4)
+
+        sampler_iter = iter(sampler)
+        self.assertEqual(next(sampler_iter), 500)

torchft/ddp.py

@@ -0,0 +1,105 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+Distributed Data Parallel
+==========================
+
+This module implements a DistributedDataParallel wrapper that works with the
+Manager to provide fault tolerance.
+"""
+
+import os
+import sys
+from typing import cast, Optional, TYPE_CHECKING
+from unittest.mock import patch
+
+import torch
+import torch.distributed as dist
+from torch import nn
+from torch.distributed.algorithms.join import Joinable
+from torch.nn import parallel
+
+from torchft.process_group import ProcessGroup, ProcessGroupDummy, ProcessGroupGloo
+
+if TYPE_CHECKING:
+    from torchft.manager import _ManagedFuture, Manager
+
+
+class DistributedDataParallel(parallel.DistributedDataParallel):
+    """
+    This is a patched DistributedDataParallel implementation that makes it
+    compatible with torchft.
+
+    Important notes:
+
+    * This requires states to be synced on step 0 using an external mechanism
+      rather than an internal broadcast (torchft.Manager will do this).
+    * Using non-basic features of the DDP may cause your model to catch fire as
+      they haven't been tested with torchft.
+    * This doesn't any sanity checks such as verifying parameter sizes are the
+      same across workers.
+    """
+
+    def __init__(self, manager: "Manager", module: nn.Module, **kwargs: object) -> None:
+        # use a dummy PG to soak up the init all reduce, actual comms will go
+        # through the comm_hook.
+        pg = ProcessGroupDummy(0, 1)
+
+        super().__init__(
+            module,
+            process_group=pg,
+            # HACK: This forces the reducer to never rebuild buckets.
+            # The reducer normally rebuilds the buckets after the first training
+            # step which can improve performance but is incompatible with
+            # torchft as it will cause the buckets to diverge for recovering
+            # replicas.
+            find_unused_parameters=True,
+            # pyre-fixme[6]: got object
+            **kwargs,
+        )
+
+        self.register_comm_hook(manager, self._comm_hook)
+
+    @staticmethod
+    def _comm_hook(
+        state: "Manager", bucket: dist.GradBucket
+    ) -> torch.futures.Future[torch.Tensor]:
+        work = state.allreduce(bucket.buffer())
+        work.wait()
+        fut = work.get_future()
+
+        # We need to return the underlying future here otherwise
+        # this can hang
+        fut = cast("_ManagedFuture[torch.Tensor]", fut)
+        assert fut._fut
+        return fut._fut
+
+
+class PureDistributedDataParallel(nn.Module):
+    """
+    A pure Python reimplementation of the DDP wrapper.
+
+    We recommend using DistributedDataParallel instead of this class.
+
+    This calls one allreduce per gradient tensor and doesn't use a reducer. This
+    may be very slow for real models.
+    """
+
+    def __init__(self, manager: "Manager", module: nn.Module) -> None:
+        super().__init__()
+
+        self.module = module
+
+        def post_grad_hook(p: torch.Tensor) -> None:
+            if p.grad is not None:
+                manager.allreduce(p.grad)
+
+        for p in module.parameters():
+            p.register_post_accumulate_grad_hook(post_grad_hook)
+
+    def forward(self, *args: object) -> object:
+        return self.module(*args)

torchft/ddp_test.py

@@ -0,0 +1,68 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from unittest import TestCase
+from unittest.mock import create_autospec
+
+import torch
+import torch.distributed as dist
+from torch import nn
+from torch.distributed.distributed_c10d import Work
+from torch.futures import Future
+
+from torchft.ddp import DistributedDataParallel, PureDistributedDataParallel
+from torchft.manager import _ManagedWork, Manager
+from torchft.process_group import ProcessGroupBabyGloo, ProcessGroupGloo
+from torchft.work import _DummyWork
+
+
+class TestDDP(TestCase):
+    def test_pure_ddp(self) -> None:
+        manager = create_autospec(Manager)
+
+        m = nn.Linear(3, 4)
+        m = PureDistributedDataParallel(manager, m)
+
+        inp = torch.rand(2, 3)
+        out = m(inp)
+        loss = out.mean()
+        loss.backward()
+
+        for p in m.parameters():
+            self.assertIsNotNone(p.grad)
+
+        self.assertEqual(manager.allreduce.call_count, len(list(m.parameters())))
+
+    def test_ddp(self) -> None:
+        manager = create_autospec(Manager)
+
+        call_count = 0
+
+        # pyre-ignore[53]: Captured variable `manager` is not annotated.
+        def allreduce(
+            tensor: torch.Tensor,
+        ) -> Work:
+            nonlocal call_count
+
+            call_count += 1
+
+            work = _DummyWork(tensor)
+            return _ManagedWork(manager, work, tensor)
+
+        manager.allreduce = allreduce
+
+        m = nn.Linear(3, 4)
+        m = DistributedDataParallel(manager, m)
+
+        inp = torch.rand(2, 3)
+        out = m(inp)
+        loss = out.mean()
+        loss.backward()
+
+        for p in m.parameters():
+            self.assertIsNotNone(p.grad)
+
+        self.assertGreaterEqual(call_count, 1)