petals/src/client/remote_model.py

# this code is in active development, interfaces may change
import os
import torch
from typing import List, Optional, Tuple, Union

import torch
import hivemind
import torch
import torch.nn as nn
from hivemind import get_logger, use_hivemind_log_handler

from src.bloom.model import (
    BloomConfig,
    BloomForCausalLM,
    BloomForSequenceClassification,
    BloomModel,
    BloomPreTrainedModel,
    LMHead,
)
from src.client.remote_sequential import RemoteSequential
from src.client.remote_generation import RemoteGenerationMixin
from src.utils.generation_algorithms import DecodingAlgorithm
from src.utils.generation_constraints import ABConstraint

use_hivemind_log_handler("in_root_logger")
logger = get_logger(__file__)


class DistributedBloomConfig(BloomConfig):
    """
    A bloom config that contains information about DHT peers.
    To create a distributed model, one must provide dht_prefix and either initial_peers or dht.
    """

    initial_peers: Tuple[str, ...] = ()  # a list of initial peers for hivemind DHT
    dht_prefix: str  # a prefix for all dht keys that correspond to this model (usually equal to model name)
    dht: Optional[hivemind.DHT] = None  # a running DHT instance, e.g. when using the same DHT for multiple models
    chunk_size_for_efficient_fp16_on_cpu: int = 10000  # a chunk size for a LM head for efficient half-precision on CPU
    num_prefix_tokens: int = 0  # a number of tokens for prompt tuning.


class DistributedBloomModel(BloomModel):
    """BloomModel, but all transformer layers are hosted by the swarm"""

    config_class = DistributedBloomConfig

    def __init__(self, config: DistributedBloomConfig):
        assert config.dht_prefix, "Could not find dht_prefix in config, please create model with dht_prefix=..."
        assert config.initial_peers or config.dht, "Please specify initial_peers=list(...) or dht=hivemind.DHT(...)"

        n_layer, config.n_layer = config.n_layer, 0  # temporarily set n_layer to 0 to prevent layer initialization
        super().__init__(config)
        assert len(self.h) == 0
        config.n_layer = n_layer

        dht = (
            config.dht
            if config.dht is not None
            else hivemind.DHT(initial_peers=config.initial_peers, client_mode=True, start=True)
        )
        assert isinstance(dht, hivemind.DHT) and dht.is_alive(), "dht must be a running hivemind.DHT instance"
        self.h = RemoteSequential(config, dht, config.dht_prefix)

        # Forbid accumulate grads for embeddings and layernorm
        self.set_requires_grad(False)

    def set_requires_grad(self, value):
        for p in self.parameters():
            p.requires_grad = value

    def forward(self, *args, use_cache=None, **kwargs):
        if use_cache:
            raise ValueError(
                "Distributed forward does not support use_cache; for efficient cache-aware generation, "
                "please use model.transformer.inference_session() or model.generate(...)"
            )
        return super().forward(*args, use_cache=False, **kwargs)


class DistributedBloomPrefix(DistributedBloomModel):
    """DistributedBloomModel with prefix tokens for prompt tuning"""

    def __init__(self, config):
        super().__init__(config)
        assert config.num_prefix_tokens > 0, "The number of prefix tokens must be > 0"
        self.prefix_length = config.num_prefix_tokens

        self.prompt_embeddings = nn.Embedding(self.prefix_length, config.hidden_size)
        self.prefix_tokens = torch.arange(self.prefix_length).long()

    def get_prompt(self, batch_size):
        prefix_tokens = self.prefix_tokens.unsqueeze(0).expand(batch_size, -1)
        prefix_tokens = prefix_tokens.to(self.word_embeddings.weight.device)
        prompts = self.prompt_embeddings(prefix_tokens)
        return prompts

    def forward(
        self,
        input_ids: Optional[torch.LongTensor],
        inputs_embeds: Optional[torch.Tensor],
        attention_mask: Optional[torch.Tensor],
        past_key_values=None,
        position_ids=None,
        head_mask=None,
        use_cache=None,
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None,
    ):
        assert (
            input_ids is None or inputs_embeds is None
        ), "You cannot specify both input_ids and inputs_embeds at the same time"
        assert input_ids is not None or inputs_embeds is not None, "You must specify either input_ids or inputs_embeds"

        if inputs_embeds is None:
            inputs_embeds = self.word_embeddings(input_ids)

        batch_size = inputs_embeds.shape[0]

        if attention_mask is not None:
            prefix_attention_mask = torch.ones(batch_size, self.prefix_length, device=attention_mask.device)
            attention_mask = torch.cat((prefix_attention_mask, attention_mask), dim=1)

        prompts = self.get_prompt(batch_size)
        inputs_embeds = torch.cat([prompts, inputs_embeds], dim=1)

        transformer_outputs = super().forward(
            inputs_embeds=inputs_embeds,
            attention_mask=attention_mask,
            past_key_values=past_key_values,
            position_ids=position_ids,
            head_mask=head_mask,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        # Remove prefix
        last_hidden_state = transformer_outputs[0][:, self.prefix_length :]
        transformer_outputs["last_hidden_state"] = last_hidden_state
        return transformer_outputs


class DistributedBloomForCausalLM(BloomForCausalLM, RemoteGenerationMixin):
    """DistributedBloomForCausalLM, but all transformer layers are hosted by the swarm"""
    config_class = DistributedBloomConfig

    def __init__(self, config: DistributedBloomConfig):
        BloomPreTrainedModel.__init__(self, config)
        if config.num_prefix_tokens > 0:
            self.transformer = DistributedBloomPrefix(config)
        else:
            self.transformer = DistributedBloomModel(config)
        self.lm_head = LMHead(config, self.transformer.word_embeddings)

        # Initialize weights and apply final processing
        self.post_init()

    def get_input_embeddings(self):
        return self.transformer.word_embeddings

    def get_output_embeddings(self):
        if self.config.tie_word_embeddings:
            return None
        return self.lm_head

    def set_input_embeddings(self, new_embeddings: nn.Embedding):
        assert isinstance(new_embeddings, nn.Embedding)
        self.transformer.word_embeddings = self.lm_head.word_embeddings = new_embeddings
        assert self.lm_head.bias is None or len(self.lm_head.bias) == new_embeddings.num_embeddings

    def set_output_embeddings(self, new_lm_head: nn.Linear):
        with torch.no_grad():
            self.lm_head.word_embeddings.weight[...] = new_lm_head.weight
            self.lm_head.bias[...] = new_lm_head.bias

    def generate(
        self,
        inputs: Optional[torch.Tensor] = None,
        do_sample: Optional[bool] = None,
        temperature: float = 1.0,
        top_k: Optional[int] = None,
        top_p: Optional[float] = None,
        eos_token_id: Optional[int] = None,
        max_new_tokens: Optional[int] = None,
        decoding_algorithm: Optional[DecodingAlgorithm] = None,
        provided_constraints: List[ABConstraint] = [],
        **model_kwargs,
    ) -> torch.Tensor:
        return RemoteGenerationMixin.generate(
            self,
            inputs,
            do_sample,
            temperature,
            top_k,
            top_p,
            eos_token_id,
            max_new_tokens,
            decoding_algorithm,
            provided_constraints,
            **model_kwargs,
        )


class DistributedBloomForSequenceClassification(BloomForSequenceClassification):
    config_class = DistributedBloomConfig

    def __init__(self, config: DistributedBloomConfig):
        super().__init__(config)
        if config.num_prefix_tokens > 0:
            self.transformer = DistributedBloomPrefix(config)
        else:
            self.transformer = DistributedBloomModel(config)

        # Initialize weights and apply final processing
        self.post_init()