File size: 951 Bytes
96fe658 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 |
import os
kwargs = {
'per_device_train_batch_size': 5,
'save_steps': 5,
'gradient_accumulation_steps': 1,
'num_train_epochs': 1,
}
def test_train_eval_loop():
os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
from swift.llm import sft_main, TrainArguments
sft_main(
TrainArguments(
model='Qwen/Qwen2.5-0.5B-Instruct',
dataset=['AI-ModelScope/alpaca-gpt4-data-zh#100'],
target_modules=['all-linear', 'all-embedding'],
modules_to_save=['all-embedding', 'all-norm'],
eval_strategy='steps',
eval_steps=5,
per_device_eval_batch_size=5,
eval_use_evalscope=True,
eval_dataset=['gsm8k'],
eval_dataset_args={'gsm8k': {
'few_shot_num': 0
}},
eval_limit=10,
report_to=['wandb'],
**kwargs))
if __name__ == '__main__':
test_train_eval_loop()
|