wehub-resource-sync
61 lines
1.5 KiB
Python
61 lines
1.5 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
import torch
|
|
|
|
from deepspeed.accelerator import get_accelerator
|
|
from deepspeed.inference.v2.allocator import on_device
|
|
from deepspeed.inference.v2.inference_parameter import InferenceParameter
|
|
from deepspeed.inference.v2.model_implementations.parameter_base import ParameterBase, ParametrizedList
|
|
|
|
|
|
class SimpleParam(ParameterBase):
|
|
"""
|
|
Parameter with single dependency.
|
|
"""
|
|
|
|
param: torch.Tensor
|
|
|
|
@on_device
|
|
def finalize(self) -> torch.Tensor:
|
|
return self.inference_model.transform(self.param)
|
|
|
|
|
|
class SimpleParametrizedList(ParametrizedList):
|
|
"""
|
|
Parameter list based on `num_dependencies` attribute.
|
|
"""
|
|
|
|
count_attr: str = "num_dependencies"
|
|
|
|
|
|
class ListParam(ParameterBase):
|
|
"""
|
|
Parameter with list dependency.
|
|
|
|
NOTE: This uses the tuple workaround for the `ParametrizedList` class
|
|
as described in the docstring of `ParametrizedList`.
|
|
"""
|
|
|
|
params: SimpleParametrizedList
|
|
|
|
@on_device
|
|
def finalize(self) -> torch.Tensor:
|
|
return self.inference_model.transform(torch.cat(tuple(self.params)))
|
|
|
|
|
|
class DummyInferenceModel:
|
|
|
|
@property
|
|
def num_dependencies(self) -> int:
|
|
return 2
|
|
|
|
def transform(self, param: torch.Tensor) -> torch.Tensor:
|
|
return InferenceParameter.initialize(param)
|
|
|
|
|
|
def validate_device(tensor: torch.Tensor):
|
|
assert tensor.device == torch.device(get_accelerator().current_device())
|