update eval and readme

README.md

@@ -10,6 +10,9 @@ Y. Wu.
 **ESFT** aims to efficiently customize Large Language Models (LLMs) with Mixture-of-Experts (MoE) architecture by adjusting only task-relevant parts, improving efficiency and performance while using fewer resources and storage. 
  
 
+## 📰 News
+
+📅 **2024.8.11:** We now release the **ESFT training code**! ✨ You can now try it with your own models and dataset!
 
 
 ## 🚀 Quick Start 
@@ -19,9 +22,9 @@ git clone https://github.com/deepseek-ai/ESFT.git
 cd esft
 ```
 
-### Install dependencies
+### Install required dependencies
 ```bash
-pip install transformers torch safetensors
+pip install transformers torch safetensors accelerate
 ```
 
 ### Download necessary adapters
@@ -32,35 +35,38 @@ bash scripts/download_adapters.sh
 
 
 ## 🔧Key Scripts
-1. **eval.py**
-This script evaluates the performance of the model on various datasets. **Usage:**
+1. **eval_multigpu.py**
+This script evaluates the performance of the model on various datasets. See **scripts/eval.sh** for detailed configs and explanations.
+
+**Usage:**
 ```bash
-python scripts/eval.py \
-    --eval_datasets=translation \
+python eval_multigpu.py \
+    --eval_dataset=translation \
     --base_model_path=deepseek-ai/ESFT-vanilla-lite \
-    --adapter_dir=all_models/adapters/token \
-    --output_dir=results/completions/token \
-    --max_new_tokens=512 \
-    --openai_api_key=REPLACE_WITH_YOUR_KEY \
-    --eval_batch_size=2
+    --adapter_dir=all_models/adapters/token/translation \
+    --output_path=results/completions/token/translation.jsonl \
+    --openai_api_key=YOUR_OPENAI_API_KEY
 ```
 
+
 2. **get_expert_scores.py**
 This script calculates the scores for each expert based on the evaluation datasets.
 **Usage:**
 ```bash
-python scripts/get_expert_scores.py \
-    --eval_datasets=intent,summary,law,translation \
+python scripts/expert/get_expert_scores.py \
+    --eval_dataset=translation \
     --base_model_path=deepseek-ai/ESFT-vanilla-lite \
-    --output_dir=results/expert_scores \
-    --n_sample_tokens=8192 # the sample size hyperparameter
+    --output_dir=results/expert_scores/translation \
+    --n_sample_tokens=131072 \
+    --world_size=4 \
+    --gpus_per_rank=2
 ```
 
 3. **generate_expert_config.py**
 This script generates the configuration to convert a MoE model with only task-relevant tasks trained based on evaluation scores.
 **Usage:**
 ```bash
-python scripts/generate_expert_config.py \
+python scripts/expert/generate_expert_config.py \
     --eval_datasets=intent,summary,law,translation \
     --expert_scores_dir=results/expert_scores \
     --output_dir=results/expert_configs \
@@ -68,13 +74,32 @@ python scripts/generate_expert_config.py \
     --top_p=0.2 # the scoring function and top_p are hyperparameters
 ```
 
+4. **train.py** and **train_ep.py**
+This script trains the model with the expert configuration generated by the previous script. The train_ep.py file uses expert parallel and has been optimized for multi-GPU training.
+**Usage:**
+```bash
+python train.py \
+    --base_model_path=deepseek-ai/ESFT-vanilla-lite \
+    --expert_config=results/expert_configs/intent.json \
+    --train_dataset=intent \
+    --train_config=configs/base.yaml \
+    --output_dir=results/checkpoints/intent
+    
+torchrun --nproc-per-node=8 train_ep.py \
+    --base_model_path=deepseek-ai/ESFT-vanilla-lite \
+    --expert_config=results/expert_configs/translation.json \
+    --train_dataset=translation \
+    --train_config=configs/base.yaml \
+    --output_dir=results/checkpoints/translation
+
+```
 
 ## Contact and Support
 For bug reports, feature requests, and general inquiries, please open an issue on our GitHub Issues page. Make sure to include as much detail as possible to help us address your issue quickly.
 
 ## 🌟Todo list
 - ☑️  📝 Update models, evaluation scripts, and expert selection scripts
-- 🔲 🔧 Update training scripts
+- ☑️ 🔧 Update training scripts
 - 🔲 🚀 More...