chore: import upstream snapshot with attribution

deepspeed/inference/v2/modules/implementations/unembed/__init__.py

@@ -0,0 +1,6 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from .ragged_unembed import DSRaggedUnembed

deepspeed/inference/v2/modules/implementations/unembed/ragged_unembed.py

@@ -0,0 +1,120 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import Any, Dict, Optional
+
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from ....allocator import empty_from
+from ....inference_utils import DtypeEnum, ActivationType
+from ....kernels.core_ops import CUDAFPLN, BlasLibLinear, CUDARMSNorm, CUDABiasActivation
+from ....kernels.ragged_ops import RaggedLogitsGather
+from ....ragged import RaggedBatchWrapper
+from ...interfaces import DSUnembedBase, DSUnembedRegistry
+from ...configs import DSUnembedConfig
+
+
+@DSUnembedRegistry.register_module
+class DSRaggedUnembed(DSUnembedBase):
+    """
+    Ragged unembedding implementation. This implementation will gather only the last token
+    of each sequence in the ragged inflight batch and calculate the logits only for those rows.
+    """
+
+    @staticmethod
+    def name():
+        return 'ragged_unembed'
+
+    @staticmethod
+    def supports_config(config: DSUnembedConfig):
+
+        if DtypeEnum(config.dtype) not in [DtypeEnum.fp16, DtypeEnum.bf16, DtypeEnum.fp32]:
+            return False
+
+        try:
+            _ = RaggedLogitsGather(config.model_dim, config.dtype)
+        except ValueError:
+            return False
+
+        if config.norm_type == 'rms_norm':
+            try:
+                _ = CUDARMSNorm(config.model_dim, config.dtype)
+            except ValueError:
+                return False
+        elif config.norm_type == 'layer_norm':
+            try:
+                _ = CUDAFPLN(config.model_dim, config.dtype)
+            except ValueError:
+                return False
+
+        return True
+
+    def __init__(self, config: DSUnembedConfig, implementation_config: Dict[str, Any]) -> None:
+        super().__init__(config, implementation_config)
+
+        self._logits_gather = RaggedLogitsGather(config.model_dim, self._config.dtype)
+
+        if self._config.norm_type == 'layer_norm':
+            self._norm = CUDAFPLN(self._config.model_dim, self._config.dtype)
+        elif self._config.norm_type == 'rms_norm':
+            self._norm = CUDARMSNorm(self._config.model_dim, self._config.dtype)
+        else:
+            self._norm = None
+
+        self._linear = BlasLibLinear(self._config.dtype)
+        # Here the activation kernel is being used to apply bias, hence the identity activation type!
+        self._act_fn = CUDABiasActivation(self._config.vocab_size, self._config.dtype, ActivationType.IDENTITY)
+
+        self._intermediate = torch.empty((self._config.max_sequences, self._config.model_dim),
+                                         dtype=self._config.dtype,
+                                         device=get_accelerator().current_device())
+
+        self._output = torch.empty((self._config.max_sequences, self._config.vocab_size),
+                                   dtype=self._config.dtype,
+                                   device=get_accelerator().current_device())
+
+    @property
+    def output(self) -> torch.Tensor:
+        return self._output
+
+    def forward(self,
+                hidden_states: torch.Tensor,
+                vocab_embedding: torch.Tensor,
+                ragged_metadata: RaggedBatchWrapper,
+                bias: Optional[torch.Tensor] = None,
+                gamma: Optional[torch.Tensor] = None,
+                beta: Optional[torch.Tensor] = None) -> torch.Tensor:
+        """
+        Return final model logits.
+
+        Args:
+            hidden_states (torch.Tensor): The hidden states from the model. This is the output of the
+                final layer of the model.
+            vocab_embedding (torch.Tensor): The vocab embedding table.
+            raged_metadata (RaggedBatchWrapper): The ragged batch metadata.
+            gamma (Optional[torch.Tensor]): The gamma tensor for normalization.
+            beta (Optional[torch.Tensor]): The beta tensor for normalization.
+        """
+
+        cut_down_hidden_states = empty_from(self._intermediate,
+                                            (ragged_metadata.current_sequences, self._config.model_dim))
+        self._logits_gather(cut_down_hidden_states, hidden_states, ragged_metadata)
+
+        if self._config.norm_type == 'rms_norm':
+            if gamma is None:
+                raise ValueError('RMS Normalization enabled but gamma not provided.')
+            self._norm(cut_down_hidden_states, cut_down_hidden_states, gamma)
+        elif self._config.norm_type == 'layer_norm':
+            if gamma is None or beta is None:
+                raise ValueError('Normalization enabled but gamma and/or beta not provided.')
+            self._norm(cut_down_hidden_states, cut_down_hidden_states, gamma, beta)
+
+        output = empty_from(self._output, (ragged_metadata.current_sequences, self._config.vocab_size))
+        self._linear(output, cut_down_hidden_states, vocab_embedding)
+        if bias is not None:
+            self._act_fn(output, bias)
+
+        return output