已合并
fix: use deterministic placeholder initialization for DeepSeek4 hash layers #4898
丁子叉创建于 28 天前
fix: use deterministic placeholder initialization for DeepSeek4 hash layers #4898
已合并
共 1 个文件变更+115-95
| @@ -14,7 +14,6 @@ | |||
| 14 | # limitations under the License. | 14 | # limitations under the License. |
| 15 | 15 | ||
丁 | |||
| 16 | from functools import wraps | 16 | from functools import wraps |
| 17 | -from functools import partial | ||
| 18 | 17 | ||
| 19 | import torch | 18 | import torch |
| 20 | import torch.nn.functional as F | 19 | import torch.nn.functional as F |
| @@ -36,11 +35,9 @@ from mindspeed_llm.core.transformer.moe.moe_utils import topk_softmax_with_capac | |||
| 36 | def group_limited_greedy_topKgating(self, logits: torch.Tensor): | 35 | def group_limited_greedy_topKgating(self, logits: torch.Tensor): |
| 37 | args = get_args() | 36 | args = get_args() |
| 38 | seq_length = logits.shape[0] | 37 | seq_length = logits.shape[0] |
| 39 | - | 38 | + |
| 40 | scores = F.softmax(logits, dim=1) | 39 | scores = F.softmax(logits, dim=1) |
| 41 | - group_scores = ( | 40 | + group_scores = scores.view(args.micro_batch_size * seq_length, self.n_group, -1).max(dim=-1).values # [n, EP] |
| 42 | - scores.view(args.micro_batch_size * seq_length, self.n_group, -1).max(dim=-1).values | ||
| 43 | - ) # [n, EP] | ||
| 44 | 41 | ||
| 45 | group_idx = torch.topk(group_scores, k=self.topk_group, dim=-1, sorted=False)[1] # [n, top_k_group] | 42 | group_idx = torch.topk(group_scores, k=self.topk_group, dim=-1, sorted=False)[1] # [n, top_k_group] |
| 46 | 43 | ||
| @@ -48,17 +45,13 @@ def group_limited_greedy_topKgating(self, logits: torch.Tensor): | |||
| 48 | group_mask.scatter_(1, group_idx, 1) # [n, EP] | 45 | group_mask.scatter_(1, group_idx, 1) # [n, EP] |
| 49 | score_mask = ( | 46 | score_mask = ( |
| 50 | group_mask.unsqueeze(-1) | 47 | group_mask.unsqueeze(-1) |
| 51 | - .expand( | 48 | + .expand(args.micro_batch_size * seq_length, self.n_group, args.num_experts // self.n_group) |
| 52 | - args.micro_batch_size * seq_length, self.n_group, args.num_experts // self.n_group | ||
| 53 | - ) | ||
| 54 | .reshape(args.micro_batch_size * seq_length, -1) | 49 | .reshape(args.micro_batch_size * seq_length, -1) |
| 55 | ) # [n, e] | 50 | ) # [n, e] |
| 56 | 51 | ||
| 57 | tmp_scores = scores.masked_fill(~score_mask.bool(), 0.0) # [n, e] | 52 | tmp_scores = scores.masked_fill(~score_mask.bool(), 0.0) # [n, e] |
| 58 | 53 | ||
| 59 | - topk_weight, topk_idx = torch.topk( | 54 | + topk_weight, topk_idx = torch.topk(tmp_scores, k=args.moe_router_topk, dim=-1, sorted=False) |
| 60 | - tmp_scores, k=args.moe_router_topk, dim=-1, sorted=False | ||
| 61 | - ) | ||
| 62 | 55 | ||
| 63 | ### norm gate to sum 1 | 56 | ### norm gate to sum 1 |
| 64 | if args.moe_router_topk > 1 and args.norm_topk_prob: | 57 | if args.moe_router_topk > 1 and args.norm_topk_prob: |
| @@ -77,7 +70,6 @@ def group_limited_greedy_topKgating(self, logits: torch.Tensor): | |||
| 77 | 70 | ||
| 78 | scores_for_aux = scores # [s*b, n_global_experts] | 71 | scores_for_aux = scores # [s*b, n_global_experts] |
| 79 | topk_idx_for_aux_loss = topk_idx.view(args.micro_batch_size, -1) # [b, s*top_k] | 72 | topk_idx_for_aux_loss = topk_idx.view(args.micro_batch_size, -1) # [b, s*top_k] |
| 80 | - topk_group_idx_for_aux_loss = group_idx.view(args.micro_batch_size, -1) # [b, s*topk_group] | ||
| 81 | fi, Pi, l_aux = None, None, 0 | 73 | fi, Pi, l_aux = None, None, 0 |
| 82 | 74 | ||
| 83 | ######################################################### | 75 | ######################################################### |
| @@ -89,9 +81,7 @@ def group_limited_greedy_topKgating(self, logits: torch.Tensor): | |||
| 89 | if args.seq_aux: | 81 | if args.seq_aux: |
| 90 | scores_for_seq_aux = scores_for_aux.view(args.micro_batch_size, seq_length, -1) | 82 | scores_for_seq_aux = scores_for_aux.view(args.micro_batch_size, seq_length, -1) |
| 91 | # [b, s, n_global_experts] | 83 | # [b, s, n_global_experts] |
| 92 | - ce = torch.zeros( | 84 | + ce = torch.zeros(args.micro_batch_size, args.num_experts, device=logits.device) # [b, n_global_experts] |
| 93 | - args.micro_batch_size, args.num_experts, device=logits.device | ||
| 94 | - ) # [b, n_global_experts] | ||
| 95 | ce.scatter_add_( | 85 | ce.scatter_add_( |
| 96 | 1, | 86 | 1, |
| 97 | topk_idx_for_aux_loss, | 87 | topk_idx_for_aux_loss, |
| @@ -105,13 +95,13 @@ def group_limited_greedy_topKgating(self, logits: torch.Tensor): | |||
| 105 | torch.distributed.all_reduce(ce, group=sequence_partition_group) | 95 | torch.distributed.all_reduce(ce, group=sequence_partition_group) |
| 106 | 96 | ||
| 107 | num_tokens = seq_length * num_sub_sequence | 97 | num_tokens = seq_length * num_sub_sequence |
| 108 | - fi = ce.div(num_sub_sequence * num_tokens * args.moe_router_topk / args.num_experts) # [b, n_global_experts] | 98 | + fi = ce.div( |
| 99 | + num_sub_sequence * num_tokens * args.moe_router_topk / args.num_experts | ||
| 100 | + ) # [b, n_global_experts] | ||
| 109 | Pi = scores_for_seq_aux.mean(dim=1) # [b, n_global_experts] | 101 | Pi = scores_for_seq_aux.mean(dim=1) # [b, n_global_experts] |
| 110 | l_expert_aux = (Pi * fi).sum(dim=1).mean() * self.config.moe_aux_loss_coeff | 102 | l_expert_aux = (Pi * fi).sum(dim=1).mean() * self.config.moe_aux_loss_coeff |
| 111 | else: | 103 | else: |
| 112 | - mask_ce = F.one_hot( | 104 | + mask_ce = F.one_hot(topk_idx_for_aux_loss.view(-1), num_classes=args.num_experts) |
| 113 | - topk_idx_for_aux_loss.view(-1), num_classes=args.num_experts | ||
| 114 | - ) | ||
| 115 | ce = mask_ce.to(logits.dtype).mean(0) | 105 | ce = mask_ce.to(logits.dtype).mean(0) |
| 116 | Pi = scores_for_aux.mean(0) | 106 | Pi = scores_for_aux.mean(0) |
| 117 | fi = ce * args.num_experts | 107 | fi = ce * args.num_experts |
| @@ -130,9 +120,7 @@ def group_limited_greedy_topKgating(self, logits: torch.Tensor): | |||
| 130 | if fi is None: | 120 | if fi is None: |
| 131 | scores_for_seq_aux = scores_for_aux.view(args.micro_batch_size, seq_length, -1) | 121 | scores_for_seq_aux = scores_for_aux.view(args.micro_batch_size, seq_length, -1) |
| 132 | 122 | ||
| 133 | - ce = torch.zeros( | 123 | + ce = torch.zeros(args.micro_batch_size, args.num_experts, device=logits.device) # [b, n_global_experts] |
| 134 | - args.micro_batch_size, args.num_experts, device=logits.device | ||
| 135 | - ) # [b, n_global_experts] | ||
| 136 | ce.scatter_add_( | 124 | ce.scatter_add_( |
| 137 | 1, | 125 | 1, |
| 138 | topk_idx_for_aux_loss, | 126 | topk_idx_for_aux_loss, |
| @@ -147,9 +135,7 @@ def group_limited_greedy_topKgating(self, logits: torch.Tensor): | |||
| 147 | 135 | ||
| 148 | else: | 136 | else: |
| 149 | if fi is None: | 137 | if fi is None: |
| 150 | - mask_ce = F.one_hot( | 138 | + mask_ce = F.one_hot(topk_idx_for_aux_loss.view(-1), num_classes=args.num_experts) |
| 151 | - topk_idx_for_aux_loss.view(-1), num_classes=args.num_experts | ||
| 152 | - ) | ||
| 153 | ce = mask_ce.to(logits.dtype).mean(0) | 139 | ce = mask_ce.to(logits.dtype).mean(0) |
| 154 | Pi = scores_for_aux.mean(0) | 140 | Pi = scores_for_aux.mean(0) |
| 155 | fi = ce * args.num_experts | 141 | fi = ce * args.num_experts |
| @@ -222,20 +208,20 @@ def group_limited_greedy_topKgating(self, logits: torch.Tensor): | |||
| 222 | class custom_multiplier(torch.autograd.Function): | 208 | class custom_multiplier(torch.autograd.Function): |
| 223 | 209 | ||
| 224 | def forward( | 210 | def forward( |
| 225 | - ctx, | 211 | + ctx, |
| 226 | - scores: torch.Tensor, | 212 | + scores: torch.Tensor, |
| 227 | - multiplier: torch.Tensor, | 213 | + multiplier: torch.Tensor, |
| 228 | - selected_experts: torch.Tensor, | 214 | + selected_experts: torch.Tensor, |
| 229 | - masked_gates: torch.Tensor, | 215 | + masked_gates: torch.Tensor, |
| 230 | - mask_for_one: torch.Tensor, | 216 | + mask_for_one: torch.Tensor, |
| 231 | ): | 217 | ): |
| 232 | ctx.save_for_backward(multiplier, selected_experts, masked_gates) | 218 | ctx.save_for_backward(multiplier, selected_experts, masked_gates) |
| 233 | return multiplier * mask_for_one | 219 | return multiplier * mask_for_one |
| 234 | 220 | ||
| 235 | 221 | ||
| 236 | def backward( | 222 | def backward( |
| 237 | - ctx, | 223 | + ctx, |
| 238 | - grad_at_output: torch.Tensor, | 224 | + grad_at_output: torch.Tensor, |
| 239 | ): | 225 | ): |
| 240 | multiplier, selected_experts, masked_gates = ctx.saved_tensors | 226 | multiplier, selected_experts, masked_gates = ctx.saved_tensors |
| 241 | 227 | ||
| @@ -272,9 +258,14 @@ def sparsemixer_top2(self, scores, jitter_eps=0.01): | |||
| 272 | masked_gates = scores.masked_fill(mask_logits_threshold, float('-inf')) | 258 | masked_gates = scores.masked_fill(mask_logits_threshold, float('-inf')) |
| 273 | if self.training: | 259 | if self.training: |
| 274 | # gumbel sampling, more robust than than the multinomial method | 260 | # gumbel sampling, more robust than than the multinomial method |
| 275 | - selected_experts = (masked_gates - torch.empty_like( | 261 | + selected_experts = ( |
| 276 | - masked_gates, memory_format=torch.legacy_contiguous_format | 262 | + ( |
| 277 | - ).exponential_().log()).max(dim=-1)[1].unsqueeze(-1) | 263 | + masked_gates |
| 264 | + - torch.empty_like(masked_gates, memory_format=torch.legacy_contiguous_format).exponential_().log() | ||
| 265 | + ) | ||
| 266 | + .max(dim=-1)[1] | ||
| 267 | + .unsqueeze(-1) | ||
| 268 | + ) | ||
| 278 | else: | 269 | else: |
| 279 | selected_experts = max_ind | 270 | selected_experts = max_ind |
| 280 | 271 | ||
| @@ -287,7 +278,7 @@ def sparsemixer_top2(self, scores, jitter_eps=0.01): | |||
| 287 | max_scores, max_ind = masked_gates.max(dim=-1, keepdim=True) | 278 | max_scores, max_ind = masked_gates.max(dim=-1, keepdim=True) |
| 288 | mask_for_one = torch.logical_or( | 279 | mask_for_one = torch.logical_or( |
| 289 | selected_experts == max_ind, | 280 | selected_experts == max_ind, |
| 290 | - torch.rand_like(max_scores) > 0.75 # Heun's third-order method: f(x) - f(0) = .25 f'(x) + .75 f'(x/3.) | 281 | + torch.rand_like(max_scores) > 0.75, # Heun's third-order method: f(x) - f(0) = .25 f'(x) + .75 f'(x/3.) |
| 291 | ).int() | 282 | ).int() |
| 292 | # 1 -> 1.0 & 0 -> 1./3: lambda x: (x + 0.5) / 1.5 | 283 | # 1 -> 1.0 & 0 -> 1./3: lambda x: (x + 0.5) / 1.5 |
| 293 | mask_for_one = torch.add(0.3333, mask_for_one, alpha=0.6667).type_as(masked_gates) | 284 | mask_for_one = torch.add(0.3333, mask_for_one, alpha=0.6667).type_as(masked_gates) |
| @@ -318,10 +309,14 @@ def sparsemixer_top2(self, scores, jitter_eps=0.01): | |||
| 318 | # apply mask | 309 | # apply mask |
| 319 | masked_gates_top2 = masked_scores.masked_fill(mask_logits_threshold, float('-inf')) | 310 | masked_gates_top2 = masked_scores.masked_fill(mask_logits_threshold, float('-inf')) |
| 320 | if self.training: | 311 | if self.training: |
| 321 | - selected_experts_top2 = (masked_gates_top2 - torch.empty_like( | 312 | + selected_experts_top2 = ( |
| 322 | - masked_gates_top2, memory_format=torch.legacy_contiguous_format | 313 | + ( |
| 323 | - ).exponential_().log() | 314 | + masked_gates_top2 |
| 324 | - ).max(dim=-1)[1].unsqueeze(-1) # gumbel sampling, more robust than than the multinomial method | 315 | + - torch.empty_like(masked_gates_top2, memory_format=torch.legacy_contiguous_format).exponential_().log() |
| 316 | + ) | ||
| 317 | + .max(dim=-1)[1] | ||
| 318 | + .unsqueeze(-1) | ||
| 319 | + ) # gumbel sampling, more robust than than the multinomial method | ||
| 325 | else: | 320 | else: |
| 326 | selected_experts_top2 = max_ind | 321 | selected_experts_top2 = max_ind |
| 327 | # compute scores for gradients | 322 | # compute scores for gradients |
| @@ -333,7 +328,7 @@ def sparsemixer_top2(self, scores, jitter_eps=0.01): | |||
| 333 | max_scores, max_ind = masked_gates_top2.max(dim=-1, keepdim=True) | 328 | max_scores, max_ind = masked_gates_top2.max(dim=-1, keepdim=True) |
| 334 | mask_for_one_top2 = torch.logical_or( | 329 | mask_for_one_top2 = torch.logical_or( |
| 335 | selected_experts_top2 == max_ind, | 330 | selected_experts_top2 == max_ind, |
| 336 | - torch.rand_like(max_scores).uniform_() > 0.75 | 331 | + torch.rand_like(max_scores).uniform_() > 0.75, |
| 337 | # Heun's third-order method: f(x) - f(0) = .25 f'(x) + .75 f'(x/3.) | 332 | # Heun's third-order method: f(x) - f(0) = .25 f'(x) + .75 f'(x/3.) |
| 338 | ).int() | 333 | ).int() |
| 339 | # 1 -> 1.0 & 0 -> 1./3: lambda x: (x + 0.5) / 1.5 | 334 | # 1 -> 1.0 & 0 -> 1./3: lambda x: (x + 0.5) / 1.5 |
| @@ -366,11 +361,16 @@ def topk_router_build_hash_module(self): | |||
| 366 | 361 | ||
| 367 | self.hash = self.layer_number <= mg_args.n_hash_layers | 362 | self.hash = self.layer_number <= mg_args.n_hash_layers |
| 368 | if self.hash: | 363 | if self.hash: |
| 369 | - # self.tid2eid hash [vocab_size, top_k] | 364 | + # DSv4-Pro provides a pre-trained tid2eid table in its inference checkpoint, but no |
| 370 | - self.tid2eid = torch.nn.Parameter( | 365 | + # public initialization recipe is available. This round-robin initialization is only |
| 371 | - torch.stack([torch.randperm(mg_args.moe_router_topk) for _ in range(mg_args.padded_vocab_size)]), | 366 | + # a placeholder to make hash layers runnable from scratch and is not suitable for |
| 372 | - requires_grad=False | 367 | + # real-world training. |
| 368 | + token_ids = torch.arange(mg_args.padded_vocab_size) | ||
| 369 | + tid2eid = torch.stack( | ||
| 370 | + [(token_ids + offset) % mg_args.num_experts for offset in range(mg_args.moe_router_topk)], | ||
| 371 | + dim=1, | ||
| 373 | ) | 372 | ) |
| 373 | + self.tid2eid = torch.nn.Parameter(tid2eid, requires_grad=False) | ||
| 374 | self.expert_bias = None | 374 | self.expert_bias = None |
| 375 | 375 | ||
| 376 | 376 | ||
| @@ -396,15 +396,16 @@ def topk_router_init_wrapper(function): | |||
| 396 | torch.zeros(self.num_experts, dtype=torch.float32), | 396 | torch.zeros(self.num_experts, dtype=torch.float32), |
| 397 | persistent=False, | 397 | persistent=False, |
| 398 | ) | 398 | ) |
| 399 | - self.register_buffer( | 399 | + self.register_buffer('expert_bias', torch.zeros(self.num_experts, dtype=torch.float32)) |
| 400 | - 'expert_bias', torch.zeros(self.num_experts, dtype=torch.float32) | ||
| 401 | - ) | ||
| 402 | else: | 400 | else: |
| 403 | self.local_tokens_per_expert = None | 401 | self.local_tokens_per_expert = None |
| 404 | self.expert_bias = None | 402 | self.expert_bias = None |
| 405 | 403 | ||
| 406 | - self.n_group = mg_args.moe_router_num_groups if mg_args.moe_router_num_groups is not None else ( | 404 | + self.n_group = ( |
| 407 | - mg_args.expert_model_parallel_size) | 405 | + mg_args.moe_router_num_groups |
| 406 | + if mg_args.moe_router_num_groups is not None | ||
| 407 | + else (mg_args.expert_model_parallel_size) | ||
| 408 | + ) | ||
| 408 | self.topk_group = mg_args.moe_router_group_topk | 409 | self.topk_group = mg_args.moe_router_group_topk |
| 409 | self.norm_topk_prob = mg_args.norm_topk_prob | 410 | self.norm_topk_prob = mg_args.norm_topk_prob |
| 410 | setattr(self.__class__, 'build_hash_module', topk_router_build_hash_module) | 411 | setattr(self.__class__, 'build_hash_module', topk_router_build_hash_module) |
| @@ -412,28 +413,30 @@ def topk_router_init_wrapper(function): | |||
| 412 | return topk_router_init | 413 | return topk_router_init |
| 413 | 414 | ||
| 414 | 415 | ||
| 415 | -def topk_router_forward_patch(self, input: torch.Tensor, input_ids: torch.Tensor = None): | 416 | +def topk_router_forward_patch( # pylint: disable=redefined-builtin |
| 416 | - """ | 417 | + self, input: torch.Tensor, input_ids: torch.Tensor = None |
| 417 | - patch for TopKRouter forward | 418 | +): |
| 419 | + """ | ||
| 420 | + patch for TopKRouter forward | ||
| 418 | 421 | ||
| 419 | - Args: | 422 | + Args: |
| 420 | - input (torch.Tensor): Input tensor. | 423 | + input (torch.Tensor): Input tensor. |
| 421 | - input_ids (torch.Tensor): Input ids. | 424 | + input_ids (torch.Tensor): Input ids. |
| 422 | - """ | 425 | + """ |
| 423 | - self._maintain_float32_expert_bias() | 426 | + self._maintain_float32_expert_bias() |
| 424 | 427 | ||
| 425 | - # Apply input jitter | 428 | + # Apply input jitter |
| 426 | - input = self.apply_input_jitter(input) | 429 | + input = self.apply_input_jitter(input) |
| 427 | - logits = self.gating(input) | 430 | + logits = self.gating(input) |
| 428 | 431 | ||
| 429 | - scores, routing_map = self.routing(logits, input_ids) | 432 | + scores, routing_map = self.routing(logits, input_ids) |
| 430 | 433 | ||
| 431 | - return scores, routing_map | 434 | + return scores, routing_map |
| 432 | 435 | ||
| 433 | 436 | ||
| 434 | def apply_seq_aux_loss(self, activation, logits, topk_idx): | 437 | def apply_seq_aux_loss(self, activation, logits, topk_idx): |
| 435 | """ | 438 | """ |
| 436 | - Apply complementary sequence-wise auxiliary loss | 439 | + Apply complementary sequence-wise auxiliary loss |
| 437 | """ | 440 | """ |
| 438 | 441 | ||
| 439 | args = get_args() | 442 | args = get_args() |
| @@ -450,7 +453,7 @@ def apply_seq_aux_loss(self, activation, logits, topk_idx): | |||
| 450 | if self.expert_bias is not None: | 453 | if self.expert_bias is not None: |
| 451 | scores = scores + self.expert_bias | 454 | scores = scores + self.expert_bias |
| 452 | scores = scores / (scores.sum(dim=-1, keepdim=True) + 1e-20) | 455 | scores = scores / (scores.sum(dim=-1, keepdim=True) + 1e-20) |
| 453 | - elif self.score_function == "sqrtsoftplus": | 456 | + elif self.score_function == "sqrtsoftplus": |
| 454 | scores = F.softplus(logits).sqrt() | 457 | scores = F.softplus(logits).sqrt() |
| 455 | if self.expert_bias is not None: | 458 | if self.expert_bias is not None: |
| 456 | scores = scores + self.expert_bias | 459 | scores = scores + self.expert_bias |
| @@ -461,8 +464,12 @@ def apply_seq_aux_loss(self, activation, logits, topk_idx): | |||
| 461 | scores_for_aux = scores # [s*b, n_global_experts] | 464 | scores_for_aux = scores # [s*b, n_global_experts] |
| 462 | topk_idx_for_aux_loss = topk_idx.view(args.micro_batch_size, -1) # [b, s*top_k] | 465 | topk_idx_for_aux_loss = topk_idx.view(args.micro_batch_size, -1) # [b, s*top_k] |
| 463 | scores_for_seq_aux = scores_for_aux.view(args.micro_batch_size, seq_length, -1) | 466 | scores_for_seq_aux = scores_for_aux.view(args.micro_batch_size, seq_length, -1) |
| 464 | - ce = torch.stack([torch.histc(x.to(torch.int32), bins=args.num_experts, min=0, max=args.num_experts) for x in | 467 | + ce = torch.stack( |
| 465 | - topk_idx_for_aux_loss]) | 468 | + [ |
| 469 | + torch.histc(x.to(torch.int32), bins=args.num_experts, min=0, max=args.num_experts) | ||
| 470 | + for x in topk_idx_for_aux_loss | ||
| 471 | + ] | ||
| 472 | + ) | ||
| 466 | 473 | ||
| 467 | num_sub_sequence = 1 | 474 | num_sub_sequence = 1 |
| 468 | sequence_partition_group = parallel_state.get_context_parallel_group() | 475 | sequence_partition_group = parallel_state.get_context_parallel_group() |
| @@ -487,7 +494,7 @@ def apply_seq_aux_loss(self, activation, logits, topk_idx): | |||
| 487 | return activation | 494 | return activation |
| 488 | 495 | ||
| 489 | 496 | ||
| 490 | -def topk_router_gating_func(self, input: torch.Tensor): | 497 | +def topk_router_gating_func(self, input: torch.Tensor): # pylint: disable=redefined-builtin |
| 491 | _args = get_args() | 498 | _args = get_args() |
| 492 | 499 | ||
| 493 | if _args.router_gating_in_fp32: | 500 | if _args.router_gating_in_fp32: |
| @@ -495,6 +502,7 @@ def topk_router_gating_func(self, input: torch.Tensor): | |||
| 495 | # if weight is not requires_grad like lora finetune, can not autograd for weight in checkpoint_manager | 502 | # if weight is not requires_grad like lora finetune, can not autograd for weight in checkpoint_manager |
| 496 | logits = F.linear(input.type(torch.float32), self.weight.type(torch.float32)) | 503 | logits = F.linear(input.type(torch.float32), self.weight.type(torch.float32)) |
| 497 | else: | 504 | else: |
| 505 | + | ||
| 498 | def to_fp32(_input, weight): | 506 | def to_fp32(_input, weight): |
| 499 | return _input.type(torch.float32), weight.type(torch.float32) | 507 | return _input.type(torch.float32), weight.type(torch.float32) |
| 500 | 508 | ||
| @@ -507,6 +515,7 @@ def topk_router_gating_func(self, input: torch.Tensor): | |||
| 507 | else: | 515 | else: |
| 508 | if self.config.moe_router_dtype == 'fp8': | 516 | if self.config.moe_router_dtype == 'fp8': |
| 509 | from mindspeed.te.pytorch.fp8.recipes import matmul_fp8 | 517 | from mindspeed.te.pytorch.fp8.recipes import matmul_fp8 |
| 518 | + | ||
| 510 | logits = matmul_fp8(input, self.weight) | 519 | logits = matmul_fp8(input, self.weight) |
| 511 | else: | 520 | else: |
| 512 | if self.config.moe_router_dtype == 'fp32': | 521 | if self.config.moe_router_dtype == 'fp32': |
| @@ -539,10 +548,8 @@ def topk_router_routing(self, logits: torch.Tensor, input_ids: torch.Tensor = No | |||
| 539 | logits = self.apply_z_loss(logits) | 548 | logits = self.apply_z_loss(logits) |
| 540 | 549 | ||
| 541 | args = get_args() | 550 | args = get_args() |
| 542 | - if ( | 551 | + dispatcher_type = self.config.moe_token_dispatcher_type |
| 543 | - self.config.tensor_model_parallel_size > 1 | 552 | + if self.config.tensor_model_parallel_size > 1 and dispatcher_type == "alltoall_seq": |
| 544 | - and self.config.moe_token_dispatcher_type == "alltoall_seq" | ||
| 545 | - ): | ||
| 546 | # Gather the logits from the TP region | 553 | # Gather the logits from the TP region |
| 547 | logits = gather_from_sequence_parallel_region(logits) | 554 | logits = gather_from_sequence_parallel_region(logits) |
| 548 | 555 | ||
| @@ -562,7 +569,7 @@ def topk_router_routing(self, logits: torch.Tensor, input_ids: torch.Tensor = No | |||
| 562 | logits_ = torch.softmax(logits, dim=-1, dtype=torch.float32) | 569 | logits_ = torch.softmax(logits, dim=-1, dtype=torch.float32) |
| 563 | else: | 570 | else: |
| 564 | logits_ = torch.softmax(logits, dim=-1, dtype=torch.float32).type_as(logits) | 571 | logits_ = torch.softmax(logits, dim=-1, dtype=torch.float32).type_as(logits) |
| 565 | - | 572 | + |
| 566 | if self.expert_bias is not None: | 573 | if self.expert_bias is not None: |
| 567 | logits_for_routing = logits_ + self.expert_bias | 574 | logits_for_routing = logits_ + self.expert_bias |
| 568 | _, indices = torch.topk(logits_for_routing, k=self.topk, dim=1) | 575 | _, indices = torch.topk(logits_for_routing, k=self.topk, dim=1) |
| @@ -618,11 +625,12 @@ def topk_router_routing(self, logits: torch.Tensor, input_ids: torch.Tensor = No | |||
| 618 | ) | 625 | ) |
| 619 | args = get_args() | 626 | args = get_args() |
| 620 | if self.training and args.seq_aux: | 627 | if self.training and args.seq_aux: |
| 621 | - scores = apply_seq_aux_loss(self, | 628 | + scores = apply_seq_aux_loss( |
| 622 | - activation=scores, | 629 | + self, |
| 623 | - logits=logits, | 630 | + activation=scores, |
| 624 | - topk_idx=routing_map, | 631 | + logits=logits, |
| 625 | - ) | 632 | + topk_idx=routing_map, |
| 633 | + ) | ||
| 626 | else: | 634 | else: |
| 627 | raise ValueError(f"Unsupported MoE routing type: {self.routing_type}") | 635 | raise ValueError(f"Unsupported MoE routing type: {self.routing_type}") |
| 628 | if args.moe_tp_extend_ep: | 636 | if args.moe_tp_extend_ep: |
| @@ -635,9 +643,14 @@ def topk_router_routing(self, logits: torch.Tensor, input_ids: torch.Tensor = No | |||
| 635 | 643 | ||
| 636 | # fix router if needed | 644 | # fix router if needed |
| 637 | if args.fix_router: | 645 | if args.fix_router: |
| 646 | + | ||
| 638 | def fix_indices(index_tensor, logits_shape, router_topk): | 647 | def fix_indices(index_tensor, logits_shape, router_topk): |
| 639 | - expert_select = torch.arange(index_tensor.shape[0] * router_topk, device=index_tensor.device, | 648 | + expert_select = ( |
| 640 | - dtype=torch.int64).view(index_tensor.shape[0], router_topk) % logits_shape[-1] | 649 | + torch.arange(index_tensor.shape[0] * router_topk, device=index_tensor.device, dtype=torch.int64).view( |
| 650 | + index_tensor.shape[0], router_topk | ||
| 651 | + ) | ||
| 652 | + % logits_shape[-1] | ||
| 653 | + ) | ||
| 641 | routing_map = torch.zeros(index_tensor.shape, device=index_tensor.device, dtype=torch.bool) | 654 | routing_map = torch.zeros(index_tensor.shape, device=index_tensor.device, dtype=torch.bool) |
| 642 | routing_map.scatter_(1, expert_select, True) | 655 | routing_map.scatter_(1, expert_select, True) |
| 643 | return routing_map | 656 | return routing_map |
| @@ -692,7 +705,9 @@ def global_aux_loss_load_balancing(self, logits: torch.Tensor): | |||
| 692 | return probs, routing_map | 705 | return probs, routing_map |
| 693 | 706 | ||
| 694 | 707 | ||
| 695 | -def global_aux_loss_topk_router_forward(self, input: torch.Tensor): | 708 | +def global_aux_loss_topk_router_forward( # pylint: disable=redefined-builtin |
| 709 | + self, input: torch.Tensor | ||
| 710 | +): | ||
| 696 | """ | 711 | """ |
| 697 | Forward pass of the router. | 712 | Forward pass of the router. |
| 698 | 713 | ||
| @@ -709,6 +724,7 @@ def global_aux_loss_topk_router_forward(self, input: torch.Tensor): | |||
| 709 | 724 | ||
| 710 | return scores, routing_map, logits.detach() | 725 | return scores, routing_map, logits.detach() |
| 711 | 726 | ||
| 727 | + | ||
| 712 | def global_load_balancing_loss_func(router_logits, attention_mask, config): | 728 | def global_load_balancing_loss_func(router_logits, attention_mask, config): |
| 713 | """ | 729 | """ |
| 714 | Computes auxiliary load balancing loss as in Switch Transformer - implemented in Pytorch. | 730 | Computes auxiliary load balancing loss as in Switch Transformer - implemented in Pytorch. |
| @@ -729,16 +745,18 @@ def global_load_balancing_loss_func(router_logits, attention_mask, config): | |||
| 729 | Returns: | 745 | Returns: |
| 730 | The auxiliary loss. | 746 | The auxiliary loss. |
| 731 | """ | 747 | """ |
| 732 | - | ||
| 733 | if router_logits is None or not isinstance(router_logits, tuple): | 748 | if router_logits is None or not isinstance(router_logits, tuple): |
| 734 | return 0 | 749 | return 0 |
| 735 | 750 | ||
| 736 | - if isinstance(router_logits, tuple): | 751 | + compute_device = router_logits[0].device |
| 737 | - compute_device = router_logits[0].device | 752 | + concatenated_gate_logits = torch.cat( |
| 738 | - concatenated_gate_logits = torch.cat( | 753 | + [ |
| 739 | - [layer_gate.to(compute_device).transpose(0, 1).reshape(-1, layer_gate.shape[2]) | 754 | + layer_gate.to(compute_device).transpose(0, 1).reshape(-1, layer_gate.shape[2]) |
| 740 | - for layer_gate in router_logits], dim=0) | 755 | + for layer_gate in router_logits |
| 741 | - | 756 | + ], |
| 757 | + dim=0, | ||
| 758 | + ) | ||
| 759 | + | ||
| 742 | top_k = config.moe_router_topk | 760 | top_k = config.moe_router_topk |
| 743 | num_experts = concatenated_gate_logits.shape[1] | 761 | num_experts = concatenated_gate_logits.shape[1] |
| 744 | 762 | ||
| @@ -788,9 +806,10 @@ def global_load_balancing_loss_func(router_logits, attention_mask, config): | |||
| 788 | 806 | ||
| 789 | return overall_loss * num_experts | 807 | return overall_loss * num_experts |
| 790 | 808 | ||
| 791 | -# TODO remove it when megatron support sqrtsoftplus | 809 | + |
| 792 | -def transformer_config_post_init_wrapper(fn): | 810 | +# Remove this wrapper when Megatron supports sqrtsoftplus. |
| 793 | - @wraps(fn) | 811 | +def transformer_config_post_init_wrapper(fn): |
| 812 | + | ||
| 794 | def wrapper(self): # | 813 | def wrapper(self): # |
| 795 | allowed_score_function = {"softmax", "sigmoid", "sqrtsoftplus"} | 814 | allowed_score_function = {"softmax", "sigmoid", "sqrtsoftplus"} |
| 796 | bypass_flag = ( | 815 | bypass_flag = ( |
| @@ -799,13 +818,14 @@ def transformer_config_post_init_wrapper(fn): | |||
| 799 | and getattr(self, "moe_router_score_function", None) != "sigmoid" | 818 | and getattr(self, "moe_router_score_function", None) != "sigmoid" |
| 800 | ) | 819 | ) |
| 801 | if not bypass_flag: | 820 | if not bypass_flag: |
| 802 | - return fn(self) | 821 | + return fn(self) |
| 803 | old_score_fn = self.moe_router_score_function | 822 | old_score_fn = self.moe_router_score_function |
| 804 | try: | 823 | try: |
| 805 | # bypass megatron's check | 824 | # bypass megatron's check |
| 806 | self.moe_router_score_function = "sigmoid" | 825 | self.moe_router_score_function = "sigmoid" |
| 807 | - return fn(self) | 826 | + return fn(self) |
| 808 | finally: | 827 | finally: |
| 809 | # restore user's config | 828 | # restore user's config |
| 810 | self.moe_router_score_function = old_score_fn | 829 | self.moe_router_score_function = old_score_fn |
| 811 | - return wrapper | 830 | + |
| 831 | + return wrapper | ||
修改hash layer参数随机初始化逻辑与开源保持一致，并说明其局限性