Modernize MosaicBERT #1943

Job	Run time
pre-commit (3.9, benchmarks/bert)	3m 34s
pre-commit (3.9, benchmarks/resnet_cifar)	3m 50s
pre-commit (3.9, benchmarks/deeplab)	3m 49s
pre-commit (3.9, benchmarks/resnet_imagenet)	3m 51s
pre-commit (3.9, end-to-end-examples/sec_10k_qa)	3m 50s
pre-commit (3.9, end-to-end-examples/support-chatbot)	2m 20s
pre-commit (3.9, end-to-end-examples/stable_diffusion)	3m 50s
pre-commit (3.9, end-to-end-examples/stable_diffusion_dreambooth)	3m 50s
pre-commit (3.9, third-party/nemo)	2m 26s
pre-commit (3.9, third-party/gpt_neox)	2m 31s
pre-commit (3.9, third-party/trlx)	2m 27s
pre-commit (3.9, inference-deployments/instructor)	3m 52s
pre-commit (3.9, inference-deployments/mosaic_diffusion)	3m 52s
pre-commit (3.9, inference-deployments/mpt)	3m 48s
	47m 50s

Provide feedback