| from __future__ import annotations |
|
|
| import os |
| from dataclasses import dataclass |
|
|
| os.environ.setdefault("PYTORCH_ENABLE_MPS_FALLBACK", "1") |
|
|
| import torch |
|
|
|
|
| DEVICE_ENV = "R1_MATH_DEVICE" |
| DTYPE_ENV = "R1_MATH_DTYPE" |
|
|
|
|
| @dataclass(frozen=True) |
| class RuntimeDevice: |
| device: torch.device |
| dtype: torch.dtype |
|
|
| @property |
| def type(self) -> str: |
| return self.device.type |
|
|
| @property |
| def label(self) -> str: |
| if self.type == "cuda": |
| name = torch.cuda.get_device_name(self.device) |
| return f"CUDA - {name}" |
| if self.type == "mps": |
| return "MPS - Apple GPU" |
| return "CPU" |
|
|
|
|
| def is_mps_available() -> bool: |
| backend = getattr(torch.backends, "mps", None) |
| return bool(backend and backend.is_available()) |
|
|
|
|
| def resolve_device(preferred: str | None = None) -> torch.device: |
| requested = (preferred or os.getenv(DEVICE_ENV) or "auto").strip().lower() |
|
|
| if requested in {"auto", ""}: |
| if torch.cuda.is_available(): |
| return torch.device("cuda") |
| if is_mps_available(): |
| return torch.device("mps") |
| return torch.device("cpu") |
|
|
| if requested.startswith("cuda"): |
| if not torch.cuda.is_available(): |
| raise RuntimeError(f"{DEVICE_ENV}={requested!r}, but CUDA is not available.") |
| return torch.device(requested) |
|
|
| if requested == "mps": |
| if not is_mps_available(): |
| raise RuntimeError(f"{DEVICE_ENV}=mps, but PyTorch MPS is not available.") |
| return torch.device("mps") |
|
|
| if requested == "cpu": |
| return torch.device("cpu") |
|
|
| raise ValueError(f"Unsupported {DEVICE_ENV}: {requested!r}. Use auto, cuda, mps, or cpu.") |
|
|
|
|
| def _dtype_from_name(name: str) -> torch.dtype: |
| normalized = name.strip().lower() |
| if normalized in {"float16", "fp16", "half"}: |
| return torch.float16 |
| if normalized in {"bfloat16", "bf16"}: |
| return torch.bfloat16 |
| if normalized in {"float32", "fp32", "full"}: |
| return torch.float32 |
| raise ValueError(f"Unsupported {DTYPE_ENV}: {name!r}. Use auto, float16, bfloat16, or float32.") |
|
|
|
|
| def default_dtype_for_device(device: torch.device) -> torch.dtype: |
| requested = (os.getenv(DTYPE_ENV) or "auto").strip().lower() |
| if requested not in {"auto", ""}: |
| return _dtype_from_name(requested) |
|
|
| if device.type == "cuda": |
| return torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16 |
| if device.type == "mps": |
| return torch.float16 |
| return torch.float32 |
|
|
|
|
| def get_runtime_device(preferred: str | None = None) -> RuntimeDevice: |
| device = resolve_device(preferred) |
| return RuntimeDevice(device=device, dtype=default_dtype_for_device(device)) |
|
|
|
|
| def dtype_name(dtype: torch.dtype) -> str: |
| return str(dtype).removeprefix("torch.") |
|
|
|
|
| def runtime_metadata(preferred: str | None = None) -> dict[str, str]: |
| runtime = get_runtime_device(preferred) |
| return { |
| "device": runtime.type, |
| "device_label": runtime.label, |
| "torch_dtype": dtype_name(runtime.dtype), |
| } |
|
|
|
|
| def mps_safe_training_defaults(device: torch.device) -> dict[str, bool | str]: |
| if device.type == "mps": |
| return { |
| "bf16": False, |
| "fp16": False, |
| "tf32": False, |
| "optim": "adamw_torch", |
| "dataloader_pin_memory": False, |
| "use_mps_device": True, |
| } |
|
|
| if device.type == "cuda": |
| return { |
| "bf16": torch.cuda.is_bf16_supported(), |
| "fp16": not torch.cuda.is_bf16_supported(), |
| "tf32": True, |
| "optim": "paged_adamw_8bit", |
| "dataloader_pin_memory": True, |
| "use_mps_device": False, |
| } |
|
|
| return { |
| "bf16": False, |
| "fp16": False, |
| "tf32": False, |
| "optim": "adamw_torch", |
| "dataloader_pin_memory": False, |
| "use_mps_device": False, |
| } |
|
|
