diff --git a/axlearn/common/array_serialization.py b/axlearn/common/array_serialization.py
index db78cfa27..2bcfb82f1 100644
--- a/axlearn/common/array_serialization.py
+++ b/axlearn/common/array_serialization.py
@@ -36,6 +36,7 @@
 from jax._src import array, typing
 from jax._src.layout import Format
 from jax.experimental.array_serialization import serialization
+from packaging import version
 
 from axlearn.common.utils import Tensor
 
@@ -75,6 +76,7 @@ def shard_coordinate(self):
 
 # Tuple (and thus hashable) representation of a slice object (start, end, step).
 _SliceTuple = tuple[Optional[int], Optional[int], Optional[int]]
+JAX_VERSION = version.parse(jax.__version__)
 
 
 def _slices_to_tuple(slices: list[slice]) -> tuple[_SliceTuple, ...]:
@@ -304,8 +306,15 @@ async def _async_serialize(
         and jax.process_count() > 1
         and arr_inp.is_fully_addressable
     )
-    # pylint: disable-next=protected-access
-    if not serialization.ts_impl._spec_has_metadata(tensorstore_spec):
+    # pylint: disable=protected-access
+    if JAX_VERSION >= version.parse("0.6.2"):
+        spec_has_metadata = serialization.ts_impl._spec_has_metadata
+    elif JAX_VERSION >= version.parse("0.5.3"):
+        spec_has_metadata = serialization._spec_has_metadata
+    else:
+        raise ValueError(f"Unsupported JAX version for spec_has_metadata: {jax.__version__}")
+
+    if not spec_has_metadata(tensorstore_spec):
         # pylint: disable-next=protected-access
         tensorstore_spec["metadata"] = serialization._get_metadata(arr_inp)
     if "dtype" not in tensorstore_spec:
@@ -486,7 +495,15 @@ async def _async_deserialize(
     async def cb(index: array.Index, device: jax.Device):
         requested_domain = ts.IndexTransform(input_shape=shape)[index].domain
         restricted_domain = t.domain.intersect(requested_domain)
-        requested_bytes = serialization.ts_impl.estimate_read_memory_footprint(t, restricted_domain)
+        if JAX_VERSION >= version.parse("0.6.2"):
+            estimate_read_memory_footprint = serialization.ts_impl.estimate_read_memory_footprint
+        elif JAX_VERSION >= version.parse("0.5.3"):
+            estimate_read_memory_footprint = serialization.estimate_read_memory_foot_print
+        else:
+            raise ValueError(
+                f"Unsupported JAX version: {JAX_VERSION}. Version must be 0.5.3 or newer"
+            )
+        requested_bytes = estimate_read_memory_footprint(t, restricted_domain)
         # Limit the bytes read for every shard.
         await byte_limiter.wait_for_bytes(requested_bytes)
         read_ts = t[restricted_domain]
@@ -563,10 +580,15 @@ async def cb(index: array.Index, device: jax.Device):
         await byte_limiter.release_bytes(requested_bytes)
         return result
 
-    # pylint: disable-next=protected-access
-    return await serialization.ts_impl._create_async_array_from_callback(
-        shape, dtype, in_sharding, cb
-    )
+    # pylint: disable=protected-access
+    if JAX_VERSION >= version.parse("0.6.2"):
+        create_async_array_from_callback = serialization.ts_impl._create_async_array_from_callback
+    elif JAX_VERSION >= version.parse("0.5.3"):
+        create_async_array_from_callback = serialization.create_async_array_from_callback
+    else:
+        raise ValueError("Unsupported JAX version: {JAX_VERSION}. Version must be 0.5.3 or newer.")
+
+    return await create_async_array_from_callback(shape, in_sharding, cb)
 
 
 # Reference:
@@ -647,6 +669,15 @@ def serialize(
 
         commit_futures = [[] for _ in range(len(tensorstore_specs))]
 
+        if JAX_VERSION >= version.parse("0.6.2"):
+            async_serialize = serialization.ts_impl.async_serialize
+        elif JAX_VERSION >= version.parse("0.5.3"):
+            async_serialize = serialization.async_serialize
+        else:
+            raise ValueError(
+                f"Unsupported JAX version: {JAX_VERSION}. Version must be 0.5.3 or newer."
+            )
+
         # pylint: disable-next=redefined-outer-name
         async def _run_serializer():
             future_writer = jax.tree.map(
diff --git a/axlearn/common/flash_attention/gpu_attention.py b/axlearn/common/flash_attention/gpu_attention.py
index c8060dce0..59c96af72 100644
--- a/axlearn/common/flash_attention/gpu_attention.py
+++ b/axlearn/common/flash_attention/gpu_attention.py
@@ -44,7 +44,6 @@
 )
 from jax.ad_checkpoint import checkpoint_name
 from jax.experimental import pallas as pl
-from jax.experimental.pallas.triton import CompilerParams as TritonCompilerParams
 
 from axlearn.common.attention_bias import (
     NEG_INF,
@@ -69,7 +68,14 @@
 from axlearn.common.kv_cache.base_kv_cache import BaseKVCache
 from axlearn.common.kv_cache.kv_cache import KVCache
 from axlearn.common.layers import get_dropout_mask
-from axlearn.common.utils import Nested, Tensor
+from axlearn.common.utils import _JAX_MEMORY_SPACE_SUPPORT, Nested, Tensor
+
+# pylint: disable=ungrouped-imports
+if _JAX_MEMORY_SPACE_SUPPORT:
+    from jax.experimental.pallas.triton import CompilerParams as TritonCompilerParams
+else:
+    from jax.experimental.pallas.triton import TritonCompilerParams
+# pylint: disable=ungrouped-imports
 
 
 def _segment_mask(
diff --git a/axlearn/common/flash_attention/gpu_decoding.py b/axlearn/common/flash_attention/gpu_decoding.py
index 73c5c15cf..c3dca5d9e 100644
--- a/axlearn/common/flash_attention/gpu_decoding.py
+++ b/axlearn/common/flash_attention/gpu_decoding.py
@@ -49,7 +49,6 @@
 from jax import lax
 from jax._src.cudnn.fused_attention_stablehlo import check_compute_capability
 from jax.experimental import pallas as pl
-from jax.experimental.pallas.triton import CompilerParams as TritonCompilerParams
 
 from axlearn.common.attention_bias import (
     NEG_INF,
@@ -61,7 +60,14 @@
 from axlearn.common.flash_attention.common import BaseSingleStepDecoding, get_gpu_dot_precision
 from axlearn.common.kv_cache.base_kv_cache import BaseKVCache
 from axlearn.common.kv_cache.kv_cache import KVCache
-from axlearn.common.utils import Nested, Tensor
+from axlearn.common.utils import _JAX_MEMORY_SPACE_SUPPORT, Nested, Tensor
+
+# pylint: disable=ungrouped-imports
+if _JAX_MEMORY_SPACE_SUPPORT:
+    from jax.experimental.pallas.triton import CompilerParams as TritonCompilerParams
+else:
+    from jax.experimental.pallas.triton import TritonCompilerParams
+# pylint: enable=ungrouped-imports
 
 
 # Note: split_k_seq_len must be a multiple of block_k.
diff --git a/axlearn/common/flash_attention/gpu_paged_attention.py b/axlearn/common/flash_attention/gpu_paged_attention.py
index fbe17bc7a..9a3dd9d0e 100644
--- a/axlearn/common/flash_attention/gpu_paged_attention.py
+++ b/axlearn/common/flash_attention/gpu_paged_attention.py
@@ -19,7 +19,6 @@
 import jax.numpy as jnp
 from jax import lax
 from jax.experimental import pallas as pl
-from jax.experimental.pallas.triton import CompilerParams as TritonCompilerParams
 
 from axlearn.common.attention_bias import (
     NEG_INF,
@@ -31,7 +30,16 @@
 from axlearn.common.flash_attention.common import BasePagedAttention, get_gpu_dot_precision
 from axlearn.common.flash_attention.gpu_decoding import _get_sm_count as get_sm_count
 from axlearn.common.kv_cache.base_kv_cache import BaseKVCache
-from axlearn.common.utils import Nested, Tensor
+from axlearn.common.utils import _JAX_MEMORY_SPACE_SUPPORT, Nested, Tensor
+
+# pylint: disable=ungrouped-imports
+if _JAX_MEMORY_SPACE_SUPPORT:
+    from jax.experimental.pallas.triton import CompilerParams as TritonCompilerParams
+else:
+    from jax.experimental.pallas.triton import (  # isort: skip
+        TritonCompilerParams,
+    )
+# pylint: enable=ungrouped-imports
 
 
 def _paged_attention_kernel(
diff --git a/axlearn/common/kv_cache/paged_kv_cache_gpu_kernel.py b/axlearn/common/kv_cache/paged_kv_cache_gpu_kernel.py
index 7483ada44..5358d908c 100644
--- a/axlearn/common/kv_cache/paged_kv_cache_gpu_kernel.py
+++ b/axlearn/common/kv_cache/paged_kv_cache_gpu_kernel.py
@@ -8,9 +8,15 @@
 import jax
 import jax.numpy as jnp
 from jax.experimental import pallas as pl
-from jax.experimental.pallas.triton import CompilerParams as TritonCompilerParams
 
-from axlearn.common.utils import Tensor
+from axlearn.common.utils import _JAX_MEMORY_SPACE_SUPPORT, Tensor
+
+# pylint: disable=ungrouped-imports
+if _JAX_MEMORY_SPACE_SUPPORT:
+    from jax.experimental.pallas.triton import CompilerParams as TritonCompilerParams
+else:
+    from jax.experimental.pallas.triton import TritonCompilerParams
+# pylint: enable=ungrouped-imports
 
 
 def _scatter_pages_kernel(
diff --git a/axlearn/common/optimizers.py b/axlearn/common/optimizers.py
index 4c618564d..6f044e18c 100644
--- a/axlearn/common/optimizers.py
+++ b/axlearn/common/optimizers.py
@@ -37,8 +37,6 @@
 import typing_extensions
 from absl import logging
 from jax import numpy as jnp
-
-# from jax._src.sharding_impls import TransferToMemoryKind
 from optax._src import numerics
 
 from axlearn.common import flax_struct, schedule
@@ -54,6 +52,8 @@
     TransformPartitionSpecFn,
 )
 from axlearn.common.utils import (
+    DEVICE_MEMORY,
+    HOST_MEMORY,
     MemoryKind,
     Nested,
     NestedTensor,
@@ -63,6 +63,7 @@
     expand_vdicts,
     flatten_items,
     register_per_param_settings,
+    transfer_to_memory_kind,
     tree_paths,
     vectorized_tree_map,
 )
@@ -2073,8 +2074,8 @@ def offload_optimizer(
     optimizer: ConfigOr[PartitionedGradientTransformation],
     *,
     pattern: Union[str, re.Pattern] = ".*",
-    offload_src: MemoryKind = "device",
-    offload_dst: MemoryKind = "pinned_host",
+    offload_src: MemoryKind = DEVICE_MEMORY,
+    offload_dst: MemoryKind = HOST_MEMORY,
 ) -> PartitionedGradientTransformation:
     """Offload the state of the wrapped optimizer that matches `pattern` to `offload_dst`.
 
@@ -2146,14 +2147,8 @@ def _move_fn(state: optax.OptState, dst: MemoryKind) -> optax.OptState:
         # released, so we have less memory pressure at that point in time.
         # memory_kind = sharding.memory_kind(dst)
         return jax.tree.map(
-            lambda path, tensor: (
-                jax.device_put(tensor, tensor.sharding.with_memory_kind(dst))
-                if re.fullmatch(pattern, path) and hasattr(tensor, "sharding")
-                else tensor
-            ),
-            tree_paths(state),
-            state,
-        )
+            transfer_to_memory_kind(tensor, dst) if re.fullmatch(pattern, path) else tensor
+            )
 
     def update_fn(updates: optax.Updates, state: optax.OptState, params: NestedOptParam):
         state = _move_fn(state, offload_src)
diff --git a/axlearn/common/optimizers_test.py b/axlearn/common/optimizers_test.py
index b142f1ff6..9cf66a886 100644
--- a/axlearn/common/optimizers_test.py
+++ b/axlearn/common/optimizers_test.py
@@ -60,6 +60,7 @@
 from axlearn.common.schedule import Schedule, adafactor_decay_rate, decay_bias_correction
 from axlearn.common.test_utils import TestCase, assert_allclose
 from axlearn.common.utils import (
+    _JAX_MEMORY_SPACE_SUPPORT,
     NestedPartitionSpec,
     PartitionSpec,
     Tensor,
@@ -428,9 +429,15 @@ def compute_loss(x):
 
         backend = jax.default_backend()
         if offload:
-            if backend and backend in ["gpu", "tpu"]:
+            jaxpr_str = str(jax.make_jaxpr(jit_fn)(params, state))
+            if _JAX_MEMORY_SPACE_SUPPORT:
                 self.assertIn(
-                    "MemoryKind='pinned_host'",
+                    "memory_kind=host",
+                    jaxpr_str,
+                )
+            else:
+                self.assertIn(
+                    "TransferToMemoryKind(memory_kind='pinned_host')",
                     str(jax.make_jaxpr(jit_fn)(params, state)),
                 )
         loss, new_loss = jit_fn(params, state)
diff --git a/axlearn/common/utils.py b/axlearn/common/utils.py
index da904fb8d..b7e165e03 100644
--- a/axlearn/common/utils.py
+++ b/axlearn/common/utils.py
@@ -54,6 +54,7 @@
 from jax.experimental import mesh_utils, multihost_utils
 from jax.extend.core import Primitive
 from jax.sharding import PartitionSpec
+from packaging import version
 
 from axlearn.common import serialization
 from axlearn.common.config import (
@@ -66,6 +67,9 @@
     register_validator,
 )
 
+# Define the version of JAX for compatibility on MemKind
+_JAX_MEMORY_SPACE_SUPPORT = version.parse(jax.__version__) >= version.parse("0.7.0")
+
 # New code should use Nested[XX] instead of NestedXX.
 # Old definitions are provided for backwards compatibility.
 _NestedT = TypeVar("_NestedT")
@@ -118,7 +122,23 @@ def __len__(self):
 # "pinned_host" = Page locked memory on CPU, which can be address directly by accelerators by
 # direct memory access (DMA). For TPU, "pinned_host" memory layout follows TPU device tile
 # layout and usually cannot be zero-copy converted to a CPU-tensor.
-MemoryKind = Literal["device", "pinned_host"]
+if _JAX_MEMORY_SPACE_SUPPORT:
+    MemoryKind = [jax.memory.Space.Device, jax.memory.Space.Host]
+    DEVICE_MEMORY = jax.memory.Space.Device
+    HOST_MEMORY = jax.memory.Space.Host
+
+    def transfer_to_memory_kind(tensor: Tensor, memory_kind: MemoryKind) -> Tensor:
+        return jax.device_put(tensor, memory_kind)
+
+else:
+    from jax._src.sharding_impls import TransferToMemoryKind  # pylint: disable=ungrouped-imports
+
+    MemoryKind = Literal["device", "pinned_host"]
+    DEVICE_MEMORY = "device"
+    HOST_MEMORY = "pinned_host"
+
+    def transfer_to_memory_kind(tensor: Tensor, memory_kind: MemoryKind) -> Tensor:
+        return jax.device_put(tensor, TransferToMemoryKind(memory_kind))
 
 
 @dataclasses.dataclass
diff --git a/axlearn/experiments/text/gpt/c4_trainer.py b/axlearn/experiments/text/gpt/c4_trainer.py
index e138a1d5a..2cc602218 100644
--- a/axlearn/experiments/text/gpt/c4_trainer.py
+++ b/axlearn/experiments/text/gpt/c4_trainer.py
@@ -49,11 +49,8 @@
 from axlearn.common.input_lm import lm_text_preprocessor
 from axlearn.common.utils import get_data_dir
 from axlearn.experiments.text.common import DataMixtureComponent, vocab
-from axlearn.experiments.text.gpt import envy, fuji, gspmd, qwen  # pytype: disable=pyi-error
-from axlearn.experiments.text.gpt.common import (  # pytype: disable=pyi-error
-    mixture_train_input_source,
-    tfds_input,
-)
+from axlearn.experiments.text.gpt import fuji, gspmd
+from axlearn.experiments.text.gpt.common import mixture_train_input_source, tfds_input
 from axlearn.experiments.text.gpt.vocabulary_fuji_v3 import FujiV3Vocabulary
 
 
@@ -114,6 +111,4 @@ def named_trainer_configs() -> dict[str, TrainerConfigFn]:
     config_map = {}
     config_map.update(fuji.trainer_configs(_train_input_source, _eval_input_sources))
     config_map.update(gspmd.trainer_configs(_train_input_source, _eval_input_sources))
-    config_map.update(envy.trainer_configs(_train_input_source, _eval_input_sources))
-    config_map.update(qwen.trainer_configs(_train_input_source, _eval_input_sources))
     return config_map