From 0cf33200fa7578d1b514873ed71bc7f118b91ef5 Mon Sep 17 00:00:00 2001 From: continuousml Date: Sun, 19 Jul 2026 19:21:09 -0700 Subject: [PATCH] Fix vocab tiling memory space mismatch under parameter host offload --- src/maxtext/utils/max_utils.py | 9 ++++ src/maxtext/utils/vocabulary_tiling.py | 7 +++ tests/unit/tiling_test.py | 68 ++++++++++++++++++++++++++ 3 files changed, 84 insertions(+) diff --git a/src/maxtext/utils/max_utils.py b/src/maxtext/utils/max_utils.py index 5058745964..5df82ec278 100644 --- a/src/maxtext/utils/max_utils.py +++ b/src/maxtext/utils/max_utils.py @@ -130,6 +130,15 @@ def device_space(): return jax._src.sharding_impls.TransferToMemoryKind("device") # pylint: disable=protected-access # pytype: disable=module-attr +def host_space(): + """Version guard for jax.memory.Space.Host.""" + # See b/436565838 for more. + if Version(jax.__version__) >= Version("0.7.1"): + return jax.memory.Space.Host # pytype: disable=module-attr + else: + return jax._src.sharding_impls.TransferToMemoryKind("pinned_host") # pylint: disable=protected-access # pytype: disable=module-attr + + def calculate_total_params_per_chip(params): """Calculate total params per chip.""" diff --git a/src/maxtext/utils/vocabulary_tiling.py b/src/maxtext/utils/vocabulary_tiling.py index 9633528d63..822c0767b3 100644 --- a/src/maxtext/utils/vocabulary_tiling.py +++ b/src/maxtext/utils/vocabulary_tiling.py @@ -232,6 +232,10 @@ def _bwd_scan_body(grad_params_acc, chunk_data): # 1.0 since total_loss is sum of all individual chunked loss (grad_params_update, grad_hidden_chunk) = vjp_fn(1.0) + # Params offloaded to host yield host cotangents. Move them to device so the + # accumulation operands share the same memory space. + if config.parameter_memory_host_offload: + grad_params_update = jax.device_put(grad_params_update, max_utils.device_space()) grad_hidden_chunk = _maybe_shard_with_name(grad_hidden_chunk, chunked_hidden_spec) grad_params_acc = jax.tree_util.tree_map( @@ -252,6 +256,9 @@ def _bwd_scan_body(grad_params_acc, chunk_data): grad_params = jax.tree_util.tree_map(lambda g: g * loss_cotangent, grad_params) # 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) + # Move cotangents back to each primal's memory space; custom_vjp requires memory space match. + if config.parameter_memory_host_offload: + grad_params = jax.device_put(grad_params, max_utils.host_space()) # Give back sharding constraint grad_reshaped_hidden_states = _reshape(grad_reshaped_hidden_states, (batch_size, seq_len, emb_dim), hidden_spec) return ( diff --git a/tests/unit/tiling_test.py b/tests/unit/tiling_test.py index b1d257c9c5..9783a7062f 100644 --- a/tests/unit/tiling_test.py +++ b/tests/unit/tiling_test.py @@ -272,6 +272,74 @@ def test_vocab_tiling_gradient_with_z_loss(self): "Gradients do not match for vocab tiling when z-loss is enabled.", ) + @pytest.mark.cpu_only + def test_vocab_tiling_gradient_param_host_offload(self): + """ + Tests loss and gradient correctness when parameters are offloaded to host, + comparing on-device computation vs. host-offloaded computation. + """ + cfg = pyconfig.initialize( + self.base_config, + run_name="grad_test_vt_no_offload", + 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_targets = jax.random.split(self.rng) + + params = model.init( + {"params": rng_model, "dropout": rng_model}, + self.dummy_inputs, + self.dummy_inputs, + ) + + 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)), + } + + loss_ref, grads_ref = self.get_grads(cfg, params, data) + + cfg_offload = pyconfig.initialize( + self.base_config, + run_name="grad_test_vt_param_host_offload", + 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, + parameter_memory_host_offload=True, + ) + params_host = jax.device_put(params, max_utils.host_space()) + loss_offload, grads_offload = self.get_grads(cfg_offload, params_host, data) + # Gradients are in the primals' host memory space; move to device to compare. + grads_offload = jax.device_put(grads_offload, max_utils.device_space()) + + # Loss correctness test + assert jnp.allclose(loss_ref, loss_offload, rtol=self.rtol), "Losses do not match with parameter host offload." + + # Gradient correctness test + self.assert_pytrees_all_close( + grads_ref, + grads_offload, + "Gradients do not match for vocab tiling with parameter host offload.", + ) + @pytest.mark.tpu_only def test_vocab_tiling_nnx_loss(self): """