t5x/examples/decoder_only/network.py

# Copyright 2022 The T5X Authors.
#
# 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.

"""Minimal decoder-only Transformer model."""

from typing import Any, Optional, Sequence

from flax import linen as nn
from flax import struct
import jax.numpy as jnp
from t5x.examples.decoder_only import layers


@struct.dataclass
class TransformerConfig:
  """Global hyperparameters used to minimize obnoxious kwarg plumbing."""
  vocab_size: int
  # Activation dtypes.
  dtype: Any = jnp.float32
  emb_dim: int = 512
  num_heads: int = 8
  num_layers: int = 6
  head_dim: int = 64
  mlp_dim: int = 2048
  # Activation functions are retrieved from Flax.
  mlp_activations: Sequence[str] = ('relu',)
  dropout_rate: float = 0.1
  # If `True`, the embedding weights are used in the decoder output layer.
  logits_via_embedding: bool = False


class DecoderLayer(nn.Module):
  """Transformer decoder layer."""
  config: TransformerConfig

  @nn.compact
  def __call__(self,
               inputs: jnp.ndarray,
               decoder_mask: Optional[jnp.ndarray] = None,
               deterministic: bool = False,
               decode: bool = False,
               max_decode_length: Optional[int] = None,
               prefill: bool = False,
               prefill_lengths: Optional[jnp.ndarray] = None):
    """Applies decoder block module."""
    cfg = self.config

    # Relative position embedding as attention biases.
    l = max_decode_length if decode and max_decode_length else inputs.shape[-2]

    # During decoding, this module will be called with `decode=True` first to
    # initialize the decoder cache, including a cached relpos bias. The prefill
    # codepath will call this once again with `decode=False`, which is slightly
    # wasteful but generally harmless. During subsequent decode steps, this will
    # be called with `decode=True` and will reuse the cached bias. This
    # significantly improves performance during decoding with many decode steps.
    decoder_bias = layers.RelativePositionBiases(
        num_buckets=32,
        max_distance=128,
        num_heads=cfg.num_heads,
        dtype=cfg.dtype,
        embedding_init=nn.initializers.variance_scaling(1.0, 'fan_avg',
                                                        'uniform'),
        name='relpos_bias')(
            l, l, False, decode=decode)

    # `inputs` is layer input with a shape [batch, length, emb_dim].
    x = layers.LayerNorm(
        dtype=cfg.dtype, name='pre_self_attention_layer_norm')(
            inputs)

    # Self-attention block
    x = layers.MultiHeadDotProductAttention(
        num_heads=cfg.num_heads,
        dtype=cfg.dtype,
        head_dim=cfg.head_dim,
        dropout_rate=cfg.dropout_rate,
        name='self_attention')(
            x,
            x,
            decoder_mask,
            decoder_bias,
            deterministic=deterministic,
            decode=decode,
            prefill=prefill,
            prefill_lengths=prefill_lengths)
    x = nn.Dropout(
        rate=cfg.dropout_rate,
        broadcast_dims=(-2,),
        name='post_self_attention_dropout')(
            x, deterministic=deterministic)
    x = x + inputs

    # MLP block.
    y = layers.LayerNorm(dtype=cfg.dtype, name='pre_mlp_layer_norm')(x)
    y = layers.MlpBlock(
        intermediate_dim=cfg.mlp_dim,
        activations=cfg.mlp_activations,
        intermediate_dropout_rate=cfg.dropout_rate,
        dtype=cfg.dtype,
        name='mlp',
    )(y, deterministic=deterministic)
    y = nn.Dropout(
        rate=cfg.dropout_rate, broadcast_dims=(-2,), name='post_mlp_dropout')(
            y, deterministic=deterministic)
    y = y + x

    return y


class Decoder(nn.Module):
  """A stack of decoder layers."""
  config: TransformerConfig

  @nn.compact
  def __call__(self,
               decoder_input_tokens: jnp.ndarray,
               decoder_target_tokens: jnp.ndarray,
               decoder_segment_ids: Optional[jnp.ndarray] = None,
               decoder_positions: Optional[jnp.ndarray] = None,
               decoder_causal_attention: Optional[jnp.ndarray] = None,
               *,
               enable_dropout: bool = True,
               decode: bool = False,
               max_decode_length: Optional[int] = None,
               prefill: Optional[bool] = None,
               prefill_lengths: Optional[jnp.ndarray] = None):
    """Applies LanguageModel on the inputs.

    For a decoder-only architecture with the notion of "prefix", e.g., a prefix
    LM where the prefix corresponds to the "inputs" of a supervised dataset, we
    perform the "prefill" operation to fill the autoregressive cache
    corresponding to the prefix region in one go. Then the autoregressive
    decoding starts after the prefix. This makes the decoding process more
    efficient. In addition, it gives an option to use bidirectional attention in
    the prefix region because the cache is filled simultaneously.

    Args:
      decoder_input_tokens: input token to the decoder.
      decoder_target_tokens: target token to the decoder.
      decoder_segment_ids: decoder segmentation info for packed examples.
      decoder_positions: decoder subsequence positions for packed examples.
      decoder_causal_attention: a binary mask indicating the portion of the
        sequence to apply bidirectional attention to instead of causal. As an
        example, useful to specify the "inputs" portion of a concatenated
        sequence for a prefix LM.
      enable_dropout: enables dropout if set to True.
      decode: whether to prepare and use an autoregressive cache as opposed to
        using teacher-forcing.
      max_decode_length: maximum sequence length to be decoded.
      prefill: whether to run a partial sequence to prefill the cache.
      prefill_lengths: an array of shape [batch] denoting the length of each
        partial sequence we are filling in the cache.

    Returns:
      logits array.
    """
    cfg = self.config
    deterministic = not enable_dropout
    assert decoder_input_tokens.ndim == 2  # [batch, len]

    if decode:
      decoder_mask = None
    else:
      decoder_mask = layers.make_decoder_mask(
          decoder_target_tokens=decoder_target_tokens,
          dtype=cfg.dtype,
          decoder_causal_attention=decoder_causal_attention,
          decoder_segment_ids=decoder_segment_ids)

    embedding = layers.Embed(
        num_embeddings=cfg.vocab_size,
        features=cfg.emb_dim,
        dtype=cfg.dtype,
        attend_dtype=jnp.float32,  # for logit training stability
        embedding_init=nn.initializers.normal(stddev=1.0),
        one_hot=True,
        name='token_embedder')
    y = embedding(decoder_input_tokens.astype('int32'))

    y = nn.Dropout(
        rate=cfg.dropout_rate, broadcast_dims=(-2,), name='input_dropout')(
            y, deterministic=deterministic)
    y = y.astype(cfg.dtype)

    for lyr in range(cfg.num_layers):
      # [batch, length, emb_dim] -> [batch, length, emb_dim]
      y = DecoderLayer(
          config=cfg, name=f'layers_{lyr}')(
              y,
              decoder_mask=decoder_mask,
              deterministic=deterministic,
              decode=decode,
              max_decode_length=max_decode_length,
              prefill=prefill,
              prefill_lengths=prefill_lengths)

    y = layers.LayerNorm(dtype=cfg.dtype, name='decoder_norm')(y)
    y = nn.Dropout(
        rate=cfg.dropout_rate, broadcast_dims=(-2,), name='output_dropout')(
            y, deterministic=deterministic)

    # [batch, length, emb_dim] -> [batch, length, vocab_size]
    if cfg.logits_via_embedding:
      # Use the transpose of embedding matrix for the logit transform.
      logits = embedding.attend(y)
      # Correctly normalize pre-softmax logits for this shared case.
      logits = logits / jnp.sqrt(y.shape[-1])
    else:
      # Use a separate dense layer for the logit transform.
      logits = layers.DenseGeneral(
          cfg.vocab_size,
          dtype=jnp.float32,  # Use float32 for stabiliity.
          kernel_axes=('embed', 'vocab'),
          name='logits_dense')(
              y)
    return logits


# TODO(hwchung): remove this after figuring out the name scope issue.
class DecoderWrapper(nn.Module):
  """Thin wrapper for the outer "decoder/" name scope."""

  config: TransformerConfig

  def setup(self):
    self.decoder = Decoder(self.config, name='decoder')

  def __call__(self, *args, **kwargs):
    return self.decoder(*args, **kwargs)