Update train.py

The official pretraining script. Use it like this
```bash
python3 train.py \
--dataset=fineweb-edu \
--n_layer=24 \
--n_head=16 \
--n_embd=1024 \
--block_size=1024 \
--batch_size=8 \
--gradient_accumulation_steps=16 \
--learning_rate=6e-4 \
--max_iters=300000 \
--eval_interval=1000 \
--eval_iters=100 \
--log_interval=5 \
--weight_decay=0.1 \
--warmup_iters=2000 \
--lr_decay_iters=300000 \
--min_lr=6e-5 \
--dtype=bfloat16 \
--compile=True \
--always_save_checkpoint=True \
--init_from=scratch \
--out_dir=/home/user/350m_fineweb
```

train.py

@@ -386,6 +386,6 @@ while True:
     # termination conditions
     if iter_num > max_iters:
         break
- 
+
 if ddp:
     destroy_process_group()