Mike0021/zonos2

1from __future__ import annotations2 3from typing import Dict4 5import torch6import torch.nn.functional as F7from zonos2.core import get_global_ctx8from zonos2.distributed import DistributedCommunicator, get_tp_info9from zonos2.utils import divide_up, nvtx_annotate10 11from .base import BaseOP12 13 14class VocabParallelEmbedding(BaseOP):15    def __init__(16        self,17        num_embeddings: int,18        embedding_dim: int,19    ):20        super().__init__()21        tp_info = get_tp_info()22        tp_rank = tp_info.rank23        self.tp_size = tp_info.size24        self.num_embeddings = num_embeddings25        self.num_embeddings_tp = divide_up(num_embeddings, self.tp_size)26        start_idx = self.num_embeddings_tp * tp_rank27        finish_idx = min(start_idx + self.num_embeddings_tp, num_embeddings)28        self.vocab_range = (start_idx, finish_idx - start_idx)29        self.weight = torch.empty(self.num_embeddings_tp, embedding_dim)30        self._comm = DistributedCommunicator()31 32    @nvtx_annotate("Embedding")33    def forward(self, x: torch.Tensor) -> torch.Tensor:34        from zonos2.kernel import indexing35 36        y = indexing(37            weights=self.weight,38            indices=x,39            vocab_range=self.vocab_range if self.tp_size > 1 else None,40        )41 42        return self._comm.all_reduce(y) if self.tp_size > 1 else y43 44 45class ParallelLMHead(VocabParallelEmbedding):46    def __init__(47        self,48        num_embeddings: int,49        embedding_dim: int,50        bias: bool = False,51        tie_word_embeddings: bool = False,52        tied_embedding: VocabParallelEmbedding | None = None,53    ):54        super().__init__(num_embeddings, embedding_dim)55        self.bias = torch.empty(self.num_embeddings_tp) if bias else None56        self.tied_embedding = tied_embedding57        assert (tied_embedding is not None) == tie_word_embeddings58 59    def load_state_dict(60        self,61        state_dict: Dict[str, torch.Tensor],62        *,63        prefix: str = "",64        _internal: bool = False,65    ) -> None:66        if not self.tied_embedding:67            return super().load_state_dict(state_dict, prefix=prefix, _internal=_internal)68        else:69            # pop the lm_head.weights and lm_head.bias if they exist70            possible_weight = f"{prefix}.weight"71            possible_bias = f"{prefix}.bias"72            if possible_weight in state_dict:73                state_dict.pop(possible_weight)74            if possible_bias in state_dict:75                state_dict.pop(possible_bias)76 77    def state_dict(78        self,79        *,80        prefix: str = "",81        result: Dict[str, torch.Tensor] | None = None,82    ) -> Dict[str, torch.Tensor]:83        if not self.tied_embedding:84            return super().state_dict(prefix=prefix, result=result)85        return {} if result is None else result86 87    @nvtx_annotate("LMHead")88    def forward(self, x: torch.Tensor) -> torch.Tensor:89        ctx = get_global_ctx()90        batch = ctx.batch91        bs = batch.size92        if batch.is_prefill:93            indices = batch.attn_metadata.get_last_indices(bs)94            x = x[indices].contiguous()95            del indices96 97        module = self.tied_embedding or self98        logits = F.linear(x, module.weight, self.bias)99        if self.tp_size == 1:100            return logits101        input_shape = logits.shape102        output_tensor = self._comm.all_gather(logits)103 104        if bs == 1:105            return output_tensor.view(1, -1)[:, : self.num_embeddings]106 107        output_tensor = output_tensor.view((self.tp_size,) + input_shape)108        output_tensor = output_tensor.movedim(0, -1)109        output_tensor = output_tensor.reshape(input_shape[:1] + (self.tp_size * input_shape[1],))110        return output_tensor[:, : self.num_embeddings]111