Skip to content

Internal change #220

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
copybara-service wants to merge 1 commit into
base: main
Choose a base branch
from
+508 −16
Open

Internal change #220

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 14 additions & 6 deletions recml/core/ops/binary_cross_entropy_ops.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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()
)
Expand Down Expand Up @@ -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:
Expand Down
78 changes: 78 additions & 0 deletions recml/core/ops/binary_cross_entropy_ops_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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()
12 changes: 7 additions & 5 deletions recml/core/ops/binary_focal_cross_entropy_ops.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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()
)
Expand Down Expand Up @@ -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:
Expand Down
64 changes: 64 additions & 0 deletions recml/core/ops/binary_focal_cross_entropy_ops_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Loading
Loading