How to load codegen models in jaxformer locally??

[harry-stark]

I've tried to extended the local master with load_ckpt function, but ran into trouble with:

assert jax.process_count() == ckpt['process_count']

Any pointers for this?? Do we've to use jax.distributed.initialize or is there another way for doing it with a single process ??

Edit: Conversion script for jax to hf model? Does one can adopt gptj script of conversion??

Setup: single node with 8 A100s

cc @enijkamp

[glicerico]

hey @harry-stark , I am having problem with the same assertion, even running in gcloud TPUs, did you solve this?