已合并
fix: Fix the issue where the LoRA A/B matrices of the DeepSeek V4 linear_o_down_proj layer are not updated. #4922
xiejiahao2333创建于 25 天前
fix: Fix the issue where the LoRA A/B matrices of the DeepSeek V4 linear_o_down_proj layer are not updated. #4922
已合并
共 1 个文件变更+46-9
| @@ -508,6 +508,51 @@ class DeepSeek4SelfAttention(MegatronModule): | |||
| 508 | avg_group=parallel_state.get_tensor_and_context_parallel_group(), | 508 | avg_group=parallel_state.get_tensor_and_context_parallel_group(), |
丁 | |||
| 509 | ) | 509 | ) |
| 510 | 510 | ||
| 511 | + def _linear_o_down_proj(self, grouped_output): | ||
| 512 | + weight_woa = rearrange( | ||
| 513 | + self.linear_o_down_proj.weight, | ||
| 514 | + '(g l) (d h)->g l (d h)', # outdim*indim | ||
| 515 | + d=self.head_dim // self.n_groups, | ||
| 516 | + l=self.o_lora_rank, | ||
| 517 | + h=self.n_heads, | ||
| 518 | + g=self.n_local_groups, | ||
| 519 | + ) | ||
| 520 | + output = torch.einsum("sbgd,gld->sbgl", grouped_output, weight_woa) | ||
| 521 | + | ||
| 522 | + linear = self.linear_o_down_proj | ||
C 26.1.0分支同步修改   | |||
| 523 | + if not (hasattr(linear, "lora_A") and hasattr(linear, "lora_B")): | ||
| 524 | + return output | ||
| 525 | + if getattr(linear, "disable_adapters", False) or getattr(linear, "merged", False): | ||
| 526 | + return output | ||
| 527 | + | ||
| 528 | + active_adapters = getattr(linear, "active_adapters", None) | ||
| 529 | + if active_adapters is None: | ||
| 530 | + active_adapter = getattr(linear, "active_adapter", None) | ||
| 531 | + active_adapters = [active_adapter] if isinstance(active_adapter, str) else active_adapter | ||
| 532 | + if not active_adapters: | ||
| 533 | + return output | ||
| 534 | + | ||
| 535 | + for active_adapter in active_adapters: | ||
| 536 | + if active_adapter not in linear.lora_A.keys() or active_adapter not in linear.lora_B.keys(): | ||
| 537 | + continue | ||
| 538 | + | ||
| 539 | + lora_A = linear.lora_A[active_adapter].weight | ||
| 540 | + lora_B = linear.lora_B[active_adapter].weight | ||
| 541 | + scaling = linear.scaling[active_adapter] | ||
| 542 | + | ||
| 543 | + lora_input = grouped_output.to(lora_A.dtype) | ||
| 544 | + lora_a_output = torch.einsum("sbgd,rd->sbgr", lora_input, lora_A) | ||
| 545 | + lora_b_weight = rearrange( | ||
| 546 | + lora_B, | ||
| 547 | + "(g l) r -> g l r", | ||
| 548 | + g=self.n_local_groups, | ||
| 549 | + l=self.o_lora_rank, | ||
| 550 | + ) | ||
| 551 | + lora_delta = torch.einsum("sbgr,glr->sbgl", lora_a_output, lora_b_weight) * scaling | ||
| 552 | + output = output + lora_delta.to(output.dtype) | ||
| 553 | + | ||
| 554 | + return output | ||
| 555 | + | ||
| 511 | def forward( | 556 | def forward( |
| 512 | self, | 557 | self, |
| 513 | hidden_states: torch.Tensor, | 558 | hidden_states: torch.Tensor, |
| @@ -787,15 +832,7 @@ class DeepSeek4SelfAttention(MegatronModule): | |||
| 787 | d=self.head_dim, | 832 | d=self.head_dim, |
| 788 | ) | 833 | ) |
| 789 | 834 | ||
| 790 | - weight_woa = rearrange( | 835 | + o = self._linear_o_down_proj(o) |
| 791 | - self.linear_o_down_proj.weight, | ||
| 792 | - '(g l) (d h)->g l (d h)', # outdim*indim | ||
| 793 | - d=self.head_dim // self.n_groups, | ||
| 794 | - l=self.o_lora_rank, | ||
| 795 | - h=self.n_heads, | ||
| 796 | - g=self.n_local_groups, | ||
| 797 | - ) | ||
| 798 | - o = torch.einsum("sbgd,gld->sbgl", o, weight_woa) | ||
| 799 | core_attn_out, bias = self.linear_o_up_proj(o.flatten(2)) | 836 | core_attn_out, bias = self.linear_o_up_proj(o.flatten(2)) |
| 800 | 837 | ||
| 801 | return core_attn_out, bias | 838 | return core_attn_out, bias |
PR名称和PR描述修改为英文描述