已合并
Add fused_linear_crocess_ectropy_loss #3110
AtomGit-Bot创建于 2025年9月3日
Add fused_linear_crocess_ectropy_loss #3110
已合并
从refs/pull/3110/head合入到master
共 3 个文件变更+108-0
| @@ -7932,6 +7932,15 @@ custom: | |||
| 7932 | - func: npu_dynamic_block_quant(Tensor x, *, float min_scale=0.0, str round_mode="rint", int dst_type=23, int row_block_size=1, int col_block_size=128) -> (Tensor y, Tensor scale) | 7932 | - func: npu_dynamic_block_quant(Tensor x, *, float min_scale=0.0, str round_mode="rint", int dst_type=23, int row_block_size=1, int col_block_size=128) -> (Tensor y, Tensor scale) |
| 7933 | op_api: all_version | 7933 | op_api: all_version |
| 7934 | 7934 | ||
| 7935 | + - func: fused_linear_online_max_sum(Tensor input, Tensor weight, Tensor target, int vocab_start_index, int vocab_end_index, bool return_logits=False) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor) | ||
| 7936 | + op_api: all_version | ||
| 7937 | + | ||
| 7938 | + - func: fused_cross_entropy_loss_with_max_sum(Tensor logits_max, Tensor sum_exp_logits, Tensor predicted_logits, *, float? label_smoothing=0.0, Tensor? input=None, Tensor? weight=None, Tensor? vocab_parallel_logits=None) -> (Tensor, Tensor) | ||
| 7939 | + op_api: all_version | ||
| 7940 | + | ||
| 7941 | + - func: fused_linear_cross_entropy_loss_with_max_sum_grad(Tensor grad, Tensor input, Tensor weight, Tensor target_mask, Tensor masked_target, float label_smoothing=0.0, Tensor? logits_max=None, Tensor? sum_exp_logits=None, Tensor? softmax=None) -> (Tensor, Tensor) | ||
| 7942 | + op_api: all_version | ||
| 7943 | + | ||
| 7935 | symint: | 7944 | symint: |
| 7936 | - func: npu_gather_sparse_index_backward(Tensor grad, SymInt[] self_sizes, Tensor index) -> Tensor | 7945 | - func: npu_gather_sparse_index_backward(Tensor grad, SymInt[] self_sizes, Tensor index) -> Tensor |
| 7937 | op_api: all_version | 7946 | op_api: all_version |
| @@ -0,0 +1,90 @@ | |||
| 1 | +// Copyright (c) 2025 Huawei Technologies Co., Ltd | ||
| 2 | +// All rights reserved. | ||
| 3 | + | ||
| 4 | + | ||
| 5 | + | ||
| 6 | + | ||
| 7 | + | ||
| 8 | +namespace op_api { | ||
| 9 | +using npu_preparation = at_npu::native::OpPreparation; | ||
| 10 | + | ||
| 11 | +::std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor> fused_linear_online_max_sum( | ||
| 12 | + const at::Tensor & input, | ||
| 13 | + const at::Tensor & weight, | ||
| 14 | + const at::Tensor & target, | ||
| 15 | + int64_t vocab_start_index, | ||
| 16 | + int64_t vocab_end_index, | ||
| 17 | + bool return_logits) | ||
| 18 | +{ | ||
| 19 | + auto output_size_0 = {input.size(0)}; | ||
| 20 | + auto output_size_1 = {(input.size(0)+7)/8}; | ||
| 21 | + auto output_dtype_0 = at::kFloat; | ||
| 22 | + auto output_dtype_1 = at::kByte; | ||
| 23 | + auto output_dtype_2 = target.scalar_type(); | ||
| 24 | + auto output_dtype_3 = input.scalar_type(); | ||
| 25 | + | ||
| 26 | + at::Tensor vocab_parallel_logits; | ||
| 27 | + if (return_logits) { | ||
| 28 | + auto output_size_2 = c10::SmallVector<int64_t, op_infer::SIZE>{input.size(0), weight.size(0)}; | ||
| 29 | + vocab_parallel_logits = npu_preparation::apply_tensor_without_format(output_size_2, input.options().dtype(output_dtype_3)); | ||
| 30 | + } else { | ||
| 31 | + vocab_parallel_logits = return_logits ? vocab_parallel_logits : at::Tensor(); | ||
| 32 | + } | ||
| 33 | + | ||
| 34 | + at::Tensor logits_max = npu_preparation::apply_tensor_without_format(output_size_0, input.options().dtype(output_dtype_0)); | ||
| 35 | + at::Tensor sum_exp_logits = npu_preparation::apply_tensor_without_format(output_size_0, input.options().dtype(output_dtype_0)); | ||
| 36 | + at::Tensor predicted_logits = npu_preparation::apply_tensor_without_format(output_size_0, input.options().dtype(output_dtype_0)); | ||
| 37 | + at::Tensor target_mask = npu_preparation::apply_tensor_without_format(output_size_1, input.options().dtype(output_dtype_1)); | ||
| 38 | + at::Tensor masked_target = npu_preparation::apply_tensor_without_format(output_size_0, input.options().dtype(output_dtype_2)); | ||
| 39 | + | ||
| 40 | + EXEC_NPU_CMD(aclnnFusedLinearOnlineMaxSum, input, weight, target, vocab_start_index, vocab_end_index, logits_max, sum_exp_logits, predicted_logits, target_mask, masked_target, vocab_parallel_logits); | ||
| 41 | + return std::make_tuple(std::move(logits_max), std::move(sum_exp_logits), std::move(predicted_logits), std::move(target_mask), std::move(masked_target), std::move(vocab_parallel_logits)); | ||
| 42 | +} | ||
| 43 | + | ||
| 44 | +::std::tuple<at::Tensor, at::Tensor> fused_cross_entropy_loss_with_max_sum( | ||
| 45 | + const at::Tensor & logits_max, | ||
| 46 | + const at::Tensor & sum_exp_logits, | ||
| 47 | + const at::Tensor & predicted_logits, | ||
| 48 | + c10::optional<double> label_smoothing, | ||
| 49 | + const c10::optional<at::Tensor> & input, | ||
| 50 | + const c10::optional<at::Tensor> & weight, | ||
| 51 | + const c10::optional<at::Tensor> & vocab_parallel_logits) | ||
| 52 | +{ | ||
| 53 | + auto label_smoothing_value = label_smoothing.value_or(0.0); | ||
| 54 | + auto output_size_0 = logits_max.sizes(); | ||
| 55 | + auto output_dtype_0 = at::kFloat; | ||
| 56 | + | ||
| 57 | + at::Tensor softmax; | ||
| 58 | + if (vocab_parallel_logits.has_value() && vocab_parallel_logits.value().defined()) { | ||
| 59 | + auto output_size_1 = vocab_parallel_logits.value().sizes(); | ||
| 60 | + softmax = npu_preparation::apply_tensor_without_format(output_size_1, logits_max.options().dtype(output_dtype_0)); | ||
| 61 | + } else { | ||
| 62 | + softmax = at::Tensor(); | ||
| 63 | + } | ||
| 64 | + | ||
| 65 | + at::Tensor loss = npu_preparation::apply_tensor_without_format(output_size_0, logits_max.options().dtype(output_dtype_0)); | ||
| 66 | + | ||
| 67 | + EXEC_NPU_CMD(aclnnFusedCrossEntropyLossWithMaxSum, logits_max, sum_exp_logits, predicted_logits, label_smoothing_value, input, weight, vocab_parallel_logits, loss, softmax); | ||
| 68 | + return std::make_tuple(std::move(loss), std::move(softmax)); | ||
| 69 | +} | ||
| 70 | + | ||
| 71 | +::std::tuple<at::Tensor, at::Tensor> fused_linear_cross_entropy_loss_with_max_sum_grad( | ||
| 72 | + const at::Tensor & grad, | ||
| 73 | + const at::Tensor & input, | ||
| 74 | + const at::Tensor & weight, | ||
| 75 | + const at::Tensor & target_mask, | ||
| 76 | + const at::Tensor & masked_target, | ||
| 77 | + double label_smoothing, | ||
| 78 | + const c10::optional<at::Tensor> & logits_max, | ||
| 79 | + const c10::optional<at::Tensor> & sum_exp_logits, | ||
| 80 | + const c10::optional<at::Tensor> & softmax) | ||
| 81 | +{ | ||
| 82 | + auto output_size_0 = {input.size(0), input.size(1)}; | ||
| 83 | + auto output_size_1 = {weight.size(0), weight.size(1)}; | ||
| 84 | + auto output_dtype_0 = input.scalar_type(); | ||
| 85 | + at::Tensor input_grad = npu_preparation::apply_tensor_without_format(output_size_0, grad.options().dtype(output_dtype_0)); | ||
| 86 | + at::Tensor weight_grad = npu_preparation::apply_tensor_without_format(output_size_1, grad.options().dtype(output_dtype_0)); | ||
| 87 | + EXEC_NPU_CMD(aclnnFusedLinearCrossEntropyLossGrad, grad, input, weight, target_mask, masked_target, label_smoothing, logits_max, sum_exp_logits, softmax, input_grad, weight_grad); | ||
| 88 | + return std::make_tuple(std::move(input_grad), std::move(weight_grad)); | ||
| 89 | +} | ||
| 90 | +} | ||
| @@ -4348,5 +4348,14 @@ | |||
| 4348 | }, | 4348 | }, |
| 4349 | "func: npu_dynamic_block_quant(Tensor x, *, float min_scale=0.0, str round_mode=\"rint\", int dst_type=23, int row_block_size=1, int col_block_size=128) -> (Tensor y, Tensor scale)": { | 4349 | "func: npu_dynamic_block_quant(Tensor x, *, float min_scale=0.0, str round_mode=\"rint\", int dst_type=23, int row_block_size=1, int col_block_size=128) -> (Tensor y, Tensor scale)": { |
| 4350 | "version": ["all_version"] | 4350 | "version": ["all_version"] |
| 4351 | + }, | ||
| 4352 | + "func: fused_linear_online_max_sum(Tensor input, Tensor weight, Tensor target, int vocab_start_index, int vocab_end_index, bool return_logits=False) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor)": { | ||
| 4353 | + "version": ["all_version"] | ||
| 4354 | + }, | ||
| 4355 | + "func: fused_cross_entropy_loss_with_max_sum(Tensor logits_max, Tensor sum_exp_logits, Tensor predicted_logits, *, float? label_smoothing=0.0, Tensor? input=None, Tensor? weight=None, Tensor? vocab_parallel_logits=None) -> (Tensor, Tensor)": { | ||
| 4356 | + "version": ["all_version"] | ||
| 4357 | + }, | ||
| 4358 | + "func: fused_linear_cross_entropy_loss_with_max_sum_grad(Tensor grad, Tensor input, Tensor weight, Tensor target_mask, Tensor masked_target, float label_smoothing=0.0, Tensor? logits_max=None, Tensor? sum_exp_logits=None, Tensor? softmax=None) -> (Tensor, Tensor)": { | ||
| 4359 | + "version": ["all_version"] | ||
| 4351 | } | 4360 | } |
| 4352 | } | 4361 | } |