Merge pull request #3503 from zhenying-liu:param-offload-fix

PiperOrigin-RevId: 889513863

Author: Google-ML-Automation

src/maxtext/trainers/pre_train/train.py

@@ -355,6 +355,11 @@ def train_step(model, config, state_mesh_shardings, params_shardings, state, dat
       lambda x: x.astype(config.grad_dtype) if x.dtype == jnp.float32 else x,
       raw_grads,
   )
+  if config.parameter_memory_host_offload:
+    raw_grads = jax.device_put(
+        raw_grads,
+        max_utils.with_memory_kind(params_shardings, "device"),
+    )
   intermediate_outputs = aux["intermediate_outputs"]
   total_weights = aux["total_weights"]
   moe_lb_loss = aux["moe_lb_loss"]

src/maxtext/utils/max_utils.py

@@ -25,6 +25,8 @@
 import time
 from typing import Any
 
+from packaging.version import Version
+
 from etils import epath
 import flax
 import jax
@@ -82,7 +84,7 @@ def calculate_num_params_from_pytree(params):
 def device_space():
   """Version guard for jax.memory.Space.Device."""
   # See b/436565838 for more.
-  if jax.__version__ >= "0.7.1":
+  if Version(jax.__version__) >= Version("0.7.1"):
     return jax.memory.Space.Device  # pytype: disable=module-attr
   else:
     return jax._src.sharding_impls.TransferToMemoryKind("device")  # pylint: disable=protected-access # pytype: disable=module-attr