Fix Bugs of TE/JAX

tests/jax/test_sharding.py

@@ -38,9 +38,13 @@ def _get_sharding_resource(mesh_names, sharding_type):
  ((4,), ("tp",), ShardingType.TP_ROW), ((2, 2), ("dp", "tp"), ShardingType.DP_TP_COL),
  ((2, 2), ("dp", "tp"), ShardingType.DP_TP_ROW)]
 
-LOGICAL_RULES = [[(('a1', None), ('a2', 'ma2')), False],
- [(('a1', None), ('a2', 'ma2'), ('a3', ('ma31', 'ma32'))), True],
- [(('a1', None), ('a2', 'ma2'), ('batch', 'batch_1200234')), True]]
+LOGICAL_RULES = [
+ [(('a1', None), ('a2', 'ma2')), False],
+ [(('a1', None), ('a2', 'ma2'), ('a3', ('ma31', 'ma32'))), True],
+ [(('a1', None), ('a2', 'ma2'), ('a3', 'ma31'), ('a3', 'ma32')), False],
+ [(('a1', None), ('a2', 'ma2'), ('batch', 'batch_1200234')), True],
+ [(('a1', None), ('a2', 'ma2'), ('a2', 'ma1'), ('batch', 'model'), ('batch', 'data')), True],
+]
 SRS = [
  ShardingResource(),
  ShardingResource('data', None),

tests/jax/utils.py

@@ -321,8 +321,9 @@ def __call__(self, inputs, deterministic: bool = False):
 
  # Take elementwise product of above intermediate activations.
  x = functools.reduce(operator.mul, activations)
+ dropout_broadcast_dims = (0,) if self.transpose_batch_sequence else (1,)
  # Apply dropout and final dense output projection.
- x = nn.Dropout(rate=self.intermediate_dropout_rate, broadcast_dims=(-2,))(
+ x = nn.Dropout(rate=self.intermediate_dropout_rate, broadcast_dims=dropout_broadcast_dims)(
  x, deterministic=deterministic) # Broadcast along length.
  if self.transpose_batch_sequence:
  x = nn_partitioning.with_sharding_constraint(x, ('length', 'batch', 'mlp'))

transformer_engine/jax/fp8.py

@@ -190,7 +190,7 @@ def update_amax_history(amax_buffers: jnp.ndarray) -> jnp.ndarray:
  Update the amax history
  """
  updated_amax_buffers = jnp.roll(amax_buffers, -1, 1)
- updated_amax_buffers.at[:, 0].set(0)
+ updated_amax_buffers = updated_amax_buffers.at[:, 0].set(0)
  return updated_amax_buffers
 
  @staticmethod

transformer_engine/jax/module.py

@@ -683,6 +683,8 @@ class LayerNormMLP(TransformerEngineBase):
  Each activation has its own transformation layer.
  intermediate_dropout_rate: float, default = 0.1
  Dropout probability for the dropout op after the :attr:`activations`.
+ intermediate_hidden_dropout_dims: Sequence[int], default = ()
+ Dimensions that will share the same dropout mask for hidden
  axis: Union[Iterable[int], int], default = -1
  An integer tuple with axes to apply the transformation on.
 
@@ -716,6 +718,7 @@ class LayerNormMLP(TransformerEngineBase):
  return_layernorm_output: bool = True
  activations: Sequence[Union[str, Callable]] = ('relu',)
  intermediate_dropout_rate: float = 0.1
+ intermediate_hidden_dropout_dims: Sequence[int] = ()
  axis: Union[Iterable[int], int] = -1
  dtype: DType = jnp.float32
  transpose_batch_sequence: bool = True
@@ -912,8 +915,9 @@ def fp8_meta_generator():
  z = functools.reduce(operator.mul, activations)
  z = jnp.reshape(z, (*z.shape[:-2], -1))
 
- z = nn.Dropout(rate=self.intermediate_dropout_rate, broadcast_dims=(-2,))(
- z, deterministic=deterministic) # Broadcast along length.
+ z = nn.Dropout(rate=self.intermediate_dropout_rate,
+ broadcast_dims=self.intermediate_hidden_dropout_dims)(
+ z, deterministic=deterministic)
 
  # DenseGeneral 2
  hidden_size = inputs.shape[-1]

transformer_engine/jax/transformer.py

@@ -53,6 +53,10 @@ def extend_logical_axis_rules(rules: LogicalRules) -> LogicalRules:
  .. warning::
  Please make sure ShardingResource is set via fp8_autocast before calling this function.
 
+ .. note::
+ This function is only needed when using TransformerLayer. For other modules, such as
+ DenseGeneral, please properly set axes of kernels and bias.
+
  Parameters
  ----------
  rules : Sequence[Tuple[str, Union[str, None]]]
@@ -73,10 +77,12 @@ def extend_logical_axis_rules(rules: LogicalRules) -> LogicalRules:
  f"Thie axis_name should be str, but got {type(key)}."
  assert isinstance(val, str) or (val is None), \
  f"Thie mesh_axis_name should be str or None, but got {type(val)}."
- rules_map[key] = val
+ if key in rules_map:
+ rules_map[key].append(val)
+ else:
+ rules_map[key] = [val]
 
  gsr = global_shard_resource()
-
  te_logical_axis_rules = (('batch', gsr.dp_resource), ('embed', None), ('mlp', gsr.tp_resource),
  ('heads', gsr.tp_resource), ('kv', None), ('qkv_dim', None),
  ('kv_dim', None), ('joined_kv', gsr.tp_resource), ('act', None),
@@ -87,7 +93,7 @@ def extend_logical_axis_rules(rules: LogicalRules) -> LogicalRules:
  key = item[0]
  val = item[1]
  if key in rules_map:
- assert rules_map[key] == val, \
+ assert len(rules_map[key]) == 1 and rules_map[key][0] == val, \
  f"The rule diverged between TE and given rule." \
  f"Axis:{key} map to {rules_map[key]} in the given" \
  f" rules, but {val} in TE's rules."
@@ -447,41 +453,38 @@ def kv_init(key, shape, dtype):
  if decode:
  is_initialized = self.has_variable('cache', 'cached_key')
 
- # TODO (Ming Huang): Check performance on GPU withou swap dimensions # pylint: disable=fixme
- def swap_dims(x):
- return x[:-3] + tuple(x[i] for i in [-2, -1, -3])
-
- 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),
+ cached_key = self.variable('cache', 'cached_key', jnp.zeros, key.shape, key.dtype)
+ cached_value = self.variable('cache', 'cached_value', jnp.zeros, value.shape,
  value.dtype)
  cache_index = self.variable('cache', 'cache_index',
  lambda: jnp.array(0, dtype=jnp.int32))
  if is_initialized:
- batch, num_heads, head_dim, length = cached_key.value.shape
+ if self.transpose_batch_sequence:
+ length, batch, num_heads, head_dim = cached_key.value.shape
+ expected_shape = (1, batch, num_heads, head_dim)
+ one_hot_indices_shape = (length, 1, 1, 1)
+ else:
+ batch, length, num_heads, head_dim = cached_key.value.shape
+ expected_shape = (batch, 1, num_heads, head_dim)
+ one_hot_indices_shape = (1, length, 1, 1)
 
  # Sanity shape check of cached key against input query.
- expected_shape = (batch, 1, num_heads, head_dim)
  if expected_shape != query.shape:
  raise ValueError(
  'Autoregressive cache shape error, '
  f"expected query shape {expected_shape} instead got {query.shape}.")
 
  cur_index = cache_index.value
  one_hot_indices = jax_nn.one_hot(cur_index, length, dtype=key.dtype)
- one_token_key = jnp.moveaxis(key, -3, -1)
- one_token_value = jnp.moveaxis(value, -3, -1)
- key = cached_key.value + one_token_key * one_hot_indices
- value = cached_value.value + one_token_value * one_hot_indices
+ one_hot_indices = jnp.reshape(one_hot_indices, one_hot_indices_shape)
+ key = cached_key.value + key * one_hot_indices
+ value = cached_value.value + value * one_hot_indices
  cached_key.value = key
  cached_value.value = value
  cache_index.value = cache_index.value + 1
 
- key = jnp.moveaxis(key, -1, -3)
- value = jnp.moveaxis(value, -1, -3)
-
  mask = combine_masks(
- mask, jnp.broadcast_to(jnp.arange(length) <= cur_index, (batch, 1, 1, length)))
+ mask, jnp.broadcast_to(jnp.arange(length) > cur_index, (batch, 1, 1, length)))
 
  if bias is not None:
  bias = dynamic_vector_slice_in_dim(jnp.squeeze(bias, axis=0),
@@ -889,10 +892,11 @@ def hidden_dropout(x, deterministic):
  assert isinstance(self.hidden_dropout_dims, Sequence)
  x_shape_len = len(x.shape)
  for dims in self.hidden_dropout_dims:
- assert -x_shape_len < dims < x_shape_len
+ assert -x_shape_len <= dims < x_shape_len
 
  return nn.Dropout(rate=self.hidden_dropout,
- broadcast_dims=self.hidden_dropout_dims)(x, deterministic)
+ broadcast_dims=self.hidden_dropout_dims)(x,
+ deterministic=deterministic)
 
  x = hidden_dropout(x, deterministic)
  if self.drop_path > 0.0:
@@ -944,6 +948,7 @@ def hidden_dropout(x, deterministic):
  intermediate_dim=self.mlp_hidden_size,
  activations=self.mlp_activations,
  intermediate_dropout_rate=self.hidden_dropout,
+ intermediate_hidden_dropout_dims=self.hidden_dropout_dims,
  dtype=self.dtype,
  scale_axes=('embed',),
  kernel_init=self.mlp_kernel_init,

tests/jax/test_sharding.py

@@ -38,9 +38,13 @@ def _get_sharding_resource(mesh_names, sharding_type):
  ((4,), ("tp",), ShardingType.TP_ROW), ((2, 2), ("dp", "tp"), ShardingType.DP_TP_COL),
  ((2, 2), ("dp", "tp"), ShardingType.DP_TP_ROW)]
 
-LOGICAL_RULES = [[(('a1', None), ('a2', 'ma2')), False],
- [(('a1', None), ('a2', 'ma2'), ('a3', ('ma31', 'ma32'))), True],
- [(('a1', None), ('a2', 'ma2'), ('batch', 'batch_1200234')), True]]
+LOGICAL_RULES = [
+ [(('a1', None), ('a2', 'ma2')), False],
+ [(('a1', None), ('a2', 'ma2'), ('a3', ('ma31', 'ma32'))), True],
+ [(('a1', None), ('a2', 'ma2'), ('a3', 'ma31'), ('a3', 'ma32')), False],
+ [(('a1', None), ('a2', 'ma2'), ('batch', 'batch_1200234')), True],
+ [(('a1', None), ('a2', 'ma2'), ('a2', 'ma1'), ('batch', 'model'), ('batch', 'data')), True],
+]
 SRS = [
  ShardingResource(),
  ShardingResource('data', None),

tests/jax/utils.py

@@ -321,8 +321,9 @@ def __call__(self, inputs, deterministic: bool = False):
 
  # Take elementwise product of above intermediate activations.
  x = functools.reduce(operator.mul, activations)
+ dropout_broadcast_dims = (0,) if self.transpose_batch_sequence else (1,)
  # Apply dropout and final dense output projection.
- x = nn.Dropout(rate=self.intermediate_dropout_rate, broadcast_dims=(-2,))(
+ x = nn.Dropout(rate=self.intermediate_dropout_rate, broadcast_dims=dropout_broadcast_dims)(
  x, deterministic=deterministic) # Broadcast along length.
  if self.transpose_batch_sequence:
  x = nn_partitioning.with_sharding_constraint(x, ('length', 'batch', 'mlp'))

transformer_engine/jax/fp8.py

@@ -190,7 +190,7 @@ def update_amax_history(amax_buffers: jnp.ndarray) -> jnp.ndarray:
  Update the amax history
  """
  updated_amax_buffers = jnp.roll(amax_buffers, -1, 1)
- updated_amax_buffers.at[:, 0].set(0)
+ updated_amax_buffers = updated_amax_buffers.at[:, 0].set(0)
  return updated_amax_buffers
 
  @staticmethod

transformer_engine/jax/module.py

@@ -683,6 +683,8 @@ class LayerNormMLP(TransformerEngineBase):
  Each activation has its own transformation layer.
  intermediate_dropout_rate: float, default = 0.1
  Dropout probability for the dropout op after the :attr:`activations`.
+ intermediate_hidden_dropout_dims: Sequence[int], default = ()
+ Dimensions that will share the same dropout mask for hidden
  axis: Union[Iterable[int], int], default = -1
  An integer tuple with axes to apply the transformation on.
 
@@ -716,6 +718,7 @@ class LayerNormMLP(TransformerEngineBase):
  return_layernorm_output: bool = True
  activations: Sequence[Union[str, Callable]] = ('relu',)
  intermediate_dropout_rate: float = 0.1
+ intermediate_hidden_dropout_dims: Sequence[int] = ()
  axis: Union[Iterable[int], int] = -1
  dtype: DType = jnp.float32
  transpose_batch_sequence: bool = True
@@ -912,8 +915,9 @@ def fp8_meta_generator():
  z = functools.reduce(operator.mul, activations)
  z = jnp.reshape(z, (*z.shape[:-2], -1))
 
- z = nn.Dropout(rate=self.intermediate_dropout_rate, broadcast_dims=(-2,))(
- z, deterministic=deterministic) # Broadcast along length.
+ z = nn.Dropout(rate=self.intermediate_dropout_rate,
+ broadcast_dims=self.intermediate_hidden_dropout_dims)(
+ z, deterministic=deterministic)
 
  # DenseGeneral 2
  hidden_size = inputs.shape[-1]

transformer_engine/jax/transformer.py

@@ -53,6 +53,10 @@ def extend_logical_axis_rules(rules: LogicalRules) -> LogicalRules:
  .. warning::
  Please make sure ShardingResource is set via fp8_autocast before calling this function.
 
+ .. note::
+ This function is only needed when using TransformerLayer. For other modules, such as
+ DenseGeneral, please properly set axes of kernels and bias.
+
  Parameters
  ----------
  rules : Sequence[Tuple[str, Union[str, None]]]
@@ -73,10 +77,12 @@ def extend_logical_axis_rules(rules: LogicalRules) -> LogicalRules:
  f"Thie axis_name should be str, but got {type(key)}."
  assert isinstance(val, str) or (val is None), \
  f"Thie mesh_axis_name should be str or None, but got {type(val)}."
- rules_map[key] = val
+ if key in rules_map:
+ rules_map[key].append(val)
+ else:
+ rules_map[key] = [val]
 
  gsr = global_shard_resource()
-
  te_logical_axis_rules = (('batch', gsr.dp_resource), ('embed', None), ('mlp', gsr.tp_resource),
  ('heads', gsr.tp_resource), ('kv', None), ('qkv_dim', None),
  ('kv_dim', None), ('joined_kv', gsr.tp_resource), ('act', None),
@@ -87,7 +93,7 @@ def extend_logical_axis_rules(rules: LogicalRules) -> LogicalRules:
  key = item[0]
  val = item[1]
  if key in rules_map:
- assert rules_map[key] == val, \
+ assert len(rules_map[key]) == 1 and rules_map[key][0] == val, \
  f"The rule diverged between TE and given rule." \
  f"Axis:{key} map to {rules_map[key]} in the given" \
  f" rules, but {val} in TE's rules."
@@ -447,41 +453,38 @@ def kv_init(key, shape, dtype):
  if decode:
  is_initialized = self.has_variable('cache', 'cached_key')
 
- # TODO (Ming Huang): Check performance on GPU withou swap dimensions # pylint: disable=fixme
- def swap_dims(x):
- return x[:-3] + tuple(x[i] for i in [-2, -1, -3])
-
- 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),
+ cached_key = self.variable('cache', 'cached_key', jnp.zeros, key.shape, key.dtype)
+ cached_value = self.variable('cache', 'cached_value', jnp.zeros, value.shape,
  value.dtype)
  cache_index = self.variable('cache', 'cache_index',
  lambda: jnp.array(0, dtype=jnp.int32))
  if is_initialized:
- batch, num_heads, head_dim, length = cached_key.value.shape
+ if self.transpose_batch_sequence:
+ length, batch, num_heads, head_dim = cached_key.value.shape
+ expected_shape = (1, batch, num_heads, head_dim)
+ one_hot_indices_shape = (length, 1, 1, 1)
+ else:
+ batch, length, num_heads, head_dim = cached_key.value.shape
+ expected_shape = (batch, 1, num_heads, head_dim)
+ one_hot_indices_shape = (1, length, 1, 1)
 
  # Sanity shape check of cached key against input query.
- expected_shape = (batch, 1, num_heads, head_dim)
  if expected_shape != query.shape:
  raise ValueError(
  'Autoregressive cache shape error, '
  f"expected query shape {expected_shape} instead got {query.shape}.")
 
  cur_index = cache_index.value
  one_hot_indices = jax_nn.one_hot(cur_index, length, dtype=key.dtype)
- one_token_key = jnp.moveaxis(key, -3, -1)
- one_token_value = jnp.moveaxis(value, -3, -1)
- key = cached_key.value + one_token_key * one_hot_indices
- value = cached_value.value + one_token_value * one_hot_indices
+ one_hot_indices = jnp.reshape(one_hot_indices, one_hot_indices_shape)
+ key = cached_key.value + key * one_hot_indices
+ value = cached_value.value + value * one_hot_indices
  cached_key.value = key
  cached_value.value = value
  cache_index.value = cache_index.value + 1
 
- key = jnp.moveaxis(key, -1, -3)
- value = jnp.moveaxis(value, -1, -3)
-
  mask = combine_masks(
- mask, jnp.broadcast_to(jnp.arange(length) <= cur_index, (batch, 1, 1, length)))
+ mask, jnp.broadcast_to(jnp.arange(length) > cur_index, (batch, 1, 1, length)))
 
  if bias is not None:
  bias = dynamic_vector_slice_in_dim(jnp.squeeze(bias, axis=0),
@@ -889,10 +892,11 @@ def hidden_dropout(x, deterministic):
  assert isinstance(self.hidden_dropout_dims, Sequence)
  x_shape_len = len(x.shape)
  for dims in self.hidden_dropout_dims:
- assert -x_shape_len < dims < x_shape_len
+ assert -x_shape_len <= dims < x_shape_len
 
  return nn.Dropout(rate=self.hidden_dropout,
- broadcast_dims=self.hidden_dropout_dims)(x, deterministic)
+ broadcast_dims=self.hidden_dropout_dims)(x,
+ deterministic=deterministic)
 
  x = hidden_dropout(x, deterministic)
  if self.drop_path > 0.0:
@@ -944,6 +948,7 @@ def hidden_dropout(x, deterministic):
  intermediate_dim=self.mlp_hidden_size,
  activations=self.mlp_activations,
  intermediate_dropout_rate=self.hidden_dropout,
+ intermediate_hidden_dropout_dims=self.hidden_dropout_dims,
  dtype=self.dtype,
  scale_axes=('embed',),
  kernel_init=self.mlp_kernel_init,

tests/jax/test_sharding.py

@@ -38,9 +38,13 @@ def _get_sharding_resource(mesh_names, sharding_type):
  ((4,), ("tp",), ShardingType.TP_ROW), ((2, 2), ("dp", "tp"), ShardingType.DP_TP_COL),
  ((2, 2), ("dp", "tp"), ShardingType.DP_TP_ROW)]
 
-LOGICAL_RULES = [[(('a1', None), ('a2', 'ma2')), False],
- [(('a1', None), ('a2', 'ma2'), ('a3', ('ma31', 'ma32'))), True],
- [(('a1', None), ('a2', 'ma2'), ('batch', 'batch_1200234')), True]]
+LOGICAL_RULES = [
+ [(('a1', None), ('a2', 'ma2')), False],
+ [(('a1', None), ('a2', 'ma2'), ('a3', ('ma31', 'ma32'))), True],
+ [(('a1', None), ('a2', 'ma2'), ('a3', 'ma31'), ('a3', 'ma32')), False],
+ [(('a1', None), ('a2', 'ma2'), ('batch', 'batch_1200234')), True],
+ [(('a1', None), ('a2', 'ma2'), ('a2', 'ma1'), ('batch', 'model'), ('batch', 'data')), True],
+]
 SRS = [
  ShardingResource(),
  ShardingResource('data', None),

tests/jax/utils.py

@@ -321,8 +321,9 @@ def __call__(self, inputs, deterministic: bool = False):
 
  # Take elementwise product of above intermediate activations.
  x = functools.reduce(operator.mul, activations)
+ dropout_broadcast_dims = (0,) if self.transpose_batch_sequence else (1,)
  # Apply dropout and final dense output projection.
- x = nn.Dropout(rate=self.intermediate_dropout_rate, broadcast_dims=(-2,))(
+ x = nn.Dropout(rate=self.intermediate_dropout_rate, broadcast_dims=dropout_broadcast_dims)(
  x, deterministic=deterministic) # Broadcast along length.
  if self.transpose_batch_sequence:
  x = nn_partitioning.with_sharding_constraint(x, ('length', 'batch', 'mlp'))

transformer_engine/jax/fp8.py

@@ -190,7 +190,7 @@ def update_amax_history(amax_buffers: jnp.ndarray) -> jnp.ndarray:
  Update the amax history
  """
  updated_amax_buffers = jnp.roll(amax_buffers, -1, 1)
- updated_amax_buffers.at[:, 0].set(0)
+ updated_amax_buffers = updated_amax_buffers.at[:, 0].set(0)
  return updated_amax_buffers
 
  @staticmethod

transformer_engine/jax/module.py

@@ -683,6 +683,8 @@ class LayerNormMLP(TransformerEngineBase):
  Each activation has its own transformation layer.
  intermediate_dropout_rate: float, default = 0.1
  Dropout probability for the dropout op after the :attr:`activations`.
+ intermediate_hidden_dropout_dims: Sequence[int], default = ()
+ Dimensions that will share the same dropout mask for hidden
  axis: Union[Iterable[int], int], default = -1
  An integer tuple with axes to apply the transformation on.
 
@@ -716,6 +718,7 @@ class LayerNormMLP(TransformerEngineBase):
  return_layernorm_output: bool = True
  activations: Sequence[Union[str, Callable]] = ('relu',)
  intermediate_dropout_rate: float = 0.1
+ intermediate_hidden_dropout_dims: Sequence[int] = ()
  axis: Union[Iterable[int], int] = -1
  dtype: DType = jnp.float32
  transpose_batch_sequence: bool = True
@@ -912,8 +915,9 @@ def fp8_meta_generator():
  z = functools.reduce(operator.mul, activations)
  z = jnp.reshape(z, (*z.shape[:-2], -1))
 
- z = nn.Dropout(rate=self.intermediate_dropout_rate, broadcast_dims=(-2,))(
- z, deterministic=deterministic) # Broadcast along length.
+ z = nn.Dropout(rate=self.intermediate_dropout_rate,
+ broadcast_dims=self.intermediate_hidden_dropout_dims)(
+ z, deterministic=deterministic)
 
  # DenseGeneral 2
  hidden_size = inputs.shape[-1]

transformer_engine/jax/transformer.py

@@ -53,6 +53,10 @@ def extend_logical_axis_rules(rules: LogicalRules) -> LogicalRules:
  .. warning::
  Please make sure ShardingResource is set via fp8_autocast before calling this function.
 
+ .. note::
+ This function is only needed when using TransformerLayer. For other modules, such as
+ DenseGeneral, please properly set axes of kernels and bias.
+
  Parameters
  ----------
  rules : Sequence[Tuple[str, Union[str, None]]]
@@ -73,10 +77,12 @@ def extend_logical_axis_rules(rules: LogicalRules) -> LogicalRules:
  f"Thie axis_name should be str, but got {type(key)}."
  assert isinstance(val, str) or (val is None), \
  f"Thie mesh_axis_name should be str or None, but got {type(val)}."
- rules_map[key] = val
+ if key in rules_map:
+ rules_map[key].append(val)
+ else:
+ rules_map[key] = [val]
 
  gsr = global_shard_resource()
-
  te_logical_axis_rules = (('batch', gsr.dp_resource), ('embed', None), ('mlp', gsr.tp_resource),
  ('heads', gsr.tp_resource), ('kv', None), ('qkv_dim', None),
  ('kv_dim', None), ('joined_kv', gsr.tp_resource), ('act', None),
@@ -87,7 +93,7 @@ def extend_logical_axis_rules(rules: LogicalRules) -> LogicalRules:
  key = item[0]
  val = item[1]
  if key in rules_map:
- assert rules_map[key] == val, \
+ assert len(rules_map[key]) == 1 and rules_map[key][0] == val, \
  f"The rule diverged between TE and given rule." \
  f"Axis:{key} map to {rules_map[key]} in the given" \
  f" rules, but {val} in TE's rules."
@@ -447,41 +453,38 @@ def kv_init(key, shape, dtype):
  if decode:
  is_initialized = self.has_variable('cache', 'cached_key')
 
- # TODO (Ming Huang): Check performance on GPU withou swap dimensions # pylint: disable=fixme
- def swap_dims(x):
- return x[:-3] + tuple(x[i] for i in [-2, -1, -3])
-
- 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),
+ cached_key = self.variable('cache', 'cached_key', jnp.zeros, key.shape, key.dtype)
+ cached_value = self.variable('cache', 'cached_value', jnp.zeros, value.shape,
  value.dtype)
  cache_index = self.variable('cache', 'cache_index',
  lambda: jnp.array(0, dtype=jnp.int32))
  if is_initialized:
- batch, num_heads, head_dim, length = cached_key.value.shape
+ if self.transpose_batch_sequence:
+ length, batch, num_heads, head_dim = cached_key.value.shape
+ expected_shape = (1, batch, num_heads, head_dim)
+ one_hot_indices_shape = (length, 1, 1, 1)
+ else:
+ batch, length, num_heads, head_dim = cached_key.value.shape
+ expected_shape = (batch, 1, num_heads, head_dim)
+ one_hot_indices_shape = (1, length, 1, 1)
 
  # Sanity shape check of cached key against input query.
- expected_shape = (batch, 1, num_heads, head_dim)
  if expected_shape != query.shape:
  raise ValueError(
  'Autoregressive cache shape error, '
  f"expected query shape {expected_shape} instead got {query.shape}.")
 
  cur_index = cache_index.value
  one_hot_indices = jax_nn.one_hot(cur_index, length, dtype=key.dtype)
- one_token_key = jnp.moveaxis(key, -3, -1)
- one_token_value = jnp.moveaxis(value, -3, -1)
- key = cached_key.value + one_token_key * one_hot_indices
- value = cached_value.value + one_token_value * one_hot_indices
+ one_hot_indices = jnp.reshape(one_hot_indices, one_hot_indices_shape)
+ key = cached_key.value + key * one_hot_indices
+ value = cached_value.value + value * one_hot_indices
  cached_key.value = key
  cached_value.value = value
  cache_index.value = cache_index.value + 1
 
- key = jnp.moveaxis(key, -1, -3)
- value = jnp.moveaxis(value, -1, -3)
-
  mask = combine_masks(
- mask, jnp.broadcast_to(jnp.arange(length) <= cur_index, (batch, 1, 1, length)))
+ mask, jnp.broadcast_to(jnp.arange(length) > cur_index, (batch, 1, 1, length)))
 
  if bias is not None:
  bias = dynamic_vector_slice_in_dim(jnp.squeeze(bias, axis=0),
@@ -889,10 +892,11 @@ def hidden_dropout(x, deterministic):
  assert isinstance(self.hidden_dropout_dims, Sequence)
  x_shape_len = len(x.shape)
  for dims in self.hidden_dropout_dims:
- assert -x_shape_len < dims < x_shape_len
+ assert -x_shape_len <= dims < x_shape_len
 
  return nn.Dropout(rate=self.hidden_dropout,
- broadcast_dims=self.hidden_dropout_dims)(x, deterministic)
+ broadcast_dims=self.hidden_dropout_dims)(x,
+ deterministic=deterministic)
 
  x = hidden_dropout(x, deterministic)
  if self.drop_path > 0.0:
@@ -944,6 +948,7 @@ def hidden_dropout(x, deterministic):
  intermediate_dim=self.mlp_hidden_size,
  activations=self.mlp_activations,
  intermediate_dropout_rate=self.hidden_dropout,
+ intermediate_hidden_dropout_dims=self.hidden_dropout_dims,
  dtype=self.dtype,
  scale_axes=('embed',),
  kernel_init=self.mlp_kernel_init,

tests/jax/test_sharding.py

@@ -38,9 +38,13 @@ def _get_sharding_resource(mesh_names, sharding_type):
  ((4,), ("tp",), ShardingType.TP_ROW), ((2, 2), ("dp", "tp"), ShardingType.DP_TP_COL),
  ((2, 2), ("dp", "tp"), ShardingType.DP_TP_ROW)]
 
-LOGICAL_RULES = [[(('a1', None), ('a2', 'ma2')), False],
- [(('a1', None), ('a2', 'ma2'), ('a3', ('ma31', 'ma32'))), True],
- [(('a1', None), ('a2', 'ma2'), ('batch', 'batch_1200234')), True]]
+LOGICAL_RULES = [
+ [(('a1', None), ('a2', 'ma2')), False],
+ [(('a1', None), ('a2', 'ma2'), ('a3', ('ma31', 'ma32'))), True],
+ [(('a1', None), ('a2', 'ma2'), ('a3', 'ma31'), ('a3', 'ma32')), False],
+ [(('a1', None), ('a2', 'ma2'), ('batch', 'batch_1200234')), True],
+ [(('a1', None), ('a2', 'ma2'), ('a2', 'ma1'), ('batch', 'model'), ('batch', 'data')), True],
+]
 SRS = [
  ShardingResource(),
  ShardingResource('data', None),

tests/jax/utils.py

@@ -321,8 +321,9 @@ def __call__(self, inputs, deterministic: bool = False):
 
  # Take elementwise product of above intermediate activations.
  x = functools.reduce(operator.mul, activations)
+ dropout_broadcast_dims = (0,) if self.transpose_batch_sequence else (1,)
  # Apply dropout and final dense output projection.
- x = nn.Dropout(rate=self.intermediate_dropout_rate, broadcast_dims=(-2,))(
+ x = nn.Dropout(rate=self.intermediate_dropout_rate, broadcast_dims=dropout_broadcast_dims)(
  x, deterministic=deterministic) # Broadcast along length.
  if self.transpose_batch_sequence:
  x = nn_partitioning.with_sharding_constraint(x, ('length', 'batch', 'mlp'))

transformer_engine/jax/fp8.py

@@ -190,7 +190,7 @@ def update_amax_history(amax_buffers: jnp.ndarray) -> jnp.ndarray:
  Update the amax history
  """
  updated_amax_buffers = jnp.roll(amax_buffers, -1, 1)
- updated_amax_buffers.at[:, 0].set(0)
+ updated_amax_buffers = updated_amax_buffers.at[:, 0].set(0)
  return updated_amax_buffers
 
  @staticmethod

transformer_engine/jax/module.py

@@ -683,6 +683,8 @@ class LayerNormMLP(TransformerEngineBase):
  Each activation has its own transformation layer.
  intermediate_dropout_rate: float, default = 0.1
  Dropout probability for the dropout op after the :attr:`activations`.
+ intermediate_hidden_dropout_dims: Sequence[int], default = ()
+ Dimensions that will share the same dropout mask for hidden
  axis: Union[Iterable[int], int], default = -1
  An integer tuple with axes to apply the transformation on.
 
@@ -716,6 +718,7 @@ class LayerNormMLP(TransformerEngineBase):
  return_layernorm_output: bool = True
  activations: Sequence[Union[str, Callable]] = ('relu',)
  intermediate_dropout_rate: float = 0.1
+ intermediate_hidden_dropout_dims: Sequence[int] = ()
  axis: Union[Iterable[int], int] = -1
  dtype: DType = jnp.float32
  transpose_batch_sequence: bool = True
@@ -912,8 +915,9 @@ def fp8_meta_generator():
  z = functools.reduce(operator.mul, activations)
  z = jnp.reshape(z, (*z.shape[:-2], -1))
 
- z = nn.Dropout(rate=self.intermediate_dropout_rate, broadcast_dims=(-2,))(
- z, deterministic=deterministic) # Broadcast along length.
+ z = nn.Dropout(rate=self.intermediate_dropout_rate,
+ broadcast_dims=self.intermediate_hidden_dropout_dims)(
+ z, deterministic=deterministic)
 
  # DenseGeneral 2
  hidden_size = inputs.shape[-1]

transformer_engine/jax/transformer.py

@@ -53,6 +53,10 @@ def extend_logical_axis_rules(rules: LogicalRules) -> LogicalRules:
  .. warning::
  Please make sure ShardingResource is set via fp8_autocast before calling this function.
 
+ .. note::
+ This function is only needed when using TransformerLayer. For other modules, such as
+ DenseGeneral, please properly set axes of kernels and bias.
+
  Parameters
  ----------
  rules : Sequence[Tuple[str, Union[str, None]]]
@@ -73,10 +77,12 @@ def extend_logical_axis_rules(rules: LogicalRules) -> LogicalRules:
  f"Thie axis_name should be str, but got {type(key)}."
  assert isinstance(val, str) or (val is None), \
  f"Thie mesh_axis_name should be str or None, but got {type(val)}."
- rules_map[key] = val
+ if key in rules_map:
+ rules_map[key].append(val)
+ else:
+ rules_map[key] = [val]
 
  gsr = global_shard_resource()
-
  te_logical_axis_rules = (('batch', gsr.dp_resource), ('embed', None), ('mlp', gsr.tp_resource),
  ('heads', gsr.tp_resource), ('kv', None), ('qkv_dim', None),
  ('kv_dim', None), ('joined_kv', gsr.tp_resource), ('act', None),
@@ -87,7 +93,7 @@ def extend_logical_axis_rules(rules: LogicalRules) -> LogicalRules:
  key = item[0]
  val = item[1]
  if key in rules_map:
- assert rules_map[key] == val, \
+ assert len(rules_map[key]) == 1 and rules_map[key][0] == val, \
  f"The rule diverged between TE and given rule." \
  f"Axis:{key} map to {rules_map[key]} in the given" \
  f" rules, but {val} in TE's rules."
@@ -447,41 +453,38 @@ def kv_init(key, shape, dtype):
  if decode:
  is_initialized = self.has_variable('cache', 'cached_key')
 
- # TODO (Ming Huang): Check performance on GPU withou swap dimensions # pylint: disable=fixme
- def swap_dims(x):
- return x[:-3] + tuple(x[i] for i in [-2, -1, -3])
-
- 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),
+ cached_key = self.variable('cache', 'cached_key', jnp.zeros, key.shape, key.dtype)
+ cached_value = self.variable('cache', 'cached_value', jnp.zeros, value.shape,
  value.dtype)
  cache_index = self.variable('cache', 'cache_index',
  lambda: jnp.array(0, dtype=jnp.int32))
  if is_initialized:
- batch, num_heads, head_dim, length = cached_key.value.shape
+ if self.transpose_batch_sequence:
+ length, batch, num_heads, head_dim = cached_key.value.shape
+ expected_shape = (1, batch, num_heads, head_dim)
+ one_hot_indices_shape = (length, 1, 1, 1)
+ else:
+ batch, length, num_heads, head_dim = cached_key.value.shape
+ expected_shape = (batch, 1, num_heads, head_dim)
+ one_hot_indices_shape = (1, length, 1, 1)
 
  # Sanity shape check of cached key against input query.
- expected_shape = (batch, 1, num_heads, head_dim)
  if expected_shape != query.shape:
  raise ValueError(
  'Autoregressive cache shape error, '
  f"expected query shape {expected_shape} instead got {query.shape}.")
 
  cur_index = cache_index.value
  one_hot_indices = jax_nn.one_hot(cur_index, length, dtype=key.dtype)
- one_token_key = jnp.moveaxis(key, -3, -1)
- one_token_value = jnp.moveaxis(value, -3, -1)
- key = cached_key.value + one_token_key * one_hot_indices
- value = cached_value.value + one_token_value * one_hot_indices
+ one_hot_indices = jnp.reshape(one_hot_indices, one_hot_indices_shape)
+ key = cached_key.value + key * one_hot_indices
+ value = cached_value.value + value * one_hot_indices
  cached_key.value = key
  cached_value.value = value
  cache_index.value = cache_index.value + 1
 
- key = jnp.moveaxis(key, -1, -3)
- value = jnp.moveaxis(value, -1, -3)
-
  mask = combine_masks(
- mask, jnp.broadcast_to(jnp.arange(length) <= cur_index, (batch, 1, 1, length)))
+ mask, jnp.broadcast_to(jnp.arange(length) > cur_index, (batch, 1, 1, length)))
 
  if bias is not None:
  bias = dynamic_vector_slice_in_dim(jnp.squeeze(bias, axis=0),
@@ -889,10 +892,11 @@ def hidden_dropout(x, deterministic):
  assert isinstance(self.hidden_dropout_dims, Sequence)
  x_shape_len = len(x.shape)
  for dims in self.hidden_dropout_dims:
- assert -x_shape_len < dims < x_shape_len
+ assert -x_shape_len <= dims < x_shape_len
 
  return nn.Dropout(rate=self.hidden_dropout,
- broadcast_dims=self.hidden_dropout_dims)(x, deterministic)
+ broadcast_dims=self.hidden_dropout_dims)(x,
+ deterministic=deterministic)
 
  x = hidden_dropout(x, deterministic)
  if self.drop_path > 0.0:
@@ -944,6 +948,7 @@ def hidden_dropout(x, deterministic):
  intermediate_dim=self.mlp_hidden_size,
  activations=self.mlp_activations,
  intermediate_dropout_rate=self.hidden_dropout,
+ intermediate_hidden_dropout_dims=self.hidden_dropout_dims,
  dtype=self.dtype,
  scale_axes=('embed',),
  kernel_init=self.mlp_kernel_init,

tests/jax/test_sharding.py

@@ -38,9 +38,13 @@ def _get_sharding_resource(mesh_names, sharding_type):
  ((4,), ("tp",), ShardingType.TP_ROW), ((2, 2), ("dp", "tp"), ShardingType.DP_TP_COL),
  ((2, 2), ("dp", "tp"), ShardingType.DP_TP_ROW)]
 
-LOGICAL_RULES = [[(('a1', None), ('a2', 'ma2')), False],
- [(('a1', None), ('a2', 'ma2'), ('a3', ('ma31', 'ma32'))), True],
- [(('a1', None), ('a2', 'ma2'), ('batch', 'batch_1200234')), True]]
+LOGICAL_RULES = [
+ [(('a1', None), ('a2', 'ma2')), False],
+ [(('a1', None), ('a2', 'ma2'), ('a3', ('ma31', 'ma32'))), True],
+ [(('a1', None), ('a2', 'ma2'), ('a3', 'ma31'), ('a3', 'ma32')), False],
+ [(('a1', None), ('a2', 'ma2'), ('batch', 'batch_1200234')), True],
+ [(('a1', None), ('a2', 'ma2'), ('a2', 'ma1'), ('batch', 'model'), ('batch', 'data')), True],
+]
 SRS = [
  ShardingResource(),
  ShardingResource('data', None),

tests/jax/utils.py

@@ -321,8 +321,9 @@ def __call__(self, inputs, deterministic: bool = False):
 
  # Take elementwise product of above intermediate activations.
  x = functools.reduce(operator.mul, activations)
+ dropout_broadcast_dims = (0,) if self.transpose_batch_sequence else (1,)
  # Apply dropout and final dense output projection.
- x = nn.Dropout(rate=self.intermediate_dropout_rate, broadcast_dims=(-2,))(
+ x = nn.Dropout(rate=self.intermediate_dropout_rate, broadcast_dims=dropout_broadcast_dims)(
  x, deterministic=deterministic) # Broadcast along length.
  if self.transpose_batch_sequence:
  x = nn_partitioning.with_sharding_constraint(x, ('length', 'batch', 'mlp'))

transformer_engine/jax/fp8.py

@@ -190,7 +190,7 @@ def update_amax_history(amax_buffers: jnp.ndarray) -> jnp.ndarray:
  Update the amax history
  """
  updated_amax_buffers = jnp.roll(amax_buffers, -1, 1)
- updated_amax_buffers.at[:, 0].set(0)
+ updated_amax_buffers = updated_amax_buffers.at[:, 0].set(0)
  return updated_amax_buffers
 
  @staticmethod

transformer_engine/jax/module.py

@@ -683,6 +683,8 @@ class LayerNormMLP(TransformerEngineBase):
  Each activation has its own transformation layer.
  intermediate_dropout_rate: float, default = 0.1
  Dropout probability for the dropout op after the :attr:`activations`.
+ intermediate_hidden_dropout_dims: Sequence[int], default = ()
+ Dimensions that will share the same dropout mask for hidden
  axis: Union[Iterable[int], int], default = -1
  An integer tuple with axes to apply the transformation on.
 
@@ -716,6 +718,7 @@ class LayerNormMLP(TransformerEngineBase):
  return_layernorm_output: bool = True
  activations: Sequence[Union[str, Callable]] = ('relu',)
  intermediate_dropout_rate: float = 0.1
+ intermediate_hidden_dropout_dims: Sequence[int] = ()
  axis: Union[Iterable[int], int] = -1
  dtype: DType = jnp.float32
  transpose_batch_sequence: bool = True
@@ -912,8 +915,9 @@ def fp8_meta_generator():
  z = functools.reduce(operator.mul, activations)
  z = jnp.reshape(z, (*z.shape[:-2], -1))
 
- z = nn.Dropout(rate=self.intermediate_dropout_rate, broadcast_dims=(-2,))(
- z, deterministic=deterministic) # Broadcast along length.
+ z = nn.Dropout(rate=self.intermediate_dropout_rate,
+ broadcast_dims=self.intermediate_hidden_dropout_dims)(
+ z, deterministic=deterministic)
 
  # DenseGeneral 2
  hidden_size = inputs.shape[-1]

transformer_engine/jax/transformer.py

@@ -53,6 +53,10 @@ def extend_logical_axis_rules(rules: LogicalRules) -> LogicalRules:
  .. warning::
  Please make sure ShardingResource is set via fp8_autocast before calling this function.
 
+ .. note::
+ This function is only needed when using TransformerLayer. For other modules, such as
+ DenseGeneral, please properly set axes of kernels and bias.
+
  Parameters
  ----------
  rules : Sequence[Tuple[str, Union[str, None]]]
@@ -73,10 +77,12 @@ def extend_logical_axis_rules(rules: LogicalRules) -> LogicalRules:
  f"Thie axis_name should be str, but got {type(key)}."
  assert isinstance(val, str) or (val is None), \
  f"Thie mesh_axis_name should be str or None, but got {type(val)}."
- rules_map[key] = val
+ if key in rules_map:
+ rules_map[key].append(val)
+ else:
+ rules_map[key] = [val]
 
  gsr = global_shard_resource()
-
  te_logical_axis_rules = (('batch', gsr.dp_resource), ('embed', None), ('mlp', gsr.tp_resource),
  ('heads', gsr.tp_resource), ('kv', None), ('qkv_dim', None),
  ('kv_dim', None), ('joined_kv', gsr.tp_resource), ('act', None),
@@ -87,7 +93,7 @@ def extend_logical_axis_rules(rules: LogicalRules) -> LogicalRules:
  key = item[0]
  val = item[1]
  if key in rules_map:
- assert rules_map[key] == val, \
+ assert len(rules_map[key]) == 1 and rules_map[key][0] == val, \
  f"The rule diverged between TE and given rule." \
  f"Axis:{key} map to {rules_map[key]} in the given" \
  f" rules, but {val} in TE's rules."
@@ -447,41 +453,38 @@ def kv_init(key, shape, dtype):
  if decode:
  is_initialized = self.has_variable('cache', 'cached_key')
 
- # TODO (Ming Huang): Check performance on GPU withou swap dimensions # pylint: disable=fixme
- def swap_dims(x):
- return x[:-3] + tuple(x[i] for i in [-2, -1, -3])
-
- 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),
+ cached_key = self.variable('cache', 'cached_key', jnp.zeros, key.shape, key.dtype)
+ cached_value = self.variable('cache', 'cached_value', jnp.zeros, value.shape,
  value.dtype)
  cache_index = self.variable('cache', 'cache_index',
  lambda: jnp.array(0, dtype=jnp.int32))
  if is_initialized:
- batch, num_heads, head_dim, length = cached_key.value.shape
+ if self.transpose_batch_sequence:
+ length, batch, num_heads, head_dim = cached_key.value.shape
+ expected_shape = (1, batch, num_heads, head_dim)
+ one_hot_indices_shape = (length, 1, 1, 1)
+ else:
+ batch, length, num_heads, head_dim = cached_key.value.shape
+ expected_shape = (batch, 1, num_heads, head_dim)
+ one_hot_indices_shape = (1, length, 1, 1)
 
  # Sanity shape check of cached key against input query.
- expected_shape = (batch, 1, num_heads, head_dim)
  if expected_shape != query.shape:
  raise ValueError(
  'Autoregressive cache shape error, '
  f"expected query shape {expected_shape} instead got {query.shape}.")
 
  cur_index = cache_index.value
  one_hot_indices = jax_nn.one_hot(cur_index, length, dtype=key.dtype)
- one_token_key = jnp.moveaxis(key, -3, -1)
- one_token_value = jnp.moveaxis(value, -3, -1)
- key = cached_key.value + one_token_key * one_hot_indices
- value = cached_value.value + one_token_value * one_hot_indices
+ one_hot_indices = jnp.reshape(one_hot_indices, one_hot_indices_shape)
+ key = cached_key.value + key * one_hot_indices
+ value = cached_value.value + value * one_hot_indices
  cached_key.value = key
  cached_value.value = value
  cache_index.value = cache_index.value + 1
 
- key = jnp.moveaxis(key, -1, -3)
- value = jnp.moveaxis(value, -1, -3)
-
  mask = combine_masks(
- mask, jnp.broadcast_to(jnp.arange(length) <= cur_index, (batch, 1, 1, length)))
+ mask, jnp.broadcast_to(jnp.arange(length) > cur_index, (batch, 1, 1, length)))
 
  if bias is not None:
  bias = dynamic_vector_slice_in_dim(jnp.squeeze(bias, axis=0),
@@ -889,10 +892,11 @@ def hidden_dropout(x, deterministic):
  assert isinstance(self.hidden_dropout_dims, Sequence)
  x_shape_len = len(x.shape)
  for dims in self.hidden_dropout_dims:
- assert -x_shape_len < dims < x_shape_len
+ assert -x_shape_len <= dims < x_shape_len
 
  return nn.Dropout(rate=self.hidden_dropout,
- broadcast_dims=self.hidden_dropout_dims)(x, deterministic)
+ broadcast_dims=self.hidden_dropout_dims)(x,
+ deterministic=deterministic)
 
  x = hidden_dropout(x, deterministic)
  if self.drop_path > 0.0:
@@ -944,6 +948,7 @@ def hidden_dropout(x, deterministic):
  intermediate_dim=self.mlp_hidden_size,
  activations=self.mlp_activations,
  intermediate_dropout_rate=self.hidden_dropout,
+ intermediate_hidden_dropout_dims=self.hidden_dropout_dims,
  dtype=self.dtype,
  scale_axes=('embed',),
  kernel_init=self.mlp_kernel_init,

tests/jax/test_sharding.py

@@ -38,9 +38,13 @@ def _get_sharding_resource(mesh_names, sharding_type):
  ((4,), ("tp",), ShardingType.TP_ROW), ((2, 2), ("dp", "tp"), ShardingType.DP_TP_COL),
  ((2, 2), ("dp", "tp"), ShardingType.DP_TP_ROW)]
 
-LOGICAL_RULES = [[(('a1', None), ('a2', 'ma2')), False],
- [(('a1', None), ('a2', 'ma2'), ('a3', ('ma31', 'ma32'))), True],
- [(('a1', None), ('a2', 'ma2'), ('batch', 'batch_1200234')), True]]
+LOGICAL_RULES = [
+ [(('a1', None), ('a2', 'ma2')), False],
+ [(('a1', None), ('a2', 'ma2'), ('a3', ('ma31', 'ma32'))), True],
+ [(('a1', None), ('a2', 'ma2'), ('a3', 'ma31'), ('a3', 'ma32')), False],
+ [(('a1', None), ('a2', 'ma2'), ('batch', 'batch_1200234')), True],
+ [(('a1', None), ('a2', 'ma2'), ('a2', 'ma1'), ('batch', 'model'), ('batch', 'data')), True],
+]
 SRS = [
  ShardingResource(),
  ShardingResource('data', None),

tests/jax/utils.py

@@ -321,8 +321,9 @@ def __call__(self, inputs, deterministic: bool = False):
 
  # Take elementwise product of above intermediate activations.
  x = functools.reduce(operator.mul, activations)
+ dropout_broadcast_dims = (0,) if self.transpose_batch_sequence else (1,)
  # Apply dropout and final dense output projection.
- x = nn.Dropout(rate=self.intermediate_dropout_rate, broadcast_dims=(-2,))(
+ x = nn.Dropout(rate=self.intermediate_dropout_rate, broadcast_dims=dropout_broadcast_dims)(
  x, deterministic=deterministic) # Broadcast along length.
  if self.transpose_batch_sequence:
  x = nn_partitioning.with_sharding_constraint(x, ('length', 'batch', 'mlp'))

transformer_engine/jax/fp8.py

@@ -190,7 +190,7 @@ def update_amax_history(amax_buffers: jnp.ndarray) -> jnp.ndarray:
  Update the amax history
  """
  updated_amax_buffers = jnp.roll(amax_buffers, -1, 1)
- updated_amax_buffers.at[:, 0].set(0)
+ updated_amax_buffers = updated_amax_buffers.at[:, 0].set(0)
  return updated_amax_buffers
 
  @staticmethod

transformer_engine/jax/module.py

@@ -683,6 +683,8 @@ class LayerNormMLP(TransformerEngineBase):
  Each activation has its own transformation layer.
  intermediate_dropout_rate: float, default = 0.1
  Dropout probability for the dropout op after the :attr:`activations`.
+ intermediate_hidden_dropout_dims: Sequence[int], default = ()
+ Dimensions that will share the same dropout mask for hidden
  axis: Union[Iterable[int], int], default = -1
  An integer tuple with axes to apply the transformation on.
 
@@ -716,6 +718,7 @@ class LayerNormMLP(TransformerEngineBase):
  return_layernorm_output: bool = True
  activations: Sequence[Union[str, Callable]] = ('relu',)
  intermediate_dropout_rate: float = 0.1
+ intermediate_hidden_dropout_dims: Sequence[int] = ()
  axis: Union[Iterable[int], int] = -1
  dtype: DType = jnp.float32
  transpose_batch_sequence: bool = True
@@ -912,8 +915,9 @@ def fp8_meta_generator():
  z = functools.reduce(operator.mul, activations)
  z = jnp.reshape(z, (*z.shape[:-2], -1))
 
- z = nn.Dropout(rate=self.intermediate_dropout_rate, broadcast_dims=(-2,))(
- z, deterministic=deterministic) # Broadcast along length.
+ z = nn.Dropout(rate=self.intermediate_dropout_rate,
+ broadcast_dims=self.intermediate_hidden_dropout_dims)(
+ z, deterministic=deterministic)
 
  # DenseGeneral 2
  hidden_size = inputs.shape[-1]

transformer_engine/jax/transformer.py

@@ -53,6 +53,10 @@ def extend_logical_axis_rules(rules: LogicalRules) -> LogicalRules:
  .. warning::
  Please make sure ShardingResource is set via fp8_autocast before calling this function.
 
+ .. note::
+ This function is only needed when using TransformerLayer. For other modules, such as
+ DenseGeneral, please properly set axes of kernels and bias.
+
  Parameters
  ----------
  rules : Sequence[Tuple[str, Union[str, None]]]
@@ -73,10 +77,12 @@ def extend_logical_axis_rules(rules: LogicalRules) -> LogicalRules:
  f"Thie axis_name should be str, but got {type(key)}."
  assert isinstance(val, str) or (val is None), \
  f"Thie mesh_axis_name should be str or None, but got {type(val)}."
- rules_map[key] = val
+ if key in rules_map:
+ rules_map[key].append(val)
+ else:
+ rules_map[key] = [val]
 
  gsr = global_shard_resource()
-
  te_logical_axis_rules = (('batch', gsr.dp_resource), ('embed', None), ('mlp', gsr.tp_resource),
  ('heads', gsr.tp_resource), ('kv', None), ('qkv_dim', None),
  ('kv_dim', None), ('joined_kv', gsr.tp_resource), ('act', None),
@@ -87,7 +93,7 @@ def extend_logical_axis_rules(rules: LogicalRules) -> LogicalRules:
  key = item[0]
  val = item[1]
  if key in rules_map:
- assert rules_map[key] == val, \
+ assert len(rules_map[key]) == 1 and rules_map[key][0] == val, \
  f"The rule diverged between TE and given rule." \
  f"Axis:{key} map to {rules_map[key]} in the given" \
  f" rules, but {val} in TE's rules."
@@ -447,41 +453,38 @@ def kv_init(key, shape, dtype):
  if decode:
  is_initialized = self.has_variable('cache', 'cached_key')
 
- # TODO (Ming Huang): Check performance on GPU withou swap dimensions # pylint: disable=fixme
- def swap_dims(x):
- return x[:-3] + tuple(x[i] for i in [-2, -1, -3])
-
- 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),
+ cached_key = self.variable('cache', 'cached_key', jnp.zeros, key.shape, key.dtype)
+ cached_value = self.variable('cache', 'cached_value', jnp.zeros, value.shape,
  value.dtype)
  cache_index = self.variable('cache', 'cache_index',
  lambda: jnp.array(0, dtype=jnp.int32))
  if is_initialized:
- batch, num_heads, head_dim, length = cached_key.value.shape
+ if self.transpose_batch_sequence:
+ length, batch, num_heads, head_dim = cached_key.value.shape
+ expected_shape = (1, batch, num_heads, head_dim)
+ one_hot_indices_shape = (length, 1, 1, 1)
+ else:
+ batch, length, num_heads, head_dim = cached_key.value.shape
+ expected_shape = (batch, 1, num_heads, head_dim)
+ one_hot_indices_shape = (1, length, 1, 1)
 
  # Sanity shape check of cached key against input query.
- expected_shape = (batch, 1, num_heads, head_dim)
  if expected_shape != query.shape:
  raise ValueError(
  'Autoregressive cache shape error, '
  f"expected query shape {expected_shape} instead got {query.shape}.")
 
  cur_index = cache_index.value
  one_hot_indices = jax_nn.one_hot(cur_index, length, dtype=key.dtype)
- one_token_key = jnp.moveaxis(key, -3, -1)
- one_token_value = jnp.moveaxis(value, -3, -1)
- key = cached_key.value + one_token_key * one_hot_indices
- value = cached_value.value + one_token_value * one_hot_indices
+ one_hot_indices = jnp.reshape(one_hot_indices, one_hot_indices_shape)
+ key = cached_key.value + key * one_hot_indices
+ value = cached_value.value + value * one_hot_indices
  cached_key.value = key
  cached_value.value = value
  cache_index.value = cache_index.value + 1
 
- key = jnp.moveaxis(key, -1, -3)
- value = jnp.moveaxis(value, -1, -3)
-
  mask = combine_masks(
- mask, jnp.broadcast_to(jnp.arange(length) <= cur_index, (batch, 1, 1, length)))
+ mask, jnp.broadcast_to(jnp.arange(length) > cur_index, (batch, 1, 1, length)))
 
  if bias is not None:
  bias = dynamic_vector_slice_in_dim(jnp.squeeze(bias, axis=0),
@@ -889,10 +892,11 @@ def hidden_dropout(x, deterministic):
  assert isinstance(self.hidden_dropout_dims, Sequence)
  x_shape_len = len(x.shape)
  for dims in self.hidden_dropout_dims:
- assert -x_shape_len < dims < x_shape_len
+ assert -x_shape_len <= dims < x_shape_len
 
  return nn.Dropout(rate=self.hidden_dropout,
- broadcast_dims=self.hidden_dropout_dims)(x, deterministic)
+ broadcast_dims=self.hidden_dropout_dims)(x,
+ deterministic=deterministic)
 
  x = hidden_dropout(x, deterministic)
  if self.drop_path > 0.0:
@@ -944,6 +948,7 @@ def hidden_dropout(x, deterministic):
  intermediate_dim=self.mlp_hidden_size,
  activations=self.mlp_activations,
  intermediate_dropout_rate=self.hidden_dropout,
+ intermediate_hidden_dropout_dims=self.hidden_dropout_dims,
  dtype=self.dtype,
  scale_axes=('embed',),
  kernel_init=self.mlp_kernel_init,

tests/jax/test_sharding.py

@@ -38,9 +38,13 @@ def _get_sharding_resource(mesh_names, sharding_type):
  ((4,), ("tp",), ShardingType.TP_ROW), ((2, 2), ("dp", "tp"), ShardingType.DP_TP_COL),
  ((2, 2), ("dp", "tp"), ShardingType.DP_TP_ROW)]
 
-LOGICAL_RULES = [[(('a1', None), ('a2', 'ma2')), False],
- [(('a1', None), ('a2', 'ma2'), ('a3', ('ma31', 'ma32'))), True],
- [(('a1', None), ('a2', 'ma2'), ('batch', 'batch_1200234')), True]]
+LOGICAL_RULES = [
+ [(('a1', None), ('a2', 'ma2')), False],
+ [(('a1', None), ('a2', 'ma2'), ('a3', ('ma31', 'ma32'))), True],
+ [(('a1', None), ('a2', 'ma2'), ('a3', 'ma31'), ('a3', 'ma32')), False],
+ [(('a1', None), ('a2', 'ma2'), ('batch', 'batch_1200234')), True],
+ [(('a1', None), ('a2', 'ma2'), ('a2', 'ma1'), ('batch', 'model'), ('batch', 'data')), True],
+]
 SRS = [
  ShardingResource(),
  ShardingResource('data', None),

tests/jax/utils.py

@@ -321,8 +321,9 @@ def __call__(self, inputs, deterministic: bool = False):
 
  # Take elementwise product of above intermediate activations.
  x = functools.reduce(operator.mul, activations)
+ dropout_broadcast_dims = (0,) if self.transpose_batch_sequence else (1,)
  # Apply dropout and final dense output projection.
- x = nn.Dropout(rate=self.intermediate_dropout_rate, broadcast_dims=(-2,))(
+ x = nn.Dropout(rate=self.intermediate_dropout_rate, broadcast_dims=dropout_broadcast_dims)(
  x, deterministic=deterministic) # Broadcast along length.
  if self.transpose_batch_sequence:
  x = nn_partitioning.with_sharding_constraint(x, ('length', 'batch', 'mlp'))

transformer_engine/jax/fp8.py

@@ -190,7 +190,7 @@ def update_amax_history(amax_buffers: jnp.ndarray) -> jnp.ndarray:
  Update the amax history
  """
  updated_amax_buffers = jnp.roll(amax_buffers, -1, 1)
- updated_amax_buffers.at[:, 0].set(0)
+ updated_amax_buffers = updated_amax_buffers.at[:, 0].set(0)
  return updated_amax_buffers
 
  @staticmethod

transformer_engine/jax/module.py

@@ -683,6 +683,8 @@ class LayerNormMLP(TransformerEngineBase):
  Each activation has its own transformation layer.
  intermediate_dropout_rate: float, default = 0.1
  Dropout probability for the dropout op after the :attr:`activations`.
+ intermediate_hidden_dropout_dims: Sequence[int], default = ()
+ Dimensions that will share the same dropout mask for hidden
  axis: Union[Iterable[int], int], default = -1
  An integer tuple with axes to apply the transformation on.
 
@@ -716,6 +718,7 @@ class LayerNormMLP(TransformerEngineBase):
  return_layernorm_output: bool = True
  activations: Sequence[Union[str, Callable]] = ('relu',)
  intermediate_dropout_rate: float = 0.1
+ intermediate_hidden_dropout_dims: Sequence[int] = ()
  axis: Union[Iterable[int], int] = -1
  dtype: DType = jnp.float32
  transpose_batch_sequence: bool = True
@@ -912,8 +915,9 @@ def fp8_meta_generator():
  z = functools.reduce(operator.mul, activations)
  z = jnp.reshape(z, (*z.shape[:-2], -1))
 
- z = nn.Dropout(rate=self.intermediate_dropout_rate, broadcast_dims=(-2,))(
- z, deterministic=deterministic) # Broadcast along length.
+ z = nn.Dropout(rate=self.intermediate_dropout_rate,
+ broadcast_dims=self.intermediate_hidden_dropout_dims)(
+ z, deterministic=deterministic)
 
  # DenseGeneral 2
  hidden_size = inputs.shape[-1]

transformer_engine/jax/transformer.py

@@ -53,6 +53,10 @@ def extend_logical_axis_rules(rules: LogicalRules) -> LogicalRules:
  .. warning::
  Please make sure ShardingResource is set via fp8_autocast before calling this function.
 
+ .. note::
+ This function is only needed when using TransformerLayer. For other modules, such as
+ DenseGeneral, please properly set axes of kernels and bias.
+
  Parameters
  ----------
  rules : Sequence[Tuple[str, Union[str, None]]]
@@ -73,10 +77,12 @@ def extend_logical_axis_rules(rules: LogicalRules) -> LogicalRules:
  f"Thie axis_name should be str, but got {type(key)}."
  assert isinstance(val, str) or (val is None), \
  f"Thie mesh_axis_name should be str or None, but got {type(val)}."
- rules_map[key] = val
+ if key in rules_map:
+ rules_map[key].append(val)
+ else:
+ rules_map[key] = [val]
 
  gsr = global_shard_resource()
-
  te_logical_axis_rules = (('batch', gsr.dp_resource), ('embed', None), ('mlp', gsr.tp_resource),
  ('heads', gsr.tp_resource), ('kv', None), ('qkv_dim', None),
  ('kv_dim', None), ('joined_kv', gsr.tp_resource), ('act', None),
@@ -87,7 +93,7 @@ def extend_logical_axis_rules(rules: LogicalRules) -> LogicalRules:
  key = item[0]
  val = item[1]
  if key in rules_map:
- assert rules_map[key] == val, \
+ assert len(rules_map[key]) == 1 and rules_map[key][0] == val, \
  f"The rule diverged between TE and given rule." \
  f"Axis:{key} map to {rules_map[key]} in the given" \
  f" rules, but {val} in TE's rules."
@@ -447,41 +453,38 @@ def kv_init(key, shape, dtype):
  if decode:
  is_initialized = self.has_variable('cache', 'cached_key')
 
- # TODO (Ming Huang): Check performance on GPU withou swap dimensions # pylint: disable=fixme
- def swap_dims(x):
- return x[:-3] + tuple(x[i] for i in [-2, -1, -3])
-
- 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),
+ cached_key = self.variable('cache', 'cached_key', jnp.zeros, key.shape, key.dtype)
+ cached_value = self.variable('cache', 'cached_value', jnp.zeros, value.shape,
  value.dtype)
  cache_index = self.variable('cache', 'cache_index',
  lambda: jnp.array(0, dtype=jnp.int32))
  if is_initialized:
- batch, num_heads, head_dim, length = cached_key.value.shape
+ if self.transpose_batch_sequence:
+ length, batch, num_heads, head_dim = cached_key.value.shape
+ expected_shape = (1, batch, num_heads, head_dim)
+ one_hot_indices_shape = (length, 1, 1, 1)
+ else:
+ batch, length, num_heads, head_dim = cached_key.value.shape
+ expected_shape = (batch, 1, num_heads, head_dim)
+ one_hot_indices_shape = (1, length, 1, 1)
 
  # Sanity shape check of cached key against input query.
- expected_shape = (batch, 1, num_heads, head_dim)
  if expected_shape != query.shape:
  raise ValueError(
  'Autoregressive cache shape error, '
  f"expected query shape {expected_shape} instead got {query.shape}.")
 
  cur_index = cache_index.value
  one_hot_indices = jax_nn.one_hot(cur_index, length, dtype=key.dtype)
- one_token_key = jnp.moveaxis(key, -3, -1)
- one_token_value = jnp.moveaxis(value, -3, -1)
- key = cached_key.value + one_token_key * one_hot_indices
- value = cached_value.value + one_token_value * one_hot_indices
+ one_hot_indices = jnp.reshape(one_hot_indices, one_hot_indices_shape)
+ key = cached_key.value + key * one_hot_indices
+ value = cached_value.value + value * one_hot_indices
  cached_key.value = key
  cached_value.value = value
  cache_index.value = cache_index.value + 1
 
- key = jnp.moveaxis(key, -1, -3)
- value = jnp.moveaxis(value, -1, -3)
-
  mask = combine_masks(
- mask, jnp.broadcast_to(jnp.arange(length) <= cur_index, (batch, 1, 1, length)))
+ mask, jnp.broadcast_to(jnp.arange(length) > cur_index, (batch, 1, 1, length)))
 
  if bias is not None:
  bias = dynamic_vector_slice_in_dim(jnp.squeeze(bias, axis=0),
@@ -889,10 +892,11 @@ def hidden_dropout(x, deterministic):
  assert isinstance(self.hidden_dropout_dims, Sequence)
  x_shape_len = len(x.shape)
  for dims in self.hidden_dropout_dims:
- assert -x_shape_len < dims < x_shape_len
+ assert -x_shape_len <= dims < x_shape_len
 
  return nn.Dropout(rate=self.hidden_dropout,
- broadcast_dims=self.hidden_dropout_dims)(x, deterministic)
+ broadcast_dims=self.hidden_dropout_dims)(x,
+ deterministic=deterministic)
 
  x = hidden_dropout(x, deterministic)
  if self.drop_path > 0.0:
@@ -944,6 +948,7 @@ def hidden_dropout(x, deterministic):
  intermediate_dim=self.mlp_hidden_size,
  activations=self.mlp_activations,
  intermediate_dropout_rate=self.hidden_dropout,
+ intermediate_hidden_dropout_dims=self.hidden_dropout_dims,
  dtype=self.dtype,
  scale_axes=('embed',),
  kernel_init=self.mlp_kernel_init,

tests/jax/test_sharding.py

@@ -38,9 +38,13 @@ def _get_sharding_resource(mesh_names, sharding_type):
  ((4,), ("tp",), ShardingType.TP_ROW), ((2, 2), ("dp", "tp"), ShardingType.DP_TP_COL),
  ((2, 2), ("dp", "tp"), ShardingType.DP_TP_ROW)]
 
-LOGICAL_RULES = [[(('a1', None), ('a2', 'ma2')), False],
- [(('a1', None), ('a2', 'ma2'), ('a3', ('ma31', 'ma32'))), True],
- [(('a1', None), ('a2', 'ma2'), ('batch', 'batch_1200234')), True]]
+LOGICAL_RULES = [
+ [(('a1', None), ('a2', 'ma2')), False],
+ [(('a1', None), ('a2', 'ma2'), ('a3', ('ma31', 'ma32'))), True],
+ [(('a1', None), ('a2', 'ma2'), ('a3', 'ma31'), ('a3', 'ma32')), False],
+ [(('a1', None), ('a2', 'ma2'), ('batch', 'batch_1200234')), True],
+ [(('a1', None), ('a2', 'ma2'), ('a2', 'ma1'), ('batch', 'model'), ('batch', 'data')), True],
+]
 SRS = [
  ShardingResource(),
  ShardingResource('data', None),

tests/jax/utils.py

@@ -321,8 +321,9 @@ def __call__(self, inputs, deterministic: bool = False):
 
  # Take elementwise product of above intermediate activations.
  x = functools.reduce(operator.mul, activations)
+ dropout_broadcast_dims = (0,) if self.transpose_batch_sequence else (1,)
  # Apply dropout and final dense output projection.
- x = nn.Dropout(rate=self.intermediate_dropout_rate, broadcast_dims=(-2,))(
+ x = nn.Dropout(rate=self.intermediate_dropout_rate, broadcast_dims=dropout_broadcast_dims)(
  x, deterministic=deterministic) # Broadcast along length.
  if self.transpose_batch_sequence:
  x = nn_partitioning.with_sharding_constraint(x, ('length', 'batch', 'mlp'))

transformer_engine/jax/fp8.py

@@ -190,7 +190,7 @@ def update_amax_history(amax_buffers: jnp.ndarray) -> jnp.ndarray:
  Update the amax history
  """
  updated_amax_buffers = jnp.roll(amax_buffers, -1, 1)
- updated_amax_buffers.at[:, 0].set(0)
+ updated_amax_buffers = updated_amax_buffers.at[:, 0].set(0)
  return updated_amax_buffers
 
  @staticmethod

transformer_engine/jax/module.py

@@ -683,6 +683,8 @@ class LayerNormMLP(TransformerEngineBase):
  Each activation has its own transformation layer.
  intermediate_dropout_rate: float, default = 0.1
  Dropout probability for the dropout op after the :attr:`activations`.
+ intermediate_hidden_dropout_dims: Sequence[int], default = ()
+ Dimensions that will share the same dropout mask for hidden
  axis: Union[Iterable[int], int], default = -1
  An integer tuple with axes to apply the transformation on.
 
@@ -716,6 +718,7 @@ class LayerNormMLP(TransformerEngineBase):
  return_layernorm_output: bool = True
  activations: Sequence[Union[str, Callable]] = ('relu',)
  intermediate_dropout_rate: float = 0.1
+ intermediate_hidden_dropout_dims: Sequence[int] = ()
  axis: Union[Iterable[int], int] = -1
  dtype: DType = jnp.float32
  transpose_batch_sequence: bool = True
@@ -912,8 +915,9 @@ def fp8_meta_generator():
  z = functools.reduce(operator.mul, activations)
  z = jnp.reshape(z, (*z.shape[:-2], -1))
 
- z = nn.Dropout(rate=self.intermediate_dropout_rate, broadcast_dims=(-2,))(
- z, deterministic=deterministic) # Broadcast along length.
+ z = nn.Dropout(rate=self.intermediate_dropout_rate,
+ broadcast_dims=self.intermediate_hidden_dropout_dims)(
+ z, deterministic=deterministic)
 
  # DenseGeneral 2
  hidden_size = inputs.shape[-1]

transformer_engine/jax/transformer.py

@@ -53,6 +53,10 @@ def extend_logical_axis_rules(rules: LogicalRules) -> LogicalRules:
  .. warning::
  Please make sure ShardingResource is set via fp8_autocast before calling this function.
 
+ .. note::
+ This function is only needed when using TransformerLayer. For other modules, such as
+ DenseGeneral, please properly set axes of kernels and bias.
+
  Parameters
  ----------
  rules : Sequence[Tuple[str, Union[str, None]]]
@@ -73,10 +77,12 @@ def extend_logical_axis_rules(rules: LogicalRules) -> LogicalRules:
  f"Thie axis_name should be str, but got {type(key)}."
  assert isinstance(val, str) or (val is None), \
  f"Thie mesh_axis_name should be str or None, but got {type(val)}."
- rules_map[key] = val
+ if key in rules_map:
+ rules_map[key].append(val)
+ else:
+ rules_map[key] = [val]
 
  gsr = global_shard_resource()
-
  te_logical_axis_rules = (('batch', gsr.dp_resource), ('embed', None), ('mlp', gsr.tp_resource),
  ('heads', gsr.tp_resource), ('kv', None), ('qkv_dim', None),
  ('kv_dim', None), ('joined_kv', gsr.tp_resource), ('act', None),
@@ -87,7 +93,7 @@ def extend_logical_axis_rules(rules: LogicalRules) -> LogicalRules:
  key = item[0]
  val = item[1]
  if key in rules_map:
- assert rules_map[key] == val, \
+ assert len(rules_map[key]) == 1 and rules_map[key][0] == val, \
  f"The rule diverged between TE and given rule." \
  f"Axis:{key} map to {rules_map[key]} in the given" \
  f" rules, but {val} in TE's rules."
@@ -447,41 +453,38 @@ def kv_init(key, shape, dtype):
  if decode:
  is_initialized = self.has_variable('cache', 'cached_key')
 
- # TODO (Ming Huang): Check performance on GPU withou swap dimensions # pylint: disable=fixme
- def swap_dims(x):
- return x[:-3] + tuple(x[i] for i in [-2, -1, -3])
-
- 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),
+ cached_key = self.variable('cache', 'cached_key', jnp.zeros, key.shape, key.dtype)
+ cached_value = self.variable('cache', 'cached_value', jnp.zeros, value.shape,
  value.dtype)
  cache_index = self.variable('cache', 'cache_index',
  lambda: jnp.array(0, dtype=jnp.int32))
  if is_initialized:
- batch, num_heads, head_dim, length = cached_key.value.shape
+ if self.transpose_batch_sequence:
+ length, batch, num_heads, head_dim = cached_key.value.shape
+ expected_shape = (1, batch, num_heads, head_dim)
+ one_hot_indices_shape = (length, 1, 1, 1)
+ else:
+ batch, length, num_heads, head_dim = cached_key.value.shape
+ expected_shape = (batch, 1, num_heads, head_dim)
+ one_hot_indices_shape = (1, length, 1, 1)
 
  # Sanity shape check of cached key against input query.
- expected_shape = (batch, 1, num_heads, head_dim)
  if expected_shape != query.shape:
  raise ValueError(
  'Autoregressive cache shape error, '
  f"expected query shape {expected_shape} instead got {query.shape}.")
 
  cur_index = cache_index.value
  one_hot_indices = jax_nn.one_hot(cur_index, length, dtype=key.dtype)
- one_token_key = jnp.moveaxis(key, -3, -1)
- one_token_value = jnp.moveaxis(value, -3, -1)
- key = cached_key.value + one_token_key * one_hot_indices
- value = cached_value.value + one_token_value * one_hot_indices
+ one_hot_indices = jnp.reshape(one_hot_indices, one_hot_indices_shape)
+ key = cached_key.value + key * one_hot_indices
+ value = cached_value.value + value * one_hot_indices
  cached_key.value = key
  cached_value.value = value
  cache_index.value = cache_index.value + 1
 
- key = jnp.moveaxis(key, -1, -3)
- value = jnp.moveaxis(value, -1, -3)
-
  mask = combine_masks(
- mask, jnp.broadcast_to(jnp.arange(length) <= cur_index, (batch, 1, 1, length)))
+ mask, jnp.broadcast_to(jnp.arange(length) > cur_index, (batch, 1, 1, length)))
 
  if bias is not None:
  bias = dynamic_vector_slice_in_dim(jnp.squeeze(bias, axis=0),
@@ -889,10 +892,11 @@ def hidden_dropout(x, deterministic):
  assert isinstance(self.hidden_dropout_dims, Sequence)
  x_shape_len = len(x.shape)
  for dims in self.hidden_dropout_dims:
- assert -x_shape_len < dims < x_shape_len
+ assert -x_shape_len <= dims < x_shape_len
 
  return nn.Dropout(rate=self.hidden_dropout,
- broadcast_dims=self.hidden_dropout_dims)(x, deterministic)
+ broadcast_dims=self.hidden_dropout_dims)(x,
+ deterministic=deterministic)
 
  x = hidden_dropout(x, deterministic)
  if self.drop_path > 0.0:
@@ -944,6 +948,7 @@ def hidden_dropout(x, deterministic):
  intermediate_dim=self.mlp_hidden_size,
  activations=self.mlp_activations,
  intermediate_dropout_rate=self.hidden_dropout,
+ intermediate_hidden_dropout_dims=self.hidden_dropout_dims,
  dtype=self.dtype,
  scale_axes=('embed',),
  kernel_init=self.mlp_kernel_init,