d9d 0.1.0__py3-none-any.whl

d9d/core/types/__init__.py

@@ -0,0 +1,12 @@
+"""
+Common type definitions used throughout the framework.
+"""
+from .data import CollateFn
+from .pytree import PyTree, ScalarTree, TensorTree
+
+__all__ = [
+    "CollateFn",
+    "PyTree",
+    "ScalarTree",
+    "TensorTree"
+]

d9d/core/types/data.py

@@ -0,0 +1,14 @@
+from collections.abc import Callable, Sequence
+from typing import TypeAlias, TypeVar
+
+from .pytree import PyTree
+
+TDataTree = TypeVar("TDataTree", bound=PyTree)
+
+CollateFn: TypeAlias = Callable[[Sequence[TDataTree]], TDataTree]
+"""
+Type alias for a function that collates a sequence of samples into a batch.
+
+The function receives a sequence of individual data point structures (PyTrees)
+and is responsible for stacking or merging them into a single batched structure.
+"""

d9d/core/types/pytree.py

@@ -0,0 +1,26 @@
+from typing import TypeAlias, TypeVar
+
+import torch
+
+TLeaf = TypeVar("TLeaf")
+
+PyTree: TypeAlias = TLeaf | list["PyTree[TLeaf]"] | dict[str, "PyTree[TLeaf]"] | tuple["PyTree[TLeaf]", ...]
+"""
+A recursive type definition representing a tree of data.
+
+This type alias covers standard Python containers (dictionaries, lists, tuples)
+nested arbitrarily deep, terminating in a leaf node of type `TLeaf`.
+
+This is commonly used for handling nested state dictionaries or arguments
+passed to functions that support recursive traversal (similar to `torch.utils._pytree`).
+"""
+
+TensorTree: TypeAlias = PyTree[torch.Tensor]
+"""
+A recursive tree structure where the leaf nodes are PyTorch Tensors.
+"""
+
+ScalarTree: TypeAlias = PyTree[str | float | int | bool]
+"""
+A recursive tree structure where the leaf nodes are python scalars (str, float, int).
+"""

d9d/dataset/__init__.py

@@ -0,0 +1,17 @@
+"""
+This package provides utilities and torch.utils.data.Dataset implementations.
+"""
+
+from .buffer_sorted import BufferSortedDataset, DatasetImplementingSortKeyProtocol
+from .padding import PaddingSide1D, pad_stack_1d
+from .sharded import ShardedDataset, ShardIndexingMode, shard_dataset_data_parallel
+
+__all__ = [
+    "BufferSortedDataset",
+    "DatasetImplementingSortKeyProtocol",
+    "PaddingSide1D",
+    "ShardIndexingMode",
+    "ShardedDataset",
+    "pad_stack_1d",
+    "shard_dataset_data_parallel"
+]

d9d/dataset/buffer_sorted.py

@@ -0,0 +1,143 @@
+import pickle  # noqa: S403
+import random
+from typing import Any, Protocol, TypeVar
+
+from torch.distributed.checkpoint.stateful import Stateful
+from torch.utils.data import Dataset
+
+_T_co = TypeVar("_T_co", covariant=True)
+
+
+class DatasetImplementingSortKeyProtocol(Protocol[_T_co]):
+    """
+    Protocol for datasets that support retrieval of a specific key for sorting purposes.
+
+    This is typically used for length-based bucketing/sorting where the dataset
+    needs to expose the length of an item without loading the full item.
+    """
+
+    def __len__(self) -> int:
+        """Returns the total number of items in the dataset."""
+        ...
+
+    def sort_key(self, index: int) -> Any:
+        """
+        Returns a value used for sorting the dataset at the given index.
+
+        Args:
+            index: The index of the item.
+
+        Returns:
+            A comparable value (e.g., int length) used for sorting.
+        """
+        ...
+
+    def __getitem__(self, item: int) -> _T_co:
+        """Retrieves the item at the specific index."""
+        ...
+
+
+class BufferSortedDataset(Dataset[_T_co], Stateful):
+    """
+    A dataset wrapper that groups items into buffers, sorts them, and yields them with local shuffling.
+
+    This prevents extreme padding in variable-length training (by grouping similar lengths)
+    while maintaining enough randomness to ensure statistical variance in updates.
+
+    Algorithm:
+
+    1. Select a range of indices (size `buffer_size`).
+    2. Sort these indices based on `base_dataset.sort_key()`.
+    3. Break the sorted list into packs of size `pack_size`.
+    4. Shuffle the order of these packs.
+    5. Flatten the list and serve items.
+    """
+
+    def __init__(
+            self,
+            base_dataset: DatasetImplementingSortKeyProtocol[_T_co],
+            buffer_size: int,
+            pack_size: int,
+            init_seed: int | None = None
+    ):
+        """
+        Constructs a BufferSortedDataset object.
+
+        Args:
+            base_dataset: The underlying dataset implementing the `DatasetImplementingSortKeyProtocol` protocol.
+            buffer_size: The number of items to load into the buffer for sorting.
+            pack_size: The size of local groups (batches/micro-batches) that remain
+                contiguous after sorting, but are shuffled relative to other packs.
+            init_seed: Seed for the random number generator used for shuffling packs.
+        """
+
+        self._base_dataset = base_dataset
+        self._buffer_size = buffer_size
+        self._pack_size = pack_size
+
+        self._rng = random.Random(init_seed ^ 0x105E7 if init_seed is not None else None)
+        self._buffer_indices: list[int] = []
+        self._buffer_idx: int = -1
+
+    def _update_buffer_idx(self, buffer_idx: int):
+        select_start = buffer_idx * self._buffer_size
+        select_end = (buffer_idx + 1) * self._buffer_size
+        select_end = min(select_end, len(self._base_dataset))
+
+        base_idx = list(range(select_start, select_end))
+        base_sort_keys = [self._base_dataset.sort_key(idx) for idx in range(select_start, select_end)]
+
+        local_idx = list(range(len(base_idx)))
+        local_idx = sorted(local_idx, key=lambda local_id: base_sort_keys[local_id])
+
+        local_idx_batch = [
+            local_idx[i: i + self._pack_size]
+            for i in range(0, len(local_idx), self._pack_size)
+        ]
+        self._rng.shuffle(local_idx_batch)
+        local_idx = [y for x in local_idx_batch for y in x]
+
+        self._buffer_indices = [base_idx[local_id] for local_id in local_idx]
+
+        self._buffer_idx = buffer_idx
+
+    def __getitem__(self, index: int) -> _T_co:
+        """
+        Retrieves an item from the locally sorted/shuffled buffer.
+
+        Args:
+            index: The global index.
+
+        Returns:
+            The dataset item.
+        """
+
+        needs_buffer_idx = index // self._buffer_size
+        if self._buffer_idx != needs_buffer_idx:
+            self._update_buffer_idx(needs_buffer_idx)
+
+        take_id = self._buffer_indices[index % self._buffer_size]
+
+        return self._base_dataset[take_id]
+
+    def __len__(self) -> int:
+        """Returns the length of the base dataset."""
+
+        return len(self._base_dataset)
+
+    def state_dict(self) -> dict[str, Any]:
+        ret = {
+            "seed": pickle.dumps(self._rng.getstate()),
+            "buffer_idx": self._buffer_idx,
+            "buffer_indices": self._buffer_indices,
+        }
+        if isinstance(self._base_dataset, Stateful):
+            ret["base_dataset"] = self._base_dataset.state_dict()
+        return ret
+
+    def load_state_dict(self, state_dict: dict[str, Any]) -> None:
+        self._rng.setstate(pickle.loads(state_dict["seed"]))  # noqa: S301
+        self._buffer_idx = state_dict["buffer_idx"]
+        self._buffer_indices = state_dict["buffer_indices"]
+        if isinstance(self._base_dataset, Stateful):
+            self._base_dataset.load_state_dict(state_dict["base_dataset"])

d9d/dataset/padding.py

@@ -0,0 +1,79 @@
+from collections.abc import Sequence
+from enum import StrEnum
+
+import torch
+import torch.nn.functional as F
+
+
+class PaddingSide1D(StrEnum):
+    """
+    Enum specifying the side for padding 1D sequences.
+
+    Attributes:
+        left: Pad on the left side.
+        right: Pad on the right side.
+    """
+
+    left = "left"
+    right = "right"
+
+
+def _padding_side_1d_to_config(side: PaddingSide1D, difference: int) -> tuple[int, ...]:
+    match side:
+        case PaddingSide1D.left:
+            return difference, 0
+        case PaddingSide1D.right:
+            return 0, difference
+        case _:
+            raise ValueError("Unknown padding side")
+
+
+def pad_stack_1d(
+        items: Sequence[torch.Tensor],
+        pad_value: int,
+        padding_side: PaddingSide1D = PaddingSide1D.right,
+        pad_to_multiple_of: int | None = None
+) -> torch.Tensor:
+    """
+    Stacks 1D tensors into a batch, applying padding.
+
+    Calculates the maximum length among the input tensors (optionally aligning to a multiple),
+    pads elements to match this length on the specified side, and stacks them.
+
+    Args:
+        items: A sequence of 1D tensors to be stacked.
+        pad_value: The value used for padding.
+        padding_side: The side on which to apply padding (left or right).
+        pad_to_multiple_of: Optional integer. If provided, ensures the target length
+            is a multiple of this value.
+
+    Returns:
+        A single stacked tensor of shape (batch, max_length).
+
+    Raises:
+        ValueError: If no items are provided or if `pad_to_multiple_of` is <= 0.
+    """
+
+    if not items:
+        raise ValueError("Cannot stack 0 items")
+    if pad_to_multiple_of is not None and pad_to_multiple_of <= 0:
+        raise ValueError("pad_to_multiple_of should be > 0")
+
+    max_len = max(x.shape[0] for x in items)
+
+    if pad_to_multiple_of is not None and (remainder := max_len % pad_to_multiple_of) != 0:
+        max_len = max_len + (pad_to_multiple_of - remainder)
+
+    padded_items = []
+
+    for x in items:
+        difference = max_len - x.shape[0]
+
+        if difference == 0:
+            padded_items.append(x)
+        else:
+            padded_items.append(
+                F.pad(x, _padding_side_1d_to_config(padding_side, difference), value=pad_value)
+            )
+
+    return torch.stack(padded_items, dim=0)

d9d/dataset/sharded.py

@@ -0,0 +1,195 @@
+import math
+from collections.abc import Sized
+from enum import StrEnum
+from typing import Any, TypeVar
+
+from torch.distributed.checkpoint.stateful import Stateful
+from torch.utils.data import Dataset
+
+from d9d.core.dist_context import BATCH_DOMAIN, DistributedContext
+
+
+class ShardIndexingMode(StrEnum):
+    """
+    Defines how the dataset is split across shards.
+
+    Modes:
+        sequential: Round-robin distribution.
+
+            shard0: 0, 4, 8, 12
+            shard1: 1, 5, 9, 13
+            shard2: 2, 6, 10
+            shard3: 3, 7, 11
+
+        chunked: Contiguous blocks.
+
+            shard0: 0, 1, 2, 3
+            shard1: 4, 5, 6, 7
+            shard2: 8, 9, 10, 11
+            shard3: 12, 13
+    """
+
+    sequential = "sequential"
+    chunked = "chunked"
+
+
+_T_co = TypeVar("_T_co", covariant=True)
+
+
+class ShardedDataset(Dataset[_T_co], Stateful):
+    """
+    A dataset wrapper that acts as a view onto a specific shard of the underlying dataset.
+
+    This is useful for Data Parallel training where each process should only see
+    a subset of the data. It supports different indexing modes and optional padding
+    to ensure all shards have equal length (preventing hangs in distributed collectives).
+    """
+
+    def __init__(
+            self,
+            dataset: Dataset[_T_co],
+            total_shards: int,
+            current_shard: int,
+            indexing_mode: ShardIndexingMode,
+            pad_to_equal_size_across_shards: bool
+    ):
+        """
+        Constructs a ShardedDataset object.
+
+        Args:
+            dataset: The underlying dataset to shard.
+            total_shards: The total number of shards (e.g., number of DP ranks).
+            current_shard: The index of the current shard (e.g., current DP rank).
+            indexing_mode: How indices are assigned to shards (sequential/round-robin or chunked).
+            pad_to_equal_size_across_shards: If True, the length of the dataset will be padded
+                so that all shards report the same length. The last standard element is repeated.
+        """
+
+        if not isinstance(dataset, Sized):
+            raise ValueError("Dataset should implement __len__ method")
+
+        self._dataset = dataset
+
+        self._total_shards = total_shards
+        self._current_shard = current_shard
+
+        self._indexing_mode = indexing_mode
+        self._pad_to_equal_size_across_shards = pad_to_equal_size_across_shards
+
+    def _compute_real_index_sequential(self, index: int) -> int:
+        return index * self._total_shards + self._current_shard
+
+    def _get_base_index_unsafe(self, index: int) -> int:
+        """
+        Calculates the underlying dataset index for a given shard index,
+        without boundary checking.
+        """
+
+        match self._indexing_mode:
+            case ShardIndexingMode.sequential:
+                base_index = index * self._total_shards + self._current_shard
+
+                return base_index
+            case ShardIndexingMode.chunked:
+                ceil_len = math.ceil(len(self._dataset) / self._total_shards)
+                shard_start_offset = ceil_len * self._current_shard
+
+                return shard_start_offset + index
+            case _:
+                raise ValueError(f"Unknown shard indexing mode: {self._indexing_mode}")
+
+    def __getitem__(self, index: int) -> _T_co:
+        """
+        Retrieves an item from the underlying dataset mapping logic shard index to physical index.
+
+        If padding is enabled and the index exceeds the valid data for this shard,
+        the last item in the dataset is returned.
+
+        Args:
+            index: The index relative to this shard.
+
+        Returns:
+            The data item.
+        """
+
+        base_index = self._get_base_index_unsafe(index)
+        if base_index >= len(self._dataset):
+            base_index = len(self._dataset) - 1
+        return self._dataset[base_index]
+
+    def __len__(self) -> int:
+        """
+        Returns the number of items in this specific shard.
+
+        If `pad_to_equal_size_across_shards` is True, this returns the ceiling
+        length (max length across all shards).
+        """
+
+        ceil_len = math.ceil(len(self._dataset) / self._total_shards)
+
+        if self._pad_to_equal_size_across_shards:
+            return ceil_len
+
+        shards_remainder = len(self._dataset) % self._total_shards
+        match self._indexing_mode:
+            case ShardIndexingMode.sequential:
+                shards_full = len(self._dataset) // self._total_shards
+                return shards_full + 1 if self._current_shard < shards_remainder else shards_full
+            case ShardIndexingMode.chunked:
+                is_shard_last = self._current_shard == self._total_shards - 1
+                if not is_shard_last or shards_remainder == 0:
+                    return ceil_len
+                else:
+                    return ceil_len - (self._total_shards - shards_remainder)
+
+    def load_state_dict(self, state_dict: dict[str, Any]) -> None:
+        if isinstance(self._dataset, Stateful):
+            self._dataset.load_state_dict(state_dict["dataset"])
+
+        # check whether env mismatched
+        if state_dict["total_shards"] != self._total_shards:
+            raise ValueError("Shard count mismatch")
+        self._total_shards = state_dict["total_shards"]
+
+        self._current_shard = state_dict["current_shard"]
+
+    def state_dict(self) -> dict[str, Any]:
+        dct: dict[str, Any] = {
+            "total_shards": self._total_shards,
+            "current_shard": self._current_shard
+        }
+        if isinstance(self._dataset, Stateful):
+            dct["dataset"] = self._dataset.state_dict()
+        return dct
+
+
+def shard_dataset_data_parallel(
+        dataset: Dataset[_T_co],
+        dist_context: DistributedContext,
+        indexing_mode: ShardIndexingMode = ShardIndexingMode.sequential,
+        pad_to_equal_size_across_shards: bool = True
+) -> Dataset[_T_co]:
+    """
+    Wraps a dataset into a ShardedDataset based on the Data Parallel dimension of the distributed context.
+
+    This is a helper function to automatically determine the correct rank and world size
+    from the 'dp' (Data Parallel) mesh dimension within the batch domain DeviceMesh.
+
+    Args:
+        dataset: The source dataset to shard.
+        dist_context: The distributed context.
+        indexing_mode: The strategy for splitting data indices (sequential/round-robin or chunked).
+        pad_to_equal_size_across_shards: If True, ensures all shards have the same length by padding.
+
+    Returns:
+        A dataset instance representing the local shard.
+    """
+
+    dp_mesh = dist_context.mesh_for(BATCH_DOMAIN)["dp"]
+    return ShardedDataset(
+        dataset=dataset,
+        total_shards=dp_mesh.size(),
+        current_shard=dp_mesh.get_local_rank(),
+        indexing_mode=indexing_mode,
+        pad_to_equal_size_across_shards=pad_to_equal_size_across_shards
+    )

d9d/internals/determinism/__init__.py

@@ -0,0 +1,10 @@
+"""
+This package provides utilities for making your distributed setup deterministic.
+"""
+
+
+from .seed import set_seeds
+
+__all__ = [
+    "set_seeds"
+]

d9d/internals/determinism/seed.py

@@ -0,0 +1,63 @@
+import os
+import random
+from typing import cast
+
+import torch
+import torch.distributed.tensor
+
+from d9d.core.dist_context import REGULAR_DOMAIN, DistributedContext
+
+
+def set_seeds(
+    dist_context: DistributedContext,
+    seed: int,
+    distinct_seed_mesh_dim: str = "pp",
+) -> None:
+    """
+    Sets random seeds for Python, NumPy, and PyTorch.
+
+    This function sets seeds deterministically based on the provided base seed and the
+    process's rank within a specific mesh dimension.
+
+    The seed is shifted by the rank in the `distinct_seed_mesh_dim` (e.g., Pipeline Parallel rank).
+    This ensures that processes in different pipeline stages operate with different random states,
+    while processes that should share randomness (like Expert Parallel peers) can be synchronized.
+
+    Args:
+        dist_context: The distributed context.
+        seed: The base random seed.
+        distinct_seed_mesh_dim: The name of the mesh dimension along which seeds should
+            be distinct (e.g., 'pp' for pipeline parallelism). Ranks along other dimensions
+            will share the seed.
+    """
+
+    # Mutate seed based on PP rank if distributed
+    if dist_context.mesh_params.is_distributed:
+        distinct_mesh = dist_context.mesh_for(REGULAR_DOMAIN)[distinct_seed_mesh_dim]
+        seed = (seed + distinct_mesh.get_local_rank()) % 2**64
+
+    dist_context.logger.info(f"Set seed {seed}")
+
+    torch.manual_seed(seed)
+    os.environ["PYTHONHASHSEED"] = str(seed % 2**32)
+    random.seed(seed)
+
+    try:
+        import numpy as np  # noqa: PLC0415
+        np.random.seed(seed)
+    except ImportError:
+        pass
+
+    # Set DTensor seeding if distributed
+    if dist_context.mesh_params.is_distributed:
+        mesh_regular = dist_context.mesh_for(REGULAR_DOMAIN)
+        duplicate_seed_mesh_dim = tuple(
+            name
+            for name
+            in cast(list[str], mesh_regular.mesh_dim_names)
+            if name != distinct_seed_mesh_dim
+        )
+        duplicate_seed_mesh = mesh_regular[duplicate_seed_mesh_dim] if len(duplicate_seed_mesh_dim) != 0 else None
+
+        if duplicate_seed_mesh and duplicate_seed_mesh.get_coordinate() is not None:
+            torch.distributed.tensor._random.manual_seed(seed, duplicate_seed_mesh)  # noqa: SLF001

d9d/internals/grad_norm/__init__.py

@@ -0,0 +1,8 @@
+from .group import ParametersForNorm, group_parameters_for_norm
+from .norm import clip_grad_norm_distributed_
+
+__all__ = [
+    "ParametersForNorm",
+    "clip_grad_norm_distributed_",
+    "group_parameters_for_norm"
+]

d9d/internals/grad_norm/group.py

@@ -0,0 +1,87 @@
+import dataclasses
+from collections import defaultdict
+from collections.abc import Iterable
+from typing import Any
+
+import torch
+from torch import nn
+from torch.distributed import DeviceMesh
+from torch.distributed.tensor import DTensor, Shard
+
+
+@dataclasses.dataclass(kw_only=True, frozen=True)
+class GradNormGroup:
+    """
+    Defines a group of parameters that share the same distributed properties.
+
+    This grouping is used to batch gradient norm reductions efficiently. Parameters
+    sharing the same device mesh shards can be reduced in a single communication collective.
+
+    Attributes:
+        shard_meshes: A tuple of device meshes where the parameters are sharded, or None if replicated/local.
+        device: The device where parameters reside.
+        grad_dtype: The data type of the gradients.
+    """
+
+    shard_meshes: tuple[DeviceMesh, ...] | None
+    device: torch.device
+    grad_dtype: torch.dtype | None
+
+
+ParametersForNorm = dict[GradNormGroup, list[nn.Parameter]]
+
+
+def _extract_shard_meshes(param: nn.Parameter) -> tuple[DeviceMesh, ...] | None:
+    data = param.data
+
+    if not isinstance(data, DTensor):
+        return None
+
+    mesh = data.device_mesh
+    mesh_dim_names = mesh.mesh_dim_names
+    if mesh_dim_names is None:
+        raise ValueError("Only named meshes are supported.")
+
+    shard_placement_dim_names: list[str] = []
+
+    for dim_i, placement in enumerate(data.placements):
+        if isinstance(placement, Shard):
+            shard_placement_dim_names.append(mesh_dim_names[dim_i])
+
+    if len(shard_placement_dim_names) == 0:
+        return None
+
+    return tuple(mesh[name] for name in shard_placement_dim_names)
+
+
+def _group_sort_key(item: tuple[GradNormGroup, list[nn.Parameter]]) -> Any:
+    # put items WITH shard_meshes on top so they are processed first so we benefit from comm-comp overlap
+    return item[0].shard_meshes is None
+
+
+def group_parameters_for_norm(parameters: Iterable[nn.Parameter]) -> ParametersForNorm:
+    """
+    Groups parameters based on their distributed tensor characteristics.
+
+    Groups parameters by their sharding meshes, device, and gradient data type.
+
+    Args:
+        parameters: The iterable of parameters to group.
+
+    Returns:
+        A dictionary mapping synchronization groups to lists of parameters.
+    """
+
+    grouped_params: ParametersForNorm = defaultdict(list)
+    for param in parameters:
+        if not param.requires_grad:
+            continue
+
+        group = GradNormGroup(
+            shard_meshes=_extract_shard_meshes(param),
+            grad_dtype=param.grad_dtype,
+            device=param.device
+        )
+        grouped_params[group].append(param)
+    # we are sure dict is ordered in python 3.11 so we can sort it...
+    return dict(sorted(grouped_params.items(), key=_group_sort_key))