Internal change

recml/core/ops/binary_cross_entropy_ops.py

@@ -249,6 +249,9 @@ class BCEConfig:
       ``embeddings.shape[0]``.
     optimize_large_l: Whether to use the target-free vocabulary pass + sparse
       positive-target correction. Defaults to False.
+    force_bwd_block_v: Whether to force the backward pass to use `block_v`
+      directly (aligned to lane width and capped by vocabulary size) instead of
+      clamping to `_max_safe_block_v`. Defaults to False.
   """
 
   block_v: int
@@ -257,6 +260,7 @@ class BCEConfig:
   use_pallas: bool = True
   global_vocab: int | None = None
   optimize_large_l: bool = False
+  force_bwd_block_v: bool = False
 
   # Sharding specs for VJP backward pass optimization
   mesh: jax.sharding.Mesh | None = None
@@ -595,30 +599,32 @@ def cut_binary_cross_entropy(
     *,
     return_per_target_losses: bool = False,
     return_metrics: bool = False,
+    return_per_target_metrics: bool = False,
     block_v: int | None = None,
     block_n: int | None = None,
     mesh: jax.sharding.Mesh | None = None,
     act_spec: jax.sharding.PartitionSpec | None = None,
     emb_spec: jax.sharding.PartitionSpec | None = None,
     use_pallas: bool = True,
     optimize_large_l: bool = False,
+    force_bwd_block_v: bool = False,
 ) -> (
     jt.Float[jt.Array, ""]
     | tuple[jt.Float[jt.Array, ""], jt.Float[jt.Array, "... B N"]]
     | tuple[
         jt.Float[jt.Array, ""],
-        jt.Float[jt.Array, ""],
-        jt.Float[jt.Array, ""],
-        jt.Float[jt.Array, ""],
-        jt.Float[jt.Array, ""],
+        jt.Float[jt.Array, "... B N"] | jt.Float[jt.Array, ""],
+        jt.Float[jt.Array, "... B N"] | jt.Float[jt.Array, ""],
+        jt.Float[jt.Array, "... B N"] | jt.Float[jt.Array, ""],
+        jt.Float[jt.Array, "... B N"] | jt.Float[jt.Array, ""],
     ]
     | tuple[
         jt.Float[jt.Array, ""],
         jt.Float[jt.Array, "... B N"],
-        jt.Float[jt.Array, ""],
-        jt.Float[jt.Array, ""],
-        jt.Float[jt.Array, ""],
-        jt.Float[jt.Array, ""],
+        jt.Float[jt.Array, "... B N"] | jt.Float[jt.Array, ""],
+        jt.Float[jt.Array, "... B N"] | jt.Float[jt.Array, ""],
+        jt.Float[jt.Array, "... B N"] | jt.Float[jt.Array, ""],
+        jt.Float[jt.Array, "... B N"] | jt.Float[jt.Array, ""],
     ]
 ):
   """Computes binary cross entropy loss over unmaterialized logits.
@@ -643,6 +649,8 @@ def cut_binary_cross_entropy(
       position's weight.
     return_per_target_losses: If True, also return the per-target loss tensor.
     return_metrics: If True, also return TP, FP, FN, TN metric counts.
+    return_per_target_metrics: If True (with ``return_metrics=True``), return
+      per-position TP, FP, FN, TN tensors instead of summing them to scalars.
     block_v: Vocab-axis block size. Auto-picked if omitted.
     block_n: Sequence-axis block size for backward kernel.
     mesh: Optional mesh to use for sharding.
@@ -652,9 +660,21 @@ def cut_binary_cross_entropy(
       implementation.
     optimize_large_l: If True, use the target-free vocabulary pass + sparse
       positive-target correction. Defaults to False.
+    force_bwd_block_v: If True, force the backward pass to use `block_v`
+      directly instead of clamping to `_max_safe_block_v`. Defaults to False.
 
   Returns:
-    Scalar loss, optionally paired with per-target losses and/or metrics.
+    - ``loss`` if ``return_per_target_losses=False`` and
+      ``return_metrics=False``.
+    - ``(loss, losses)`` if ``return_per_target_losses=True`` and
+      ``return_metrics=False``.
+    - ``(loss, tp, fp, fn, tn)`` if ``return_per_target_losses=False`` and
+      ``return_metrics=True``.
+    - ``(loss, losses, tp, fp, fn, tn)`` if ``return_per_target_losses=True``
+      and ``return_metrics=True``.
+    Here ``loss`` is a scalar ``[]``, ``losses`` has shape ``[..., B, N]``, and
+    ``(tp, fp, fn, tn)`` each have shape ``[..., B, N]`` when
+    ``return_per_target_metrics=True`` or scalar shape ``[]`` when ``False``.
 
   Raises:
     ValueError: If ``activations`` has rank 5 or above, or if ``weights`` does
@@ -715,6 +735,7 @@ def cut_binary_cross_entropy(
       use_pallas=use_pallas,
       global_vocab=vocab_size,
       optimize_large_l=optimize_large_l,
+      force_bwd_block_v=force_bwd_block_v,
   )
 
   res = _cut_binary_cross_entropy(
@@ -730,22 +751,24 @@ def cut_binary_cross_entropy(
     if weights is not None:
       losses = losses * weights
       weight_sum = jnp.sum(weights)
-      tp_sum = jnp.sum(tp * weights)
-      fp_sum = jnp.sum(fp * weights)
-      fn_sum = jnp.sum(fn * weights)
-      tn_sum = jnp.sum(tn * weights)
+      tp = tp * weights
+      fp = fp * weights
+      fn = fn * weights
+      tn = tn * weights
     else:
       weight_sum = np.prod(losses_shape)
-      tp_sum = jnp.sum(tp)
-      fp_sum = jnp.sum(fp)
-      fn_sum = jnp.sum(fn)
-      tn_sum = jnp.sum(tn)
 
     loss = jnp.sum(losses) / (weight_sum + EPS)
 
+    if not return_per_target_metrics:
+      tp = jnp.sum(tp)
+      fp = jnp.sum(fp)
+      fn = jnp.sum(fn)
+      tn = jnp.sum(tn)
+
     if return_per_target_losses:
-      return loss, losses, tp_sum, fp_sum, fn_sum, tn_sum
-    return loss, tp_sum, fp_sum, fn_sum, tn_sum
+      return loss, losses, tp, fp, fn, tn
+    return loss, tp, fp, fn, tn
   else:
     losses = res
 
@@ -871,6 +894,9 @@ def _effective_block_v(
   """
   lane = _pallas_lane()
   padded_d = pl.align_to(hidden, lane)
+  if config.force_bwd_block_v:
+    block_v = pl.align_to(min(config.block_v, vocab), lane)
+    return padded_d, min(block_v, pl.align_to(vocab, lane))
   max_safe = _max_safe_block_v(_pallas_vmem_budget(), padded_d, config.block_n)
   block_v = pl.align_to(min(config.block_v, vocab, max_safe), lane)
   return padded_d, min(block_v, pl.align_to(vocab, lane), max_safe)

recml/core/ops/binary_cross_entropy_ops_benchmark_test.py

@@ -96,7 +96,7 @@ def test_mini_benchmark_cut_bce(self):
 
     def run_cut(act, emb, bv=_BLOCK_V):
       return binary_cross_entropy_ops.cut_binary_cross_entropy(
-          act, emb, targets, block_v=bv
+          act, emb, targets, block_v=bv, optimize_large_l=True
       )
 
     grad_cut = jax.jit(jax.grad(run_cut, argnums=(0, 1)))

recml/core/ops/binary_cross_entropy_ops_test.py

@@ -854,6 +854,72 @@ def test_cut_bce_metrics(self):
     self.assertIsNotNone(fn)
     self.assertIsNotNone(tn)
 
+  @parameterized.named_parameters(
+      ('unweighted_no_per_target_losses', False, False),
+      ('weighted_no_per_target_losses', True, False),
+      ('weighted_with_per_target_losses', True, True),
+  )
+  def test_cut_bce_per_target_metrics(
+      self, use_weights: bool, return_per_target_losses: bool
+  ):
+    batch, seq_len, hidden_dim, vocab_size, num_labels = 2, 16, 32, 128, 2
+    block_v = 64
+    key_act, key_emb, key_tgt, key_w = jax.random.split(
+        jax.random.PRNGKey(17), 4
+    )
+    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
+    )
+    weights = None
+    if use_weights:
+      weights = jax.random.bernoulli(key_w, 0.7, (batch, seq_len)).astype(
+          jnp.float32
+      )
+
+    res = binary_cross_entropy_ops.cut_binary_cross_entropy(
+        activations,
+        embeddings,
+        targets,
+        weights,
+        block_v=block_v,
+        return_metrics=True,
+        return_per_target_metrics=True,
+        return_per_target_losses=return_per_target_losses,
+        use_pallas=False,
+    )
+    if return_per_target_losses:
+      loss, losses, tp, fp, fn, tn = res
+    else:
+      loss, tp, fp, fn, tn = res
+      losses = None
+
+    loss_ref, losses_ref = _keras_bce(
+        activations, embeddings, targets, weights=weights
+    )
+    np.testing.assert_allclose(loss, loss_ref, atol=1e-5, rtol=1e-5)
+    if losses is not None:
+      np.testing.assert_allclose(losses, losses_ref, atol=1e-5, rtol=1e-5)
+
+    logits = jnp.matmul(activations, embeddings.T)
+    multi_hot = jnp.max(jax.nn.one_hot(targets, vocab_size), axis=-2) > 0
+    preds = logits > 0.0
+    tp_ref = jnp.sum(multi_hot & preds, axis=-1, dtype=jnp.float32)
+    fp_ref = jnp.sum(~multi_hot & preds, axis=-1, dtype=jnp.float32)
+    fn_ref = jnp.sum(multi_hot & ~preds, axis=-1, dtype=jnp.float32)
+    tn_ref = jnp.sum(~multi_hot & ~preds, axis=-1, dtype=jnp.float32)
+    if weights is not None:
+      tp_ref = tp_ref * weights
+      fp_ref = fp_ref * weights
+      fn_ref = fn_ref * weights
+      tn_ref = tn_ref * weights
+
+    np.testing.assert_allclose(tp, tp_ref, rtol=1e-6)
+    np.testing.assert_allclose(fp, fp_ref, rtol=1e-6)
+    np.testing.assert_allclose(fn, fn_ref, rtol=1e-6)
+    np.testing.assert_allclose(tn, tn_ref, rtol=1e-6)
+
   def test_cut_bce_block_v_capped_at_vocab(self):
     activations = jnp.ones((2, 64, 32))
     embeddings = jnp.ones((200, 32))
@@ -1088,6 +1154,62 @@ def loss_fn(act, emb, opt):
     np.testing.assert_allclose(g_act_opt, g_act_ref, atol=1e-5, rtol=1e-5)
     np.testing.assert_allclose(g_emb_opt, g_emb_ref, atol=1e-5, rtol=1e-5)
 
+  def test_force_bwd_block_v(self):
+    config_clamped = binary_cross_entropy_ops.BCEConfig(
+        block_v=11008, block_n=256, force_bwd_block_v=False
+    )
+    config_forced = binary_cross_entropy_ops.BCEConfig(
+        block_v=11008, block_n=256, force_bwd_block_v=True
+    )
+    _, bwd_v_clamped = binary_cross_entropy_ops._effective_block_v(
+        config_clamped, hidden=256, vocab=16384
+    )
+    _, bwd_v_forced = binary_cross_entropy_ops._effective_block_v(
+        config_forced, hidden=256, vocab=16384
+    )
+    self.assertLess(bwd_v_clamped, 11008)
+    self.assertEqual(bwd_v_forced, 11008)
+
+    batch, seq_len, hidden_dim, vocab_size, num_labels = 2, 16, 128, 2048, 4
+    key_act, key_emb, key_tgt = jax.random.split(jax.random.PRNGKey(99), 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
+    )
+
+    def loss_fn(act, emb, force_bwd):
+      return binary_cross_entropy_ops.cut_binary_cross_entropy(
+          act,
+          emb,
+          targets,
+          block_v=1024,
+          block_n=128,
+          force_bwd_block_v=force_bwd,
+      )
+
+    with mock.patch.object(
+        binary_cross_entropy_ops, '_max_safe_block_v', return_value=512
+    ):
+      loss_clamped = loss_fn(activations, embeddings, False)
+      loss_forced = loss_fn(activations, embeddings, True)
+      np.testing.assert_allclose(
+          loss_forced, loss_clamped, atol=1e-6, rtol=1e-6
+      )
+
+      g_act_clamped, g_emb_clamped = jax.jit(
+          jax.grad(lambda a, e: loss_fn(a, e, False), argnums=(0, 1))
+      )(activations, embeddings)
+      g_act_forced, g_emb_forced = jax.jit(
+          jax.grad(lambda a, e: loss_fn(a, e, True), argnums=(0, 1))
+      )(activations, embeddings)
+    np.testing.assert_allclose(
+        g_act_forced, g_act_clamped, atol=1e-5, rtol=1e-5
+    )
+    np.testing.assert_allclose(
+        g_emb_forced, g_emb_clamped, atol=1e-5, rtol=1e-5
+    )
+
 
 if __name__ == '__main__':
   absltest.main()