feat: add comprehensive XPU and MPS support for PyTorch models

examples/benchmarks/LSTM/workflow_config_lstm_Alpha158.yaml

@@ -62,11 +62,11 @@ task:
             n_epochs: 200
             lr: 1e-3
             early_stop: 10
-            batch_size: 800
+            batch_size: 2000
             metric: loss
             loss: mse
-            n_jobs: 20
-            GPU: 0
+            n_jobs: 6
+            GPU: "xpu"
     dataset:
         class: TSDatasetH
         module_path: qlib.data.dataset

qlib/_version.py

@@ -0,0 +1,34 @@
+# file generated by setuptools-scm
+# don't change, don't track in version control
+
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
+
+TYPE_CHECKING = False
+if TYPE_CHECKING:
+    from typing import Tuple
+    from typing import Union
+
+    VERSION_TUPLE = Tuple[Union[int, str], ...]
+    COMMIT_ID = Union[str, None]
+else:
+    VERSION_TUPLE = object
+    COMMIT_ID = object
+
+version: str
+__version__: str
+__version_tuple__: VERSION_TUPLE
+version_tuple: VERSION_TUPLE
+commit_id: COMMIT_ID
+__commit_id__: COMMIT_ID
+
+__version__ = version = '0.9.8.dev18'
+__version_tuple__ = version_tuple = (0, 9, 8, 'dev18')
+
+__commit_id__ = commit_id = 'g477160e4a'

qlib/contrib/model/pytorch_adarnn.py

@@ -12,7 +12,7 @@
 import torch.nn.functional as F
 import torch.optim as optim
 from torch.autograd import Function
-from qlib.contrib.model.pytorch_utils import count_parameters
+from qlib.contrib.model.pytorch_utils import count_parameters, empty_cache
 from qlib.data.dataset import DatasetH
 from qlib.data.dataset.handler import DataHandlerLP
 from qlib.log import get_module_logger
@@ -81,7 +81,10 @@ def __init__(
         self.optimizer = optimizer.lower()
         self.loss = loss
         self.n_splits = n_splits
-        self.device = torch.device("cuda:%d" % GPU if torch.cuda.is_available() and GPU >= 0 else "cpu")
+        if isinstance(GPU, str):
+            self.device = torch.device(GPU)
+        else:
+            self.device = torch.device("cuda:%d" % GPU if torch.cuda.is_available() and GPU >= 0 else "cpu")
         self.seed = seed
 
         self.logger.info(
@@ -298,7 +301,7 @@ def fit(
         torch.save(best_param, save_path)
 
         if self.use_gpu:
-            torch.cuda.empty_cache()
+            empty_cache(self.device)
         return best_score
 
     def predict(self, dataset: DatasetH, segment: Union[Text, slice] = "test"):
@@ -396,7 +399,10 @@ def __init__(
         self.model_type = model_type
         self.trans_loss = trans_loss
         self.len_seq = len_seq
-        self.device = torch.device("cuda:%d" % GPU if torch.cuda.is_available() and GPU >= 0 else "cpu")
+        if isinstance(GPU, str):
+            self.device = torch.device(GPU)
+        else:
+            self.device = torch.device("cuda:%d" % GPU if torch.cuda.is_available() and GPU >= 0 else "cpu")
         in_size = self.n_input
 
         features = nn.ModuleList()
@@ -558,7 +564,10 @@ def __init__(self, loss_type="cosine", input_dim=512, GPU=0):
         """
         self.loss_type = loss_type
         self.input_dim = input_dim
-        self.device = torch.device("cuda:%d" % GPU if torch.cuda.is_available() and GPU >= 0 else "cpu")
+        if isinstance(GPU, str):
+            self.device = torch.device(GPU)
+        else:
+            self.device = torch.device("cuda:%d" % GPU if torch.cuda.is_available() and GPU >= 0 else "cpu")
 
     def compute(self, X, Y):
         """Compute adaptation loss

qlib/contrib/model/pytorch_add.py

@@ -83,7 +83,10 @@ def __init__(
         self.optimizer = optimizer.lower()
         self.base_model = base_model
         self.model_path = model_path
-        self.device = torch.device("cuda:%d" % (GPU) if torch.cuda.is_available() and GPU >= 0 else "cpu")
+        if isinstance(GPU, str):
+            self.device = torch.device(GPU)
+        else:
+            self.device = torch.device("cuda:%d" % (GPU) if torch.cuda.is_available() and GPU >= 0 else "cpu")
         self.seed = seed
 
         self.gamma = gamma
@@ -414,7 +417,7 @@ def fit(
         save_path = get_or_create_path(save_path)
         torch.save(best_param, save_path)
         if self.use_gpu:
-            torch.cuda.empty_cache()
+            empty_cache(self.device)
 
     def predict(self, dataset: DatasetH, segment: Union[Text, slice] = "test"):
         x_test = dataset.prepare(segment, col_set="feature", data_key=DataHandlerLP.DK_I)

qlib/contrib/model/pytorch_alstm.py

@@ -11,6 +11,7 @@
 import copy
 from ...utils import get_or_create_path
 from ...log import get_module_logger
+from .pytorch_utils import empty_cache
 
 import torch
 import torch.nn as nn
@@ -70,7 +71,10 @@ def __init__(
         self.early_stop = early_stop
         self.optimizer = optimizer.lower()
         self.loss = loss
-        self.device = torch.device("cuda:%d" % (GPU) if torch.cuda.is_available() and GPU >= 0 else "cpu")
+        if isinstance(GPU, str):
+            self.device = torch.device(GPU)
+        else:
+            self.device = torch.device("cuda:%d" % (GPU) if torch.cuda.is_available() and GPU >= 0 else "cpu")
         self.seed = seed
 
         self.logger.info(
@@ -262,7 +266,7 @@ def fit(
         torch.save(best_param, save_path)
 
         if self.use_gpu:
-            torch.cuda.empty_cache()
+            empty_cache(self.device)
 
     def predict(self, dataset: DatasetH, segment: Union[Text, slice] = "test"):
         if not self.fitted:

qlib/contrib/model/pytorch_alstm_ts.py

@@ -11,6 +11,7 @@
 import copy
 from ...utils import get_or_create_path
 from ...log import get_module_logger
+from .pytorch_utils import empty_cache
 
 import torch
 import torch.nn as nn
@@ -74,7 +75,10 @@ def __init__(
         self.early_stop = early_stop
         self.optimizer = optimizer.lower()
         self.loss = loss
-        self.device = torch.device("cuda:%d" % (GPU) if torch.cuda.is_available() and GPU >= 0 else "cpu")
+        if isinstance(GPU, str):
+            self.device = torch.device(GPU)
+        else:
+            self.device = torch.device("cuda:%d" % (GPU) if torch.cuda.is_available() and GPU >= 0 else "cpu")
         self.n_jobs = n_jobs
         self.seed = seed
 
@@ -282,7 +286,7 @@ def fit(
         torch.save(best_param, save_path)
 
         if self.use_gpu:
-            torch.cuda.empty_cache()
+            empty_cache(self.device)
 
     def predict(self, dataset: DatasetH, segment: Union[Text, slice] = "test"):
         if not self.fitted:

qlib/contrib/model/pytorch_gats.py

@@ -11,6 +11,7 @@
 import copy
 from ...utils import get_or_create_path
 from ...log import get_module_logger
+from .pytorch_utils import empty_cache
 import torch
 import torch.nn as nn
 import torch.optim as optim
@@ -75,7 +76,10 @@ def __init__(
         self.loss = loss
         self.base_model = base_model
         self.model_path = model_path
-        self.device = torch.device("cuda:%d" % (GPU) if torch.cuda.is_available() and GPU >= 0 else "cpu")
+        if isinstance(GPU, str):
+            self.device = torch.device(GPU)
+        else:
+            self.device = torch.device("cuda:%d" % (GPU) if torch.cuda.is_available() and GPU >= 0 else "cpu")
         self.seed = seed
 
         self.logger.info(
@@ -296,7 +300,7 @@ def fit(
         torch.save(best_param, save_path)
 
         if self.use_gpu:
-            torch.cuda.empty_cache()
+            empty_cache(self.device)
 
     def predict(self, dataset: DatasetH, segment: Union[Text, slice] = "test"):
         if not self.fitted:

qlib/contrib/model/pytorch_gats_ts.py

@@ -10,6 +10,7 @@
 import copy
 from ...utils import get_or_create_path
 from ...log import get_module_logger
+from .pytorch_utils import empty_cache
 import torch
 import torch.nn as nn
 import torch.optim as optim
@@ -94,7 +95,10 @@ def __init__(
         self.loss = loss
         self.base_model = base_model
         self.model_path = model_path
-        self.device = torch.device("cuda:%d" % (GPU) if torch.cuda.is_available() and GPU >= 0 else "cpu")
+        if isinstance(GPU, str):
+            self.device = torch.device(GPU)
+        else:
+            self.device = torch.device("cuda:%d" % (GPU) if torch.cuda.is_available() and GPU >= 0 else "cpu")
         self.n_jobs = n_jobs
         self.seed = seed
 
@@ -310,7 +314,7 @@ def fit(
         torch.save(best_param, save_path)
 
         if self.use_gpu:
-            torch.cuda.empty_cache()
+            empty_cache(self.device)
 
     def predict(self, dataset):
         if not self.fitted:

qlib/contrib/model/pytorch_general_nn.py

@@ -26,6 +26,7 @@
     get_or_create_path,
 )
 from ...log import get_module_logger
+from .pytorch_utils import empty_cache
 
 from ...model.utils import ConcatDataset
 
@@ -83,7 +84,10 @@ def __init__(
         self.optimizer = optimizer.lower()
         self.loss = loss
         self.weight_decay = weight_decay
-        self.device = torch.device("cuda:%d" % (GPU) if torch.cuda.is_available() and GPU >= 0 else "cpu")
+        if isinstance(GPU, str):
+            self.device = torch.device(GPU)
+        else:
+            self.device = torch.device("cuda:%d" % (GPU) if torch.cuda.is_available() and GPU >= 0 else "cpu")
         self.n_jobs = n_jobs
         self.seed = seed
 
@@ -329,7 +333,7 @@ def fit(
         torch.save(best_param, save_path)
 
         if self.use_gpu:
-            torch.cuda.empty_cache()
+            empty_cache(self.device)
 
     def predict(
         self,

qlib/contrib/model/pytorch_gru.py

@@ -17,6 +17,7 @@
 from ...data.dataset import DatasetH
 from ...data.dataset.handler import DataHandlerLP
 from ...log import get_module_logger
+from .pytorch_utils import empty_cache
 from ...model.base import Model
 from ...utils import get_or_create_path
 from .pytorch_utils import count_parameters
@@ -70,7 +71,10 @@ def __init__(
         self.early_stop = early_stop
         self.optimizer = optimizer.lower()
         self.loss = loss
-        self.device = torch.device("cuda:%d" % (GPU) if torch.cuda.is_available() and GPU >= 0 else "cpu")
+        if isinstance(GPU, str):
+            self.device = torch.device(GPU)
+        else:
+            self.device = torch.device("cuda:%d" % (GPU) if torch.cuda.is_available() and GPU >= 0 else "cpu")
         self.seed = seed
 
         self.logger.info(
@@ -287,7 +291,7 @@ def fit(
                 rec.log_metrics(step=i, **{k: v})
 
         if self.use_gpu:
-            torch.cuda.empty_cache()
+            empty_cache(self.device)
 
     def predict(self, dataset: DatasetH, segment: Union[Text, slice] = "test"):
         if not self.fitted:

qlib/contrib/model/pytorch_gru_ts.py

@@ -10,6 +10,7 @@
 import copy
 from ...utils import get_or_create_path
 from ...log import get_module_logger
+from .pytorch_utils import empty_cache
 
 import torch
 import torch.nn as nn
@@ -72,7 +73,10 @@ def __init__(
         self.early_stop = early_stop
         self.optimizer = optimizer.lower()
         self.loss = loss
-        self.device = torch.device("cuda:%d" % (GPU) if torch.cuda.is_available() and GPU >= 0 else "cpu")
+        if isinstance(GPU, str):
+            self.device = torch.device(GPU)
+        else:
+            self.device = torch.device("cuda:%d" % (GPU) if torch.cuda.is_available() and GPU >= 0 else "cpu")
         self.n_jobs = n_jobs
         self.seed = seed
 
@@ -276,7 +280,7 @@ def fit(
         torch.save(best_param, save_path)
 
         if self.use_gpu:
-            torch.cuda.empty_cache()
+            empty_cache(self.device)
 
     def predict(self, dataset):
         if not self.fitted:

qlib/contrib/model/pytorch_hist.py

@@ -13,6 +13,7 @@
 import copy
 from ...utils import get_or_create_path
 from ...log import get_module_logger
+from .pytorch_utils import empty_cache
 import torch
 import torch.nn as nn
 import torch.optim as optim
@@ -80,7 +81,10 @@ def __init__(
         self.model_path = model_path
         self.stock2concept = stock2concept
         self.stock_index = stock_index
-        self.device = torch.device("cuda:%d" % (GPU) if torch.cuda.is_available() and GPU >= 0 else "cpu")
+        if isinstance(GPU, str):
+            self.device = torch.device(GPU)
+        else:
+            self.device = torch.device("cuda:%d" % (GPU) if torch.cuda.is_available() and GPU >= 0 else "cpu")
         self.seed = seed
 
         self.logger.info(
@@ -326,6 +330,8 @@ def fit(
         self.logger.info("best score: %.6lf @ %d" % (best_score, best_epoch))
         self.HIST_model.load_state_dict(best_param)
         torch.save(best_param, save_path)
+        if self.use_gpu:
+            empty_cache(self.device)
 
     def predict(self, dataset: DatasetH, segment: Union[Text, slice] = "test"):
         if not self.fitted:

qlib/contrib/model/pytorch_igmtf.py

@@ -11,6 +11,7 @@
 import copy
 from ...utils import get_or_create_path
 from ...log import get_module_logger
+from .pytorch_utils import empty_cache
 
 import torch
 import torch.nn as nn
@@ -74,7 +75,10 @@ def __init__(
         self.loss = loss
         self.base_model = base_model
         self.model_path = model_path
-        self.device = torch.device("cuda:%d" % (GPU) if torch.cuda.is_available() and GPU >= 0 else "cpu")
+        if isinstance(GPU, str):
+            self.device = torch.device(GPU)
+        else:
+            self.device = torch.device("cuda:%d" % (GPU) if torch.cuda.is_available() and GPU >= 0 else "cpu")
         self.seed = seed
 
         self.logger.info(
@@ -322,7 +326,7 @@ def fit(
         torch.save(best_param, save_path)
 
         if self.use_gpu:
-            torch.cuda.empty_cache()
+            empty_cache(self.device)
 
     def predict(self, dataset: DatasetH, segment: Union[Text, slice] = "test"):
         if not self.fitted: