|
| 1 | +# Copyright (c) Meta Platforms, Inc. and affiliates. |
| 2 | +# All rights reserved. |
| 3 | +# |
| 4 | +# This source code is licensed under the BSD-style license found in the |
| 5 | +# LICENSE file in the root directory of this source tree. |
| 6 | + |
| 7 | +"""Full DTensor infrastructure for SPMD-style parallelization. |
| 8 | +
|
| 9 | +When ``training.full_dtensor`` is enabled, all model parameters, buffers, |
| 10 | +and inputs become DTensors on a multi-dimensional SPMD mesh. FSDP uses |
| 11 | +``DataParallelMeshDims`` to identify which mesh dimensions are |
| 12 | +data-parallel. |
| 13 | +
|
| 14 | +TP sharding is handled by ``Module.parallelize(spmd_mesh, full_dtensor=True)`` |
| 15 | +using config-based ``ShardingSpec``. |
| 16 | +""" |
| 17 | + |
| 18 | +from typing import Any |
| 19 | + |
| 20 | +import torch |
| 21 | +import torch.nn as nn |
| 22 | +from torch.distributed.device_mesh import DeviceMesh |
| 23 | +from torch.distributed.fsdp import DataParallelMeshDims |
| 24 | +from torch.distributed.tensor import DTensor, Replicate, Shard |
| 25 | +from torch.distributed.tensor.placement_types import Placement |
| 26 | + |
| 27 | +from torchtitan.distributed.parallel_dims import ParallelDims |
| 28 | + |
| 29 | + |
| 30 | +def validate_config(parallel_dims: ParallelDims, model_config: Any) -> None: |
| 31 | + """Validate that the current configuration is compatible with full DTensor. |
| 32 | +
|
| 33 | + Raises NotImplementedError with a clear message if incompatible. |
| 34 | + """ |
| 35 | + if parallel_dims.ep_enabled: |
| 36 | + raise NotImplementedError( |
| 37 | + "full_dtensor is not supported with Expert Parallel. " |
| 38 | + "Disable EP or disable full_dtensor." |
| 39 | + ) |
| 40 | + |
| 41 | + layers = getattr(model_config, "layers", None) |
| 42 | + layer = layers[0] if layers else None |
| 43 | + attn_config = getattr(layer, "attention", None) if layer else None |
| 44 | + attn_backend = getattr(attn_config, "attn_backend", "sdpa") |
| 45 | + if attn_backend in ("flex", "varlen"): |
| 46 | + raise NotImplementedError( |
| 47 | + f"full_dtensor is not supported with {attn_backend} attention. " |
| 48 | + "Flex/varlen attention does not support DTensor dispatch. " |
| 49 | + "Use sdpa attention or disable full_dtensor." |
| 50 | + ) |
| 51 | + |
| 52 | + |
| 53 | +def get_dense_spmd_mesh(parallel_dims: ParallelDims) -> DeviceMesh: |
| 54 | + """Get the dense SPMD mesh for full DTensor parallelization. |
| 55 | +
|
| 56 | + Uses the ``full_dtensor_dense`` global mesh which has separate ``dp_shard`` |
| 57 | + and ``cp`` dimensions. Disabled dimensions (degree 1) are filtered out. |
| 58 | +
|
| 59 | + The result is cached on the ``ParallelDims`` object so that all callers |
| 60 | + share the exact same ``DeviceMesh`` object — FSDP requires object identity. |
| 61 | + """ |
| 62 | + if hasattr(parallel_dims, "_spmd_mesh"): |
| 63 | + return parallel_dims._spmd_mesh # type: ignore[return-value] |
| 64 | + |
| 65 | + mesh_names = [ |
| 66 | + n |
| 67 | + for n in ["dp_replicate", "dp_shard", "cp", "tp"] |
| 68 | + if parallel_dims.get_optional_mesh(n) |
| 69 | + ] |
| 70 | + assert mesh_names, "full_dtensor requires at least one mesh dimension" |
| 71 | + mesh = parallel_dims.get_mesh(mesh_names) |
| 72 | + parallel_dims._spmd_mesh = mesh # type: ignore[attr-defined] |
| 73 | + return mesh |
| 74 | + |
| 75 | + |
| 76 | +def get_dp_mesh_dims(parallel_dims: ParallelDims) -> DataParallelMeshDims: |
| 77 | + """Build DataParallelMeshDims for dense (non-MoE) parameters. |
| 78 | +
|
| 79 | + Uses ``dp_shard`` and ``cp`` as separate shard dimensions rather than |
| 80 | + the flattened ``fsdp`` from the non-full-dtensor path. |
| 81 | + """ |
| 82 | + shard_dims: list[str] = [] |
| 83 | + if parallel_dims.dp_shard_enabled: |
| 84 | + shard_dims.append("dp_shard") |
| 85 | + if parallel_dims.cp_enabled: |
| 86 | + shard_dims.append("cp") |
| 87 | + |
| 88 | + if len(shard_dims) > 1: |
| 89 | + shard: str | tuple[str, ...] | None = tuple(shard_dims) |
| 90 | + elif shard_dims: |
| 91 | + shard = shard_dims[0] |
| 92 | + else: |
| 93 | + shard = None |
| 94 | + |
| 95 | + replicate: str | None = None |
| 96 | + if parallel_dims.dp_replicate_enabled: |
| 97 | + replicate = "dp_replicate" |
| 98 | + |
| 99 | + return DataParallelMeshDims(shard=shard, replicate=replicate) |
| 100 | + |
| 101 | + |
| 102 | +def resolve_fsdp_mesh( |
| 103 | + model: nn.Module, |
| 104 | + parallel_dims: ParallelDims, |
| 105 | + full_dtensor: bool, |
| 106 | +) -> tuple[DeviceMesh, DataParallelMeshDims | None]: |
| 107 | + """Select the FSDP mesh and optional DataParallelMeshDims. |
| 108 | +
|
| 109 | + In full DTensor mode, returns the SPMD mesh and DataParallelMeshDims. |
| 110 | + In standard mode, returns the conventional dp_mesh and None. |
| 111 | + """ |
| 112 | + if full_dtensor: |
| 113 | + spmd_mesh = get_dense_spmd_mesh(parallel_dims) |
| 114 | + _remove_sdpa_math_backend(model) |
| 115 | + dp_mesh_dims = get_dp_mesh_dims(parallel_dims) |
| 116 | + return spmd_mesh, dp_mesh_dims |
| 117 | + else: |
| 118 | + names = ( |
| 119 | + ["dp_replicate", "fsdp"] if parallel_dims.dp_replicate_enabled else ["fsdp"] |
| 120 | + ) |
| 121 | + return parallel_dims.get_mesh(names), None |
| 122 | + |
| 123 | + |
| 124 | +def _remove_sdpa_math_backend(model: nn.Module) -> None: |
| 125 | + """Remove MATH backend from SDPA modules. |
| 126 | +
|
| 127 | + SDPA MATH backend decomposes into primitive ops that mix plain tensors |
| 128 | + with DTensors, causing errors. Flash and efficient backends have proper |
| 129 | + DTensor dispatch rules. |
| 130 | + """ |
| 131 | + from torch.nn.attention import SDPBackend |
| 132 | + |
| 133 | + from torchtitan.models.common.attention import ScaledDotProductAttention |
| 134 | + |
| 135 | + for module in model.modules(): |
| 136 | + if isinstance(module, ScaledDotProductAttention): |
| 137 | + if SDPBackend.MATH in module.sdpa_backends: |
| 138 | + module.sdpa_backends = [ |
| 139 | + b for b in module.sdpa_backends if b != SDPBackend.MATH |
| 140 | + ] |
| 141 | + |
| 142 | + |
| 143 | +def parallelize_inputs( |
| 144 | + parallel_dims: ParallelDims, |
| 145 | + inputs: torch.Tensor, |
| 146 | + labels: torch.Tensor, |
| 147 | + extra_kwargs: dict[str, torch.Tensor | None] | None = None, |
| 148 | +) -> tuple[DTensor, DTensor]: |
| 149 | + """Convert inputs, labels, and extra kwargs to DTensors on the SPMD mesh. |
| 150 | +
|
| 151 | + DP dims get Shard(0) (batch), CP gets Shard(1) (sequence), TP gets Replicate. |
| 152 | + Tensor values in extra_kwargs (e.g. positions) use the same placements. |
| 153 | + """ |
| 154 | + mesh = get_dense_spmd_mesh(parallel_dims) |
| 155 | + placements: list[Placement] = [] |
| 156 | + if parallel_dims.dp_replicate_enabled: |
| 157 | + placements.append(Shard(0)) |
| 158 | + if parallel_dims.dp_shard_enabled: |
| 159 | + placements.append(Shard(0)) |
| 160 | + if parallel_dims.cp_enabled: |
| 161 | + placements.append(Shard(1)) |
| 162 | + if parallel_dims.tp_enabled: |
| 163 | + placements.append(Replicate()) |
| 164 | + |
| 165 | + assert mesh.ndim == len(placements) |
| 166 | + |
| 167 | + # Convert extra_kwargs tensors (e.g. positions) to DTensors |
| 168 | + if extra_kwargs is not None: |
| 169 | + for key, value in extra_kwargs.items(): |
| 170 | + if isinstance(value, torch.Tensor) and not isinstance(value, DTensor): |
| 171 | + extra_kwargs[key] = DTensor.from_local(value, mesh, placements) |
| 172 | + |
| 173 | + return DTensor.from_local(inputs, mesh, placements), DTensor.from_local( |
| 174 | + labels, mesh, placements |
| 175 | + ) |
0 commit comments