Add DeepSpeed bf16 configuration

[dashstander]

For main DeepSpeed, setting the precision configuration parameter is insufficient to enable bf16 training. DeepSpeed uses a special bf16 dict configuration to enable using bf16, as documented here. This adds that configuration to NeoX and main DeepSpeed does the rest of the work.

[StellaAthena]

Since it seems like this is a dict with a single value, instead of making the user specify it wouldn’t it make more sense to have the code set the configs that DeepSpeed expects automatically when the precision is set to bf16?

[CLAassistant]

All committers have signed the CLA.

[StellaAthena]

I think at it’s core my question is whether this is actually more intuitive than:

precision = 'bfloat16'
precision_args = {whatever you need goes here}

[dashstander]

As things currently work we can't just have a precision_args argument that gets turned into a bf16 or fp16 dict for DeepSpeed because every argument that gets passed to DeepSpeed needs to exist in the megatron/neox_args/deepspeed_args.py file.

The changes I've made (and didn't push all of last night, whoops) make it so that a user only has to set the precision argument and the bare-bones of everything else gets set. It also raises an error if those dictionaries are set but the precision argument disagrees.

I'm open to other approaches though.

[dashstander]

I still need to test this to confirm there aren't any stupid bugs, but otherwise @StellaAthena and @Quentin-Anthony I'd appreciate if you took a look at this when you have a moment.

If people still don't like this--and honestly I don't love it--then I think we should open another issue where we figure out a way for people pass stuff straight into DeepSpeed without getting checked by NeoX or needing to have it be checked into our code. I know we've talked about it before. In that scenario we'd be able to just handle the default cases when people set precision and for advanced users who want to mess with the fp16 args dict they can do so.

[Quentin-Anthony]

What's this change for?

[Quentin-Anthony]

Why'd we remove these details?

[dashstander]

They're tied to the specific bf16 config that you wanted removed. I can move them to the section on the precision config?