[Fix]fix multiprocessing fails to launch on Ascend NPU

Author: zhouzaida

.github/workflows/merge_stage_test.yml

@@ -12,6 +12,7 @@ on:
       - "CONTRIBUTING_zh-CN.md"
       - ".pre-commit-config.yaml"
       - ".pre-commit-config-zh-cn.yaml"
+      - "examples/**"
     branches:
       - main
 

.github/workflows/pr_stage_test.yml

@@ -12,6 +12,7 @@ on:
       - "CONTRIBUTING_zh-CN.md"
       - ".pre-commit-config.yaml"
       - ".pre-commit-config-zh-cn.yaml"
+      - "examples/**"
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.ref }}

mmengine/device/utils.py

@@ -1,5 +1,4 @@
 # Copyright (c) OpenMMLab. All rights reserved.
-import os
 from typing import Optional
 
 import torch
@@ -8,10 +7,6 @@
     import torch_npu  # noqa: F401
     import torch_npu.npu.utils as npu_utils
 
-    # Enable operator support for dynamic shape and
-    # binary operator support on the NPU.
-    npu_jit_compile = bool(os.getenv('NPUJITCompile', False))
-    torch.npu.set_compile_mode(jit_compile=npu_jit_compile)
     IS_NPU_AVAILABLE = hasattr(torch, 'npu') and torch.npu.is_available()
 except Exception:
     IS_NPU_AVAILABLE = False

mmengine/model/base_model/base_model.py

@@ -1,4 +1,5 @@
 # Copyright (c) OpenMMLab. All rights reserved.
+import os
 from abc import abstractmethod
 from collections import OrderedDict
 from typing import Dict, Optional, Tuple, Union
@@ -267,6 +268,12 @@ def _set_device(self, device: torch.device) -> None:
                 buffers in this module.
         """
 
+        if device.type == 'npu':
+            # Enable operator support for dynamic shape and
+            # binary operator support on the NPU.
+            npu_jit_compile = bool(os.getenv('NPUJITCompile', False))
+            torch.npu.set_compile_mode(jit_compile=npu_jit_compile)
+
         def apply_fn(module):
             if not isinstance(module, BaseDataPreprocessor):
                 return