Skip to content
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

Jump to bottom

[Bugfix] Fix bias for 0-dim tensors in gemm #1246

Merged
merged 5 commits into from
Oct 17, 2024
Merged

[Bugfix] Fix bias for 0-dim tensors in gemm #1246

Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
3fe8543
fix bias for 0-dim tensor
yaox12 Oct 12, 2024
8035717
add check
yaox12 Oct 12, 2024
b5fa1c3
use numel() instead of nullptr
yaox12 Oct 15, 2024
55cb727
Merge branch 'main' into xiny/fix_bias
yaox12 Oct 16, 2024
256b7ff
Merge branch 'main' into xiny/fix_bias
yaox12 Oct 17, 2024
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
30 changes: 22 additions & 8 deletions transformer_engine/pytorch/csrc/extensions/gemm.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,10 +15,16 @@ void te_gemm(at::Tensor A, at::Tensor A_scale_inverse, transformer_engine::DType
at::Tensor workspace, size_t workspaceSize, bool accumulate,
bool use_split_accumulator, int math_sm_count) {
using namespace transformer_engine;
if (A.data_ptr() == nullptr || B.data_ptr() == nullptr) {
if (D.data_ptr() != nullptr && !accumulate) D.zero_();
if (bias.data_ptr() != nullptr) bias.zero_();
if (pre_gelu_out.data_ptr() != nullptr) pre_gelu_out.zero_();
if (A.numel() == 0 || B.numel() == 0) {
if (D.numel() != 0 && !accumulate) D.zero_();
if (bias.numel() != 0 && grad) {
if (B.numel() == 0) {
bias.zero_();
} else {
bias.copy_(B.sum(0));
}
}
if (pre_gelu_out.numel() != 0) pre_gelu_out.zero_();
return;
}

Expand Down Expand Up @@ -109,10 +115,16 @@ void te_grouped_gemm(std::vector<at::Tensor> A, at::Tensor A_scale_inverse, int
return tensor_wrappers.back().data();
};
for (size_t i = 0; i < A.size(); i++) {
if (A[i].data_ptr() == nullptr || B[i].data_ptr() == nullptr) {
if (D[i].data_ptr() != nullptr && !accumulate) D[i].zero_();
if (bias[i].data_ptr() != nullptr) bias[i].zero_();
if (pre_gelu_out[i].data_ptr() != nullptr) pre_gelu_out[i].zero_();
if (A[i].numel() == 0 || B[i].numel() == 0) {
if (D[i].numel() != 0 && !accumulate) D[i].zero_();
if (bias[i].numel() != 0 && grad) {
if (B[i].numel() == 0) {
bias[i].zero_();
} else {
bias[i].copy_(B[i].sum(0));
}
}
if (pre_gelu_out[i].numel() != 0) pre_gelu_out[i].zero_();
continue;
}

Expand Down Expand Up @@ -175,6 +187,8 @@ void te_grouped_gemm_single_output(
void* d_i_ptr = reinterpret_cast<void*>(D.data_ptr());
for (size_t i = 0; i < A.size(); i++) {
if (m_splits[i] == 0) continue;
NVTE_CHECK(A[i].data_ptr() != nullptr, "A[", i, "] must not be nullptr.");
NVTE_CHECK(B[i].data_ptr() != nullptr, "B[", i, "] must not be nullptr.");
NVTE_CHECK(A[i].is_contiguous(), "A[", i, "] must be contiguous.");
NVTE_CHECK(B[i].is_contiguous(), "B[", i, "] must be contiguous.");
te_A.emplace_back(make_tensor(
Expand Down
Loading