10-20x faster load checkpoint (for critic/reward model) (#675)

Co-authored-by: Lev Kurilenko <[email protected]>

Author: awan-10

applications/DeepSpeed-Chat/training/step1_supervised_finetuning/training_scripts/single_node/run_llama_7b.sh renamed to applications/DeepSpeed-Chat/training/step1_supervised_finetuning/training_scripts/single_node/llama2/run_llama2_7b.sh

@@ -6,17 +6,17 @@
 OUTPUT=$1
 ZERO_STAGE=$2
 if [ "$OUTPUT" == "" ]; then
-    OUTPUT=./output_step1_llama_7b_epoch4_lr9.65e-6_test
+    OUTPUT=./output_step1_llama2_7b
 fi
 if [ "$ZERO_STAGE" == "" ]; then
     ZERO_STAGE=3
 fi
 mkdir -p $OUTPUT
 
-deepspeed --include="worker-1" main.py \
+deepspeed main.py \
    --data_path Dahoas/rm-static Dahoas/full-hh-rlhf Dahoas/synthetic-instruct-gptj-pairwise yitingxie/rlhf-reward-datasets \
    --data_split 2,4,4 \
-   --model_name_or_path decapoda-research/llama-7b-hf \
+   --model_name_or_path meta-llama/Llama-2-7b-hf \
    --per_device_train_batch_size 4 \
    --per_device_eval_batch_size 4 \
    --max_seq_len 512 \

applications/DeepSpeed-Chat/training/step1_supervised_finetuning/training_scripts/single_node/llama2/run_llama2_7b_lora.sh

@@ -0,0 +1,36 @@
+#!/bin/bash
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+OUTPUT=$1
+ZERO_STAGE=$2
+if [ "$OUTPUT" == "" ]; then
+    OUTPUT=./output_step1_llama2_7b_lora
+fi
+if [ "$ZERO_STAGE" == "" ]; then
+    ZERO_STAGE=3
+fi
+mkdir -p $OUTPUT
+
+deepspeed main.py \
+   --data_path Dahoas/rm-static Dahoas/full-hh-rlhf Dahoas/synthetic-instruct-gptj-pairwise yitingxie/rlhf-reward-datasets \
+   --data_split 2,4,4 \
+   --model_name_or_path meta-llama/Llama-2-7b-hf \
+   --per_device_train_batch_size 4 \
+   --per_device_eval_batch_size 4 \
+   --max_seq_len 512 \
+   --learning_rate 9.65e-6 \
+   --weight_decay 0. \
+   --num_train_epochs 4  \
+   --gradient_accumulation_steps 1 \
+   --lr_scheduler_type cosine \
+   --num_warmup_steps 0 \
+   --seed 1234 \
+   --gradient_checkpointing \
+   --zero_stage $ZERO_STAGE \
+   --deepspeed \
+   --lora_dim 128 \
+   --lora_module_name "layers." \
+   --output_dir $OUTPUT \
+   &> $OUTPUT/training.log

applications/DeepSpeed-Chat/training/step1_supervised_finetuning/training_scripts/single_node/run_1.3b.sh renamed to applications/DeepSpeed-Chat/training/step1_supervised_finetuning/training_scripts/single_node/opt/run_1.3b.sh

@@ -14,9 +14,9 @@ fi
 mkdir -p $OUTPUT
 
 deepspeed main.py \
-   --data_path Dahoas/rm-static Dahoas/full-hh-rlhf Dahoas/synthetic-instruct-gptj-pairwise yitingxie/rlhf-reward-datasets \
+   --data_path Dahoas/rm-static \
    --data_split 2,4,4 \
-   --model_name_or_path decapoda-research/llama-7b-hf \
+   --model_name_or_path meta-llama/Llama-2-7b-hf \
    --per_device_train_batch_size 8 \
    --per_device_eval_batch_size 8 \
    --max_seq_len 512 \
@@ -31,5 +31,6 @@ deepspeed main.py \
    --gradient_checkpointing \
    --zero_stage $ZERO_STAGE \
    --deepspeed \
+   --offload \
    --output_dir $OUTPUT \
    &> $OUTPUT/training.log

applications/DeepSpeed-Chat/training/step1_supervised_finetuning/training_scripts/single_node/run_1.3b_lora.sh renamed to applications/DeepSpeed-Chat/training/step1_supervised_finetuning/training_scripts/single_node/opt/run_1.3b_lora.sh

@@ -0,0 +1,38 @@
+#!/bin/bash
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+OUTPUT=$1
+ZERO_STAGE=$2
+if [ "$OUTPUT" == "" ]; then
+    OUTPUT=./output_step2_llama_7b_epoch1_lr9.65e-6
+fi
+if [ "$ZERO_STAGE" == "" ]; then
+    ZERO_STAGE=3
+fi
+mkdir -p $OUTPUT
+
+deepspeed main.py \
+   --data_path Dahoas/rm-static \
+   --data_split 2,4,4 \
+   --model_name_or_path meta-llama/Llama-2-7b-hf \
+   --per_device_train_batch_size 8 \
+   --per_device_eval_batch_size 8 \
+   --max_seq_len 512 \
+   --learning_rate 9.65e-6 \
+   --weight_decay 0.1 \
+   --num_padding_at_beginning 0 \
+   --num_train_epochs 1  \
+   --gradient_accumulation_steps 1 \
+   --lr_scheduler_type cosine \
+   --num_warmup_steps 0 \
+   --seed 1234 \
+   --gradient_checkpointing \
+   --zero_stage $ZERO_STAGE \
+   --deepspeed \
+   --offload \
+   --lora_dim 128 \
+   --lora_module_name "layers." \
+   --output_dir $OUTPUT \
+   &> $OUTPUT/training.log