Checkpoint Scripts
Utilities
inspect_checkpoints.py
Reports information about a saved checkpoint.
usage: inspect_checkpoints.py [-h] [--attributes [ATTRIBUTES ...]] [--interactive] [--compare] [--diff] dir
positional arguments:
dir The checkpoint dir to inspect. Must be either: - a directory containing pickle binaries saved with 'torch.save' ending in .pt or .ckpt - a single path to a .pt or .ckpt file - two comma separated directories -
in which case the script will *compare* the two checkpoints
options:
-h, --help show this help message and exit
--attributes [ATTRIBUTES ...]
Name of one or several attributes to query. To access an attribute within a nested structure, use '/' as separator.
--interactive, -i Drops into interactive shell after printing the summary.
--compare, -c If true, script will compare two directories separated by commas
--diff, -d In compare mode, only print diffs
HuggingFace Scripts
convert_hf_to_sequential.py
A script for converting publicly available Huggingface (HF) checkpoints to NeoX format.
Note that this script requires access to corresponding config files for equivalent NeoX models to those found in Hugging face.
Example usage: (Converts the 70M Pythia model to NeoX format)
================================================================
OMPI_COMM_WORLD_RANK=0 CUDA_VISIBLE_DEVICES=0 python tools/ckpts/convert_hf_to_sequential.py \
--hf-model-name pythia-70m-v0 \
--revision 143000 \
--output-dir checkpoints/neox_converted/pythia/70m \
--cache-dir checkpoints/HF \
--config configs/pythia/70M.yml configs/local_setup.yml \
--test
For multi-gpu support we must initialize deepspeed:
NOTE: This requires manually changing the arguments below.
================================================================
CUDA_VISIBLE_DEVICES=0,1,2,3 python ./deepy.py tools/ckpts/convert_hf_to_sequential.py \
-d configs pythia/70M.yml local_setup.yml
convert_module_to_hf.py
Converts a NeoX model with pipeline parallelism greater than 1 to a HuggingFace transformers GPTNeoXForCausalLM model
Note that this script does not support all NeoX features. Please investigate carefully whether your model is compatible with all architectures supported by the GPTNeoXForCausalLM class in HF.
(e.g. position embeddings such as AliBi may not be supported by Huggingface's GPT-NeoX architecture)
usage: convert_module_to_hf.py [-h] [--input_dir INPUT_DIR] [--config_file CONFIG_FILE] [--output_dir OUTPUT_DIR] [--upload]
Merge MP partitions and convert to HF Model.
options:
-h, --help show this help message and exit
--input_dir INPUT_DIR
Path to NeoX checkpoint, e.g. /path/to/model/global_step143000
--config_file CONFIG_FILE
Path to config file for the input NeoX checkpoint.
--output_dir OUTPUT_DIR
Output dir, where to save the HF Model, tokenizer, and configs
--upload Set to true in order to upload to the HF Hub directly.
convert_sequential_to_hf.py
Converts a NeoX model without pipeline parallelism to a HuggingFace transformers GPTNeoXForCausalLM model.
usage: convert_sequential_to_hf.py [-h] [--input_dir INPUT_DIR] [--config_file CONFIG_FILE] [--output_dir OUTPUT_DIR] [--upload]
Merge MP partitions and convert to HF Model.
options:
-h, --help show this help message and exit
--input_dir INPUT_DIR
Path to NeoX checkpoint, e.g. /path/to/model/global_step143000
--config_file CONFIG_FILE
Path to config file for the input NeoX checkpoint.
--output_dir OUTPUT_DIR
Output dir, where to save the HF Model, tokenizer, and configs
--upload Set to true in order to upload to the HF Hub directly.
upload.py
Uploads a converted checkpoint to the HuggingFace hub.
python upload.py <converted-ckpt-dir> <repo-name> <branch-name>
NeoX-20B Scripts
merge20b.py
Reduces model and pipeline parallelism of a 20B checkpoint to 1 and 1.
usage: merge20b.py [-h] [--input_dir INPUT_DIR] [--output_dir OUTPUT_DIR]
Merge 20B checkpoint.
options:
-h, --help show this help message and exit
--input_dir INPUT_DIR
Checkpoint dir, which should contain (e.g. a folder named "global_step150000")
--output_dir OUTPUT_DIR
Output dir, to save the 1-GPU weights configs
Llama Scripts
convert_raw_llama_weights_to_neox.py
Takes a Llama checkpoint and puts it into a NeoX-compatible format.
usage: convert_raw_llama_weights_to_neox.py [-h] [--input_dir INPUT_DIR] [--model_size {7B,13B,30B,65B,tokenizer_only}] [--output_dir OUTPUT_DIR] [--num_output_shards NUM_OUTPUT_SHARDS] [--pipeline_parallel]
Convert raw LLaMA checkpoints to GPT-NeoX format.
options:
-h, --help show this help message and exit
--input_dir INPUT_DIR
Location of LLaMA weights, which contains tokenizer.model and model folders
--model_size {7B,13B,30B,65B,tokenizer_only}
--output_dir OUTPUT_DIR
Location to write GPT-NeoX mode
--num_output_shards NUM_OUTPUT_SHARDS
--pipeline_parallel Only use if PP>1
convert_hf_llama_to_neox.py
Takes an HF Llama checkpoint and puts it into a NeoX-compatible format.
Note that this does not support pipeline parallelism!
usage: convert_hf_llama_to_neox.py [-h] [--tp TP] [--pp PP] [--model MODEL] [--model_path MODEL_PATH]
options:
-h, --help show this help message and exit
--tp TP Number of tensor parallelism ranks
--pp PP Number of pipeline parallelism stages
--model MODEL HF model name
--model_path MODEL_PATH
Path to save model