Automatically configure device name for different accelerators. For example, on Ascend NPU, this function defaults the trainer device to "npu" unless explicitly set to "cpu". Args: config: Configuration object with trainer.device attribute.
(config)
| 143 | |
| 144 | |
| 145 | def auto_set_device(config) -> None: |
| 146 | """Automatically configure device name for different accelerators. |
| 147 | |
| 148 | For example, on Ascend NPU, this function defaults the trainer device to "npu" |
| 149 | unless explicitly set to "cpu". |
| 150 | |
| 151 | Args: |
| 152 | config: Configuration object with trainer.device attribute. |
| 153 | """ |
| 154 | if config and hasattr(config, "trainer") and hasattr(config.trainer, "device"): |
| 155 | if is_torch_npu_available(): |
| 156 | if config.trainer.device not in ["cpu", "npu"]: |
| 157 | logger.warning( |
| 158 | f"Detect setting config.trainer.device to {config.trainer.device} for Ascend NPU, maybe" |
| 159 | f"from default value in config file, automatically set to `npu` instead." |
| 160 | ) |
| 161 | |
| 162 | config.trainer.device = "npu" |
| 163 | # Other cases: set device to "cuda" via config file, no need to change. |
| 164 | |
| 165 | |
| 166 | def get_device_capability(device_id: int = 0) -> tuple[int | None, int | None]: |