Spaces:
Sleeping
Sleeping
File size: 1,715 Bytes
17ff0d8 |
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 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 |
CMD="
python -m sdlm.train_reward_model \
--dataset_name argilla/ultrafeedback-binarized-preferences-cleaned \
\
--num_train_epochs 1 \
--per_device_train_batch_size 1 \
--remove_unused_columns=False \
--gradient_checkpointing=True \
--warmup_ratio 0.03 \
--learning_rate=2e-5 \
--report_to="tensorboard" \
--logging_steps=50 \
--save_total_limit 1 \
--optim adamw_torch_fused \
--evaluation_strategy="steps" \
--max_length=512 \
--gradient_checkpointing \
--bf16 \
--include_padding=False \
--use_tulu_chat_template=True \
--use_flash_attention2=True \
"
# on beaker, load from niklas' trained mistral model.
if [ ! -z "${BEAKER}" ]; then
gantry run -y -n mistral_rm_train -t mistral_rm_train --allow-dirty \
--workspace ai2/tess2 \
--gpus 1 \
--priority preemptible \
--budget ai2/allennlp \
--cluster ai2/jupiter-cirrascale-2 \
--env 'HF_HOME=/net/nfs.cirrascale/allennlp/jaket/.hf' \
--env 'PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python' \
--dataset '01J0PF0NKZP7SD8TMRH2PD0NFK:/model' \
--beaker-image 'ai2/pytorch2.0.0-cuda11.8-python3.10' \
--env-secret HF_TOKEN=HF_TOKEN \
--venv 'base' \
--pip requirements.txt \
-- ${CMD} \
--model_name_or_path /model \
--eval_steps 200 \
--save_steps 400 \
--gradient_accumulation_steps 128 \
--output_dir /results
else
${CMD} \
--model_name_or_path mistralai/Mistral-7B-v0.1 \
--eval_steps 1 \
--eval_steps 5 \
--save_steps 5 \
--gradient_accumulation_steps 1 \
--output_dir outputs/test
fi
|