# 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, Partial, 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. Both one-dimensional and multi-dimensional ``DeviceMesh`` are supported. On a multi-dimensional mesh, exactly one mesh dim must shard the class dimension — that dim is treated as the "TP" dim and can be in any position (not necessarily last). The remaining mesh dims may be ``Shard`` (e.g. batch-sharded for DP/CP) or ``Replicate``. Args: input (:class:`DTensor`): Input logits. Assumed to be sharded on the class dimension on exactly one mesh dim. target (Union[:class:`torch.Tensor`, :class:`DTensor`]): Must be ground truth class indices (class probabilities currently not supported). Its placements are derived from ``input``'s placements by (1) replacing the TP dim's ``Shard(class_dim)`` with ``Replicate`` and (2) shifting any other ``Shard(d)`` with ``d > class_dim`` down by one (since the class tensor dim is removed from the output):: input tensor shape input placements target shape target placements ---------------------- -------------------------- --------------- ------------------------ (batch, class) (Shard(0), Shard(1)) (batch,) (Shard(0), Replicate()) (batch, class) (Replicate(), Shard(1)) (batch,) (Replicate(), Replicate()) (batch, class, seq) (Shard(2), Shard(1)) (batch, seq) (Shard(1), Replicate()) (batch, class, seq) (Shard(0), Shard(1)) (batch, seq) (Shard(0), Replicate()) A plain :class:`torch.Tensor` is only accepted when the derived target placements are all ``Replicate`` (e.g. on a one-dimensional mesh); otherwise a :class:`DTensor` must be passed explicitly so that the intended sharding is unambiguous. weight (Union[:class:`torch.Tensor`, :class:`DTensor`], optional): If given, assumed to be replicated across the ``DeviceMesh``. label_smoothing: Currently not supported. Returns: A :class:`DTensor` whose placements depend on ``reduction``: - ``reduction="none"`` — per-sample loss; inherits the ``target`` placements (``Shard`` on batch-parallel dims, ``Replicate`` on TP). - ``reduction="sum"`` — scalar; ``Replicate`` on the TP dim (the custom all-reduce makes each rank's local value globally correct over TP), ``Partial("sum")`` on every other ``Shard`` mesh dim so that the cross-rank reduction happens lazily on materialization/redistribution, and ``Replicate`` on ``Replicate`` mesh dims. - ``reduction="mean"`` — only supported on a one-dimensional mesh; returns a fully replicated :class:`DTensor`. On a one-dimensional mesh all three cases simplify to a fully replicated :class:`DTensor` (the original behavior). .. note:: ``reduction="mean"`` is only supported on a one-dimensional ``DeviceMesh``. On a multi-dimensional mesh the per-rank division by ``total_weight`` uses a local count, so aggregating across non-TP dims does not yield the correct global mean; use ``reduction="sum"`` or ``"none"`` instead and divide by the global count yourself if needed. 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() # Return the mesh_dim whose placement is Shard(dim). Exactly one such mesh dim # (the "TP" dim) is required; other mesh dims must be Shard on a different # tensor dim or Replicate. Partial on a non-TP dim is rejected so that # reduction="none" cannot silently accept garbage (replicate_reduction_dims # would otherwise rewrite Partial -> Replicate in target_placements). def _find_all_reduce_mesh_dim(placements: tuple[Placement, ...], dim: int) -> int: shard_mesh_dims = [i for i, p in enumerate(placements) if p.is_shard(dim)] if len(shard_mesh_dims) != 1: raise ValueError( f"loss_parallel() requires exactly one mesh dim to shard tensor " f"dimension {dim}, but got placements {placements} with " f"{len(shard_mesh_dims)} such mesh dim(s)." ) mesh_dim = shard_mesh_dims[0] for i, p in enumerate(placements): if i == mesh_dim: continue if not (p.is_shard() or p.is_replicate()): raise ValueError( f"loss_parallel() expects non-TP mesh dims to be " f"Shard or Replicate, got {p} at dim {i}." ) return mesh_dim 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): if any(p.is_shard() for p in placements): # A plain torch.Tensor with Shard placements is ambiguous: it could # be a full global tensor the user expects us to shard, or an # already-local slice. Require a DTensor so the intent is explicit. raise ValueError( f"loss_parallel() requires a DTensor (not a plain torch.Tensor) " f"when the derived placements {placements} contain Shard — this " f"happens on a multi-dimensional mesh with a batch-sharded non-TP " f"dim. Wrap the tensor with DTensor.from_local or distribute_tensor " f"to make the sharding explicit." ) 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) # On an N-D mesh, result.sum() is a per-rank sum over the local batch; # the caller emits Partial("sum") on non-TP Shard dims so the cross-rank # reduction happens when the output is later materialized/redistributed. # reduction="mean" is rejected in the handler for N-D mesh because # total_weight is a per-rank local count. 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: # reduction="mean" is not yet correct on multi-dimensional mesh: the # local division by total_weight uses a per-rank count, so summing # across non-TP dims does not yield the global mean. See the NOTE in # _nll_loss_forward. if reduction == Reduction.MEAN.value and spec.mesh.ndim > 1: raise NotImplementedError( "loss_parallel() with reduction='mean' is only supported on " "one-dimensional DeviceMesh; got mesh with ndim=" f"{spec.mesh.ndim}. Use reduction='sum' or 'none' instead." ) # On the shard dim (TP), the custom all-reduce produces a globally # correct result -> Replicate. On other dims (e.g. DP/CP), each rank # only reduced its local batch -> Partial(sum). out_placements_list: list[Placement] = [] for i, p in enumerate(spec.placements): if i == mesh_dim: out_placements_list.append(Replicate()) elif p.is_shard(): out_placements_list.append(Partial()) else: out_placements_list.append(p) output_placements = tuple(out_placements_list) # 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) # For reduction="none", autograd may deliver grad_output with different # placements than the forward output (e.g. Replicate on a batch-sharded dim), # causing a shape mismatch with the local batch. Redistribute it to match # target_placements — which is what the forward output uses for reduction="none". # For reduction="sum"/"mean" the forward output is a 0-D scalar, so a placement # mismatch (e.g. Replicate vs. Partial("sum")) cannot produce a shape mismatch # against x._local_tensor and no redistribute is needed here. if reduction == Reduction.NONE.value: grad_output = grad_output.redistribute(spec.mesh, target_placements) # tensor inputs to _propagate_tensor_meta need to be DTensors # pyrefly: ignore [bad-assignment] args = list(args) # pyrefly: ignore [unsupported-operation] args[0] = grad_output # 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)