Fix vocab tiling hidden-state cotangent scaling

src/maxtext/utils/vocabulary_tiling.py

@@ -253,6 +253,7 @@ def _bwd_scan_body(grad_params_acc, chunk_data):
     grad_reshaped_hidden_states = _maybe_shard_with_name(grad_reshaped_hidden_states, reshaped_hidden_spec)
     # Chain-rule to accumulate gradients
     grad_params = jax.tree_util.tree_map(lambda g: g * loss_cotangent, grad_params)
+    grad_reshaped_hidden_states *= loss_cotangent
     # Cast cotangents back to each primal's dtype; custom_vjp requires dtype match.
     grad_params = jax.tree_util.tree_map(lambda x, y: y.astype(x.dtype), gathered_params, grad_params)
     # Give back sharding constraint
@@ -469,6 +470,7 @@ def _bwd_scan_body(grad_head_acc, chunk_data):
     )
     grad_reshaped_hidden_states = _maybe_shard_with_name(grad_reshaped_hidden_states, reshaped_hidden_spec)
     grad_head = jax.tree_util.tree_map(lambda g: g * loss_cotangent, grad_head)
+    grad_reshaped_hidden_states *= loss_cotangent
     grad_head = jax.tree_util.tree_map(lambda x, y: y.astype(x.dtype), chunk_head_params, grad_head)
     grad_reshaped_hidden_states = _reshape(grad_reshaped_hidden_states, (batch_size, seq_len, emb_dim), hidden_spec)
 

tests/unit/attention_test.py

@@ -3480,7 +3480,9 @@ def test_tpu_dot_product_context_parallel_with_indexer(
           "indexer_topk": 32,
       },
   )
-  @pytest.mark.skip(reason="Indexer with all-gather context parallelism diverges from the dot_product reference; fix tracked in #4947.")
+  @pytest.mark.skip(
+      reason="Indexer with all-gather context parallelism diverges from the dot_product reference; fix tracked in #4947."
+  )
   @pytest.mark.tpu_only
   def test_tpu_flash_attention_context_parallel_with_indexer(
       self, context_parallel_load_balance, ici_context_parallelism=2, indexer_topk=256

tests/unit/tiling_test.py

@@ -317,6 +317,50 @@ def test_vocab_tiling_nnx_loss(self):
         xent_sum_tiled, xent_sum_ref, rtol=self.rtol, atol=self.atol
     ), f"NNX vocab tiling loss {xent_sum_tiled} does not match non-tiled reference {xent_sum_ref}."
 
+  @pytest.mark.cpu_only
+  def test_linen_vocab_tiling_hidden_gradient_respects_outer_loss_scale(self):
+    """The tiled loss's hidden-state gradient obeys the chain rule."""
+    cfg = pyconfig.initialize(
+        self.base_config,
+        run_name="linen_vocab_tiling_outer_loss_scale",
+        enable_checkpointing=False,
+        enable_dropout=False,
+        max_target_length=self.seq_len,
+        per_device_batch_size=self.batch_size,
+        logits_via_embedding=False,
+        base_num_decoder_layers=0,
+        dtype="float32",
+        matmul_precision="high",
+        num_vocab_tiling=4,
+    )
+    quant = quantizations.configure_quantization(cfg)
+    devices_array = maxtext_utils.create_device_mesh(cfg)
+    mesh = Mesh(devices_array, cfg.mesh_axes)
+    model = models.transformer_as_linen(cfg, mesh=mesh, quant=quant, model_mode=MODEL_MODE_TRAIN)
+    rng_model, rng_hidden, rng_targets = jax.random.split(self.rng, 3)
+    params = model.init(
+        {"params": rng_model, "dropout": rng_model},
+        self.dummy_inputs,
+        self.dummy_inputs,
+    )
+    hidden_states = jax.random.normal(rng_hidden, (self.batch_size, self.seq_len, cfg.emb_dim), dtype=jnp.float32)
+    data = {
+        "targets": jax.random.randint(rng_targets, (self.batch_size, self.seq_len), 0, cfg.vocab_size),
+        "targets_segmentation": jnp.ones((self.batch_size, self.seq_len)),
+    }
+
+    def scaled_loss(h, outer_loss_scale):
+      total_loss, _ = vocab_tiling_linen_loss(h, data, cfg, model, params, is_train=True)
+      return outer_loss_scale * total_loss
+
+    grad_with_scale = jax.jit(jax.grad(scaled_loss, argnums=0))
+    base_grad = grad_with_scale(hidden_states, 1.0)
+    doubled_grad = grad_with_scale(hidden_states, 2.0)
+
+    assert jnp.all(jnp.isfinite(base_grad))
+    assert jnp.any(base_grad != 0)
+    assert jnp.allclose(doubled_grad, 2.0 * base_grad, rtol=1e-5, atol=1e-6)
+
   @pytest.mark.tpu_only
   def test_vocab_tiling_gradient_non_tied_embedding(self):
     """
@@ -888,6 +932,26 @@ def test_nnx_vocab_tiling_grad_over_hidden_states(self):
     assert tile_grad_h.dtype == hidden_states.dtype
     assert jnp.allclose(ref_grad_h, tile_grad_h, rtol=self.rtol, atol=self.atol), "grad_hidden_states diverged"
 
+  @pytest.mark.cpu_only
+  def test_nnx_vocab_tiling_hidden_gradient_respects_outer_loss_scale(self):
+    """The tiled loss's hidden-state gradient obeys the chain rule."""
+    cfg, model = self._build_cfg_and_model(num_vocab_tiling=4)
+    hidden_states, labels, segmentation = self._make_inputs(cfg)
+    graphdef, params, rest = self._split_and_axes(cfg, model)
+    tile_loss_fn = self._tiled_loss_fn(cfg, graphdef, rest, hidden_states, labels, segmentation)
+
+    def scaled_loss(p, h, outer_loss_scale):
+      return outer_loss_scale * tile_loss_fn(p, h)
+
+    grad_with_scale = self._g(scaled_loss, argnums=1)
+    with nn_partitioning.axis_rules(cfg.logical_axis_rules):
+      base_grad = grad_with_scale(params, hidden_states, 1.0)
+      doubled_grad = grad_with_scale(params, hidden_states, 2.0)
+
+    assert jnp.all(jnp.isfinite(base_grad))
+    assert jnp.any(base_grad != 0)
+    assert jnp.allclose(doubled_grad, 2.0 * base_grad, rtol=1e-5, atol=1e-6)
+
   @pytest.mark.tpu_only
   def test_nnx_vocab_tiling_bf16_hidden_states(self):
     """bf16 hidden_states: loss/grad parity holds and the grad keeps the bf16 dtype."""