已合并
fix gelu inductor decomp to match eager Gelu/GeluV2 and add erf decomp for fusion #43650
wangzili121创建于 8月3日
fix gelu inductor decomp to match eager Gelu/GeluV2 and add erf decomp for fusion #43650
已合并
共 2 个文件变更+365-105
| @@ -220,42 +220,136 @@ def _register_mlir_dvm_decompositions(): | |||
| 220 | 220 | ||
| 221 | return grad_t, grad_cos, grad_sin | 221 | return grad_t, grad_cos, grad_sin |
| 222 | 222 | ||
| 223 | + # Constants from cann/ops-nn gelu / gelu_grad | ||
| 224 | + _GELU_BETA = 1.595769121605730711759 # sqrt(8/pi) | ||
| 225 | + _GELU_KAPPA = 0.044715 | ||
| 226 | + _GELU_AN = -0.0713548162726002527220 # -BETA * KAPPA | ||
| 227 | + _GELU_A3 = 0.2140644488178007 # BETA * 3 * KAPPA | ||
| 228 | + _M_SQRT1_2 = 0.70710678118654752440 | ||
| 229 | + _INV_SQRT_2PI = 0.3989422804 | ||
| 230 | + | ||
| 231 | + # AscendC Erf PADE (adv_api/detail/math/erf/erf_common_impl.h) | ||
| 232 | + _ERF_CLIP = 3.92 | ||
| 233 | + _ERF_P5 = 0.053443748819 | ||
| 234 | + _ERF_P4 = 0.75517016694e1 | ||
| 235 | + _ERF_P3 = 0.10162808918e3 | ||
| 236 | + _ERF_P2 = 0.13938061484e4 | ||
| 237 | + _ERF_P1 = 0.50637915060e4 | ||
| 238 | + _ERF_P0 = 0.29639384698e5 | ||
| 239 | + _ERF_Q4 = 0.31212858877e2 | ||
| 240 | + _ERF_Q3 = 0.39856963806e3 | ||
| 241 | + _ERF_Q2 = 0.30231248150e4 | ||
| 242 | + _ERF_Q1 = 0.13243365831e5 | ||
| 243 | + _ERF_Q0 = 0.26267224157e5 | ||
| 244 | + | ||
| 245 | + def _npu_use_compatible_gelu_v2() -> bool: | ||
| 246 | + try: | ||
| 247 | + from torch_npu.npu import are_compatible_impl_enabled | ||
| 248 | + | ||
| 249 | + return are_compatible_impl_enabled() | ||
| 250 | + except Exception: | ||
| 251 | + return False | ||
| 252 | + | ||
| 253 | + def _gelu_use_tanh_approx(approximate: str) -> bool: | ||
| 254 | + if approximate == "tanh": | ||
| 255 | + return True | ||
| 256 | + if approximate == "none": | ||
| 257 | + return not _npu_use_compatible_gelu_v2() | ||
| 258 | + raise RuntimeError( | ||
| 259 | + f"approximate argument must be either none or tanh, but got {approximate!r}." | ||
| 260 | + ) | ||
| 261 | + | ||
| 262 | + def _erf_clip_fp32(x: Tensor) -> Tensor: | ||
| 263 | + """AscendC ErfClip: Mins(x, 3.92) then Maxs(..., -3.92).""" | ||
| 264 | + x = torch.clamp_max(x, _ERF_CLIP) | ||
| 265 | + x = torch.clamp_min(x, -_ERF_CLIP) | ||
| 266 | + return x | ||
| 267 | + | ||
| 268 | + def _erf_compute_p_fp32(x: Tensor, x2: Tensor) -> Tensor: | ||
| 269 | + """AscendC ErfComputeP: P(x) Horner (Muls/Adds/Mul).""" | ||
| 270 | + t = x2 * _ERF_P5 | ||
| 271 | + t = t + _ERF_P4 | ||
| 272 | + t = x2 * t | ||
| 273 | + t = t + _ERF_P3 | ||
| 274 | + t = x2 * t | ||
| 275 | + t = t + _ERF_P2 | ||
| 276 | + t = x2 * t | ||
| 277 | + t = t + _ERF_P1 | ||
| 278 | + t = x2 * t | ||
| 279 | + t = t + _ERF_P0 | ||
| 280 | + return x * t | ||
| 281 | + | ||
| 282 | + def _erf_compute_q_fp32(x2: Tensor) -> Tensor: | ||
| 283 | + """AscendC ErfComputeQ: Q(x) Horner (Adds/Mul).""" | ||
| 284 | + t = x2 + _ERF_Q4 | ||
| 285 | + t = x2 * t | ||
| 286 | + t = t + _ERF_Q3 | ||
| 287 | + t = x2 * t | ||
| 288 | + t = t + _ERF_Q2 | ||
| 289 | + t = x2 * t | ||
| 290 | + t = t + _ERF_Q1 | ||
| 291 | + t = x2 * t | ||
| 292 | + t = t + _ERF_Q0 | ||
| 293 | + return t | ||
| 294 | + | ||
| 295 | + def _erf_pade_fp32(x: Tensor) -> Tensor: | ||
| 296 | + """AscendC ErfCompute: clip then P(x)/Q(x).""" | ||
| 297 | + x = _erf_clip_fp32(x) | ||
| 298 | + x2 = x * x | ||
| 299 | + p = _erf_compute_p_fp32(x, x2) | ||
| 300 | + q = _erf_compute_q_fp32(x2) | ||
| 301 | + return p / q | ||
| 302 | + | ||
| 303 | + def erf(a: Tensor) -> Tensor: | ||
| 304 | + """Match AscendC ErfImpl; fp16: cast up CAST_NONE, compute in fp32, cast down.""" | ||
| 305 | + orig_dtype = a.dtype | ||
| 306 | + if orig_dtype != torch.float32: | ||
| 307 | + a = a.to(torch.float32) | ||
| 308 | + out = _erf_pade_fp32(a) | ||
| 309 | + if orig_dtype != torch.float32: | ||
| 310 | + out = out.to(orig_dtype) | ||
| 311 | + return out | ||
| 312 | + | ||
| 223 | def gelu(a, approximate: str = "none"): | 313 | def gelu(a, approximate: str = "none"): |
| 224 | - """ | 314 | + """Match eager gelu; see cann/ops-nn gelu decomp.""" |
| 225 | - Reference implementation of torch.nn.functional.gelu | 315 | + orig_dtype = a.dtype |
| 226 | - """ | 316 | + if orig_dtype != torch.float32: |
| 227 | - M_SQRT2 = 1.41421356237309504880 | 317 | + a = a.to(torch.float32) |
| 228 | - M_2_SQRTPI = 1.12837916709551257390 | 318 | + if _gelu_use_tanh_approx(approximate): |
| 229 | - kBeta = M_SQRT2 * M_2_SQRTPI * 0.5 | 319 | + a_cube = a * a * a |
| 230 | - kKappa = 0.044715 | 320 | + out = a / (1.0 + torch.exp(-_GELU_BETA * (a + _GELU_KAPPA * a_cube))) |
| 231 | - a_cube = a * a * a | 321 | + else: |
| 232 | - inner = kBeta * (a + kKappa * a_cube) | 322 | + # GeluV2ErfPost: (1 + erf) * (0.5 * x) |
| 233 | - return 0.5 * a * (1 + torch.tanh(inner)) | 323 | + out = (1.0 + _erf_pade_fp32(a * _M_SQRT1_2)) * (0.5 * a) |
| 324 | + if orig_dtype != torch.float32: | ||
| 325 | + out = out.to(orig_dtype) | ||
| 326 | + return out | ||
| 234 | 327 | ||
| 235 | def gelu_backward(grad: Tensor, self: Tensor, approximate: str = "none"): | 328 | def gelu_backward(grad: Tensor, self: Tensor, approximate: str = "none"): |
| 236 | - M_SQRT2 = 1.41421356237309504880 | 329 | + """Match eager gelu_backward; see cann/ops-nn gelu_grad decomp.""" |
| 237 | - M_SQRT1_2 = 0.70710678118654752440 | 330 | + orig_dtype = grad.dtype |
| 238 | - M_2_SQRTPI = 1.12837916709551257390 | 331 | + if orig_dtype != torch.float32: |
| 239 | - kBeta = M_SQRT2 * M_2_SQRTPI * 0.5 | 332 | + grad = grad.to(torch.float32) |
| 240 | - kKappa = 0.044715 | 333 | + self = self.to(torch.float32) |
| 241 | - x_sq = self * self | 334 | + if _gelu_use_tanh_approx(approximate): |
| 242 | - x_cube = x_sq * self | 335 | + x_sq = self * self |
| 243 | - inner = kBeta * (self + kKappa * x_cube) | 336 | + px = torch.exp((-_GELU_BETA + _GELU_AN * x_sq) * self) |
| 244 | - tanh_inner = torch.tanh(inner) | 337 | + res0 = (_GELU_BETA + _GELU_A3 * x_sq) * self |
| 245 | - | 338 | + div = 1.0 / (1.0 + px) |
| 246 | - left = 0.5 * self | 339 | + resp = px * div * res0 * div |
| 247 | - right = 1.0 + tanh_inner | 340 | + resp = torch.where(torch.isnan(resp), torch.zeros_like(resp), resp) |
| 248 | - | 341 | + out = grad * (resp + div) |
| 249 | - left_derivative = 0.5 * right | 342 | + else: |
| 250 | - | 343 | + cdf = 0.5 * (1.0 + _erf_pade_fp32(self * _M_SQRT1_2)) |
| 251 | - tanh_derivative = (tanh_inner * tanh_inner) * -1.0 + 1.0 | 344 | + pdf = _INV_SQRT_2PI * torch.exp(self * self * -0.5) |
| 252 | - inner_derivative = kBeta * (1.0 + 3.0 * kKappa * x_sq) | 345 | + out = grad * (cdf + self * pdf) |
| 253 | - right_derivative = left * tanh_derivative * inner_derivative | 346 | + if orig_dtype != torch.float32: |
| 254 | - | 347 | + out = out.to(orig_dtype) |
| 255 | - return grad * (left_derivative + right_derivative) | 348 | + return out |
| 256 | 349 | ||
| 257 | register_decomposition(torch.ops.aten.convolution_backward)(npu_convolution_backward) | 350 | register_decomposition(torch.ops.aten.convolution_backward)(npu_convolution_backward) |
| 258 | register_decomposition(torch.ops.aten._softmax_backward_data.default)(npu__softmax_backward_data) | 351 | register_decomposition(torch.ops.aten._softmax_backward_data.default)(npu__softmax_backward_data) |
| 352 | + register_decomposition(torch.ops.aten.erf.default)(erf) | ||
| 259 | register_decomposition(torch.ops.aten.gelu.default)(gelu) | 353 | register_decomposition(torch.ops.aten.gelu.default)(gelu) |
| 260 | register_decomposition(torch.ops.aten.gelu_backward.default)(gelu_backward) | 354 | register_decomposition(torch.ops.aten.gelu_backward.default)(gelu_backward) |
| 261 | # register_decomposition(torch.ops.npu.npu_rms_norm.default)(npu_rms_norm) | 355 | # register_decomposition(torch.ops.npu.npu_rms_norm.default)(npu_rms_norm) |
| @@ -442,35 +536,135 @@ def _override_softmax_backward_decomp_no_fma(): | |||
| 442 | 536 | ||
| 443 | 537 | ||
| 444 | def _override_gelu_decomp(): | 538 | def _override_gelu_decomp(): |
| 539 | + # Align default Inductor gelu with cann/ops-nn gelu decomp | ||
| 445 | from torch._inductor.decomposition import decompositions as _ind_decomps | 540 | from torch._inductor.decomposition import decompositions as _ind_decomps |
| 446 | 541 | ||
| 542 | + _GELU_BETA = 1.595769121605730711759 | ||
| 543 | + _GELU_KAPPA = 0.044715 | ||
| 544 | + _GELU_AN = -0.0713548162726002527220 | ||
| 545 | + _GELU_A3 = 0.2140644488178007 | ||
| 546 | + _M_SQRT1_2 = 0.70710678118654752440 | ||
| 547 | + _INV_SQRT_2PI = 0.3989422804 | ||
| 548 | + | ||
| 549 | + _ERF_CLIP = 3.92 | ||
| 550 | + _ERF_P5 = 0.053443748819 | ||
| 551 | + _ERF_P4 = 0.75517016694e1 | ||
| 552 | + _ERF_P3 = 0.10162808918e3 | ||
| 553 | + _ERF_P2 = 0.13938061484e4 | ||
| 554 | + _ERF_P1 = 0.50637915060e4 | ||
| 555 | + _ERF_P0 = 0.29639384698e5 | ||
| 556 | + _ERF_Q4 = 0.31212858877e2 | ||
| 557 | + _ERF_Q3 = 0.39856963806e3 | ||
| 558 | + _ERF_Q2 = 0.30231248150e4 | ||
| 559 | + _ERF_Q1 = 0.13243365831e5 | ||
| 560 | + _ERF_Q0 = 0.26267224157e5 | ||
| 561 | + | ||
| 562 | + def _use_compatible_gelu_v2() -> bool: | ||
| 563 | + try: | ||
| 564 | + from torch_npu.npu import are_compatible_impl_enabled | ||
| 565 | + | ||
| 566 | + return are_compatible_impl_enabled() | ||
| 567 | + except Exception: | ||
| 568 | + return False | ||
| 569 | + | ||
| 570 | + def _use_tanh_approx(approximate: str) -> bool: | ||
| 571 | + if approximate == "tanh": | ||
| 572 | + return True | ||
| 573 | + if approximate == "none": | ||
| 574 | + return not _use_compatible_gelu_v2() | ||
| 575 | + raise RuntimeError( | ||
| 576 | + f"approximate argument must be either none or tanh, but got {approximate!r}." | ||
| 577 | + ) | ||
| 578 | + | ||
| 579 | + def _erf_clip_fp32(x: torch.Tensor) -> torch.Tensor: | ||
| 580 | + """AscendC ErfClip: Mins(x, 3.92) then Maxs(..., -3.92).""" | ||
| 581 | + x = torch.clamp_max(x, _ERF_CLIP) | ||
| 582 | + x = torch.clamp_min(x, -_ERF_CLIP) | ||
| 583 | + return x | ||
| 584 | + | ||
| 585 | + def _erf_compute_p_fp32(x: torch.Tensor, x2: torch.Tensor) -> torch.Tensor: | ||
| 586 | + """AscendC ErfComputeP: P(x) Horner (Muls/Adds/Mul).""" | ||
| 587 | + t = x2 * _ERF_P5 | ||
| 588 | + t = t + _ERF_P4 | ||
| 589 | + t = x2 * t | ||
| 590 | + t = t + _ERF_P3 | ||
| 591 | + t = x2 * t | ||
| 592 | + t = t + _ERF_P2 | ||
| 593 | + t = x2 * t | ||
| 594 | + t = t + _ERF_P1 | ||
| 595 | + t = x2 * t | ||
| 596 | + t = t + _ERF_P0 | ||
| 597 | + return x * t | ||
| 598 | + | ||
| 599 | + def _erf_compute_q_fp32(x2: torch.Tensor) -> torch.Tensor: | ||
| 600 | + """AscendC ErfComputeQ: Q(x) Horner (Adds/Mul).""" | ||
| 601 | + t = x2 + _ERF_Q4 | ||
| 602 | + t = x2 * t | ||
| 603 | + t = t + _ERF_Q3 | ||
| 604 | + t = x2 * t | ||
| 605 | + t = t + _ERF_Q2 | ||
| 606 | + t = x2 * t | ||
| 607 | + t = t + _ERF_Q1 | ||
| 608 | + t = x2 * t | ||
| 609 | + t = t + _ERF_Q0 | ||
| 610 | + return t | ||
| 611 | + | ||
| 612 | + def _erf_pade_fp32(x: torch.Tensor) -> torch.Tensor: | ||
| 613 | + """AscendC ErfCompute: clip then P(x)/Q(x).""" | ||
| 614 | + x = _erf_clip_fp32(x) | ||
| 615 | + x2 = x * x | ||
| 616 | + p = _erf_compute_p_fp32(x, x2) | ||
| 617 | + q = _erf_compute_q_fp32(x2) | ||
| 618 | + return p / q | ||
| 619 | + | ||
| 620 | + def _erf(a: torch.Tensor) -> torch.Tensor: | ||
| 621 | + orig_dtype = a.dtype | ||
| 622 | + if orig_dtype != torch.float32: | ||
| 623 | + a = a.to(torch.float32) | ||
| 624 | + out = _erf_pade_fp32(a) | ||
| 625 | + if orig_dtype != torch.float32: | ||
| 626 | + out = out.to(orig_dtype) | ||
| 627 | + return out | ||
| 628 | + | ||
| 447 | def _gelu(x: torch.Tensor, approximate: str = "none"): | 629 | def _gelu(x: torch.Tensor, approximate: str = "none"): |
| 448 | - two_sqrt_2_over_pi = 1.5957691216057308 | 630 | + orig_dtype = x.dtype |
| 449 | - coeff = 0.044715 | 631 | + if orig_dtype != torch.float32: |
| 450 | - x_cubed = x * x * x | 632 | + x = x.to(torch.float32) |
| 451 | - z = two_sqrt_2_over_pi * (x + coeff * x_cubed) | 633 | + if _use_tanh_approx(approximate): |
| 452 | - sigmoid_z = torch.sigmoid(z) | 634 | + x_cubed = x * x * x |
| 453 | - result = x * sigmoid_z | 635 | + out = x / (1.0 + torch.exp(-_GELU_BETA * (x + _GELU_KAPPA * x_cubed))) |
| 454 | - return result | 636 | + else: |
| 637 | + # GeluV2ErfPost: (1 + erf) * (0.5 * x) | ||
| 638 | + out = (1.0 + _erf_pade_fp32(x * _M_SQRT1_2)) * (0.5 * x) | ||
| 639 | + if orig_dtype != torch.float32: | ||
| 640 | + out = out.to(orig_dtype) | ||
| 641 | + return out | ||
| 455 | 642 | ||
| 456 | def _gelu_backward(grad: torch.Tensor, self: torch.Tensor, approximate: str = "none"): | 643 | def _gelu_backward(grad: torch.Tensor, self: torch.Tensor, approximate: str = "none"): |
| 457 | - two_sqrt_2_over_pi = 1.5957691216057308 | 644 | + orig_dtype = grad.dtype |
| 458 | - coeff = 0.044715 | 645 | + if orig_dtype != torch.float32: |
| 459 | - x_sq = self * self | 646 | + grad = grad.to(torch.float32) |
| 460 | - x_cubed = x_sq * self | 647 | + self = self.to(torch.float32) |
| 461 | - z = two_sqrt_2_over_pi * (self + coeff * x_cubed) | 648 | + if _use_tanh_approx(approximate): |
| 462 | - sigmoid_z = torch.sigmoid(z) | 649 | + x_sq = self * self |
| 463 | - # dz/dx = two_sqrt_2_over_pi * (1 + 3*coeff*x^2) | 650 | + px = torch.exp((-_GELU_BETA + _GELU_AN * x_sq) * self) |
| 464 | - z_derivative = two_sqrt_2_over_pi * (1.0 + 3.0 * coeff * x_sq) | 651 | + res0 = (_GELU_BETA + _GELU_A3 * x_sq) * self |
| 465 | - # d/dx [x * sigmoid(z)] = sigmoid(z) + x * sigmoid(z)*(1 - sigmoid(z)) * dz/dx | 652 | + div = 1.0 / (1.0 + px) |
| 466 | - sigmoid_derivative = sigmoid_z * (1.0 - sigmoid_z) | 653 | + resp = px * div * res0 * div |
| 467 | - result_derivative = sigmoid_z + self * sigmoid_derivative * z_derivative | 654 | + resp = torch.where(torch.isnan(resp), torch.zeros_like(resp), resp) |
| 468 | - return grad * result_derivative | 655 | + out = grad * (resp + div) |
| 656 | + else: | ||
| 657 | + cdf = 0.5 * (1.0 + _erf_pade_fp32(self * _M_SQRT1_2)) | ||
| 658 | + pdf = _INV_SQRT_2PI * torch.exp(self * self * -0.5) | ||
| 659 | + out = grad * (cdf + self * pdf) | ||
| 660 | + if orig_dtype != torch.float32: | ||
| 661 | + out = out.to(orig_dtype) | ||
| 662 | + return out | ||
| 469 | 663 | ||
| 664 | + _ind_decomps[aten.erf.default] = _erf | ||
| 470 | _ind_decomps[aten.gelu.default] = _gelu | 665 | _ind_decomps[aten.gelu.default] = _gelu |
| 471 | _ind_decomps[aten.gelu_backward.default] = _gelu_backward | 666 | _ind_decomps[aten.gelu_backward.default] = _gelu_backward |
| 472 | 667 | ||
| 473 | - | ||
| 474 | def _override_rms_norm_decomp(): | 668 | def _override_rms_norm_decomp(): |
| 475 | from torch._inductor.decomposition import decompositions as _ind_decomps | 669 | from torch._inductor.decomposition import decompositions as _ind_decomps |
| 476 | 670 | ||
| @@ -1,5 +1,3 @@ | |||
| 1 | -import math | ||
| 2 | - | ||
| 3 | import torch | 1 | import torch |
| 4 | from torch._decomp import remove_decompositions | 2 | from torch._decomp import remove_decompositions |
| 5 | from torch._inductor import decomposition as inductor_decomp | 3 | from torch._inductor import decomposition as inductor_decomp |
| @@ -72,59 +70,6 @@ def tanh(a): | |||
| 72 | return out | 70 | return out |
| 73 | 71 | ||
| 74 | 72 | ||
| 75 | -def gelu(a: torch.Tensor, approximate: str = "none"): | ||
| 76 | - """ | ||
| 77 | - y = -sqrt(8/pi) * (x + 0.044715 * x^3) | ||
| 78 | - out = x / (1 + exp(y)) | ||
| 79 | - """ | ||
| 80 | - orig_dtype = a.dtype | ||
| 81 | - if orig_dtype != torch.float32: | ||
| 82 | - a = a.to(torch.float32) | ||
| 83 | - | ||
| 84 | - M_SQRT2 = math.sqrt(2) | ||
| 85 | - M_2_SQRTPI = 2.0 / math.sqrt(math.pi) | ||
| 86 | - kBeta = M_SQRT2 * M_2_SQRTPI | ||
| 87 | - kKappa = 0.044715 | ||
| 88 | - | ||
| 89 | - a_cube = a * a * a | ||
| 90 | - inner = a + kKappa * a_cube | ||
| 91 | - y = -kBeta * inner | ||
| 92 | - out = a / (1.0 + torch.exp(y)) | ||
| 93 | - | ||
| 94 | - if orig_dtype != torch.float32: | ||
| 95 | - out = out.to(orig_dtype) | ||
| 96 | - return out | ||
| 97 | - | ||
| 98 | - | ||
| 99 | -def gelu_backward(grad, self, approximate: str = "none"): | ||
| 100 | - orig_dtype = grad.dtype | ||
| 101 | - if orig_dtype != torch.float32: | ||
| 102 | - grad = grad.to(torch.float32) | ||
| 103 | - self = self.to(torch.float32) | ||
| 104 | - M_SQRT2 = math.sqrt(2) | ||
| 105 | - M_2_SQRTPI = 2.0 / math.sqrt(math.pi) | ||
| 106 | - kBeta = M_SQRT2 * M_2_SQRTPI * 0.5 | ||
| 107 | - kKappa = 0.044715 | ||
| 108 | - x_sq = self * self | ||
| 109 | - x_cube = x_sq * self | ||
| 110 | - inner = kBeta * (self + kKappa * x_cube) | ||
| 111 | - tanh_inner = torch.tanh(inner) | ||
| 112 | - | ||
| 113 | - left = 0.5 * self | ||
| 114 | - right = 1.0 + tanh_inner | ||
| 115 | - | ||
| 116 | - left_derivative = 0.5 * right | ||
| 117 | - | ||
| 118 | - tanh_derivative = (tanh_inner * tanh_inner) * -1.0 + 1.0 | ||
| 119 | - inner_derivative = kBeta * (1.0 + 3.0 * kKappa * x_sq) | ||
| 120 | - right_derivative = left * tanh_derivative * inner_derivative | ||
| 121 | - out = grad * (left_derivative + right_derivative) | ||
| 122 | - | ||
| 123 | - if orig_dtype != torch.float32: | ||
| 124 | - out = out.to(orig_dtype) | ||
| 125 | - return out | ||
| 126 | - | ||
| 127 | - | ||
| 128 | def sigmoid(a: torch.Tensor) -> torch.Tensor: | 73 | def sigmoid(a: torch.Tensor) -> torch.Tensor: |
| 129 | orig_dtype = a.dtype | 74 | orig_dtype = a.dtype |
| 130 | if orig_dtype != torch.float32: | 75 | if orig_dtype != torch.float32: |
| @@ -135,6 +80,127 @@ def sigmoid(a: torch.Tensor) -> torch.Tensor: | |||
| 135 | return out | 80 | return out |
| 136 | 81 | ||
| 137 | 82 | ||
| 83 | +# Constants from cann/ops-nn gelu / gelu_grad | ||
| 84 | +_GELU_BETA = 1.595769121605730711759 # sqrt(8/pi) | ||
| 85 | +_GELU_KAPPA = 0.044715 | ||
| 86 | +_GELU_AN = -0.0713548162726002527220 # -BETA * KAPPA | ||
| 87 | +_GELU_A3 = 0.2140644488178007 # BETA * 3 * KAPPA | ||
| 88 | +_M_SQRT1_2 = 0.70710678118654752440 | ||
| 89 | +_INV_SQRT_2PI = 0.3989422804 | ||
| 90 | + | ||
| 91 | +# AscendC Erf PADE | ||
| 92 | +_ERF_CLIP = 3.92 | ||
| 93 | +_ERF_P5 = 0.053443748819 | ||
| 94 | +_ERF_P4 = 0.75517016694e1 | ||
| 95 | +_ERF_P3 = 0.10162808918e3 | ||
| 96 | +_ERF_P2 = 0.13938061484e4 | ||
| 97 | +_ERF_P1 = 0.50637915060e4 | ||
| 98 | +_ERF_P0 = 0.29639384698e5 | ||
| 99 | +_ERF_Q4 = 0.31212858877e2 | ||
| 100 | +_ERF_Q3 = 0.39856963806e3 | ||
| 101 | +_ERF_Q2 = 0.30231248150e4 | ||
| 102 | +_ERF_Q1 = 0.13243365831e5 | ||
| 103 | +_ERF_Q0 = 0.26267224157e5 | ||
| 104 | + | ||
| 105 | + | ||
| 106 | +def _npu_use_compatible_gelu_v2() -> bool: | ||
| 107 | + try: | ||
| 108 | + from torch_npu.npu import are_compatible_impl_enabled | ||
| 109 | + | ||
| 110 | + return are_compatible_impl_enabled() | ||
| 111 | + except Exception: | ||
| 112 | + return False | ||
| 113 | + | ||
| 114 | + | ||
| 115 | +def _gelu_use_tanh_approx(approximate: str) -> bool: | ||
| 116 | + if approximate == "tanh": | ||
| 117 | + return True | ||
| 118 | + if approximate == "none": | ||
| 119 | + return not _npu_use_compatible_gelu_v2() | ||
| 120 | + raise RuntimeError( | ||
| 121 | + f"approximate argument must be either none or tanh, but got {approximate!r}." | ||
| 122 | + ) | ||
| 123 | + | ||
| 124 | + | ||
| 125 | +def _erf_clip_fp32(x: torch.Tensor) -> torch.Tensor: | ||
| 126 | + x = torch.clamp_max(x, _ERF_CLIP) | ||
| 127 | + x = torch.clamp_min(x, -_ERF_CLIP) | ||
| 128 | + return x | ||
| 129 | + | ||
| 130 | + | ||
| 131 | +def _erf_compute_p_fp32(x: torch.Tensor, x2: torch.Tensor) -> torch.Tensor: | ||
| 132 | + t = x2 * _ERF_P5 | ||
| 133 | + t = t + _ERF_P4 | ||
| 134 | + t = x2 * t | ||
| 135 | + t = t + _ERF_P3 | ||
| 136 | + t = x2 * t | ||
| 137 | + t = t + _ERF_P2 | ||
| 138 | + t = x2 * t | ||
| 139 | + t = t + _ERF_P1 | ||
| 140 | + t = x2 * t | ||
| 141 | + t = t + _ERF_P0 | ||
| 142 | + return x * t | ||
| 143 | + | ||
| 144 | + | ||
| 145 | +def _erf_compute_q_fp32(x2: torch.Tensor) -> torch.Tensor: | ||
| 146 | + t = x2 + _ERF_Q4 | ||
| 147 | + t = x2 * t | ||
| 148 | + t = t + _ERF_Q3 | ||
| 149 | + t = x2 * t | ||
| 150 | + t = t + _ERF_Q2 | ||
| 151 | + t = x2 * t | ||
| 152 | + t = t + _ERF_Q1 | ||
| 153 | + t = x2 * t | ||
| 154 | + t = t + _ERF_Q0 | ||
| 155 | + return t | ||
| 156 | + | ||
| 157 | + | ||
| 158 | +def _erf_pade_fp32(x: torch.Tensor) -> torch.Tensor: | ||
| 159 | + x = _erf_clip_fp32(x) | ||
| 160 | + x2 = x * x | ||
| 161 | + p = _erf_compute_p_fp32(x, x2) | ||
| 162 | + q = _erf_compute_q_fp32(x2) | ||
| 163 | + return p / q | ||
| 164 | + | ||
| 165 | + | ||
| 166 | +def gelu(a: torch.Tensor, approximate: str = "none"): | ||
| 167 | + """Match eager gelu; DVM path uses resp==resp NaN clear in backward.""" | ||
| 168 | + orig_dtype = a.dtype | ||
| 169 | + if orig_dtype != torch.float32: | ||
| 170 | + a = a.to(torch.float32) | ||
| 171 | + if _gelu_use_tanh_approx(approximate): | ||
| 172 | + a_cube = a * a * a | ||
| 173 | + out = a / (1.0 + torch.exp(-_GELU_BETA * (a + _GELU_KAPPA * a_cube))) | ||
| 174 | + else: | ||
| 175 | + out = (1.0 + _erf_pade_fp32(a * _M_SQRT1_2)) * (0.5 * a) | ||
| 176 | + if orig_dtype != torch.float32: | ||
| 177 | + out = out.to(orig_dtype) | ||
| 178 | + return out | ||
| 179 | + | ||
| 180 | + | ||
| 181 | +def gelu_backward(grad: torch.Tensor, self: torch.Tensor, approximate: str = "none"): | ||
| 182 | + """Match eager gelu_backward; NaN clear via resp==resp (CANN Compare EQ).""" | ||
| 183 | + orig_dtype = grad.dtype | ||
| 184 | + if orig_dtype != torch.float32: | ||
| 185 | + grad = grad.to(torch.float32) | ||
| 186 | + self = self.to(torch.float32) | ||
| 187 | + if _gelu_use_tanh_approx(approximate): | ||
| 188 | + x_sq = self * self | ||
| 189 | + px = torch.exp((-_GELU_BETA + _GELU_AN * x_sq) * self) | ||
| 190 | + res0 = (_GELU_BETA + _GELU_A3 * x_sq) * self | ||
| 191 | + div = 1.0 / (1.0 + px) | ||
| 192 | + resp = px * div * res0 * div | ||
| 193 | + resp = torch.where(resp == resp, resp, torch.zeros_like(resp)) | ||
| 194 | + out = grad * (resp + div) | ||
| 195 | + else: | ||
| 196 | + cdf = 0.5 * (1.0 + _erf_pade_fp32(self * _M_SQRT1_2)) | ||
| 197 | + pdf = _INV_SQRT_2PI * torch.exp(self * self * -0.5) | ||
| 198 | + out = grad * (cdf + self * pdf) | ||
| 199 | + if orig_dtype != torch.float32: | ||
| 200 | + out = out.to(orig_dtype) | ||
| 201 | + return out | ||
| 202 | + | ||
| 203 | + | ||
| 138 | _dvm_inductor_decomp_patched = False | 204 | _dvm_inductor_decomp_patched = False |
| 139 | 205 | ||
| 140 | 206 | ||