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()