Initial import: grid-bot — grid trading bot for BTC-USDT on Cifra Markets
This commit is contained in:
Kolp
@@ -0,0 +1,96 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates
|
||||
|
||||
import torch
|
||||
import torch.distributed.tensor._ops # force import all built-in dtensor ops
|
||||
from torch.distributed.device_mesh import DeviceMesh, init_device_mesh # noqa: F401
|
||||
from torch.distributed.tensor._api import (
|
||||
distribute_module,
|
||||
distribute_tensor,
|
||||
DTensor,
|
||||
empty,
|
||||
full,
|
||||
ones,
|
||||
rand,
|
||||
randn,
|
||||
zeros,
|
||||
)
|
||||
from torch.distributed.tensor.placement_types import (
|
||||
_StridedShard,
|
||||
Partial,
|
||||
Placement,
|
||||
Replicate,
|
||||
Shard,
|
||||
)
|
||||
from torch.optim.optimizer import (
|
||||
_foreach_supported_types as _optim_foreach_supported_types,
|
||||
)
|
||||
from torch.utils._foreach_utils import (
|
||||
_foreach_supported_types as _util_foreach_supported_types,
|
||||
)
|
||||
|
||||
|
||||
# All public APIs from dtensor package
|
||||
__all__ = [
|
||||
"DTensor",
|
||||
"distribute_tensor",
|
||||
"distribute_module",
|
||||
"Shard",
|
||||
"Replicate",
|
||||
"Partial",
|
||||
"Placement",
|
||||
"ones",
|
||||
"empty",
|
||||
"full",
|
||||
"rand",
|
||||
"randn",
|
||||
"zeros",
|
||||
]
|
||||
|
||||
# For weights_only torch.load
|
||||
from ._dtensor_spec import (
|
||||
DTensorSpec as _DTensorSpec,
|
||||
ShardOrderEntry as _ShardOrderEntry,
|
||||
TensorMeta as _TensorMeta,
|
||||
)
|
||||
|
||||
|
||||
torch.serialization.add_safe_globals(
|
||||
[
|
||||
DeviceMesh,
|
||||
_DTensorSpec,
|
||||
_TensorMeta,
|
||||
_ShardOrderEntry,
|
||||
DTensor,
|
||||
Partial,
|
||||
Replicate,
|
||||
Shard,
|
||||
_StridedShard,
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
# Append DTensor to the list of supported types for foreach implementation for optimizer
|
||||
# and clip_grad_norm_ so that we will try to use foreach over the for-loop implementation on CUDA.
|
||||
if DTensor not in _optim_foreach_supported_types:
|
||||
_optim_foreach_supported_types.append(DTensor)
|
||||
|
||||
if DTensor not in _util_foreach_supported_types:
|
||||
_util_foreach_supported_types.append(DTensor) # type: ignore[arg-type]
|
||||
|
||||
|
||||
# Set namespace for exposed private names
|
||||
DTensor.__module__ = "torch.distributed.tensor"
|
||||
distribute_tensor.__module__ = "torch.distributed.tensor"
|
||||
distribute_module.__module__ = "torch.distributed.tensor"
|
||||
ones.__module__ = "torch.distributed.tensor"
|
||||
empty.__module__ = "torch.distributed.tensor"
|
||||
full.__module__ = "torch.distributed.tensor"
|
||||
rand.__module__ = "torch.distributed.tensor"
|
||||
randn.__module__ = "torch.distributed.tensor"
|
||||
zeros.__module__ = "torch.distributed.tensor"
|
||||
|
||||
# Register DTensor dispatch for higher order operators
|
||||
from torch._higher_order_ops.print import _register_dtensor_impl
|
||||
|
||||
|
||||
_register_dtensor_impl()
|
||||
File diff suppressed because it is too large
Load Diff
+566
@@ -0,0 +1,566 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import logging
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
from functools import lru_cache
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
import torch.distributed._functional_collectives as funcol
|
||||
import torch.distributed.tensor._dtensor_spec as dtensor_spec
|
||||
from torch._C._distributed_c10d import _resolve_process_group
|
||||
from torch._logging import warning_once
|
||||
from torch.distributed._functional_collectives import _are_we_tracing
|
||||
from torch.distributed._local_tensor import (
|
||||
local_tensor_mode,
|
||||
maybe_run_for_local_tensor,
|
||||
)
|
||||
from torch.distributed.device_mesh import _mesh_resources, DeviceMesh
|
||||
from torch.distributed.distributed_c10d import (
|
||||
broadcast,
|
||||
get_group_rank,
|
||||
get_rank,
|
||||
GroupName,
|
||||
ProcessGroup,
|
||||
scatter,
|
||||
Work,
|
||||
)
|
||||
from torch.fx.experimental.symbolic_shapes import guard_or_false
|
||||
from torch.types import IntLikeType
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Opaque types must be registered before defining schemas that reference them,
|
||||
# so the schema parser recognizes the type names and uses PyObjectType (which
|
||||
# wraps Python objects as ConcretePyObjectHolder) instead of AnyType (which
|
||||
# calls toTypeInferredIValue and fails for Python-only opaque types).
|
||||
from torch.distributed.device_mesh import _register_distributed_opaque_types
|
||||
|
||||
|
||||
_register_distributed_opaque_types()
|
||||
|
||||
_dtensor_lib = torch.library.Library("_dtensor", "FRAGMENT")
|
||||
_dtensor_lib.define(
|
||||
"mesh_get_process_group("
|
||||
"torch.distributed.device_mesh.DeviceMesh mesh, int dim"
|
||||
") -> torch.distributed.distributed_c10d.ProcessGroup"
|
||||
)
|
||||
|
||||
|
||||
@torch.library.impl("_dtensor::mesh_get_process_group", "CompositeExplicitAutograd")
|
||||
def _mesh_get_process_group_impl(mesh, dim):
|
||||
return mesh.get_group(dim)
|
||||
|
||||
|
||||
@torch.library.register_fake("_dtensor::mesh_get_process_group")
|
||||
def _mesh_get_process_group_fake(mesh, dim):
|
||||
from torch._library.fake_class_registry import maybe_unwrap_fake_script_object
|
||||
|
||||
real_mesh = maybe_unwrap_fake_script_object(mesh)
|
||||
return real_mesh.get_group(dim)
|
||||
|
||||
|
||||
@torch.library.register_fake("_dtensor::shard_dim_alltoall")
|
||||
def _shard_dim_alltoall_meta(
|
||||
input, gather_dim, shard_dim, group_name: GroupName | ProcessGroup
|
||||
):
|
||||
if isinstance(group_name, str):
|
||||
# pyrefly: ignore[bad-argument-type] # pyrefly bug
|
||||
group_name = _resolve_process_group(group_name)
|
||||
group_size = group_name.size()
|
||||
stacked_list = [torch.empty_like(input) for _ in range(group_size)]
|
||||
group_rank = get_group_rank(group_name, get_rank())
|
||||
|
||||
cat_tensor = torch.cat(stacked_list, dim=gather_dim)
|
||||
# pyrefly: ignore [unsupported-operation]
|
||||
chunk_size = cat_tensor.size(shard_dim) // group_size
|
||||
chunk = torch.narrow(cat_tensor, shard_dim, group_rank * chunk_size, chunk_size)
|
||||
return chunk.contiguous()
|
||||
|
||||
|
||||
def shard_dim_alltoall(input, gather_dim, shard_dim, mesh, mesh_dim):
|
||||
if mesh.device_type == "cpu" and local_tensor_mode() is None:
|
||||
# Gloo does not support alltoall, so falling back to allgather + chunk
|
||||
warning_once(
|
||||
logger,
|
||||
"CPU process group does not support alltoall yet, falling back with allgather + chunk!",
|
||||
)
|
||||
out = funcol.all_gather_tensor(input, gather_dim, (mesh, mesh_dim))
|
||||
if isinstance(out, funcol.AsyncCollectiveTensor):
|
||||
# stick to the same behavior for the alltoall case, remove this once we enable alltoall async
|
||||
out = out.wait()
|
||||
from torch.distributed.tensor.placement_types import Shard
|
||||
|
||||
out = Shard._custom_chunk(out, mesh.size(mesh_dim), dim=shard_dim)[
|
||||
mesh.get_local_rank(mesh_dim)
|
||||
]
|
||||
return out.contiguous()
|
||||
|
||||
group = funcol._resolve_group((mesh, mesh_dim))
|
||||
# TODO: enable async op for shard_dim_alltoall
|
||||
return torch.ops._dtensor.shard_dim_alltoall(
|
||||
input, gather_dim, shard_dim, funcol._group_or_group_name(group)
|
||||
)
|
||||
|
||||
|
||||
def mesh_scatter(
|
||||
output: torch.Tensor,
|
||||
scatter_list: list[torch.Tensor],
|
||||
mesh: DeviceMesh,
|
||||
mesh_dim: int = 0,
|
||||
async_op: bool = False,
|
||||
*,
|
||||
group_src: int = 0,
|
||||
) -> Work | None:
|
||||
"""
|
||||
scatter a list of tensors to a device mesh dimension. We by default
|
||||
use the first rank of the mesh dimension as the source of truth, i.e
|
||||
for a 2d mesh [[0, 1], [2, 3]], if we scatter on mesh_dim = 1, we will
|
||||
scatter the tensor list on rank 0 to rank 0/1, and tensor list on rank
|
||||
2 to rank 2/3.
|
||||
|
||||
Args:
|
||||
output (torch.Tensor): the tensor to receive the scattered list.
|
||||
scatter_list (List[torch.Tensor]): the tensor list to be scattered.
|
||||
mesh_dim (int, optional): indicate which mesh dimension we want
|
||||
to scatter on, we by default choose the first rank on the
|
||||
mesh dimension as source of truth.
|
||||
|
||||
Keyword args:
|
||||
group_src (int, optional): the group rank of the source data for the
|
||||
logical/global tensor, on the specific mesh dimension. By default, we
|
||||
use ``group_rank=0`` on each DeviceMesh dimension as the source data
|
||||
to preserve the single-device semantic. If passing ``None`` explicitly,
|
||||
this method simply uses its local data with no communication.
|
||||
|
||||
Returns:
|
||||
A :class:`Work` object
|
||||
"""
|
||||
# TODO: Ideally we should use the meta tensor way
|
||||
# (to register a meta kernel for the collective op)
|
||||
# so that it would avoid the communication. Need to
|
||||
# remove the check below once that is done.
|
||||
if output.is_meta:
|
||||
return None
|
||||
dim_group = mesh.get_group(mesh_dim)
|
||||
if not isinstance(dim_group, ProcessGroup):
|
||||
raise AssertionError
|
||||
|
||||
if group_src == get_rank(dim_group):
|
||||
fut = scatter(
|
||||
output,
|
||||
scatter_list=scatter_list,
|
||||
group=dim_group,
|
||||
async_op=async_op,
|
||||
group_src=group_src,
|
||||
)
|
||||
else:
|
||||
fut = scatter(
|
||||
output,
|
||||
scatter_list=None,
|
||||
group=dim_group,
|
||||
async_op=async_op,
|
||||
group_src=group_src,
|
||||
)
|
||||
|
||||
return fut
|
||||
|
||||
|
||||
def mesh_broadcast(
|
||||
tensor: torch.Tensor,
|
||||
mesh: DeviceMesh,
|
||||
mesh_dim: int = 0,
|
||||
async_op: bool = False,
|
||||
*,
|
||||
group_src: int = 0,
|
||||
) -> Work | None:
|
||||
"""
|
||||
broadcast the tensor to a device mesh dimension. We by default
|
||||
use the first rank of the mesh dimension as the source of truth, i.e
|
||||
for a 2d mesh [[0, 1], [2, 3]], if we broadcast on mesh_dim = 1, we will
|
||||
broadcast the tensor on rank 0 to rank 0/1, and tensor on rank 2
|
||||
to rank 2/3.
|
||||
|
||||
Args:
|
||||
tensor (torch.Tensor): tensor to broadcast.
|
||||
mesh_dim (int, optional): indicate which mesh dimension we want
|
||||
to scatter on, we by default choose the first rank on the
|
||||
mesh dimension as source of truth.
|
||||
|
||||
Keyword args:
|
||||
group_src (int, optional): the group rank of the source data for the
|
||||
logical/global tensor, on the specific mesh dimension. By default, we
|
||||
use ``group_rank=0`` on each DeviceMesh dimension as the source data
|
||||
to preserve the single-device semantic. If passing ``None`` explicitly,
|
||||
this method simply uses its local data with no communication.
|
||||
|
||||
Returns:
|
||||
A :class:`Work` object
|
||||
"""
|
||||
# TODO: Ideally we should use the meta tensor way
|
||||
# (to register a meta kernel for the collective op)
|
||||
# so that it would avoid the communication. Need to
|
||||
# remove the check below once that is done.
|
||||
if tensor.is_meta:
|
||||
return None
|
||||
dim_group = mesh.get_group(mesh_dim)
|
||||
if not isinstance(dim_group, ProcessGroup):
|
||||
raise AssertionError
|
||||
|
||||
return broadcast(tensor, group=dim_group, async_op=async_op, group_src=group_src)
|
||||
|
||||
|
||||
@maybe_run_for_local_tensor
|
||||
def pad_tensor(
|
||||
tensor: torch.Tensor, pad_dim: int, pad_size: IntLikeType
|
||||
) -> torch.Tensor:
|
||||
# During tracing, always emit the pad op even when pad_size=0 so all
|
||||
# ranks produce identical FX graph structure (SPMD).
|
||||
# guard_or_false returns False for symbolic sizes, so the pad is always
|
||||
# emitted during tracing. In eager with concrete pad_size=0, it returns
|
||||
# True and we skip the no-op pad.
|
||||
if guard_or_false(pad_size == 0) and not _are_we_tracing():
|
||||
return tensor
|
||||
pad = [0, 0] * (tensor.ndim - pad_dim)
|
||||
pad[-1] = pad_size # pyrefly: ignore[unsupported-operation]
|
||||
return torch.nn.functional.pad(tensor, pad)
|
||||
|
||||
|
||||
@maybe_run_for_local_tensor
|
||||
def unpad_tensor(
|
||||
tensor: torch.Tensor, pad_dim: int, pad_size: IntLikeType
|
||||
) -> torch.Tensor:
|
||||
# During tracing, always emit the narrow op even when pad_size=0 so all
|
||||
# ranks produce identical FX graph structure (SPMD).
|
||||
if guard_or_false(pad_size == 0) and not _are_we_tracing():
|
||||
return tensor
|
||||
return tensor.narrow(
|
||||
pad_dim,
|
||||
start=0,
|
||||
length=tensor.size(pad_dim) - pad_size,
|
||||
)
|
||||
|
||||
|
||||
def fill_empty_tensor_to_shards(
|
||||
shards: list[torch.Tensor], shard_dim: int, num_empty_tensors: int
|
||||
) -> list[torch.Tensor]:
|
||||
if num_empty_tensors == 0:
|
||||
return shards
|
||||
tensor_size = list(shards[0].size())
|
||||
tensor_size[shard_dim] = 0
|
||||
tensor = shards[0].new_zeros(tensor_size)
|
||||
shards.extend(tensor for _ in range(num_empty_tensors))
|
||||
return shards
|
||||
|
||||
|
||||
def check_tensor_meta(
|
||||
local_tensor, check_shape_stride=False
|
||||
) -> Optional["dtensor_spec.TensorMeta"]:
|
||||
local_metadata = {
|
||||
"dtype": local_tensor.dtype,
|
||||
"requires_grad": local_tensor.requires_grad,
|
||||
}
|
||||
|
||||
if check_shape_stride:
|
||||
local_metadata.update(
|
||||
{"shape": local_tensor.shape, "stride": local_tensor.stride()}
|
||||
)
|
||||
|
||||
gathered_metadata = [None for _ in range(torch.distributed.get_world_size())]
|
||||
torch.distributed.all_gather_object(gathered_metadata, local_metadata)
|
||||
|
||||
# Check if metadata is consistent across ranks
|
||||
if not all(meta == local_metadata for meta in gathered_metadata):
|
||||
raise ValueError(
|
||||
"Inconsistent tensor metadata (including shape and stride) across ranks."
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
def spec_to_bytes(spec: "dtensor_spec.DTensorSpec") -> int:
|
||||
if spec.tensor_meta is None:
|
||||
raise AssertionError("spec should have tensor meta defined!")
|
||||
return spec.tensor_meta.dtype.itemsize * math.prod(spec.shape)
|
||||
|
||||
|
||||
@dataclass
|
||||
class MeshTopoInfo:
|
||||
"""
|
||||
Mesh information for collective cost estimation
|
||||
"""
|
||||
|
||||
mesh: DeviceMesh
|
||||
mesh_dim_devices: list[int]
|
||||
mesh_dim_bandwidth: list[float]
|
||||
mesh_dim_latency: list[float]
|
||||
|
||||
@staticmethod
|
||||
@lru_cache(None)
|
||||
def build_from_mesh(mesh: DeviceMesh) -> "MeshTopoInfo":
|
||||
# Generate mesh topology info for intra-host/inter-host communication pattern
|
||||
# Note that we made bunch of assumptions for simplicity:
|
||||
# 1. we assume the mesh is homogeneous, and it's gpu/nccl model
|
||||
# 2. we assume gpu arch is Ampere or Hopper
|
||||
# 3. we assume collectives are all ring base algo for now
|
||||
num_devices_per_host = _mesh_resources.num_devices_per_host(mesh.device_type)
|
||||
# the base bw number (intra-node), GB/s
|
||||
base_bw = 87.7
|
||||
mesh_dim_bandwidth = [base_bw] * mesh.ndim
|
||||
# the latency in terms of us (intra-node, nv-link)
|
||||
mesh_dim_latency = [0.6] * mesh.ndim
|
||||
mesh_dim_devices = [1] * mesh.ndim
|
||||
|
||||
total_num_devices = 1
|
||||
for mesh_dim in reversed(range(mesh.ndim)):
|
||||
num_devices = mesh.size(mesh_dim)
|
||||
mesh_dim_devices[mesh_dim] = num_devices
|
||||
total_num_devices *= num_devices
|
||||
if total_num_devices > num_devices_per_host:
|
||||
# magic number for inter-host communication bandwidth/latency factor
|
||||
# This number assumes latest GPU arch, i.e. Ampere or Hopper
|
||||
# TODO: see if we need to tweak this or offer a way for user
|
||||
# to specify the bandwidths/latency
|
||||
mesh_dim_bandwidth[mesh_dim] *= 0.22
|
||||
# set to ethernet latency for inter-host
|
||||
mesh_dim_latency[mesh_dim] = 2.7
|
||||
|
||||
return MeshTopoInfo(
|
||||
mesh, mesh_dim_devices, mesh_dim_bandwidth, mesh_dim_latency
|
||||
)
|
||||
|
||||
|
||||
def allgather_cost(bytes_gb: float, mesh_topo: MeshTopoInfo, mesh_dim: int) -> float:
|
||||
num_devices_on_mesh_dim = mesh_topo.mesh_dim_devices[mesh_dim]
|
||||
mesh_dim_bandwidth = mesh_topo.mesh_dim_bandwidth[mesh_dim]
|
||||
num_hops = num_devices_on_mesh_dim - 1
|
||||
# base latency + comm latency
|
||||
latency = 6.6 + num_hops * mesh_topo.mesh_dim_latency[mesh_dim] # us
|
||||
bw = (bytes_gb * num_hops / num_devices_on_mesh_dim) / mesh_dim_bandwidth # s
|
||||
return latency + bw * 1e6 # rescale to us
|
||||
|
||||
|
||||
def allreduce_cost(bytes_gb: float, mesh_topo: MeshTopoInfo, mesh_dim: int) -> float:
|
||||
num_devices_on_mesh_dim = mesh_topo.mesh_dim_devices[mesh_dim]
|
||||
mesh_dim_bandwidth = mesh_topo.mesh_dim_bandwidth[mesh_dim]
|
||||
# allreduce have almost 2x comm bytes compare to allgather/reduce_scatter
|
||||
num_hops = 2 * (num_devices_on_mesh_dim - 1)
|
||||
|
||||
latency = 6.6 + num_hops * mesh_topo.mesh_dim_latency[mesh_dim]
|
||||
bw = (bytes_gb * num_hops / num_devices_on_mesh_dim) / mesh_dim_bandwidth
|
||||
return latency + bw * 1e6
|
||||
|
||||
|
||||
def reduce_scatter_cost(
|
||||
bytes_gb: float,
|
||||
mesh_topo: MeshTopoInfo,
|
||||
mesh_dim: int,
|
||||
) -> float:
|
||||
num_devices_on_mesh_dim = mesh_topo.mesh_dim_devices[mesh_dim]
|
||||
mesh_dim_bandwidth = mesh_topo.mesh_dim_bandwidth[mesh_dim]
|
||||
num_hops = num_devices_on_mesh_dim - 1
|
||||
# base latency + comm latency
|
||||
latency = 6.6 + num_hops * mesh_topo.mesh_dim_latency[mesh_dim]
|
||||
bw = (bytes_gb * num_hops / num_devices_on_mesh_dim) / mesh_dim_bandwidth
|
||||
return latency + bw * 1e6
|
||||
|
||||
|
||||
def _compute_placement_transition_cost(
|
||||
current_placement: "dtensor_spec.Placement",
|
||||
target_placement: "dtensor_spec.Placement",
|
||||
mesh_topo: MeshTopoInfo,
|
||||
mesh_dim: int,
|
||||
comm_bytes_gb: float,
|
||||
) -> tuple[float, float]:
|
||||
"""
|
||||
Compute the cost of transitioning from one placement to another on a single mesh dimension.
|
||||
|
||||
Args:
|
||||
current_placement: The current placement on the mesh dimension.
|
||||
target_placement: The target placement on the mesh dimension.
|
||||
mesh_topo: Mesh topology information for cost estimation.
|
||||
mesh_dim: The mesh dimension where the transition happens.
|
||||
comm_bytes_gb: The communication bytes in GB for this step.
|
||||
|
||||
Returns:
|
||||
A tuple of (cost, updated_comm_bytes_gb):
|
||||
- cost: The communication cost for this transition (float("inf") if invalid).
|
||||
- updated_comm_bytes_gb: The updated communication bytes after this step.
|
||||
"""
|
||||
if current_placement == target_placement:
|
||||
return 0.0, comm_bytes_gb
|
||||
|
||||
num_devices_on_mesh_dim = mesh_topo.mesh_dim_devices[mesh_dim]
|
||||
|
||||
# NOTE: is_shard() does not match _StridedShard; see _is_shard_like().
|
||||
# Safe today: redistribute_cost bails with inf when shard_order is None.
|
||||
if current_placement.is_shard() and target_placement.is_replicate():
|
||||
# allgather gives larger comm bytes
|
||||
comm_bytes_gb *= num_devices_on_mesh_dim
|
||||
return allgather_cost(comm_bytes_gb, mesh_topo, mesh_dim), comm_bytes_gb
|
||||
elif current_placement.is_shard() and target_placement.is_shard():
|
||||
# should be alltoall comm, since we haven't implement it yet, add 1.0 as penalty
|
||||
# to favor allgather instead
|
||||
# TODO: add alltoall_cost
|
||||
return allgather_cost(comm_bytes_gb, mesh_topo, mesh_dim) + 1.0, comm_bytes_gb
|
||||
elif current_placement.is_partial() and target_placement.is_replicate():
|
||||
return allreduce_cost(comm_bytes_gb, mesh_topo, mesh_dim), comm_bytes_gb
|
||||
elif current_placement.is_partial() and target_placement.is_shard():
|
||||
cost = reduce_scatter_cost(comm_bytes_gb, mesh_topo, mesh_dim)
|
||||
# after reduce_scatter the comm bytes for further collectives halved.
|
||||
comm_bytes_gb /= num_devices_on_mesh_dim
|
||||
return cost, comm_bytes_gb
|
||||
elif current_placement.is_shard() and target_placement.is_partial():
|
||||
# ban shard -> partial as it does not make sense to perform
|
||||
# this redistribute
|
||||
return float("inf"), comm_bytes_gb
|
||||
elif current_placement.is_partial() and target_placement.is_partial():
|
||||
# we already handled the == case at the top, and we ban converting between partial types.
|
||||
return float("inf"), comm_bytes_gb
|
||||
elif current_placement.is_replicate() and target_placement.is_shard():
|
||||
comm_bytes_gb /= num_devices_on_mesh_dim
|
||||
return 0.0, comm_bytes_gb
|
||||
|
||||
return 0.0, comm_bytes_gb
|
||||
|
||||
|
||||
def one_step_redistribute_cost(
|
||||
current_spec: "dtensor_spec.DTensorSpec",
|
||||
target_spec: "dtensor_spec.DTensorSpec",
|
||||
) -> float:
|
||||
"""
|
||||
Calculate the cost of a single redistribution step between two DTensorSpecs.
|
||||
|
||||
This function computes the communication cost for a one-step redistribution
|
||||
where the current and target specs differ by exactly one placement on one
|
||||
mesh dimension.
|
||||
|
||||
Args:
|
||||
current_spec: The current DTensorSpec.
|
||||
target_spec: The target DTensorSpec.
|
||||
|
||||
Returns:
|
||||
The communication cost for this step (float("inf") if invalid).
|
||||
"""
|
||||
if current_spec.mesh != target_spec.mesh:
|
||||
return float("inf")
|
||||
|
||||
if current_spec.placements == target_spec.placements:
|
||||
return 0.0
|
||||
|
||||
# Find the mesh dimension that differs
|
||||
mesh_dim = -1
|
||||
current_placement = None
|
||||
target_placement = None
|
||||
for i, (cur, tgt) in enumerate(
|
||||
zip(current_spec.placements, target_spec.placements)
|
||||
):
|
||||
if cur != tgt:
|
||||
if mesh_dim != -1:
|
||||
# More than one dimension differs - not a single step
|
||||
raise ValueError(
|
||||
"one_step_redistribute_cost expects specs that differ by exactly one placement"
|
||||
)
|
||||
mesh_dim = i
|
||||
current_placement = cur
|
||||
target_placement = tgt
|
||||
|
||||
if mesh_dim == -1:
|
||||
return 0.0
|
||||
|
||||
if current_placement is None or target_placement is None:
|
||||
raise AssertionError
|
||||
|
||||
mesh_topo = MeshTopoInfo.build_from_mesh(current_spec.mesh)
|
||||
comm_bytes_gb = (
|
||||
spec_to_bytes(current_spec) / current_spec.num_shards / 1024 / 1024 / 1024
|
||||
)
|
||||
|
||||
cost, _ = _compute_placement_transition_cost(
|
||||
current_placement, target_placement, mesh_topo, mesh_dim, comm_bytes_gb
|
||||
)
|
||||
return cost
|
||||
|
||||
|
||||
def redistribute_cost(
|
||||
current_spec: "dtensor_spec.DTensorSpec",
|
||||
target_spec: "dtensor_spec.DTensorSpec",
|
||||
) -> float:
|
||||
"""
|
||||
This function returns the cost of redistribute from current to target DTensorSpec.
|
||||
|
||||
NOTE:
|
||||
1. Only consider communication cost here, since computation costs for redistribute
|
||||
are quite trivial (i.e. we only need to narrow or simple division)
|
||||
2. Only consider redistribute cost on same mesh, cross mesh communication cost is
|
||||
not quite needed for operator strategy estimation/selection.
|
||||
"""
|
||||
if current_spec.mesh != target_spec.mesh:
|
||||
# make infinite cost if meshes are not same
|
||||
# TODO: see if we want to support this once there's cross mesh communication
|
||||
return float("inf")
|
||||
if current_spec.is_replicated():
|
||||
# short-cut: comm cost is 0 if current spec is already full replication
|
||||
return 0.0
|
||||
|
||||
# TODO(zpcore): test placements with _StridedShard if we replace shard_order
|
||||
# with _StridedShard.
|
||||
if (
|
||||
current_spec.placements == target_spec.placements
|
||||
and current_spec.shard_order == target_spec.shard_order
|
||||
):
|
||||
return 0.0
|
||||
|
||||
# For sub-meshes, ranks not participating in the mesh should not compute
|
||||
# redistribution costs. Return 0 since they won't actually participate.
|
||||
if not current_spec.mesh._is_current_rank_part_of_mesh():
|
||||
return 0.0
|
||||
|
||||
mesh_topo = MeshTopoInfo.build_from_mesh(current_spec.mesh)
|
||||
cost = 0.0
|
||||
comm_bytes_gb = (
|
||||
spec_to_bytes(current_spec) / current_spec.num_shards / 1024 / 1024 / 1024
|
||||
)
|
||||
# Transformation that considered for redistribute cost:
|
||||
# 1. allgather 2. alltoall
|
||||
# 3. allreduce 4. reduce_scatter
|
||||
from torch.distributed._functional_collectives import _are_we_tracing
|
||||
from torch.distributed.tensor._redistribute import (
|
||||
_gen_transform_infos,
|
||||
_gen_transform_infos_non_cached,
|
||||
)
|
||||
|
||||
# TODO(zpcore): Support _StridedShard redistribution. Remove the temporary
|
||||
# fix, which is to prevent StridedShard erroring out.
|
||||
if current_spec.shard_order is None or target_spec.shard_order is None:
|
||||
return float("inf")
|
||||
|
||||
# No redistribution needed when placements are already identical.
|
||||
# This also prevents potential failures in _gen_transform_infos for certain configurations
|
||||
# (e.g., sub-meshes) where finding a transform path between identical states may error out.
|
||||
# TODO(zpcore): test placements with _StridedShard if we replace shard_order
|
||||
# with _StridedShard.
|
||||
if (
|
||||
current_spec.placements == target_spec.placements
|
||||
and current_spec.shard_order == target_spec.shard_order
|
||||
):
|
||||
return cost
|
||||
|
||||
if _are_we_tracing():
|
||||
transform_infos = _gen_transform_infos_non_cached(current_spec, target_spec)
|
||||
else:
|
||||
transform_infos = _gen_transform_infos(current_spec, target_spec)
|
||||
for transform_info in transform_infos:
|
||||
if current_spec.tensor_meta is None:
|
||||
raise AssertionError("spec should have tensor meta defined!")
|
||||
current = transform_info.src_dst_placements[0]
|
||||
target = transform_info.src_dst_placements[1]
|
||||
mesh_dim = transform_info.mesh_dim
|
||||
step_cost, comm_bytes_gb = _compute_placement_transition_cost(
|
||||
current, target, mesh_topo, mesh_dim, comm_bytes_gb
|
||||
)
|
||||
if step_cost == float("inf"):
|
||||
return float("inf")
|
||||
cost += step_cost
|
||||
|
||||
return cost
|
||||
@@ -0,0 +1,330 @@
|
||||
# mypy: allow-untyped-defs
|
||||
"""
|
||||
Decomposition-based sharding propagation for DTensor.
|
||||
|
||||
When an operator doesn't have a registered sharding strategy, we derive one by
|
||||
tracing through its decomposition. The decomposed ops (which do have strategies)
|
||||
determine how placements propagate through the original op.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import itertools
|
||||
from typing import Any, TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
from torch._decomp import decomposition_table
|
||||
from torch.distributed._functional_collectives import _are_we_tracing
|
||||
from torch.distributed.device_mesh import DeviceMesh
|
||||
from torch.distributed.tensor._dtensor_spec import DTensorSpec
|
||||
from torch.distributed.tensor._op_schema import OpSchema, OpStrategy, RuntimeSchemaInfo
|
||||
from torch.distributed.tensor._utils import try_find_mesh_from_args
|
||||
from torch.distributed.tensor.placement_types import (
|
||||
_StridedShard,
|
||||
Placement,
|
||||
Replicate,
|
||||
Shard,
|
||||
)
|
||||
from torch.fx.experimental.symbolic_shapes import GuardOnDataDependentSymNode
|
||||
from torch.utils._python_dispatch import TorchDispatchMode
|
||||
|
||||
|
||||
def _infer_schema_info_from_op(op: OpOverload) -> RuntimeSchemaInfo:
|
||||
"""Infer RuntimeSchemaInfo from an operator's schema for decomposition ops"""
|
||||
schema = op._schema
|
||||
|
||||
# Find first non-tensor positional arg index
|
||||
static_argnum = None
|
||||
for i, arg in enumerate(schema.arguments):
|
||||
if arg.kwarg_only:
|
||||
break
|
||||
if arg.type.kind() != "TensorType" and static_argnum is None:
|
||||
static_argnum = i
|
||||
break
|
||||
|
||||
# Find keyword-only args that aren't tensors
|
||||
kwarg_only_names = []
|
||||
for arg in schema.arguments:
|
||||
if arg.kwarg_only and arg.type.kind() != "TensorType":
|
||||
kwarg_only_names.append(arg.name)
|
||||
|
||||
kwargs = {}
|
||||
if static_argnum is not None:
|
||||
kwargs["static_argnum"] = static_argnum
|
||||
if kwarg_only_names:
|
||||
# pyrefly: ignore [unsupported-operation]
|
||||
kwargs["static_kwargkey"] = kwarg_only_names
|
||||
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
return RuntimeSchemaInfo(**kwargs)
|
||||
|
||||
|
||||
from torch.utils._pytree import tree_any, tree_flatten, tree_map, tree_map_only
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from torch._ops import OpOverload
|
||||
from torch.distributed.tensor._sharding_prop import ShardingPropagator
|
||||
|
||||
|
||||
def _extract_input_specs(op_schema: OpSchema) -> tuple[DTensorSpec | object, ...]:
|
||||
return op_schema.args_schema + tuple(op_schema.kwargs_schema.values())
|
||||
|
||||
|
||||
class PlacementTrackingMode(TorchDispatchMode):
|
||||
"""
|
||||
TorchDispatchMode that tracks DTensor placements through op execution.
|
||||
|
||||
Used during decomposition tracing: intercepts each op, propagates sharding
|
||||
via the ShardingPropagator, and records output placements on the result tensors.
|
||||
"""
|
||||
|
||||
def __init__(self, sharding_prop: ShardingPropagator, mesh: DeviceMesh):
|
||||
super().__init__()
|
||||
self.sharding_prop = sharding_prop
|
||||
self.mesh = mesh
|
||||
|
||||
def __torch_dispatch__(self, func, types, args=(), kwargs=None):
|
||||
args_schema, kwargs_schema = tree_map(
|
||||
lambda x: getattr(x, "_spec", x) if isinstance(x, torch.Tensor) else x,
|
||||
(args, kwargs or {}),
|
||||
)
|
||||
|
||||
if not tree_any(
|
||||
lambda x: isinstance(x, DTensorSpec), (args_schema, kwargs_schema)
|
||||
):
|
||||
raise NotImplementedError(f"No DTensorSpec found in args/kwargs for {func}")
|
||||
|
||||
# Set schema_info so the LRU cache key includes static args
|
||||
op_schema = OpSchema(func, args_schema, kwargs_schema)
|
||||
schema_info = self.sharding_prop.op_to_schema_info.get(func)
|
||||
if schema_info is None:
|
||||
schema_info = (
|
||||
self.sharding_prop.op_to_schema_info_for_single_dim_strategy.get(func)
|
||||
)
|
||||
if schema_info is not None:
|
||||
op_schema.schema_info = schema_info
|
||||
op_schema._recompute_comparison_key()
|
||||
|
||||
if _are_we_tracing():
|
||||
output_sharding = self.sharding_prop.propagate_op_sharding_non_cached(
|
||||
op_schema
|
||||
)
|
||||
else:
|
||||
output_sharding = self.sharding_prop.propagate_op_sharding(op_schema)
|
||||
|
||||
if (
|
||||
output_sharding.needs_redistribute # pyrefly: ignore [missing-attribute]
|
||||
and (
|
||||
redistribute_schema
|
||||
:= output_sharding.redistribute_schema # pyrefly: ignore [missing-attribute]
|
||||
)
|
||||
is not None
|
||||
):
|
||||
# a pure .needs_redistribute check is too broad; we want to ban redistribution,
|
||||
# but this flag is set for view ops that convert global shape -> local shape args.
|
||||
# During decomposition tracing on meta tensors at global shape, the shape adjustment
|
||||
# is irrelevant — only reject true redistribution.
|
||||
for orig, desired in zip(
|
||||
op_schema.args_spec,
|
||||
redistribute_schema.args_spec, # pyrefly: ignore [missing-attribute]
|
||||
):
|
||||
if orig.placements != desired.placements:
|
||||
raise RuntimeError(
|
||||
f"Decomposition requires redistribution for {func}"
|
||||
)
|
||||
|
||||
out = func(*args, **kwargs)
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
self._record_output_specs(out, output_sharding.output_spec)
|
||||
return out
|
||||
|
||||
def _record_output_specs(self, output: Any, output_spec: DTensorSpec | Any) -> None:
|
||||
if isinstance(output, torch.Tensor) and output_spec is not None:
|
||||
output._spec = output_spec # pyrefly: ignore [missing-attribute]
|
||||
elif isinstance(output, (tuple, list)) and isinstance(
|
||||
output_spec, (tuple, list)
|
||||
):
|
||||
for t, s in zip(output, output_spec):
|
||||
self._record_output_specs(t, s)
|
||||
|
||||
|
||||
class DecompShardingStrategy:
|
||||
"""
|
||||
Generates sharding strategies for ops by tracing through their decompositions.
|
||||
|
||||
For each candidate input placement combination, runs the decomposition on meta
|
||||
tensors under PlacementTrackingMode to determine the output placement. These
|
||||
single-dimension strategies are then expanded to the full mesh.
|
||||
"""
|
||||
|
||||
def __init__(self, sharding_prop: ShardingPropagator):
|
||||
self.sharding_prop = sharding_prop
|
||||
# Cache fake meshes per device type to avoid repeated allocation.
|
||||
# A fake size-1 mesh ensures identical strategy computation across all ranks
|
||||
# during decomposition tracing, avoiding potential SPMD divergence.
|
||||
# False negatives are avoided (all sizes % 1 == 0), while false positives
|
||||
# are caught on expansion to the real, multi-dim device mesh.
|
||||
self._fake_meshes: dict[str, DeviceMesh] = {}
|
||||
|
||||
def _get_fake_mesh(self, device_type: str) -> DeviceMesh:
|
||||
fake_mesh = self._fake_meshes.get(device_type)
|
||||
if fake_mesh is None:
|
||||
fake_mesh = DeviceMesh(device_type, [0], _init_backend=False, _rank=0)
|
||||
self._fake_meshes[device_type] = fake_mesh
|
||||
return fake_mesh
|
||||
|
||||
@staticmethod
|
||||
def has_decomp(op: OpOverload) -> bool:
|
||||
# Check if op has a decomposition (explicit or CIA)
|
||||
return op in decomposition_table or op._can_decompose()
|
||||
|
||||
def ensure_schema_info(self, op: OpOverload) -> None:
|
||||
"""
|
||||
Register schema_info for decomposition op on first invocation.
|
||||
Needed for correct shard prop cache key.
|
||||
"""
|
||||
if op not in self.sharding_prop.op_to_schema_info:
|
||||
schema_info = _infer_schema_info_from_op(op)
|
||||
self.sharding_prop.op_to_schema_info[op] = schema_info
|
||||
|
||||
def propagate_strategy(
|
||||
self,
|
||||
op_schema: OpSchema,
|
||||
) -> OpStrategy | None:
|
||||
if not tree_any(
|
||||
lambda x: isinstance(x, DTensorSpec),
|
||||
(op_schema.args_schema, op_schema.kwargs_schema),
|
||||
):
|
||||
return None
|
||||
|
||||
candidate_placements = self._get_candidate_placements(op_schema)
|
||||
mesh = try_find_mesh_from_args(
|
||||
op_schema.op,
|
||||
op_schema.args_schema + tuple(op_schema.kwargs_schema.values()),
|
||||
)
|
||||
|
||||
fake_mesh = self._get_fake_mesh(mesh.device_type)
|
||||
single_dim_strategies = []
|
||||
output_placements: list[Placement | tuple[Placement, ...]] = []
|
||||
for input_placements in candidate_placements:
|
||||
try:
|
||||
output = self._propagate_through_decomp(
|
||||
op_schema,
|
||||
input_placements,
|
||||
fake_mesh,
|
||||
)
|
||||
except NotImplementedError:
|
||||
return None
|
||||
except GuardOnDataDependentSymNode:
|
||||
return None
|
||||
except (RuntimeError, KeyError, IndexError):
|
||||
# TODO(pianpwk): RuntimeError is raised when redistribution is detected; switch to a custom error type
|
||||
# Runtime/KeyError/IndexError can also occur in view ops
|
||||
continue
|
||||
|
||||
output_placements = (
|
||||
[output] if not isinstance(output, tuple) else list(output)
|
||||
)
|
||||
single_dim_strategies.append(output_placements + list(input_placements))
|
||||
|
||||
if not single_dim_strategies:
|
||||
raise AssertionError(
|
||||
"Sharding propagation should have produced at least Replicate() strategy"
|
||||
)
|
||||
|
||||
n_outputs = len(output_placements)
|
||||
strategy_schema = self.sharding_prop._wrap_with_op_strategy(op_schema)
|
||||
# Import here to avoid circular import at module load time
|
||||
from torch.distributed.tensor._ops.utils import ( # noqa: F811
|
||||
expand_to_full_mesh_op_strategy,
|
||||
)
|
||||
|
||||
return expand_to_full_mesh_op_strategy(
|
||||
mesh, strategy_schema, single_dim_strategies, input_index=n_outputs
|
||||
)
|
||||
|
||||
def _propagate_through_decomp(
|
||||
self,
|
||||
op_schema: OpSchema,
|
||||
placement: tuple[Placement | None],
|
||||
mesh: DeviceMesh,
|
||||
) -> Placement | tuple[Placement, ...]:
|
||||
op = op_schema.op
|
||||
if op in decomposition_table:
|
||||
decomp_fn = decomposition_table[op]
|
||||
elif op._can_decompose():
|
||||
decomp_fn = op.decompose
|
||||
else:
|
||||
raise NotImplementedError(f"No decomposition found for {op}")
|
||||
|
||||
placement_iter = iter(placement)
|
||||
|
||||
def to_meta(x):
|
||||
p = next(placement_iter)
|
||||
if isinstance(x, DTensorSpec):
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
meta = torch.empty(x.shape, dtype=x.tensor_meta.dtype, device="meta")
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
meta._spec = DTensorSpec(mesh, (p,), tensor_meta=x.tensor_meta)
|
||||
return meta
|
||||
return x
|
||||
|
||||
# Disable LocalTensorMode during decomposition tracing to prevent
|
||||
# interference with meta tensor operations
|
||||
from torch.distributed._local_tensor import maybe_disable_local_tensor_mode
|
||||
|
||||
with maybe_disable_local_tensor_mode():
|
||||
# Create meta tensors and run decomposition outside LocalTensorMode
|
||||
args_meta = tree_map(to_meta, op_schema.args_schema)
|
||||
kwargs_meta = tree_map(to_meta, op_schema.kwargs_schema)
|
||||
|
||||
with PlacementTrackingMode(self.sharding_prop, mesh):
|
||||
output = decomp_fn(*args_meta, **kwargs_meta)
|
||||
|
||||
def get_placement(t):
|
||||
if isinstance(t, torch.Tensor):
|
||||
spec = getattr(t, "_spec", None)
|
||||
return spec.placements[0] if spec else None
|
||||
return None
|
||||
|
||||
result = tree_map(get_placement, output)
|
||||
if isinstance(result, (tuple, list)):
|
||||
flat = [p for p in result if p is not None]
|
||||
return flat[0] if len(flat) == 1 else tuple(flat)
|
||||
return result
|
||||
|
||||
@staticmethod
|
||||
def _get_candidate_placements(
|
||||
op_schema: OpSchema,
|
||||
) -> list[tuple[Placement | None]]:
|
||||
tensor_specs = _extract_input_specs(op_schema)
|
||||
flat_specs, _ = tree_flatten(list(tensor_specs))
|
||||
|
||||
# Step 1: Collect unique placements across all DTensorSpec inputs
|
||||
all_placements: set[Placement] = {Replicate()}
|
||||
tree_map_only(
|
||||
DTensorSpec,
|
||||
lambda spec: all_placements.update(spec.placements),
|
||||
flat_specs,
|
||||
)
|
||||
|
||||
# Step 2: For each input, use the placement set, but expand Shard/StridedShard to all tensor dims
|
||||
candidates: list[list[Placement | None]] = []
|
||||
for spec in flat_specs:
|
||||
if not isinstance(spec, DTensorSpec):
|
||||
candidates.append([None])
|
||||
else:
|
||||
options = set(all_placements)
|
||||
for p in all_placements:
|
||||
if isinstance(p, _StridedShard):
|
||||
options |= {
|
||||
_StridedShard(i, split_factor=p.split_factor)
|
||||
for i in range(spec.ndim)
|
||||
}
|
||||
elif isinstance(p, Shard):
|
||||
options |= {Shard(i) for i in range(spec.ndim)}
|
||||
candidates.append(list(options))
|
||||
|
||||
# pyrefly: ignore [bad-argument-type, no-matching-overload]
|
||||
return list(itertools.product(*candidates))
|
||||
@@ -0,0 +1,802 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates
|
||||
import contextlib
|
||||
import logging
|
||||
import warnings
|
||||
from collections.abc import Sequence
|
||||
from typing import cast
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
import torch.distributed.tensor._api as dtensor
|
||||
import torch.distributed.tensor._random as random
|
||||
from torch._library.utils import fill_defaults
|
||||
from torch._logging import LazyString
|
||||
from torch._prims.rng_prims import run_dtensor_rng_op
|
||||
from torch.distributed._functional_collectives import _are_we_tracing
|
||||
from torch.distributed.device_mesh import DeviceMesh
|
||||
from torch.distributed.tensor._dtensor_spec import DTensorSpec, TensorMeta
|
||||
from torch.distributed.tensor._nonlinear_redux import argminmax_handler
|
||||
from torch.distributed.tensor._op_schema import (
|
||||
OpInfo,
|
||||
OpSchema,
|
||||
OutputSharding,
|
||||
OutputSpecType,
|
||||
)
|
||||
from torch.distributed.tensor._random import is_rng_supported_mesh
|
||||
from torch.distributed.tensor._redistribute import redistribute_local_tensor
|
||||
from torch.distributed.tensor._sharding_prop import ShardingPropagator
|
||||
from torch.distributed.tensor._tp_conv import (
|
||||
convolution_backward_handler,
|
||||
convolution_handler,
|
||||
)
|
||||
from torch.distributed.tensor._utils import (
|
||||
_format_implicit_redistribution_msg,
|
||||
ExplicitRedistributionContext,
|
||||
try_find_mesh_from_args,
|
||||
)
|
||||
from torch.distributed.tensor.placement_types import Partial, Placement, Replicate
|
||||
from torch.utils._debug_mode import get_active_debug_mode
|
||||
from torch.utils._python_dispatch import return_and_correct_aliasing
|
||||
|
||||
|
||||
try:
|
||||
from torch.utils import _cxx_pytree as pytree
|
||||
except ImportError:
|
||||
from torch.utils import _pytree as pytree # type: ignore[no-redef]
|
||||
|
||||
aten = torch.ops.aten
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# The C++ DTensor dispatch fast path caches whether debug logging is
|
||||
# enabled. Wrap setLevel so the cached flag is reset automatically.
|
||||
_orig_setLevel = logger.setLevel
|
||||
|
||||
|
||||
def _setLevel_and_reinit(level: int) -> None:
|
||||
_orig_setLevel(level)
|
||||
torch._C._reinit_DTensor_dispatch_logger()
|
||||
|
||||
|
||||
logger.setLevel = _setLevel_and_reinit # type: ignore[method-assign]
|
||||
|
||||
|
||||
def as_strided_handler(
|
||||
op_call: torch._ops.OpOverload,
|
||||
args: tuple[object, ...],
|
||||
kwargs: dict[str, object],
|
||||
):
|
||||
args, kwargs = fill_defaults(op_call._schema, args, kwargs)
|
||||
if kwargs:
|
||||
raise AssertionError
|
||||
tensor, size, stride, storage_offset = args
|
||||
if (
|
||||
tensor.size() == tuple(size)
|
||||
and tensor.stride() == tuple(stride)
|
||||
and (storage_offset is None or tensor.storage_offset() == storage_offset)
|
||||
):
|
||||
return torch.ops.aten.alias.default(tensor)
|
||||
raise RuntimeError("as_strided not supported with DTensor")
|
||||
|
||||
|
||||
def is_same_size_handler(
|
||||
op_call: torch._ops.OpOverload,
|
||||
args: tuple[object, ...],
|
||||
kwargs: dict[str, object],
|
||||
) -> bool:
|
||||
lhs = cast(torch.Tensor, args[0])
|
||||
rhs = cast(torch.Tensor, args[1])
|
||||
return lhs.shape == rhs.shape
|
||||
|
||||
|
||||
def is_pinned_handler(
|
||||
op_call: torch._ops.OpOverload,
|
||||
args: tuple[object, ...],
|
||||
kwargs: dict[str, object],
|
||||
) -> bool:
|
||||
tensor = cast(dtensor.DTensor, args[0])
|
||||
return tensor._local_tensor.is_pinned()
|
||||
|
||||
|
||||
def found_inf_reduce_handler(
|
||||
op_call: torch._ops.OpOverload,
|
||||
args: tuple[object, ...],
|
||||
kwargs: dict[str, object],
|
||||
) -> None:
|
||||
op_info = dtensor.DTensor._op_dispatcher.unwrap_to_op_info(op_call, args, kwargs)
|
||||
local_tensor_args = pytree.tree_unflatten(
|
||||
cast(list[object], op_info.local_args),
|
||||
op_info.args_tree_spec, # type: ignore[arg-type]
|
||||
)
|
||||
local_tensor_args = cast(tuple[object, ...], local_tensor_args)
|
||||
op_call(*local_tensor_args, **op_info.local_kwargs)
|
||||
|
||||
grad_dtensor = cast(list[dtensor.DTensor], args[0])[0]
|
||||
grad_placements = grad_dtensor.placements
|
||||
mesh = grad_dtensor.device_mesh
|
||||
|
||||
found_inf_placements: list[Placement] = []
|
||||
for placement in grad_placements:
|
||||
if isinstance(placement, Replicate):
|
||||
found_inf_placements.append(placement)
|
||||
else:
|
||||
found_inf_placements.append(Partial("max"))
|
||||
|
||||
target_tensor = cast(torch.Tensor, args[1])
|
||||
spec = DTensorSpec(
|
||||
mesh=mesh,
|
||||
placements=tuple(found_inf_placements),
|
||||
tensor_meta=TensorMeta(
|
||||
shape=target_tensor.size(),
|
||||
stride=target_tensor.stride(),
|
||||
dtype=target_tensor.dtype,
|
||||
),
|
||||
)
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
found_inf_dtensor = dtensor.DTensor(
|
||||
local_tensor=target_tensor, # pyrefly: ignore [unexpected-keyword]
|
||||
spec=spec, # pyrefly: ignore [unexpected-keyword]
|
||||
requires_grad=False, # pyrefly: ignore [unexpected-keyword]
|
||||
)
|
||||
found_inf = found_inf_dtensor.full_tensor()
|
||||
target_tensor.copy_(found_inf)
|
||||
|
||||
|
||||
class OpDispatcher:
|
||||
"""
|
||||
Op dispatching class instance to handle args/kwargs pre-processing (un-wrapping), sharding
|
||||
propagation, redistribute local args, local compute, and post-processing (re-wrapping). It
|
||||
also handles any op specific logic if necessary.
|
||||
|
||||
NOTE: Given the runtime overhead of Tensor subclass (__torch_dispatch__), the OpDispatcher
|
||||
is designed to minimize the CPU overhead by using the tricks of proper unflattening, faster
|
||||
pytree if needed, and leveraging various caching mechanisms implemented in the sharding
|
||||
propagation and redistribute modules. The CPU overhead is critical to eager mode performance,
|
||||
one need to carefully measure the CPU overhead when making significant changes to the
|
||||
OpDispatcher and ShardingPropagator.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.sharding_propagator = ShardingPropagator()
|
||||
# NOTE: must stay in sync with is_random_op in
|
||||
# torch/csrc/autograd/python_variable.cpp
|
||||
self._random_ops = {
|
||||
aten.native_dropout.default,
|
||||
aten.normal_.default,
|
||||
aten.rand.default,
|
||||
aten.rand_like.default,
|
||||
aten.randn.default,
|
||||
aten.randn_like.default,
|
||||
aten.randint_like.default,
|
||||
aten.randint_like.low_dtype,
|
||||
aten.randint_like.low_dtype_out,
|
||||
aten.uniform_.default,
|
||||
aten.bernoulli.default,
|
||||
aten.bernoulli_.float,
|
||||
}
|
||||
self._squeeze_inplace_ops = {
|
||||
aten.squeeze_.dim,
|
||||
aten.squeeze_.default,
|
||||
aten.squeeze_.dims,
|
||||
}
|
||||
self._custom_op_handlers = {
|
||||
aten.is_same_size.default: is_same_size_handler,
|
||||
aten.is_pinned.default: is_pinned_handler,
|
||||
aten.convolution.default: convolution_handler,
|
||||
aten.convolution_backward.default: convolution_backward_handler,
|
||||
aten._amp_foreach_non_finite_check_and_unscale_.default: found_inf_reduce_handler,
|
||||
aten.as_strided.default: as_strided_handler,
|
||||
aten.argmin.default: argminmax_handler,
|
||||
aten.argmax.default: argminmax_handler,
|
||||
}
|
||||
|
||||
# ********************************************************************************************
|
||||
# def dispatch(...)
|
||||
#
|
||||
# NOTE: this class no longer contains the top-level dispatch entrypoint!
|
||||
# See #167051 for details
|
||||
#
|
||||
# The entrypoint has been moved to C++, and it handles common cases and then calls back into
|
||||
# OpDispatcher python to handle corner cases.
|
||||
# See dispatchDTensorOp() defined in python_variable.cpp and called from python_arg_parser.cpp
|
||||
# ********************************************************************************************
|
||||
|
||||
# This flag is used internally to control whether we treat the torch.Tensor(non-DTensor)
|
||||
# as implicitly replicated or we throw error to user.
|
||||
# NOTE: It is EXTREMELY UNSAFE to turn this flag on by default so we intentionally leave
|
||||
# it as False by default.
|
||||
@property
|
||||
def _allow_implicit_replication(self) -> bool:
|
||||
return torch._C._get_dtensor_allow_implicit_replication()
|
||||
|
||||
@_allow_implicit_replication.setter
|
||||
def _allow_implicit_replication(self, value: bool) -> None:
|
||||
return torch._C._set_dtensor_allow_implicit_replication(value)
|
||||
|
||||
def _propagate_op_sharding_dispatch_slow_path(
|
||||
self,
|
||||
op_call: torch._ops.OpOverload,
|
||||
args: tuple[object, ...],
|
||||
kwargs: dict[str, object],
|
||||
op_info: OpInfo,
|
||||
# The logic here is a bit messy. There are several reasons why the
|
||||
# C++ fastpath may have bailed out. If we just cache missed, we will
|
||||
# come here because we need to actually calculate the real thing.
|
||||
# There's no need to have a SECOND Python cache lookup; the C++ native
|
||||
# cache completely subsumes it. But sometimes, we will have failed
|
||||
# to compute the cache key in C++ entirely. In this case, we DO need
|
||||
# to do a cache lookup in Python, as the missing cache key in C++
|
||||
# means we don't have access to it all. Furthermore, without duping
|
||||
# this function, we need to do the try_cache test inside of the
|
||||
# try-except block so that either case hits the inference mode /
|
||||
# exception rewrapping case.
|
||||
#
|
||||
# This should be cleaned up. First, ensuring the C++ codepath can
|
||||
# always compute a key will be a big help. Second, we should properly
|
||||
# fastpath inference mode composite implicit autograd so that you
|
||||
# don't have to throw an exception even in "fastpath".
|
||||
try_cache: bool,
|
||||
) -> object:
|
||||
# NOTE: schema should always be populated when calling this function,
|
||||
# as it's only called from C++ after unwrap_to_op_info (create_schema=True).
|
||||
# See dispatchDTensorOp in python_variable.cpp line 1453-1460.
|
||||
if op_info.schema is None:
|
||||
raise AssertionError(
|
||||
"op_info.schema should not be None in sharding propagation. "
|
||||
"This function should only be called after unwrap_to_op_info."
|
||||
)
|
||||
try:
|
||||
# We have basically inlined propagate() here, but WITHOUT the
|
||||
# output_sharding assignment
|
||||
if try_cache and not _are_we_tracing():
|
||||
result = self.sharding_propagator.propagate_op_sharding(op_info.schema)
|
||||
else:
|
||||
result = self.sharding_propagator.propagate_op_sharding_non_cached(
|
||||
op_info.schema
|
||||
)
|
||||
if logger.handlers and logger.isEnabledFor(logging.DEBUG):
|
||||
logger.debug(
|
||||
"sharding_prop MISS (C++ fast path): %s -> %s",
|
||||
op_info.schema,
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
result.output_spec,
|
||||
)
|
||||
return result
|
||||
except NotImplementedError:
|
||||
if torch._C._dispatch_has_kernel_for_dispatch_key(
|
||||
op_call.name(), torch._C.DispatchKey.CompositeImplicitAutograd
|
||||
):
|
||||
# When running under inference mode, CompositeImplicitAutograd ops show up in __torch_dispatch__,
|
||||
# so we manually decompose them, here
|
||||
out = op_call.decompose(*args, **kwargs)
|
||||
if out is NotImplemented:
|
||||
raise AssertionError from None
|
||||
return out
|
||||
else:
|
||||
raise
|
||||
except Exception as e:
|
||||
raise RuntimeError(
|
||||
f"{e}\n\nSharding propagation failed for {op_info.schema or op_call}"
|
||||
) from e
|
||||
|
||||
def _dispatch_get_local_results_slow_path(
|
||||
self,
|
||||
op_call: torch._ops.OpOverload,
|
||||
args: tuple[object, ...],
|
||||
op_info: OpInfo,
|
||||
) -> object:
|
||||
output_sharding = op_info.output_sharding
|
||||
if output_sharding is None:
|
||||
raise AssertionError("output sharding should not be None")
|
||||
if op_info is None:
|
||||
raise AssertionError("op_info should never be None")
|
||||
|
||||
# Record output placements for debugging
|
||||
debug_mode = get_active_debug_mode()
|
||||
if debug_mode is not None and output_sharding.output_spec is not None:
|
||||
debug_mode.record_output_placements(output_sharding.output_spec)
|
||||
|
||||
mesh = op_info.compute_mesh
|
||||
participating = mesh._is_current_rank_part_of_mesh()
|
||||
local_results = None
|
||||
if participating:
|
||||
# computation that happens in the current rank of the mesh, normal case
|
||||
if output_sharding.needs_redistribute:
|
||||
# If sharding propagation decision needs redistribute, perform redistribute
|
||||
# on args first, which could potentially modify args (i.e. allgather certain arg)
|
||||
if output_sharding.redistribute_schema is None:
|
||||
raise AssertionError
|
||||
self.redistribute_local_args(
|
||||
op_info,
|
||||
output_sharding.redistribute_schema,
|
||||
output_sharding.use_val_from_redistribute_schema,
|
||||
)
|
||||
|
||||
local_tensor_args = (
|
||||
pytree.tree_unflatten(
|
||||
cast(list[object], op_info.local_args),
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
op_info.args_tree_spec,
|
||||
)
|
||||
if op_info.args_tree_spec
|
||||
else op_info.local_args
|
||||
)
|
||||
|
||||
# run local op computation with potentially modified args/kwargs
|
||||
local_tensor_args = cast(tuple[object, ...], local_tensor_args)
|
||||
if op_call in self._random_ops:
|
||||
if not random._rng_tracker and is_rng_supported_mesh(mesh):
|
||||
# Default to `OffsetBasedRNGTracker` if the parallelism API did not already construct one
|
||||
# Skip RNG state sync during tracing to avoid lazily initializing real RNG state under fake mode.
|
||||
run_state_sync = not _are_we_tracing()
|
||||
if not run_state_sync:
|
||||
logger.info(
|
||||
"DTensor RNG tracker is being lazily initialized during tracing. "
|
||||
"RNG states may not be synchronized across ranks, which can lead "
|
||||
"to silent incorrectness. Please call `torch.manual_seed()` with "
|
||||
"the same seed on all ranks before compiling DTensor random ops.",
|
||||
stacklevel=2,
|
||||
)
|
||||
random._rng_tracker = random.OffsetBasedRNGTracker(
|
||||
mesh, run_state_sync
|
||||
)
|
||||
|
||||
first_arg, first_local_arg = (
|
||||
cast(dtensor.DTensor, args[0]),
|
||||
cast(torch.Tensor, local_tensor_args[0]),
|
||||
)
|
||||
|
||||
# If the user provided a generator, we hook it up to our RNG manager, but we also pop it from kwargs
|
||||
# so the op_call does not directly use it (we want op_call to fall back to the 'default' which is
|
||||
# our RNG manager)
|
||||
maybe_user_generator = op_info.local_kwargs.pop("generator", None)
|
||||
if not (
|
||||
maybe_user_generator is None
|
||||
or isinstance(maybe_user_generator, torch.Generator)
|
||||
):
|
||||
raise AssertionError
|
||||
|
||||
if (
|
||||
random._rng_tracker
|
||||
and not first_local_arg.is_meta
|
||||
and random._rng_tracker.distribute_region_enabled
|
||||
):
|
||||
if (
|
||||
maybe_user_generator is not None
|
||||
or first_local_arg.device.type != "cuda"
|
||||
or (
|
||||
not _are_we_tracing()
|
||||
and type(first_local_arg) is not torch.Tensor
|
||||
)
|
||||
):
|
||||
with random._rng_tracker._distribute_region(
|
||||
first_arg._spec, generator=maybe_user_generator
|
||||
):
|
||||
local_results = op_call(
|
||||
*local_tensor_args, **op_info.local_kwargs
|
||||
)
|
||||
else:
|
||||
# CUDA device without user generator, use HOP for traceability
|
||||
if not isinstance(
|
||||
random._rng_tracker, random.OffsetBasedRNGTracker
|
||||
):
|
||||
raise AssertionError
|
||||
start_offset_incr, end_offset_incr = (
|
||||
random._rng_tracker._compute_rng_offsets(first_arg._spec)
|
||||
)
|
||||
local_results = run_dtensor_rng_op(
|
||||
start_offset_incr,
|
||||
end_offset_incr,
|
||||
op_call,
|
||||
*local_tensor_args,
|
||||
**op_info.local_kwargs,
|
||||
)
|
||||
else:
|
||||
# No rng_tracker, meta tensor, or distribute_region disabled
|
||||
local_results = op_call(*local_tensor_args, **op_info.local_kwargs)
|
||||
else:
|
||||
# normal case, run local sharded op computation
|
||||
if (
|
||||
output_sharding.needs_redistribute
|
||||
and output_sharding.redistribute_schema is not None
|
||||
and output_sharding.redistribute_schema.op != op_call
|
||||
):
|
||||
# Op was rewritten (e.g., squeeze.default → squeeze.dims)
|
||||
local_results = output_sharding.redistribute_schema.op(
|
||||
*local_tensor_args, **op_info.local_kwargs
|
||||
)
|
||||
else:
|
||||
local_results = op_call(*local_tensor_args, **op_info.local_kwargs)
|
||||
|
||||
else:
|
||||
# For a non-participating device (happens on rank that does not belong to
|
||||
# the device mesh), we do:
|
||||
# 1. if the return type is scalar, set the local result to None.
|
||||
# 2. if the return type is Tensor or List[Tensor], return empty
|
||||
# tensor(s) with correct dtype.
|
||||
spec = output_sharding.output_spec
|
||||
ret_list = op_call._schema.returns
|
||||
|
||||
if spec is None:
|
||||
# For a scalar return type, the non-participating device has None
|
||||
# as its local result
|
||||
local_results = None
|
||||
else:
|
||||
|
||||
def default_tensor(spec: DTensorSpec) -> torch.Tensor:
|
||||
if spec.tensor_meta is not None:
|
||||
shape = spec.tensor_meta.shape
|
||||
dtype = spec.tensor_meta.dtype
|
||||
if len(shape) == 0:
|
||||
# scalar tensor
|
||||
return torch.zeros((), dtype=dtype)
|
||||
else:
|
||||
# non-scalar tensor
|
||||
return torch.tensor([], dtype=dtype)
|
||||
else:
|
||||
raise RuntimeError(f"{spec} has no tensor metadata.")
|
||||
|
||||
if isinstance(spec, DTensorSpec):
|
||||
# return a Tensor value
|
||||
local_results = default_tensor(spec)
|
||||
elif isinstance(spec, Sequence):
|
||||
# return a List[Tensor] value
|
||||
local_results = [
|
||||
default_tensor(s) if s is not None else None for s in spec
|
||||
]
|
||||
if not isinstance(local_results, list):
|
||||
raise AssertionError
|
||||
if None in local_results:
|
||||
ret_type = str(ret_list[0].type)
|
||||
raise NotImplementedError(
|
||||
f"return type {ret_type} in DTensor op is not supported"
|
||||
)
|
||||
return local_results
|
||||
|
||||
def _dispatch_fast_path_python_tail(
|
||||
self,
|
||||
op_call: torch._ops.OpOverload,
|
||||
args: tuple[object, ...],
|
||||
kwargs: dict[str, object],
|
||||
compute_mesh: DeviceMesh,
|
||||
output_sharding: OutputSharding,
|
||||
local_results: object,
|
||||
participating: bool,
|
||||
is_inplace_op: bool,
|
||||
is_out_variant_op: bool,
|
||||
) -> object:
|
||||
"""
|
||||
Tail of main dispatching logic, called from C++ fast path.
|
||||
"""
|
||||
|
||||
# Record output placements for debugging
|
||||
debug_mode = get_active_debug_mode()
|
||||
if debug_mode is not None and output_sharding.output_spec is not None:
|
||||
debug_mode.record_output_placements(output_sharding.output_spec)
|
||||
|
||||
if output_sharding.output_spec is None:
|
||||
if op_call == aten.equal.default:
|
||||
# The output of the equal op is a bool, by converting it into a
|
||||
# a single value tensor, we can use all-reduce with min reduce op
|
||||
# to simulate logical and.
|
||||
if not (local_results is None or isinstance(local_results, bool)):
|
||||
raise AssertionError
|
||||
r = torch.tensor(
|
||||
int(local_results) if local_results is not None else 1,
|
||||
device=compute_mesh.device_type,
|
||||
)
|
||||
dist.all_reduce(r, op=dist.ReduceOp.MIN)
|
||||
local_results = bool(r.item())
|
||||
|
||||
if is_inplace_op:
|
||||
# inplace op should return self instead of re-wrapping
|
||||
if output_sharding.output_spec is not None:
|
||||
output_spec = output_sharding.output_spec
|
||||
if not isinstance(output_spec, DTensorSpec):
|
||||
raise AssertionError
|
||||
if not isinstance(args[0], dtensor.DTensor):
|
||||
raise AssertionError
|
||||
|
||||
# NOTE: squeeze_ inplace ops may change the tensor's metadata
|
||||
# (shape/strides). We special-case them to update the spec.
|
||||
if op_call in self._squeeze_inplace_ops:
|
||||
# update the spec to handle tensor meta changes
|
||||
args[0]._spec = output_spec
|
||||
# use return_and_correct_aliasing to match the outer and the inner
|
||||
# aliasing. See https://github.com/pytorch/pytorch/pull/158954
|
||||
return return_and_correct_aliasing(op_call, args, kwargs, args[0])
|
||||
else:
|
||||
# For all other inplace ops, check if placement changes are required
|
||||
# Inplace operations that change placement are not supported because
|
||||
# they would require redistribution, which breaks aliasing semantics.
|
||||
# If there are views into the tensor, the views would not be updated.
|
||||
if args[0]._spec.placements != output_spec.placements:
|
||||
raise RuntimeError(
|
||||
f"{op_call}: in-place operations that require placement changes "
|
||||
f"are not supported. The operation would change placement from "
|
||||
f"{args[0]._spec.placements} to {output_spec.placements}, "
|
||||
f"which requires redistribution and breaks aliasing semantics. "
|
||||
f"Please use the out-of-place version of this operation instead."
|
||||
)
|
||||
# Most inplace ops don't change tensor meta, so no spec update needed
|
||||
return args[0]
|
||||
else:
|
||||
return None
|
||||
elif is_out_variant_op:
|
||||
# out variant could possibly have multiple out args (i.e. lu_unpack.out)
|
||||
output_specs = (
|
||||
(output_sharding.output_spec,)
|
||||
if not isinstance(output_sharding.output_spec, tuple)
|
||||
else output_sharding.output_spec
|
||||
)
|
||||
out_dts = []
|
||||
spec_idx = 0
|
||||
for argument in op_call._schema.arguments:
|
||||
if argument.is_out:
|
||||
out_dt = cast(dtensor.DTensor, kwargs[argument.name])
|
||||
out_dt._spec = cast(DTensorSpec, output_specs[spec_idx])
|
||||
out_dts.append(out_dt)
|
||||
spec_idx += 1
|
||||
|
||||
if len(out_dts) < 1:
|
||||
raise AssertionError("out variant should have at least one out arg")
|
||||
return tuple(out_dts) if len(out_dts) > 1 else out_dts[0]
|
||||
else:
|
||||
if op_call != aten.equal.default:
|
||||
raise AssertionError(op_call)
|
||||
ret = self.wrap(local_results, output_sharding.output_spec) # type: ignore[possibly-undefined]
|
||||
if participating and op_call._schema._is_view_op():
|
||||
return return_and_correct_aliasing(op_call, args, kwargs, ret)
|
||||
else:
|
||||
return ret
|
||||
|
||||
@staticmethod
|
||||
def redistribute_local_args(
|
||||
op_info: OpInfo,
|
||||
suggested_input_schema: OpSchema,
|
||||
use_val_from_redistribute_schema: bool,
|
||||
) -> None:
|
||||
debug_mode = get_active_debug_mode()
|
||||
|
||||
# NOTE: it's very rare that we need to reshard kwargs so we intentionally skip it
|
||||
if op_info.args_tree_spec is not None:
|
||||
flatten_args_schema_to_reshard = tuple(
|
||||
pytree.tree_leaves(suggested_input_schema.args_schema)
|
||||
)
|
||||
else:
|
||||
flatten_args_schema_to_reshard = suggested_input_schema.args_schema
|
||||
|
||||
new_local_args: list[object] = []
|
||||
for i, arg_spec in enumerate(op_info.flat_args_schema):
|
||||
reshard_arg_spec = flatten_args_schema_to_reshard[i]
|
||||
if isinstance(arg_spec, DTensorSpec):
|
||||
local_tensor = cast(torch.Tensor, op_info.local_args[i])
|
||||
if arg_spec != reshard_arg_spec:
|
||||
redistribute_context = (
|
||||
debug_mode.record_redistribute_calls( # type: ignore[union-attr]
|
||||
i, arg_spec, reshard_arg_spec
|
||||
)
|
||||
if debug_mode is not None
|
||||
else contextlib.nullcontext()
|
||||
)
|
||||
|
||||
ExplicitRedistributionContext.observe_redistribution(
|
||||
arg_spec,
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
reshard_arg_spec,
|
||||
LazyString(
|
||||
_format_implicit_redistribution_msg,
|
||||
op_info.schema or suggested_input_schema.op,
|
||||
),
|
||||
)
|
||||
with redistribute_context:
|
||||
resharded_local_tensor = redistribute_local_tensor(
|
||||
local_tensor,
|
||||
arg_spec,
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
reshard_arg_spec,
|
||||
)
|
||||
new_local_args.append(resharded_local_tensor)
|
||||
else:
|
||||
new_local_args.append(local_tensor)
|
||||
else:
|
||||
if use_val_from_redistribute_schema:
|
||||
# args can be updated for view related ops, we refer to the
|
||||
# update in redistribute_schema.
|
||||
new_local_args.append(reshard_arg_spec)
|
||||
else:
|
||||
new_local_args.append(arg_spec)
|
||||
|
||||
# Append extra non-tensor args from rewritten schema (e.g., dims tuple).
|
||||
if use_val_from_redistribute_schema:
|
||||
for i in range(
|
||||
len(op_info.flat_args_schema), len(flatten_args_schema_to_reshard)
|
||||
):
|
||||
new_local_args.append(flatten_args_schema_to_reshard[i])
|
||||
|
||||
op_info.local_args = tuple(new_local_args)
|
||||
|
||||
def unwrap_to_op_info(
|
||||
self,
|
||||
op_call: torch._ops.OpOverload,
|
||||
args: tuple[object, ...],
|
||||
kwargs: dict[str, object],
|
||||
) -> OpInfo:
|
||||
return self._unwrap_to_op_info_impl(op_call, args, kwargs, True)
|
||||
|
||||
def _unwrap_to_op_info_impl(
|
||||
self,
|
||||
op_call: torch._ops.OpOverload,
|
||||
args: tuple[object, ...],
|
||||
kwargs: dict[str, object],
|
||||
create_schema: bool,
|
||||
) -> OpInfo:
|
||||
# get runtime schema info to determine whether to use pytree to flatten inputs
|
||||
runtime_schema_info = self.sharding_propagator.op_to_schema_info.get(
|
||||
op_call, None
|
||||
)
|
||||
if runtime_schema_info is None:
|
||||
runtime_schema_info = (
|
||||
self.sharding_propagator.op_to_schema_info_for_single_dim_strategy.get(
|
||||
op_call, None
|
||||
)
|
||||
)
|
||||
|
||||
# Auto-detect needs_pytree if any arg is a list/tuple containing tensors
|
||||
def _contains_tensor(arg: object) -> bool:
|
||||
if isinstance(arg, (list, tuple)):
|
||||
return any(isinstance(item, torch.Tensor) for item in arg)
|
||||
return False
|
||||
|
||||
needs_pytree = (
|
||||
runtime_schema_info is not None and runtime_schema_info.needs_pytree
|
||||
) or any(_contains_tensor(arg) for arg in args)
|
||||
|
||||
if needs_pytree:
|
||||
# flatten args/kwargs when op says necessary or args contain lists/tuples
|
||||
tree_args, args_spec = pytree.tree_flatten(args)
|
||||
args_list: Sequence[object] = tree_args
|
||||
else:
|
||||
args_list, args_spec = args, None
|
||||
|
||||
args_schema: list[object] = []
|
||||
kwargs_schema: dict[str, object] = {}
|
||||
local_args: list[object] = []
|
||||
local_kwargs: dict[str, object] = {}
|
||||
compute_mesh: DeviceMesh | None = None
|
||||
|
||||
for arg in args_list:
|
||||
if isinstance(arg, dtensor.DTensor):
|
||||
local_args.append(arg._local_tensor)
|
||||
args_schema.append(arg._spec)
|
||||
if compute_mesh is None:
|
||||
# record the first compute device mesh from args
|
||||
compute_mesh = arg.device_mesh
|
||||
elif isinstance(arg, torch.Tensor):
|
||||
compute_mesh = compute_mesh or try_find_mesh_from_args(
|
||||
op_call, args_list
|
||||
)
|
||||
args_schema.append(
|
||||
self._try_replicate_spec_for_scalar_tensor(
|
||||
op_call, arg, compute_mesh
|
||||
)
|
||||
)
|
||||
local_args.append(arg)
|
||||
else:
|
||||
# non DTensor/Tensor args (i.e. int/float/bool), just add to args_schema/local_args
|
||||
args_schema.append(arg)
|
||||
local_args.append(arg)
|
||||
|
||||
for k, v in kwargs.items():
|
||||
if isinstance(v, dtensor.DTensor):
|
||||
local_kwargs[k] = v._local_tensor
|
||||
kwargs_schema[k] = v._spec
|
||||
if compute_mesh is None:
|
||||
# record the first compute device mesh from kwargs
|
||||
compute_mesh = v.device_mesh
|
||||
elif isinstance(v, torch.Tensor):
|
||||
compute_mesh = compute_mesh or try_find_mesh_from_args(
|
||||
op_call, args_list
|
||||
)
|
||||
kwargs_schema[k] = self._try_replicate_spec_for_scalar_tensor(
|
||||
op_call,
|
||||
v,
|
||||
compute_mesh,
|
||||
)
|
||||
local_kwargs[k] = v
|
||||
else:
|
||||
# non DTensor/Tensor args (i.e. int/float/bool), just add to args_schema/local_args
|
||||
kwargs_schema[k] = v
|
||||
local_kwargs[k] = v
|
||||
|
||||
if compute_mesh is None:
|
||||
raise AssertionError(
|
||||
f"found no DeviceMesh from dtensor args for {op_call}!"
|
||||
)
|
||||
op_info = OpInfo(
|
||||
compute_mesh,
|
||||
OpSchema(
|
||||
op_call,
|
||||
(
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
pytree.tree_unflatten(args_schema, args_spec)
|
||||
if args_spec
|
||||
else tuple(args_schema)
|
||||
),
|
||||
kwargs_schema,
|
||||
schema_info=runtime_schema_info,
|
||||
)
|
||||
if create_schema
|
||||
else None, # type: ignore[arg-type]
|
||||
args_schema,
|
||||
tuple(local_args),
|
||||
local_kwargs,
|
||||
args_spec,
|
||||
)
|
||||
return op_info
|
||||
|
||||
@staticmethod
|
||||
def wrap(res: object, spec: OutputSpecType) -> object:
|
||||
if isinstance(res, torch.Tensor):
|
||||
if spec is not None:
|
||||
if not isinstance(spec, DTensorSpec):
|
||||
raise AssertionError(
|
||||
f"output spec does not match with output! Expected DTensorSpec, got {spec}."
|
||||
)
|
||||
# pyrefly: ignore [bad-argument-type, bad-argument-count, unexpected-keyword]
|
||||
return dtensor.DTensor(res, spec, requires_grad=res.requires_grad)
|
||||
else:
|
||||
# if output does not have a DTensorSpec due to specific ops, it must be a scalar tensor
|
||||
if res.ndim != 0:
|
||||
raise AssertionError("output tensor should be scalar!")
|
||||
return res
|
||||
elif isinstance(res, (list, tuple)):
|
||||
if not (spec is not None and isinstance(spec, (list, tuple))):
|
||||
raise AssertionError(
|
||||
f"output spec does not match with output! Expected list/tuple, got {spec}."
|
||||
)
|
||||
res_list = []
|
||||
for e, s in zip(res, spec):
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
res_list.append(OpDispatcher.wrap(e, s))
|
||||
|
||||
return tuple(res_list) if isinstance(res, tuple) else res_list
|
||||
else:
|
||||
# if the res contains only non tensor values (i.e. int/float/none), we simply return it
|
||||
# without rewrapping to DTensor.
|
||||
return res
|
||||
|
||||
def _try_replicate_spec_for_scalar_tensor(
|
||||
self,
|
||||
op_call: torch._ops.OpOverload,
|
||||
tensor_arg: torch.Tensor,
|
||||
compute_mesh: DeviceMesh,
|
||||
) -> DTensorSpec:
|
||||
# util function to produce a replicate spec for a scalar tensor arg/kwarg
|
||||
if tensor_arg.numel() == 1 and tensor_arg.ndim == 1:
|
||||
warnings.warn(
|
||||
"Found a non-scalar tensor with numel=1 and ndim!=0, "
|
||||
"we are implicitly creating a replicated DTensor for it. "
|
||||
"However, please consider changing it to a scalar tensor "
|
||||
"or explicitly create a DTensor under distributed environment.",
|
||||
stacklevel=2,
|
||||
)
|
||||
|
||||
if tensor_arg.numel() == 1 or self._allow_implicit_replication:
|
||||
# scalar tensor can be safely treated as replicated
|
||||
replication_spec = DTensorSpec(
|
||||
compute_mesh,
|
||||
(Replicate(),) * compute_mesh.ndim,
|
||||
tensor_meta=TensorMeta(
|
||||
shape=tensor_arg.shape,
|
||||
stride=tensor_arg.stride(),
|
||||
dtype=tensor_arg.dtype,
|
||||
),
|
||||
)
|
||||
else:
|
||||
raise RuntimeError(
|
||||
f"{op_call}: got mixed torch.Tensor and DTensor, need to convert all"
|
||||
" torch.Tensor to DTensor before calling distributed operators!"
|
||||
" Please see https://docs.pytorch.org/docs/main/distributed.tensor.html#mixed-tensor-and-dtensor-operations"
|
||||
" for more details."
|
||||
)
|
||||
return replication_spec
|
||||
@@ -0,0 +1,763 @@
|
||||
import hashlib
|
||||
import itertools
|
||||
import math
|
||||
from collections import defaultdict
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, cast, NamedTuple
|
||||
|
||||
import torch
|
||||
from torch.distributed.device_mesh import DeviceMesh
|
||||
from torch.distributed.tensor.placement_types import (
|
||||
_is_shard_like,
|
||||
_MaskPartial,
|
||||
_StridedShard,
|
||||
Partial,
|
||||
Placement,
|
||||
Replicate,
|
||||
Shard,
|
||||
)
|
||||
from torch.utils._debug_mode import _stringify_shape
|
||||
from torch.utils._dtype_abbrs import dtype_abbrs
|
||||
|
||||
|
||||
# Defined here (not in placement_types.py) because decoding split_factor into
|
||||
# a shard order is a DTensorSpec concern — placement_types doesn't know about
|
||||
# shard orders.
|
||||
class _StridedShardNotDecodableError(ValueError):
|
||||
"""Raised when _StridedShard split_factor cannot be decoded into a shard order."""
|
||||
|
||||
|
||||
class ShardOrderEntry(NamedTuple):
|
||||
"""
|
||||
Represents how a single tensor dimension is sharded across mesh dimensions.
|
||||
|
||||
Attributes:
|
||||
tensor_dim: The tensor dimension being sharded (e.g., 0, 1, 2 for a 3D tensor).
|
||||
mesh_dims: Tuple of mesh dimensions across which this tensor dimension is sharded,
|
||||
in execution order. The first mesh dim is applied first, second is applied
|
||||
second, etc. This tuple is guaranteed to be non-empty.
|
||||
|
||||
Examples:
|
||||
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_DISTRIBUTED)
|
||||
>>> # Tensor dim 1 sharded across mesh dim 2, then mesh dim 0
|
||||
>>> ShardOrderEntry(tensor_dim=1, mesh_dims=(2, 0))
|
||||
|
||||
>>> # Tensor dim 0 sharded only on mesh dim 1
|
||||
>>> ShardOrderEntry(tensor_dim=0, mesh_dims=(1,))
|
||||
"""
|
||||
|
||||
tensor_dim: int
|
||||
mesh_dims: tuple[int, ...] # guaranteed to be non-empty
|
||||
|
||||
|
||||
# Type alias for the complete shard order specification
|
||||
# A tuple of ShardOrderEntry, one per sharded tensor dimension
|
||||
#
|
||||
# Example:
|
||||
# shard_order = (
|
||||
# ShardOrderEntry(tensor_dim=0, mesh_dims=(1,)),
|
||||
# ShardOrderEntry(tensor_dim=2, mesh_dims=(0, 3)),
|
||||
# )
|
||||
# This means:
|
||||
# - Tensor dimension 0 is sharded on mesh dimension 1
|
||||
# - Tensor dimension 2 is sharded on mesh dimension 0 first, then mesh dimension 3
|
||||
ShardOrder = tuple[ShardOrderEntry, ...]
|
||||
|
||||
|
||||
class TensorMeta(NamedTuple):
|
||||
# simple named tuple to represent tensor metadata
|
||||
# intentionally to stay simple only for sharding
|
||||
# propagation purposes.
|
||||
shape: torch.Size
|
||||
stride: tuple[int, ...]
|
||||
dtype: torch.dtype
|
||||
|
||||
|
||||
# used internally to propagate the placements
|
||||
@dataclass
|
||||
class DTensorSpec:
|
||||
mesh: DeviceMesh
|
||||
placements: tuple[Placement, ...]
|
||||
|
||||
# tensor meta will only be set during sharding propagation
|
||||
tensor_meta: TensorMeta | None = None
|
||||
|
||||
# When a tensor dimension is sharded across multiple mesh axes,
|
||||
# `shard_order` specifies the sequence in which these shardings are applied.
|
||||
# This order determines how tensor shards are mapped and distributed across
|
||||
# devices.
|
||||
#
|
||||
# Example:
|
||||
# For a tensor of shape [8, 16] and a 3D device mesh, if dim 0 is sharded over
|
||||
# mesh dim 1, and dim 1 is sharded over mesh dim 0 and then mesh dim 2,
|
||||
# the shard_order would be:
|
||||
# shard_order = (
|
||||
# ShardOrderEntry(tensor_dim=0, mesh_dims=(1,)),
|
||||
# ShardOrderEntry(tensor_dim=1, mesh_dims=(0, 2)),
|
||||
# )
|
||||
shard_order: ShardOrder = None # type: ignore[assignment]
|
||||
|
||||
# When True, _StridedShard placements encode the shard order, and the
|
||||
# shard_order field must be left as None (it will be derived on demand).
|
||||
# Set explicitly to False to treat _StridedShard as a regular Shard,
|
||||
# e.g., in view propagation.
|
||||
use_strided_shard_as_shard_order: bool | None = None
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if not isinstance(self.placements, tuple):
|
||||
self.placements = tuple(self.placements)
|
||||
if self.use_strided_shard_as_shard_order is None:
|
||||
if any(isinstance(p, _StridedShard) for p in self.placements):
|
||||
self.use_strided_shard_as_shard_order = True
|
||||
else:
|
||||
self.use_strided_shard_as_shard_order = False
|
||||
if self.use_strided_shard_as_shard_order:
|
||||
if self.shard_order is not None:
|
||||
raise ValueError(
|
||||
"DTensorSpec doesn't allow specify shard_order when "
|
||||
"use_strided_shard_as_shard_order is True. This may result "
|
||||
"in conflicting shard order."
|
||||
)
|
||||
else:
|
||||
if self.shard_order is None:
|
||||
self.shard_order = self.compute_default_shard_order(self.placements)
|
||||
|
||||
self._hash: int | None = None
|
||||
|
||||
@staticmethod
|
||||
def _normalize_placements_into_shard_order(
|
||||
placements: tuple[Placement, ...],
|
||||
mesh: DeviceMesh,
|
||||
use_strided_shard_as_shard_order: bool = True,
|
||||
) -> tuple[tuple[Placement, ...], ShardOrder]:
|
||||
# If use_strided_shard_as_shard_order, it means the StridedShard/Shard
|
||||
# combinations should be interpreted as shard order.
|
||||
if use_strided_shard_as_shard_order:
|
||||
# _StridedShard in placements, try check if it can be decoded as shard order
|
||||
shard_order = DTensorSpec._maybe_convert_StridedShard_to_shard_order(
|
||||
placements, mesh
|
||||
)
|
||||
if shard_order is None:
|
||||
raise _StridedShardNotDecodableError(
|
||||
f"_StridedShard placements {placements} cannot be decoded "
|
||||
"into a corresponding shard_order"
|
||||
)
|
||||
normalized_placements = tuple(
|
||||
[
|
||||
p if not isinstance(p, _StridedShard) else Shard(p.dim)
|
||||
for p in placements
|
||||
]
|
||||
)
|
||||
return normalized_placements, shard_order
|
||||
else:
|
||||
return placements, DTensorSpec.compute_default_shard_order(placements)
|
||||
|
||||
@staticmethod
|
||||
def compute_default_shard_order(
|
||||
placements: tuple[Placement, ...],
|
||||
) -> ShardOrder:
|
||||
"""
|
||||
Compute the default shard order from placements.
|
||||
|
||||
Returns a ShardOrder where each ShardOrderEntry maps a tensor dimension
|
||||
to the mesh dimensions it's sharded on, in left-to-right order.
|
||||
|
||||
Args:
|
||||
placements: Tuple of Placement objects representing how a tensor is
|
||||
distributed across mesh dimensions.
|
||||
"""
|
||||
# follow default left-to-right device order if shard_order is not specified
|
||||
tensor_dim_to_mesh_dims: defaultdict[int, list[int]] = defaultdict(list)
|
||||
mesh_ndim = len(placements)
|
||||
for mesh_dim in range(mesh_ndim):
|
||||
if _is_shard_like(placements[mesh_dim]):
|
||||
placement = placements[mesh_dim]
|
||||
shard_dim = placement.dim # pyrefly: ignore [missing-attribute]
|
||||
if shard_dim < 0:
|
||||
raise AssertionError(
|
||||
f"Shard dim {shard_dim} in placements {placements} must be normalized"
|
||||
)
|
||||
tensor_dim_to_mesh_dims[shard_dim].append(mesh_dim)
|
||||
|
||||
# Convert dict into ShardOrderEntry tuples
|
||||
default_shard_order = tuple(
|
||||
ShardOrderEntry(tensor_dim=key, mesh_dims=tuple(value))
|
||||
for key, value in sorted(tensor_dim_to_mesh_dims.items())
|
||||
if value
|
||||
)
|
||||
return default_shard_order
|
||||
|
||||
@staticmethod
|
||||
def _convert_shard_order_to_StridedShard(
|
||||
shard_order: ShardOrder, placements: tuple[Placement, ...], mesh: DeviceMesh
|
||||
) -> tuple[Placement, ...]:
|
||||
"""
|
||||
Convert ShardOrder to placements with _StridedShard.
|
||||
|
||||
This function converts a ShardOrder specification into a tuple of Placement objects,
|
||||
using _StridedShard when a tensor dimension is sharded across multiple mesh dimensions
|
||||
in a non-default order. The split_factor of each _StridedShard is determined by the
|
||||
product of mesh dimension sizes that appear earlier in the shard order but later in
|
||||
the placement tuple.
|
||||
|
||||
Args:
|
||||
shard_order: ShardOrder specification indicating which tensor dimensions are
|
||||
sharded on which mesh dimensions and in what execution order.
|
||||
placements: Tuple of Placement objects that does not contain _StridedShard.
|
||||
mesh: DeviceMesh containing the size information for each mesh dimension.
|
||||
|
||||
Returns:
|
||||
Updated tuple of Placement objects with Shard or _StridedShard placements.
|
||||
|
||||
Algorithm:
|
||||
For each ShardOrderEntry in shard_order:
|
||||
- For each mesh dimension in the entry's mesh_dims (in order):
|
||||
- Calculate split_factor as the product of mesh sizes for all mesh dimensions
|
||||
that appear:
|
||||
1. Earlier in the shard order (lower index in mesh_dims), and
|
||||
2. Later in the placement tuple (higher mesh dimension index)
|
||||
- If split_factor == 1: use normal Shard
|
||||
- Otherwise: use _StridedShard with the calculated split_factor
|
||||
|
||||
Example:
|
||||
>>> # xdoctest: +SKIP("Requires DeviceMesh")
|
||||
>>> # Tensor dimension 0 sharded on mesh dims [2, 0, 1] in that order
|
||||
>>> # mesh = DeviceMesh([4, 3, 2]) # sizes: mesh[0]=4, mesh[1]=3, mesh[2]=2
|
||||
>>> shard_order = (ShardOrderEntry(tensor_dim=0, mesh_dims=(2, 0, 1)),)
|
||||
>>> placements = (Shard(0), Shard(0), Shard(0))
|
||||
>>> # For mesh_dim=2 (index 0 in mesh_dims): no earlier dims, split_factor=1
|
||||
>>> # -> placements[2] = Shard(0)
|
||||
>>> # For mesh_dim=0 (index 1 in mesh_dims): mesh_dim=2 is earlier and has index 2>0
|
||||
>>> # -> split_factor = mesh.size(2) = 2
|
||||
>>> # -> placements[0] = _StridedShard(0, split_factor=2)
|
||||
>>> # For mesh_dim=1 (index 2 in mesh_dims): mesh_dim=2 is earlier and has index 2>1
|
||||
>>> # -> split_factor = mesh.size(2) = 2
|
||||
>>> # -> placements[1] = _StridedShard(0, split_factor=2)
|
||||
>>> # Result: (_StridedShard(0, sf=2), _StridedShard(0, sf=2), Shard(0))
|
||||
"""
|
||||
placements_list = list(placements)
|
||||
for entry in shard_order:
|
||||
tensor_dim = entry.tensor_dim
|
||||
mesh_dims = entry.mesh_dims
|
||||
for idx in range(len(mesh_dims)):
|
||||
# TODO(zpcore): split_factor from `view` and `shard order`
|
||||
# should be able to be multiplied into one. Need to loosen the
|
||||
# condition here.
|
||||
mesh_dim = mesh_dims[idx]
|
||||
if type(placements[mesh_dim]) is not Shard:
|
||||
raise ValueError(
|
||||
f"Only Shard placement can be converted to _StridedShard, "
|
||||
f"found {placements[mesh_dim]} in {placements=}."
|
||||
)
|
||||
split_factor = math.prod(
|
||||
mesh.size(i) for i in mesh_dims[:idx] if i > mesh_dim
|
||||
)
|
||||
if split_factor == 1:
|
||||
# use normal Shard
|
||||
placements_list[mesh_dim] = Shard(tensor_dim)
|
||||
else:
|
||||
placements_list[mesh_dim] = _StridedShard(
|
||||
tensor_dim, split_factor=split_factor
|
||||
)
|
||||
return tuple(placements_list)
|
||||
|
||||
@staticmethod
|
||||
def _maybe_convert_StridedShard_to_shard_order(
|
||||
placements: tuple[Placement, ...], mesh: DeviceMesh
|
||||
) -> ShardOrder | None:
|
||||
"""
|
||||
Try to convert _StridedShard placements to ShardOrder.
|
||||
|
||||
This is the inverse of `_convert_shard_order_to_StridedShard`. It reconstructs the shard
|
||||
order by examining the split_factor of each _StridedShard and determining its position
|
||||
in the execution order. If the _StridedShard configuration cannot be represented as a
|
||||
valid ShardOrder (i.e., there's no shard order that produces the observed split_factors),
|
||||
this function returns None.
|
||||
|
||||
Args:
|
||||
placements: Tuple of Placement objects that may contain _StridedShard.
|
||||
mesh: DeviceMesh containing the size information for each mesh dimension.
|
||||
|
||||
Returns:
|
||||
ShardOrder if conversion is possible, None otherwise. For placements without
|
||||
_StridedShard, returns the default shard order.
|
||||
|
||||
Algorithm:
|
||||
1. If no _StridedShard in placements, return default shard order
|
||||
2. Create an empty list for each tensor dimension to represent mesh dim ordering
|
||||
3. Iterate through placements in reverse order (right to left):
|
||||
- For each Shard/_StridedShard on a tensor dimension:
|
||||
- Extract its split_factor (1 for Shard, split_factor for _StridedShard)
|
||||
- Find the position in mesh_dims_order where accumulated_sf equals split_factor
|
||||
- accumulated_sf is the product of mesh sizes of mesh dimensions that appear
|
||||
earlier in mesh_dims_order (lower indices)
|
||||
- Insert mesh_dim at the found position
|
||||
4. If no valid position found for any split_factor, return None (unable to convert)
|
||||
5. Construct ShardOrderEntry for each tensor dimension from mesh_dims_order
|
||||
|
||||
Example:
|
||||
>>> # xdoctest: +SKIP("Requires DeviceMesh")
|
||||
>>> # mesh = DeviceMesh([4, 3, 2]) # sizes: mesh[0]=4, mesh[1]=3, mesh[2]=2
|
||||
>>> # placements = (_StridedShard(0, sf=2), _StridedShard(0, sf=2), Shard(0))
|
||||
>>> # Process tensor_dim=0 from right to left:
|
||||
>>> # - mesh_dim=2: Shard(0) with sf=1
|
||||
>>> # Try position 0: accumulated_sf=1, matches! Insert at position 0
|
||||
>>> # Current mesh_dims_order order: [2]
|
||||
>>> # - mesh_dim=1: _StridedShard(0, sf=2) with sf=2
|
||||
>>> # Try position 0: accumulated_sf=1, no match
|
||||
>>> # Try position 1: accumulated_sf=1*mesh.size(2)=2, matches! Insert at position 1
|
||||
>>> # Current mesh_dims_order order: [2, 1]
|
||||
>>> # - mesh_dim=0: _StridedShard(0, sf=2) with sf=2
|
||||
>>> # Try position 0: accumulated_sf=1, no match
|
||||
>>> # Try position 1: accumulated_sf=1*mesh.size(2)=2, matches! Insert at position 1
|
||||
>>> # Final mesh_dims_order order: [2, 0, 1]
|
||||
>>> # Result: ShardOrder((ShardOrderEntry(tensor_dim=0, mesh_dims=(2, 0, 1)),))
|
||||
>>> # This means: first shard on mesh_dim=2, then mesh_dim=0, then mesh_dim=1
|
||||
|
||||
Note:
|
||||
This function validates that _StridedShard can be represented as a ShardOrder.
|
||||
Not all _StridedShard configurations are valid - the split_factor must match
|
||||
the product of mesh sizes in some execution order.
|
||||
"""
|
||||
if not any(isinstance(p, _StridedShard) for p in placements):
|
||||
return DTensorSpec.compute_default_shard_order(placements)
|
||||
max_tensor_dim = max([i.dim for i in placements if _is_shard_like(i)]) + 1
|
||||
shard_order = []
|
||||
|
||||
tensor_dim_to_mesh_dims_order: list[list[int]] = [
|
||||
[] for i in range(max_tensor_dim)
|
||||
]
|
||||
for mesh_dim in reversed(range(len(placements))):
|
||||
cur_placement = placements[mesh_dim]
|
||||
if _is_shard_like(cur_placement):
|
||||
tensor_dim = cur_placement.dim
|
||||
mesh_dims_order = tensor_dim_to_mesh_dims_order[tensor_dim]
|
||||
cur_sf = 1
|
||||
if isinstance(cur_placement, _StridedShard):
|
||||
cur_sf = cur_placement.split_factor
|
||||
accumulated_sf = 1
|
||||
find_order = False
|
||||
for i in range(len(mesh_dims_order) + 1):
|
||||
if accumulated_sf == cur_sf:
|
||||
mesh_dims_order.insert(i, mesh_dim)
|
||||
find_order = True
|
||||
break
|
||||
if i < len(mesh_dims_order):
|
||||
accumulated_sf *= mesh.size(mesh_dims_order[i])
|
||||
if not find_order:
|
||||
# _StridedShard is not convertible to ShardOrder
|
||||
return None
|
||||
else:
|
||||
if not isinstance(cur_placement, Replicate | Partial | _MaskPartial):
|
||||
raise ValueError(
|
||||
f"Unsupported placement type {type(cur_placement)} encountered in "
|
||||
f"{placements}; expected Replicate, Partial, or _MaskPartial."
|
||||
)
|
||||
for tensor_dim in range(max_tensor_dim):
|
||||
if len(tensor_dim_to_mesh_dims_order[tensor_dim]) > 0:
|
||||
shard_order.append(
|
||||
ShardOrderEntry(
|
||||
tensor_dim=tensor_dim,
|
||||
mesh_dims=tuple(tensor_dim_to_mesh_dims_order[tensor_dim]),
|
||||
)
|
||||
)
|
||||
return tuple(shard_order)
|
||||
|
||||
def _verify_shard_order(self, shard_order: ShardOrder) -> None:
|
||||
"""Verify that the shard_order is valid and matches the placements."""
|
||||
total_shard = 0
|
||||
if any(isinstance(p, _StridedShard) for p in self.placements):
|
||||
# _StridedShard shard_order validation not yet supported;
|
||||
# the Shard-only checks below (line 390, 394) would fail.
|
||||
return
|
||||
prev_tensor_dim = -1
|
||||
for entry in shard_order:
|
||||
tensor_dim = entry.tensor_dim
|
||||
mesh_dims = entry.mesh_dims
|
||||
if len(mesh_dims) <= 0:
|
||||
raise AssertionError(f"shard_order {shard_order} has empty mesh dim")
|
||||
if tensor_dim < 0:
|
||||
raise AssertionError(
|
||||
f"shard_order {shard_order} has invalid tensor dim {tensor_dim}"
|
||||
)
|
||||
if tensor_dim <= prev_tensor_dim:
|
||||
raise AssertionError("tensor dim should be sorted in shard_order")
|
||||
prev_tensor_dim = tensor_dim
|
||||
total_shard += len(mesh_dims)
|
||||
for mesh_dim in mesh_dims:
|
||||
if not (0 <= mesh_dim < len(self.placements)):
|
||||
raise AssertionError(
|
||||
f"shard_order {shard_order} has invalid mesh dim {mesh_dims}"
|
||||
)
|
||||
if self.placements[mesh_dim] != Shard(tensor_dim):
|
||||
raise AssertionError(
|
||||
f"placement[{mesh_dim}] doesn't have a matching shard in shard_order"
|
||||
)
|
||||
if total_shard != sum(1 for p in self.placements if isinstance(p, Shard)):
|
||||
raise AssertionError
|
||||
|
||||
def __setattr__(self, attr: str, value: Any) -> None:
|
||||
if attr == "shard_order" and value is not None:
|
||||
self._verify_shard_order(value)
|
||||
super().__setattr__(attr, value)
|
||||
# Make sure to recompute the hash in case any of the hashed attributes
|
||||
# change (though we do not expect `mesh`, `placements` or `shard_order`
|
||||
# to change)
|
||||
if hasattr(self, "_hash") and attr in (
|
||||
"mesh",
|
||||
"placements",
|
||||
"tensor_meta",
|
||||
"shard_order",
|
||||
):
|
||||
self._hash = None
|
||||
# This assert was triggered by buggy handling for dict outputs in some
|
||||
# FX passes, where you accidentally iterate over a dict and try to put
|
||||
# keys into TensorMeta. See https://github.com/pytorch/pytorch/issues/157919
|
||||
if attr == "tensor_meta" and value is not None:
|
||||
from torch.fx.passes.shape_prop import TensorMetadata
|
||||
|
||||
# TODO: the TensorMetadata arises from
|
||||
# test/distributed/tensor/experimental/test_tp_transform.py::TensorParallelTest::test_tp_transform_e2e
|
||||
# but I actually can't reproduce it, maybe it is also a bug!
|
||||
if not isinstance(value, TensorMeta | TensorMetadata):
|
||||
raise AssertionError(repr(value))
|
||||
|
||||
def _hash_key(self) -> tuple[Any, ...]:
|
||||
"""Return the tuple used for hashing. Used by both __hash__ and _stable_hash."""
|
||||
if self.tensor_meta is not None:
|
||||
return (
|
||||
self.mesh,
|
||||
self.placements,
|
||||
self.shard_order,
|
||||
self.tensor_meta.shape,
|
||||
self.tensor_meta.stride,
|
||||
self.tensor_meta.dtype,
|
||||
)
|
||||
return (self.mesh, self.placements, self.shard_order)
|
||||
|
||||
def _hash_impl(self) -> int:
|
||||
# hashing and equality check for DTensorSpec are used to cache the sharding
|
||||
# propagation results. We only need to consider the mesh, placements, shape
|
||||
# dtype and stride.
|
||||
# Caveat: we need to keep this in mind and sync hash and eq if we add more
|
||||
# fields to them.
|
||||
return hash(self._hash_key())
|
||||
|
||||
def __hash__(self) -> int:
|
||||
# We lazily cache the spec to avoid recomputing the hash upon each
|
||||
# use, where we make sure to update the hash when the `tensor_meta`
|
||||
# changes by overriding `__setattr__`. This must be lazy so that Dynamo
|
||||
# does not try to hash non-singleton `SymInt`s for the stride.
|
||||
if self._hash is None:
|
||||
self._hash = self._hash_impl()
|
||||
return self._hash
|
||||
|
||||
def _stable_hash(self) -> str:
|
||||
"""
|
||||
Return a stable hash for AOT autograd caching.
|
||||
[See note: Tensor subclass stable hashing for AOT autograd cache]
|
||||
"""
|
||||
# Get hash key, but replace mesh with its stable hash
|
||||
key = self._hash_key()
|
||||
# First element is mesh, replace with its stable hash
|
||||
stable_key = (self.mesh._stable_hash(),) + key[1:]
|
||||
return hashlib.blake2b(repr(stable_key).encode(), digest_size=16).hexdigest()
|
||||
|
||||
def _check_equals(self, other: object, skip_shapes: bool = False) -> bool:
|
||||
if not (
|
||||
isinstance(other, DTensorSpec)
|
||||
and self.mesh == other.mesh
|
||||
and self.placements == other.placements
|
||||
and self.shard_order == other.shard_order
|
||||
):
|
||||
return False
|
||||
if self.tensor_meta is None or other.tensor_meta is None:
|
||||
return self.tensor_meta == other.tensor_meta
|
||||
|
||||
if skip_shapes:
|
||||
return self.tensor_meta.dtype == other.tensor_meta.dtype
|
||||
return (
|
||||
self.tensor_meta.shape == other.tensor_meta.shape # type: ignore[union-attr]
|
||||
and self.tensor_meta.stride == other.tensor_meta.stride # type: ignore[union-attr]
|
||||
and self.tensor_meta.dtype == other.tensor_meta.dtype # type: ignore[union-attr]
|
||||
)
|
||||
|
||||
def __eq__(self, other: object, /) -> bool:
|
||||
return self._check_equals(other)
|
||||
|
||||
def __str__(self) -> str:
|
||||
"""
|
||||
human readable representation of the DTensorSpec
|
||||
"""
|
||||
placement_str = self.format_shard_order_str(self.placements, self.shard_order)
|
||||
if self.tensor_meta is not None:
|
||||
tensor_shape = _stringify_shape(self.tensor_meta.shape)
|
||||
tensor_dtype = dtype_abbrs[self.tensor_meta.dtype]
|
||||
else:
|
||||
tensor_shape = "unknown shape"
|
||||
tensor_dtype = "unknown dtype"
|
||||
|
||||
return f"Spec({tensor_dtype}{tensor_shape}({placement_str}))"
|
||||
|
||||
@staticmethod
|
||||
def is_default_device_order(shard_order: ShardOrder) -> bool:
|
||||
"""
|
||||
Check if the device order is the default left-to-right order.
|
||||
"""
|
||||
if shard_order is None:
|
||||
# Missing `shard_order` attribute, possibly due to _StridedShard in
|
||||
# placements.
|
||||
return False
|
||||
for entry in shard_order:
|
||||
mesh_dims = entry.mesh_dims
|
||||
is_increasing = all(
|
||||
prev < nxt for prev, nxt in itertools.pairwise(mesh_dims)
|
||||
)
|
||||
if not is_increasing:
|
||||
return False
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def format_shard_order_str(
|
||||
placements: tuple[Placement, ...],
|
||||
shard_order: ShardOrder | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
Format DTensor sharding information as a human-readable string.
|
||||
|
||||
This method formats the sharding pattern in mesh-centric order, showing the placement
|
||||
for each mesh dimension sequentially. When a tensor dimension is sharded across multiple
|
||||
mesh dimensions, the order index indicates the execution sequence of the sharding operations.
|
||||
|
||||
Args:
|
||||
placements: Tuple of placement objects for each mesh dimension.
|
||||
shard_order: Optional ShardOrder specifying the sharding order.
|
||||
|
||||
Returns:
|
||||
String representation of the sharding pattern in mesh-centric format.
|
||||
|
||||
Example:
|
||||
For a 3D tensor on a 2x2x2x2 mesh (16 devices) with::
|
||||
|
||||
placements = [Partial(), Shard(1), Shard(1), Replicate()]
|
||||
shard_order = (ShardOrderEntry(tensor_dim=1, mesh_dims=(2, 1)),)
|
||||
|
||||
Mesh configuration:
|
||||
- mesh_dim_0: Partial reduction (sum)
|
||||
- mesh_dim_1: Shard tensor dimension 1 (executed second, order index 1)
|
||||
- mesh_dim_2: Shard tensor dimension 1 (executed first, order index 0)
|
||||
- mesh_dim_3: Replicate
|
||||
|
||||
Output: ``"PS(1)[1]S(1)[0]R"``
|
||||
|
||||
Explanation:
|
||||
- ``P``: mesh dimension 0 has partial reduction
|
||||
- ``S(1)[1]``: mesh dimension 1 shards tensor dimension 1 (order index 1 means second)
|
||||
- ``S(1)[0]``: mesh dimension 2 shards tensor dimension 1 (order index 0 means first)
|
||||
- ``R``: mesh dimension 3 replicates
|
||||
|
||||
The format follows mesh dimension order (0, 1, 2, 3), and when a tensor dimension
|
||||
is sharded across multiple mesh dimensions, the bracketed index shows the execution
|
||||
order: ``[0]`` is executed first, ``[1]`` is executed second, etc.
|
||||
"""
|
||||
out_str = ""
|
||||
# native dtensor-style sharding representation: map from mesh
|
||||
# dim to tensor dim
|
||||
for mesh_dim, placement in enumerate(placements):
|
||||
if _is_shard_like(placement):
|
||||
if shard_order is not None:
|
||||
for entry in shard_order:
|
||||
tensor_dim = entry.tensor_dim
|
||||
mesh_dims = entry.mesh_dims
|
||||
|
||||
if placement.dim == tensor_dim:
|
||||
if mesh_dim not in mesh_dims:
|
||||
raise AssertionError
|
||||
if len(mesh_dims) > 1:
|
||||
out_str += f"{placement}[{mesh_dims.index(mesh_dim)}]"
|
||||
else:
|
||||
# no need to show device order if the tensor dim is
|
||||
# only sharded in one mesh dim
|
||||
out_str += str(placement)
|
||||
break
|
||||
else:
|
||||
out_str += str(placement)
|
||||
else:
|
||||
out_str += str(placement)
|
||||
return out_str
|
||||
|
||||
@property
|
||||
def shape(self) -> torch.Size:
|
||||
if self.tensor_meta is None:
|
||||
raise ValueError("tensor_meta is not set")
|
||||
return self.tensor_meta.shape
|
||||
|
||||
@property
|
||||
def stride(self) -> tuple[int, ...]:
|
||||
if self.tensor_meta is None:
|
||||
raise ValueError("tensor_meta is not set")
|
||||
return self.tensor_meta.stride
|
||||
|
||||
@property
|
||||
def ndim(self) -> int:
|
||||
if self.tensor_meta is None:
|
||||
raise ValueError("tensor_meta is not set")
|
||||
return len(self.tensor_meta.shape)
|
||||
|
||||
@property
|
||||
def num_shards(self) -> int:
|
||||
num_shards = 1
|
||||
for i, placement in enumerate(self.placements):
|
||||
if _is_shard_like(placement):
|
||||
num_shards *= self.mesh.size(i)
|
||||
return num_shards
|
||||
|
||||
@property
|
||||
def device_mesh(self) -> DeviceMesh:
|
||||
# simple aliasing for the mesh field, make some
|
||||
# checks that mixes DTensor/DTensorSpec easier
|
||||
return self.mesh
|
||||
|
||||
@property
|
||||
def dim_map(self) -> list[int]:
|
||||
"""
|
||||
dim_map is a property we derive from `placements` of
|
||||
the distributed tensor. It simply return a list of ints
|
||||
where dim_map[i] denotes the sharding mapping to the mesh
|
||||
dimension, and len(dim_map) == dist_tensor.ndim
|
||||
dim_map[i] = -1: means tensor dim i replicate on mesh
|
||||
dim_map[i] = j: means tensor dim i shard on mesh dim j
|
||||
|
||||
For example, we have a dist tensor that have the shape of
|
||||
[18, 20, 30], and device_mesh([0, 1, 2, 3]), placements:
|
||||
[Shard(1)], the dim_map of this placement would be:
|
||||
[-1, 0, -1]. This representation is pretty helpful during
|
||||
sharding propagation where we could know exactly each
|
||||
tensor dimension is sharded or not.
|
||||
|
||||
Note that if placements contains `_Partial`, we have to
|
||||
explicitly deal with it, so that when we create a DTensorSpec
|
||||
with dim_map, we could properly record the pending sums.
|
||||
"""
|
||||
# dims mapping of dist tensor sharding
|
||||
# return size of tensor ndim, -1 represent replicate
|
||||
# and int >=0 represent shard on that device mesh dim
|
||||
r = [-1] * self.ndim
|
||||
for i, placement in enumerate(self.placements):
|
||||
if _is_shard_like(placement):
|
||||
shard_dim = placement.dim
|
||||
if r[shard_dim] > -1:
|
||||
raise ValueError(
|
||||
f"Tensor dim {shard_dim} is already sharded on mesh dim {r[shard_dim]},"
|
||||
" DTensor operator implementation does not support things like hybrid"
|
||||
" sharding strategies yet (i.e. [Shard(0), Shard(0)])"
|
||||
)
|
||||
r[shard_dim] = i
|
||||
return r
|
||||
|
||||
@property
|
||||
def num_shards_map(self) -> list[int]:
|
||||
"""
|
||||
dim_map is a property we derive from `placements` of
|
||||
the distributed tensor. Unlike `dim_map`, `num_shards_map`
|
||||
denotes how many shards each tensor dim has. Like `dim_map`:
|
||||
len(num_shards_map) == dist_tensor.ndim
|
||||
num_shards_map[i] = 1: means tensor dim i is not sharded
|
||||
num_shards_map[i] = j: means tensor dim i has j shards in total
|
||||
|
||||
For example, we have a dist tensor of shape [18, 20, 30],
|
||||
a device_mesh ([[0, 1, 2, 3], [4, 5, 6, 7]]), and placements
|
||||
([Shard(1), Shard(0)]), the num_shards_map of this distributed tensor
|
||||
would be: [4, 2, 1].
|
||||
"""
|
||||
r = [1] * self.ndim
|
||||
for i, placement in enumerate(self.placements):
|
||||
if _is_shard_like(placement):
|
||||
r[placement.dim] *= self.mesh.size(i)
|
||||
|
||||
return r
|
||||
|
||||
@property
|
||||
def sums(self) -> list[int]:
|
||||
"""
|
||||
sums is a property we derive from `placements` of the
|
||||
distributed tensor. It simply return a list of ints where
|
||||
sums[i] denotes the pending sum (partial) on mesh dim i
|
||||
"""
|
||||
return [
|
||||
idx
|
||||
for idx, placement in enumerate(self.placements)
|
||||
if placement.is_partial()
|
||||
]
|
||||
|
||||
@classmethod
|
||||
def from_dim_map(
|
||||
cls,
|
||||
mesh: DeviceMesh,
|
||||
dim_map: list[int],
|
||||
sums: list[int],
|
||||
tensor_meta: TensorMeta | None = None,
|
||||
) -> "DTensorSpec":
|
||||
"""
|
||||
Construct a DTensorSpec from dim_map list and pending sum.
|
||||
|
||||
Args:
|
||||
mesh (class:`DeviceMesh`): device mesh to be used in the DTensorSpec
|
||||
dim_map (List[int]): a list of integer that represents sharding on each
|
||||
tensor dimension, see `dim_map` property doc for details
|
||||
sums (List[int]): a list of integer that represents the dist tensor have
|
||||
pending sum on which device mesh dimension.
|
||||
tensor meta (TensorMeta): DTensor metadata
|
||||
|
||||
Return:
|
||||
a class:`DTensorSpec` object
|
||||
"""
|
||||
# by default replicate on device mesh dims
|
||||
placements: list[Placement] = [Replicate() for _ in range(mesh.ndim)]
|
||||
|
||||
# find all mesh dims that need pending reductions
|
||||
for s in sums:
|
||||
placements[s] = Partial()
|
||||
|
||||
for i, m in enumerate(dim_map):
|
||||
if m >= 0:
|
||||
placement = placements[m]
|
||||
if placement.is_shard(): # dim_map only produces Shard placements
|
||||
placement = cast(Shard, placement)
|
||||
raise RuntimeError(
|
||||
f"DeviceMesh dimension can't be mapped to two dimension of the same tensor: {i} and {placement.dim}"
|
||||
)
|
||||
elif placement.is_partial():
|
||||
raise RuntimeError(
|
||||
f"DeviceMesh dimension {m} cannot be both shard and partial!"
|
||||
)
|
||||
placements[m] = Shard(i)
|
||||
|
||||
return cls(mesh, tuple(placements), tensor_meta=tensor_meta)
|
||||
|
||||
def is_replicated(self) -> bool:
|
||||
"""
|
||||
return True if the current DTensorSpec replicates on all mesh dims (devices)
|
||||
"""
|
||||
return all(placement.is_replicate() for placement in self.placements)
|
||||
|
||||
def is_sharded(self) -> bool:
|
||||
"""
|
||||
return True if the current DTensorSpec uses Shard() or _StridedShard() placement on any mesh dims (devices)
|
||||
"""
|
||||
return any(_is_shard_like(placement) for placement in self.placements)
|
||||
|
||||
def shallow_copy_with_tensor_meta(
|
||||
self, tensor_meta: TensorMeta | None
|
||||
) -> "DTensorSpec":
|
||||
"""
|
||||
Shallow copy the DTensorSpec with a new tensor_meta.
|
||||
"""
|
||||
if tensor_meta is None:
|
||||
raise AssertionError("shallow copy with no tensor_meta!")
|
||||
return DTensorSpec(
|
||||
self.mesh,
|
||||
self.placements,
|
||||
tensor_meta=tensor_meta,
|
||||
use_strided_shard_as_shard_order=self.use_strided_shard_as_shard_order,
|
||||
)
|
||||
+305
@@ -0,0 +1,305 @@
|
||||
import operator
|
||||
from functools import reduce
|
||||
from typing import cast
|
||||
|
||||
import torch
|
||||
import torch.distributed._functional_collectives as funcol
|
||||
import torch.distributed.tensor._api as dtensor
|
||||
from torch.distributed.tensor._op_schema import OutputSharding
|
||||
from torch.distributed.tensor._utils import compute_local_shape_and_global_offset
|
||||
from torch.distributed.tensor.placement_types import (
|
||||
_StridedShard,
|
||||
Partial,
|
||||
Placement,
|
||||
Replicate,
|
||||
Shard,
|
||||
)
|
||||
|
||||
|
||||
# Mapping from argmin/argmax ops to their corresponding value ops (min/max)
|
||||
_ARGMINMAX_REDUCTION_OPS = {
|
||||
torch.ops.aten.argmax.default: torch.max,
|
||||
torch.ops.aten.argmin.default: torch.min,
|
||||
}
|
||||
|
||||
|
||||
def _get_output_sharding(
|
||||
op_call: torch._ops.OpOverload,
|
||||
args: tuple[object, ...],
|
||||
kwargs: dict[str, object],
|
||||
) -> OutputSharding:
|
||||
"""Get the output sharding for the given op."""
|
||||
op_info = dtensor.DTensor._op_dispatcher.unwrap_to_op_info(op_call, args, kwargs)
|
||||
dtensor.DTensor._op_dispatcher.sharding_propagator.propagate(op_info)
|
||||
output_sharding = op_info.output_sharding
|
||||
if output_sharding is None:
|
||||
raise AssertionError("output sharding should not be None")
|
||||
return output_sharding
|
||||
|
||||
|
||||
def _prep_arguments(
|
||||
op_call_repr: str,
|
||||
args: tuple[object, ...],
|
||||
kwargs: dict[str, object] | None,
|
||||
) -> tuple[
|
||||
torch.Tensor,
|
||||
torch.Size,
|
||||
"torch.distributed.device_mesh.DeviceMesh",
|
||||
tuple[Placement, ...],
|
||||
int | None,
|
||||
bool,
|
||||
]:
|
||||
"""
|
||||
Prepare arguments for nonlinear reduction ops.
|
||||
|
||||
Returns:
|
||||
local_tensor: The local tensor to operate on
|
||||
global_shape: The global shape of the DTensor
|
||||
device_mesh: The device mesh
|
||||
placements: The placements tuple
|
||||
dim: The reduction dimension (can be None)
|
||||
keepdim: Whether to keep the reduced dimension
|
||||
"""
|
||||
input_dtensor = cast(dtensor.DTensor, args[0])
|
||||
dim: int | None = None
|
||||
keepdim: bool = False
|
||||
|
||||
if not isinstance(input_dtensor, dtensor.DTensor):
|
||||
raise NotImplementedError
|
||||
if len(args) > 1:
|
||||
dim = cast(int, args[1])
|
||||
if len(args) > 2:
|
||||
keepdim = cast(bool, args[2])
|
||||
if kwargs:
|
||||
if "dim" in kwargs:
|
||||
dim = cast(int, kwargs["dim"])
|
||||
if "keepdim" in kwargs:
|
||||
keepdim = cast(bool, kwargs["keepdim"])
|
||||
device_mesh = input_dtensor.device_mesh
|
||||
placements = input_dtensor.placements
|
||||
|
||||
# check for partial placements and handle it as a replicate.
|
||||
if any(isinstance(p, Partial) for p in placements):
|
||||
target_placements = [
|
||||
Replicate() if isinstance(p, Partial) else p for p in placements
|
||||
]
|
||||
input_dtensor = input_dtensor.redistribute(
|
||||
device_mesh=device_mesh, placements=target_placements
|
||||
)
|
||||
placements = input_dtensor.placements
|
||||
local_tensor = input_dtensor.to_local()
|
||||
global_shape = input_dtensor.shape
|
||||
|
||||
return local_tensor, global_shape, device_mesh, placements, dim, keepdim
|
||||
|
||||
|
||||
def _get_expected_shape(
|
||||
local_tensor: torch.Tensor, dim: int | None, keepdim: bool
|
||||
) -> torch.Size:
|
||||
"""Compute the expected output shape after reduction."""
|
||||
input_shape = list(local_tensor.shape)
|
||||
if dim is None:
|
||||
expected_shape = (
|
||||
torch.Size([1] * len(input_shape)) if keepdim else torch.Size([])
|
||||
)
|
||||
elif keepdim:
|
||||
if input_shape:
|
||||
input_shape[dim] = 1
|
||||
expected_shape = torch.Size(input_shape)
|
||||
else:
|
||||
if input_shape:
|
||||
input_shape.pop(dim)
|
||||
expected_shape = torch.Size(input_shape)
|
||||
|
||||
return expected_shape
|
||||
|
||||
|
||||
def _collect_shard_mesh_dims(
|
||||
op_call_repr: str,
|
||||
local_tensor: torch.Tensor,
|
||||
placements: tuple[Placement, ...],
|
||||
dim: int | None,
|
||||
) -> list[int]:
|
||||
"""Collect mesh dimensions that are sharded along the reduction dimension."""
|
||||
shard_mesh_dims: list[int] = []
|
||||
for mesh_dim, p in enumerate(placements):
|
||||
if isinstance(p, Shard):
|
||||
if dim is None or p.dim == (dim if dim >= 0 else local_tensor.ndim + dim):
|
||||
shard_mesh_dims.append(mesh_dim)
|
||||
elif isinstance(p, _StridedShard):
|
||||
raise NotImplementedError(f"{op_call_repr} does not support _StridedShard!")
|
||||
return shard_mesh_dims
|
||||
|
||||
|
||||
def _convert_to_global_idxs(
|
||||
local_idx: torch.Tensor,
|
||||
global_shape: torch.Size,
|
||||
device_mesh: "torch.distributed.device_mesh.DeviceMesh",
|
||||
placements: tuple[Placement, ...],
|
||||
dim: int | None,
|
||||
) -> tuple[int, torch.Tensor]:
|
||||
"""Convert local indices to global indices."""
|
||||
local_shape, global_offset = compute_local_shape_and_global_offset(
|
||||
global_shape, device_mesh, placements
|
||||
)
|
||||
|
||||
if dim is None:
|
||||
# Convert flat local index → flat global index using arithmetic ops
|
||||
# instead of torch.unravel_index, which doesn't support SymInt shapes.
|
||||
gathered_idxs = torch.zeros_like(local_idx)
|
||||
remaining = local_idx
|
||||
for i in range(len(local_shape)):
|
||||
local_stride = reduce(operator.mul, local_shape[i + 1 :], 1)
|
||||
global_stride = reduce(operator.mul, global_shape[i + 1 :], 1)
|
||||
coord = remaining // local_stride
|
||||
remaining = remaining % local_stride
|
||||
gathered_idxs = gathered_idxs + (coord + global_offset[i]) * global_stride
|
||||
gather_dim = 0
|
||||
else:
|
||||
gather_dim = dim
|
||||
gathered_idxs = local_idx + global_offset[dim]
|
||||
return gather_dim, gathered_idxs
|
||||
|
||||
|
||||
def _gather_tensors(
|
||||
gather_dim: int,
|
||||
gathered_idxs: torch.Tensor,
|
||||
local_redux: torch.Tensor,
|
||||
device_mesh: "torch.distributed.device_mesh.DeviceMesh",
|
||||
shard_mesh_dims: list[int],
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Gather the min or max of the tensors and their corresponding indices.
|
||||
|
||||
Args:
|
||||
gather_dim: The dim to stack the collected min/max tensors.
|
||||
gathered_idxs: The local tensor holding the corresponding indices.
|
||||
local_redux: The local tensor holding the operator's value i.e. min/max.
|
||||
device_mesh: Device mesh of the DTensor.
|
||||
shard_mesh_dims: List of mesh dimensions that are sharded.
|
||||
|
||||
Returns:
|
||||
All gathered tensors (gathered_redux, gathered_idxs) of the reducing operator.
|
||||
"""
|
||||
gathered_redux = local_redux
|
||||
for mesh_dim in shard_mesh_dims:
|
||||
gathered_redux = funcol.all_gather_tensor(
|
||||
gathered_redux,
|
||||
gather_dim=gather_dim,
|
||||
group=(device_mesh, mesh_dim),
|
||||
)
|
||||
gathered_idxs = funcol.all_gather_tensor(
|
||||
gathered_idxs,
|
||||
gather_dim=gather_dim,
|
||||
group=(device_mesh, mesh_dim),
|
||||
)
|
||||
return gathered_redux, gathered_idxs
|
||||
|
||||
|
||||
def argminmax_handler(
|
||||
op_call: torch._ops.OpOverload,
|
||||
args: tuple[object, ...],
|
||||
kwargs: dict[str, object],
|
||||
) -> object:
|
||||
"""
|
||||
Handler for aten.argmin.default and aten.argmax.default ops.
|
||||
|
||||
This is a pure function handler that doesn't require instantiation.
|
||||
"""
|
||||
if op_call not in _ARGMINMAX_REDUCTION_OPS:
|
||||
raise NotImplementedError(f"Unsupported reduction op: {op_call}")
|
||||
|
||||
local_tensor, global_shape, device_mesh, placements, dim, keepdim = _prep_arguments(
|
||||
str(op_call), args, kwargs
|
||||
)
|
||||
output_sharding = _get_output_sharding(op_call, args, kwargs)
|
||||
|
||||
expected_shape = _get_expected_shape(local_tensor, dim, keepdim)
|
||||
shard_mesh_dims = _collect_shard_mesh_dims(
|
||||
str(op_call), local_tensor, placements, dim
|
||||
)
|
||||
|
||||
# Compute local reduction
|
||||
if dim is None:
|
||||
val_op = _ARGMINMAX_REDUCTION_OPS[op_call]
|
||||
# unsqueeze scalars to 1-d so they can be allgathered
|
||||
local_redux = val_op(local_tensor).unsqueeze(0)
|
||||
local_idx = op_call(local_tensor).unsqueeze(0)
|
||||
else:
|
||||
val_op = _ARGMINMAX_REDUCTION_OPS[op_call]
|
||||
local_redux, local_idx = val_op(local_tensor, dim=dim, keepdim=True)
|
||||
|
||||
if not shard_mesh_dims:
|
||||
return dtensor.DTensor._op_dispatcher.wrap(
|
||||
local_idx.reshape(expected_shape), output_sharding.output_spec
|
||||
)
|
||||
|
||||
gather_dim, gathered_idxs = _convert_to_global_idxs(
|
||||
local_idx, global_shape, device_mesh, placements, dim
|
||||
)
|
||||
gathered_redux, gather_idxs = _gather_tensors(
|
||||
gather_dim, gathered_idxs, local_redux, device_mesh, shard_mesh_dims
|
||||
)
|
||||
# Select the rank with the best value; use dim=0 when dim was None since
|
||||
# the scalars were unsqueezed to 1-d for gathering
|
||||
select_dim = 0 if dim is None else dim
|
||||
rank_winner = op_call(gathered_redux, select_dim, True)
|
||||
final_idx = torch.gather(gather_idxs, dim=gather_dim, index=rank_winner)
|
||||
|
||||
return dtensor.DTensor._op_dispatcher.wrap(
|
||||
final_idx.reshape(expected_shape), output_sharding.output_spec
|
||||
)
|
||||
|
||||
|
||||
def minmax_dim_handler(
|
||||
op_call: torch._ops.OpOverload,
|
||||
args: tuple[object, ...],
|
||||
kwargs: dict[str, object],
|
||||
) -> object:
|
||||
"""
|
||||
Handler for aten.min.dim and aten.max.dim ops.
|
||||
|
||||
This is a pure function handler that doesn't require instantiation.
|
||||
"""
|
||||
local_tensor, global_shape, device_mesh, placements, dim, keepdim = _prep_arguments(
|
||||
str(op_call), args, kwargs
|
||||
)
|
||||
output_sharding = _get_output_sharding(op_call, args, kwargs)
|
||||
|
||||
expected_shape = _get_expected_shape(local_tensor, dim, keepdim)
|
||||
shard_mesh_dims = _collect_shard_mesh_dims(
|
||||
str(op_call), local_tensor, placements, dim
|
||||
)
|
||||
|
||||
# Compute local reduction - min/max with dim always requires dim
|
||||
if dim is None:
|
||||
raise AssertionError
|
||||
local_redux, local_idx = op_call(local_tensor, dim=dim, keepdim=True)
|
||||
|
||||
if not shard_mesh_dims:
|
||||
return dtensor.DTensor._op_dispatcher.wrap(
|
||||
(
|
||||
local_redux.reshape(expected_shape),
|
||||
local_idx.reshape(expected_shape),
|
||||
),
|
||||
output_sharding.output_spec,
|
||||
)
|
||||
|
||||
gather_dim, gathered_idxs = _convert_to_global_idxs(
|
||||
local_idx, global_shape, device_mesh, placements, dim
|
||||
)
|
||||
|
||||
gathered_redux, gather_idxs = _gather_tensors(
|
||||
gather_dim, gathered_idxs, local_redux, device_mesh, shard_mesh_dims
|
||||
)
|
||||
# The op_call here is min/max with dim which returns (values, indices)
|
||||
final_redux, rank_winner = op_call(gathered_redux, dim, True)
|
||||
final_idx = torch.gather(gather_idxs, dim=gather_dim, index=rank_winner)
|
||||
|
||||
return dtensor.DTensor._op_dispatcher.wrap(
|
||||
(
|
||||
final_redux.reshape(expected_shape),
|
||||
final_idx.reshape(expected_shape),
|
||||
),
|
||||
output_sharding.output_spec,
|
||||
)
|
||||
@@ -0,0 +1,719 @@
|
||||
# mypy: allow-untyped-defs
|
||||
"""
|
||||
DTensor operator schema definitions and utilities.
|
||||
|
||||
This module defines the core data structures and utilities for describing and managing
|
||||
distributed tensor operations in PyTorch's DTensor system. It provides the foundational
|
||||
schema types used for sharding propagation, operator strategy selection, and distributed
|
||||
execution planning.
|
||||
|
||||
Key components:
|
||||
- OpSpec: Describes acceptable sharding placements for operations
|
||||
- OpStrategy: Represents the possible sharding strategies for an operator
|
||||
- TupleStrategy: Container for multiple strategies when ops have tuple/list of tensors input
|
||||
- OpSchema: Describes operator input/output schemas with DTensorSpecs
|
||||
- OutputSharding: Manages output sharding specifications and redistribution
|
||||
- RuntimeSchemaInfo: Runtime execution metadata for operators
|
||||
- OpInfo: Complete runtime operator execution information
|
||||
|
||||
These schema definitions enable the DTensor system to:
|
||||
1. Propagate tensor sharding information to the operator outputs
|
||||
2. Greedily select sharding strategies for distributed operations
|
||||
3. Plan and execute tensor redistributions when needed
|
||||
4. Cache sharding decisions for performance optimization
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
from dataclasses import dataclass
|
||||
from functools import cached_property
|
||||
from typing import Any
|
||||
from typing_extensions import deprecated
|
||||
|
||||
import torch
|
||||
from torch._C import (
|
||||
_DTensor_OpSchema_post_init,
|
||||
_DTensor_OpSchema_recompute_comparison_key,
|
||||
)
|
||||
from torch._ops import OpOverload
|
||||
from torch.distributed.device_mesh import DeviceMesh
|
||||
from torch.distributed.tensor._dtensor_spec import DTensorSpec, TensorMeta
|
||||
from torch.distributed.tensor.placement_types import Placement
|
||||
|
||||
|
||||
try:
|
||||
from torch.utils._cxx_pytree import (
|
||||
register_pytree_node,
|
||||
tree_leaves,
|
||||
tree_map_only,
|
||||
TreeSpec,
|
||||
)
|
||||
except ImportError:
|
||||
from torch.utils._pytree import ( # type: ignore[no-redef, assignment]
|
||||
register_pytree_node,
|
||||
tree_leaves,
|
||||
tree_map_only,
|
||||
TreeSpec,
|
||||
)
|
||||
|
||||
|
||||
# Common type aliases
|
||||
ArgsType = tuple[object, ...]
|
||||
KwargsType = dict[str, object]
|
||||
|
||||
PlacementList = list[Placement | None]
|
||||
|
||||
# ATen op schemas could have Tensor, Tuple[Tensor] and List[Tensor], so output type should
|
||||
# be the same set of possibilities.
|
||||
OutputSpecType = DTensorSpec | Sequence[DTensorSpec | None] | None
|
||||
|
||||
|
||||
def _rebuild_tensor_from_dtensor_meta(arg) -> object:
|
||||
"""
|
||||
This is used to propagate tensor metadata, must be under fake mode
|
||||
"""
|
||||
if arg.tensor_meta is None:
|
||||
raise AssertionError("DTensorSpec does not contain tensor_meta.")
|
||||
return torch.empty_strided(
|
||||
arg.tensor_meta.shape,
|
||||
arg.tensor_meta.stride,
|
||||
dtype=arg.tensor_meta.dtype,
|
||||
)
|
||||
|
||||
|
||||
def _pretty_print_spec(spec: object) -> str:
|
||||
if spec is None:
|
||||
return "None"
|
||||
elif isinstance(spec, DTensorSpec):
|
||||
return "".join([str(p) for p in spec.placements])
|
||||
elif isinstance(spec, Sequence):
|
||||
return "(" + ", ".join([_pretty_print_spec(s) for s in spec]) + ")"
|
||||
else:
|
||||
raise RuntimeError(f"Unknown spec type to print: spec={spec}")
|
||||
|
||||
|
||||
@dataclass
|
||||
class OpSpec:
|
||||
"""
|
||||
An OpSpec describes an acceptable sharding placements of an operation, with the
|
||||
specified DTensorSpecs for both the output and the inputs.
|
||||
|
||||
note: when the op return value is a single DTensor object, output_specs is
|
||||
DTensorSpec; when the return value is a tuple of Optional[DTensor],
|
||||
output_specs is a tuple of Optional[DTensorSpec].
|
||||
|
||||
note: we MUST produce an DTensorSpec for every output that is a Tensor. None
|
||||
entries only occur for non-Tensor outputs (e.g., operators that return Optional[Tensor],
|
||||
or non-Tensor outputs.)
|
||||
|
||||
invariant: the DeviceMesh on all DTensorSpec must be the same
|
||||
"""
|
||||
|
||||
# output_specs and input_specs are related: for this op, given these input_specs,
|
||||
# this is the way the output would look
|
||||
# Note: output_specs can be None for ops that don't return tensors (e.g., _linalg_check_errors)
|
||||
output_specs: DTensorSpec | tuple[DTensorSpec | None, ...] | None
|
||||
input_specs: Sequence[DTensorSpec] | None = None
|
||||
|
||||
"""
|
||||
redistribute_cost tells how expensive it is to redistribute a given input into the
|
||||
placement specified in this OpSpec.
|
||||
|
||||
outer list: one entry (list) per (tensor) input in the op's arg schema
|
||||
inner list: one entry (cost value) per possible sharding spec for that input
|
||||
|
||||
Example:
|
||||
-------
|
||||
another_op() -> tensor_a # another_op produces the output that becomes our first input
|
||||
my_op(tensor_a)
|
||||
|
||||
Let's assume this OpSpec's input_specs are [Replicate()],
|
||||
but another_op() supports 2 strategies (OpSpecs) which produce outputs of
|
||||
Replicate()
|
||||
Shard(0)
|
||||
|
||||
In this example, redistribute_costs would look like this
|
||||
[
|
||||
# one row representing "my_op's first input" (tensor_a)
|
||||
[
|
||||
# two entries, one for each strategies supported by another_op
|
||||
0.0, # cost of redistributing tensor_a from 'Replicate()'
|
||||
K, # cost of redistributing tensor_a from 'Shard(0)'
|
||||
],
|
||||
"""
|
||||
redistribute_cost: list[list[float]] | None = None
|
||||
|
||||
@cached_property
|
||||
def output_spec(self) -> DTensorSpec:
|
||||
"""
|
||||
This function requires that the strategy have exactly one DTensorSpec as the
|
||||
output spec. If the output_specs is a tuple, we throw an exception.
|
||||
"""
|
||||
if isinstance(self.output_specs, DTensorSpec):
|
||||
return self.output_specs
|
||||
else:
|
||||
raise ValueError(
|
||||
f"function output_spec expects a single DTensorSpec but got: {self.output_specs}"
|
||||
)
|
||||
|
||||
@cached_property
|
||||
def mesh(self):
|
||||
if isinstance(self.output_specs, DTensorSpec):
|
||||
return self.output_specs.mesh
|
||||
elif isinstance(self.output_specs, tuple):
|
||||
out_spec = self.output_specs[0]
|
||||
if not isinstance(out_spec, DTensorSpec):
|
||||
raise AssertionError
|
||||
return out_spec.mesh
|
||||
elif self.output_specs is None:
|
||||
# For no-output ops, get mesh from input_specs
|
||||
if self.input_specs is None or len(self.input_specs) <= 0:
|
||||
raise AssertionError(
|
||||
"Cannot determine mesh: output_specs is None and input_specs is empty"
|
||||
)
|
||||
return self.input_specs[0].mesh
|
||||
else:
|
||||
raise ValueError(
|
||||
f"function output_spec expects a single DTensorSpec or a tuple of DTensorSpec but got: {self.output_specs}"
|
||||
)
|
||||
|
||||
def input_spec(self, index: int = 0) -> DTensorSpec:
|
||||
if self.input_specs is None:
|
||||
raise AssertionError("input_specs of OpSpec is None!")
|
||||
if len(self.input_specs) <= index:
|
||||
raise AssertionError(
|
||||
f"Invalid index {index} for input_specs of length "
|
||||
f"{len(self.input_specs)}: {self.input_specs}"
|
||||
)
|
||||
return self.input_specs[index]
|
||||
|
||||
def __str__(self) -> str:
|
||||
if self.input_specs is not None:
|
||||
input_specs_str = f"{_pretty_print_spec(self.input_specs)} -> "
|
||||
else:
|
||||
input_specs_str = ""
|
||||
output_spec_str = _pretty_print_spec(self.output_specs)
|
||||
return f"{input_specs_str}{output_spec_str}"
|
||||
|
||||
def __hash__(self) -> int:
|
||||
if self.output_specs is None:
|
||||
output_hash = hash(None)
|
||||
elif isinstance(self.output_specs, DTensorSpec):
|
||||
output_hash = hash(self.output_specs)
|
||||
else:
|
||||
output_hash = hash(tuple(self.output_specs))
|
||||
input_hash = hash(tuple(self.input_specs)) if self.input_specs else 0
|
||||
return hash((output_hash, input_hash))
|
||||
|
||||
def __eq__(self, other: object) -> bool:
|
||||
if not isinstance(other, OpSpec):
|
||||
return False
|
||||
return (
|
||||
self.output_specs == other.output_specs
|
||||
and self.input_specs == other.input_specs
|
||||
)
|
||||
|
||||
|
||||
class StrategyType:
|
||||
"""
|
||||
Base class type for op strategy, We have two StrategyType:
|
||||
OpStrategy and TupleStrategy
|
||||
"""
|
||||
|
||||
|
||||
class OpStrategy(StrategyType):
|
||||
"""
|
||||
OpStrategy that consists of a list of sharding strategies associated with the op,
|
||||
where each strategy is an OpSpec that describes the acceptable input/output sharding.
|
||||
|
||||
invariant: the DeviceMesh on all OpSpec must be the same
|
||||
"""
|
||||
|
||||
def __init__(self, strategies: list[OpSpec]) -> None:
|
||||
super().__init__()
|
||||
self.strategies: list[OpSpec] = strategies
|
||||
|
||||
def __str__(self) -> str:
|
||||
strategy_list_str = ", ".join([str(strategy) for strategy in self.strategies])
|
||||
mesh_shape = self.mesh_shape
|
||||
return f"OpStrategy[{strategy_list_str}] @ mesh: {mesh_shape}"
|
||||
|
||||
def max_num_shards(self) -> int:
|
||||
"""
|
||||
Returns the max number of shards across all OpSpecs
|
||||
"""
|
||||
return max(strategy.output_spec.num_shards for strategy in self.strategies)
|
||||
|
||||
@property
|
||||
def mesh(self):
|
||||
return self.strategies[0].mesh
|
||||
|
||||
@property
|
||||
def mesh_shape(self):
|
||||
return self.strategies[0].mesh.shape
|
||||
|
||||
@property
|
||||
def ndim(self):
|
||||
return self.strategies[0].output_spec.ndim
|
||||
|
||||
@property
|
||||
def shape(self):
|
||||
return self.strategies[0].output_spec.shape
|
||||
|
||||
@property
|
||||
def tensor_meta(self) -> TensorMeta:
|
||||
# TODO upstream this assert to DTensorSpec itself and fill any missing TensorMetas
|
||||
if self.strategies[0].output_spec.tensor_meta is None:
|
||||
raise AssertionError
|
||||
return self.strategies[0].output_spec.tensor_meta
|
||||
|
||||
def __hash__(self) -> int:
|
||||
return hash(tuple(self.strategies))
|
||||
|
||||
def __eq__(self, other: object) -> bool:
|
||||
return isinstance(other, OpStrategy) and self.strategies == other.strategies
|
||||
|
||||
|
||||
class TupleStrategy(StrategyType):
|
||||
"""
|
||||
TupleStrategy is a special case for operators that are fundamentally compound or batched such that some subset
|
||||
of the inputs and outputs are completely unrelated to some other subset.
|
||||
|
||||
Generally, foreach_* ops are the most common use-case for TupleStrategy, because they accept lists of inputs,
|
||||
but operate independently on each input or tuple of zipped inputs.
|
||||
|
||||
For example, [out_a, out_b] = torch.foreach_add([a, b], scalar): input a's sharding only affects out_a's sharding,
|
||||
independent of b and out_b.
|
||||
|
||||
An example of an operator that should NOT use TupleStrategy is torch.split. It produces a List[Tensor]
|
||||
as its output, but the sharding decision of one output is bound together with the decision
|
||||
of each other output and the common input.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
children: Sequence[StrategyType],
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.children: Sequence[StrategyType] = children
|
||||
|
||||
@property
|
||||
@deprecated(
|
||||
"TupleStrategy.childs is deprecated, use TupleStrategy.children instead.", # codespell:ignore childs
|
||||
category=FutureWarning,
|
||||
)
|
||||
def childs(self) -> Sequence[StrategyType]: # codespell:ignore childs
|
||||
"""
|
||||
Alias for children, to maintain backward compatibility.
|
||||
"""
|
||||
return self.children
|
||||
|
||||
def child_mesh(self, index: int) -> DeviceMesh:
|
||||
op_strategy = self.children[index]
|
||||
if not isinstance(op_strategy, OpStrategy):
|
||||
raise AssertionError
|
||||
return op_strategy.mesh
|
||||
|
||||
def __str__(self) -> str:
|
||||
child_strategies_str = ", ".join(
|
||||
[f"{str(strat)}" for idx, strat in enumerate(self.children)]
|
||||
)
|
||||
return f"TupleStrategy({child_strategies_str})"
|
||||
|
||||
def __hash__(self) -> int:
|
||||
return hash(tuple(self.children))
|
||||
|
||||
def __eq__(self, other: object) -> bool:
|
||||
return isinstance(other, TupleStrategy) and self.children == other.children
|
||||
|
||||
|
||||
try:
|
||||
register_pytree_node(
|
||||
TupleStrategy,
|
||||
lambda node: (node.children, None),
|
||||
lambda children, _: TupleStrategy(tuple(children)),
|
||||
)
|
||||
except ValueError:
|
||||
# already registered TupleStrategy, skip
|
||||
pass
|
||||
|
||||
|
||||
@dataclass
|
||||
class RuntimeSchemaInfo:
|
||||
"""
|
||||
RuntimeSchemaInfo stores the operator schema related information for runtime (eager)
|
||||
execution. This is mainly used for two ways: 1. to generate hash for args to determine
|
||||
whether to re-run sharding prop or not 2. to determine if we need pytree
|
||||
"""
|
||||
|
||||
# This static_argnum records static arg "starting index" for ops that have non-tensor
|
||||
# args/kwargs which would affect sharding propagation results. All args starting from
|
||||
# this index would be hashed to our sharding cache.
|
||||
# Note that only a few ops need this information, e.g. view, transpose, var.dim, etc.
|
||||
static_argnum: int = 100
|
||||
# This static_kwargkey records static kwarg names which would affect sharding prop
|
||||
static_kwargkey: list[str] | None = None
|
||||
# each op can decide if it wants to use pytree flatten/unflatten during operator
|
||||
# eager execution, by default we don't need to do flatten/unflatten, only if the
|
||||
# op indicate it needs to, this is to accelerate eager performance.
|
||||
needs_pytree: bool = False
|
||||
|
||||
|
||||
@dataclass
|
||||
class OpSchema:
|
||||
"""
|
||||
OpSchema is a data class that describes an operator input schemas, it includes
|
||||
DTensorSpecs/OpStrategies (instead of DTensor) and non-tensor args/kwargs (positional
|
||||
order preserved). It is mainly used by the DTensor's dispatching logic to perform various
|
||||
actions (i.e. sharding propagation, caching sharding decisions, redistribute, etc.)
|
||||
|
||||
NOTE: this must be used as a read only data class
|
||||
TODO: make this a frozen dataclass
|
||||
|
||||
Args:
|
||||
op: the operator overload we are intercepting
|
||||
args_schema: contains args except that the DTensor args have been replaced
|
||||
with its DTensorSpec or OpStrategy
|
||||
kwargs_schema: contains kwargs except that the DTensor kwargs have been replaced
|
||||
with its DTensorSpec or OpStrategy
|
||||
"""
|
||||
|
||||
op: OpOverload
|
||||
args_schema: ArgsType
|
||||
kwargs_schema: KwargsType
|
||||
|
||||
schema_info: RuntimeSchemaInfo | None = None
|
||||
|
||||
_comparison_key: tuple[object, ...] | None = None
|
||||
|
||||
@property
|
||||
def args_spec(self) -> tuple[DTensorSpec, ...]:
|
||||
"""
|
||||
args_spec: Tuple[DTensorSpec, ...]: contains a clean list of args spec list
|
||||
with NO non-DTensor positional arguments (i.e. int/float/tuple, etc)
|
||||
mainly used by sharding propagation to propagate the output spec
|
||||
"""
|
||||
args = (
|
||||
tree_leaves(self.args_schema)
|
||||
if self.schema_info is not None and self.schema_info.needs_pytree
|
||||
else self.args_schema
|
||||
)
|
||||
return tuple(item for item in args if isinstance(item, DTensorSpec))
|
||||
|
||||
@property
|
||||
def args_strategy(self) -> tuple[OpStrategy, ...]:
|
||||
# filter out non-relevant values from args schema to get a clean OpStrategy list
|
||||
# separate with args_spec for the ease of type annotation
|
||||
# TODO: see if we should merge this with args_spec
|
||||
args = (
|
||||
tree_leaves(self.args_schema)
|
||||
if self.schema_info is not None and self.schema_info.needs_pytree
|
||||
else self.args_schema
|
||||
)
|
||||
return tuple(item for item in args if isinstance(item, OpStrategy))
|
||||
|
||||
@property
|
||||
def kwargs_strategy(self) -> tuple[OpStrategy, ...]:
|
||||
# returns OpStrategy items from kwargs_schema.
|
||||
kwargs_vals = (
|
||||
tree_leaves(self.kwargs_schema)
|
||||
if self.schema_info is not None and self.schema_info.needs_pytree
|
||||
else self.kwargs_schema.values()
|
||||
)
|
||||
return tuple(item for item in kwargs_vals if isinstance(item, OpStrategy))
|
||||
|
||||
@property
|
||||
def args_meta(self) -> tuple[TensorMeta | Any, ...]:
|
||||
# Used for calling single_dim strategy functions, which aren't allowed to see DTensorSpecs/Meshes
|
||||
# like args_spec, but has OpStrategy replaced with corresponding TensorMeta,
|
||||
# and TupleStrategy replaced with tuple of TensorMeta
|
||||
# preserves the original pytree structure
|
||||
# example:
|
||||
# args_schema = (OpStrategy1, TupleStrategy([OpStrategy2, OpStrategy3]), OpStrategy4)
|
||||
# args_meta: (TensorMeta1, (TensorMeta2, TensorMeta3), TensorMeta4)
|
||||
|
||||
def convert_to_meta(item):
|
||||
if isinstance(item, OpStrategy):
|
||||
return item.tensor_meta
|
||||
elif isinstance(item, TupleStrategy):
|
||||
return tuple(convert_to_meta(child) for child in item.children)
|
||||
elif isinstance(item, (list, tuple)):
|
||||
converted = [convert_to_meta(child) for child in item]
|
||||
return type(item)(converted)
|
||||
else:
|
||||
return item
|
||||
|
||||
return tuple(convert_to_meta(arg) for arg in self.args_schema)
|
||||
|
||||
@property
|
||||
def kwargs_meta(self) -> dict[str, object]:
|
||||
# like args_meta, but for kwargs
|
||||
|
||||
def convert_to_meta(item):
|
||||
if isinstance(item, OpStrategy):
|
||||
return item.tensor_meta
|
||||
elif isinstance(item, TupleStrategy):
|
||||
return tuple(convert_to_meta(child) for child in item.children)
|
||||
elif isinstance(item, (list, tuple)):
|
||||
converted = [convert_to_meta(child) for child in item]
|
||||
return type(item)(converted)
|
||||
else:
|
||||
return item
|
||||
|
||||
return {
|
||||
key: convert_to_meta(value) for key, value in self.kwargs_schema.items()
|
||||
}
|
||||
|
||||
def __repr__(self) -> str:
|
||||
args_schema = ", ".join([str(arg_schema) for arg_schema in self.args_schema])
|
||||
return (
|
||||
f"OpSchema(op={self.op},"
|
||||
f" args_schema=({args_schema}),"
|
||||
f" kwargs_schema={self.kwargs_schema})"
|
||||
)
|
||||
|
||||
def __str__(self) -> str:
|
||||
args_schema: list[str] = []
|
||||
device_mesh = None
|
||||
|
||||
for arg in self.args_schema:
|
||||
if isinstance(arg, DTensorSpec):
|
||||
args_schema.append(str(arg))
|
||||
device_mesh = arg.mesh
|
||||
elif isinstance(arg, OpStrategy):
|
||||
if len(arg.strategies) != 1:
|
||||
raise AssertionError
|
||||
args_schema.append(_pretty_print_spec(arg.strategies[0].output_specs))
|
||||
device_mesh = arg.mesh
|
||||
elif isinstance(arg, TupleStrategy):
|
||||
first_op_strategy = arg.children[0]
|
||||
if not isinstance(first_op_strategy, OpStrategy):
|
||||
raise AssertionError
|
||||
device_mesh = first_op_strategy.mesh
|
||||
args_schema.append(str(arg))
|
||||
else:
|
||||
args_schema.append(str(arg))
|
||||
|
||||
return f"{self.op}({', '.join(args_schema)}) on {device_mesh})"
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
_DTensor_OpSchema_post_init(self)
|
||||
|
||||
def arg_type_tensor_or_tensor_list_like(self, arg: object) -> bool:
|
||||
is_tensor = isinstance(arg, DTensorSpec)
|
||||
if is_tensor:
|
||||
return True
|
||||
|
||||
if not isinstance(arg, list):
|
||||
return False
|
||||
|
||||
return all(isinstance(e, DTensorSpec) or e is None for e in arg)
|
||||
|
||||
def return_type_tuple_tensor_like(self) -> bool:
|
||||
# all dispatch ops could only return Tuple[Tensor] or have None/ints/floats
|
||||
# in the tuple, but the first element must be a Tensor, so this check is enough
|
||||
return_types = self.op._schema.returns
|
||||
return len(return_types) > 1 and isinstance(
|
||||
return_types[0].type, torch.TensorType
|
||||
)
|
||||
|
||||
def return_type_list_tensor_like(self) -> bool:
|
||||
# returns True if the return type is a List
|
||||
return_types = self.op._schema.returns
|
||||
return len(return_types) == 1 and isinstance(
|
||||
return_types[0].type, torch.ListType
|
||||
)
|
||||
|
||||
def return_type_tensor(self) -> bool:
|
||||
return_types = self.op._schema.returns
|
||||
# all dispatch ops only return Tensor or Tuple[Tensor] for tensor like
|
||||
# return types, so this check is enough for tensor like types
|
||||
return len(return_types) > 0 and isinstance(
|
||||
return_types[0].type, torch.TensorType
|
||||
)
|
||||
|
||||
def get_mesh_from_args(self, validate: bool = True) -> DeviceMesh:
|
||||
"""
|
||||
This util can be used to get a mesh from the OpSchema that contains multiple
|
||||
DTensors as arguments. When `validate` is True, it will try to validate that all the
|
||||
arguments have the same mesh to avoid unexpected cross mesh errors.
|
||||
|
||||
NOTE: this util currently does not handle TupleStrategy when `validate=True`,
|
||||
this is because for TupleStrategy there could be different types of checks, i.e.:
|
||||
- for stack and cat like op, we need to check within a TupleStrategy is every
|
||||
input is on the same mesh
|
||||
- for foreach like ops we need to check "zipped" inputs are on the same mesh
|
||||
for each index.
|
||||
"""
|
||||
mesh = None
|
||||
# Scan all args to find the first DTensorSpec/OpStrategy (not just the first arg)
|
||||
for arg in self.args_schema:
|
||||
if isinstance(arg, (DTensorSpec, OpStrategy)):
|
||||
mesh = arg.mesh
|
||||
break
|
||||
elif isinstance(arg, (list, tuple, TupleStrategy)):
|
||||
# Scan all elements in the list/tuple, not just the first one,
|
||||
# to handle cases like List[Optional[Tensor]] where first elem may be None
|
||||
elems = arg.children if isinstance(arg, TupleStrategy) else arg
|
||||
for elem in elems:
|
||||
if isinstance(elem, (DTensorSpec, OpStrategy)):
|
||||
mesh = elem.mesh
|
||||
break
|
||||
if mesh is not None:
|
||||
break
|
||||
if mesh is None:
|
||||
raise ValueError(f"Cannot find device mesh from args for op : {self.op}.")
|
||||
|
||||
if validate:
|
||||
for arg in self.args_schema[1:]:
|
||||
if isinstance(arg, (DTensorSpec, OpStrategy)) and arg.mesh != mesh:
|
||||
raise RuntimeError(
|
||||
f"DTensor does not support cross-mesh operation on {self.op}! "
|
||||
f"Got meshes: {mesh} {arg.mesh}. "
|
||||
f"Please make sure all the arguments have the same DeviceMesh."
|
||||
)
|
||||
|
||||
return mesh
|
||||
|
||||
def is_inplace_op(self) -> bool:
|
||||
# simple analysis of function schema to determine
|
||||
# if this is an inplace variant, it might not
|
||||
# be entirely correct, but it's good enough for now.
|
||||
return self.op._schema.name[-1] == "_"
|
||||
|
||||
def is_out_variant_op(self) -> bool:
|
||||
# simple analysis of function schema to determine
|
||||
# if this is an out variant, it might not
|
||||
# be entirely correct, but it's good enough for now.
|
||||
return "out" in self.op._schema.overload_name
|
||||
|
||||
def is_view_op(self) -> bool:
|
||||
return self.op._schema._is_view_op()
|
||||
|
||||
def _recompute_comparison_key(self) -> None:
|
||||
_DTensor_OpSchema_recompute_comparison_key(self)
|
||||
|
||||
def __hash__(self) -> int:
|
||||
return hash(self._comparison_key)
|
||||
|
||||
def __eq__(self, other: object) -> bool:
|
||||
# early return checks
|
||||
if not isinstance(other, OpSchema):
|
||||
return False
|
||||
|
||||
if self.op != other.op:
|
||||
return False
|
||||
|
||||
if len(self.args_schema) != len(other.args_schema):
|
||||
return False
|
||||
|
||||
return self._comparison_key == other._comparison_key
|
||||
|
||||
def gen_fake_args(self) -> ArgsType:
|
||||
"""
|
||||
gen_fake_args: generate fake args for the operator, this is mainly used
|
||||
by sharding propagation rules to generate fake args for the operator
|
||||
to run the local tensor operator and get the output spec.
|
||||
"""
|
||||
return tree_map_only(
|
||||
DTensorSpec,
|
||||
_rebuild_tensor_from_dtensor_meta,
|
||||
self.args_schema,
|
||||
is_leaf=lambda x: isinstance(x, DTensorSpec),
|
||||
)
|
||||
|
||||
def gen_fake_kwargs(self) -> KwargsType:
|
||||
"""
|
||||
gen_fake_kwargs: generate fake kwargs for the operator, this is mainly used
|
||||
by sharding propagation rules to generate fake kwargs for the operator
|
||||
to run the local tensor operator and get the output spec.
|
||||
"""
|
||||
return tree_map_only(
|
||||
DTensorSpec,
|
||||
_rebuild_tensor_from_dtensor_meta,
|
||||
self.kwargs_schema,
|
||||
is_leaf=lambda x: isinstance(x, DTensorSpec),
|
||||
)
|
||||
|
||||
def _inplace_rewrap_schema_suggestion(self, origin_schema: "OpSchema") -> None:
|
||||
suggestion_args_spec = self.args_spec
|
||||
new_arg_schema: list[object] = []
|
||||
idx_of_args_spec = 0
|
||||
if (
|
||||
origin_schema.schema_info is not None
|
||||
and origin_schema.schema_info.needs_pytree
|
||||
):
|
||||
args_schema: Sequence[Any] = tree_leaves(origin_schema.args_schema)
|
||||
else:
|
||||
args_schema = origin_schema.args_schema
|
||||
for arg in args_schema:
|
||||
if isinstance(arg, DTensorSpec):
|
||||
new_arg_schema.append(suggestion_args_spec[idx_of_args_spec])
|
||||
idx_of_args_spec += 1
|
||||
else:
|
||||
new_arg_schema.append(arg)
|
||||
self.args_schema = tuple(new_arg_schema)
|
||||
self.kwargs_schema = origin_schema.kwargs_schema
|
||||
self._recompute_comparison_key()
|
||||
|
||||
|
||||
@dataclass
|
||||
class OutputSharding:
|
||||
"""
|
||||
OutputSharding is a data class that is used by the sharding propagation,
|
||||
it could set the output_spec upon successful propagation. If needs_redistribute
|
||||
is set to True, a redistribute_schema would be returned together to indicate
|
||||
the input arguments needs to be redistributed before the op execution.
|
||||
|
||||
NOTE: the redistribute_schema generated by sharding propagation should be
|
||||
exactly the same as the operator OpSchema, except the DTensorSpecs
|
||||
"""
|
||||
|
||||
# specifies the output sharding pattern
|
||||
output_spec: OutputSpecType
|
||||
# schema for redistribution if needed
|
||||
redistribute_schema: OpSchema | None = None
|
||||
# flag indicating if inputs need redistribution
|
||||
needs_redistribute: bool = False
|
||||
# flag to use values from `redistribute_schema`
|
||||
use_val_from_redistribute_schema: bool = False
|
||||
|
||||
@cached_property
|
||||
def mesh(self):
|
||||
if isinstance(self.output_spec, DTensorSpec):
|
||||
return self.output_spec.mesh
|
||||
elif isinstance(self.output_spec, tuple):
|
||||
out_spec = self.output_spec[0]
|
||||
if isinstance(out_spec, DTensorSpec):
|
||||
return out_spec.mesh
|
||||
else:
|
||||
raise ValueError(f"Unknown output spec type: {type(out_spec)}")
|
||||
else:
|
||||
raise ValueError(f"Unknown output spec type: {type(self.output_spec)}")
|
||||
|
||||
|
||||
@dataclass
|
||||
class OpInfo:
|
||||
"""
|
||||
All Runtime Op execution info are packed here
|
||||
"""
|
||||
|
||||
# The first compute device mesh recorded from args
|
||||
# NOTE: one op could have multiple meshes from its args. We just record the first
|
||||
# mesh here to check if current rank should participate in computation or not.
|
||||
compute_mesh: DeviceMesh
|
||||
|
||||
# compete runtime operator infos
|
||||
# NOTE: schema can be None due to C++ fast path optimization. When the C++
|
||||
# dispatch layer (dispatchDTensorOp in python_variable.cpp) finds a cached
|
||||
# sharding decision, it skips creating the full OpSchema to reduce CPU overhead.
|
||||
# In this case, OpInfo is created with create_schema=False, setting schema to None.
|
||||
# The operator information is still available through output_sharding.redistribute_schema
|
||||
# when redistribution is needed.
|
||||
schema: OpSchema | None
|
||||
flat_args_schema: list[object]
|
||||
local_args: Sequence[object]
|
||||
local_kwargs: dict[str, object]
|
||||
args_tree_spec: TreeSpec | None = None
|
||||
|
||||
# the output sharding info
|
||||
output_sharding: OutputSharding | None = None
|
||||
@@ -0,0 +1,9 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates
|
||||
from ._conv_ops import * # noqa: F403
|
||||
from ._embedding_ops import * # noqa: F403
|
||||
from ._math_ops import * # noqa: F403
|
||||
from ._matrix_ops import * # noqa: F403
|
||||
from ._pointwise_ops import * # noqa: F403
|
||||
from ._random_ops import * # noqa: F403
|
||||
from ._tensor_ops import * # noqa: F403
|
||||
from ._view_ops import * # noqa: F403
|
||||
+288
@@ -0,0 +1,288 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates
|
||||
import string
|
||||
from typing import cast
|
||||
|
||||
import torch
|
||||
from torch.distributed.tensor._dtensor_spec import DTensorSpec, TensorMeta
|
||||
from torch.distributed.tensor._op_schema import OpSchema, OutputSharding
|
||||
from torch.distributed.tensor._ops.utils import prod
|
||||
from torch.distributed.tensor._utils import compute_local_shape_and_global_offset
|
||||
|
||||
|
||||
def _replace_char_in_str(string: str, new_char: str, idx: int) -> str:
|
||||
return string[:idx] + new_char + string[idx + 1 :]
|
||||
|
||||
|
||||
def _gen_reshard_suggestions(
|
||||
op_schema: OpSchema,
|
||||
input_dims: list[str],
|
||||
input_specs: tuple[DTensorSpec, ...],
|
||||
dim_to_sharding: dict[str, int],
|
||||
pending_sum: list[int],
|
||||
) -> OutputSharding:
|
||||
suggested_arg_specs: list[DTensorSpec] = []
|
||||
for input_dim, input_spec in zip(input_dims, input_specs):
|
||||
dim_map = [dim_to_sharding[dim] for dim in input_dim]
|
||||
suggested_arg_specs.append(
|
||||
DTensorSpec.from_dim_map(
|
||||
mesh=input_spec.mesh,
|
||||
dim_map=dim_map,
|
||||
sums=pending_sum,
|
||||
tensor_meta=input_spec.tensor_meta,
|
||||
)
|
||||
)
|
||||
suggested_schema = OpSchema(op_schema.op, tuple(suggested_arg_specs), {})
|
||||
suggested_schema._inplace_rewrap_schema_suggestion(op_schema)
|
||||
return OutputSharding(
|
||||
None,
|
||||
redistribute_schema=suggested_schema,
|
||||
)
|
||||
|
||||
|
||||
def einop_rule(
|
||||
equation: str,
|
||||
op_schema: OpSchema,
|
||||
*,
|
||||
linearity: bool = False,
|
||||
enforce_sharding: dict[str, int] | None = None,
|
||||
) -> OutputSharding:
|
||||
"""
|
||||
Propagate the sharding of inputs to output for ops whose data moves according to einsum notation.
|
||||
|
||||
This is mostly borrowed from @zdevito's sharding simulator. Examples:
|
||||
mk,kn->mn - einsum
|
||||
ij,ij->ij - addition
|
||||
ij,j->ij - broadcasted addition
|
||||
ij->i - reduction
|
||||
Other ops could use this propagation algorithm when applied, note
|
||||
that einsum propagation only deal with list of specs (DTensor specs)
|
||||
as it only works on list of tensors!
|
||||
|
||||
linearity in einop_rule means that the calling op `f` follows this rule:
|
||||
f(a + b) = f(a) + f(b)
|
||||
|
||||
In this case we can propagate the partial sum, note that linearity in einop
|
||||
only applies to partial sum, not other operations like min/max (which are
|
||||
associative but not linear).
|
||||
"""
|
||||
# parse einop equation and extract arg specs
|
||||
inputs, outputs = equation.split("->")
|
||||
input_dims, output_dims = inputs.split(","), outputs.split(",")
|
||||
input_specs = op_schema.args_spec
|
||||
# NOTE: only support single output unless needed in future
|
||||
output_dim = output_dims[0]
|
||||
|
||||
dim_to_sharding: dict[str, int] = {}
|
||||
dim_to_size: dict[str, int] = {}
|
||||
# record pending sum, key is mesh dimension, value is pending sum
|
||||
# counter across input specs
|
||||
pending_sums_counter: dict[int, int] = {}
|
||||
seen_shardings: dict[int, str] = {}
|
||||
needs_reshard = False
|
||||
|
||||
def merge_sharding(dim: str, a: int, b: int) -> int:
|
||||
# merge the sharding of inputs if it's able to merge, i.e. we can merge
|
||||
# replicate and shard to shard, but this will trigger an reshard operation
|
||||
if a != b:
|
||||
if a == -1 or b == -1:
|
||||
# reshard the replicate to match the sharded one
|
||||
nonlocal needs_reshard
|
||||
needs_reshard = True
|
||||
return a if a != -1 else b
|
||||
else:
|
||||
# TODO: further merge the sharding properly (i.e. reshard one input to replicate)
|
||||
raise RuntimeError(
|
||||
f"{equation}: dim {dim} sharded two different ways: {a} and {b}"
|
||||
)
|
||||
else:
|
||||
return a
|
||||
|
||||
for input_dim, input_spec in zip(input_dims, input_specs):
|
||||
# deal with partial sums
|
||||
input_sums = input_spec.sums
|
||||
for sum_dim in input_sums:
|
||||
if sum_dim not in pending_sums_counter:
|
||||
seen_shardings[sum_dim] = "+"
|
||||
# update pending sum counter for pending sum mesh
|
||||
# dimension with the occurrence from each input
|
||||
pending_sums_counter[sum_dim] = pending_sums_counter.get(sum_dim, 0) + 1
|
||||
|
||||
for idx, (dim, mesh_dim) in enumerate(zip(input_dim, input_spec.dim_map)):
|
||||
if enforce_sharding and dim in enforce_sharding:
|
||||
if enforce_sharding[dim] != mesh_dim:
|
||||
needs_reshard = True
|
||||
dim_to_sharding[dim] = enforce_sharding[dim]
|
||||
dim_to_size[dim] = input_spec.shape[idx]
|
||||
elif dim not in dim_to_sharding:
|
||||
dim_to_sharding[dim] = mesh_dim
|
||||
dim_to_size[dim] = input_spec.shape[idx]
|
||||
else:
|
||||
dim_to_sharding[dim] = merge_sharding(
|
||||
dim, dim_to_sharding[dim], mesh_dim
|
||||
)
|
||||
if dim_to_size[dim] != input_spec.shape[idx]:
|
||||
raise AssertionError
|
||||
|
||||
# after merging sharding, we check if there're multiple
|
||||
# sharding on the same mesh dim.
|
||||
merged_sharding_for_dim = dim_to_sharding[dim]
|
||||
if merged_sharding_for_dim != -1:
|
||||
if (
|
||||
merged_sharding_for_dim in seen_shardings
|
||||
and dim != seen_shardings[merged_sharding_for_dim]
|
||||
):
|
||||
needs_reshard = True
|
||||
seen_shardings[merged_sharding_for_dim] += dim
|
||||
else:
|
||||
seen_shardings[merged_sharding_for_dim] = dim
|
||||
|
||||
if pending_sums_counter and not linearity:
|
||||
# return reshard suggestion with no pending sum, because we already properly
|
||||
# merge the sharding, this reshard suggestion is legit to use
|
||||
return _gen_reshard_suggestions(
|
||||
op_schema, input_dims, input_specs, dim_to_sharding, []
|
||||
)
|
||||
else:
|
||||
# It's a op that support linearity, but not all input arguments are partial
|
||||
# we fail the sharding propagation with suggestion to make all inputs be
|
||||
# partial on the corresponding mesh dim (all inputs should be partial for
|
||||
# the mesh dims in order to execute locally and delay the sum reduction)
|
||||
for value in pending_sums_counter.values():
|
||||
if value != len(input_specs):
|
||||
needs_reshard = True
|
||||
|
||||
for mesh_dim, dims in seen_shardings.items():
|
||||
if len(dims) > 1:
|
||||
# we found different input dims are being sharded on the same mesh dim
|
||||
# in order to perform local op computation, we need to reshard inputs
|
||||
# base on some simple heuristics, now we simply pick the one with least comm
|
||||
# volume. (i.e. the input with least size)
|
||||
# TODO: consider a more advanced heuristic to pick the best sharding
|
||||
costs = []
|
||||
for d in dims:
|
||||
cost = 0
|
||||
for input_dim, input_spec in zip(input_dims, input_specs):
|
||||
if (
|
||||
d in input_dim
|
||||
and input_spec.dim_map[input_dim.index(d)] == mesh_dim
|
||||
):
|
||||
if input_spec.tensor_meta is None:
|
||||
raise AssertionError
|
||||
global_shape = input_spec.tensor_meta.shape
|
||||
local_shape, _ = compute_local_shape_and_global_offset(
|
||||
global_shape,
|
||||
input_spec.mesh,
|
||||
input_spec.placements,
|
||||
skip_offset=True,
|
||||
)
|
||||
cost += prod(local_shape) * input_spec.mesh.size(mesh_dim)
|
||||
|
||||
costs.append(cost)
|
||||
d_to_keep_sharding = dims[costs.index(max(costs))]
|
||||
for d in dims:
|
||||
# update dim_to_sharding to keep the sharding of the dim with
|
||||
# highest comm and make the rest of the dims to replicate
|
||||
if d != d_to_keep_sharding:
|
||||
dim_to_sharding[d] = -1
|
||||
|
||||
pending_sums = list(pending_sums_counter.keys())
|
||||
if needs_reshard:
|
||||
return _gen_reshard_suggestions(
|
||||
op_schema, input_dims, input_specs, dim_to_sharding, pending_sums
|
||||
)
|
||||
|
||||
# generate output pending sum if a dim is sharded, and it appears in input
|
||||
# but not output
|
||||
for dim, shard_on_mesh in dim_to_sharding.items():
|
||||
if dim not in output_dims[0] and shard_on_mesh != -1:
|
||||
pending_sums.append(shard_on_mesh)
|
||||
|
||||
# if no need to reshard, we directly generate the output sharding
|
||||
output_dim_map = []
|
||||
output_shape = []
|
||||
for dim in output_dim:
|
||||
if dim == "1":
|
||||
# find output dim that is a singleton dimension, mark sharding and shape
|
||||
output_dim_map.append(-1)
|
||||
output_shape.append(1)
|
||||
else:
|
||||
output_dim_map.append(dim_to_sharding[dim])
|
||||
output_shape.append(dim_to_size[dim])
|
||||
|
||||
# XXX: since we still need to have intermediate shape calculation, we need
|
||||
# to pass in the shape here. We should remove this once sharding decomp works
|
||||
# for ops like addmm
|
||||
if input_specs[0].tensor_meta is None:
|
||||
raise AssertionError
|
||||
tensor_meta = TensorMeta(
|
||||
torch.Size(output_shape),
|
||||
input_specs[0].tensor_meta.stride,
|
||||
input_specs[0].tensor_meta.dtype,
|
||||
)
|
||||
return OutputSharding(
|
||||
DTensorSpec.from_dim_map(
|
||||
input_specs[0].mesh,
|
||||
output_dim_map,
|
||||
pending_sums,
|
||||
tensor_meta=tensor_meta,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def pointwise_rule(op_schema: OpSchema, linearity: bool = False) -> OutputSharding:
|
||||
"""
|
||||
Propagate the sharding for pointwise operations.
|
||||
|
||||
Examples:
|
||||
ij,ij->ij - addition/mul
|
||||
ij,j->ij - broadcasted addition
|
||||
"""
|
||||
alphabet = string.ascii_lowercase
|
||||
# find the max_dim first in case we need to broadcasting
|
||||
input_specs = op_schema.args_spec
|
||||
max_dim = max(input.ndim for input in input_specs)
|
||||
dimchars = []
|
||||
singleton_counter: list[int] = [0] * max_dim
|
||||
for input in input_specs:
|
||||
start_dim = max_dim - input.ndim
|
||||
p = alphabet[start_dim:max_dim]
|
||||
# handle the "broadcasting to a common shape case"
|
||||
# see https://pytorch.org/docs/stable/notes/broadcasting.html
|
||||
# If any of the dimensions is singleton dimension (i.e. 1).
|
||||
# we mark the dim char as a special "1" to distinguish with
|
||||
# the non-singleton dimension, so that sharding propagation
|
||||
# should just ignore the singleton dimension.
|
||||
if len(input_specs) > 1:
|
||||
for i in range(max_dim):
|
||||
if i < start_dim:
|
||||
# treat the leading miss dim chars as singleton
|
||||
singleton_counter[i] += 1
|
||||
elif input.shape[i - start_dim] == 1:
|
||||
# mark singleton dim char as a special "1" in einop rule
|
||||
singleton_counter[i] += 1
|
||||
p = _replace_char_in_str(p, "1", (i - start_dim))
|
||||
|
||||
dimchars.append(p)
|
||||
out_dimchars = alphabet[:max_dim]
|
||||
# check if we replace the all inputs dim char with singleton dimension,
|
||||
# if we replace all inputs, we also need to replace the output dimension.
|
||||
for output_dim_idx in range(len(out_dimchars)):
|
||||
if singleton_counter[output_dim_idx] == len(input_specs):
|
||||
out_dimchars = _replace_char_in_str(out_dimchars, "1", output_dim_idx)
|
||||
|
||||
fmt = f"{','.join(p for p in dimchars)}->{out_dimchars}"
|
||||
|
||||
enforce_sharding: dict[str, int] = {}
|
||||
if op_schema.is_inplace_op():
|
||||
follow_spec = op_schema.args_spec[0]
|
||||
enforce_sharding.update(zip(out_dimchars, follow_spec.dim_map))
|
||||
elif op_schema.is_out_variant_op():
|
||||
follow_spec = cast(DTensorSpec, op_schema.kwargs_schema["out"])
|
||||
enforce_sharding.update(zip(out_dimchars, follow_spec.dim_map))
|
||||
|
||||
return einop_rule(
|
||||
fmt,
|
||||
op_schema,
|
||||
linearity=linearity,
|
||||
enforce_sharding=enforce_sharding,
|
||||
)
|
||||
@@ -0,0 +1,202 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates
|
||||
# implement matrix related ops for distributed tensor
|
||||
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from torch.distributed.tensor._dtensor_spec import DTensorSpec, TensorMeta
|
||||
from torch.distributed.tensor._op_schema import (
|
||||
OpSchema,
|
||||
OutputSharding,
|
||||
RuntimeSchemaInfo,
|
||||
)
|
||||
from torch.distributed.tensor._ops.single_dim_strategy import (
|
||||
_ShardingPlaceholder,
|
||||
register_single_dim_strategy,
|
||||
)
|
||||
from torch.distributed.tensor._ops.utils import register_prop_rule
|
||||
from torch.distributed.tensor.placement_types import Partial, Placement, Replicate
|
||||
|
||||
|
||||
aten = torch.ops.aten
|
||||
|
||||
|
||||
@register_prop_rule(aten.convolution.default)
|
||||
def convolution_rules(op_schema: OpSchema) -> OutputSharding:
|
||||
(
|
||||
input_spec,
|
||||
weight_spec,
|
||||
bias_spec,
|
||||
stride,
|
||||
padding,
|
||||
dilation,
|
||||
_transposed,
|
||||
_output_padding,
|
||||
_groups,
|
||||
) = op_schema.args_schema
|
||||
|
||||
if not isinstance(input_spec, DTensorSpec):
|
||||
raise AssertionError
|
||||
if not isinstance(weight_spec, DTensorSpec):
|
||||
raise AssertionError
|
||||
# bias_spec can be None (optional parameter in aten.convolution schema)
|
||||
if bias_spec is not None:
|
||||
if not isinstance(bias_spec, DTensorSpec):
|
||||
raise AssertionError
|
||||
if input_spec.tensor_meta is None:
|
||||
raise AssertionError
|
||||
if weight_spec.tensor_meta is None:
|
||||
raise AssertionError
|
||||
in_shape = input_spec.tensor_meta.shape
|
||||
weight_shape = weight_spec.tensor_meta.shape
|
||||
if not isinstance(stride, list):
|
||||
raise AssertionError(f"stride must be list, got {type(stride)}")
|
||||
if not isinstance(padding, list):
|
||||
raise AssertionError(f"padding must be list, got {type(padding)}")
|
||||
if not isinstance(dilation, list):
|
||||
raise AssertionError(f"dilation must be list, got {type(dilation)}")
|
||||
# weight_shape might not be torch.Size in all cases (e.g., SymIntArrayRef during tracing)
|
||||
# so we don't assert its type, just use it
|
||||
out_conv_shape = [
|
||||
(d + 2 * padding[i] - dilation[i] * (weight_shape[i + 1] - 1) - 1) // stride[i]
|
||||
+ 1
|
||||
for (i, d) in enumerate(in_shape[2:])
|
||||
]
|
||||
output_shape = [in_shape[0], weight_shape[0]] + out_conv_shape
|
||||
output_stride = [1]
|
||||
for i in range(1, len(output_shape)):
|
||||
output_stride.insert(0, output_stride[0] * output_shape[-i])
|
||||
output_dim_map = input_spec.dim_map
|
||||
pending_sums = input_spec.sums
|
||||
|
||||
tensor_meta = TensorMeta(
|
||||
torch.Size(output_shape),
|
||||
tuple(output_stride),
|
||||
input_spec.tensor_meta.dtype,
|
||||
)
|
||||
return OutputSharding(
|
||||
DTensorSpec.from_dim_map(
|
||||
input_spec.mesh,
|
||||
output_dim_map,
|
||||
pending_sums,
|
||||
tensor_meta=tensor_meta,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@register_prop_rule(aten.convolution_backward.default)
|
||||
def convolution_backward_rules(op_schema: OpSchema) -> OutputSharding:
|
||||
input_spec = op_schema.args_schema[0]
|
||||
(
|
||||
grad_output_spec,
|
||||
input_spec,
|
||||
weight_spec,
|
||||
bias_shape_opt,
|
||||
_stride,
|
||||
_padding,
|
||||
_dilation,
|
||||
_transposed,
|
||||
_output_padding,
|
||||
_groups,
|
||||
_output_mask,
|
||||
) = op_schema.args_schema
|
||||
|
||||
if not isinstance(grad_output_spec, DTensorSpec):
|
||||
raise AssertionError
|
||||
if not isinstance(input_spec, DTensorSpec):
|
||||
raise AssertionError
|
||||
if not isinstance(weight_spec, DTensorSpec):
|
||||
raise AssertionError
|
||||
# bias_shape_opt can be None (optional parameter in aten.convolution_backward schema)
|
||||
if bias_shape_opt is not None:
|
||||
if not isinstance(bias_shape_opt, list):
|
||||
raise AssertionError
|
||||
if input_spec.tensor_meta is None:
|
||||
raise AssertionError
|
||||
weight_tensor_meta = weight_spec.tensor_meta
|
||||
|
||||
# Only create bias_tensor_meta if bias_shape_opt is not None
|
||||
if bias_shape_opt is not None:
|
||||
bias_tensor_meta = TensorMeta(
|
||||
torch.Size(bias_shape_opt),
|
||||
(1,),
|
||||
input_spec.tensor_meta.dtype,
|
||||
)
|
||||
else:
|
||||
bias_tensor_meta = None
|
||||
|
||||
grad_input_spec = input_spec
|
||||
grad_weight_spec = DTensorSpec.from_dim_map(
|
||||
input_spec.mesh,
|
||||
[-1, -1, -1, -1],
|
||||
[0],
|
||||
tensor_meta=weight_tensor_meta,
|
||||
)
|
||||
|
||||
# Only create grad_bias_spec if we have bias_tensor_meta
|
||||
if bias_tensor_meta is not None:
|
||||
grad_bias_spec = DTensorSpec.from_dim_map(
|
||||
input_spec.mesh,
|
||||
[-1],
|
||||
[0],
|
||||
tensor_meta=bias_tensor_meta,
|
||||
)
|
||||
else:
|
||||
grad_bias_spec = None
|
||||
|
||||
# TODO: actually the output_mask is not respected here, we should
|
||||
# set the corresponding spec to `None` if the output_mask is not `False`
|
||||
# for a certain output Tensor. This also applies to the conv handler
|
||||
# in torch/distributed/tensor/_tp_conv.py
|
||||
return OutputSharding([grad_input_spec, grad_weight_spec, grad_bias_spec])
|
||||
|
||||
|
||||
# Single-dim strategies for autoparallel optimizer support.
|
||||
# These coexist with the prop_rules above — strategies take precedence
|
||||
# in the propagation path, while the prop_rules + custom handlers in
|
||||
# _tp_conv.py continue to handle runtime dispatch.
|
||||
|
||||
|
||||
@register_single_dim_strategy(
|
||||
[aten.convolution.default],
|
||||
schema_info=RuntimeSchemaInfo(2),
|
||||
)
|
||||
def convolution_single_dim_strategy(
|
||||
op: torch._ops.OpOverload,
|
||||
args_schema: tuple[Any, ...],
|
||||
kwargs_schema: dict[str, Any],
|
||||
) -> list[list[Placement | _ShardingPlaceholder]]:
|
||||
bias_meta = args_schema[2]
|
||||
# [output, input, weight, (bias)]
|
||||
rule: list[Placement | _ShardingPlaceholder] = [
|
||||
_ShardingPlaceholder(0), # output
|
||||
_ShardingPlaceholder(0), # input
|
||||
Replicate(), # weight
|
||||
]
|
||||
if bias_meta is not None:
|
||||
rule.append(Replicate()) # bias
|
||||
return [rule]
|
||||
|
||||
|
||||
@register_single_dim_strategy(
|
||||
[aten.convolution_backward.default],
|
||||
schema_info=RuntimeSchemaInfo(3),
|
||||
)
|
||||
def convolution_backward_single_dim_strategy(
|
||||
op: torch._ops.OpOverload,
|
||||
args_schema: tuple[Any, ...],
|
||||
kwargs_schema: dict[str, Any],
|
||||
) -> list[list[Placement | _ShardingPlaceholder | None]]:
|
||||
bias_sizes = args_schema[3]
|
||||
has_bias = bias_sizes is not None
|
||||
# outputs: [grad_input, grad_weight, grad_bias]
|
||||
# inputs: [grad_output, input, weight]
|
||||
rule: list[Placement | _ShardingPlaceholder | None] = [
|
||||
_ShardingPlaceholder(0), # grad_input
|
||||
Partial("sum"), # grad_weight
|
||||
Partial("sum") if has_bias else None, # grad_bias
|
||||
_ShardingPlaceholder(0), # grad_output
|
||||
_ShardingPlaceholder(0), # input
|
||||
Replicate(), # weight
|
||||
]
|
||||
return [rule]
|
||||
+195
@@ -0,0 +1,195 @@
|
||||
import itertools
|
||||
from dataclasses import dataclass
|
||||
|
||||
from torch.distributed.device_mesh import DeviceMesh
|
||||
from torch.distributed.tensor._dtensor_spec import DTensorSpec
|
||||
from torch.distributed.tensor._op_schema import OpSpec, OpStrategy
|
||||
from torch.distributed.tensor.placement_types import (
|
||||
Partial,
|
||||
Placement,
|
||||
Replicate,
|
||||
Shard,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class EinsumDims:
|
||||
contracting_dims: list[str]
|
||||
batch_dims: list[str]
|
||||
lhs_out_only_dims: list[str]
|
||||
rhs_out_only_dims: list[str]
|
||||
|
||||
@classmethod
|
||||
def parse_equation(cls, equation: str) -> tuple[list[str], str]:
|
||||
# parse einop equation and extract arg specs
|
||||
"""
|
||||
Parse the einsum equation str to input dim chars and output dim char
|
||||
"""
|
||||
inputs, outputs = equation.split("->")
|
||||
input_dims, output_dims = inputs.split(","), outputs.split(",")
|
||||
|
||||
# NOTE: only support at most two inputs, and single output
|
||||
# extend to support more inputs if needed in future
|
||||
if len(input_dims) > 2:
|
||||
raise AssertionError("Only support at most two inputs")
|
||||
if len(output_dims) != 1:
|
||||
raise AssertionError("Only support single output")
|
||||
output_dim = output_dims[0]
|
||||
return input_dims, output_dim
|
||||
|
||||
@classmethod
|
||||
def parse_dims(cls, input_dims: list[str], output_dim: str) -> "EinsumDims":
|
||||
"""
|
||||
Parse the dims and extract the contracting, batch, and free dimensions
|
||||
for the left and right hand sides.
|
||||
"""
|
||||
dim_char_set: set[str] = set()
|
||||
for input_dim in input_dims:
|
||||
dim_char_set.update(input_dim)
|
||||
|
||||
# get a deterministic order of all dim chars
|
||||
all_dim_chars = sorted(dim_char_set)
|
||||
|
||||
# parse input and output dimensions
|
||||
lhs_out_only_dims, rhs_out_only_dims = [], []
|
||||
batch_dims, contracting_dims = [], []
|
||||
|
||||
for dim_char in all_dim_chars:
|
||||
if dim_char not in output_dim:
|
||||
contracting_dims.append(dim_char)
|
||||
else:
|
||||
is_batch_dim = True
|
||||
for input_dim in input_dims:
|
||||
is_batch_dim = is_batch_dim and dim_char in input_dim
|
||||
|
||||
if is_batch_dim:
|
||||
batch_dims.append(dim_char)
|
||||
else:
|
||||
if len(input_dims) != 2:
|
||||
raise AssertionError(
|
||||
"free dimension only supported for two inputs!"
|
||||
)
|
||||
lhs, rhs = input_dims
|
||||
if dim_char in lhs:
|
||||
lhs_out_only_dims.append(dim_char)
|
||||
elif dim_char in rhs:
|
||||
rhs_out_only_dims.append(dim_char)
|
||||
else:
|
||||
raise RuntimeError("Invalid dimension character")
|
||||
|
||||
return cls(
|
||||
contracting_dims=contracting_dims,
|
||||
batch_dims=batch_dims,
|
||||
lhs_out_only_dims=lhs_out_only_dims,
|
||||
rhs_out_only_dims=rhs_out_only_dims,
|
||||
)
|
||||
|
||||
|
||||
def gen_einsum_strategies(
|
||||
equation: str,
|
||||
mesh: DeviceMesh,
|
||||
*,
|
||||
linearity: bool = False,
|
||||
) -> OpStrategy:
|
||||
"""
|
||||
Generate a strategy list for the ops that follow einsum style notation.
|
||||
|
||||
In principle, each mesh dim is independent of other device mesh dim when we
|
||||
generate strategies. So we generate strategy over each device mesh dim and
|
||||
do product combination on all mesh dims. We basically follow the below rule
|
||||
for each device mesh dim:
|
||||
|
||||
1. Shard on contracting dim: When both inputs shard on contracting dim over
|
||||
the same device dim. The result will be Partial over that device dim.
|
||||
|
||||
2. Shard on noncontracting dim:
|
||||
2.1: Shard on batch dim: output, both inputs all should shard on batch
|
||||
dim.
|
||||
2.2: Shard on lhs only dim or rhs only dim: both output and lhs or rhs
|
||||
input should shard on this free dim.
|
||||
|
||||
3. Linearity (Partial): If enabled, set Partial on output and inputs over
|
||||
the same device mesh dim.
|
||||
"""
|
||||
# parse einop equation and extract dims
|
||||
input_dims, output_dim = EinsumDims.parse_equation(equation)
|
||||
edims = EinsumDims.parse_dims(input_dims, output_dim)
|
||||
all_mesh_dim_strategies = []
|
||||
|
||||
# generate strategies for each mesh dim and do cartesian product for final strategy. E.g., for a 2D mesh, we can have [P(),R,R]
|
||||
strategies_over_one_mesh_dim = []
|
||||
|
||||
# placement list stores placements of [output, input1, input2, ...]
|
||||
# first we always have replicate all for inputs and output
|
||||
placement_list: list[Placement] = [Replicate()] * (len(input_dims) + 1)
|
||||
strategies_over_one_mesh_dim.append(placement_list)
|
||||
|
||||
# split batch dim
|
||||
for batch_dim in edims.batch_dims:
|
||||
output_batch_dim = output_dim.index(batch_dim)
|
||||
placement_list = [Shard(output_batch_dim)]
|
||||
for input_dim in input_dims:
|
||||
input_batch_dim = input_dim.index(batch_dim)
|
||||
placement_list.append(Shard(input_batch_dim))
|
||||
|
||||
strategies_over_one_mesh_dim.append(placement_list)
|
||||
|
||||
# split contracting dim
|
||||
# NOTE: This is the only strategy that produces a Partial output, and it
|
||||
# hardcodes Partial("sum"). No strategy accepts Partial as an input
|
||||
# placement, so Partial inputs are always redistributed to Shard or
|
||||
# Replicate. This means Partial("avg") inputs cannot be preserved through
|
||||
# the op. Use gen_single_dim_einsum_strategies with per-input linearity
|
||||
# for proper Partial support.
|
||||
for contracting_dim in edims.contracting_dims:
|
||||
# Contracting dim can shard on same device axis for both inputs. This
|
||||
# results in the output being Partial on that device axis. For example:
|
||||
# bmk_{x},k_{x}n -> bmn{Ux} (becomes partial over device axis x)
|
||||
placement_list = [Partial()]
|
||||
for input_dim in input_dims:
|
||||
input_contracting_dim = input_dim.index(contracting_dim)
|
||||
placement_list.append(Shard(input_contracting_dim))
|
||||
|
||||
strategies_over_one_mesh_dim.append(placement_list)
|
||||
|
||||
# split lhs free dim
|
||||
for lhs_dim in edims.lhs_out_only_dims:
|
||||
lhs_free_dim_output = output_dim.index(lhs_dim)
|
||||
lhs_free_dim_input = input_dims[0].index(lhs_dim)
|
||||
# this means split the lhs input and output
|
||||
# i.e. S(0), R -> S(0)
|
||||
lhs_placement_list: list[Placement] = [
|
||||
Shard(lhs_free_dim_output),
|
||||
Shard(lhs_free_dim_input),
|
||||
Replicate(),
|
||||
]
|
||||
strategies_over_one_mesh_dim.append(lhs_placement_list)
|
||||
|
||||
# split rhs free dim
|
||||
for rhs_dim in edims.rhs_out_only_dims:
|
||||
rhs_free_dim_output = output_dim.index(rhs_dim)
|
||||
rhs_free_dim_input = input_dims[1].index(rhs_dim)
|
||||
rhs_placement_list: list[Placement] = [
|
||||
Shard(rhs_free_dim_output),
|
||||
Replicate(),
|
||||
Shard(rhs_free_dim_input),
|
||||
]
|
||||
strategies_over_one_mesh_dim.append(rhs_placement_list)
|
||||
|
||||
# linearity strategy
|
||||
if linearity:
|
||||
linearity_placement_list: list[Placement] = [Partial()]
|
||||
for _ in input_dims:
|
||||
linearity_placement_list.append(Partial())
|
||||
strategies_over_one_mesh_dim.append(linearity_placement_list)
|
||||
|
||||
# generate strategies for entire mesh
|
||||
all_mesh_dim_strategies = [strategies_over_one_mesh_dim] * mesh.ndim
|
||||
strategy_combs = itertools.product(*all_mesh_dim_strategies)
|
||||
all_strategies = []
|
||||
for strategy_comb in strategy_combs:
|
||||
spec_list = [DTensorSpec(mesh, tuple(specs)) for specs in zip(*strategy_comb)]
|
||||
strat = OpSpec(output_specs=spec_list[0], input_specs=spec_list[1:])
|
||||
all_strategies.append(strat)
|
||||
|
||||
return OpStrategy(all_strategies)
|
||||
+117
@@ -0,0 +1,117 @@
|
||||
# mypy: allow-untyped-defs
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates
|
||||
# implement matrix related ops for distributed tensor
|
||||
from typing import cast
|
||||
|
||||
import torch
|
||||
from torch.distributed.tensor._op_schema import (
|
||||
OpSchema,
|
||||
OpStrategy,
|
||||
PlacementList,
|
||||
RuntimeSchemaInfo,
|
||||
StrategyType,
|
||||
)
|
||||
from torch.distributed.tensor._ops.utils import (
|
||||
expand_to_full_mesh_op_strategy,
|
||||
register_op_strategy,
|
||||
)
|
||||
from torch.distributed.tensor.placement_types import (
|
||||
_MaskPartial,
|
||||
Partial,
|
||||
Replicate,
|
||||
Shard,
|
||||
)
|
||||
|
||||
|
||||
aten = torch.ops.aten
|
||||
|
||||
|
||||
@register_op_strategy(aten.embedding.default)
|
||||
def embedding_strategy(op_schema: OpSchema) -> StrategyType:
|
||||
"""
|
||||
This strategy handles embedding op. We have two possible embedding shardings:
|
||||
rowwise and colwise
|
||||
"""
|
||||
weight_strategy = cast(OpStrategy, op_schema.args_schema[0])
|
||||
indices_strategy = cast(OpStrategy, op_schema.args_schema[1])
|
||||
mesh = op_schema.get_mesh_from_args()
|
||||
|
||||
weight_shape = weight_strategy.shape
|
||||
indices_shape = indices_strategy.shape
|
||||
output_emd_dim = len(indices_shape)
|
||||
|
||||
single_mesh_dim_strategies = []
|
||||
|
||||
# placement list stores placements of [output, weight, input_indices]
|
||||
# first we always have replicate all for inputs and output
|
||||
all_replicate: PlacementList = [Replicate()] * 3
|
||||
single_mesh_dim_strategies.append(all_replicate)
|
||||
|
||||
# colwise sharding, output shard on last dim, weight shard on dim 1, input replicate
|
||||
colwise_sharding: PlacementList = [Shard(output_emd_dim), Shard(1), Replicate()]
|
||||
single_mesh_dim_strategies.append(colwise_sharding)
|
||||
|
||||
# rowwise sharding, output is embedding partial, weight shard on dim 0, input accepts embedding partial
|
||||
embedding_partial_placement = _MaskPartial(offset_shape=weight_shape, offset_dim=0)
|
||||
|
||||
# NOTE we want to reuse the same mask partial placement so that we can reuse the same mask that generates
|
||||
# from the input indices and use it for output reduction
|
||||
rowwise_sharding: PlacementList = [
|
||||
embedding_partial_placement,
|
||||
Shard(0),
|
||||
embedding_partial_placement,
|
||||
]
|
||||
single_mesh_dim_strategies.append(rowwise_sharding)
|
||||
|
||||
# batch dim sharding, weight replicated, input can shard on any dim, output follows input
|
||||
for input_dim in range(len(indices_shape)):
|
||||
batch_sharding: PlacementList = [
|
||||
Shard(input_dim),
|
||||
Replicate(),
|
||||
Shard(input_dim),
|
||||
]
|
||||
single_mesh_dim_strategies.append(batch_sharding)
|
||||
|
||||
return expand_to_full_mesh_op_strategy(mesh, op_schema, single_mesh_dim_strategies)
|
||||
|
||||
|
||||
@register_op_strategy(
|
||||
aten.embedding_dense_backward.default,
|
||||
schema_info=RuntimeSchemaInfo(static_argnum=2),
|
||||
)
|
||||
def embedding_dense_backward_strategy(op_schema: OpSchema) -> StrategyType:
|
||||
"""
|
||||
This strategy handles embedding op. We have two possible embedding shardings:
|
||||
rowwise and colwise
|
||||
"""
|
||||
grad_out_strategy = cast(OpStrategy, op_schema.args_schema[0])
|
||||
indices_strategy = cast(OpStrategy, op_schema.args_schema[1])
|
||||
mesh = op_schema.get_mesh_from_args()
|
||||
|
||||
grad_out_shape = grad_out_strategy.shape
|
||||
indices_shape = indices_strategy.shape
|
||||
grad_out_ndim = len(grad_out_shape)
|
||||
|
||||
single_mesh_dim_strategies = []
|
||||
|
||||
# placement list stores placements of [output, weight, input_indices]
|
||||
# first we always have replicate all for inputs and output
|
||||
all_replicate: PlacementList = [Replicate()] * 3
|
||||
single_mesh_dim_strategies.append(all_replicate)
|
||||
|
||||
# colwise sharding backward, grad_out shard on last dim, input replicate,
|
||||
# weight grad shard colwise
|
||||
colwise_sharding: PlacementList = [Shard(1), Shard(grad_out_ndim - 1), Replicate()]
|
||||
single_mesh_dim_strategies.append(colwise_sharding)
|
||||
|
||||
# batch dim sharding, weight replicated, grad_out/input have same sharding
|
||||
# that can shard on any dim, weight grad partial
|
||||
for input_dim in range(len(indices_shape)):
|
||||
batch_sharding: PlacementList = [Partial(), Shard(input_dim), Shard(input_dim)]
|
||||
single_mesh_dim_strategies.append(batch_sharding)
|
||||
|
||||
# grad_out partial, input replicate, weight grad keep partial
|
||||
partial_sharding: PlacementList = [Partial(), Partial(), Replicate()]
|
||||
single_mesh_dim_strategies.append(partial_sharding)
|
||||
|
||||
return expand_to_full_mesh_op_strategy(mesh, op_schema, single_mesh_dim_strategies)
|
||||
+44
@@ -0,0 +1,44 @@
|
||||
# mypy: allow-untyped-defs
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates
|
||||
from dataclasses import dataclass
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
@dataclass
|
||||
class MaskBuffer:
|
||||
data: torch.Tensor | None = None
|
||||
# refcount allows shared usage of the MaskBuffer, as long as all users have the same data
|
||||
refcount: int = 0
|
||||
|
||||
def materialize_mask(self, mask):
|
||||
if self.refcount == 0:
|
||||
self.data = mask
|
||||
else:
|
||||
if self.data is None:
|
||||
raise AssertionError
|
||||
if not torch.equal(self.data, mask):
|
||||
raise RuntimeError(
|
||||
"MaskBuffer has been materialized with conflicting data"
|
||||
)
|
||||
self.refcount += 1
|
||||
|
||||
def release_mask(self):
|
||||
if self.refcount == 0 or self.data is None:
|
||||
raise RuntimeError("MaskBuffer has not been materialized")
|
||||
self.refcount -= 1
|
||||
if self.refcount == 0:
|
||||
self.data = None
|
||||
|
||||
def apply_mask(self, tensor):
|
||||
if self.refcount == 0 or self.data is None:
|
||||
raise RuntimeError("MaskBuffer has not been materialized")
|
||||
|
||||
# NOTE: MaskPartial is being used by the embedding op and the gather op.
|
||||
# For gather, the mask has the same dimension as the output tensor, whereas
|
||||
# the output of the embedding op has an additional dimension compare to the input,
|
||||
# hence the output masking logic below having two different cases.
|
||||
if tensor.ndim == self.data.ndim:
|
||||
tensor[self.data] = 0.0
|
||||
else:
|
||||
tensor[self.data, :] = 0.0
|
||||
File diff suppressed because it is too large
Load Diff
+1327
File diff suppressed because it is too large
Load Diff
+592
@@ -0,0 +1,592 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates
|
||||
from collections.abc import Callable
|
||||
|
||||
import torch
|
||||
from torch._ops import OpOverload
|
||||
from torch.distributed.tensor._dtensor_spec import TensorMeta
|
||||
from torch.distributed.tensor._op_schema import ArgsType, KwargsType, RuntimeSchemaInfo
|
||||
from torch.distributed.tensor._ops.single_dim_strategy import (
|
||||
_ShardingPlaceholder,
|
||||
register_single_dim_strategy,
|
||||
)
|
||||
from torch.distributed.tensor._ops.utils import infer_broadcast_dims_map
|
||||
from torch.distributed.tensor.placement_types import Partial, Placement, Replicate
|
||||
|
||||
|
||||
aten = torch.ops.aten
|
||||
prims = torch.ops.prims
|
||||
|
||||
# Linear pointwise ops, split by linearity type.
|
||||
unary_linear_ops = [aten.to.dtype]
|
||||
|
||||
|
||||
def _common_pointwise_single_dim_strategy(
|
||||
partial_extra_rules: list[list[Placement | _ShardingPlaceholder]] | None = None,
|
||||
) -> Callable[
|
||||
[OpOverload, ArgsType, KwargsType], list[list[Placement | _ShardingPlaceholder]]
|
||||
]:
|
||||
"""Factory for single-dim strategies that add partial placement rules.
|
||||
|
||||
Returns strategies shaped [output, *args] only. Tensor kwarg placements
|
||||
(e.g. ``out``, ``lr``) are appended by the wrapper in
|
||||
``_register_single_dim_pointwise``.
|
||||
"""
|
||||
|
||||
def strategy(
|
||||
op: OpOverload,
|
||||
args_schema: ArgsType,
|
||||
kwargs_schema: KwargsType,
|
||||
) -> list[list[Placement | _ShardingPlaceholder]]:
|
||||
tensor_arg_metas: list[TensorMeta] = [
|
||||
arg for arg in args_schema if isinstance(arg, TensorMeta)
|
||||
]
|
||||
common_shape = torch.broadcast_shapes(
|
||||
*[arg.shape for arg in args_schema if isinstance(arg, TensorMeta)]
|
||||
)
|
||||
# For multi-output ops (e.g. frexp), all outputs share the same
|
||||
# pointwise sharding, so replicate the output placement.
|
||||
num_outputs = sum(1 for r in op._schema.returns if "Tensor" in str(r.type))
|
||||
placements: list[list[Placement | _ShardingPlaceholder]] = []
|
||||
for i in range(len(common_shape)):
|
||||
shard_placements: list[Placement | _ShardingPlaceholder] = [
|
||||
_ShardingPlaceholder(i)
|
||||
] * num_outputs
|
||||
for arg in tensor_arg_metas:
|
||||
common_dim_to_arg_dim = infer_broadcast_dims_map(
|
||||
common_shape, arg.shape
|
||||
)
|
||||
# If the output shard dim maps to an input dim, shard that
|
||||
# input dim; otherwise it was broadcast, so replicate.
|
||||
if common_dim_to_arg_dim[i] >= 0:
|
||||
shard_placements.append(
|
||||
_ShardingPlaceholder(common_dim_to_arg_dim[i])
|
||||
)
|
||||
else:
|
||||
shard_placements.append(Replicate())
|
||||
placements.append(shard_placements)
|
||||
if partial_extra_rules:
|
||||
n_tensors = len(tensor_arg_metas)
|
||||
expected_len = num_outputs + n_tensors
|
||||
for rule in partial_extra_rules:
|
||||
# Filter rather than assert: some ops (e.g. mul.Tensor) mix
|
||||
# unary rules (len 2, for scalar promotion) and binary rules
|
||||
# (len 3, for tensor-tensor), so mismatched lengths are expected.
|
||||
# see _MUL_RULES to see how _UNARY_LINEAR_RULES handles the
|
||||
# scalar promotion case
|
||||
if len(rule) == expected_len:
|
||||
placements.append(rule)
|
||||
return placements
|
||||
|
||||
return strategy
|
||||
|
||||
|
||||
def _is_list_op(op: OpOverload) -> bool:
|
||||
"""Returns True if op is a foreach, amp_foreach, or fused op."""
|
||||
name = op.name()
|
||||
return name.startswith(("aten::_foreach_", "aten::_amp_foreach_", "aten::_fused_"))
|
||||
|
||||
|
||||
# The state_steps arg of fused adam / adamw is a Replicate scalar tensor, which will be put on
|
||||
# the compute_mesh of an op across all parameter groups, even when not all parameter groups
|
||||
# are on the same device mesh. This idx will help avoid hitting exceptions or unnecessary
|
||||
# redistribute during sharding propagation.
|
||||
_FUSED_OP_SCALAR_IDX = 5
|
||||
|
||||
# Ops registered with extra Partial rules; populated by _register_single_dim_pointwise
|
||||
# when partial_extra_rules is not None, to avoid double-registration from tag discovery.
|
||||
_specially_registered_ops: set[OpOverload] = set()
|
||||
|
||||
|
||||
def _register_single_dim_pointwise(
|
||||
op: OpOverload,
|
||||
partial_extra_rules: list[list[Placement]] | None = None,
|
||||
static_argnum: int = 0,
|
||||
) -> None:
|
||||
if partial_extra_rules is not None:
|
||||
_specially_registered_ops.add(op)
|
||||
inner_fn = _common_pointwise_single_dim_strategy(
|
||||
partial_extra_rules=partial_extra_rules # pyrefly: ignore[bad-argument-type]
|
||||
)
|
||||
|
||||
# Wrap to append tensor kwarg placements in schema declaration order.
|
||||
# out = output placement (s[0]); everything else (e.g. lr) = Replicate.
|
||||
# TODO: move kwargs handling upstream if this works
|
||||
def strategy_fn(
|
||||
op: OpOverload,
|
||||
args: ArgsType,
|
||||
kwargs: KwargsType,
|
||||
_fn: Callable = inner_fn,
|
||||
) -> list[list[Placement | _ShardingPlaceholder]]:
|
||||
strategies = _fn(op, args, kwargs)
|
||||
kw_names = [k for k, v in kwargs.items() if isinstance(v, TensorMeta)]
|
||||
if not kw_names:
|
||||
return strategies
|
||||
return [
|
||||
s + [s[0] if name == "out" else Replicate() for name in kw_names]
|
||||
for s in strategies
|
||||
]
|
||||
|
||||
if _is_list_op(op):
|
||||
schema_info = RuntimeSchemaInfo(needs_pytree=True)
|
||||
else:
|
||||
schema_info = RuntimeSchemaInfo(static_argnum, static_kwargkey=["out"])
|
||||
# Fused ops (e.g. _fused_adam_) have state_steps on a potentially different
|
||||
# mesh; see the note in expand_to_full_mesh_op_strategy for details.
|
||||
different_mesh_args: list[int] | None = None
|
||||
if op.name().startswith("aten::_fused_"):
|
||||
different_mesh_args = [_FUSED_OP_SCALAR_IDX]
|
||||
register_single_dim_strategy(
|
||||
op,
|
||||
schema_info=schema_info,
|
||||
allow_uneven_sharding=True,
|
||||
allow_unbacked_sharding=True,
|
||||
different_mesh_args=different_mesh_args,
|
||||
)(strategy_fn)
|
||||
|
||||
|
||||
_UNARY_LINEAR_RULES: list[list[Placement]] = [
|
||||
[Partial("sum"), Partial("sum")],
|
||||
[Partial("avg"), Partial("avg")],
|
||||
]
|
||||
|
||||
binary_additive_ops = [
|
||||
aten.add.Tensor,
|
||||
aten.add_.Tensor,
|
||||
aten.add.out,
|
||||
aten.sub.Tensor,
|
||||
aten.sub_.Tensor,
|
||||
aten.sub.out,
|
||||
# foreach variants
|
||||
aten._foreach_add.List,
|
||||
aten._foreach_add_.List,
|
||||
aten._foreach_sub.List,
|
||||
aten._foreach_sub_.List,
|
||||
]
|
||||
|
||||
_BINARY_ADDITIVE_RULES: list[list[Placement]] = [
|
||||
[Partial("sum"), Partial("sum"), Partial("sum")],
|
||||
[Partial("avg"), Partial("avg"), Partial("avg")],
|
||||
# P(x), R -> P(x): adding/subtracting a replicated value preserves partial types
|
||||
# avg, max, min. sum would result in R being added n times, n = num_ranks
|
||||
# (the replicated value is constant across ranks, so reduce order is unaffected)
|
||||
[Partial("avg"), Partial("avg"), Replicate()],
|
||||
[Partial("max"), Partial("max"), Replicate()],
|
||||
[Partial("min"), Partial("min"), Replicate()],
|
||||
# R, P(avg) -> P(avg): avg is linear so this holds for any alpha
|
||||
# (R, P(max/min) excluded: negative alpha would flip the ordering)
|
||||
[Partial("avg"), Replicate(), Partial("avg")],
|
||||
]
|
||||
|
||||
for op in binary_additive_ops:
|
||||
_register_single_dim_pointwise(op, _BINARY_ADDITIVE_RULES)
|
||||
|
||||
# mul: partials propagate through either arg. div: only through numerator.
|
||||
binary_mul_ops = [
|
||||
aten.mul.Tensor,
|
||||
aten.mul_.Tensor,
|
||||
aten.mul.out,
|
||||
# foreach variants
|
||||
aten._foreach_mul.List,
|
||||
aten._foreach_mul_.List,
|
||||
aten._foreach_mul.Tensor,
|
||||
aten._foreach_mul_.Tensor,
|
||||
]
|
||||
binary_div_ops = [
|
||||
aten.div.Tensor,
|
||||
aten.div_.Tensor,
|
||||
aten.div.out,
|
||||
# foreach variants
|
||||
aten._foreach_div.List,
|
||||
aten._foreach_div_.List,
|
||||
aten._foreach_div.Tensor,
|
||||
aten._foreach_div_.Tensor,
|
||||
]
|
||||
|
||||
# _UNARY_LINEAR_RULES handles the scalar promotion case: Python's __mul__/__truediv__
|
||||
# promote scalars to 0-dim tensors, so aten.mul.Scalar dispatches as aten.mul.Tensor
|
||||
# with n_tensors=1, matching the length-2 unary rules.
|
||||
_MUL_RULES: list[list[Placement]] = [
|
||||
[Partial("sum"), Partial("sum"), Replicate()],
|
||||
[Partial("avg"), Partial("avg"), Replicate()],
|
||||
[Partial("sum"), Replicate(), Partial("sum")],
|
||||
[Partial("avg"), Replicate(), Partial("avg")],
|
||||
]
|
||||
|
||||
_DIV_RULES: list[list[Placement]] = [
|
||||
[Partial("sum"), Partial("sum"), Replicate()],
|
||||
[Partial("avg"), Partial("avg"), Replicate()],
|
||||
]
|
||||
|
||||
for op in binary_mul_ops:
|
||||
_register_single_dim_pointwise(op, _UNARY_LINEAR_RULES + _MUL_RULES)
|
||||
|
||||
for op in binary_div_ops:
|
||||
_register_single_dim_pointwise(op, _UNARY_LINEAR_RULES + _DIV_RULES)
|
||||
|
||||
scalar_linear_ops = [
|
||||
aten.div.Scalar,
|
||||
aten.div_.Scalar,
|
||||
aten.mul.Scalar,
|
||||
aten.mul_.Scalar,
|
||||
# foreach variants
|
||||
aten._foreach_div.Scalar,
|
||||
aten._foreach_div_.Scalar,
|
||||
aten._foreach_mul.Scalar,
|
||||
aten._foreach_mul_.Scalar,
|
||||
aten._foreach_div.ScalarList,
|
||||
aten._foreach_div_.ScalarList,
|
||||
aten._foreach_mul.ScalarList,
|
||||
aten._foreach_mul_.ScalarList,
|
||||
]
|
||||
|
||||
for op in scalar_linear_ops:
|
||||
_register_single_dim_pointwise(op, _UNARY_LINEAR_RULES, static_argnum=1)
|
||||
|
||||
|
||||
# Non-decreasing unary ops: f(max(a,b)) = max(f(a),f(b)).
|
||||
# Only ops that are non-decreasing on their ENTIRE domain belong here.
|
||||
# Ops with restricted domains (e.g. log on (0,∞), asin on [-1,1]) do NOT qualify
|
||||
# because P(max) offsets can push inputs outside the valid domain.
|
||||
non_decreasing_unary_ops = [
|
||||
aten.asinh.default,
|
||||
aten.asinh_.default,
|
||||
aten.asinh.out,
|
||||
aten.atan.default,
|
||||
aten.atan_.default,
|
||||
aten.atan.out,
|
||||
aten.ceil.default,
|
||||
aten.ceil_.default,
|
||||
aten.ceil.out,
|
||||
aten.deg2rad.default,
|
||||
aten.deg2rad_.default,
|
||||
aten.deg2rad.out,
|
||||
aten.erf.default,
|
||||
aten.erf_.default,
|
||||
aten.erf.out,
|
||||
aten.exp.default,
|
||||
aten.exp_.default,
|
||||
aten.exp.out,
|
||||
aten.exp2.default,
|
||||
aten.exp2_.default,
|
||||
aten.exp2.out,
|
||||
aten.expm1.default,
|
||||
aten.expm1_.default,
|
||||
aten.expm1.out,
|
||||
aten.floor.default,
|
||||
aten.floor_.default,
|
||||
aten.floor.out,
|
||||
aten.rad2deg.default,
|
||||
aten.rad2deg_.default,
|
||||
aten.rad2deg.out,
|
||||
aten.relu.default,
|
||||
aten.relu_.default,
|
||||
aten.round.decimals,
|
||||
aten.round.default,
|
||||
aten.round_.decimals,
|
||||
aten.round_.default,
|
||||
aten.round.decimals_out,
|
||||
aten.round.out,
|
||||
aten.sgn.default,
|
||||
aten.sgn_.default,
|
||||
aten.sgn.out,
|
||||
aten.sigmoid.default,
|
||||
aten.sigmoid_.default,
|
||||
aten.sigmoid.out,
|
||||
aten.sign.default,
|
||||
aten.sign_.default,
|
||||
aten.sign.out,
|
||||
aten.sinh.default,
|
||||
aten.sinh_.default,
|
||||
aten.sinh.out,
|
||||
aten.tanh.default,
|
||||
aten.tanh_.default,
|
||||
aten.tanh.out,
|
||||
aten.trunc.default,
|
||||
aten.trunc_.default,
|
||||
aten.trunc.out,
|
||||
# nan_to_num is non-decreasing on its entire domain (including nan/inf):
|
||||
# it maps -inf→min, nan→0, inf→max, and is identity elsewhere.
|
||||
aten.nan_to_num.default,
|
||||
aten.nan_to_num_.default,
|
||||
aten.nan_to_num.out,
|
||||
# hardshrink: x if |x|>lambd else 0. Non-decreasing on entire domain.
|
||||
aten.hardshrink.default,
|
||||
# I1(x) is monotonically non-decreasing for all real x.
|
||||
aten.special_modified_bessel_i1.default,
|
||||
# threshold(x, t, v): x if x > t else v. Non-decreasing for v <= t (the
|
||||
# common case, including the default v=0, t=0).
|
||||
aten.threshold.default,
|
||||
# foreach variants
|
||||
aten._foreach_exp.default,
|
||||
aten._foreach_exp_.default,
|
||||
aten._foreach_clamp_max_.Scalar,
|
||||
aten._foreach_clamp_min_.Scalar,
|
||||
]
|
||||
|
||||
_NON_DECREASING_RULES: list[list[Placement]] = [
|
||||
[Partial("max"), Partial("max")],
|
||||
[Partial("min"), Partial("min")],
|
||||
]
|
||||
|
||||
for op in non_decreasing_unary_ops:
|
||||
_register_single_dim_pointwise(op, _NON_DECREASING_RULES)
|
||||
|
||||
# Non-increasing unary ops: f(max(a,b)) = min(f(a),f(b)).
|
||||
# Note: acos excluded due to domain constraints [-1,1] causing validation failures
|
||||
non_increasing_unary_ops: list[OpOverload] = [
|
||||
aten.erfc.default,
|
||||
aten.erfc_.default,
|
||||
aten.erfc.out,
|
||||
aten.special_erfcx.default,
|
||||
aten.special_erfcx.out,
|
||||
]
|
||||
|
||||
_NON_INCREASING_RULES: list[list[Placement]] = [
|
||||
[Partial("min"), Partial("max")],
|
||||
[Partial("max"), Partial("min")],
|
||||
]
|
||||
|
||||
for op in non_increasing_unary_ops:
|
||||
_register_single_dim_pointwise(op, _NON_INCREASING_RULES)
|
||||
|
||||
# Bessel K functions are strictly decreasing for x > 0 but undefined at x <= 0.
|
||||
# Only P(min)->P(max) is safe: P(min) offsets add positive values to the
|
||||
# non-holding rank, keeping all inputs positive. P(max) offsets subtract,
|
||||
# which can push inputs to x <= 0 producing NaN.
|
||||
_POSITIVE_DOMAIN_NON_INCREASING_RULES: list[list[Placement]] = [
|
||||
[Partial("max"), Partial("min")],
|
||||
]
|
||||
|
||||
for op in [
|
||||
aten.special_modified_bessel_k0.default,
|
||||
aten.special_modified_bessel_k1.default,
|
||||
aten.special_scaled_modified_bessel_k0.default,
|
||||
aten.special_scaled_modified_bessel_k1.default,
|
||||
]:
|
||||
_register_single_dim_pointwise(op, _POSITIVE_DOMAIN_NON_INCREASING_RULES)
|
||||
|
||||
# neg is linear: -(A1 + A2) = -A1 + -A2
|
||||
neg_ops = [
|
||||
aten.neg.default,
|
||||
aten.neg_.default,
|
||||
aten.neg.out,
|
||||
# foreach variants
|
||||
aten._foreach_neg.default,
|
||||
aten._foreach_neg_.default,
|
||||
]
|
||||
|
||||
_NEG_RULES: list[list[Placement]] = _UNARY_LINEAR_RULES + _NON_INCREASING_RULES
|
||||
|
||||
for op in neg_ops:
|
||||
_register_single_dim_pointwise(op, _NEG_RULES)
|
||||
|
||||
# xlog1py(x, y) = x * log1p(y). Linear in x with y replicated:
|
||||
# (a+b)*log1p(y) = a*log1p(y) + b*log1p(y).
|
||||
_XLOG1PY_RULES: list[list[Placement]] = [
|
||||
[Partial("sum"), Partial("sum"), Replicate()],
|
||||
[Partial("avg"), Partial("avg"), Replicate()],
|
||||
]
|
||||
|
||||
for op in [aten.special_xlog1py.default, aten.special_xlog1py.other_scalar]:
|
||||
_register_single_dim_pointwise(op, _XLOG1PY_RULES)
|
||||
|
||||
|
||||
# All-partial-preserving unary ops: P(x)->P(x) for all x.
|
||||
# TODO: positive should be removed once CIA (Copy Is All) optimizes it away.
|
||||
all_partial_preserving_unary_ops = [
|
||||
aten.to.dtype,
|
||||
aten.positive.default,
|
||||
]
|
||||
|
||||
_ALL_PARTIAL_PRESERVING_RULES: list[list[Placement]] = [
|
||||
[Partial(r), Partial(r)] for r in ("sum", "avg", "max", "min")
|
||||
]
|
||||
|
||||
for op in all_partial_preserving_unary_ops:
|
||||
_register_single_dim_pointwise(op, _ALL_PARTIAL_PRESERVING_RULES)
|
||||
|
||||
all_partial_preserving_binary_ops = [
|
||||
aten.copy_.default,
|
||||
prims.copy_to.default,
|
||||
]
|
||||
|
||||
_ALL_PARTIAL_BINARY_PRESERVING_RULES: list[list[Placement]] = [
|
||||
[Partial(r), Partial(r), Partial(r)] for r in ("sum", "avg", "max", "min")
|
||||
]
|
||||
|
||||
for op in all_partial_preserving_binary_ops:
|
||||
_register_single_dim_pointwise(op, _ALL_PARTIAL_BINARY_PRESERVING_RULES)
|
||||
|
||||
# Monotonic increasing in both args but don't preserve any specific partial type.
|
||||
monotonic_binary_ops = [
|
||||
aten.logaddexp.default,
|
||||
aten.logaddexp.out,
|
||||
aten.logaddexp2.default,
|
||||
aten.logaddexp2.out,
|
||||
]
|
||||
|
||||
_MONOTONE_BINARY_BASE_RULES: list[list[Placement]] = [
|
||||
[Partial("max"), Partial("max"), Replicate()],
|
||||
[Partial("max"), Replicate(), Partial("max")],
|
||||
[Partial("min"), Partial("min"), Replicate()],
|
||||
[Partial("min"), Replicate(), Partial("min")],
|
||||
]
|
||||
|
||||
for op in monotonic_binary_ops:
|
||||
_register_single_dim_pointwise(op, _MONOTONE_BINARY_BASE_RULES)
|
||||
|
||||
# Binary ops monotonically increasing in both arguments.
|
||||
# max-preserving: P(max)+P(max)->P(max) because max(max(a),max(b)) = max(a,b)
|
||||
monotonic_max_preserving_binary_ops = [
|
||||
aten.clamp_min.Tensor,
|
||||
aten.fmax.default,
|
||||
aten.fmax.out,
|
||||
aten.maximum.default,
|
||||
aten.maximum.out,
|
||||
prims.fmax.default,
|
||||
# foreach variants
|
||||
aten._foreach_maximum_.List,
|
||||
]
|
||||
|
||||
_MONOTONE_MAX_PRESERVING_BINARY_BASE_RULES: list[list[Placement]] = [
|
||||
*_MONOTONE_BINARY_BASE_RULES,
|
||||
[Partial("max"), Partial("max"), Partial("max")],
|
||||
]
|
||||
|
||||
for op in monotonic_max_preserving_binary_ops:
|
||||
_register_single_dim_pointwise(op, _MONOTONE_MAX_PRESERVING_BINARY_BASE_RULES)
|
||||
|
||||
# min-preserving: P(min)+P(min)->P(min) because min(min(a),min(b)) = min(a,b)
|
||||
monotonic_min_preserving_binary_ops = [
|
||||
aten.clamp_max.Tensor,
|
||||
aten.fmin.default,
|
||||
aten.fmin.out,
|
||||
aten.minimum.default,
|
||||
aten.minimum.out,
|
||||
prims.fmin.default,
|
||||
]
|
||||
|
||||
_MONOTONE_MIN_PRESERVING_BINARY_BASE_RULES: list[list[Placement]] = [
|
||||
*_MONOTONE_BINARY_BASE_RULES,
|
||||
[Partial("min"), Partial("min"), Partial("min")],
|
||||
]
|
||||
|
||||
for op in monotonic_min_preserving_binary_ops:
|
||||
_register_single_dim_pointwise(op, _MONOTONE_MIN_PRESERVING_BINARY_BASE_RULES)
|
||||
|
||||
|
||||
# Ops that are pointwise for DTensor purposes but lack torch.Tag.pointwise.
|
||||
# TODO(pianpwk): add torch.Tag.pointwise to these ops in native_functions.yaml
|
||||
# so this list can be removed.
|
||||
_extra_pointwise_ops: list[OpOverload] = [
|
||||
aten.__irshift__.Scalar,
|
||||
aten.__irshift__.Tensor,
|
||||
aten._conj.default,
|
||||
aten.abs_.default,
|
||||
aten.copysign_.Scalar,
|
||||
aten.copysign_.Tensor,
|
||||
aten.exponential_.default,
|
||||
aten.float_power.Scalar,
|
||||
aten.float_power.Scalar_out,
|
||||
aten.float_power.Tensor_Scalar,
|
||||
aten.float_power.Tensor_Scalar_out,
|
||||
aten.float_power.Tensor_Tensor,
|
||||
aten.float_power.Tensor_Tensor_out,
|
||||
aten.masked_fill_.Scalar,
|
||||
aten.native_dropout_backward.out,
|
||||
aten.polygamma_.default,
|
||||
aten.rrelu_with_noise.default,
|
||||
aten.where.self_out,
|
||||
aten.xlogy_.Scalar_Other,
|
||||
prims.bessel_i0e.default,
|
||||
prims.bessel_i1.default,
|
||||
prims.bessel_i1e.default,
|
||||
prims.bessel_j0.default,
|
||||
prims.bessel_j1.default,
|
||||
prims.div.default,
|
||||
prims.erfcx.default,
|
||||
prims.frexp.default,
|
||||
prims.gcd.default,
|
||||
prims.ndtri.default,
|
||||
prims.ne.default,
|
||||
prims.spherical_bessel_j0.default,
|
||||
prims.zeta.default,
|
||||
# foreach variants
|
||||
aten._foreach_abs.default,
|
||||
aten._foreach_abs_.default,
|
||||
aten._foreach_addcdiv_.Scalar,
|
||||
aten._foreach_addcdiv_.ScalarList,
|
||||
aten._foreach_addcdiv_.Tensor,
|
||||
aten._foreach_addcmul.Scalar,
|
||||
aten._foreach_addcmul_.Scalar,
|
||||
aten._foreach_addcmul_.ScalarList,
|
||||
aten._foreach_addcmul_.Tensor,
|
||||
aten._foreach_lerp_.Scalar,
|
||||
aten._foreach_pow.List,
|
||||
aten._foreach_pow.ScalarList,
|
||||
aten._foreach_reciprocal_.default,
|
||||
aten._foreach_sub.Scalar,
|
||||
aten._foreach_sub_.Scalar,
|
||||
aten._foreach_sub.ScalarList,
|
||||
aten._foreach_sub_.ScalarList,
|
||||
aten._foreach_sqrt.default,
|
||||
aten._foreach_sqrt_.default,
|
||||
aten._foreach_zero_.default,
|
||||
aten._foreach_cos.default,
|
||||
aten._foreach_cos_.default,
|
||||
aten._foreach_log.default,
|
||||
aten._foreach_log_.default,
|
||||
aten._amp_foreach_non_finite_check_and_unscale_.default,
|
||||
# foreach linearity variants
|
||||
aten._foreach_add.Scalar,
|
||||
aten._foreach_add_.Scalar,
|
||||
aten._foreach_add_.ScalarList,
|
||||
# fused optimizer ops
|
||||
aten._fused_adam_.default,
|
||||
aten._fused_adam.default,
|
||||
aten._fused_adam.tensor_lr,
|
||||
aten._fused_adam_.tensor_lr,
|
||||
aten._fused_adamw_.default,
|
||||
aten._fused_adamw.default,
|
||||
aten._fused_adamw.tensor_lr,
|
||||
aten._fused_adamw_.tensor_lr,
|
||||
]
|
||||
|
||||
|
||||
def _get_pointwise_ops_from_tag() -> list[OpOverload]:
|
||||
"""
|
||||
Auto-discover pointwise ops via torch.Tag.pointwise, from ops.aten, ops.prims.
|
||||
"""
|
||||
ops = []
|
||||
for ns in [torch.ops.aten, torch.ops.prims]:
|
||||
for attr_name in dir(ns):
|
||||
attr = getattr(ns, attr_name)
|
||||
if isinstance(attr, torch._ops.OpOverloadPacket):
|
||||
for overload_name in attr.overloads():
|
||||
op = getattr(attr, overload_name)
|
||||
if torch.Tag.pointwise in op.tags:
|
||||
ops.append(op)
|
||||
return ops
|
||||
|
||||
|
||||
pointwise_ops = [
|
||||
op
|
||||
for op in _get_pointwise_ops_from_tag() + _extra_pointwise_ops
|
||||
if op not in _specially_registered_ops
|
||||
]
|
||||
|
||||
|
||||
for op in pointwise_ops:
|
||||
_register_single_dim_pointwise(op)
|
||||
|
||||
|
||||
def register_inductor_prims() -> None:
|
||||
"""Register DTensor sharding strategies for inductor prims ops.
|
||||
|
||||
Called lazily because inductor prims are created via make_prim() in
|
||||
torch._inductor.inductor_prims, which is imported after this module.
|
||||
"""
|
||||
# TODO: handle other inductor prims ops that may need DTensor sharding
|
||||
# strategies (e.g. mul_rn, div_rn). Those are more complicated and not
|
||||
# necessarily pointwise.
|
||||
_register_single_dim_pointwise(prims.fma.default)
|
||||
@@ -0,0 +1,43 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates
|
||||
import torch
|
||||
from torch.distributed.tensor._op_schema import (
|
||||
OpSchema,
|
||||
OpSpec,
|
||||
OpStrategy,
|
||||
StrategyType,
|
||||
)
|
||||
from torch.distributed.tensor._ops.utils import is_tensor_partial, register_op_strategy
|
||||
|
||||
|
||||
aten = torch.ops.aten
|
||||
|
||||
|
||||
@register_op_strategy(
|
||||
[
|
||||
aten.normal_.default,
|
||||
aten.uniform_.default,
|
||||
aten.native_dropout.default,
|
||||
aten.bernoulli_.float,
|
||||
aten.bernoulli.default,
|
||||
]
|
||||
)
|
||||
def random_op_strategy(op_schema: OpSchema) -> StrategyType:
|
||||
self_strategy = op_schema.args_schema[0]
|
||||
if not isinstance(self_strategy, OpStrategy):
|
||||
raise AssertionError
|
||||
|
||||
random_strategy = OpStrategy([])
|
||||
for arg_strategy in self_strategy.strategies:
|
||||
arg_spec = arg_strategy.output_spec
|
||||
if is_tensor_partial(arg_spec):
|
||||
# TODO: figure out how inplace random op should behave when it's partial
|
||||
raise RuntimeError(f"{op_schema.op} with Partial is not supported yet!")
|
||||
random_strategy.strategies.append(
|
||||
OpSpec(
|
||||
output_specs=arg_spec,
|
||||
input_specs=(arg_spec,),
|
||||
redistribute_cost=[[0.0] * len(self_strategy.strategies)],
|
||||
)
|
||||
)
|
||||
|
||||
return random_strategy
|
||||
+1531
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
+1084
File diff suppressed because it is too large
Load Diff
+2304
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,671 @@
|
||||
# mypy: allow-untyped-defs
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates
|
||||
import functools
|
||||
import itertools
|
||||
import operator
|
||||
from collections.abc import Callable, Iterable, Sequence
|
||||
from typing import TypeAlias, TypeVar
|
||||
|
||||
import torch
|
||||
from torch._prims_common import DimsSequenceType, DimsType
|
||||
from torch.distributed.tensor._api import DTensor
|
||||
from torch.distributed.tensor._collective_utils import redistribute_cost
|
||||
from torch.distributed.tensor._dtensor_spec import DTensorSpec, TensorMeta
|
||||
from torch.distributed.tensor._op_schema import (
|
||||
OpSchema,
|
||||
OpSpec,
|
||||
OpStrategy,
|
||||
OutputSharding,
|
||||
PlacementList,
|
||||
RuntimeSchemaInfo,
|
||||
StrategyType,
|
||||
)
|
||||
from torch.distributed.tensor.device_mesh import DeviceMesh
|
||||
from torch.distributed.tensor.placement_types import (
|
||||
_is_shard_like,
|
||||
_StridedShard,
|
||||
Partial,
|
||||
Placement,
|
||||
Replicate,
|
||||
Shard,
|
||||
)
|
||||
|
||||
|
||||
def _get_registration_wrapper(
|
||||
registration_fn,
|
||||
op: torch._ops.OpOverload | list[torch._ops.OpOverload],
|
||||
schema_info: RuntimeSchemaInfo | None,
|
||||
arg_names_that_require_specializing_cache_strategy: list[str] | None,
|
||||
):
|
||||
def wrapper(impl):
|
||||
overloads = op if isinstance(op, list) else [op]
|
||||
for overload in overloads:
|
||||
curr_schema_info = None
|
||||
if (
|
||||
schema_info is None
|
||||
and arg_names_that_require_specializing_cache_strategy is not None
|
||||
):
|
||||
specialized_args = [
|
||||
a.name
|
||||
for a in overload._schema.arguments
|
||||
if a.name in arg_names_that_require_specializing_cache_strategy
|
||||
]
|
||||
if any(specialized_args):
|
||||
curr_schema_info = RuntimeSchemaInfo(
|
||||
static_kwargkey=specialized_args
|
||||
)
|
||||
else:
|
||||
curr_schema_info = schema_info
|
||||
registration_fn(overload, impl, curr_schema_info)
|
||||
return impl
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
# convenient wrapper to register sharding propagation rules
|
||||
def register_prop_rule(
|
||||
op: torch._ops.OpOverload | list[torch._ops.OpOverload],
|
||||
schema_info: RuntimeSchemaInfo | None = None,
|
||||
) -> Callable[
|
||||
[Callable[[OpSchema], OutputSharding]], Callable[[OpSchema], OutputSharding]
|
||||
]:
|
||||
return _get_registration_wrapper(
|
||||
DTensor._op_dispatcher.sharding_propagator.register_sharding_prop_rule,
|
||||
op,
|
||||
schema_info,
|
||||
arg_names_that_require_specializing_cache_strategy=None,
|
||||
)
|
||||
|
||||
|
||||
# Note:
|
||||
# using TypeVar here allows the registration decorator to preserve the specific type info of the wrapped strategy,
|
||||
# while hardcoding the typing on the wrapper (e.g. Callable[[OpSchema], StrategyType]) would mean mypy would treat
|
||||
# the return value of the wrapped strategy as always being a `StrategyType` even if it were a derived class like
|
||||
# MyStrategyType(StrategyType).
|
||||
_OpSchemaT = TypeVar("_OpSchemaT", bound=OpSchema)
|
||||
_StrategyTypeT = TypeVar("_StrategyTypeT", bound=StrategyType)
|
||||
_ShardingStrategyFunc: TypeAlias = Callable[[_OpSchemaT], _StrategyTypeT]
|
||||
|
||||
|
||||
def register_op_strategy(
|
||||
op: torch._ops.OpOverload | list[torch._ops.OpOverload],
|
||||
schema_info: RuntimeSchemaInfo | None = None,
|
||||
) -> Callable[[_ShardingStrategyFunc], _ShardingStrategyFunc]:
|
||||
# For every ATen op that accepts any args in this list,
|
||||
# the arg itself can impact the strides (and potentially the sharding strategy)
|
||||
# of the output tensor.
|
||||
# thus, we will detect ATen schemas with any of these args and ensure
|
||||
# that they get specialized here.
|
||||
arg_names_that_require_specializing_cache_strategy = [
|
||||
"memory_format",
|
||||
]
|
||||
return _get_registration_wrapper(
|
||||
DTensor._op_dispatcher.sharding_propagator.register_op_strategy,
|
||||
op,
|
||||
schema_info,
|
||||
arg_names_that_require_specializing_cache_strategy,
|
||||
)
|
||||
|
||||
|
||||
def replicate_op_strategy(op_schema: OpSchema) -> StrategyType:
|
||||
"""
|
||||
Fallback strategy all use Replication()
|
||||
"""
|
||||
args_strategy = op_schema.args_strategy
|
||||
kwargs_strategy = op_schema.kwargs_strategy
|
||||
inputs_strategy = args_strategy + kwargs_strategy
|
||||
|
||||
output_type = [str(ret.type) for ret in op_schema.op._schema.returns]
|
||||
output_len = output_type.count("Tensor")
|
||||
# TODO(zpcore): Confirm if view op can be handle properly or not. Prevent
|
||||
# handling view ops until confirmed.
|
||||
if op_schema.op.is_view:
|
||||
raise RuntimeError(
|
||||
"fallback strategy is unable to handle view ops until confirmed"
|
||||
)
|
||||
if "List[Tensor]" in output_type:
|
||||
raise RuntimeError(
|
||||
"fallback strategy is unable to handle ops with List[Tensor] output "
|
||||
"because size of the list may depend on the op's input value"
|
||||
)
|
||||
|
||||
mesh = inputs_strategy[0].mesh
|
||||
|
||||
dim_sharding: PlacementList = [Replicate()] * (output_len + len(inputs_strategy))
|
||||
single_dim_placement = [dim_sharding]
|
||||
return expand_to_full_mesh_op_strategy(
|
||||
mesh, op_schema, single_dim_placement, input_index=output_len
|
||||
)
|
||||
|
||||
|
||||
def as_list(
|
||||
x: list[object] | object,
|
||||
# pyre-fixme[11]: Annotation `immutable_list` is not defined as a type.
|
||||
) -> list[object] | torch.fx.immutable_collections.immutable_list: # type: ignore[valid-type]
|
||||
# During tracing, `aten.sum.dim_IntList` uses `immutable_list` for its args,
|
||||
# which is an object but treated as a list by the tracer. Therefore, keep
|
||||
# `immutable_list` intact here as well.
|
||||
if type(x) is list or isinstance(x, torch.fx.immutable_collections.immutable_list):
|
||||
return x
|
||||
else:
|
||||
return [x]
|
||||
|
||||
|
||||
def normalize_dim(dim: int, ndim: int) -> int:
|
||||
return dim if dim >= 0 else dim + ndim
|
||||
|
||||
|
||||
def normalize_dims(dims: DimsType, ndim: int) -> DimsSequenceType:
|
||||
"""Normalize a dim or a sequence of dims, so that they are all positive."""
|
||||
if isinstance(dims, int):
|
||||
dims = (normalize_dim(dims, ndim),)
|
||||
elif isinstance(dims, list):
|
||||
dims = [normalize_dim(dim, ndim) for dim in dims]
|
||||
elif isinstance(dims, tuple):
|
||||
dims = tuple(normalize_dim(dim, ndim) for dim in dims)
|
||||
return dims
|
||||
|
||||
|
||||
def prod(xs: Iterable[int]) -> int:
|
||||
return functools.reduce(operator.mul, xs, 1)
|
||||
|
||||
|
||||
def is_tensor_shardable(
|
||||
shape: Sequence[int],
|
||||
spec: DTensorSpec,
|
||||
allow_unbacked_sharding: bool | None = None,
|
||||
) -> bool:
|
||||
"""
|
||||
Check if the shape is shardable according to the spec.
|
||||
|
||||
This function handles both `Shard` and `_StridedShard` placements:
|
||||
- For `Shard`: checks if the tensor dimension size >= number of shards
|
||||
- For `_StridedShard`: additionally checks if the dimension is shardable after
|
||||
splitting with the placement's `split_factor`
|
||||
|
||||
allow_unbacked_sharding: determines the fallback value if unbacked shapes are involved,
|
||||
and the queried shape properties are not statically known.
|
||||
|
||||
e.g. when asking if u0 is shardable on num_shards, and u0 has generic bounds [0, inf],
|
||||
the behavior of allow_unbacked_sharding is:
|
||||
|
||||
None: will data-dependent error
|
||||
True: assumes shardability; we return True, allowing zero-size shards at runtime when u0 < num_shards.
|
||||
False: returns False, and lower-bounding u0, e.g. torch._check(u0 >= num_shards), is needed to enable sharding.
|
||||
"""
|
||||
from torch.fx.experimental.symbolic_shapes import guard_or_false, guard_or_true
|
||||
|
||||
if allow_unbacked_sharding not in [None, True, False]:
|
||||
raise AssertionError
|
||||
guard_fn = {
|
||||
None: bool,
|
||||
True: guard_or_false,
|
||||
False: guard_or_true,
|
||||
}[allow_unbacked_sharding]
|
||||
|
||||
# number of shards in each tensor dimension
|
||||
num_shards = [1] * len(shape)
|
||||
for i, placement in enumerate(spec.placements):
|
||||
if _is_shard_like(placement):
|
||||
shard_dim = placement.dim
|
||||
if shard_dim >= len(shape):
|
||||
return False
|
||||
num_shards[shard_dim] *= spec.mesh.size(i)
|
||||
if isinstance(placement, _StridedShard):
|
||||
# make sure tensor dim `shard_dim` is shardable after splitting
|
||||
# with split_factor
|
||||
if guard_fn(
|
||||
shape[shard_dim] < num_shards[shard_dim] * placement.split_factor
|
||||
):
|
||||
return False
|
||||
else:
|
||||
if guard_fn(shape[shard_dim] < num_shards[shard_dim]):
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
|
||||
def is_tensor_evenly_shardable(shape: Sequence[int], spec: DTensorSpec) -> bool:
|
||||
"""Check if the shape is evenly shardable according to the spec."""
|
||||
# number of shards in each tensor dimension
|
||||
num_shards = [1] * len(shape)
|
||||
for i, placement in enumerate(spec.placements):
|
||||
if _is_shard_like(placement):
|
||||
shard_dim = placement.dim
|
||||
if shard_dim >= len(shape):
|
||||
return False
|
||||
num_shards[shard_dim] *= spec.mesh.size(i)
|
||||
if isinstance(placement, _StridedShard):
|
||||
if (
|
||||
shape[shard_dim] % (placement.split_factor * num_shards[shard_dim])
|
||||
!= 0
|
||||
):
|
||||
return False
|
||||
else:
|
||||
if shape[shard_dim] % num_shards[shard_dim] != 0:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def is_tensor_evenly_shardable_on_dim(
|
||||
shape: Sequence[int], spec: DTensorSpec, dim: int
|
||||
) -> bool:
|
||||
"""Check if the shape is evenly shardable according to the spec on dim."""
|
||||
dim = normalize_dim(dim, len(shape))
|
||||
|
||||
num_shards = 1
|
||||
for i, placement in enumerate(spec.placements):
|
||||
if _is_shard_like(placement) and placement.dim == dim:
|
||||
num_shards *= spec.mesh.size(i)
|
||||
if isinstance(placement, _StridedShard):
|
||||
# _StridedShard._split_tensor first chunks into split_factor
|
||||
# groups, then into num_shards within each group, so the dim
|
||||
# must be divisible by the product of both. This is stricter
|
||||
# than the final num_shards check and implies it. Note:
|
||||
# num_shards already includes spec.mesh.size(i) from this
|
||||
# iteration, so the check covers the full shard count.
|
||||
if shape[dim] % (placement.split_factor * num_shards) != 0:
|
||||
return False
|
||||
|
||||
return shape[dim] % num_shards == 0
|
||||
|
||||
|
||||
def is_tensor_dim_sharded(spec: DTensorSpec, dim: int) -> bool:
|
||||
"""Return True if tensor dim is sharded."""
|
||||
return any(_is_shard_like(p) and p.dim == dim for p in spec.placements)
|
||||
|
||||
|
||||
def is_tensor_partial(spec: DTensorSpec) -> bool:
|
||||
"""Return True if tensor is partial on the mesh."""
|
||||
return any(p.is_partial() for p in spec.placements)
|
||||
|
||||
|
||||
def infer_broadcast_dims_map(
|
||||
common_shape: torch.Size, input_shape: torch.Size
|
||||
) -> list[int]:
|
||||
# infer the broadcast dims map, where it maps from the common shape dim to the input shape dim
|
||||
# this is aligned with the broadcast semantics
|
||||
# e.g. if common_shape = [1, 2, 3, 4] and input_shape = [2, 3, 4],
|
||||
# broadcast_dims_map will be [-1, 0, 1, 2]
|
||||
# meaning that dim 0 in the output has no mapping to the input, and dim 1 in the output maps to dim 0 in the input
|
||||
from torch.fx.experimental.symbolic_shapes import guard_or_false
|
||||
|
||||
common_ndim = len(common_shape)
|
||||
input_ndim = len(input_shape)
|
||||
broadcast_dims_map = [-1] * common_ndim
|
||||
for idx in range(-1, -1 - input_ndim, -1):
|
||||
if guard_or_false(input_shape[idx] == common_shape[idx]):
|
||||
broadcast_dims_map[common_ndim + idx] = input_ndim + idx
|
||||
return broadcast_dims_map
|
||||
|
||||
|
||||
def map_placements_after_broadcast(
|
||||
placements: tuple[Placement, ...],
|
||||
shape: torch.Size,
|
||||
broadcast_dims_map: list[int],
|
||||
partial_to_replicate: bool = False,
|
||||
) -> tuple[Placement, ...]:
|
||||
"""Map each placement based on the output shape after broadcast."""
|
||||
new_placements: list[Placement] = []
|
||||
for placement in placements:
|
||||
if isinstance(placement, Partial):
|
||||
if partial_to_replicate:
|
||||
# map the partial placement to replicate
|
||||
new_placements.append(Replicate())
|
||||
else:
|
||||
new_placements.append(placement)
|
||||
elif isinstance(placement, Replicate):
|
||||
new_placements.append(placement)
|
||||
else:
|
||||
if not _is_shard_like(placement):
|
||||
raise AssertionError
|
||||
shard_dim = normalize_dim(placement.dim, len(shape))
|
||||
new_shard_dim = broadcast_dims_map[shard_dim]
|
||||
if new_shard_dim != -1:
|
||||
# there's a map from the common shape shard dim to
|
||||
# the input shape shard dim before broadcasting,
|
||||
# use that instead
|
||||
if isinstance(placement, _StridedShard):
|
||||
new_placements.append(
|
||||
_StridedShard(
|
||||
new_shard_dim, split_factor=placement.split_factor
|
||||
)
|
||||
)
|
||||
else:
|
||||
new_placements.append(Shard(new_shard_dim))
|
||||
else:
|
||||
# there's no map between common shape shard dim and
|
||||
# the input shape shard dim before broadcasting,
|
||||
# in this case it means implicit broadcasting happen
|
||||
# in this dim, so we can just mark it as replicate
|
||||
# and implicit broadcast will broadcast automatically
|
||||
# to the sharded shape
|
||||
new_placements.append(Replicate())
|
||||
|
||||
return tuple(new_placements)
|
||||
|
||||
|
||||
def generate_redistribute_costs(
|
||||
src_strategy: OpStrategy, dst_spec: DTensorSpec
|
||||
) -> list[float]:
|
||||
"""Generates one row in the 'redistribute_costs' matrix in an OpSpec
|
||||
The length of the returned list will match the number of strategies in 'src_strategy'.
|
||||
|
||||
Each value in the row is the cost of redistributing from a particular src_strategy to dst_spec.
|
||||
"""
|
||||
redistribute_costs: list[float] = [
|
||||
redistribute_cost(strat.output_spec, dst_spec)
|
||||
for strat in src_strategy.strategies
|
||||
]
|
||||
|
||||
return redistribute_costs
|
||||
|
||||
|
||||
def expand_to_full_mesh_op_strategy(
|
||||
mesh: DeviceMesh,
|
||||
op_schema: OpSchema,
|
||||
single_mesh_dim_strategies: list[PlacementList],
|
||||
*,
|
||||
output_tensor_meta: TensorMeta | Sequence[TensorMeta | None] | None = None,
|
||||
input_index: int = 1,
|
||||
inplace_op: bool = False,
|
||||
allow_unbacked_sharding: bool | None = None,
|
||||
allow_uneven_sharding: bool = False,
|
||||
is_valid_strategy_cb: Callable[
|
||||
[list[DTensorSpec], DTensorSpec | tuple[DTensorSpec | None, ...]], bool
|
||||
]
|
||||
| None = None,
|
||||
different_mesh_args: list[int] | None = None,
|
||||
) -> OpStrategy:
|
||||
"""
|
||||
Convenience function to allow writing a sharding strategy considering only a single mesh dimension,
|
||||
and have it expanded combinatorially to all mesh dimensions.
|
||||
|
||||
Args:
|
||||
mesh (DeviceMesh): the device mesh to expand the strategy to
|
||||
op_schema (OpSchema): the op schema
|
||||
single_mesh_dim_strategies (list[PlacementList]): the sharding strategies to expand. The outer list is over
|
||||
different strategies. The inner PlacementList is over the outputs and inputs of the op. If input_index is 1,
|
||||
a PlacementList looks like [output_placement, input_placement1, input_placement2, ...].
|
||||
output_tensor_meta: tensor metadata for the output(s), used to populate DTensorSpec.tensor_meta field
|
||||
input_index: the number of outputs of the op, defaults to 1
|
||||
inplace_op: whether the op is inplace or not, defaults to False
|
||||
is_valid_strategy_cb: a callback function to filter out invalid sharding rules, defaults to None.
|
||||
|
||||
Example: Let's say `my_op(tensor_x, tensor_y) - > output_tensor` can support sharding or replicating tensor_x,
|
||||
but always requires tensor_y to be replicated. We can specify these valid combinations ignoring mesh dims.
|
||||
Then, we can rely on `expand_to_full_mesh_op_strategy` to create every possible combination of these shardings
|
||||
over multiple mesh dimensions, filtering out any combinations that are invalid based on the actual mesh dim size.
|
||||
|
||||
single_mesh_dim_strategies = [
|
||||
# first strategy: return output sharded on first dim, shard tensor_x on its first dim, replicate tensor_y
|
||||
[Shard(0), Shard(0), Replicate()]
|
||||
# second strategy: replicate output, and both inputs
|
||||
[Replicate(), Replicate(), Replicate()]
|
||||
]
|
||||
"""
|
||||
# Expand the single_mesh_dim_strategies to full mesh dim strategies.
|
||||
all_mesh_dim_strategies = [single_mesh_dim_strategies] * mesh.ndim
|
||||
|
||||
strategy_combs = itertools.product(*all_mesh_dim_strategies)
|
||||
|
||||
args_strategy = op_schema.args_strategy
|
||||
kwargs_strategy = op_schema.kwargs_strategy
|
||||
input_args_strategy = args_strategy + kwargs_strategy
|
||||
|
||||
# Propagate use_strided_shard_as_shard_order from inputs so that
|
||||
# strategy specs with _StridedShard get the correct flag (and thus
|
||||
# correct shard_order) at construction time, avoiding shard_order
|
||||
# mismatches in redistribute_cost computation.
|
||||
_input_use_strided: bool | None = None
|
||||
for input_strat in input_args_strategy:
|
||||
input_spec = input_strat.strategies[0].output_spec
|
||||
if any(isinstance(p, _StridedShard) for p in input_spec.placements):
|
||||
_input_use_strided = input_spec.use_strided_shard_as_shard_order
|
||||
break
|
||||
|
||||
all_strategies = []
|
||||
# Track input placements if we skip strategies due to inplace placement mismatch
|
||||
blocking_inplace_input_placements: tuple[Placement, ...] | None = None
|
||||
for strategy_comb in strategy_combs:
|
||||
spec_list: list[DTensorSpec | None] = []
|
||||
# Track how many non-None output specs we've seen (for output_tensor_meta indexing).
|
||||
# This is needed because output_tensor_meta may contain only non-None entries,
|
||||
# so we can't use position directly when there are None entries in the output.
|
||||
output_spec_count = 0
|
||||
# Track input args separately since not all tensor inputs have OpStrategy
|
||||
# (e.g., philox_seed/offset in SDPA are scalar tensors without OpStrategy)
|
||||
input_strategy_counter = 0
|
||||
for position, specs in enumerate(zip(*strategy_comb, strict=True)):
|
||||
if specs[0] is not None:
|
||||
# Populate tensor_meta field for both output and input specs,
|
||||
# including for tuple output cases
|
||||
tensor_meta = None
|
||||
# Use position to determine output vs input territory
|
||||
# (position includes None entries, unlike the old spec_index)
|
||||
if position < input_index:
|
||||
# This is an output position
|
||||
if output_tensor_meta is not None:
|
||||
if isinstance(output_tensor_meta, TensorMeta):
|
||||
tensor_meta = output_tensor_meta
|
||||
elif isinstance(output_tensor_meta, (tuple, list)):
|
||||
if output_spec_count < len(output_tensor_meta):
|
||||
tensor_meta = output_tensor_meta[output_spec_count]
|
||||
output_spec_count += 1
|
||||
else:
|
||||
# This is an input position
|
||||
# Only get tensor_meta if we have a corresponding input_args_strategy entry
|
||||
if input_strategy_counter < len(input_args_strategy):
|
||||
tensor_meta = input_args_strategy[
|
||||
input_strategy_counter
|
||||
].tensor_meta
|
||||
input_strategy_counter += 1
|
||||
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
use_strided = (
|
||||
_input_use_strided
|
||||
if _input_use_strided is not None
|
||||
and any(isinstance(p, _StridedShard) for p in specs)
|
||||
else None
|
||||
)
|
||||
spec_list.append(
|
||||
DTensorSpec(
|
||||
mesh,
|
||||
specs,
|
||||
tensor_meta=tensor_meta,
|
||||
use_strided_shard_as_shard_order=use_strided,
|
||||
)
|
||||
)
|
||||
else:
|
||||
spec_list.append(None)
|
||||
|
||||
# Skip strategy combinations that would create mixed partial types
|
||||
# (except sum+avg which commute with each other).
|
||||
# We check (type, reduce_op) pairs rather than just reduce_op because
|
||||
# Partial subclasses like _MaskPartial have different reduction semantics
|
||||
# even when they share the same reduce_op string.
|
||||
has_mixed_partial = False
|
||||
for spec in spec_list:
|
||||
if spec is not None:
|
||||
partial_kinds = {
|
||||
(type(p), p.reduce_op)
|
||||
for p in spec.placements
|
||||
if isinstance(p, Partial)
|
||||
}
|
||||
if len(partial_kinds) > 1:
|
||||
reduce_ops = {ro for _, ro in partial_kinds}
|
||||
types = {t for t, _ in partial_kinds}
|
||||
if not (len(types) == 1 and reduce_ops == {"sum", "avg"}):
|
||||
has_mixed_partial = True
|
||||
break
|
||||
if has_mixed_partial:
|
||||
continue
|
||||
|
||||
input_specs: list[DTensorSpec] = [
|
||||
s for s in spec_list[input_index:] if isinstance(s, DTensorSpec)
|
||||
]
|
||||
|
||||
if len(input_specs) != len(input_args_strategy):
|
||||
raise AssertionError(
|
||||
f"input_specs({len(input_specs)}) != strategies({len(input_args_strategy)}: "
|
||||
f"{len(args_strategy)} args + {len(kwargs_strategy)} kwargs)"
|
||||
)
|
||||
|
||||
# Note [Multi-mesh args]
|
||||
#
|
||||
# Some ops accept args whose DTensor lives on a different DeviceMesh
|
||||
# than the op's primary compute mesh. We call these "multi-mesh
|
||||
# args". They arise in fused optimizer ops (e.g. _fused_adam_)
|
||||
# where *state_steps* is a per-rank scalar counter allocated on a
|
||||
# smaller sub-mesh (e.g. 1-D DP) while params and grads live on a
|
||||
# larger mesh (e.g. 2-D DP × TP).
|
||||
#
|
||||
# Why must these args be Replicate?
|
||||
# Sharding implies a specific partitioning of a tensor's data
|
||||
# across the ranks of a mesh. If a tensor doesn't even *exist*
|
||||
# on the compute mesh, there is no meaningful way to interpret a
|
||||
# Shard placement for it. Replicate, on the other hand, is
|
||||
# mesh-agnostic: every rank already holds the full data, so the
|
||||
# op can simply read the value regardless of which mesh owns it.
|
||||
#
|
||||
# What we do here:
|
||||
# We preserve the original mesh and Replicate placement for these
|
||||
# args so the propagator does not try to redistribute them onto
|
||||
# the compute mesh (which would fail or produce wrong results).
|
||||
#
|
||||
# This is distinct from the *element_mesh* handling in
|
||||
# single_dim_strategy.py, which deals with foreach ops where
|
||||
# different *elements* in a tensor list may live on different
|
||||
# sub-meshes (e.g. param group A on 2-D mesh, param group B on
|
||||
# 1-D mesh).
|
||||
# TODO: refactor fused_ops handling so that there are no longer
|
||||
# args on different meshes
|
||||
if different_mesh_args is not None:
|
||||
for idx in different_mesh_args:
|
||||
if idx < len(input_args_strategy):
|
||||
cross_mesh_input = input_args_strategy[idx]
|
||||
original_spec = cross_mesh_input.strategies[0].output_spec
|
||||
if original_spec.mesh != mesh:
|
||||
if not all(p == Replicate() for p in original_spec.placements):
|
||||
raise RuntimeError(
|
||||
f"Cross-mesh input at index {idx} must be Replicate, "
|
||||
f"but got {original_spec.placements}"
|
||||
)
|
||||
input_specs[idx] = DTensorSpec(
|
||||
mesh=original_spec.mesh,
|
||||
placements=original_spec.placements,
|
||||
tensor_meta=original_spec.tensor_meta,
|
||||
)
|
||||
self_spec = input_args_strategy[0].strategies[0].output_spec
|
||||
|
||||
redistribute_input = self_spec.placements != input_specs[0].placements
|
||||
mismatching_input_output = (
|
||||
spec_list[0] is not None and spec_list[0].placements != self_spec.placements
|
||||
)
|
||||
if inplace_op and (redistribute_input or mismatching_input_output):
|
||||
# For inplace ops, both the proposed input[0] and the output must
|
||||
# match self's runtime placement: input[0] because self can't be
|
||||
# redistributed, output because the result IS self.
|
||||
if blocking_inplace_input_placements is None:
|
||||
blocking_inplace_input_placements = self_spec.placements
|
||||
continue
|
||||
|
||||
# For out= variant ops, output placement must match the "out" kwarg's placement
|
||||
if (
|
||||
op_schema.is_out_variant_op()
|
||||
and "out" in op_schema.kwargs_schema
|
||||
and isinstance(op_schema.kwargs_schema["out"], OpStrategy)
|
||||
):
|
||||
out_kwarg_spec = op_schema.kwargs_schema["out"].strategies[0].output_spec
|
||||
# spec_list[0] is the output spec for this strategy combination
|
||||
if spec_list[0] is not None:
|
||||
if spec_list[0].placements != out_kwarg_spec.placements:
|
||||
continue
|
||||
|
||||
output_specs: tuple[DTensorSpec | None, ...] | DTensorSpec | None
|
||||
if input_index == 0:
|
||||
# No outputs (e.g., _linalg_check_errors)
|
||||
output_specs = None
|
||||
elif input_index > 1:
|
||||
output_specs = tuple(spec_list[:input_index])
|
||||
else:
|
||||
if spec_list[0] is not None:
|
||||
output_specs = spec_list[0]
|
||||
else:
|
||||
raise RuntimeError("output spec is None")
|
||||
|
||||
# check all inputs are shardable
|
||||
if not all(
|
||||
is_tensor_shardable(
|
||||
inp.shape, s, allow_unbacked_sharding=allow_unbacked_sharding
|
||||
)
|
||||
or (
|
||||
allow_uneven_sharding
|
||||
and inp.strategies[0].output_spec.placements == s.placements
|
||||
)
|
||||
for inp, s in zip(input_args_strategy, input_specs)
|
||||
):
|
||||
continue
|
||||
|
||||
# perform additional op-specific filtering
|
||||
# Skip callback for no-output ops (output_specs is None)
|
||||
if is_valid_strategy_cb is not None and output_specs is not None:
|
||||
if not is_valid_strategy_cb(input_specs, output_specs):
|
||||
continue
|
||||
|
||||
redistribute_cost = [
|
||||
generate_redistribute_costs(input_strategy, input_spec)
|
||||
for input_strategy, input_spec in zip(input_args_strategy, input_specs)
|
||||
]
|
||||
|
||||
strategy = OpSpec(
|
||||
output_specs=output_specs,
|
||||
input_specs=input_specs,
|
||||
redistribute_cost=redistribute_cost,
|
||||
)
|
||||
all_strategies.append(strategy)
|
||||
|
||||
# If all strategies were filtered out due to inplace placement mismatch,
|
||||
# raise a clear error message instead of returning an empty OpStrategy
|
||||
# (which would later cause a cryptic "min() arg is an empty sequence" error)
|
||||
if not all_strategies and blocking_inplace_input_placements is not None:
|
||||
raise RuntimeError(
|
||||
f"{op_schema.op}: in-place operations that require placement changes "
|
||||
f"are not supported. The input has placement {blocking_inplace_input_placements}, "
|
||||
f"but no valid strategy preserves this placement. "
|
||||
f"Please use the out-of-place version of this operation instead."
|
||||
)
|
||||
|
||||
return OpStrategy(all_strategies)
|
||||
|
||||
|
||||
def shift_shard_dims_after_insert(
|
||||
placements: Sequence[Placement], insert_dim: int = 0
|
||||
) -> Sequence[Placement]:
|
||||
normalized_placements: list[Placement] = []
|
||||
for placement in placements:
|
||||
if isinstance(placement, _StridedShard) and placement.dim >= insert_dim:
|
||||
normalized_placements.append(
|
||||
_StridedShard(placement.dim + 1, split_factor=placement.split_factor)
|
||||
)
|
||||
elif isinstance(placement, Shard) and placement.dim >= insert_dim:
|
||||
normalized_placements.append(Shard(placement.dim + 1))
|
||||
else:
|
||||
normalized_placements.append(placement)
|
||||
return normalized_placements
|
||||
|
||||
|
||||
def shift_shard_dims_after_remove(
|
||||
placements: Sequence[Placement], remove_dim: int = 0
|
||||
) -> Sequence[Placement]:
|
||||
normalized_placements: list[Placement] = []
|
||||
for placement in placements:
|
||||
if isinstance(placement, _StridedShard) and placement.dim > remove_dim:
|
||||
normalized_placements.append(
|
||||
_StridedShard(placement.dim - 1, split_factor=placement.split_factor)
|
||||
)
|
||||
elif isinstance(placement, Shard) and placement.dim > remove_dim:
|
||||
normalized_placements.append(Shard(placement.dim - 1))
|
||||
else:
|
||||
normalized_placements.append(placement)
|
||||
return normalized_placements
|
||||
@@ -0,0 +1,480 @@
|
||||
# mypy: allow-untyped-defs
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates
|
||||
import contextlib
|
||||
import warnings
|
||||
from collections.abc import Sequence
|
||||
from logging import getLogger
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
from torch.distributed._local_tensor import maybe_run_for_local_tensor
|
||||
from torch.distributed.device_mesh import _get_device_handle, DeviceMesh
|
||||
from torch.distributed.tensor._dtensor_spec import DTensorSpec
|
||||
from torch.distributed.tensor.placement_types import _StridedShard, Shard
|
||||
from torch.types import IntLikeType
|
||||
|
||||
|
||||
logger = getLogger(__name__)
|
||||
|
||||
__all__ = [
|
||||
"is_rng_supported_mesh",
|
||||
"manual_seed",
|
||||
"OffsetBasedRNGTracker",
|
||||
]
|
||||
|
||||
_rng_tracker: Optional["_RNGStateTracker"] = None
|
||||
|
||||
|
||||
def is_rng_supported_mesh(device_mesh: DeviceMesh) -> bool:
|
||||
"""Checks if the current device of ``device_mesh`` supports DTensor's random APIs.
|
||||
Currently DTensor Random APIs only supports cuda/cuda-like devices. We suggest
|
||||
users call this API to test the availability before using our random APIs.
|
||||
|
||||
Args:
|
||||
device_mesh (:class:`DeviceMesh`): The device mesh on which we check if the
|
||||
random ops APIs are supported.
|
||||
|
||||
Returns:
|
||||
A bool value. True if ``device_mesh`` supports DTensor Random APIs; False otherwise.
|
||||
|
||||
.. warning::
|
||||
Currently we only support correct RNG on cuda/cuda-like devices.
|
||||
"""
|
||||
device_handle = _get_device_handle(device_mesh.device_type)
|
||||
if device_handle and hasattr(device_handle, "set_rng_state"):
|
||||
return True
|
||||
else:
|
||||
# TODO: Logs way too much
|
||||
warnings.warn(
|
||||
f"DTensor random operators may not have complete support on {device_mesh.device_type} device mesh",
|
||||
stacklevel=2,
|
||||
)
|
||||
return False
|
||||
|
||||
|
||||
def manual_seed(seed: int, device_mesh: DeviceMesh) -> None:
|
||||
"""Sets the seed for generating random numbers for the calling rank.
|
||||
|
||||
Args:
|
||||
seed (int): The desired seed.
|
||||
device_mesh (:class:`DeviceMesh`): The device mesh to set the seed. It is
|
||||
required that the ``device_mesh`` include the calling rank. This is
|
||||
to ensure that the SPMD region maintains a synchronous RNG state, which
|
||||
means no ranks should be initialized with values other than ``seed``.
|
||||
|
||||
Returns:
|
||||
None
|
||||
|
||||
.. warning::
|
||||
:func:`manual_seed` does not check the ``seed`` value correctness. Users must
|
||||
ensure on their own that the value passed in is the desired ``seed`` for ranks
|
||||
within ``device_mesh``.
|
||||
If ``device_mesh`` is a sub-mesh and the calling rank is not a part of it,
|
||||
``manual_seed`` will throw an error.
|
||||
Current implementation only supports a GPU device mesh.
|
||||
"""
|
||||
if not is_rng_supported_mesh(device_mesh):
|
||||
warnings.warn(
|
||||
"DTensor manual_seed() may not have complete support "
|
||||
f"on {device_mesh.device_type} device mesh",
|
||||
stacklevel=2,
|
||||
)
|
||||
return
|
||||
|
||||
# TODO: deprecate this API, but also need to ensure we disable broadcast for PP case, and that's currently
|
||||
# bundled together with this API. See torchtitan/distributed/utils.py:set_determinism
|
||||
# warnings.warn(
|
||||
# "DTensor manual_seed() is deprecated, since DTensor no longer maintains a separate copy of generator state. "
|
||||
# "Use `torch.manual_seed` instead"
|
||||
# )
|
||||
# Note: we still need to ensure setting `run_state_sync=False` to support the pp case
|
||||
|
||||
# instantiate a RNG tracker if haven't. By default DTensor uses an
|
||||
# OffsetBasedRNGTracker to perform random operators.
|
||||
global _rng_tracker
|
||||
if not _rng_tracker:
|
||||
_rng_tracker = OffsetBasedRNGTracker(device_mesh, run_state_sync=False)
|
||||
|
||||
if device_mesh.get_coordinate() is None:
|
||||
raise RuntimeError(
|
||||
"manual_seed requires the current rank to be a part of the device mesh "
|
||||
"otherwise DTensor RNG state on the rank will not be initialized and "
|
||||
"the behavior of DTensor random ops is undefined."
|
||||
)
|
||||
|
||||
# DTensor no longer maintains a copy of rng state. manual seed on dtensor is the same thing
|
||||
# as manual seed on torch.
|
||||
#
|
||||
# torch.manual_seed will handle LocalTensor mode correctly by
|
||||
# iterating through all ranks if seed is a LocalIntNode.
|
||||
torch.manual_seed(seed)
|
||||
|
||||
|
||||
class _PhiloxState:
|
||||
"""
|
||||
Convenience accessor for interpreting the packed bits of (seed: uint64, offset: uint64) in the philox state,
|
||||
which for some reason is actually exposed as a size-16 uint8 tensor.
|
||||
|
||||
The state is always moved to .cpu since it is necessary for it to be on CPU before applying it back to a generator.
|
||||
"""
|
||||
|
||||
def __init__(self, state: torch.Tensor):
|
||||
self._state = state.to("cpu")
|
||||
|
||||
@property
|
||||
def state(self):
|
||||
return self._state
|
||||
|
||||
@property
|
||||
def offset(self) -> torch.Tensor:
|
||||
return self._state[8:].view(dtype=torch.int64)
|
||||
|
||||
@offset.setter
|
||||
def offset(self, offset: torch.Tensor) -> None:
|
||||
if offset.numel() != 1:
|
||||
raise AssertionError
|
||||
self._state[8:] = offset.view(torch.uint8)
|
||||
|
||||
@property
|
||||
def seed(self) -> torch.Tensor:
|
||||
return self._state[:8].view(dtype=torch.uint64)
|
||||
|
||||
@seed.setter
|
||||
def seed(self, seed: torch.Tensor) -> None:
|
||||
if seed.numel() != 1:
|
||||
raise AssertionError
|
||||
self._state[:8] = seed.view(torch.uint8)
|
||||
|
||||
|
||||
class _RNGStateTracker:
|
||||
"""
|
||||
_RNGStateTracker stores Random Number Generator (RNG) state (a ByteTensor object)
|
||||
in a dict, mapping from a corresponding tag to each state tensor. It also provides
|
||||
a set of convenient utility methods to help access/modify the state tensors. The most
|
||||
important interface is _distribute_region which will be used when DTensor executes
|
||||
a random op (an operator that calls RNG).
|
||||
"""
|
||||
|
||||
def __init__(self, device: torch.device):
|
||||
self._device = device
|
||||
self._device_handle = _get_device_handle(self._device.type)
|
||||
if not (self._device_handle and self._device_handle.is_available()):
|
||||
raise RuntimeError(
|
||||
f"{self.__class__.__name__} instantiation requires the presence of "
|
||||
f"{device.type} device but couldn't find."
|
||||
)
|
||||
self._use_distribute_region = True
|
||||
|
||||
@property
|
||||
def distribute_region_enabled(self) -> bool:
|
||||
return self._use_distribute_region
|
||||
|
||||
@distribute_region_enabled.setter
|
||||
def distribute_region_enabled(self, value) -> None:
|
||||
self._use_distribute_region = value
|
||||
|
||||
def _distribute_region(
|
||||
self, spec: DTensorSpec, generator: torch.Generator | None = None
|
||||
):
|
||||
pass
|
||||
|
||||
def _manual_seed(self, parallel_seed: int) -> None:
|
||||
pass
|
||||
|
||||
|
||||
class OffsetBasedRNGTracker(_RNGStateTracker):
|
||||
"""
|
||||
This subclass of ``_RNGStateTracker`` defines the default policy of how RNG states
|
||||
should be shared and synchronized among all ranks to respect the semantics of DTensor
|
||||
random operators.
|
||||
|
||||
note: _RNGStateTracker only supports cuda/cuda-like device.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
device_mesh: DeviceMesh,
|
||||
run_state_sync: bool = True,
|
||||
):
|
||||
super().__init__(_resolve_device(device_mesh=device_mesh))
|
||||
if self._device_handle is None:
|
||||
raise AssertionError
|
||||
# DTensor RNG tracker so far only supports CUDA/CUDA-like devices
|
||||
if self._device.type == "cpu":
|
||||
raise RuntimeError(
|
||||
f"{self.__class__.__name__} instantiation requires the presence of "
|
||||
f"CUDA/CUDA-like/XPU device. Got {self._device.type} instead."
|
||||
)
|
||||
|
||||
if run_state_sync:
|
||||
rng_state = self._get_device_state()
|
||||
# synchronize RNG state using rank 0's current one
|
||||
torch.distributed.broadcast(rng_state, 0)
|
||||
my_rng_state = self._get_device_state()
|
||||
if not all(my_rng_state == rng_state):
|
||||
logger.warning(
|
||||
"DTensor is synchronizing RNG states of every rank with the state from rank 0. "
|
||||
"This behavior is deprecated. "
|
||||
"Please call `torch.manual_seed()` on every rank that participates in SPMD DTensor Operations with "
|
||||
"the same seed. If using Pipeline Parallelism, each pipeling state would use a different seed, "
|
||||
"but all ranks belonging to one pipeline stage would use the same seed."
|
||||
)
|
||||
self._set_device_state(rng_state)
|
||||
|
||||
def _get_device_state(self) -> torch.Tensor:
|
||||
if self._device.type == "hpu":
|
||||
self._device_handle.set_rng_ctx("philox")
|
||||
rng_state = self._device_handle.get_rng_state().to(self._device)
|
||||
if self._device.type == "hpu":
|
||||
self._device_handle.unset_rng_ctx("philox")
|
||||
return rng_state
|
||||
|
||||
def _set_device_state(self, state: torch.Tensor):
|
||||
# It seems that the underlying generator wants a cpu tensor but the dtensor code expects `_get_device_state`
|
||||
# to convert to a 'device' tensor, probably because we may use it with our backend comms for sync/debug
|
||||
# for now, we just convert back to cpu here to make sure it always works.
|
||||
if self._device.type == "hpu":
|
||||
self._device_handle.set_rng_ctx("philox")
|
||||
self._device_handle.set_rng_state(state.to("cpu"))
|
||||
if self._device.type == "hpu":
|
||||
self._device_handle.unset_rng_ctx("philox")
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _distribute_region(
|
||||
self, spec: DTensorSpec, generator: torch.Generator | None = None
|
||||
):
|
||||
from torch.distributed._local_tensor import maybe_enable_local_tracker
|
||||
|
||||
if local_tracker_context := maybe_enable_local_tracker(
|
||||
self._device.type, self.distribute_region_enabled, spec, generator
|
||||
):
|
||||
with local_tracker_context:
|
||||
yield
|
||||
return
|
||||
|
||||
# regular (non-LocalTensor) mode
|
||||
if generator is not None:
|
||||
# This is a little hacky, but for any user-passed generator, we store its state under a unique key,
|
||||
# not because we need to keep a copy of it but because its the easiest way to make it work with the
|
||||
# existing set/get APIs. We also ensure we remove it from rng_states after each _distribute_region.
|
||||
state = _PhiloxState(generator.get_state())
|
||||
else:
|
||||
state = _PhiloxState(self._get_device_state())
|
||||
|
||||
if self.distribute_region_enabled:
|
||||
if self._device.type == "hpu":
|
||||
self._device_handle.set_rng_ctx("philox")
|
||||
old_offset = state.offset.clone()
|
||||
self._set_pre_op_offset(state, spec)
|
||||
with torch.random.fork_rng(
|
||||
devices=[self._device], device_type=self._device.type
|
||||
):
|
||||
if self._device_handle is None:
|
||||
raise AssertionError
|
||||
self._device_handle.set_rng_state(state.state)
|
||||
try:
|
||||
yield # execute the region code
|
||||
finally:
|
||||
# update offset to synchronize among ranks
|
||||
self._set_post_op_offset(state, spec, old_offset)
|
||||
if self._device.type == "hpu":
|
||||
self._device_handle.unset_rng_ctx("philox")
|
||||
else:
|
||||
yield
|
||||
|
||||
if generator is not None:
|
||||
# ensure we (a) propagate the state advancement back to the user's RNG so its visible and impacts any future
|
||||
# usage of that RNG (dtensor or non-dtensor), (b) drop it from our own cache so that if the user updates
|
||||
# the seed value in their rng and uses it with DTensor again, we always use the latest value
|
||||
generator.set_state(state.state)
|
||||
else:
|
||||
self._set_device_state(state.state)
|
||||
|
||||
def _set_pre_op_offset(self, state: _PhiloxState, spec: DTensorSpec) -> None:
|
||||
"""Set the starting RNG offset for current device's local shard before actual
|
||||
op execution. The pre_op_offset value should start from the current RNG offset
|
||||
and increment by the size of local shard until it reaches the size of the whole
|
||||
DTensor. For different ranks that hold the same DTensor shard, their pre_op_offset
|
||||
will be the same.
|
||||
|
||||
Args:
|
||||
state (:class:`Tensor`): The generator state to modify
|
||||
spec (:class:`DTensorSpec`): the spec of the DTensor object on which
|
||||
we prepare the offset for running random ops.
|
||||
|
||||
Returns:
|
||||
None
|
||||
|
||||
.. warning::
|
||||
Note that, current implementation does not consider DTensor's continguity.
|
||||
|
||||
Example:
|
||||
take a DTensor of shape [8, 16] as an example. Assume that the DTensor
|
||||
is placed on a device mesh with placements ([Shard(1), Replicate(), Shard(0)]),
|
||||
and the mesh is:
|
||||
[[[0, 1], [2, 3]], [[4, 5], [6, 7]]]
|
||||
``spec.mesh.get_coordinate()`` provides the coordinate of the current rank
|
||||
in the mesh. For example, the coordinate of rank 5 is (1, 0, 1).
|
||||
|
||||
Another concept to introduce besides rank coordinate is shard coordinate.
|
||||
Each rank holds a local shard of the DTensor. In the example, the DTensor
|
||||
is partitioned into 4 [4, 8] shards. The first shard has 2 replicas and
|
||||
rank 0 (coord (0, 0, 0)) and rank 2 (coord (0, 1, 0)) have 1 replica each.
|
||||
That being said, the local shard on rank 0 and rank 2 correspond to the same
|
||||
shard of the DTensor. To denote each DTensor shard, we use a shard coordinate
|
||||
(in the example, it will be a tuple (i, j) where shard (i, j) has the slice
|
||||
DTensor[4 * i : 4 * (i + 1), 8 * j : 8 * (j + 1)], 0 <= i < 2, 0 <= j < 2).
|
||||
|
||||
Once we have rank coordinate and shard coordinate, we can calculate on each rank
|
||||
what shard of the DTensor the rank holds, with the help of dim_map. The dim_map
|
||||
of the above DTensor is [2, 0] so the shard coordinate of a rank with rank coord
|
||||
(x, y, z) is simply (z, x) by taking(rank_coord[dim_map[0]],rank_coord[dim_map[1]]).
|
||||
Following this calculation,
|
||||
rank 0 and rank 2 holds the shard of coord (0, 0);
|
||||
rank 1 and rank 3 holds the shard of coord (0, 1);
|
||||
rank 4 and rank 6 holds the shard of coord (1, 0);
|
||||
rank 5 and rank 7 holds the shard of coord (1, 1);
|
||||
|
||||
The last value to calculate before obtaining the starting offset is the shard linear index.
|
||||
The starting offset for each rank will be its shard_linear_index * local_tensor_numel.
|
||||
"""
|
||||
start_offset_incr, _ = self._compute_rng_offsets(spec)
|
||||
state.offset = state.offset + start_offset_incr
|
||||
|
||||
def _set_post_op_offset(
|
||||
self, state: _PhiloxState, spec: DTensorSpec, old_offset: torch.Tensor
|
||||
) -> None:
|
||||
"""Sets the RNG to a synchronized state after running the local random op. Every
|
||||
rank should set its RNG offset to `old_offset + DTensor.numel()` where old_offset is
|
||||
the offset before calling `set_pre_op_offset` i.e. the offset before running DTensor
|
||||
random ops.
|
||||
|
||||
Args:
|
||||
state (:class:`Tensor`): The generator state to modify.
|
||||
spec (:class:`DTensorSpec`): the spec of the DTensor object on which
|
||||
we post-process the offset for running random ops.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
_, end_offset_incr = self._compute_rng_offsets(spec)
|
||||
state.offset = old_offset + end_offset_incr
|
||||
|
||||
def _compute_rng_offsets(self, spec: DTensorSpec) -> tuple[int, int]:
|
||||
"""Compute the RNG offset increments for a distributed random op.
|
||||
|
||||
These values are derived from mesh topology, placements, and tensor shape,
|
||||
and are static for a given compiled graph. They can be burned into the graph
|
||||
as integer constants rather than keeping the DTensorSpec around at runtime.
|
||||
|
||||
Returns:
|
||||
(start_offset_incr, end_offset_incr) — both aligned to multiples of 4.
|
||||
"""
|
||||
from torch.distributed.tensor._ops.utils import prod
|
||||
|
||||
mesh = spec.mesh
|
||||
mesh_coordinate = [mesh._sym_get_coordinate(i) for i in range(mesh.ndim)]
|
||||
|
||||
shard_idx_by_dim, total_num_shards_by_dim = _calc_shard_info(
|
||||
mesh_coordinate, spec
|
||||
)
|
||||
shard_linear_idx = self._calc_shard_linear_idx(
|
||||
shard_idx_by_dim, total_num_shards_by_dim
|
||||
)
|
||||
local_size = prod(_calc_first_shard_size(spec))
|
||||
# pytorch: offset must be multiple of 4
|
||||
# source: aten/src/ATen/cuda/CUDAGeneratorImpl.cpp
|
||||
start_offset_incr = (shard_linear_idx * local_size + 3) // 4 * 4
|
||||
end_offset_incr = (prod(spec.shape) + 3) // 4 * 4
|
||||
|
||||
return start_offset_incr, end_offset_incr
|
||||
|
||||
def _calc_shard_linear_idx(
|
||||
self, shard_coord: Sequence[IntLikeType], shard_size: Sequence[IntLikeType]
|
||||
) -> IntLikeType:
|
||||
return _calc_shard_linear_idx(shard_coord, shard_size)
|
||||
|
||||
|
||||
def _calc_first_shard_size(spec: DTensorSpec) -> list[int]:
|
||||
local_size_on_rank_0 = list(spec.shape)
|
||||
for idx, placement in enumerate(spec.placements):
|
||||
if isinstance(placement, Shard | _StridedShard):
|
||||
mesh_dim_size = spec.mesh.size(idx)
|
||||
shard_dim = placement.dim
|
||||
local_size_on_rank_0[shard_dim], _ = placement._local_shard_size_and_offset(
|
||||
spec.shape[shard_dim],
|
||||
mesh_dim_size,
|
||||
0,
|
||||
)
|
||||
return local_size_on_rank_0
|
||||
|
||||
|
||||
def _calc_shard_info(
|
||||
mesh_coordinate: Sequence[IntLikeType], spec: DTensorSpec
|
||||
) -> tuple[list[IntLikeType], list[IntLikeType]]:
|
||||
mesh = spec.mesh
|
||||
# note: dim_map does not allow double sharding which is the FSDP(fully_shard)+TP
|
||||
# case. Replace the custom logic with dim_map once we support it.
|
||||
dim_map: list[int | list[int]] = [-1] * spec.ndim
|
||||
for i, placement in enumerate(spec.placements):
|
||||
if isinstance(placement, Shard | _StridedShard):
|
||||
shard_dim = placement.dim
|
||||
if dim_map[shard_dim] == -1:
|
||||
dim_map[shard_dim] = [i]
|
||||
else:
|
||||
mesh_dim_list = dim_map[shard_dim]
|
||||
if not isinstance(mesh_dim_list, list):
|
||||
raise AssertionError
|
||||
mesh_dim_list.append(i)
|
||||
|
||||
# Compute shard coordinate:
|
||||
# The coordinate on each tensor dim is a tuple (idx, range)
|
||||
# If a DTensor is partitioned on its dim i into n shards, and the current rank
|
||||
# holds the j-th, then its shard coordinate will be (idx=j, range=n) on dim i
|
||||
mesh_size = mesh.shape
|
||||
shard_idx_by_dim = []
|
||||
total_num_shards_by_dim: list[
|
||||
IntLikeType
|
||||
] = [] # total number of shards on each tensor dim
|
||||
for mesh_dim in dim_map:
|
||||
shard_idx: IntLikeType = 0
|
||||
total_num_shards: IntLikeType = 1
|
||||
# the tensor dim is sharded on more than 1 mesh dim
|
||||
if isinstance(mesh_dim, list):
|
||||
rank_coord = [mesh_coordinate[d] for d in mesh_dim]
|
||||
num_shards = [mesh_size[d] for d in mesh_dim]
|
||||
# compute the shard idx and total number of shards
|
||||
for idx, size in zip(rank_coord, num_shards):
|
||||
shard_idx = shard_idx * size + idx
|
||||
total_num_shards *= size
|
||||
|
||||
shard_idx_by_dim.append(shard_idx)
|
||||
total_num_shards_by_dim.append(total_num_shards)
|
||||
return shard_idx_by_dim, total_num_shards_by_dim
|
||||
|
||||
|
||||
def _calc_shard_linear_idx(
|
||||
shard_coord: Sequence[IntLikeType], shard_size: Sequence[IntLikeType]
|
||||
) -> IntLikeType:
|
||||
# compute shard linear index
|
||||
shard_linear_idx: IntLikeType = 0
|
||||
shard_coord_stride: IntLikeType = 1
|
||||
for idx, size in zip(reversed(shard_coord), reversed(shard_size)):
|
||||
shard_linear_idx += idx * shard_coord_stride
|
||||
shard_coord_stride *= size
|
||||
|
||||
return shard_linear_idx
|
||||
|
||||
|
||||
def _resolve_device(device_mesh: DeviceMesh) -> torch.device:
|
||||
device_type = device_mesh.device_type
|
||||
device_handle = _get_device_handle(device_type)
|
||||
if device_handle is None:
|
||||
raise AssertionError
|
||||
device_idx = device_mesh.get_rank() % device_handle.device_count()
|
||||
|
||||
@maybe_run_for_local_tensor
|
||||
def get_device(device_idx):
|
||||
return torch.device(f"{device_type}:{device_idx:d}")
|
||||
|
||||
return get_device(device_idx)
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,362 @@
|
||||
# mypy: allow-untyped-defs
|
||||
# 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 typing import Any
|
||||
|
||||
import torch
|
||||
from torch.distributed.checkpoint.metadata import (
|
||||
ChunkStorageMetadata,
|
||||
MetadataIndex,
|
||||
TensorProperties,
|
||||
TensorStorageMetadata,
|
||||
)
|
||||
from torch.distributed.checkpoint.planner import (
|
||||
TensorWriteData,
|
||||
WriteItem,
|
||||
WriteItemType,
|
||||
)
|
||||
|
||||
|
||||
aten = torch.ops.aten
|
||||
|
||||
|
||||
class LocalShardsWrapper(torch.Tensor):
|
||||
"""
|
||||
A wrapper class to hold local shards of a DTensor.
|
||||
This class is used largely for checkpointing purposes and implicitly subtypes
|
||||
the _Checkpointable protocol.
|
||||
"""
|
||||
|
||||
__slots__ = ["_local_shards", "_storage_meta"]
|
||||
_local_shards: list[torch.Tensor]
|
||||
_storage_meta: TensorStorageMetadata
|
||||
|
||||
@staticmethod
|
||||
def __new__(
|
||||
cls, local_shards: list[torch.Tensor], local_offsets: list[tuple[int, ...]]
|
||||
) -> "LocalShardsWrapper":
|
||||
if not all(
|
||||
tensor.device == local_shards[0].device for tensor in local_shards[1:]
|
||||
):
|
||||
raise AssertionError
|
||||
|
||||
# if empty shard, we create a empty tensor
|
||||
if len(local_shards) == 0:
|
||||
r = torch.Tensor._make_wrapper_subclass(
|
||||
cls,
|
||||
torch.Size([0, 0]),
|
||||
)
|
||||
r._local_shards = []
|
||||
r._storage_meta = TensorStorageMetadata(
|
||||
properties=TensorProperties(),
|
||||
size=torch.Size([0, 0]),
|
||||
chunks=[
|
||||
ChunkStorageMetadata(
|
||||
offsets=torch.Size([0, 0]), sizes=torch.Size([0, 0])
|
||||
)
|
||||
],
|
||||
)
|
||||
return r
|
||||
|
||||
# we calculate the total tensor size by "concat" on second tensor dimension
|
||||
cat_tensor_shape = list(local_shards[0].size())
|
||||
if len(local_shards) > 1 and local_shards[0].ndim == 2: # column-wise sharding
|
||||
for shard in local_shards[1:]:
|
||||
cat_tensor_shape[1] += shard.size()[1]
|
||||
|
||||
# in cases of sharding optimizer rowwise, we calculate total tensor size by "concat" on first tensor dimension
|
||||
if len(local_shards) > 1 and local_shards[0].ndim == 1: # column-wise sharding
|
||||
for shard in local_shards[1:]:
|
||||
cat_tensor_shape[0] += shard.size()[0]
|
||||
|
||||
wrapper_properties = TensorProperties.create_from_tensor(local_shards[0])
|
||||
wrapper_shape = torch.Size(cat_tensor_shape)
|
||||
chunks_meta = [
|
||||
ChunkStorageMetadata(
|
||||
offsets=torch.Size(offset),
|
||||
sizes=shard.size(),
|
||||
)
|
||||
for shard, offset in zip(local_shards, local_offsets)
|
||||
]
|
||||
|
||||
r = torch.Tensor._make_wrapper_subclass(
|
||||
cls,
|
||||
torch.Size(cat_tensor_shape),
|
||||
)
|
||||
r._local_shards = local_shards
|
||||
r._storage_meta = TensorStorageMetadata(
|
||||
properties=wrapper_properties,
|
||||
size=wrapper_shape,
|
||||
chunks=chunks_meta,
|
||||
)
|
||||
|
||||
return r
|
||||
|
||||
# necessary for ops dispatching from this subclass to its local shards
|
||||
@classmethod
|
||||
def __torch_dispatch__(cls, func, types, args=(), kwargs=None): # type: ignore[override]
|
||||
kwargs = kwargs or {}
|
||||
|
||||
dispatcher = {
|
||||
torch.ops._c10d_functional.all_gather_into_tensor.default: cls.handle_all_gather_into_tensor,
|
||||
torch.ops._c10d_functional.wait_tensor.default: cls.handle_wait_tensor,
|
||||
aten._to_copy.default: cls.handle_to_copy,
|
||||
aten.view.default: cls.handle_view,
|
||||
aten.equal.default: cls.handle_equal,
|
||||
aten.detach.default: cls.handle_detach,
|
||||
aten.clone.default: cls.handle_clone,
|
||||
aten.new_empty.default: cls.handle_new_empty,
|
||||
}
|
||||
|
||||
if func in dispatcher:
|
||||
return dispatcher[func](args, kwargs)
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
f"{func} is not supported for LocalShardsWrapper!"
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def handle_all_gather_into_tensor(args, kwargs) -> torch.Tensor:
|
||||
dim = args[0].local_sizes()[0][1]
|
||||
cat_tensor = torch.cat(
|
||||
[t.view(-1) for t in args[0].local_shards()], dim=0
|
||||
).view(-1, dim)
|
||||
return torch.ops._c10d_functional.all_gather_into_tensor.default(
|
||||
cat_tensor, *args[1:], **kwargs
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def handle_wait_tensor(args, kwargs) -> torch.Tensor:
|
||||
return torch.ops._c10d_functional.wait_tensor(args[0])
|
||||
|
||||
@staticmethod
|
||||
def handle_to_copy(args, kwargs) -> torch.Tensor:
|
||||
res_shards_list = [
|
||||
aten._to_copy.default(shard, *args[1:], **kwargs)
|
||||
for shard in args[0].local_shards()
|
||||
]
|
||||
return LocalShardsWrapper(res_shards_list, args[0].local_offsets())
|
||||
|
||||
@staticmethod
|
||||
def handle_view(args, kwargs) -> "LocalShardsWrapper":
|
||||
view_shape = args[1]
|
||||
res_shards_list = []
|
||||
if len(args[0].local_shards()) > 1:
|
||||
if args[0].local_shards()[0].ndim == 2:
|
||||
if not (
|
||||
args[0].storage_metadata().size[0] == view_shape[0]
|
||||
and args[0].storage_metadata().size[1] == view_shape[1]
|
||||
):
|
||||
raise AssertionError
|
||||
# This accounts for a DTensor quirk, when multiple shards are present on a rank, DTensor on
|
||||
# init calls view_as() on the global tensor shape
|
||||
# will fail because the view shape is not applicable to individual shards.
|
||||
res_shards_list = [
|
||||
aten.view.default(shard, shard.shape, **kwargs)
|
||||
for shard in args[0].local_shards()
|
||||
]
|
||||
elif args[0].local_shards()[0].ndim == 1:
|
||||
if args[0].storage_metadata().size[0] != view_shape[0]:
|
||||
raise AssertionError
|
||||
# This case is for optimizer sharding as regardless of sharding type, optimizer state is row wise sharded
|
||||
res_shards_list = [
|
||||
aten.view.default(shard, shard.shape, **kwargs)
|
||||
for shard in args[0].local_shards()
|
||||
]
|
||||
else:
|
||||
raise NotImplementedError("No support for view on tensors ndim > 2")
|
||||
else:
|
||||
# view is called per shard
|
||||
res_shards_list = [
|
||||
aten.view.default(shard, args[1], **kwargs)
|
||||
for shard in args[0].local_shards()
|
||||
]
|
||||
return LocalShardsWrapper(res_shards_list, args[0].local_offsets())
|
||||
|
||||
@staticmethod
|
||||
def handle_equal(args, kwargs) -> bool:
|
||||
"""
|
||||
LocalShardsWrapper equal impl also checks for equality of storage metadata
|
||||
and the order of shards
|
||||
"""
|
||||
a, b = args[0], args[1]
|
||||
if len(a.local_shards()) != len(b.local_shards()):
|
||||
return False
|
||||
if not all(
|
||||
aten.equal.default(x, y) for x, y in zip(a.local_shards(), b.local_shards())
|
||||
):
|
||||
return False
|
||||
if a.storage_metadata() != b.storage_metadata():
|
||||
return False
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def handle_detach(args, kwargs) -> "LocalShardsWrapper":
|
||||
self_ls = args[0]
|
||||
deatched_local_shards = [
|
||||
aten.detach.default(shard) for shard in self_ls.local_shards()
|
||||
]
|
||||
self_ls._local_shards = deatched_local_shards
|
||||
self_ls._storage_meta.properties.requires_grad = False
|
||||
return self_ls
|
||||
|
||||
@staticmethod
|
||||
def handle_clone(args, kwargs) -> "LocalShardsWrapper":
|
||||
self_ls = args[0]
|
||||
desired_memory_format = kwargs.get("memory_format", None)
|
||||
if desired_memory_format and desired_memory_format != torch.preserve_format:
|
||||
raise NotImplementedError(
|
||||
f"{desired_memory_format} is not supported for LocalShardsWrapper!"
|
||||
)
|
||||
cloned_local_shards = [
|
||||
shard.clone(memory_format=desired_memory_format)
|
||||
for shard in self_ls._local_shards
|
||||
]
|
||||
return LocalShardsWrapper(cloned_local_shards, self_ls.local_offsets())
|
||||
|
||||
@staticmethod
|
||||
def handle_new_empty(args, kwargs) -> "LocalShardsWrapper":
|
||||
self_ls = args[0]
|
||||
return LocalShardsWrapper(
|
||||
[torch.empty_like(shard) for shard in self_ls._local_shards],
|
||||
self_ls.local_offsets(),
|
||||
)
|
||||
|
||||
@property
|
||||
def device(self) -> torch._C.device: # type: ignore[override]
|
||||
return (
|
||||
self._local_shards[0].device if self._local_shards else torch.device("meta")
|
||||
)
|
||||
|
||||
@property
|
||||
def is_meta(self) -> bool: # type: ignore[override]
|
||||
return self._local_shards[0].is_meta if self._local_shards else True
|
||||
|
||||
def is_pinned(self) -> bool: # type: ignore[override]
|
||||
return self._storage_meta.properties.pin_memory
|
||||
|
||||
def requires_grad_(self, requires_grad: bool = True) -> "LocalShardsWrapper":
|
||||
self._storage_meta.properties.requires_grad = requires_grad
|
||||
[shard.requires_grad_(requires_grad) for shard in self._local_shards]
|
||||
return self
|
||||
|
||||
def local_shards(self) -> list[torch.Tensor]:
|
||||
"""
|
||||
Returns a list of :class:`torch.Tensor' corresponding to the
|
||||
local shards for this rank. Returns an empty list if the current rank
|
||||
does not host any shards for this Tensor.
|
||||
"""
|
||||
return self._local_shards
|
||||
|
||||
def local_sizes(self) -> list[torch.Size]:
|
||||
"""
|
||||
Returns a list of :class:`torch.Size' corresponding to the
|
||||
local sizes for the shards on this rank. Returns an empty list if the current rank
|
||||
does not host any shards for this Tensor.
|
||||
"""
|
||||
return [chunk.sizes for chunk in self._storage_meta.chunks]
|
||||
|
||||
def local_offsets(self) -> list[torch.Size]:
|
||||
"""
|
||||
Returns a list of :class:`torch.Size' corresponding to the
|
||||
local offsets for the shards on this rank. Returns an empty list if the current rank
|
||||
does not host any shards for this Tensor.
|
||||
"""
|
||||
return [chunk.offsets for chunk in self._storage_meta.chunks]
|
||||
|
||||
@property
|
||||
def local_chunks(self) -> list[ChunkStorageMetadata]:
|
||||
"""
|
||||
Returns a :class:`list[ChunkStorageMetadata]` object corresponding to the
|
||||
metadata for each tensor shard
|
||||
"""
|
||||
return self._storage_meta.chunks
|
||||
|
||||
def storage_metadata(self) -> TensorStorageMetadata:
|
||||
"""
|
||||
Returns a :class:`TensorStorageMetadata` object corresponding to the
|
||||
metadata for the local tensor on current rank
|
||||
"""
|
||||
return self._storage_meta
|
||||
|
||||
def is_empty_shard(self) -> bool:
|
||||
"""
|
||||
Returns a :class:`bool` object indicating if the local tensor on current rank
|
||||
is an empty tensor
|
||||
"""
|
||||
return self._storage_meta.size[0] == 0 and self._storage_meta.size[1] == 0
|
||||
|
||||
def __create_write_items__(self, fqn: str, object: Any) -> list[WriteItem]:
|
||||
"""
|
||||
For compatibility with DCP, we support creation of WriteItems
|
||||
such that they can be saved properly.
|
||||
"""
|
||||
return [
|
||||
WriteItem(
|
||||
index=MetadataIndex(fqn, chunks.offsets),
|
||||
type=WriteItemType.SHARD,
|
||||
tensor_data=TensorWriteData(
|
||||
chunk=ChunkStorageMetadata(
|
||||
offsets=chunks.offsets,
|
||||
sizes=chunks.sizes,
|
||||
),
|
||||
properties=self._storage_meta.properties,
|
||||
size=object.size(),
|
||||
),
|
||||
)
|
||||
for tensor, chunks in zip(self.local_shards(), self.local_chunks)
|
||||
]
|
||||
|
||||
def __create_chunk_list__(self) -> list[ChunkStorageMetadata]:
|
||||
"""
|
||||
For compatibility with DCP, we support creation of chunk lists
|
||||
such that they can be saved properly.
|
||||
"""
|
||||
return self._storage_meta.chunks
|
||||
|
||||
def __get_tensor_shard__(self, index: MetadataIndex) -> torch.Tensor:
|
||||
"""
|
||||
For compatibility with DCP, we support finding shard based on index
|
||||
Return a 'torch.Tensor' shard based on 'MetadataIndex'.
|
||||
"""
|
||||
# Fast lookup path
|
||||
if index.index is not None:
|
||||
if (
|
||||
len(self._local_shards) > index.index
|
||||
and self._storage_meta.chunks[index.index].offsets == index.offset
|
||||
):
|
||||
return self._local_shards[index.index]
|
||||
|
||||
if index.offset is not None:
|
||||
for shard, chunk in zip(self._local_shards, self._storage_meta.chunks):
|
||||
if chunk.offsets == index.offset:
|
||||
return shard
|
||||
|
||||
# Empty shard case
|
||||
if len(self._local_shards) == 0 and self._storage_meta.chunks[
|
||||
0
|
||||
].sizes == torch.Size([0, 0]):
|
||||
return torch.empty(0)
|
||||
|
||||
raise ValueError(
|
||||
f"Could not find shard at '{index.offset}' for FQN: '{index.fqn}'"
|
||||
)
|
||||
|
||||
def _get_tensor_size_bytes(self) -> int:
|
||||
object_size = 0
|
||||
for shard in self.local_shards():
|
||||
object_size += shard.nelement() * shard.element_size()
|
||||
return object_size
|
||||
|
||||
def __hash__(self) -> int:
|
||||
return id(self)
|
||||
|
||||
def __repr__(self) -> str: # type: ignore[override]
|
||||
return f"LocalShardsWrapper:{self._local_shards} {self._storage_meta}"
|
||||
|
||||
def __str__(self) -> str:
|
||||
return f"LocalShardsWrapper:{self._local_shards} {self._storage_meta}"
|
||||
@@ -0,0 +1,314 @@
|
||||
# mypy: allow-untyped-defs
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates
|
||||
# implement matrix related ops for distributed tensor
|
||||
from typing import cast
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
import torch.distributed.tensor._api as dtensor
|
||||
|
||||
|
||||
aten = torch.ops.aten
|
||||
|
||||
|
||||
def _requires_data_exchange(padding, dim_map) -> bool:
|
||||
# Data exchange is not need if only sharded across batch dim
|
||||
if all(x == -1 for x in dim_map[1:]):
|
||||
return False
|
||||
# TODO: whether there requires data exchange is currently determined by padding
|
||||
return padding[-1] != 0
|
||||
|
||||
|
||||
def _is_supported(input_size, kernel_size, stride, padding, dilation):
|
||||
if dilation[-1] != 1:
|
||||
raise RuntimeError("Dilation must be 1 for tensor parallel convolution.")
|
||||
if padding[-1] != 0:
|
||||
if stride[-1] != 1:
|
||||
raise RuntimeError(
|
||||
"Stride must be 1 when there is padding for tensor parallel convolution."
|
||||
)
|
||||
if kernel_size[-1] // 2 > input_size[-1]:
|
||||
raise RuntimeError(
|
||||
"kernel_size[-1] // 2 should be less than or equal to input_size[-1] for tensor parallel convolution."
|
||||
)
|
||||
else:
|
||||
if not (input_size[-1] % stride[-1] == 0 and stride[-1] == kernel_size[-1]):
|
||||
raise RuntimeError(
|
||||
"It requires that input_size[-1] is divisible by stride[-1] and stride[-1] equals kernel_size[-1] "
|
||||
"when there is padding for tensor parallel convolution."
|
||||
)
|
||||
return True
|
||||
|
||||
|
||||
def _ring_send_recv_construct(in_tensor, d1, d2, left, right, rank, size):
|
||||
# dist comms and reconstruct local input tensor
|
||||
send_to_right = in_tensor[..., -d1:].contiguous()
|
||||
send_to_left = in_tensor[..., :d2].contiguous()
|
||||
recv_from_right = torch.zeros_like(send_to_left)
|
||||
recv_from_left = torch.zeros_like(send_to_right)
|
||||
|
||||
send_op_right = dist.P2POp(dist.isend, send_to_right, right)
|
||||
send_op_left = dist.P2POp(dist.isend, send_to_left, left)
|
||||
recv_op_right = dist.P2POp(dist.irecv, recv_from_right, right)
|
||||
recv_op_left = dist.P2POp(dist.irecv, recv_from_left, left)
|
||||
|
||||
reqs = dist.batch_isend_irecv(
|
||||
[send_op_right, send_op_left, recv_op_left, recv_op_right]
|
||||
)
|
||||
for req in reqs:
|
||||
req.wait()
|
||||
|
||||
if rank == 0:
|
||||
in_tensor = torch.cat([in_tensor, recv_from_right], dim=-1)
|
||||
elif rank == size - 1:
|
||||
in_tensor = torch.cat([recv_from_left, in_tensor], dim=-1)
|
||||
else:
|
||||
in_tensor = torch.cat([recv_from_left, in_tensor, recv_from_right], dim=-1)
|
||||
|
||||
return in_tensor
|
||||
|
||||
|
||||
def _ring_send_recv_aggregate(grad_in_tensor, d1, d2, left, right, rank, size):
|
||||
# dist comms and aggregate gradients for edge pixels
|
||||
send_to_right = grad_in_tensor[:, :, :, -d2:].contiguous()
|
||||
send_to_left = grad_in_tensor[:, :, :, :d1].contiguous()
|
||||
recv_from_right = torch.zeros_like(send_to_left)
|
||||
recv_from_left = torch.zeros_like(send_to_right)
|
||||
|
||||
send_op_right = dist.P2POp(dist.isend, send_to_right, right)
|
||||
send_op_left = dist.P2POp(dist.isend, send_to_left, left)
|
||||
recv_op_right = dist.P2POp(dist.irecv, recv_from_right, right)
|
||||
recv_op_left = dist.P2POp(dist.irecv, recv_from_left, left)
|
||||
|
||||
reqs = dist.batch_isend_irecv(
|
||||
[send_op_right, send_op_left, recv_op_left, recv_op_right]
|
||||
)
|
||||
for req in reqs:
|
||||
req.wait()
|
||||
|
||||
if rank == 0:
|
||||
grad_in_tensor = grad_in_tensor[:, :, :, :-d2]
|
||||
grad_in_tensor[:, :, :, -d1:] = torch.add(
|
||||
grad_in_tensor[:, :, :, -d1:], recv_from_right
|
||||
)
|
||||
elif rank == size - 1:
|
||||
grad_in_tensor = grad_in_tensor[:, :, :, d1:]
|
||||
grad_in_tensor[:, :, :, :d2] = torch.add(
|
||||
grad_in_tensor[:, :, :, :d2], recv_from_left
|
||||
)
|
||||
else:
|
||||
grad_in_tensor = grad_in_tensor[:, :, :, d1:-d2]
|
||||
grad_in_tensor[:, :, :, -d1:] = torch.add(
|
||||
grad_in_tensor[:, :, :, -d1:], recv_from_right
|
||||
)
|
||||
grad_in_tensor[:, :, :, :d2] = torch.add(
|
||||
grad_in_tensor[:, :, :, :d2], recv_from_left
|
||||
)
|
||||
|
||||
|
||||
def tp_convolution(
|
||||
op_call: torch._ops.OpOverload,
|
||||
local_tensor_args: tuple[object, ...],
|
||||
local_tensor_kwargs: dict[str, object],
|
||||
dim_map: list[int],
|
||||
) -> object:
|
||||
if op_call != aten.convolution.default:
|
||||
raise AssertionError
|
||||
if len(local_tensor_args) != 9:
|
||||
raise AssertionError
|
||||
|
||||
rank = dist.get_rank()
|
||||
size = dist.get_world_size()
|
||||
in_tensor = cast(torch.Tensor, local_tensor_args[0])
|
||||
weight = cast(torch.Tensor, local_tensor_args[1])
|
||||
stride, padding, dilation = local_tensor_args[3:6]
|
||||
|
||||
if not isinstance(padding, list):
|
||||
raise AssertionError
|
||||
|
||||
if not _requires_data_exchange(padding, dim_map):
|
||||
local_results = op_call(*local_tensor_args, **local_tensor_kwargs)
|
||||
return local_results
|
||||
else:
|
||||
if not _is_supported(in_tensor.shape, weight.shape, stride, padding, dilation):
|
||||
raise AssertionError(
|
||||
"tp_convolution data exchange requires supported stride/padding/dilation"
|
||||
)
|
||||
# step 0 compute the overlap pixels of the input tensor
|
||||
d = weight.shape[-1] - 1
|
||||
d1 = d // 2
|
||||
d2 = d - d1
|
||||
if d1 + d2 != d:
|
||||
raise AssertionError
|
||||
right = (rank + 1) % size
|
||||
left = (rank - 1 + size) % size
|
||||
|
||||
# step1 reconstruct local input tensor
|
||||
in_tensor = _ring_send_recv_construct(
|
||||
in_tensor, d1, d2, left, right, rank, size
|
||||
)
|
||||
|
||||
# step2 feed local input tensor to op_call
|
||||
local_tensor_args_list = list(local_tensor_args)
|
||||
local_tensor_args_list[0] = in_tensor
|
||||
local_tensor_args = cast(tuple[object, ...], local_tensor_args_list)
|
||||
local_results = op_call(*local_tensor_args, **local_tensor_kwargs)
|
||||
|
||||
# step3 remove extra outputs from the results
|
||||
padding_w = padding[-1]
|
||||
w = local_results.size(-1)
|
||||
if rank == 0:
|
||||
local_results = local_results[..., : w - padding_w]
|
||||
elif rank == size - 1:
|
||||
local_results = local_results[..., padding_w:]
|
||||
else:
|
||||
local_results = local_results[..., padding_w : w - padding_w]
|
||||
|
||||
return local_results
|
||||
|
||||
|
||||
def tp_convolution_backward(
|
||||
op_call: torch._ops.OpOverload,
|
||||
local_tensor_args: tuple[object, ...],
|
||||
local_tensor_kwargs: dict[str, object],
|
||||
dim_map: list[int],
|
||||
) -> object:
|
||||
if op_call != aten.convolution_backward.default:
|
||||
raise AssertionError
|
||||
if len(local_tensor_args) != 11:
|
||||
raise AssertionError
|
||||
|
||||
rank = dist.get_rank()
|
||||
size = dist.get_world_size()
|
||||
grad_out_tensor = cast(torch.Tensor, local_tensor_args[0])
|
||||
in_tensor = cast(torch.Tensor, local_tensor_args[1])
|
||||
weight = cast(torch.Tensor, local_tensor_args[2])
|
||||
stride, padding, dilation = local_tensor_args[4:7]
|
||||
|
||||
if not isinstance(padding, list):
|
||||
raise AssertionError
|
||||
|
||||
if not _requires_data_exchange(padding, dim_map):
|
||||
local_results = op_call(*local_tensor_args, **local_tensor_kwargs)
|
||||
return local_results
|
||||
else:
|
||||
if not _is_supported(in_tensor.shape, weight.shape, stride, padding, dilation):
|
||||
raise AssertionError(
|
||||
"tp_convolution_backward data exchange requires supported stride/padding/dilation"
|
||||
)
|
||||
# step 0 compute the overlap pixels of the input tensor
|
||||
d = weight.shape[3] - 1
|
||||
d1 = d // 2
|
||||
d2 = d - d1
|
||||
if d1 + d2 != d:
|
||||
raise AssertionError
|
||||
right = (rank + 1) % size
|
||||
left = (rank - 1 + size) % size
|
||||
|
||||
# step1 reconstruct local input tensor
|
||||
in_tensor = _ring_send_recv_construct(
|
||||
in_tensor, d1, d2, left, right, rank, size
|
||||
)
|
||||
|
||||
# step2 reconstruct local gradient output tensor
|
||||
padding_w = padding[1]
|
||||
if rank == 0:
|
||||
grad_out_tensor = torch.nn.functional.pad(
|
||||
grad_out_tensor, (0, padding_w), "constant", 0
|
||||
)
|
||||
elif rank == size - 1:
|
||||
grad_out_tensor = torch.nn.functional.pad(
|
||||
grad_out_tensor, (padding_w, 0), "constant", 0
|
||||
)
|
||||
else:
|
||||
grad_out_tensor = torch.nn.functional.pad(
|
||||
grad_out_tensor, (padding_w, padding_w), "constant", 0
|
||||
)
|
||||
|
||||
# step3 feed local input tensor to op_call
|
||||
local_tensor_args_list = list(local_tensor_args)
|
||||
local_tensor_args_list[0] = grad_out_tensor
|
||||
local_tensor_args_list[1] = in_tensor
|
||||
local_tensor_args = cast(tuple[object, ...], local_tensor_args_list)
|
||||
local_results = op_call(*local_tensor_args, **local_tensor_kwargs)
|
||||
|
||||
# step4 aggregate gradients for edge pixels
|
||||
grad_in_tensor = local_results[0]
|
||||
if grad_in_tensor is not None:
|
||||
grad_in_tensor = _ring_send_recv_aggregate(
|
||||
grad_in_tensor, d1, d2, left, right, rank, size
|
||||
)
|
||||
local_results = list(local_results)
|
||||
local_results[0] = grad_in_tensor
|
||||
|
||||
local_results = cast(tuple[object, ...], local_results)
|
||||
|
||||
return local_results
|
||||
|
||||
|
||||
def convolution_handler(
|
||||
op_call: torch._ops.OpOverload,
|
||||
args: tuple[object, ...],
|
||||
kwargs: dict[str, object],
|
||||
) -> object:
|
||||
# extract local tensor and sharding infos to a OpInfo
|
||||
op_info = dtensor.DTensor._op_dispatcher.unwrap_to_op_info(op_call, args, kwargs)
|
||||
|
||||
# sharding propagation
|
||||
dtensor.DTensor._op_dispatcher.sharding_propagator.propagate(op_info)
|
||||
output_sharding = op_info.output_sharding
|
||||
if output_sharding is None:
|
||||
raise AssertionError("output sharding should not be None")
|
||||
output_spec = output_sharding.output_spec
|
||||
if not isinstance(output_spec, dtensor.DTensorSpec):
|
||||
raise AssertionError
|
||||
|
||||
# local propagation
|
||||
local_results = tp_convolution(
|
||||
op_call,
|
||||
tuple(op_info.local_args),
|
||||
op_info.local_kwargs,
|
||||
output_spec.dim_map,
|
||||
)
|
||||
|
||||
return dtensor.DTensor._op_dispatcher.wrap(local_results, output_spec)
|
||||
|
||||
|
||||
def convolution_backward_handler(
|
||||
op_call: torch._ops.OpOverload,
|
||||
args: tuple[object, ...],
|
||||
kwargs: dict[str, object],
|
||||
) -> object:
|
||||
# Redistribute grad_output tensor to the same placement as input tensor
|
||||
# pyrefly: ignore [bad-assignment]
|
||||
args = list(args)
|
||||
if not (
|
||||
isinstance(args[0], dtensor.DTensor) and isinstance(args[1], dtensor.DTensor)
|
||||
):
|
||||
raise AssertionError
|
||||
# pyrefly: ignore [unsupported-operation]
|
||||
args[0] = args[0].redistribute(args[1].device_mesh, args[1].placements)
|
||||
args = tuple(args)
|
||||
|
||||
# extract local tensor and sharding infos to a OpInfo
|
||||
op_info = dtensor.DTensor._op_dispatcher.unwrap_to_op_info(op_call, args, kwargs)
|
||||
|
||||
# sharding propagation
|
||||
dtensor.DTensor._op_dispatcher.sharding_propagator.propagate(op_info)
|
||||
output_sharding = op_info.output_sharding
|
||||
if output_sharding is None:
|
||||
raise AssertionError("output sharding should not be None")
|
||||
if not isinstance(op_info.flat_args_schema[0], dtensor.DTensorSpec):
|
||||
raise AssertionError
|
||||
|
||||
# local propagation
|
||||
local_results = tp_convolution_backward(
|
||||
op_call,
|
||||
tuple(op_info.local_args),
|
||||
op_info.local_kwargs,
|
||||
op_info.flat_args_schema[0].dim_map,
|
||||
)
|
||||
|
||||
return dtensor.DTensor._op_dispatcher.wrap(
|
||||
local_results, output_sharding.output_spec
|
||||
)
|
||||
@@ -0,0 +1,525 @@
|
||||
import logging
|
||||
import threading
|
||||
from collections.abc import Callable, Sequence
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import torch.distributed._functional_collectives as funcol
|
||||
import torch.distributed.tensor._api as dtensor
|
||||
from torch._logging import LazyString
|
||||
from torch._prims_common import ShapeType
|
||||
from torch.distributed import RankType
|
||||
from torch.distributed._local_tensor import maybe_run_for_local_tensor
|
||||
from torch.distributed.device_mesh import DeviceMesh
|
||||
from torch.distributed.tensor._collective_utils import redistribute_cost
|
||||
from torch.distributed.tensor._dtensor_spec import DTensorSpec
|
||||
from torch.distributed.tensor._op_schema import OpSchema
|
||||
from torch.distributed.tensor.placement_types import (
|
||||
_is_shard_like,
|
||||
_StridedShard,
|
||||
Partial,
|
||||
Placement,
|
||||
Replicate,
|
||||
Shard,
|
||||
)
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _format_implicit_redistribution_msg(schema: OpSchema) -> str:
|
||||
return f"Implicit redistribution occurred for {schema} while ExplicitRedistributionContext was active"
|
||||
|
||||
|
||||
class ExplicitRedistributionContext:
|
||||
"""
|
||||
Within this context manager, DTensor will refuse to perform implicit redistribution,
|
||||
instead raising an error. Manual calls to ``redistribute()`` are required wherever a redistribution
|
||||
must occur to avoid erroring. This can be used to ensure that the user is aware of all redistribution.
|
||||
|
||||
Note: it is easier to use this mode on just the forward pass of a typical DTensor program, as the backwards pass
|
||||
may contain implicit redistribution calls that are not visible to the user and difficult to replace with manual
|
||||
calls. Redistribution during backward can be made explicit by writing `autograd.Function`s that are no-op
|
||||
during forward and perform a manual redistribution during backwards.
|
||||
|
||||
enable (bool) if False, disables the context manager. Can be used nested inside an enabled region.
|
||||
|
||||
strict (bool) if True, triggers on any redistribution. If False, only triggers on redistributions that perform
|
||||
communication.
|
||||
|
||||
mode (str) Determines what happens when ExplicitRedistributionContext triggers:
|
||||
"raise": raises an exception, "warn" issues a warning
|
||||
"""
|
||||
|
||||
_local = threading.local()
|
||||
|
||||
def __init__(self, enable: bool = True, strict: bool = False, mode="raise"):
|
||||
self._enable = enable
|
||||
self._strict = strict
|
||||
if mode not in ("raise", "warn"):
|
||||
raise RuntimeError(f"Invalid mode {mode}")
|
||||
self._raise_on_redistribution = mode == "raise"
|
||||
|
||||
@classmethod
|
||||
def observe_redistribution(
|
||||
cls,
|
||||
src_spec: DTensorSpec,
|
||||
dst_spec: DTensorSpec,
|
||||
redistribution_msg: LazyString,
|
||||
):
|
||||
if instance := getattr(cls._local, "_active", None):
|
||||
allowed = True
|
||||
if instance._enable:
|
||||
if instance._strict:
|
||||
allowed = False
|
||||
else:
|
||||
allowed = redistribute_cost(src_spec, dst_spec) <= 0
|
||||
if not allowed:
|
||||
if instance._raise_on_redistribution:
|
||||
raise RuntimeError(redistribution_msg)
|
||||
else:
|
||||
logger.warning(redistribution_msg)
|
||||
|
||||
def __enter__(self):
|
||||
self._prev = getattr(ExplicitRedistributionContext._local, "_active", None)
|
||||
ExplicitRedistributionContext._local._active = self
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
ExplicitRedistributionContext._local._active = self._prev
|
||||
|
||||
|
||||
def compute_local_shape_and_global_offset(
|
||||
global_shape: ShapeType,
|
||||
mesh: DeviceMesh,
|
||||
placements: Sequence[Placement],
|
||||
skip_offset: bool = False,
|
||||
) -> tuple[tuple[int, ...], tuple[int, ...]]:
|
||||
"""
|
||||
Compute the local tensor shape and the global offsets into the original tensor
|
||||
of a DTensor on its current global rank. This is useful for checkpointing purpose.
|
||||
|
||||
Example:
|
||||
global_tensor = [[0, 1, 2, 3, 4], sharded on mesh (DP=2, TP=2) with (Shard(1), Shard(1))
|
||||
[10, 11, 12, 13, 14]]
|
||||
|
||||
This table shows the return value of local_shape and global_offset for each rank.
|
||||
(`local_tensor` is for illustration only).
|
||||
|
||||
Note how the first coordinate of global_offset is always 0, corresponding to tensor dim 0 being replicated.
|
||||
|
||||
Rank local_tensor local_shape global_offset
|
||||
-------------------------------------------------------------
|
||||
0 [[0, 1], (2, 2) (0, 0)
|
||||
[10, 11]]
|
||||
|
||||
1 [[2], (2, 1) (0, 2)
|
||||
[12]]
|
||||
|
||||
2 [[3], (2, 1) (0, 3)
|
||||
[13]]
|
||||
|
||||
3 [[4], (2, 1) (0, 4)
|
||||
[14]]
|
||||
|
||||
Args:
|
||||
global_shape (ShapeType): The global shape of the DTensor.
|
||||
mesh (:class:`DeviceMesh`): The device mesh this DTensor is distributed on.
|
||||
placements (Sequence[:class:`Placement`]]): The placements of the DTensor.
|
||||
skip_offset (bool): If True, skip computing the global offsets and return an empty
|
||||
tuple for global_offset. This can improve performance when only the local shape
|
||||
is needed. Defaults to False.
|
||||
|
||||
Return:
|
||||
local_shape: the shape of the DTensor's _local_tensor on the current rank.
|
||||
global_offset: a tuple of offsets for each dimension of the global tensor shape,
|
||||
identifying how this shard fits into the global tensor in each dimension. If
|
||||
skip_offset is True, this will be an empty tuple.
|
||||
|
||||
"""
|
||||
empty_offset = ()
|
||||
if not mesh._is_current_rank_part_of_mesh():
|
||||
# if rank not in the mesh, return empty offset
|
||||
return ((0,), empty_offset)
|
||||
|
||||
return _compute_local_shape_and_global_offset(
|
||||
global_shape, mesh.shape, mesh._sym_get_coordinate, placements, skip_offset
|
||||
)
|
||||
|
||||
|
||||
@maybe_run_for_local_tensor
|
||||
def _get_shard_size_and_offsets(
|
||||
curr_local_size: int,
|
||||
mesh_dim_size: int,
|
||||
rank: RankType,
|
||||
placement: Shard | _StridedShard,
|
||||
previous_offsets,
|
||||
zero_global_offset: int,
|
||||
skip_offset: bool,
|
||||
) -> tuple[int, torch.Tensor | None]:
|
||||
kwargs: dict[str, Any] = {
|
||||
"curr_local_size": curr_local_size,
|
||||
"num_chunks": mesh_dim_size,
|
||||
"rank": rank,
|
||||
}
|
||||
if isinstance(placement, _StridedShard):
|
||||
kwargs["return_first_offset"] = False
|
||||
shard_size, shard_offsets = placement._local_shard_size_and_offset(**kwargs)
|
||||
if skip_offset:
|
||||
return shard_size, None
|
||||
if shard_size == 0:
|
||||
return shard_size, torch.arange(zero_global_offset, zero_global_offset + 1)
|
||||
if isinstance(placement, Shard) and not isinstance(placement, _StridedShard):
|
||||
if not isinstance(shard_offsets, int):
|
||||
raise AssertionError
|
||||
index = torch.arange(shard_offsets, shard_offsets + shard_size)
|
||||
else:
|
||||
if not isinstance(shard_offsets, list):
|
||||
raise AssertionError
|
||||
index = torch.tensor(shard_offsets)
|
||||
if previous_offsets is None:
|
||||
return shard_size, index
|
||||
else:
|
||||
return shard_size, previous_offsets[index]
|
||||
|
||||
|
||||
@maybe_run_for_local_tensor
|
||||
def _get_first_offset(offsets: torch.Tensor) -> int:
|
||||
return int(offsets[0])
|
||||
|
||||
|
||||
# accept 'plain data types' to enable simpler unit testing without creating device mesh
|
||||
def _compute_local_shape_and_global_offset(
|
||||
global_shape: ShapeType,
|
||||
mesh_shape: ShapeType,
|
||||
my_coordinate: list[int] | Callable[[int], RankType] | None,
|
||||
placements: Sequence[Placement],
|
||||
skip_offset: bool = False,
|
||||
) -> tuple[tuple[int, ...], tuple[int, ...]]:
|
||||
"""
|
||||
Suppose you have a full tensor with size global_shape, and you have sharded
|
||||
it according to placements for mesh_shape. This function returns, for a
|
||||
specific coordinate my_coordinate in the device mesh:
|
||||
|
||||
- The size of your local shard WITHOUT padding (i.e., if you have
|
||||
an uneven split, your size might be smaller than the other entries
|
||||
in your dim), and
|
||||
|
||||
- Where the data for your shard begins, in the full tensor.
|
||||
|
||||
This function is fairly simple if your tensor is evenly sharded; the complication
|
||||
is around uneven splits. There is also some complication for handling StridedShard,
|
||||
which changes the order you should apply sharding.
|
||||
|
||||
Args:
|
||||
global_shape (ShapeType): The global shape of the tensor.
|
||||
mesh_shape (ShapeType): The shape of the device mesh.
|
||||
my_coordinate (Optional[list[int]]): The coordinate of the current rank in the device mesh.
|
||||
placements (Sequence[Placement]): The placements of the DTensor.
|
||||
skip_offset (bool): If True, skip computing the global offsets and return an empty
|
||||
tuple for global_offset. This can improve performance when only the local shape
|
||||
is needed. Defaults to False.
|
||||
|
||||
Returns:
|
||||
tuple: A tuple containing:
|
||||
- local_shape (tuple[int, ...]): The shape of the local shard on the current rank.
|
||||
- global_offset (tuple[int, ...]): The offsets for each dimension identifying where
|
||||
this shard begins in the global tensor. If skip_offset is True, this will be an
|
||||
empty tuple.
|
||||
"""
|
||||
|
||||
if isinstance(my_coordinate, (list, tuple)):
|
||||
_coord: list | tuple = my_coordinate
|
||||
|
||||
def coordinate_lookup(dim: int) -> RankType:
|
||||
return _coord[dim]
|
||||
else:
|
||||
if my_coordinate is None:
|
||||
raise AssertionError
|
||||
coordinate_lookup = my_coordinate
|
||||
|
||||
local_shape = list(global_shape)
|
||||
# Perform shard from left to right. For example,
|
||||
# global tensor: [0, 1, 2, 3, 4, 5, 6, 7]
|
||||
# placements: S(0), SS(0, split_factor=2)
|
||||
# mesh_shape: (2, 2)
|
||||
# After S(0), shard_dim_to_global_offsets are
|
||||
# {0: [0, 1, 2, 3]} on my_coordinate [0, 0] [0, 1]
|
||||
# {0: [4, 5, 6, 7]} on my_coordinate [1, 0] [1, 1]
|
||||
# After SS(0, split_factor=2), shard_dim_to_global_offsets are
|
||||
# {0: [0, 2]} on my_coordinate [0, 0]
|
||||
# {0: [1, 3]} on my_coordinate [0, 1]
|
||||
# {0: [4, 6]} on my_coordinate [1, 0]
|
||||
# {0: [5, 7]} on my_coordinate [1, 1]
|
||||
shard_dim_to_global_offsets = {}
|
||||
for mesh_dim, placement in enumerate(placements):
|
||||
if not isinstance(placement, (Shard, _StridedShard)):
|
||||
continue
|
||||
shard_dim = placement.dim
|
||||
zero_global_offset = global_shape[shard_dim]
|
||||
if shard_dim >= len(local_shape):
|
||||
raise AssertionError(
|
||||
f"Sharding dim {shard_dim} greater than tensor ndim {len(local_shape)}"
|
||||
)
|
||||
previous_offsets = shard_dim_to_global_offsets.get(shard_dim)
|
||||
shard_size, shard_offsets = _get_shard_size_and_offsets(
|
||||
local_shape[shard_dim],
|
||||
mesh_shape[mesh_dim],
|
||||
coordinate_lookup(mesh_dim),
|
||||
placement,
|
||||
previous_offsets,
|
||||
zero_global_offset,
|
||||
skip_offset,
|
||||
)
|
||||
local_shape[shard_dim] = shard_size
|
||||
shard_dim_to_global_offsets[shard_dim] = shard_offsets
|
||||
if skip_offset:
|
||||
return tuple(local_shape), ()
|
||||
global_offset = [0] * len(global_shape)
|
||||
for shard_dim, global_offsets in shard_dim_to_global_offsets.items():
|
||||
global_offset[shard_dim] = _get_first_offset(global_offsets)
|
||||
return tuple(local_shape), tuple(global_offset)
|
||||
|
||||
|
||||
compute_global_tensor_info = torch._C._DTensor_compute_global_tensor_info
|
||||
|
||||
|
||||
def compute_local_tensor_info(
|
||||
global_tensor: torch.Tensor,
|
||||
mesh: DeviceMesh,
|
||||
placements: Sequence[Placement],
|
||||
) -> tuple[list[int], list[int]]:
|
||||
"""
|
||||
Compute the local size and stride of a DTensor from the given global tensor info.
|
||||
|
||||
For example, if we have a global tensor with size (4, 8, 4) and stride (32, 1, 8).
|
||||
If the DTensor placements are [Shard(2)] and world_size is 2;
|
||||
then the local size is (4, 8, 2) and stride is (16, 1, 8).
|
||||
|
||||
Args:
|
||||
tensor (:class:`torch.Tensor`):
|
||||
Global tensor which DTensor will distribute
|
||||
mesh (:class:`DeviceMesh`):
|
||||
Object which describes the mesh topology
|
||||
of devices for the DTensor.
|
||||
placements (Sequence[:class:`Placement`]):
|
||||
The attribute of the DTensor that describes its layout
|
||||
on the mesh topology.
|
||||
|
||||
Returns:
|
||||
local_shape: A List of int which specifies the size of the local tensor.
|
||||
local_stride: A List of int which specifies the stride of the local tensor.
|
||||
"""
|
||||
local_shape = list(global_tensor.size())
|
||||
local_stride = list(global_tensor.stride())
|
||||
|
||||
for idx, placement in enumerate(placements):
|
||||
mesh_dim_size = mesh.size(idx)
|
||||
if _is_shard_like(placement):
|
||||
if placement.dim < 0:
|
||||
raise AssertionError(
|
||||
"Shard placements should have negative dims normalized in "
|
||||
f"the user-facing APIs: {placement}"
|
||||
)
|
||||
shard_dim = placement.dim
|
||||
if shard_dim >= len(local_shape):
|
||||
raise AssertionError(
|
||||
f"Sharding dim {shard_dim} greater than tensor ndim {len(local_shape)} "
|
||||
f"for placement number {idx}."
|
||||
)
|
||||
|
||||
global_dim_size = local_shape[shard_dim]
|
||||
if global_dim_size % mesh_dim_size != 0:
|
||||
raise AssertionError(
|
||||
f"Global dim {global_dim_size} not divisible by mesh size {mesh_dim_size}"
|
||||
)
|
||||
local_shape[shard_dim] = global_dim_size // mesh_dim_size
|
||||
|
||||
# shrink strides that were scaled up globally
|
||||
for i in range(len(local_stride)):
|
||||
if (
|
||||
i != shard_dim
|
||||
and local_stride[i] >= local_stride[shard_dim] * mesh_dim_size
|
||||
):
|
||||
local_stride[i] = local_stride[i] // mesh_dim_size
|
||||
|
||||
elif not isinstance(placement, (Replicate, Partial)):
|
||||
raise RuntimeError(f"placement type {type(placement)} not supported!")
|
||||
|
||||
return local_shape, local_stride
|
||||
|
||||
|
||||
def compute_global_tensor_shape(
|
||||
shape: torch.Size, mesh: DeviceMesh, placements: Sequence[Placement]
|
||||
) -> torch.Size:
|
||||
"""
|
||||
Compute the global size of a DTensor from the given local tensor shape,
|
||||
the mesh and placements. Different from `compute_global_tensor_info`,
|
||||
which assumes sharding is even, this util allgathers local shards' shapes
|
||||
from all ranks and thus can support uneven sharding.
|
||||
NOTE: Currently this function only supports 1D mesh.
|
||||
|
||||
Args:
|
||||
shape (:class:`torch.Size`):
|
||||
Shape of the local tensor
|
||||
mesh (:class:`DeviceMesh`):
|
||||
Object which describes the mesh topology
|
||||
of devices for the DTensor.
|
||||
placements (Sequence[:class:`Placement`]]):
|
||||
The attribute of the DTensor that describes its layout
|
||||
on the mesh topology.
|
||||
|
||||
Return:
|
||||
tensor_shape: Shape of the global DTensor.
|
||||
"""
|
||||
if len(placements) != 1:
|
||||
raise NotImplementedError(
|
||||
"compute_global_tensor_shape only supports 1 placement for now."
|
||||
)
|
||||
|
||||
if len(placements) != mesh.ndim:
|
||||
raise RuntimeError(
|
||||
"Expected one placement per mesh dim, "
|
||||
f"but found {len(placements)} placements and {mesh.ndim} mesh dims."
|
||||
)
|
||||
|
||||
if isinstance(placements[0], Replicate):
|
||||
return shape
|
||||
# NOTE: isinstance(_, Shard) does not match _StridedShard; see _is_shard_like().
|
||||
elif isinstance(placements[0], Shard):
|
||||
|
||||
@maybe_run_for_local_tensor
|
||||
def _create_local_shape_tensor(shape):
|
||||
return torch.tensor(list(shape), device=mesh.device_type)
|
||||
|
||||
local_shape = _create_local_shape_tensor(shape)
|
||||
gathered_shaped_tensors = [
|
||||
torch.empty_like(local_shape, device=local_shape.device)
|
||||
for _ in range(mesh.size())
|
||||
]
|
||||
funcol.all_gather_inplace(gathered_shaped_tensors, local_shape, mesh)
|
||||
|
||||
@maybe_run_for_local_tensor
|
||||
def _validate_and_compute_global_shape(local_shape, gathered_shaped_tensors):
|
||||
sharded_dim_sum = 0
|
||||
shard_dim = placements[0].dim # type: ignore[union-attr]
|
||||
other_dims = [d for d in range(len(shape)) if d != shard_dim]
|
||||
for shape_tensor in gathered_shaped_tensors:
|
||||
if not torch.equal(local_shape[other_dims], shape_tensor[other_dims]):
|
||||
raise RuntimeError(
|
||||
"Non-sharded dimensions should have identical size across ranks."
|
||||
)
|
||||
shape_tensor_list = shape_tensor.tolist()
|
||||
sharded_dim_sum += shape_tensor_list[shard_dim]
|
||||
return sharded_dim_sum
|
||||
|
||||
sharded_dim_sum = _validate_and_compute_global_shape(
|
||||
local_shape, gathered_shaped_tensors
|
||||
)
|
||||
global_shape = list(shape)
|
||||
global_shape[placements[0].dim] = sharded_dim_sum
|
||||
return torch.Size(global_shape)
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
f"Placement type {type(placements[0])} not supported."
|
||||
)
|
||||
|
||||
|
||||
def try_find_mesh_from_args(
|
||||
op_call: torch._ops.OpOverload, args: Sequence[object]
|
||||
) -> DeviceMesh:
|
||||
"""
|
||||
Find the device mesh object from args.
|
||||
It returns None if no mesh is found.
|
||||
NOTE: we can optimize this search if needed
|
||||
"""
|
||||
for arg in args:
|
||||
if isinstance(arg, (dtensor.DTensor, DTensorSpec)):
|
||||
return arg.device_mesh
|
||||
elif (
|
||||
isinstance(arg, (list, tuple))
|
||||
and len(arg) > 0
|
||||
and isinstance(arg[0], (dtensor.DTensor, DTensorSpec))
|
||||
):
|
||||
return arg[0].device_mesh
|
||||
|
||||
raise ValueError(f"Cannot find device mesh from args for op : {op_call}.")
|
||||
|
||||
|
||||
def compute_local_stride(
|
||||
global_stride: ShapeType, local_shape: ShapeType
|
||||
) -> tuple[int, ...]:
|
||||
"""
|
||||
Compute the stride of a local tensor shard, given the global stride and local shape.
|
||||
|
||||
Derives strides by preserving the memory layout (dimension ordering) implied
|
||||
by the global strides, then computing contiguous strides for the local shape
|
||||
in that order. Assumes the global tensor is non-overlapping and dense.
|
||||
"""
|
||||
ndim = len(global_stride)
|
||||
# Sort dims by global stride descending to recover memory layout order.
|
||||
# Stable sort preserves original dim order for ties, which only occur
|
||||
# on size-1 dims where the stride value is semantically irrelevant.
|
||||
perm = sorted(range(ndim), key=lambda d: global_stride[d], reverse=True)
|
||||
local_strides = [0] * ndim
|
||||
s = 1
|
||||
for d in reversed(perm):
|
||||
local_strides[d] = s
|
||||
s *= local_shape[d]
|
||||
return tuple(local_strides)
|
||||
|
||||
|
||||
def normalize_to_torch_size(size) -> torch.Size: # type: ignore[no-untyped-def]
|
||||
"""
|
||||
Unify variable types of size argument to torch.Size
|
||||
Acceptable types include:
|
||||
int, Sequence[int], Tuple[int], Tuple[Sequence[int]],
|
||||
or torch.Size
|
||||
"""
|
||||
if isinstance(size, torch.Size):
|
||||
return size
|
||||
|
||||
if isinstance(size, int):
|
||||
torch_size = [size]
|
||||
elif len(size) == 1 and isinstance(size[0], Sequence):
|
||||
torch_size = list(size[0])
|
||||
else:
|
||||
torch_size = list(size)
|
||||
return torch.Size(torch_size)
|
||||
|
||||
|
||||
def assert_no_mixed_partial_types(placements: Sequence[Placement]) -> None:
|
||||
"""
|
||||
Assert that a placement list doesn't contain mixed Partial reduce types.
|
||||
|
||||
Mixed Partial types (e.g., ``Partial("sum")`` and ``Partial("max")`` together in the
|
||||
same placement list) are not supported and will raise a ``ValueError``. This restriction
|
||||
exists because nonlinear reductions (e.g., max) don't commute with linear reductions
|
||||
(e.g., sum), which means the relative ordering of different partial types would be
|
||||
semantically critical during redistribution. Rather than introducing complex ordering
|
||||
constraints, we prohibit mixing different Partial reduce types.
|
||||
|
||||
Note: Partial("sum") and Partial("avg") DO commute with each other, so they can be ordered
|
||||
arbitrarily, and we allow this.
|
||||
|
||||
This function is called internally by public APIs like :meth:`DTensor.from_local` and
|
||||
:func:`distribute_tensor` to validate placements early, before DTensor construction.
|
||||
|
||||
Args:
|
||||
placements (Sequence[:class:`Placement`]): A sequence of placement specifications
|
||||
to validate.
|
||||
|
||||
Raises:
|
||||
ValueError: If the placements contain more than one distinct Partial reduce type.
|
||||
"""
|
||||
partial_reduce_ops: set[str] = set()
|
||||
for p in placements:
|
||||
if isinstance(p, Partial):
|
||||
partial_reduce_ops.add(p.reduce_op)
|
||||
|
||||
if len(partial_reduce_ops) > 1 and partial_reduce_ops != {"sum", "avg"}:
|
||||
raise ValueError(
|
||||
f"Mixed Partial reduce types are not supported in the same placement list. "
|
||||
f"Found reduce ops: {partial_reduce_ops}. "
|
||||
f"Please ensure all Partial placements use the same reduce operation."
|
||||
)
|
||||
@@ -0,0 +1,70 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import torch._C
|
||||
from torch.distributed.tensor.debug._comm_mode import CommDebugMode
|
||||
from torch.distributed.tensor.debug._visualize_sharding import visualize_sharding
|
||||
|
||||
|
||||
__all__ = ["CommDebugMode", "visualize_sharding"]
|
||||
|
||||
|
||||
def _get_python_sharding_prop_cache_info():
|
||||
"""
|
||||
Get the cache info for the Python sharding propagation cache, used for debugging purpose only.
|
||||
This would return a named tuple showing hits, misses, maxsize and cursize of the sharding
|
||||
propagator cache. Note that directly calling into the sharding propagator does not share cache
|
||||
state with the DTensor dispatch fast path!
|
||||
"""
|
||||
from torch.distributed.tensor._api import DTensor
|
||||
|
||||
return (
|
||||
DTensor._op_dispatcher.sharding_propagator.propagate_op_sharding.cache_info() # type:ignore[attr-defined]
|
||||
)
|
||||
|
||||
|
||||
def _get_fast_path_sharding_prop_cache_stats():
|
||||
"""
|
||||
Get a tuple (hits, misses) for the fast path sharding propagation cache, used for debugging
|
||||
only.
|
||||
"""
|
||||
return torch._C._get_DTensor_sharding_propagator_cache_stats()
|
||||
|
||||
|
||||
def _clear_python_sharding_prop_cache():
|
||||
"""
|
||||
Clears the cache for the Python sharding propagation cache, used for debugging purpose only.
|
||||
"""
|
||||
from torch.distributed.tensor._api import DTensor
|
||||
|
||||
return (
|
||||
DTensor._op_dispatcher.sharding_propagator.propagate_op_sharding.cache_clear() # type:ignore[attr-defined]
|
||||
)
|
||||
|
||||
|
||||
def _clear_fast_path_sharding_prop_cache():
|
||||
"""
|
||||
Clears the cache for the fast path sharding propagation cache, used for debugging purpose only.
|
||||
"""
|
||||
torch._C._clear_DTensor_sharding_propagator_cache()
|
||||
|
||||
|
||||
def _clear_sharding_prop_cache():
|
||||
"""
|
||||
Clears both the Python and fast path sharding propagation caches, used for debugging purpose only.
|
||||
This is the recommended way to clear all sharding propagation caches.
|
||||
"""
|
||||
_clear_python_sharding_prop_cache()
|
||||
_clear_fast_path_sharding_prop_cache()
|
||||
|
||||
|
||||
def _reinit_dispatch_logger():
|
||||
"""
|
||||
Resets the cached DTensor dispatch logger state so that the next DTensor
|
||||
dispatch re-checks whether debug logging is enabled. Call this after
|
||||
changing the log level on the ``torch.distributed.tensor._dispatch`` logger.
|
||||
"""
|
||||
torch._C._reinit_DTensor_dispatch_logger()
|
||||
|
||||
|
||||
# Set namespace for exposed private names
|
||||
CommDebugMode.__module__ = "torch.distributed.tensor.debug"
|
||||
visualize_sharding.__module__ = "torch.distributed.tensor.debug"
|
||||
+747
@@ -0,0 +1,747 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import copy
|
||||
import json
|
||||
import re
|
||||
import weakref
|
||||
from collections import defaultdict
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import torch.nn
|
||||
from torch._guards import detect_fake_mode
|
||||
from torch.autograd.graph import register_multi_grad_hook
|
||||
from torch.distributed._tools.mod_tracker import ModTracker
|
||||
from torch.distributed.tensor._api import DTensor
|
||||
from torch.nn.modules.module import (
|
||||
register_module_forward_hook,
|
||||
register_module_forward_pre_hook,
|
||||
register_module_full_backward_pre_hook,
|
||||
)
|
||||
from torch.utils._python_dispatch import TorchDispatchMode
|
||||
from torch.utils._pytree import tree_flatten
|
||||
|
||||
|
||||
__all__ = ["CommDebugMode"]
|
||||
|
||||
funcol_native = torch.ops._c10d_functional
|
||||
funcol_py = torch.ops.c10d_functional
|
||||
funcol_autograd = torch.ops._c10d_functional_autograd
|
||||
c10d_ops = torch.ops.c10d
|
||||
|
||||
NATIVE_TO_PY_MAPPING = {
|
||||
funcol_native.all_gather_into_tensor: funcol_py.all_gather_into_tensor,
|
||||
funcol_native.all_gather_into_tensor_coalesced: funcol_py.all_gather_into_tensor_coalesced,
|
||||
funcol_native.all_reduce: funcol_py.all_reduce,
|
||||
funcol_native.all_reduce_coalesced: funcol_py.all_reduce_coalesced,
|
||||
funcol_native.all_to_all_single: funcol_py.all_to_all_single,
|
||||
funcol_native.broadcast: funcol_py.broadcast,
|
||||
funcol_native.reduce_scatter_tensor: funcol_py.reduce_scatter_tensor,
|
||||
funcol_native.reduce_scatter_tensor_coalesced: funcol_py.reduce_scatter_tensor_coalesced,
|
||||
# functional ops
|
||||
funcol_autograd.all_to_all_single: funcol_py.all_to_all_single,
|
||||
}
|
||||
|
||||
c10d_collective_ops = {
|
||||
c10d_ops._allgather_base_,
|
||||
c10d_ops._reduce_scatter_base_,
|
||||
c10d_ops.allgather_,
|
||||
c10d_ops.allgather_coalesced_,
|
||||
c10d_ops.allgather_into_tensor_coalesced_,
|
||||
c10d_ops.allreduce_,
|
||||
c10d_ops.allreduce_coalesced_,
|
||||
c10d_ops.alltoall_,
|
||||
c10d_ops.alltoall_base_,
|
||||
c10d_ops.broadcast_,
|
||||
c10d_ops.gather_,
|
||||
c10d_ops.scatter_,
|
||||
c10d_ops.reduce_,
|
||||
c10d_ops.reduce_scatter_,
|
||||
c10d_ops.reduce_scatter_tensor_coalesced_,
|
||||
}
|
||||
|
||||
trivial_ops = {
|
||||
"aten.detach.default",
|
||||
"aten.t.default",
|
||||
"aten.view.default",
|
||||
"aten._to_copy.default",
|
||||
"aten.as_strided.default",
|
||||
"aten.transpose.int",
|
||||
}
|
||||
|
||||
|
||||
class _CommModeModuleTracker(ModTracker):
|
||||
"""
|
||||
Inherits ModuleTracker and expands on its functionality to track the
|
||||
parameters and sharding information of a model at a module-level
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.module_helper_dict = {}
|
||||
self.module_parameters_dict = {}
|
||||
self.module_parents_dict = {}
|
||||
self.register_forward_hook_handles = {}
|
||||
self.parent_dict = {}
|
||||
self.parent_list = []
|
||||
self.sharding_dict = {}
|
||||
self.activation_checkpointing = False
|
||||
self.name = ""
|
||||
|
||||
def _fw_set_module_hook(self, mod, input, output):
|
||||
"""
|
||||
Updates the current module after module finishes running and
|
||||
all other hooks are resolved
|
||||
"""
|
||||
|
||||
if self.is_bw:
|
||||
self.activation_checkpointing = True
|
||||
else:
|
||||
self.activation_checkpointing = False
|
||||
|
||||
if not self.activation_checkpointing:
|
||||
# module is no longer parent of next modules
|
||||
self.parent_list.pop()
|
||||
|
||||
# set current module to previous parent module
|
||||
self.name = self.parent_list[-1]
|
||||
|
||||
def _fw_pre_hook(self, mod, input):
|
||||
"""
|
||||
This function is called before the forward pass of a module. It
|
||||
collects the parameters and sharding information of a module and
|
||||
stores it in a dictionary.
|
||||
"""
|
||||
if self.is_bw:
|
||||
self.activation_checkpointing = True
|
||||
else:
|
||||
self.activation_checkpointing = False
|
||||
|
||||
self.name = super()._get_mod_name(mod)
|
||||
w_mod = weakref.ref(mod)
|
||||
|
||||
# adds current sub-module to module tracker parent class
|
||||
super()._get_append_fn(w_mod, self.name, False)()
|
||||
|
||||
args, _ = tree_flatten(input)
|
||||
tensors = [a for a in args if isinstance(a, torch.Tensor) and a.requires_grad]
|
||||
if not self.is_bw and tensors:
|
||||
register_multi_grad_hook(
|
||||
tensors, super()._get_pop_fn(w_mod, self.name, True)
|
||||
)
|
||||
|
||||
if not self.activation_checkpointing:
|
||||
# contains information about module ordering and depth in the module tree
|
||||
if self.name not in self.module_helper_dict:
|
||||
self.module_helper_dict[self.name] = {}
|
||||
|
||||
self.module_helper_dict[self.name]["module_type"] = (
|
||||
str(type(mod)).replace("<", "").replace(">", "")
|
||||
)
|
||||
self.module_helper_dict[self.name]["depth"] = len(self.parents) - 1
|
||||
|
||||
for param_name, param in mod.named_parameters(recurse=False):
|
||||
if self.name not in self.module_parameters_dict:
|
||||
self.module_parameters_dict[self.name] = {}
|
||||
|
||||
self.module_parameters_dict[self.name][param_name] = param.data
|
||||
|
||||
if isinstance(param.data, DTensor):
|
||||
key_name = self.name + "." + param_name
|
||||
self.sharding_dict[key_name] = param.data.placements
|
||||
|
||||
if "parameters" not in self.module_helper_dict[self.name]:
|
||||
self.module_helper_dict[self.name]["parameters"] = {}
|
||||
|
||||
self.module_helper_dict[self.name]["parameters"][param_name] = str(
|
||||
param.data.placements
|
||||
)
|
||||
|
||||
# used to store module's parents to ensure correctness in backward pass/checkpointing
|
||||
if self.name not in self.module_parents_dict:
|
||||
self.module_parents_dict[self.name] = copy.deepcopy(self.parents)
|
||||
|
||||
# used to create parent-child module associations for json dumps
|
||||
parent = self.parent_list[-1]
|
||||
if parent not in self.parent_dict:
|
||||
self.parent_dict[parent] = []
|
||||
|
||||
self.parent_dict[parent].append(self.name)
|
||||
self.parent_list.append(self.name)
|
||||
|
||||
self.register_forward_hook_handles[self.name] = mod.register_forward_hook(
|
||||
self._fw_set_module_hook
|
||||
)
|
||||
|
||||
def _fw_post_hook(self, mod, input, output): # pylint: disable=useless-parent-delegation
|
||||
"""
|
||||
This function is called when the forward pass of a module is called.
|
||||
It updates the module tracker and removes the module from parent data
|
||||
"""
|
||||
|
||||
super()._fw_post_hook(mod, input, output)
|
||||
|
||||
def _bw_hook(self, mod, output):
|
||||
"""
|
||||
This function is called when the backward pass of a module is called. It
|
||||
updates the current module for backward passes
|
||||
"""
|
||||
self.activation_checkpointing = False
|
||||
self.name = super()._get_mod_name(mod)
|
||||
|
||||
def __enter__(self):
|
||||
self.activation_checkpointing = False
|
||||
self.module_parameters_dict.clear()
|
||||
self.sharding_dict.clear()
|
||||
self.parent_dict.clear()
|
||||
self.parent_list = ["Global"]
|
||||
self.module_helper_dict.clear()
|
||||
self.module_helper_dict["Global"] = {"depth": 0}
|
||||
self.module_parents_dict.clear()
|
||||
self.module_parents_dict["Global"] = set()
|
||||
self._fw_pre_handle = register_module_forward_pre_hook(self._fw_pre_hook)
|
||||
self._fw_post_handle = register_module_forward_hook(self._fw_post_hook)
|
||||
self.register_forward_hook_handles.clear()
|
||||
self._bw_handle = register_module_full_backward_pre_hook(self._bw_hook)
|
||||
self.name = "Global"
|
||||
|
||||
def __exit__(self, *args):
|
||||
super().__exit__(*args)
|
||||
self._bw_handle.remove()
|
||||
|
||||
# removes all forward_hook handles added in the pre-hook
|
||||
for handle in self.register_forward_hook_handles.values():
|
||||
handle.remove()
|
||||
|
||||
def print_paramater_info(self):
|
||||
print(self.module_parameters_dict)
|
||||
|
||||
def print_sharding_info(self):
|
||||
for key, value in self.sharding_dict.items():
|
||||
print(key + ": " + str(value))
|
||||
|
||||
|
||||
class CommDebugMode(TorchDispatchMode):
|
||||
"""
|
||||
:class:`CommDebugMode` is a context manager that counts the number of
|
||||
functional collectives within its context. It does this using a
|
||||
``TorchDispatchMode``.
|
||||
|
||||
.. note:: Not all collectives are supported yet.
|
||||
|
||||
Example usage
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
mod = ...
|
||||
comm_mode = CommDebugMode()
|
||||
with comm_mode:
|
||||
mod.sum().backward()
|
||||
print(comm_mode.get_comm_counts())
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.supports_higher_order_operators = True
|
||||
self.comm_counts: dict[Any, int] = defaultdict(int)
|
||||
self.comm_module_counts = {}
|
||||
self.comm_module_operation_counts = {}
|
||||
self.comm_registry = set()
|
||||
for native_op, py_op in NATIVE_TO_PY_MAPPING.items():
|
||||
self.comm_registry.add(native_op)
|
||||
self.comm_registry.add(py_op)
|
||||
|
||||
self.comm_registry.add(torch.ops._dtensor.shard_dim_alltoall)
|
||||
self.advanced_module_tracker = _CommModeModuleTracker()
|
||||
|
||||
def generate_json_dump(self, file_name="comm_mode_log.json", noise_level=3):
|
||||
"""
|
||||
Creates json file used to build browser visual
|
||||
0. prints module-level collective counts
|
||||
1. prints dTensor operations not included in trivial operations
|
||||
2. prints operations not included in trivial operations
|
||||
3. prints all operations
|
||||
"""
|
||||
|
||||
(
|
||||
include_DTensor_ops,
|
||||
include_module_data,
|
||||
include_ops,
|
||||
include_trivial_ops,
|
||||
) = self._set_noise_parameters(noise_level)
|
||||
|
||||
# recursively builds json data
|
||||
def add_json_information(json_dict, fqn):
|
||||
json_dict["fqn"] = fqn
|
||||
json_dict["module_type"] = ""
|
||||
json_dict["parameters"] = []
|
||||
json_dict["children"] = []
|
||||
json_dict["collectives_forward"] = []
|
||||
json_dict["collectives_backward"] = []
|
||||
json_dict["operations_forward"] = []
|
||||
json_dict["operations_backward"] = []
|
||||
|
||||
# adds module layer type and parameters, and their sharding
|
||||
if (
|
||||
"module_type" in self.advanced_module_tracker.module_helper_dict[fqn]
|
||||
and include_module_data
|
||||
):
|
||||
json_dict["module_type"] = (
|
||||
self.advanced_module_tracker.module_helper_dict[fqn]["module_type"]
|
||||
)
|
||||
|
||||
if "parameters" in self.advanced_module_tracker.module_helper_dict[fqn]:
|
||||
for (
|
||||
param_name,
|
||||
placement,
|
||||
) in self.advanced_module_tracker.module_helper_dict[fqn][
|
||||
"parameters"
|
||||
].items():
|
||||
json_dict["parameters"].append((param_name, placement))
|
||||
|
||||
# adds module collective information
|
||||
if fqn in self.comm_module_counts:
|
||||
for collective, count in self.comm_module_counts[fqn][
|
||||
"forward"
|
||||
].items():
|
||||
json_dict["collectives_forward"].append((str(collective), count))
|
||||
|
||||
for collective, count in self.comm_module_counts[fqn][
|
||||
"backward"
|
||||
].items():
|
||||
json_dict["collectives_backward"].append((str(collective), count))
|
||||
|
||||
# adds module operation information
|
||||
forward_operations = []
|
||||
backward_operations = []
|
||||
checkpointing_operations = []
|
||||
|
||||
# only get operations if the minimum operation noise level is set to true
|
||||
if include_DTensor_ops:
|
||||
if fqn in self.comm_module_operation_counts:
|
||||
(
|
||||
forward_operations,
|
||||
backward_operations,
|
||||
checkpointing_operations,
|
||||
) = self._get_operations_list(
|
||||
self.comm_module_operation_counts[fqn]
|
||||
)
|
||||
|
||||
# remove all operations who don't have DTensor inputs
|
||||
if not include_ops:
|
||||
forward_operations = [
|
||||
op for op in forward_operations if len(op["input_sharding"])
|
||||
]
|
||||
backward_operations = [
|
||||
op for op in backward_operations if len(op["input_sharding"])
|
||||
]
|
||||
checkpointing_operations = [
|
||||
op for op in checkpointing_operations if len(op["input_sharding"])
|
||||
]
|
||||
|
||||
# remove all operations in trivial operations set
|
||||
if not include_trivial_ops:
|
||||
forward_operations = [
|
||||
op
|
||||
for op in forward_operations
|
||||
if str(op["name"]) not in trivial_ops
|
||||
]
|
||||
backward_operations = [
|
||||
op
|
||||
for op in backward_operations
|
||||
if str(op["name"]) not in trivial_ops
|
||||
]
|
||||
checkpointing_operations = [
|
||||
op
|
||||
for op in checkpointing_operations
|
||||
if str(op["name"]) not in trivial_ops
|
||||
]
|
||||
|
||||
# converts operation information into string format for json.dumps()
|
||||
forward_operations = copy.deepcopy(forward_operations)
|
||||
for op in forward_operations:
|
||||
op["name"] = str(op["name"])
|
||||
|
||||
for i in range(len(op["input_sharding"])):
|
||||
op["input_sharding"][i] = str(op["input_sharding"][i])
|
||||
op["input_shape"][i] = str(op["input_shape"][i])
|
||||
|
||||
backward_operations = copy.deepcopy(backward_operations)
|
||||
for op in backward_operations:
|
||||
op["name"] = str(op["name"])
|
||||
|
||||
for i in range(len(op["input_sharding"])):
|
||||
op["input_sharding"][i] = str(op["input_sharding"][i])
|
||||
op["input_shape"][i] = str(op["input_shape"][i])
|
||||
|
||||
checkpointing_operations = copy.deepcopy(checkpointing_operations)
|
||||
for op in checkpointing_operations:
|
||||
op["name"] = str(op["name"])
|
||||
|
||||
for i in range(len(op["input_sharding"])):
|
||||
op["input_sharding"][i] = str(op["input_sharding"][i])
|
||||
op["input_shape"][i] = str(op["input_shape"][i])
|
||||
|
||||
json_dict["operations_forward"] = forward_operations
|
||||
json_dict["operations_backward"] = backward_operations
|
||||
json_dict["operations_checkpointing"] = checkpointing_operations
|
||||
|
||||
if fqn not in self.advanced_module_tracker.parent_dict:
|
||||
return json_dict
|
||||
|
||||
# recursively adds module's children
|
||||
for ele in self.advanced_module_tracker.parent_dict[fqn]:
|
||||
json_dict["children"].append(add_json_information({}, ele))
|
||||
|
||||
return json_dict
|
||||
|
||||
json_dict: dict[str, Any] = {}
|
||||
add_json_information(json_dict, "Global")
|
||||
|
||||
# converts dictionary into json file
|
||||
with open(file_name, "w") as json_file:
|
||||
json.dump(json_dict, json_file, indent=4)
|
||||
|
||||
def generate_comm_debug_tracing_table(self, noise_level=3):
|
||||
"""
|
||||
Generates detailed table displaying operations and collective tracing information
|
||||
on a module level. Amount of information is dependent on noise_level
|
||||
|
||||
0. prints module-level collective counts
|
||||
1. prints dTensor operations not included in trivial operations, module information
|
||||
2. prints operations not included in trivial operations
|
||||
3. prints all operations
|
||||
"""
|
||||
|
||||
(
|
||||
include_DTensor_ops,
|
||||
include_module_data,
|
||||
include_ops,
|
||||
include_trivial_ops,
|
||||
) = self._set_noise_parameters(noise_level)
|
||||
|
||||
table = ""
|
||||
for fqn in self.advanced_module_tracker.module_helper_dict:
|
||||
# setting up indentations for table formatting
|
||||
indent = " " * (
|
||||
2 * self.advanced_module_tracker.module_helper_dict[fqn]["depth"]
|
||||
)
|
||||
table += f"{indent}{fqn}\n"
|
||||
|
||||
if include_module_data:
|
||||
if (
|
||||
"module_type"
|
||||
in self.advanced_module_tracker.module_helper_dict[fqn]
|
||||
):
|
||||
module_type = self.advanced_module_tracker.module_helper_dict[fqn][
|
||||
"module_type"
|
||||
]
|
||||
table += f"{indent}*module type: {module_type}\n"
|
||||
|
||||
if "parameters" in self.advanced_module_tracker.module_helper_dict[fqn]:
|
||||
table += f"{indent}*Parameter List\n"
|
||||
for (
|
||||
param_name,
|
||||
placement,
|
||||
) in self.advanced_module_tracker.module_helper_dict[fqn][
|
||||
"parameters"
|
||||
].items():
|
||||
table += f"{indent} *{param_name}: {placement}\n"
|
||||
|
||||
indent += " "
|
||||
collective_indent = " " * (
|
||||
2 * self.advanced_module_tracker.module_helper_dict[fqn]["depth"] + 2
|
||||
)
|
||||
operation_indent = " " * (
|
||||
2 * self.advanced_module_tracker.module_helper_dict[fqn]["depth"] + 3
|
||||
)
|
||||
|
||||
# separate the module's collective and operations by forward and backward
|
||||
forward_collectives = {}
|
||||
backward_collectives = {}
|
||||
if fqn in self.comm_module_counts:
|
||||
forward_collectives = self.comm_module_counts[fqn]["forward"]
|
||||
backward_collectives = self.comm_module_counts[fqn]["backward"]
|
||||
|
||||
forward_operations = []
|
||||
backward_operations = []
|
||||
checkpointing_operations = []
|
||||
|
||||
if include_DTensor_ops:
|
||||
if fqn in self.comm_module_operation_counts:
|
||||
(
|
||||
forward_operations,
|
||||
backward_operations,
|
||||
checkpointing_operations,
|
||||
) = self._get_operations_list(
|
||||
self.comm_module_operation_counts[fqn]
|
||||
)
|
||||
|
||||
def add_tracing_information(table, collectives_dict, operation_list):
|
||||
"""
|
||||
adds tracing information for module's forward or backward
|
||||
"""
|
||||
for collective, count in collectives_dict.items():
|
||||
table += (
|
||||
f"\033[1;33m{collective_indent}*{collective}: {count}\033[0m\n"
|
||||
)
|
||||
|
||||
def add_operations(
|
||||
table, operation, collective_indent, operation_indent
|
||||
):
|
||||
"""
|
||||
adds operation information to the table
|
||||
"""
|
||||
table += f"\033[1;33m{collective_indent}**{operation_name}\033[0m\n"
|
||||
|
||||
if len(operation["input_shape"]):
|
||||
operation_shape = operation["input_shape"]
|
||||
operation_sharding = operation["input_sharding"]
|
||||
operation_device_mesh = operation["device_mesh"]
|
||||
|
||||
table += f"\033[1;31m{operation_indent}shape: {operation_shape}\033[0m\n"
|
||||
table += f"\033[1;31m{operation_indent}sharding: {operation_sharding}\033[0m\n"
|
||||
table += f"\033[1;31m{operation_indent}device mesh: {operation_device_mesh}\033[0m\n"
|
||||
|
||||
return table
|
||||
|
||||
for operation in operation_list:
|
||||
operation_name = str(operation["name"])
|
||||
|
||||
# include all operations
|
||||
if include_trivial_ops:
|
||||
table = add_operations(
|
||||
table, operation, collective_indent, operation_indent
|
||||
)
|
||||
|
||||
# include all operations not in trivial operations
|
||||
elif include_ops and operation_name not in trivial_ops:
|
||||
table = add_operations(
|
||||
table, operation, collective_indent, operation_indent
|
||||
)
|
||||
|
||||
# only include dTensor operations not in trivial set
|
||||
elif (
|
||||
include_DTensor_ops
|
||||
and (operation_name not in trivial_ops)
|
||||
and len(operation["input_shape"])
|
||||
):
|
||||
table = add_operations(
|
||||
table, operation, collective_indent, operation_indent
|
||||
)
|
||||
|
||||
return table
|
||||
|
||||
if len(forward_collectives) or len(forward_operations):
|
||||
table += f"{indent}FORWARD PASS\n"
|
||||
table = add_tracing_information(
|
||||
table, forward_collectives, forward_operations
|
||||
)
|
||||
|
||||
if len(backward_collectives) or len(backward_operations):
|
||||
table += f"{indent}BACKWARD PASS\n"
|
||||
table = add_tracing_information(
|
||||
table, backward_collectives, backward_operations
|
||||
)
|
||||
|
||||
if len(checkpointing_operations):
|
||||
table += f"{indent}ACTIVATION CHECKPOINTING\n"
|
||||
table = add_tracing_information(table, {}, checkpointing_operations)
|
||||
|
||||
return table
|
||||
|
||||
def _get_operations_list(self, module_operation_counts):
|
||||
forward_operations = [
|
||||
op for op in module_operation_counts["operations_list"] if not op["is_bw"]
|
||||
]
|
||||
backward_operations = [
|
||||
op
|
||||
for op in module_operation_counts["operations_list"]
|
||||
if op["is_bw"] and not op["is_activation_checkpointing"]
|
||||
]
|
||||
checkpointing_operations = [
|
||||
op
|
||||
for op in module_operation_counts["operations_list"]
|
||||
if op["is_activation_checkpointing"]
|
||||
]
|
||||
|
||||
return forward_operations, backward_operations, checkpointing_operations
|
||||
|
||||
def get_total_counts(self) -> int:
|
||||
return sum(self.comm_counts.values())
|
||||
|
||||
def get_comm_counts(self) -> dict[Any, int]:
|
||||
"""Returns the communication counts as a dictionary.
|
||||
|
||||
Returns:
|
||||
Dict[Any, int]: The communication counts as a dictionary.
|
||||
"""
|
||||
return self.comm_counts
|
||||
|
||||
def get_parameter_info(self) -> dict[str, dict[str, Any]]:
|
||||
return self.advanced_module_tracker.module_parameters_dict
|
||||
|
||||
def get_sharding_info(self) -> dict[str, dict[str, Any]]:
|
||||
return self.advanced_module_tracker.sharding_dict
|
||||
|
||||
def __enter__(self):
|
||||
self.comm_counts.clear()
|
||||
self.comm_module_counts.clear()
|
||||
self.comm_module_counts["Global"] = {}
|
||||
self.comm_module_counts["Global"]["forward"] = defaultdict(int)
|
||||
self.comm_module_counts["Global"]["backward"] = defaultdict(int)
|
||||
|
||||
self.comm_module_operation_counts.clear()
|
||||
|
||||
super().__enter__()
|
||||
self.advanced_module_tracker.__enter__()
|
||||
return self
|
||||
|
||||
# pyrefly: ignore [bad-override]
|
||||
def __exit__(self, *args):
|
||||
self.advanced_module_tracker.__exit__()
|
||||
super().__exit__(*args)
|
||||
|
||||
def log_comm_debug_tracing_table_to_file(
|
||||
self, file_name="comm_mode_log.txt", noise_level=3
|
||||
):
|
||||
"""
|
||||
Alternative to console CommDebugMode output, writes to file specified by the user
|
||||
"""
|
||||
ansi_escape = re.compile(r"\x1B\[[0-?]*[ -/]*[@-~]")
|
||||
table = ansi_escape.sub("", self.generate_comm_debug_tracing_table(noise_level))
|
||||
|
||||
with open(file_name, "w") as log_file:
|
||||
log_file.write(table)
|
||||
|
||||
def _set_noise_parameters(self, noise_level):
|
||||
"""
|
||||
sets variables controlling what information displays based on noise level
|
||||
"""
|
||||
include_DTensor_ops = False
|
||||
include_module_data = False
|
||||
include_ops = False
|
||||
include_trivial_ops = False
|
||||
|
||||
if noise_level > 0:
|
||||
include_DTensor_ops = True
|
||||
include_module_data = True
|
||||
|
||||
if noise_level > 1:
|
||||
include_ops = True
|
||||
|
||||
if noise_level > 2:
|
||||
include_trivial_ops = True
|
||||
|
||||
return (
|
||||
include_DTensor_ops,
|
||||
include_module_data,
|
||||
include_ops,
|
||||
include_trivial_ops,
|
||||
)
|
||||
|
||||
def __torch_dispatch__(self, func, types, args=(), kwargs=None):
|
||||
# When running this mode with DTensor, ordinarily all modes will
|
||||
# run **before** subclasses get a chance to run.
|
||||
# Returning NotImplemented here gives us a chance to let DTensor
|
||||
# run and desugar into comms ops, before CommDebugMode sees them.
|
||||
|
||||
# Higher-order operators (e.g. run_dtensor_rng_op) don't have
|
||||
# _overloadpacket and aren't collectives — just redispatch.
|
||||
if isinstance(func, torch._ops.HigherOrderOperator):
|
||||
kwargs = kwargs if kwargs else {}
|
||||
return func(*args, **kwargs)
|
||||
|
||||
# sets up operation-level collective count
|
||||
if self.advanced_module_tracker.name not in self.comm_module_operation_counts:
|
||||
# dictionary should hold module input and output shape, operations list and collective counter
|
||||
self.comm_module_operation_counts[self.advanced_module_tracker.name] = {
|
||||
"operations_list": []
|
||||
}
|
||||
operation_dict = {}
|
||||
operation_dict["name"] = func
|
||||
|
||||
operation_dict["input_shape"] = []
|
||||
operation_dict["input_sharding"] = []
|
||||
operation_dict["device_mesh"] = ""
|
||||
|
||||
# tracks if the operation is part of the backward pass
|
||||
operation_dict["is_bw"] = self.advanced_module_tracker.is_bw
|
||||
|
||||
# tracks if the operation is part of activation checkpointing
|
||||
operation_dict["is_activation_checkpointing"] = (
|
||||
self.advanced_module_tracker.activation_checkpointing
|
||||
)
|
||||
|
||||
if any(t == DTensor for t in types):
|
||||
for ele in args:
|
||||
if isinstance(ele, DTensor):
|
||||
# saves shapes and placements of all DTensor args
|
||||
operation_dict["input_shape"].append(ele.shape)
|
||||
operation_dict["input_sharding"].append(ele.placements)
|
||||
operation_dict["device_mesh"] = str(ele.device_mesh)
|
||||
|
||||
self.comm_module_operation_counts[self.advanced_module_tracker.name][
|
||||
"operations_list"
|
||||
].append(operation_dict)
|
||||
|
||||
return NotImplemented
|
||||
|
||||
kwargs = kwargs if kwargs else {}
|
||||
out = func(*args, **kwargs)
|
||||
func_packet = func._overloadpacket
|
||||
|
||||
# We have many tests that use CommDebugMode to verify the occurrence of
|
||||
# collectives. These tests do so by querying comm_counts with legacy
|
||||
# funcol ops as key. For the purpose of native funcol migration, we
|
||||
# need these tests to work for both legacy and native funcol. To avoid
|
||||
# the need to modify all tests to accommodate the two implementations,
|
||||
# we make CommDebugMode translate native funcol ops into legacy funcol
|
||||
# ops until the migration finishes.
|
||||
|
||||
if func_packet in self.comm_registry or func_packet in c10d_collective_ops:
|
||||
if func_packet in NATIVE_TO_PY_MAPPING:
|
||||
func_packet = NATIVE_TO_PY_MAPPING[func_packet]
|
||||
self.comm_counts[func_packet] += 1
|
||||
|
||||
key = "forward"
|
||||
if self.advanced_module_tracker.is_bw:
|
||||
key = "backward"
|
||||
|
||||
# adds collective count to current module
|
||||
if self.advanced_module_tracker.name not in self.comm_module_counts:
|
||||
self.comm_module_counts[self.advanced_module_tracker.name] = {}
|
||||
self.comm_module_counts[self.advanced_module_tracker.name][
|
||||
"forward"
|
||||
] = defaultdict(int)
|
||||
self.comm_module_counts[self.advanced_module_tracker.name][
|
||||
"backward"
|
||||
] = defaultdict(int)
|
||||
self.comm_module_counts[self.advanced_module_tracker.name][key][
|
||||
func_packet
|
||||
] += 1
|
||||
|
||||
# adds collective count to parent modules
|
||||
for par in self.advanced_module_tracker.module_parents_dict[
|
||||
self.advanced_module_tracker.name
|
||||
]:
|
||||
# makes sure we aren't double counting when current sub-module hasn't been removed from parents
|
||||
if par != self.advanced_module_tracker.name:
|
||||
if par not in self.comm_module_counts:
|
||||
self.comm_module_counts[par] = {}
|
||||
self.comm_module_counts[par]["forward"] = defaultdict(int)
|
||||
self.comm_module_counts[par]["backward"] = defaultdict(int)
|
||||
self.comm_module_counts[par][key][func_packet] += 1
|
||||
|
||||
# if tensor op uses fake tensors, return
|
||||
if detect_fake_mode(args):
|
||||
return out
|
||||
|
||||
# add tensor operation to module operation list
|
||||
self.comm_module_operation_counts[self.advanced_module_tracker.name][
|
||||
"operations_list"
|
||||
].append(operation_dict)
|
||||
|
||||
return out
|
||||
|
||||
def __repr__(self):
|
||||
return f"CommDebugMode(get_total_counts()={self.get_total_counts()})"
|
||||
+105
@@ -0,0 +1,105 @@
|
||||
# mypy: allow-untyped-defs
|
||||
from operator import itemgetter
|
||||
|
||||
import torch
|
||||
import torch.fx
|
||||
import torch.nn as nn
|
||||
from functorch.compile import make_boxed_func
|
||||
from torch._functorch.compilers import aot_module
|
||||
from torch._inductor.decomposition import select_decomp_table
|
||||
from torch.distributed.tensor import DTensor
|
||||
|
||||
|
||||
inductor_decomps = select_decomp_table()
|
||||
|
||||
graphs: list[torch.fx.GraphModule] = []
|
||||
|
||||
|
||||
def fwd_bwd_compiler(fx_g, _):
|
||||
graphs.append(fx_g)
|
||||
return make_boxed_func(fx_g)
|
||||
|
||||
|
||||
def get_inductor_decomp_graphs(model: nn.Module, args, kwargs):
|
||||
"""
|
||||
Obtain forward and backward graphs of a model with inductor decompositions using tracing and aot_module.
|
||||
|
||||
Convenient util to get the fwd and bwd graphs of an arbitrary model
|
||||
with inductor decompositions. Note that this would simply do tracing
|
||||
with aot_module and don't ensure correctness. This is useful to track
|
||||
the ops needed in DTensor.
|
||||
"""
|
||||
compiled_mod = aot_module(
|
||||
model, fw_compiler=fwd_bwd_compiler, decompositions=inductor_decomps
|
||||
)
|
||||
output = compiled_mod(*args, **kwargs)
|
||||
|
||||
if output.ndim != 0:
|
||||
# if output is not a scalar tensor, by default sum it in order to
|
||||
# run backward
|
||||
output = output.sum()
|
||||
|
||||
output.backward()
|
||||
|
||||
# one fwd, one bwd graph
|
||||
if len(graphs) != 2:
|
||||
raise AssertionError
|
||||
return graphs
|
||||
|
||||
|
||||
def print_op_coverage_summary(model: nn.Module, args, kwargs, *, output_csv=False):
|
||||
"""
|
||||
Util to print the operator coverage summary of a certain model with tabulute.
|
||||
|
||||
Must have tabulate module installed.
|
||||
"""
|
||||
# python module required for summary
|
||||
import csv
|
||||
|
||||
from tabulate import tabulate
|
||||
|
||||
fwd_graph, bwd_graph = get_inductor_decomp_graphs(model, args, kwargs)
|
||||
|
||||
op_counts = {}
|
||||
|
||||
for node in fwd_graph.graph.nodes:
|
||||
if node.op == "call_function" and isinstance(
|
||||
node.target, torch._ops.OpOverload
|
||||
):
|
||||
if node.target not in op_counts:
|
||||
op_counts[node.target] = 0
|
||||
|
||||
op_counts[node.target] += 1
|
||||
|
||||
for node in bwd_graph.graph.nodes:
|
||||
if node.op == "call_function" and isinstance(
|
||||
node.target, torch._ops.OpOverload
|
||||
):
|
||||
if node.target not in op_counts:
|
||||
op_counts[node.target] = 0
|
||||
|
||||
op_counts[node.target] += 1
|
||||
|
||||
op_infos = []
|
||||
|
||||
for op, count in op_counts.items():
|
||||
supported = op in DTensor._op_dispatcher.sharding_propagator.op_to_rules
|
||||
op_infos.append([op, str(op._schema), count, supported])
|
||||
|
||||
# sort the op info base on the total count index
|
||||
count_idx = 2
|
||||
op_infos.sort(key=itemgetter(count_idx), reverse=True)
|
||||
|
||||
headers = ["Operator", "Schema", "Total Count", "Supported"]
|
||||
print(tabulate(op_infos, headers=headers))
|
||||
|
||||
if output_csv:
|
||||
# Open a CSV file for writing
|
||||
with open("op_summary.csv", "w", newline="") as csv_file:
|
||||
# Create a CSV writer object
|
||||
csv_writer = csv.writer(csv_file)
|
||||
|
||||
csv_writer.writerow(headers)
|
||||
# Write each table row to the CSV file
|
||||
for row in op_infos:
|
||||
csv_writer.writerow(row)
|
||||
+227
@@ -0,0 +1,227 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import importlib.util
|
||||
|
||||
import numpy as np
|
||||
|
||||
from torch._prims_common import ShapeType
|
||||
from torch.distributed.tensor._utils import _compute_local_shape_and_global_offset
|
||||
|
||||
|
||||
__all__ = ["visualize_sharding"]
|
||||
|
||||
Color = tuple[float, float, float]
|
||||
|
||||
|
||||
def _create_table(
|
||||
shards: list[tuple[tuple[int, int], tuple[int, int], int]], device_kind: str = ""
|
||||
):
|
||||
"""
|
||||
Creates a tabulate table given row and column ranges with device name
|
||||
"""
|
||||
from tabulate import tabulate
|
||||
|
||||
# Extract unique row and column ranges
|
||||
row_ranges = sorted({block[0] for block in shards})
|
||||
col_ranges = sorted({block[1] for block in shards})
|
||||
|
||||
# Create a matrix initialized with empty strings
|
||||
matrix = [["" for _ in col_ranges] for _ in row_ranges]
|
||||
|
||||
# Fill the matrix with values
|
||||
for block in shards:
|
||||
row_index = row_ranges.index(block[0])
|
||||
col_index = col_ranges.index(block[1])
|
||||
if matrix[row_index][col_index] == "":
|
||||
matrix[row_index][col_index] = device_kind + ":" + str(block[2])
|
||||
else:
|
||||
matrix[row_index][col_index] += "," + str(block[2])
|
||||
|
||||
# Prepare headers
|
||||
row_headers = [f"Row {r[0]}-{r[1]}" for r in row_ranges]
|
||||
col_headers = [f"Col {c[0]}-{c[1]}" for c in col_ranges]
|
||||
|
||||
return tabulate(matrix, headers=col_headers, showindex=row_headers)
|
||||
|
||||
|
||||
def make_color_iter(color_map, num_rows, num_cols):
|
||||
num_colors = num_rows * num_cols
|
||||
for idx in range(num_colors):
|
||||
yield color_map(idx)
|
||||
|
||||
|
||||
def _canonicalize_color(color: Color) -> str:
|
||||
if isinstance(color, str):
|
||||
return color
|
||||
r, g, b = (int(a * 255) for a in color)
|
||||
return f"#{r:02X}{g:02X}{b:02X}"
|
||||
|
||||
|
||||
def _get_text_color(color: str) -> str:
|
||||
r, g, b = map(lambda x: int(x, 16), (color[1:3], color[3:5], color[5:7])) # noqa: C417
|
||||
if (r * 0.299 + g * 0.587 + b * 0.114) > 186:
|
||||
return "#000000"
|
||||
return "#ffffff"
|
||||
|
||||
|
||||
def _create_rich_table(
|
||||
shape: ShapeType,
|
||||
shards: list[tuple[tuple[int, int], tuple[int, int], int]],
|
||||
device_kind: str = "",
|
||||
scale: float = 1.0,
|
||||
min_width: int = 9,
|
||||
max_width: int = 80,
|
||||
):
|
||||
import matplotlib
|
||||
import rich.align
|
||||
import rich.box
|
||||
import rich.console
|
||||
import rich.padding
|
||||
import rich.style
|
||||
import rich.table
|
||||
|
||||
dtensor_height = shape[0]
|
||||
dtensor_width = shape[1] if len(shape) == 2 else 1
|
||||
|
||||
row_ranges = sorted({s[0] for s in shards})
|
||||
col_ranges = sorted({s[1] for s in shards})
|
||||
num_rows, num_cols = len(row_ranges), len(col_ranges)
|
||||
|
||||
console = rich.console.Console(width=max_width)
|
||||
use_color = console.color_system
|
||||
color_iter = make_color_iter(matplotlib.colormaps["tab20b"], num_rows, num_cols)
|
||||
|
||||
base_height = int(10 * scale)
|
||||
aspect_ratio = (shape[1] if len(shape) == 2 else 1) / shape[0]
|
||||
base_width = int(base_height * aspect_ratio)
|
||||
height_to_width_ratio = 2.5
|
||||
|
||||
table = rich.table.Table(
|
||||
show_header=False,
|
||||
show_lines=not use_color,
|
||||
padding=0,
|
||||
highlight=not use_color,
|
||||
pad_edge=False,
|
||||
box=rich.box.SQUARE if not use_color else None,
|
||||
)
|
||||
for row in range(num_rows):
|
||||
table_row = []
|
||||
for col in range(num_cols):
|
||||
entry = (
|
||||
device_kind
|
||||
+ ":"
|
||||
+ ",".join(
|
||||
[
|
||||
str(device_id)
|
||||
for row_range, col_range, device_id in shards
|
||||
if row_range == row_ranges[row] and col_range == col_ranges[col]
|
||||
]
|
||||
)
|
||||
)
|
||||
width = (col_ranges[col][1] - col_ranges[col][0]) / dtensor_width
|
||||
width = int(width * base_width * height_to_width_ratio)
|
||||
height = (row_ranges[row][1] - row_ranges[row][0]) / dtensor_height
|
||||
height = int(height * base_height)
|
||||
left_padding, remainder = divmod(width - len(entry) - 2, 2)
|
||||
right_padding = left_padding + remainder
|
||||
top_padding, remainder = divmod(height - 2, 2)
|
||||
bottom_padding = top_padding + remainder
|
||||
if use_color:
|
||||
color = _canonicalize_color(next(color_iter)[:3])
|
||||
text_color = _get_text_color(color)
|
||||
top_padding += 1
|
||||
bottom_padding += 1
|
||||
left_padding += 1
|
||||
right_padding += 1
|
||||
else:
|
||||
color = None
|
||||
text_color = None
|
||||
padding = (
|
||||
max(top_padding, 0),
|
||||
max(right_padding, 0),
|
||||
max(bottom_padding, 0),
|
||||
max(left_padding, 0),
|
||||
)
|
||||
table_row.append(
|
||||
rich.padding.Padding(
|
||||
rich.align.Align(entry, "center", vertical="middle"),
|
||||
padding,
|
||||
style=rich.style.Style(bgcolor=color, color=text_color),
|
||||
)
|
||||
)
|
||||
table.add_row(*table_row)
|
||||
console.print(table, end="\n\n")
|
||||
|
||||
|
||||
def visualize_sharding(dtensor, header="", use_rich: bool = False):
|
||||
"""
|
||||
Visualizes sharding in the terminal for :class:`DTensor` that are 1D or 2D.
|
||||
|
||||
.. note:: This requires the ``tabulate`` package, or ``rich`` and ``matplotlib``.
|
||||
No sharding info will be printed for empty tensors
|
||||
"""
|
||||
if dtensor.numel() == 0: # Do not print empty dtensors.
|
||||
return
|
||||
|
||||
if len(dtensor.shape) >= 3:
|
||||
raise RuntimeError("visualize sharding supports only 1D or 2D DTensor")
|
||||
|
||||
if dtensor.device_mesh.get_coordinate() is None: # current rank is not in the mesh
|
||||
return
|
||||
|
||||
# Only display the visualization once for each DTensor, on the rank whose
|
||||
# coordinate is 0 on all dimensions. For example, if the mesh is a full mesh,
|
||||
# we will only print on rank 0.
|
||||
local_rank_zero_on_all_dim = all(
|
||||
dtensor.device_mesh.get_local_rank(mesh_dim=dim) == 0
|
||||
for dim in range(dtensor.device_mesh.ndim)
|
||||
)
|
||||
if not local_rank_zero_on_all_dim:
|
||||
return
|
||||
|
||||
device_coords = {
|
||||
int(device_index.item()): list(coord)
|
||||
for coord, device_index in np.ndenumerate(
|
||||
np.array(dtensor.device_mesh.mesh.tolist())
|
||||
)
|
||||
}
|
||||
|
||||
device_shard_shape_and_offsets = {
|
||||
device_index: _compute_local_shape_and_global_offset(
|
||||
dtensor.shape,
|
||||
dtensor.device_mesh.shape,
|
||||
lambda i: device_coords[device_index][i],
|
||||
dtensor.placements,
|
||||
)
|
||||
for device_index in device_coords
|
||||
}
|
||||
|
||||
# Extend shards in a 1D tensor to 2D
|
||||
device_shard_shape_and_offsets = {
|
||||
device_index: (
|
||||
shape if len(shape) == 2 else (shape[0], 1),
|
||||
offset if len(offset) == 2 else (offset[0], 0),
|
||||
)
|
||||
for device_index, (shape, offset) in device_shard_shape_and_offsets.items()
|
||||
}
|
||||
|
||||
shards = [
|
||||
(
|
||||
(offset[0], offset[0] + shape[0] - 1),
|
||||
(offset[1], offset[1] + shape[1] - 1),
|
||||
device_index,
|
||||
)
|
||||
for device_index, (shape, offset) in device_shard_shape_and_offsets.items()
|
||||
]
|
||||
|
||||
if (
|
||||
importlib.util.find_spec("rich")
|
||||
and importlib.util.find_spec("matplotlib")
|
||||
and use_rich
|
||||
):
|
||||
_create_rich_table(
|
||||
dtensor.shape, shards, device_kind=dtensor.device_mesh.device_type
|
||||
)
|
||||
elif importlib.util.find_spec("tabulate"):
|
||||
print(_create_table(shards, device_kind=dtensor.device_mesh.device_type))
|
||||
else:
|
||||
raise ValueError("`visualize_sharding` requires either `rich` or `tabulate`.")
|
||||
@@ -0,0 +1,9 @@
|
||||
from torch.distributed.device_mesh import ( # noqa: F401
|
||||
_get_device_handle,
|
||||
_mesh_resources,
|
||||
DeviceMesh,
|
||||
init_device_mesh,
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["init_device_mesh", "DeviceMesh"]
|
||||
+34
@@ -0,0 +1,34 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates
|
||||
from collections.abc import Iterator
|
||||
from contextlib import contextmanager
|
||||
|
||||
from torch.distributed.tensor._api import DTensor
|
||||
from torch.distributed.tensor.experimental._attention import context_parallel
|
||||
from torch.distributed.tensor.experimental._func_map import local_map
|
||||
from torch.distributed.tensor.experimental._register_sharding import register_sharding
|
||||
|
||||
|
||||
__all__ = ["context_parallel", "implicit_replication", "local_map", "register_sharding"]
|
||||
|
||||
|
||||
@contextmanager
|
||||
def implicit_replication() -> Iterator[None]:
|
||||
"""
|
||||
This context manager allows :class:`DTensor` to implicitly treat all non-DTensors (``torch.Tensor``)
|
||||
in the program be replicate :class:`DTensor` s during the operator computation.
|
||||
|
||||
.. warning:: This might possible lead to incorrect results if ``torch.Tensor`` s are not replicated
|
||||
in practice, please use it at your discretion.
|
||||
"""
|
||||
try:
|
||||
DTensor._op_dispatcher._allow_implicit_replication = True
|
||||
yield
|
||||
finally:
|
||||
DTensor._op_dispatcher._allow_implicit_replication = False
|
||||
|
||||
|
||||
# Set namespace for exposed private names
|
||||
context_parallel.__module__ = "torch.distributed.tensor.experimental"
|
||||
implicit_replication.__module__ = "torch.distributed.tensor.experimental"
|
||||
local_map.__module__ = "torch.distributed.tensor.experimental"
|
||||
register_sharding.__module__ = "torch.distributed.tensor.experimental"
|
||||
+44
@@ -0,0 +1,44 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates
|
||||
# Backward compatibility stub - this module has been moved to _context_parallel/_attention.py
|
||||
|
||||
from ._context_parallel._attention import (
|
||||
_CausalBehavior,
|
||||
_context_parallel_shard,
|
||||
_ContextParallel,
|
||||
_cp_options,
|
||||
_disable_context_parallel_dispatcher,
|
||||
_enable_context_parallel_dispatcher,
|
||||
_is_causal_behavior,
|
||||
_RotateMethod,
|
||||
_templated_ring_attention,
|
||||
context_parallel,
|
||||
context_parallel_unshard,
|
||||
set_rotate_method,
|
||||
)
|
||||
from ._context_parallel._load_balancer import (
|
||||
_HeadTailLoadBalancer,
|
||||
_LoadBalancer,
|
||||
_PerDocumentHeadTailLoadBalancer,
|
||||
_PTRRLoadBalancer,
|
||||
)
|
||||
|
||||
|
||||
# TODO(fegin): add deprecation message once the final interfaces are concluded.
|
||||
__all__ = [
|
||||
"_CausalBehavior",
|
||||
"_context_parallel_shard",
|
||||
"_ContextParallel",
|
||||
"_cp_options",
|
||||
"_disable_context_parallel_dispatcher",
|
||||
"_enable_context_parallel_dispatcher",
|
||||
"_is_causal_behavior",
|
||||
"_RotateMethod",
|
||||
"_templated_ring_attention",
|
||||
"context_parallel",
|
||||
"context_parallel_unshard",
|
||||
"set_rotate_method",
|
||||
"_HeadTailLoadBalancer",
|
||||
"_LoadBalancer",
|
||||
"_PerDocumentHeadTailLoadBalancer",
|
||||
"_PTRRLoadBalancer",
|
||||
]
|
||||
+46
@@ -0,0 +1,46 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates
|
||||
# Context Parallel components
|
||||
|
||||
from ._attention import (
|
||||
_CausalBehavior,
|
||||
_context_parallel_shard,
|
||||
_ContextParallel,
|
||||
_cp_options,
|
||||
_disable_context_parallel_dispatcher,
|
||||
_enable_context_parallel_dispatcher,
|
||||
_is_causal_behavior,
|
||||
_RotateMethod,
|
||||
context_parallel,
|
||||
context_parallel_unshard,
|
||||
set_rotate_method,
|
||||
)
|
||||
from ._cp_custom_ops import flex_cp_allgather
|
||||
from ._load_balancer import (
|
||||
_HeadTailLoadBalancer,
|
||||
_LoadBalancer,
|
||||
_PerDocumentHeadTailLoadBalancer,
|
||||
_PTRRLoadBalancer,
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
# From _attention
|
||||
"_CausalBehavior",
|
||||
"_context_parallel_shard",
|
||||
"_ContextParallel",
|
||||
"_cp_options",
|
||||
"_disable_context_parallel_dispatcher",
|
||||
"_enable_context_parallel_dispatcher",
|
||||
"_is_causal_behavior",
|
||||
"_RotateMethod",
|
||||
"context_parallel",
|
||||
"context_parallel_unshard",
|
||||
"set_rotate_method",
|
||||
# From _cp_custom_ops
|
||||
"flex_cp_allgather",
|
||||
# From _load_balancer
|
||||
"_HeadTailLoadBalancer",
|
||||
"_LoadBalancer",
|
||||
"_PerDocumentHeadTailLoadBalancer",
|
||||
"_PTRRLoadBalancer",
|
||||
]
|
||||
+1713
File diff suppressed because it is too large
Load Diff
+88
@@ -0,0 +1,88 @@
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import torch.distributed._functional_collectives as funcol
|
||||
import torch.distributed.distributed_c10d as c10d
|
||||
|
||||
|
||||
@torch.library.custom_op("cplib::flex_cp_allgather", mutates_args=())
|
||||
def flex_cp_allgather(
|
||||
k: torch.Tensor, v: torch.Tensor, seq_dim: int, pg_name: c10d.GroupName
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
k = k.contiguous()
|
||||
v = v.contiguous()
|
||||
k = funcol.all_gather_tensor(k, seq_dim, pg_name)
|
||||
v = funcol.all_gather_tensor(v, seq_dim, pg_name)
|
||||
if isinstance(k, funcol.AsyncCollectiveTensor):
|
||||
k = k.wait()
|
||||
if isinstance(v, funcol.AsyncCollectiveTensor):
|
||||
v = v.wait()
|
||||
return k, v
|
||||
|
||||
|
||||
@flex_cp_allgather.register_fake
|
||||
def _(
|
||||
k: torch.Tensor, v: torch.Tensor, seq_dim: int, pg_name: c10d.GroupName
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
shape_k = list(k.shape)
|
||||
shape_v = list(v.shape)
|
||||
shape_k[seq_dim] *= c10d._get_group_size_by_name(pg_name)
|
||||
shape_v[seq_dim] *= c10d._get_group_size_by_name(pg_name)
|
||||
new_k = torch.empty(shape_k, dtype=k.dtype, device=k.device)
|
||||
new_v = torch.empty(shape_v, dtype=v.dtype, device=v.device)
|
||||
return new_k, new_v
|
||||
|
||||
|
||||
@torch.library.custom_op("cplib::flex_cp_allgather_backward", mutates_args=())
|
||||
def flex_cp_allgather_backward(
|
||||
grad_full_k: torch.Tensor,
|
||||
grad_full_v: torch.Tensor,
|
||||
seq_dim: int,
|
||||
pg_name: c10d.GroupName,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
grad_k = funcol.reduce_scatter_tensor(grad_full_k, "sum", seq_dim, pg_name)
|
||||
if isinstance(grad_k, funcol.AsyncCollectiveTensor):
|
||||
grad_k = grad_k.wait()
|
||||
grad_v = funcol.reduce_scatter_tensor(grad_full_v, "sum", seq_dim, pg_name)
|
||||
if isinstance(grad_v, funcol.AsyncCollectiveTensor):
|
||||
grad_v = grad_v.wait()
|
||||
|
||||
return grad_k, grad_v
|
||||
|
||||
|
||||
@flex_cp_allgather_backward.register_fake
|
||||
def _(
|
||||
grad_full_k: torch.Tensor,
|
||||
grad_full_v: torch.Tensor,
|
||||
seq_dim: int,
|
||||
pg_name: c10d.GroupName,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
shape_k = list(grad_full_k.shape)
|
||||
shape_v = list(grad_full_v.shape)
|
||||
shape_k[seq_dim] //= c10d._get_group_size_by_name(pg_name)
|
||||
shape_v[seq_dim] //= c10d._get_group_size_by_name(pg_name)
|
||||
new_grad_k = torch.empty(
|
||||
shape_k, dtype=grad_full_k.dtype, device=grad_full_k.device
|
||||
)
|
||||
new_grad_v = torch.empty(
|
||||
shape_v, dtype=grad_full_v.dtype, device=grad_full_v.device
|
||||
)
|
||||
return new_grad_k, new_grad_v
|
||||
|
||||
|
||||
def _flex_cp_allgather_backward(
|
||||
ctx: Any, grad_full_k: torch.Tensor, grad_full_v: torch.Tensor
|
||||
) -> tuple[torch.Tensor, torch.Tensor, None, None]:
|
||||
grad_k, grad_v = flex_cp_allgather_backward(
|
||||
grad_full_k, grad_full_v, ctx.seq_dim, ctx.pg_name
|
||||
)
|
||||
return grad_k, grad_v, None, None
|
||||
|
||||
|
||||
def _flex_cp_setup_context(ctx: Any, inputs: Any, output: Any) -> None:
|
||||
_, _, ctx.seq_dim, ctx.pg_name = inputs
|
||||
|
||||
|
||||
flex_cp_allgather.register_autograd(
|
||||
_flex_cp_allgather_backward, setup_context=_flex_cp_setup_context
|
||||
)
|
||||
+478
@@ -0,0 +1,478 @@
|
||||
# this file contains the `_LoadBalancer` class and its family of implementation
|
||||
# for different load-balancing strategies in tensor sharding.
|
||||
import functools
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
from torch.nn.attention.flex_attention import BlockMask
|
||||
|
||||
|
||||
# make it private since it's still a prototype
|
||||
class _LoadBalancer(ABC):
|
||||
@abstractmethod
|
||||
def _generate_indices(self, restore: bool = False) -> Tensor | None:
|
||||
"""
|
||||
Generate indices for load balancing.
|
||||
Args:
|
||||
restore (bool):
|
||||
|
||||
Returns:
|
||||
The generated indices of shape `(1, seq_len)` if the load-balancing is
|
||||
identical within the batch, or `(batch_size, seq_len)` if the load-balancing
|
||||
should vary within the batch.
|
||||
|
||||
Warning:
|
||||
For Multi-Head Attention, we require the masks over the head dimension are identical
|
||||
(i.e. the return value of `_generate_indices()` does not have `heads` dimension).
|
||||
|
||||
Example:
|
||||
Here is the causal mask for attention where q_len == kv_len == 8:
|
||||
KV_index
|
||||
[1, 0, 0, 0, 0, 0, 0, 0]
|
||||
[1, 1, 0, 0, 0, 0, 0, 0]
|
||||
[1, 1, 1, 0, 0, 0, 0, 0]
|
||||
Q_index [1, 1, 1, 1, 0, 0, 0, 0]
|
||||
[1, 1, 1, 1, 1, 0, 0, 0]
|
||||
[1, 1, 1, 1, 1, 1, 0, 0]
|
||||
[1, 1, 1, 1, 1, 1, 1, 0]
|
||||
[1, 1, 1, 1, 1, 1, 1, 1]
|
||||
|
||||
This mask matrix also represents the computation required to compute
|
||||
the masked Q @ K^T by:
|
||||
- mask[i, j] == 1: the computation of Q[i, :] dot K[j, :] is required
|
||||
- mask[i, j] == 0: the computation should be skipped
|
||||
|
||||
Therefore the number of 1s in matrix represents the amount of computation
|
||||
required.
|
||||
|
||||
Assume we want to distribute this Q @ K^T computation to 2 devices, then
|
||||
the matrix is also distributed as:
|
||||
KV_index
|
||||
[1, 0, 0, 0, 0, 0, 0, 0]
|
||||
[1, 1, 0, 0, 0, 0, 0, 0]
|
||||
[1, 1, 1, 0, 0, 0, 0, 0] rank 0
|
||||
[1, 1, 1, 1, 0, 0, 0, 0]
|
||||
Q_index ------------------------
|
||||
[1, 1, 1, 1, 1, 0, 0, 0]
|
||||
[1, 1, 1, 1, 1, 1, 0, 0] rank 1
|
||||
[1, 1, 1, 1, 1, 1, 1, 0]
|
||||
[1, 1, 1, 1, 1, 1, 1, 1]
|
||||
|
||||
An imbalance of computation is observed on these 2 ranks and this could make
|
||||
rank 1 the straggler when performing Context Parallel. In order to balance
|
||||
the computation, we need to rearrange the QKV tensors before sharding in such a
|
||||
way that the result mask matrix is evenly distributed over devices and each
|
||||
rank has the number of 1s as close as possible.
|
||||
|
||||
This method defines the strategy of how to rearrange the QKV tensor for better
|
||||
load-balance:
|
||||
- when `restore == False`, this method returns an indices tensor `rearrange_idx`
|
||||
such that Q[rearrange_idx] is the desired Q tensor after rearranging.
|
||||
- when `restore == True`, this method returns an indices tensor `restore_idx`
|
||||
such that Q[rearrange_idx][restore_idx] == Q, i.e. restoring the rearranged tensor
|
||||
back to the original status before rearranging.
|
||||
"""
|
||||
|
||||
|
||||
class _HeadTailLoadBalancer(_LoadBalancer):
|
||||
def __init__(self, seq_length: int, world_size: int, device: str | torch.device):
|
||||
self.seq_length = seq_length
|
||||
self.world_size = world_size
|
||||
self.device = device
|
||||
|
||||
def _generate_indices(self, restore: bool = False) -> Tensor:
|
||||
"""
|
||||
Generate head-and-tail load balancing indices or restore indices.
|
||||
Args:
|
||||
restore:
|
||||
If True, generate restore indices that map head-and-tail rearranged
|
||||
positions back to original positions. If False, generate load
|
||||
balance indices that rearrange original positions to head-and-tail pattern.
|
||||
|
||||
Returns:
|
||||
The generated indices of shape `(1, seq_len)` because the load-balancing is
|
||||
identical within the batch.
|
||||
|
||||
Warning:
|
||||
For Multi-Head Attention, we require the masks over the head dimension are identical
|
||||
(i.e. the return value of `_generate_indices()` does not have `heads` dimension).
|
||||
|
||||
Example:
|
||||
Here is the causal mask for attention where q_len == kv_len == 8:
|
||||
KV_index
|
||||
[1, 0, 0, 0, 0, 0, 0, 0]
|
||||
[1, 1, 0, 0, 0, 0, 0, 0]
|
||||
[1, 1, 1, 0, 0, 0, 0, 0]
|
||||
Q_index [1, 1, 1, 1, 0, 0, 0, 0]
|
||||
[1, 1, 1, 1, 1, 0, 0, 0]
|
||||
[1, 1, 1, 1, 1, 1, 0, 0]
|
||||
[1, 1, 1, 1, 1, 1, 1, 0]
|
||||
[1, 1, 1, 1, 1, 1, 1, 1]
|
||||
|
||||
Head-tail load-balance strategy rearranges the Q tensor by combining
|
||||
Q[0:k] (on seq dim) and Q[-k:] for rank 0, Q[k:2k] and Q[-2k:-k] for
|
||||
rank 1, and so on. In python code it looks like:
|
||||
|
||||
k = Q.size(0) // (2 * cp_world_size)
|
||||
for rank in range(cp_world_size):
|
||||
reordered_Q[rank * 2 * k : (rank + 1) * 2 * k] = torch.cat(
|
||||
(Q[rank * k : (rank + 1) * k], Q[-(rank + 1) * k : -rank * k])
|
||||
)
|
||||
|
||||
This can also be done by tensor slicing. For the above example, the indices
|
||||
tensor for slicing is:
|
||||
slice_indices = Tensor([0, 7, 1, 6, 2, 5, 3, 4])
|
||||
|
||||
After reordering QKV using the `slice_indices`, the corresponding mask matrix
|
||||
distributing over 2 devices becomes well-balanced:
|
||||
KV_index
|
||||
[1, 0, 0, 0, 0, 0, 0, 0]
|
||||
[1, 1, 1, 1, 1, 1, 1, 1]
|
||||
[1, 1, 0, 0, 0, 0, 0, 0] rank 0
|
||||
[1, 1, 1, 1, 1, 1, 1, 0]
|
||||
Q_index ------------------------
|
||||
[1, 1, 1, 0, 0, 0, 0, 0]
|
||||
[1, 1, 1, 1, 1, 1, 0, 0] rank 1
|
||||
[1, 1, 1, 1, 0, 0, 0, 0]
|
||||
[1, 1, 1, 1, 1, 0, 0, 0]
|
||||
|
||||
To restore the reordering and putting the tensor back, slicing op can do the
|
||||
trick with a `restore_indices` such that:
|
||||
slice_indices[restore_indices] == Tensor([0, 1, 2, ...])
|
||||
|
||||
In this way, `reordered_Q[restore_indices]` will just be the original Q.
|
||||
"""
|
||||
seq_length = self.seq_length
|
||||
world_size = self.world_size
|
||||
if seq_length % (world_size * 2) != 0:
|
||||
raise AssertionError
|
||||
chunk_size = seq_length // (world_size * 2)
|
||||
|
||||
# Split sequence into 2*world_size chunks, then pair chunk r with
|
||||
# chunk (2*world_size - 1 - r) for each rank.
|
||||
indices = torch.arange(seq_length, dtype=torch.int, device=self.device)
|
||||
chunks = indices.view(world_size * 2, chunk_size)
|
||||
head_idx = torch.arange(world_size, device=self.device)
|
||||
tail_idx = 2 * world_size - 1 - head_idx
|
||||
paired = torch.stack([chunks[head_idx], chunks[tail_idx]], dim=1)
|
||||
all_indices_tensor = paired.reshape(-1)
|
||||
|
||||
if restore:
|
||||
all_indices_tensor = torch.argsort(all_indices_tensor)
|
||||
|
||||
return all_indices_tensor.unsqueeze(0) # add batch dim
|
||||
|
||||
|
||||
class _PerDocumentHeadTailLoadBalancer(_LoadBalancer):
|
||||
def __init__(
|
||||
self,
|
||||
seq_length_per_doc: list[list[int]],
|
||||
world_size: int,
|
||||
device: str | torch.device,
|
||||
):
|
||||
"""
|
||||
`seq_length_per_doc` has size (B, seq_len) if the load-balancing should vary
|
||||
within the batch. Otherwise `seq_length_per_doc` should have size (1, seq_len).
|
||||
"""
|
||||
self.seq_length_per_doc = seq_length_per_doc
|
||||
self.world_size = world_size
|
||||
self.device = device
|
||||
|
||||
def _generate_indices(self, restore: bool = False) -> Tensor:
|
||||
"""
|
||||
Generate the per-document head-and-tail rearrange indices so that after rearranging
|
||||
the input is load-balanced in per-document head-and-tail style.
|
||||
|
||||
Args:
|
||||
restore:
|
||||
If True, generate restore indices that map per-document head-and-tail
|
||||
rearranged positions back to original positions. If False, generate load
|
||||
balance indices that rearrange original positions to per-document
|
||||
head-and-tail pattern.
|
||||
|
||||
Returns:
|
||||
The generated indices of shape `(batch_size, seq_len)` if the load-balancing
|
||||
should vary within the batch. Otherwise, it should have shape `(1, seq_len)`.
|
||||
|
||||
Warning:
|
||||
For Multi-Head Attention, we require the masks over the head dimension are identical
|
||||
(i.e. `seq_length_per_doc` must have size (B, seq_len) or (1, seq_len)).
|
||||
|
||||
Example:
|
||||
Here is the document causal mask for attention where q_len == kv_len == 16:
|
||||
KV_index
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
|
||||
[1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
|
||||
[1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
|
||||
[1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
|
||||
[0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
|
||||
[0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
|
||||
[0, 0, 0, 0, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0]
|
||||
Q_index [0, 0, 0, 0, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0]
|
||||
[0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0]
|
||||
[0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0]
|
||||
[0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0]
|
||||
[0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0]
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0]
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0]
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 0]
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1]
|
||||
|
||||
The per-document head-and-tail load-balancer will apply head-and-tail
|
||||
reordering within each document. After load-balancing for context-parallel
|
||||
on 2 devices, the above mask matrix will look like this:
|
||||
KV_index
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
|
||||
[1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
|
||||
[0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
|
||||
[0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
|
||||
[0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0]
|
||||
[0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0]
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0]
|
||||
Q_index [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1]
|
||||
------------------------------------------------
|
||||
[1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
|
||||
[1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
|
||||
[0, 0, 0, 0, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0]
|
||||
[0, 0, 0, 0, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0]
|
||||
[0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0]
|
||||
[0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0]
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0]
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 0]
|
||||
"""
|
||||
return torch.stack(
|
||||
[
|
||||
self._generate_indices_for_batch(seq_lengths, restore)
|
||||
for seq_lengths in self.seq_length_per_doc
|
||||
]
|
||||
)
|
||||
|
||||
def _generate_indices_for_batch(self, seq_length_per_doc, restore) -> Tensor: # type: ignore[no-untyped-def]
|
||||
world_size = self.world_size
|
||||
device = self.device
|
||||
if not all(
|
||||
seq_length % (2 * world_size) == 0 for seq_length in seq_length_per_doc
|
||||
):
|
||||
raise AssertionError
|
||||
chunk_length_per_doc = [
|
||||
seq_length // (2 * world_size) for seq_length in seq_length_per_doc
|
||||
]
|
||||
|
||||
indices = []
|
||||
document_start_idx = 0
|
||||
for seq_length, chunk_length in zip(seq_length_per_doc, chunk_length_per_doc):
|
||||
# Generate the indices for the current document
|
||||
for rank in range(world_size):
|
||||
head_chunk_start_idx = document_start_idx + chunk_length * rank
|
||||
tail_chunk_end_idx = document_start_idx + chunk_length * (
|
||||
2 * world_size - rank
|
||||
)
|
||||
indices.append(
|
||||
torch.arange(
|
||||
head_chunk_start_idx,
|
||||
head_chunk_start_idx + chunk_length,
|
||||
device=device,
|
||||
)
|
||||
)
|
||||
indices.append(
|
||||
torch.arange(
|
||||
tail_chunk_end_idx - chunk_length,
|
||||
tail_chunk_end_idx,
|
||||
device=device,
|
||||
)
|
||||
)
|
||||
|
||||
document_start_idx += seq_length
|
||||
|
||||
indices_tensor = torch.cat(indices)
|
||||
if restore:
|
||||
indices_tensor = torch.argsort(indices_tensor)
|
||||
|
||||
return indices_tensor
|
||||
|
||||
|
||||
class _PTRRLoadBalancer(_LoadBalancer):
|
||||
"""
|
||||
Processing-Time based Round-Robin (PTRR) load balancer. This load balancer should
|
||||
only be used for flex_attention() since it leverages `BlockMask`.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
block_mask: BlockMask,
|
||||
world_size: int,
|
||||
):
|
||||
"""
|
||||
`block_mask` must have shape (B, 1, seq_len, seq_len) or (1, 1, seq_len, seq_len).
|
||||
"""
|
||||
self.block_mask = block_mask
|
||||
self.world_size = world_size
|
||||
|
||||
@staticmethod
|
||||
def ptrr_scheduling(process_time: Tensor, group_size: int) -> Tensor:
|
||||
"""
|
||||
Separate the tasks into `group_size` groups using PTRR scheduling.
|
||||
process_time:
|
||||
1D tensor of size n, where n is the number of tasks. The value
|
||||
is the process time of the task. Size `n` must be divisible by
|
||||
`group_size`.
|
||||
group_size:
|
||||
the number of groups
|
||||
|
||||
Returns:
|
||||
tasks_in_group (list[list[int]]):
|
||||
A collection of list[int] and each list should have size `n // group_size`
|
||||
(`group_size` lists in total). Each element is an index in the input
|
||||
`process_time` (i.e. [0, len(process_time) - 1]).
|
||||
|
||||
Example:
|
||||
process_time = [9, 14, 2, 20, 10, 15, 8, 14, 16, 19, 15, 3, 12, 1, 12, 10]
|
||||
tasks_in_group = [
|
||||
[3, 12, 13, 14], # values = [1, 12, 12, 20], sum = 45
|
||||
[2, 4, 7, 9], # values = [2, 10, 14, 19], sum = 45
|
||||
[1, 8, 11, 15], # values = [14, 16, 3, 10], sum = 43
|
||||
[0, 5, 6, 10] # values = [9, 15, 8, 15], sum = 47
|
||||
]
|
||||
"""
|
||||
if process_time.ndim != 1:
|
||||
raise AssertionError
|
||||
|
||||
num_tasks = process_time.size(0)
|
||||
|
||||
if num_tasks % group_size != 0:
|
||||
raise NotImplementedError(
|
||||
f"num_tasks {num_tasks} must be divisible by group_size {group_size}"
|
||||
)
|
||||
|
||||
device = process_time.device
|
||||
_, sorted_indices_descending = torch.sort(
|
||||
process_time, descending=True, stable=True
|
||||
) # if process time is tied, the order is preserved
|
||||
sorted_indices_descending_reversed = torch.flip(
|
||||
sorted_indices_descending.view(-1, group_size), dims=[1]
|
||||
).view(-1)
|
||||
tasks_in_group = torch.where(
|
||||
torch.arange(num_tasks, device=device) // group_size % 2 == 0,
|
||||
sorted_indices_descending,
|
||||
sorted_indices_descending_reversed,
|
||||
)
|
||||
tasks_in_group = tasks_in_group.view(-1, group_size).transpose(
|
||||
0, 1
|
||||
) # (group_size, n // group_size)
|
||||
|
||||
# sort each group. This step should not have impact on correctness
|
||||
# nor execution run time, but it helps users visualize the mask
|
||||
tasks_in_group, _ = torch.sort(tasks_in_group, dim=1)
|
||||
return tasks_in_group
|
||||
|
||||
def _generate_indices(self, restore: bool = False) -> Tensor:
|
||||
"""
|
||||
Generate the PTRR reorder indices of shape `(1, seq_len)` or `(batch_size, seq_len)`.
|
||||
|
||||
Args:
|
||||
restore:
|
||||
If True, generate restore indices that map Processing-Time based Round-Robin
|
||||
(PTRR) rearranged positions back to original positions. If False, generate
|
||||
load balance indices that rearrange original positions to PTRR pattern.
|
||||
|
||||
Returns:
|
||||
The generated indices of shape `(1, seq_len)` if the load-balancing is
|
||||
identical within the batch (i.e. `BlockMask.shape[0] == 1`), or
|
||||
`(batch_size, seq_len)` if the load-balancing should vary within the batch.
|
||||
|
||||
Warning:
|
||||
For Multi-Head Attention, we require the masks over the head dimension are identical
|
||||
(i.e. `self.block_mask` must have shape (B, 1, seq_len, seq_len) or (1, 1, seq_len, seq_len)).
|
||||
|
||||
Example:
|
||||
Here is the document causal mask for attention whereq_len == kv_len == 16 * BLOCK_SIZE
|
||||
(each entry is a block):
|
||||
KV_index
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] -> row value = 1
|
||||
[1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] -> row value = 2
|
||||
[1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] -> row value = 3
|
||||
[1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] -> row value = 4
|
||||
[0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] -> row value = 1
|
||||
[0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] -> row value = 2
|
||||
[0, 0, 0, 0, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0] -> row value = 3
|
||||
Q_index [0, 0, 0, 0, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0] -> row value = 4
|
||||
[0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0] -> row value = 5
|
||||
[0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0] -> row value = 6
|
||||
[0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0] -> row value = 7
|
||||
[0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0] -> row value = 8
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0] -> row value = 1
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0] -> row value = 2
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 0] -> row value = 3
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1] -> row value = 4
|
||||
|
||||
The reorder indices will be: [2, 3, 5, 6, 8, 11, 12, 13, 0, 1, 4, 7, 9, 10, 14, 15] and
|
||||
the mask matrix will look like:
|
||||
KV_index
|
||||
[1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] -> row value = 3
|
||||
[1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] -> row value = 4
|
||||
[0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] -> row value = 2
|
||||
[0, 0, 0, 0, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0] -> row value = 3
|
||||
[0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0] -> row value = 5 rank 0 (sum=28)
|
||||
[0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0] -> row value = 8
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0] -> row value = 1
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0] -> row value = 2
|
||||
------------------------------------------------
|
||||
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] -> row value = 1
|
||||
[1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] -> row value = 2
|
||||
[0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] -> row value = 1
|
||||
[0, 0, 0, 0, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0] -> row value = 4
|
||||
[0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0] -> row value = 6 rank 1 (sum=28)
|
||||
[0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0] -> row value = 7
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 0] -> row value = 3
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1] -> row value = 4
|
||||
"""
|
||||
block_mask = self.block_mask
|
||||
kv_num_blocks = block_mask.kv_num_blocks
|
||||
full_kv_num_blocks = block_mask.full_kv_num_blocks
|
||||
non_sparse_kv_num_blocks = (
|
||||
kv_num_blocks + full_kv_num_blocks
|
||||
if full_kv_num_blocks is not None
|
||||
else kv_num_blocks
|
||||
)
|
||||
B, H, Q = non_sparse_kv_num_blocks.shape
|
||||
# requirement: the masking is identical across heads (i.e. H == 1 in BlockMask)
|
||||
non_sparse_kv_num_blocks = non_sparse_kv_num_blocks.view(-1, Q) # (B, Q_BLK)
|
||||
|
||||
batch_ptrr = torch.vmap(
|
||||
functools.partial(
|
||||
_PTRRLoadBalancer.ptrr_scheduling,
|
||||
group_size=self.world_size,
|
||||
)
|
||||
)
|
||||
ptrr_indices = batch_ptrr(
|
||||
non_sparse_kv_num_blocks
|
||||
) # (B, group_size, num_blks_in_group)
|
||||
ptrr_indices = ptrr_indices.reshape(B, -1) # (B, num_blocks)
|
||||
|
||||
# NOTE: only support the case where the qkv block size are equal
|
||||
q_blk_size, kv_blk_size = block_mask.BLOCK_SIZE
|
||||
if q_blk_size != kv_blk_size:
|
||||
raise AssertionError("for now only support q_blk_size == kv_blk_size")
|
||||
|
||||
indices = torch.arange(
|
||||
q_blk_size * ptrr_indices.size(1), device=ptrr_indices.device
|
||||
).view(-1, q_blk_size) # (NUM_BLOCKS, BLOCK_SIZE)
|
||||
indices = indices[ptrr_indices].view(B, -1) # (B, qkv_size)
|
||||
|
||||
if restore:
|
||||
# pyrefly: ignore[missing-argument]
|
||||
indices = torch.vmap(torch.argsort)(indices)
|
||||
|
||||
return indices
|
||||
|
||||
|
||||
def _create_default_load_balancer(
|
||||
seq_length: int, world_size: int, device: str | torch.device
|
||||
) -> _LoadBalancer | None:
|
||||
from ._attention import _cp_options
|
||||
|
||||
if _cp_options.enable_load_balance:
|
||||
return _HeadTailLoadBalancer(seq_length, world_size, device)
|
||||
else:
|
||||
return None
|
||||
+408
@@ -0,0 +1,408 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates
|
||||
"""
|
||||
Context Parallelism sharding rules for scaled_dot_product attention operators.
|
||||
|
||||
The sharding rules for CP cannot be embedded by default because Shard(2) is not
|
||||
a valid sharding for SDPA without CP enabled. This module provides utilities to
|
||||
dynamically install Shard(2) sharding rules when CP is activated.
|
||||
"""
|
||||
|
||||
from contextlib import contextmanager
|
||||
|
||||
import torch
|
||||
from torch.distributed.tensor._op_schema import (
|
||||
OpSchema,
|
||||
OpStrategy,
|
||||
PlacementList,
|
||||
RuntimeSchemaInfo,
|
||||
)
|
||||
from torch.distributed.tensor._ops.utils import (
|
||||
expand_to_full_mesh_op_strategy,
|
||||
register_op_strategy,
|
||||
)
|
||||
from torch.distributed.tensor.debug import (
|
||||
_clear_fast_path_sharding_prop_cache,
|
||||
_clear_python_sharding_prop_cache,
|
||||
)
|
||||
from torch.distributed.tensor.placement_types import Replicate, Shard
|
||||
|
||||
|
||||
aten = torch.ops.aten
|
||||
|
||||
SEQ_DIM = 2
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _op_strategy_context(op_overload, strategy_func, schema_info=None):
|
||||
"""
|
||||
Context manager for setting and clearing op strategies for Context Parallelism.
|
||||
|
||||
Args:
|
||||
op_overload: The operator overload to set or clear the strategy for.
|
||||
strategy_func: The strategy function to set for the operator overload.
|
||||
schema_info: Optional schema information for the operator overload.
|
||||
|
||||
Yields:
|
||||
None
|
||||
"""
|
||||
from torch.distributed.tensor import DTensor
|
||||
|
||||
propagator = DTensor._op_dispatcher.sharding_propagator
|
||||
_origin_op_strategy_funcs = None
|
||||
_origin_op_strategy_schema = None
|
||||
try:
|
||||
# Save original strategy if exists
|
||||
if op_overload in propagator.op_strategy_funcs:
|
||||
_origin_op_strategy_funcs = propagator.op_strategy_funcs[op_overload]
|
||||
if op_overload in propagator.op_to_schema_info:
|
||||
_origin_op_strategy_schema = propagator.op_to_schema_info[op_overload]
|
||||
|
||||
# Register the new op strategy
|
||||
register_op_strategy(op_overload, schema_info=schema_info)(strategy_func)
|
||||
yield (_origin_op_strategy_funcs, _origin_op_strategy_schema)
|
||||
finally:
|
||||
# Restore original strategy
|
||||
if _origin_op_strategy_funcs is None:
|
||||
if op_overload in propagator.op_strategy_funcs:
|
||||
del propagator.op_strategy_funcs[op_overload]
|
||||
else:
|
||||
propagator.op_strategy_funcs[op_overload] = _origin_op_strategy_funcs
|
||||
|
||||
if _origin_op_strategy_schema is None:
|
||||
if op_overload in propagator.op_to_schema_info:
|
||||
del propagator.op_to_schema_info[op_overload]
|
||||
else:
|
||||
propagator.op_to_schema_info[op_overload] = _origin_op_strategy_schema
|
||||
|
||||
# Ideally, we should clear the cache, but it is too expensive.
|
||||
# _clear_python_sharding_prop_cache()
|
||||
# _clear_fast_path_sharding_prop_cache()
|
||||
|
||||
|
||||
# ==================== Flash Attention Strategies ====================
|
||||
|
||||
|
||||
def _scaled_dot_product_flash_attention_cp_strategy(op_schema: OpSchema) -> OpStrategy:
|
||||
"""
|
||||
Strategy for flash attention forward with Context Parallelism support.
|
||||
This includes the base strategies plus CP-specific sequence dimension sharding.
|
||||
"""
|
||||
# Import here to avoid circular dependency
|
||||
from torch.distributed.tensor._ops._matrix_ops import (
|
||||
_scaled_dot_product_flash_attention_base_strategies,
|
||||
)
|
||||
|
||||
# Get the base strategies (without CP modifications)
|
||||
mesh = op_schema.get_mesh_from_args()
|
||||
single_mesh_dim_strategies = _scaled_dot_product_flash_attention_base_strategies(
|
||||
op_schema
|
||||
)
|
||||
|
||||
# Add Context Parallelism strategy: shards on the sequence dim
|
||||
return_debug_mask = len(op_schema.args_schema) >= 6 and op_schema.args_schema[5]
|
||||
debug_attn_mask_sharding = Shard(SEQ_DIM) if return_debug_mask else Replicate()
|
||||
|
||||
cp_strategy: PlacementList = [
|
||||
Shard(SEQ_DIM), # output
|
||||
Shard(SEQ_DIM), # logsumexp
|
||||
None, # cum_seq_q
|
||||
None, # cum_seq_k
|
||||
None, # max_q
|
||||
None, # max_k
|
||||
Replicate(), # rng_state
|
||||
None, # unused
|
||||
debug_attn_mask_sharding, # debugattn
|
||||
Shard(SEQ_DIM), # q
|
||||
Shard(SEQ_DIM), # k
|
||||
Shard(SEQ_DIM), # v
|
||||
]
|
||||
single_mesh_dim_strategies.append(cp_strategy)
|
||||
|
||||
return expand_to_full_mesh_op_strategy(
|
||||
mesh, op_schema, single_mesh_dim_strategies, input_index=9
|
||||
)
|
||||
|
||||
|
||||
def _scaled_dot_product_flash_attention_backward_cp_strategy(
|
||||
op_schema: OpSchema,
|
||||
) -> OpStrategy:
|
||||
"""
|
||||
Strategy for flash attention backward with Context Parallelism support.
|
||||
"""
|
||||
from torch.distributed.tensor._ops._matrix_ops import (
|
||||
_scaled_dot_product_flash_attention_backward_base_strategies,
|
||||
)
|
||||
|
||||
mesh = op_schema.get_mesh_from_args(validate=False)
|
||||
single_mesh_dim_strategies = (
|
||||
_scaled_dot_product_flash_attention_backward_base_strategies(op_schema)
|
||||
)
|
||||
|
||||
tensor_input_indices = [
|
||||
i
|
||||
for i, arg_spec in enumerate(op_schema.args_schema)
|
||||
if isinstance(arg_spec, OpStrategy)
|
||||
]
|
||||
num_tensor_inputs = len(tensor_input_indices)
|
||||
|
||||
# Context Parallelism: shards on the sequence dim
|
||||
cp_strategy: PlacementList = [
|
||||
Shard(SEQ_DIM), # grad_q
|
||||
Shard(SEQ_DIM), # grad_k
|
||||
Shard(SEQ_DIM), # grad_v
|
||||
Shard(SEQ_DIM), # grad_output
|
||||
Shard(SEQ_DIM), # q
|
||||
Shard(SEQ_DIM), # k
|
||||
Shard(SEQ_DIM), # v
|
||||
Shard(SEQ_DIM), # output
|
||||
Shard(SEQ_DIM), # logsumexp
|
||||
]
|
||||
cp_strategy.extend([Replicate()] * (num_tensor_inputs - 6))
|
||||
single_mesh_dim_strategies.append(cp_strategy)
|
||||
|
||||
return expand_to_full_mesh_op_strategy(
|
||||
mesh, op_schema, single_mesh_dim_strategies, input_index=3
|
||||
)
|
||||
|
||||
|
||||
# ==================== Efficient Attention Strategies ====================
|
||||
|
||||
|
||||
def _scaled_dot_product_efficient_attention_cp_strategy(
|
||||
op_schema: OpSchema,
|
||||
) -> OpStrategy:
|
||||
"""
|
||||
Strategy for efficient attention forward with Context Parallelism support.
|
||||
"""
|
||||
from torch.distributed.tensor._ops._matrix_ops import (
|
||||
_scaled_dot_product_efficient_attention_base_strategies,
|
||||
)
|
||||
|
||||
mesh = op_schema.get_mesh_from_args()
|
||||
single_mesh_dim_strategies = (
|
||||
_scaled_dot_product_efficient_attention_base_strategies(op_schema)
|
||||
)
|
||||
|
||||
# Add Context Parallelism strategy
|
||||
has_attn_bias = op_schema.args_schema[3] is not None
|
||||
|
||||
cp_strategy: PlacementList = [
|
||||
Shard(SEQ_DIM), # output
|
||||
Shard(SEQ_DIM), # logsumexp
|
||||
None, # philox_seed
|
||||
None, # philox_offset
|
||||
Shard(SEQ_DIM), # q
|
||||
Shard(SEQ_DIM), # k
|
||||
Shard(SEQ_DIM), # v
|
||||
]
|
||||
if has_attn_bias:
|
||||
cp_strategy.append(Replicate()) # attn bias - not sharded for CP
|
||||
single_mesh_dim_strategies.append(cp_strategy)
|
||||
|
||||
return expand_to_full_mesh_op_strategy(
|
||||
mesh, op_schema, single_mesh_dim_strategies, input_index=4
|
||||
)
|
||||
|
||||
|
||||
def _scaled_dot_product_efficient_attention_backward_cp_strategy(
|
||||
op_schema: OpSchema,
|
||||
) -> OpStrategy:
|
||||
"""
|
||||
Strategy for efficient attention backward with Context Parallelism support.
|
||||
"""
|
||||
from torch.distributed.tensor._ops._matrix_ops import (
|
||||
_scaled_dot_product_efficient_attention_backward_base_strategies,
|
||||
)
|
||||
|
||||
mesh = op_schema.get_mesh_from_args(validate=False)
|
||||
single_mesh_dim_strategies = (
|
||||
_scaled_dot_product_efficient_attention_backward_base_strategies(op_schema)
|
||||
)
|
||||
|
||||
has_attn_bias = op_schema.args_schema[4] is not None
|
||||
|
||||
# Context Parallelism: shards on the sequence dim
|
||||
cp_strategy: PlacementList = [
|
||||
Shard(SEQ_DIM), # grad_q
|
||||
Shard(SEQ_DIM), # grad_k
|
||||
Shard(SEQ_DIM), # grad_v
|
||||
Shard(1) if has_attn_bias else None, # grad_bias
|
||||
Shard(SEQ_DIM), # grad_output
|
||||
Shard(SEQ_DIM), # q
|
||||
Shard(SEQ_DIM), # k
|
||||
Shard(SEQ_DIM), # v
|
||||
Shard(SEQ_DIM), # output
|
||||
Shard(SEQ_DIM), # logsumexp
|
||||
]
|
||||
if has_attn_bias:
|
||||
cp_strategy.insert(8, Shard(1)) # attn_bias input
|
||||
cp_strategy.extend([Replicate(), Replicate()])
|
||||
single_mesh_dim_strategies.append(cp_strategy)
|
||||
|
||||
return expand_to_full_mesh_op_strategy(
|
||||
mesh, op_schema, single_mesh_dim_strategies, input_index=4
|
||||
)
|
||||
|
||||
|
||||
# ==================== cuDNN Attention Strategies ====================
|
||||
|
||||
|
||||
def _scaled_dot_product_cudnn_attention_cp_strategy(op_schema: OpSchema) -> OpStrategy:
|
||||
"""
|
||||
Strategy for cudnn attention forward with Context Parallelism support.
|
||||
"""
|
||||
from torch.distributed.tensor._ops._matrix_ops import (
|
||||
_scaled_dot_product_cudnn_attention_base_strategies,
|
||||
)
|
||||
|
||||
mesh = op_schema.get_mesh_from_args()
|
||||
single_mesh_dim_strategies = _scaled_dot_product_cudnn_attention_base_strategies(
|
||||
op_schema
|
||||
)
|
||||
|
||||
(
|
||||
query_strategy,
|
||||
_,
|
||||
_,
|
||||
attn_bias_strategy,
|
||||
compute_log_sumexp,
|
||||
*rest_args,
|
||||
) = op_schema.args_schema
|
||||
return_debug_mask = len(op_schema.args_schema) >= 8 and rest_args[2]
|
||||
has_attn_bias = attn_bias_strategy is not None
|
||||
|
||||
# Context Parallelism: shards on the sequence dim
|
||||
logsumexp_sharding = Shard(SEQ_DIM) if compute_log_sumexp else Replicate()
|
||||
debug_attn_mask_sharding = Shard(SEQ_DIM) if return_debug_mask else None
|
||||
|
||||
cp_strategy: PlacementList = [
|
||||
Shard(SEQ_DIM), # output
|
||||
logsumexp_sharding, # logsumexp
|
||||
None, # cum_seq_q
|
||||
None, # cum_seq_k
|
||||
None, # max_q
|
||||
None, # max_k
|
||||
None, # philox_seed
|
||||
None, # philox_offset
|
||||
debug_attn_mask_sharding, # debug_attn_mask
|
||||
Shard(SEQ_DIM), # q
|
||||
Shard(SEQ_DIM), # k
|
||||
Shard(SEQ_DIM), # v
|
||||
]
|
||||
if has_attn_bias:
|
||||
cp_strategy.append(Replicate()) # attn_bias - not sharded for CP
|
||||
single_mesh_dim_strategies.append(cp_strategy)
|
||||
|
||||
return expand_to_full_mesh_op_strategy(
|
||||
mesh, op_schema, single_mesh_dim_strategies, input_index=9
|
||||
)
|
||||
|
||||
|
||||
def _scaled_dot_product_cudnn_attention_backward_cp_strategy(
|
||||
op_schema: OpSchema,
|
||||
) -> OpStrategy:
|
||||
"""
|
||||
Strategy for cudnn attention backward with Context Parallelism support.
|
||||
"""
|
||||
from torch.distributed.tensor._ops._matrix_ops import (
|
||||
_scaled_dot_product_cudnn_attention_backward_base_strategies,
|
||||
)
|
||||
|
||||
mesh = op_schema.get_mesh_from_args(validate=False)
|
||||
single_mesh_dim_strategies = (
|
||||
_scaled_dot_product_cudnn_attention_backward_base_strategies(op_schema)
|
||||
)
|
||||
|
||||
has_attn_bias = op_schema.args_schema[8] is not None
|
||||
has_scale = len(op_schema.args_schema) >= 16 and False
|
||||
|
||||
# Context Parallelism: shards on the sequence dim
|
||||
cp_sharding_gout: PlacementList = [Shard(SEQ_DIM)] * 3 # grad_q, grad_k, grad_v
|
||||
cp_sharding_ginp: PlacementList = [
|
||||
Shard(SEQ_DIM)
|
||||
] * 6 # grad_output, q, k, v, output, logsumexp
|
||||
cp_sharding_ginp += [Replicate()] * 2 # philox_seed, philox_offset
|
||||
cp_sharding_ginp += [Shard(SEQ_DIM) if has_attn_bias else None] # attn_bias
|
||||
cp_sharding_ginp += [
|
||||
None
|
||||
] * 6 # cum_seq_q, cum_seq_k, max_q, max_k, dropout_p, is_causal
|
||||
if has_scale:
|
||||
cp_sharding_ginp.append(None)
|
||||
|
||||
cp_sharding = cp_sharding_gout + cp_sharding_ginp
|
||||
single_mesh_dim_strategies.append(cp_sharding)
|
||||
|
||||
return expand_to_full_mesh_op_strategy(
|
||||
mesh, op_schema, single_mesh_dim_strategies, input_index=3
|
||||
)
|
||||
|
||||
|
||||
# Store context managers and original strategies
|
||||
_cp_strategy_contexts = {}
|
||||
_original_strategies = {}
|
||||
|
||||
|
||||
def register_cp_sharding_rules():
|
||||
"""Register Context Parallelism sharding rules for all scaled_dot_product ops."""
|
||||
global _cp_strategy_contexts, _original_strategies
|
||||
|
||||
# If already registered, don't register again
|
||||
if _cp_strategy_contexts:
|
||||
return
|
||||
|
||||
# Define ops and their corresponding CP strategy functions
|
||||
cp_strategies = [
|
||||
(
|
||||
aten._scaled_dot_product_flash_attention.default,
|
||||
_scaled_dot_product_flash_attention_cp_strategy,
|
||||
RuntimeSchemaInfo(5),
|
||||
),
|
||||
(
|
||||
aten._scaled_dot_product_flash_attention_backward.default,
|
||||
_scaled_dot_product_flash_attention_backward_cp_strategy,
|
||||
None,
|
||||
),
|
||||
(
|
||||
aten._scaled_dot_product_efficient_attention.default,
|
||||
_scaled_dot_product_efficient_attention_cp_strategy,
|
||||
RuntimeSchemaInfo(4),
|
||||
),
|
||||
(
|
||||
aten._scaled_dot_product_efficient_attention_backward.default,
|
||||
_scaled_dot_product_efficient_attention_backward_cp_strategy,
|
||||
None,
|
||||
),
|
||||
(
|
||||
aten._scaled_dot_product_cudnn_attention.default,
|
||||
_scaled_dot_product_cudnn_attention_cp_strategy,
|
||||
RuntimeSchemaInfo(4),
|
||||
),
|
||||
(
|
||||
aten._scaled_dot_product_cudnn_attention_backward.default,
|
||||
_scaled_dot_product_cudnn_attention_backward_cp_strategy,
|
||||
None,
|
||||
),
|
||||
]
|
||||
|
||||
# Register each strategy
|
||||
for op_overload, strategy_func, schema_info in cp_strategies:
|
||||
ctx = _op_strategy_context(op_overload, strategy_func, schema_info)
|
||||
orig_funcs, orig_schema = ctx.__enter__()
|
||||
_cp_strategy_contexts[op_overload] = ctx
|
||||
_original_strategies[op_overload] = (orig_funcs, orig_schema)
|
||||
|
||||
|
||||
def unregister_cp_sharding_rules(clear_the_cache=False):
|
||||
"""Unregister Context Parallelism sharding rules and restore original strategies."""
|
||||
global _cp_strategy_contexts, _original_strategies
|
||||
|
||||
# Exit all context managers
|
||||
for ctx in _cp_strategy_contexts.values():
|
||||
ctx.__exit__(None, None, None)
|
||||
|
||||
if clear_the_cache:
|
||||
_clear_fast_path_sharding_prop_cache()
|
||||
_clear_python_sharding_prop_cache()
|
||||
|
||||
_cp_strategy_contexts = {}
|
||||
_original_strategies = {}
|
||||
+285
@@ -0,0 +1,285 @@
|
||||
# mypy: allow-untyped-defs
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates
|
||||
import functools
|
||||
from collections.abc import Callable, Sequence
|
||||
|
||||
import torch
|
||||
from torch.distributed._functional_collectives import AsyncCollectiveTensor
|
||||
from torch.distributed.tensor import DeviceMesh, DTensor
|
||||
from torch.distributed.tensor.placement_types import Placement
|
||||
|
||||
|
||||
try:
|
||||
from torch.utils import _cxx_pytree as pytree
|
||||
except ImportError:
|
||||
from torch.utils import _pytree as pytree # type: ignore[no-redef]
|
||||
|
||||
|
||||
__all__ = ["local_map"]
|
||||
|
||||
PlacementType = Sequence[Placement] | None
|
||||
InputPlacements = tuple[PlacementType, ...] | None
|
||||
OutputPlacements = PlacementType | tuple[PlacementType, ...]
|
||||
|
||||
|
||||
def local_map(
|
||||
func: Callable | None = None,
|
||||
out_placements: OutputPlacements = None,
|
||||
in_placements: InputPlacements = None,
|
||||
in_grad_placements: InputPlacements = None,
|
||||
device_mesh: DeviceMesh | None = None,
|
||||
*,
|
||||
redistribute_inputs: bool = False,
|
||||
):
|
||||
"""
|
||||
:meth:`local_map` is an experimental API that allows users to pass :class:`DTensor` s
|
||||
to a function that is written to be applied on ``torch.Tensor`` s. It is done by extracting
|
||||
the local components of :class:`DTensor`, call the function, and wrap the outputs to
|
||||
:class:`DTensor` according to the ``out_placements``.
|
||||
|
||||
Args:
|
||||
func (Callable): the function to be applied on each local shard of
|
||||
:class:`DTensor` s.
|
||||
out_placements (Union[`PlacementType`, Tuple[`PlacementType`, ...]]):
|
||||
the desired placements of the :class:`DTensor` s in ``func``'s flattened output.
|
||||
If the flattened ``output`` is a single value, the ``out_placements`` should be
|
||||
of type `PlacementType`. Otherwise if the flattened ``output`` has multiple
|
||||
values, the ``out_placements`` should be a tuple of `PlacementType` values 1:1
|
||||
mapping to the flattened ``output``.
|
||||
Besides, for :class:`Tensor` output, we use `PlacementType` as its
|
||||
placements (a `Tuple[Placement]` value). For non-Tensor output, the `PlacementType`
|
||||
should be `None`.
|
||||
Note that the only exception is when no :class:`DTensor` argument is passed
|
||||
in. In this case, even if `out_placements` is not `None`, the result function
|
||||
should ignore the desired placements because the function is not running with
|
||||
:class:`DTensor` s.
|
||||
in_placements (Tuple[`PlacementType`, ...], optional):
|
||||
the required placements of the :class:`DTensor` s in the flattened inputs of ``func``.
|
||||
If ``in_placements`` is specified, :meth:`local_map` would examine whether the
|
||||
placements of each :class:`DTensor` argument is the same as the required
|
||||
placements or not. If the placements are not the same and
|
||||
``redistribute_inputs`` is ``False``, an exception will be raised. Otherwise if
|
||||
``redistribute_inputs`` is ``True``, the argument will be first redistributed to
|
||||
the required sharding placements before passing its local tensor to ``func``.
|
||||
The only exception is when required placements are not ``None`` and the
|
||||
argument is a :class:`torch.Tensor`. In this case, the placements examination
|
||||
will be skipped and the argument will be directly passed to ``func``.
|
||||
If ``in_placements`` is ``None``, no placements examination will be performed.
|
||||
Default: None
|
||||
in_grad_placements (Tuple[`PlacementType`, ...], optional):
|
||||
the placements hint of the :class:`DTensor` s gradient corresponds
|
||||
to the flattened input DTensor. This argument is the hint that user
|
||||
can give to :meth:`to_local` in case the gradient layout of the
|
||||
local tensor input does not match its :class:`DTensor` input layout.
|
||||
If not specified, we will assume the gradient layout of the local
|
||||
tensor input remains the same as the original :class:`DTensor` input
|
||||
and use that for gradient computation. Default: None.
|
||||
device_mesh (:class:`DeviceMesh`, optional):
|
||||
the device mesh that the output :class:`DTensor` s are placed on. If not
|
||||
specified, this will be inferred from the first input :class:`DTensor`'s device
|
||||
mesh. Default: None.
|
||||
|
||||
Keyword Args:
|
||||
redistribute_inputs (bool, optional):
|
||||
the bool value indicating whether to reshard the input :class:`DTensor` s when
|
||||
their placements are different from the required input placements. If this
|
||||
value is ``False`` and some :class:`DTensor` input has a different placement,
|
||||
an exception will be raised. Default: False.
|
||||
|
||||
Returns:
|
||||
A ``Callable`` that applies ``func`` to each local shard of the input :class:`DTensor`
|
||||
and returns a :class:`DTensor` constructed from the return value of ``func``.
|
||||
|
||||
Raises:
|
||||
AssertionError: For any non-DTensor output, we require its corresponding
|
||||
output placement in ``out_placements`` be None. An AssertionError will be raised
|
||||
if this is not the case.
|
||||
|
||||
ValueError: If ``redistribute_inputs=False`` but the input :class:`DTensor` needs
|
||||
a redistribution according to ``in_placements``.
|
||||
|
||||
Example:
|
||||
>>> # xdoctest: +SKIP("distributed")
|
||||
>>> def mm_allreduce_forward(device_mesh, W, X):
|
||||
>>> partial_sum_tensor = torch.mm(W, X)
|
||||
>>> reduced_tensor = funcol.all_reduce(partial_sum_tensor, "sum", device_mesh)
|
||||
>>> return reduced_tensor
|
||||
>>>
|
||||
>>> W = torch.randn(12, 8, requires_grad=False)
|
||||
>>> X = torch.randn(8, 16, requires_grad=False)
|
||||
>>> Y = torch.mm(W, X)
|
||||
>>> row_wise = [Shard(0)] # row-wise sharding placements on 1-d mesh
|
||||
>>> col_wise = [Shard(1)] # col-wise sharding placements on 1-d mesh
|
||||
>>>
|
||||
>>> # local_mm_allreduce_forward is the function wrapped with DTensor/Tensor conversion
|
||||
>>> local_mm_allreduce_forward = local_map(
|
||||
>>> mm_allreduce_forward,
|
||||
>>> out_placements=[Replicate()],
|
||||
>>> in_placements=[col_wise, row_wise],
|
||||
>>> device_mesh=device_mesh,
|
||||
>>> )
|
||||
>>>
|
||||
>>> W_dt = distribute_tensor(
|
||||
... W, device_mesh, (col_wise)
|
||||
... ) # col-wisely sharded W tensor
|
||||
>>> X_dt = distribute_tensor(
|
||||
... X, device_mesh, (row_wise)
|
||||
... ) # row-wisely sharded X tensor
|
||||
>>> Y_dt = local_mm_allreduce_forward(
|
||||
... device_mesh, W_dt, X_dt
|
||||
... ) # apply local_mm_allreduce_forward to DTensors
|
||||
|
||||
.. note:: This API is currently experimental and subject to change
|
||||
"""
|
||||
|
||||
if func is None:
|
||||
# decorator mode
|
||||
def decorated(func):
|
||||
return local_map(
|
||||
func=func,
|
||||
out_placements=out_placements,
|
||||
in_placements=in_placements,
|
||||
in_grad_placements=in_grad_placements,
|
||||
device_mesh=device_mesh,
|
||||
redistribute_inputs=redistribute_inputs,
|
||||
)
|
||||
|
||||
return decorated
|
||||
|
||||
return functools.partial(
|
||||
_local_map_wrapped,
|
||||
func,
|
||||
out_placements,
|
||||
in_placements,
|
||||
in_grad_placements,
|
||||
device_mesh,
|
||||
redistribute_inputs,
|
||||
)
|
||||
|
||||
|
||||
def _local_map_wrapped(
|
||||
func: Callable,
|
||||
out_placements: OutputPlacements,
|
||||
in_placements: InputPlacements,
|
||||
in_grad_placements: InputPlacements,
|
||||
device_mesh: DeviceMesh | None,
|
||||
redistribute_inputs: bool,
|
||||
*args,
|
||||
**kwargs,
|
||||
):
|
||||
# process input args
|
||||
flat_args, args_spec = pytree.tree_flatten(args)
|
||||
if in_placements is not None:
|
||||
if len(in_placements) != len(flat_args):
|
||||
raise AssertionError(
|
||||
f"in_placements length {len(in_placements)} does not match the number "
|
||||
f"of input args {len(flat_args)}!"
|
||||
)
|
||||
|
||||
# we assume every DTensor object is placed on the same device mesh
|
||||
flat_local_args = []
|
||||
seen_dtensor_arg = False
|
||||
for idx, arg in enumerate(flat_args):
|
||||
if isinstance(arg, DTensor):
|
||||
# TODO: the current code doesn't consider the uneven sharding case
|
||||
# Need to think about what the consequence is when the input DTensor
|
||||
# is uneven sharded.
|
||||
if device_mesh is None: # infer device mesh from the DTensor arg
|
||||
device_mesh = arg.device_mesh
|
||||
|
||||
# this function is applied to at least one DTensor argument
|
||||
seen_dtensor_arg = True
|
||||
|
||||
if in_placements is not None:
|
||||
spec = in_placements[idx]
|
||||
if spec is None:
|
||||
raise AssertionError(
|
||||
f"DTensor input {arg} expects placements but received {spec}!"
|
||||
)
|
||||
|
||||
if not isinstance(spec, tuple):
|
||||
spec = tuple(spec)
|
||||
|
||||
if arg.placements != spec:
|
||||
if redistribute_inputs:
|
||||
# redistribute to input placements
|
||||
arg = arg.redistribute(placements=spec)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"arg {arg} in local_map has a mismatched placements: "
|
||||
f"arg placements is {arg.placements} but the input "
|
||||
f"placements is {spec}! "
|
||||
"If redistribute_inputs is wanted, set "
|
||||
"redistribute_inputs=True to local_map."
|
||||
)
|
||||
|
||||
if in_grad_placements is not None:
|
||||
spec = in_grad_placements[idx]
|
||||
if spec is None:
|
||||
raise AssertionError(
|
||||
f"DTensor input {arg} expects in grad placements but received {spec}!"
|
||||
)
|
||||
if not isinstance(spec, tuple):
|
||||
spec = tuple(spec)
|
||||
local_arg = arg.to_local(grad_placements=spec)
|
||||
else:
|
||||
local_arg = arg.to_local()
|
||||
|
||||
if isinstance(local_arg, AsyncCollectiveTensor):
|
||||
local_arg = local_arg.wait()
|
||||
|
||||
flat_local_args.append(local_arg)
|
||||
else:
|
||||
# Non-Tensor input must have None in `in_placements`
|
||||
if in_placements is not None and not isinstance(arg, torch.Tensor):
|
||||
spec = in_placements[idx]
|
||||
if spec is not None:
|
||||
raise AssertionError(
|
||||
f"Non-Tensor input {arg} expects None placements "
|
||||
f"but received {spec}!"
|
||||
)
|
||||
|
||||
flat_local_args.append(arg)
|
||||
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
local_args = pytree.tree_unflatten(flat_local_args, args_spec)
|
||||
|
||||
out = func(*local_args, **kwargs)
|
||||
|
||||
if seen_dtensor_arg:
|
||||
# process output to be DTensor if we've seen DTensor inputs
|
||||
flat_out, out_spec = pytree.tree_flatten(out)
|
||||
|
||||
flat_dist_out = []
|
||||
out_placements_tuple = (
|
||||
out_placements if isinstance(out_placements, tuple) else (out_placements,)
|
||||
)
|
||||
if len(flat_out) != len(out_placements_tuple):
|
||||
raise AssertionError(
|
||||
"local_map requires one PlacementType be provided for each output value,"
|
||||
f" received {len(out_placements_tuple)} out_placements but"
|
||||
f" {len(flat_out)} is expected!"
|
||||
)
|
||||
for out, spec in zip(flat_out, out_placements_tuple):
|
||||
if isinstance(out, torch.Tensor):
|
||||
if isinstance(out, DTensor):
|
||||
raise AssertionError(
|
||||
f"torch.Tensor output expected but received {type(out)}: {out}"
|
||||
)
|
||||
|
||||
flat_dist_out.append(
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
DTensor.from_local(out, device_mesh, spec, run_check=False)
|
||||
)
|
||||
else:
|
||||
if spec is not None:
|
||||
raise AssertionError(
|
||||
f"Non-tensor output {out} expects None placements but received {spec}!"
|
||||
)
|
||||
|
||||
flat_dist_out.append(out)
|
||||
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
return pytree.tree_unflatten(flat_dist_out, out_spec)
|
||||
else:
|
||||
return out
|
||||
+136
@@ -0,0 +1,136 @@
|
||||
# mypy: allow-untyped-defs
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates
|
||||
from collections.abc import Callable, Sequence
|
||||
from functools import partial
|
||||
|
||||
import torch
|
||||
from torch._ops import OpOverload
|
||||
from torch.distributed.tensor import DTensor
|
||||
from torch.distributed.tensor._op_schema import (
|
||||
OpSchema,
|
||||
OpStrategy,
|
||||
PlacementList,
|
||||
RuntimeSchemaInfo,
|
||||
StrategyType,
|
||||
TupleStrategy,
|
||||
)
|
||||
from torch.distributed.tensor._ops.utils import expand_to_full_mesh_op_strategy
|
||||
|
||||
|
||||
__all__ = ["register_sharding"]
|
||||
|
||||
|
||||
def register_sharding(op: OpOverload | list[OpOverload]):
|
||||
"""
|
||||
:meth:`register_sharding` is an experimental API that allows users to register sharding
|
||||
strategies for an operator when the tensor inputs and outputs are DTensor.
|
||||
It can be useful when: (1) there doesn't exist a default sharding strategy for ``op``,
|
||||
e.g. when ``op`` is a custom operator that is not supported by :class:`DTensor`; (2)
|
||||
when users would like to overwrite default sharding strategies of existing operators.
|
||||
|
||||
Args:
|
||||
op (Union[OpOverload, List[OpOverload]]):
|
||||
An op or a list of ops to register the customized sharding function.
|
||||
|
||||
Returns:
|
||||
A function decorator which can be used to wrap a function that defines the sharding
|
||||
strategy for the operator specified in ``op``. The defined sharding strategy will be
|
||||
registered to DTensor and will override the default sharding strategy if DTensor has
|
||||
already implemented the operator. The customized sharding function takes the same inputs
|
||||
as the original op (except that if an arg is a :class:`torch.Tensor`, it will be
|
||||
replaced by a tensor-like object that DTensor uses internally). The function should
|
||||
return a sequence of 2-tuples, each specifying acceptable output placements and its
|
||||
corresponding input placements.
|
||||
|
||||
Example:
|
||||
>>> # xdoctest: +SKIP("distributed")
|
||||
>>> @register_sharding(aten._softmax.default)
|
||||
>>> def custom_softmax_sharding(x, dim, half_to_float):
|
||||
>>> softmax_dim = dim if dim >= 0 else dim + x.ndim
|
||||
>>> acceptable_shardings = []
|
||||
>>>
|
||||
>>> all_replicate = ([Replicate()], [Replicate(), None, None])
|
||||
>>> acceptable_shardings.append(all_replicate)
|
||||
>>>
|
||||
>>> for sharding_dim in range(x.ndim):
|
||||
>>> if sharding_dim != softmax_dim:
|
||||
>>> all_sharded = (
|
||||
>>> [Shard(sharding_dim)],
|
||||
>>> [Shard(sharding_dim), None, None],
|
||||
>>> )
|
||||
>>> acceptable_shardings.append(all_sharded)
|
||||
>>>
|
||||
>>> return acceptable_shardings
|
||||
|
||||
.. note:: This API is currently experimental and subject to change
|
||||
"""
|
||||
|
||||
def custom_strategy(
|
||||
custom_sharding_fn: Callable[
|
||||
..., Sequence[tuple[PlacementList, PlacementList]]
|
||||
],
|
||||
op_schema: OpSchema,
|
||||
) -> StrategyType:
|
||||
def strategy_to_spec(strategy: object) -> object:
|
||||
if isinstance(strategy, OpStrategy):
|
||||
# take the output spec from the first strategy
|
||||
return strategy.strategies[0].output_spec
|
||||
elif isinstance(strategy, TupleStrategy):
|
||||
return tuple(strategy_to_spec(s) for s in strategy.children)
|
||||
else:
|
||||
return strategy
|
||||
|
||||
mesh = op_schema.get_mesh_from_args()
|
||||
|
||||
args_schema = tuple(strategy_to_spec(i) for i in op_schema.args_schema)
|
||||
kwargs_schema = {
|
||||
k: strategy_to_spec(v) for k, v in op_schema.kwargs_schema.items()
|
||||
}
|
||||
|
||||
acceptable_shardings = custom_sharding_fn(*args_schema, **kwargs_schema)
|
||||
|
||||
single_mesh_dim_strategies: list[PlacementList] = []
|
||||
for output_specs, input_specs in acceptable_shardings:
|
||||
single_mesh_dim_strategies.append(output_specs + input_specs)
|
||||
|
||||
# TODO: handle out variant ops
|
||||
return expand_to_full_mesh_op_strategy(
|
||||
mesh,
|
||||
op_schema,
|
||||
single_mesh_dim_strategies,
|
||||
input_index=len(op_schema.op._schema.returns),
|
||||
inplace_op=op_schema.is_inplace_op(),
|
||||
)
|
||||
|
||||
def wrapper(custom_sharding_fn):
|
||||
def derive_schema_info(op):
|
||||
# NOTE: without user directly providing RuntimeSchemaInfo, for now
|
||||
# we create it in a conservative fashion as follows:
|
||||
# 1. let static_argnum be the first int argument
|
||||
# 2. let static_kwargkey include all the int type kwargs
|
||||
# 3. always set needs_pytree=True
|
||||
static_argnum = 100
|
||||
static_kwargkey: list[str] = []
|
||||
for i, arg in enumerate(op._schema.arguments):
|
||||
if isinstance(arg.type, torch.IntType) or (
|
||||
isinstance(arg.type, torch.OptionalType)
|
||||
and isinstance(arg.type.getElementType(), torch.IntType)
|
||||
):
|
||||
static_argnum = min(i, static_argnum)
|
||||
if arg.kwarg_only:
|
||||
static_kwargkey.append(arg.name)
|
||||
return RuntimeSchemaInfo(
|
||||
static_argnum, static_kwargkey or None, needs_pytree=True
|
||||
)
|
||||
|
||||
overloads = op if isinstance(op, list) else [op]
|
||||
for overload in overloads:
|
||||
DTensor._op_dispatcher.sharding_propagator.register_op_strategy(
|
||||
overload,
|
||||
partial(custom_strategy, custom_sharding_fn),
|
||||
derive_schema_info(overload),
|
||||
)
|
||||
|
||||
return custom_sharding_fn
|
||||
|
||||
return wrapper
|
||||
+576
@@ -0,0 +1,576 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import copy
|
||||
import operator
|
||||
from collections.abc import Sequence
|
||||
from typing import Any, cast
|
||||
|
||||
import torch
|
||||
from torch._subclasses.fake_tensor import FakeTensor
|
||||
from torch.distributed.tensor import DeviceMesh, distribute_tensor, DTensor
|
||||
from torch.distributed.tensor._dtensor_spec import DTensorSpec, TensorMeta
|
||||
from torch.distributed.tensor._op_schema import (
|
||||
OpSchema,
|
||||
OpSpec,
|
||||
OutputSharding,
|
||||
OutputSpecType,
|
||||
)
|
||||
from torch.distributed.tensor._redistribute import redistribute_local_tensor
|
||||
from torch.distributed.tensor.parallel.style import ColwiseParallel, ParallelStyle
|
||||
from torch.distributed.tensor.placement_types import Placement, Replicate, Shard
|
||||
from torch.export import ExportedProgram
|
||||
from torch.export.exported_program import ExportGraphSignature
|
||||
from torch.fx import GraphModule
|
||||
from torch.fx.experimental.proxy_tensor import make_fx
|
||||
from torch.fx.node import Node
|
||||
from torch.fx.passes.infra.pass_base import PassBase, PassResult
|
||||
from torch.fx.passes.shape_prop import _extract_tensor_metadata
|
||||
from torch.utils import _pytree as pytree
|
||||
|
||||
|
||||
__all__ = ["tensor_parallel_transformation"]
|
||||
|
||||
aten = torch.ops.aten
|
||||
|
||||
|
||||
def tensor_parallel_transformation(
|
||||
exported_program: ExportedProgram,
|
||||
rank: int,
|
||||
world_size: int,
|
||||
device_type: str,
|
||||
parallel_strategies: dict[str, ParallelStyle],
|
||||
) -> ExportedProgram:
|
||||
"""
|
||||
The entry point function to perform graph transformations on an exported program
|
||||
to transform a single-device graph into a tensor parallel graph.
|
||||
|
||||
.. warning::
|
||||
This API is experimental and subject to change.
|
||||
"""
|
||||
|
||||
gm = exported_program.graph_module
|
||||
sig = copy.deepcopy(exported_program.graph_signature)
|
||||
state_dict = copy.copy(exported_program.state_dict)
|
||||
|
||||
with gm._set_replace_hook(sig.get_replace_hook()):
|
||||
res = _TensorParallelTransformPass(
|
||||
rank,
|
||||
world_size,
|
||||
device_type,
|
||||
state_dict,
|
||||
exported_program.graph_signature,
|
||||
parallel_strategies,
|
||||
)(gm)
|
||||
if res is None:
|
||||
raise AssertionError
|
||||
gm = res.graph_module
|
||||
|
||||
return exported_program._update(gm, sig, state_dict=state_dict)
|
||||
|
||||
|
||||
class _TensorParallelTransformPass(PassBase):
|
||||
"""
|
||||
This pass is responsible for transforming a single-device graph into a tensor parallel
|
||||
graph. It will mark the OpSpec of each node in the graph, partition the graph into
|
||||
distributed graph, then shard the parameters/buffers accordingly.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
rank: int,
|
||||
world_size: int,
|
||||
device_type: str,
|
||||
state_dict: dict[str, torch.Tensor],
|
||||
graph_signature: ExportGraphSignature,
|
||||
parallel_strategies: dict[str, ParallelStyle],
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.rank = rank
|
||||
self.mesh = DeviceMesh(device_type, torch.arange(world_size))
|
||||
self.state_dict: dict[str, torch.Tensor] = state_dict
|
||||
self.graph_signature = graph_signature
|
||||
self.parallel_strategies = parallel_strategies
|
||||
|
||||
def call(self, graph_module) -> PassResult:
|
||||
gm = copy.deepcopy(graph_module)
|
||||
|
||||
parameter_placements = _generate_parameter_and_buffer_placements(
|
||||
list(self.state_dict.keys()), self.parallel_strategies
|
||||
)
|
||||
placement_strategies = _mark_sharding(
|
||||
gm, self.graph_signature, self.mesh, parameter_placements
|
||||
)
|
||||
_partitioner(gm)
|
||||
_shard_state_dict(
|
||||
self.state_dict, placement_strategies, self.graph_signature, self.mesh
|
||||
)
|
||||
return PassResult(gm, True)
|
||||
|
||||
|
||||
def _generate_parameter_and_buffer_placements(
|
||||
params_and_buffers: list[str],
|
||||
parallel_strategies: dict[str, ParallelStyle],
|
||||
) -> dict[str, Placement]:
|
||||
"""
|
||||
Build parameter placements based on the give parallel style of linear layers.
|
||||
"""
|
||||
parameter_placements: dict[str, Placement] = {}
|
||||
for linear_fqn, parallel_style in parallel_strategies.items():
|
||||
weight_fqn = f"{linear_fqn}.weight"
|
||||
bias_fqn = f"{linear_fqn}.bias"
|
||||
if weight_fqn not in params_and_buffers:
|
||||
raise AssertionError
|
||||
parameter_placements[weight_fqn] = (
|
||||
Shard(0) if parallel_style == ColwiseParallel else Shard(1)
|
||||
)
|
||||
if bias_fqn in params_and_buffers:
|
||||
parameter_placements[bias_fqn] = (
|
||||
Shard(0) if parallel_style == ColwiseParallel else Replicate()
|
||||
)
|
||||
return parameter_placements
|
||||
|
||||
|
||||
def _mark_tensor_parallel_shardings(
|
||||
gm: GraphModule,
|
||||
graph_signature: ExportGraphSignature,
|
||||
mesh: DeviceMesh,
|
||||
parameter_placements: dict[str, Placement],
|
||||
) -> dict[Node, OpSpec]:
|
||||
"""
|
||||
Mark the placement strategies of the parameter and buffer placeholder nodes.
|
||||
"""
|
||||
placement_strategies: dict[Node, OpSpec] = {}
|
||||
num_params_and_buffers = len(graph_signature.inputs_to_parameters) + len(
|
||||
graph_signature.inputs_to_buffers
|
||||
)
|
||||
placeholder_idx: int = 0
|
||||
for node in gm.graph.nodes:
|
||||
if node.op == "placeholder":
|
||||
if placeholder_idx < num_params_and_buffers:
|
||||
fqn: str = _get_input_node_fqn(node.name, graph_signature)
|
||||
placement: Placement = (
|
||||
parameter_placements[fqn]
|
||||
if fqn in parameter_placements
|
||||
else Replicate()
|
||||
)
|
||||
placement_strategies[node] = _create_placement_strategy(
|
||||
node,
|
||||
mesh,
|
||||
placements=(placement,),
|
||||
)
|
||||
placeholder_idx += 1
|
||||
else:
|
||||
placement_strategies[node] = _create_placement_strategy(
|
||||
node,
|
||||
mesh,
|
||||
placements=(Replicate(),),
|
||||
)
|
||||
return placement_strategies
|
||||
|
||||
|
||||
def _get_input_node_fqn(input_name: str, graph_signature: ExportGraphSignature) -> str:
|
||||
"""
|
||||
Return the FQN of an input node.
|
||||
"""
|
||||
if input_name in graph_signature.inputs_to_parameters:
|
||||
return graph_signature.inputs_to_parameters[input_name]
|
||||
elif input_name in graph_signature.inputs_to_buffers:
|
||||
return graph_signature.inputs_to_buffers[input_name]
|
||||
else:
|
||||
raise ValueError(
|
||||
f"{input_name} not found in inputs_to_parameters or inputs_to_buffers"
|
||||
)
|
||||
|
||||
|
||||
def _mark_sharding(
|
||||
gm: GraphModule,
|
||||
graph_signature: ExportGraphSignature,
|
||||
mesh: DeviceMesh,
|
||||
parameter_placements: dict[str, Placement],
|
||||
) -> dict[Node, OpSpec]:
|
||||
"""
|
||||
Mark the sharding strategy for each node in the graph module.
|
||||
"""
|
||||
placement_strategies: dict[Node, OpSpec] = _mark_tensor_parallel_shardings(
|
||||
gm,
|
||||
graph_signature,
|
||||
mesh,
|
||||
parameter_placements,
|
||||
)
|
||||
|
||||
for node in gm.graph.nodes:
|
||||
if node.op == "placeholder":
|
||||
if node not in placement_strategies:
|
||||
placement_strategies[node] = _create_placement_strategy(
|
||||
node, mesh, placements=(Replicate(),)
|
||||
)
|
||||
node.meta["sharding"] = placement_strategies[node]
|
||||
elif node.op == "call_function":
|
||||
if node.target is operator.getitem:
|
||||
input_nodes = node.all_input_nodes
|
||||
if len(input_nodes) != 1:
|
||||
raise AssertionError(
|
||||
f"non-compute op only support one input now, found node: {node} "
|
||||
f"with length of inputs: {len(node.args)}"
|
||||
)
|
||||
arg_strategy = placement_strategies[input_nodes[0]]
|
||||
placement_strategies[node] = _create_placement_strategy(
|
||||
node,
|
||||
mesh,
|
||||
placements=arg_strategy.output_spec.placements,
|
||||
input_specs=_get_input_node_specs(node, placement_strategies),
|
||||
)
|
||||
node.meta["sharding"] = placement_strategies[node]
|
||||
else:
|
||||
op_schema = _get_op_schema(node, placement_strategies)
|
||||
|
||||
# get DTensor specs for inputs and outputs
|
||||
sharding_propagator = DTensor._op_dispatcher.sharding_propagator
|
||||
if (
|
||||
op_schema.op not in sharding_propagator.op_strategy_funcs
|
||||
and op_schema.op not in sharding_propagator.op_to_rules
|
||||
and op_schema.op
|
||||
not in sharding_propagator.op_single_dim_strategy_funcs
|
||||
):
|
||||
# Mark all as replicated
|
||||
output_sharding = _generate_default_output_sharding(
|
||||
node,
|
||||
mesh,
|
||||
op_schema,
|
||||
)
|
||||
else:
|
||||
output_sharding = DTensor._op_dispatcher.sharding_propagator.propagate_op_sharding( # type: ignore[assignment]
|
||||
op_schema,
|
||||
)
|
||||
placement_strategies[node] = OpSpec(
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
output_specs=_get_output_spec_from_output_sharding(output_sharding),
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
input_specs=output_sharding.redistribute_schema.args_spec
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
if output_sharding.redistribute_schema is not None
|
||||
else _get_input_node_specs(node, placement_strategies),
|
||||
)
|
||||
node.meta["sharding"] = placement_strategies[node]
|
||||
elif node.op == "output":
|
||||
node.meta["sharding"] = None
|
||||
else:
|
||||
raise RuntimeError(f"op code {node.op} not supported")
|
||||
return placement_strategies
|
||||
|
||||
|
||||
def _get_output_spec_from_output_sharding(
|
||||
output_sharding: OutputSharding,
|
||||
) -> DTensorSpec:
|
||||
"""
|
||||
Util function to extract output spec from output sharding.
|
||||
"""
|
||||
if isinstance(output_sharding.output_spec, DTensorSpec):
|
||||
return output_sharding.output_spec
|
||||
else:
|
||||
# For ops that return multiple outputs, the outputs should have the same output spec
|
||||
if not isinstance(output_sharding.output_spec, Sequence):
|
||||
raise AssertionError
|
||||
if output_sharding.output_spec[0] is None:
|
||||
raise AssertionError
|
||||
output_sharding.output_spec[0].tensor_meta = None
|
||||
return output_sharding.output_spec[0]
|
||||
|
||||
|
||||
def _create_placement_strategy(
|
||||
node: Node,
|
||||
mesh: DeviceMesh,
|
||||
placements: tuple[Placement, ...],
|
||||
input_specs: Sequence[DTensorSpec] | None = None,
|
||||
) -> OpSpec:
|
||||
"""
|
||||
Util function to construct an OpSpec for a given node.
|
||||
"""
|
||||
placement = OpSpec(
|
||||
input_specs=input_specs,
|
||||
output_specs=DTensorSpec(
|
||||
mesh=mesh,
|
||||
placements=placements,
|
||||
),
|
||||
)
|
||||
_populate_tensor_meta(node, placement.output_specs)
|
||||
return placement
|
||||
|
||||
|
||||
def _populate_tensor_meta(node: Node, output_spec: OutputSpecType) -> None:
|
||||
"""
|
||||
Util function to populate tensor meta of output_spec based on node metadata.
|
||||
"""
|
||||
if isinstance(node.meta["val"], Sequence):
|
||||
if not isinstance(output_spec, Sequence):
|
||||
raise AssertionError
|
||||
for spec, fake_tensor in zip(output_spec, node.meta["val"]):
|
||||
if spec is None:
|
||||
raise AssertionError
|
||||
spec.tensor_meta = TensorMeta(
|
||||
shape=fake_tensor.shape,
|
||||
stride=fake_tensor.stride(),
|
||||
dtype=fake_tensor.dtype,
|
||||
)
|
||||
else:
|
||||
if not isinstance(output_spec, DTensorSpec):
|
||||
raise AssertionError
|
||||
output_spec.tensor_meta = TensorMeta(
|
||||
shape=node.meta["val"].shape,
|
||||
stride=node.meta["val"].stride(),
|
||||
dtype=node.meta["val"].dtype,
|
||||
)
|
||||
|
||||
|
||||
def _generate_default_output_sharding(
|
||||
node: Node,
|
||||
mesh: DeviceMesh,
|
||||
op_schema: OpSchema,
|
||||
) -> OutputSharding:
|
||||
"""
|
||||
Util function to create a default output sharding that suggests Replicate placement for both args and outputs.
|
||||
"""
|
||||
|
||||
def update_arg_spec(arg_spec: DTensorSpec) -> DTensorSpec:
|
||||
return DTensorSpec(
|
||||
mesh=arg_spec.mesh,
|
||||
placements=(Replicate(),),
|
||||
tensor_meta=arg_spec.tensor_meta,
|
||||
)
|
||||
|
||||
new_op_schema = OpSchema(
|
||||
op=op_schema.op,
|
||||
args_schema=pytree.tree_map_only(
|
||||
DTensorSpec, update_arg_spec, op_schema.args_schema
|
||||
),
|
||||
kwargs_schema=op_schema.kwargs_schema,
|
||||
)
|
||||
|
||||
def create_output_spec(tensor: FakeTensor) -> DTensorSpec:
|
||||
return DTensorSpec(
|
||||
mesh=mesh,
|
||||
placements=(Replicate(),),
|
||||
tensor_meta=TensorMeta(
|
||||
shape=tensor.shape,
|
||||
stride=tensor.stride(),
|
||||
dtype=tensor.dtype,
|
||||
),
|
||||
)
|
||||
|
||||
return OutputSharding(
|
||||
output_spec=pytree.tree_map_only(
|
||||
FakeTensor, create_output_spec, node.meta["val"]
|
||||
),
|
||||
redistribute_schema=new_op_schema,
|
||||
needs_redistribute=True,
|
||||
)
|
||||
|
||||
|
||||
def _partitioner(gm: torch.fx.GraphModule) -> torch.fx.GraphModule:
|
||||
"""
|
||||
Graph partitioner that partitions the single device graph
|
||||
to distributed graph
|
||||
"""
|
||||
for node in gm.graph.nodes:
|
||||
node_sharding = node.meta["sharding"]
|
||||
if node.op == "placeholder":
|
||||
out_spec = node_sharding.output_spec
|
||||
local_val = _partition_val(node.meta["val"], out_spec)
|
||||
# update node value
|
||||
node.meta["val"] = local_val
|
||||
elif node.op == "call_function":
|
||||
out_spec = node_sharding.output_spec
|
||||
# check if there's misaligned sharding, insert reshard if there is
|
||||
expected_input_specs = node_sharding.input_specs
|
||||
for idx, input_arg in enumerate(node.all_input_nodes):
|
||||
input_arg_sharding = input_arg.meta["sharding"]
|
||||
input_arg_spec = input_arg_sharding.output_spec
|
||||
desired_spec = (
|
||||
out_spec
|
||||
if expected_input_specs is None
|
||||
else expected_input_specs[idx]
|
||||
)
|
||||
if input_arg_spec != desired_spec:
|
||||
_insert_reshard_gm(
|
||||
gm, node, input_arg, input_arg_spec, desired_spec
|
||||
)
|
||||
# convert output val to its local component
|
||||
output_val = node.meta["val"]
|
||||
node.meta["val"] = _partition_val(output_val, out_spec)
|
||||
elif node.op == "output":
|
||||
for input_arg in node.all_input_nodes:
|
||||
# input args of output should be Replicate, otherwise redistribution is needed.
|
||||
input_args_to_check: Sequence[Node] = (
|
||||
input_arg if isinstance(input_arg, Sequence) else [input_arg]
|
||||
)
|
||||
for arg in input_args_to_check:
|
||||
arg_sharding = arg.meta["sharding"]
|
||||
arg_spec = arg_sharding.output_spec
|
||||
desired_spec = copy.copy(arg_spec)
|
||||
desired_spec.placements = (Replicate(),)
|
||||
if arg_spec != desired_spec:
|
||||
_insert_reshard_gm(gm, node, arg, arg_spec, desired_spec)
|
||||
else:
|
||||
raise RuntimeError(f"op code {node} not supported")
|
||||
|
||||
_clean_up_graph_metadata(gm)
|
||||
gm.graph.lint()
|
||||
gm.recompile()
|
||||
return gm
|
||||
|
||||
|
||||
def _partition_val(val: Any, spec: DTensorSpec) -> Any:
|
||||
"""
|
||||
util function to convert a full tensor val to its local component
|
||||
"""
|
||||
if isinstance(val, torch.Tensor):
|
||||
local_shard = val
|
||||
if val.ndim == 0:
|
||||
# If it's already a scalar tensor, it is already local, we don't
|
||||
# need to do anything
|
||||
return local_shard
|
||||
|
||||
for idx, placement in enumerate(spec.placements):
|
||||
# NOTE: is_shard() does not match _StridedShard; see _is_shard_like().
|
||||
if placement.is_shard():
|
||||
placement = cast(Shard, placement)
|
||||
num_chunks = spec.mesh.size(mesh_dim=idx)
|
||||
my_coord = spec.mesh.get_coordinate()
|
||||
if my_coord is None:
|
||||
raise AssertionError("current rank not in mesh!")
|
||||
my_coord_on_mesh_dim = my_coord[idx]
|
||||
local_shard = placement._select_split_tensor(
|
||||
local_shard,
|
||||
num_chunks,
|
||||
my_coord_on_mesh_dim,
|
||||
with_padding=False,
|
||||
contiguous=True,
|
||||
clone=False,
|
||||
)
|
||||
return local_shard
|
||||
elif isinstance(val, (list, tuple)):
|
||||
return val.__class__(_partition_val(v, spec) for v in val)
|
||||
else:
|
||||
raise RuntimeError(f"val type {type(val)} not supported")
|
||||
|
||||
|
||||
def _insert_reshard_gm(
|
||||
gm: torch.fx.GraphModule,
|
||||
node: Node,
|
||||
input_arg: Node,
|
||||
input_arg_spec: DTensorSpec,
|
||||
desired_spec: DTensorSpec,
|
||||
) -> None:
|
||||
"""
|
||||
Transform the graph for tensor redistribution.
|
||||
"""
|
||||
input_arg_spec.tensor_meta = input_arg.meta["tensor_meta"]
|
||||
desired_spec.tensor_meta = input_arg.meta["tensor_meta"]
|
||||
input_arg_tensor = input_arg.meta["val"]
|
||||
|
||||
# insert reshard operation
|
||||
def reshard_fn(local_tensor: torch.Tensor) -> torch.Tensor:
|
||||
return redistribute_local_tensor(
|
||||
local_tensor,
|
||||
input_arg_spec,
|
||||
desired_spec,
|
||||
)
|
||||
|
||||
reshard_gm = make_fx(reshard_fn)(input_arg_tensor)
|
||||
reshard_gm_nodes = list(reshard_gm.graph.nodes)
|
||||
input_node = reshard_gm_nodes[0]
|
||||
with gm.graph.inserting_before(node):
|
||||
# copy nn_module_stack metadata for output, all-reduce nodes
|
||||
for reshard_node in reshard_gm.graph.nodes:
|
||||
if reshard_node.op not in ["placeholder", "output"]:
|
||||
reshard_node.meta["nn_module_stack"] = (
|
||||
copy.copy(input_arg.meta["nn_module_stack"])
|
||||
if input_arg.op != "placeholder"
|
||||
else copy.copy(node.meta["nn_module_stack"])
|
||||
)
|
||||
output_node = gm.graph.graph_copy(
|
||||
reshard_gm.graph,
|
||||
val_map={
|
||||
input_node: input_arg,
|
||||
},
|
||||
)
|
||||
node.replace_input_with(input_arg, output_node) # type: ignore[arg-type]
|
||||
|
||||
|
||||
def _clean_up_graph_metadata(gm: torch.fx.GraphModule) -> None:
|
||||
"""
|
||||
Clean up the graph by removing sharding and partitioning related metadata
|
||||
"""
|
||||
for node in gm.graph.nodes:
|
||||
if "sharding" in node.meta:
|
||||
del node.meta["sharding"]
|
||||
if "val" in node.meta and isinstance(node.meta["val"], torch.Tensor):
|
||||
local_tensor_meta = _extract_tensor_metadata(node.meta["val"])
|
||||
node.meta["tensor_meta"] = local_tensor_meta
|
||||
|
||||
|
||||
def _get_input_node_specs(
|
||||
node: Node, placement_strategies: dict[Node, OpSpec]
|
||||
) -> tuple[DTensorSpec, ...]:
|
||||
"""
|
||||
Get the input specs of a node.
|
||||
"""
|
||||
input_specs_list: list[DTensorSpec] = []
|
||||
for input_arg in node.all_input_nodes:
|
||||
if input_arg in placement_strategies:
|
||||
output_spec = placement_strategies[input_arg].output_specs
|
||||
if not isinstance(output_spec, DTensorSpec):
|
||||
raise AssertionError
|
||||
input_specs_list.append(output_spec)
|
||||
else:
|
||||
raise ValueError(f"{input_arg} does not have output_spec populated.")
|
||||
return tuple(input_specs_list)
|
||||
|
||||
|
||||
def _get_op_schema(node: Node, placement_strategies: dict[Node, OpSpec]) -> OpSchema:
|
||||
"""
|
||||
Util function to construct the operator schema of a node.
|
||||
"""
|
||||
args_schema_list = pytree.tree_map_only(
|
||||
Node, lambda arg: placement_strategies[arg].output_specs, node.args
|
||||
)
|
||||
op_schema = OpSchema(
|
||||
op=cast(torch._ops.OpOverload, node.target),
|
||||
args_schema=tuple(args_schema_list),
|
||||
kwargs_schema=cast(dict[str, object], node.kwargs),
|
||||
)
|
||||
return op_schema
|
||||
|
||||
|
||||
def _shard_state_dict(
|
||||
state_dict: dict[str, torch.Tensor],
|
||||
placement_strategies: dict[Node, OpSpec],
|
||||
graph_signature: ExportGraphSignature,
|
||||
mesh: DeviceMesh,
|
||||
) -> None:
|
||||
"""
|
||||
Inplace partition the weights based on the OpSpec
|
||||
"""
|
||||
for node, op_spec in placement_strategies.items():
|
||||
if node.op != "placeholder":
|
||||
continue
|
||||
if node.name in graph_signature.inputs_to_parameters:
|
||||
fqn = graph_signature.inputs_to_parameters[node.name]
|
||||
elif node.name in graph_signature.inputs_to_buffers:
|
||||
fqn = graph_signature.inputs_to_buffers[node.name]
|
||||
else:
|
||||
continue
|
||||
if fqn not in state_dict:
|
||||
raise AssertionError(f"{fqn} not found in state dict: {state_dict.keys()}")
|
||||
|
||||
original_param = state_dict[fqn]
|
||||
dtensor_param = distribute_tensor(
|
||||
original_param,
|
||||
mesh,
|
||||
op_spec.output_spec.placements,
|
||||
)
|
||||
local_param = dtensor_param.to_local()
|
||||
state_dict[fqn] = (
|
||||
torch.nn.Parameter(local_param)
|
||||
if isinstance(original_param, torch.nn.Parameter)
|
||||
else local_param
|
||||
)
|
||||
+25
@@ -0,0 +1,25 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates
|
||||
from torch.distributed.tensor.parallel.api import parallelize_module
|
||||
from torch.distributed.tensor.parallel.loss import loss_parallel
|
||||
from torch.distributed.tensor.parallel.style import (
|
||||
ColwiseParallel,
|
||||
ParallelStyle,
|
||||
PrepareModuleInput,
|
||||
PrepareModuleInputOutput,
|
||||
PrepareModuleOutput,
|
||||
RowwiseParallel,
|
||||
SequenceParallel,
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"ColwiseParallel",
|
||||
"ParallelStyle",
|
||||
"PrepareModuleInput",
|
||||
"PrepareModuleInputOutput",
|
||||
"PrepareModuleOutput",
|
||||
"RowwiseParallel",
|
||||
"SequenceParallel",
|
||||
"parallelize_module",
|
||||
"loss_parallel",
|
||||
]
|
||||
+51
@@ -0,0 +1,51 @@
|
||||
from functools import partial
|
||||
from typing import no_type_check
|
||||
|
||||
import torch
|
||||
from torch.distributed._functional_collectives import AsyncCollectiveTensor
|
||||
from torch.distributed.tensor import DTensor
|
||||
from torch.distributed.tensor._dtensor_spec import DTensorSpec
|
||||
|
||||
|
||||
@no_type_check
|
||||
def sync_grad_hook(grad, *, device_handle=None, compute_stream=None):
|
||||
if isinstance(grad, AsyncCollectiveTensor):
|
||||
if compute_stream is not None:
|
||||
with device_handle.stream(compute_stream):
|
||||
grad = grad.wait()
|
||||
else:
|
||||
grad = grad.wait()
|
||||
|
||||
return grad
|
||||
|
||||
|
||||
def _flatten_tensor(
|
||||
tensor: torch.Tensor,
|
||||
) -> tuple[torch.Tensor, DTensorSpec | None]:
|
||||
if isinstance(tensor, DTensor):
|
||||
tensor._local_tensor.requires_grad_()
|
||||
return tensor._local_tensor, tensor._spec
|
||||
return tensor, None
|
||||
|
||||
|
||||
@no_type_check
|
||||
def _unflatten_tensor(tensor, spec, *, device_handle=None, compute_stream=None):
|
||||
# unflatten would mainly be called every time FSDP allgather parameters.
|
||||
result = DTensor.from_local(
|
||||
tensor,
|
||||
spec.mesh,
|
||||
spec.placements,
|
||||
run_check=False,
|
||||
shape=spec.shape,
|
||||
stride=spec.stride,
|
||||
)
|
||||
if tensor.requires_grad:
|
||||
# only register the hook if the tensor requires grad
|
||||
tensor.register_hook(
|
||||
partial(
|
||||
sync_grad_hook,
|
||||
device_handle=device_handle,
|
||||
compute_stream=compute_stream,
|
||||
)
|
||||
)
|
||||
return result
|
||||
@@ -0,0 +1,142 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates
|
||||
import warnings
|
||||
from fnmatch import fnmatch
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.distributed.device_mesh import _mesh_resources, DeviceMesh
|
||||
from torch.distributed.tensor.parallel.style import ParallelStyle
|
||||
|
||||
|
||||
__all__ = ["parallelize_module"]
|
||||
|
||||
|
||||
def parallelize_module( # type: ignore[return]
|
||||
module: nn.Module,
|
||||
device_mesh: DeviceMesh | None = None,
|
||||
parallelize_plan: ParallelStyle | dict[str, ParallelStyle] | None = None,
|
||||
*,
|
||||
src_data_rank: int | None = 0,
|
||||
) -> nn.Module:
|
||||
"""
|
||||
Apply Tensor Parallelism in PyTorch by parallelizing modules or sub-modules based on a user-specified plan.
|
||||
|
||||
We parallelize module or sub_modules based on a parallelize_plan. The parallelize_plan contains
|
||||
:class:`ParallelStyle`, which indicates how user wants the module or sub_module
|
||||
to be parallelized.
|
||||
|
||||
User can also specify different parallel style per module fully qualified name (FQN).
|
||||
|
||||
Note that ``parallelize_module`` only accepts a 1-D :class:`DeviceMesh`, if you have a 2-D or N-D :class:`DeviceMesh`,
|
||||
slice the DeviceMesh to a 1-D sub DeviceMesh first then pass to this API(i.e. ``device_mesh[\"tp\"]``)
|
||||
|
||||
Args:
|
||||
module (:class:`nn.Module`):
|
||||
Module to be parallelized.
|
||||
device_mesh (:class:`DeviceMesh`, optional):
|
||||
Object which describes the mesh topology of devices for the DTensor.
|
||||
If not specified, the call must be under a DeviceMesh context.
|
||||
parallelize_plan (Union[:class:`ParallelStyle`, Dict[str, :class:`ParallelStyle`]], optional):
|
||||
The plan used to parallelize the module. It can be either a
|
||||
:class:`ParallelStyle` object which contains how we prepare
|
||||
input/output for Tensor Parallelism or it can be a dict of module
|
||||
FQN and its corresponding :class:`ParallelStyle` object. If not
|
||||
specified, the call will do nothing at the moment.
|
||||
Keyword args:
|
||||
src_data_rank (int, optional): the rank of the source data for the logical/global tensor, it is used by
|
||||
:meth:`distribute_tensor` to scatter/broadcast the shards/replicas to other ranks. By default,
|
||||
we use ``group_rank=0`` on each DeviceMesh dimension as the source data to preserve the single-device
|
||||
semantic. If passing ``None`` explicitly, :meth:`parallelize_module` simply uses its local data instead
|
||||
of trying to preserve the single-device semantic via scatter/broadcast. Default: 0
|
||||
Return:
|
||||
A :class:`nn.Module` object parallelized.
|
||||
|
||||
Example::
|
||||
>>> # xdoctest: +SKIP("distributed")
|
||||
>>> from torch.distributed.tensor.parallel import parallelize_module, ColwiseParallel
|
||||
>>> from torch.distributed.device_mesh import init_device_mesh
|
||||
>>>
|
||||
>>> # Define the module.
|
||||
>>> m = Model(...)
|
||||
>>> tp_mesh = init_device_mesh("cuda", (8,))
|
||||
>>> m = parallelize_module(m, tp_mesh, {"w1": ColwiseParallel(), "w2": RowwiseParallel()})
|
||||
>>>
|
||||
|
||||
.. note:: For complex module architecture like Attention, MLP layers, we recommend composing
|
||||
different ParallelStyles together (i.e. ``ColwiseParallel`` and ``RowwiseParallel``) and pass
|
||||
as a parallelize_plan, to achieves the desired sharding computation.
|
||||
"""
|
||||
torch._C._log_api_usage_once("torch.distributed.tensor.parallel.parallelize_module")
|
||||
|
||||
device_mesh = device_mesh or _mesh_resources.get_current_mesh()
|
||||
|
||||
if parallelize_plan is None:
|
||||
warnings.warn(
|
||||
"No parallelize_plan is provided and auto-parallel is not supported "
|
||||
"at the moment, so this parallelize_module call will do nothing.",
|
||||
stacklevel=2,
|
||||
)
|
||||
return module
|
||||
|
||||
# note: The RNG tracker will be initialized in distribute_tensor() call if it hasn't
|
||||
# been initialized.
|
||||
|
||||
if isinstance(parallelize_plan, ParallelStyle):
|
||||
parallelize_plan.src_data_rank = src_data_rank
|
||||
return parallelize_plan._apply(module, device_mesh)
|
||||
elif isinstance(parallelize_plan, dict):
|
||||
for module_path, parallelize_style in parallelize_plan.items():
|
||||
if module_path == "":
|
||||
# shortcut: empty string means to apply the plan to the current module
|
||||
parallelize_module(module, device_mesh, parallelize_style)
|
||||
continue
|
||||
|
||||
path_splits = module_path.split(".")
|
||||
# Instead of blindly popping tokens, first check the match,
|
||||
# we only consume/pop the token if we found a match.
|
||||
token = path_splits[0]
|
||||
|
||||
matched_children = list(
|
||||
filter(
|
||||
# `t[0]` is child name
|
||||
lambda t: fnmatch(t[0], token),
|
||||
module.named_children(),
|
||||
)
|
||||
)
|
||||
if not matched_children:
|
||||
# No match at this level. Log a warning and process next plan entry.
|
||||
warnings.warn(
|
||||
f"Parallelize plan key '{module_path}' could not be resolved: "
|
||||
f"no submodule matching token '{token}' in module {module}, "
|
||||
f"skipping this plan entry.",
|
||||
stacklevel=2,
|
||||
)
|
||||
continue
|
||||
|
||||
# Now that we have a match, we can consume the token.
|
||||
path_splits.pop(0)
|
||||
# apply the plan to all matched submodules
|
||||
for _, submodule in matched_children:
|
||||
if path_splits:
|
||||
# we haven't reached the leaf, apply in dict style
|
||||
leaf_path = ".".join(path_splits) # rest of the path after `token`
|
||||
parallelize_module(
|
||||
submodule,
|
||||
device_mesh,
|
||||
{leaf_path: parallelize_style},
|
||||
src_data_rank=src_data_rank,
|
||||
)
|
||||
else:
|
||||
# otherwise, directly apply style to this submodule
|
||||
parallelize_module(
|
||||
submodule,
|
||||
device_mesh,
|
||||
parallelize_style,
|
||||
src_data_rank=src_data_rank,
|
||||
)
|
||||
return module
|
||||
else:
|
||||
raise TypeError( # pyre-ignore[7]
|
||||
"Expect Union[ParallelStyle, Dict[str, ParallelStyle]] for"
|
||||
f" parallelize_plan, {type(parallelize_plan)} found!"
|
||||
)
|
||||
@@ -0,0 +1,105 @@
|
||||
# mypy: allow-untyped-defs
|
||||
from typing import Any
|
||||
|
||||
import torch.nn as nn
|
||||
from torch.distributed.tensor.parallel._data_parallel_utils import (
|
||||
_flatten_tensor,
|
||||
_unflatten_tensor,
|
||||
)
|
||||
|
||||
|
||||
__all__ = [] # type: ignore[var-annotated]
|
||||
|
||||
|
||||
def _get_submodule_n_params(module: nn.Module, path: str):
|
||||
"""
|
||||
Get submodule and the direct path of parameter from the module
|
||||
"""
|
||||
if "." in path:
|
||||
path_list = path.split(".")
|
||||
parent_module_path = ".".join(path_list[:-1])
|
||||
module = module.get_submodule(parent_module_path)
|
||||
path = path_list[-1]
|
||||
return module, path
|
||||
|
||||
|
||||
def _update_module_param(param_list: list[tuple[nn.Module, str, nn.Parameter]]):
|
||||
"""
|
||||
Update parameters within the module
|
||||
"""
|
||||
for item in param_list:
|
||||
parent_module, module_path, t = item
|
||||
if not hasattr(parent_module, module_path):
|
||||
raise AssertionError
|
||||
delattr(parent_module, module_path)
|
||||
setattr(parent_module, module_path, t)
|
||||
|
||||
|
||||
def _reconstruct_dtensor(module: nn.Module, _input: Any):
|
||||
"""
|
||||
Reconstruct DTensor parameters from local tensors
|
||||
"""
|
||||
param_list = []
|
||||
# TODO: To add perf optimizations to this iterations
|
||||
for name, t in module.named_parameters():
|
||||
if hasattr(t, "_st_info"):
|
||||
dtensor = _unflatten_tensor(t, t._st_info)
|
||||
param_list.append((*_get_submodule_n_params(module, name), dtensor))
|
||||
_update_module_param(param_list) # type: ignore[arg-type]
|
||||
|
||||
|
||||
def _localize_dtensor(
|
||||
module: nn.Module, *_: Any, ignored_params: set[nn.Parameter] | None = None
|
||||
):
|
||||
"""
|
||||
Convert DTensor parameters to local tensors
|
||||
"""
|
||||
if ignored_params is None:
|
||||
ignored_params = set()
|
||||
param_list = []
|
||||
for name, param in module.named_parameters():
|
||||
if param in ignored_params:
|
||||
continue
|
||||
t, sharding_info = _flatten_tensor(param)
|
||||
if sharding_info is not None:
|
||||
t = nn.Parameter(t)
|
||||
t._st_info = sharding_info # type: ignore[attr-defined]
|
||||
param_list.append((*_get_submodule_n_params(module, name), t))
|
||||
_update_module_param(param_list) # type: ignore[arg-type]
|
||||
|
||||
|
||||
def _pre_dp_module_transform(module: nn.Module):
|
||||
"""
|
||||
Enable the composability between Tensor Parallelism (TP) and Data
|
||||
Parallelism(DP) in PyTorch when using DDP. We need to convert Parameters which
|
||||
are DTensors to local tensors before wrapping with data parallelism API.
|
||||
We then register two hooks, one for converting local tensors back to DTensor
|
||||
preforward and one to convert DTensors back to tensors after Forward. By
|
||||
integrating this way, we avoid any special handling of DTensor parameters by DDP
|
||||
and get DTensor's gradients propagated back to DP, e.g. gradient buckets of DDP.
|
||||
|
||||
For now, this API only works with ``DistributedDataParallel``. It will later support
|
||||
other DP methods such as FSDP.
|
||||
|
||||
Args:
|
||||
module (:class:`nn.Module`):
|
||||
Module which has been applied TP on.
|
||||
|
||||
Example::
|
||||
>>> # xdoctest: +SKIP("distributed")
|
||||
>>> from torch.distributed.tensor.parallel import parallelize_module, PairwiseParallel
|
||||
>>> from torch.nn.parallel import DistributedDataParallel as DDP
|
||||
>>> from torch.distributed.tensor.parallel.ddp import pre_dp_module_transform
|
||||
>>>
|
||||
>>> # Define the module.
|
||||
>>> m = module(...)
|
||||
>>> parallelize_module(m, PairwiseParallel())
|
||||
>>> m = pre_dp_module_transform(m)
|
||||
>>> m = DDP(m)
|
||||
>>>
|
||||
"""
|
||||
|
||||
_localize_dtensor(module, None, None)
|
||||
# TODO: To add test cases and ensure that it works for nested modules
|
||||
module.register_forward_pre_hook(_reconstruct_dtensor)
|
||||
module.register_forward_hook(_localize_dtensor)
|
||||
@@ -0,0 +1,400 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import copy
|
||||
from typing import Any, cast
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
import torch.distributed._shard.sharding_spec as shard_spec
|
||||
import torch.distributed.distributed_c10d as c10d
|
||||
from torch.distributed._shard.sharded_tensor import (
|
||||
Shard,
|
||||
ShardedTensor,
|
||||
ShardedTensorMetadata,
|
||||
TensorProperties,
|
||||
)
|
||||
from torch.distributed._shard.sharding_spec import ShardMetadata
|
||||
from torch.distributed._shard.sharding_spec.chunk_sharding_spec import ChunkShardingSpec
|
||||
from torch.distributed.fsdp._common_utils import _set_fsdp_flattened
|
||||
from torch.distributed.fsdp._fsdp_extensions import FSDPExtensions
|
||||
from torch.distributed.fsdp._shard_utils import _create_chunk_sharded_tensor
|
||||
from torch.distributed.remote_device import _remote_device
|
||||
from torch.distributed.tensor import DeviceMesh, DTensor, Replicate, Shard as DShard
|
||||
from torch.distributed.tensor.parallel._data_parallel_utils import (
|
||||
_flatten_tensor,
|
||||
_unflatten_tensor,
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["DTensorExtensions"]
|
||||
|
||||
|
||||
def _get_box(tensor: DTensor) -> tuple[torch.Size, torch.Size]:
|
||||
device_mesh = tensor.device_mesh
|
||||
if device_mesh.ndim != 1:
|
||||
raise AssertionError("Only 1D DeviceMeshes currently handled")
|
||||
|
||||
placement = tensor.placements[0]
|
||||
offsets = [0] * len(tensor.size())
|
||||
num_chunks = device_mesh.size(mesh_dim=0)
|
||||
|
||||
# NOTE: is_shard() does not match _StridedShard; see _is_shard_like().
|
||||
if tensor.placements[0].is_shard():
|
||||
shard_dim = cast(DShard, placement).dim
|
||||
chunk_size = tensor.size(shard_dim) // num_chunks
|
||||
offsets[shard_dim] = chunk_size
|
||||
|
||||
return (torch.Size(offsets), tensor._local_tensor.size())
|
||||
|
||||
|
||||
def _get_box_for(tensor: DTensor, idx: int) -> tuple[torch.Size, torch.Size]:
|
||||
offsets, size = _get_box(tensor)
|
||||
return (torch.Size([val * idx for val in offsets]), size)
|
||||
|
||||
|
||||
def _get_local_box(tensor: DTensor) -> tuple[torch.Size, torch.Size]:
|
||||
device_mesh = tensor.device_mesh
|
||||
coord = device_mesh.get_coordinate()
|
||||
if coord is None:
|
||||
raise AssertionError
|
||||
return _get_box_for(tensor, coord[0])
|
||||
|
||||
|
||||
def _create_shard_md_from_dt(dt: DTensor, current_rank: int) -> ShardMetadata:
|
||||
mesh = dt.device_mesh
|
||||
if mesh.ndim != 1:
|
||||
raise AssertionError("Only 1D DeviceMeshes currently handled")
|
||||
|
||||
offsets, sizes = _get_local_box(dt)
|
||||
return ShardMetadata(
|
||||
shard_offsets=list(offsets),
|
||||
shard_sizes=list(sizes),
|
||||
placement=f"rank:{current_rank}/{dt._local_tensor.device}",
|
||||
)
|
||||
|
||||
|
||||
def _create_sharded_tensor_md_from_dt(
|
||||
dt: DTensor, dt_pg: c10d.ProcessGroup
|
||||
) -> ShardedTensorMetadata:
|
||||
# This is where it gets tricky, we have to produce a ShardedTensor that has full coverage
|
||||
# and yet has only one valid shard for the current rank.
|
||||
|
||||
shards_md = []
|
||||
my_rank = dist.get_rank(dt_pg)
|
||||
scapegoat_rank = 0 if my_rank > 0 else 1
|
||||
|
||||
# NOTE: is_shard() does not match _StridedShard; see _is_shard_like().
|
||||
if dt.placements[0].is_shard():
|
||||
shard_count = dt_pg.size()
|
||||
else:
|
||||
shard_count = 1
|
||||
|
||||
for i in range(shard_count):
|
||||
offsets, sizes = _get_box_for(dt, i)
|
||||
shards_md.append(
|
||||
ShardMetadata(
|
||||
shard_offsets=list(offsets),
|
||||
shard_sizes=list(sizes),
|
||||
placement=(
|
||||
f"rank:{scapegoat_rank if i > 0 else my_rank}/{dt._local_tensor.device}"
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
return ShardedTensorMetadata(
|
||||
shards_metadata=shards_md,
|
||||
size=dt.size(),
|
||||
tensor_properties=TensorProperties(
|
||||
dtype=dt.dtype,
|
||||
layout=dt.layout,
|
||||
requires_grad=dt.requires_grad,
|
||||
# ignore memory_format and pin_memory as those are not supported by DT
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _get_dt_pg(dt: DTensor) -> c10d.ProcessGroup:
|
||||
mesh = dt.device_mesh
|
||||
if mesh.ndim != 1:
|
||||
raise AssertionError("Only 1D DeviceMeshes currently handled")
|
||||
return mesh.get_group()
|
||||
|
||||
|
||||
def _rewrite_spec_if_needed(
|
||||
spec: shard_spec.ShardingSpec, tensor: torch.Tensor, rank: int
|
||||
) -> shard_spec.ShardingSpec:
|
||||
"""
|
||||
Rewrite ``spec`` to match the device of ``tensor``.
|
||||
|
||||
FSDP.sharded_optim_state_dict sneakly ships optimizer state to CPU so if the original ShardingSpec
|
||||
produces CUDA metadata, ST construction bombs.
|
||||
"""
|
||||
if not isinstance(spec, ChunkShardingSpec):
|
||||
return spec
|
||||
|
||||
# let's see if we need
|
||||
rewrite = False
|
||||
for p in spec.placements:
|
||||
p = cast(_remote_device, p)
|
||||
if p.rank() == rank and p.device() != tensor.device:
|
||||
rewrite = True
|
||||
break
|
||||
if rewrite:
|
||||
spec = copy.deepcopy(spec)
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
for i, placement in enumerate(spec.placements):
|
||||
placement = cast(_remote_device, placement)
|
||||
if placement.rank() == rank and placement.device() != tensor.device:
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
spec.placements[i] = _remote_device(f"rank:{rank}/{tensor.device}")
|
||||
|
||||
return spec
|
||||
|
||||
|
||||
def _chunk_tensor(
|
||||
tensor: torch.Tensor,
|
||||
rank: int,
|
||||
world_size: int,
|
||||
num_devices_per_node: int,
|
||||
pg: dist.ProcessGroup,
|
||||
) -> torch.Tensor:
|
||||
if type(tensor) is ShardedTensor:
|
||||
if len(tensor.local_shards()) != 1:
|
||||
raise AssertionError
|
||||
|
||||
inner_param = tensor.local_tensor()
|
||||
inner_st = _create_chunk_sharded_tensor(
|
||||
inner_param,
|
||||
rank,
|
||||
world_size,
|
||||
num_devices_per_node,
|
||||
pg,
|
||||
)
|
||||
|
||||
outer_local_shard = tensor.local_shards()[0]
|
||||
shards: list[Shard] = [
|
||||
Shard(inner_st, copy.deepcopy(outer_local_shard.metadata))
|
||||
]
|
||||
st_meta = copy.deepcopy(tensor.metadata())
|
||||
st_meta.tensor_properties.requires_grad = False
|
||||
|
||||
st_outer = ShardedTensor._init_from_local_shards_and_global_metadata(
|
||||
shards,
|
||||
sharded_tensor_metadata=st_meta,
|
||||
process_group=tensor._process_group,
|
||||
init_rrefs=False,
|
||||
)
|
||||
return st_outer
|
||||
elif type(tensor) is DTensor:
|
||||
device_mesh = tensor.device_mesh
|
||||
if device_mesh.ndim != 1:
|
||||
raise AssertionError("Only 1D DeviceMeshes currently handled")
|
||||
|
||||
inner_param = tensor._local_tensor
|
||||
|
||||
inner_st = _create_chunk_sharded_tensor(
|
||||
inner_param,
|
||||
rank,
|
||||
world_size,
|
||||
torch.accelerator.device_count(),
|
||||
pg,
|
||||
)
|
||||
|
||||
dt_pg = _get_dt_pg(tensor)
|
||||
# We do this differently here, we create a ST with no local shards then patch it
|
||||
shards = [
|
||||
Shard(inner_st, _create_shard_md_from_dt(tensor, dist.get_rank(dt_pg)))
|
||||
]
|
||||
|
||||
st_meta = _create_sharded_tensor_md_from_dt(tensor, dt_pg)
|
||||
st_meta.tensor_properties.requires_grad = False
|
||||
|
||||
st_outer = ShardedTensor._init_from_local_shards_and_global_metadata(
|
||||
shards,
|
||||
sharded_tensor_metadata=st_meta,
|
||||
process_group=dt_pg,
|
||||
init_rrefs=False,
|
||||
)
|
||||
|
||||
return st_outer
|
||||
else:
|
||||
return _create_chunk_sharded_tensor(
|
||||
tensor,
|
||||
rank,
|
||||
world_size,
|
||||
num_devices_per_node,
|
||||
pg,
|
||||
)
|
||||
|
||||
|
||||
def _chunk_dtensor(
|
||||
tensor: torch.Tensor,
|
||||
rank: int,
|
||||
device_mesh: DeviceMesh,
|
||||
) -> DTensor:
|
||||
"""
|
||||
Shard a tensor to chunks along the first dimension.
|
||||
|
||||
The local rank will gets its corresponding chunk as the local tensor to create a DTensor.
|
||||
"""
|
||||
root_mesh = device_mesh._get_root_mesh() if device_mesh is not None else None
|
||||
if root_mesh is None:
|
||||
raise RuntimeError("No parent device_mesh is found for FSDP device_mesh.")
|
||||
if root_mesh.ndim < 2:
|
||||
raise RuntimeError(
|
||||
f"Found parent device_mesh of ndim={root_mesh.ndim},",
|
||||
"but meshes must be at least 2D.",
|
||||
)
|
||||
|
||||
# We need to explicitly call .detach() to return a new tensor detached from the current graph.
|
||||
tensor = tensor.detach().clone()
|
||||
|
||||
# When a layer is not involved in TP, then the tensor will not be a DTensor.
|
||||
# e.g. When a layer is not sppecified in the parallelize_plan, TP will have no effect on the layer.
|
||||
# e.g. When you do PairwiseParallel on a 3 layer model, TP will have no effect on the third layer.
|
||||
if isinstance(tensor, torch.Tensor) and not isinstance(tensor, DTensor):
|
||||
# For tensors, it is replicated across tp dimension and sharded across FSDP dimension.
|
||||
# TP is the inner dimension and FSDP is the outer dimension.
|
||||
# Therefore, shard placements for tensor is (Shard(0), Replicate()).
|
||||
replicate_placements = [Replicate() for _ in range(root_mesh.ndim)]
|
||||
shard_placements = [Replicate() for _ in range(root_mesh.ndim)]
|
||||
shard_placements[0] = DShard(0) # type: ignore[call-overload]
|
||||
|
||||
return DTensor.from_local(
|
||||
tensor, root_mesh, replicate_placements, run_check=False
|
||||
).redistribute(
|
||||
device_mesh=root_mesh,
|
||||
placements=shard_placements,
|
||||
)
|
||||
|
||||
else:
|
||||
tp_placements = tensor.placements
|
||||
tp_placement = tp_placements[0]
|
||||
|
||||
tensor = tensor.to_local()
|
||||
|
||||
# For DTensors, it is sharded across tp dimension first and then sharded across FSDP dimension.
|
||||
# TP is the inner dimension and FSDP is the outer dimension.
|
||||
# Therefore, shard placements for tensor is (Shard(0), tp_placement).
|
||||
# For higher dimensional meshes, it is replicated across other dimensions. For example, with
|
||||
# HSDP the shard placements for tensor is (Replicate, Shard(0), tp_placement).
|
||||
replicate_placements = [Replicate() for _ in range(root_mesh.ndim)]
|
||||
replicate_placements[-1] = tp_placement # type: ignore[call-overload]
|
||||
shard_placements = [Replicate() for i in range(root_mesh.ndim)] # type: ignore[misc]
|
||||
shard_placements[-2] = DShard(0) # type: ignore[call-overload]
|
||||
shard_placements[-1] = tp_placement # type: ignore[call-overload]
|
||||
|
||||
return DTensor.from_local(
|
||||
tensor, root_mesh, replicate_placements, run_check=False
|
||||
).redistribute(
|
||||
device_mesh=root_mesh,
|
||||
placements=shard_placements,
|
||||
)
|
||||
|
||||
|
||||
def _pre_load_state_dict(
|
||||
tensor: torch.Tensor,
|
||||
) -> tuple[torch.Tensor, list[Shard]]:
|
||||
shards = cast(ShardedTensor, tensor).local_shards()
|
||||
if len(shards) == 1 and type(shards[0].tensor) is ShardedTensor:
|
||||
inner_tensor = shards[0].tensor
|
||||
shards = inner_tensor.local_shards() # pyre-ignore[16]
|
||||
tensor = inner_tensor
|
||||
|
||||
return (tensor, shards if len(shards) > 0 else [])
|
||||
|
||||
|
||||
def _all_gather_dtensor(
|
||||
tensor: DTensor,
|
||||
parent_mesh: DeviceMesh | None,
|
||||
) -> torch.Tensor:
|
||||
"""All gather a DTensor in its FSDP dimension and return the local tensor."""
|
||||
if parent_mesh != tensor.device_mesh:
|
||||
raise AssertionError
|
||||
|
||||
placements = list(copy.deepcopy(tensor.placements))
|
||||
# FSDP + TP: [Shard(0), tp_placement] -> [Replicate(), tp_placement]
|
||||
# HSDP + TP: [Replicate(), Shard(0), tp_placement] -> [Replicate(), Replicate(), tp_placement]
|
||||
for i in range(len(placements) - 1):
|
||||
placements[i] = Replicate()
|
||||
tensor = tensor.redistribute(
|
||||
device_mesh=tensor.device_mesh,
|
||||
placements=placements,
|
||||
)
|
||||
|
||||
return tensor.to_local()
|
||||
|
||||
|
||||
class DTensorExtensions(FSDPExtensions):
|
||||
"""
|
||||
DTensorExtension is the TensorFlattener extension needed for 2D FSDP + TP.
|
||||
|
||||
This is the implementation for FSDPExtensions defined in
|
||||
https://github.com/pytorch/pytorch/blob/main/torch/distributed/fsdp/_fsdp_extensions.py
|
||||
"""
|
||||
|
||||
def __init__(self, device_handle) -> None:
|
||||
super().__init__()
|
||||
self.compute_stream = None
|
||||
self.device_handle = device_handle
|
||||
# we have to use the dynamo disable this way to disable dynamo as the decorator way would
|
||||
# trigger build failure with torch deploy...
|
||||
self.post_unflatten_transform = torch._dynamo.disable( # type: ignore[method-assign]
|
||||
self.post_unflatten_transform
|
||||
)
|
||||
|
||||
def pre_flatten_transform(
|
||||
self,
|
||||
tensor: torch.Tensor,
|
||||
) -> tuple[torch.Tensor, Any | None]:
|
||||
return _flatten_tensor(tensor)
|
||||
|
||||
def post_unflatten_transform(
|
||||
self, tensor: torch.Tensor, param_extension: Any
|
||||
) -> torch.Tensor:
|
||||
stream = self.compute_stream or self.device_handle.current_stream()
|
||||
with self.device_handle.stream(stream):
|
||||
# runtime we put the unflattened tensor call on the compute stream since
|
||||
# the unflattened tensor might contain computations in fwd/bwd where we
|
||||
# need to sync properly.
|
||||
# TODO: this is a short term fix and we should make the get_unflat_views
|
||||
# directly happen in the compute stream.
|
||||
result = _unflatten_tensor(
|
||||
tensor,
|
||||
param_extension,
|
||||
device_handle=self.device_handle,
|
||||
compute_stream=self.compute_stream,
|
||||
)
|
||||
_set_fsdp_flattened(result)
|
||||
return result
|
||||
|
||||
def chunk_tensor(
|
||||
self,
|
||||
tensor: torch.Tensor,
|
||||
rank: int,
|
||||
world_size: int,
|
||||
num_devices_per_node: int,
|
||||
pg: dist.ProcessGroup,
|
||||
device: torch.device | None = None,
|
||||
) -> torch.Tensor:
|
||||
return _chunk_tensor(tensor, rank, world_size, num_devices_per_node, pg)
|
||||
|
||||
def chunk_dtensor(
|
||||
self,
|
||||
tensor: torch.Tensor,
|
||||
rank: int,
|
||||
device_mesh: DeviceMesh,
|
||||
) -> torch.Tensor:
|
||||
return _chunk_dtensor(tensor, rank, device_mesh)
|
||||
|
||||
def pre_load_state_dict_transform(
|
||||
self,
|
||||
tensor: torch.Tensor,
|
||||
) -> tuple[torch.Tensor, list[Shard]]:
|
||||
return _pre_load_state_dict(tensor)
|
||||
|
||||
def all_gather_dtensor(
|
||||
self,
|
||||
tensor: DTensor,
|
||||
parent_mesh: DeviceMesh | None,
|
||||
) -> torch.Tensor:
|
||||
return _all_gather_dtensor(tensor, parent_mesh)
|
||||
+107
@@ -0,0 +1,107 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates
|
||||
from functools import partial
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from torch.distributed.tensor import DeviceMesh, DTensor, Replicate, Shard
|
||||
from torch.distributed.tensor.placement_types import _is_shard_like
|
||||
|
||||
|
||||
__all__ = [
|
||||
"input_reshard",
|
||||
]
|
||||
|
||||
|
||||
def input_reshard(
|
||||
module: torch.nn.Module,
|
||||
tp_device_mesh: DeviceMesh,
|
||||
input_reshard_dim: int | None = None,
|
||||
) -> torch.nn.Module:
|
||||
"""
|
||||
Register hooks to an nn.Module for input resharding, enabling sharding and restoration during backward computation.
|
||||
|
||||
Register hooks to an nn.Module with input resharding so that we can shard
|
||||
per the given `tp_device_mesh` and `input_reshard_dim` and restore the
|
||||
input back when recomputing the activations in the backward. The reason
|
||||
why we can do this is that for Tensor Parallel(TP), the input are same
|
||||
across all TP ranks.
|
||||
|
||||
Args:
|
||||
module (:class:`nn.Module`):
|
||||
Module to be registered with input resharding.
|
||||
tp_device_mesh (:class:`DeviceMesh`):
|
||||
Object which describes the mesh topology
|
||||
of devices for Tensor Parallel.
|
||||
input_reshard_dim (Optional[int]):
|
||||
The dimension of where we perform the sharding
|
||||
of input. If set None, there is no sharding of input.
|
||||
Default: None
|
||||
|
||||
Return:
|
||||
A :class:`nn.Module` object registered with TP input resharding.
|
||||
"""
|
||||
if input_reshard_dim is None:
|
||||
return module
|
||||
|
||||
cx: torch.autograd.graph.saved_tensors_hooks | None = None
|
||||
|
||||
def input_reshard_forward_pre_hook(_: torch.nn.Module, _i: tuple[Any, ...]) -> None:
|
||||
saved_tensor_hooks = torch.autograd.graph.saved_tensors_hooks(
|
||||
partial(_pack_hook_tp, tp_device_mesh, input_reshard_dim),
|
||||
partial(_unpack_hook_tp, tp_device_mesh, input_reshard_dim),
|
||||
)
|
||||
saved_tensor_hooks.__enter__()
|
||||
nonlocal cx
|
||||
cx = saved_tensor_hooks # type: ignore[name-defined]
|
||||
|
||||
def input_reshard_backward_hook(
|
||||
_: torch.nn.Module, _i: tuple[Any, ...], _o: Any
|
||||
) -> Any:
|
||||
nonlocal cx
|
||||
cx.__exit__() # type: ignore[name-defined, union-attr]
|
||||
|
||||
module.register_forward_pre_hook(input_reshard_forward_pre_hook)
|
||||
module.register_forward_hook(input_reshard_backward_hook)
|
||||
return module
|
||||
|
||||
|
||||
def _pack_hook_tp(mesh: DeviceMesh, input_reshard_dim: int, x: torch.Tensor) -> Any: # noqa: D401
|
||||
"""Hook function called after FWD to shard input."""
|
||||
if isinstance(x, DTensor) and all(p.is_replicate() for p in x._spec.placements):
|
||||
return x.redistribute(device_mesh=mesh, placements=[Shard(input_reshard_dim)])
|
||||
elif (
|
||||
not isinstance(x, DTensor)
|
||||
and isinstance(x, torch.Tensor)
|
||||
and x.numel() >= mesh.size()
|
||||
):
|
||||
return (
|
||||
DTensor.from_local(x, device_mesh=mesh)
|
||||
.redistribute(device_mesh=mesh, placements=[Shard(input_reshard_dim)])
|
||||
.to_local()
|
||||
)
|
||||
else:
|
||||
return x
|
||||
|
||||
|
||||
def _unpack_hook_tp(mesh: DeviceMesh, input_reshard_dim: int, x: Any) -> torch.Tensor: # noqa: D401
|
||||
"""Hook function called before activation recomputing in BWD to restore input."""
|
||||
if (
|
||||
isinstance(x, DTensor)
|
||||
and len(x._spec.placements) == 1
|
||||
and _is_shard_like(x._spec.placements[0])
|
||||
):
|
||||
return x.redistribute(device_mesh=mesh, placements=[Replicate()])
|
||||
elif (
|
||||
not isinstance(x, DTensor)
|
||||
and isinstance(x, torch.Tensor)
|
||||
and x.numel() >= mesh.size()
|
||||
):
|
||||
return (
|
||||
DTensor.from_local(
|
||||
x, device_mesh=mesh, placements=[Shard(input_reshard_dim)]
|
||||
)
|
||||
.redistribute(device_mesh=mesh, placements=[Replicate()])
|
||||
.to_local()
|
||||
)
|
||||
else:
|
||||
return x
|
||||
@@ -0,0 +1,514 @@
|
||||
# mypy: allow-untyped-defs
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates
|
||||
import contextlib
|
||||
from typing import cast
|
||||
|
||||
import torch
|
||||
import torch._prims_common as utils
|
||||
import torch.distributed._functional_collectives as funcol
|
||||
import torch.distributed.distributed_c10d as c10d
|
||||
from torch import Tensor
|
||||
from torch.distributed.device_mesh import DeviceMesh
|
||||
from torch.distributed.tensor import DTensor, Replicate, Shard
|
||||
from torch.distributed.tensor._dtensor_spec import DTensorSpec, TensorMeta
|
||||
from torch.distributed.tensor._ops._embedding_ops import _MaskPartial
|
||||
from torch.distributed.tensor._ops._math_ops import (
|
||||
_skip_dim,
|
||||
Reduction,
|
||||
replicate_reduction_dims,
|
||||
)
|
||||
from torch.distributed.tensor._ops.utils import normalize_dim
|
||||
from torch.distributed.tensor.placement_types import Placement
|
||||
|
||||
|
||||
aten = torch.ops.aten
|
||||
|
||||
|
||||
__all__ = ["loss_parallel"]
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def loss_parallel():
|
||||
"""
|
||||
A context manager that enables loss parallelism, where efficient parallelized loss computation
|
||||
can be performed when the input is sharded on the class dimension. Currently only the cross-entropy
|
||||
loss is supported.
|
||||
|
||||
Within this context manager, one can use :func:`~torch.nn.functional.cross_entropy` or
|
||||
:class:`~torch.nn.CrossEntropyLoss` as usual, with the following assumptions on the input parameters.
|
||||
The corresponding ``backward()`` call, if any, also needs to happen under this context manager.
|
||||
|
||||
Args:
|
||||
input (:class:`DTensor`):
|
||||
Input logits. Assumed to be sharded on the class dimension.
|
||||
target (Union[:class:`torch.Tensor`, :class:`DTensor`]):
|
||||
Must be ground truth class indices (class probabilities currently not supported).
|
||||
Assumed to be replicated across the ``DeviceMesh``.
|
||||
weight (Union[:class:`torch.Tensor`, :class:`DTensor`], optional):
|
||||
If given, assumed to be replicated across the ``DeviceMesh``.
|
||||
label_smoothing:
|
||||
Currently not supported.
|
||||
|
||||
Returns:
|
||||
A replicated :class:`DTensor`.
|
||||
|
||||
Example:
|
||||
A sharded DTensor is manually created here to showcase the usage.
|
||||
In practice, it is usually the output of a TP module.
|
||||
|
||||
>>> # xdoctest: +SKIP("distributed")
|
||||
>>> from torch.distributed.tensor.parallel import loss_parallel
|
||||
>>> from torch.distributed.device_mesh import init_device_mesh
|
||||
>>> ...
|
||||
>>> device_mesh = init_device_mesh("cuda", (8,))
|
||||
>>> input = torch.randn(4, 16, device="cuda", requires_grad=True)
|
||||
>>> dist_input = distribute_tensor(input, device_mesh, placements=[Shard(1)])
|
||||
>>> target = torch.randint(16, (4,), device="cuda")
|
||||
>>> with loss_parallel():
|
||||
>>> loss = F.cross_entropy(dist_input, target, reduction="mean")
|
||||
>>> loss.backward()
|
||||
>>> ...
|
||||
"""
|
||||
_enable_custom_loss_ops()
|
||||
|
||||
yield
|
||||
|
||||
_disable_custom_loss_ops()
|
||||
|
||||
|
||||
# Currently only needs to support one dimensional DeviceMesh; in general return
|
||||
# the mesh_dim with placements[mesh_dim].is_shard(dim)
|
||||
def _find_all_reduce_mesh_dim(placements: tuple[Placement, ...], dim: int) -> int:
|
||||
if not len(placements) == 1:
|
||||
raise ValueError(
|
||||
"Currently loss_parallel() only supports input on one-dimensional DeviceMesh."
|
||||
)
|
||||
if not placements[0].is_shard(dim):
|
||||
raise ValueError(
|
||||
f"loss_parallel() should be enabled only when the input tensor is sharded on dimension {dim}."
|
||||
)
|
||||
return 0
|
||||
|
||||
|
||||
def _cast_to_dtensor(
|
||||
tensor, placements: tuple[Placement, ...], mesh: DeviceMesh
|
||||
) -> DTensor:
|
||||
if isinstance(tensor, DTensor):
|
||||
if tensor.placements == placements:
|
||||
return tensor
|
||||
else:
|
||||
raise RuntimeError(f"Expected {placements} but got {tensor.placements}.")
|
||||
elif isinstance(tensor, torch.Tensor):
|
||||
return DTensor.from_local(
|
||||
tensor, device_mesh=mesh, placements=placements, run_check=False
|
||||
)
|
||||
else:
|
||||
raise TypeError(f"Unsupported type {type(tensor)}")
|
||||
|
||||
|
||||
def _propagate_tensor_meta(
|
||||
op_call: torch._ops.OpOverload,
|
||||
args: tuple[object, ...],
|
||||
kwargs: dict[str, object],
|
||||
) -> TensorMeta:
|
||||
op_info = DTensor._op_dispatcher.unwrap_to_op_info(op_call, args, kwargs)
|
||||
if op_info.schema is None:
|
||||
raise AssertionError(
|
||||
"op_info.schema should not be None after unwrap_to_op_info"
|
||||
)
|
||||
tensor_meta = DTensor._op_dispatcher.sharding_propagator._propagate_tensor_meta(
|
||||
op_info.schema
|
||||
)
|
||||
if isinstance(tensor_meta, TensorMeta):
|
||||
return tensor_meta
|
||||
elif isinstance(tensor_meta, tuple):
|
||||
# pyrefly: ignore [bad-return]
|
||||
return tensor_meta[0]
|
||||
else:
|
||||
raise RuntimeError(f"Unexpected tensor meta type: {type(tensor_meta)}.")
|
||||
|
||||
|
||||
# NOTE: The implementation follows torch._decomp.decomposition._log_softmax,
|
||||
# with all_reduce manually inserted to perform distributed computation.
|
||||
def _log_softmax(x, dim, half_to_float, mesh, mesh_dim):
|
||||
if half_to_float:
|
||||
if x.dtype != torch.half:
|
||||
raise AssertionError
|
||||
computation_dtype, result_dtype = utils.elementwise_dtypes(
|
||||
x, type_promotion_kind=utils.ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT
|
||||
)
|
||||
x = x.to(dtype=computation_dtype, memory_format=torch.contiguous_format)
|
||||
if x.numel() == 0:
|
||||
shifted = x
|
||||
else:
|
||||
x_max = torch.amax(x, dim, keepdim=True)
|
||||
x_max = funcol.all_reduce(
|
||||
x_max, reduceOp=c10d.ReduceOp.MAX.name, group=(mesh, mesh_dim)
|
||||
)
|
||||
shifted = x - x_max
|
||||
shifted_sumexp = torch.sum(torch.exp(shifted), dim, keepdim=True)
|
||||
shifted_sumexp = funcol.all_reduce(
|
||||
shifted_sumexp, reduceOp=c10d.ReduceOp.SUM.name, group=(mesh, mesh_dim)
|
||||
)
|
||||
shifted_logsumexp = torch.log(shifted_sumexp)
|
||||
result = shifted - shifted_logsumexp
|
||||
if not half_to_float:
|
||||
result = result.to(result_dtype)
|
||||
return result
|
||||
|
||||
|
||||
def _log_softmax_handler(
|
||||
op_call: torch._ops.OpOverload,
|
||||
args: tuple[object, ...],
|
||||
kwargs: dict[str, object],
|
||||
) -> object:
|
||||
x = cast(DTensor, args[0])
|
||||
dim = cast(int, args[1])
|
||||
half_to_float = cast(bool, args[2])
|
||||
|
||||
spec = x._spec
|
||||
dim = normalize_dim(dim, x.dim())
|
||||
mesh_dim = _find_all_reduce_mesh_dim(spec.placements, dim)
|
||||
|
||||
output_tensor_meta = _propagate_tensor_meta(op_call, args, kwargs)
|
||||
|
||||
res = _log_softmax(x._local_tensor, dim, half_to_float, spec.mesh, mesh_dim)
|
||||
|
||||
res_spec = DTensorSpec(
|
||||
spec.mesh,
|
||||
spec.placements,
|
||||
tensor_meta=output_tensor_meta,
|
||||
)
|
||||
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
return DTensor(
|
||||
# pyrefly: ignore [bad-argument-count]
|
||||
res,
|
||||
res_spec,
|
||||
# pyrefly: ignore [unexpected-keyword]
|
||||
requires_grad=res.requires_grad,
|
||||
)
|
||||
|
||||
|
||||
# NOTE: As explained below at _nll_loss_and_log_softmax_backward, the
|
||||
# _log_softmax_backward_handler does not actually do any computation.
|
||||
def _log_softmax_backward_handler(
|
||||
op_call: torch._ops.OpOverload,
|
||||
args: tuple[object, ...],
|
||||
kwargs: dict[str, object],
|
||||
) -> object:
|
||||
grad_output = cast(DTensor, args[0])
|
||||
input_dtype = cast(torch.dtype, args[3])
|
||||
return grad_output.to(input_dtype)
|
||||
|
||||
|
||||
# NOTE: The implementation follows torch._decomp.decomposition._nll_loss_forward,
|
||||
# with customized communication inserted to perform distributed computation.
|
||||
def _nll_loss_forward(
|
||||
x: Tensor,
|
||||
target: Tensor,
|
||||
weight: Tensor | None,
|
||||
local_weight: Tensor | None,
|
||||
reduction: int,
|
||||
ignore_index: int,
|
||||
input_shape: torch.Size,
|
||||
channel_dim: int,
|
||||
mesh: DeviceMesh,
|
||||
mesh_dim: int,
|
||||
) -> tuple[Tensor, Tensor]:
|
||||
n_dims = x.dim()
|
||||
channel_dim = 1
|
||||
if n_dims < 2:
|
||||
channel_dim = 0
|
||||
|
||||
def _weight_view(weight: Tensor) -> Tensor:
|
||||
if n_dims > 1:
|
||||
shape = [
|
||||
1,
|
||||
] * n_dims
|
||||
shape[channel_dim] = weight.shape[0]
|
||||
w = weight.view(shape)
|
||||
else:
|
||||
w = weight
|
||||
return w
|
||||
|
||||
if weight is not None:
|
||||
w = _weight_view(weight)
|
||||
if local_weight is None:
|
||||
raise AssertionError
|
||||
local_w = _weight_view(local_weight)
|
||||
x = x * local_w
|
||||
safe_target = torch.where(target != ignore_index, target, 0)
|
||||
safe_target_ = safe_target.unsqueeze(channel_dim)
|
||||
|
||||
# The following code block is a distributed version of
|
||||
# result = -torch.gather(self, channel_dim, safe_target_).squeeze(channel_dim)
|
||||
partial_placement = _MaskPartial(offset_shape=input_shape, offset_dim=channel_dim)
|
||||
safe_target_partial_ = partial_placement._partition_value(
|
||||
safe_target_, mesh, mesh_dim
|
||||
)
|
||||
result_partial = torch.gather(x, channel_dim, safe_target_partial_)
|
||||
# an all_reduce happens here
|
||||
result_reduced = partial_placement._reduce_value(result_partial, mesh, mesh_dim)
|
||||
result = -result_reduced.squeeze(channel_dim)
|
||||
|
||||
result = torch.where(target != ignore_index, result, 0)
|
||||
|
||||
if reduction == Reduction.NONE.value and n_dims > 1:
|
||||
total_weight = x.new_full((), 0.0)
|
||||
return result, total_weight
|
||||
|
||||
if weight is not None:
|
||||
new_shape = list(x.shape)
|
||||
new_shape[channel_dim] = -1
|
||||
# pyrefly: ignore [unbound-name]
|
||||
w = w.expand(new_shape)
|
||||
wsum = torch.gather(w, channel_dim, safe_target_).squeeze(channel_dim)
|
||||
wsum = torch.where(target != ignore_index, wsum, 0)
|
||||
total_weight = wsum.sum()
|
||||
else:
|
||||
total_weight = (target != ignore_index).sum().to(x)
|
||||
|
||||
# NOTE: this is correct only on 1D DeviceMesh; o/w additional
|
||||
# all-reduce on result and total_weight is needed
|
||||
if reduction == Reduction.SUM.value:
|
||||
result = result.sum()
|
||||
elif reduction == Reduction.MEAN.value:
|
||||
result = result.sum() / total_weight
|
||||
|
||||
return result, total_weight
|
||||
|
||||
|
||||
def _nll_loss_forward_handler(
|
||||
op_call: torch._ops.OpOverload,
|
||||
args: tuple[object, ...],
|
||||
kwargs: dict[str, object],
|
||||
) -> object:
|
||||
x = cast(DTensor, args[0])
|
||||
target = args[1]
|
||||
weight = args[2]
|
||||
reduction = cast(int, args[3])
|
||||
ignore_index = cast(int, args[4])
|
||||
|
||||
channel_dim = 1 if x.dim() >= 2 else 0
|
||||
spec = x._spec
|
||||
mesh_dim = _find_all_reduce_mesh_dim(spec.placements, channel_dim)
|
||||
|
||||
# Check user input: if target and weight are not DTensors, convert them to DTensors;
|
||||
# if they are DTensors, check that they have the desired placements.
|
||||
target_placements = _skip_dim(
|
||||
replicate_reduction_dims(spec.placements, [channel_dim]), channel_dim
|
||||
)
|
||||
all_replicate_placements = (Replicate(),) * spec.mesh.ndim
|
||||
target = _cast_to_dtensor(target, target_placements, spec.mesh)
|
||||
local_weight = None
|
||||
if weight is not None:
|
||||
weight = _cast_to_dtensor(weight, all_replicate_placements, spec.mesh)
|
||||
# For local computation, both (replicated) weight and (sharded) local_weight
|
||||
# are needed in _nll_loss_forward(). local_weight is generated here using
|
||||
# DTensor API, without incurring any communication.
|
||||
sharded_placements = [
|
||||
Shard(0) if i == mesh_dim else Replicate() for i in range(spec.mesh.ndim)
|
||||
]
|
||||
local_weight = weight.redistribute(spec.mesh, sharded_placements)._local_tensor
|
||||
if local_weight.shape[0] != x._local_tensor.shape[channel_dim]:
|
||||
raise AssertionError
|
||||
|
||||
if reduction == Reduction.NONE.value:
|
||||
output_placements = target_placements
|
||||
else:
|
||||
output_placements = all_replicate_placements
|
||||
|
||||
# tensor inputs to _propagate_tensor_meta need to be DTensors
|
||||
# pyrefly: ignore [bad-assignment]
|
||||
args = list(args)
|
||||
# pyrefly: ignore [unsupported-operation]
|
||||
args[1], args[2] = target, weight
|
||||
output_tensor_meta = _propagate_tensor_meta(op_call, tuple(args), kwargs)
|
||||
|
||||
result, total_weight = _nll_loss_forward(
|
||||
x._local_tensor,
|
||||
target._local_tensor,
|
||||
weight._local_tensor if weight is not None else None,
|
||||
local_weight,
|
||||
reduction,
|
||||
ignore_index,
|
||||
x.shape,
|
||||
channel_dim,
|
||||
spec.mesh,
|
||||
mesh_dim,
|
||||
)
|
||||
out_spec = DTensorSpec(spec.mesh, output_placements, tensor_meta=output_tensor_meta)
|
||||
|
||||
return (
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
DTensor(
|
||||
# pyrefly: ignore [bad-argument-count]
|
||||
result,
|
||||
out_spec,
|
||||
# pyrefly: ignore [unexpected-keyword]
|
||||
requires_grad=result.requires_grad,
|
||||
),
|
||||
total_weight,
|
||||
)
|
||||
|
||||
|
||||
# NOTE: The backward computation of cross_entropy goes through two steps:
|
||||
# backward for nll_loss and then backward for log_softmax. In loss parallel,
|
||||
# the two steps are fused into the following function (called by _nll_loss_backward_handler)
|
||||
# to avoid communication when target contains class indices not class probabilities.
|
||||
# Also note that the _log_softmax_backward_handler does not perform computation.
|
||||
# The implementation resembles _nll_loss_backward and _log_softmax_backward_data
|
||||
# from torch._decomp.decomposition.
|
||||
def _nll_loss_and_log_softmax_backward(
|
||||
grad_output: Tensor,
|
||||
x: Tensor,
|
||||
target: Tensor,
|
||||
weight: Tensor | None,
|
||||
reduction: int,
|
||||
ignore_index: int,
|
||||
total_weight: Tensor,
|
||||
input_shape: torch.Size,
|
||||
channel_dim: int,
|
||||
mesh: DeviceMesh,
|
||||
mesh_dim: int,
|
||||
) -> Tensor:
|
||||
channel_dim = 0 if x.dim() < 2 else 1
|
||||
if reduction == Reduction.MEAN.value:
|
||||
grad_output = grad_output / total_weight
|
||||
|
||||
target = target.unsqueeze(channel_dim)
|
||||
safe_target = torch.where(target != ignore_index, target, 0)
|
||||
grad_input = torch.zeros_like(x)
|
||||
|
||||
# The following code block is a distributed version of
|
||||
# grad_input = torch.scatter(grad_input, channel_dim, safe_target, -1.0)
|
||||
partial_placement = _MaskPartial(offset_shape=input_shape, offset_dim=channel_dim)
|
||||
safe_target = safe_target.squeeze(channel_dim).flatten()
|
||||
masked_safe_target = partial_placement._partition_value(safe_target, mesh, mesh_dim)
|
||||
# only update grad_input to -1 if not masked
|
||||
if partial_placement.mask_buffer.data is None:
|
||||
raise AssertionError
|
||||
grad_update = partial_placement.mask_buffer.data.to(grad_input.dtype) - 1.0
|
||||
arange_1d = torch.arange(
|
||||
masked_safe_target.shape[0], device=masked_safe_target.device
|
||||
)
|
||||
# The first two cases with x.dim() <= 2 are for aten.nll_loss_backward.default;
|
||||
# the last case is for aten.nll_loss2d_backward.default.
|
||||
if x.dim() == 1:
|
||||
grad_input[masked_safe_target] = grad_update
|
||||
elif x.dim() == 2:
|
||||
grad_input[arange_1d, masked_safe_target] = grad_update
|
||||
else:
|
||||
grad_input_t = grad_input.transpose(channel_dim, -1)
|
||||
intermidate_shape = grad_input_t.shape
|
||||
grad_input_2d = grad_input_t.reshape(-1, x.shape[channel_dim])
|
||||
grad_input_2d[arange_1d, masked_safe_target] = grad_update
|
||||
grad_input = grad_input_2d.view(intermidate_shape).transpose(channel_dim, -1)
|
||||
|
||||
if grad_input.dim() > grad_output.dim() > 0:
|
||||
grad_output = grad_output.unsqueeze(channel_dim)
|
||||
|
||||
if weight is not None:
|
||||
new_shape = [1 for _ in range(x.dim())]
|
||||
new_shape[channel_dim] = weight.shape[0]
|
||||
weight = weight.reshape(new_shape)
|
||||
# In order for fused computation to work, the following line is rewritten.
|
||||
# grad_output = grad_output * weight
|
||||
new_shape = list(x.shape)
|
||||
new_shape[channel_dim] = -1
|
||||
w = weight.expand(new_shape)
|
||||
w_target = torch.gather(w, channel_dim, target)
|
||||
grad_output = grad_output * w_target
|
||||
|
||||
grad_output = torch.where(target != ignore_index, grad_output, 0)
|
||||
|
||||
# NOTE: Instead of directly returning the grad_input as grad_output for log_softmax,
|
||||
# here we perform backward computation for log_softmax altogether to avoid the
|
||||
# otherwise extra all_gather communication.
|
||||
# return grad_input * grad_output
|
||||
return (grad_input + torch.exp(x)) * grad_output
|
||||
|
||||
|
||||
def _nll_loss_backward_handler(
|
||||
op_call: torch._ops.OpOverload,
|
||||
args: tuple[object, ...],
|
||||
kwargs: dict[str, object],
|
||||
) -> object:
|
||||
grad_output = cast(DTensor, args[0])
|
||||
x = cast(DTensor, args[1])
|
||||
target = args[2]
|
||||
weight = args[3]
|
||||
reduction = cast(int, args[4])
|
||||
ignore_index = cast(int, args[5])
|
||||
total_weight = cast(Tensor, args[6])
|
||||
|
||||
channel_dim = 1 if x.dim() >= 2 else 0
|
||||
spec = x._spec
|
||||
mesh_dim = _find_all_reduce_mesh_dim(spec.placements, channel_dim)
|
||||
|
||||
# if target and weight are not DTensors, convert them to DTensors
|
||||
target_placements = _skip_dim(
|
||||
replicate_reduction_dims(spec.placements, [channel_dim]), channel_dim
|
||||
)
|
||||
all_replicate_placements = (Replicate(),) * spec.mesh.ndim
|
||||
target = _cast_to_dtensor(target, target_placements, spec.mesh)
|
||||
if weight is not None:
|
||||
weight = _cast_to_dtensor(weight, all_replicate_placements, spec.mesh)
|
||||
|
||||
# tensor inputs to _propagate_tensor_meta need to be DTensors
|
||||
# pyrefly: ignore [bad-assignment]
|
||||
args = list(args)
|
||||
# pyrefly: ignore [unsupported-operation]
|
||||
args[2], args[3] = target, weight
|
||||
# pyrefly: ignore [unsupported-operation]
|
||||
args[6] = _cast_to_dtensor(total_weight, all_replicate_placements, spec.mesh)
|
||||
output_tensor_meta = _propagate_tensor_meta(op_call, tuple(args), kwargs)
|
||||
|
||||
result = _nll_loss_and_log_softmax_backward(
|
||||
grad_output._local_tensor,
|
||||
x._local_tensor,
|
||||
target._local_tensor,
|
||||
weight._local_tensor if weight is not None else None,
|
||||
reduction,
|
||||
ignore_index,
|
||||
total_weight,
|
||||
x.shape,
|
||||
channel_dim,
|
||||
spec.mesh,
|
||||
mesh_dim,
|
||||
)
|
||||
# the output sharding is the same as input sharding: Shard(channel_dim) on mesh_dim
|
||||
out_spec = DTensorSpec(
|
||||
spec.mesh,
|
||||
spec.placements,
|
||||
tensor_meta=output_tensor_meta,
|
||||
)
|
||||
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
return DTensor(
|
||||
# pyrefly: ignore [bad-argument-count]
|
||||
result,
|
||||
out_spec,
|
||||
# pyrefly: ignore [unexpected-keyword]
|
||||
requires_grad=result.requires_grad,
|
||||
)
|
||||
|
||||
|
||||
customized_loss_ops = {
|
||||
aten._log_softmax.default: _log_softmax_handler,
|
||||
aten._log_softmax_backward_data.default: _log_softmax_backward_handler,
|
||||
aten.nll_loss_forward.default: _nll_loss_forward_handler,
|
||||
aten.nll_loss2d_forward.default: _nll_loss_forward_handler,
|
||||
aten.nll_loss_backward.default: _nll_loss_backward_handler,
|
||||
aten.nll_loss2d_backward.default: _nll_loss_backward_handler,
|
||||
}
|
||||
|
||||
|
||||
def _enable_custom_loss_ops():
|
||||
DTensor._op_dispatcher._custom_op_handlers.update(customized_loss_ops)
|
||||
|
||||
|
||||
def _disable_custom_loss_ops():
|
||||
for custom_op in customized_loss_ops:
|
||||
DTensor._op_dispatcher._custom_op_handlers.pop(custom_op)
|
||||
@@ -0,0 +1,823 @@
|
||||
# mypy: allow-untyped-defs
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates
|
||||
from abc import ABC, abstractmethod
|
||||
from functools import partial
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.distributed.tensor import (
|
||||
DeviceMesh,
|
||||
distribute_module,
|
||||
distribute_tensor,
|
||||
DTensor,
|
||||
Replicate,
|
||||
Shard,
|
||||
)
|
||||
from torch.distributed.tensor.placement_types import Placement
|
||||
|
||||
|
||||
__all__ = [
|
||||
"ParallelStyle",
|
||||
"RowwiseParallel",
|
||||
"SequenceParallel",
|
||||
"ColwiseParallel",
|
||||
"PrepareModuleInput",
|
||||
"PrepareModuleInputOutput",
|
||||
"PrepareModuleOutput",
|
||||
]
|
||||
|
||||
|
||||
class ParallelStyle(ABC):
|
||||
"""
|
||||
The parallel style contract defines how the module or submodule should be parallelized.
|
||||
|
||||
It only defines the ``apply`` method for ``parallelize_module`` to use, this allows maximum
|
||||
flexibility for different kind of style implementations.
|
||||
"""
|
||||
|
||||
src_data_rank: int | None = 0
|
||||
|
||||
@abstractmethod
|
||||
def _apply(self, module: nn.Module, device_mesh: DeviceMesh) -> nn.Module: ...
|
||||
|
||||
|
||||
class ColwiseParallel(ParallelStyle):
|
||||
"""
|
||||
Partition a compatible nn.Module in a column-wise fashion. Currently supports nn.Linear and nn.Embedding.
|
||||
Users can compose it together with RowwiseParallel to achieve the sharding of more complicated modules.
|
||||
(i.e. MLP, Attention)
|
||||
|
||||
Keyword Args:
|
||||
input_layouts (Placement, optional):
|
||||
The DTensor layout of input tensor for the nn.Module, this is used to annotate the input tensor to
|
||||
become a DTensor. If not specified, we assume the input tensor to be replicated.
|
||||
output_layouts (Placement, optional):
|
||||
The DTensor layout of the output for the nn.Module, this is used to ensure the output of the nn.Module
|
||||
with the user desired layout. If not specified, the output tensor is sharded on the last dimension.
|
||||
use_local_output (bool, optional):
|
||||
Whether to use local :class:`torch.Tensor` instead of :class:`DTensor` for the module output, default: True.
|
||||
Returns:
|
||||
A :class:`ParallelStyle` object that represents Colwise sharding of the nn.Module.
|
||||
|
||||
Example::
|
||||
>>> # xdoctest: +SKIP(failing)
|
||||
>>> from torch.distributed.tensor.parallel import parallelize_module, ColwiseParallel
|
||||
>>> from torch.distributed.device_mesh import init_device_mesh
|
||||
>>> ...
|
||||
>>> m = Model(...) # m is a nn.Module that contains a "w1" nn.Linear submodule
|
||||
>>> tp_mesh = init_device_mesh("cuda", (8,))
|
||||
>>>
|
||||
>>> # By default, the input of the "w1" Linear will be converted to Replicated DTensor
|
||||
>>> # and the output of "w1" will return :class:`torch.Tensor` that shards on the last dim.
|
||||
>>>
|
||||
>>> sharded_mod = parallelize_module(m, tp_mesh, {"w1": ColwiseParallel()})
|
||||
>>> ...
|
||||
|
||||
.. note:: By default ``ColwiseParallel`` output is sharded on the last dimension if the ``output_layouts`` not
|
||||
specified, if there're operators that require specific tensor shape (i.e. before the paired ``RowwiseParallel``),
|
||||
keep in mind that if the output is sharded the operator might need to be adjusted to the sharded size.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
input_layouts: Placement | None = None,
|
||||
output_layouts: Placement | None = None,
|
||||
use_local_output: bool = True,
|
||||
):
|
||||
super().__init__()
|
||||
self.input_layouts = (input_layouts or Replicate(),)
|
||||
self.output_layouts = (output_layouts or Shard(-1),)
|
||||
# colwise linear runtime sharding (desired sharding):
|
||||
# 1. requires replicate input
|
||||
# 2. shard output on last dim
|
||||
self.desired_input_layouts = (Replicate(),)
|
||||
self.use_local_output = use_local_output
|
||||
|
||||
@staticmethod
|
||||
def _prepare_input_fn(
|
||||
input_layouts, desired_input_layouts, mod, inputs, device_mesh
|
||||
):
|
||||
# TODO: figure out dynamo support for instance method and switch this to instance method
|
||||
|
||||
# annotate module input placements/sharding with input_layouts
|
||||
input_tensor = inputs[0]
|
||||
if not isinstance(input_tensor, DTensor):
|
||||
input_tensor = DTensor.from_local(
|
||||
input_tensor,
|
||||
device_mesh,
|
||||
input_layouts,
|
||||
run_check=False,
|
||||
)
|
||||
|
||||
# transform the input layouts to the desired layouts of ColwiseParallel
|
||||
if input_layouts != desired_input_layouts:
|
||||
input_tensor = input_tensor.redistribute(
|
||||
placements=desired_input_layouts, async_op=True
|
||||
)
|
||||
return input_tensor
|
||||
|
||||
def _partition_linear_fn(self, name, module, device_mesh):
|
||||
# colwise shard weight/bias to Shard(0), weight be Shard(0)
|
||||
# means Colwise as Linear is input * weight^T + bias, where
|
||||
# weight would become Shard(1)
|
||||
for name, param in module.named_parameters():
|
||||
dist_param = nn.Parameter(
|
||||
distribute_tensor(
|
||||
param, device_mesh, [Shard(0)], src_data_rank=self.src_data_rank
|
||||
),
|
||||
requires_grad=param.requires_grad,
|
||||
)
|
||||
module.register_parameter(name, dist_param)
|
||||
|
||||
def _partition_embedding_fn(self, name, module, device_mesh):
|
||||
# colwise shard embedding.weight is straight forward as Shard(1)
|
||||
for name, param in module.named_parameters():
|
||||
dist_param = nn.Parameter(
|
||||
distribute_tensor(
|
||||
param, device_mesh, [Shard(1)], src_data_rank=self.src_data_rank
|
||||
),
|
||||
requires_grad=param.requires_grad,
|
||||
)
|
||||
module.register_parameter(name, dist_param)
|
||||
|
||||
@staticmethod
|
||||
def _prepare_output_fn(output_layouts, use_local_output, mod, outputs, device_mesh):
|
||||
# outputs is a shard on last dimension DTensor, i.e. Shard(-1)
|
||||
if outputs.placements != output_layouts:
|
||||
outputs = outputs.redistribute(placements=output_layouts, async_op=True)
|
||||
# back to local tensor
|
||||
return outputs.to_local() if use_local_output else outputs
|
||||
|
||||
def _apply(self, module: nn.Module, device_mesh: DeviceMesh) -> nn.Module:
|
||||
if isinstance(module, nn.Linear):
|
||||
partition_fn = self._partition_linear_fn
|
||||
elif isinstance(module, nn.Embedding):
|
||||
partition_fn = self._partition_embedding_fn
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
"ColwiseParallel currently only support nn.Linear and nn.Embedding!"
|
||||
)
|
||||
|
||||
return distribute_module(
|
||||
module,
|
||||
device_mesh,
|
||||
partition_fn,
|
||||
partial(
|
||||
self._prepare_input_fn,
|
||||
self.input_layouts,
|
||||
self.desired_input_layouts,
|
||||
),
|
||||
partial(
|
||||
self._prepare_output_fn, self.output_layouts, self.use_local_output
|
||||
),
|
||||
)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
tmpstr = self.__class__.__name__ + "("
|
||||
tmpstr += f"input_layouts={self.input_layouts}, "
|
||||
tmpstr += f"output_layouts={self.output_layouts}, "
|
||||
tmpstr += f"use_local_output={self.use_local_output}"
|
||||
tmpstr += ")"
|
||||
return tmpstr
|
||||
|
||||
|
||||
class RowwiseParallel(ParallelStyle):
|
||||
"""
|
||||
Partition a compatible nn.Module in a row-wise fashion. Currently supports nn.Linear and nn.Embedding.
|
||||
Users can compose it with ColwiseParallel to achieve the sharding of more complicated modules.
|
||||
(i.e. MLP, Attention)
|
||||
|
||||
Keyword Args:
|
||||
input_layouts (Placement, optional):
|
||||
The DTensor layout of input tensor for the nn.Module, this is used to annotate the input tensor to
|
||||
become a DTensor. If not specified, we assume the input tensor to be sharded on the last dimension.
|
||||
output_layouts (Placement, optional):
|
||||
The DTensor layout of the output for the nn.Module, this is used to ensure the output of the nn.Module
|
||||
with the user desired layout. If not specified, the output tensor is replicated.
|
||||
use_local_output (bool, optional):
|
||||
Whether to use local :class:`torch.Tensor` instead of :class:`DTensor` for the module output, default: True.
|
||||
Returns:
|
||||
A :class:`ParallelStyle` object that represents Rowwise sharding of the nn.Module.
|
||||
|
||||
Example::
|
||||
>>> # xdoctest: +SKIP(failing)
|
||||
>>> from torch.distributed.tensor.parallel import parallelize_module, RowwiseParallel
|
||||
>>> from torch.distributed.device_mesh import init_device_mesh
|
||||
>>> ...
|
||||
>>> m = Model(...) # m is a nn.Module that contains a "w2" nn.Linear submodule
|
||||
>>> tp_mesh = init_device_mesh("cuda", (8,))
|
||||
>>>
|
||||
>>> # By default, the input of the "w2" Linear will be converted to DTensor that shards on the last dim
|
||||
>>> # and the output of "w2" will return a replicated :class:`torch.Tensor`.
|
||||
>>>
|
||||
>>> sharded_mod = parallelize_module(m, tp_mesh, {"w2": RowwiseParallel()}),
|
||||
>>> ...
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
input_layouts: Placement | None = None,
|
||||
output_layouts: Placement | None = None,
|
||||
use_local_output: bool = True,
|
||||
):
|
||||
super().__init__()
|
||||
self.input_layouts = (input_layouts or Shard(-1),)
|
||||
self.output_layouts = (output_layouts or Replicate(),)
|
||||
self.use_local_output = use_local_output
|
||||
|
||||
@staticmethod
|
||||
def _prepare_input_fn(
|
||||
input_layouts, desired_input_layouts, mod, inputs, device_mesh
|
||||
):
|
||||
input_tensor = inputs[0]
|
||||
if not isinstance(input_tensor, DTensor):
|
||||
input_tensor = DTensor.from_local(
|
||||
input_tensor,
|
||||
device_mesh,
|
||||
input_layouts,
|
||||
run_check=False,
|
||||
)
|
||||
|
||||
if input_layouts != desired_input_layouts:
|
||||
input_tensor = input_tensor.redistribute(
|
||||
placements=desired_input_layouts, async_op=True
|
||||
)
|
||||
return input_tensor
|
||||
|
||||
def _partition_linear_fn(self, name, module, device_mesh):
|
||||
# Rowwise shard weight to Shard(1), bias to Replicate(), weight be Shard(1)
|
||||
# means Rowwise as nn.Linear is input * weight^T + bias, where
|
||||
# weight would become Shard(0)
|
||||
module.register_parameter(
|
||||
"weight",
|
||||
nn.Parameter(
|
||||
distribute_tensor(
|
||||
module.weight,
|
||||
device_mesh,
|
||||
[Shard(1)],
|
||||
src_data_rank=self.src_data_rank,
|
||||
),
|
||||
requires_grad=module.weight.requires_grad,
|
||||
),
|
||||
)
|
||||
if getattr(module, "bias", None) is not None:
|
||||
# The Linear module has bias
|
||||
module.register_parameter(
|
||||
"bias",
|
||||
nn.Parameter(
|
||||
distribute_tensor(
|
||||
module.bias,
|
||||
device_mesh,
|
||||
[Replicate()],
|
||||
src_data_rank=self.src_data_rank,
|
||||
),
|
||||
requires_grad=module.bias.requires_grad,
|
||||
),
|
||||
)
|
||||
|
||||
def _partition_embedding_fn(self, name, module, device_mesh):
|
||||
# rowwise shard embedding.weight is Shard(0)
|
||||
for name, param in module.named_parameters():
|
||||
dist_param = nn.Parameter(
|
||||
distribute_tensor(
|
||||
param, device_mesh, [Shard(0)], src_data_rank=self.src_data_rank
|
||||
),
|
||||
requires_grad=param.requires_grad,
|
||||
)
|
||||
module.register_parameter(name, dist_param)
|
||||
|
||||
@staticmethod
|
||||
def _prepare_output_fn(output_layouts, use_local_output, mod, outputs, device_mesh):
|
||||
# Rowwise sharding produces partial output, depending on output layouts:
|
||||
# 1. to replicate -> allreduce
|
||||
# 2. to shard -> reduce_scatter
|
||||
if outputs.placements != output_layouts:
|
||||
outputs = outputs.redistribute(placements=output_layouts, async_op=True)
|
||||
# back to local tensor if use_local_output is True
|
||||
return outputs.to_local() if use_local_output else outputs
|
||||
|
||||
def _apply(self, module: nn.Module, device_mesh: DeviceMesh) -> nn.Module:
|
||||
if isinstance(module, nn.Linear):
|
||||
partition_fn = self._partition_linear_fn
|
||||
# rowwise linear runtime sharding requires input tensor shard on last dim
|
||||
self.desired_input_layouts: tuple[Placement, ...] = (Shard(-1),)
|
||||
elif isinstance(module, nn.Embedding):
|
||||
partition_fn = self._partition_embedding_fn
|
||||
# rowwise embedding runtime sharding requires input tensor replicated
|
||||
self.desired_input_layouts = (Replicate(),)
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
"RowwiseParallel currently only support nn.Linear and nn.Embedding!"
|
||||
)
|
||||
|
||||
return distribute_module(
|
||||
module,
|
||||
device_mesh,
|
||||
partition_fn,
|
||||
partial(
|
||||
self._prepare_input_fn,
|
||||
self.input_layouts,
|
||||
self.desired_input_layouts,
|
||||
),
|
||||
partial(
|
||||
self._prepare_output_fn, self.output_layouts, self.use_local_output
|
||||
),
|
||||
)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
tmpstr = self.__class__.__name__ + "("
|
||||
tmpstr += f"input_layouts={self.input_layouts}, "
|
||||
tmpstr += f"output_layouts={self.output_layouts}, "
|
||||
tmpstr += f"use_local_output={self.use_local_output}"
|
||||
tmpstr += ")"
|
||||
return tmpstr
|
||||
|
||||
|
||||
class SequenceParallel(ParallelStyle):
|
||||
"""
|
||||
SequenceParallel replicates a compatible ``nn.Module`` parameters and runs the sharded computation with
|
||||
input sharded on the sequence dimension. This currently supports ``nn.LayerNorm``, ``nn.Dropout``, and the
|
||||
`RMSNorm python implementation <https://github.com/facebookresearch/llama/blob/main/llama/model.py#L34>`__
|
||||
|
||||
This style implements the operation that is described in the paper
|
||||
`Reducing Activation Recomputation in Large Transformer Models <https://arxiv.org/abs/2205.05198>`__
|
||||
|
||||
If the input passed in to this ``nn.Module`` is a :class:`torch.Tensor`, it assumes that the input is already sharded
|
||||
on the sequence dimension and converts the input to a :class:`DTensor` sharded on the sequence dimension. If the input
|
||||
passed in to this ``nn.Module`` is already a :class:`DTensor` but is not sharded on the sequence dimension, it would
|
||||
redistribute the input to be sharded on the sequence dimension.
|
||||
|
||||
The output of the ``nn.Module`` will be sharded on the sequence dimension.
|
||||
|
||||
Keyword Args:
|
||||
sequence_dim (int, optional):
|
||||
The sequence dimension of the input tensor for the ``nn.Module``, this is used to annotate the input tensor to
|
||||
become a DTensor that is sharded on the sequence dimension, default: 1.
|
||||
use_local_output (bool, optional):
|
||||
Whether to use local :class:`torch.Tensor` instead of :class:`DTensor` for the module output, default: False.
|
||||
Returns:
|
||||
A :class:`ParallelStyle` object that represents Sequence Parallel of the ``nn.Module``.
|
||||
|
||||
Example::
|
||||
>>> # xdoctest: +SKIP(failing)
|
||||
>>> from torch.distributed.tensor.parallel import parallelize_module, SequenceParallel
|
||||
>>> from torch.distributed.device_mesh import init_device_mesh
|
||||
>>> ...
|
||||
>>> m = Model(...) # m is a nn.Module that contains a "norm" nn.LayerNorm submodule
|
||||
>>> tp_mesh = init_device_mesh("cuda", (8,))
|
||||
>>>
|
||||
>>> # By default, the input of the "norm" will be converted to DTensor that shards on the sequence dim
|
||||
>>> # and the output of "norm" will return a sharded on sequence dimension :class:`DTensor`.
|
||||
>>>
|
||||
>>> sharded_mod = parallelize_module(m, tp_mesh, {"norm": SequenceParallel()}),
|
||||
>>> ...
|
||||
|
||||
.. note:: SequenceParallel style assumes ones initialization if there are weights in the nn.Module (i.e.
|
||||
``nn.LayerNorm`` or ``RMSNorm``, and they by default have ones initialization). If you have custom
|
||||
inits for the weights on those modules, you need to broadcast the weights before/after parallelizing
|
||||
to ensure that they are replicated.
|
||||
"""
|
||||
|
||||
def __init__(self, *, sequence_dim: int = 1, use_local_output: bool = False):
|
||||
super().__init__()
|
||||
self.sequence_sharding = (Shard(sequence_dim),)
|
||||
self.use_local_output = use_local_output
|
||||
|
||||
def _replicate_module_fn(
|
||||
self, name: str, module: nn.Module, device_mesh: DeviceMesh
|
||||
):
|
||||
for p_name, param in module.named_parameters():
|
||||
# simple replication with fixed ones_ init from LayerNorm/RMSNorm, which allow
|
||||
# us to simply just use from_local
|
||||
replicated_param = torch.nn.Parameter(
|
||||
DTensor.from_local(param, device_mesh, [Replicate()], run_check=False)
|
||||
)
|
||||
module.register_parameter(p_name, replicated_param)
|
||||
|
||||
@staticmethod
|
||||
def _prepare_input_fn(sequence_sharding, mod, inputs, device_mesh):
|
||||
input_tensor = inputs[0]
|
||||
if isinstance(input_tensor, DTensor):
|
||||
# if the passed in input DTensor is not sharded on the sequence dim, we need to redistribute it
|
||||
if input_tensor.placements != sequence_sharding:
|
||||
input_tensor = input_tensor.redistribute(
|
||||
placements=sequence_sharding, async_op=True
|
||||
)
|
||||
return input_tensor
|
||||
elif isinstance(input_tensor, torch.Tensor):
|
||||
# assume the input passed in already sharded on the sequence dim and create the DTensor
|
||||
return DTensor.from_local(
|
||||
input_tensor, device_mesh, sequence_sharding, run_check=False
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"expecting input of {mod} to be a torch.Tensor or DTensor, but got {input_tensor}"
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _prepare_output_fn(use_local_output, mod, outputs, device_mesh):
|
||||
return outputs.to_local() if use_local_output else outputs
|
||||
|
||||
def _apply(self, module: nn.Module, device_mesh: DeviceMesh) -> nn.Module:
|
||||
return distribute_module(
|
||||
module,
|
||||
device_mesh,
|
||||
self._replicate_module_fn,
|
||||
partial(self._prepare_input_fn, self.sequence_sharding),
|
||||
partial(self._prepare_output_fn, self.use_local_output),
|
||||
)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
tmpstr = self.__class__.__name__ + "("
|
||||
if len(self.sequence_sharding) == 1:
|
||||
tmpstr += f"sequence_dim={self.sequence_sharding[0].dim}, "
|
||||
tmpstr += f"use_local_output={self.use_local_output}"
|
||||
tmpstr += ")"
|
||||
return tmpstr
|
||||
|
||||
|
||||
class PrepareModuleInput(ParallelStyle):
|
||||
"""
|
||||
Configure the nn.Module's inputs to convert the input tensors of the nn.Module to DTensors at runtime according to
|
||||
``input_layouts``, and perform layout redistribution according to the ``desired_input_layouts``.
|
||||
|
||||
Keyword Args:
|
||||
input_layouts (Union[Placement, Tuple[Optional[Placement]]]):
|
||||
The DTensor layouts of input tensors for the nn.Module, this is used to convert the input tensors to
|
||||
DTensors. If some inputs are not torch.Tensor or no need to convert to DTensors, ``None`` need to be specified
|
||||
as a placeholder. default: None.
|
||||
desired_input_layouts (Union[Placement, Tuple[Optional[Placement]]]):
|
||||
The desired DTensor layout of input tensors for the nn.Module, this is used to ensure the inputs of the nn.Module
|
||||
have the desired DTensor layouts. This argument needs to have the same length with ``input_layouts``. default: None.
|
||||
input_kwarg_layouts (Dict[str, Placement]):
|
||||
The DTensor layouts of input kwargs for the nn.Module, this is used to convert the input kwarg tensors to DTensors.
|
||||
default: None
|
||||
desired_input_kwarg_layouts: (Dict[str, Placement]):
|
||||
The desired DTensor layout of input kwargs for the nn.Module, this is used to ensure the inputs of the nn.Module
|
||||
have the desired DTensor layouts. default: None.
|
||||
use_local_output (bool, optional):
|
||||
Whether to use local :class:`torch.Tensor` instead of :class:`DTensor` for the module inputs, default: False.
|
||||
Returns:
|
||||
A :class:`ParallelStyle` object that prepares the sharding layouts of the nn.Module's inputs.
|
||||
|
||||
Example::
|
||||
>>> # xdoctest: +SKIP(failing)
|
||||
>>> from torch.distributed.tensor.parallel import parallelize_module, PrepareModuleInput
|
||||
>>> from torch.distributed.device_mesh import init_device_mesh
|
||||
>>> ...
|
||||
>>> block = TransformerBlock(...) # block is a nn.Module that contains an "attn" Attention submodule
|
||||
>>> tp_mesh = init_device_mesh("cuda", (8,))
|
||||
>>>
|
||||
>>> # According to the style specified below, the first input of attn will be annotated to Sharded DTensor
|
||||
>>> # and then redistributed to Replicated DTensor.
|
||||
>>> parallelize_module(
|
||||
>>> block, # this can be a submodule or module
|
||||
>>> tp_mesh,
|
||||
>>> parallelize_plan={
|
||||
>>> "attn": PrepareModuleInput(
|
||||
>>> input_layouts=(Shard(0), None, None, ...),
|
||||
>>> desired_input_layouts=(Replicate(), None, None, ...)
|
||||
>>> ),
|
||||
>>> }
|
||||
>>> )
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
input_layouts: Placement | tuple[Placement | None, ...] | None = None,
|
||||
desired_input_layouts: Placement | tuple[Placement | None, ...] | None = None,
|
||||
input_kwarg_layouts: dict[str, Placement] | None = None,
|
||||
desired_input_kwarg_layouts: dict[str, Placement] | None = None,
|
||||
use_local_output: bool = False,
|
||||
):
|
||||
self.input_layouts = (
|
||||
(input_layouts,) if isinstance(input_layouts, Placement) else input_layouts
|
||||
)
|
||||
self.desired_input_layouts = (
|
||||
(desired_input_layouts,)
|
||||
if isinstance(desired_input_layouts, Placement)
|
||||
else desired_input_layouts
|
||||
)
|
||||
self.use_local_output = use_local_output
|
||||
if self.input_layouts is not None:
|
||||
if self.desired_input_layouts is None:
|
||||
raise AssertionError("desired module inputs should not be None!")
|
||||
if len(self.input_layouts) != len(self.desired_input_layouts):
|
||||
raise AssertionError(
|
||||
"input_layouts and desired_input_layouts should have same length!"
|
||||
)
|
||||
self.with_kwargs = input_kwarg_layouts is not None
|
||||
self.input_kwarg_layouts = input_kwarg_layouts or {}
|
||||
self.desired_input_kwarg_layouts = desired_input_kwarg_layouts or {}
|
||||
if self.with_kwargs:
|
||||
if len(self.input_kwarg_layouts) != len(self.desired_input_kwarg_layouts):
|
||||
raise AssertionError(
|
||||
"input_kwarg_layouts and desired_input_kwarg_layouts should have same length!"
|
||||
)
|
||||
|
||||
def _prepare_input_arg(
|
||||
self,
|
||||
input: Any,
|
||||
mesh: DeviceMesh,
|
||||
input_layout: Placement | None,
|
||||
desired_layout: Placement | None,
|
||||
):
|
||||
if input_layout is not None:
|
||||
if isinstance(input, DTensor):
|
||||
# TODO: re-enable the check once we fix the compile path
|
||||
# assert inp.placements[0] == input_layout
|
||||
dt_inp = input
|
||||
else:
|
||||
if not isinstance(input, torch.Tensor):
|
||||
raise AssertionError("expecting input to be a torch.Tensor!")
|
||||
dt_inp = DTensor.from_local(
|
||||
input, mesh, (input_layout,), run_check=False
|
||||
)
|
||||
|
||||
if desired_layout is not None and input_layout != desired_layout:
|
||||
dt_inp = dt_inp.redistribute(placements=(desired_layout,))
|
||||
|
||||
return dt_inp.to_local() if self.use_local_output else dt_inp
|
||||
else:
|
||||
return input
|
||||
|
||||
def _prepare_input_fn(self, inputs, device_mesh):
|
||||
if self.input_layouts is None:
|
||||
return inputs
|
||||
prepared_inputs = []
|
||||
if not isinstance(inputs, tuple):
|
||||
inputs = (inputs,)
|
||||
if len(inputs) != len(self.input_layouts):
|
||||
raise ValueError("module inputs and input_layouts should have same length!")
|
||||
|
||||
if self.desired_input_layouts is None:
|
||||
raise AssertionError("desired module inputs should not be None!")
|
||||
|
||||
for inp, input_layout, desired_layout in zip(
|
||||
inputs, self.input_layouts, self.desired_input_layouts
|
||||
):
|
||||
prepared_inputs.append(
|
||||
self._prepare_input_arg(inp, device_mesh, input_layout, desired_layout)
|
||||
)
|
||||
return tuple(prepared_inputs)
|
||||
|
||||
def _prepare_input_kwarg_fn(self, inputs, kwarg_inputs, device_mesh):
|
||||
prepared_arg_inputs = self._prepare_input_fn(inputs, device_mesh)
|
||||
prepared_kwarg_inputs = {}
|
||||
for kwarg_key in kwarg_inputs:
|
||||
kwarg_val = kwarg_inputs[kwarg_key]
|
||||
input_layout = self.input_kwarg_layouts.get(kwarg_key)
|
||||
desired_input_layout = self.desired_input_kwarg_layouts.get(kwarg_key)
|
||||
|
||||
prepared_kwarg_inputs[kwarg_key] = self._prepare_input_arg(
|
||||
kwarg_val, device_mesh, input_layout, desired_input_layout
|
||||
)
|
||||
|
||||
return (prepared_arg_inputs, prepared_kwarg_inputs)
|
||||
|
||||
def _apply(self, module: nn.Module, device_mesh: DeviceMesh) -> nn.Module:
|
||||
if self.with_kwargs:
|
||||
module.register_forward_pre_hook(
|
||||
lambda _, inputs, kwargs: self._prepare_input_kwarg_fn(
|
||||
inputs, kwargs, device_mesh
|
||||
),
|
||||
with_kwargs=True,
|
||||
) # type: ignore[misc]
|
||||
else:
|
||||
module.register_forward_pre_hook(
|
||||
lambda _, inputs: self._prepare_input_fn(inputs, device_mesh)
|
||||
) # type: ignore[misc, call-arg]
|
||||
return module
|
||||
|
||||
def __repr__(self) -> str:
|
||||
tmpstr = self.__class__.__name__ + "("
|
||||
tmpstr += f"input_layouts={self.input_layouts}, "
|
||||
tmpstr += f"desired_input_layouts={self.desired_input_layouts}, "
|
||||
tmpstr += f"input_kwarg_layouts={self.input_kwarg_layouts}, "
|
||||
tmpstr += f"desired_input_kwarg_layouts={self.desired_input_kwarg_layouts}, "
|
||||
tmpstr += f"use_local_output={self.use_local_output}"
|
||||
tmpstr += ")"
|
||||
return tmpstr
|
||||
|
||||
|
||||
class PrepareModuleOutput(ParallelStyle):
|
||||
"""
|
||||
Configure the nn.Module's outputs to convert the output tensors of the nn.Module to DTensors at runtime according to
|
||||
``output_layouts``, and perform layout redistribution according to the ``desired_output_layouts``.
|
||||
|
||||
Keyword Args:
|
||||
output_layouts (Union[Placement, Tuple[Placement]]):
|
||||
The DTensor layouts of output tensors for the nn.Module, this is used to convert the output tensors to
|
||||
DTensors if they are :class:`torch.Tensor`. If some outputs are not torch.Tensor or no need to convert to DTensors,
|
||||
``None`` need to be specified as a placeholder.
|
||||
desired_output_layouts (Union[Placement, Tuple[Placement]]):
|
||||
The desired DTensor layouts of output tensors for the nn.Module, this is used to ensure the outputs of the nn.Module
|
||||
have the desired DTensor layouts.
|
||||
use_local_output (bool, optional):
|
||||
Whether to use local :class:`torch.Tensor` instead of :class:`DTensor` for the module outputs, default: True.
|
||||
Returns:
|
||||
A ParallelStyle object that prepares the sharding layouts of the nn.Module's outputs.
|
||||
|
||||
Example::
|
||||
>>> # xdoctest: +SKIP(failing)
|
||||
>>> from torch.distributed.tensor.parallel import parallelize_module, PrepareModuleOutput
|
||||
>>> from torch.distributed.device_mesh import init_device_mesh
|
||||
>>> ...
|
||||
>>> block = TransformerBlock(...) # block is a nn.Module that contains an "attn" Attention submodule
|
||||
>>> tp_mesh = init_device_mesh("cuda", (8,))
|
||||
>>>
|
||||
>>> # According to the style specified below, the output of the TransformerBlock will be converted to Replicated DTensor
|
||||
>>> # and then redistributed to Sharded DTensor.
|
||||
>>> parallelize_module(
|
||||
>>> block, # this can be a submodule or module
|
||||
>>> tp_mesh,
|
||||
>>> parallelize_plan = PrepareModuleOutput(
|
||||
>>> output_layouts=Replicate(),
|
||||
>>> desired_output_layouts=Shard(0)
|
||||
>>> )
|
||||
>>> )
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
output_layouts: Placement | tuple[Placement | None, ...],
|
||||
desired_output_layouts: Placement | tuple[Placement, ...],
|
||||
use_local_output: bool = True,
|
||||
):
|
||||
self.output_layouts = (
|
||||
(output_layouts,)
|
||||
if isinstance(output_layouts, Placement)
|
||||
else output_layouts
|
||||
)
|
||||
self.desired_output_layouts = (
|
||||
(desired_output_layouts,)
|
||||
if isinstance(desired_output_layouts, Placement)
|
||||
else desired_output_layouts
|
||||
)
|
||||
self.use_local_output = use_local_output
|
||||
if len(self.output_layouts) != len(self.desired_output_layouts):
|
||||
raise AssertionError(
|
||||
"output_layouts and desired_output_layouts should have same length!"
|
||||
)
|
||||
|
||||
def _prepare_out_fn(self, outputs, device_mesh):
|
||||
prepared_outputs = []
|
||||
if not isinstance(outputs, tuple):
|
||||
outputs = (outputs,)
|
||||
if len(outputs) != len(self.output_layouts):
|
||||
raise ValueError(
|
||||
"module outputs and output_layouts should have same length!"
|
||||
)
|
||||
|
||||
for out, out_layout, desired_out_layout in zip(
|
||||
outputs, self.output_layouts, self.desired_output_layouts
|
||||
):
|
||||
if out_layout is not None:
|
||||
if isinstance(out, DTensor):
|
||||
# TODO: re-enable the check once we fix the compile path
|
||||
# assert out.placements[0] == out_layout
|
||||
dt_out = out
|
||||
else:
|
||||
dt_out = DTensor.from_local(
|
||||
out, device_mesh, (out_layout,), run_check=False
|
||||
)
|
||||
|
||||
if out_layout != desired_out_layout:
|
||||
dt_out = dt_out.redistribute(placements=(desired_out_layout,))
|
||||
prepared_outputs.append(
|
||||
dt_out.to_local() if self.use_local_output else dt_out
|
||||
)
|
||||
else:
|
||||
prepared_outputs.append(out)
|
||||
if len(prepared_outputs) == 1:
|
||||
return prepared_outputs[0]
|
||||
else:
|
||||
return tuple(prepared_outputs)
|
||||
|
||||
def _apply(self, module: nn.Module, device_mesh: DeviceMesh) -> nn.Module:
|
||||
module.register_forward_hook(
|
||||
lambda _, inputs, outputs: self._prepare_out_fn(outputs, device_mesh)
|
||||
) # type: ignore[misc, call-arg]
|
||||
return module
|
||||
|
||||
def __repr__(self) -> str:
|
||||
tmpstr = self.__class__.__name__ + "("
|
||||
tmpstr += f"output_layouts={self.output_layouts}, "
|
||||
tmpstr += f"desired_output_layouts={self.desired_output_layouts}, "
|
||||
tmpstr += f"use_local_output={self.use_local_output}"
|
||||
tmpstr += ")"
|
||||
return tmpstr
|
||||
|
||||
|
||||
class PrepareModuleInputOutput(ParallelStyle):
|
||||
"""
|
||||
Configure the nn.Module's inputs (and outputs) to convert the input tensors (and output tensors, respectively) of the nn.Module
|
||||
to DTensors at runtime according to ``input_layouts`` (and output_layouts, respectively), and perform layout redistribution
|
||||
according to the ``desired_input_layouts`` (and ``desired_output_layouts``, respectively). This is a combination of
|
||||
:class:`PrepareModuleInput` and :class:`PrepareModuleOutput`.
|
||||
|
||||
Keyword Args:
|
||||
input_layouts (Union[Placement, Tuple[Optional[Placement]]]):
|
||||
The DTensor layouts of input tensors for the nn.Module, this is used to convert the input tensors to
|
||||
DTensors. If some inputs are not torch.Tensor or no need to convert to DTensors, ``None`` need to be specified
|
||||
as a placeholder. default: None.
|
||||
desired_input_layouts (Union[Placement, Tuple[Optional[Placement]]]):
|
||||
The desired DTensor layout of input tensors for the nn.Module, this is used to ensure the inputs of the nn.Module
|
||||
have the desired DTensor layouts. This argument needs to have the same length with ``input_layouts``. default: None.
|
||||
input_kwarg_layouts (Dict[str, Placement]):
|
||||
The DTensor layouts of input kwargs for the nn.Module, this is used to convert the input kwarg tensors to DTensors.
|
||||
default: None
|
||||
desired_input_kwarg_layouts: (Dict[str, Placement]):
|
||||
The desired DTensor layout of input kwargs for the nn.Module, this is used to ensure the inputs of the nn.Module
|
||||
have the desired DTensor layouts. default: None.
|
||||
use_local_input (bool, optional):
|
||||
Whether to use local :class:`torch.Tensor` instead of :class:`DTensor` for the module inputs, default: False.
|
||||
output_layouts (Union[Placement, Tuple[Placement]]):
|
||||
The DTensor layouts of output tensors for the nn.Module, this is used to convert the output tensors to
|
||||
DTensors if they are :class:`torch.Tensor`. If some outputs are not torch.Tensor or no need to convert to DTensors,
|
||||
``None`` need to be specified as a placeholder.
|
||||
desired_output_layouts (Union[Placement, Tuple[Placement]]):
|
||||
The desired DTensor layouts of output tensors for the nn.Module, this is used to ensure the outputs of the nn.Module
|
||||
have the desired DTensor layouts.
|
||||
use_local_output (bool, optional):
|
||||
Whether to use local :class:`torch.Tensor` instead of :class:`DTensor` for the module outputs, default: True.
|
||||
Returns:
|
||||
A :class:`ParallelStyle` object that prepares the sharding layouts of the nn.Module's inputs and outputs.
|
||||
|
||||
Example::
|
||||
>>> # xdoctest: +SKIP(failing)
|
||||
>>> from torch.distributed.tensor.parallel import parallelize_module, PrepareModuleInputOutput
|
||||
>>> from torch.distributed.device_mesh import init_device_mesh
|
||||
>>> ...
|
||||
>>> block = TransformerBlock(...) # block is a nn.Module that contains an "attn" Attention submodule
|
||||
>>> tp_mesh = init_device_mesh("cuda", (8,))
|
||||
>>>
|
||||
>>> # According to the style specified below, the first input of attn will be annotated as Sharded DTensor
|
||||
>>> # and then redistributed to Replicated DTensor, and the output of the TransformerBlock will be annotated
|
||||
>>> # as Replicated DTensor and then redistributed to Sharded DTensor.
|
||||
>>> parallelize_module(
|
||||
>>> block, # this can be a submodule or module
|
||||
>>> tp_mesh,
|
||||
>>> parallelize_plan={
|
||||
>>> "attn": PrepareModuleInputOutput(
|
||||
>>> input_layouts=(Shard(0), None, None, ...),
|
||||
>>> desired_input_layouts=(Replicate(), None, None, ...),
|
||||
>>> output_layouts=Replicate(),
|
||||
>>> desired_output_layouts=Shard(0),
|
||||
>>> ),
|
||||
>>> }
|
||||
>>> )
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
input_layouts: Placement | tuple[Placement | None, ...] | None = None,
|
||||
desired_input_layouts: Placement | tuple[Placement | None, ...] | None = None,
|
||||
input_kwarg_layouts: dict[str, Placement] | None = None,
|
||||
desired_input_kwarg_layouts: dict[str, Placement] | None = None,
|
||||
use_local_input: bool = False,
|
||||
output_layouts: Placement | tuple[Placement | None, ...],
|
||||
desired_output_layouts: Placement | tuple[Placement, ...],
|
||||
use_local_output: bool = True,
|
||||
):
|
||||
self.prepare_module_input = PrepareModuleInput(
|
||||
input_layouts=input_layouts,
|
||||
desired_input_layouts=desired_input_layouts,
|
||||
input_kwarg_layouts=input_kwarg_layouts,
|
||||
desired_input_kwarg_layouts=desired_input_kwarg_layouts,
|
||||
use_local_output=use_local_input,
|
||||
)
|
||||
self.prepare_module_output = PrepareModuleOutput(
|
||||
output_layouts=output_layouts,
|
||||
desired_output_layouts=desired_output_layouts,
|
||||
use_local_output=use_local_output,
|
||||
)
|
||||
|
||||
def _apply(self, module: nn.Module, device_mesh: DeviceMesh) -> nn.Module:
|
||||
self.prepare_module_input._apply(module, device_mesh)
|
||||
self.prepare_module_output._apply(module, device_mesh)
|
||||
|
||||
return module
|
||||
|
||||
def __repr__(self) -> str:
|
||||
tmpstr = self.__class__.__name__ + "("
|
||||
tmpstr += f"input_layouts={self.prepare_module_input.input_layouts}, "
|
||||
tmpstr += (
|
||||
f"desired_input_layouts={self.prepare_module_input.desired_input_layouts}, "
|
||||
)
|
||||
tmpstr += (
|
||||
f"input_kwarg_layouts={self.prepare_module_input.input_kwarg_layouts}, "
|
||||
)
|
||||
tmpstr += f"desired_input_kwarg_layouts={self.prepare_module_input.desired_input_kwarg_layouts}, "
|
||||
tmpstr += f"use_local_input={self.prepare_module_input.use_local_output}, "
|
||||
tmpstr += f"output_layouts={self.prepare_module_output.output_layouts}, "
|
||||
tmpstr += f"desired_output_layouts={self.prepare_module_output.desired_output_layouts}, "
|
||||
tmpstr += f"use_local_output={self.prepare_module_output.use_local_output}"
|
||||
tmpstr += ")"
|
||||
return tmpstr
|
||||
+1575
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user