Fix vocab tiling memory space mismatch under parameter host offload

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."""
 

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 (

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):
     """