train

scripts/TRAIN.md

@@ -47,11 +47,10 @@ cp config.json out/converted_model/
 
 ```python
 import torch
-from transformers import AutoModel
+from safetensors.torch import save_file
 
-state_dict = torch.load('out/converted_model/model.pth')
-model = AutoModel.from_pretrained('TinyLlama/TinyLlama_v1.1', state_dict=state_dict, ignore_mismatched_sizes=True)
-model.save_pretrained('out/converted_model/')
+state_dict = torch.load('out/converted_model/model.pth', map_location='cpu')
+save_file(state_dict, 'out/converted_model/model.safetensors')
 ```
 
 ## Evaluate

scripts/requirements-lit.in

@@ -1,11 +0,0 @@
-# pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
-tqdm
-datasets
-jinja2
-transformers
-bitsandbytes
-wandb
-# litgpt[all]
-litgpt[all] @ git+https://github.com/Lightning-AI/litgpt.git
-litdata
-grokadamw

scripts/requirements.in

@@ -4,4 +4,8 @@ datasets
 jinja2
 transformers
 bitsandbytes
-wandb
+wandb
+# litgpt[all]
+litgpt[all] @ git+https://github.com/Lightning-AI/litgpt.git
+litdata
+grokadamw