fix(nnx): convert NNX state to checkpoint dict layout before saving in gpt3 paxml conversion

Author: hsuan-lun-chiang

src/maxtext/checkpoint_conversion/standalone_scripts/convert_gpt3_ckpt_from_paxml.py

@@ -316,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.