From d60bb24e9244736e01a8d1f3623631a70b81d305 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Juliusz=20Straszy=C5=84ski?= Date: Thu, 20 Aug 2026 15:43:33 +0000 Subject: [PATCH 1/2] Fix vocab tiling hidden-state cotangent scaling --- src/maxtext/utils/vocabulary_tiling.py | 2 + tests/unit/tiling_test.py | 64 ++++++++++++++++++++++++++ 2 files changed, 66 insertions(+) diff --git a/src/maxtext/utils/vocabulary_tiling.py b/src/maxtext/utils/vocabulary_tiling.py index e88e43833f..a8f7452ebb 100644 --- a/src/maxtext/utils/vocabulary_tiling.py +++ b/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) diff --git a/tests/unit/tiling_test.py b/tests/unit/tiling_test.py index e518729449..89f04d394f 100644 --- a/tests/unit/tiling_test.py +++ b/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.""" From b566596241743109eb22694853ba173ac8509f99 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Juliusz=20Straszy=C5=84ski?= Date: Fri, 21 Aug 2026 19:11:26 +0000 Subject: [PATCH 2/2] Format attention test with pyink --- tests/unit/attention_test.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/unit/attention_test.py b/tests/unit/attention_test.py index 2cdf98e9af..632efbfa31 100644 --- a/tests/unit/attention_test.py +++ b/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