import multiprocessing as mp import multiprocessing.pool from functools import partial from typing import Tuple, List, Dict, Optional import numpy as np import torch import torch.nn as nn from torch.autograd.function import once_differentiable from .expert import RemoteExpert, _RemoteModuleCall from ..utils import nested_map, check_numpy, run_and_await_k, nested_pack, nested_flatten, DUMMY, run_in_background from ..utils import run_isolated_forward, EmulatedAutogradContext, run_isolated_backward, map_with_parallel_backward class RemoteMixtureOfExperts(nn.Module): """ A torch module that performs mixture of experts inference with a local gating function and multiple remote experts. Natively supports pytorch autograd. :note: By default, not all experts are guaranteed to perform forward pass. Moreover, not all of those who ran forward pass are guaranteed to perform backward pass. In the latter case, gradient will be averaged without the missing experts :param in_features: common input size for experts and gating function :param grid_size: tesseract dimensions that form expert uid (see below) :param uid_prefix: common prefix for all expert uids expert uid follows the pattern {uid_prefix}.{0...grid_size[0]}.{0...grid_size[1]}...{0...grid_size[-1]} :param network: TesseractNetwork where the experts reside :param num_workers: number of threads for parallel network operation :param k_best: queries this many experts with highest scores :param k_min: makes sure at least this many experts returned output :param timeout_after_k_min: waits for this many seconds after k_min experts returned results. Any expert that didn't manage to return output after that delay is considered unavailable :param expert_padding: internal value used to denote "absent expert". Should not coincide with any expert uid. :param allow_broadcasting: if RemoteMixtureOfExperts if fed with input dimension above 2, allow_broadcasting=True will flatten first d-1 input dimensions, apply RemoteMixtureOfExperts and un-flatten again allow_broadcasting=False will raise an error """ def __init__(self, *, in_features, grid_size: Tuple[int], network, k_best, k_min=1, forward_timeout=None, timeout_after_k_min=1.0, backward_k_min=1, backward_timeout=None, uid_prefix='', expert_padding=None, allow_broadcasting=True): super().__init__() self.network, self.grid_size = network, grid_size self.uid_prefix, self.expert_padding = uid_prefix, expert_padding self.k_best, self.k_min, self.backward_k_min = k_best, k_min, backward_k_min self.forward_timeout, self.timeout_after_k_min, self.backward_timeout = forward_timeout, timeout_after_k_min, backward_timeout self.allow_broadcasting = allow_broadcasting self.proj = nn.Linear(in_features, sum(grid_size)) # jointly predict logits for all grid dimensions self._outputs_schema = None def forward(self, input: torch.Tensor, *args: torch.Tensor, **kwargs: torch.Tensor): """ Choose k best experts with beam search, then call chosen experts and average their outputs. :param input: a tensor of values that are used to estimate gating function, batch-first :param args: extra positional parameters that will be passed to each expert after input, batch-first :param kwargs: extra keyword parameters that will be passed to each expert, batch-first :returns: averaged predictions of all experts that delivered result on time, nested structure of batch-first """ if self.allow_broadcasting and input.ndim != 2: # flatten extra dimensions, apply the function and then un-flatten them back to normal like nn.Linear does flattened_dims = input.shape[:-1] input_flat = input.view(-1, input.shape[-1]) args_flat = [tensor.view(-1, tensor.shape[len(flattened_dims):]) for tensor in args] kwargs_flat = {key: tensor.view(-1, tensor.shape[len(flattened_dims):]) for key, tensor in kwargs.items()} out_flat = self.forward(input_flat, *args_flat, **kwargs_flat) return nested_map(lambda tensor: tensor.view(flattened_dims, tensor.shape[len(flattened_dims):]), out_flat) # 1. compute scores and find most appropriate experts with beam search grid_scores = self.proj(input).split_with_sizes(self.grid_size, dim=-1) chosen_experts = self.beam_search(grid_scores, self.k_best) # ^-- List[batch_size] of List[RemoteExpert] chosen for every input in batch expert_logits = self.compute_expert_scores(grid_scores, chosen_experts) expert_inputs = ((input, *args), kwargs) input_schema = nested_map(lambda x: None, expert_inputs) flat_inputs_per_expert = tuple(zip(*[tensor.split(1, dim=0) for tensor in nested_flatten(expert_inputs)])) batch_jobs_args = tuple( (expert_logits[i, :len(chosen_experts)], chosen_experts[i], self.k_min, self.timeout_after_k_min, self.backward_k_min, self.forward_timeout, self.backward_timeout, input_schema, *flat_inputs_per_expert[i]) for i in range(len(input)) ) averaged_outputs_flat = map_with_parallel_backward(_RemoteMoECall, *batch_jobs_args) return nested_pack(averaged_outputs_flat, self.outputs_schema) def beam_search(self, grid_scores: List[torch.Tensor], k_best: int, **kwargs) -> List[List[RemoteExpert]]: """ Find and return k best experts in the grid using (exact) beam search of the product space :param grid_scores: scores predicted for each dimension in the grid, :type grid_scores: a sequence of tensors of shape[batch_size, self.grid_size[i]] :param k_best: how many of the top experts participate in the computation :param kwargs: extra keyword parameters passed to self.network.first_k_active :returns: a list of *batch_size* lists that contain chosen experts for one sample each inner list contains \ RemoteExpert instances for *up to* k_best experts """ assert len(grid_scores) == len(self.grid_size) assert all(len(dim_scores.shape) == 2 for dim_scores in grid_scores) batch_size = len(grid_scores[0]) beam = np.array([[self.uid_prefix]] * batch_size, dtype=object) # [batch_size, up_to_beam_size] scores = np.zeros([batch_size, 1], dtype=np.float64) delimeters = np.array(self.network.UID_DELIMETER)[None, None, None] # pre-compute numpy array for fast concat for dim_index, dim_scores in enumerate(grid_scores): dim_scores = check_numpy(dim_scores) assert dim_scores.shape[-1] == self.grid_size[dim_index] # create all possible successsors from current beam dim_indices = np.arange(dim_scores.shape[1]).astype(str) new_candidates = beam[:, :, None] + delimeters + dim_indices[None, None, :] new_candidates = new_candidates.reshape([batch_size, -1]) new_scores = scores[:, :, None] + dim_scores[:, None, :] new_scores = new_scores.reshape([batch_size, -1]) # select k best candidates according to scores but only those that are still active new_order = np.argsort(- new_scores, axis=-1) top_alive_lookups = [ run_in_background(self.network.first_k_active, cands[order], k_best, **kwargs) for cands, order in zip(new_candidates, new_order)] batch_cand_to_score = [ dict(zip(cands, cand_scores)) for cands, cand_scores in zip(new_candidates, new_scores)] top_alive_prefixes = [result.result() for result in top_alive_lookups] top_alive_scores = [list(map(cand_to_score.get, top_cands)) for cand_to_score, top_cands in zip(batch_cand_to_score, top_alive_prefixes)] # pad up to beam size beam = np.array([row + [self.expert_padding] * (k_best - len(row)) for row in top_alive_prefixes], dtype='object') scores = np.array([row + [-float('inf')] * (k_best - len(row)) for row in top_alive_scores], dtype='float32') unique_experts = self.network.get_experts(list(set( uid for row in beam for uid in row if uid != self.expert_padding))) if self._outputs_schema is None: self._outputs_schema = next(iter(unique_experts)).info['outputs_schema'] unique_experts_by_uid = {expert.uid: expert for expert in unique_experts if expert != self.expert_padding} return [[unique_experts_by_uid[uid] for uid in row if uid in unique_experts_by_uid] for row in beam] def compute_expert_scores( self, grid_scores: List[torch.Tensor], batch_experts: List[List[RemoteExpert]]) -> torch.Tensor: """ TODO docstring here """ expert_counts = list(map(len, batch_experts)) batch_size = len(batch_experts) max_num_experts = max(expert_counts) total_num_experts = sum(expert_counts) expert_index_in_batch = torch.arange(total_num_experts, device=grid_scores[0].device) expert_strides = torch.cumsum(torch.as_tensor([0] + expert_counts, device=grid_scores[0].device), dim=-1)[:-1] flat_batch_indices = (expert_index_in_batch >= expert_strides[:, None]).to(torch.int32).sum(0) - 1 flat_local_indices = expert_index_in_batch - expert_strides[flat_batch_indices] flat_experts = [expert for row in batch_experts for expert in row] grid_indices = np.zeros([len(flat_experts), len(grid_scores)], dtype=np.int64) for i, expert in enumerate(flat_experts): expert_indices = expert.uid[len(self.uid_prefix) + len(self.network.UID_DELIMETER):] expert_indices = list(map(int, expert_indices.split(self.network.UID_DELIMETER))) grid_indices[i] = expert_indices scores_per_dim = [ dim_scores[flat_batch_indices, dim_indices] if len(flat_batch_indices) else torch.zeros(0) for dim_scores, dim_indices in zip(grid_scores, grid_indices.T)] flat_scores = torch.sum(torch.stack(scores_per_dim, dim=0), dim=0) scores = torch.full((batch_size, max_num_experts), fill_value=-float('inf'), device=grid_scores[0].device) scores[flat_batch_indices, flat_local_indices] = flat_scores # backprop-able w.r.t. flat_scores return scores @property def outputs_schema(self): if self._outputs_schema is None: # grab some expert to set ensemble output shape dummy_scores = self.proj(torch.randn(1, self.proj.in_features)).split_with_sizes(self.grid_size, dim=-1) self._outputs_schema = self.beam_search(dummy_scores, k_best=1)[0][0].info['outputs_schema'] return self._outputs_schema class _RemoteMoECall(torch.autograd.Function): """ Internal autograd-friendly function that calls multiple experts on the same input and averages their outputs. This function that can recover from individual failures during forward and/or backward passes. For user-friendly version of this function, use RemoteMixtureOfExperts module. """ @classmethod def forward(cls, ctx, expert_logits: torch.Tensor, experts: List[RemoteExpert], k_min: int, timeout_after_k_min: float, backward_k_min: int, timeout_total: Optional[float], backward_timeout: Optional[float], input_schema, *flat_inputs: torch.Tensor) -> Tuple[torch.Tensor]: expert_args, expert_kwargs = nested_pack(flat_inputs, structure=input_schema) assert expert_logits.ndim == 1 and len(expert_logits) == len(experts) # 1. call experts and await results jobs = [partial(cls._run_expert_forward, expert, *expert_args, **expert_kwargs) for expert in experts] results = run_and_await_k(jobs, k=k_min, timeout_after_k=timeout_after_k_min, timeout_total=timeout_total) alive_contexts, alive_outputs, alive_ix = zip(*[(result[0], result[1], ix) for ix, result in enumerate(results) if not isinstance(result, BaseException)]) # ^ ^ ^-- a list of indices of experts that returned outputs in time # \ \-- list of outputs of every expert that didn't die on us # \-- a list of autograd contexts, used for parallel backward # 2. compute softmax weights for alive experts and average outputs alive_ix = torch.as_tensor(alive_ix, device=expert_logits.device) alive_expert_probs = torch.softmax(expert_logits[alive_ix], dim=0) stacked_alive_outputs = tuple(map(torch.stack, zip(*alive_outputs))) flat_average_outputs = tuple((alive_expert_probs @ stacked_out.flatten(1)).view(*stacked_out.shape[1:]) for stacked_out in stacked_alive_outputs) # 3. save individual outputs for backward pass ctx.save_for_backward(expert_logits, alive_ix, alive_expert_probs, *stacked_alive_outputs) ctx._saved_non_tensors = alive_contexts, backward_k_min, backward_timeout return tuple(map(torch.Tensor.detach, flat_average_outputs)) @classmethod @once_differentiable def backward(cls, ctx, *grad_outputs_flat: torch.Tensor) -> Tuple[Optional[torch.Tensor], ...]: """ Like normal backward, but we ignore any experts that failed during backward pass """ expert_logits, alive_ix, alive_expert_probas, *stacked_alive_outputs = ctx.saved_tensors alive_contexts, backward_k_min, backward_timeout = ctx._saved_non_tensors jobs = [partial(cls._run_expert_backward, ctx, prob, *grad_outputs_flat) for ctx, prob in zip(alive_contexts, alive_expert_probas.split(1))] results = run_and_await_k(jobs, k=backward_k_min, timeout_after_k=None, timeout_total=backward_timeout) backward_survivors_in_alive_ix, survived_grad_inputs = zip(*((i, grads) for i, grads in enumerate(results))) backward_survivors_in_alive_ix = torch.as_tensor(backward_survivors_in_alive_ix, device=expert_logits.device) backward_survivors_ix = alive_ix[backward_survivors_in_alive_ix] survived_probas = torch.softmax(expert_logits[backward_survivors_ix], dim=0) weight_ratios = survived_probas / alive_expert_probas[backward_survivors_in_alive_ix] flat_grad_inputs = tuple((weight_ratios @ stacked_grad_inp.flatten(1)).view(stacked_grad_inp.shape[1:]) for stacked_grad_inp in map(torch.stack, zip(*survived_grad_inputs))) # compute grad w.r.t. logits grad_wrt_probs = sum(tuple( torch.sum(grad_out[None, ...] * stacked_avive_out[backward_survivors_in_alive_ix], dim=tuple(range(1, stacked_avive_out.ndim))) for grad_out, stacked_avive_out in zip(grad_outputs_flat, stacked_alive_outputs) )) softmax_jacobian = torch.diagflat(survived_probas) - torch.ger(survived_probas, survived_probas) grad_wrt_logits = grad_wrt_probs @ softmax_jacobian return grad_wrt_logits, None, None, None, None, None, None, None, *flat_grad_inputs @staticmethod def _run_expert_forward(expert: RemoteExpert, *args: torch.Tensor, **kwargs: torch.Tensor): """ Call remote expert and return flattened outputs. Compatible with concurrent autograd. """ flat_inputs = nested_flatten((args, kwargs)) return run_isolated_forward(_RemoteModuleCall, DUMMY, expert.uid, expert.host, expert.port, *flat_inputs) @staticmethod def _run_expert_backward(ctx: EmulatedAutogradContext, weight: torch.Tensor, *grad_outputs: torch.Tensor): backward_result = run_isolated_backward(_RemoteModuleCall, ctx, *(grad * weight for grad in grad_outputs)) grad_dummy, no_grad_uid, no_grad_hostname, no_grad_port, *grad_inputs = backward_result return grad_inputs