Missing ShardingHelper class when using opt-350m to evaluate

[Harry-mic]

Hi! Thanks for your awesome work! I meet a problem when running the following scripts:

python -m coh.models.opt.opt_serve \
    --load_opt_config='json::/root/autodl-tmp/chain-of-hindsight/datasets/opt-350m/config.json' \
    --load_checkpoint='/root/autodl-tmp/chain-of-hindsight/coh_output/ba56a63584e34c28b29679f43aae2ea1/wandb/latest-run/run-ba56a63584e34c28b29679f43aae2ea1.wandb' \
    --mp_mesh_dim=-1 \
    --dtype='bf16' \
    --input_length=1024 \
    --seq_length=2048 \
    --do_sample=True \
    --temperature=1.0 \
    --lm_server.port=5007 \
    --lm_server.pre_compile='all' \
    --lm_server.chat_prepend_text='' \
    --lm_server.chat_lm_prefix='An helpful answer:' \
    --lm_server.chat_lm_suffix='</s>' \
    --lm_server.chat_user_prefix='User: ' \
    --lm_server.chat_user_suffix=' '

As you suggest, I change your example(gptj-6B) :

--load_gptj_config='Your checkpoint path' \
--load_checkpoint='Your checkpoint path' \

to opt-350m:

--load_opt_config='json::/root/autodl-tmp/chain-of-hindsight/datasets/opt-350m/config.json' \
--load_checkpoint='/root/autodl-tmp/chain-of-hindsight/coh_output/ba56a63584e34c28b29679f43aae2ea1/wandb/latest-run/run-ba56a63584e34c28b29679f43aae2ea1.wandb' \

However, there exist a bug:
ImportError: cannot import name 'ShardingHelper' from 'coh.jax_utils' (/root/autodl-tmp/chain-of-hindsight/coh/jax_utils.py)
I search the 'ShardingHelper' class in the entire folders but find no result. Therefore, it seems that the class is missing. Could you please give some suggestions on where is the 'ShardingHelper' class?

[forhaoliu]

Thanks for the interest and report.
We have updated code with better sharding and training process. Please note that we dropped support to OPT because we found LLaMA and GPT-J are consistently better in CoH and other things.