utils.py 56.2 KB
Newer Older
1
# Copyright (c) 2022-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
3
#
# See LICENSE for license information.
zlsh80826's avatar
zlsh80826 committed
4
"""Utility for the TE layer tests"""
5

6
import functools
7
import math
8
import operator
9
from typing import Any, Callable, Dict, Tuple, Sequence, Union, Iterable, Optional
10

11
import jax
12
13
import jax.numpy as jnp
import numpy as np
14
15
from flax import linen as nn
from flax.linen import partitioning as nn_partitioning
16
from flax.linen.attention import combine_masks
17
18
19
from jax import lax, vmap
from jax import nn as jax_nn
from jax import random as jax_random
20

21
from transformer_engine.jax.attention import AttnMaskType, make_swa_mask
22
23
from transformer_engine.jax.fp8 import DType as TEDType

24
25
26
27
PRNGKey = Any
Shape = Tuple[int, ...]
DType = jnp.dtype
Array = Any
28
29
30
PrecisionLike = Union[
    None, str, lax.Precision, Tuple[str, str], Tuple[lax.Precision, lax.Precision]
]
31
32
Initializer = Callable[[PRNGKey, Shape, DType], Array]

33

34
def is_devices_enough(required):
zlsh80826's avatar
zlsh80826 committed
35
36
37
    """
    Check if the available GPUs is enough
    """
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
    return len(jax.devices()) >= required


def _generate_drop_path_shape(shape: Sequence[int], batch_dim: int) -> Sequence[int]:
    # Generate broadcast dims for drop_path.
    drop_path_shape = list(range(0, len(shape)))
    drop_path_shape.pop(batch_dim)
    return drop_path_shape


def _normalize_axes(axes: Iterable[int], ndim: int) -> Tuple[int]:
    # A tuple by convention. len(axes_tuple) then also gives the rank efficiently.
    return tuple(ax if ax >= 0 else ndim + ax for ax in axes)


def _canonicalize_tuple(x):
    if isinstance(x, Iterable):
        return tuple(x)
    return (x,)


def _convert_to_activation_function(fn_or_string: Union[str, Callable]) -> Callable:
    """Convert a string to an activation function."""
61
    if fn_or_string == "linear":
62
63
64
65
66
67
68
69
70
71
72
        return lambda x: x
    if isinstance(fn_or_string, str):
        return getattr(nn, fn_or_string)
    if callable(fn_or_string):
        return fn_or_string
    raise ValueError(f"don't know how to convert {fn_or_string} to an activation function")


def combine_biases(*masks: Optional[Array]):
    """Combine attention biases.

73
74
    Args:
      *masks: set of attention bias arguments to combine, some can be None.
75

76
77
78
    Returns:
      Combined mask, reduced by summation, returns None if no masks given.
    """
79
80
81
    masks = [m for m in masks if m is not None]
    if not masks:
        return None
82
83
84
    assert all(
        map(lambda x: x.ndim == masks[0].ndim, masks)
    ), f"masks must have same rank: {tuple(map(lambda x: x.ndim, masks))}"
85
86
87
88
89
90
    mask, *other_masks = masks
    for other_mask in other_masks:
        mask = mask + other_mask
    return mask


91
92
93
class DotProductAttention(nn.Module):
    transpose_batch_sequence: bool = True
    scale_attn_logits: bool = True
94
    dropout_rate: float = 0.0
95
96
    dtype: DType = jnp.float32
    float32_logits: bool = False
97
98
    """Computes dot-product attention given query, key, and value.

99
100
101
    This is the core function for applying attention based on
    https://arxiv.org/abs/1706.03762. It calculates the attention weights given
    query and key and combines the values using the attention weights.
102

103
104
105
106
107
108
    Args:
        dropout_rate: dropout rate
        dtype: the dtype of the computation (default: float32)
        float32_logits: bool, if True then compute logits in float32 to avoid
        numerical issues with bfloat16.
    """
109

110
    @nn.compact
111
112
113
114
115
116
117
118
    def __call__(
        self,
        query: Array,
        key: Array,
        value: Array,
        bias: Optional[Array] = None,
        deterministic: bool = False,
    ):
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
        """
        Args:
            query: queries for calculating attention with shape of `[batch, q_length,
            num_heads, qk_depth_per_head]`.
            key: keys for calculating attention with shape of `[batch, kv_length,
            num_gqa_groups, qk_depth_per_head]`.
            value: values to be used in attention with shape of `[batch, kv_length,
            num_gqa_groups, v_depth_per_head]`.
            bias: bias for the attention weights. This should be broadcastable to the
            shape `[batch, num_heads, q_length, kv_length]` This can be used for
            incorporating causal masks, padding masks, proximity bias, etc.
            dropout_rng: JAX PRNGKey: to be used for dropout
            deterministic: bool, deterministic or not (to apply dropout)
        Returns:
            Output of shape `[batch, length, num_heads, v_depth_per_head]`.
        """
135
        assert key.ndim == query.ndim == value.ndim, "q, k, v must have same rank."
136
        batch_dim = 1 if self.transpose_batch_sequence else 0
137
138
139
        assert (
            query.shape[batch_dim] == key.shape[batch_dim] == value.shape[batch_dim]
        ), "q, k, v batch dims must match."
140
        sequence_dim = 0 if self.transpose_batch_sequence else 1
141
142
143
        assert key.shape[sequence_dim] == value.shape[sequence_dim], "k, v lengths must match."
        assert key.shape[-2] == value.shape[-2], "k, v num_heads must match."
        assert query.shape[-1] == key.shape[-1], "q, k head_dim must match."
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161

        if self.scale_attn_logits:
            head_dim = query.shape[-1]
            depth_scaling = jnp.sqrt(head_dim).astype(self.dtype)
            query = query / depth_scaling

        # Casting logits and softmax computation for float32 for model stability.
        if self.float32_logits:
            query = query.astype(jnp.float32)
            key = key.astype(jnp.float32)

        # `attn_weights`: [batch, num_heads, groups, q_length, kv_length]
        h_q, h_kv = query.shape[-2], key.shape[-2]
        assert (h_q % h_kv == 0) and (h_q >= h_kv)
        group_size = h_q // h_kv
        grouped_query = query.reshape((*query.shape[:2], h_kv, group_size, query.shape[-1]))

        if self.transpose_batch_sequence:
162
            attn_weights = jnp.einsum("qbhgd,kbhd->bhgqk", grouped_query, key)
163
        else:
164
            attn_weights = jnp.einsum("bqhgd,bkhd->bhgqk", grouped_query, key)
165
166
167
168
169
170
171
172
173
174
175
176
177
178

        # reshape back to normal DPA shape for bias/softmax/dropout
        b, h, g, q, k = attn_weights_with_groups_shape = attn_weights.shape
        attn_weights_without_groups_shape = (b, h * g, q, k)
        attn_weights = attn_weights.reshape(attn_weights_without_groups_shape)

        # Apply attention bias: masking, dropout, proximity bias, etc.
        if bias is not None:
            attn_weights = attn_weights + bias.astype(attn_weights.dtype)

        # Normalize the attention weights across `kv_length` dimension.
        attn_weights = jax_nn.softmax(attn_weights).astype(self.dtype)

        # Apply attention dropout.
179
        if not deterministic and self.dropout_rate > 0.0:
180
181
182
183
            keep_prob = 1.0 - self.dropout_rate
            # T5 broadcasts along the "length" dim, but unclear which one that
            # corresponds to in positional dimensions here, assuming query dim.
            dropout_shape = list(attn_weights.shape)
184
            dropout_rng = self.make_rng("dropout")
185
            keep = jax_random.bernoulli(dropout_rng, keep_prob, dropout_shape)
186
            multiplier = keep.astype(attn_weights.dtype) / jnp.asarray(keep_prob, dtype=self.dtype)
187
188
189
190
191
192
            attn_weights = attn_weights * multiplier

        attn_weights = attn_weights.reshape(attn_weights_with_groups_shape)

        # Take the linear combination of `value`.
        if self.transpose_batch_sequence:
193
            return jnp.einsum("bhgqk,kbhd->qbhgd", attn_weights, value).reshape(query.shape)
194

195
        return jnp.einsum("bhgqk,bkhd->bqhgd", attn_weights, value).reshape(query.shape)
196
197
198
199
200


class DenseGeneral(nn.Module):
    """A linear transformation with flexible axes and FP8 support.

201
202
203
204
205
206
207
    Attributes:
    features: tuple with numbers of output features.
    axis: tuple with axes to apply the transformation on.
    dtype: the dtype of the computation (default: float32).
    kernel_init: initializer function for the weight matrix.
    use_bias: whether to add a bias to the output (default: False).
    bias_init: initializer function for the bias vector.
208
    """
209

210
211
212
213
214
215
216
217
218
219
220
    features: Union[Iterable[int], int]
    axis: Union[Iterable[int], int] = -1
    dtype: DType = jnp.float32
    kernel_init: Initializer = None
    kernel_axes: Tuple[str, ...] = ()
    use_bias: bool = False
    bias_init: Initializer = nn.initializers.zeros
    bias_axes: Tuple[str, ...] = ()

    def __post_init__(self):
        if self.kernel_init is None:
221
            self.kernel_init = nn.initializers.variance_scaling(1.0, "fan_in", "truncated_normal")
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
        super().__post_init__()

    @nn.compact
    def __call__(self, inputs: Array) -> Array:
        """Applies a linear transformation to the inputs along multiple dimensions.

        Args:
        inputs: The nd-array to be transformed.

        Returns:
        The transformed input.
        """
        features = _canonicalize_tuple(self.features)
        axis = _canonicalize_tuple(self.axis)

        inputs = jnp.asarray(inputs, self.dtype)
        axis = _normalize_axes(axis, inputs.ndim)

        kernel_shape = tuple(inputs.shape[ax] for ax in axis) + features
        kernel_param_shape = (np.prod([inputs.shape[ax] for ax in axis]), np.prod(features))
242
243
244
        kernel = nn_partitioning.param_with_axes(
            "kernel", self.kernel_init, kernel_param_shape, jnp.float32, axes=self.kernel_axes
        )
245
246
247
248
249

        kernel = jnp.asarray(kernel, self.dtype)
        kernel = jnp.reshape(kernel, kernel_shape)

        if self.use_bias:
250
251
252
            bias = nn_partitioning.param_with_axes(
                "bias", self.bias_init, self.features, jnp.float32, axes=self.bias_axes
            )
253
            bias = bias.astype(self.dtype)
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
        else:
            bias = None

        contract_ind = tuple(range(0, len(axis)))

        y = lax.dot_general(inputs, kernel, ((axis, contract_ind), ((), ())))

        if bias is not None:
            y += jnp.reshape(bias, (1,) * (y.ndim - 1) + (-1,))
        return y


class MlpBlock(nn.Module):
    """Transformer MLP / feed-forward block.

269
270
271
272
273
274
275
276
277
278
    Attributes:
      intermediate_dim: Shared dimension of hidden layers.
      activations: Type of activations for each layer.  Each element is either
        'linear', a string function name in flax.linen, or a function.
      kernel_init: Kernel function, passed to the dense layers.
      deterministic: Whether the dropout layers should be deterministic.
      intermediate_dropout_rate: Dropout rate used after the intermediate layers.
      dtype: Type for the dense layer.
    """

279
280
    transpose_batch_sequence: bool
    intermediate_dim: int = 2048
281
    activations: Sequence[Union[str, Callable]] = ("relu",)
282
283
    kernel_init: Initializer = None
    intermediate_dropout_rate: float = 0.1
284
285
    intermediate_dropout_dims: Sequence[int] = ()
    use_bias: bool = False
286
    dtype: Any = jnp.float32
287
    fuse_wi: bool = True
288
289
290

    def __post_init__(self):
        if self.kernel_init is None:
291
            self.kernel_init = nn.initializers.variance_scaling(1.0, "fan_in", "truncated_normal")
292
293
294
295
296
297
298
299
300
301
        super().__post_init__()

    @nn.compact
    def __call__(self, inputs, deterministic: bool = False):
        """Applies Transformer MlpBlock module."""
        # Iterate over specified MLP input activation functions.
        # e.g. ('relu',) or ('gelu', 'linear') for gated-gelu.

        activations = []
        if self.fuse_wi:
302
            dense_name = "wi"
303
            num_activations = len(self.activations)
304
305
306
307
308
309
310
311
312
            x = DenseGeneral(
                self.intermediate_dim * num_activations,
                dtype=self.dtype,
                kernel_init=self.kernel_init,
                kernel_axes=("embed", "mlp"),
                use_bias=self.use_bias,
                bias_axes="mlp",
                name=dense_name,
            )(inputs)
313
314
315
316
317
318
            x = jnp.split(x, num_activations, axis=-1)
            for idx, act_fn in enumerate(self.activations):
                x_i = _convert_to_activation_function(act_fn)(x[idx])
                activations.append(x_i)
        else:
            for idx, act_fn in enumerate(self.activations):
319
320
321
322
323
324
325
326
327
328
                dense_name = "wi" if len(self.activations) == 1 else f"wi_{idx}"
                x = DenseGeneral(
                    self.intermediate_dim,
                    dtype=self.dtype,
                    kernel_init=self.kernel_init,
                    kernel_axes=("embed", "mlp"),
                    use_bias=self.use_bias,
                    bias_axes="mlp",
                    name=dense_name,
                )(inputs)
329
330
331
332
333
334
                x = _convert_to_activation_function(act_fn)(x)
                activations.append(x)

        # Take elementwise product of above intermediate activations.
        x = functools.reduce(operator.mul, activations)
        # Apply dropout and final dense output projection.
335
336
337
338
339
        x = nn.Dropout(
            rate=self.intermediate_dropout_rate, broadcast_dims=self.intermediate_dropout_dims
        )(
            x, deterministic=deterministic
        )  # Broadcast along length.
340
        if self.transpose_batch_sequence:
341
            x = nn_partitioning.with_sharding_constraint(x, ("length", "batch", "mlp"))
342
        else:
343
344
345
346
347
348
349
350
351
352
            x = nn_partitioning.with_sharding_constraint(x, ("batch", "length", "mlp"))
        output = DenseGeneral(
            inputs.shape[-1],
            dtype=self.dtype,
            kernel_init=self.kernel_init,
            kernel_axes=("mlp", "embed"),
            use_bias=self.use_bias,
            bias_axes="embed",
            name="wo",
        )(x)
353
354
355
        return output


356
def apply_rotary_pos_emb_alternate(
357
358
359
360
361
362
363
364
    inputs: jnp.ndarray,
    position: jnp.ndarray,
    min_timescale: int = 1,
    max_timescale: int = 10000,
):
    embedding_dim = inputs.shape[-1]
    half_embedding_dim = embedding_dim // 2
    fraction = 2 * jnp.arange(0, half_embedding_dim) / embedding_dim
365
    timescale = min_timescale * (max_timescale / min_timescale) ** fraction
366
367
368
369
370
371
372
373
374
375
376
377
378
    timescale = jnp.expand_dims(timescale, axis=tuple(range(inputs.ndim - 1)))
    position = jnp.expand_dims(position, axis=tuple(range(2, inputs.ndim)))
    sinusoid_inp = position / timescale
    sin = jnp.sin(sinusoid_inp)
    cos = jnp.cos(sinusoid_inp)
    first_half, second_half = jnp.split(inputs, 2, axis=-1)
    first_part = first_half * cos - second_half * sin
    second_part = second_half * cos + first_half * sin
    first_part = first_part.astype(inputs.dtype)
    second_part = second_part.astype(inputs.dtype)
    return jnp.concatenate([first_part, second_part], axis=-1)


379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
def apply_rotary_pos_emb_consecutive(
    inputs: jnp.ndarray,
    position: jnp.ndarray,
    min_timescale: int = 1,
    max_timescale: int = 10000,
):
    embedding_dim = inputs.shape[-1]
    half_embedding_dim = embedding_dim // 2
    fraction = 2 * jnp.arange(0, half_embedding_dim) / embedding_dim

    inputs_shifted_left = jnp.concatenate([inputs[..., 1:], inputs[..., :1]], axis=-1)
    inputs_shifted_right = jnp.concatenate([inputs[..., -1:], inputs[..., :-1]], axis=-1)
    inputs_shifted = jax.lax.select(
        jnp.tile(
            jnp.mod(jnp.arange(embedding_dim, dtype=jnp.int32), 2),
            inputs.shape[:-1] + (1,),
        ),
        inputs_shifted_right,
        inputs_shifted_left,
    )
    fraction = jnp.repeat(fraction, 2)
400
    timescale = min_timescale * (max_timescale / min_timescale) ** fraction
401
402
403
404
405
406
407
408
409
410
411
412

    position = jnp.expand_dims(position, axis=tuple(range(2, inputs.ndim)))

    sinusoid_inp = position / timescale
    sin = jnp.sin(sinusoid_inp)
    cos = jnp.cos(sinusoid_inp)
    sign = jnp.sign(jnp.mod(jnp.arange(embedding_dim, dtype=jnp.int32), 2) - 0.5)
    outputs = inputs * cos + inputs_shifted * sin * sign

    return outputs


413
414
415
416
417
418
419
420
421
dynamic_vector_slice_in_dim = vmap(lax.dynamic_slice_in_dim, in_axes=(None, 0, None, None))


class MultiHeadAttention(nn.Module):
    """Multi-head dot-product attention.

    Attributes:
      num_heads: number of attention heads. Features (i.e. inputs_q.shape[-1])
        should be divisible by the number of heads.
zlsh80826's avatar
zlsh80826 committed
422
      num_gqa_groups: number of kv attention heads
423
424
425
426
427
428
      head_dim: dimension of each head.
      dtype: the dtype of the computation.
      dropout_rate: dropout rate
      kernel_init: initializer for the kernel of the Dense layers.
      float32_logits: bool, if True then compute logits in float32 to avoid
        numerical issues with bfloat16.
429
    """
430

zlsh80826's avatar
zlsh80826 committed
431
432
433
434
    num_heads: int = 8
    num_gqa_groups: int | None = None
    head_dim: int = 64
    transpose_batch_sequence: bool = True
435
    dtype: DType = jnp.float32
436
    dropout_rate: float = 0.0
437
    kernel_init: Initializer = None
438
    float32_logits: bool = False  # computes logits in float32 for stability.
439
440
    scale_attn_logits: bool = False
    scaled_query_init: bool = True
441
    enable_rotary_pos_emb: bool = False
442
    rotary_pos_emb_group_method: str = "consecutive"
443
    fuse_qkv: bool = True
444
    use_bias: bool = False
445
446
447

    def __post_init__(self):
        if self.kernel_init is None:
448
            self.kernel_init = nn.initializers.variance_scaling(1.0, "fan_in", "normal")
zlsh80826's avatar
zlsh80826 committed
449
450
        if self.num_gqa_groups is None:
            self.num_gqa_groups = self.num_attention_heads
451
452
453
        super().__post_init__()

    @nn.compact
454
455
456
457
458
459
460
461
462
463
    def __call__(
        self,
        inputs_q: Array,
        inputs_kv: Array,
        mask: Optional[Array] = None,
        bias: Optional[Array] = None,
        *,
        decode: bool = False,
        deterministic: bool = False,
    ) -> Array:
464
465
        """Applies multi-head dot product attention on the input data.

466
467
        Projects the inputs into multi-headed query, key, and value vectors,
        applies dot-product attention and project the results to an output vector.
468

469
470
471
472
473
474
        There are two modes: decoding and non-decoding (e.g., training). The mode is
        determined by `decode` argument. For decoding, this method is called twice,
        first to initialize the cache and then for an actual decoding process. The
        two calls are differentiated by the presence of 'cached_key' in the variable
        dict. In the cache initialization stage, the cache variables are initialized
        as zeros and will be filled in the subsequent decoding process.
475

476
477
478
479
        In the cache initialization call, `inputs_q` has a shape [batch, length,
        q_features] and `inputs_kv`: [batch, length, kv_features]. During the
        incremental decoding stage, query, key and value all have the shape [batch,
        1, qkv_features] corresponding to a single step.
480

481
482
483
484
485
486
487
        Args:
          inputs_q: input queries of shape `[batch, q_length, q_features]`.
          inputs_kv: key/values of shape `[batch, kv_length, kv_features]`.
          mask: attention mask of shape `[batch, num_heads, q_length, kv_length]`.
          bias: attention bias of shape `[batch, num_heads, q_length, kv_length]`.
          decode: Whether to prepare and use an autoregressive cache.
          deterministic: Disables dropout if set to True.
488

489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
        Returns:
          output of shape `[batch, length, q_features]`.
        """
        q_projection = functools.partial(
            DenseGeneral,
            axis=-1,
            features=self.num_heads * self.head_dim,
            kernel_axes=("embed", "joined_kv"),
            use_bias=self.use_bias,
            bias_axes="joined_kv",
            dtype=self.dtype,
        )

        kv_projection = functools.partial(
            DenseGeneral,
            axis=-1,
            features=self.num_gqa_groups * self.head_dim,
            kernel_axes=("embed", "joined_kv"),
            use_bias=self.use_bias,
            bias_axes="joined_kv",
            dtype=self.dtype,
        )
511
512
513
514
515

        # NOTE: T5 does not explicitly rescale the attention logits by
        #       1/sqrt(depth_kq)!  This is folded into the initializers of the
        #       linear transformations, which is equivalent under Adafactor
        depth_scaling = jnp.sqrt(self.head_dim).astype(self.dtype)
516
517
518
        query_init = lambda *args: self.kernel_init(*args) / (
            depth_scaling if self.scaled_query_init else 1.0
        )
519
520
521
522
523
524
525
526
527
528
529
530

        # Project inputs_q to multi-headed q/k/v
        # dimensions are then [batch, length, num_heads, head_dim]

        def qkv_init(key, shape, dtype):
            assert shape[-1] % 3 == 0

            q_shape = (shape[0], shape[1] // 3)
            k_shape = (shape[0], shape[1] // 3)
            v_shape = (shape[0], shape[1] // 3)

            q_kernel = query_init(key, q_shape, dtype)
zlsh80826's avatar
zlsh80826 committed
531
532
            k_kernel = self.kernel_init(key, k_shape, dtype)
            v_kernel = self.kernel_init(key, v_shape, dtype)
533
534
535

            return jnp.concatenate([q_kernel, k_kernel, v_kernel], axis=-1, dtype=dtype)

536
537
538
        is_self_attn = inputs_q is inputs_kv
        is_gqa = self.num_heads != self.num_gqa_groups
        is_qkvpack = is_self_attn and not is_gqa
zlsh80826's avatar
zlsh80826 committed
539

540
        if self.fuse_qkv:
zlsh80826's avatar
zlsh80826 committed
541
            if is_qkvpack:
542
543
544
545
546
547
548
549
550
551
                qkv_proj = DenseGeneral(
                    axis=-1,
                    features=self.num_heads * self.head_dim * 3,
                    kernel_axes=("embed", "joined_kv"),
                    kernel_init=qkv_init,
                    use_bias=self.use_bias,
                    bias_axes="joined_kv",
                    name="qkv",
                    dtype=self.dtype,
                )(inputs_kv)
552
                query, key, value = jnp.split(
553
554
555
556
                    qkv_proj,
                    [self.num_heads * self.head_dim, self.num_heads * self.head_dim * 2],
                    axis=-1,
                )
557
            else:
558
559
560
561
562
563
564
565
566
567
568
569
                query = q_projection(kernel_init=query_init, name="query")(inputs_q)

                kv_proj = DenseGeneral(
                    axis=-1,
                    features=self.num_gqa_groups * self.head_dim * 2,
                    kernel_axes=("embed", "joined_kv"),
                    kernel_init=self.kernel_init,
                    use_bias=self.use_bias,
                    bias_axes="joined_kv",
                    name="kv",
                    dtype=self.dtype,
                )(inputs_kv)
zlsh80826's avatar
zlsh80826 committed
570
                key, value = jnp.split(kv_proj, [self.num_gqa_groups * self.head_dim], axis=-1)
571
        else:
572
573
574
            query = q_projection(kernel_init=query_init, name="query")(inputs_q)
            key = kv_projection(kernel_init=self.kernel_init, name="key")(inputs_kv)
            value = kv_projection(kernel_init=self.kernel_init, name="value")(inputs_kv)
575

576
577
578
579
        if self.enable_rotary_pos_emb:
            batch_dim = 1 if self.transpose_batch_sequence else 0
            seq_dim = 1 - batch_dim

580
581
            q_position = jnp.expand_dims(jnp.arange(query.shape[seq_dim]), axis=batch_dim)
            k_position = jnp.expand_dims(jnp.arange(query.shape[seq_dim]), axis=batch_dim)
582

583
            if self.rotary_pos_emb_group_method == "alternate":
584
585
586
587
                apply_rotary_pos_emb = apply_rotary_pos_emb_alternate
            else:
                apply_rotary_pos_emb = apply_rotary_pos_emb_consecutive

588
589
590
591
            query = query.reshape((*query.shape[:2], self.num_heads, self.head_dim))
            key = key.reshape((*key.shape[:2], self.num_gqa_groups, self.head_dim))
            query = apply_rotary_pos_emb(query, q_position)
            key = apply_rotary_pos_emb(key, k_position)
592

zlsh80826's avatar
zlsh80826 committed
593
594
595
        query = query.reshape((*query.shape[:2], self.num_heads, self.head_dim))
        key = key.reshape((*key.shape[:2], self.num_gqa_groups, self.head_dim))
        value = value.reshape((*value.shape[:2], self.num_gqa_groups, self.head_dim))
596
597

        if self.transpose_batch_sequence:
598
599
600
601
602
603
604
            query = nn_partitioning.with_sharding_constraint(
                query, ("length", "batch", "heads", "kv")
            )
            key = nn_partitioning.with_sharding_constraint(key, ("length", "batch", "heads", "kv"))
            value = nn_partitioning.with_sharding_constraint(
                value, ("length", "batch", "heads", "kv")
            )
605
        else:
606
607
608
609
610
611
612
            query = nn_partitioning.with_sharding_constraint(
                query, ("batch", "length", "heads", "kv")
            )
            key = nn_partitioning.with_sharding_constraint(key, ("batch", "length", "heads", "kv"))
            value = nn_partitioning.with_sharding_constraint(
                value, ("batch", "length", "heads", "kv")
            )
613
614
615

        if decode:
            # Detect if we're initializing by absence of existing cache data.
616
            is_initialized = self.has_variable("cache", "cached_key")
617
618
619
620
621
            # The key and value have dimension [batch, length, num_heads, head_dim],
            # but we cache them as [batch, num_heads, head_dim, length] as a TPU
            # fusion optimization. This also enables the "scatter via one-hot
            # broadcast" trick, which means we do a one-hot broadcast instead of a
            # scatter/gather operations, resulting in a 3-4x speedup in practice.
zlsh80826's avatar
zlsh80826 committed
622
            swap_dims = lambda x: x[:-3] + tuple(x[i] for i in [-2, -1, -3])
623
624
625
626
627
628
629
630
631
            cached_key = self.variable(
                "cache", "cached_key", jnp.zeros, swap_dims(key.shape), key.dtype
            )
            cached_value = self.variable(
                "cache", "cached_value", jnp.zeros, swap_dims(value.shape), value.dtype
            )
            cache_index = self.variable(
                "cache", "cache_index", lambda: jnp.array(0, dtype=jnp.int32)
            )
632
633
634
635
636
637
638
639
            if is_initialized:
                batch, num_heads, head_dim, length = cached_key.value.shape
                # During fast autoregressive decoding, we feed one position at a time,
                # and cache the keys and values step by step.
                # Sanity shape check of cached key against input query.
                expected_shape = (batch, 1, num_heads, head_dim)
                if expected_shape != query.shape:
                    raise ValueError(
640
641
642
                        "Autoregressive cache shape error, "
                        f"expected query shape {expected_shape} instead got {query.shape}."
                    )
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669

                # Create a OHE of the current index. NOTE: the index is increased below.
                cur_index = cache_index.value
                one_hot_indices = jax_nn.one_hot(cur_index, length, dtype=key.dtype)
                # In order to update the key, value caches with the current key and
                # value, we move the length axis to the back, similar to what we did for
                # the cached ones above.
                # Note these are currently the key and value of a single position, since
                # we feed one position at a time.
                one_token_key = jnp.moveaxis(key, -3, -1)
                one_token_value = jnp.moveaxis(value, -3, -1)
                # Update key, value caches with our new 1d spatial slices.
                # We implement an efficient scatter into the cache via one-hot
                # broadcast and addition.
                key = cached_key.value + one_token_key * one_hot_indices
                value = cached_value.value + one_token_value * one_hot_indices
                cached_key.value = key
                cached_value.value = value
                cache_index.value = cache_index.value + 1
                # Move the keys and values back to their original shapes.
                key = jnp.moveaxis(key, -1, -3)
                value = jnp.moveaxis(value, -1, -3)

                # Causal mask for cached decoder self-attention: our single query
                # position should only attend to those key positions that have already
                # been generated and cached, not the remaining zero elements.
                mask = combine_masks(
670
                    jnp.logical_not(mask),
671
672
                    jnp.broadcast_to(
                        jnp.arange(length) <= cur_index,
673
674
675
676
677
678
679
                        # (1, 1, length) represent (head dim, query length, key length)
                        # query length is 1 because during decoding we deal with one
                        # index.
                        # The same mask is applied to all batch elements and heads.
                        (batch, 1, 1, length),
                    ),
                )
680
681
682
683
684
685
686

                # Grab the correct relative attention bias during decoding. This is
                # only required during single step decoding.
                if bias is not None:
                    # The bias is a full attention matrix, but during decoding we only
                    # have to take a slice of it.
                    # This is equivalent to bias[..., cur_index:cur_index+1, :].
687
688
689
                    bias = dynamic_vector_slice_in_dim(
                        jnp.squeeze(bias, axis=0), jnp.reshape(cur_index, (-1)), 1, -2
                    )
690
691
692
693

        # Convert the boolean attention mask to an attention bias.
        if mask is not None:
            # attention mask in the form of attention bias
694
695
696
697
698
            attention_bias = lax.select(
                mask > 0,
                jnp.full(mask.shape, 0.0).astype(self.dtype),
                jnp.full(mask.shape, -1e10).astype(self.dtype),
            )
699
700
701
702
703
704
705
706
        else:
            attention_bias = None

        # Add provided bias term (e.g. relative position embedding).
        if bias is not None:
            attention_bias = combine_biases(attention_bias, bias)

        # Apply attention.
707
708
709
710
711
712
713
        x = DotProductAttention(
            transpose_batch_sequence=self.transpose_batch_sequence,
            scale_attn_logits=self.scale_attn_logits,
            dropout_rate=self.dropout_rate,
            dtype=self.dtype,
            float32_logits=self.float32_logits,
        )(query, key, value, bias=attention_bias, deterministic=deterministic)
714
715
716
717

        x = x.reshape((x.shape[0], x.shape[1], x.shape[2] * x.shape[3]))

        if self.transpose_batch_sequence:
718
            x = nn_partitioning.with_sharding_constraint(x, ("length", "batch", "joined_kv"))
719
        else:
720
            x = nn_partitioning.with_sharding_constraint(x, ("batch", "length", "joined_kv"))
721
722
723

        # Back to the original inputs dimensions.
        out = DenseGeneral(
724
            features=inputs_q.shape[-1],  # output dim is set to the input dim.
725
726
            axis=-1,
            kernel_init=self.kernel_init,
727
            kernel_axes=("joined_kv", "embed"),
728
            use_bias=self.use_bias,
729
            bias_axes="embed",
730
            dtype=self.dtype,
731
732
            name="out",
        )(x)
733
734
735
736
737
        return out


class LayerNorm(nn.Module):
    """T5 Layer normalization operating on the last axis of the input data."""
738

739
740
    epsilon: float = 1e-6
    dtype: Any = jnp.float32
741
    layernorm_type: str = "layernorm"
742
743
    zero_centered_gamma: bool = False
    scale_init: Initializer = None
744
745
    bias_init: Initializer = nn.initializers.zeros

746
747
748
749
750
751
752
753
    def __post_init__(self):
        if self.scale_init is None:
            if not self.zero_centered_gamma:
                self.scale_init = nn.initializers.ones
            else:
                self.scale_init = nn.initializers.zeros
        super().__post_init__()

754
755
756
757
758
759
760
    @nn.compact
    def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
        """Applies layer normalization on the input."""

        x = jnp.asarray(x, jnp.float32)
        features = x.shape[-1]

761
762
763
        scale = nn_partitioning.param_with_axes(
            "scale", self.scale_init, (features,), jnp.float32, axes=("embed",)
        )
764
765
        scale = jnp.asarray(scale, self.dtype)

766
        if self.layernorm_type == "layernorm":
767
768
769
770
            mean = jnp.mean(x, axis=-1, keepdims=True)
            var = jnp.mean(jnp.square(x - mean), axis=-1, keepdims=True)
            y = (x - mean) * lax.rsqrt(var + self.epsilon)

771
772
773
            bias = nn_partitioning.param_with_axes(
                "ln_bias", self.bias_init, (features,), jnp.float32, axes=("embed",)
            )
774
775
            bias = jnp.asarray(bias, self.dtype)

776
777
778
            if not self.zero_centered_gamma:
                z = y * scale + bias
            else:
779
                z = y * (scale + 1.0) + bias
780
        else:
781
            assert self.layernorm_type == "rmsnorm"
782
            assert not self.zero_centered_gamma
783
            mean2 = jnp.mean(lax.square(x), axis=-1, keepdims=True)
784
            y = x * lax.rsqrt(mean2 + self.epsilon)
785
786
            z = y * scale

787
        return jnp.asarray(z, self.dtype)
788
789
790
791
792


class RelativePositionBiases(nn.Module):
    """Adds T5-style relative positional embeddings to the attention logits.

793
794
795
796
797
798
799
800
801
802
803
    Attributes:
      num_buckets: Number of buckets to bucket distances between key and query
        positions into.
      max_distance: Maximum distance before everything is lumped into the last
        distance bucket.
      num_heads: Number of heads in the attention layer. Each head will get a
        different relative position weighting.
      dtype: Type of arrays through this module.
      embedding_init: initializer for relative embedding table.
    """

804
805
806
807
808
809
810
    num_buckets: int
    max_distance: int
    num_heads: int
    dtype: Any
    embedding_init: Callable[..., Array] = nn.linear.default_embed_init

    @staticmethod
811
812
813
    def _relative_position_bucket(
        relative_position, bidirectional=True, num_buckets=32, max_distance=128
    ):
814
815
        """Translate relative position to a bucket number for relative attention.

816
817
818
819
820
821
822
823
824
825
        The relative position is defined as memory_position - query_position, i.e.
        the distance in tokens from the attending position to the attended-to
        position.  If bidirectional=False, then positive relative positions are
        invalid.
        We use smaller buckets for small absolute relative_position and larger
        buckets for larger absolute relative_positions.  All relative
        positions >=max_distance  map to the same bucket.  All relative
        positions <=-max_distance map to the same bucket.  This should allow for
        more graceful generalization to longer sequences than the model has been
        trained on.
826

827
828
829
830
831
        Args:
          relative_position: an int32 array
          bidirectional: a boolean - whether the attention is bidirectional
          num_buckets: an integer
          max_distance: an integer
832

833
834
835
836
        Returns:
          a Tensor with the same shape as relative_position, containing int32
            values in the range [0, num_buckets)
        """
837
838
839
840
841
842
843
844
845
846
847
848
        ret = 0
        n = -relative_position
        if bidirectional:
            num_buckets //= 2
            ret += (n < 0).astype(np.int32) * num_buckets
            n = np.abs(n)
        else:
            n = np.maximum(n, 0)
        # now n is in the range [0, inf)
        max_exact = num_buckets // 2
        is_small = n < max_exact
        val_if_large = max_exact + (
849
850
851
852
            np.log(n.astype(np.float32) / max_exact + np.finfo(np.float32).eps)
            / np.log(max_distance / max_exact)
            * (num_buckets - max_exact)
        ).astype(np.int32)
853
854
855
856
857
858
859
860
        val_if_large = np.minimum(val_if_large, num_buckets - 1)
        ret += np.where(is_small, n, val_if_large)
        return ret

    @nn.compact
    def __call__(self, qlen, klen, bidirectional=True):
        """Produce relative position embedding attention biases.

861
862
863
864
865
        Args:
          qlen: attention query length.
          klen: attention key length.
          bidirectional: whether to allow positive memory-query relative position
            embeddings.
866

867
868
869
        Returns:
          output: `(1, len, q_len, k_len)` attention bias
        """
870
871
        context_position = np.arange(qlen, dtype=jnp.int32)[:, None]
        memory_position = np.arange(klen, dtype=jnp.int32)[None, :]
872
873
874
875
876
877
878
        relative_position = memory_position - context_position  # shape (qlen, klen)
        rp_bucket = self._relative_position_bucket(
            relative_position,
            bidirectional=bidirectional,
            num_buckets=self.num_buckets,
            max_distance=self.max_distance,
        )
879
        relative_attention_bias = nn_partitioning.param_with_axes(
880
881
882
            "rel_embedding",
            self.embedding_init,
            (self.num_heads, self.num_buckets),
883
            jnp.float32,
884
885
            axes=("heads", "relpos_buckets"),
        )
886
887
888
889
890
891
892
893
894
895
896
897
898

        relative_attention_bias = jnp.asarray(relative_attention_bias, self.dtype)
        # Instead of using a slow gather, we create a leading-dimension one-hot
        # array from rp_bucket and use it to perform the gather-equivalent via a
        # contraction, i.e.:
        # (num_head, num_buckets) x (num_buckets one-hot, qlen, klen).
        # This is equivalent to relative_attention_bias[:, rp_bucket]
        bcast_iota = lax.broadcasted_iota(jnp.int32, (self.num_buckets, 1, 1), 0)
        rp_bucket_one_hot = jnp.array(rp_bucket[jnp.newaxis, ...] == bcast_iota, dtype=self.dtype)
        # --> shape (qlen, klen, num_heads)
        values = lax.dot_general(
            relative_attention_bias,
            rp_bucket_one_hot,
899
900
            (((1,), (0,)), ((), ())),  # rhs, lhs contracting dims
        )  # no batched dims
901
902
903
904
905
        # Add a singleton batch dimension.
        # --> shape (1, num_heads, qlen, klen)
        return values[jnp.newaxis, ...]


906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
def apply_swa_mask(
    attn_mask_type: str,
    original_mask: Array,
    window_size: Tuple[int, int] = (-1, -1),
) -> Array:
    """Apply the sliding window mask to a given mask"""
    mask_map = {
        "no_mask": AttnMaskType.NO_MASK,
        "padding": AttnMaskType.PADDING_MASK,
        "causal": AttnMaskType.CAUSAL_MASK,
        "padding_causal": AttnMaskType.PADDING_CAUSAL_MASK,
        "causal_bottom_right": AttnMaskType.CAUSAL_BOTTOM_RIGHT_MASK,
        "padding_causal_bottom_right": AttnMaskType.PADDING_CAUSAL_BOTTOM_RIGHT_MASK,
    }
    _attn_mask_type = mask_map.get(attn_mask_type, None)
    assert _attn_mask_type is not None
    max_seqlen_q = original_mask.shape[-2]
    max_seqlen_kv = original_mask.shape[-1]
    swa_mask = make_swa_mask(
        max_seqlen_q, max_seqlen_kv, window_size, _attn_mask_type, dtype=original_mask.dtype
    )
    # In swa_mask and original_mask 0 is masked out
    swa_mask_bcast = jnp.broadcast_to(swa_mask, original_mask.shape)
    new_mask = jnp.where(original_mask == 1, swa_mask_bcast, original_mask)
    return new_mask


933
934
class EncoderLayer(nn.Module):
    """Transformer encoder layer."""
935

936
    enable_relative_embedding: bool = True
937
    relative_embedding: nn.Module = None
zlsh80826's avatar
zlsh80826 committed
938
939
    num_attention_heads: int = 8
    num_gqa_groups: int | None = None
940
    head_dim: int = 64
941
942
943
944
945
    hidden_dropout: float = 0.1
    hidden_dropout_dims: Sequence[int] = ()
    attention_dropout: float = 0.1
    intermediate_dropout: float = 0.1
    intermediate_dropout_dims: Sequence[int] = ()
946
947
948
949
950
    transpose_batch_sequence: bool = True
    float32_attention_logits: bool = False
    scale_attn_logits: bool = False
    scaled_query_init: bool = True
    mlp_dim: int = 2048
951
    mlp_activations: Sequence[str] = ("relu",)
952
    use_bias: bool = False
953
954
    dtype: Any = jnp.float32
    apply_residual_connection_post_layernorm: bool = False
955
    layernorm_type: str = "layernorm"
956
    layernorm_epsilon: float = 1e-6
957
    zero_centered_gamma: bool = False
958
959
    output_layernorm: bool = False
    drop_path: float = 0.0
960
    enable_rotary_pos_emb: bool = False
961
    rotary_pos_emb_group_method: str = "consecutive"
962
    fuse_qkv_params: bool = True
963
964
    fuse_mlp_wi: bool = True
    self_attn_bias_type: Any = None
965
966
    self_attn_mask_type: str = "no_mask"
    window_size: Tuple[int, int] = (-1, -1)
967

zlsh80826's avatar
zlsh80826 committed
968
969
970
971
972
    def __post_init__(self):
        if self.num_gqa_groups is None:
            self.num_gqa_groups = self.num_attention_heads
        super().__post_init__()

973
974
    @nn.compact
    def __call__(self, inputs, encoder_mask=None, deterministic=False):
975
976
977
978
979
980
981
        # Currently cuDNN backend only supports SWA for causal/padding_causal, follow this
        encoder_mask = apply_swa_mask(
            self.self_attn_mask_type,
            encoder_mask,
            self.window_size,
        )

982
983
984
985
        # Relative position embedding as attention biases.
        sequence_dim = 0 if self.transpose_batch_sequence else 1
        batch_dim = 1 - sequence_dim

986
987
        if self.enable_relative_embedding:
            if self.relative_embedding is None:
988
989
990
991
992
993
994
995
                rel_emb = RelativePositionBiases(
                    num_buckets=32,
                    max_distance=128,
                    num_heads=self.num_attention_heads,
                    dtype=self.dtype,
                    embedding_init=nn.initializers.variance_scaling(1.0, "fan_avg", "uniform"),
                    name="relpos_bias",
                )
996
997
998
            else:
                rel_emb = self.relative_embedding
            encoder_bias = rel_emb(inputs.shape[sequence_dim], inputs.shape[sequence_dim], True)
999
        else:
1000
            encoder_bias = None
1001
1002
1003
1004
1005
1006

        # Attention block.
        residual = inputs

        if not self.output_layernorm:
            # Attention block.
1007
1008
1009
1010
1011
1012
1013
            x = LayerNorm(
                layernorm_type=self.layernorm_type,
                epsilon=self.layernorm_epsilon,
                zero_centered_gamma=self.zero_centered_gamma,
                dtype=self.dtype,
                name="pre_attention_layer_norm",
            )(inputs)
1014
1015
1016
1017
1018
1019
1020

            if self.apply_residual_connection_post_layernorm:
                residual = x
        else:
            x = inputs

        # [batch, length, emb_dim] -> [batch, length, emb_dim]
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
        x = MultiHeadAttention(
            num_heads=self.num_attention_heads,
            num_gqa_groups=self.num_gqa_groups,
            dtype=self.dtype,
            head_dim=self.head_dim,
            transpose_batch_sequence=self.transpose_batch_sequence,
            dropout_rate=self.attention_dropout,
            float32_logits=self.float32_attention_logits,
            scale_attn_logits=self.scale_attn_logits,
            scaled_query_init=self.scaled_query_init,
            fuse_qkv=self.fuse_qkv_params,
            enable_rotary_pos_emb=self.enable_rotary_pos_emb,
            rotary_pos_emb_group_method=self.rotary_pos_emb_group_method,
            use_bias=self.use_bias,
            name="attention",
        )(x, x, encoder_mask, encoder_bias, deterministic=deterministic)
        x = nn.Dropout(rate=self.hidden_dropout, broadcast_dims=self.hidden_dropout_dims)(
            x, deterministic=deterministic
        )
1040
1041
        if self.drop_path > 0.0:
            drop_path_shape = _generate_drop_path_shape(x.shape, batch_dim)
1042
1043
1044
            x = nn.Dropout(rate=self.drop_path, broadcast_dims=drop_path_shape)(
                x, deterministic=deterministic
            )
1045
1046
1047
1048
        x = x + residual

        # MLP block.
        residual = x
1049
1050
1051
1052
1053
1054
1055
        y = LayerNorm(
            layernorm_type=self.layernorm_type,
            epsilon=self.layernorm_epsilon,
            zero_centered_gamma=self.zero_centered_gamma,
            dtype=self.dtype,
            name="pre_mlp_layer_norm",
        )(x)
1056
1057
1058
1059
1060
1061
1062
1063
1064

        if self.apply_residual_connection_post_layernorm:
            residual = y

        # [batch, length, emb_dim] -> [batch, length, emb_dim]
        y = MlpBlock(
            transpose_batch_sequence=self.transpose_batch_sequence,
            intermediate_dim=self.mlp_dim,
            activations=self.mlp_activations,
1065
1066
1067
            intermediate_dropout_rate=self.intermediate_dropout,
            intermediate_dropout_dims=self.intermediate_dropout_dims,
            use_bias=self.use_bias,
1068
1069
            dtype=self.dtype,
            fuse_wi=self.fuse_mlp_wi,
1070
            name="mlp",
1071
        )(y, deterministic=deterministic)
1072
1073
1074
        y = nn.Dropout(rate=self.hidden_dropout, broadcast_dims=self.hidden_dropout_dims)(
            y, deterministic=deterministic
        )
1075
1076
        if self.drop_path > 0.0:
            drop_path_shape = _generate_drop_path_shape(y.shape, batch_dim)
1077
1078
1079
            y = nn.Dropout(rate=self.drop_path, broadcast_dims=drop_path_shape)(
                y, deterministic=deterministic
            )
1080
1081
1082
        y = y + residual

        if self.output_layernorm:
1083
1084
1085
1086
1087
1088
1089
            y = LayerNorm(
                layernorm_type=self.layernorm_type,
                epsilon=self.layernorm_epsilon,
                zero_centered_gamma=self.zero_centered_gamma,
                dtype=self.dtype,
                name="output_layernorm",
            )(y)
1090
1091
1092
1093
1094
        return y


class DecoderLayer(nn.Module):
    """Transformer decoder layer that attends to the encoder."""
1095

1096
    enable_relative_embedding: bool = True
1097
    relative_embedding: nn.Module = None
zlsh80826's avatar
zlsh80826 committed
1098
1099
    num_attention_heads: int = 8
    num_gqa_groups: int | None = None
1100
    head_dim: int = 64
1101
1102
1103
1104
1105
    hidden_dropout: float = 0.1
    hidden_dropout_dims: Sequence[int] = ()
    attention_dropout: float = 0.1
    intermediate_dropout: float = 0.1
    intermediate_dropout_dims: Sequence[int] = ()
1106
1107
1108
1109
1110
    transpose_batch_sequence: bool = True
    float32_attention_logits: bool = False
    scale_attn_logits: bool = False
    scaled_query_init: bool = True
    mlp_dim: int = 2048
1111
    mlp_activations: Sequence[str] = ("relu",)
1112
    use_bias: bool = False
1113
1114
1115
    dtype: Any = jnp.float32
    apply_residual_connection_post_layernorm: bool = False
    output_layernorm: bool = False
1116
    layernorm_type: str = "layernorm"
1117
    layernorm_epsilon: float = 1e-6
1118
    zero_centered_gamma: bool = False
1119
    drop_path: float = 0.0
1120
    enable_rotary_pos_emb: bool = False
1121
    rotary_pos_emb_group_method: str = "consecutive"
1122
    fuse_qkv_params: bool = True
1123
1124
    fuse_mlp_wi: bool = True
    self_attn_bias_type: Any = None
1125
1126
    self_attn_mask_type: str = "no_mask"
    window_size: Tuple[int, int] = (-1, -1)
1127

zlsh80826's avatar
zlsh80826 committed
1128
1129
1130
1131
1132
    def __post_init__(self):
        if self.num_gqa_groups is None:
            self.num_gqa_groups = self.num_attention_heads
        super().__post_init__()

1133
    @nn.compact
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
    def __call__(
        self,
        inputs,
        encoded,
        decoder_mask=None,
        encoder_decoder_mask=None,
        deterministic=False,
        decode=False,
        max_decode_length=None,
    ):
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
        decoder_mask = apply_swa_mask(
            self.self_attn_mask_type,
            decoder_mask,
            self.window_size,
        )

        encoder_decoder_mask = apply_swa_mask(
            "padding",
            encoder_decoder_mask,
            self.window_size,
        )

1156
1157
1158
        # Relative position embedding as attention biases.
        sequence_dim = 0 if self.transpose_batch_sequence else 1
        batch_dim = 1 - sequence_dim
1159
1160
1161
1162

        if self.enable_relative_embedding:
            l = max_decode_length if decode and max_decode_length else inputs.shape[sequence_dim]
            if self.relative_embedding is None:
1163
1164
1165
1166
1167
1168
1169
1170
                rel_emb = RelativePositionBiases(
                    num_buckets=32,
                    max_distance=128,
                    num_heads=self.num_attention_heads,
                    dtype=self.dtype,
                    embedding_init=nn.initializers.variance_scaling(1.0, "fan_avg", "uniform"),
                    name="relpos_bias",
                )
1171
1172
1173
            else:
                rel_emb = self.relative_embedding
            decoder_bias = rel_emb(l, l, False)
1174
        else:
1175
            decoder_bias = None
1176
1177
1178
1179
1180
1181

        # inputs: embedded inputs to the decoder with shape [batch, length, emb_dim]
        residual = inputs

        if not self.output_layernorm:
            # Attention block.
1182
1183
1184
1185
1186
1187
1188
            x = LayerNorm(
                layernorm_type=self.layernorm_type,
                epsilon=self.layernorm_epsilon,
                zero_centered_gamma=self.zero_centered_gamma,
                dtype=self.dtype,
                name="pre_self_attention_layer_norm",
            )(inputs)
1189
1190
1191
1192
1193
1194
1195

            if self.apply_residual_connection_post_layernorm:
                residual = x
        else:
            x = inputs

        # Self-attention block
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
        x = MultiHeadAttention(
            num_heads=self.num_attention_heads,
            num_gqa_groups=self.num_gqa_groups,
            dtype=self.dtype,
            head_dim=self.head_dim,
            transpose_batch_sequence=self.transpose_batch_sequence,
            dropout_rate=self.attention_dropout,
            float32_logits=self.float32_attention_logits,
            scale_attn_logits=self.scale_attn_logits,
            scaled_query_init=self.scaled_query_init,
            enable_rotary_pos_emb=self.enable_rotary_pos_emb,
            rotary_pos_emb_group_method=self.rotary_pos_emb_group_method,
            fuse_qkv=self.fuse_qkv_params,
            use_bias=self.use_bias,
            name="self_attention",
        )(x, x, decoder_mask, decoder_bias, deterministic=deterministic, decode=decode)
        x = nn.Dropout(rate=self.hidden_dropout, broadcast_dims=self.hidden_dropout_dims)(
            x, deterministic=deterministic
        )
1215
1216
        if self.drop_path > 0.0:
            drop_path_shape = _generate_drop_path_shape(x.shape, batch_dim)
1217
1218
1219
            x = nn.Dropout(rate=self.drop_path, broadcast_dims=drop_path_shape)(
                x, deterministic=deterministic
            )
1220
1221
1222
1223
        x = x + residual

        # Encoder-Decoder block.
        residual = x
1224
1225
1226
1227
1228
1229
1230
        y = LayerNorm(
            layernorm_type=self.layernorm_type,
            epsilon=self.layernorm_epsilon,
            zero_centered_gamma=self.zero_centered_gamma,
            dtype=self.dtype,
            name="pre_cross_attention_layer_norm",
        )(x)
1231
1232
1233

        if self.apply_residual_connection_post_layernorm:
            residual = y
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
        y = MultiHeadAttention(
            num_heads=self.num_attention_heads,
            num_gqa_groups=self.num_gqa_groups,
            dtype=self.dtype,
            head_dim=self.head_dim,
            transpose_batch_sequence=self.transpose_batch_sequence,
            dropout_rate=self.attention_dropout,
            float32_logits=self.float32_attention_logits,
            scale_attn_logits=self.scale_attn_logits,
            scaled_query_init=self.scaled_query_init,
            enable_rotary_pos_emb=self.enable_rotary_pos_emb,
            rotary_pos_emb_group_method=self.rotary_pos_emb_group_method,
            fuse_qkv=self.fuse_qkv_params,
            use_bias=self.use_bias,
            name="encoder_decoder_attention",
        )(y, encoded, encoder_decoder_mask, deterministic=deterministic)
        y = nn.Dropout(rate=self.hidden_dropout, broadcast_dims=self.hidden_dropout_dims)(
            y, deterministic=deterministic
        )
1253
1254
1255
1256
        y = y + residual

        # MLP block.
        residual = y
1257
1258
1259
1260
1261
1262
1263
        z = LayerNorm(
            layernorm_type=self.layernorm_type,
            epsilon=self.layernorm_epsilon,
            zero_centered_gamma=self.zero_centered_gamma,
            dtype=self.dtype,
            name="pre_mlp_layer_norm",
        )(y)
1264
1265
1266
1267
1268
1269
        if self.apply_residual_connection_post_layernorm:
            residual = z
        z = MlpBlock(
            transpose_batch_sequence=self.transpose_batch_sequence,
            intermediate_dim=self.mlp_dim,
            activations=self.mlp_activations,
1270
1271
1272
            intermediate_dropout_rate=self.intermediate_dropout,
            intermediate_dropout_dims=self.intermediate_dropout_dims,
            use_bias=self.use_bias,
1273
1274
            dtype=self.dtype,
            fuse_wi=self.fuse_mlp_wi,
1275
            name="mlp",
1276
        )(z, deterministic=deterministic)
1277
1278
1279
        z = nn.Dropout(rate=self.hidden_dropout, broadcast_dims=self.hidden_dropout_dims)(
            z, deterministic=deterministic
        )
1280
1281
        if self.drop_path > 0.0:
            drop_path_shape = _generate_drop_path_shape(z.shape, batch_dim)
1282
1283
1284
            z = nn.Dropout(rate=self.drop_path, broadcast_dims=drop_path_shape)(
                z, deterministic=deterministic
            )
1285
1286
1287
        z = z + residual

        if self.output_layernorm:
1288
1289
1290
1291
1292
1293
1294
            z = LayerNorm(
                layernorm_type=self.layernorm_type,
                epsilon=self.layernorm_epsilon,
                zero_centered_gamma=self.zero_centered_gamma,
                dtype=self.dtype,
                name="output_layernorm",
            )(z)
1295
1296
1297
1298

        return z


1299
def make_causal_mask(batch, seqlen, dtype=jnp.uint8):
zlsh80826's avatar
zlsh80826 committed
1300
1301
1302
    """
    Generate causal mask
    """
1303
1304
1305
1306
1307
1308
1309
1310
1311
1312
    shape = (batch, seqlen)
    idxs = jnp.broadcast_to(jnp.arange(shape[-1], dtype=jnp.int32), shape)

    mask = jnp.greater_equal(jnp.expand_dims(idxs, axis=-1), jnp.expand_dims(idxs, axis=-2))
    mask = jnp.expand_dims(mask, axis=-3)
    mask = 1 - mask
    return mask.astype(dtype)


def make_self_mask(batch, seqlen, dtype=jnp.uint8):
zlsh80826's avatar
zlsh80826 committed
1313
1314
1315
    """
    Generate attention mask
    """
1316
1317
1318
1319
1320
1321
1322
    shape = (batch, seqlen)
    mask = jnp.ones((*shape, shape[-1]))
    mask = jnp.expand_dims(mask, axis=-3)
    mask = 1 - mask
    return mask.astype(dtype)


1323
1324
1325
1326
1327
1328
1329
1330
1331
1332
1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
1343
1344
1345
1346
1347
1348
1349
1350
def assert_allclose(
    actual: Array,
    desired: Array,
    rtol: Optional[float] = None,
    atol: Optional[float] = None,
    dtype: Optional[Union[DType, TEDType, np.dtype, str]] = None,
    **kwargs,
) -> None:
    """Check if two tensors are close.

    Args:
      actual: test tensor.
      desired: reference tensor.
      dtype: data type or data type name (default: inferred from
        `actual`).
      rtol: relative tolerance (default: based on `dtype`).
      atol: absolute tolerance (default: based on `dtype`).
      **kwargs: keyword arguments to pass to np.testing.assert_allclose.
    """

    # Infer data type if needed
    if dtype is None:
        if isinstance(actual, float):
            dtype = "float32"
        else:
            dtype = actual.dtype

    # Determine tolerances
zlsh80826's avatar
zlsh80826 committed
1351
    tols = {}
1352
1353
1354
1355
1356
1357
1358
1359
    if rtol is None or atol is None:
        tols = dtype_tols(dtype)
    if rtol is not None:
        tols["rtol"] = rtol
    if atol is not None:
        tols["atol"] = atol

    # Cast tensors to fp32
1360
1361
1362
1363
    if not isinstance(actual, float):
        actual = actual.astype(jnp.float32)
    if not isinstance(desired, float):
        desired = desired.astype(jnp.float32)
1364
1365
1366
1367
1368

    # Check if tensors are close
    np.testing.assert_allclose(actual, desired, **tols, **kwargs)


1369
1370
1371
1372
def assert_tree_like_allclose(expected, actual, rtol=1e-05, atol=1e-08):
    flatten_expected, _ = jax.tree_util.tree_flatten_with_path(expected)
    flatten_actual, _ = jax.tree_util.tree_flatten_with_path(actual)

1373
1374
1375
    for (expected_path, expected_value), (actual_path, actual_value) in zip(
        flatten_expected, flatten_actual
    ):
1376
1377
        assert expected_path == actual_path
        key_str = jax.tree_util.keystr(expected_path)
1378
1379
1380
1381
1382
1383
1384
        assert_allclose(
            expected_value,
            actual_value,
            rtol=rtol,
            atol=atol,
            err_msg=f"Value of expected{key_str} and actual{key_str} is not close",
        )
1385
1386


1387
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407
1408
1409
1410
1411
1412
1413
1414
1415
1416
1417
1418
1419
1420
1421
1422
def dtype_tols(
    dtype: Union[DType, TEDType, np.dtype],
    reference_value: float = 1.0,
) -> Dict[str, float]:
    """Expected numerical tolerance for a data type.

    Args:
      dtype: data type.
      reference_value: reference value (default: 1).

    Returns:
      Dictionary with "rtol" and "atol" as keys

    """

    # Convert to JAX dtype if needed
    if isinstance(dtype, TEDType):
        dtype = {
            TEDType.kByte: jnp.uint8,
            TEDType.kInt32: jnp.int32,
            TEDType.kInt64: jnp.int64,
            TEDType.kFloat32: jnp.float32,
            TEDType.kFloat16: jnp.float16,
            TEDType.kBFloat16: jnp.bfloat16,
            TEDType.kFloat8E4M3: jnp.float8_e4m3fn,
            TEDType.kFloat8E5M2: jnp.float8_e5m2,
        }[dtype]
    elif isinstance(dtype, np.dtype):
        dtype = jnp.dtype(dtype)

    # Expect bit-wise accuracy for integer dtypes
    if not jnp.issubdtype(dtype, jnp.floating):
        return dict(rtol=0, atol=0)

    # Estimate floating-point error
    finfo = jnp.finfo(dtype)
1423
    eps_relaxed = math.pow(finfo.eps, 2 / 3)
1424
1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
1435
    with jax.default_device(jax.devices("cpu")[0]):
        if isinstance(reference_value, (float, int)):
            reference_value = jnp.array(reference_value, dtype=dtype)
        else:
            reference_value = reference_value.astype(dtype)
        spacing_high = jnp.nextafter(reference_value, finfo.max) - reference_value
        spacing_low = reference_value - jnp.nextafter(reference_value, finfo.min)
        ulp = max(spacing_high.item(), spacing_low.item())
    return dict(
        rtol=eps_relaxed,
        atol=max(ulp, eps_relaxed),
    )
1436
1437


1438
def sync_params_values(dst, src, transformations, sep="/"):
1439
1440
1441
1442
1443
1444
1445
1446
1447
1448
1449
1450
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
1462
1463
1464
1465
1466
1467
1468
    """
    This function will reconstuct a tree with dst's tree_def/shape and src's value.
    transformations is a map that records the key mappings between dst and src.
    If no dst key found in the transformerations, it will fall back to src key = dst key.
    transformations = {
        dst key map 0: src key map 0,
        dst key map 1: src key map 1,
        ...
        # if dst key = src key, we don't need to add it
    }
    """
    src_values = {}
    for key, value in jax.tree_util.tree_leaves_with_path(src):
        normalized_key = sep.join(x.key for x in key)
        src_values[normalized_key] = value

    flatten_dst, dst_tree_def = jax.tree_util.tree_flatten_with_path(dst)
    synced_dst_values = []

    for key, value in flatten_dst:
        normalized_key = sep.join(x.key for x in key)
        if normalized_key in transformations:
            corresponding_src_key = transformations[normalized_key]
        else:
            corresponding_src_key = normalized_key
        synced_dst_values.append(src_values[corresponding_src_key])

    synced_dst = jax.tree_util.tree_unflatten(dst_tree_def, synced_dst_values)

    return jax.tree_util.tree_map(lambda x, y: x.reshape(y.shape), synced_dst, dst)