[bert] update the zero pretraining and finetuning with new zero

[FrankLeeeee]

Fixed the training script such that len(dataloader) works fine. This script is updated with the new zero api as well.