chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,258 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
import weakref
|
||||
from abc import abstractmethod
|
||||
from typing import Type
|
||||
|
||||
import torch
|
||||
|
||||
from deepspeed.compat import get_annotations_from_namespace, get_annotations
|
||||
|
||||
# Currently have dependency loops for the type hints.
|
||||
InferenceModel = Type["InferenceModel"]
|
||||
LayerContainer = Type["LayerContainer"]
|
||||
|
||||
MAPPING_KEY = "PARAM_MAPPING"
|
||||
|
||||
|
||||
def make_param_getter(clsname, param):
|
||||
"""
|
||||
Normal getter implementation for a property.
|
||||
"""
|
||||
|
||||
def param_getter(self):
|
||||
return getattr(self, f"__{clsname}__{param}")
|
||||
|
||||
return param_getter
|
||||
|
||||
|
||||
def make_param_setter(clsname, param):
|
||||
"""
|
||||
Setter implementation that will call complete component to potentially
|
||||
finalize the parameter.
|
||||
"""
|
||||
|
||||
def param_setter(self, value):
|
||||
setattr(self, f"__{clsname}__{param}", value)
|
||||
self.dtype = value.dtype
|
||||
self.complete_component()
|
||||
|
||||
return param_setter
|
||||
|
||||
|
||||
def make_readonly_setter():
|
||||
"""
|
||||
Setter implementation that will raise an error if called.
|
||||
"""
|
||||
|
||||
def paramlist_setter(self, value):
|
||||
raise ValueError("Cannot set a ParametrizedList directly.")
|
||||
|
||||
return paramlist_setter
|
||||
|
||||
|
||||
class ParameterMetaclass(type):
|
||||
"""
|
||||
MetaClass for the ParameterBase base class. This class will parse the `src_params`
|
||||
attribute and create properties for each of the dependencies. A dependency can either
|
||||
be represented as a string, which is interpreted as a named Tensor, or a `ParametrizedList`
|
||||
subclass.
|
||||
"""
|
||||
|
||||
def __new__(cls, clsname, bases, attrs):
|
||||
|
||||
annotations = get_annotations_from_namespace(attrs)
|
||||
dependencies = {
|
||||
name: annotation
|
||||
for name, annotation in annotations.items() if issubclass(annotation, (torch.Tensor, ParametrizedList))
|
||||
}
|
||||
n_dependencies = len(dependencies)
|
||||
|
||||
# Create properties for each of our dependencies
|
||||
for d_name, d_type in dependencies.items():
|
||||
if issubclass(d_type, ParametrizedList):
|
||||
assert hasattr(
|
||||
d_type, "count_attr"
|
||||
), "ParametrizedList must have a count_attr attribute to access on the inference module."
|
||||
attrs[d_name] = property(make_param_getter(clsname, d_name), make_readonly_setter())
|
||||
else: # torch.Tensor
|
||||
attrs[d_name] = property(make_param_getter(clsname, d_name), make_param_setter(clsname, d_name))
|
||||
|
||||
new_cls = super().__new__(cls, clsname, bases, attrs)
|
||||
new_cls.n_dependencies = n_dependencies
|
||||
|
||||
return new_cls
|
||||
|
||||
def __call__(cls, *args, **kwargs):
|
||||
new_obj = super().__call__(*args, **kwargs)
|
||||
new_obj.__init__(*args, **kwargs)
|
||||
|
||||
setattr(new_obj, "dest_param", None)
|
||||
|
||||
# Initialize our dependences to None/empty `ParametrizedList`s
|
||||
for name, annotation in get_annotations(new_obj).items():
|
||||
if issubclass(annotation, ParametrizedList):
|
||||
#TODO(jeff): update assert with this, model implementation attribute does not align or missing wrt the ParametrizedList attributes
|
||||
assert hasattr(
|
||||
new_obj.inference_model, annotation.count_attr
|
||||
), f"new_obj={new_obj.__class__.__name__}, name={name}, annotation.count_attr={annotation.count_attr}"
|
||||
param_list = annotation(new_obj, getattr(new_obj.inference_model, annotation.count_attr))
|
||||
setattr(new_obj, f"__{new_obj.__class__.__name__}__{name}", param_list)
|
||||
else: # torch.Tensor
|
||||
setattr(new_obj, f"__{new_obj.__class__.__name__}__{name}", None)
|
||||
|
||||
return new_obj
|
||||
|
||||
|
||||
class ParameterBase(metaclass=ParameterMetaclass):
|
||||
"""
|
||||
A ParameterBase allows us to consolidate tracking the dependencies of loading a parameter from
|
||||
a checkpoint into a single object. This class should not be used directly, but rather subclassed
|
||||
and the `src_params` attribute set to a list of strings and/or `ParametrizedList`s.
|
||||
"""
|
||||
|
||||
# inference_model: InferenceModel
|
||||
"""
|
||||
Inference model that will provide context on how to shard and transform the parameter.
|
||||
"""
|
||||
|
||||
#completed_components: int
|
||||
"""
|
||||
How many of the layer dependencies have been met. This is used to determine when the parameter
|
||||
is ready to be finalized. A ParametrizedList counts as a single dependency for the purposes
|
||||
of this counter.
|
||||
"""
|
||||
|
||||
def __init__(self, model: InferenceModel, parent_container: LayerContainer) -> None:
|
||||
"""
|
||||
Direct constructor. This should not be called from client code.
|
||||
|
||||
Args:
|
||||
model (InferenceModel): Inference model that will be used to shard and transform the
|
||||
parameter in `finalize`.
|
||||
parent_container (LayerContainer): The parent container that this parameter is a member
|
||||
of. We will build a weakref to this container to call the finalization callback.
|
||||
"""
|
||||
self.inference_model = model
|
||||
self.completed_components = 0
|
||||
self.parent_container = weakref.ref(parent_container)
|
||||
|
||||
@abstractmethod
|
||||
def finalize(self) -> torch.Tensor:
|
||||
"""
|
||||
Finalize the parameter after all of its source parameters have been set. This method
|
||||
will be automatically called when all inputs have been set. It should return the Tensor
|
||||
with all transformations performed on it.
|
||||
"""
|
||||
pass
|
||||
|
||||
def complete_component(self) -> None:
|
||||
"""
|
||||
Mark a component as completed. This should be called by the relevant setter of a direct
|
||||
property or a ParametrizedList. This method will automatically call `finalize` when all
|
||||
dependencies have been met and then call the finalization callback on the parent container.
|
||||
|
||||
Once the finalization callback has been called, the parameter will be replaced with the
|
||||
`dst_param` attribute on the parent container, and this instance will be destroyed.
|
||||
"""
|
||||
self.completed_components += 1
|
||||
|
||||
if self.completed_components != self.n_dependencies:
|
||||
return
|
||||
|
||||
finalized_param = self.finalize()
|
||||
self.parent_container().finalization_callback(self, finalized_param)
|
||||
|
||||
|
||||
class ParametrizedList:
|
||||
"""
|
||||
A ParametrizedList is a list of parameters that are dependencies
|
||||
of a `ParameterBase` but may vary in length depending on the model
|
||||
configuration (rather than architecture). For example, a MoE layer
|
||||
may have different number of experts depending on the size of the model.
|
||||
|
||||
This class is used to manage these lists and provide integer indexing
|
||||
of a single component rather than accessing names directly. For example,
|
||||
it tends to be more natural to access the 8th expert with `experts[8]`
|
||||
rather than a name like `expert_8`, especially as an attribute.
|
||||
|
||||
To inherit from this class, set static variables `name` and `count_attr`.
|
||||
|
||||
```python
|
||||
class MyParametrizedList(ParametrizedList):
|
||||
count_attr: str = "my_list_count"
|
||||
```
|
||||
|
||||
In the above example, `my_list_count` should be an accessible attribute
|
||||
of the inference model (i.e. via `self.inference_model.my_list_count`).
|
||||
|
||||
NOTE: There are some APIs in which this type cannot be used as if it is
|
||||
just a list of Tensors. For example, `torch.cat(param_list)` will not work.
|
||||
However, you can make it compatible with a tuple wrapper:
|
||||
`torch.cat(tuple(param_list))`
|
||||
"""
|
||||
|
||||
n_params: int
|
||||
"""
|
||||
Number of params this list contains.
|
||||
"""
|
||||
|
||||
param: ParameterBase
|
||||
"""
|
||||
WeakRef to the owning parameter.
|
||||
"""
|
||||
|
||||
def __init__(self, param: ParameterBase, n_params: int) -> None:
|
||||
"""
|
||||
Constructor. Should not be called from client code.
|
||||
|
||||
Args:
|
||||
param (ParameterBase): The owning parameter.
|
||||
n_params (int): The number of parameters this list contains. This should be
|
||||
"""
|
||||
self.n_params = n_params
|
||||
self.set_params = 0
|
||||
self.param = weakref.ref(param)
|
||||
self._params = [None] * n_params
|
||||
|
||||
def __getitem__(self, index):
|
||||
return self._params[index]
|
||||
|
||||
def __setitem__(self, index, value):
|
||||
if self._params[index] is not None:
|
||||
raise ValueError("Cannot set a parameter twice.")
|
||||
|
||||
self._params[index] = value
|
||||
self.set_params += 1
|
||||
|
||||
if self.set_params != self.n_params:
|
||||
return
|
||||
|
||||
self.param().complete_component()
|
||||
|
||||
def __iter__(self):
|
||||
return iter(self._params)
|
||||
|
||||
|
||||
def ParamList(attr: str):
|
||||
"""
|
||||
Helper to create a subclass of ParametrizedList with the desired `count_attr`.
|
||||
|
||||
In this manner, we can annotate the type of a Parameter dependency with the
|
||||
following:
|
||||
|
||||
```python
|
||||
class CustomParameter(ParameterBase):
|
||||
dependency_list: ParamList("dependencies_count_name")
|
||||
```
|
||||
|
||||
where "dependencies_count_name" is the name of the attribute on the inference model.
|
||||
"""
|
||||
|
||||
class ParametrizedListInstance(ParametrizedList):
|
||||
count_attr: str = attr
|
||||
|
||||
return ParametrizedListInstance
|
||||
Reference in New Issue
Block a user