diff --git a/src/maxtext/checkpoint_conversion/standalone_scripts/convert_gpt3_ckpt_from_paxml.py b/src/maxtext/checkpoint_conversion/standalone_scripts/convert_gpt3_ckpt_from_paxml.py
index 12bc48c676..9ee71b0197 100644
--- a/src/maxtext/checkpoint_conversion/standalone_scripts/convert_gpt3_ckpt_from_paxml.py
+++ b/src/maxtext/checkpoint_conversion/standalone_scripts/convert_gpt3_ckpt_from_paxml.py
@@ -213,6 +213,7 @@ def init_state_fn():
     state_map = {
         "['optimizer']['step'].value": ("step", None),
         "['optimizer']['opt_state'][0]['count'].value": ("opt_states_0.no_prefix_0.count", None),
+        "['optimizer']['opt_state']['count'].value": ("opt_states_0.no_prefix_0.count", None),
     }
   else:
     state_map = {
@@ -237,10 +238,18 @@ def get_layer_prefix(keystr_pax):
           f"opt_states_0.{prefix_pax_opt_state}.m{keystr_pax}",
           transform_fn,
       )
+      state_map[f"['optimizer']['opt_state']['mu']{keystr_maxtext}.value"] = (
+          f"opt_states_0.{prefix_pax_opt_state}.m{keystr_pax}",
+          transform_fn,
+      )
       state_map[f"['optimizer']['opt_state'][0]['nu']{keystr_maxtext}.value"] = (
           f"opt_states_0.{prefix_pax_opt_state}.v{keystr_pax}",
           transform_fn,
       )
+      state_map[f"['optimizer']['opt_state']['nu']{keystr_maxtext}.value"] = (
+          f"opt_states_0.{prefix_pax_opt_state}.v{keystr_pax}",
+          transform_fn,
+      )
     else:
       state_map[f".params['params']{keystr_maxtext}"] = (f"mdl_vars{keystr_pax}", transform_fn)
       state_map[f".opt_state.mu['params']{keystr_maxtext}"] = (
@@ -254,6 +263,8 @@ def get_layer_prefix(keystr_pax):
 
   def verify_fn(key_path, _):
     keystr = jax.tree_util.keystr(key_path)
+    if "['rngs']" in keystr:
+      return
     assert keystr in state_map, f"{keystr} not found"
 
   jax.tree_util.tree_map_with_path(verify_fn, state)
@@ -264,6 +275,8 @@ def verify_fn(key_path, _):
 
   def map_fn(key_path, value):
     key_path_str = jax.tree_util.keystr(key_path)
+    if "['rngs']" in key_path_str:
+      return value
     file_path, transform_fn = state_map[key_path_str]
     full_path = os.path.join(paxml_ckpt_prefix, file_path)
     spec = {"driver": "zarr", "metadata_key": ".zarray", "kvstore": {}}
@@ -303,6 +316,8 @@ def map_fn(key_path, value):
   max_utils.print_mem_stats("converted state finished")
 
   step_value = int(converted_state.optimizer.step.value) if cfg.pure_nnx else converted_state.step
+  if cfg.pure_nnx:
+    converted_state = train_state_nnx.to_checkpoint_dict(converted_state)
   if checkpointing.save_checkpoint(checkpoint_manager, step_value, converted_state):
     max_logging.log(f"saved a checkpoint at step {step_value}")
   # Upon preemption, exit when and only when all ongoing saves are complete.