This commit is contained in:
Gang huang 2024-03-20 11:26:44 +08:00 committed by GitHub
commit 90861dc077
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 21 additions and 0 deletions

View File

@ -13,6 +13,16 @@ python run.py
to test the code. to test the code.
If you are using Conda, run
```shell
conda env create -f environment.yml
conda activate grok
python run.py
```
to test the code.
The script loads the checkpoint and samples from the model on a test input. The script loads the checkpoint and samples from the model on a test input.
Due to the large size of the model (314B parameters), a machine with enough GPU memory is required to test the model with the example code. Due to the large size of the model (314B parameters), a machine with enough GPU memory is required to test the model with the example code.

11
environment.yml Normal file
View File

@ -0,0 +1,11 @@
name: grok
channels:
- conda-forge
- nvidia
- defaults
dependencies:
- python=3.9
- numpy==1.26.4
- dm-haiku==0.0.12
- sentencepiece==0.2.0
- jax[cuda12-pip]==0.4.25