Skip to content

GitLab

Unverified Commit 5a495a39 authored by Przemyslaw Tredak's avatar Przemyslaw Tredak Committed by GitHub
Browse files

Fix the use-after-free bug in unfused normalization (#2002) 


Signed-off-by: default avatarPrzemek Tredak <ptredak@nvidia.com>
parent aac74427
Show whitespace changes
Inline Side-by-side
Showing with 2 additions and 2 deletions
transformer_engine/pytorch/csrc/extensions/normalization.cpp
View file @ 5a495a39
...@@ -108,9 +108,9 @@ std::vector<py::object> layernorm_fwd(py::handle input, py::handle weight, Maybe ...@@ -108,9 +108,9 @@ std::vector<py::object> layernorm_fwd(py::handle input, py::handle weight, Maybe
} }
} }
TensorWrapper unquantized_out_cu; TensorWrapper unquantized_out_cu;
py::object unquantized_out;
if (force_unfused_kernel) { if (force_unfused_kernel) {
NoneQuantizer q{none}; NoneQuantizer q{none};
py::object unquantized_out;
std::tie(unquantized_out_cu, unquantized_out) = q.create_tensor(size, out_dtype); std::tie(unquantized_out_cu, unquantized_out) = q.create_tensor(size, out_dtype);
} }
TensorWrapper &kernel_out_cu = force_unfused_kernel ? unquantized_out_cu : out_cu; TensorWrapper &kernel_out_cu = force_unfused_kernel ? unquantized_out_cu : out_cu;
...@@ -269,9 +269,9 @@ std::vector<py::object> rmsnorm_fwd(const py::handle &input, const py::handle &w ...@@ -269,9 +269,9 @@ std::vector<py::object> rmsnorm_fwd(const py::handle &input, const py::handle &w
} }
} }
TensorWrapper unquantized_out_cu; TensorWrapper unquantized_out_cu;
py::object unquantized_out;
if (force_unfused_kernel) { if (force_unfused_kernel) {
NoneQuantizer q{none}; NoneQuantizer q{none};
py::object unquantized_out;
std::tie(unquantized_out_cu, unquantized_out) = q.create_tensor(size, out_dtype); std::tie(unquantized_out_cu, unquantized_out) = q.create_tensor(size, out_dtype);
} }
TensorWrapper &kernel_out_cu = force_unfused_kernel ? unquantized_out_cu : out_cu; TensorWrapper &kernel_out_cu = force_unfused_kernel ? unquantized_out_cu : out_cu;
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment