orbax to reg checkpointer conversion

[lkolluru05]

Command

Orbax Emergency test:

`axlearn gcp bundle --name=$NAME
--bundler_spec=allow_dirty=True
--bundler_type=artifactregistry
--bundler_spec=dockerfile=Dockerfile
--bundler_spec=image=tpu
--bundler_spec=target=tpu

axlearn gcp launch run --cluster=stoelinga-axlearn
--runner_name gke_tpu_single
--name=$NAME
--instance_type=tpu-v6e-16
--host_mount_spec=name=tmp,host_path=/tmp,mount_path=/host-tmp
--num_replicas=3
--bundler_spec=allow_dirty=True
--bundler_type=artifactregistry --bundler_spec=image=tpu
--bundler_spec=dockerfile=Dockerfile --bundler_spec=target=tpu
-- python3 -m axlearn.common.launch_trainer_main
--init_module=axlearn.common.checkpointer_orbax_emergency:local_ckpt_dir=/host-tmp/checkpoints
--module=text.gpt.c4_trainer
--config=fuji-7B-v2-flash-orbaxem
--trainer_dir=$OUTPUT_DIR
--data_dir=gs://axlearn-public/tensorflow_datasets
--jax_backend=tpu
--mesh_selector=tpu-v6e-16
--trace_at_steps=3`

Orbax test

`axlearn gcp bundle --name=$NAME
--bundler_spec=allow_dirty=True
--bundler_type=artifactregistry
--bundler_spec=dockerfile=Dockerfile
--bundler_spec=image=tpu
--bundler_spec=target=tpu

axlearn gcp launch run --cluster=lkolluru-axlearn
--runner_name=gke_tpu_pathways
--name=$NAME
--instance_type=tpu-v6e-16
--num_replicas=1
--bundler_spec=allow_dirty=True
--bundler_type=artifactregistry --bundler_spec=image=tpu
--bundler_spec=dockerfile=Dockerfile --bundler_spec=target=tpu
-- python3 -m axlearn.common.launch_trainer_main
--init_module=axlearn.common.checkpointer_orbax
--module=text.gpt.c4_trainer
--config=fuji-7B-v2-flash-orbax
--trainer_dir=$OUTPUT_DIR
--data_dir=gs://axlearn-public/tensorflow_datasets
--jax_backend=proxy
--mesh_selector=tpu-v6e-16
--trace_at_steps=3`

[findmyway]

If I understand it correctly, here we're still using an online approach to do the checkpoint conversion. This means that we have to allocate the same resource (or at least a slice) of training stage for each checkpoint conversion.

I'm wondering if we can do the conversion offline on a CPU only node with a large memory.