Bf16

[sdtblck]

Currently working with everything but partitioning checkpointed activations + pipeline parallel - but those fixes will be in deepspeed anyway, so imo this should mostly be ready to merge.

• gets bf16 working for data parallel + model parallel training
• adds new layernorm from https://github.com/NVIDIA/Megatron-LM that works with bf16
• adds bf16 support for fused attention kernels

All nccl communication for bf16 now happens in fp32. So it's a little slower if you're communication bound. Untested with zero optimizer as of right now.

Will need to install https://github.com/EleutherAI/DeeperSpeed/tree/bf16 for it to work, currently.

TODO:

• merge deeperspeed bf16 branch to main
• test with zero optimizer
• get partition activations working
• get pipe parallel working

[sdtblck]

Update: It doesn't work with zero 😆

[sdtblck]

Ok, now working for zero 1 + 2.

[sdtblck]

aaand now pipeline parallel + activation checkpointing too!
Aside from Zero 3, everything works. I'll wait till we have zero 3 working in our codebase to take a look at that anyway, since i suspect it will be a fair bit of work.