Muennighoff
commited on
Commit
•
ff75c09
1
Parent(s):
922f03b
Create sbatch_4b284b6bc4.sh
Browse files- sbatch_4b284b6bc4.sh +162 -0
sbatch_4b284b6bc4.sh
ADDED
@@ -0,0 +1,162 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
#!/bin/bash
|
2 |
+
#SBATCH --nodes=32
|
3 |
+
#SBATCH --ntasks-per-node=1
|
4 |
+
#SBATCH --cpus-per-task=40 ##############
|
5 |
+
#SBATCH --mem=256G ##############
|
6 |
+
#SBATCH -p standard-g ##############
|
7 |
+
#SBATCH -t 48:00:00 ##############
|
8 |
+
#SBATCH --gpus-per-node=mi250:8 ##############
|
9 |
+
#SBATCH --exclusive=user
|
10 |
+
#SBATCH --hint=nomultithread ##############
|
11 |
+
#SBATCH --account=project_462000119 ##############
|
12 |
+
#SBATCH -o %j.out ##############
|
13 |
+
#SBATCH -e %j.err ##############
|
14 |
+
|
15 |
+
VARIANT=4b284b6bc4
|
16 |
+
|
17 |
+
# if run without sbatch, invoke here
|
18 |
+
if [ -z $SLURM_JOB_ID ]; then
|
19 |
+
mkdir -p logs
|
20 |
+
sbatch "$0"
|
21 |
+
exit
|
22 |
+
fi
|
23 |
+
|
24 |
+
set -euo pipefail
|
25 |
+
|
26 |
+
# symlink logs/latest.out and logs/latest.err
|
27 |
+
ln -f -s $SLURM_JOB_ID.out logs/latest.out
|
28 |
+
ln -f -s $SLURM_JOB_ID.err logs/latest.err
|
29 |
+
|
30 |
+
KILL_SWITCH_PATH=kill-switch-$VARIANT
|
31 |
+
CHECKPOINT_PATH=checkpoints_$VARIANT
|
32 |
+
TENSORBOARD_PATH=tensorboard_$VARIANT
|
33 |
+
|
34 |
+
# Data
|
35 |
+
VOCAB_FILE="gpt2/vocab.json"
|
36 |
+
MERGE_FILE="gpt2/merges.txt"
|
37 |
+
|
38 |
+
TRAIN_DATA_PATH=train6b.txt
|
39 |
+
# "train: 1.0 0:1 gpt2tok_c4_en_6B_text_document"
|
40 |
+
VALID_DATA_PATH=val.txt
|
41 |
+
# "validation: 1.0 0:1 gpt2tok_c4validation_text_document"
|
42 |
+
|
43 |
+
PP_SIZE=1
|
44 |
+
TP_SIZE=2
|
45 |
+
|
46 |
+
MICRO_BATCH_SIZE=2
|
47 |
+
GRADIENT_ACCUMULATION_STEPS=1
|
48 |
+
WORLD_SIZE=$((SLURM_GPUS_ON_NODE*SLURM_JOB_NUM_NODES))
|
49 |
+
GLOBAL_BATCH_SIZE=$((MICRO_BATCH_SIZE*WORLD_SIZE*GRADIENT_ACCUMULATION_STEPS))
|
50 |
+
|
51 |
+
# Model parameters
|
52 |
+
source model_params.sh
|
53 |
+
MODEL_PARAM=("${PARAM_4516M[@]}")
|
54 |
+
NHIDDEN=${MODEL_PARAM[0]}
|
55 |
+
FFN_HIDDEN_SIZE=${MODEL_PARAM[1]}
|
56 |
+
KV_SIZE=${MODEL_PARAM[2]}
|
57 |
+
NHEADS=${MODEL_PARAM[3]}
|
58 |
+
NLAYERS=${MODEL_PARAM[4]}
|
59 |
+
SEQ_LEN=2048
|
60 |
+
|
61 |
+
echo "Model parameters: d_model $NHIDDEN ffw_size $FFN_HIDDEN_SIZE kv_size $KV_SIZE n_heads $NHEADS n_layers $NLAYERS"
|
62 |
+
|
63 |
+
SAVE_INTERVAL=10000
|
64 |
+
|
65 |
+
# Tokens: 84_000_000_000
|
66 |
+
# -> Samples: 41_015_625.0
|
67 |
+
TRAIN_SAMPLES=41_015_625
|
68 |
+
|
69 |
+
OPTIMIZER_ARGS=" \
|
70 |
+
--optimizer adam \
|
71 |
+
--adam-beta1 0.9 \
|
72 |
+
--adam-beta2 0.95 \
|
73 |
+
--adam-eps 1e-8 \
|
74 |
+
--lr 2e-4 \
|
75 |
+
--min-lr 2e-5 \
|
76 |
+
--lr-decay-style cosine \
|
77 |
+
--lr-decay-samples $TRAIN_SAMPLES \
|
78 |
+
--lr-warmup-samples 410_156 \
|
79 |
+
--clip-grad 1.0 \
|
80 |
+
--weight-decay 1e-1 \
|
81 |
+
"
|
82 |
+
|
83 |
+
GPT_ARGS=" \
|
84 |
+
--num-layers $NLAYERS \
|
85 |
+
--hidden-size $NHIDDEN \
|
86 |
+
--num-attention-heads $NHEADS \
|
87 |
+
--kv-channels $KV_SIZE \
|
88 |
+
--ffn-hidden-size $FFN_HIDDEN_SIZE \
|
89 |
+
--seq-length $SEQ_LEN \
|
90 |
+
--max-position-embeddings $SEQ_LEN \
|
91 |
+
--micro-batch-size $MICRO_BATCH_SIZE \
|
92 |
+
--global-batch-size $GLOBAL_BATCH_SIZE \
|
93 |
+
--train-samples $TRAIN_SAMPLES \
|
94 |
+
--vocab-file $VOCAB_FILE \
|
95 |
+
--merge-file $MERGE_FILE \
|
96 |
+
--clip-grad 1.0 \
|
97 |
+
--kill-switch-path $KILL_SWITCH_PATH \
|
98 |
+
--bf16 \
|
99 |
+
$OPTIMIZER_ARGS \
|
100 |
+
"
|
101 |
+
|
102 |
+
OUTPUT_ARGS=" \
|
103 |
+
--log-interval 10 \
|
104 |
+
--save-interval $SAVE_INTERVAL \
|
105 |
+
--eval-interval 1000 \
|
106 |
+
--eval-iters 100 \
|
107 |
+
--tensorboard-dir $TENSORBOARD_PATH \
|
108 |
+
--tensorboard-queue-size 5 \
|
109 |
+
--log-timers-to-tensorboard \
|
110 |
+
--log-batch-size-to-tensorboard \
|
111 |
+
--log-validation-ppl-to-tensorboard \
|
112 |
+
"
|
113 |
+
|
114 |
+
ZERO_STAGE=0
|
115 |
+
|
116 |
+
mkdir -p ds_configs
|
117 |
+
DS_CONFIG_PATH="ds_configs/$SLURM_JOB_ID.json"
|
118 |
+
|
119 |
+
cat <<EOF > $DS_CONFIG_PATH
|
120 |
+
{
|
121 |
+
"train_micro_batch_size_per_gpu": $MICRO_BATCH_SIZE,
|
122 |
+
"train_batch_size": $GLOBAL_BATCH_SIZE,
|
123 |
+
"gradient_clipping": 1.0,
|
124 |
+
"zero_optimization": {
|
125 |
+
"stage": $ZERO_STAGE
|
126 |
+
},
|
127 |
+
"bf16": {
|
128 |
+
"enabled": true
|
129 |
+
},
|
130 |
+
"steps_per_print": 2000,
|
131 |
+
"wall_clock_breakdown": false
|
132 |
+
}
|
133 |
+
EOF
|
134 |
+
|
135 |
+
DEEPSPEED_ARGS=" \
|
136 |
+
--deepspeed \
|
137 |
+
--deepspeed_config $DS_CONFIG_PATH \
|
138 |
+
--zero-stage $ZERO_STAGE \
|
139 |
+
"
|
140 |
+
|
141 |
+
CMD=" \
|
142 |
+
Megatron-DeepSpeed/pretrain_gpt.py \
|
143 |
+
--tensor-model-parallel-size $TP_SIZE \
|
144 |
+
--pipeline-model-parallel-size $PP_SIZE \
|
145 |
+
$GPT_ARGS \
|
146 |
+
$OUTPUT_ARGS \
|
147 |
+
--save $CHECKPOINT_PATH \
|
148 |
+
--load $CHECKPOINT_PATH \
|
149 |
+
--train-weighted-split-paths-path $TRAIN_DATA_PATH \
|
150 |
+
--valid-weighted-split-paths-path $VALID_DATA_PATH \
|
151 |
+
--data-impl mmap \
|
152 |
+
$DEEPSPEED_ARGS \
|
153 |
+
"
|
154 |
+
|
155 |
+
echo $CMD
|
156 |
+
|
157 |
+
echo "START $SLURM_JOBID: $(date)"
|
158 |
+
|
159 |
+
# bash launch_srun.sh $CMD
|
160 |
+
srun --label launch.sh $CMD
|
161 |
+
|
162 |
+
echo "END $SLURM_JOBID: $(date)"
|