torchft-nightly 2026.1.3__cp310-cp310-manylinux_2_24_x86_64.whl

torchft/_test_utils.py

@@ -0,0 +1,111 @@
+# 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 torch
+
+
+def any_nan(ts: list[torch.Tensor]) -> bool:
+    """
+    Check if any tensor in the list contains NaN values.
+
+    Args:
+        ts: List of tensors to check for NaN values
+
+    Returns:
+        True if any tensor contains NaN values, False otherwise
+    """
+    for t in ts:
+        if torch.isnan(t).any():
+            return True
+    return False
+
+
+def combine_views(
+    views: list[list[tuple[int, ...]]],
+    combinations: list[list[tuple[int, ...]]],
+    tmp: list[tuple[int, ...]],
+    i: int,
+) -> None:
+    """
+    Recursively generate all possible combinations of views from a list of
+    lists of views.
+
+    This function uses backtracking to generate all possible combinations by
+    selecting each list in the input. The results are stored in the
+    combinations list.
+
+    Args:
+        views: A list of lists, where each inner list contains possible view
+            shapes (tuples)
+        combinations: Output list where all combinations will be stored
+        tmp: Temporary list to build the current combination
+        i: Current index in the views list being processed
+
+    Returns:
+        None. Results are stored in the combinations list passed as
+        an argument.
+    """
+    if i == len(views):
+        combinations.append(tmp.copy())
+        return
+
+    for j in range(len(views[i])):
+        tmp.append(views[i][j])
+        combine_views(views, combinations, tmp, i + 1)
+        tmp.pop()
+
+
+def gen_views(inp: torch.Tensor) -> list[tuple[int, ...]]:
+    """
+    Generate all possible 2D views (shapes) for a tensor with a given number
+    of elements.
+
+    This function finds all pairs of integers (m, n) such that m * n equals the
+    total number of elements in the input tensor. These pairs represent possible
+    2D shapes that the tensor can be reshaped into.
+
+    Args:
+        inp: Input tensor
+
+    Returns:
+        A list of tuples, where each tuple (m, n) represents a possible 2D shape
+        such that m * n equals the total number of elements in the input tensor
+    """
+    size = inp.numel()
+
+    views = []
+    for m in range(1 if size % 2 == 0 else 2, size):
+        if size % m == 0:
+            views.append((m, size // m))
+
+    return views
+
+
+def gen_splits(inp: torch.Tensor, split_size: int) -> list[list[tuple[int, ...]]]:
+    """
+    Split a tensor into chunks and generate all possible combinations of views.
+
+    This function first splits the input tensor into chunks of the specified size,
+    then generates all possible 2D views for each chunk, and finally computes all
+    possible combinations of these views across all chunks.
+
+    Args:
+        inp: Input tensor to be split
+        split_size: Size of each chunk
+
+    Returns:
+        A list of lists, where each inner list contains a combination of view
+        shapes, one for each chunk of the input tensor
+    """
+    views = []
+
+    for split in torch.split(inp, split_size):
+        views.append(gen_views(split))
+
+    combinations = []
+    combine_views(views, combinations, [], 0)
+
+    return combinations

torchft/_torchft.pyi

@@ -0,0 +1,116 @@
+# 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 dataclasses import dataclass
+from datetime import timedelta
+from typing import Hashable, List, Optional
+
+class ManagerClient:
+    def __init__(self, addr: str, connect_timeout: timedelta) -> None: ...
+    def _quorum(
+        self,
+        group_rank: int,
+        step: int,
+        checkpoint_metadata: str,
+        shrink_only: bool,
+        timeout: timedelta,
+        commit_failures: int,
+        init_sync: bool = True,
+    ) -> QuorumResult: ...
+    def _checkpoint_metadata(self, rank: int, timeout: timedelta) -> str: ...
+    def should_commit(
+        self,
+        group_rank: int,
+        step: int,
+        should_commit: bool,
+        timeout: timedelta,
+    ) -> bool: ...
+
+class QuorumResult:
+    quorum_id: int
+    replica_rank: int
+    replica_world_size: int
+    recover_src_manager_address: str
+    recover_src_replica_rank: Optional[int]
+    recover_dst_replica_ranks: List[int]
+    store_address: str
+    max_step: int
+    max_replica_rank: Optional[int]
+    max_world_size: int
+    heal: bool
+    commit_failures: int
+    replica_ids: list[str]
+
+class ManagerServer:
+    def __init__(
+        self,
+        replica_id: str,
+        lighthouse_addr: str,
+        hostname: str,
+        bind: str,
+        store_addr: str,
+        world_size: int,
+        heartbeat_interval: timedelta,
+        connect_timeout: timedelta,
+        quorum_retries: int,
+    ) -> None: ...
+    def address(self) -> str: ...
+    def shutdown(self) -> None: ...
+
+class LighthouseServer:
+    def __init__(
+        self,
+        bind: str,
+        min_replicas: int,
+        join_timeout_ms: Optional[int] = None,
+        quorum_tick_ms: Optional[int] = None,
+        heartbeat_timeout_ms: Optional[int] = None,
+    ) -> None: ...
+    def address(self) -> str: ...
+    def shutdown(self) -> None: ...
+
+@dataclass
+class QuorumMember:
+    replica_id: str
+    address: str
+    store_address: str
+    step: int
+    world_size: int
+    shrink_only: bool
+    data: Optional[dict[Hashable, object]] = None
+
+@dataclass
+class Timestamp:
+    seconds: int
+    nanos: int
+
+@dataclass
+class Quorum:
+    quorum_id: str
+    participants: List[QuorumMember]
+    created: Timestamp
+
+@dataclass
+class LighthouseClient:
+    addr: str
+    connect_timeout: timedelta
+
+    def quorum(
+        self,
+        replica_id: str,
+        timeout: timedelta,
+        address: Optional[str] = None,
+        store_address: Optional[str] = None,
+        step: Optional[int] = None,
+        world_size: Optional[int] = None,
+        shrink_only: Optional[bool] = None,
+        data: Optional[dict[Hashable, object]] = None,
+    ) -> Quorum: ...
+    def heartbeat(
+        self,
+        replica_id: str,
+        timeout: timedelta = timedelta(seconds=5),
+    ) -> None: ...

torchft/checkpointing/__init__.py

@@ -0,0 +1,20 @@
+# 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.
+
+"""
+Checkpointing
+==============
+
+This module implements methods for checkpointing and resuming training from a checkpoint.
+"""
+
+from torchft.checkpointing.http_transport import HTTPTransport
+from torchft.checkpointing.transport import CheckpointTransport
+
+__all__ = [
+    "HTTPTransport",
+    "CheckpointTransport",
+]

torchft/checkpointing/_rwlock.py

@@ -0,0 +1,136 @@
+# 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.
+
+"""rwlock.py
+
+Adapted from: https://github.com/tylerneylon/rwlock/blob/main/rwlock.py
+
+A class to implement read-write locks on top of the standard threading
+library.
+
+This is implemented with two mutexes (threading.Lock instances) as per this
+wikipedia pseudocode:
+
+https://en.wikipedia.org/wiki/Readers%E2%80%93writer_lock#Using_two_mutexes
+
+__________________________
+License info (MIT):
+
+*******
+
+Copyright 2023 Tyler Neylon and contributors
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated
+documentation files (the "Software"), to deal in the Software without restriction, including without limitation the
+rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit
+persons to whom the Software is furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE
+WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
+COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
+OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+
+*******
+"""
+
+from contextlib import contextmanager
+from threading import Lock
+from typing import Generator
+
+
+class RWLock(object):
+    """RWLock class; this is meant to allow an object to be read from by
+    multiple threads, but only written to by a single thread at a time. See:
+    https://en.wikipedia.org/wiki/Readers%E2%80%93writer_lock
+
+    All operations are timed and will throw TimeoutError if the timeout is
+    exceeded.
+
+    Usage:
+
+        from rwlock import RWLock
+
+        my_obj_rwlock = RWLock(timeout=60.0)
+
+        # When reading from my_obj:
+        with my_obj_rwlock.r_lock():
+            do_read_only_things_with(my_obj)
+
+        # When writing to my_obj:
+        with my_obj_rwlock.w_lock():
+            mutate(my_obj)
+    """
+
+    def __init__(self, timeout: float = -1) -> None:
+        self.timeout = timeout
+
+        self._w_lock = Lock()
+        self._num_r_lock = Lock()
+        self._num_r = 0
+
+    # ___________________________________________________________________
+    # Reading methods.
+
+    def r_acquire(self) -> None:
+        if not self._num_r_lock.acquire(timeout=self.timeout):
+            raise TimeoutError(
+                f"Timed out waiting for rlock after {self.timeout} seconds"
+            )
+
+        self._num_r += 1
+        if self._num_r == 1:
+            if not self._w_lock.acquire(timeout=self.timeout):
+                self._num_r -= 1
+                self._num_r_lock.release()
+                raise TimeoutError(
+                    f"Timed out waiting for wlock after {self.timeout} seconds"
+                )
+
+        self._num_r_lock.release()
+
+    def r_release(self) -> None:
+        assert self._num_r > 0
+        self._num_r_lock.acquire()
+        self._num_r -= 1
+        if self._num_r == 0:
+            self._w_lock.release()
+        self._num_r_lock.release()
+
+    @contextmanager
+    def r_lock(self) -> Generator[None, None, None]:
+        """This method is designed to be used via the `with` statement."""
+        self.r_acquire()
+        try:
+            yield
+        finally:
+            self.r_release()
+
+    # ___________________________________________________________________
+    # Writing methods.
+
+    def w_acquire(self) -> None:
+        if not self._w_lock.acquire(timeout=self.timeout):
+            raise TimeoutError(
+                f"Timed out waiting for wlock after {self.timeout} seconds"
+            )
+
+    def w_release(self) -> None:
+        self._w_lock.release()
+
+    @contextmanager
+    def w_lock(self) -> Generator[None, None, None]:
+        """This method is designed to be used via the `with` statement."""
+        self.w_acquire()
+        try:
+            yield
+        finally:
+            self.w_release()
+
+    def w_locked(self) -> bool:
+        """Returns True if the lock is currently locked for reading."""
+        return self._w_lock.locked()

torchft/checkpointing/_serialization.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.
+
+import io
+import warnings
+from typing import IO
+
+import torch
+
+
+def _fallback_save(obj: object, f: IO[bytes]) -> None:
+    warnings.warn(
+        "using slow fallback torch.save implementation, please upgrade to PT 2.7+ for fast streaming saves"
+    )
+
+    torch.save(obj, f)
+
+
+def _fallback_load(f: IO[bytes], weights_only: bool = True) -> object:
+    warnings.warn(
+        "using slow fallback torch.load implementation, please upgrade to PT 2.7+ for fast streaming loads"
+    )
+
+    # torch.load requires a seekable file object
+    buf = f.read()
+    reader = io.BytesIO(buf)
+
+    return torch.load(reader, weights_only=weights_only)
+
+
+try:
+    # upgrade to PT 2.7 once released
+    from torch.distributed._serialization import _streaming_load, _streaming_save
+except ImportError:
+    _streaming_load = _fallback_load
+    _streaming_save = _fallback_save