423A35C7/LLaMA-Factory
1
0
Fork 0
mirror of https://github.com/hiyouga/LLaMA-Factory.git synced 2026-11-06 06:59:59 +08:00
Code Issues Packages Projects Releases Wiki Activity

[npu] Skip unnecessary gradients in NPU fused MoE backward (#10918)

Browse Source
This commit is contained in:
xvxuopop
2026-11-04 21:32:27 +08:00
committed by GitHub
parent 84024cefa6
commit d13f1632e6
1 changed files with 65 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files

107
src/llamafactory/v1/plugins/model_plugins/kernels/ops/mlp/npu_fused_moe.py
View File

@@ -60,7 +60,11 @@ class GmmFunction(torch.autograd.Function):
Returns: Returns:
Tensor: The result of the grouped matrix multiplication. Tensor: The result of the grouped matrix multiplication.
""" """
ctx.save_for_backward(x, weight) # Input gradients need weights; weight gradients need inputs.
ctx.save_for_backward(
x if ctx.needs_input_grad[1] else None,
weight if ctx.needs_input_grad[0] else None,
)
ctx.group_list = group_list ctx.group_list = group_list
fwd_output = torch_npu.npu_grouped_matmul( fwd_output = torch_npu.npu_grouped_matmul(
@@ -82,19 +86,28 @@ class GmmFunction(torch.autograd.Function):
input_tensor, weight = ctx.saved_tensors input_tensor, weight = ctx.saved_tensors
group_list = ctx.group_list group_list = ctx.group_list
weight = torch.transpose(weight, 1, 2) grad_input = grad_weight = None
grad_input = torch_npu.npu_grouped_matmul( # Frozen experts still propagate input gradients to upstream LoRA layers.
[grad_output], [weight], bias=None, group_list=group_list, split_item=2, group_type=0, group_list_type=1 if ctx.needs_input_grad[0]:
)[0] grad_input = torch_npu.npu_grouped_matmul(
grad_weight = torch_npu.npu_grouped_matmul( [grad_output],
[input_tensor.T], [weight.transpose(1, 2)],
[grad_output], bias=None,
bias=None, group_list=group_list,
group_list=group_list, split_item=2,
split_item=3, group_type=0,
group_type=2, group_list_type=1,
group_list_type=1, )[0]
)[0] if ctx.needs_input_grad[1]:
grad_weight = torch_npu.npu_grouped_matmul(
[input_tensor.T],
[grad_output],
bias=None,
group_list=group_list,
split_item=3,
group_type=2,
group_list_type=1,
)[0]
return grad_input, grad_weight, None return grad_input, grad_weight, None
@@ -120,7 +133,10 @@ class HybridGmmFunction(torch.autograd.Function):
ctx.split_sizes = split_sizes ctx.split_sizes = split_sizes
ctx.num_experts = num_experts ctx.num_experts = num_experts
ctx.save_for_backward(*args) ctx.save_for_backward(
*(x if ctx.needs_input_grad[1 + num_experts + i] else None for i, x in enumerate(x_list)),
*(w if ctx.needs_input_grad[1 + i] else None for i, w in enumerate(weight_list)),
)
outputs = torch_npu.npu_grouped_matmul( outputs = torch_npu.npu_grouped_matmul(
x_list, weight_list, bias=None, group_list=None, split_item=0, group_type=-1 x_list, weight_list, bias=None, group_list=None, split_item=0, group_type=-1
@@ -147,35 +163,42 @@ class HybridGmmFunction(torch.autograd.Function):
grad_outputs_contiguous = [g.contiguous() for g in grad_outputs] grad_outputs_contiguous = [g.contiguous() for g in grad_outputs]
w_t_list = [w.t() for w in weight_list] grad_x_list = [None] * num_experts
grad_x_list = torch_npu.npu_grouped_matmul( input_indices = [i for i in range(num_experts) if ctx.needs_input_grad[1 + i]]
grad_outputs_contiguous, # List[Tensor], 每个 [M_i, N] if input_indices:
w_t_list, # List[Tensor], 每个 [N, K] (view) grad_inputs = torch_npu.npu_grouped_matmul(
bias=None, [grad_outputs_contiguous[i] for i in input_indices],
group_list=None, [weight_list[i].t() for i in input_indices],
split_item=0, bias=None,
group_type=-1, group_list=None,
) split_item=0,
group_type=-1,
)
for i, grad in zip(input_indices, grad_inputs, strict=True):
grad_x_list[i] = grad
x_concat = torch.cat(x_list, dim=0) grad_w_list = [None] * num_experts
dy_concat = torch.cat(grad_outputs_contiguous, dim=0) # [Total_M, N] weight_indices = [i for i in range(num_experts) if ctx.needs_input_grad[1 + num_experts + i]]
if weight_indices:
x_concat = torch.cat([x_list[i] for i in weight_indices], dim=0)
dy_concat = torch.cat([grad_outputs_contiguous[i] for i in weight_indices], dim=0)
group_list = torch.tensor(
[split_sizes[i] for i in weight_indices], device=x_concat.device, dtype=torch.int64
)
grad_w_stack = torch_npu.npu_grouped_matmul(
[x_concat.t()],
[dy_concat],
bias=None,
group_list=group_list,
split_item=3,
group_type=2,
group_list_type=1,
)[0]
group_list = torch.tensor(split_sizes, device=x_concat.device, dtype=torch.int64) if grad_w_stack.dim() != 3:
raise RuntimeError(f"Unexpected grad_w_stack shape: {grad_w_stack.shape}")
grad_w_stack = torch_npu.npu_grouped_matmul( for i, grad in zip(weight_indices, grad_w_stack.unbind(0), strict=True):
[x_concat.t()], grad_w_list[i] = grad
[dy_concat],
bias=None,
group_list=group_list,
split_item=3,
group_type=2,
group_list_type=1,
)[0]
if grad_w_stack.dim() == 3:
grad_w_list = list(torch.unbind(grad_w_stack, dim=0))
else:
raise RuntimeError(f"Unexpected grad_w_stack shape: {grad_w_stack.shape}")
return (None, *grad_x_list, *grad_w_list) return (None, *grad_x_list, *grad_w_list)