!deepspeed kto.py
--save_path '/root/chai/kto'
--save_steps 50
--logging_steps 10
--eval_steps -1
--train_batch_size 256
--micro_train_batch_size 1
--pretrain '/root/chai/mistralai/Mistral-Nemo-Instruct-2407'
--bf16
--max_epochs 1
--max_len 1348
--zero_stage 2
--learning_rate 5e-7
--dataset '/root/chai/1_kto'
--input_key instruction
--output_key response
--label_key score
--flash_attn
--beta 0.1
--max_samples 100000
--gradient_checkpointing
--lora_rank 16
--lora_alpha 32
--lora_dropout 0.05