|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
import torch |
|
|
from typing import Callable |
|
|
from torch import Tensor |
|
|
from packaging import version as pkg_version |
|
|
|
|
|
|
|
|
class OnDevice(object): |
|
|
""" |
|
|
Create modules/tensors w. specific devices and dtypes. Examples: |
|
|
|
|
|
Create MyModule which consists of many different sub-modules and parameters. In this case we can create |
|
|
MyModule as a collection of 'meta' tensors by passing `device='meta'` or we can create the module _directly_ |
|
|
on a CUDA device by passing `device=f'cuda:{local_rank}'` (where `local_rank` is the local GPU id. |
|
|
|
|
|
with OnDevice(dtype=torch.float16, device='meta'): |
|
|
model = MyModel() |
|
|
|
|
|
with OnDevice(dtype=torch.float16, device=f'cuda:{local_rank}'): |
|
|
model = MyModel() |
|
|
|
|
|
""" |
|
|
|
|
|
_orig_torch_empty = torch.empty |
|
|
_orig_torch_zeros = torch.zeros |
|
|
_orig_torch_ones = torch.ones |
|
|
_orig_torch_full = torch.full |
|
|
|
|
|
def __init__(self, dtype, device="meta", enabled=True): |
|
|
self.dtype = dtype |
|
|
self.enabled = enabled |
|
|
self.device = device |
|
|
|
|
|
if device == "meta": |
|
|
if pkg_version.parse('1.10') > pkg_version.parse(torch.__version__): |
|
|
raise NotImplementedError("Meta tensor support is not available, please upgrade to torch 1.10+") |
|
|
|
|
|
def fp_tensor_constructor(self, fn: Callable, target_fp_dtype: torch.dtype) -> Callable: |
|
|
|
|
|
def wrapped_fn(*args, **kwargs) -> Tensor: |
|
|
if kwargs.get("device", None) is None: |
|
|
kwargs['device'] = self.device |
|
|
tensor: Tensor = fn(*args, **kwargs) |
|
|
if tensor.is_floating_point(): |
|
|
tensor = tensor.to(target_fp_dtype) |
|
|
return tensor |
|
|
|
|
|
return wrapped_fn |
|
|
|
|
|
def get_new_tensor_fn_for_dtype(self, dtype: torch.dtype) -> Callable: |
|
|
|
|
|
def new_tensor(cls, *args) -> Tensor: |
|
|
tensor = OnDevice._orig_torch_empty(0, device=self.device).new_empty(*args) |
|
|
if tensor.is_floating_point(): |
|
|
tensor = tensor.to(dtype) |
|
|
return tensor |
|
|
|
|
|
return new_tensor |
|
|
|
|
|
def __enter__(self): |
|
|
if not self.enabled: |
|
|
return |
|
|
torch.Tensor.__old_new__ = torch.Tensor.__new__ |
|
|
torch.Tensor.__new__ = self.get_new_tensor_fn_for_dtype(self.dtype) |
|
|
torch.empty = self.fp_tensor_constructor(self._orig_torch_empty, self.dtype) |
|
|
torch.zeros = self.fp_tensor_constructor(self._orig_torch_zeros, self.dtype) |
|
|
torch.ones = self.fp_tensor_constructor(self._orig_torch_ones, self.dtype) |
|
|
torch.full = self.fp_tensor_constructor(self._orig_torch_full, self.dtype) |
|
|
|
|
|
def __exit__(self, exc_type, exc_value, traceback): |
|
|
if not self.enabled: |
|
|
return |
|
|
torch.Tensor.__new__ = torch.Tensor.__old_new__ |
|
|
torch.empty = self._orig_torch_empty |
|
|
torch.zeros = self._orig_torch_zeros |
|
|
torch.ones = self._orig_torch_ones |
|
|
torch.full = self._orig_torch_full |
|
|
|
