已合并
feat: kvallgather cp thd load balance #3820
m0_50947149创建于 7月26日
feat: kvallgather cp thd load balance #3820
已合并
共 4 个文件变更+694-363
| @@ -5,25 +5,28 @@ from mindspeed.core.context_parallel import get_args | |||
| 5 | from mindspeed.core.context_parallel import mpu | 5 | from mindspeed.core.context_parallel import mpu |
| 6 | from mindspeed.core.tensor_parallel_y_union_cp import TensorParallelYUnionCP | 6 | from mindspeed.core.tensor_parallel_y_union_cp import TensorParallelYUnionCP |
| 7 | from mindspeed.model.transformer import set_attention_mask, get_attention_mask | 7 | from mindspeed.model.transformer import set_attention_mask, get_attention_mask |
| 8 | -from mindspeed.core.context_parallel.model_parallel_utils import (get_context_parallel_for_hybrid_ulysses_world_size, | 8 | +from mindspeed.core.context_parallel.model_parallel_utils import ( |
| 9 | - get_context_parallel_for_hybrid_ulysses_rank, | 9 | + get_context_parallel_for_hybrid_ulysses_world_size, |
| 10 | - get_context_parallel_for_hybrid_ring_world_size, | 10 | + get_context_parallel_for_hybrid_ulysses_rank, |
| 11 | - get_context_parallel_for_hybrid_ring_rank) | 11 | + get_context_parallel_for_hybrid_ring_world_size, |
| 12 | -from mindspeed.core.context_parallel.utils import (set_scheduling_info, | 12 | + get_context_parallel_for_hybrid_ring_rank, |
| 13 | - set_remapped_seq_order, | 13 | +) |
| 14 | - adaptive_reschedule_task, | 14 | +from mindspeed.core.context_parallel.utils import ( |
| 15 | - get_adaptive_cp_mask_list_by_user, | 15 | + set_scheduling_info, |
| 16 | - get_adaptive_cp_grid_mask_by_user, | 16 | + set_remapped_seq_order, |
| 17 | - generate_adaptive_cp_mask_list_by_user, | 17 | + adaptive_reschedule_task, |
| 18 | - generate_adaptive_cp_grid_mask_by_user, | 18 | + get_adaptive_cp_mask_list_by_user, |
| 19 | - pad_data) | 19 | + get_adaptive_cp_grid_mask_by_user, |
| 20 | + generate_adaptive_cp_mask_list_by_user, | ||
| 21 | + generate_adaptive_cp_grid_mask_by_user, | ||
| 22 | + pad_data, | ||
| 23 | +) | ||
| 20 | 24 | ||
| 21 | _ACTUAL_SEQ_LEN = None | 25 | _ACTUAL_SEQ_LEN = None |
| 22 | _REARRANGE_IDX_TENSOR = None | 26 | _REARRANGE_IDX_TENSOR = None |
| 23 | 27 | ||
| 24 | 28 | ||
| 25 | def get_actual_seq_len(): | 29 | def get_actual_seq_len(): |
| 26 | - global _ACTUAL_SEQ_LEN | ||
| 27 | return _ACTUAL_SEQ_LEN | 30 | return _ACTUAL_SEQ_LEN |
| 28 | 31 | ||
| 29 | 32 | ||
| @@ -50,7 +53,9 @@ def get_ring_degree(): | |||
| 50 | 53 | ||
| 51 | def _broadcast(item): | 54 | def _broadcast(item): |
| 52 | if item is not None: | 55 | if item is not None: |
| 53 | - torch.distributed.broadcast(item, mpu.get_tensor_model_parallel_src_rank(), group=mpu.get_tensor_model_parallel_group()) | 56 | + torch.distributed.broadcast( |
| 57 | + item, mpu.get_tensor_model_parallel_src_rank(), group=mpu.get_tensor_model_parallel_group() | ||
| 58 | + ) | ||
| 54 | 59 | ||
| 55 | 60 | ||
| 56 | def broadcast_dynamic(item): | 61 | def broadcast_dynamic(item): |
| @@ -74,21 +79,24 @@ def generate_rearrange_idx_tensor(tp_y_cp_size): | |||
| 74 | rearrange_index = [] | 79 | rearrange_index = [] |
| 75 | for i in range(tp_y_cp_size): | 80 | for i in range(tp_y_cp_size): |
| 76 | rearrange_index.extend([i, 2 * tp_y_cp_size - 1 - i]) | 81 | rearrange_index.extend([i, 2 * tp_y_cp_size - 1 - i]) |
| 77 | - _REARRANGE_IDX_TENSOR = torch.tensor(rearrange_index, device='cpu', pin_memory=True).to(device='npu', non_blocking=True) | 82 | + _REARRANGE_IDX_TENSOR = torch.tensor(rearrange_index, device='cpu', pin_memory=True).to( |
| 83 | + device='npu', non_blocking=True | ||
| 84 | + ) | ||
| 78 | return _REARRANGE_IDX_TENSOR | 85 | return _REARRANGE_IDX_TENSOR |
| 79 | 86 | ||
| 80 | 87 | ||
| 81 | def batch_index(seq1d, seq_len): | 88 | def batch_index(seq1d, seq_len): |
| 82 | from bisect import bisect_right | 89 | from bisect import bisect_right |
| 90 | + | ||
| 83 | end_points = list(range(seq_len, seq1d[-1] + 1, seq_len)) | 91 | end_points = list(range(seq_len, seq1d[-1] + 1, seq_len)) |
| 84 | indexes = [0] + [bisect_right(seq1d, p) for p in end_points] | 92 | indexes = [0] + [bisect_right(seq1d, p) for p in end_points] |
| 85 | - seq_batch = [seq1d[indexes[i]:indexes[i + 1]] for i in range(len(indexes) - 1)] | 93 | + seq_batch = [seq1d[indexes[i] : indexes[i + 1]] for i in range(len(indexes) - 1)] |
| 86 | return [[elem - i * seq_len for elem in seq] for i, seq in enumerate(seq_batch)] | 94 | return [[elem - i * seq_len for elem in seq] for i, seq in enumerate(seq_batch)] |
| 87 | 95 | ||
| 88 | 96 | ||
| 89 | def get_batch_on_this_cp_rank(batch): | 97 | def get_batch_on_this_cp_rank(batch): |
| 90 | - """ Slice batch input along sequence dimension into multiple chunks, | 98 | + """Slice batch input along sequence dimension into multiple chunks, |
| 91 | - which are parallelized across GPUs in a context parallel group. | 99 | + which are parallelized across GPUs in a context parallel group. |
| 92 | """ | 100 | """ |
| 93 | 101 | ||
| 94 | # With causal masking, each token only attends to its prior tokens. Simply split | 102 | # With causal masking, each token only attends to its prior tokens. Simply split |
| @@ -104,13 +112,15 @@ def get_batch_on_this_cp_rank(batch): | |||
| 104 | if cp_size == 1: | 112 | if cp_size == 1: |
| 105 | return batch | 113 | return batch |
| 106 | 114 | ||
| 107 | - tp_y_cp_size = TensorParallelYUnionCP().get_parallel_group_world_size() if args.tp_2d else args.context_parallel_size | 115 | + tp_y_cp_size = ( |
| 116 | + TensorParallelYUnionCP().get_parallel_group_world_size() if args.tp_2d else args.context_parallel_size | ||
| 117 | + ) | ||
| 108 | if not tp_y_cp_size > 1: | 118 | if not tp_y_cp_size > 1: |
| 109 | return batch | 119 | return batch |
| 110 | 120 | ||
| 111 | cp_expanded_by_2d_tp = args.tp_y > 1 | 121 | cp_expanded_by_2d_tp = args.tp_y > 1 |
| 112 | if args.reset_attention_mask and args.attention_mask_type == 'causal': | 122 | if args.reset_attention_mask and args.attention_mask_type == 'causal': |
| 113 | - if args.context_parallel_algo in ['ulysses_cp_algo', 'kvallgather_cp_algo']: | 123 | + if args.context_parallel_algo == 'ulysses_cp_algo': |
| 114 | batch = _get_batch_on_this_cp_rank_in_ulysses_cp(batch) | 124 | batch = _get_batch_on_this_cp_rank_in_ulysses_cp(batch) |
| 115 | else: | 125 | else: |
| 116 | batch = _get_batch_on_this_cp_rank_in_megatron_cp_eod_padding(batch, get_actual_seq_len()) | 126 | batch = _get_batch_on_this_cp_rank_in_megatron_cp_eod_padding(batch, get_actual_seq_len()) |
| @@ -138,10 +148,18 @@ def get_batch_on_this_cp_rank(batch): | |||
| 138 | 148 | ||
| 139 | 149 | ||
| 140 | def _get_batch_on_this_cp_rank_in_megatron_cp_eod_padding(batch, actual_seq_len): | 150 | def _get_batch_on_this_cp_rank_in_megatron_cp_eod_padding(batch, actual_seq_len): |
| 151 | + if actual_seq_len is None or actual_seq_len.numel() == 0: | ||
| 152 | + raise AssertionError("actual_seq_len must be provided for EOD load-balanced context parallelism.") | ||
| 153 | + | ||
| 141 | def get_index(actual_seq_len_cpu, cp_rank, cp_size): | 154 | def get_index(actual_seq_len_cpu, cp_rank, cp_size): |
| 142 | starts = torch.cat([torch.tensor([0]), actual_seq_len_cpu[:-1]]) | 155 | starts = torch.cat([torch.tensor([0]), actual_seq_len_cpu[:-1]]) |
| 143 | ends = actual_seq_len_cpu | 156 | ends = actual_seq_len_cpu |
| 144 | - chunk_sizes = (ends - starts) // (2 * cp_size) | 157 | + seq_lens = ends - starts |
| 158 | + if torch.any(seq_lens <= 0): | ||
| 159 | + raise AssertionError("Each EOD subsequence must have a positive sequence length.") | ||
| 160 | + if torch.any(seq_lens % (2 * cp_size) != 0): | ||
| 161 | + raise AssertionError("Each EOD subsequence length must be divisible by 2 * context_parallel_size.") | ||
| 162 | + chunk_sizes = seq_lens // (2 * cp_size) | ||
| 145 | 163 | ||
| 146 | first_starts = starts + cp_rank * chunk_sizes | 164 | first_starts = starts + cp_rank * chunk_sizes |
| 147 | first_ends = first_starts + chunk_sizes | 165 | first_ends = first_starts + chunk_sizes |
| @@ -168,9 +186,9 @@ def _get_batch_on_this_cp_rank_in_megatron_cp_eod_padding(batch, actual_seq_len) | |||
| 168 | if val is not None: | 186 | if val is not None: |
| 169 | seq_dim = 1 if key != 'attention_mask' else 2 | 187 | seq_dim = 1 if key != 'attention_mask' else 2 |
| 170 | bsz = val.shape[0] | 188 | bsz = val.shape[0] |
| 171 | - val = val.view(-1, *val.shape[seq_dim + 1:]) | 189 | + val = val.view(-1, *val.shape[seq_dim + 1 :]) |
| 172 | val = val.index_select(0, index) | 190 | val = val.index_select(0, index) |
| 173 | - val = val.view(bsz, -1, *val.shape[seq_dim + 1:]) | 191 | + val = val.view(bsz, -1, *val.shape[seq_dim + 1 :]) |
| 174 | 192 | ||
| 175 | batch[key] = val | 193 | batch[key] = val |
| 176 | 194 | ||
| @@ -222,7 +240,7 @@ def _get_batch_on_this_tp_y_cp_rank_in_megatron_cp(batch): | |||
| 222 | *val.shape[0:seq_dim], | 240 | *val.shape[0:seq_dim], |
| 223 | 2 * tp_y_cp_size, | 241 | 2 * tp_y_cp_size, |
| 224 | val.shape[seq_dim] // (2 * tp_y_cp_size), | 242 | val.shape[seq_dim] // (2 * tp_y_cp_size), |
| 225 | - *val.shape[(seq_dim + 1):], | 243 | + *val.shape[(seq_dim + 1) :], |
| 226 | ) | 244 | ) |
| 227 | 245 | ||
| 228 | val = val.index_select(seq_dim, index=rearrange_idx_tensor) | 246 | val = val.index_select(seq_dim, index=rearrange_idx_tensor) |
| @@ -232,7 +250,7 @@ def _get_batch_on_this_tp_y_cp_rank_in_megatron_cp(batch): | |||
| 232 | *val.shape[0:seq_dim], | 250 | *val.shape[0:seq_dim], |
| 233 | cp_size, | 251 | cp_size, |
| 234 | val.shape[seq_dim] // cp_size, | 252 | val.shape[seq_dim] // cp_size, |
| 235 | - *val.shape[(seq_dim + 1):], | 253 | + *val.shape[(seq_dim + 1) :], |
| 236 | ) | 254 | ) |
| 237 | # [b, 1, s/cp] -> [b, s/cp] | 255 | # [b, 1, s/cp] -> [b, s/cp] |
| 238 | val = val[:, cp_rank].view(b, -1) | 256 | val = val[:, cp_rank].view(b, -1) |
| @@ -253,11 +271,11 @@ def _get_batch_on_this_cp_rank_in_megatron_cp(batch): | |||
| 253 | *val.shape[0:seq_dim], | 271 | *val.shape[0:seq_dim], |
| 254 | 2 * cp_size, | 272 | 2 * cp_size, |
| 255 | val.shape[seq_dim] // (2 * cp_size), | 273 | val.shape[seq_dim] // (2 * cp_size), |
| 256 | - *val.shape[(seq_dim + 1):], | 274 | + *val.shape[(seq_dim + 1) :], |
| 257 | ) | 275 | ) |
| 258 | index = torch.tensor([cp_rank, (2 * cp_size - cp_rank - 1)], device=val.device) | 276 | index = torch.tensor([cp_rank, (2 * cp_size - cp_rank - 1)], device=val.device) |
| 259 | val = val.index_select(seq_dim, index) | 277 | val = val.index_select(seq_dim, index) |
| 260 | - val = val.view(*val.shape[0:seq_dim], -1, *val.shape[(seq_dim + 2):]) | 278 | + val = val.view(*val.shape[0:seq_dim], -1, *val.shape[(seq_dim + 2) :]) |
| 261 | batch[key] = val | 279 | batch[key] = val |
| 262 | 280 | ||
| 263 | return batch | 281 | return batch |
| @@ -323,11 +341,11 @@ def _get_batch_on_this_cp_rank_in_hybrid_cp(batch): | |||
| 323 | *val.shape[0:seq_dim], | 341 | *val.shape[0:seq_dim], |
| 324 | 2 * r_size, | 342 | 2 * r_size, |
| 325 | val.shape[seq_dim] // (2 * r_size), | 343 | val.shape[seq_dim] // (2 * r_size), |
| 326 | - *val.shape[(seq_dim + 1):], | 344 | + *val.shape[(seq_dim + 1) :], |
| 327 | ) | 345 | ) |
| 328 | index = torch.tensor([r_rank, (2 * r_size - r_rank - 1)], device=val.device) | 346 | index = torch.tensor([r_rank, (2 * r_size - r_rank - 1)], device=val.device) |
| 329 | val = val.index_select(seq_dim, index) | 347 | val = val.index_select(seq_dim, index) |
| 330 | - val = val.view(*val.shape[0:seq_dim], -1, *val.shape[(seq_dim + 2):]) | 348 | + val = val.view(*val.shape[0:seq_dim], -1, *val.shape[(seq_dim + 2) :]) |
| 331 | val = val.chunk(u_size, dim=seq_dim)[u_rank].contiguous() | 349 | val = val.chunk(u_size, dim=seq_dim)[u_rank].contiguous() |
| 332 | batch[key] = val | 350 | batch[key] = val |
| 333 | 351 | ||
| @@ -353,6 +371,7 @@ def _get_batch_on_this_cp_rank_in_adaptive_cp(batch): | |||
| 353 | if len(attention_mask.shape) != 2: | 371 | if len(attention_mask.shape) != 2: |
| 354 | raise AssertionError("The fusion attention operator currently only support 2D attention mask.") | 372 | raise AssertionError("The fusion attention operator currently only support 2D attention mask.") |
| 355 | from mindspeed.core.context_parallel.utils import AdaptiveCpOps | 373 | from mindspeed.core.context_parallel.utils import AdaptiveCpOps |
| 374 | + | ||
| 356 | adaptive_cp_ops = AdaptiveCpOps() | 375 | adaptive_cp_ops = AdaptiveCpOps() |
| 357 | remapped_seq_order, scheduling = adaptive_cp_ops.get_adaptive_cp_info(attention_mask, cp_size) | 376 | remapped_seq_order, scheduling = adaptive_cp_ops.get_adaptive_cp_info(attention_mask, cp_size) |
| 358 | mask_list = adaptive_cp_ops.get_mask_list(attention_mask, scheduling, remapped_seq_order, cp_rank, cp_size) | 377 | mask_list = adaptive_cp_ops.get_mask_list(attention_mask, scheduling, remapped_seq_order, cp_rank, cp_size) |
| @@ -366,8 +385,9 @@ def _get_batch_on_this_cp_rank_in_adaptive_cp(batch): | |||
| 366 | if key != 'attention_mask' and val is not None: | 385 | if key != 'attention_mask' and val is not None: |
| 367 | seq_dim = 1 | 386 | seq_dim = 1 |
| 368 | per = val.shape[seq_dim] // cp_size | 387 | per = val.shape[seq_dim] // cp_size |
| 369 | - index = torch.tensor(remapped_seq_order[cp_rank * per:(cp_rank + 1) * per], device=val.device, | 388 | + index = torch.tensor( |
| 370 | - dtype=torch.int) | 389 | + remapped_seq_order[cp_rank * per : (cp_rank + 1) * per], device=val.device, dtype=torch.int |
| 390 | + ) | ||
| 371 | val = val.index_select(seq_dim, index) | 391 | val = val.index_select(seq_dim, index) |
| 372 | batch[key] = val | 392 | batch[key] = val |
| 373 | return batch | 393 | return batch |
| @@ -394,6 +414,7 @@ def _get_batch_on_this_cp_rank_in_hybrid_adaptive_cp(batch): | |||
| 394 | if len(attention_mask.shape) != 2: | 414 | if len(attention_mask.shape) != 2: |
| 395 | raise AssertionError("The fusion attention operator currently only support 2D attention mask.") | 415 | raise AssertionError("The fusion attention operator currently only support 2D attention mask.") |
| 396 | from mindspeed.core.context_parallel.utils import AdaptiveCpOps | 416 | from mindspeed.core.context_parallel.utils import AdaptiveCpOps |
| 417 | + | ||
| 397 | adaptive_cp_ops = AdaptiveCpOps() | 418 | adaptive_cp_ops = AdaptiveCpOps() |
| 398 | remapped_seq_order, scheduling = adaptive_cp_ops.get_adaptive_cp_info(attention_mask, adap_size) | 419 | remapped_seq_order, scheduling = adaptive_cp_ops.get_adaptive_cp_info(attention_mask, adap_size) |
| 399 | mask_list = adaptive_cp_ops.get_mask_list(attention_mask, scheduling, remapped_seq_order, adap_rank, adap_size) | 420 | mask_list = adaptive_cp_ops.get_mask_list(attention_mask, scheduling, remapped_seq_order, adap_rank, adap_size) |
| @@ -408,7 +429,7 @@ def _get_batch_on_this_cp_rank_in_hybrid_adaptive_cp(batch): | |||
| 408 | seq_dim = 1 | 429 | seq_dim = 1 |
| 409 | per = val.shape[seq_dim] // adap_size // ulys_size | 430 | per = val.shape[seq_dim] // adap_size // ulys_size |
| 410 | which_per = adap_rank * ulys_size + ulys_rank | 431 | which_per = adap_rank * ulys_size + ulys_rank |
| 411 | - index = torch.tensor(remapped_seq_order[which_per * per:(which_per + 1) * per], device=val.device) | 432 | + index = torch.tensor(remapped_seq_order[which_per * per : (which_per + 1) * per], device=val.device) |
| 412 | val = val.index_select(seq_dim, index) | 433 | val = val.index_select(seq_dim, index) |
| 413 | batch[key] = val | 434 | batch[key] = val |
| 414 | 435 | ||
| @@ -429,7 +450,7 @@ def get_batch_on_this_tp_rank(data_iterator): | |||
| 429 | 'labels': data["labels"].cuda(non_blocking=True), | 450 | 'labels': data["labels"].cuda(non_blocking=True), |
| 430 | 'loss_mask': data["loss_mask"].cuda(non_blocking=True), | 451 | 'loss_mask': data["loss_mask"].cuda(non_blocking=True), |
| 431 | 'attention_mask': None if "attention_mask" not in data else data["attention_mask"].cuda(non_blocking=True), | 452 | 'attention_mask': None if "attention_mask" not in data else data["attention_mask"].cuda(non_blocking=True), |
| 432 | - 'position_ids': data["position_ids"].cuda(non_blocking=True) | 453 | + 'position_ids': data["position_ids"].cuda(non_blocking=True), |
| 433 | } | 454 | } |
| 434 | 455 | ||
| 435 | if args.pipeline_model_parallel_size == 1: | 456 | if args.pipeline_model_parallel_size == 1: |
| @@ -456,29 +477,38 @@ def get_batch_on_this_tp_rank(data_iterator): | |||
| 456 | 477 | ||
| 457 | if args.reset_attention_mask: | 478 | if args.reset_attention_mask: |
| 458 | actual_seq_len = broadcast_dynamic(data['actual_seq_len']) | 479 | actual_seq_len = broadcast_dynamic(data['actual_seq_len']) |
| 459 | - if args.attention_mask_type == 'causal' \ | 480 | + if ( |
| 460 | - and args.context_parallel_size > 1 \ | 481 | + args.attention_mask_type == 'causal' |
| 461 | - and args.context_parallel_algo == 'megatron_cp_algo': | 482 | + and args.context_parallel_size > 1 |
| 462 | - actual_seq_len = pad_data(actual_seq_len, batch, args.context_parallel_size, args.tensor_model_parallel_size) | 483 | + and args.context_parallel_algo in ['megatron_cp_algo', 'kvallgather_cp_algo'] |
| 463 | - actual_seq_len /= get_ring_degree() | 484 | + ): |
| 485 | + actual_seq_len = pad_data( | ||
| 486 | + actual_seq_len, batch, args.context_parallel_size, args.tensor_model_parallel_size | ||
| 487 | + ) | ||
| 488 | + actual_seq_len = actual_seq_len // get_ring_degree() | ||
| 464 | set_actual_seq_len(actual_seq_len) | 489 | set_actual_seq_len(actual_seq_len) |
| 465 | 490 | ||
| 466 | else: | 491 | else: |
| 467 | - tokens = torch.empty((args.micro_batch_size, args.seq_length), dtype=torch.int64, | 492 | + tokens = torch.empty( |
| 468 | - device=torch.cuda.current_device()) | 493 | + (args.micro_batch_size, args.seq_length), dtype=torch.int64, device=torch.cuda.current_device() |
| 469 | - labels = torch.empty((args.micro_batch_size, args.seq_length), dtype=torch.int64, | 494 | + ) |
| 470 | - device=torch.cuda.current_device()) | 495 | + labels = torch.empty( |
| 471 | - loss_mask = torch.empty((args.micro_batch_size, args.seq_length), dtype=torch.float32, | 496 | + (args.micro_batch_size, args.seq_length), dtype=torch.int64, device=torch.cuda.current_device() |
| 472 | - device=torch.cuda.current_device()) | 497 | + ) |
| 498 | + loss_mask = torch.empty( | ||
| 499 | + (args.micro_batch_size, args.seq_length), dtype=torch.float32, device=torch.cuda.current_device() | ||
| 500 | + ) | ||
| 473 | if getattr(args, 'create_attention_mask_in_dataloader', False): | 501 | if getattr(args, 'create_attention_mask_in_dataloader', False): |
| 474 | attention_mask = torch.empty( | 502 | attention_mask = torch.empty( |
| 475 | - (args.micro_batch_size, 1, args.seq_length, args.seq_length), dtype=torch.bool, | 503 | + (args.micro_batch_size, 1, args.seq_length, args.seq_length), |
| 476 | - device=torch.cuda.current_device() | 504 | + dtype=torch.bool, |
| 505 | + device=torch.cuda.current_device(), | ||
| 477 | ) | 506 | ) |
| 478 | else: | 507 | else: |
| 479 | attention_mask = None | 508 | attention_mask = None |
| 480 | - position_ids = torch.empty((args.micro_batch_size, args.seq_length), dtype=torch.int64, | 509 | + position_ids = torch.empty( |
| 481 | - device=torch.cuda.current_device()) | 510 | + (args.micro_batch_size, args.seq_length), dtype=torch.int64, device=torch.cuda.current_device() |
| 511 | + ) | ||
| 482 | 512 | ||
| 483 | if args.pipeline_model_parallel_size == 1: | 513 | if args.pipeline_model_parallel_size == 1: |
| 484 | _broadcast(tokens) | 514 | _broadcast(tokens) |
| @@ -514,17 +544,20 @@ def get_batch_on_this_tp_rank(data_iterator): | |||
| 514 | 'labels': labels, | 544 | 'labels': labels, |
| 515 | 'loss_mask': loss_mask, | 545 | 'loss_mask': loss_mask, |
| 516 | 'attention_mask': attention_mask, | 546 | 'attention_mask': attention_mask, |
| 517 | - 'position_ids': position_ids | 547 | + 'position_ids': position_ids, |
| 518 | } | 548 | } |
| 519 | 549 | ||
| 520 | if args.reset_attention_mask: | 550 | if args.reset_attention_mask: |
| 521 | actual_seq_len = broadcast_dynamic(None) | 551 | actual_seq_len = broadcast_dynamic(None) |
| 522 | - if args.attention_mask_type == 'causal' \ | 552 | + if ( |
| 523 | - and args.context_parallel_size > 1 \ | 553 | + args.attention_mask_type == 'causal' |
| 524 | - and args.context_parallel_algo == 'megatron_cp_algo': | 554 | + and args.context_parallel_size > 1 |
| 525 | - actual_seq_len = pad_data(actual_seq_len, batch, args.context_parallel_size, args.tensor_model_parallel_size) | 555 | + and args.context_parallel_algo in ['megatron_cp_algo', 'kvallgather_cp_algo'] |
| 526 | - actual_seq_len /= get_ring_degree() | 556 | + ): |
| 557 | + actual_seq_len = pad_data( | ||
| 558 | + actual_seq_len, batch, args.context_parallel_size, args.tensor_model_parallel_size | ||
| 559 | + ) | ||
| 560 | + actual_seq_len = actual_seq_len // get_ring_degree() | ||
| 527 | set_actual_seq_len(actual_seq_len) | 561 | set_actual_seq_len(actual_seq_len) |
| 528 | 562 | ||
| 529 | return batch | 563 | return batch |
| 530 | - | ||
| @@ -1,19 +1,15 @@ | |||
| 1 | # Copyright (c) 2025, Huawei Technologies Co., Ltd. All rights reserved. | 1 | # Copyright (c) 2025, Huawei Technologies Co., Ltd. All rights reserved. |
| 2 | 2 | ||
| 3 | from abc import ABC, abstractmethod | 3 | from abc import ABC, abstractmethod |
| 4 | -from typing import Optional, List, Any | 4 | +from typing import List, Any |
| 5 | 5 | ||
| 6 | import torch | 6 | import torch |
| 7 | -from einops import rearrange | ||
| 8 | 7 | ||
| 9 | from .kvallgather_context_parallel import ( | 8 | from .kvallgather_context_parallel import ( |
| 10 | AttnFuncWithCPAndKVAllGatherForSBHD, | 9 | AttnFuncWithCPAndKVAllGatherForSBHD, |
| 11 | AttnFuncWithCPAndKVAllGatherForTHD, | 10 | AttnFuncWithCPAndKVAllGatherForTHD, |
| 12 | ) | 11 | ) |
| 13 | -from .hamilton_context_parallel import ( | 12 | +from .hamilton_context_parallel import AttnFuncWithCPAndKVA2AForSBHD, AttnFuncWithCPAndKVA2AForTHD |
| 14 | - AttnFuncWithCPAndKVA2AForSBHD, | ||
| 15 | - AttnFuncWithCPAndKVA2AForTHD | ||
| 16 | -) | ||
| 17 | 13 | ||
| 18 | from .ulysses_context_parallel import AttnFuncWithCPAndQKVOA2A | 14 | from .ulysses_context_parallel import AttnFuncWithCPAndQKVOA2A |
| 19 | from .utils import prepare_sbhd_format, prepare_thd_format | 15 | from .utils import prepare_sbhd_format, prepare_thd_format |
| @@ -21,12 +17,14 @@ from .utils import prepare_sbhd_format, prepare_thd_format | |||
| 21 | 17 | ||
| 22 | class BaseCPStrategy(torch.nn.Module, ABC): | 18 | class BaseCPStrategy(torch.nn.Module, ABC): |
| 23 | """Base class for all Context Parallelism strategies""" | 19 | """Base class for all Context Parallelism strategies""" |
| 24 | - | 20 | + |
| 25 | - def __init__(self, | 21 | + def __init__( |
| 26 | - softmax_scale: float, | 22 | + self, |
| 27 | - attention_dropout: float = 0.0, | 23 | + softmax_scale: float, |
| 28 | - attention_type: str = "self", | 24 | + attention_dropout: float = 0.0, |
| 29 | - deterministic: bool = False): | 25 | + attention_type: str = "self", |
| 26 | + deterministic: bool = False, | ||
| 27 | + ): | ||
| 30 | super().__init__() | 28 | super().__init__() |
| 31 | self.softmax_scale = softmax_scale | 29 | self.softmax_scale = softmax_scale |
| 32 | self.attention_dropout = attention_dropout | 30 | self.attention_dropout = attention_dropout |
| @@ -35,21 +33,21 @@ class BaseCPStrategy(torch.nn.Module, ABC): | |||
| 35 | 33 | ||
| 36 | 34 | ||
| 37 | def forward( | 35 | def forward( |
| 38 | - self, | 36 | + self, |
| 39 | - query_layer, | 37 | + query_layer, |
| 40 | - key_layer, | 38 | + key_layer, |
| 41 | - value_layer, | 39 | + value_layer, |
| 42 | - attention_mask, | 40 | + attention_mask, |
| 43 | - qkv_format, | 41 | + qkv_format, |
| 44 | - cu_seqlens_q, | 42 | + cu_seqlens_q, |
| 45 | - cu_seqlens_kv, | 43 | + cu_seqlens_kv, |
| 46 | - attn_mask_type, | 44 | + attn_mask_type, |
| 47 | - max_seqlen_q, | 45 | + max_seqlen_q, |
| 48 | - max_seqlen_kv, | 46 | + max_seqlen_kv, |
| 49 | - cp_group, | 47 | + cp_group, |
| 50 | - cp_global_ranks, | 48 | + cp_global_ranks, |
| 51 | - cp_stream, | 49 | + cp_stream, |
| 52 | - **kwargs | 50 | + **kwargs, |
| 53 | ) -> torch.Tensor: | 51 | ) -> torch.Tensor: |
| 54 | pass | 52 | pass |
| 55 | 53 | ||
| @@ -57,38 +55,39 @@ class BaseCPStrategy(torch.nn.Module, ABC): | |||
| 57 | class UlyssesCPStrategy(BaseCPStrategy): | 55 | class UlyssesCPStrategy(BaseCPStrategy): |
| 58 | """Ulysses-based Context Parallelism Strategy""" | 56 | """Ulysses-based Context Parallelism Strategy""" |
| 59 | 57 | ||
| 60 | - def __init__(self, | 58 | + def __init__( |
| 61 | - softmax_scale: float, | 59 | + self, |
| 62 | - attention_dropout: float = 0.0, | 60 | + softmax_scale: float, |
| 63 | - attention_type: str = "self", | 61 | + attention_dropout: float = 0.0, |
| 64 | - deterministic: bool = False, | 62 | + attention_type: str = "self", |
| 65 | - scatter_idx: int = 2, | 63 | + deterministic: bool = False, |
| 66 | - gather_idx: int = 0, | 64 | + scatter_idx: int = 2, |
| 67 | - gather_size_calculator=None, # Injected dependency | 65 | + gather_idx: int = 0, |
| 66 | + gather_size_calculator=None, # Injected dependency | ||
| 68 | ): | 67 | ): |
| 69 | super().__init__(softmax_scale, attention_dropout, attention_type, deterministic) | 68 | super().__init__(softmax_scale, attention_dropout, attention_type, deterministic) |
| 70 | self.ulysses_comm_para = { | 69 | self.ulysses_comm_para = { |
| 71 | 'scatter_idx': scatter_idx, | 70 | 'scatter_idx': scatter_idx, |
| 72 | 'gather_idx': gather_idx, | 71 | 'gather_idx': gather_idx, |
| 73 | - 'gather_size_calculator': gather_size_calculator | 72 | + 'gather_size_calculator': gather_size_calculator, |
| 74 | } | 73 | } |
| 75 | 74 | ||
| 76 | def forward( | 75 | def forward( |
| 77 | - self, | 76 | + self, |
| 78 | - query_layer, | 77 | + query_layer, |
| 79 | - key_layer, | 78 | + key_layer, |
| 80 | - value_layer, | 79 | + value_layer, |
| 81 | - attention_mask, | 80 | + attention_mask, |
| 82 | - qkv_format, | 81 | + qkv_format, |
| 83 | - cu_seqlens_q, | 82 | + cu_seqlens_q, |
| 84 | - cu_seqlens_kv, | 83 | + cu_seqlens_kv, |
| 85 | - attn_mask_type, | 84 | + attn_mask_type, |
| 86 | - max_seqlen_q, | 85 | + max_seqlen_q, |
| 87 | - max_seqlen_kv, | 86 | + max_seqlen_kv, |
| 88 | - cp_group, | 87 | + cp_group, |
| 89 | - cp_global_ranks, | 88 | + cp_global_ranks, |
| 90 | - cp_stream, | 89 | + cp_stream, |
| 91 | - **kwargs | 90 | + **kwargs, |
| 92 | ): | 91 | ): |
| 93 | return AttnFuncWithCPAndQKVOA2A( | 92 | return AttnFuncWithCPAndQKVOA2A( |
| 94 | query_layer, | 93 | query_layer, |
| @@ -104,7 +103,7 @@ class UlyssesCPStrategy(BaseCPStrategy): | |||
| 104 | self.deterministic, | 103 | self.deterministic, |
| 105 | cp_group, | 104 | cp_group, |
| 106 | cp_stream, | 105 | cp_stream, |
| 107 | - self.ulysses_comm_para | 106 | + self.ulysses_comm_para, |
| 108 | ) | 107 | ) |
| 109 | 108 | ||
| 110 | 109 | ||
| @@ -112,28 +111,25 @@ class KVAllGatherCPStrategy(BaseCPStrategy): | |||
| 112 | """AllGather-based Context Parallelism Strategy""" | 111 | """AllGather-based Context Parallelism Strategy""" |
| 113 | 112 | ||
| 114 | def forward( | 113 | def forward( |
| 115 | - self, | 114 | + self, |
| 116 | - query_layer, | 115 | + query_layer, |
| 117 | - key_layer, | 116 | + key_layer, |
| 118 | - value_layer, | 117 | + value_layer, |
| 119 | - attention_mask, | 118 | + attention_mask, |
| 120 | - qkv_format, | 119 | + qkv_format, |
| 121 | - cu_seqlens_q, | 120 | + cu_seqlens_q, |
| 122 | - cu_seqlens_kv, | 121 | + cu_seqlens_kv, |
| 123 | - attn_mask_type, | 122 | + attn_mask_type, |
| 124 | - max_seqlen_q, | 123 | + max_seqlen_q, |
| 125 | - max_seqlen_kv, | 124 | + max_seqlen_kv, |
| 126 | - cp_group, | 125 | + cp_group, |
| 127 | - cp_global_ranks, | 126 | + cp_global_ranks, |
| 128 | - cp_stream, | 127 | + cp_stream, |
| 129 | - **kwargs | 128 | + **kwargs, |
| 130 | - ): | 129 | + ): |
| 131 | - | ||
| 132 | # Prepare tensors based on format | 130 | # Prepare tensors based on format |
| 133 | if qkv_format == 'sbhd': | 131 | if qkv_format == 'sbhd': |
| 134 | - query_layer, key_layer, value_layer, n_head = prepare_sbhd_format( | 132 | + query_layer, key_layer, value_layer, n_head = prepare_sbhd_format(query_layer, key_layer, value_layer) |
| 135 | - query_layer, key_layer, value_layer | ||
| 136 | - ) | ||
| 137 | 133 | ||
| 138 | return AttnFuncWithCPAndKVAllGatherForSBHD.apply( | 134 | return AttnFuncWithCPAndKVAllGatherForSBHD.apply( |
| 139 | query_layer, | 135 | query_layer, |
| @@ -147,13 +143,11 @@ class KVAllGatherCPStrategy(BaseCPStrategy): | |||
| 147 | self.softmax_scale, | 143 | self.softmax_scale, |
| 148 | self.deterministic, | 144 | self.deterministic, |
| 149 | cp_group, | 145 | cp_group, |
| 150 | - cp_stream | 146 | + cp_stream, |
| 151 | ) | 147 | ) |
| 152 | - | 148 | + |
| 153 | elif qkv_format == 'thd': | 149 | elif qkv_format == 'thd': |
| 154 | - n_head, cu_seqlens_q, cu_seqlens_kv = prepare_thd_format( | 150 | + _, n_head, _ = query_layer.shape |
| 155 | - query_layer, cu_seqlens_q, cu_seqlens_kv | ||
| 156 | - ) | ||
| 157 | 151 | ||
| 158 | return AttnFuncWithCPAndKVAllGatherForTHD.apply( | 152 | return AttnFuncWithCPAndKVAllGatherForTHD.apply( |
| 159 | query_layer, | 153 | query_layer, |
| @@ -168,9 +162,9 @@ class KVAllGatherCPStrategy(BaseCPStrategy): | |||
| 168 | self.deterministic, | 162 | self.deterministic, |
| 169 | cp_group, | 163 | cp_group, |
| 170 | cu_seqlens_q, | 164 | cu_seqlens_q, |
| 171 | - cu_seqlens_kv | 165 | + cu_seqlens_kv, |
| 172 | ) | 166 | ) |
| 173 | - | 167 | + |
| 174 | else: | 168 | else: |
| 175 | raise ValueError(f"Unsupported qkv_format: {qkv_format}") | 169 | raise ValueError(f"Unsupported qkv_format: {qkv_format}") |
| 176 | 170 | ||
| @@ -184,48 +178,48 @@ class HamiltonCPStrategy(BaseCPStrategy): | |||
| 184 | For more details, please refer to `HamiltonAttention <https://arxiv.org/pdf/2509.26541>`_. | 178 | For more details, please refer to `HamiltonAttention <https://arxiv.org/pdf/2509.26541>`_. |
| 185 | """ | 179 | """ |
| 186 | 180 | ||
| 187 | - def __init__(self, | 181 | + def __init__( |
| 188 | - softmax_scale: float, | 182 | + self, |
| 189 | - attention_dropout: float = 0.0, | 183 | + softmax_scale: float, |
| 190 | - attention_type: str = "self", | 184 | + attention_dropout: float = 0.0, |
| 191 | - deterministic: bool = False, | 185 | + attention_type: str = "self", |
| 192 | - path_num: int = 0, | 186 | + deterministic: bool = False, |
| 193 | - out_mapping: Any = None, | 187 | + path_num: int = 0, |
| 194 | - in_mapping: Any = None, | 188 | + out_mapping: Any = None, |
| 195 | - permute_index: Any = None, | 189 | + in_mapping: Any = None, |
| 196 | - restore_index: Any = None | 190 | + permute_index: Any = None, |
| 197 | - ): | 191 | + restore_index: Any = None, |
| 192 | + ): | ||
| 198 | super().__init__(softmax_scale, attention_dropout, attention_type, deterministic) | 193 | super().__init__(softmax_scale, attention_dropout, attention_type, deterministic) |
| 199 | self.ha_comm_para = { | 194 | self.ha_comm_para = { |
| 200 | 'path_num': path_num, | 195 | 'path_num': path_num, |
| 201 | 'out_mapping': out_mapping, | 196 | 'out_mapping': out_mapping, |
| 202 | 'in_mapping': in_mapping, | 197 | 'in_mapping': in_mapping, |
| 203 | 'permute_index': permute_index, | 198 | 'permute_index': permute_index, |
| 204 | - 'restore_index': restore_index | 199 | + 'restore_index': restore_index, |
| 205 | } | 200 | } |
| 206 | 201 | ||
| 207 | def forward( | 202 | def forward( |
| 208 | - self, | 203 | + self, |
| 209 | - query_layer, | 204 | + query_layer, |
| 210 | - key_layer, | 205 | + key_layer, |
| 211 | - value_layer, | 206 | + value_layer, |
| 212 | - attention_mask, | 207 | + attention_mask, |
| 213 | - qkv_format, | 208 | + qkv_format, |
| 214 | - cu_seqlens_q, | 209 | + cu_seqlens_q, |
| 215 | - cu_seqlens_kv, | 210 | + cu_seqlens_kv, |
| 216 | - attn_mask_type, | 211 | + attn_mask_type, |
| 217 | - max_seqlen_q, | 212 | + max_seqlen_q, |
| 218 | - max_seqlen_kv, | 213 | + max_seqlen_kv, |
| 219 | - cp_group, | 214 | + cp_group, |
| 220 | - cp_global_ranks, | 215 | + cp_global_ranks, |
| 221 | - **kwargs | 216 | + cp_stream, |
| 217 | + **kwargs, | ||
| 222 | ): | 218 | ): |
| 223 | # Prepare tensors based on format | 219 | # Prepare tensors based on format |
| 224 | if qkv_format == 'sbhd': | 220 | if qkv_format == 'sbhd': |
| 225 | # transform the qkv format into 3 dimensions | 221 | # transform the qkv format into 3 dimensions |
| 226 | - query_layer, key_layer, value_layer, n_head = prepare_sbhd_format( | 222 | + query_layer, key_layer, value_layer, n_head = prepare_sbhd_format(query_layer, key_layer, value_layer) |
| 227 | - query_layer, key_layer, value_layer | ||
| 228 | - ) | ||
| 229 | 223 | ||
| 230 | return AttnFuncWithCPAndKVA2AForSBHD.apply( | 224 | return AttnFuncWithCPAndKVA2AForSBHD.apply( |
| 231 | query_layer, | 225 | query_layer, |
| @@ -239,13 +233,11 @@ class HamiltonCPStrategy(BaseCPStrategy): | |||
| 239 | self.softmax_scale, | 233 | self.softmax_scale, |
| 240 | self.deterministic, | 234 | self.deterministic, |
| 241 | cp_group, | 235 | cp_group, |
| 242 | - self.ha_comm_para | 236 | + self.ha_comm_para, |
| 243 | ) | 237 | ) |
| 244 | 238 | ||
| 245 | elif qkv_format == 'thd': | 239 | elif qkv_format == 'thd': |
| 246 | - n_head, cu_seqlens_q, cu_seqlens_kv = prepare_thd_format( | 240 | + n_head, cu_seqlens_q, cu_seqlens_kv = prepare_thd_format(query_layer, cu_seqlens_q, cu_seqlens_kv) |
| 247 | - query_layer, cu_seqlens_q, cu_seqlens_kv | ||
| 248 | - ) | ||
| 249 | 241 | ||
| 250 | return AttnFuncWithCPAndKVA2AForTHD.apply( | 242 | return AttnFuncWithCPAndKVA2AForTHD.apply( |
| 251 | query_layer, | 243 | query_layer, |
| @@ -261,7 +253,7 @@ class HamiltonCPStrategy(BaseCPStrategy): | |||
| 261 | cp_group, | 253 | cp_group, |
| 262 | cu_seqlens_q, | 254 | cu_seqlens_q, |
| 263 | cu_seqlens_kv, | 255 | cu_seqlens_kv, |
| 264 | - self.ha_comm_para | 256 | + self.ha_comm_para, |
| 265 | ) | 257 | ) |
| 266 | 258 | ||
| 267 | else: | 259 | else: |
| @@ -270,41 +262,35 @@ class HamiltonCPStrategy(BaseCPStrategy): | |||
| 270 | 262 | ||
| 271 | class CPStrategyFactory: | 263 | class CPStrategyFactory: |
| 272 | """Factory for creating Context Parallelism strategies""" | 264 | """Factory for creating Context Parallelism strategies""" |
| 273 | - | 265 | + |
| 274 | _STRATEGY_MAP = { | 266 | _STRATEGY_MAP = { |
| 275 | "kvallgather_cp_algo": KVAllGatherCPStrategy, | 267 | "kvallgather_cp_algo": KVAllGatherCPStrategy, |
| 276 | "ulysses_cp_algo": UlyssesCPStrategy, | 268 | "ulysses_cp_algo": UlyssesCPStrategy, |
| 277 | - "hamilton_cp_algo": HamiltonCPStrategy | 269 | + "hamilton_cp_algo": HamiltonCPStrategy, |
| 278 | } | 270 | } |
| 279 | - | 271 | + |
| 280 | 272 | ||
| 281 | - def create_strategy( | 273 | + def create_strategy(strategy_type: str, **strategy_params) -> BaseCPStrategy: |
| 282 | - strategy_type: str, | ||
| 283 | - **strategy_params | ||
| 284 | - ) -> BaseCPStrategy: | ||
| 285 | """Create a CP strategy instance | 274 | """Create a CP strategy instance |
| 286 | - | 275 | + |
| 287 | Args: | 276 | Args: |
| 288 | strategy_type: Type of strategy to create | 277 | strategy_type: Type of strategy to create |
| 289 | **strategy_params: Parameters for the strategy constructor | 278 | **strategy_params: Parameters for the strategy constructor |
| 290 | - | 279 | + |
| 291 | Returns: | 280 | Returns: |
| 292 | Instance of the requested strategy | 281 | Instance of the requested strategy |
| 293 | - | 282 | + |
| 294 | Raises: | 283 | Raises: |
| 295 | ValueError: If strategy_type is not supported | 284 | ValueError: If strategy_type is not supported |
| 296 | """ | 285 | """ |
| 297 | if strategy_type not in CPStrategyFactory._STRATEGY_MAP: | 286 | if strategy_type not in CPStrategyFactory._STRATEGY_MAP: |
| 298 | available = CPStrategyFactory.get_available_strategies() | 287 | available = CPStrategyFactory.get_available_strategies() |
| 299 | - raise ValueError( | 288 | + raise ValueError(f"Unsupported CP strategy: {strategy_type}. Available strategies: {available}") |
| 300 | - f"Unsupported CP strategy: {strategy_type}. " | 289 | + |
| 301 | - f"Available strategies: {available}" | ||
| 302 | - ) | ||
| 303 | - | ||
| 304 | strategy_class = CPStrategyFactory._STRATEGY_MAP[strategy_type] | 290 | strategy_class = CPStrategyFactory._STRATEGY_MAP[strategy_type] |
| 305 | return strategy_class(**strategy_params) | 291 | return strategy_class(**strategy_params) |
| 306 | - | 292 | + |
| 307 | 293 | ||
| 308 | def get_available_strategies() -> List[str]: | 294 | def get_available_strategies() -> List[str]: |
| 309 | """Get list of available strategy names""" | 295 | """Get list of available strategy names""" |
| 310 | - return list(CPStrategyFactory._STRATEGY_MAP.keys()) | 296 | + return list(CPStrategyFactory._STRATEGY_MAP.keys()) |
| @@ -1,6 +1,9 @@ | |||
| 1 | # Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. | 1 | # Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. |
| 2 | +# Copyright (c) 2022-2025; NVIDIA CORPORATION AFFILIATES. | ||
| 3 | +# Copyright c) 2022-2025 Advanced Micro Devices Inc. | ||
| 2 | # Copyright (c) 2025, Huawei Technologies Co., Ltd. All rights reserved. | 4 | # Copyright (c) 2025, Huawei Technologies Co., Ltd. All rights reserved. |
| 3 | 5 | ||
| 6 | +from functools import lru_cache | ||
| 4 | from typing import Optional, Tuple | 7 | from typing import Optional, Tuple |
| 5 | 8 | ||
| 6 | import torch | 9 | import torch |
| @@ -10,12 +13,14 @@ from .utils import get_distributed_rank, get_distributed_world_size | |||
| 10 | 13 | ||
| 11 | _seq_chunk_ids_cache_for_reordering_before_attn = {} | 14 | _seq_chunk_ids_cache_for_reordering_before_attn = {} |
| 12 | _seq_chunk_ids_cache_for_reordering_after_attn = {} | 15 | _seq_chunk_ids_cache_for_reordering_after_attn = {} |
| 16 | +_THD_LOAD_BALANCED_CP_METADATA_CACHE_SIZE = 8 | ||
| 17 | +_thd_load_balanced_cp_metadata_cache = {} | ||
| 13 | 18 | ||
| 14 | 19 | ||
| 15 | def gather_along_first_dim( | 20 | def gather_along_first_dim( |
| 16 | - inp: torch.Tensor, | 21 | + inp: torch.Tensor, |
| 17 | - process_group, | 22 | + process_group, |
| 18 | - async_op: bool = False, | 23 | + async_op: bool = False, |
| 19 | ) -> Tuple[torch.Tensor, Optional[torch.distributed.Work]]: | 24 | ) -> Tuple[torch.Tensor, Optional[torch.distributed.Work]]: |
| 20 | """ | 25 | """ |
| 21 | All-gather tensors and concatenate along first dimension. | 26 | All-gather tensors and concatenate along first dimension. |
| @@ -44,9 +49,7 @@ def gather_along_first_dim( | |||
| 44 | 49 | ||
| 45 | 50 | ||
| 46 | def reduce_scatter_along_first_dim( | 51 | def reduce_scatter_along_first_dim( |
| 47 | - inp: torch.Tensor, | 52 | + inp: torch.Tensor, process_group, async_op: bool = False |
| 48 | - process_group, | ||
| 49 | - async_op: bool = False | ||
| 50 | ) -> Tuple[torch.Tensor, Optional[torch.distributed.Work]]: | 53 | ) -> Tuple[torch.Tensor, Optional[torch.distributed.Work]]: |
| 51 | """Reduce-scatter the input tensor across model parallel group.""" | 54 | """Reduce-scatter the input tensor across model parallel group.""" |
| 52 | world_size = get_distributed_world_size(process_group) | 55 | world_size = get_distributed_world_size(process_group) |
| @@ -60,9 +63,7 @@ def reduce_scatter_along_first_dim( | |||
| 60 | dim_size[0] = dim_size[0] // world_size | 63 | dim_size[0] = dim_size[0] // world_size |
| 61 | 64 | ||
| 62 | output = torch.empty(dim_size, dtype=inp.dtype, device=inp.device) | 65 | output = torch.empty(dim_size, dtype=inp.dtype, device=inp.device) |
| 63 | - handle = torch.distributed.reduce_scatter_tensor( | 66 | + handle = torch.distributed.reduce_scatter_tensor(output, inp.contiguous(), group=process_group, async_op=async_op) |
| 64 | - output, inp.contiguous(), group=process_group, async_op=async_op | ||
| 65 | - ) | ||
| 66 | return output, handle | 67 | return output, handle |
| 67 | 68 | ||
| 68 | 69 | ||
| @@ -73,7 +74,6 @@ def get_seq_chunk_ids_for_reordering_before_attn(cp_size, device): | |||
| 73 | be contigupus before attention compute. This function is to compute sequence chunk ids for | 74 | be contigupus before attention compute. This function is to compute sequence chunk ids for |
| 74 | reordering. | 75 | reordering. |
| 75 | """ | 76 | """ |
| 76 | - global _seq_chunk_ids_cache_for_reordering_before_attn | ||
L | |||
| 77 | if (cp_size, device) not in _seq_chunk_ids_cache_for_reordering_before_attn: | 77 | if (cp_size, device) not in _seq_chunk_ids_cache_for_reordering_before_attn: |
| 78 | chunk_ids = torch.empty(2 * cp_size, dtype=torch.int32, device=device) | 78 | chunk_ids = torch.empty(2 * cp_size, dtype=torch.int32, device=device) |
| 79 | for rank in range(cp_size): | 79 | for rank in range(cp_size): |
| @@ -89,7 +89,6 @@ def get_seq_chunk_ids_for_reordering_after_attn(cp_size, device): | |||
| 89 | We need to reorder sequence chunks back to discontiguous after attention compute. This function | 89 | We need to reorder sequence chunks back to discontiguous after attention compute. This function |
| 90 | is to compute sequence chunk ids for reordering. | 90 | is to compute sequence chunk ids for reordering. |
| 91 | """ | 91 | """ |
| 92 | - global _seq_chunk_ids_cache_for_reordering_after_attn | ||
| 93 | if (cp_size, device) not in _seq_chunk_ids_cache_for_reordering_after_attn: | 92 | if (cp_size, device) not in _seq_chunk_ids_cache_for_reordering_after_attn: |
| 94 | chunk_ids = torch.empty(2 * cp_size, dtype=torch.int32, device=device) | 93 | chunk_ids = torch.empty(2 * cp_size, dtype=torch.int32, device=device) |
| 95 | for rank in range(cp_size): | 94 | for rank in range(cp_size): |
| @@ -99,9 +98,7 @@ def get_seq_chunk_ids_for_reordering_after_attn(cp_size, device): | |||
| 99 | return _seq_chunk_ids_cache_for_reordering_after_attn[(cp_size, device)] | 98 | return _seq_chunk_ids_cache_for_reordering_after_attn[(cp_size, device)] |
| 100 | 99 | ||
| 101 | 100 | ||
| 102 | -def get_kv_seq_info_after_all_gather( | 101 | +def get_kv_seq_info_after_all_gather(local_chunk_id, cp_size, max_seqlen_kv, causal): |
| 103 | - local_chunk_id, cp_size, max_seqlen_kv, causal | ||
| 104 | -): | ||
| 105 | """Compute KV sequence index range and update window size after all-gather.""" | 102 | """Compute KV sequence index range and update window size after all-gather.""" |
| 106 | local_chunk_end_idx = (local_chunk_id + 1) * max_seqlen_kv | 103 | local_chunk_end_idx = (local_chunk_id + 1) * max_seqlen_kv |
| 107 | full_seq_end_idx = max_seqlen_kv * cp_size * 2 | 104 | full_seq_end_idx = max_seqlen_kv * cp_size * 2 |
| @@ -124,19 +121,19 @@ class AttnFuncWithCPAndKVAllGatherForSBHD(torch.autograd.Function): | |||
| 124 | 121 | ||
| 125 | 122 | ||
| 126 | def forward( | 123 | def forward( |
| 127 | - ctx, | 124 | + ctx, |
| 128 | - q, | 125 | + q, |
| 129 | - k, | 126 | + k, |
| 130 | - v, | 127 | + v, |
| 131 | - n_head, | 128 | + n_head, |
| 132 | - attention_mask, | 129 | + attention_mask, |
| 133 | - qkv_format, | 130 | + qkv_format, |
| 134 | - attn_mask_type, | 131 | + attn_mask_type, |
| 135 | - attention_dropout, | 132 | + attention_dropout, |
| 136 | - softmax_scale, | 133 | + softmax_scale, |
| 137 | - deterministic, | 134 | + deterministic, |
| 138 | - cp_group, | 135 | + cp_group, |
| 139 | - cp_stream | 136 | + cp_stream, |
| 140 | ): | 137 | ): |
| 141 | if softmax_scale is None: | 138 | if softmax_scale is None: |
| 142 | softmax_scale = q.shape[-1] ** (-0.5) | 139 | softmax_scale = q.shape[-1] ** (-0.5) |
| @@ -153,7 +150,7 @@ class AttnFuncWithCPAndKVAllGatherForSBHD(torch.autograd.Function): | |||
| 153 | raise AssertionError("Sequence length per GPU needs to be divisible by 2!") | 150 | raise AssertionError("Sequence length per GPU needs to be divisible by 2!") |
| 154 | 151 | ||
| 155 | # [s, b, h] -> [2, s//2, b, h] | 152 | # [s, b, h] -> [2, s//2, b, h] |
| 156 | - q = q.view(2, q.shape[seq_dim] // 2, *q.shape[(seq_dim + 1):]) | 153 | + q = q.view(2, q.shape[seq_dim] // 2, *q.shape[(seq_dim + 1) :]) |
| 157 | 154 | ||
| 158 | # [s, b, h] -> [cp, s, b, h] | 155 | # [s, b, h] -> [cp, s, b, h] |
| 159 | k_ag, _ = gather_along_first_dim(k, cp_group) | 156 | k_ag, _ = gather_along_first_dim(k, cp_group) |
| @@ -190,13 +187,11 @@ class AttnFuncWithCPAndKVAllGatherForSBHD(torch.autograd.Function): | |||
| 190 | with torch.npu.stream(flash_attn_streams[i]): | 187 | with torch.npu.stream(flash_attn_streams[i]): |
| 191 | # [2, sq//2, b, h] -> [sq//2, b, h] | 188 | # [2, sq//2, b, h] -> [sq//2, b, h] |
| 192 | q_ = q.select(seq_dim, i).contiguous() | 189 | q_ = q.select(seq_dim, i).contiguous() |
| 193 | - kv_seq_range_per_step[i] = ( | 190 | + kv_seq_range_per_step[i] = get_kv_seq_info_after_all_gather( |
| 194 | - get_kv_seq_info_after_all_gather( | 191 | + local_seq_chunk_ids[i], |
| 195 | - local_seq_chunk_ids[i], | 192 | + cp_size, |
| 196 | - cp_size, | 193 | + kv_len, |
| 197 | - kv_len, | 194 | + causal, |
| 198 | - causal, | ||
| 199 | - ) | ||
| 200 | ) | 195 | ) |
| 201 | seq_start_idx, seq_end_idx = ( | 196 | seq_start_idx, seq_end_idx = ( |
| 202 | kv_seq_range_per_step[i][0], | 197 | kv_seq_range_per_step[i][0], |
| @@ -219,7 +214,7 @@ class AttnFuncWithCPAndKVAllGatherForSBHD(torch.autograd.Function): | |||
| 219 | next_tockens=0, | 214 | next_tockens=0, |
| 220 | keep_prob=1 - attention_dropout, | 215 | keep_prob=1 - attention_dropout, |
| 221 | inner_precise=0, | 216 | inner_precise=0, |
| 222 | - sparse_mode=3 | 217 | + sparse_mode=3, |
| 223 | ) | 218 | ) |
| 224 | out_per_step[i] = attn_outs[0] | 219 | out_per_step[i] = attn_outs[0] |
| 225 | softmax_max[i] = attn_outs[1] | 220 | softmax_max[i] = attn_outs[1] |
| @@ -236,14 +231,7 @@ class AttnFuncWithCPAndKVAllGatherForSBHD(torch.autograd.Function): | |||
| 236 | # [2, s//2, b, h] -> [s, b, h] | 231 | # [2, s//2, b, h] -> [s, b, h] |
| 237 | q = q.view(-1, *q.shape[2:]) | 232 | q = q.view(-1, *q.shape[2:]) |
| 238 | 233 | ||
| 239 | - ctx.save_for_backward( | 234 | + ctx.save_for_backward(q, k, v, *out_per_step, *softmax_max, *softmax_sum) |
| 240 | - q, | ||
| 241 | - k, | ||
| 242 | - v, | ||
| 243 | - *out_per_step, | ||
| 244 | - *softmax_max, | ||
| 245 | - *softmax_sum | ||
| 246 | - ) | ||
| 247 | 235 | ||
| 248 | ctx.kv_seq_range_per_step = kv_seq_range_per_step | 236 | ctx.kv_seq_range_per_step = kv_seq_range_per_step |
| 249 | ctx.cp_group = cp_group | 237 | ctx.cp_group = cp_group |
| @@ -270,7 +258,7 @@ class AttnFuncWithCPAndKVAllGatherForSBHD(torch.autograd.Function): | |||
| 270 | seq_dim = ctx.qkv_format.index("s") | 258 | seq_dim = ctx.qkv_format.index("s") |
| 271 | 259 | ||
| 272 | # [s, b, h] -> [2, s//2, b, h] | 260 | # [s, b, h] -> [2, s//2, b, h] |
| 273 | - q = q.view(2, q.shape[seq_dim] // 2, *q.shape[(seq_dim + 1):]) | 261 | + q = q.view(2, q.shape[seq_dim] // 2, *q.shape[(seq_dim + 1) :]) |
| 274 | # [s, b, h_v] -> [2, s//2, b, h_v] | 262 | # [s, b, h_v] -> [2, s//2, b, h_v] |
| 275 | if k.shape[-1] == v.shape[-1]: | 263 | if k.shape[-1] == v.shape[-1]: |
| 276 | dout = dout.view(q.shape) | 264 | dout = dout.view(q.shape) |
| @@ -340,7 +328,7 @@ class AttnFuncWithCPAndKVAllGatherForSBHD(torch.autograd.Function): | |||
| 340 | pre_tockens=65536, | 328 | pre_tockens=65536, |
| 341 | next_tockens=0, | 329 | next_tockens=0, |
| 342 | sparse_mode=3, | 330 | sparse_mode=3, |
| 343 | - keep_prob=1 - ctx.attention_dropout | 331 | + keep_prob=1 - ctx.attention_dropout, |
| 344 | ) | 332 | ) |
| 345 | dq_per_step[i] = attn_grad_outs[0] | 333 | dq_per_step[i] = attn_grad_outs[0] |
| 346 | dk_per_step[i] = attn_grad_outs[1] | 334 | dk_per_step[i] = attn_grad_outs[1] |
| @@ -399,83 +387,151 @@ class AttnFuncWithCPAndKVAllGatherForSBHD(torch.autograd.Function): | |||
| 399 | ) | 387 | ) |
| 400 | 388 | ||
| 401 | 389 | ||
| 402 | -def get_cu_seqlens_qkv_before_attn(cu_seqlens_q, cp_size, rank): | 390 | +def get_thd_load_balanced_cp_metadata(cu_seqlens, cp_size, rank, device): |
| 403 | - """ | 391 | + """Build token indices for EOD load-balanced KVAllGather THD attention. |
| 404 | - get splited cu_seqlens_q, cu_seqlens_kv, and kv_seq_range for local rank before attention | 392 | + |
| 393 | + Each packed subsequence is divided into ``2 * cp_size`` chunks. Rank ``r`` owns | ||
| 394 | + chunk ``r`` and chunk ``2 * cp_size - r - 1``. K/V all-gather produces a | ||
| 395 | + rank-major tensor, while TND attention requires each logical sequence to be | ||
| 396 | + contiguous, so this helper builds direct rank-major indices for the two | ||
| 397 | + sequence-major causal KV prefixes used by the local Q chunks. | ||
| 405 | """ | 398 | """ |
| 399 | + if cp_size < 1: | ||
| 400 | + raise AssertionError("context_parallel_size must be greater than zero.") | ||
| 401 | + if rank < 0 or rank >= cp_size: | ||
| 402 | + raise AssertionError("context parallel rank must be in [0, context_parallel_size).") | ||
| 406 | 403 | ||
| 407 | - total_length = cu_seqlens_q[-1] | 404 | + cu_seqlens = tuple(int(seq_len) for seq_len in cu_seqlens) |
| 408 | - chunk_size = total_length // cp_size | 405 | + if not cu_seqlens: |
| 409 | - start = rank * chunk_size | 406 | + raise AssertionError("cu_seqlens must contain at least one sequence endpoint.") |
| 410 | - end = start + chunk_size | ||
| 411 | 407 | ||
| 412 | - rank_cu_seqlens_q = [] | 408 | + cache_key = (cp_size, rank, device) |
| 413 | - rank_cu_seqlens_kv = [] | 409 | + cached_metadata = _thd_load_balanced_cp_metadata_cache.get(cache_key) |
| 414 | - cur_len_q = 0 | 410 | + if cached_metadata is not None and cached_metadata["cu_seqlens"] == cu_seqlens: |
| 415 | - cur_len_kv = 0 | 411 | + return cached_metadata["metadata"] |
| 416 | 412 | ||
| 417 | - kv_seq_range = [start, end] | 413 | + seq_starts = (0,) + cu_seqlens[:-1] |
| 414 | + seq_lens = [end - start for start, end in zip(seq_starts, cu_seqlens)] | ||
| 415 | + if any(seq_len <= 0 for seq_len in seq_lens): | ||
| 416 | + raise AssertionError("Each packed subsequence must have a positive sequence length.") | ||
| 417 | + if any(seq_len % (2 * cp_size) != 0 for seq_len in seq_lens): | ||
| 418 | + raise AssertionError("Each packed subsequence length must be divisible by 2 * context_parallel_size.") | ||
| 418 | 419 | ||
| 419 | - for i in range(len(cu_seqlens_q)): | 420 | + chunk_lens = [seq_len // (2 * cp_size) for seq_len in seq_lens] |
| 420 | - if i == 0: | 421 | + local_seq_starts = [] |
| 421 | - seq_start = 0 | 422 | + local_total_len = 0 |
| 422 | - else: | 423 | + for chunk_len in chunk_lens: |
| 423 | - seq_start = cu_seqlens_q[i - 1] | 424 | + local_seq_starts.append(local_total_len) |
| 424 | - seq_end = cu_seqlens_q[i] | 425 | + local_total_len += 2 * chunk_len |
| 425 | 426 | ||
| 426 | - # skip non-overlap sequences | 427 | + full_total_len = cu_seqlens[-1] |
| 427 | - if seq_end <= start: | 428 | + if local_total_len * cp_size != full_total_len: |
| 428 | - continue | 429 | + raise AssertionError("Packed sequence length is inconsistent with context parallel size.") |
| 429 | - if seq_start >= end: | ||
| 430 | - break | ||
| 431 | 430 | ||
| 432 | - # calculate overlapped seqlens_q | 431 | + kv_index_in_rank_major = [] |
| 433 | - overlap_start_q = max(seq_start, start) | 432 | + actual_seq_qlen = [] |
| 434 | - overlap_end_q = min(seq_end, end) | 433 | + actual_seq_kvlen = [] |
| 435 | - overlap_len_q = overlap_end_q - overlap_start_q | 434 | + local_chunk_ids = (rank, 2 * cp_size - rank - 1) |
| 435 | + q_end = 0 | ||
| 436 | + kv_end = 0 | ||
| 436 | 437 | ||
| 437 | - if overlap_len_q > 0: | 438 | + for local_seq_start, chunk_len in zip(local_seq_starts, chunk_lens): |
| 438 | - cur_len_q += overlap_len_q | 439 | + for chunk_id in local_chunk_ids: |
| 439 | - rank_cu_seqlens_q.append(cur_len_q) | 440 | + kv_prefix_len = (chunk_id + 1) * chunk_len |
| 441 | + for prefix_chunk_id in range(chunk_id + 1): | ||
| 442 | + if prefix_chunk_id < cp_size: | ||
| 443 | + source_rank = prefix_chunk_id | ||
| 444 | + source_slot = 0 | ||
| 445 | + else: | ||
| 446 | + source_rank = 2 * cp_size - prefix_chunk_id - 1 | ||
| 447 | + source_slot = 1 | ||
| 448 | + source_start = source_rank * local_total_len + local_seq_start + source_slot * chunk_len | ||
| 449 | + kv_index_in_rank_major.extend(range(source_start, source_start + chunk_len)) | ||
| 440 | 450 | ||
| 441 | - # calculate overlapped seqlens_kv | 451 | + q_end += chunk_len |
| 442 | - overlap_start_kv = seq_start | 452 | + kv_end += kv_prefix_len |
| 443 | - overlap_end_kv = min(seq_end, end) | 453 | + actual_seq_qlen.append(q_end) |
| 444 | - overlap_len_kv = overlap_end_kv - overlap_start_kv | 454 | + actual_seq_kvlen.append(kv_end) |
| 445 | 455 | ||
| 446 | - if overlap_len_kv > 0: | 456 | + kv_index_in_rank_major = torch.tensor(kv_index_in_rank_major, dtype=torch.long, device=device) |
| 447 | - cur_len_kv += overlap_len_kv | ||
| 448 | - rank_cu_seqlens_kv.append(cur_len_kv) | ||
| 449 | 457 | ||
| 450 | - # get kv seq range after all gather | 458 | + metadata = { |
| 451 | - if seq_start < start: | 459 | + "cu_seqlens": cu_seqlens, |
| 452 | - kv_seq_range[0] = seq_start | 460 | + "local_total_len": local_total_len, |
| 461 | + "full_total_len": full_total_len, | ||
| 462 | + "kv_index_in_rank_major": kv_index_in_rank_major, | ||
| 463 | + "actual_seq_qlen": actual_seq_qlen, | ||
| 464 | + "actual_seq_kvlen": actual_seq_kvlen, | ||
| 465 | + } | ||
| 466 | + _thd_load_balanced_cp_metadata_cache[cache_key] = { | ||
| 467 | + "cu_seqlens": cu_seqlens, | ||
| 468 | + "metadata": metadata, | ||
| 469 | + } | ||
| 470 | + return metadata | ||
| 453 | 471 | ||
| 454 | - return rank_cu_seqlens_q, rank_cu_seqlens_kv, kv_seq_range | 472 | + |
| 473 | + | ||
| 474 | +def _get_thd_load_balanced_cp_metadata_for_tensor(cu_seqlens, _tensor_version, cp_size, rank, device): | ||
| 475 | + """Reuse metadata across layers that share the same cu_seqlens tensor.""" | ||
| 476 | + return get_thd_load_balanced_cp_metadata(cu_seqlens.tolist(), cp_size, rank, device) | ||
| 477 | + | ||
| 478 | + | ||
| 479 | +def clear_thd_load_balanced_cp_metadata_cache(): | ||
| 480 | + """Clear tensor and metadata caches, primarily for tests.""" | ||
| 481 | + _get_thd_load_balanced_cp_metadata_for_tensor.cache_clear() | ||
| 482 | + _thd_load_balanced_cp_metadata_cache.clear() | ||
| 483 | + | ||
| 484 | + | ||
| 485 | +def _get_thd_load_balanced_cp_metadata_cached(cu_seqlens, cp_size, rank, device): | ||
| 486 | + """Avoid repeated cu_seqlens conversion for Transformer layers in one microbatch.""" | ||
| 487 | + if isinstance(cu_seqlens, torch.Tensor): | ||
| 488 | + try: | ||
| 489 | + tensor_version = cu_seqlens._version | ||
| 490 | + except RuntimeError: | ||
| 491 | + tensor_version = None | ||
| 492 | + if tensor_version is not None: | ||
| 493 | + return _get_thd_load_balanced_cp_metadata_for_tensor(cu_seqlens, tensor_version, cp_size, rank, device) | ||
| 494 | + cu_seqlens = cu_seqlens.tolist() | ||
| 495 | + return get_thd_load_balanced_cp_metadata(cu_seqlens, cp_size, rank, device) | ||
| 496 | + | ||
| 497 | + | ||
| 498 | +def _validate_thd_attention_shapes(q, k, v, n_head): | ||
| 499 | + """Validate TND shapes for MHA, GQA and expanded MLA attention.""" | ||
| 500 | + if q.ndim != 3 or k.ndim != 3 or v.ndim != 3: | ||
| 501 | + raise AssertionError("TND Q, K and V must all be three-dimensional tensors.") | ||
| 502 | + if q.dtype != k.dtype or q.dtype != v.dtype: | ||
| 503 | + raise AssertionError("Q, K and V must have the same dtype.") | ||
| 504 | + if q.shape[1] != n_head: | ||
| 505 | + raise AssertionError("n_head must match the number of query heads for TND attention.") | ||
| 506 | + if k.shape[1] != v.shape[1]: | ||
| 507 | + raise AssertionError("K and V must have the same number of heads.") | ||
| 508 | + if q.shape[1] == 0 or k.shape[1] == 0 or q.shape[1] % k.shape[1] != 0: | ||
| 509 | + raise AssertionError("The number of query heads must be divisible by the number of KV heads.") | ||
| 510 | + if q.shape[-1] != k.shape[-1]: | ||
| 511 | + raise AssertionError("Q and K must have the same head dimension.") | ||
| 512 | + if k.shape[-1] < v.shape[-1]: | ||
| 513 | + raise AssertionError("The K head dimension must be greater than or equal to the V head dimension.") | ||
| 455 | 514 | ||
| 456 | 515 | ||
| 457 | class AttnFuncWithCPAndKVAllGatherForTHD(torch.autograd.Function): | 516 | class AttnFuncWithCPAndKVAllGatherForTHD(torch.autograd.Function): |
| 458 | - """ | 517 | + """KVAllGather THD attention with EOD-aware causal load balancing.""" |
| 459 | - Attention implementation with context parallelism. KV all-gather between CP ranks is exposed. | ||
| 460 | - For THD format (TND shape_order) | ||
| 461 | - """ | ||
| 462 | 518 | ||
| 463 | 519 | ||
| 464 | def forward( | 520 | def forward( |
| 465 | - ctx, | 521 | + ctx, |
| 466 | - q, | 522 | + q, |
| 467 | - k, | 523 | + k, |
| 468 | - v, | 524 | + v, |
| 469 | - n_head, | 525 | + n_head, |
| 470 | - attention_mask, | 526 | + attention_mask, |
| 471 | - qkv_format, | 527 | + qkv_format, |
| 472 | - attn_mask_type, | 528 | + attn_mask_type, |
| 473 | - attention_dropout, | 529 | + attention_dropout, |
| 474 | - softmax_scale, | 530 | + softmax_scale, |
| 475 | - deterministic, | 531 | + deterministic, |
| 476 | - cp_group, | 532 | + cp_group, |
| 477 | - cu_seqlens_q, | 533 | + cu_seqlens_q, |
| 478 | - cu_seqlens_kv | 534 | + cu_seqlens_kv, |
| 479 | ): | 535 | ): |
| 480 | if softmax_scale is None: | 536 | if softmax_scale is None: |
| 481 | softmax_scale = q.shape[-1] ** (-0.5) | 537 | softmax_scale = q.shape[-1] ** (-0.5) |
| @@ -483,37 +539,30 @@ class AttnFuncWithCPAndKVAllGatherForTHD(torch.autograd.Function): | |||
| 483 | cp_size = get_distributed_world_size(cp_group) | 539 | cp_size = get_distributed_world_size(cp_group) |
| 484 | rank = get_distributed_rank(cp_group) | 540 | rank = get_distributed_rank(cp_group) |
| 485 | 541 | ||
| 486 | - causal = 'causal' in attn_mask_type | 542 | + if 'causal' not in attn_mask_type: |
| 487 | - if not causal: | ||
| 488 | raise AssertionError("Only causal mask is supported for AllGatherContextParallel.") | 543 | raise AssertionError("Only causal mask is supported for AllGatherContextParallel.") |
| 544 | + _validate_thd_attention_shapes(q, k, v, n_head) | ||
| 545 | + if q.shape[0] != k.shape[0] or q.shape[0] != v.shape[0]: | ||
| 546 | + raise AssertionError("Q, K and V must have the same local token count for THD self-attention.") | ||
| 547 | + | ||
| 548 | + metadata = _get_thd_load_balanced_cp_metadata_cached(cu_seqlens_q, cp_size, rank, q.device) | ||
| 549 | + if cu_seqlens_q is not cu_seqlens_kv: | ||
| 550 | + kv_metadata = _get_thd_load_balanced_cp_metadata_cached(cu_seqlens_kv, cp_size, rank, q.device) | ||
| 551 | + if kv_metadata["cu_seqlens"] != metadata["cu_seqlens"]: | ||
| 552 | + raise AssertionError("cu_seqlens_q and cu_seqlens_kv must be the same for THD format.") | ||
| 553 | + if q.shape[0] != metadata["local_total_len"]: | ||
| 554 | + raise AssertionError("The local THD token count does not match the EOD load-balanced CP partition.") | ||
| 489 | 555 | ||
| 490 | - if cu_seqlens_q != cu_seqlens_kv: | ||
| 491 | - raise AssertionError("cu_seqlens_q and cu_seqlens_kv must be the same for THD format.") | ||
| 492 | - | ||
| 493 | - # [t, n, d] -> [cp, t, n, d] | ||
| 494 | k_ag, _ = gather_along_first_dim(k, cp_group) | 556 | k_ag, _ = gather_along_first_dim(k, cp_group) |
| 495 | v_ag, _ = gather_along_first_dim(v, cp_group) | 557 | v_ag, _ = gather_along_first_dim(v, cp_group) |
| 496 | - # [cp, t, n, d] -> [cp*t, n, d] | 558 | + kv_index = metadata["kv_index_in_rank_major"] |
| 497 | - k_ag = k_ag.view(-1, *k.shape[1:]) | 559 | + k_attn = k_ag.index_select(0, kv_index) |
| 498 | - v_ag = v_ag.view(-1, *v.shape[1:]) | 560 | + v_attn = v_ag.index_select(0, kv_index) |
| 499 | - | ||
| 500 | - # get cu_seqlens_q and cu_seqlens_kv | ||
| 501 | - actual_seq_qlen, actual_seq_kvlen, kv_seq_range = get_cu_seqlens_qkv_before_attn( | ||
| 502 | - cu_seqlens_q, | ||
| 503 | - cp_size, | ||
| 504 | - rank | ||
| 505 | - ) | ||
| 506 | - seq_start_idx, seq_end_idx = ( | ||
| 507 | - kv_seq_range[0], | ||
| 508 | - kv_seq_range[1], | ||
| 509 | - ) | ||
| 510 | - # [cp*s, n, d] -> [s_range, n, d] | ||
| 511 | - k_, v_ = [x[seq_start_idx:seq_end_idx] for x in [k_ag, v_ag]] | ||
| 512 | 561 | ||
| 513 | attn_outs = torch_npu.npu_fusion_attention( | 562 | attn_outs = torch_npu.npu_fusion_attention( |
| 514 | q, | 563 | q, |
| 515 | - k_, | 564 | + k_attn, |
| 516 | - v_, | 565 | + v_attn, |
| 517 | n_head, | 566 | n_head, |
| 518 | 'TND', | 567 | 'TND', |
| 519 | pse=None, | 568 | pse=None, |
| @@ -525,64 +574,41 @@ class AttnFuncWithCPAndKVAllGatherForTHD(torch.autograd.Function): | |||
| 525 | keep_prob=1 - attention_dropout, | 574 | keep_prob=1 - attention_dropout, |
| 526 | inner_precise=0, | 575 | inner_precise=0, |
| 527 | sparse_mode=3, | 576 | sparse_mode=3, |
| 528 | - actual_seq_qlen=actual_seq_qlen, | 577 | + actual_seq_qlen=metadata["actual_seq_qlen"], |
| 529 | - actual_seq_kvlen=actual_seq_kvlen | 578 | + actual_seq_kvlen=metadata["actual_seq_kvlen"], |
| 530 | ) | 579 | ) |
| 531 | - out = attn_outs[0] | 580 | + out, softmax_max, softmax_sum = attn_outs[:3] |
| 532 | - softmax_max = attn_outs[1] | ||
| 533 | - softmax_sum = attn_outs[2] | ||
| 534 | 581 | ||
| 535 | - ctx.save_for_backward( | 582 | + ctx.save_for_backward(q, k, v, out, softmax_max, softmax_sum) |
| 536 | - q, | ||
| 537 | - k, | ||
| 538 | - v, | ||
| 539 | - out, | ||
| 540 | - softmax_max, | ||
| 541 | - softmax_sum | ||
| 542 | - ) | ||
| 543 | - | ||
| 544 | - ctx.kv_seq_range = kv_seq_range | ||
| 545 | - ctx.actual_seq_qlen = actual_seq_qlen | ||
| 546 | - ctx.actual_seq_kvlen = actual_seq_kvlen | ||
| 547 | ctx.cp_group = cp_group | 583 | ctx.cp_group = cp_group |
| 548 | ctx.n_head = n_head | 584 | ctx.n_head = n_head |
| 549 | ctx.attention_dropout = attention_dropout | 585 | ctx.attention_dropout = attention_dropout |
| 550 | ctx.softmax_scale = softmax_scale | 586 | ctx.softmax_scale = softmax_scale |
| 551 | ctx.attention_mask = attention_mask | 587 | ctx.attention_mask = attention_mask |
| 588 | + ctx.kv_index_in_rank_major = metadata["kv_index_in_rank_major"] | ||
| 589 | + ctx.actual_seq_qlen = metadata["actual_seq_qlen"] | ||
| 590 | + ctx.actual_seq_kvlen = metadata["actual_seq_kvlen"] | ||
| 552 | return out | 591 | return out |
| 553 | 592 | ||
| 554 | 593 | ||
| 555 | def backward(ctx, dout): | 594 | def backward(ctx, dout): |
| 556 | cp_size = get_distributed_world_size(ctx.cp_group) | 595 | cp_size = get_distributed_world_size(ctx.cp_group) |
| 596 | + saved_tensors = ctx.saved_tensors | ||
| 597 | + q, k, v, out, softmax_max, softmax_sum = saved_tensors | ||
| 557 | 598 | ||
| 558 | - q, k, v, out, softmax_max, softmax_sum = ctx.saved_tensors | 599 | + k_ag, _ = gather_along_first_dim(k, ctx.cp_group) |
| 559 | - kv_seq_range = ctx.kv_seq_range | 600 | + v_ag, _ = gather_along_first_dim(v, ctx.cp_group) |
| 560 | - actual_seq_qlen = ctx.actual_seq_qlen | 601 | + kv_index = ctx.kv_index_in_rank_major |
| 561 | - actual_seq_kvlen = ctx.actual_seq_kvlen | 602 | + k_attn = k_ag.index_select(0, kv_index) |
| 603 | + v_attn = v_ag.index_select(0, kv_index) | ||
| 562 | 604 | ||
| 563 | - # [t, n, d] | ||
| 564 | - dq = torch.empty_like(q) | ||
| 565 | - # [cp*t, n, d] | ||
| 566 | dk = torch.zeros((k.shape[0] * cp_size, *k.shape[1:]), dtype=k.dtype, device=k.device) | 605 | dk = torch.zeros((k.shape[0] * cp_size, *k.shape[1:]), dtype=k.dtype, device=k.device) |
| 567 | dv = torch.zeros((v.shape[0] * cp_size, *v.shape[1:]), dtype=v.dtype, device=v.device) | 606 | dv = torch.zeros((v.shape[0] * cp_size, *v.shape[1:]), dtype=v.dtype, device=v.device) |
| 568 | 607 | ||
| 569 | - # [t, n, d] -> [cp, t, n, d] | ||
| 570 | - k_ag, _ = gather_along_first_dim(k, ctx.cp_group) | ||
| 571 | - v_ag, _ = gather_along_first_dim(v, ctx.cp_group) | ||
| 572 | - # [cp, t, n, d] -> [cp*t, n, d] | ||
| 573 | - k_ag = k_ag.view(-1, *k.shape[1:]) | ||
| 574 | - v_ag = v_ag.view(-1, *v.shape[1:]) | ||
| 575 | - | ||
| 576 | - seq_start_idx, seq_end_idx = ( | ||
| 577 | - kv_seq_range[0], | ||
| 578 | - kv_seq_range[1], | ||
| 579 | - ) | ||
| 580 | - k_, v_ = [x[seq_start_idx:seq_end_idx] for x in [k_ag, v_ag]] | ||
| 581 | - | ||
| 582 | attn_grad_outs = torch_npu.npu_fusion_attention_grad( | 608 | attn_grad_outs = torch_npu.npu_fusion_attention_grad( |
| 583 | q, | 609 | q, |
| 584 | - k_, | 610 | + k_attn, |
| 585 | - v_, | 611 | + v_attn, |
| 586 | dout, | 612 | dout, |
| 587 | ctx.n_head, | 613 | ctx.n_head, |
| 588 | "TND", | 614 | "TND", |
| @@ -597,14 +623,13 @@ class AttnFuncWithCPAndKVAllGatherForTHD(torch.autograd.Function): | |||
| 597 | next_tockens=0, | 623 | next_tockens=0, |
| 598 | sparse_mode=3, | 624 | sparse_mode=3, |
| 599 | keep_prob=1 - ctx.attention_dropout, | 625 | keep_prob=1 - ctx.attention_dropout, |
| 600 | - actual_seq_qlen=actual_seq_qlen, | 626 | + actual_seq_qlen=ctx.actual_seq_qlen, |
| 601 | - actual_seq_kvlen=actual_seq_kvlen | 627 | + actual_seq_kvlen=ctx.actual_seq_kvlen, |
| 602 | ) | 628 | ) |
| 603 | - dq.copy_(attn_grad_outs[0]) | 629 | + dq = attn_grad_outs[0] |
| 604 | - dk[seq_start_idx:seq_end_idx].copy_(attn_grad_outs[1]) | 630 | + dk.index_add_(0, kv_index, attn_grad_outs[1]) |
| 605 | - dv[seq_start_idx:seq_end_idx].copy_(attn_grad_outs[2]) | 631 | + dv.index_add_(0, kv_index, attn_grad_outs[2]) |
| 606 | 632 | ||
| 607 | - # communication for dk and dv | ||
| 608 | dk, _ = reduce_scatter_along_first_dim(dk, ctx.cp_group) | 633 | dk, _ = reduce_scatter_along_first_dim(dk, ctx.cp_group) |
| 609 | dv, _ = reduce_scatter_along_first_dim(dv, ctx.cp_group) | 634 | dv, _ = reduce_scatter_along_first_dim(dv, ctx.cp_group) |
| 610 | 635 | ||
| @@ -622,7 +647,4 @@ class AttnFuncWithCPAndKVAllGatherForTHD(torch.autograd.Function): | |||
| 622 | None, | 647 | None, |
| 623 | None, | 648 | None, |
| 624 | None, | 649 | None, |
| 625 | - None, | 650 | + ) |
| 626 | - None, | ||
| 627 | - None, | ||
| 628 | - ) | ||
| @@ -0,0 +1,290 @@ | |||
| 1 | +# Copyright (c) 2026, Huawei Technologies Co., Ltd. All rights reserved. | ||
| 2 | + | ||
| 3 | +import math | ||
L  m0_50947149 13 天前 评论： 13 天前 评论： | |||
| 4 | + | ||
| 5 | +import pytest | ||
| 6 | +import torch | ||
| 7 | +import torch.distributed as dist | ||
| 8 | +import torch_npu | ||
| 9 | + | ||
| 10 | +from megatron.core import mpu | ||
| 11 | + | ||
| 12 | +from mindspeed.core.context_parallel import get_batch_utils | ||
| 13 | +from mindspeed.te.pytorch.attention.dot_product_attention import kvallgather_context_parallel as kvallgather_cp | ||
| 14 | +from mindspeed.te.pytorch.attention.dot_product_attention.kvallgather_context_parallel import ( | ||
| 15 | + AttnFuncWithCPAndKVAllGatherForTHD, | ||
| 16 | + clear_thd_load_balanced_cp_metadata_cache, | ||
| 17 | + get_thd_load_balanced_cp_metadata, | ||
| 18 | +) | ||
| 19 | +from tests_extend.commons import initialize_model_parallel, set_random_seed | ||
| 20 | +from tests_extend.unit_tests.common import DistributedTest | ||
| 21 | + | ||
| 22 | + | ||
| 23 | +def _get_local_token_index(cu_seqlens, cp_size, rank, device): | ||
| 24 | + ranges = [] | ||
| 25 | + seq_start = 0 | ||
| 26 | + for seq_end in cu_seqlens: | ||
| 27 | + chunk_len = (seq_end - seq_start) // (2 * cp_size) | ||
| 28 | + ranges.append( | ||
| 29 | + torch.arange( | ||
| 30 | + seq_start + rank * chunk_len, | ||
| 31 | + seq_start + (rank + 1) * chunk_len, | ||
| 32 | + device=device, | ||
| 33 | + ) | ||
| 34 | + ) | ||
| 35 | + ranges.append( | ||
| 36 | + torch.arange( | ||
| 37 | + seq_end - (rank + 1) * chunk_len, | ||
| 38 | + seq_end - rank * chunk_len, | ||
| 39 | + device=device, | ||
| 40 | + ) | ||
| 41 | + ) | ||
| 42 | + seq_start = seq_end | ||
| 43 | + return torch.cat(ranges).long() | ||
| 44 | + | ||
| 45 | + | ||
| 46 | +def _diff_cu_seqlens(cu_seqlens): | ||
| 47 | + return [cu_seqlens[0]] + [cu_seqlens[idx] - cu_seqlens[idx - 1] for idx in range(1, len(cu_seqlens))] | ||
| 48 | + | ||
| 49 | + | ||
| 50 | +def _get_combined_kv_prefix_index(cu_seqlens, cp_size, rank, device): | ||
| 51 | + chunk_ids = (rank, 2 * cp_size - rank - 1) | ||
| 52 | + ranges = [] | ||
| 53 | + seq_start = 0 | ||
| 54 | + for seq_end in cu_seqlens: | ||
| 55 | + chunk_len = (seq_end - seq_start) // (2 * cp_size) | ||
| 56 | + for chunk_id in chunk_ids: | ||
| 57 | + ranges.append( | ||
| 58 | + torch.arange( | ||
| 59 | + seq_start, | ||
| 60 | + seq_start + (chunk_id + 1) * chunk_len, | ||
| 61 | + device=device, | ||
| 62 | + ) | ||
| 63 | + ) | ||
| 64 | + seq_start = seq_end | ||
| 65 | + return torch.cat(ranges).long() | ||
| 66 | + | ||
| 67 | + | ||
| 68 | + | ||
| 69 | + "cp_size,cu_seqlens", | ||
| 70 | + [ | ||
| 71 | + (2, [16, 40, 72]), | ||
| 72 | + (4, [32, 80, 144]), | ||
| 73 | + ], | ||
| 74 | +) | ||
| 75 | +def test_thd_load_balanced_cp_metadata(cp_size, cu_seqlens): | ||
| 76 | + full_token_ids = torch.arange(cu_seqlens[-1]) | ||
| 77 | + rank_major_token_ids = torch.cat( | ||
| 78 | + [ | ||
| 79 | + full_token_ids.index_select( | ||
| 80 | + 0, | ||
| 81 | + _get_local_token_index(cu_seqlens, cp_size, rank, torch.device("cpu")), | ||
| 82 | + ) | ||
| 83 | + for rank in range(cp_size) | ||
| 84 | + ] | ||
| 85 | + ) | ||
| 86 | + | ||
| 87 | + rank_workloads = [] | ||
| 88 | + for rank in range(cp_size): | ||
| 89 | + metadata = get_thd_load_balanced_cp_metadata(cu_seqlens, cp_size, rank, torch.device("cpu")) | ||
| 90 | + | ||
| 91 | + q_lens = _diff_cu_seqlens(metadata["actual_seq_qlen"]) | ||
| 92 | + kv_lens = _diff_cu_seqlens(metadata["actual_seq_kvlen"]) | ||
| 93 | + expected_q_lens = [] | ||
| 94 | + expected_kv_lens = [] | ||
| 95 | + seq_start = 0 | ||
| 96 | + for seq_end in cu_seqlens: | ||
| 97 | + chunk_len = (seq_end - seq_start) // (2 * cp_size) | ||
| 98 | + expected_q_lens.extend((chunk_len, chunk_len)) | ||
| 99 | + expected_kv_lens.extend( | ||
| 100 | + ( | ||
| 101 | + (rank + 1) * chunk_len, | ||
| 102 | + (2 * cp_size - rank) * chunk_len, | ||
| 103 | + ) | ||
| 104 | + ) | ||
| 105 | + seq_start = seq_end | ||
| 106 | + assert q_lens == expected_q_lens | ||
| 107 | + assert kv_lens == expected_kv_lens | ||
| 108 | + assert sum(q_lens) == metadata["local_total_len"] | ||
| 109 | + | ||
| 110 | + kv_index = metadata["kv_index_in_rank_major"] | ||
| 111 | + expected_kv_index = _get_combined_kv_prefix_index(cu_seqlens, cp_size, rank, torch.device("cpu")) | ||
| 112 | + assert torch.equal( | ||
| 113 | + rank_major_token_ids.index_select(0, kv_index), | ||
| 114 | + full_token_ids.index_select(0, expected_kv_index), | ||
| 115 | + ) | ||
| 116 | + | ||
| 117 | + combined_grad = torch.arange(1, kv_index.numel() + 1, dtype=torch.float32) | ||
| 118 | + rank_major_grad = torch.zeros(metadata["full_total_len"]) | ||
| 119 | + reference_grad = torch.zeros(metadata["full_total_len"]) | ||
| 120 | + rank_major_grad.index_add_(0, kv_index, combined_grad) | ||
| 121 | + reference_grad.index_add_(0, expected_kv_index, combined_grad) | ||
| 122 | + assert torch.equal( | ||
| 123 | + rank_major_grad, | ||
| 124 | + reference_grad.index_select(0, rank_major_token_ids), | ||
| 125 | + ) | ||
| 126 | + | ||
| 127 | + rank_workloads.append(sum(q_len * kv_len for q_len, kv_len in zip(q_lens, kv_lens))) | ||
| 128 | + | ||
| 129 | + assert len(set(rank_workloads)) == 1 | ||
| 130 | + | ||
| 131 | + | ||
| 132 | +def test_thd_metadata_tensor_identity_cache_avoids_reconversion(): | ||
| 133 | + clear_thd_load_balanced_cp_metadata_cache() | ||
| 134 | + cu_seqlens = torch.tensor([16, 40, 72]) | ||
| 135 | + metadata = kvallgather_cp._get_thd_load_balanced_cp_metadata_cached(cu_seqlens, 2, 0, torch.device("cpu")) | ||
| 136 | + | ||
| 137 | + cache_info = kvallgather_cp._get_thd_load_balanced_cp_metadata_for_tensor.cache_info() # pylint: disable=no-value-for-parameter | ||
| 138 | + cached_metadata = kvallgather_cp._get_thd_load_balanced_cp_metadata_cached(cu_seqlens, 2, 0, torch.device("cpu")) | ||
| 139 | + updated_cache_info = kvallgather_cp._get_thd_load_balanced_cp_metadata_for_tensor.cache_info() # pylint: disable=no-value-for-parameter | ||
| 140 | + assert cached_metadata is metadata | ||
| 141 | + assert updated_cache_info.hits == cache_info.hits + 1 | ||
| 142 | + | ||
| 143 | + | ||
| 144 | +def test_thd_metadata_tensor_mutation_invalidates_identity_cache(): | ||
| 145 | + clear_thd_load_balanced_cp_metadata_cache() | ||
| 146 | + cu_seqlens = torch.tensor([16, 40, 72]) | ||
| 147 | + metadata = kvallgather_cp._get_thd_load_balanced_cp_metadata_cached(cu_seqlens, 2, 0, torch.device("cpu")) | ||
| 148 | + | ||
| 149 | + cu_seqlens[-1] = 80 | ||
| 150 | + updated_metadata = kvallgather_cp._get_thd_load_balanced_cp_metadata_cached(cu_seqlens, 2, 0, torch.device("cpu")) | ||
| 151 | + assert updated_metadata is not metadata | ||
| 152 | + assert updated_metadata["full_total_len"] == 80 | ||
| 153 | + | ||
| 154 | + | ||
| 155 | +class TestKVAllGatherContextParallelTHD(DistributedTest): | ||
| 156 | + world_size = 2 | ||
| 157 | + | ||
| 158 | + def test_eod_load_balanced_batch_partition(self, monkeypatch): | ||
| 159 | + initialize_model_parallel(context_parallel_size=self.world_size) | ||
| 160 | + rank = dist.get_rank() | ||
| 161 | + cu_seqlens = [16, 40, 72] | ||
| 162 | + full_tokens = torch.arange(cu_seqlens[-1], device="npu").view(1, -1) | ||
| 163 | + batch = { | ||
| 164 | + "tokens": full_tokens.clone(), | ||
| 165 | + "labels": full_tokens.clone(), | ||
| 166 | + "loss_mask": torch.ones_like(full_tokens, dtype=torch.float32), | ||
| 167 | + "attention_mask": None, | ||
| 168 | + "position_ids": full_tokens.clone(), | ||
| 169 | + } | ||
| 170 | + | ||
| 171 | + monkeypatch.setattr(get_batch_utils, "get_ring_degree", lambda: 1) | ||
| 172 | + get_batch_utils._get_batch_on_this_cp_rank_in_megatron_cp_eod_padding( | ||
| 173 | + batch, torch.tensor(cu_seqlens, device="npu") | ||
| 174 | + ) | ||
| 175 | + | ||
| 176 | + local_index = _get_local_token_index(cu_seqlens, self.world_size, rank, torch.device("npu")) | ||
| 177 | + expected_tokens = full_tokens.index_select(1, local_index) | ||
| 178 | + assert torch.equal(batch["tokens"], expected_tokens) | ||
| 179 | + assert torch.equal(batch["labels"], expected_tokens) | ||
| 180 | + assert torch.equal(batch["position_ids"], expected_tokens) | ||
| 181 | + | ||
| 182 | + | ||
| 183 | + | ||
| 184 | + "query_heads,kv_heads,qk_head_dim,v_head_dim", | ||
| 185 | + [ | ||
| 186 | + pytest.param(4, 4, 128, 128, id="mha"), | ||
| 187 | + pytest.param(8, 2, 128, 128, id="gqa"), | ||
| 188 | + pytest.param(4, 4, 192, 128, id="mla"), | ||
| 189 | + ], | ||
| 190 | + ) | ||
| 191 | + def test_forward_and_backward_match_non_cp_reference( | ||
| 192 | + self, | ||
| 193 | + dtype, | ||
| 194 | + query_heads, | ||
| 195 | + kv_heads, | ||
| 196 | + qk_head_dim, | ||
| 197 | + v_head_dim, | ||
| 198 | + ): | ||
| 199 | + initialize_model_parallel(context_parallel_size=self.world_size) | ||
| 200 | + set_random_seed(1234) | ||
| 201 | + | ||
| 202 | + rank = dist.get_rank() | ||
| 203 | + cu_seqlens = [16, 40] | ||
| 204 | + total_tokens = cu_seqlens[-1] | ||
| 205 | + softmax_scale = 1.0 / math.sqrt(qk_head_dim) | ||
| 206 | + | ||
| 207 | + q_ref = torch.randn( | ||
| 208 | + total_tokens, | ||
| 209 | + query_heads, | ||
| 210 | + qk_head_dim, | ||
| 211 | + dtype=dtype, | ||
| 212 | + device="npu", | ||
| 213 | + requires_grad=True, | ||
| 214 | + ) | ||
| 215 | + k_ref = torch.randn( | ||
| 216 | + total_tokens, | ||
| 217 | + kv_heads, | ||
| 218 | + qk_head_dim, | ||
| 219 | + dtype=dtype, | ||
| 220 | + device="npu", | ||
| 221 | + requires_grad=True, | ||
| 222 | + ) | ||
| 223 | + v_ref = torch.randn( | ||
| 224 | + total_tokens, | ||
| 225 | + kv_heads, | ||
| 226 | + v_head_dim, | ||
| 227 | + dtype=dtype, | ||
| 228 | + device="npu", | ||
| 229 | + requires_grad=True, | ||
| 230 | + ) | ||
| 231 | + dout = torch.randn( | ||
| 232 | + total_tokens, | ||
| 233 | + query_heads, | ||
| 234 | + v_head_dim, | ||
| 235 | + dtype=dtype, | ||
| 236 | + device="npu", | ||
| 237 | + ) | ||
| 238 | + attention_mask = ~torch.tril(torch.ones((2048, 2048), dtype=torch.bool, device="npu")) | ||
| 239 | + | ||
| 240 | + out_ref = torch_npu.npu_fusion_attention( | ||
| 241 | + q_ref, | ||
| 242 | + k_ref, | ||
| 243 | + v_ref, | ||
| 244 | + query_heads, | ||
| 245 | + "TND", | ||
| 246 | + pse=None, | ||
| 247 | + padding_mask=None, | ||
| 248 | + atten_mask=attention_mask, | ||
| 249 | + scale=softmax_scale, | ||
| 250 | + pre_tockens=65536, | ||
| 251 | + next_tockens=0, | ||
| 252 | + keep_prob=1.0, | ||
| 253 | + inner_precise=0, | ||
| 254 | + sparse_mode=3, | ||
| 255 | + actual_seq_qlen=cu_seqlens, | ||
| 256 | + actual_seq_kvlen=cu_seqlens, | ||
| 257 | + )[0] | ||
| 258 | + out_ref.backward(dout) | ||
| 259 | + | ||
| 260 | + local_index = _get_local_token_index(cu_seqlens, self.world_size, rank, torch.device("npu")) | ||
| 261 | + q = q_ref.detach().index_select(0, local_index).requires_grad_(True) | ||
| 262 | + k = k_ref.detach().index_select(0, local_index).requires_grad_(True) | ||
| 263 | + v = v_ref.detach().index_select(0, local_index).requires_grad_(True) | ||
| 264 | + | ||
| 265 | + out = AttnFuncWithCPAndKVAllGatherForTHD.apply( | ||
| 266 | + q, | ||
| 267 | + k, | ||
| 268 | + v, | ||
| 269 | + query_heads, | ||
| 270 | + attention_mask, | ||
| 271 | + "thd", | ||
| 272 | + "causal", | ||
| 273 | + 0.0, | ||
| 274 | + softmax_scale, | ||
| 275 | + False, | ||
| 276 | + mpu.get_context_parallel_group(), | ||
| 277 | + cu_seqlens, | ||
| 278 | + cu_seqlens, | ||
| 279 | + ) | ||
| 280 | + assert out.shape == (q.shape[0], query_heads, v_head_dim) | ||
| 281 | + out.backward(dout.index_select(0, local_index)) | ||
| 282 | + | ||
| 283 | + tolerances = {"atol": 5e-3, "rtol": 5e-3} | ||
| 284 | + if dtype == torch.bfloat16: | ||
| 285 | + tolerances = {"atol": 2.5e-2, "rtol": 2.5e-2} | ||
| 286 | + | ||
| 287 | + assert torch.allclose(out, out_ref.detach().index_select(0, local_index), **tolerances) | ||
| 288 | + assert torch.allclose(q.grad, q_ref.grad.index_select(0, local_index), **tolerances) | ||
| 289 | + assert torch.allclose(k.grad, k_ref.grad.index_select(0, local_index), **tolerances) | ||
| 290 | + assert torch.allclose(v.grad, v_ref.grad.index_select(0, local_index), **tolerances) | ||
这里对全局变量引用的定义删除，是否会对存量逻辑产生影响？