feat: add mps support

Author: mjun0812

config/__base__/config.yaml

@@ -27,6 +27,7 @@ seed: 42
 
 # Device Params
 gpu:
+  device: "cuda"
   use: 0
   multi: false
   use_cudnn: true
@@ -35,7 +36,6 @@ gpu:
   fsdp:
     min_num_params: 100000000
     use_cpu_offload: false
-use_cpu: false
 
 # AutoMixedPrecision
 use_amp: false

script/edit_configs.py

@@ -8,7 +8,7 @@ def main() -> None:
     parser = argparse.ArgumentParser(description="Edit Config")
     parser.add_argument("config_path_or_dir", type=str, help="Path to the config file or dir")
     parser.add_argument(
-        "override", type=str, help="Override the config file ex: batch=2,use_cpu=True"
+        "override", type=str, help="Override the config file ex: batch=2,gpu.device=cpu"
     )
     args = parser.parse_args()
 

script/show_model.py

@@ -7,16 +7,19 @@
     build_batched_transform,
     build_transforms,
 )
+from src.utils import setup_device
 
 
 @ConfigManager.argparse
 def main(cfg: ExperimentConfig) -> None:
     cfg.model.checkpoint = None
 
-    if cfg.use_cpu or not torch.cuda.is_available():
-        device = torch.device("cpu")
-    else:
-        device = torch.device("cuda:0")
+    device = setup_device(
+        device_type=cfg.gpu.device,
+        device_index=cfg.gpu.use,
+        use_cudnn=cfg.gpu.use_cudnn,
+        verbose=False,
+    )
 
     model = build_model(cfg.model, phase="train")
     model = model.to(device)

script/show_transform.py

@@ -21,7 +21,7 @@
 def main(cfg: ExperimentConfig) -> None:
     phase = cfg.get("phase", "train")
     print(f"Phase: {phase}")
-    cfg.use_cpu = True
+    cfg.gpu.device = "cpu"
 
     cfg_dataset = cfg.dataset.get(phase)
     transform = build_transforms(cfg_dataset.transforms)

src/__init__.py

@@ -28,6 +28,8 @@
     build_dataloader,
     build_dataset,
     build_sampler,
+    is_ram_cache_supported,
+    validate_ram_cache_config,
 )
 
 # Evaluation
@@ -104,7 +106,7 @@
     save_model,
     save_model_info,
     save_optimizer,
-    set_device,
+    setup_device,
     time_synchronized,
     worker_init_fn,
 )

src/config/config.py

@@ -187,14 +187,16 @@ class GPUConfig:
     """Configuration for GPU usage and distributed training.
 
     Attributes:
-        use: GPU device(s) to use (string or int).
+        device: Device type to use ("cuda", "mps", "cpu").
+        use: GPU device(s) to use (string or int). Can be "mps" for MPS.
         multi: Whether to use multiple GPUs.
         use_cudnn: Whether to use cuDNN backend.
         use_tf32: Whether to use TensorFloat-32 format.
         multi_strategy: Strategy for multi-GPU training (dp, ddp, fsdp).
         fsdp: Configuration for FSDP when using fsdp strategy.
     """
 
+    device: str = "cuda"
     use: str | int = 0
     multi: bool = False
     use_cudnn: bool = True
@@ -265,7 +267,6 @@ class ExperimentConfig:
         ram_cache_size_gb: Size of RAM cache in GB.
         seed: Random seed for reproducibility.
         gpu: GPU and distributed training configuration.
-        use_cpu: Whether to force CPU usage.
         use_amp: Whether to use Automatic Mixed Precision.
         amp_dtype: Data type for AMP (fp16 or bf16).
         amp_init_scale: Initial scale for AMP gradient scaler.
@@ -311,7 +312,6 @@ class ExperimentConfig:
 
     # device config
     gpu: GPUConfig = field(default_factory=GPUConfig)
-    use_cpu: bool = False
 
     # AutoMixedPrecision
     use_amp: bool = False

src/dataloaders/__init__.py

@@ -2,7 +2,14 @@
 
 from ..utils import import_submodules
 from .base import BaseDataset
-from .build import DATASET_REGISTRY, build_dataloader, build_dataset, build_sampler
+from .build import (
+    DATASET_REGISTRY,
+    build_dataloader,
+    build_dataset,
+    build_sampler,
+    is_ram_cache_supported,
+    validate_ram_cache_config,
+)
 from .tensor_cache import TensorCache
 from .types import DatasetOutput
 

src/dataloaders/build.py

@@ -1,3 +1,5 @@
+import platform
+
 import torch
 from loguru import logger
 from torch.utils.data import BatchSampler, DataLoader, Dataset
@@ -18,6 +20,35 @@
 """Registry for dataset classes."""
 
 
+def is_ram_cache_supported() -> bool:
+    """Check if RAM cache is supported on the current platform.
+
+    Returns:
+        bool: True if RAM cache is supported (Linux only), False otherwise.
+    """
+    return platform.system() == "Linux"
+
+
+def validate_ram_cache_config(use_ram_cache: bool, verbose: bool = True) -> bool:
+    """Validate RAM cache configuration for current platform.
+
+    Args:
+        use_ram_cache: Whether RAM cache is requested.
+        verbose: Whether to log warnings.
+
+    Returns:
+        bool: True if RAM cache can be used, False otherwise.
+    """
+    if use_ram_cache and not is_ram_cache_supported():
+        if verbose:
+            logger.warning(
+                f"RAM cache is only supported on Linux. Current OS: {platform.system()}. "
+                "Disabling RAM cache."
+            )
+        return False
+    return use_ram_cache
+
+
 def build_dataset(
     cfg: DatasetConfig,
     transforms: Compose | None,
@@ -37,11 +68,14 @@ def build_dataset(
 
     Raises:
         AssertionError: If RAM cache size exceeds available shared memory.
+        RuntimeError: If RAM cache is requested on non-Linux systems.
     """
-    if use_ram_cache:
-        assert ram_cache_size_gb <= get_free_shm_size() / BYTES_PER_GIB, (
-            "RAM Cache size is too large"
-        )
+    # Validate RAM cache configuration
+    use_ram_cache = validate_ram_cache_config(use_ram_cache)
+
+    if use_ram_cache and ram_cache_size_gb is not None:
+        if ram_cache_size_gb > get_free_shm_size() / BYTES_PER_GIB:
+            raise RuntimeError("RAM Cache size is too large")
         cache = TensorCache(size_limit_gb=ram_cache_size_gb)
         logger.info(f"Use RAM Cache: {ram_cache_size_gb}GB")
     else:

src/utils/__init__.py

@@ -18,7 +18,7 @@
     save_model,
     save_model_info,
     save_optimizer,
-    set_device,
+    setup_device,
     time_synchronized,
     worker_init_fn,
 )

src/utils/registry.py

@@ -52,9 +52,10 @@ def _do_register(self, name: str, obj: object) -> None:
         Raises:
             AssertionError: If an object with the same name is already registered.
         """
-        assert name not in self._obj_map, (
-            f"An object named '{name}' was already registered in '{self._name}' registry!"
-        )
+        if name in self._obj_map:
+            raise RuntimeError(
+                f"An object named '{name}' was already registered in '{self._name}' registry!"
+            )
         self._obj_map[name] = obj
 
     def register(self, obj: object | None = None, name: str | None = None) -> object | None: