|
|
--- |
|
|
license: apache-2.0 |
|
|
base_model: |
|
|
- xai-org/grok-2 |
|
|
--- |
|
|
### huihui-ai/grok-2 |
|
|
This Python [script](https://huggingface.co/huihui-ai/grok-2/blob/main/convert_safetensors.py) is designed to process and merge sharded weight files |
|
|
(in safetensors format) for a machine learning model, specifically targeting the [xai-org/grok-2](https://huggingface.co/xai-org/grok-2) model. The main functionalities include: |
|
|
|
|
|
Just a simple merge, without any inference code, and does not indicate whether the final model is reasonable or correct. |
|
|
|
|
|
Now, do we need a custom MixtralForCausalLM? |
|
|
|
|
|
1. **Collecting safetensors files**: Locates all `pytorch_model-*.safetensors` files in the specified model directory. |
|
|
2. **Loading files into cache**: Loads all safetensors files into memory and builds a key-to-file mapping. |
|
|
3. **Merging Tensor Parallel (TP) shards**: Merges shards for tensor parallelism (TP=8) along specific dimensions and verifies the merged tensor shapes. |
|
|
4. **Grouping weights by layer**: Organizes weights by model layer, with special weights (e.g., `lm_head.weight`, `model.embed_tokens.weight`, and `model.norm.weight`) handled separately. |
|
|
5. **Saving merged weights**: Saves the grouped weights as new safetensors files and generates a new index file [pytorch_model.bin.index.json](https://huggingface.co/huihui-ai/grok-2/blob/main/pytorch_model.bin.index.json). |
|
|
|
|
|
### Features |
|
|
- **Input**: Safetensors files in the `xai-org/grok-2` model directory. |
|
|
- **Output**: Layer-organized safetensors files and an index file in the `huihui-ai/grok-2` directory. |
|
|
- **Tensor Parallelism Support**: Handles TP=8 shards, merging tensors along specific dimensions (`w1.weight` and `w3.weight` along dim=0, `w2.weight` along dim=1). |
|
|
- **Error Handling**: Includes warnings and handling for missing files, shape mismatches, and other exceptions. |
|
|
- **Shape Validation**: Verifies shapes for specific weights (e.g., MoE layer weights), ensuring merged tensors match expected shapes (e.g., `(16384, 8192)` or `(8192, 16384)`). |
|
|
|
|
|
### Usage |
|
|
1. Install the required Python libraries: |
|
|
```bash |
|
|
pip install torch safetensors |
|
|
``` |
|
|
2. Place the script in an environment with the `xai-org/grok-2` model directory. |
|
|
3. Run the script: |
|
|
```bash |
|
|
python convert_safetensors.py |
|
|
``` |
|
|
4. Output files will be saved in the `huihui-ai/grok-2` directory, including layer-organized safetensors files and an index file. |
|
|
|
|
|
### Notes |
|
|
- Ensure the input directory `xai-org/grok-2` contains valid `pytorch_model-*.safetensors` files. |
|
|
- The script assumes a tensor parallelism degree of 8 (`tp_count = 8`). Modify the `tp_count` value in the script if needed. |
|
|
- Memory requirements may be high; run on a machine with sufficient memory. |
|
|
- If shards are missing or shapes mismatch, the script will print warnings and attempt to proceed. |
