Skip to content

[NPU] align fused linear JSD forward with accum_dtype API #1346

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
zheliuyu wants to merge 1 commit into
base: main
Choose a base branch
from
Open

[NPU] align fused linear JSD forward with accum_dtype API #1346

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 17 additions & 3 deletions src/liger_kernel/ops/backends/_ascend/ops/fused_linear_jsd.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,7 @@ def fused_linear_jsd_forward(
ignore_index,
has_label,
temperature,
accum_dtype=None,
):
device = student_input.device
dtype = student_input.dtype
Expand All @@ -63,7 +64,12 @@ def fused_linear_jsd_forward(
chunk_size = triton.next_power_of_2(triton.cdiv(BT, inc_factor)) # (BT + inc_factor - 1) // inc_factor
num_chunks = triton.cdiv(BT, chunk_size) # (BT + chunk_size - 1) // chunk_size

grad_weight = torch.zeros_like(student_weight, device=device) if student_weight.requires_grad else None
if accum_dtype is None:
grad_weight = torch.zeros_like(student_weight, device=device) if student_weight.requires_grad else None
else:
grad_weight = (
torch.zeros_like(student_weight, dtype=accum_dtype, device=device) if student_weight.requires_grad else None
)
grad_input = torch.zeros_like(student_input)
# we use fp32 for loss accumulator
loss_1d = torch.zeros((BT, V), dtype=torch.float32, device=device)
Expand Down Expand Up @@ -139,9 +145,15 @@ def fused_linear_jsd_forward(
grad_input[start_idx:end_idx] = student_logits_chunk @ student_weight

if grad_weight is not None:
grad_weight.add_(student_logits_chunk.t() @ student_input_chunk)
if accum_dtype is None:
grad_weight.add_(student_logits_chunk.t() @ student_input_chunk)
else:
grad_weight.add_(torch.mm(student_logits_chunk.t(), student_input_chunk).float())

loss = torch.sum(loss_1d)
grad_weight = (
grad_weight.to(student_weight.dtype) if grad_weight is not None and accum_dtype is not None else grad_weight
)
return loss, grad_input, grad_weight


Expand Down Expand Up @@ -202,6 +214,7 @@ def forward(
jsd_beta: float = 0.5,
ignore_index: int = -100,
temperature: float = 1.0,
accum_dtype: Optional[torch.dtype] = None,
):
"""
Args:
Expand Down Expand Up @@ -236,6 +249,7 @@ def forward(
ignore_index,
has_label,
temperature,
accum_dtype,
)
# downcast to dtype and store for backward
ctx.save_for_backward(
Expand All @@ -249,4 +263,4 @@ def forward(
def backward(ctx, grad_output):
(grad_input, grad_weight) = ctx.saved_tensors
grad_input, grad_weight = fused_linear_jsd_backward(grad_output, grad_input, grad_weight)
return (grad_input, grad_weight, None, None, None, None, None, None)
return (grad_input, grad_weight, None, None, None, None, None, None, None)
Loading