fix: Fix the issue where the LoRA A/B matrices of the DeepSeek V4 linear_o_down_proj layer are not updated.-MindSpeed-LLM-AtomGit

已合并

fix: Fix the issue where the LoRA A/B matrices of the DeepSeek V4 linear_o_down_proj layer are not updated. #4922

xiejiahao2333创建于 5 天前从xiejiahao2333/MindSpeed-LLM-Lora-bugfix： fix-lora-down-proj 合入到 Ascend/MindSpeed-LLM： master

fix: Fix the issue where the LoRA A/B matrices of the DeepSeek V4 linear_o_down_proj layer are not updated. #4922

已合并

xiejiahao2333创建于 5 天前

从xiejiahao2333/MindSpeed-LLM-Lora-bugfix fix-lora-down-proj合入到Ascend/MindSpeed-LLM master

xiejiahao2333

5 天前

https://gitcode.com/Ascend/MindSpeed-LLM/issues/1772

What this PR does / why we need it?

Please describe the background and detailed changes of the PR. If it is a bugfix, please attach the related issue.
Fix the issue where the A/B matrices of the DeepSeek V4 linear_o_down_proj layer are not updated.
一、Issue:：The B matrix of linear_o_down_proj remains all zeros and is not updated. The A matrix also shows cases where it is not updated.

二、Root cause:：
The original implementation at /opt/MindSpeed-LLM/mindspeed_llm/tasks/models/transformer/deepseek4/g2_attention.py:572 directly reads self.linear_o_down_proj_weight for the einsum computation. This is equivalent to using only the original weight of this layer for computation, without adding the LoRA A/B matrices. As a result, the A/B matrices of linear_o_down_proj do not receive gradients and are not updated.

Does this PR introduce any user-facing change?

Please describe whether the PR will result in any user-facing usage changes. If there is related documentation, please specify its path.
Solution:：
The newly added logic, self._linear_o_down_proj_lora(), works as follows:

a. First, check whether linear_o_down_proj has LoRA A/B matrices.

b. If LoRA is not enabled, directly return the original output, so normal training is not affected.

c. If LoRA is enabled, retrieve lora_A, lora_B, and scaling.

d. Use the input grouped_o to compute an additional LoRA branch:

lora_delta = grouped_o * lora_A * lora_B * scaling

e. Add the LoRA delta back to the original result:

output = base_output + lora_delta

How was this patch tested?

Please explain how to verify the correctness and effectiveness of this feature, as well as its usage constraints and limitations.
Commit ID: 1ed31e0e36aefdd2e868ec8e865164fd0b030480
Training configuration:

{
    "task_type": "sft",
    "model_type": "deepseek4",
    "peft_type": "full",
    "custom": {
        "use_mp_training": true,
        "train_dir": "/dsV4/train_chatml.json"
    },
    "train": {
        "tokenizer-name-or-path": "/models/deepseek-ai/DeepSeek-V4-Flash-BF16",
        "tokenizer-type": "PretrainedFromHF",
        "load": "/dsV4/weight/DeepSeek-V4-Flash-0731-Megatron-pp4tp1ep4",
        "save": "./save_lora",
        "prompt_type": "deepseek4",
        "finetune": true,
        "no-load-optim": true,
        "no-load-rng": true,
        "seed": 1234,
        "split": "100,0,0",
        "lr": 1e-5,
        "min-lr": "1.0e-7",
        "lr-decay-style": "cosine",
        "lr-warmup-iters": 5,
        "bf16": true,
        "weight-decay": "1e-2",
        "clip-grad": "1.0",
        "adam-beta1": "0.9",
        "adam-beta2": "0.999",
        "train-iters": 2000,
        "micro-batch-size": 1,
        "global-batch-size": 128,
        "tensor-model-parallel-size": 1,
        "pipeline-model-parallel-size": 4,
        "expert-model-parallel-size": 4,
        "expert-tensor-parallel-size": 1,
        "compress-rope-theta": 160000.0,
        "context-parallel-size": 1,
        "rope_factor": 16.0,
        "sequence-parallel": true,
        "seq-length": 4096,
        "mtp-num-layers": 0,
        "initial-loss-scale": 65536.0,
        "log-interval": 1,
        "eval-interval": 1000,
        "eval-iters": 1000,
        "save-interval": 50,
        "use-flash-attn": true,
        "use-distributed-optimizer": true,
        "tensorboard-queue-size": 1000,
        "tensorboard-dir": "./log_dir",
        "recompute-granularity": "full",
        "swap-optimizer": true,
        "use-mcore-models": true,
        "is-instruction-dataset": true,
        "use-triton-mhc": false,
        "ckpt-format": "torch",
        "moe-permute-fusion": true,
        "moe-grouped-gemm": true,
        "moe-shared-expert-overlap": false,
        "recompute-method": "uniform",
        "recompute-num-layers": 1,
        "no-save-optim": true,
        "no-save-rng": true,
        "log-throughput": true,
        "fix-router": false,
        "lora-r": 16,
        "lora-alpha": 32,
        "lora-fusion": true,
        "lora-target-modules": ["linear_q", "linear_kv", "linear_fc1", "linear_fc2", "linear_q_up_proj", "linear_o_down_proj", "linear_o_up_proj"],
        "moe-token-dispatcher-type": "alltoall_seq"
   }
}

Loss curve:

Update status of the A/B matrices：

vLLM test result after training:

Pull Request已成功合入, 合并人@ascend-robot

（感谢 xiejiahao2333 的贡献）

X xiejiahao2333

5 天前创建了 pull request，commit 9111ce94

atomgit-bot

5 天前评论：

变更摘要

此 PR 修复了 DeepSeekV4 模型中 linear_o_down_proj 层 LoRA A/B 矩阵梯度不更新的问题。根本原因是原实现直接使用 self.linear_o_down_proj_weight 进行 einsum 计算，仅依赖原始权重，未将 LoRA A/B 矩阵纳入计算路径，导致 LoRA 参数无法获得梯度更新。修复方案是在 DeepSeek4SelfAttention 中新增 _linear_o_down_proj_lora 方法，显式计算 LoRA 分支的增量并叠加到基础输出上。

主要改动

新增 _linear_o_down_proj_lora 方法：在 DeepSeek4SelfAttention 类中新增该方法，依次检查 linear_o_down_proj 是否存在 lora_A/lora_B 属性、适配器是否被禁用或已合并、以及活跃适配器列表；对每个活跃适配器，通过 lora_A.weight 和 lora_B.weight 计算 LoRA 增量（lora_delta = grouped_o * lora_A * lora_B * scaling），并叠加到基础输出上返回。
forward 方法中保留原始输入 grouped_o：在调用 weight_woa 的 einsum 之前，将未变换的 o 保存为 grouped_o，供后续 LoRA 分支使用，确保 LoRA 计算的输入与基础路径一致。
在 einsum 后追加 LoRA 分支调用：在 o = torch.einsum("sbgd,gld->sbgl", grouped_o, weight_woa) 之后，立即调用 o = self._linear_o_down_proj_lora(grouped_o, o)，使 LoRA A/B 矩阵参与前向计算并能够接收梯度更新。

ascend-robot成员

5 天前评论：

openLiBing可接管PR, PR详情链接:
https://www.openlibing.com/apps/transitionPrDetail?owner=Ascend&repo=MindSpeed-LLM&number=4922&platform=gitcode&codeHostingPlatformFlag=gitcode

atomgit-bot

5 天前评论：

代码审查

✅ 未发现问题

ascend-robot成员

5 天前添加了label：ascend-cla/yes

此处折叠了82条消息查看更多

ChenJingyi成员

1 天前评论：

/approve

ascend-robot成员

1 天前添加了label：approved

ascend-robot成员

1 天前合入了pull request

ascend-robot成员

1 天前评论：

Pull Request 已合并或已关闭。

If you want to solve this problem, you can click here to do it in the FAQs.

ascend-robot成员

1 天前评论：

Pull Request 已合并或已关闭。

If you want to solve this problem, you can click here to do it in the FAQs.