Update train.py

train.py

@@ -17,17 +17,17 @@ EPOCHS = 1
 LEARNING_RATE = 2e-4
 FACTOR = 12 ** 3 // 3
 MAX_SEQ_LENGTH = 128
-VOCAB_SIZE = 32000
+VOCAB_SIZE = 52000
 INPUT_DATASET = "HuggingFaceTB/smollm-corpus"
 INSTRUCT_DATASET = "nroggendorff/elephant"
 OUTPUT_REPO = "nroggendorff/smallama"
-INSTRUCT_FINETUNE_BOOL = True
-INIT = 6
-SHARD_SIZE = int(15e+5)
+INSTRUCT_FINETUNE_BOOL = False
+INIT = 0
+SHARD_SIZE = int(2e+6)
 FP16 = True
-WARMUP_STEPS = 0
-WEIGHT_DECAY = 0.
-GRADIENT_ACCUMULATION_STEPS = 1
+WARMUP_STEPS = 50
+WEIGHT_DECAY = 1e-3
+GRADIENT_ACCUMULATION_STEPS = 2
 PUSH_TO_HUB = True
 
 class Space: