[DESIGN] user experience about petals HOT 8 CLOSED

import os
import torch
from src import DistributedBloom

os.environ['BLOOM_TOKEN'] = input("Paste your account token here: ").strip()
os.environ['INITIAL_PEERS'] = "/ip4/123.123.123.123/tcp/1337 /ip6/something/udp/42/quic"
model = DistributedBloom.from_pretrained("bigscience/distributed-bloom-176b")
# or MyFinetunedDistributedBloom.from_pretrained("myname/finetuned-bloom")
# this loads only embeddings, logits, and config (hyperparameters)


# inference: same as most models from HF hub
generated_tokens = model.generate(tokenized_prefix, max_length=10, max_points_per_token=0.1)

# training (e.g. prompt-tuning)
opt = torch.optim.Adam(model.parameters(), lr=1e-5)
budget = dict(max_points_per_token=0.1, max_backward_points_per_token=0.25)  # all params are optional
for data_batch in my_dataloader:
    opt.zero_grad()
    loss = compute_loss(model.forward(**data_batch, **budget))
    loss.backward()
    opt.step()  # note: this step will only update local parameters: embeddings, lm_head biases, optional adapters

class DistribtuedBloom:
    emb: nn.Embedding                         # local modules initialized in from_pretrained
    lm_head: nn.Linear
    remote_model: RemoteSequential            # networking, compression and fault tolerance magic happens here
        
    def forward(self, tokens):
        embs = self.emb(tokens)                                   # trains same as regular model
        last_hid = self.remote_model.forward(embs)                # model can do backprop, but does not update frozen transformer blocks
        # ^-- note 1: this module hides any network errors from the user, automatically finds alternative servers in case of errors
        # ^-- note 2: some remote_model knock-offs may contain local trainable adapters that *are* trainable
        return self.lm_head(last_hid)                             # trains same as regular model
    
    def generate(self, prefix: torch.LongTensor, max_length: int **generation_params_including_budget):
        outputs = []
        with self.remote_model.begin_inference_session(**generation_params_including_budget) as sess:
            # ^-- find out generation params, budget, etc -- do not talk to servers YET
            past_inputs = prefix
            for i in range(max_length - prefix.shape[1]):
                hidden_states = self.emb(past_inputs)

                hidden_states = sess.step(hidden_states, **optional_beam_search_hypo_indices)
                # then send 1 token through pipeline, update attention caches, return last layer activities
                # ^-- hidden from the user: find servers, establish connections, broker deals, recover from failures

                logits = self.lm_head(hidden_state)
                next_token = argmax_or_some_algorithm(logits)
                outputs.append(next_token)
                past_inputs = next_token
            return outputs

class RemoteSequential:
    dht: DHT
    metadata_about_servers
    optional_trainable_parts: nn.Module

    def forward(self, *inputs: torch.Tensor, **budget_kwargs) -> torch.Tensor:
        """Split inputs along batch dimension, send to remote servers, retry in case of failure/timeout"""

    def begin_inference_session(self, **generation_params) -> RemoteInferenceSession:
        """Create a fault-tolerant inference session that supports generic incremental inference"""

class RemoteInferenceSession:
    """
    represents a single inference session on client side
    finds best remote servers from across DHT, dynamically finds replacements for failed servers
    keeps local caches for intermediate hidden states between servers
    """
    def step(self, hidden_states: torch.Tensor, **additional_step_params) -> torch.Tensor:
        """
        Run a chunk of hidden states through all transformer blocks, return last layer hidden states
        :param hidden_states: tensor [batch_size_usually_1, chunk_length_usually_1, hidden_size]
        :param additional_step_params: beam search hypothesis indices, backtrack steps, etc
        :returns: hidden_states, tensor of same shape as hidden_states
        """