<commit_before># Copyright 2022 Google LLC.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Core Sparse MoE utils using pjit.

Many thanks to <NAME> and <NAME>, for helping with the einsum
implementation, and to <NAME> for helping writing the lift transform.

The following abbreviations are sometimes used to name the size of different
axes in the arrays.

G = num_groups. It must be a multiple of num_experts.
S = group_size.
E = num_experts.
C = capacity.
K = num_selected_experts. It must be <= num_experts.
"""
import abc
import math
from typing import Any, Callable, Dict, Mapping, Optional, Tuple

import flax.core.lift
import flax.linen.partitioning
import flax.linen.transforms
import flax.struct
import jax
from jax.experimental import pjit
import jax.numpy as jnp
import vmoe.partitioning


Array = jnp.ndarray
PartitionSpec = pjit.PartitionSpec
with_sharding_constraint = vmoe.partitioning.with_sharding_constraint
_add_axis_to_metadata = flax.linen.partitioning._add_axis_to_metadata  # pylint: disable=protected-access


class BaseDispatcher(abc.ABC):
  """Base class for different dispatcher implementations.

  Dispatchers are in charge of preparing the data to be dispatched to the
  different experts, and then combining the outputs of each expert for each
  item. There are different ways of doing so with different memory / flops /
  runtime implications when running on actual hardware.

  In all cases, when dispatching data, they take an array of shape (G, S, ...).
  The groups (G) are dispatched independently of each other. The items in each
  group (S) will take place in the buffer (of capacity C) of items to be
  processed by each expert (E). The output is an array of shape (E, G * C, ...)
  with the elements to be processed by each expert.

  When combining data, they take an array of shape (E, G * C, ...) and output
  an array of shape (G, S, ...). Notice that the trailing dimensions (...) at
  combine might not be the same as the ones at dispatch (e.g. if the expert
  changes the shape of the data).
  """

  @abc.abstractmethod
  def dispatch(self, data: Array) -> Array:
    """Dispatches data to experts.

    Args:
      data: (G, S, ...) array with the data to dispatch to the experts.

    Returns:
      (E, G * C, ...) array with the data to be processed by each expert.
    """

  @abc.abstractmethod
  def combine(self, data: Array) -> Array:
    """Combines outputs from multiple experts.

    Args:
      data: (E, G * C, ...) array with the output data from each expert.

    Returns:
      (G, S, ...) array with the combined outputs from each expert.
    """


@flax.struct.dataclass
class EinsumDispatcher(BaseDispatcher):
  """Dispatcher using Einsum.

  Attributes:
    combine_weights: (G, S, E, C) array with the combine weights for each item
      (G, S) for each expert (E) and buffer position (C).
    dispatch_weights: Optional. (G, S, E, C) array with the dispatch weights of
      each item (G, S) for each expert (E) and buffer position (C).
    partition_spec: Optional. PartitionSpec used to constrain the sharding of
      the data arrays. By default (None), no sharding constraint is specified.
    einsum_precision: Optional. Precision used in all the einsums (e.g.
      combining the outputs of different experts).
  """
  combine_weights: Array
  dispatch_weights: Optional[Array] = None
  partition_spec: Optional[PartitionSpec] = flax.struct.field(
      pytree_node=False, default=None)
  einsum_precision: jax.lax.Precision = flax.struct.field(
      pytree_node=False, default=jax.lax.Precision.DEFAULT)

  def dispatch(self, data: Array) -> Array:
    dispatch_weights = (
        self.combine_weights > 0
        if self.dispatch_weights is None else self.dispatch_weights)
    data = jnp.einsum("GSEC,GS...->GEC...", dispatch_weights, data,
                      precision=self.einsum_precision)
    return _dispatch(data, self.partition_spec)

  def combine(self, data: Array) -> Array:
    """Combines data from experts according to combine_weights."""
    num_groups, _, _, _ = self.combine_weights.shape
    data = _receive(data, num_groups, self.partition_spec)
    return jnp.einsum("GSEC,GEC...->GS...", self.combine_weights, data,
                      precision=self.einsum_precision)


@flax.struct.dataclass
class ExpertIndicesDispatcher(BaseDispatcher):
  """Dispatcher using scatter/gather with (expert, buffer) indices.

  Attributes:
    indices: (G, S, K, 2) integer array with the (expert, buffer) indices of
      each item (G, S) and their K-selected experts. The tuple (expert, buffer)
      for each item is represented in the last dimension (of size 2).
    combine_weights: (G, S, K) array with the combine weights of each item
      (G, S) and their K-selected experts.
    num_experts: Number of experts.
    capacity: Capacity of each expert's buffer per group.
    partition_spec: Optional. PartitionSpec used to constrain the sharding of
      the data arrays. By default (None), no sharding constraint is specified.
    einsum_precision: Optional. Precision used in all the einsums (e.g.
      combining the outputs of different experts).
  """
  indices: Array  # (G, S, K, 2).
  combine_weights: Array   # (G, S, K).
  num_experts: int = flax.struct.field(pytree_node=False)
  capacity: int = flax.struct.field(pytree_node=False)
  partition_spec: Optional[PartitionSpec] = flax.struct.field(
      pytree_node=False, default=None)
  einsum_precision: jax.lax.Precision = flax.struct.field(
      pytree_node=False, default=jax.lax.Precision.DEFAULT)

  def dispatch(self, data: Array) -> Array:
    num_groups, _, num_selected_experts, _ = self.indices.shape
    _, _, *item_shape = data.shape
    data = jnp.repeat(data, num_selected_experts, axis=1)
    indices = self.indices.reshape(num_groups, -1, 2)
    shape = (self.num_experts, self.capacity, *item_shape)
    data = jax.vmap(lambda x, i: _scatter_nd(i, x, shape))(data, indices)
    return _dispatch(data, self.partition_spec)

  def combine(self, data: Array) -> Array:
    num_groups, _, _ = self.combine_weights.shape
    data = _receive(data, num_groups, self.partition_spec)
    data = jax.vmap(lambda x, i: x[i[:, :, 0], i[:, :, 1]])(data, self.indices)
    # Mask invalid gathered data.
    mask = jnp.logical_and(self.indices[..., 0] < self.num_experts,
                           self.indices[..., 1] < self.capacity)
    data = data * mask.reshape(mask.shape + (1,) * (data.ndim - 3))
    # Weighted sum of the outputs of the K-selected experts for each item.
    return jnp.einsum("GSK...,GSK->GS...", data, self.combine_weights,
                      precision=jax.lax.Precision.HIGHEST)


@flax.struct.dataclass
class Bfloat16Dispatcher(BaseDispatcher):
  """Dispatcher wrapper converting data to bfloat16 to save bandwidth."""
  dispatcher: BaseDispatcher

  def dispatch(self, data: Array) -> Array:
    dtype = data.dtype
    data = _cast_to_bfloat16(data)
    data = self.dispatcher.dispatch(data)
    return data.astype(dtype)

  def combine(self, data: Array) -> Array:
    dtype = data.dtype
    data = _cast_to_bfloat16(data)
    data = self.dispatcher.combine(data)
    return data.astype(dtype)<commit_msg>internal<commit_after># Copyright 2022 Google LLC.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Core Sparse MoE utils using pjit.

Many thanks to <NAME> and <NAME>, for helping with the einsum
implementation, and to <NAME> for helping writing the lift transform.

The following abbreviations are sometimes used to name the size of different
axes in the arrays.

G = num_groups. It must be a multiple of num_experts.
S = group_size.
E = num_experts.
C = capacity.
K = num_selected_experts. It must be <= num_experts.
"""
import abc
import math
from typing import Any, Callable, Dict, Mapping, Optional, Tuple

import flax.core.lift
import flax.linen.partitioning
import flax.linen.transforms
import flax.struct
import jax
from jax.experimental import pjit
import jax.numpy as jnp
import vmoe.partitioning


Array = jnp.ndarray
PartitionSpec = pjit.PartitionSpec
with_sharding_constraint = vmoe.partitioning.with_sharding_constraint
_add_axis_to_metadata = flax.linen.partitioning._add_axis_to_metadata  # pylint: disable=protected-access


class BaseDispatcher(abc.ABC):
  """Base class for different dispatcher implementations.

  Dispatchers are in charge of preparing the data to be dispatched to the
  different experts, and then combining the outputs of each expert for each
  item. There are different ways of doing so with different memory / flops /
  runtime implications when running on actual hardware.

  In all cases, when dispatching data, they take an array of shape (G, S, ...).
  The groups (G) are dispatched independently of each other. The items in each
  group (S) will take place in the buffer (of capacity C) of items to be
  processed by each expert (E). The output is an array of shape (E, G * C, ...)
  with the elements to be processed by each expert.

  When combining data, they take an array of shape (E, G * C, ...) and output
  an array of shape (G, S, ...). Notice that the trailing dimensions (...) at
  combine might not be the same as the ones at dispatch (e.g. if the expert
  changes the shape of the data).
  """

  @abc.abstractmethod
  def dispatch(self, data: Array) -> Array:
    """Dispatches data to experts.

    Args:
      data: (G, S, ...) array with the data to dispatch to the experts.

    Returns:
      (E, G * C, ...) array with the data to be processed by each expert.
    """

  @abc.abstractmethod
  def combine(self, data: Array) -> Array:
    """Combines outputs from multiple experts.

    Args:
      data: (E, G * C, ...) array with the output data from each expert.

    Returns:
      (G, S, ...) array with the combined outputs from each expert.
    """


@flax.struct.dataclass
class DenseEinsumDispatcher(BaseDispatcher):
  """Dispatcher using Einsum, dispatching data to all experts.

  This is similar to EinsumDispatcher, but with the assumption that C = S.

  Attributes:
    combine_weights: (G, S, E) array with the combine weights for each item
      (G, S) for each expert (E).
    partition_spec: Optional. PartitionSpec used to constrain the sharding of
      the data arrays. By default (None), no sharding constraint is specified.
    einsum_precision: Optional. Precision used in all the einsums (e.g.
      combining the outputs of different experts).
  """
  combine_weights: Array
  partition_spec: Optional[PartitionSpec] = flax.struct.field(
      pytree_node=False, default=None)
  einsum_precision: jax.lax.Precision = flax.struct.field(
      pytree_node=False, default=jax.lax.Precision.DEFAULT)

  def dispatch(self, data: Array) -> Array:
    dispatch_weights = jnp.ones_like(self.combine_weights, dtype=jnp.bool_)
    data = jnp.einsum("GSE,GS...->GES...", dispatch_weights, data,
                      precision=self.einsum_precision)
    return _dispatch(data, self.partition_spec)

  def combine(self, data: Array) -> Array:
    """Combines data from experts according to combine_weights."""
    num_groups, _, _ = self.combine_weights.shape
    data = _receive(data, num_groups, self.partition_spec)
    return jnp.einsum("GSE,GES...->GS...", self.combine_weights, data,
                      precision=self.einsum_precision)


@flax.struct.dataclass
class EinsumDispatcher(BaseDispatcher):
  """Dispatcher using Einsum.

  Attributes:
    combine_weights: (G, S, E, C) array with the combine weights for each item
      (G, S) for each expert (E) and buffer position (C).
    dispatch_weights: Optional. (G, S, E, C) array with the dispatch weights of
      each item (G, S) for each expert (E) and buffer position (C).
    partition_spec: Optional. PartitionSpec used to constrain the sharding of
      the data arrays. By default (None), no sharding constraint is specified.
    einsum_precision: Optional. Precision used in all the einsums (e.g.
      combining the outputs of different experts).
  """
  combine_weights: Array
  dispatch_weights: Optional[Array] = None
  partition_spec: Optional[PartitionSpec] = flax.struct.field(
      pytree_node=False, default=None)
  einsum_precision: jax.lax.Precision = flax.struct.field(
      pytree_node=False, default=jax.lax.Precision.DEFAULT)

  def dispatch(self, data: Array) -> Array:
    dispatch_weights = (
        self.combine_weights > 0
        if self.dispatch_weights is None else self.dispatch_weights)
    data = jnp.einsum("GSEC,GS...->GEC...", dispatch_weights, data,
                      precision=self.einsum_precision)
    return _dispatch(data, self.partition_spec)

  def combine(self, data: Array) -> Array:
    """Combines data from experts according to combine_weights."""
    num_groups, _, _, _ = self.combine_weights.shape
    data = _receive(data, num_groups, self.partition_spec)
    return jnp.einsum("GSEC,GEC...->GS...", self.combine_weights, data,
                      precision=self.einsum_precision)


@flax.struct.dataclass
class ExpertIndicesDispatcher(BaseDispatcher):
  """Dispatcher using scatter/gather with (expert, buffer) indices.

  Attributes:
    indices: (G, S, K, 2) integer array with the (expert, buffer) indices of
      each item (G, S) and their K-selected experts. The tuple (expert, buffer)
      for each item is represented in the last dimension (of size 2).
    combine_weights: (G, S, K) array with the combine weights of each item
      (G, S) and their K-selected experts.
    num_experts: Number of experts.
    capacity: Capacity of each expert's buffer per group.
    partition_spec: Optional. PartitionSpec used to constrain the sharding of
      the data arrays. By default (None), no sharding constraint is specified.
    einsum_precision: Optional. Precision used in all the einsums (e.g.
      combining the outputs of different experts).
  """
  indices: Array  # (G, S, K, 2).
  combine_weights: Array   # (G, S, K).
  num_experts: int = flax.struct.field(pytree_node=False)
  capacity: int = flax.struct.field(pytree_node=False)
  partition_spec: Optional[PartitionSpec] = flax.struct.field(
      pytree_node=False, default=None)
  einsum_precision: jax.lax.Precision = flax.struct.field(
      pytree_node=False, default=jax.lax.Precision.DEFAULT)

  def dispatch(self, data: Array) -> Array:
    num_groups, _, num_selected_experts, _ = self.indices.shape
    _, _, *item_shape = data.shape
    data = jnp.repeat(data, num_selected_experts, axis=1)
    indices = self.indices.reshape(num_groups, -1, 2)
    shape = (self.num_experts, self.capacity, *item_shape)
    data = jax.vmap(lambda x, i: _scatter_nd(i, x, shape))(data, indices)
    return _dispatch(data, self.partition_spec)

  def combine(self, data: Array) -> Array:
    num_groups, _, _ = self.combine_weights.shape
    data = _receive(data, num_groups, self.partition_spec)
    data = jax.vmap(lambda x, i: x[i[:, :, 0], i[:, :, 1]])(data, self.indices)
    # Mask invalid gathered data.
    mask = jnp.logical_and(self.indices[..., 0] < self.num_experts,
                           self.indices[..., 1] < self.capacity)
    data = data * mask.reshape(mask.shape + (1,) * (data.ndim - 3))
    # Weighted sum of the outputs of the K-selected experts for each item.
    return jnp.einsum("GSK...,GSK->GS...", data, self.combine_weights,
                      precision=jax.lax.Precision.HIGHEST)


@flax.struct.dataclass
class Bfloat16Dispatcher(BaseDispatcher):
  """Dispatcher wrapper converting data to bfloat16 to save bandwidth."""
  dispatcher: BaseDispatcher

  def dispatch(self, data: Array) -> Array:
    dtype = data.dtype
    data = _cast_to_bfloat16(data)
    data = self.dispatcher.dispatch(data)
    return data.astype(dtype)

  def combine(self, data: Array) -> Array:
    dtype = data.dtype
    data = _cast_to_bfloat16(data)
    data = self.dispatcher.combine(data)
    return data.astype(dtype)


def get_dense_einsum_dispatcher(gates,
                                **dispatcher_kwargs) -> DenseEinsumDispatcher:
  # The dispatching algorithm is trivial, because all tokens are sent to
  # all experts this is coded implicitly in the DenseEinsumDispatcher class.
  return DenseEinsumDispatcher(combine_weights=gates, **dispatcher_kwargs)
{'fraction_non_alphanumeric': 0.09776876267748479, 'fraction_numerical': 0.0033613445378151263, 'mean_word_length': 3.51255230125523, 'pattern_counts': {'":': 0, '<': 15, '<?xml version=': 0, '>': 38, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3566374', 'n_tokens_mistral': 5348, 'n_tokens_neox': 5179, 'n_words': 1841}