wehub-resource-sync
90 lines
2.7 KiB
Python
90 lines
2.7 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
from typing import Dict
|
|
|
|
import torch
|
|
|
|
CORE_PARAM = "_ds_core_param_key"
|
|
|
|
STR_TO_DTYPE = {
|
|
"torch.float32": torch.float32,
|
|
"torch.float64": torch.float64,
|
|
"torch.float16": torch.float16,
|
|
"torch.bfloat16": torch.bfloat16,
|
|
"torch.int64": torch.int64,
|
|
"torch.int32": torch.int32,
|
|
"torch.int16": torch.int16,
|
|
"torch.int8": torch.int8,
|
|
"torch.uint8": torch.uint8,
|
|
"torch.bool": torch.bool,
|
|
}
|
|
|
|
|
|
class InferenceParameter(torch.Tensor):
|
|
"""
|
|
An extension of the torch.Tensor class to support our inference focused features. One important
|
|
thing to note here is that an InferenceParam can be used a torch.Tensor, but outputs of
|
|
torch.Tensor operations will not be InferenceParams.
|
|
"""
|
|
|
|
@staticmethod
|
|
def __new__(cls, tensor, *args, **kwargs):
|
|
new_tensor = super().__new__(cls, tensor, *args, **kwargs)
|
|
if hasattr(tensor, "_aux_attrs"):
|
|
setattr(new_tensor, "_aux_attrs", tensor.aux_attrs)
|
|
return new_tensor
|
|
|
|
def to(self, *args, **kwargs):
|
|
new_tensor = super().to(*args, **kwargs)
|
|
if hasattr(self, "_aux_attrs"):
|
|
setattr(new_tensor, "_aux_attrs", self.aux_attrs)
|
|
try:
|
|
_ = torch.device(args[0])
|
|
for name, attr in new_tensor.aux_attrs.items():
|
|
new_attr = attr.to(*args, **kwargs)
|
|
setattr(new_tensor, name, new_attr)
|
|
new_tensor.aux_attrs[name] = new_attr
|
|
except Exception:
|
|
pass
|
|
|
|
return new_tensor
|
|
|
|
@classmethod
|
|
def initialize(cls, core_param: torch.Tensor, **kwargs) -> 'InferenceParameter':
|
|
"""
|
|
Create the inference parameter.
|
|
"""
|
|
param = InferenceParameter(core_param)
|
|
setattr(param, "_aux_attrs", kwargs)
|
|
|
|
for attr_name, attr in kwargs.items():
|
|
if hasattr(param, attr_name):
|
|
raise ValueError(f"Attribute {attr_name} already exists on param.")
|
|
|
|
if not isinstance(attr, torch.Tensor):
|
|
raise ValueError(f"Attribute {attr_name} must be a tensor.")
|
|
|
|
setattr(param, attr_name, attr)
|
|
|
|
return param
|
|
|
|
@classmethod
|
|
def initialize_raw(self, **kwargs) -> 'InferenceParameter':
|
|
"""
|
|
All kwargs must be torch.Tensors and must include the core parameter.
|
|
"""
|
|
if CORE_PARAM not in kwargs:
|
|
raise ValueError(f"Must provide core parameter, with key {CORE_PARAM}.")
|
|
|
|
return InferenceParameter.initialize(kwargs[CORE_PARAM], **kwargs)
|
|
|
|
@property
|
|
def aux_attrs(self) -> Dict[str, torch.Tensor]:
|
|
"""
|
|
Dictionary of auxiliary attributes.
|
|
"""
|
|
return self._aux_attrs
|