已合并
feat: support triton_experimental in TorchBench #43627
rmch创建于 8月3日
feat: support triton_experimental in TorchBench #43627
已合并
共 3 个文件变更+236-87
| @@ -171,10 +171,13 @@ | |||
| 171 | 171 | ||
| 172 | 4. NPU图模式后端指定 | 172 | 4. NPU图模式后端指定 |
| 173 | 173 | ||
| 174 | - 当前NPU图模式后端通过`--npu-backend`参数指定，支持`mlir`、`dvm`、`akg`、`triton`四种模式，不显示指定会默认选择四种模式中端到端时间加速比最大的图模式后端，使用示例如下 | 174 | + 当前NPU图模式后端通过`--npu-backend`参数指定，支持`mlir`、`dvm`、`akg`、`triton`和`triton_experimental`五种模式；不显式指定时，会默认选择端到端时间加速比最大的图模式后端。其中`triton`对应默认 Triton 后端，`triton_experimental`对应独立的实验性 Triton 后端。使用示例如下 |
| 175 | 175 | ||
| 176 | ```shell | 176 | ```shell |
| 177 | python3 torchbench.py --accuracy --cold-start-latency --train --float32 --backend inductor --npu-backend mlir --only BERT_pytorch --iterations 50 | 177 | python3 torchbench.py --accuracy --cold-start-latency --train --float32 --backend inductor --npu-backend mlir --only BERT_pytorch --iterations 50 |
| 178 | + | ||
| 179 | + # 使用triton_experimental后端 | ||
| 180 | + python3 torchbench.py --accuracy --cold-start-latency --train --float32 --backend inductor --npu-backend triton_experimental --only nvidia_deeprecommender --iterations 50 --disable-aclgraph --dynamic-shapes | ||
| 178 | ``` | 181 | ``` |
| 179 | 182 | ||
| 180 | 当前可通过 `--mfusion` 参数开启 MFusion 图算融合优化功能, 配合不同的NPU图模式后端, 进一步提升模型的性能，使用示例如下 | 183 | 当前可通过 `--mfusion` 参数开启 MFusion 图算融合优化功能, 配合不同的NPU图模式后端, 进一步提升模型的性能，使用示例如下 |
| @@ -302,7 +305,7 @@ | |||
| 302 | 305 | ||
| 303 | 1. NPU图模式后端指定 | 306 | 1. NPU图模式后端指定 |
| 304 | 307 | ||
| 305 | - 当前NPU图模式后端通过`--npu-backend`参数指定，支持`mlir`、`dvm`、`akg`、`triton`四种模式，使用示例如下 | 308 | + 当前NPU图模式后端通过`--npu-backend`参数指定，支持`mlir`、`dvm`、`akg`、`triton`和`triton_experimental`五种模式。其中`triton`对应默认 Triton 后端，`triton_experimental`对应独立的实验性 Triton 后端。使用示例如下 |
| 306 | 309 | ||
| 307 | ```shell | 310 | ```shell |
| 308 | python3 huggingface.py --accuracy --cold-start-latency --train --float32 --backend inductor --npu-backend mlir --only AlbertForMaskedLM --iterations 50 | 311 | python3 huggingface.py --accuracy --cold-start-latency --train --float32 --backend inductor --npu-backend mlir --only AlbertForMaskedLM --iterations 50 |
| @@ -844,6 +844,16 @@ class BenchmarkRunner: | |||
| 844 | 844 | ||
| 845 | def run_n_iterations(self, mod, inputs, run_mode=None): | 845 | def run_n_iterations(self, mod, inputs, run_mode=None): |
| 846 | n = self.args.iterations | 846 | n = self.args.iterations |
| 847 | + | ||
| 848 | + if is_npu_available and self.args.enable_profiler is not None: | ||
| 849 | + # Run the first iteration before starting the outer profiler. | ||
| 850 | + # Inductor compilation and Triton autotuning can create their own | ||
| 851 | + # profiler sessions; nesting those inside the benchmark profiler | ||
| 852 | + # can corrupt the CANN trace (for example, by leaving it without a | ||
| 853 | + # TASK table). | ||
| 854 | + self.model_iter_fn(mod, inputs, collect_outputs=False) | ||
R | |||
| 855 | + synchronize() | ||
| 856 | + | ||
| 847 | if run_mode is None: | 857 | if run_mode is None: |
| 848 | for _ in range(n - 1): | 858 | for _ in range(n - 1): |
| 849 | self.model_iter_fn(mod, inputs, collect_outputs=False) | 859 | self.model_iter_fn(mod, inputs, collect_outputs=False) |
| @@ -866,7 +876,7 @@ class BenchmarkRunner: | |||
| 866 | enable=self.args.enable_profiler is not None, | 876 | enable=self.args.enable_profiler is not None, |
| 867 | level=prof_level, | 877 | level=prof_level, |
| 868 | warmup=10, | 878 | warmup=10, |
| 869 | - active=n, | 879 | + active=max(1, n - 10), |
| 870 | save_path=prof_output_dir, | 880 | save_path=prof_output_dir, |
| 871 | ) | 881 | ) |
| 872 | else: | 882 | else: |
| @@ -1581,7 +1591,14 @@ def parse_args(args=None): | |||
| 1581 | ) | 1591 | ) |
| 1582 | parser.add_argument( | 1592 | parser.add_argument( |
| 1583 | "--npu-backend", | 1593 | "--npu-backend", |
| 1584 | - choices=["mlir", "dvm", "akg", "triton", "default"], | 1594 | + choices=[ |
| 1595 | + "mlir", | ||
| 1596 | + "dvm", | ||
| 1597 | + "akg", | ||
| 1598 | + "triton", | ||
| 1599 | + "triton_experimental", | ||
| 1600 | + "default", | ||
| 1601 | + ], | ||
| 1585 | default="default", | 1602 | default="default", |
| 1586 | help="Specify NPU backend (only effective when --backend is inductor)", | 1603 | help="Specify NPU backend (only effective when --backend is inductor)", |
| 1587 | ) | 1604 | ) |
| @@ -2077,8 +2094,14 @@ def configure_compile_options(args, runner): | |||
| 2077 | if backend == "inductor" and hasattr(args, "npu_backend"): | 2094 | if backend == "inductor" and hasattr(args, "npu_backend"): |
| 2078 | if npu_backend == "default": | 2095 | if npu_backend == "default": |
| 2079 | npu_backend = get_npu_backend(args) | 2096 | npu_backend = get_npu_backend(args) |
| 2080 | - if npu_backend in ["mlir", "dvm"]: | 2097 | + backend_env = { |
| 2081 | - os.environ["TORCHINDUCTOR_NPU_BACKEND"] = npu_backend | 2098 | + "mlir": "mlir", |
| 2099 | + "dvm": "dvm", | ||
| 2100 | + "triton": "default", | ||
| 2101 | + "triton_experimental": "triton_experimental", | ||
| 2102 | + } | ||
| 2103 | + if npu_backend in backend_env: | ||
| 2104 | + os.environ["TORCHINDUCTOR_NPU_BACKEND"] = backend_env[npu_backend] | ||
| 2082 | if npu_backend == "akg": | 2105 | if npu_backend == "akg": |
| 2083 | os.environ["TORCHINDUCTOR_NPU_BACKEND"] = "mlir" | 2106 | os.environ["TORCHINDUCTOR_NPU_BACKEND"] = "mlir" |
| 2084 | os.environ["TORCHINDUCTOR_USE_AKG"] = "1" | 2107 | os.environ["TORCHINDUCTOR_USE_AKG"] = "1" |
| @@ -1,3 +1,4 @@ | |||
| 1 | +import gc | ||
| 1 | import logging | 2 | import logging |
| 2 | import os | 3 | import os |
| 3 | import sys | 4 | import sys |
| @@ -56,6 +57,21 @@ def use_aclnn(): | |||
| 56 | os.environ["USE_ACLOP"] = "0" | 57 | os.environ["USE_ACLOP"] = "0" |
| 57 | 58 | ||
| 58 | 59 | ||
| 60 | +def _is_triton_experimental_backend(): | ||
| 61 | + """Return True when the active NPU inductor backend is triton_experimental. | ||
| 62 | + | ||
| 63 | + configure_compile_options() sets TORCHINDUCTOR_NPU_BACKEND before | ||
| 64 | + patch_model() runs, so this reflects the backend selected for the run. | ||
| 65 | + | ||
| 66 | + The triton_experimental backend does not consume the ascend_npu_ir config | ||
| 67 | + (GENERATE_LIST / force_fallback_kernel_names / decomps_to_exclude_npu), and | ||
| 68 | + it relies on the stock inductor decomposition table. The compiler-tuning | ||
| 69 | + patches below are therefore either meaningless or actively harmful for it, | ||
| 70 | + so callers skip them when this returns True. | ||
| 71 | + """ | ||
| 72 | + return os.environ.get("TORCHINDUCTOR_NPU_BACKEND") == "triton_experimental" | ||
| 73 | + | ||
| 74 | + | ||
| 59 | def _hf_t5_mt5_conditionalgeneration_forward_new( | 75 | def _hf_t5_mt5_conditionalgeneration_forward_new( |
| 60 | self, | 76 | self, |
| 61 | hidden_states, | 77 | hidden_states, |
| @@ -316,52 +332,121 @@ def _patch_model_7(): | |||
| 316 | 332 | ||
| 317 | def _patch_model_10(): | 333 | def _patch_model_10(): |
| 318 | try: | 334 | try: |
| 319 | - import torch | 335 | + from torchbenchmark.models.nvidia_deeprecommender import ( |
| 320 | - | 336 | + nvinfer, |
| 321 | - import torch_npu._inductor # noqa: F401 | 337 | + nvtrain, |
| 322 | - # from torch_npu.contrib import transfer_to_npu # noqa: F401 | ||
| 323 | - except ImportError: | ||
| 324 | - log.warning("NPU_FlAG is False!") | ||
| 325 | - return | ||
| 326 | - | ||
| 327 | - try: | ||
| 328 | - from torchbenchmark.models.nvidia_deeprecommender.nvtrain import ( | ||
| 329 | - DeepRecommenderTrainBenchmark, | ||
| 330 | ) | 338 | ) |
| 331 | - from torchbenchmark.models.nvidia_deeprecommender.reco_encoder.model import model | ||
| 332 | except ImportError: | 339 | except ImportError: |
| 333 | log.warning( | 340 | log.warning( |
| 334 | - "Import nvidia_deeprecommender failed or could not get DeepRecommenderTrainBenchmark" | 341 | + "Import nvidia_deeprecommender failed; the NPU compatibility patch " |
| 335 | - "from module torchbenchmark.models.nvidia_deeprecommender.nvtrain.DeepRecommenderTrainBenchmark" | 342 | + "was not applied" |
| 336 | ) | 343 | ) |
| 337 | return | 344 | return |
| 338 | 345 | ||
| 339 | - def new_init( | 346 | + train_cls = nvtrain.DeepRecommenderTrainBenchmark |
| 340 | - self, device="cpu", jit=False, batch_size=256, process_command_line=False | 347 | + inference_cls = nvinfer.DeepRecommenderInferenceBenchmark |
| 348 | + if getattr(train_cls, "_npu_patch_applied", False): | ||
| 349 | + return | ||
| 350 | + | ||
| 351 | + original_train_init = train_cls.__init__ | ||
| 352 | + original_inference_init = inference_cls.__init__ | ||
| 353 | + | ||
| 354 | + def reset_optimizer(self): | ||
| 355 | + if self.args.optimizer == "adam": | ||
| 356 | + self.optimizer = nvtrain.optim.Adam( | ||
| 357 | + self.rencoder.parameters(), | ||
| 358 | + lr=self.args.lr, | ||
| 359 | + weight_decay=self.args.weight_decay, | ||
| 360 | + ) | ||
| 361 | + elif self.args.optimizer == "adagrad": | ||
| 362 | + self.optimizer = nvtrain.optim.Adagrad( | ||
| 363 | + self.rencoder.parameters(), | ||
| 364 | + lr=self.args.lr, | ||
| 365 | + weight_decay=self.args.weight_decay, | ||
| 366 | + ) | ||
| 367 | + elif self.args.optimizer == "momentum": | ||
| 368 | + self.optimizer = nvtrain.optim.SGD( | ||
| 369 | + self.rencoder.parameters(), | ||
| 370 | + lr=self.args.lr, | ||
| 371 | + momentum=0.9, | ||
| 372 | + weight_decay=self.args.weight_decay, | ||
| 373 | + ) | ||
| 374 | + self.scheduler = nvtrain.MultiStepLR( | ||
| 375 | + self.optimizer, | ||
| 376 | + milestones=[24, 36, 48, 66, 72], | ||
| 377 | + gamma=0.5, | ||
| 378 | + ) | ||
| 379 | + elif self.args.optimizer == "rmsprop": | ||
| 380 | + self.optimizer = nvtrain.optim.RMSprop( | ||
| 381 | + self.rencoder.parameters(), | ||
| 382 | + lr=self.args.lr, | ||
| 383 | + momentum=0.9, | ||
| 384 | + weight_decay=self.args.weight_decay, | ||
| 385 | + ) | ||
| 386 | + else: | ||
| 387 | + raise ValueError(f"Unknown optimizer kind: {self.args.optimizer}") | ||
| 388 | + | ||
| 389 | + def train_init( | ||
| 390 | + self, device="cpu", jit=False, batch_size=256, processCommandLine=False | ||
| 341 | ): | 391 | ): |
| 342 | - self.TrainInit("cuda", jit, batch_size, process_command_line) | 392 | + target_device = torch.device(device) |
| 393 | + if target_device.type != "npu": | ||
| 394 | + return original_train_init( | ||
| 395 | + self, device, jit, batch_size, processCommandLine | ||
| 396 | + ) | ||
| 397 | + | ||
| 398 | + # The upstream benchmark constructor only accepts CPU and CUDA. Build | ||
| 399 | + # its state on CPU first, then move the model and inputs to the actual | ||
| 400 | + # NPU device without relying on transfer_to_npu or CUDA API rewriting. | ||
| 401 | + original_train_init(self, "cpu", jit, batch_size, processCommandLine) | ||
| 402 | + self.device = target_device | ||
| 403 | + | ||
| 404 | + # Keep the existing DVM-friendly shape workaround while making it | ||
| 405 | + # backend-independent. The outer TorchBench runner also uses this model | ||
| 406 | + # with triton_experimental. | ||
| 407 | + if hasattr(self, "scheduler"): | ||
| 408 | + del self.scheduler | ||
| 409 | + del self.optimizer | ||
| 410 | + del self.rencoder | ||
| 411 | + del self.toyinputs | ||
| 412 | + gc.collect() | ||
| 343 | 413 | ||
| 344 | self.toyvocab = 197952 | 414 | self.toyvocab = 197952 |
| 345 | self.toyinputs = torch.randn(self.toybatch, self.toyvocab) | 415 | self.toyinputs = torch.randn(self.toybatch, self.toyvocab) |
| 346 | if self.toytest: | 416 | if self.toytest: |
| 347 | - self.rencoder = model.AutoEncoder( | 417 | + self.rencoder = nvtrain.model.AutoEncoder( |
| 348 | - layer_sizes=[self.toyvocab] + [int(l) for l in self.args.hidden_layers.split(',')], | 418 | + layer_sizes=[self.toyvocab] |
| 419 | + + [int(layer) for layer in self.args.hidden_layers.split(",")], | ||
| 349 | nl_type=self.args.non_linearity_type, | 420 | nl_type=self.args.non_linearity_type, |
| 350 | is_constrained=self.args.constrained, | 421 | is_constrained=self.args.constrained, |
| 351 | dp_drop_prob=self.args.drop_prob, | 422 | dp_drop_prob=self.args.drop_prob, |
| 352 | last_layer_activations=not self.args.skip_last_layer_nl, | 423 | last_layer_activations=not self.args.skip_last_layer_nl, |
| 353 | ) | 424 | ) |
| 354 | 425 | ||
| 355 | - if hasattr(self, "args"): | 426 | + self.args.use_cuda = False |
| 356 | - self.args.use_cuda = True | 427 | + self.rencoder = self.rencoder.to(target_device) |
| 428 | + self.toyinputs = self.toyinputs.to(target_device) | ||
| 429 | + reset_optimizer(self) | ||
| 357 | 430 | ||
| 358 | - if hasattr(self, "rencoder"): | 431 | + def inference_init( |
| 359 | - self.rencoder = self.rencoder.npu() | 432 | + self, device="cpu", jit=False, batch_size=256, usecommandlineargs=False |
| 433 | + ): | ||
| 434 | + target_device = torch.device(device) | ||
| 435 | + if target_device.type != "npu": | ||
| 436 | + return original_inference_init( | ||
| 437 | + self, device, jit, batch_size, usecommandlineargs | ||
| 438 | + ) | ||
| 360 | 439 | ||
| 361 | - if hasattr(self, "toyinputs"): | 440 | + original_inference_init(self, "cpu", jit, batch_size, usecommandlineargs) |
| 362 | - self.toyinputs = self.toyinputs.to("npu") | 441 | + self.device = target_device |
| 442 | + self.args.use_cuda = False | ||
| 443 | + self.rencoder = self.rencoder.to(target_device) | ||
| 444 | + self.toyinputs = self.toyinputs.to(target_device) | ||
| 363 | 445 | ||
| 364 | - DeepRecommenderTrainBenchmark.__init__ = new_init | 446 | + train_cls.__init__ = train_init |
| 447 | + inference_cls.__init__ = inference_init | ||
| 448 | + train_cls._npu_patch_applied = True | ||
| 449 | + inference_cls._npu_patch_applied = True | ||
| 365 | 450 | ||
| 366 | 451 | ||
| 367 | 452 | ||
| @@ -466,6 +551,12 @@ def _patch_model_19(): | |||
| 466 | 551 | ||
| 467 | 552 | ||
| 468 | def patch_remove_ops_from_generate_list(op_names=None): | 553 | def patch_remove_ops_from_generate_list(op_names=None): |
| 554 | + if _is_triton_experimental_backend(): | ||
| 555 | + print( | ||
| 556 | + "[patch] triton_experimental backend: skip GENERATE_LIST tuning " | ||
| 557 | + "(ascend_npu_ir only)." | ||
| 558 | + ) | ||
| 559 | + return | ||
| 469 | try: | 560 | try: |
| 470 | import torch | 561 | import torch |
| 471 | 562 | ||
| @@ -496,6 +587,12 @@ def patch_remove_ops_from_generate_list(op_names=None): | |||
| 496 | 587 | ||
| 497 | 588 | ||
| 498 | def patch_remove_decomposition(op_names=None): | 589 | def patch_remove_decomposition(op_names=None): |
| 590 | + if _is_triton_experimental_backend(): | ||
| 591 | + print( | ||
| 592 | + "[patch] triton_experimental backend: keep stock inductor " | ||
| 593 | + "decompositions (skip removal)." | ||
| 594 | + ) | ||
| 595 | + return | ||
| 499 | try: | 596 | try: |
| 500 | from torch._decomp import remove_decompositions | 597 | from torch._decomp import remove_decompositions |
| 501 | from torch._inductor import decomposition as inductor_decomp | 598 | from torch._inductor import decomposition as inductor_decomp |
| @@ -518,6 +615,78 @@ def patch_remove_decomposition(op_names=None): | |||
| 518 | print(f"[patch] Failed to remove decompositions from inductor: {e}") | 615 | print(f"[patch] Failed to remove decompositions from inductor: {e}") |
| 519 | 616 | ||
| 520 | 617 | ||
| 618 | +def patch_force_fallback_kernels(kernel_names=None): | ||
| 619 | + """Force ascend_npu_ir to fall back on the given fused kernels. | ||
| 620 | + | ||
| 621 | + Only meaningful for the mlir/dvm (ascend_npu_ir) backend; the | ||
| 622 | + triton_experimental backend does not consume this config, so skip it there. | ||
| 623 | + """ | ||
| 624 | + if _is_triton_experimental_backend(): | ||
| 625 | + print( | ||
| 626 | + "[patch] triton_experimental backend: skip force_fallback_kernel_names " | ||
| 627 | + "(ascend_npu_ir only)." | ||
| 628 | + ) | ||
| 629 | + return | ||
| 630 | + if not kernel_names: | ||
| 631 | + print("[patch] No kernel names provided, nothing to do.") | ||
| 632 | + return | ||
| 633 | + try: | ||
| 634 | + from torch_npu._inductor.ascend_npu_ir.ascend_npu_ir import ( | ||
| 635 | + config as anir_config, | ||
| 636 | + ) | ||
| 637 | + | ||
| 638 | + for name in kernel_names: | ||
| 639 | + anir_config.force_fallback_kernel_names[name] = True | ||
| 640 | + print(f"[patch] Forced fallback for {len(kernel_names)} kernel(s).") | ||
| 641 | + except ImportError: | ||
| 642 | + log.warning("import ascend_npu_ir config failed for force_fallback patch") | ||
| 643 | + | ||
| 644 | + | ||
| 645 | +def patch_exclude_decomps_npu(op_names=None): | ||
| 646 | + """Exclude ops from ascend_npu_ir decomposition and the inductor table. | ||
| 647 | + | ||
| 648 | + The decomps_to_exclude_npu list is ascend_npu_ir-specific, and the | ||
| 649 | + accompanying remove_decompositions() call mutates the global inductor | ||
| 650 | + decomposition table that triton_experimental relies on, so skip both there. | ||
| 651 | + """ | ||
| 652 | + if _is_triton_experimental_backend(): | ||
| 653 | + print( | ||
| 654 | + "[patch] triton_experimental backend: keep stock decompositions " | ||
| 655 | + "(skip decomps_to_exclude_npu)." | ||
| 656 | + ) | ||
| 657 | + return | ||
| 658 | + if not op_names: | ||
| 659 | + print("[patch] No op names provided, nothing to do.") | ||
| 660 | + return | ||
| 661 | + try: | ||
| 662 | + from torch._decomp import remove_decompositions | ||
| 663 | + from torch._inductor import decomposition as inductor_decomp | ||
| 664 | + | ||
| 665 | + from torch_npu._inductor.ascend_npu_ir.ascend_npu_ir import ( | ||
| 666 | + config as anir_config, | ||
| 667 | + ) | ||
| 668 | + | ||
| 669 | + ops = [] | ||
| 670 | + for name in op_names: | ||
| 671 | + op = torch.ops | ||
| 672 | + for p in name.split("."): | ||
| 673 | + op = getattr(op, p) | ||
| 674 | + ops.append(op) | ||
| 675 | + | ||
| 676 | + if hasattr(anir_config, "decomps_to_exclude_npu") and isinstance( | ||
| 677 | + anir_config.decomps_to_exclude_npu, list | ||
| 678 | + ): | ||
| 679 | + for op in ops: | ||
| 680 | + if op not in anir_config.decomps_to_exclude_npu: | ||
| 681 | + anir_config.decomps_to_exclude_npu.append(op) | ||
| 682 | + remove_decompositions(inductor_decomp.decompositions, ops) | ||
| 683 | + print(f"[patch] Excluded {len(ops)} decomposition(s) for NPU.") | ||
| 684 | + except Exception: | ||
| 685 | + log.warning( | ||
| 686 | + "import config failed for decomps_to_exclude_npu patch", exc_info=True | ||
| 687 | + ) | ||
| 688 | + | ||
| 689 | + | ||
| 521 | 690 | ||
| 522 | def _patch_model_20(): | 691 | def _patch_model_20(): |
| 523 | import numpy as np | 692 | import numpy as np |
| @@ -570,14 +739,7 @@ def _patch_model_21(): | |||
| 570 | # The current operator suffers from severe performance degradation. | 739 | # The current operator suffers from severe performance degradation. |
| 571 | # This patch will be removed after the issue is fixed in the future. | 740 | # This patch will be removed after the issue is fixed in the future. |
| 572 | patch_remove_decomposition(["aten._softmax"]) | 741 | patch_remove_decomposition(["aten._softmax"]) |
| 573 | - try: | 742 | + patch_force_fallback_kernels(["mlir_fused_add_lt_neg_where_16"]) |
| 574 | - from torch_npu._inductor.ascend_npu_ir.ascend_npu_ir import ( | ||
| 575 | - config as anir_config, | ||
| 576 | - ) | ||
| 577 | - | ||
| 578 | - anir_config.force_fallback_kernel_names["mlir_fused_add_lt_neg_where_16"] = True | ||
| 579 | - except ImportError: | ||
| 580 | - log.warning("import config failed for hf_T5_base patch") | ||
| 581 | 743 | ||
| 582 | 744 | ||
| 583 | 745 | ||
| @@ -585,14 +747,7 @@ def _patch_model_22(): | |||
| 585 | # The current operator suffers from severe performance degradation. | 747 | # The current operator suffers from severe performance degradation. |
| 586 | # This patch will be removed after the issue is fixed in the future. | 748 | # This patch will be removed after the issue is fixed in the future. |
| 587 | patch_remove_decomposition(["aten._softmax"]) | 749 | patch_remove_decomposition(["aten._softmax"]) |
| 588 | - try: | 750 | + patch_force_fallback_kernels(["mlir_fused_add_lt_neg_where_16"]) |
| 589 | - from torch_npu._inductor.ascend_npu_ir.ascend_npu_ir import ( | ||
| 590 | - config as anir_config, | ||
| 591 | - ) | ||
| 592 | - | ||
| 593 | - anir_config.force_fallback_kernel_names["mlir_fused_add_lt_neg_where_16"] = True | ||
| 594 | - except ImportError: | ||
| 595 | - log.warning("import config failed for hf_T5_large patch") | ||
| 596 | 751 | ||
| 597 | 752 | ||
| 598 | 753 | ||
| @@ -611,26 +766,12 @@ def _patch_squeezenet1_1(): | |||
| 611 | fallbackdiv | 766 | fallbackdiv |
| 612 | """ | 767 | """ |
| 613 | patch_remove_ops_from_generate_list(["aten.div"]) | 768 | patch_remove_ops_from_generate_list(["aten.div"]) |
| 614 | - try: | 769 | + patch_force_fallback_kernels(["mlir_fused_relu_threshold_backward_2"]) |
| 615 | - from torch_npu._inductor.ascend_npu_ir.ascend_npu_ir import ( | ||
| 616 | - config as anir_config, | ||
| 617 | - ) | ||
| 618 | - | ||
| 619 | - anir_config.force_fallback_kernel_names["mlir_fused_relu_threshold_backward_2"] = True | ||
| 620 | - except ImportError: | ||
| 621 | - log.warning("import config failed for squeezenet1_1 patch") | ||
| 622 | 770 | ||
| 623 | 771 | ||
| 624 | 772 | ||
| 625 | def _patch_model_24(): | 773 | def _patch_model_24(): |
| 626 | - try: | 774 | + patch_force_fallback_kernels(["mlir_fused_add_lt_neg_where_13"]) |
| 627 | - from torch_npu._inductor.ascend_npu_ir.ascend_npu_ir import ( | ||
| 628 | - config as anir_config, | ||
| 629 | - ) | ||
| 630 | - | ||
| 631 | - anir_config.force_fallback_kernel_names["mlir_fused_add_lt_neg_where_13"] = True | ||
| 632 | - except ImportError: | ||
| 633 | - log.warning("import config failed for T5ForConditionalGeneration patch") | ||
| 634 | from torch._higher_order_ops.effects import ( | 775 | from torch._higher_order_ops.effects import ( |
| 635 | _EffectType, | 776 | _EffectType, |
| 636 | _register_effectful_op, | 777 | _register_effectful_op, |
| @@ -644,30 +785,12 @@ def _patch_model_24(): | |||
| 644 | 785 | ||
| 645 | 786 | ||
| 646 | def _patch_model_25(): | 787 | def _patch_model_25(): |
| 647 | - try: | 788 | + patch_exclude_decomps_npu( |
| 648 | - import torch | 789 | + [ |
| 649 | - from torch._decomp import remove_decompositions | 790 | + "aten.native_layer_norm", |
| 650 | - from torch._inductor import decomposition as inductor_decomp | 791 | + "aten.native_layer_norm_backward", |
| 651 | - | ||
| 652 | - from torch_npu._inductor.ascend_npu_ir.ascend_npu_ir import ( | ||
| 653 | - config as anir_config, | ||
| 654 | - ) | ||
| 655 | - | ||
| 656 | - aten = torch.ops.aten | ||
| 657 | - ops_to_add = [ | ||
| 658 | - aten.native_layer_norm, | ||
| 659 | - aten.native_layer_norm_backward, | ||
| 660 | ] | 792 | ] |
| 661 | - | 793 | + ) |
| 662 | - if hasattr(anir_config, "decomps_to_exclude_npu") and isinstance( | ||
| 663 | - anir_config.decomps_to_exclude_npu, list | ||
| 664 | - ): | ||
| 665 | - for op in ops_to_add: | ||
| 666 | - if op not in anir_config.decomps_to_exclude_npu: | ||
| 667 | - anir_config.decomps_to_exclude_npu.append(op) | ||
| 668 | - remove_decompositions(inductor_decomp.decompositions, ops_to_add) | ||
| 669 | - except Exception: | ||
| 670 | - log.warning("import config failed for BartForCausalLM patch", exc_info=True) | ||
| 671 | 794 | ||
| 672 | 795 | ||
| 673 | 796 | ||
[P1] 这里的预热调用没有区分 run_mode。run_mode is None 的 accuracy/普通执行随后仍会进入 range(n - 1) 并再执行一次 collect_outputs=True，因此总共执行 n+1 次；训练模式还会额外更新一次权重，改变准确率和最终状态。请只在 profiler 性能分支预热，或相应扣减后续迭代，并补 accuracy/train 覆盖。