d9d 0.1.0__py3-none-any.whl

d9d/loop/component/data_loader_factory.py

@@ -0,0 +1,258 @@
+from collections.abc import Callable, Iterable, Iterator
+from typing import Any, Self, TypedDict, Unpack
+
+import torch
+import torch.utils._pytree as pytree  # noqa: PLC2701
+from torch.utils.data import Dataset, Sampler
+from torchdata.stateful_dataloader import StatefulDataLoader
+
+from d9d.core.dist_context import BATCH_DOMAIN, DistributedContext
+from d9d.core.types import CollateFn, PyTree
+from d9d.loop.config import DataLoadingConfig
+from d9d.loop.control import DatasetProvider, InitializeDatasetContext
+
+from .batch_maths import BatchMaths
+
+
+class DataLoaderKwargs(TypedDict, total=False):
+    """
+    Type definition for arguments accepted by the PyTorch DataLoader.
+    """
+
+    batch_size: int | None
+    shuffle: bool | None
+    sampler: Sampler | Iterable | None
+    batch_sampler: Sampler[list] | Iterable[list] | None
+    num_workers: int
+    collate_fn: CollateFn
+    pin_memory: bool
+    drop_last: bool
+    timeout: float
+    worker_init_fn: Callable | None
+    multiprocessing_context: Any
+    generator: Any
+    prefetch_factor: int | None
+    persistent_workers: bool
+    pin_memory_device: str
+
+
+def _move_to_device(data: PyTree, device: torch.types.Device) -> PyTree:
+    return pytree.tree_map(lambda x: x.to(device), data)
+
+
+class IteratorBatchGroup(Iterator):
+    """
+    An iterator that groups items from a base iterator into sub-streams.
+
+    This class is utilized for gradient accumulation where
+        a single optimizer step consumes multiple micro-batches (the group).
+
+    It also moves the data to the specified device immediately upon access.
+    """
+
+    def __init__(
+            self,
+            base: Iterator,
+            device: torch.types.Device,
+            batch_group_size: int
+    ):
+        """
+        Constructs an IteratorBatchGroup object.
+
+        Args:
+            base: The underlying data iterator (usually from a DataLoader).
+            device: The target device to move tensors to.
+            batch_group_size: The number of micro-batches to yield within one group.
+        """
+
+        self._base = base
+        self._device = device
+
+        self._batch_group_size = batch_group_size
+
+        self._is_end = False
+
+    def __next__(self) -> PyTree:
+        """
+        Advances the iterator.
+
+        Returns:
+            A generator that yields `batch_group_size` items (micro-batches),
+            with each item already moved to the configured device.
+
+        Raises:
+            StopIteration: If the underlying iterator is exhausted.
+        """
+
+        if self._is_end:
+            raise StopIteration()
+
+        try:
+            sample_item = next(self._base)
+        except StopIteration:
+            self._is_end = True
+            raise StopIteration() from None
+
+        def _iter_inside_group():
+            yield _move_to_device(sample_item, self._device)
+
+            for _ in range(self._batch_group_size - 1):
+                try:
+                    item = next(self._base)
+                    yield _move_to_device(item, self._device)
+                except StopIteration:
+                    self._is_end = True
+                    break
+
+        return _iter_inside_group()
+
+    def __iter__(self) -> Self:
+        """Returns self."""
+
+        return self
+
+
+class StatefulDataLoaderDataParallelAware(StatefulDataLoader):
+    """
+    A stateful data loader that is aware of data parallel ranks.
+
+    This loader extends the standard torchdata StatefulDataLoader to ensure
+    that checkpoints are saved with rank-specific keys.
+
+    It also wraps the iterator to support batch grouping for gradient accumulation and
+    automatically transfer data to bound device.
+    """
+
+    def __init__(
+            self,
+            dataset: Dataset,
+            dp_rank: int,
+            device: torch.types.Device,
+            group_size: int,
+            **kwargs: Unpack[DataLoaderKwargs]
+    ):
+        """
+        Constructs a StatefulDataLoaderDataParallelAware object.
+
+        Args:
+            dataset: The dataset to load from.
+            dp_rank: The Data Parallel rank of the current process (used for state checkpointing).
+            device: The device to move data to.
+            group_size: The number of batches to group together (e.g., for gradient accumulation).
+            **kwargs: Standard arguments passed to the parent DataLoader.
+        """
+
+        super().__init__(dataset, **kwargs)
+        self._dp_rank = dp_rank
+        self._device = device
+        self._group_size = group_size
+
+    def state_dict(self) -> dict[str, Any]:
+        return {
+            f"dp_{self._dp_rank}": super().state_dict()
+        }
+
+    def load_state_dict(self, state_dict: dict[str, Any]) -> None:
+        super().load_state_dict(state_dict[f"dp_{self._dp_rank}"])
+
+    def __iter__(self) -> Iterator:
+        return IteratorBatchGroup(
+            super().__iter__(),
+            device=self._device,
+            batch_group_size=self._group_size
+        )
+
+
+class DataLoaderFactory:
+    """
+    Factory class for creating configured DataLoaders.
+
+    This class centralizes the creation logic for training and inference
+    data loaders, applying configurations.
+    """
+
+    def __init__(
+            self,
+            dist_context: DistributedContext,
+            provider: DatasetProvider,
+            config_data_loading: DataLoadingConfig,
+            batch_maths: BatchMaths
+    ):
+        """
+        Constructs a DataLoaderFactory object.
+
+        Args:
+            dist_context: The distributed context containing mesh and rank information.
+            provider: The provider callable that initializes the dataset and collator.
+            config_data_loading: Specific configuration for data loading.
+            batch_maths: BatchMaths object.
+        """
+
+        self._dist_context = dist_context
+        self._provider = provider
+
+        self._config_data_loading = config_data_loading
+
+        self._batch_maths = batch_maths
+
+    def _build_dataloader(
+            self,
+            provider: DatasetProvider,
+            batch_size: int,
+            group_size: int,
+            drop_last: bool
+    ) -> StatefulDataLoader:
+        result = provider(InitializeDatasetContext(
+            dist_context=self._dist_context,
+            batch_maths=self._batch_maths
+        ))
+
+        return StatefulDataLoaderDataParallelAware(
+            result.dataset,
+            collate_fn=result.collator,
+            group_size=group_size,
+            num_workers=self._config_data_loading.num_workers,
+            persistent_workers=self._config_data_loading.persistent_workers,
+            pin_memory=self._config_data_loading.pin_memory,
+            batch_size=batch_size,
+            dp_rank=self._dist_context.mesh_for(BATCH_DOMAIN)["dp"].size(),
+            device="cuda",
+            drop_last=drop_last
+        )
+
+    def build_dataloader_for_train_job(self) -> StatefulDataLoader:
+        """
+        Builds and returns a StatefulDataLoader configured for training.
+
+        This loader is configured to drop the last incomplete batch and group
+        batches according to the gradient accumulation settings defined in
+        BatchMaths.
+
+        Returns:
+            A configured StatefulDataLoader instance.
+        """
+
+        return self._build_dataloader(
+            self._provider,
+            batch_size=self._batch_maths.data_loader_batch_size,
+            group_size=self._batch_maths.num_microbatches_gradient_accumulation,
+            drop_last=True
+        )
+
+    def build_dataloader_for_infer_job(self) -> StatefulDataLoader:
+        """
+        Builds and returns a StatefulDataLoader configured for inference.
+
+        This loader processes batches one by one (group size of 1) and does
+        not drop the last batch.
+
+        Returns:
+            A configured StatefulDataLoader instance.
+        """
+
+        return self._build_dataloader(
+            self._provider,
+            batch_size=self._batch_maths.data_loader_batch_size,
+            group_size=1,
+            drop_last=False
+        )

d9d/loop/component/garbage_collector.py

@@ -0,0 +1,94 @@
+import gc
+import time
+from contextlib import AbstractContextManager
+from types import TracebackType
+from typing import Self
+
+from d9d.core.dist_context import DistributedContext
+from d9d.loop.config import GarbageCollectionConfig
+
+from .stepper import Stepper
+
+
+class ManualGarbageCollector(AbstractContextManager):
+    """
+    Manages efficient Python garbage collection during the training loop.
+
+    This context manager disables automatic garbage collection upon entry to prevent
+    unpredictable latency spikes during training steps. It allows performing
+    manual collection at specific intervals (periodic) or specific points (forced).
+    """
+
+    def __init__(
+            self,
+            dist_ctx: DistributedContext,
+            config: GarbageCollectionConfig,
+            step: Stepper
+    ):
+        """
+        Constructs the garbage collector manager.
+
+        Args:
+            dist_ctx: The distributed context.
+            config: Configuration determining how often GC should run.
+            step: Stepper instance used to track the current training step.
+        """
+        self._dist_ctx = dist_ctx
+        self._config = config
+        self._step = step
+
+    def __enter__(self) -> Self:
+        """
+        Disables automatic garbage collection and performs an initial full collection.
+
+        Returns:
+            The calling instance.
+        """
+
+        gc.disable()
+        self._collect(generation=2)
+
+        return self
+
+    def __exit__(
+            self,
+            exc_type: type[BaseException] | None,
+            exc_value: BaseException | None,
+            traceback: TracebackType | None, /
+    ) -> None:
+        """
+        Re-enables automatic garbage collection and performs a final full collection.
+
+        Args:
+            exc_type: The type of the exception raised (if any).
+            exc_value: The exception instance raised (if any).
+            traceback: The traceback object (if any).
+        """
+
+        gc.enable()
+        self._collect(generation=2)
+
+    def collect_periodic(self):
+        """
+        Triggers garbage collection if the current step matches the configured period.
+
+        This typically performs a faster (generation 1) collection rather than a full sweep.
+        """
+
+        if self._step.should_do_action(self._config.period_steps, enable_on_last_step_if_periodic=False):
+            self._collect(generation=1)
+
+    def collect_forced(self):
+        """
+        Forces a full garbage collection run regardless of the step count.
+
+        This performs a generation 2 collection.
+        """
+
+        self._collect(generation=2)
+
+    def _collect(self, generation: int):
+        begin = time.monotonic()
+        gc.collect(generation)
+        end = time.monotonic()
+        self._dist_ctx.logger.info(f"[GC] Garbage collection for generation {generation} took {end - begin}s")

d9d/loop/component/gradient_clipper.py

@@ -0,0 +1,89 @@
+from contextlib import contextmanager
+
+from d9d.core.dist_context import REGULAR_DOMAIN, DistributedContext
+from d9d.internals.grad_norm import ParametersForNorm, clip_grad_norm_distributed_, group_parameters_for_norm
+from d9d.loop.config import GradientClippingConfig
+from d9d.tracker import BaseTrackerRun
+
+from .model_stage_factory import TrackedModules
+from .stepper import Stepper
+
+
+class GradientClipper:
+    """
+    Manages gradient clipping and logging of gradient norms in a distributed execution environment.
+    """
+
+    def __init__(
+            self,
+            dist_context: DistributedContext,
+            tracked_modules: TrackedModules,
+            config: GradientClippingConfig,
+            stepper: Stepper
+    ):
+        """
+        Constructs the gradient clipper.
+
+        Args:
+            dist_context: The distributed context.
+            tracked_modules: Container of model modules whose parameters need clipping.
+            config: Configuration defining max norm and logging frequency.
+            stepper: Stepper instance used to track the current training step.
+        """
+
+        self._dist_context = dist_context
+        self._tracked_modules = tracked_modules
+        self._config = config
+        self._stepper = stepper
+
+        self._parameter_groups: ParametersForNorm | None = None
+
+    def _all_parameters(self):
+        for model in self._tracked_modules.modules:
+            yield from model.parameters()
+
+    @contextmanager
+    def install(self):
+        """
+        Context manager that prepares and groups parameters for efficient norm calculation.
+
+        It calculates necessary metadata (such as segregating shared parameters) to ensure
+        correct global norm calculation across the pipeline parallel mesh.
+        """
+
+        self._parameter_groups = group_parameters_for_norm(self._all_parameters())
+        yield
+        self._parameter_groups = None
+
+    def clip_and_log(self, run: BaseTrackerRun):
+        """
+        Clips gradients to the configured maximum norm and logs the total L2 norm.
+
+        This method performs an in-place modification of parameter gradients if a
+        maximum norm is configured. It calculates the global gradient norm across
+        distributed ranks.
+
+        Args:
+            run: The tracker run instance used for logging the norm scalar.
+
+        Raises:
+            ValueError: If called outside the ``install`` context manager scope.
+        """
+
+        should_log = self._stepper.should_do_action(self._config.log_total_steps)
+
+        if not self._config.max_norm and not should_log:
+            return
+
+        if self._parameter_groups is None:
+            raise ValueError("Parameter groups are not configured")
+
+        grad_norm = clip_grad_norm_distributed_(
+            parameter_groups=self._parameter_groups,
+            max_norm=self._config.max_norm,
+            norm_type=2.0,
+            pp_mesh=self._dist_context.mesh_for(REGULAR_DOMAIN)["pp"],
+        )
+
+        if should_log:
+            run.scalar(name="l2_grad_norm_total", value=grad_norm.item())

d9d/loop/component/gradient_manager.py

@@ -0,0 +1,149 @@
+from contextlib import contextmanager
+
+import torch
+from torch.distributed.tensor import DTensor
+
+from d9d.core.dist_context import DistributedContext
+from d9d.internals.grad_sync import GradientSynchronizer
+from d9d.loop.config import GradientManagerConfig
+from d9d.metric.impl import WeightedMeanMetric
+
+from .batch_maths import BatchMaths
+from .model_stage_factory import TrackedModules
+
+
+class GradientManager:
+    """
+    Manages the lifecycle of gradients during the training loop.
+
+    This class handles gradient synchronization across ranks,
+    gradient data type configuration, and loss scaling based on accumulated weights.
+    It orchestrates the `GradientSynchronizer` and ensures gradients are correctly
+    prepared before the optimizer step.
+    """
+
+    def __init__(
+            self,
+            dist_context: DistributedContext,
+            tracked_modules: TrackedModules,
+            batch_maths: BatchMaths,
+            config: GradientManagerConfig
+    ):
+        """
+        Constructs the GradientManager and initializes the internal synchronizer.
+
+        Args:
+            dist_context: The distributed context.
+            tracked_modules: Container of model modules to manage gradients for.
+            batch_maths: Calculation utility for batch sizes and accumulation steps.
+            config: Configuration for gradient handling.
+        """
+
+        self._dist_context = dist_context
+        self._tracked_modules = tracked_modules
+        self._batch_maths = batch_maths
+        self._config = config
+        self._loss = WeightedMeanMetric()
+        self._loss.to("cuda")
+
+        self._grad_sync = GradientSynchronizer(
+            [list(module.parameters()) for module in self._tracked_modules.modules],
+            bucket_size_mb=self._config.bucket_size_mb,
+            require_accumulations=self._batch_maths.num_backward_calls
+        )
+        self._grads_to_scale: list[torch.Tensor] | None = None
+
+    def _setup_grad_dtype(self):
+        if self._config.grad_dtype is None:
+            return
+
+        for mod in self._tracked_modules.modules:
+            for param in mod.parameters():
+                if param.requires_grad:
+                    param.grad_dtype = getattr(torch, self._config.grad_dtype)
+
+    def _bind_grads_to_scale(self):
+        grads_to_scale: list[torch.Tensor] = []
+
+        for mod in self._tracked_modules.modules:
+            for param in mod.parameters():
+                if param.grad is None:
+                    continue
+                grad = param.grad.to_local() if isinstance(param.grad, DTensor) else param.grad
+                grads_to_scale.append(grad)
+
+        self._grads_to_scale = grads_to_scale
+
+    def _unbind_grads_to_scale(self):
+        self._grads_to_scale = None
+
+    def _scale_grads(self):
+        scale_factor = 1.0 / self._loss.accumulated_weight
+        torch._foreach_mul_(self._grads_to_scale, scale_factor)
+
+    @contextmanager
+    def install(self):
+        """
+        Context manager to activate gradient handling for a forward/backward pass.
+
+        This sets up gradient dtypes, install backward hooks for synchronization via
+        the `GradientSynchronizer`, and binds gradients for later scaling. It acts
+        as the boundary for the accumulation phase.
+        """
+
+        self._setup_grad_dtype()
+        self._grad_sync.bind()
+        self._bind_grads_to_scale()
+        yield
+        self._unbind_grads_to_scale()
+        self._grad_sync.unbind()
+
+    def add_loss_with_weight(self, loss: torch.Tensor, loss_weight: torch.Tensor):
+        """
+        Accumulates a loss value and its corresponding weight into the internal metric.
+
+        Args:
+            loss: The computed loss scalar.
+            loss_weight: The weight asscociated with this loss.
+        """
+
+        self._loss.update(loss, loss_weight)
+
+    def sync_and_scale(self):
+        """
+        Finalizes gradients to be ready for the optimizer step.
+
+        This method performs the following operations:
+
+        1. Waits for all gradient synchronization hooks to complete.
+        2. Synchronizes the accumulated loss/weights across the distributed context.
+        3. Scales the gradients by the inverse of the total accumulated weight to
+           normalize them.
+        """
+
+        self._grad_sync.wait()
+
+        self._loss.trigger_sync(self._dist_context)
+        self._loss.wait_sync(self._dist_context)
+        self._scale_grads()
+
+    def compute_global_loss(self) -> torch.Tensor:
+        """
+        Calculates the final weighted mean loss.
+
+        Returns:
+            The averaged loss scalar across all accumulation steps and ranks.
+        """
+
+        return self._loss.compute()
+
+    def zero_grad(self):
+        """
+        Resets the internal state for the next training step.
+
+        This clears the accumulated gradients in the synchronizer and resets the
+        loss metrics.
+        """
+
+        self._grad_sync.zero_grad()
+        self._loss.reset()