| |
| |
|
|
| |
|
|
| from functools import reduce |
| from typing import Iterable |
| from collections import defaultdict |
| import torch |
|
|
| from deepspeed.accelerator import get_accelerator |
|
|
|
|
| class Allocator: |
| cache = defaultdict(dict) |
|
|
| def empty_from(tensor: torch.Tensor, shape: Iterable[int]) -> torch.Tensor: |
| try: |
| return Allocator.cache[tensor][shape] |
| except KeyError: |
| shape_size = reduce(lambda x, y: x * y, shape) |
| if shape_size == 0: |
| raise ValueError("Cannot create empty tensor with size 0") |
| Allocator.cache[tensor][shape] = tensor.flatten()[:shape_size].view(shape) |
| return Allocator.cache[tensor][shape] |
|
|
|
|
| empty_from = Allocator.empty_from |
|
|
|
|
| def on_device(method) -> torch.Tensor: |
| """ |
| Wraps a method to ensure the returned tensor is on the current device. |
| """ |
|
|
| def wrapped(self, *args, **kwargs): |
| tensor = method(self, *args, **kwargs) |
| if isinstance(tensor, torch.Tensor): |
| return tensor.to(get_accelerator().current_device()) |
| return tensor |
|
|
| return wrapped |
|
|
