d9d 0.1.0__py3-none-any.whl

d9d/internals/grad_norm/norm.py

@@ -0,0 +1,169 @@
+import math
+
+import torch.distributed as dist
+import torch.nn.utils
+from torch import nn
+from torch.autograd.profiler import record_function
+from torch.distributed import DeviceMesh
+from torch.distributed.tensor import DTensor
+
+from d9d.internals.grad_norm.group import ParametersForNorm
+
+
+def _reduce_op_from_norm_type(norm_type: float) -> dist.ReduceOp.RedOpType:
+    if math.isinf(norm_type):
+        return dist.ReduceOp.MAX
+    else:
+        return dist.ReduceOp.SUM
+
+
+def _parameter_to_local_grad(parameter: nn.Parameter) -> torch.Tensor:
+    grad = parameter.grad
+
+    if grad is None:
+        raise ValueError("None grad detected")
+
+    if isinstance(grad, DTensor):
+        return grad.to_local()
+    else:
+        return grad
+
+
+def _get_local_norm_pow(
+        parameters: list[nn.Parameter],
+        norm_type: float
+) -> torch.Tensor:
+    # calculates for local
+
+    if len(parameters) == 0:
+        return torch.tensor(0.0, device="cuda")
+
+    norm_val = torch.nn.utils.get_total_norm(
+        [_parameter_to_local_grad(x) for x in parameters],
+        norm_type=norm_type,
+        foreach=True,
+        error_if_nonfinite=False
+    )
+
+    if math.isinf(norm_type):
+        return norm_val
+    else:
+        return norm_val ** norm_type
+
+
+def _get_global_norm_pow_horizontal(
+        parameter_groups: ParametersForNorm,
+        norm_type: float
+) -> torch.Tensor:
+    # calculates for horizontal parallelism
+    if len(parameter_groups) == 0:
+        return torch.tensor(0.0, device="cuda")
+
+    norms: list[torch.Tensor] = []
+    works: list[dist.Work] = []
+    for group, group_params in parameter_groups.items():
+        local_norm_pow = _get_local_norm_pow(group_params, norm_type=norm_type)
+        if group.shard_meshes is not None:
+            if len(group.shard_meshes) != 1:
+                raise ValueError(
+                    "Currently we do not support calculating norm for tensors that are sharded on multiple dims - feel "
+                    "free to file an issue if you need it."
+                )
+            process_group = group.shard_meshes[0].get_group()
+            work = dist.all_reduce(
+                local_norm_pow,
+                op=_reduce_op_from_norm_type(norm_type),
+                group=process_group,
+                async_op=True
+            )
+            works.append(work)
+        norms.append(local_norm_pow)
+
+    for work in works:
+        work.wait()
+
+    norms_total = torch.stack(norms, dim=0)
+
+    if math.isinf(norm_type):
+        return norms_total.max()
+    else:
+        return norms_total.sum()
+
+
+def _get_global_norm_pow_pp(
+        parameter_groups: ParametersForNorm,
+        norm_type: float,
+        pp_mesh: DeviceMesh | None
+) -> torch.Tensor:
+    norm = _get_global_norm_pow_horizontal(
+        parameter_groups=parameter_groups,
+        norm_type=norm_type
+    )
+    if pp_mesh is not None:
+        dist.all_reduce(norm, op=_reduce_op_from_norm_type(norm_type), group=pp_mesh.get_group())
+    return norm
+
+
+def _clip_grad_with_norm_(
+        parameter_groups: ParametersForNorm,
+        max_norm: float,
+        total_norm: torch.Tensor
+):
+    clip_coef = max_norm / (total_norm + 1e-6)
+    clip_coef_clamped = torch.clamp(clip_coef, max=1.0)
+
+    for group in parameter_groups.values():
+        grads = [_parameter_to_local_grad(x) for x in group]
+        torch._foreach_mul_(grads, clip_coef_clamped)
+
+
+def clip_grad_norm_distributed_(
+        parameter_groups: ParametersForNorm,
+        max_norm: float | None,
+        norm_type: float,
+        pp_mesh: DeviceMesh | None
+) -> torch.Tensor:
+    """
+    Clips gradient norms in a fully distributed environment.
+
+    This function calculates the global gradient norm across all dimensions of parallelism
+    (Horizontal - DP/CP/TP/EP/..., and Pipeline) and scales the gradients in-place to ensure the norm
+    does not exceed max_norm.
+
+    It accurately handles DTensors by identifying their sharding placements and performing
+    reductions only on the necessary process groups.
+
+    Overlaps communication and computation if possible.
+
+    Args:
+        parameter_groups: Dictionary grouping parameters by synchronization requirements,
+            typically created by `group_parameters_for_norm`.
+        max_norm: The maximum allowed norm of the gradients. If None, the function
+            calculates and returns the global norm without modifying the gradients.
+        norm_type: The type of the norm to calculate (e.g., 2.0 for L2 norm, inf for max norm).
+        pp_mesh: The device mesh representing the pipeline parallel dimension, needed
+            to reduce norms across pipeline stages.
+
+    Returns:
+        The calculated global gradient norm.
+    """
+
+    with record_function("Gradient Clipping"):
+        global_norm_pow = _get_global_norm_pow_pp(
+            parameter_groups=parameter_groups,
+            norm_type=norm_type,
+            pp_mesh=pp_mesh
+        )
+        if math.isinf(norm_type):
+            global_norm = global_norm_pow
+        else:
+            global_norm = global_norm_pow ** (1.0 / norm_type)
+
+        if max_norm:
+            _clip_grad_with_norm_(
+                parameter_groups,
+                max_norm=max_norm,
+                total_norm=global_norm
+            )
+
+        return global_norm

d9d/internals/grad_sync/__init__.py

@@ -0,0 +1,14 @@
+"""
+Gradient synchronization utilities.
+
+This package provides the infrastructure for manual gradient bucketing and
+asynchronous reduction, similar to DistributedDataParallel but exposed
+for internal framework usage with DTensors.
+"""
+
+
+from .synchronizer import GradientSynchronizer
+
+__all__ = [
+    "GradientSynchronizer"
+]

d9d/internals/grad_sync/bucket.py

@@ -0,0 +1,317 @@
+import abc
+from typing import cast
+
+import torch
+import torch.distributed as dist
+from torch import Tensor, nn
+from torch.autograd.profiler import record_function
+from torch.distributed import DeviceMesh
+from torch.distributed.tensor import DTensor
+from torch.utils.hooks import RemovableHandle
+
+from .placement_helper import dist_grad_from_local
+
+
+class AbstractGradientBucket(abc.ABC):
+    """
+    Interface for a bucket containing a subset of model parameters.
+
+    A bucket manages the memory layout and synchronization lifecycle of the
+    gradients associated with its parameters.
+    """
+
+    @abc.abstractmethod
+    def bind(self):
+        """
+        Initializes the bucket state.
+
+        This involves allocating contiguous memory buffers (if applicable),
+        registering backward hooks, and preparing the gradients for accumulation.
+        """
+
+    @abc.abstractmethod
+    def unbind(self):
+        """
+        Cleans up the bucket state.
+
+        Removes hooks, deallocates buffers, and detaches gradients.
+        """
+
+    @abc.abstractmethod
+    def zero_grad(self):
+        """
+        Zeros out the gradients and resets accumulation counters.
+        """
+
+    @abc.abstractmethod
+    def mark_sync(self):
+        """
+        Marks this bucket as synchronized.
+        """
+
+
+class LocalGradientBucket(AbstractGradientBucket):
+    """
+    A bucket for parameters that do not require distributed synchronization.
+    """
+
+    def __init__(self, params: list[nn.Parameter]):
+        """
+        Constructs a LocalGradientBucket.
+
+        Args:
+            params: List of parameters to manage.
+        """
+
+        self._params = params
+
+    def bind(self):
+        """
+        No-op for local buckets as they do not require special buffering.
+        """
+
+    def unbind(self):
+        """
+        No-op for local buckets.
+        """
+
+    def wait(self):
+        """
+        No-op as no async communication is performed.
+        """
+
+    @torch.no_grad()
+    def zero_grad(self):
+        """
+        Directly zeros the grad attribute of the parameters.
+        """
+
+        for param in self._params:
+            param.grad = None
+
+    def mark_sync(self):
+        """
+        No-op for local buckets.
+        """
+
+
+class AccumulationCounter:
+    """
+    Tracks the number of gradient accumulation steps for a set of parameters.
+    """
+
+    def __init__(self, require_accumulations: int, parameters: list[nn.Parameter]):
+        """
+        Constructs an AccumulationCounter.
+
+        Args:
+            require_accumulations: Number of accumulations required before sync.
+            parameters: List of parameters to track.
+        """
+
+        self._require_accumulations = require_accumulations
+        self._param_to_sync_count = {param: 0 for param in parameters}
+
+    def reset(self):
+        """
+        Resets all counters to zero.
+        """
+
+        self._param_to_sync_count = {param: 0 for param in self._param_to_sync_count}
+
+    def update(self, param: nn.Parameter):
+        """
+        Increments the counter for a specific parameter.
+
+        Args:
+            param: The parameter that finished a backward step.
+        """
+
+        self._param_to_sync_count[param] += 1
+
+    def is_ready(self) -> bool:
+        """
+        Checks if all parameters have reached the required number of accumulations.
+
+        Returns:
+            True if synchronization can proceed.
+        """
+
+        return all(x == self._require_accumulations for x in self._param_to_sync_count.values())
+
+
+class SyncGradientBucket(AbstractGradientBucket):
+    """
+    A bucket that manages a contiguous memory buffer for gradients and performs async reduction.
+
+    This bucket flattens the gradients of its parameters into a single contiguous
+    Tensor to enable efficient batched all-reduce operations.
+    """
+
+    def __init__(
+            self,
+            parameters: list[nn.Parameter],
+            require_accumulations: int,
+            device: torch.device,
+            grad_dtype: torch.dtype,
+            reduce_mesh: DeviceMesh,
+            communicate_stream: torch.cuda.Stream
+    ):
+        """
+        Constructs a SyncGradientBucket.
+
+        Args:
+            parameters: List of parameters to manage.
+            require_accumulations: Number of accumulations before triggering reduce.
+            device: Device where parameters reside.
+            grad_dtype: Data type for the gradients.
+            reduce_mesh: DeviceMesh on which reduction happens.
+            communicate_stream: Stream where all the asynchronous communications will be scheduled
+        """
+
+        if not all(isinstance(x.data, DTensor) for x in parameters):
+            raise ValueError("All parameters passed in synchronizable bucket should contain DTensor data")
+
+        self._params = parameters
+        self._accum_counter = AccumulationCounter(require_accumulations, parameters)
+        self._device = device
+        self._grad_dtype = grad_dtype
+        # iterate from innermost to outermost group
+        self._reduce_groups: list[dist.ProcessGroup] = reduce_mesh.get_all_groups()[::-1]
+
+        self._buffer: Tensor | None = None
+        self._hooks: list[RemovableHandle] | None = None
+
+        self._communicate_stream = communicate_stream
+        self._ready_to_sync = False
+
+    def _bind_buffer(self):
+        """
+        Allocates the flat buffer and redirects parameter gradients to view into it.
+        """
+
+        buffer_size = sum(cast(DTensor, param.data).to_local().numel() for param in self._params)
+
+        self._buffer = torch.zeros(
+            (buffer_size,),
+            dtype=self._grad_dtype,
+            device=self._device
+        )
+
+        offset = 0
+
+        for param in self._params:
+            data = cast(DTensor, param.data)
+            local_param = data.to_local()
+
+            local_grad = self._buffer[offset:offset + local_param.numel()].view(local_param.shape)
+
+            param.grad = dist_grad_from_local(data, local_grad)
+
+            offset += local_param.numel()
+
+    @torch.no_grad()
+    def _post_accumulation_hook(self, param: nn.Parameter):
+        """
+        Hook executed after backward pass for a parameter.
+
+        Updates the accumulation counter and triggers the asynchronous all-reduce
+        if the bucket is ready.
+
+        Args:
+            param: The parameter that finished backward pass.
+        """
+
+        self._accum_counter.update(param)
+
+        if not self._accum_counter.is_ready():
+            return
+
+        if self._ready_to_sync:
+            raise ValueError("Tried to accumulate, but synchronization was not performed")
+
+        with record_function("Gradient Sync"):
+            # wait for backward operation is complete
+            self._communicate_stream.wait_stream(torch.cuda.current_stream())
+            # execute all sync operations in sequential order (to ensure
+            # data safety), but in a DIFFERENT stream
+            with torch.cuda.stream(self._communicate_stream):
+                for group in self._reduce_groups:
+                    dist.all_reduce(
+                        self._buffer,
+                        op=dist.ReduceOp.SUM,
+                        group=group
+                    )
+            self._ready_to_sync = True
+
+    def _bind_hooks(self):
+        """
+        Registers post-accumulate hooks on all parameters.
+        """
+
+        hooks = []
+        for param in self._params:
+            hooks.append(param.register_post_accumulate_grad_hook(self._post_accumulation_hook))
+        self._hooks = hooks
+
+    @torch.no_grad()
+    def bind(self):
+        """
+        Allocates the contiguous buffer and registers hooks.
+        """
+
+        self._bind_buffer()
+        self._bind_hooks()
+
+    def _unbind_buffer(self):
+        """
+        Deallocates the buffer and clears parameter gradients.
+        """
+
+        self._buffer = None
+
+        for param in self._params:
+            param.grad = None
+
+    def _unbind_hooks(self):
+        """
+        Removes all registered hooks.
+        """
+
+        if self._hooks is None:
+            return
+
+        for hook in self._hooks:
+            hook.remove()
+        self._hooks = None
+
+    @torch.no_grad()
+    def unbind(self):
+        """
+        Cleans up buffer and hooks.
+        """
+
+        self._unbind_buffer()
+        self._unbind_hooks()
+
+    @torch.no_grad()
+    def zero_grad(self):
+        """
+        Zeros the contiguous buffer, resets counters, and marks params as awaiting sync.
+
+        Raises:
+            ValueError: If the buffer is not initialized (call bind first).
+        """
+
+        buffer = self._buffer
+        if buffer is None:
+            raise ValueError("Buffer is not initialized")
+
+        buffer.zero_()
+        self._accum_counter.reset()
+
+    def mark_sync(self):
+        if not self._ready_to_sync:
+            raise ValueError("This bucket is not ready for sync.")
+
+        self._ready_to_sync = False

d9d/internals/grad_sync/placement_helper.py

@@ -0,0 +1,23 @@
+from torch import Tensor
+from torch.distributed.tensor import DTensor
+
+
+def dist_grad_from_local(data: DTensor, local_grad: Tensor) -> DTensor:
+    """
+    Constructs a DTensor gradient from a local tensor using data placement info.
+
+    Args:
+        data: The original parameter DTensor (source of metadata).
+        local_grad: The local tensor containing gradient data.
+
+    Returns:
+        A new DTensor wrapping the local gradient.
+    """
+
+    return DTensor.from_local(
+        local_grad,
+        shape=data.shape,
+        stride=data.stride(),
+        device_mesh=data.device_mesh,
+        placements=data.placements
+    )