已合并
[feature]分布式支持context parallel #35982
TrHan创建于 5月18日
[feature]分布式支持context parallel #35982
已合并
共 8 个文件变更+1137-3
| @@ -32,6 +32,8 @@ PyTorch Ascend Adapter named torch_npu is committed to maintaining compatibility | |||
| 32 | 32 | ||
| 33 | - **[torch.distributed.optim](#torchdistributedoptim)** | 33 | - **[torch.distributed.optim](#torchdistributedoptim)** |
| 34 | 34 | ||
| 35 | +- **[torch.distributed.tensor.experimental](#torchdistributedtensorexperimental)** | ||
| 36 | + | ||
| 35 | - **[torch.distributed.tensor.parallel](#torchdistributedtensorparallel)** | 37 | - **[torch.distributed.tensor.parallel](#torchdistributedtensorparallel)** |
| 36 | 38 | ||
| 37 | - **[torch.distributed.checkpoint](#torchdistributedcheckpoint)** | 39 | - **[torch.distributed.checkpoint](#torchdistributedcheckpoint)** |
| @@ -122,9 +124,10 @@ PyTorch Ascend Adapter named torch_npu is committed to maintaining compatibility | |||
| 122 | 124 | ||
| 123 | - **[Understanding CUDA Memory Usage](#understanding-cuda-memory-usage)** | 125 | - **[Understanding CUDA Memory Usage](#understanding-cuda-memory-usage)** |
| 124 | 126 | ||
| 125 | -> **NOTE:** | 127 | +>  **NOTE:** |
| 126 | ->1. Marked as compatible, please note if there are any limitations. | 128 | +> |
| 127 | ->2. Not indicating whether it is compatible means that the compatibility of this API has not been tested and will be improved in the future. | 129 | +> 1. Marked as compatible, please note if there are any limitations. |
| 130 | +> 2. Not indicating whether it is compatible means that the compatibility of this API has not been tested and will be improved in the future. | ||
| 128 | 131 | ||
| 129 | ## torch | 132 | ## torch |
| 130 | 133 | ||
| @@ -11891,6 +11894,27 @@ Just like torch.cuda above, you only need to replace ``torch.cuda.amp.xxx`` with | |||
| 11891 | </tbody> | 11894 | </tbody> |
| 11892 | </table> | 11895 | </table> |
| 11893 | 11896 | ||
| 11897 | +## torch.distributed.tensor.experimental | ||
 | |||
| 11898 | + | ||
| 11899 | +<a name="table-context-parallel-experimental"></a> | ||
| 11900 | +<table><thead align="left"><tr id="row-context-parallel-experimental-head"><th class="cellrowborder" valign="top" width="44.84%" id="mcps-context-parallel-experimental-api"><p id="p-context-parallel-experimental-api"><a name="p-context-parallel-experimental-api"></a><a name="p-context-parallel-experimental-api"></a>PyTorch API</p> | ||
| 11901 | +</th> | ||
| 11902 | +<th class="cellrowborder" valign="top" width="27.58%" id="mcps-context-parallel-experimental-compat"><p id="p-context-parallel-experimental-compat"><a name="p-context-parallel-experimental-compat"></a><a name="p-context-parallel-experimental-compat"></a>Compatibility</p> | ||
| 11903 | +</th> | ||
| 11904 | +<th class="cellrowborder" valign="top" width="27.58%" id="mcps-context-parallel-experimental-limit"><p id="p-context-parallel-experimental-limit"><a name="p-context-parallel-experimental-limit"></a><a name="p-context-parallel-experimental-limit"></a>Limitations</p> | ||
| 11905 | +</th> | ||
| 11906 | +</tr> | ||
| 11907 | +</thead> | ||
| 11908 | +<tbody><tr id="row-context-parallel-experimental-context-parallel"><td class="cellrowborder" valign="top" width="44.84%" headers="mcps-context-parallel-experimental-api "><p id="p-context-parallel-experimental-context-parallel"><a name="p-context-parallel-experimental-context-parallel"></a><a name="p-context-parallel-experimental-context-parallel"></a>torch.distributed.tensor.experimental.context_parallel</p> | ||
| 11909 | +</td> | ||
| 11910 | +<td class="cellrowborder" valign="top" width="27.58%" headers="mcps-context-parallel-experimental-compat "><p id="p-context-parallel-experimental-context-parallel-compat"><a name="p-context-parallel-experimental-context-parallel-compat"></a><a name="p-context-parallel-experimental-context-parallel-compat"></a>Y</p> | ||
| 11911 | +</td> | ||
| 11912 | +<td class="cellrowborder" valign="top" width="27.58%" headers="mcps-context-parallel-experimental-limit "><p id="p-context-parallel-experimental-context-parallel-limit"><a name="p-context-parallel-experimental-context-parallel-limit"></a><a name="p-context-parallel-experimental-context-parallel-limit"></a>Only supports the NPU fused SDPA path. The q/k/v tensors must use BNSD layout. pse, padding_mask, prefix, actual_seq_qlen, actual_seq_kvlen, sink, and arbitrary non-causal attention masks are not supported. Load balancing requires causal attention.</p> | ||
| 11913 | +</td> | ||
| 11914 | +</tr> | ||
| 11915 | +</tbody> | ||
| 11916 | +</table> | ||
| 11917 | + | ||
| 11894 | ## torch.distributed.tensor.parallel | 11918 | ## torch.distributed.tensor.parallel |
| 11895 | 11919 | ||
| 11896 | <a name="table42392735619"></a> | 11920 | <a name="table42392735619"></a> |
| @@ -0,0 +1,316 @@ | |||
| 1 | +# Copyright (c) Meta Platforms, Inc. and affiliates | ||
| 2 | +# Owner(s): ["oncall: distributed"] | ||
| 3 | + | ||
| 4 | +""" | ||
| 5 | +NPU context parallel SDPA regression tests. | ||
| 6 | + | ||
| 7 | +The main scenario intentionally mirrors PyTorch's | ||
| 8 | +RingAttentionTest.test_ring_attention_sdpa so torch_npu validates the same | ||
| 9 | +context-parallel user behavior as native DTensor. NPU-specific assertions also | ||
| 10 | +cover the fused npu_fusion_attention_v3 dispatcher path, ring softmax merge, | ||
| 11 | +BNSD layout handling, communication counts, bf16 tolerance, and gradient | ||
| 12 | +unsharding behavior. | ||
| 13 | +""" | ||
| 14 | + | ||
| 15 | +from collections.abc import Callable | ||
| 16 | + | ||
| 17 | +import torch | ||
| 18 | +import torch.distributed as dist | ||
| 19 | +import torch.nn.functional as F | ||
| 20 | +import torch_npu | ||
| 21 | +import torch_npu.distributed.tensor.experimental._context_parallel._attention | ||
| 22 | +from torch.distributed.tensor import DeviceMesh | ||
| 23 | +from torch.distributed.tensor.debug import CommDebugMode | ||
| 24 | +from torch.distributed.tensor.experimental._attention import ( | ||
| 25 | + _context_parallel_shard, | ||
| 26 | + _ContextParallel, | ||
| 27 | + _cp_options, | ||
| 28 | + _disable_context_parallel_dispatcher, | ||
| 29 | + _enable_context_parallel_dispatcher, | ||
| 30 | + _HeadTailLoadBalancer, | ||
| 31 | + _RotateMethod, | ||
| 32 | + context_parallel, | ||
| 33 | + context_parallel_unshard, | ||
| 34 | + set_rotate_method, | ||
| 35 | +) | ||
| 36 | +from torch.distributed.tensor.parallel import parallelize_module | ||
| 37 | +from torch.nn.attention import sdpa_kernel, SDPBackend | ||
| 38 | + | ||
| 39 | +from torch_npu.testing._internal.common_dtensor import NPUDTensorTestBase | ||
| 40 | +from torch_npu.testing.common_distributed import skipIfUnsupportMultiNPU, with_comms | ||
| 41 | +from torch_npu.testing.common_utils import SupportedDevices | ||
| 42 | +from torch_npu.testing.testcase import run_tests | ||
| 43 | + | ||
| 44 | + | ||
| 45 | +c10d_functional = torch.ops.c10d_functional | ||
| 46 | + | ||
| 47 | +ROTATER_ENUM_TO_STR = { | ||
| 48 | + _RotateMethod.ALL_GATHER: "allgather", | ||
| 49 | + _RotateMethod.ALL_TO_ALL: "alltoall", | ||
| 50 | +} | ||
| 51 | + | ||
| 52 | +ATTENTION_TOLERANCES = { | ||
| 53 | + torch.bfloat16: (1e-2, 8e-3), | ||
| 54 | + torch.float32: (2e-6, 1e-5), | ||
| 55 | +} | ||
| 56 | + | ||
| 57 | +def _ordered_bf16_bits(x: torch.Tensor) -> torch.Tensor: | ||
| 58 | + bits = x.detach().to(torch.bfloat16).cpu().contiguous().view(torch.int16).to(torch.int32) | ||
| 59 | + bits = torch.where(bits < 0, bits + 65536, bits) | ||
| 60 | + sign = (bits & 0x8000) != 0 | ||
| 61 | + return torch.where(sign, 0xFFFF - bits, bits + 0x8000) | ||
| 62 | + | ||
| 63 | + | ||
| 64 | +def _bf16_ulp_diff(actual: torch.Tensor, expected: torch.Tensor) -> torch.Tensor: | ||
| 65 | + return (_ordered_bf16_bits(actual) - _ordered_bf16_bits(expected)).abs() | ||
| 66 | + | ||
| 67 | + | ||
| 68 | +class SDPAWrapper(torch.nn.Module): | ||
| 69 | + def __init__(self, compiled: bool, backend: SDPBackend) -> None: | ||
| 70 | + super().__init__() | ||
| 71 | + self.compiled = compiled | ||
| 72 | + self.backend = backend | ||
| 73 | + if compiled: | ||
| 74 | + self._compiled_sdpa = torch.compile( | ||
| 75 | + F.scaled_dot_product_attention, | ||
| 76 | + fullgraph=True, | ||
| 77 | + backend="aot_eager", | ||
| 78 | + ) | ||
| 79 | + | ||
| 80 | + def forward(self, *args: object, **kwargs: object) -> torch.Tensor: | ||
| 81 | + with sdpa_kernel(self.backend): | ||
| 82 | + if self.compiled: | ||
| 83 | + return self._compiled_sdpa(*args, **kwargs) | ||
| 84 | + return F.scaled_dot_product_attention(*args, **kwargs) | ||
| 85 | + | ||
| 86 | + | ||
| 87 | +class TestContextParallelAttention(NPUDTensorTestBase): | ||
| 88 | + | ||
| 89 | + def world_size(self) -> int: | ||
| 90 | + device_count = torch.npu.device_count() if torch.npu.is_available() else 0 | ||
| 91 | + return min(8, device_count) if device_count >= 2 else 2 | ||
| 92 | + | ||
| 93 | + def _make_load_balancer( | ||
| 94 | + self, | ||
| 95 | + load_balance: bool, | ||
| 96 | + seq_length: int, | ||
| 97 | + ) -> _HeadTailLoadBalancer | None: | ||
| 98 | + if not load_balance: | ||
| 99 | + return None | ||
| 100 | + return _HeadTailLoadBalancer(seq_length, self.world_size, torch.device(self.device_type)) | ||
| 101 | + | ||
| 102 | + def _ring_attention_sdpa( | ||
| 103 | + self, | ||
| 104 | + cp_q: torch.Tensor, | ||
| 105 | + cp_k: torch.Tensor, | ||
| 106 | + cp_v: torch.Tensor, | ||
| 107 | + *, | ||
| 108 | + fn_eval: Callable, | ||
| 109 | + mesh: DeviceMesh, | ||
| 110 | + seq_dim: int, | ||
| 111 | + is_causal: bool, | ||
| 112 | + compiled: bool, | ||
| 113 | + backend: SDPBackend, | ||
| 114 | + rotater: _RotateMethod, | ||
| 115 | + test_forward_only: bool, | ||
| 116 | + load_balance: bool, | ||
| 117 | + use_context: bool, | ||
| 118 | + ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None, int]: | ||
| 119 | + ring_update_call_count = 0 | ||
| 120 | + original_ring_update = getattr(torch_npu, "npu_ring_attention_update", None) | ||
| 121 | + | ||
| 122 | + if callable(original_ring_update): | ||
| 123 | + def counted_ring_update(*args, **kwargs): | ||
| 124 | + nonlocal ring_update_call_count | ||
| 125 | + ring_update_call_count += 1 | ||
| 126 | + return original_ring_update(*args, **kwargs) | ||
| 127 | + | ||
| 128 | + torch_npu.npu_ring_attention_update = counted_ring_update | ||
| 129 | + | ||
| 130 | + cp_context = None | ||
| 131 | + try: | ||
| 132 | + if not use_context: | ||
| 133 | + cp_plan = _ContextParallel( | ||
| 134 | + seq_dim=seq_dim, | ||
| 135 | + attention_type=_ContextParallel.AttentionType.SDPA, | ||
| 136 | + ) | ||
| 137 | + attention = parallelize_module(SDPAWrapper(compiled=compiled, backend=backend), mesh, cp_plan) | ||
| 138 | + load_balancer = self._make_load_balancer(load_balance, cp_q.size(seq_dim)) | ||
| 139 | + cp_q, cp_k, cp_v = _context_parallel_shard( | ||
| 140 | + mesh, | ||
| 141 | + (cp_q, cp_k, cp_v), | ||
| 142 | + (seq_dim,) * 3, | ||
| 143 | + load_balancer=load_balancer, | ||
| 144 | + ) | ||
| 145 | + _enable_context_parallel_dispatcher() | ||
| 146 | + else: | ||
| 147 | + _cp_options.enable_load_balance = load_balance | ||
| 148 | + cp_context = context_parallel( | ||
| 149 | + mesh, | ||
| 150 | + buffers=(cp_q, cp_k, cp_v), | ||
| 151 | + buffer_seq_dims=(seq_dim,) * 3, | ||
| 152 | + ) | ||
| 153 | + cp_context.__enter__() | ||
| 154 | + attention = F.scaled_dot_product_attention | ||
| 155 | + if compiled: | ||
| 156 | + attention = torch.compile(attention, fullgraph=True, backend="aot_eager") | ||
| 157 | + | ||
| 158 | + for target in (cp_q, cp_k, cp_v): | ||
| 159 | + target.requires_grad = True | ||
| 160 | + | ||
| 161 | + with CommDebugMode() as comm_mode: | ||
| 162 | + with sdpa_kernel(backend): | ||
| 163 | + cp_out = fn_eval( | ||
| 164 | + attention, | ||
| 165 | + cp_q, | ||
| 166 | + cp_k, | ||
| 167 | + cp_v, | ||
| 168 | + is_causal=is_causal, | ||
| 169 | + ) | ||
| 170 | + | ||
| 171 | + if not compiled and rotater == _RotateMethod.ALL_TO_ALL: | ||
| 172 | + expected_all2all = ( | ||
| 173 | + self.world_size - 1 | ||
| 174 | + if test_forward_only | ||
| 175 | + else self.world_size * 3 - 2 | ||
| 176 | + ) | ||
| 177 | + self.assertDictEqual( | ||
| 178 | + comm_mode.get_comm_counts(), | ||
| 179 | + {c10d_functional.all_to_all_single: expected_all2all}, | ||
| 180 | + ) | ||
| 181 | + | ||
| 182 | + cp_dq, cp_dk, cp_dv = cp_q.grad, cp_k.grad, cp_v.grad | ||
| 183 | + for target in (cp_q, cp_k, cp_v): | ||
| 184 | + target.requires_grad = False | ||
| 185 | + return cp_out, cp_dq, cp_dk, cp_dv, ring_update_call_count | ||
| 186 | + finally: | ||
| 187 | + if not use_context: | ||
| 188 | + _disable_context_parallel_dispatcher() | ||
| 189 | + elif cp_context is not None: | ||
| 190 | + cp_context.__exit__(None, None, None) | ||
| 191 | + if callable(original_ring_update): | ||
| 192 | + torch_npu.npu_ring_attention_update = original_ring_update | ||
| 193 | + | ||
| 194 | + | ||
| 195 | + | ||
| 196 | + | ||
| 197 | + def test_ring_attention_sdpa(self) -> None: | ||
| 198 | + self.run_subtests( | ||
| 199 | + { | ||
| 200 | + "is_causal": [True, False], | ||
| 201 | + "compiled": [False, True], | ||
| 202 | + "backend": [SDPBackend.OVERRIDEABLE], | ||
| 203 | + "load_balance": [False, True], | ||
| 204 | + "rotater": [_RotateMethod.ALL_TO_ALL, _RotateMethod.ALL_GATHER], | ||
| 205 | + "test_forward_only": [False, True], | ||
| 206 | + "use_context": [False, True], | ||
| 207 | + "dtype": [torch.bfloat16, torch.float32], | ||
| 208 | + }, | ||
| 209 | + self._test_ring_attention_sdpa, | ||
| 210 | + ) | ||
| 211 | + | ||
| 212 | + def _test_ring_attention_sdpa( | ||
| 213 | + self, | ||
| 214 | + is_causal: bool, | ||
| 215 | + compiled: bool, | ||
| 216 | + backend: SDPBackend, | ||
| 217 | + load_balance: bool, | ||
| 218 | + rotater: _RotateMethod, | ||
| 219 | + test_forward_only: bool, | ||
| 220 | + use_context: bool, | ||
| 221 | + dtype: torch.dtype, | ||
| 222 | + ) -> None: | ||
| 223 | + if load_balance and not is_causal: | ||
| 224 | + return | ||
| 225 | + | ||
| 226 | + set_rotate_method(ROTATER_ENUM_TO_STR[rotater]) | ||
| 227 | + self.assertEqual(_cp_options.rotate_method, rotater) | ||
| 228 | + | ||
| 229 | + device_mesh = DeviceMesh(self.device_type, torch.arange(0, self.world_size)) | ||
| 230 | + bs = 8 | ||
| 231 | + seq_length = 1024 | ||
| 232 | + seq_dim = 2 | ||
| 233 | + dim = 32 | ||
| 234 | + nheads = 8 | ||
| 235 | + | ||
| 236 | + torch.manual_seed(10) | ||
| 237 | + q, k, v = [ | ||
| 238 | + torch.rand( | ||
| 239 | + (bs, nheads, seq_length * self.world_size, dim), | ||
| 240 | + device=self.device_type, | ||
| 241 | + dtype=dtype, | ||
| 242 | + requires_grad=True, | ||
| 243 | + ) | ||
| 244 | + for _ in range(3) | ||
| 245 | + ] | ||
| 246 | + | ||
| 247 | + with torch.no_grad(): | ||
| 248 | + dist.broadcast(q, src=0) | ||
| 249 | + dist.broadcast(k, src=0) | ||
| 250 | + dist.broadcast(v, src=0) | ||
| 251 | + | ||
| 252 | + def fn_eval(fn, *args, **kwargs): | ||
| 253 | + if test_forward_only: | ||
| 254 | + with torch.no_grad(): | ||
| 255 | + return fn(*args, **kwargs) | ||
| 256 | + out = fn(*args, **kwargs) | ||
| 257 | + out.sum().backward() | ||
| 258 | + return out | ||
| 259 | + | ||
| 260 | + with sdpa_kernel(backend): | ||
| 261 | + out = fn_eval(F.scaled_dot_product_attention, q, k, v, is_causal=is_causal) | ||
| 262 | + | ||
| 263 | + cp_q, cp_k, cp_v = [target.detach().clone() for target in (q, k, v)] | ||
| 264 | + cp_out, cp_dq, cp_dk, cp_dv, ring_update_call_count = self._ring_attention_sdpa( | ||
| 265 | + cp_q, | ||
| 266 | + cp_k, | ||
| 267 | + cp_v, | ||
| 268 | + fn_eval=fn_eval, | ||
| 269 | + mesh=device_mesh, | ||
| 270 | + seq_dim=seq_dim, | ||
| 271 | + is_causal=is_causal, | ||
| 272 | + compiled=compiled, | ||
| 273 | + backend=backend, | ||
| 274 | + rotater=rotater, | ||
| 275 | + test_forward_only=test_forward_only, | ||
| 276 | + load_balance=load_balance, | ||
| 277 | + use_context=use_context, | ||
| 278 | + ) | ||
| 279 | + | ||
| 280 | + call_count = torch.tensor([ring_update_call_count], device=self.device_type) | ||
| 281 | + dist.all_reduce(call_count) | ||
| 282 | + if callable(getattr(torch_npu, "npu_ring_attention_update", None)): | ||
| 283 | + self.assertGreater(call_count.item(), 0) | ||
| 284 | + | ||
| 285 | + load_balancer = self._make_load_balancer(load_balance, q.size(seq_dim)) | ||
| 286 | + (cp_out,) = context_parallel_unshard( | ||
| 287 | + device_mesh, | ||
| 288 | + [cp_out], | ||
| 289 | + [seq_dim], | ||
| 290 | + load_balancer=load_balancer, | ||
| 291 | + ) | ||
| 292 | + | ||
| 293 | + atol, rtol = ATTENTION_TOLERANCES[dtype] | ||
| 294 | + torch.testing.assert_close(out, cp_out, atol=atol, rtol=rtol) | ||
| 295 | + if dtype == torch.bfloat16: | ||
| 296 | + self.assertLessEqual(int(_bf16_ulp_diff(cp_out, out).max().item()), 1) | ||
| 297 | + | ||
| 298 | + if test_forward_only: | ||
| 299 | + return | ||
| 300 | + | ||
| 301 | + cp_dq, cp_dk, cp_dv = context_parallel_unshard( | ||
| 302 | + device_mesh, | ||
| 303 | + [cp_dq, cp_dk, cp_dv], | ||
| 304 | + [seq_dim] * 3, | ||
| 305 | + load_balancer=load_balancer, | ||
| 306 | + ) | ||
| 307 | + | ||
| 308 | + torch.testing.assert_close(q.grad, cp_dq, atol=atol, rtol=rtol) | ||
| 309 | + torch.testing.assert_close(k.grad, cp_dk, atol=atol, rtol=rtol) | ||
| 310 | + torch.testing.assert_close(v.grad, cp_dv, atol=atol, rtol=rtol) | ||
| 311 | + if dtype == torch.bfloat16: | ||
| 312 | + self.assertLessEqual(int(_bf16_ulp_diff(cp_dv, v.grad).max().item()), 1) | ||
| 313 | + | ||
| 314 | + | ||
| 315 | +if __name__ == "__main__": | ||
| 316 | + run_tests() | ||
| @@ -569,6 +569,9 @@ | |||
| 569 | "torch_npu.distributed.reinit_process_group": { | 569 | "torch_npu.distributed.reinit_process_group": { |
| 570 | "signature": "(group=None, rebuild_link=True)" | 570 | "signature": "(group=None, rebuild_link=True)" |
| 571 | }, | 571 | }, |
| 572 | + "torch_npu.distributed.tensor.experimental.context_parallel": { | ||
| 573 | + "signature": "(mesh: torch.distributed.device_mesh.DeviceMesh, *, buffers: list[torch.Tensor] | None = None, buffer_seq_dims: list[int] | None = None, no_restore_buffers: set[torch.Tensor] | None = None) -> collections.abc.Generator[None, None, None]" | ||
| 574 | + }, | ||
| 572 | "torch_npu.distributed.rpc.options.NPUTensorPipeRpcBackendOptions": { | 575 | "torch_npu.distributed.rpc.options.NPUTensorPipeRpcBackendOptions": { |
| 573 | "signature": "(*, num_worker_threads: int = 16, rpc_timeout: float = 60.0, init_method: str = 'env://', device_maps: Optional[Dict[str, Dict[Union[int, str, torch.device], Union[int, str, torch.device]]]] = None, devices: Optional[List[Union[int, str, torch.device]]] = None, _transports: Optional[List] = None, _channels: Optional[List] = None)" | 576 | "signature": "(*, num_worker_threads: int = 16, rpc_timeout: float = 60.0, init_method: str = 'env://', device_maps: Optional[Dict[str, Dict[Union[int, str, torch.device], Union[int, str, torch.device]]]] = None, devices: Optional[List[Union[int, str, torch.device]]] = None, _transports: Optional[List] = None, _channels: Optional[List] = None)" |
| 574 | }, | 577 | }, |
| @@ -6,3 +6,4 @@ import torch_npu.distributed.tensor._moe_ops | |||
| 6 | import torch_npu.distributed.tensor._pointwise_ops | 6 | import torch_npu.distributed.tensor._pointwise_ops |
| 7 | import torch_npu.distributed.tensor._sharded_tensor_patch | 7 | import torch_npu.distributed.tensor._sharded_tensor_patch |
| 8 | import torch_npu.distributed.tensor._view_ops | 8 | import torch_npu.distributed.tensor._view_ops |
| 9 | +import torch_npu.distributed.tensor.experimental | ||
| @@ -0,0 +1,7 @@ | |||
| 1 | +from torch_npu.distributed.tensor.experimental._context_parallel import ( | ||
| 2 | + context_parallel, | ||
| 3 | +) | ||
| 4 | + | ||
| 5 | +__all__ = ["context_parallel"] | ||
| 6 | + | ||
| 7 | +context_parallel.__module__ = __name__ | ||
| @@ -0,0 +1,35 @@ | |||
| 1 | +import functools | ||
| 2 | +from importlib import import_module | ||
| 3 | + | ||
| 4 | +# Import NPU _attention module to trigger side-effect injection (replaces native dispatcher functions) | ||
| 5 | +import_module("torch_npu.distributed.tensor.experimental._context_parallel._attention") | ||
| 6 | + | ||
| 7 | +# Re-export public API from native module (_enable_cp_dtensor_dispatcher is now the NPU-injected version) | ||
| 8 | +from torch.distributed.tensor.experimental._context_parallel._attention import ( | ||
| 9 | + _ContextParallel, | ||
| 10 | + _context_parallel_shard, | ||
| 11 | + context_parallel as _torch_context_parallel, | ||
| 12 | + context_parallel_unshard as _torch_context_parallel_unshard, | ||
| 13 | + set_rotate_method as _torch_set_rotate_method, | ||
| 14 | +) | ||
| 15 | +from torch.distributed.tensor.experimental._context_parallel._load_balancer import ( | ||
| 16 | + _HeadTailLoadBalancer, | ||
| 17 | + _LoadBalancer, | ||
| 18 | +) | ||
| 19 | + | ||
| 20 | +_WRAPPER_ASSIGNMENTS = ("__name__", "__qualname__", "__doc__", "__annotations__") | ||
| 21 | + | ||
| 22 | + | ||
| 23 | + | ||
| 24 | +def context_parallel(*args, **kwargs): | ||
| 25 | + return _torch_context_parallel(*args, **kwargs) | ||
| 26 | + | ||
| 27 | + | ||
| 28 | + | ||
| 29 | +def context_parallel_unshard(*args, **kwargs): | ||
| 30 | + return _torch_context_parallel_unshard(*args, **kwargs) | ||
| 31 | + | ||
| 32 | + | ||
| 33 | + | ||
| 34 | +def set_rotate_method(*args, **kwargs): | ||
| 35 | + return _torch_set_rotate_method(*args, **kwargs) | ||
| @@ -0,0 +1,12 @@ | |||
| 1 | +import torch.nn.functional as F | ||
| 2 | +import torch.distributed.tensor.experimental._context_parallel._attention as _native | ||
| 3 | + | ||
| 4 | +from ._npu_attention import ( | ||
| 5 | + npu_disable_cp_dtensor_dispatcher, | ||
| 6 | + npu_enable_cp_dtensor_dispatcher, | ||
| 7 | +) | ||
| 8 | + | ||
| 9 | + | ||
| 10 | +# Replace dispatcher registration functions in the native module | ||
| 11 | +_native._enable_cp_dtensor_dispatcher = npu_enable_cp_dtensor_dispatcher | ||
| 12 | +_native._disable_cp_dtensor_dispatcher = npu_disable_cp_dtensor_dispatcher | ||
| @@ -0,0 +1,736 @@ | |||
| 1 | +import os | ||
| 2 | +import logging | ||
| 3 | + | ||
| 4 | +import torch | ||
| 5 | +import torch_npu | ||
| 6 | +from torch.distributed.tensor import DTensor, Shard | ||
| 7 | + | ||
| 8 | +from torch.distributed.tensor.experimental._context_parallel._attention import ( | ||
| 9 | + _cp_options, | ||
| 10 | + _templated_ring_attention, | ||
| 11 | + _templated_ring_attention_backward, | ||
| 12 | +) | ||
| 13 | + | ||
| 14 | +logger = logging.getLogger("torch.distributed._context_parallel") | ||
| 15 | + | ||
| 16 | +# ============================================================================ | ||
| 17 | +# npu_fusion_attention_v3 param index | ||
| 18 | +# | ||
| 19 | +# Forward: | ||
| 20 | +# 0:query 1:key 2:value 3:head_num 4:input_layout 5:pse 6:padding_mask | ||
| 21 | +# 7:atten_mask 8:scale 9:keep_prob 10:pre_tockens 11:next_tockens | ||
| 22 | +# 12:inner_precise 13:prefix 14:actual_seq_qlen 15:actual_seq_kvlen | ||
| 23 | +# 16:sparse_mode 17:gen_mask_parallel 18:sync 19:softmax_layout 20:sink | ||
| 24 | +# | ||
| 25 | +# Backward: | ||
| 26 | +# 0:query 1:key 2:value 3:dy 4:head_num 5:input_layout 6:pse 7:padding_mask | ||
| 27 | +# 8:atten_mask 9:softmax_max 10:softmax_sum 11:softmax_in 12:attention_in | ||
| 28 | +# 13:scale_value 14:keep_prob 15:pre_tockens 16:next_tockens 17:inner_precise | ||
| 29 | +# 18:seed 19:offset 20:prefix 21:actual_seq_qlen 22:actual_seq_kvlen | ||
| 30 | +# 23:sparse_mode 24:gen_mask_parallel 25:sync 26:softmax_layout 27:sink | ||
| 31 | +# ============================================================================ | ||
| 32 | + | ||
| 33 | +_FWD_IX = dict( | ||
| 34 | + head_num=3, input_layout=4, pse=5, padding_mask=6, atten_mask=7, | ||
| 35 | + scale=8, keep_prob=9, pre_tockens=10, next_tockens=11, inner_precise=12, | ||
| 36 | + prefix=13, actual_seq_qlen=14, actual_seq_kvlen=15, sparse_mode=16, | ||
| 37 | + gen_mask_parallel=17, sync=18, softmax_layout=19, sink=20, | ||
| 38 | +) | ||
| 39 | + | ||
| 40 | +_BWD_IX = dict( | ||
| 41 | + head_num=4, input_layout=5, pse=6, padding_mask=7, atten_mask=8, | ||
| 42 | + softmax_max=9, softmax_sum=10, softmax_in=11, attention_in=12, | ||
| 43 | + scale_value=13, keep_prob=14, pre_tockens=15, next_tockens=16, | ||
| 44 | + inner_precise=17, seed=18, offset=19, prefix=20, actual_seq_qlen=21, | ||
| 45 | + actual_seq_kvlen=22, sparse_mode=23, gen_mask_parallel=24, sync=25, | ||
| 46 | + softmax_layout=26, sink=27, | ||
| 47 | +) | ||
| 48 | + | ||
| 49 | +def _get(args, ix: dict, name: str, default=None): | ||
| 50 | + """Get value from args by name via index table; returns default if out of bounds.""" | ||
| 51 | + i = ix[name] | ||
| 52 | + return args[i] if len(args) > i else default | ||
| 53 | + | ||
| 54 | +def _validate_bnsd_layout( | ||
| 55 | + query: torch.Tensor, | ||
| 56 | + key: torch.Tensor, | ||
| 57 | + value: torch.Tensor, | ||
| 58 | + input_layout, | ||
| 59 | + *, | ||
| 60 | + op_name: str, | ||
| 61 | +) -> None: | ||
| 62 | + """Fail fast for layouts unsupported by the current CP ring path.""" | ||
| 63 | + layout = input_layout.upper() if isinstance(input_layout, str) else input_layout | ||
| 64 | + if layout != "BNSD": | ||
| 65 | + raise NotImplementedError( | ||
| 66 | + f"{op_name} currently supports BNSD q/k/v only in NPU context parallel, " | ||
| 67 | + f"got input_layout={input_layout!r} with " | ||
| 68 | + f"q={tuple(query.shape)} k={tuple(key.shape)} v={tuple(value.shape)}" | ||
| 69 | + ) | ||
| 70 | + if query.dim() != 4 or key.dim() != 4 or value.dim() != 4: | ||
| 71 | + raise NotImplementedError( | ||
| 72 | + f"{op_name} currently expects 4D BNSD q/k/v in NPU context parallel, " | ||
| 73 | + f"got q.dim={query.dim()} k.dim={key.dim()} v.dim={value.dim()} with " | ||
| 74 | + f"q={tuple(query.shape)} k={tuple(key.shape)} v={tuple(value.shape)}" | ||
| 75 | + ) | ||
| 76 | + | ||
| 77 | +# Passthrough param names: not controlled by ring attention, forwarded to every step | ||
| 78 | +_PASSTHROUGH_NAMES = [ | ||
| 79 | + "head_num", "input_layout", "pre_tockens", "next_tockens", | ||
| 80 | + "inner_precise", "gen_mask_parallel", "sync", "softmax_layout", | ||
| 81 | +] | ||
| 82 | + | ||
| 83 | +def _extract_passthrough(args, ix: dict) -> dict: | ||
| 84 | + """Extract passthrough params from args by index table; skips None values.""" | ||
| 85 | + pt = {} | ||
| 86 | + for name in _PASSTHROUGH_NAMES: | ||
| 87 | + v = _get(args, ix, name) | ||
| 88 | + if v is not None: | ||
| 89 | + pt[name] = v | ||
| 90 | + return pt | ||
| 91 | + | ||
| 92 | +_UNSUPPORTED_CP_PASSTHROUGH_NAMES = ( | ||
| 93 | + "pse", | ||
| 94 | + "padding_mask", | ||
| 95 | + "prefix", | ||
| 96 | + "actual_seq_qlen", | ||
| 97 | + "actual_seq_kvlen", | ||
| 98 | + "sink", | ||
| 99 | +) | ||
| 100 | + | ||
| 101 | +def _is_present(value) -> bool: | ||
| 102 | + if value is None: | ||
| 103 | + return False | ||
| 104 | + if isinstance(value, (list, tuple)) and len(value) == 0: | ||
| 105 | + return False | ||
| 106 | + return True | ||
| 107 | + | ||
| 108 | +def _validate_cp_passthrough_args(args, ix: dict, *, op_name: str) -> None: | ||
| 109 | + """Reject per-sequence/per-logit inputs that are not ring-step transformed yet.""" | ||
| 110 | + unsupported = [ | ||
| 111 | + name | ||
| 112 | + for name in _UNSUPPORTED_CP_PASSTHROUGH_NAMES | ||
| 113 | + if _is_present(_get(args, ix, name)) | ||
| 114 | + ] | ||
| 115 | + if unsupported: | ||
| 116 | + raise NotImplementedError( | ||
| 117 | + f"{op_name} in NPU context parallel does not support " | ||
| 118 | + f"{', '.join(unsupported)} yet. These inputs are tied to global " | ||
| 119 | + "sequence positions or attention logits, so they must be sliced and/or " | ||
| 120 | + "rotated together with q/k/v for each ring step." | ||
| 121 | + ) | ||
| 122 | + | ||
| 123 | + softmax_layout = _get(args, ix, "softmax_layout", "") | ||
| 124 | + if softmax_layout not in (None, ""): | ||
| 125 | + raise NotImplementedError( | ||
| 126 | + f"{op_name} in NPU context parallel currently supports the default " | ||
| 127 | + f"BNSD softmax layout only, got softmax_layout={softmax_layout!r}." | ||
| 128 | + ) | ||
| 129 | + | ||
| 130 | +def _validate_cp_sparse_args(args, ix: dict, *, op_name: str) -> None: | ||
| 131 | + """Keep mask semantics limited to the ring path we actually transform.""" | ||
| 132 | + sparse_mode = _get(args, ix, "sparse_mode", 0) | ||
| 133 | + atten_mask = _get(args, ix, "atten_mask") | ||
| 134 | + if sparse_mode not in (0, 1, 2, 3): | ||
| 135 | + raise NotImplementedError( | ||
| 136 | + f"{op_name} in NPU context parallel currently supports only full " | ||
| 137 | + f"attention and causal sparse modes 1/2/3, got sparse_mode={sparse_mode!r}." | ||
| 138 | + ) | ||
| 139 | + if _is_present(atten_mask) and sparse_mode not in (1, 2, 3): | ||
| 140 | + raise NotImplementedError( | ||
| 141 | + f"{op_name} in NPU context parallel does not support arbitrary " | ||
| 142 | + f"atten_mask yet. Pass causal sparse_mode 1/2/3 for causal attention; " | ||
| 143 | + f"got sparse_mode={sparse_mode!r}." | ||
| 144 | + ) | ||
| 145 | + | ||
| 146 | +# ============================================================================ | ||
| 147 | +# Global stack: forward pushes step_caches, backward pops them. | ||
| 148 | +# | ||
| 149 | +# C++ autograd engine runs backward on different threads (verified via tid mismatch), | ||
| 150 | +# so threading.local() cannot be used. A module-level list works because: | ||
| 151 | +# - Python GIL guarantees thread safety | ||
| 152 | +# - LIFO order matches autograd reverse order (last forward → first backward) | ||
| 153 | +# - Each rank is an independent process | ||
| 154 | +# ============================================================================ | ||
| 155 | + | ||
| 156 | +_step_cache_stack: list = [] | ||
| 157 | + | ||
| 158 | +# ============================================================================ | ||
| 159 | +# Format Conversion: softmax_max/sum → logsumexp | ||
| 160 | +# ============================================================================ | ||
| 161 | + | ||
| 162 | +def _convert_softmax_to_logsumexp( | ||
| 163 | + softmax_max: torch.Tensor, | ||
| 164 | + softmax_sum: torch.Tensor, | ||
| 165 | +) -> torch.Tensor: | ||
| 166 | + """npu_fusion_attention softmax_max/sum [B,N,S,8] → logsumexp [B,N,S]。 | ||
| 167 | + | ||
| 168 | + slot 0: lse = max + log(sum(exp(x - max)))。 | ||
| 169 | + """ | ||
| 170 | + sm_max = softmax_max[:, :, :, 0].float() | ||
| 171 | + sm_sum = softmax_sum[:, :, :, 0].float() | ||
| 172 | + return sm_max + torch.log(sm_sum + 1e-10) | ||
| 173 | + | ||
| 174 | +def _get_ring_attention_update(): | ||
| 175 | + ring_update = getattr(torch_npu, "npu_ring_attention_update", None) | ||
| 176 | + return ring_update if callable(ring_update) else None | ||
| 177 | + | ||
| 178 | +def _get_softmax_merge_impl() -> str: | ||
| 179 | + if _get_ring_attention_update() is not None: | ||
| 180 | + return "op" | ||
| 181 | + return "python" | ||
| 182 | + | ||
| 183 | +def _bnsd_to_sbh(attn_out: torch.Tensor) -> torch.Tensor: | ||
| 184 | + """Map BNSD attention output to the SBH layout expected by ring_update.""" | ||
| 185 | + B, N, S, D = attn_out.shape | ||
| 186 | + return attn_out.permute(2, 0, 1, 3).contiguous().view(S, B, N * D) | ||
| 187 | + | ||
| 188 | +def _sbh_to_bnsd(attn_out: torch.Tensor, *, head_num: int) -> torch.Tensor: | ||
| 189 | + """Map SBH attention output back to the CP main-path BNSD layout.""" | ||
| 190 | + S, B, H = attn_out.shape | ||
| 191 | + if H % head_num != 0: | ||
| 192 | + raise RuntimeError( | ||
| 193 | + f"Cannot convert SBH attention back to BNSD: hidden={H} is not divisible " | ||
| 194 | + f"by head_num={head_num}" | ||
| 195 | + ) | ||
| 196 | + D = H // head_num | ||
| 197 | + return attn_out.view(S, B, head_num, D).permute(1, 2, 0, 3).contiguous() | ||
| 198 | + | ||
| 199 | +def _merge_softmax_stats_python( | ||
| 200 | + prev_softmax_max: torch.Tensor, | ||
| 201 | + prev_softmax_sum: torch.Tensor, | ||
| 202 | + cur_softmax_max: torch.Tensor, | ||
| 203 | + cur_softmax_sum: torch.Tensor, | ||
| 204 | +) -> tuple[torch.Tensor, torch.Tensor]: | ||
| 205 | + """Reference merge in Python, keeping the current BNSD/BNS8 main path.""" | ||
| 206 | + prev_max = prev_softmax_max.float() if _cp_options.convert_to_f32 else prev_softmax_max | ||
| 207 | + prev_sum = prev_softmax_sum.float() if _cp_options.convert_to_f32 else prev_softmax_sum | ||
| 208 | + cur_max = cur_softmax_max.float() if _cp_options.convert_to_f32 else cur_softmax_max | ||
| 209 | + cur_sum = cur_softmax_sum.float() if _cp_options.convert_to_f32 else cur_softmax_sum | ||
| 210 | + | ||
| 211 | + new_max = torch.maximum(prev_max, cur_max) | ||
| 212 | + prev_scale = torch.exp(prev_max.float() - new_max.float()) | ||
| 213 | + cur_scale = torch.exp(cur_max.float() - new_max.float()) | ||
| 214 | + new_sum = prev_sum.float() * prev_scale + cur_sum.float() * cur_scale | ||
| 215 | + return new_max.to(prev_max.dtype), new_sum.to(prev_sum.dtype) | ||
| 216 | + | ||
| 217 | +def _merge_softmax_stats_with_ring_update( | ||
| 218 | + prev_attn_out: torch.Tensor, | ||
| 219 | + prev_softmax_max: torch.Tensor, | ||
| 220 | + prev_softmax_sum: torch.Tensor, | ||
| 221 | + cur_attn_out: torch.Tensor, | ||
| 222 | + cur_softmax_max: torch.Tensor, | ||
| 223 | + cur_softmax_sum: torch.Tensor, | ||
| 224 | +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: | ||
| 225 | + """Use torch_npu.npu_ring_attention_update with a thin BNSD <-> SBH adapter.""" | ||
| 226 | + ring_update = _get_ring_attention_update() | ||
| 227 | + if ring_update is None: | ||
| 228 | + raise RuntimeError( | ||
| 229 | + "The op softmax merge path requires torch_npu.npu_ring_attention_update, " | ||
| 230 | + "but the API is not available in the current torch_npu build." | ||
| 231 | + ) | ||
| 232 | + | ||
| 233 | + merged_attn_sbh, merged_max, merged_sum = ring_update( | ||
| 234 | + _bnsd_to_sbh(prev_attn_out), | ||
| 235 | + prev_softmax_max.float(), | ||
| 236 | + prev_softmax_sum.float(), | ||
| 237 | + _bnsd_to_sbh(cur_attn_out), | ||
| 238 | + cur_softmax_max.float(), | ||
| 239 | + cur_softmax_sum.float(), | ||
| 240 | + input_layout="SBH", | ||
| 241 | + ) | ||
| 242 | + merged_attn = _sbh_to_bnsd(merged_attn_sbh, head_num=prev_attn_out.size(1)) | ||
| 243 | + return merged_attn, merged_max, merged_sum | ||
| 244 | + | ||
| 245 | +def _make_forward_op(step_caches: list, *, pt: dict): | ||
| 246 | + """Create a per-step forward op closure for ring attention. | ||
| 247 | + | ||
| 248 | + step_caches (out param): list mutated in-place — each call appends | ||
| 249 | + (merged_max, merged_sum, seed, offset) for the current step. | ||
| 250 | + pt: passthrough dict of params that do not vary across ring steps | ||
| 251 | + (head_num, input_layout, pre_tockens, next_tockens, etc.). | ||
| 252 | + """ | ||
| 253 | + | ||
| 254 | + merged_max = None | ||
| 255 | + merged_sum = None | ||
| 256 | + merged_attn = None | ||
| 257 | + merge_impl = _get_softmax_merge_impl() | ||
| 258 | + | ||
| 259 | + def _op(query, key, value, *, is_causal=False, dropout_p=0.0, scale=None): | ||
| 260 | + nonlocal merged_max, merged_sum, merged_attn | ||
| 261 | + | ||
| 262 | + B, N, S, D = query.shape | ||
| 263 | + softmax_scale = scale if scale is not None else (1.0 / D**0.5) | ||
| 264 | + | ||
| 265 | + atten_mask = None | ||
| 266 | + sparse_mode = 0 | ||
| 267 | + if is_causal: | ||
| 268 | + atten_mask = torch.triu( | ||
| 269 | + torch.ones(S, S, dtype=torch.uint8, device=query.device), diagonal=1 | ||
| 270 | + ).unsqueeze(0).unsqueeze(0) | ||
| 271 | + sparse_mode = 1 | ||
| 272 | + | ||
| 273 | + (attention_score, softmax_max, softmax_sum, _softmax_out, | ||
| 274 | + seed, offset, _numels) = torch_npu.npu_fusion_attention( | ||
| 275 | + query, key, value, | ||
| 276 | + head_num=pt.get("head_num", N), | ||
| 277 | + input_layout=pt.get("input_layout", "BNSD"), | ||
| 278 | + scale=softmax_scale, | ||
| 279 | + atten_mask=atten_mask, | ||
| 280 | + sparse_mode=sparse_mode, | ||
| 281 | + keep_prob=1.0 - dropout_p, | ||
| 282 | + pse=None, | ||
| 283 | + padding_mask=None, | ||
| 284 | + pre_tockens=pt.get("pre_tockens", 2147483647), | ||
| 285 | + next_tockens=pt.get("next_tockens", 2147483647), | ||
| 286 | + inner_precise=pt.get("inner_precise", 0), | ||
| 287 | + prefix=None, | ||
| 288 | + actual_seq_qlen=None, | ||
| 289 | + actual_seq_kvlen=None, | ||
| 290 | + gen_mask_parallel=pt.get("gen_mask_parallel", True), | ||
| 291 | + sync=pt.get("sync", False), | ||
| 292 | + softmax_layout=pt.get("softmax_layout", ""), | ||
| 293 | + sink=None, | ||
| 294 | + ) | ||
| 295 | + | ||
| 296 | + # Online softmax merge: fold current-step stats into the running accumulator. | ||
| 297 | + # When _cp_options.convert_to_f32 is set, merge runs entirely in fp32. | ||
| 298 | + if merged_max is None: | ||
| 299 | + merged_attn = attention_score.detach() | ||
| 300 | + merged_max = softmax_max.float() if _cp_options.convert_to_f32 else softmax_max | ||
| 301 | + merged_sum = softmax_sum.float() if _cp_options.convert_to_f32 else softmax_sum | ||
| 302 | + elif softmax_max.shape[2] != merged_max.shape[2]: | ||
| 303 | + # IS_CAUSAL partial: _templated_ring_attention uses query.chunk(2, dim=2)[1], | ||
| 304 | + # so the current step only handles the trailing half of Q positions. | ||
| 305 | + S_half = softmax_max.shape[2] | ||
| 306 | + prev_half_max = merged_max[:, :, S_half:, :] | ||
| 307 | + prev_half_sum = merged_sum[:, :, S_half:, :] | ||
| 308 | + if merge_impl == "op": | ||
| 309 | + if merged_attn is None: | ||
| 310 | + raise RuntimeError( | ||
| 311 | + "merged_attn is None in IS_CAUSAL partial merge with " | ||
| 312 | + "merge_impl='op'; this indicates a logic error in forward step merging." | ||
| 313 | + ) | ||
| 314 | + prev_half_attn = merged_attn[:, :, S_half:, :] | ||
| 315 | + merged_half_attn, new_half_max, new_half_sum = _merge_softmax_stats_with_ring_update( | ||
| 316 | + prev_half_attn, | ||
| 317 | + prev_half_max, | ||
| 318 | + prev_half_sum, | ||
| 319 | + attention_score.detach(), | ||
| 320 | + softmax_max, | ||
| 321 | + softmax_sum, | ||
| 322 | + ) | ||
| 323 | + merged_attn = torch.cat([ | ||
| 324 | + merged_attn[:, :, :S_half, :], | ||
| 325 | + merged_half_attn.detach(), | ||
| 326 | + ], dim=2) | ||
| 327 | + else: | ||
| 328 | + new_half_max, new_half_sum = _merge_softmax_stats_python( | ||
| 329 | + prev_half_max, | ||
| 330 | + prev_half_sum, | ||
| 331 | + softmax_max, | ||
| 332 | + softmax_sum, | ||
| 333 | + ) | ||
| 334 | + | ||
| 335 | + merged_max = torch.cat([ | ||
| 336 | + merged_max[:, :, :S_half, :], | ||
| 337 | + new_half_max, | ||
| 338 | + ], dim=2) | ||
| 339 | + merged_sum = torch.cat([ | ||
| 340 | + merged_sum[:, :, :S_half, :], | ||
| 341 | + new_half_sum, | ||
| 342 | + ], dim=2) | ||
| 343 | + else: | ||
| 344 | + if merge_impl == "op": | ||
| 345 | + # merged_attn not used in ring attention | ||
| 346 | + if merged_attn is None: | ||
| 347 | + raise RuntimeError( | ||
| 348 | + "merged_attn is None with merge_impl='op'; " | ||
| 349 | + "this indicates a logic error in forward step merging." | ||
| 350 | + ) | ||
| 351 | + merged_attn, merged_max, merged_sum = _merge_softmax_stats_with_ring_update( | ||
| 352 | + merged_attn, | ||
| 353 | + merged_max, | ||
| 354 | + merged_sum, | ||
| 355 | + attention_score.detach(), | ||
| 356 | + softmax_max, | ||
| 357 | + softmax_sum, | ||
| 358 | + ) | ||
| 359 | + merged_attn = merged_attn.detach() | ||
| 360 | + else: | ||
| 361 | + merged_max, merged_sum = _merge_softmax_stats_python( | ||
| 362 | + merged_max, | ||
| 363 | + merged_sum, | ||
| 364 | + softmax_max, | ||
| 365 | + softmax_sum, | ||
| 366 | + ) | ||
| 367 | + | ||
| 368 | + # Push merged softmax stats (cloned to guard against later steps), seed, offset. | ||
| 369 | + # Per-step attention_score is not saved; backward uses the merged output via _bop(out). | ||
| 370 | + step_caches.append((merged_max.clone(), merged_sum.clone(), seed, offset)) | ||
| 371 | + | ||
| 372 | + # Return per-step logsumexp (fed to the native merger for attention output), not the merged version. | ||
| 373 | + logsumexp = _convert_softmax_to_logsumexp(softmax_max, softmax_sum) | ||
| 374 | + | ||
| 375 | + return ( | ||
| 376 | + attention_score, | ||
| 377 | + logsumexp, | ||
| 378 | + None, # cum_seq_q | ||
| 379 | + None, # cum_seq_k | ||
| 380 | + S, # max_q | ||
| 381 | + S, # max_k | ||
| 382 | + torch.tensor(seed, dtype=torch.int64, device=query.device), | ||
| 383 | + None, # unused | ||
| 384 | + torch.empty(0, device=query.device), # debug_attn_mask | ||
| 385 | + ) | ||
| 386 | + | ||
| 387 | + return _op | ||
| 388 | + | ||
| 389 | + | ||
| 390 | +def _make_backward_op(step_caches: list, *, pt: dict): | ||
| 391 | + """Create a per-step backward op closure for ring attention. | ||
| 392 | + | ||
| 393 | + step_caches: the same list that _make_forward_op populated — consumed in LIFO order. | ||
| 394 | + step_caches[i] = (merged_max, merged_sum, seed, offset) for step i; | ||
| 395 | + step_caches[-1][:2] holds the final globally-normalized stats for the backward kernel. | ||
| 396 | + pt: passthrough dict (same semantics as _make_forward_op). | ||
| 397 | + """ | ||
| 398 | + | ||
| 399 | + idx_box = [0] | ||
| 400 | + | ||
| 401 | + def _bop(grad_out, query, key, value, out, logsumexp, | ||
| 402 | + cum_seq_q, cum_seq_k, max_q, max_k, | ||
| 403 | + dropout_p, is_causal, | ||
| 404 | + philox_seed, philox_offset, | ||
| 405 | + *, scale=None): | ||
| 406 | + N = query.size(1) | ||
| 407 | + D = query.size(-1) | ||
| 408 | + S = query.size(2) | ||
| 409 | + softmax_scale = scale if scale is not None else (1.0 / D**0.5) | ||
| 410 | + | ||
| 411 | + i = idx_box[0] | ||
| 412 | + # step_caches[i] holds the running merged state at step i (grows over ring rounds). | ||
| 413 | + # Backward needs the final globally-normalized stats, so take step_caches[-1][:2]. | ||
| 414 | + _, _, seed, offset = step_caches[i] | ||
| 415 | + merged_max, merged_sum = step_caches[-1][:2] | ||
| 416 | + idx_box[0] += 1 | ||
| 417 | + | ||
| 418 | + # IS_CAUSAL partial: native backward already sliced query/out/grad_out to chunk(2)[1]. | ||
| 419 | + # Slice merged_max/sum to the corresponding trailing half; out is already sliced. | ||
| 420 | + if query.size(2) != merged_max.size(2): | ||
| 421 | + S_half = query.size(2) | ||
| 422 | + merged_max = merged_max[:, :, S_half:, :] | ||
| 423 | + merged_sum = merged_sum[:, :, S_half:, :] | ||
| 424 | + | ||
| 425 | + atten_mask = None | ||
| 426 | + sparse_mode = 0 | ||
| 427 | + if is_causal: | ||
| 428 | + atten_mask = torch.triu( | ||
| 429 | + torch.ones(S, S, dtype=torch.uint8, device=query.device), diagonal=1 | ||
| 430 | + ).unsqueeze(0).unsqueeze(0) | ||
| 431 | + sparse_mode = 1 | ||
| 432 | + | ||
| 433 | + grads = torch_npu.npu_fusion_attention_grad( | ||
| 434 | + query, key, value, | ||
| 435 | + dy=grad_out, | ||
| 436 | + head_num=pt.get("head_num", N), | ||
| 437 | + input_layout=pt.get("input_layout", "BNSD"), | ||
| 438 | + softmax_max=merged_max, | ||
| 439 | + softmax_sum=merged_sum, | ||
| 440 | + attention_in=out, | ||
| 441 | + scale_value=softmax_scale, | ||
| 442 | + keep_prob=1.0 - dropout_p, | ||
| 443 | + atten_mask=atten_mask, | ||
| 444 | + sparse_mode=sparse_mode, | ||
| 445 | + seed=seed, | ||
| 446 | + offset=offset, | ||
| 447 | + pse=None, | ||
| 448 | + padding_mask=None, | ||
| 449 | + pre_tockens=pt.get("pre_tockens", 2147483647), | ||
| 450 | + next_tockens=pt.get("next_tockens", 2147483647), | ||
| 451 | + inner_precise=pt.get("inner_precise", 0), | ||
| 452 | + numels=0, | ||
| 453 | + ) | ||
| 454 | + return grads[0], grads[1], grads[2] | ||
| 455 | + | ||
| 456 | + return _bop | ||
| 457 | + | ||
| 458 | + | ||
| 459 | +# ============================================================================ | ||
| 460 | +# Common: DTensor unwrap helper | ||
| 461 | +# ============================================================================ | ||
| 462 | + | ||
| 463 | +def _unwrap_args(args): | ||
| 464 | + """Unwrap all DTensors in args to _local_tensor; return (local_args, mesh).""" | ||
| 465 | + local_args = [] | ||
| 466 | + mesh = None | ||
| 467 | + for i, a in enumerate(args): | ||
| 468 | + if isinstance(a, DTensor): | ||
| 469 | + local_args.append(a._local_tensor) | ||
| 470 | + if mesh is None: | ||
| 471 | + mesh = a.device_mesh | ||
| 472 | + elif isinstance(a, torch.Tensor): | ||
| 473 | + local_args.append(a) | ||
| 474 | + else: | ||
| 475 | + local_args.append(a) | ||
| 476 | + return local_args, mesh | ||
| 477 | + | ||
| 478 | + | ||
| 479 | +# ============================================================================ | ||
| 480 | +# Forward DTensor Handler — intercepts npu_fusion_attention_v3 | ||
| 481 | +# ============================================================================ | ||
| 482 | +def _npu_fa_v3_handler(op_call, args, kwargs): | ||
| 483 | + """Intercept npu_fusion_attention_v3, run ring attention, return a v3 6-tuple. | ||
| 484 | + | ||
| 485 | + Unwraps DTensor args manually (standard unwrap_to_op_info rejects mixed | ||
| 486 | + DTensor/plain-tensor args). Delegates directly to _templated_ring_attention. | ||
| 487 | + Pushes per-step caches onto the global stack for backward consumption. | ||
| 488 | + """ | ||
| 489 | + local_args, mesh = _unwrap_args(args) | ||
| 490 | + if mesh is None: | ||
| 491 | + raise RuntimeError("No DTensor found in _npu_fa_v3_handler args") | ||
| 492 | + | ||
| 493 | + query, key, value = local_args[0], local_args[1], local_args[2] | ||
| 494 | + | ||
| 495 | + # Ring-attention-controlled params: extract causal semantics and scale from args | ||
| 496 | + sparse_mode = _get(local_args, _FWD_IX, "sparse_mode", 0) | ||
| 497 | + scale = _get(local_args, _FWD_IX, "scale") | ||
| 498 | + keep_prob = _get(local_args, _FWD_IX, "keep_prob", 1.0) | ||
| 499 | + input_layout = _get(local_args, _FWD_IX, "input_layout", "BNSD") | ||
| 500 | + | ||
| 501 | + logger.debug( | ||
| 502 | + "CP forward handler: q=%s k=%s v=%s mesh=%s sparse_mode=%s scale=%s keep_prob=%s", | ||
| 503 | + tuple(query.shape), tuple(key.shape), tuple(value.shape), | ||
| 504 | + mesh, sparse_mode, scale, keep_prob, | ||
| 505 | + ) | ||
| 506 | + | ||
| 507 | + _validate_bnsd_layout( | ||
| 508 | + query, | ||
| 509 | + key, | ||
| 510 | + value, | ||
| 511 | + input_layout, | ||
| 512 | + op_name="npu_fusion_attention_v3", | ||
| 513 | + ) | ||
| 514 | + _validate_cp_passthrough_args( | ||
| 515 | + local_args, | ||
| 516 | + _FWD_IX, | ||
| 517 | + op_name="npu_fusion_attention_v3", | ||
| 518 | + ) | ||
| 519 | + _validate_cp_sparse_args( | ||
| 520 | + local_args, | ||
| 521 | + _FWD_IX, | ||
| 522 | + op_name="npu_fusion_attention_v3", | ||
| 523 | + ) | ||
| 524 | + | ||
| 525 | + is_causal = sparse_mode in (1, 2, 3) | ||
| 526 | + dropout_p = (1.0 - keep_prob) if isinstance(keep_prob, (int, float)) else 0.0 | ||
| 527 | + softmax_scale = scale if scale is not None else (1.0 / query.shape[-1] ** 0.5) | ||
| 528 | + pt = _extract_passthrough(local_args, _FWD_IX) | ||
| 529 | + | ||
| 530 | + step_caches: list = [] | ||
| 531 | + op = _make_forward_op(step_caches, pt=pt) | ||
| 532 | + group = mesh.get_group() | ||
| 533 | + result = _templated_ring_attention( | ||
| 534 | + group, | ||
| 535 | + seq_dim=2, | ||
| 536 | + op=op, | ||
| 537 | + query=query, | ||
| 538 | + key=key, | ||
| 539 | + value=value, | ||
| 540 | + is_causal=is_causal, | ||
| 541 | + dropout_p=dropout_p, | ||
| 542 | + scale=softmax_scale, | ||
| 543 | + ) | ||
| 544 | + attn_output = result[0] | ||
| 545 | + merged_lse = result[1] | ||
| 546 | + | ||
| 547 | + # Push (step_caches, is_causal, merged_out, merged_lse) for backward handler. | ||
| 548 | + # Backward needs merged_out to compute D = sum(dout * O_merged); per-step raw output is not enough. | ||
| 549 | + _step_cache_stack.append( | ||
| 550 | + (step_caches, is_causal, attn_output.detach(), merged_lse.detach()) | ||
| 551 | + ) | ||
| 552 | + | ||
| 553 | + # Build v3 6-tuple. softmax_max/sum come from the final step; op_plugin typically only uses attn_output. | ||
| 554 | + B, N, S, D = attn_output.shape | ||
| 555 | + dev = attn_output.device | ||
| 556 | + if step_caches: | ||
| 557 | + sm_max, sm_sum, _, _ = step_caches[-1] | ||
| 558 | + softmax_max = sm_max | ||
| 559 | + softmax_sum = sm_sum | ||
| 560 | + else: | ||
| 561 | + softmax_max = torch.zeros(B, N, S, 8, dtype=torch.float32, device=dev) | ||
| 562 | + softmax_sum = torch.zeros(B, N, S, 8, dtype=torch.float32, device=dev) | ||
| 563 | + | ||
| 564 | + return ( | ||
| 565 | + attn_output, | ||
| 566 | + softmax_max, | ||
| 567 | + softmax_sum, | ||
| 568 | + torch.zeros(0, device=dev), # softmax_out | ||
| 569 | + torch.tensor([0], dtype=torch.int64, device="cpu"), # seed | ||
| 570 | + torch.tensor([0], dtype=torch.int64, device="cpu"), # offset | ||
| 571 | + ) | ||
| 572 | + | ||
| 573 | + | ||
| 574 | +# ============================================================================ | ||
| 575 | +# Backward DTensor Handler — intercepts npu_fusion_attention_grad_v3 | ||
| 576 | +# ============================================================================ | ||
| 577 | + | ||
| 578 | +def _npu_fa_grad_v3_handler(op_call, args, kwargs): | ||
| 579 | + """Intercept npu_fusion_attention_grad_v3, run ring attention backward, return DTensor grads. | ||
| 580 | + | ||
| 581 | + Pops per-step caches from the global stack (pushed by the forward handler). | ||
| 582 | + Delegates to _templated_ring_attention_backward, then wraps grad_q/k/v as | ||
| 583 | + DTensors with Shard(2). The softmax stats saved by AutogradNPU's backward | ||
| 584 | + node are placeholders; the real per-step caches come from _step_cache_stack. | ||
| 585 | + """ | ||
| 586 | + local_args, mesh = _unwrap_args(args) | ||
| 587 | + if mesh is None: | ||
| 588 | + raise RuntimeError("No DTensor found in npu_fusion_attention_grad_v3 args") | ||
| 589 | + | ||
| 590 | + query, key, value, dy = local_args[0], local_args[1], local_args[2], local_args[3] | ||
| 591 | + input_layout = _get(local_args, _BWD_IX, "input_layout", "BNSD") | ||
| 592 | + _validate_bnsd_layout( | ||
| 593 | + query, | ||
| 594 | + key, | ||
| 595 | + value, | ||
| 596 | + input_layout, | ||
| 597 | + op_name="npu_fusion_attention_grad_v3", | ||
| 598 | + ) | ||
| 599 | + _validate_cp_passthrough_args( | ||
| 600 | + local_args, | ||
| 601 | + _BWD_IX, | ||
| 602 | + op_name="npu_fusion_attention_grad_v3", | ||
| 603 | + ) | ||
| 604 | + _validate_cp_sparse_args( | ||
| 605 | + local_args, | ||
| 606 | + _BWD_IX, | ||
| 607 | + op_name="npu_fusion_attention_grad_v3", | ||
| 608 | + ) | ||
| 609 | + scale_value = _get(local_args, _BWD_IX, "scale_value") | ||
| 610 | + keep_prob = _get(local_args, _BWD_IX, "keep_prob", 1.0) | ||
| 611 | + | ||
| 612 | + dropout_p = (1.0 - keep_prob) if isinstance(keep_prob, (int, float)) else 0.0 | ||
| 613 | + if scale_value is not None: | ||
| 614 | + softmax_scale = float(scale_value) if isinstance(scale_value, (int, float)) else scale_value | ||
| 615 | + else: | ||
| 616 | + softmax_scale = 1.0 / query.shape[-1] ** 0.5 | ||
| 617 | + | ||
| 618 | + stack = _step_cache_stack | ||
| 619 | + if not stack: | ||
| 620 | + raise RuntimeError( | ||
| 621 | + "step_cache_stack is empty in backward handler; " | ||
| 622 | + "forward caches were never pushed or already consumed." | ||
| 623 | + ) | ||
| 624 | + step_caches, is_causal, merged_out, merged_lse = stack.pop() | ||
| 625 | + | ||
| 626 | + pt = _extract_passthrough(local_args, _BWD_IX) | ||
| 627 | + | ||
| 628 | + bop = _make_backward_op(step_caches, pt=pt) | ||
| 629 | + group = mesh.get_group() | ||
| 630 | + zero = torch.zeros(0, device=query.device) | ||
| 631 | + result = _templated_ring_attention_backward( | ||
| 632 | + group, | ||
| 633 | + seq_dim=2, | ||
| 634 | + op=bop, | ||
| 635 | + grad_out=dy, | ||
| 636 | + grad_out_name="grad_out", | ||
| 637 | + query=query, | ||
| 638 | + key=key, | ||
| 639 | + value=value, | ||
| 640 | + out=merged_out, | ||
| 641 | + logsumexp=merged_lse, | ||
| 642 | + is_causal=is_causal, | ||
| 643 | + cum_seq_q=zero, | ||
| 644 | + cum_seq_k=zero, | ||
| 645 | + max_q=query.size(2), | ||
| 646 | + max_k=key.size(2), | ||
| 647 | + dropout_p=dropout_p, | ||
| 648 | + philox_seed=zero, | ||
| 649 | + philox_offset=zero, | ||
| 650 | + scale=softmax_scale, | ||
| 651 | + ) | ||
| 652 | + grad_q, grad_k, grad_v = result[0], result[1], result[2] | ||
| 653 | + | ||
| 654 | + grad_q_dt = DTensor.from_local(grad_q, mesh, [Shard(2)], run_check=False) | ||
| 655 | + grad_k_dt = DTensor.from_local(grad_k, mesh, [Shard(2)], run_check=False) | ||
| 656 | + grad_v_dt = DTensor.from_local(grad_v, mesh, [Shard(2)], run_check=False) | ||
| 657 | + | ||
| 658 | + dev = query.device | ||
| 659 | + grad_pse = torch.zeros(0, device=dev) | ||
| 660 | + grad_sink = torch.zeros(0, device=dev) | ||
| 661 | + return (grad_q_dt, grad_k_dt, grad_v_dt, grad_pse, grad_sink) | ||
| 662 | + | ||
| 663 | + | ||
| 664 | +# ============================================================================ | ||
| 665 | +# CP Sharding Rule | ||
| 666 | +# ============================================================================ | ||
| 667 | + | ||
| 668 | +def _scaled_dot_product_attention_cp_strategy(op_schema): | ||
| 669 | + """CP strategy: Shard(2) on q/k/v and output. Hardcoded dim=2 assumes BNSD layout.""" | ||
| 670 | + from torch.distributed.tensor._ops.utils import ( | ||
| 671 | + expand_to_full_mesh_op_strategy as _expand, | ||
| 672 | + ) | ||
| 673 | + | ||
| 674 | + mesh = op_schema.get_mesh_from_args() | ||
| 675 | + cp_strategy = [ | ||
| 676 | + Shard(2), # output | ||
| 677 | + Shard(2), # query | ||
| 678 | + Shard(2), # key | ||
| 679 | + Shard(2), # value | ||
| 680 | + ] | ||
| 681 | + return _expand(mesh, op_schema, [cp_strategy], input_index=1) | ||
| 682 | + | ||
| 683 | + | ||
| 684 | +_npu_fa = torch.ops.npu.npu_fusion_attention_v3.default | ||
| 685 | +_npu_fa_grad = torch.ops.npu.npu_fusion_attention_grad_v3.default | ||
| 686 | + | ||
| 687 | +_npu_custom_ops = { | ||
| 688 | + _npu_fa: _npu_fa_v3_handler, | ||
| 689 | + _npu_fa_grad: _npu_fa_grad_v3_handler, | ||
| 690 | +} | ||
| 691 | + | ||
| 692 | +# ============================================================================ | ||
| 693 | +# CP Dispatcher Enable/Disable | ||
| 694 | +# ============================================================================ | ||
| 695 | +def npu_enable_cp_dtensor_dispatcher() -> None: | ||
| 696 | + """Register NPU SDPA forward/backward handlers and CP sharding rules.""" | ||
| 697 | + logger.info(f"registering handler keys={[str(k) for k in _npu_custom_ops.keys()]}") | ||
| 698 | + | ||
| 699 | + existing = DTensor._op_dispatcher._custom_op_handlers.copy() | ||
| 700 | + DTensor._op_dispatcher._custom_op_handlers = {**existing, **_npu_custom_ops} | ||
| 701 | + | ||
| 702 | + from torch.distributed.tensor.experimental._context_parallel._sharding_rules import ( | ||
| 703 | + register_cp_sharding_rules, | ||
| 704 | + ) | ||
| 705 | + register_cp_sharding_rules() | ||
| 706 | + | ||
| 707 | + from torch.distributed.tensor._ops.registration import register_op_strategy | ||
| 708 | + from torch.distributed.tensor._op_schema import RuntimeSchemaInfo | ||
| 709 | + register_op_strategy( | ||
| 710 | + _npu_fa, | ||
| 711 | + schema_info=RuntimeSchemaInfo(1), | ||
| 712 | + )(_scaled_dot_product_attention_cp_strategy) | ||
| 713 | + | ||
| 714 | + | ||
| 715 | +def npu_disable_cp_dtensor_dispatcher() -> None: | ||
| 716 | + """Remove NPU handlers and unregister CP sharding rules.""" | ||
| 717 | + logger.info(f"removing handler keys={[str(k) for k in _npu_custom_ops.keys()]}") | ||
| 718 | + | ||
| 719 | + DTensor._op_dispatcher._custom_op_handlers = { | ||
| 720 | + k: v | ||
| 721 | + for k, v in DTensor._op_dispatcher._custom_op_handlers.items() | ||
| 722 | + if k not in _npu_custom_ops | ||
| 723 | + } | ||
| 724 | + | ||
| 725 | + from torch.distributed.tensor.experimental._context_parallel._sharding_rules import ( | ||
| 726 | + unregister_cp_sharding_rules, | ||
| 727 | + ) | ||
| 728 | + unregister_cp_sharding_rules(clear_the_cache=False) | ||
| 729 | + | ||
| 730 | + if _step_cache_stack: | ||
| 731 | + logger.warning( | ||
| 732 | + "CP dispatcher disabled with step_cache_stack depth=%d; " | ||
| 733 | + "some forward caches were not consumed by backward.", | ||
| 734 | + len(_step_cache_stack), | ||
| 735 | + ) | ||
| 736 | + _step_cache_stack.clear() | ||
当前分支删除改资料，统一在2.7.1 docs中维护 资料问题由如下issue跟踪闭环 https://atomgit.com/Ascend/pytorch/issues/2035