Class MethodRunner

axlearn/common/inference.py:46–137

Class for running a pjit-lowered model's method on a batch of data.

# pylint: disable-next=too-few-public-methods
class MethodRunner:
    """Class for running a pjit-lowered model's method on a batch of data."""

    def __init__(
        self,
        *,
        prng_key: Tensor,
        mesh: jax.sharding.Mesh,
        input_batch_partition_spec: DataPartitionType,
        jit_run_on_batch: Callable[
            [Tensor, NestedTensor],
            tuple[Tensor, NestedTensor, NestedTensor, NestedTensor],
        ],
    ):
        """Initializes MethodRunner object.

        Args:
            prng_key: the random key used for the first run, then updated by each call.
            mesh: mesh to be used during method running, same as the one used for pjit.
            input_batch_partition_spec: partition spec for input batches.
            jit_run_on_batch: callable which takes prng key, input batch and outputs
                updated prng key, outputs, summaries, and module outputs.
        """
        self._prng_key = prng_key
        self._mesh = mesh
        self._input_batch_partition_spec = input_batch_partition_spec
        self._jit_run_on_batch = jit_run_on_batch

    @dataclass(frozen=True)
    class Output:
        """Output class of MethodRunner."""

        # Output batch as a partitioned global array.
        output_batch: NestedTensor
        # Input batch as a partitioned global array.
        input_batch: NestedTensor
        # Summaries.
        summaries: NestedTensor
        # Module outputs.
        module_outputs: NestedTensor

    def __call__(self, input_batch: NestedTensor) -> Output:
        """Computes outputs and summaries for the given input.

        The convention is for input_batch to be global arrays.
        Output batches are global arrays.
        This symmetry in global arrays for both input and output batches allows
        users to chain multiple `MethodRunner`s together without extra host-device transfer.
        If the input_batch is host-local, it will be automatically converted to
            global input batch for ease-of-use.

        Args:
            input_batch: An input batch of data. By convention, these are global arrays.
                Host-local input batches are also accepted and converted to global input batches.

        Returns:
            An Output object containing global batch inputs, outputs and summaries.
            N.B. the returned input and output batches will have the same partitioning.