mirror of https://github.com/hpcaitech/ColossalAI
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
1.5 KiB
1.5 KiB
Grok-1 Inference
An easy-to-use Python + PyTorch + HuggingFace version of 314B Grok-1. [code] [blog] [HuggingFace Grok-1 PyTorch model weights]
Install
# Make sure you install colossalai from the latest source code
git clone https://github.com/hpcaitech/ColossalAI.git
cd ColossalAI
pip install .
cd examples/language/grok-1
pip install -r requirements.txt
Tokenizer preparation
You should download the tokenizer from the official grok-1 repository.
wget https://github.com/xai-org/grok-1/raw/main/tokenizer.model
Inference
You need 8x A100 80GB or equivalent GPUs to run the inference.
We provide two scripts for inference. run_inference_fast.sh uses tensor parallelism provided by ColossalAI, and it is faster. run_inference_slow.sh uses auto device provided by transformers, and it is slower.
Command format:
./run_inference_fast.sh <model_name_or_path> <tokenizer_path>
./run_inference_slow.sh <model_name_or_path> <tokenizer_path>
model_name_or_path can be a local path or a model name from Hugging Face model hub. We provided weights on model hub, named hpcaitech/grok-1.
Command example:
./run_inference_fast.sh hpcaitech/grok-1 tokenizer.model
It will take 5-10 minutes to load checkpoints. Don't worry, it's not stuck.