[NPU]:Added support for the layer_norm operator in npu - #1113
Conversation
8c96875 to
0ca405e
Compare
|
@Tcc0403 would you mind having a preview? |
| def _numerically_stable_sum(tensor, dim=0): | ||
| if tensor.shape[dim] <= 1: | ||
| return tensor.squeeze(dim) | ||
|
|
||
| if tensor.shape[dim] <= 8: | ||
| return tensor.sum(dim=dim, dtype=torch.float32) | ||
|
|
||
| num_programs = tensor.shape[dim] | ||
| num_groups = min(4, (num_programs + 3) // 4) | ||
|
|
||
| if num_groups > 1: | ||
| group_size = (num_programs + num_groups - 1) // num_groups | ||
| groups = [] | ||
|
|
||
| for i in range(num_groups): | ||
| start_idx = i * group_size | ||
| end_idx = min((i + 1) * group_size, num_programs) | ||
| if start_idx < end_idx: | ||
| group_sum = tensor[start_idx:end_idx].sum(dim=dim, dtype=torch.float32) | ||
| groups.append(group_sum) | ||
|
|
||
| if len(groups) > 0: | ||
| return torch.stack(groups).sum(dim=0, dtype=torch.float32) |
There was a problem hiding this comment.
Is it hard to pass correctness tests without this helper function?
There was a problem hiding this comment.
When the NPU device is performing accumulation, if there are too many accumulation items, errors may occur.
There was a problem hiding this comment.
I should've asked earlier. Could you provide detailed test cases it wasn't able to pass with using normal torch.sum only?
There was a problem hiding this comment.
Sorry, when I was trying to reproduce this issue just now, I found that there was no such error problem anymore. It might be caused by my previous incorrect operation. Currently, I have submitted a new modification to remove this part. Please have a look.
0ab972b to
992b998
Compare
|
@Tcc0403 Hello, are you a bit busy recently? May I ask if you have time to take a look at this PR as well as the others I submitted recently? |
Sure, I'll review your recent PRs. If you have any new PRs or get any updates on existing ones, feel free to ping me. I wasn't checking all PRs regularly, so I might miss your updates. |
OK. Thank you for your explanation. |
|
|
||
| X_block = tl.load(X_row_ptr + col_offsets, mask=mask, other=0.0, cache_modifier=".cg").to(tl.float32) | ||
|
|
||
| row_sum += tl.sum(tl.where(mask, X_block, 0.0)) |
There was a problem hiding this comment.
We don't have to do tl.where since out of bound values are masked to 0.0 in the previous tl.load
| _layer_norm_backward_kernel_no_tiling[(grid_size,)]( | ||
| X, | ||
| X.stride(0), | ||
| X.dtype if hasattr(X.dtype, "name") else tl.float32, |
There was a problem hiding this comment.
In what scenario, would X.dtype have attribute name?
|
|
||
| # Store dX with coalesced memory access | ||
| DX_block_ptr = DX_ptr + row_idx[:, None] * DX_row_stride + col_offsets[None, :] | ||
| tl.store(DX_block_ptr, DX_f32.to(X_dtype), mask=block_mask) |
There was a problem hiding this comment.
I'm not sure about triton code would be compiled for npu device, but tl.store should implicitly typecast DX_f32 to DX.dtype
value is implicitly broadcast to pointer.shape and typecast to pointer.dtype.element_ty.
Ref: https://triton-lang.org/main/python-api/generated/triton.language.store.html
| def _layer_norm_backward_kernel_no_tiling( | ||
| X_ptr, | ||
| X_row_stride, | ||
| X_dtype: tl.constexpr, |
There was a problem hiding this comment.
You can also retrieve dtype by X_ptr.dtype.element_ty in triton code, so that we don't have to pass X_dtype in host side.
There was a problem hiding this comment.
Thank you for your explanation. I didn't know this could be done before.
|
@Tcc0403 I have fixed the issues related to masks and type conversions. Are there any other questions? |
2 participants
Summary
Added NPU-friendly operator implementations. Previously, only some unsupported writing methods were addressed based on the GPU implementation. However, when the n_col is large, the problem of UB overflow still occurs. This modification has solved this issue, and the performance has improved by 2-3 times compared to the original implementation. However, due to device limitations, the performance is still somewhat lower than that of Hugging Face. We will continue to follow up on the performance issues in the future.
Testing Done
make testto ensure correctnessmake checkstyleto ensure code stylemake test-convergenceto ensure convergence