Update README.md

README.md

@@ -1,3 +1,36 @@
----
-license: apache-2.0
----
+---
+license: apache-2.0
+---
+### huihui-ai/grok-2
+This Python script is designed to process and merge sharded weight files (in safetensors format) for a machine learning model, specifically targeting the `xai-org/grok-2` model. The main functionalities include:
+
+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`.
+
+### 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 script.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.