已合并
fix sharding strategy for npu_fusion_attention #27780
jizewei创建于 2025年12月10日
fix sharding strategy for npu_fusion_attention #27780
已合并
共 5 个文件变更+785-156
| @@ -0,0 +1,264 @@ | |||
| 1 | +import numpy as np | ||
| 2 | +import torch | ||
| 3 | +from torch.distributed._tensor import distribute_tensor, Replicate, Shard | ||
| 4 | +from torch.testing._internal.common_utils import ( | ||
| 5 | + instantiate_parametrized_tests, | ||
| 6 | + parametrize | ||
| 7 | +) | ||
| 8 | + | ||
| 9 | +import torch_npu | ||
| 10 | +from torch_npu.testing.testcase import run_tests | ||
| 11 | +from torch_npu.testing._internal.common_dtensor import NPUDTensorTestBase | ||
| 12 | +from torch_npu.testing.common_distributed import with_comms, skipIfUnsupportMultiNPU | ||
| 13 | + | ||
| 14 | + | ||
| 15 | +def get_atten_mask(shape, sparse_mode=0, pre_tokens=65536, next_tokens=65536): | ||
| 16 | + atten_mask = None | ||
| 17 | + if sparse_mode == 0: | ||
| 18 | + atten_mask_u = np.triu(np.ones(shape), k=pre_tokens + 1) | ||
| 19 | + atten_mask_l = np.tril(np.ones(shape), k=-next_tokens - 1) | ||
| 20 | + atten_masks = atten_mask_u + atten_mask_l | ||
| 21 | + atten_mask = torch.tensor(atten_masks).to(torch.float16).bool().npu() | ||
| 22 | + elif sparse_mode in [2, 3, 4]: | ||
| 23 | + atten_masks = torch.from_numpy(np.triu(np.ones([2048, 2048]), k=1)) | ||
| 24 | + atten_mask = torch.tensor(atten_masks).to(torch.float16).bool().npu() | ||
| 25 | + | ||
| 26 | + return atten_mask | ||
| 27 | + | ||
| 28 | + | ||
| 29 | +class TestAttentionOps(NPUDTensorTestBase): | ||
| 30 | + | ||
| 31 | + | ||
| 32 | + | ||
| 33 | + "sparse_mode,pre_tokens,next_tokens", | ||
| 34 | + [ | ||
| 35 | + (0, 128, 128), | ||
| 36 | + (1, 65536, 65536), | ||
| 37 | + (2, 65536, 0), | ||
| 38 | + (3, 65536, 0), | ||
| 39 | + (4, 128, 128) | ||
| 40 | + ] | ||
| 41 | + ) | ||
| 42 | + def test_npu_fusion_attention_forward_bnsd(self, sparse_mode, pre_tokens, next_tokens): | ||
| 43 | + device_mesh = self.build_device_mesh() | ||
| 44 | + | ||
| 45 | + B, N, S, D = 4, 8, 32, 32 | ||
| 46 | + shape = (B, N, S, D) | ||
| 47 | + query = torch.randn(shape, dtype=torch.float32, device="npu") | ||
| 48 | + key = torch.randn(shape, dtype=torch.float32, device="npu") | ||
| 49 | + value = torch.randn(shape, dtype=torch.float32, device="npu") | ||
| 50 | + | ||
| 51 | + scale = 0.08838 | ||
| 52 | + | ||
| 53 | + atten_mask = get_atten_mask(shape, sparse_mode, pre_tokens, next_tokens) | ||
| 54 | + result = torch_npu.npu_fusion_attention( | ||
| 55 | + query, key, value, head_num=N, input_layout="BNSD", scale=scale, sparse_mode=sparse_mode, | ||
| 56 | + atten_mask=atten_mask, pre_tockens=pre_tokens, next_tockens=next_tokens | ||
| 57 | + ) | ||
| 58 | + | ||
| 59 | + def test_placement_comb(query_placements, key_placements, value_placements, atten_mask_placements): | ||
| 60 | + dist_query = distribute_tensor(query, device_mesh, query_placements) | ||
| 61 | + dist_key = distribute_tensor(key, device_mesh, key_placements) | ||
| 62 | + dist_value = distribute_tensor(value, device_mesh, value_placements) | ||
| 63 | + dist_atten_mask = distribute_tensor( | ||
| 64 | + atten_mask, device_mesh, atten_mask_placements | ||
| 65 | + ) if atten_mask is not None else None | ||
| 66 | + dist_result = torch_npu.npu_fusion_attention( | ||
| 67 | + dist_query, dist_key, dist_value, head_num=N, input_layout="BNSD", scale=scale, | ||
| 68 | + sparse_mode=sparse_mode, atten_mask=dist_atten_mask, | ||
| 69 | + pre_tockens=pre_tokens, next_tockens=next_tokens | ||
| 70 | + ) | ||
| 71 | + self.assertEqual(dist_result[0].full_tensor(), result[0]) | ||
| 72 | + self.assertEqual(dist_result[1].full_tensor(), result[1]) | ||
| 73 | + self.assertEqual(dist_result[2].full_tensor(), result[2]) | ||
| 74 | + | ||
| 75 | + placements = [Shard(0), Shard(1), Shard(2), Shard(3), Replicate()] | ||
| 76 | + for placement in placements: | ||
| 77 | + if atten_mask is None or (isinstance(placement, Shard) and atten_mask.ndim <= placement.dim): | ||
| 78 | + test_placement_comb([placement], [placement], [placement], [Replicate()]) | ||
| 79 | + else: | ||
| 80 | + test_placement_comb([placement], [placement], [placement], [placement]) | ||
| 81 | + | ||
| 82 | + | ||
| 83 | + | ||
| 84 | + | ||
| 85 | + "sparse_mode,pre_tokens,next_tokens", | ||
| 86 | + [ | ||
| 87 | + (0, 128, 128), | ||
| 88 | + (1, 65536, 65536), | ||
| 89 | + (2, 65536, 0), | ||
| 90 | + (3, 65536, 0), | ||
| 91 | + (4, 128, 128) | ||
| 92 | + ] | ||
| 93 | + ) | ||
| 94 | + def test_npu_fusion_attention_backward_bnsd(self, sparse_mode, pre_tokens, next_tokens): | ||
| 95 | + device_mesh = self.build_device_mesh() | ||
| 96 | + | ||
| 97 | + B, N, S, D = 4, 8, 32, 32 | ||
| 98 | + shape = (B, N, S, D) | ||
| 99 | + query = torch.randn(shape, dtype=torch.float32, device="npu", requires_grad=True) | ||
| 100 | + key = torch.randn(shape, dtype=torch.float32, device="npu", requires_grad=True) | ||
| 101 | + value = torch.randn(shape, dtype=torch.float32, device="npu", requires_grad=True) | ||
| 102 | + | ||
| 103 | + scale = 0.08838 | ||
| 104 | + | ||
| 105 | + atten_mask = get_atten_mask(shape, sparse_mode, pre_tokens, next_tokens) | ||
| 106 | + result = torch_npu.npu_fusion_attention( | ||
| 107 | + query, key, value, head_num=N, input_layout="BNSD", scale=scale, sparse_mode=sparse_mode, | ||
| 108 | + atten_mask=atten_mask, pre_tockens=pre_tokens, next_tockens=next_tokens | ||
| 109 | + ) | ||
| 110 | + gard_y = torch.ones_like(result[0]) | ||
| 111 | + result[0].backward(gard_y) | ||
| 112 | + | ||
| 113 | + def test_placement_comb(query_placements, key_placements, value_placements, atten_mask_placements): | ||
| 114 | + dist_query = distribute_tensor(query, device_mesh, query_placements) | ||
| 115 | + dist_key = distribute_tensor(key, device_mesh, key_placements) | ||
| 116 | + dist_value = distribute_tensor(value, device_mesh, value_placements) | ||
| 117 | + dist_atten_mask = distribute_tensor( | ||
| 118 | + atten_mask, device_mesh, atten_mask_placements | ||
| 119 | + ) if atten_mask is not None else None | ||
| 120 | + dist_result = torch_npu.npu_fusion_attention( | ||
| 121 | + dist_query, dist_key, dist_value, head_num=N, input_layout="BNSD", scale=scale, | ||
| 122 | + sparse_mode=sparse_mode, atten_mask=dist_atten_mask, | ||
| 123 | + pre_tockens=pre_tokens, next_tockens=next_tokens | ||
| 124 | + ) | ||
| 125 | + dist_grad_y = distribute_tensor(gard_y, device_mesh, dist_result[0].placements) | ||
| 126 | + dist_result[0].backward(dist_grad_y) | ||
| 127 | + self.assertEqual(dist_result[0].full_tensor(), result[0]) | ||
| 128 | + self.assertEqual(dist_result[1].full_tensor(), result[1]) | ||
| 129 | + self.assertEqual(dist_result[2].full_tensor(), result[2]) | ||
| 130 | + self.assertEqual(dist_query.grad.full_tensor(), query.grad) | ||
| 131 | + self.assertEqual(dist_key.grad.full_tensor(), key.grad) | ||
| 132 | + self.assertEqual(dist_value.grad.full_tensor(), value.grad) | ||
| 133 | + | ||
| 134 | + placements = [Shard(0), Shard(1), Shard(2), Shard(3), Replicate()] | ||
| 135 | + for placement in placements: | ||
| 136 | + if atten_mask is None or (isinstance(placement, Shard) and atten_mask.ndim <= placement.dim): | ||
| 137 | + test_placement_comb([placement], [placement], [placement], [Replicate()]) | ||
| 138 | + else: | ||
| 139 | + test_placement_comb([placement], [placement], [placement], [placement]) | ||
| 140 | + | ||
| 141 | + | ||
| 142 | + | ||
| 143 | + def test_npu_fusion_attention_bsnd(self): | ||
| 144 | + device_mesh = self.build_device_mesh() | ||
| 145 | + | ||
| 146 | + B, N, S, D = 4, 8, 32, 32 | ||
| 147 | + shape = (B, S, N, D) | ||
| 148 | + query = torch.randn(shape, dtype=torch.float32, device="npu", requires_grad=True) | ||
| 149 | + key = torch.randn(shape, dtype=torch.float32, device="npu", requires_grad=True) | ||
| 150 | + value = torch.randn(shape, dtype=torch.float32, device="npu", requires_grad=True) | ||
| 151 | + scale = 0.08838 | ||
| 152 | + | ||
| 153 | + result = torch_npu.npu_fusion_attention(query, key, value, head_num=N, input_layout="BSND", scale=scale) | ||
| 154 | + gard_y = torch.ones_like(result[0]) | ||
| 155 | + result[0].backward(gard_y) | ||
| 156 | + | ||
| 157 | + def test_placement_comb(query_placements, key_placements, value_placements): | ||
| 158 | + dist_query = distribute_tensor(query, device_mesh, query_placements) | ||
| 159 | + dist_key = distribute_tensor(key, device_mesh, key_placements) | ||
| 160 | + dist_value = distribute_tensor(value, device_mesh, value_placements) | ||
| 161 | + dist_result = torch_npu.npu_fusion_attention( | ||
| 162 | + dist_query, dist_key, dist_value, head_num=N, input_layout="BSND", scale=scale | ||
| 163 | + ) | ||
| 164 | + dist_grad_y = distribute_tensor(gard_y, device_mesh, dist_result[0].placements) | ||
| 165 | + dist_result[0].backward(dist_grad_y) | ||
| 166 | + self.assertEqual(dist_result[0].full_tensor(), result[0]) | ||
| 167 | + self.assertEqual(dist_result[1].full_tensor(), result[1]) | ||
| 168 | + self.assertEqual(dist_result[2].full_tensor(), result[2]) | ||
| 169 | + self.assertEqual(dist_query.grad.full_tensor(), query.grad) | ||
| 170 | + self.assertEqual(dist_key.grad.full_tensor(), key.grad) | ||
| 171 | + self.assertEqual(dist_value.grad.full_tensor(), value.grad) | ||
| 172 | + | ||
| 173 | + placements = [Shard(0), Shard(1), Shard(2), Shard(3), Replicate()] | ||
| 174 | + for placement in placements: | ||
| 175 | + test_placement_comb([placement], [placement], [placement]) | ||
| 176 | + | ||
| 177 | + | ||
| 178 | + | ||
| 179 | + def test_npu_fusion_attention_bsh(self): | ||
| 180 | + device_mesh = self.build_device_mesh() | ||
| 181 | + | ||
| 182 | + B, N, S, D = 4, 8, 32, 32 | ||
| 183 | + shape = (B, S, N * D) | ||
| 184 | + query = torch.randn(shape, dtype=torch.float32, device="npu", requires_grad=True) | ||
| 185 | + key = torch.randn(shape, dtype=torch.float32, device="npu", requires_grad=True) | ||
| 186 | + value = torch.randn(shape, dtype=torch.float32, device="npu", requires_grad=True) | ||
| 187 | + scale = 0.08838 | ||
| 188 | + | ||
| 189 | + result = torch_npu.npu_fusion_attention(query, key, value, head_num=N, input_layout="BSH", scale=scale) | ||
| 190 | + gard_y = torch.ones_like(result[0]) | ||
| 191 | + result[0].backward(gard_y) | ||
| 192 | + | ||
| 193 | + def test_placement_comb(query_placements, key_placements, value_placements): | ||
| 194 | + dist_query = distribute_tensor(query, device_mesh, query_placements) | ||
| 195 | + dist_key = distribute_tensor(key, device_mesh, key_placements) | ||
| 196 | + dist_value = distribute_tensor(value, device_mesh, value_placements) | ||
| 197 | + dist_result = torch_npu.npu_fusion_attention( | ||
| 198 | + dist_query, dist_key, dist_value, head_num=N, input_layout="BSH", scale=scale | ||
| 199 | + ) | ||
| 200 | + dist_grad_y = distribute_tensor(gard_y, device_mesh, dist_result[0].placements) | ||
| 201 | + dist_result[0].backward(dist_grad_y) | ||
| 202 | + self.assertEqual(dist_result[0].full_tensor(), result[0]) | ||
| 203 | + self.assertEqual(dist_result[1].full_tensor(), result[1]) | ||
| 204 | + self.assertEqual(dist_result[2].full_tensor(), result[2]) | ||
| 205 | + self.assertEqual(dist_query.grad.full_tensor(), query.grad) | ||
| 206 | + self.assertEqual(dist_key.grad.full_tensor(), key.grad) | ||
| 207 | + self.assertEqual(dist_value.grad.full_tensor(), value.grad) | ||
| 208 | + | ||
| 209 | + placements = [Shard(0), Shard(1), Shard(2), Replicate()] | ||
| 210 | + for placement in placements: | ||
| 211 | + test_placement_comb([placement], [placement], [placement]) | ||
| 212 | + | ||
| 213 | + | ||
| 214 | + | ||
| 215 | + def test_npu_fusion_attention_tnd(self): | ||
| 216 | + device_mesh = self.build_device_mesh() | ||
| 217 | + | ||
| 218 | + B, Nq, Nkv, D = 3, 8, 2, 32 | ||
| 219 | + seq_qlen_list = np.array([1, 2, 3]) | ||
| 220 | + actual_seq_qlen = np.cumsum(seq_qlen_list) | ||
| 221 | + Sq = seq_qlen_list.sum() | ||
| 222 | + seq_kvlen_list = np.array([3, 4, 5]) | ||
| 223 | + actual_seq_kvlen = np.cumsum(seq_kvlen_list) | ||
| 224 | + Skv = seq_kvlen_list.sum() | ||
| 225 | + query = torch.randn(Sq, Nq, D, dtype=torch.float32, device="npu", requires_grad=True) | ||
| 226 | + key = torch.randn(Skv, Nkv, D, dtype=torch.float32, device="npu", requires_grad=True) | ||
| 227 | + value = torch.randn(Skv, Nkv, D, dtype=torch.float32, device="npu", requires_grad=True) | ||
| 228 | + scale = 1 / (D ** 0.5) | ||
| 229 | + | ||
| 230 | + result = torch_npu.npu_fusion_attention( | ||
| 231 | + query, key, value, head_num=Nq, input_layout="TND", scale=scale, | ||
| 232 | + actual_seq_qlen=actual_seq_qlen.tolist(), actual_seq_kvlen=actual_seq_kvlen.tolist(), softmax_layout="TND" | ||
| 233 | + ) | ||
| 234 | + grad_y = torch.ones_like(result[0]) | ||
| 235 | + result[0].backward(grad_y) | ||
| 236 | + | ||
| 237 | + def test_placement_comb(query_placements, key_placements, value_placements): | ||
| 238 | + dist_query = distribute_tensor(query, device_mesh, query_placements) | ||
| 239 | + dist_key = distribute_tensor(key, device_mesh, key_placements) | ||
| 240 | + dist_value = distribute_tensor(value, device_mesh, value_placements) | ||
| 241 | + dist_result = torch_npu.npu_fusion_attention( | ||
| 242 | + dist_query, dist_key, dist_value, head_num=Nq, input_layout="TND", scale=scale, | ||
| 243 | + actual_seq_qlen=actual_seq_qlen.tolist(), actual_seq_kvlen=actual_seq_kvlen.tolist(), | ||
| 244 | + softmax_layout="TND" | ||
| 245 | + ) | ||
| 246 | + dist_grad_y = distribute_tensor(grad_y, device_mesh, dist_result[0].placements) | ||
| 247 | + dist_result[0].backward(dist_grad_y) | ||
| 248 | + self.assertEqual(dist_result[0].full_tensor(), result[0]) | ||
| 249 | + self.assertEqual(dist_result[1].full_tensor(), result[1]) | ||
| 250 | + self.assertEqual(dist_result[2].full_tensor(), result[2]) | ||
| 251 | + self.assertEqual(dist_query.grad.full_tensor(), query.grad) | ||
| 252 | + self.assertEqual(dist_key.grad.full_tensor(), key.grad) | ||
| 253 | + self.assertEqual(dist_value.grad.full_tensor(), value.grad) | ||
| 254 | + | ||
| 255 | + placements = [Shard(0), Shard(1), Shard(2), Replicate()] | ||
| 256 | + for placement in placements: | ||
| 257 | + test_placement_comb([placement], [placement], [placement]) | ||
| 258 | + | ||
| 259 | + | ||
| 260 | +instantiate_parametrized_tests(TestAttentionOps) | ||
| 261 | + | ||
| 262 | + | ||
| 263 | +if __name__ == "__main__": | ||
| 264 | + run_tests() | ||
| @@ -1,3 +1,4 @@ | |||
| 1 | +import torch_npu.distributed.tensor._dtensor_patch # patch before register strategy | ||
| 1 | import torch_npu.distributed.tensor._matrix_ops | 2 | import torch_npu.distributed.tensor._matrix_ops |
| 2 | import torch_npu.distributed.tensor._attention | 3 | import torch_npu.distributed.tensor._attention |
| 3 | import torch_npu.distributed.tensor._math_ops | 4 | import torch_npu.distributed.tensor._math_ops |
| @@ -1,15 +1,344 @@ | |||
| 1 | import itertools | 1 | import itertools |
| 2 | -from typing import Any, Dict, Tuple | 2 | +from typing import cast, Any, Dict, Tuple |
| 3 | 3 | ||
| 4 | import torch | 4 | import torch |
| 5 | from torch.distributed.device_mesh import DeviceMesh | 5 | from torch.distributed.device_mesh import DeviceMesh |
| 6 | -from torch.distributed.tensor import distribute_module, DTensor, Replicate | 6 | +from torch.distributed._tensor.experimental import register_sharding |
| 7 | +from torch.distributed._tensor.placement_types import DTensorSpec | ||
| 8 | +from torch.distributed.tensor import DTensor, Replicate, Shard | ||
| 9 | +from torch.distributed.tensor._op_schema import ( | ||
| 10 | + OpInfo, | ||
| 11 | + OpSchema, | ||
| 12 | + OutputSharding | ||
| 13 | +) | ||
| 14 | +from torch.distributed.tensor._redistribute import redistribute_local_tensor | ||
| 7 | 15 | ||
| 8 | import torch_npu | 16 | import torch_npu |
| 9 | 17 | ||
| 10 | npu = torch.ops.npu | 18 | npu = torch.ops.npu |
| 11 | 19 | ||
| 12 | 20 | ||
| 21 | + | ||
| 22 | +# pylint:disable=huawei-too-many-arguments | ||
| 23 | +def npu_fusion_attention_strategy(query, key, value, head_num, input_layout, pse=None, padding_mask=None, | ||
| 24 | + atten_mask=None, scale=1.0, keep_prob=1.0, pre_tockens=2147483647, | ||
| 25 | + next_tockens=2147483647, inner_precise=0, prefix=None, actual_seq_qlen=None, | ||
| 26 | + actual_seq_kvlen=None, sparse_mode=0, gen_mask_parallel=True, sync=False, | ||
| 27 | + softmax_layout=""): | ||
| 28 | + # func: npu_fusion_attention(Tensor query, Tensor key, Tensor value, int head_num, str input_layout, | ||
| 29 | + # Tensor? pse=None, Tensor? padding_mask=None, Tensor? atten_mask=None, float scale=1., | ||
| 30 | + # float keep_prob=1., int pre_tockens=2147483647, int next_tockens=2147483647, | ||
| 31 | + # int inner_precise=0, int[]? prefix=None, int[]? actual_seq_qlen=None, | ||
| 32 | + # int[]? actual_seq_kvlen=None, int sparse_mode=0, bool gen_mask_parallel=True, | ||
| 33 | + # bool sync=False, str softmax_layout="") | ||
| 34 | + # -> (Tensor, Tensor, Tensor, Tensor, int, int, int) | ||
| 35 | + strategies = [] | ||
| 36 | + | ||
| 37 | + # all replicate strategy | ||
| 38 | + replicate_strategy = ( | ||
| 39 | + [ | ||
| 40 | + Replicate(), # attention_out | ||
| 41 | + Replicate(), # softmax_max | ||
| 42 | + Replicate(), # softmax_sum | ||
| 43 | + Replicate(), # softmax_out(reserve, unused now) | ||
| 44 | + None, None, None # others | ||
| 45 | + ], | ||
| 46 | + [ | ||
| 47 | + Replicate(), # query | ||
| 48 | + Replicate(), # key | ||
| 49 | + Replicate(), # value | ||
| 50 | + None, # head_num | ||
| 51 | + None, # input_layout | ||
| 52 | + None if pse is None else Replicate(), # pse | ||
| 53 | + None if padding_mask is None else Replicate(), # padding_mask | ||
| 54 | + None if atten_mask is None else Replicate(), # atten_mask | ||
| 55 | + None, None, None, None, None, None, None, None, None, None, None, None # others | ||
| 56 | + ] | ||
| 57 | + ) | ||
| 58 | + strategies.append(replicate_strategy) | ||
| 59 | + | ||
| 60 | + # only support sharding for sdpa currently, in which pse and padding_mask are not used | ||
| 61 | + # keep_prob < 1.0 may effect different results under sharding | ||
| 62 | + unused_args_in_sdpa = [pse, padding_mask, prefix, actual_seq_qlen, actual_seq_kvlen] | ||
| 63 | + if not all(arg is None for arg in unused_args_in_sdpa) or keep_prob < 1.0: | ||
| 64 | + return strategies | ||
| 65 | + | ||
| 66 | + # input layout: BSH, SBH, BSND, BNSD, TND | ||
| 67 | + # atten_mask layout: BNSS, B1SS, 11SS, SS | ||
| 68 | + # dp sharding strategy | ||
| 69 | + if 'B' in input_layout: | ||
| 70 | + batch_dim = input_layout.index('B') | ||
| 71 | + atten_mask_sharding = None | ||
| 72 | + if atten_mask is not None: | ||
| 73 | + if atten_mask.ndim == 4 and atten_mask.shape[0] != 1: # BNSS, B1SS | ||
| 74 | + atten_mask_sharding = Shard(0) | ||
| 75 | + else: # 11SS, SS | ||
| 76 | + atten_mask_sharding = Replicate() | ||
| 77 | + dp_sharding_strategy = ( | ||
| 78 | + [ | ||
| 79 | + Shard(batch_dim), # attention_out | ||
| 80 | + Shard(0), # softmax_max layout: BNS8 | ||
| 81 | + Shard(0), # softmax_sum layout: BNS8 | ||
| 82 | + Replicate(), # softmax_out(reserve, unused now) | ||
| 83 | + None, None, None # others | ||
| 84 | + ], | ||
| 85 | + [ | ||
| 86 | + Shard(batch_dim), # query | ||
| 87 | + Shard(batch_dim), # key | ||
| 88 | + Shard(batch_dim), # value | ||
| 89 | + None, # head_num | ||
| 90 | + None, # input_layout | ||
| 91 | + None, # pse | ||
| 92 | + None, # padding_mask | ||
| 93 | + atten_mask_sharding, # atten_mask | ||
| 94 | + None, None, None, None, None, None, None, None, None, None, None, None # others | ||
| 95 | + ] | ||
| 96 | + ) | ||
| 97 | + strategies.append(dp_sharding_strategy) | ||
| 98 | + | ||
| 99 | + # add tp sharding strategy | ||
| 100 | + if 'N' in input_layout: | ||
| 101 | + head_dim = input_layout.index('N') | ||
| 102 | + atten_mask_sharding = None | ||
| 103 | + if atten_mask is not None: | ||
| 104 | + if atten_mask.ndim == 4 and atten_mask.shape[1] != 1: # BNSS | ||
| 105 | + atten_mask_sharding = Shard(1) | ||
| 106 | + else: | ||
| 107 | + atten_mask_sharding = Replicate() # B1SS, 11SS, SS | ||
| 108 | + tp_sharding_strategy = ( | ||
| 109 | + [ | ||
| 110 | + Shard(head_dim), # attention_out | ||
| 111 | + Shard(1), # softmax_max layout: BNS8 | ||
| 112 | + Shard(1), # softmax_sum layout: BNS8 | ||
| 113 | + Replicate(), # softmax_out(reserve, unused now) | ||
| 114 | + None, None, None # others | ||
| 115 | + ], | ||
| 116 | + [ | ||
| 117 | + Shard(head_dim), # query | ||
| 118 | + Shard(head_dim), # key | ||
| 119 | + Shard(head_dim), # value | ||
| 120 | + None, # head_num | ||
| 121 | + None, # input_layout | ||
| 122 | + None, # pse | ||
| 123 | + None, # padding_mask | ||
| 124 | + atten_mask_sharding, # atten_mask | ||
| 125 | + None, None, None, None, None, None, None, None, None, None, None, None # others | ||
| 126 | + ] | ||
| 127 | + ) | ||
| 128 | + strategies.append(tp_sharding_strategy) | ||
| 129 | + | ||
| 130 | + return strategies | ||
| 131 | + | ||
| 132 | + | ||
| 133 | + | ||
| 134 | +def npu_fusion_attention_grad_strategy(query, key, value, dy, head_num, input_layout, pse=None, padding_mask=None, | ||
| 135 | + atten_mask=None, softmax_max=None, softmax_sum=None, softmax_in=None, | ||
| 136 | + attention_in=None, scale_value=1., keep_prob=1., pre_tockens=2147483647, | ||
| 137 | + next_tockens=2147483647, inner_precise=0, seed=0, offset=0, numels=0, | ||
| 138 | + prefix=None, actual_seq_qlen=None, actual_seq_kvlen=None, sparse_mode=0, | ||
| 139 | + gen_mask_parallel=True, sync=False, softmax_layout=""): | ||
| 140 | + # npu_fusion_attention_grad(Tensor query, Tensor key, Tensor value, Tensor dy, int head_num, str input_layout, *, | ||
| 141 | + # Tensor? pse=None, Tensor? padding_mask=None, Tensor? atten_mask=None, | ||
| 142 | + # Tensor? softmax_max=None, Tensor? softmax_sum=None, Tensor? softmax_in=None, | ||
| 143 | + # Tensor? attention_in=None, float scale_value=1., float keep_prob=1., | ||
| 144 | + # int pre_tockens=2147483647, int next_tockens=2147483647, int inner_precise=0, | ||
| 145 | + # int seed=0, int offset=0, int numels=0, int[]? prefix=None, | ||
| 146 | + # int[]? actual_seq_qlen=None, int[]? actual_seq_kvlen=None, int sparse_mode=0, | ||
| 147 | + # bool gen_mask_parallel=True, bool sync=False, str softmax_layout="") | ||
| 148 | + # -> (Tensor, Tensor, Tensor, Tensor) | ||
| 149 | + strategies = [] | ||
| 150 | + | ||
| 151 | + # all replicate strategy | ||
| 152 | + replicate_strategy = ( | ||
| 153 | + [ | ||
| 154 | + Replicate(), # grad_query | ||
| 155 | + Replicate(), # grad_key | ||
| 156 | + Replicate(), # grad_value | ||
| 157 | + Replicate() # grad_pse(reserve, unused now) | ||
| 158 | + ], | ||
| 159 | + [ | ||
| 160 | + Replicate(), # query | ||
| 161 | + Replicate(), # key | ||
| 162 | + Replicate(), # value | ||
| 163 | + Replicate(), # dy | ||
| 164 | + None, # head_num | ||
| 165 | + None, # input_layout | ||
| 166 | + None if pse is None else Replicate(), # pse | ||
| 167 | + None if padding_mask is None else Replicate(), # padding_mask | ||
| 168 | + None if atten_mask is None else Replicate(), # atten_mask | ||
| 169 | + None if softmax_max is None else Replicate(), # softmax_max | ||
| 170 | + None if softmax_sum is None else Replicate(), # softmax_sum | ||
| 171 | + None if softmax_in is None else Replicate(), # softmax_in(reserve, unused now) | ||
| 172 | + None if attention_in is None else Replicate(), # attention_in | ||
| 173 | + None, None, None, None, None, None, None, None, None, None, None, None, None, None, None # others | ||
| 174 | + ] | ||
| 175 | + ) | ||
| 176 | + strategies.append(replicate_strategy) | ||
| 177 | + | ||
| 178 | + # only support sharding for sdpa currently, in which pse and padding_mask are not used | ||
| 179 | + # keep_prob < 1.0 may effect different results under sharding | ||
| 180 | + unused_args_in_sdpa = [pse, padding_mask, prefix, actual_seq_qlen, actual_seq_kvlen] | ||
| 181 | + if not all(arg is None for arg in unused_args_in_sdpa) or keep_prob < 1.0: | ||
| 182 | + return strategies | ||
| 183 | + | ||
| 184 | + # input layout: BSH, SBH, BSND, BNSD, TND | ||
| 185 | + # atten_mask layout: BNSS, B1SS, 11SS, SS | ||
| 186 | + # dp sharding strategy | ||
| 187 | + if 'B' in input_layout: | ||
| 188 | + batch_dim = input_layout.index('B') | ||
| 189 | + atten_mask_sharding = None | ||
| 190 | + if atten_mask is not None: | ||
| 191 | + if atten_mask.ndim == 4 and atten_mask.shape[0] != 1: # BNSS, B1SS | ||
| 192 | + atten_mask_sharding = Shard(0) | ||
| 193 | + else: # 11SS, SS | ||
| 194 | + atten_mask_sharding = Replicate() | ||
| 195 | + dp_sharding_strategy = ( | ||
| 196 | + [ | ||
| 197 | + Shard(batch_dim), # grad_query | ||
| 198 | + Shard(batch_dim), # grad_key | ||
| 199 | + Shard(batch_dim), # grad_value | ||
| 200 | + Replicate() # grad_pse(reserve, unused now) | ||
| 201 | + ], | ||
| 202 | + [ | ||
| 203 | + Shard(batch_dim), # query | ||
| 204 | + Shard(batch_dim), # key | ||
| 205 | + Shard(batch_dim), # value | ||
| 206 | + Shard(batch_dim), # dy | ||
| 207 | + None, # head_num | ||
| 208 | + None, # input_layout | ||
| 209 | + None, # pse | ||
| 210 | + None, # padding_mask | ||
| 211 | + atten_mask_sharding, # atten_mask | ||
| 212 | + Shard(0) if softmax_max is not None else None, # softmax_max layout: BNS8 | ||
| 213 | + Shard(0) if softmax_sum is not None else None, # softmax_sum layout: BNS8 | ||
| 214 | + None if softmax_in is None else Replicate(), # softmax_in(reserve, unused now) | ||
| 215 | + Shard(batch_dim) if attention_in is not None else None, # attention_in | ||
| 216 | + None, None, None, None, None, None, None, None, None, None, None, None, None, None, None # others | ||
| 217 | + ] | ||
| 218 | + ) | ||
| 219 | + strategies.append(dp_sharding_strategy) | ||
| 220 | + | ||
| 221 | + # add tp sharding strategy | ||
| 222 | + if 'N' in input_layout: | ||
| 223 | + head_dim = input_layout.index('N') | ||
| 224 | + atten_mask_sharding = None | ||
| 225 | + if atten_mask is not None: | ||
| 226 | + if atten_mask.ndim == 4 and atten_mask.shape[1] != 1: # BNSS | ||
| 227 | + atten_mask_sharding = Shard(1) | ||
| 228 | + else: | ||
| 229 | + atten_mask_sharding = Replicate() # B1SS, 11SS, SS | ||
| 230 | + tp_sharding_strategy = ( | ||
| 231 | + [ | ||
| 232 | + Shard(head_dim), # grad_query | ||
| 233 | + Shard(head_dim), # grad_key | ||
| 234 | + Shard(head_dim), # grad_value | ||
| 235 | + Replicate() # grad_pse(reserve, unused now) | ||
| 236 | + ], | ||
| 237 | + [ | ||
| 238 | + Shard(head_dim), # query | ||
| 239 | + Shard(head_dim), # key | ||
| 240 | + Shard(head_dim), # value | ||
| 241 | + Shard(head_dim), # dy | ||
| 242 | + None, # head_num | ||
| 243 | + None, # input_layout | ||
| 244 | + None, # pse | ||
| 245 | + None, # padding_mask | ||
| 246 | + atten_mask_sharding, # atten_mask | ||
| 247 | + Shard(1) if softmax_max is not None else None, # softmax_max layout: BNS8 | ||
| 248 | + Shard(1) if softmax_sum is not None else None, # softmax_sum layout: BNS8 | ||
| 249 | + None if softmax_in is None else Replicate(), # softmax_in(reserve, unused now) | ||
| 250 | + Shard(head_dim) if attention_in is not None else None, # attention_in | ||
| 251 | + None, None, None, None, None, None, None, None, None, None, None, None, None, None, None # others | ||
| 252 | + ] | ||
| 253 | + ) | ||
| 254 | + strategies.append(tp_sharding_strategy) | ||
| 255 | + | ||
| 256 | + return strategies | ||
| 257 | + | ||
| 258 | + | ||
| 259 | +def _infer_target_kwargs_spec(op_schema: OpSchema, output_sharding: OutputSharding) -> Dict[str, DTensorSpec]: | ||
| 260 | + # in ShardingPropagator.propagate, only the redistribution of args is considered: | ||
| 261 | + # 1. if args do not need redistribute, output_sharding.redistribute_schema is None | ||
| 262 | + # 2. if args need redistribute, kwargs_schema in output_sharding.redistribute_schema is still from the source input | ||
| 263 | + # schema rather than the selected strategy | ||
| 264 | + # therefore we need infer the correct kwargs spec from output spec here | ||
| 265 | + input_layout = op_schema.args_schema[5] | ||
| 266 | + batch_dim = input_layout.index('B') if 'B' in input_layout else None | ||
| 267 | + dp_shard = Shard(batch_dim) if batch_dim is not None else None | ||
| 268 | + head_dim = input_layout.index('N') if 'N' in input_layout else None | ||
| 269 | + tp_shard = Shard(head_dim) if head_dim is not None else None | ||
| 270 | + output_spec = output_sharding.output_spec[0] | ||
| 271 | + kwargs_spec = {} | ||
| 272 | + for key, spec in op_schema.kwargs_schema.items(): | ||
| 273 | + if not isinstance(spec, DTensorSpec): | ||
| 274 | + kwargs_spec[key] = spec | ||
| 275 | + continue | ||
| 276 | + | ||
| 277 | + target_placement = [] | ||
| 278 | + for placement in output_spec.placements: | ||
| 279 | + if placement == Replicate(): | ||
| 280 | + target_placement.append(Replicate()) | ||
| 281 | + elif placement == dp_shard: | ||
| 282 | + if key == 'atten_mask': | ||
| 283 | + atten_mask = op_schema.kwargs_schema[key] | ||
| 284 | + if atten_mask.ndim == 4 and atten_mask.shape[0] != 1: # BNSS, B1SS | ||
| 285 | + target_placement.append(dp_shard) | ||
| 286 | + else: # 11SS, SS | ||
| 287 | + target_placement.append(Replicate()) | ||
| 288 | + elif key == 'softmax_max' or key == 'softmax_sum': | ||
| 289 | + target_placement.append(Shard(0)) | ||
| 290 | + elif key == 'attention_in': | ||
| 291 | + target_placement.append(dp_shard) | ||
| 292 | + else: # softmax_in | ||
| 293 | + target_placement.append(Replicate()) | ||
| 294 | + elif placement == tp_shard: | ||
| 295 | + if key == 'atten_mask': | ||
| 296 | + atten_mask = op_schema.kwargs_schema[key] | ||
| 297 | + if atten_mask.ndim == 4 and atten_mask.shape[1] != 1: # BNSS | ||
| 298 | + target_placement.append(tp_shard) | ||
| 299 | + else: # B1SS, 11SS, SS | ||
| 300 | + target_placement.append(Replicate()) | ||
| 301 | + elif key == 'softmax_max' or key == 'softmax_sum': | ||
| 302 | + target_placement.append(Shard(1)) | ||
| 303 | + elif key == 'attention_in': | ||
| 304 | + target_placement.append(tp_shard) | ||
| 305 | + else: # softmax_in | ||
| 306 | + target_placement.append(Replicate()) | ||
| 307 | + else: | ||
| 308 | + raise ValueError( | ||
| 309 | + f"Unexpected placement {placement} for npu_fusion_attention_grad in layout {input_layout}." | ||
| 310 | + ) | ||
| 311 | + | ||
| 312 | + kwargs_spec[key] = DTensorSpec( | ||
| 313 | + mesh=spec.mesh, | ||
| 314 | + placements=target_placement, | ||
| 315 | + tensor_meta=spec.tensor_meta | ||
| 316 | + ) | ||
| 317 | + | ||
| 318 | + return kwargs_spec | ||
| 319 | + | ||
| 320 | + | ||
| 321 | +def _redistribute_local_kwargs(op_info: OpInfo): | ||
| 322 | + src_kwargs_spec = op_info.schema.kwargs_schema | ||
| 323 | + target_kwargs_spec = _infer_target_kwargs_spec(op_info.schema, op_info.output_sharding) | ||
| 324 | + new_local_kwargs = {} | ||
| 325 | + for key, target_spec in target_kwargs_spec.items(): | ||
| 326 | + local_tensor = op_info.local_kwargs[key] | ||
| 327 | + if isinstance(target_spec, DTensorSpec): | ||
| 328 | + src_spec = src_kwargs_spec[key] | ||
| 329 | + if src_spec.placements != target_spec.placements: | ||
| 330 | + resharded_local_tensor = redistribute_local_tensor( | ||
| 331 | + local_tensor, src_spec, target_spec | ||
| 332 | + ) | ||
| 333 | + new_local_kwargs[key] = resharded_local_tensor | ||
| 334 | + else: | ||
| 335 | + new_local_kwargs[key] = local_tensor | ||
| 336 | + else: | ||
| 337 | + new_local_kwargs[key] = local_tensor | ||
| 338 | + | ||
| 339 | + op_info.local_kwargs = new_local_kwargs | ||
| 340 | + | ||
| 341 | + | ||
| 13 | def _npu_fusion_attention_handler( | 342 | def _npu_fusion_attention_handler( |
| 14 | op_call: torch._ops.OpOverload, | 343 | op_call: torch._ops.OpOverload, |
| 15 | args: Tuple[object, ...], | 344 | args: Tuple[object, ...], |
| @@ -56,18 +385,77 @@ def _npu_fusion_attention_handler( | |||
| 56 | DTensor._op_dispatcher.sharding_propagator.propagate(op_info) | 385 | DTensor._op_dispatcher.sharding_propagator.propagate(op_info) |
| 57 | output_sharding = op_info.output_sharding | 386 | output_sharding = op_info.output_sharding |
| 58 | 387 | ||
| 59 | - if op_call == npu.npu_fusion_attention.default: | 388 | + mesh = op_info.compute_mesh |
| 60 | - local_results = torch_npu.npu_fusion_attention( | 389 | + participating = mesh.get_coordinate() is not None |
| 61 | - *op_info.local_args, **op_info.local_kwargs | 390 | + if participating: |
| 62 | - ) | 391 | + # computation that happens in the current rank of the mesh, normal case |
| 63 | - elif op_call == npu.npu_fusion_attention_grad.default: | 392 | + if output_sharding.needs_redistribute: |
| 64 | - local_results = torch_npu.npu_fusion_attention_grad( | 393 | + DTensor._op_dispatcher.redistribute_local_args( |
| 65 | - *op_info.local_args, **op_info.local_kwargs | 394 | + op_info, |
| 395 | + output_sharding.redistribute_schema, | ||
| 396 | + output_sharding.use_val_from_redistribute_schema, | ||
| 397 | + ) | ||
| 398 | + local_args = ( | ||
| 399 | + pytree.tree_unflatten( | ||
| 400 | + cast(list[object], op_info.local_args), op_info.args_tree_spec | ||
| 401 | + ) | ||
| 402 | + if op_info.args_tree_spec | ||
| 403 | + else op_info.local_args | ||
| 404 | + ) | ||
| 405 | + | ||
| 406 | + if op_call == npu.npu_fusion_attention.default: | ||
| 407 | + # if sharding head_dim in qkv, need recalculate head_num in local args | ||
| 408 | + input_layout = op_info.local_args[4] | ||
| 409 | + if 'N' in input_layout: | ||
| 410 | + head_dim = input_layout.index('N') | ||
| 411 | + local_args = list(local_args) | ||
| 412 | + local_query = local_args[0] | ||
| 413 | + local_args[3] = local_query.size(head_dim) | ||
| 414 | + local_args = tuple(local_args) | ||
| 415 | + | ||
| 416 | + # run local op computation with potentially modified args/kwargs | ||
| 417 | + local_args = cast(tuple[object, ...], local_args) | ||
| 418 | + local_results = torch_npu.npu_fusion_attention( | ||
| 419 | + *local_args, **op_info.local_kwargs | ||
| 420 | + ) | ||
| 421 | + elif op_call == npu.npu_fusion_attention_grad.default: | ||
| 422 | + _redistribute_local_kwargs(op_info) | ||
| 423 | + # if sharding head_dim in qkv, need recalculate head_num in local args | ||
| 424 | + input_layout = op_info.local_args[5] | ||
| 425 | + if 'N' in input_layout: | ||
| 426 | + head_dim = input_layout.index('N') | ||
| 427 | + local_args = list(local_args) | ||
| 428 | + local_query = local_args[0] | ||
| 429 | + local_args[4] = local_query.size(head_dim) | ||
| 430 | + local_args = tuple(local_args) | ||
| 431 | + local_args = cast(tuple[object, ...], local_args) | ||
| 432 | + local_results = torch_npu.npu_fusion_attention_grad( | ||
| 433 | + *local_args, **op_info.local_kwargs | ||
| 434 | + ) | ||
| 435 | + else: | ||
| 436 | + raise NotImplementedError( | ||
| 437 | + "_npu_fusion_attention_handler only supports npu_fusion_attention and npu_fusion_attention_grad now." | ||
| 66 | ) | 438 | ) |
| 67 | else: | 439 | else: |
| 68 | - raise NotImplementedError( | 440 | + # For a non-participating device (happens on rank that does not belong to the device mesh), |
| 69 | - "_npu_fusion_attention_handler only supports npu_fusion_attention and npu_fusion_attention_grad now." | 441 | + # return empty tensor(s) with correct dtype. |
| 70 | - ) | 442 | + spec = output_sharding.output_spec |
| 443 | + | ||
| 444 | + def default_tensor(spec: DTensorSpec) -> torch.Tensor: | ||
| 445 | + if spec.tensor_meta is not None: | ||
| 446 | + shape = spec.tensor_meta.shape | ||
| 447 | + dtype = spec.tensor_meta.dtype | ||
| 448 | + if len(shape) == 0: | ||
| 449 | + # scalar tensor | ||
| 450 | + return torch.zeros((), dtype=dtype) | ||
| 451 | + else: | ||
| 452 | + # non-scalar tensor | ||
| 453 | + return torch.tensor([], dtype=dtype) | ||
| 454 | + else: | ||
| 455 | + raise RuntimeError(f"{spec} has no tensor metadata.") | ||
| 456 | + | ||
| 457 | + # only have Tensor and int outputs here | ||
| 458 | + local_results = [default_tensor(s) if s is not None else 0 for s in spec] | ||
| 71 | 459 | ||
| 72 | return DTensor._op_dispatcher.wrap(local_results, output_sharding.output_spec) | 460 | return DTensor._op_dispatcher.wrap(local_results, output_sharding.output_spec) |
| 73 | 461 | ||
| @@ -0,0 +1,120 @@ | |||
| 1 | +import itertools | ||
| 2 | +from typing import Callable, Optional | ||
| 3 | + | ||
| 4 | +import torch | ||
| 5 | +from torch.distributed.tensor._dtensor_spec import DTensorSpec | ||
| 6 | +from torch.distributed.tensor._op_schema import ( | ||
| 7 | + OpSchema, | ||
| 8 | + OpSpec, | ||
| 9 | + OpStrategy, | ||
| 10 | + PlacementList, | ||
| 11 | +) | ||
| 12 | +from torch.distributed.tensor._ops.utils import ( | ||
| 13 | + generate_redistribute_costs, | ||
| 14 | + is_tensor_shardable | ||
| 15 | +) | ||
| 16 | +from torch.distributed.tensor.device_mesh import DeviceMesh | ||
| 17 | + | ||
| 18 | +try: | ||
| 19 | + from torch.utils._cxx_pytree import tree_leaves | ||
| 20 | +except ImportError: | ||
| 21 | + from torch.utils._pytree import tree_leaves | ||
| 22 | + | ||
| 23 | + | ||
| 24 | +def _patched_kwargs_strategy(self) -> tuple[OpStrategy, ...]: | ||
| 25 | + kwargs_vals = ( | ||
| 26 | + tree_leaves(self.kwargs_schema) | ||
| 27 | + if self.schema_info is not None and self.schema_info.needs_pytree | ||
| 28 | + else self.kwargs_schema.values() | ||
| 29 | + ) | ||
| 30 | + return tuple(item for item in kwargs_vals if isinstance(item, OpStrategy)) | ||
| 31 | + | ||
| 32 | + | ||
| 33 | +def _patched_expand_to_full_mesh_op_strategy( | ||
| 34 | + mesh: DeviceMesh, | ||
| 35 | + op_schema: OpSchema, | ||
| 36 | + single_mesh_dim_strategies: list[PlacementList], | ||
| 37 | + *, | ||
| 38 | + input_index: int = 1, | ||
| 39 | + inplace_op: bool = False, | ||
| 40 | + is_valid_strategy_cb: Optional[ | ||
| 41 | + Callable[[list[DTensorSpec], tuple[Optional[DTensorSpec], ...]], bool] | ||
| 42 | + ] = None, | ||
| 43 | +) -> OpStrategy: | ||
| 44 | + # Expand the single_mesh_dim_strategies to full mesh dim strategies. | ||
| 45 | + all_mesh_dim_strategies = [single_mesh_dim_strategies] * mesh.ndim | ||
| 46 | + | ||
| 47 | + strategy_combs = itertools.product(*all_mesh_dim_strategies) | ||
| 48 | + | ||
| 49 | + all_strategies = [] | ||
| 50 | + for strategy_comb in strategy_combs: | ||
| 51 | + spec_list: list[Optional[DTensorSpec]] = [] | ||
| 52 | + for specs in zip(*strategy_comb): | ||
| 53 | + if specs[0] is not None: | ||
| 54 | + spec_list.append(DTensorSpec(mesh, specs)) | ||
| 55 | + else: | ||
| 56 | + spec_list.append(None) | ||
| 57 | + | ||
| 58 | + input_specs: list[DTensorSpec] = [s for s in spec_list[input_index:] if isinstance(s, DTensorSpec)] | ||
| 59 | + | ||
| 60 | + args_strategy = op_schema.args_strategy | ||
| 61 | + kwargs_strategy = op_schema.kwargs_strategy | ||
| 62 | + input_args_strategy = args_strategy + kwargs_strategy | ||
| 63 | + | ||
| 64 | + if len(input_specs) != len(input_args_strategy): | ||
| 65 | + raise AssertionError( | ||
| 66 | + f"input_specs({len(input_specs)}) != strategies({len(input_args_strategy)}: " | ||
| 67 | + f"{len(args_strategy)} args + {len(kwargs_strategy)} kwargs)" | ||
| 68 | + ) | ||
| 69 | + self_spec = input_args_strategy[0].strategies[0].output_spec | ||
| 70 | + | ||
| 71 | + if inplace_op and self_spec.placements != input_specs[0].placements: | ||
| 72 | + # if it's inplace op, we would only allow the OpSpec to be added when the | ||
| 73 | + # input_spec matches the first argument's runtime sharding, otherwise we skip | ||
| 74 | + continue | ||
| 75 | + | ||
| 76 | + output_specs: tuple[Optional[DTensorSpec], ...] | ||
| 77 | + if input_index > 1: | ||
| 78 | + output_specs = tuple(spec_list[:input_index]) | ||
| 79 | + else: | ||
| 80 | + if spec_list[0] is not None: | ||
| 81 | + output_specs = spec_list[0] # type: ignore[assignment] | ||
| 82 | + else: | ||
| 83 | + raise RuntimeError("output spec is None") | ||
| 84 | + | ||
| 85 | + # check all inputs are shardable | ||
| 86 | + if not all( | ||
| 87 | + is_tensor_shardable(inp.shape, s) | ||
| 88 | + for inp, s in zip(input_args_strategy, input_specs) | ||
| 89 | + ): | ||
| 90 | + continue | ||
| 91 | + | ||
| 92 | + # perform additional op-specific filtering | ||
| 93 | + if is_valid_strategy_cb is not None: | ||
| 94 | + if not is_valid_strategy_cb(input_specs, output_specs): | ||
| 95 | + continue | ||
| 96 | + | ||
| 97 | + redistribute_cost = [ | ||
| 98 | + generate_redistribute_costs(input_strategy, input_spec) | ||
| 99 | + for input_strategy, input_spec in zip(input_args_strategy, input_specs) | ||
| 100 | + ] | ||
| 101 | + | ||
| 102 | + strategy = OpSpec( | ||
| 103 | + output_specs=output_specs, | ||
| 104 | + input_specs=input_specs, | ||
| 105 | + redistribute_cost=redistribute_cost, | ||
| 106 | + ) | ||
| 107 | + all_strategies.append(strategy) | ||
| 108 | + return OpStrategy(all_strategies) | ||
| 109 | + | ||
| 110 | + | ||
| 111 | +def _apply_dtensor_patch(): | ||
| 112 | + # adding kwarg inputs handling in register sharding for previous pytorch version | ||
| 113 | + # See pytorch/pytorch/pull/168249 | ||
| 114 | + if torch.__version__ < "2.10": | ||
| 115 | + if not hasattr(OpSchema, "kwargs_strategy"): | ||
| 116 | + OpSchema.kwargs_strategy = property(_patched_kwargs_strategy) | ||
| 117 | + torch.distributed.tensor._ops.utils.expand_to_full_mesh_op_strategy = _patched_expand_to_full_mesh_op_strategy | ||
| 118 | + | ||
| 119 | + | ||
| 120 | +_apply_dtensor_patch() | ||
| @@ -131,147 +131,3 @@ def custom_matmul_backward_sharding( | |||
| 131 | 131 | ||
| 132 | acceptable_shardings.append(strategy) | 132 | acceptable_shardings.append(strategy) |
| 133 | return acceptable_shardings | 133 | return acceptable_shardings |
| 134 | - | ||
| 135 | - | ||
| 136 | - | ||
| 137 | -# pylint:disable=huawei-too-many-arguments | ||
| 138 | -def custom_npu_fusion_attention_sharding(query, key, value, head_num, input_layout, pse=None, padding_mask=None, | ||
| 139 | - atten_mask=None, scale=1.0, keep_prob=1.0, pre_tockens=2147483647, | ||
| 140 | - inner_precise=0, prefix=None, actual_seq_qlen=None, actual_seq_kvlen=None, | ||
| 141 | - sparse_mode=0, gen_mask_parallel=True, sync=False): | ||
| 142 | - acceptable_shardings = [] | ||
| 143 | - | ||
| 144 | - # add all replicate strategy | ||
| 145 | - replcate_strategy = ( | ||
| 146 | - [ | ||
| 147 | - Replicate(), # Tensor attention_score | ||
| 148 | - Replicate(), # Tensor softmax_max | ||
| 149 | - Replicate(), # Tensor softmax_sum | ||
| 150 | - Replicate(), # Tensor softmax_out | ||
| 151 | - None, # int seed | ||
| 152 | - None, # int offset | ||
| 153 | - None # int numels | ||
| 154 | - ], | ||
| 155 | - [ | ||
| 156 | - Replicate(), # Tensor query | ||
| 157 | - Replicate(), # Tensor key | ||
| 158 | - Replicate(), # Tensor value | ||
| 159 | - None, # int head_num | ||
| 160 | - None, # str input_layout | ||
| 161 | - None if pse is None else Replicate(), # Tensor? pse | ||
| 162 | - None if padding_mask is None else Replicate(), # Tensor? padding_mask | ||
| 163 | - None if atten_mask is None else Replicate(), # Tensor? atten_mask | ||
| 164 | - None # other | ||
| 165 | - ] | ||
| 166 | - ) | ||
| 167 | - | ||
| 168 | - # add sharding strategy | ||
| 169 | - for strategy_index, default_sharding in enumerate(query.placements): | ||
| 170 | - pse_sharding = None if pse is None else pse.placements[strategy_index] | ||
| 171 | - padding_mask_sharding = None if padding_mask is None else padding_mask.placements[strategy_index] | ||
| 172 | - atten_mask_sharding = None if atten_mask is None else atten_mask.placements[strategy_index] | ||
| 173 | - | ||
| 174 | - sharding_strategy = ( | ||
| 175 | - [ | ||
| 176 | - default_sharding, # Tensor attention_score | ||
| 177 | - default_sharding, # Tensor softmax_max | ||
| 178 | - default_sharding, # Tensor softmax_sum | ||
| 179 | - default_sharding, # Tensor softmax_out | ||
| 180 | - None, # int seed | ||
| 181 | - None, # int offset | ||
| 182 | - None # int numels | ||
| 183 | - ], | ||
| 184 | - [ | ||
| 185 | - default_sharding, # Tensor query | ||
| 186 | - default_sharding, # Tensor key | ||
| 187 | - default_sharding, # Tensor value | ||
| 188 | - None, # int head_num | ||
| 189 | - None, # str input_layout | ||
| 190 | - pse_sharding, # Tensor? pse | ||
| 191 | - padding_mask_sharding, # Tensor? padding_mask | ||
| 192 | - atten_mask_sharding, # Tensor? atten_mask | ||
| 193 | - None # other | ||
| 194 | - ] | ||
| 195 | - ) | ||
| 196 | - | ||
| 197 | - acceptable_shardings.append(sharding_strategy) | ||
| 198 | - | ||
| 199 | - acceptable_shardings.append(replcate_strategy) | ||
| 200 | - | ||
| 201 | - return acceptable_shardings | ||
| 202 | - | ||
| 203 | - | ||
| 204 | - | ||
| 205 | -# pylint:disable=huawei-too-many-arguments | ||
| 206 | -def custom_npu_fusion_attention_grad_sharding(query, key, value, dy, head_num, input_layout, *, pse=None, | ||
| 207 | - padding_mask=None, atten_mask=None, softmax_max=None, softmax_sum=None, | ||
| 208 | - softmax_in=None, attention_in=None, scale_value=1.0, keep_prob=1.0, | ||
| 209 | - pre_tockens=2147483647, next_tockens=2147483647, inner_precise=0, seed=0, | ||
| 210 | - offset=0, numels=0, prefix=None, actual_seq_qlen=None, | ||
| 211 | - actual_seq_kvlen=None, sparse_mode=0, gen_mask_parallel=True, sync=False): | ||
| 212 | - acceptable_shardings = [] | ||
| 213 | - | ||
| 214 | - # add all replicate strategy | ||
| 215 | - replcate_strategy = ( | ||
| 216 | - [ | ||
| 217 | - Replicate(), # Tensor grad_query | ||
| 218 | - Replicate(), # Tensor grad_key | ||
| 219 | - Replicate(), # Tensor grad_value | ||
| 220 | - Replicate(), # Tensor grad_dy | ||
| 221 | - ], | ||
| 222 | - [ | ||
| 223 | - Replicate(), # Tensor query | ||
| 224 | - Replicate(), # Tensor key | ||
| 225 | - Replicate(), # Tensor value | ||
| 226 | - Replicate(), # Tensor dy | ||
| 227 | - None, # int head_num | ||
| 228 | - None, # str input_layout | ||
| 229 | - None if pse is None else Replicate(), # Tensor? pse | ||
| 230 | - None if padding_mask is None else Replicate(), # Tensor? padding_mask | ||
| 231 | - None if atten_mask is None else Replicate(), # Tensor? atten_mask | ||
| 232 | - None if softmax_max is None else Replicate(), # Tensor? softmax_max | ||
| 233 | - None if softmax_sum is None else Replicate(), # Tensor? softmax_sum | ||
| 234 | - None if softmax_in is None else Replicate(), # Tensor? softmax_in | ||
| 235 | - None if attention_in is None else Replicate(), # Tensor? attention_in | ||
| 236 | - None # other | ||
| 237 | - ] | ||
| 238 | - ) | ||
| 239 | - acceptable_shardings.append(replcate_strategy) | ||
| 240 | - | ||
| 241 | - # add sharding strategy | ||
| 242 | - for strategy_index, default_sharding in enumerate(query.placements): | ||
| 243 | - pse_sharding = None if pse is None else pse.placements[strategy_index] | ||
| 244 | - padding_mask_sharding = None if padding_mask is None else padding_mask.placements[strategy_index] | ||
| 245 | - atten_mask_sharding = None if atten_mask is None else atten_mask.placements[strategy_index] | ||
| 246 | - asoftmax_max_sharding = None if softmax_max is None else softmax_max.placements[strategy_index] | ||
| 247 | - softmax_sum_sharding = None if softmax_sum is None else softmax_sum.placements[strategy_index] | ||
| 248 | - softmax_in_sharding = None if softmax_in is None else softmax_in.placements[strategy_index] | ||
| 249 | - attention_in_sharding = None if attention_in is None else attention_in.placements[strategy_index] | ||
| 250 | - | ||
| 251 | - sharding_strategy = ( | ||
| 252 | - [ | ||
| 253 | - default_sharding, # Tensor grad_query | ||
| 254 | - default_sharding, # Tensor grad_key | ||
| 255 | - default_sharding, # Tensor grad_value | ||
| 256 | - default_sharding, # Tensor grad_dy | ||
| 257 | - ], | ||
| 258 | - [ | ||
| 259 | - default_sharding, # Tensor query | ||
| 260 | - default_sharding, # Tensor key | ||
| 261 | - default_sharding, # Tensor value | ||
| 262 | - default_sharding, # Tensor dy | ||
| 263 | - None, # int head_num | ||
| 264 | - None, # str input_layout | ||
| 265 | - pse_sharding, # Tensor? pse | ||
| 266 | - padding_mask_sharding, # Tensor? padding_mask | ||
| 267 | - atten_mask_sharding, # Tensor? atten_mask | ||
| 268 | - asoftmax_max_sharding, # Tensor? softmax_max | ||
| 269 | - softmax_sum_sharding, # Tensor? softmax_sum | ||
| 270 | - softmax_in_sharding, # Tensor? softmax_sum | ||
| 271 | - attention_in_sharding, # Tensor? attention_in | ||
| 272 | - None # other | ||
| 273 | - ] | ||
| 274 | - ) | ||
| 275 | - acceptable_shardings.append(sharding_strategy) | ||
| 276 | - | ||
| 277 | - return acceptable_shardings | ||