Star973
1910
代码
代码
Issues260
Pull Requests351
Actions
项目讨论
Wiki
分析
Star973
1910
已合并
fix: 放宽 swiglu_group_grad 的 weight 形态校验为 element num #8541
sinobr创建于 8月11日
fix: 放宽 swiglu_group_grad 的 weight 形态校验为 element num #8541
已合并
sinobr创建于 8月11日
从qq_51545127/ops-nn_9782fix/swiglu-group-grad-weight-numel-check合入到cann/ops-nnmaster
共 10 个文件变更+59-92
Mquant/swiglu_group_grad/op_api/aclnn_swiglu_group_grad.cpp+11-12
@@ -19,16 +19,20 @@
19#include "opdev/make_op_executor.h"19#include "opdev/make_op_executor.h"
20#include "opdev/op_dfx.h"20#include "opdev/op_dfx.h"
21#include "opdev/op_log.h"21#include "opdev/op_log.h"
22+#include "log/log.h"
22#include "opdev/shape_utils.h"23#include "opdev/shape_utils.h"
23#include "opdev/tensor_view_utils.h"24#include "opdev/tensor_view_utils.h"
24#include "opdev/platform.h"25#include "opdev/platform.h"
25#include "op_api/aclnn_util.h"26#include "op_api/aclnn_util.h"
27+#include <string>
26 28 
27using namespace op;29using namespace op;
28#ifdef __cplusplus30#ifdef __cplusplus
29extern "C" {31extern "C" {
30#endif32#endif
31 33 
34+static constexpr const char* ACLNN_SWIGLU_GROUP_GRAD_NAME = "aclnnSwigluGroupGrad";
35+ 
32// ── Supported dtype list ───────────────────────────────────────────────────36// ── Supported dtype list ───────────────────────────────────────────────────
33static const std::initializer_list<DataType> DTYPE_SUPPORT_LIST = {DataType::DT_FLOAT16, DataType::DT_FLOAT,37static const std::initializer_list<DataType> DTYPE_SUPPORT_LIST = {DataType::DT_FLOAT16, DataType::DT_FLOAT,
34 DataType::DT_BF16};38 DataType::DT_BF16};
@@ -113,20 +117,15 @@ static inline bool CheckShape(const aclTensor* gradY, const aclTensor* x, const
113 117 
114 if (weightOptional != nullptr) {118 if (weightOptional != nullptr) {
115 auto weightShape = weightOptional->GetViewShape();119 auto weightShape = weightOptional->GetViewShape();
116- if (weightShape.GetDimNum() != inputRank) {120+ int64_t totalRows = 1;
117- OP_LOGE(ACLNN_ERR_PARAM_INVALID, "weightOptional rank(%zu) must equal gradY rank(%zu).",
118- weightShape.GetDimNum(), inputRank);
119- return false;
120- }
121 for (size_t i = 0; i < lastDim; ++i) {121 for (size_t i = 0; i < lastDim; ++i) {
122- if (weightShape.GetDim(i) != gradYShape.GetDim(i)) {122+ totalRows *= gradYShape.GetDim(i);
123- OP_LOGE(ACLNN_ERR_PARAM_INVALID, "weightOptional.shape[%zu]=%ld != gradY.shape[%zu]=%ld", i,
124- weightShape.GetDim(i), i, gradYShape.GetDim(i));
125- return false;
126- }
127 }123 }
128- if (weightShape.GetDim(lastDim) != 1) {124+ int64_t weightElementNum = weightShape.GetShapeSize();
129- OP_LOGE(ACLNN_ERR_PARAM_INVALID, "weightOptional.shape[-1]=%ld != 1", weightShape.GetDim(lastDim));125+ if (weightElementNum != totalRows) {
126+ OP_LOGE_FOR_INVALID_SHAPESIZE_WITH_REASON(
127+ ACLNN_SWIGLU_GROUP_GRAD_NAME, "weightOptional", std::to_string(weightElementNum).c_str(),
128+ "The element num of weightOptional must be equal to the product of gradY leading dims.");
130 return false;129 return false;
131 }130 }
132 OP_CHECK_SHAPE_NOT_EQUAL(gradWeightOutOptional, weightOptional, return false);131 OP_CHECK_SHAPE_NOT_EQUAL(gradWeightOutOptional, weightOptional, return false);
Mquant/swiglu_group_grad/op_host/arch35/swiglu_group_grad_regbase_tiling.cpp+6-18
@@ -168,24 +168,12 @@ ge::graphStatus SwigluGroupGradArch35Tiling::ParseOptionalInputs()
168 168 
169 if (isWeight_ == 1) {169 if (isWeight_ == 1) {
170 const gert::Shape& weightShape = weightStorageShape->GetStorageShape();170 const gert::Shape& weightShape = weightStorageShape->GetStorageShape();
171- const gert::Shape& gradYShape = tilingContext->GetInputShape(GRAD_Y_INDEX)->GetStorageShape();171+ auto weightElementNum = weightShape.GetShapeSize();
172- OP_CHECK_IF(weightShape.GetDimNum() != gradYShape.GetDimNum(),172+ if (weightElementNum != totalRows_) {
173- OP_LOGE(tilingContext->GetNodeName(), "weight dims=%ld must match grad_y dims=%ld",173+ OP_LOGE_FOR_INVALID_SHAPESIZE_WITH_REASON(
174- weightShape.GetDimNum(), gradYShape.GetDimNum()),174+ tilingContext->GetNodeName(), "weight", std::to_string(weightElementNum).c_str(),
175- return ge::GRAPH_FAILED);175+ "The element num of weight must be equal to the product of grad_y leading dims.");
176- OP_CHECK_IF(weightShape.GetDim(weightShape.GetDimNum() - 1) != 1,176+ return ge::GRAPH_FAILED;
177- OP_LOGE(tilingContext->GetNodeName(), "weight.shape[-1]=%ld must be 1",
178- weightShape.GetDim(weightShape.GetDimNum() - 1)),
179- return ge::GRAPH_FAILED);
180- if (gradYShape.GetDimNum() == DIM_THREE) {
181- OP_CHECK_IF(weightShape.GetDim(0) != gradYShape.GetDim(0) || weightShape.GetDim(1) != gradYShape.GetDim(1),
182- OP_LOGE(tilingContext->GetNodeName(), "weight [B,S] must match grad_y [B,S]"),
183- return ge::GRAPH_FAILED);
184- } else {
185- OP_CHECK_IF(weightShape.GetDim(0) != totalRows_,
186- OP_LOGE(tilingContext->GetNodeName(), "weight.shape[0]=%ld != totalRows=%ld",
187- weightShape.GetDim(0), totalRows_),
188- return ge::GRAPH_FAILED);
189 }177 }
190 auto gradWeightOut = tilingContext->GetOutputShape(1);178 auto gradWeightOut = tilingContext->GetOutputShape(1);
191 OP_CHECK_IF(gradWeightOut == nullptr,179 OP_CHECK_IF(gradWeightOut == nullptr,
Mquant/swiglu_group_grad/op_host/arch35/swiglu_group_grad_tiling_base.cpp+6-17
@@ -190,23 +190,12 @@ static ge::graphStatus ParseOptionalInputs(gert::TilingContext* context, SwigluG
190 OP_LOGE(context->GetNodeName(), "y_origin rows mismatch"), return ge::GRAPH_FAILED);190 OP_LOGE(context->GetNodeName(), "y_origin rows mismatch"), return ge::GRAPH_FAILED);
191 }191 }
192 192 
193- OP_CHECK_IF(weightShape->GetDimNum() != gYShape.GetDimNum(),193+ auto weightElementNum = weightShape->GetShapeSize();
194- OP_LOGE(context->GetNodeName(), "weight dims=%ld must match grad_y dims=%ld",194+ if (weightElementNum != inputData.totalRows) {
195- weightShape->GetDimNum(), gYShape.GetDimNum()),195+ OP_LOGE_FOR_INVALID_SHAPESIZE_WITH_REASON(
196- return ge::GRAPH_FAILED);196+ context->GetNodeName(), "weight", std::to_string(weightElementNum).c_str(),
197- OP_CHECK_IF(weightShape->GetDim(weightShape->GetDimNum() - 1) != 1,197+ "The element num of weight must be equal to the product of grad_y leading dims.");
198- OP_LOGE(context->GetNodeName(), "weight.shape[-1]=%ld must be 1",198+ return ge::GRAPH_FAILED;
199- weightShape->GetDim(weightShape->GetDimNum() - 1)),
200- return ge::GRAPH_FAILED);
201- if (gYShape.GetDimNum() == DIM_THREE) {
202- OP_CHECK_IF(weightShape->GetDim(0) != gYShape.GetDim(0) || weightShape->GetDim(1) != gYShape.GetDim(1),
203- OP_LOGE(context->GetNodeName(), "weight [B,S] must match grad_y [B,S]"),
204- return ge::GRAPH_FAILED);
205- } else {
206- OP_CHECK_IF(weightShape->GetDim(0) != inputData.totalRows,
207- OP_LOGE(context->GetNodeName(), "weight.shape[0]=%ld != totalRows=%ld", weightShape->GetDim(0),
208- inputData.totalRows),
209- return ge::GRAPH_FAILED);
210 }199 }
211 }200 }
212 201 
Mquant/swiglu_group_grad/op_host/swiglu_group_grad_infershape.cpp+7-23
@@ -168,31 +168,15 @@ static ge::graphStatus InferShapeForSwigluGroupGrad(gert::InferShapeContext* con
168 const bool y_origin_unknown_rank = Ops::Base::IsUnknownRank(*y_origin_shape);168 const bool y_origin_unknown_rank = Ops::Base::IsUnknownRank(*y_origin_shape);
169 169 
L
Lliuchuangdev8月12日

保留必要校验

likedislike
170 if (!weight_unknown_rank) {170 if (!weight_unknown_rank) {
171- const size_t weight_rank = weight_shape->GetDimNum();171+ const int64_t weightElementNum = weight_shape->GetShapeSize();
172- if (weight_rank != grad_y_rank) {172+ int64_t totalRows = 1;
173- OP_LOGE(context->GetNodeName(),
174- "Invalid weight rank: weight rank(%zu) must be equal to "
175- "grad_y rank(%zu).",
176- weight_rank, grad_y_rank);
177- return ge::GRAPH_FAILED;
178- }
179- 
180 for (size_t i = 0; i + 1U < grad_y_rank; ++i) {173 for (size_t i = 0; i + 1U < grad_y_rank; ++i) {
181- const int64_t weight_dim = weight_shape->GetDim(i);174+ totalRows *= grad_y_shape->GetDim(i);
182- const int64_t grad_y_dim = grad_y_shape->GetDim(i);
183- if (!IsCompatibleDim(weight_dim, grad_y_dim)) {
184- OP_LOGE(context->GetNodeName(),
185- "Shape mismatch: weight.shape[%zu](%lld) must be equal to "
186- "grad_y.shape[%zu](%lld).",
187- i, static_cast<long long>(weight_dim), i, static_cast<long long>(grad_y_dim));
188- return ge::GRAPH_FAILED;
189- }
190 }175 }
191- 176+ if (weightElementNum != totalRows) {
192- const int64_t weight_last_dim = weight_shape->GetDim(last_dim_index);177+ OP_LOGE_FOR_INVALID_SHAPESIZE_WITH_REASON(
193- if (!IsUnknownDim(weight_last_dim) && weight_last_dim != 1) {178+ context->GetNodeName(), "weight", std::to_string(weightElementNum).c_str(),
194- OP_LOGE(context->GetNodeName(), "Invalid weight last dimension: weight.shape[-1](%lld) must be 1.",179+ "The element num of weight must be equal to the product of grad_y leading dims.");
195- static_cast<long long>(weight_last_dim));
196 return ge::GRAPH_FAILED;180 return ge::GRAPH_FAILED;
197 }181 }
198 }182 }
Mquant/swiglu_group_grad/tests/ut/op_host/arch35/test_swiglu_group_grad_tiling.cpp+7-2
@@ -37,6 +37,7 @@ struct TilingCase {
37 bool hasWeight = false;37 bool hasWeight = false;
38 bool hasYOrigin = false;38 bool hasYOrigin = false;
39 bool hasGroupIndex = false;39 bool hasGroupIndex = false;
40+ gert::Shape weightShape = {};
40 ge::graphStatus expectedStatus = ge::GRAPH_SUCCESS;41 ge::graphStatus expectedStatus = ge::GRAPH_SUCCESS;
41};42};
42 43 
@@ -103,8 +104,11 @@ void ExecuteTilingCase(const TilingCase& testCase)
103 xStorageShape.MutableStorageShape() = testCase.xShape;104 xStorageShape.MutableStorageShape() = testCase.xShape;
104 xStorageShape.MutableOriginShape() = xStorageShape.MutableStorageShape();105 xStorageShape.MutableOriginShape() = xStorageShape.MutableStorageShape();
105 gert::StorageShape weightStorageShape;106 gert::StorageShape weightStorageShape;
106- weightStorageShape.MutableStorageShape() = testCase.gradYShape;107+ weightStorageShape.MutableStorageShape() = testCase.weightShape;
107- weightStorageShape.MutableStorageShape().SetDim(weightStorageShape.MutableStorageShape().GetDimNum() - 1, 1);108+ if (weightStorageShape.MutableStorageShape().GetDimNum() == 0) {
109+ weightStorageShape.MutableStorageShape() = testCase.gradYShape;
110+ weightStorageShape.MutableStorageShape().SetDim(weightStorageShape.MutableStorageShape().GetDimNum() - 1, 1);
111+ }
108 weightStorageShape.MutableOriginShape() = weightStorageShape.MutableStorageShape();112 weightStorageShape.MutableOriginShape() = weightStorageShape.MutableStorageShape();
109 gert::StorageShape yOriginStorageShape;113 gert::StorageShape yOriginStorageShape;
110 yOriginStorageShape.MutableStorageShape() = testCase.gradYShape;114 yOriginStorageShape.MutableStorageShape() = testCase.gradYShape;
@@ -268,6 +272,7 @@ TEST_F(SwigluGroupGradTilingTest, tiling_accepts_3d_input_with_all_optional_inpu
268 testCase.hasWeight = true;272 testCase.hasWeight = true;
269 testCase.hasYOrigin = true;273 testCase.hasYOrigin = true;
270 testCase.hasGroupIndex = true;274 testCase.hasGroupIndex = true;
275+ testCase.weightShape = {8};
271 ExecuteTilingCase(testCase);276 ExecuteTilingCase(testCase);
272}277}
273 278 
Mquant/swiglu_group_grad/tests/ut/op_host/op_api/test_aclnn_swiglu_group_grad.cpp+2-2
@@ -78,10 +78,10 @@ TEST_F(l2_swiglu_group_grad_test, l2_normal_FLOAT_ND_with_topk_weight)
78{78{
79 auto dyDesc = TensorDesc({4, 16}, ACL_FLOAT, ACL_FORMAT_ND);79 auto dyDesc = TensorDesc({4, 16}, ACL_FLOAT, ACL_FORMAT_ND);
80 auto xDesc = TensorDesc({4, 32}, ACL_FLOAT, ACL_FORMAT_ND);80 auto xDesc = TensorDesc({4, 32}, ACL_FLOAT, ACL_FORMAT_ND);
81- auto weightDesc = TensorDesc({4, 1}, ACL_FLOAT, ACL_FORMAT_ND);81+ auto weightDesc = TensorDesc({4}, ACL_FLOAT, ACL_FORMAT_ND);
82 auto yOriginDesc = TensorDesc({4, 16}, ACL_FLOAT, ACL_FORMAT_ND);82 auto yOriginDesc = TensorDesc({4, 16}, ACL_FLOAT, ACL_FORMAT_ND);
83 auto dxOutDesc = TensorDesc({4, 32}, ACL_FLOAT, ACL_FORMAT_ND);83 auto dxOutDesc = TensorDesc({4, 32}, ACL_FLOAT, ACL_FORMAT_ND);
84- auto dWeightDesc = TensorDesc({4, 1}, ACL_FLOAT, ACL_FORMAT_ND);84+ auto dWeightDesc = TensorDesc({4}, ACL_FLOAT, ACL_FORMAT_ND);
85 auto ut = OP_API_UT(aclnnSwigluGroupGrad, INPUT(dyDesc, xDesc, weightDesc, yOriginDesc, nullptr, 0.0f),85 auto ut = OP_API_UT(aclnnSwigluGroupGrad, INPUT(dyDesc, xDesc, weightDesc, yOriginDesc, nullptr, 0.0f),
86 OUTPUT(dxOutDesc, dWeightDesc));86 OUTPUT(dxOutDesc, dWeightDesc));
87 uint64_t workspaceSize = 0;87 uint64_t workspaceSize = 0;
Mquant/swiglu_group_grad/tests/ut/op_host/test_swiglu_group_grad_infershape.cpp+7-2
@@ -29,6 +29,7 @@ struct InferShapeCase {
29 bool hasWeight = false;29 bool hasWeight = false;
30 bool hasYOrigin = false;30 bool hasYOrigin = false;
31 bool hasGroupIndex = false;31 bool hasGroupIndex = false;
32+ gert::Shape weightShape = {};
32 gert::Shape groupIndexShape = {2};33 gert::Shape groupIndexShape = {2};
33 ge::graphStatus expectedStatus = ge::GRAPH_SUCCESS;34 ge::graphStatus expectedStatus = ge::GRAPH_SUCCESS;
34};35};
@@ -54,8 +55,11 @@ void ExecuteInferShapeCase(const InferShapeCase& testCase)
54 55 
55 gert::Shape gradYShape = testCase.gradYShape;56 gert::Shape gradYShape = testCase.gradYShape;
56 gert::Shape xShape = testCase.xShape;57 gert::Shape xShape = testCase.xShape;
57- gert::Shape weightShape = gradYShape;58+ gert::Shape weightShape = testCase.weightShape;
58- weightShape.SetDim(weightShape.GetDimNum() - 1, 1);59+ if (weightShape.GetDimNum() == 0) {
60+ weightShape = gradYShape;
61+ weightShape.SetDim(weightShape.GetDimNum() - 1, 1);
62+ }
59 gert::Shape yOriginShape = gradYShape;63 gert::Shape yOriginShape = gradYShape;
60 gert::Shape groupIndexShape = testCase.groupIndexShape;64 gert::Shape groupIndexShape = testCase.groupIndexShape;
61 gert::Shape gradXShape = {};65 gert::Shape gradXShape = {};
@@ -145,6 +149,7 @@ TEST_F(SwigluGroupGradInferShapeTest, infershape_all_optional_inputs)
145 testCase.hasWeight = true;149 testCase.hasWeight = true;
146 testCase.hasYOrigin = true;150 testCase.hasYOrigin = true;
147 testCase.hasGroupIndex = true;151 testCase.hasGroupIndex = true;
152+ testCase.weightShape = {4};
148 ExecuteInferShapeCase(testCase);153 ExecuteInferShapeCase(testCase);
149}154}
150 155 
Mquant/swiglu_group_grad/tests/ut/torch_extension/test_swiglu_group_backward.py+1-1
@@ -126,7 +126,7 @@ def test_meta_accepts_3d_input_with_all_optional_inputs():
126 126 
127 grad_output = torch.empty((2, 4, 16), dtype=torch.float32, device="meta")127 grad_output = torch.empty((2, 4, 16), dtype=torch.float32, device="meta")
128 x = torch.empty((2, 4, 32), dtype=torch.float32, device="meta")128 x = torch.empty((2, 4, 32), dtype=torch.float32, device="meta")
129- weight = torch.empty((2, 4, 1), dtype=torch.float32, device="meta")129+ weight = torch.empty((8,), dtype=torch.float32, device="meta")
130 y_origin = torch.empty((2, 4, 16), dtype=torch.float32, device="meta")130 y_origin = torch.empty((2, 4, 16), dtype=torch.float32, device="meta")
131 group_index = torch.empty((2,), dtype=torch.int64, device="meta")131 group_index = torch.empty((2,), dtype=torch.int64, device="meta")
132 132 
Mquant/swiglu_group_grad/torch_extension/csrc/swiglu_group_grad.cpp+6-5
@@ -71,12 +71,13 @@ std::tuple<at::Tensor, c10::optional<at::Tensor>> swiglu_group_backward(
71 at::Tensor grad_weight = at::empty({0}, grad_output.options().dtype(at::kFloat));71 at::Tensor grad_weight = at::empty({0}, grad_output.options().dtype(at::kFloat));
72 72 
73 if (has_weight) {73 if (has_weight) {
74- TORCH_CHECK(weight.value().dim() == grad_output.dim(), "weight rank must equal grad_output rank");74+ int64_t weightElementNum = weight.value().numel();
75- for (int64_t dim = 0; dim < grad_output.dim() - 1; ++dim) {75+ int64_t totalRows = 1;
76- TORCH_CHECK(weight.value().size(dim) == grad_output.size(dim), "weight.shape[", dim,76+ for (int64_t i = 0; i < grad_output.dim() - 1; ++i) {
77- "] must equal grad_output.shape[", dim, "]");77+ totalRows *= grad_output.size(i);
78 }78 }
79- TORCH_CHECK(weight.value().size(-1) == 1, "weight.shape[-1] must be 1");79+ TORCH_CHECK(weightElementNum == totalRows, "weight element num must equal total rows (", totalRows,
80+ "), but got ", weightElementNum);
80 TORCH_CHECK(weight.value().scalar_type() == at::kFloat, "weight dtype must be FLOAT");81 TORCH_CHECK(weight.value().scalar_type() == at::kFloat, "weight dtype must be FLOAT");
81 grad_weight = at::empty(weight.value().sizes(), weight.value().options().dtype(at::kFloat));82 grad_weight = at::empty(weight.value().sizes(), weight.value().options().dtype(at::kFloat));
82 }83 }
Mquant/swiglu_group_grad/torch_extension/swiglu_group_grad.py+6-10
@@ -54,18 +54,14 @@ def _check_swiglu_group_backward_inputs(
54 if (weight is None) != (y_origin is None):54 if (weight is None) != (y_origin is None):
55 raise RuntimeError("weight and y_origin must be provided together")55 raise RuntimeError("weight and y_origin must be provided together")
56 if weight is not None:56 if weight is not None:
57+ weight_element_num = weight.numel()
58+ total_rows = 1
59+ for s in grad_output.shape[:-1]:
60+ total_rows *= s
57 torch._check(61 torch._check(
58- weight.dim() == grad_output.dim(),62+ weight_element_num == total_rows,
59- lambda: "weight rank must equal grad_output rank",63+ lambda: f"weight element num must equal total rows ({total_rows}), but got {weight_element_num}",
60 )64 )
61- for dim in range(grad_output.dim() - 1):
62- torch._check(
63- weight.shape[dim] == grad_output.shape[dim],
64- lambda dim=dim: (
65- f"weight.shape[{dim}] must equal grad_output.shape[{dim}]"
66- ),
67- )
68- torch._check(weight.shape[-1] == 1, lambda: "weight.shape[-1] must be 1")
69 torch._check(65 torch._check(
70 weight.dtype == torch.float32, lambda: "weight dtype must be FLOAT"66 weight.dtype == torch.float32, lambda: "weight dtype must be FLOAT"
71 )67 )