Muennighoff commited on
Commit
ff75c09
1 Parent(s): 922f03b

Create sbatch_4b284b6bc4.sh

Browse files
Files changed (1) hide show
  1. 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)"