From 4ea11cfcf45749bbe1e62aa6c115aae943162641 Mon Sep 17 00:00:00 2001 From: RecML authors Date: Tue, 3 Nov 2026 15:09:01 -0800 Subject: [PATCH] Internal change PiperOrigin-RevId: 1013324345 --- recml/core/ops/binary_cross_entropy_ops.py | 20 +- .../core/ops/binary_cross_entropy_ops_test.py | 78 +++++++ .../ops/binary_focal_cross_entropy_ops.py | 12 +- .../binary_focal_cross_entropy_ops_test.py | 64 ++++++ recml/core/utils/distribution_utils.py | 155 ++++++++++++++ recml/core/utils/distribution_utils_test.py | 195 +++++++++++++++++- 6 files changed, 508 insertions(+), 16 deletions(-) diff --git a/recml/core/ops/binary_cross_entropy_ops.py b/recml/core/ops/binary_cross_entropy_ops.py index eb12af0..ad1823e 100644 --- a/recml/core/ops/binary_cross_entropy_ops.py +++ b/recml/core/ops/binary_cross_entropy_ops.py @@ -158,6 +158,13 @@ def _get_sharding(x: jt.ArrayLike) -> jax.sharding.Sharding | None: return None +def _has_non_manual_axes( + mesh: jax.sharding.Mesh | jax.sharding.AbstractMesh, +) -> bool: + """Returns whether `mesh` has at least one non-manual axis.""" + return not mesh.empty and len(mesh.manual_axes) < len(mesh.axis_names) + + def _token_shard_count( x: jt.Float[jt.Array, "... D"], mesh: jax.sharding.Mesh | jax.sharding.AbstractMesh | None = None, @@ -171,6 +178,8 @@ def _token_shard_count( mesh, spec = sharding.mesh, sharding.spec if mesh.empty: return None + if not _has_non_manual_axes(mesh): + return 1 # A PartitionSpec may be shorter than `x.ndim`; trailing axes are unsharded. factor = 1 @@ -207,7 +216,7 @@ def _replicate_hidden_dim( sharding = _get_sharding(x) if isinstance(sharding, jax.sharding.NamedSharding): mesh = sharding.mesh - if not mesh.empty: + if _has_non_manual_axes(mesh): spec = sharding.spec new_spec_list = list(spec) + [None] * (x.ndim - len(spec)) if new_spec_list: @@ -706,10 +715,9 @@ def cut_binary_cross_entropy( # Prevent collective communications inside the loop by forcing replication of # weights. sharding = _get_sharding(embeddings) - if ( - isinstance(sharding, jax.sharding.NamedSharding) - and not sharding.mesh.empty - ): + if isinstance( + sharding, jax.sharding.NamedSharding + ) and _has_non_manual_axes(sharding.mesh): replicated_sharding = jax.sharding.NamedSharding( sharding.mesh, jax.sharding.PartitionSpec() ) @@ -1595,7 +1603,7 @@ def _bce_bwd_sharded( emb_sharding = _get_sharding(embeddings) is_sharded = ( isinstance(act_sharding, jax.sharding.NamedSharding) - and not act_sharding.mesh.empty + and _has_non_manual_axes(act_sharding.mesh) and isinstance(emb_sharding, jax.sharding.NamedSharding) ) if is_sharded: diff --git a/recml/core/ops/binary_cross_entropy_ops_test.py b/recml/core/ops/binary_cross_entropy_ops_test.py index 8e735aa..b285ab8 100644 --- a/recml/core/ops/binary_cross_entropy_ops_test.py +++ b/recml/core/ops/binary_cross_entropy_ops_test.py @@ -264,6 +264,24 @@ def test_local_tokens_with_real_mesh(self): binary_cross_entropy_ops._local_tokens(activations), 8 * 64 ) + def test_local_tokens_inside_shard_map(self): + num_devices = jax.device_count() + mesh = jax.sharding.Mesh(np.array(jax.devices()), ('batch',)) + local_counts = [] + + @jax.jit + def f(x): + def per_shard(local_x): + local_counts.append(binary_cross_entropy_ops._local_tokens(local_x)) + return local_x + + return jax.shard_map( + per_shard, mesh=mesh, in_specs=P('batch'), out_specs=P('batch') + )(x) + + f(jnp.ones((2 * num_devices, 64, 32))) + self.assertEqual(local_counts[-1], 2 * 64) + # --- `block_v` selection: the memory budget itself. ------------------------ def test_auto_block_v_is_clamped_to_vocab_size(self): @@ -1210,6 +1228,66 @@ def loss_fn(act, emb, force_bwd): g_emb_forced, g_emb_clamped, atol=1e-5, rtol=1e-5 ) + @parameterized.named_parameters( + ('pure_jax', False), + ('pallas', True), + ) + def test_cut_bce_inside_shard_map_jit_grad(self, use_pallas): + num_devices = jax.device_count() + batch, seq_len, hidden_dim = 2 * num_devices, 16, 128 + vocab_size, block_v, num_labels = 256, 128, 3 + key_act, key_emb, key_tgt = jax.random.split(jax.random.PRNGKey(17), 3) + activations = jax.random.normal(key_act, (batch, seq_len, hidden_dim)) + embeddings = jax.random.normal(key_emb, (vocab_size, hidden_dim)) + targets = jax.random.randint( + key_tgt, (batch, seq_len, num_labels), 0, vocab_size + ) + mesh = jax.sharding.Mesh(np.array(jax.devices()), ('batch',)) + + def sharded_loss(act, emb): + tiled_emb = jnp.tile(emb, (num_devices, 1)) + + def per_shard(local_act, local_emb, local_tgt): + _, per_pos = binary_cross_entropy_ops.cut_binary_cross_entropy( + local_act, + local_emb, + local_tgt, + block_v=block_v, + use_pallas=use_pallas, + return_per_target_losses=True, + ) + return per_pos + + losses = jax.shard_map( + per_shard, + mesh=mesh, + in_specs=(P('batch'), P('batch'), P('batch')), + out_specs=P('batch'), + check_vma=False, + )(act, tiled_emb, targets) + return jnp.mean(losses) + + def ref_loss(act, emb): + _, per_pos = binary_cross_entropy_ops.cut_binary_cross_entropy( + act, + emb, + targets, + block_v=block_v, + use_pallas=use_pallas, + return_per_target_losses=True, + ) + return jnp.mean(per_pos) + + loss_got, (g_act_got, g_emb_got) = jax.jit( + jax.value_and_grad(sharded_loss, argnums=(0, 1)) + )(activations, embeddings) + loss_ref, (g_act_ref, g_emb_ref) = jax.jit( + jax.value_and_grad(ref_loss, argnums=(0, 1)) + )(activations, embeddings) + np.testing.assert_allclose(loss_got, loss_ref, atol=1e-5, rtol=1e-5) + np.testing.assert_allclose(g_act_got, g_act_ref, atol=1e-5, rtol=1e-5) + np.testing.assert_allclose(g_emb_got, g_emb_ref, atol=1e-5, rtol=1e-5) + if __name__ == '__main__': absltest.main() diff --git a/recml/core/ops/binary_focal_cross_entropy_ops.py b/recml/core/ops/binary_focal_cross_entropy_ops.py index 311af8e..572a067 100644 --- a/recml/core/ops/binary_focal_cross_entropy_ops.py +++ b/recml/core/ops/binary_focal_cross_entropy_ops.py @@ -43,6 +43,9 @@ bce_ops._effective_block_v # pylint: disable=protected-access ) _get_sharding = bce_ops._get_sharding # pylint: disable=protected-access +_has_non_manual_axes = ( + bce_ops._has_non_manual_axes # pylint: disable=protected-access +) _max_safe_block_v = bce_ops._max_safe_block_v # pylint: disable=protected-access _max_safe_chunk_n = bce_ops._max_safe_chunk_n # pylint: disable=protected-access _pallas_interpret = bce_ops._pallas_interpret # pylint: disable=protected-access @@ -522,10 +525,9 @@ def cut_binary_focal_cross_entropy( ) sharding = _get_sharding(embeddings) - if ( - isinstance(sharding, jax.sharding.NamedSharding) - and not sharding.mesh.empty - ): + if isinstance( + sharding, jax.sharding.NamedSharding + ) and _has_non_manual_axes(sharding.mesh): replicated_sharding = jax.sharding.NamedSharding( sharding.mesh, jax.sharding.PartitionSpec() ) @@ -1164,7 +1166,7 @@ def _focal_bce_bwd_sharded( emb_sharding = _get_sharding(embeddings) is_sharded = ( isinstance(act_sharding, jax.sharding.NamedSharding) - and not act_sharding.mesh.empty + and _has_non_manual_axes(act_sharding.mesh) and isinstance(emb_sharding, jax.sharding.NamedSharding) ) if is_sharded: diff --git a/recml/core/ops/binary_focal_cross_entropy_ops_test.py b/recml/core/ops/binary_focal_cross_entropy_ops_test.py index b04a9a8..5b2dd92 100644 --- a/recml/core/ops/binary_focal_cross_entropy_ops_test.py +++ b/recml/core/ops/binary_focal_cross_entropy_ops_test.py @@ -1160,6 +1160,70 @@ def loss_fn(act, emb, force_bwd): g_emb_forced, g_emb_clamped, atol=1e-4, rtol=1e-4 ) + @parameterized.named_parameters( + ('pure_jax', False), + ('pallas', True), + ) + def test_cut_focal_bce_inside_shard_map_jit_grad(self, use_pallas): + num_devices = jax.device_count() + batch, seq_len, hidden_dim = 2 * num_devices, 16, 128 + vocab_size, block_v, num_labels = 256, 128, 3 + key_act, key_emb, key_tgt = jax.random.split(jax.random.PRNGKey(17), 3) + activations = jax.random.normal(key_act, (batch, seq_len, hidden_dim)) + embeddings = jax.random.normal(key_emb, (vocab_size, hidden_dim)) + targets = jax.random.randint( + key_tgt, (batch, seq_len, num_labels), 0, vocab_size + ) + mesh = jax.sharding.Mesh(np.array(jax.devices()), ('batch',)) + + def sharded_loss(act, emb): + tiled_emb = jnp.tile(emb, (num_devices, 1)) + + def per_shard(local_act, local_emb, local_tgt): + _, per_pos = ( + binary_focal_cross_entropy_ops.cut_binary_focal_cross_entropy( + local_act, + local_emb, + local_tgt, + block_v=block_v, + use_pallas=use_pallas, + return_per_target_losses=True, + ) + ) + return per_pos + + losses = jax.shard_map( + per_shard, + mesh=mesh, + in_specs=(P('batch'), P('batch'), P('batch')), + out_specs=P('batch'), + check_vma=False, + )(act, tiled_emb, targets) + return jnp.mean(losses) + + def ref_loss(act, emb): + _, per_pos = ( + binary_focal_cross_entropy_ops.cut_binary_focal_cross_entropy( + act, + emb, + targets, + block_v=block_v, + use_pallas=use_pallas, + return_per_target_losses=True, + ) + ) + return jnp.mean(per_pos) + + loss_got, (g_act_got, g_emb_got) = jax.jit( + jax.value_and_grad(sharded_loss, argnums=(0, 1)) + )(activations, embeddings) + loss_ref, (g_act_ref, g_emb_ref) = jax.jit( + jax.value_and_grad(ref_loss, argnums=(0, 1)) + )(activations, embeddings) + np.testing.assert_allclose(loss_got, loss_ref, atol=1e-5, rtol=1e-5) + np.testing.assert_allclose(g_act_got, g_act_ref, atol=1e-5, rtol=1e-5) + np.testing.assert_allclose(g_emb_got, g_emb_ref, atol=1e-5, rtol=1e-5) + if __name__ == '__main__': absltest.main() diff --git a/recml/core/utils/distribution_utils.py b/recml/core/utils/distribution_utils.py index e079588..8f6c3a6 100644 --- a/recml/core/utils/distribution_utils.py +++ b/recml/core/utils/distribution_utils.py @@ -13,8 +13,15 @@ # limitations under the License. """Utilities for inspecting the active JAX/Keras distribution mesh.""" +from collections.abc import Callable +import math +from typing import Any +import weakref + import jax +import jax.numpy as jnp import keras +from keras.src.backend.common import variables as keras_variables # Conventional name of the mesh axis used for model/tensor parallelism. Any # other axis of the mesh is assumed to shard the batch. @@ -120,3 +127,151 @@ def resolve_batch_axis( # somewhere; fall back to the full set of axis names. batch_axes = tuple(abstract_mesh.axis_names) return normalize_axis(batch_axes) + + +def fuse_table_grads_enabled() -> bool: + """Returns whether the active distribution fuses table gradients. + + A distribution opts in by setting a truthy `fuse_table_grads` attribute. + Layers that read a large table with a batch (embedding lookups, tied output + projections) then run that read through `fused_table_grad_call`. + + Returns: + True if a distribution is active and its `fuse_table_grads` is truthy. + """ + return bool( + getattr(keras.distribution.distribution(), 'fuse_table_grads', False) + ) + + +class FusedTableGradDataParallel(keras.distribution.DataParallel): + """`DataParallel` that fuses gradients across reads of each replicated table. + + Under plain `DataParallel`, XLA partitions every embedding lookup's gradient + (a scatter-add into a replicated table) as a local scatter plus an + all-reduce, and it seeds the next scatter of the same table with that + all-reduce's result. A table read by N lookups then pays a chain of up to N + table-sized all-reduces per step. With this distribution, every read of a + table runs under `jax.shard_map` on a shared per-device copy of the table. + Each device adds up its reads' table gradients locally, and one all-reduce + per table sums the result across devices. The forward computation is + unchanged; only the order in which the table gradient is summed changes. + """ + + fuse_table_grads = True + + +# Shared per-device copies of replicated tables. Maps id(table) to a weak +# reference to the table, the table's value in the trace that read it last, and +# the copies of that value. See `_shared_table_copies`. +_shared_table_copies_cache: dict[ + int, tuple[weakref.ref[keras.Variable], jax.Array, jax.Array] +] = {} + + +def _shared_table_copies( + table: keras.Variable, table_value: jax.Array, num_copies: int +) -> jax.Array: + """Returns `num_copies` copies of `table_value` stacked on dim 0. + + Every call on `table` with the same traced `table_value` gets the same + copies. JAX then adds up the table gradients of all those calls on each + device before it reaches the transpose of the tile, which sums the copies + across devices once. + + The copies are a tracer, so all calls on a table must run at the same trace + level. If the first call runs inside a nested transformation such as + `jax.checkpoint` and a later one outside it, JAX raises an escaped-tracer + error. + + Args: + table: The table variable. + table_value: The table's value in the current trace. + num_copies: The number of devices that share the table. + + Returns: + An array of shape `(num_copies * table.shape[0], *table.shape[1:])`. + """ + key = id(table) + cached = _shared_table_copies_cache.get(key) + if cached is not None and cached[0]() is table and cached[1] is table_value: + return cached[2] + copies = jnp.tile(table_value, (num_copies,) + (1,) * (table_value.ndim - 1)) + _shared_table_copies_cache[key] = ( + # Drops the entry, and the tracers in it, when the table goes away. + weakref.ref(table, lambda _: _shared_table_copies_cache.pop(key, None)), + table_value, + copies, + ) + return copies + + +def fused_table_grad_call( + fn: Callable[..., Any], + table: keras.Variable, + *batch_args: Any, +) -> Any: + """Calls `fn(table, *batch_args)` on each device and fuses table gradients. + + The table must be replicated. Each device gets its own copy, and all calls on + the table in the same trace share that copy. So each device adds up the table + gradients of all calls, and one all-reduce per table sums them across + devices. All calls on a table must run at the same trace level (see + `_shared_table_copies`). + + Under a Keras autocast scope, `fn` gets the table in the scope's dtype, but + the gradient is still all-reduced in the variable dtype. + + Args: + fn: A function of the full table and the local batch shards. Each output row + must depend only on the same row of the batch args. + table: A replicated table variable. + *batch_args: Arrays whose dim 0 is the global batch. + + Returns: + The output of `fn`, sharded on dim 0 over the batch axes. + + Raises: + NotImplementedError: If the active distribution shards the table. + """ + # TODO(jinchi): Support sharded (FSDP) tables. + if any(keras.distribution.distribution().get_variable_layout(table).axes): + raise NotImplementedError( + 'fused_table_grad_call supports only replicated tables.' + ) + abstract_mesh = get_abstract_mesh() + batch_axis = resolve_batch_axis(abstract_mesh) + + # Read the table in its variable dtype and cast inside `fn`, so the gradient + # is all-reduced in that dtype rather than in bf16. + autocast_scope = keras_variables.get_autocast_scope() + with keras_variables.AutocastScope(None): + table_value = keras.ops.convert_to_tensor(table) + + if abstract_mesh is not None and isinstance(table_value, jax.core.Tracer): + mesh = abstract_mesh + batch_axes = batch_axis if isinstance(batch_axis, tuple) else (batch_axis,) + axis_sizes = dict(zip(abstract_mesh.axis_names, abstract_mesh.axis_sizes)) + table_value = _shared_table_copies( + table, table_value, math.prod(axis_sizes[axis] for axis in batch_axes) + ) + table_spec = jax.P(batch_axis) + else: + # Eager `shard_map` on an `AbstractMesh` rejects inputs without a + # `NamedSharding`, so eager calls use the concrete mesh and one table. + mesh = keras.distribution.distribution().device_mesh.backend_mesh + table_spec = jax.P() + + def local_fn(full_table, *local_batch_args): + if autocast_scope is not None: + full_table = autocast_scope.maybe_cast(full_table) + return fn(full_table, *local_batch_args) + + # The transpose of the tile sums the copies' gradients across devices. + return jax.shard_map( + local_fn, + mesh=mesh, + in_specs=(table_spec,) + (jax.P(batch_axis),) * len(batch_args), + out_specs=jax.P(batch_axis), + check_vma=False, + )(table_value, *batch_args) diff --git a/recml/core/utils/distribution_utils_test.py b/recml/core/utils/distribution_utils_test.py index ce622a5..c9e5e63 100644 --- a/recml/core/utils/distribution_utils_test.py +++ b/recml/core/utils/distribution_utils_test.py @@ -13,14 +13,22 @@ # limitations under the License. """Tests mesh and batch-axis resolution in distribution_utils.""" +import re + from absl.testing import absltest from absl.testing import parameterized import jax +import jax.numpy as jnp import keras from keras.src.distribution import distribution_lib +import numpy as np from recml.core.utils import distribution_utils -_DEVICES = [f'cpu:{i}' for i in range(8)] +# Virtual CPU devices for the test meshes. Must be even for the 2D meshes. +_NUM_DEVICES = 8 +# Must run before JAX starts. +jax.config.update('jax_num_cpu_devices', _NUM_DEVICES) +_DEVICES = [f'cpu:{i}' for i in range(_NUM_DEVICES)] class _StubDistribution: @@ -45,6 +53,16 @@ def _abstract_mesh( ) +def _all_reduced_elements(hlo_text: str, shape: str) -> int: + """Counts all-reduced arrays of `shape` (e.g. 'f32[24,8]') in HLO text.""" + count = 0 + for line in hlo_text.splitlines(): + match = re.search(r'=\s*(.*?)\s+all-reduce(?:-start)?\(', line) + if match: + count += match.group(1).count(shape) + return count + + class DistributionUtilsTest(parameterized.TestCase): def setUp(self): @@ -83,7 +101,7 @@ def test_resolve_batch_axis_without_mesh_or_distribution(self): def test_resolve_batch_axis_prefers_distribution_over_mesh(self): device_mesh = distribution_lib.DeviceMesh( - (2, 4), ('replica', 'fsdp'), _DEVICES + (2, _NUM_DEVICES // 2), ('replica', 'fsdp'), _DEVICES ) keras.distribution.set_distribution( _StubDistribution(device_mesh, ('replica', 'fsdp')) @@ -95,7 +113,9 @@ def test_resolve_batch_axis_prefers_distribution_over_mesh(self): ) def test_resolve_batch_axis_from_data_parallel_distribution(self): - device_mesh = distribution_lib.DeviceMesh((8,), ('batch',), _DEVICES) + device_mesh = distribution_lib.DeviceMesh( + (_NUM_DEVICES,), ('batch',), _DEVICES + ) keras.distribution.set_distribution( keras.distribution.DataParallel(device_mesh=device_mesh) ) @@ -116,7 +136,7 @@ def test_get_abstract_mesh_from_global_jax_mesh(self): def test_get_abstract_mesh_from_distribution(self): device_mesh = distribution_lib.DeviceMesh( - (2, 4), ('replica', 'fsdp'), _DEVICES + (2, _NUM_DEVICES // 2), ('replica', 'fsdp'), _DEVICES ) keras.distribution.set_distribution(_StubDistribution(device_mesh, None)) @@ -124,7 +144,7 @@ def test_get_abstract_mesh_from_distribution(self): self.assertIsNotNone(mesh) self.assertEqual(mesh.axis_names, ('replica', 'fsdp')) - self.assertEqual(mesh.axis_sizes, (2, 4)) + self.assertEqual(mesh.axis_sizes, (2, _NUM_DEVICES // 2)) # Regression: `axis_types` must have one entry per axis. Passing a bare # `AxisType` raises `ValueError` on a mesh with more than one axis. self.assertEqual( @@ -133,5 +153,170 @@ def test_get_abstract_mesh_from_distribution(self): ) +# TODO(jinchi): Add sharded-table (FSDP) tests once fused table gradients +# support FSDP. +class FusedTableGradTest(parameterized.TestCase): + + def setUp(self): + super().setUp() + keras.distribution.set_distribution(None) + self.addCleanup(keras.distribution.set_distribution, None) + + def test_fuse_table_grads_disabled_without_distribution(self): + self.assertFalse(distribution_utils.fuse_table_grads_enabled()) + + def test_fused_table_grad_data_parallel_opts_in(self): + device_mesh = distribution_lib.DeviceMesh( + (_NUM_DEVICES,), ('batch',), _DEVICES + ) + keras.distribution.set_distribution( + distribution_utils.FusedTableGradDataParallel(device_mesh=device_mesh) + ) + self.assertTrue(distribution_utils.fuse_table_grads_enabled()) + + def test_fused_table_grad_call_under_data_parallel(self): + """Several reads of one replicated table, as in GUM.""" + device_mesh = distribution_lib.DeviceMesh( + (_NUM_DEVICES,), ('batch',), _DEVICES + ) + keras.distribution.set_distribution( + distribution_utils.FusedTableGradDataParallel(device_mesh=device_mesh) + ) + vocab, dim, batch, tokens = 16, 8, _NUM_DEVICES, 3 + rng = np.random.default_rng(0) + table = keras.Variable( + rng.normal(size=(vocab, dim)).astype(np.float32), name='table' + ) + ids = [ + rng.integers(0, vocab, size=(batch, tokens), dtype=np.int32) + for _ in range(3) + ] + h = rng.normal(size=(batch, dim)).astype(np.float32) + out_weights = rng.normal(size=(batch, 3 * dim + vocab)).astype(np.float32) + + def lookup(full_table, ids): + return jnp.take(full_table, ids, axis=0).sum(axis=1) + + def tied_logits(full_table, h): + return h @ full_table.T + + def loss_fn(table_value, h, fuse_grads): + with keras.StatelessScope(state_mapping=[(table, table_value)]): + outs = [] + for i in ids: + if fuse_grads: + outs.append( + distribution_utils.fused_table_grad_call(lookup, table, i) + ) + else: + outs.append(lookup(keras.ops.convert_to_tensor(table), i)) + if fuse_grads: + outs.append( + distribution_utils.fused_table_grad_call(tied_logits, table, h) + ) + else: + outs.append(tied_logits(keras.ops.convert_to_tensor(table), h)) + return jnp.sum(jnp.concatenate(outs, axis=-1) * out_weights) + + grad_fn = jax.jit( + jax.value_and_grad(loss_fn, argnums=(0, 1)), static_argnums=2 + ) + expected_loss, expected_grads = grad_fn(table.value, h, False) + loss, grads = grad_fn(table.value, h, True) + + np.testing.assert_allclose(loss, expected_loss, rtol=1e-5) + for grad, expected_grad in zip(grads, expected_grads): + np.testing.assert_allclose(grad, expected_grad, rtol=1e-5, atol=1e-5) + + def test_fused_table_grad_call_reduces_replicated_table_once(self): + """All reads of a replicated table share one gradient all-reduce.""" + device_mesh = distribution_lib.DeviceMesh( + (_NUM_DEVICES,), ('batch',), _DEVICES + ) + keras.distribution.set_distribution( + distribution_utils.FusedTableGradDataParallel(device_mesh=device_mesh) + ) + vocab, dim, batch, tokens = 24, 8, _NUM_DEVICES, 3 + rng = np.random.default_rng(0) + table = keras.Variable( + rng.normal(size=(vocab, dim)).astype(np.float32), name='table' + ) + ids = [ + rng.integers(0, vocab, size=(batch, tokens), dtype=np.int32) + for _ in range(3) + ] + h = rng.normal(size=(batch, dim)).astype(np.float32) + + def lookup(full_table, ids): + return jnp.take(full_table, ids, axis=0).sum(axis=1) + + def tied_logits(full_table, h): + return h @ full_table.T + + def loss_fn(table_value, h): + with keras.StatelessScope(state_mapping=[(table, table_value)]): + outs = [ + distribution_utils.fused_table_grad_call(lookup, table, i) + for i in ids + ] + outs.append( + distribution_utils.fused_table_grad_call(tied_logits, table, h) + ) + return sum(jnp.sum(out) for out in outs) + + hlo = ( + jax.jit(jax.grad(loss_fn, argnums=(0, 1))) + .lower(table.value, h) + .compile() + .as_text() + ) + + # 4 calls read the table, but its gradient is all-reduced once. + self.assertEqual(_all_reduced_elements(hlo, f'f32[{vocab},{dim}]'), 1) + + def test_fused_table_grad_call_in_separate_traces(self): + """Each trace makes its own table copies.""" + device_mesh = distribution_lib.DeviceMesh( + (_NUM_DEVICES,), ('batch',), _DEVICES + ) + keras.distribution.set_distribution( + distribution_utils.FusedTableGradDataParallel(device_mesh=device_mesh) + ) + vocab, dim, batch = 16, 8, _NUM_DEVICES + rng = np.random.default_rng(0) + table = keras.Variable( + rng.normal(size=(vocab, dim)).astype(np.float32), name='table' + ) + h = rng.normal(size=(batch, dim)).astype(np.float32) + + def tied_logits(full_table, h): + return h @ full_table.T + + def loss_fn(table_value, h): + with keras.StatelessScope(state_mapping=[(table, table_value)]): + logits = distribution_utils.fused_table_grad_call(tied_logits, table, h) + return jnp.sum(jnp.square(logits)) + + # Reusing the first trace's copies would raise an escaped-tracer error. + loss = jax.jit(loss_fn)(table.value, h) + table_grad, _ = jax.jit(jax.grad(loss_fn, argnums=(0, 1)))(table.value, h) + + logits = h @ np.asarray(table.value).T + np.testing.assert_allclose(loss, np.sum(np.square(logits)), rtol=1e-5) + np.testing.assert_allclose( + table_grad, 2 * logits.T @ h, rtol=1e-4, atol=1e-4 + ) + + @parameterized.named_parameters(('off', False), ('on', True)) + def test_fuse_table_grads_reads_distribution(self, value): + device_mesh = distribution_lib.DeviceMesh( + (_NUM_DEVICES,), ('batch',), _DEVICES + ) + distribution = keras.distribution.DataParallel(device_mesh=device_mesh) + distribution.fuse_table_grads = value + keras.distribution.set_distribution(distribution) + self.assertEqual(distribution_utils.fuse_table_grads_enabled(), value) + + if __name__ == '__main__': absltest.main()