wehub-resource-sync
27 lines
651 B
Python
27 lines
651 B
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
import torch
|
|
|
|
from ...model_implementations.parameter_base import ParameterBase
|
|
"""
|
|
Unembedding containers.
|
|
"""
|
|
|
|
|
|
class UnembedParameter(ParameterBase):
|
|
"""
|
|
Unembedding parameter. This will likely be mapped to the same original weight in the model as the
|
|
embedding, but we have a different preferred sharding approach.
|
|
"""
|
|
|
|
params: torch.Tensor
|
|
"""
|
|
Unembedding parameter of shape [vocab_size, model_dim].
|
|
"""
|
|
|
|
def finalize(self) -> torch.Tensor:
|
|
return self.inference_model.transform_unembed_param(self.params)
|