Add print statements to `pretrained/GPT-NeoX-20B/prepare.py` to show progress

[csris]

Describe the bug pretrained/GPT-NeoX-20B/prepare.py can take a long time to prepare the base model. It should print progress as it's converting.

To Reproduce Steps to reproduce the behavior:

1. run python pretrained/GPT-NeoX-20B/prepare.py from the root of the repo.

Expected behavior The script should print progress.

[csris]

@LorrinWWW, this is a version of your script for sharding the base model. Could you please take a look?

[LorrinWWW]

Sure!