update fsdp mixed precision

[vchiley]

What does this PR do?

Our MixedPrecision setting 'DEFAULT' should be updated to:

# If mixed_precision = 'default'
mixed_precision = MixedPrecision(
  param_dtype=autocast_precision,
  reduce_dtype=torch.float32,
  buffer_dtype=autocast_precision,
)

This will make 'DEFAULT' emulate AMP behavior.

use FULL if you need
use PURE if you want to live dangerously
standard AMP should be the 'DEFAULT' MixedPrecision setting (see above)

The current 'DEFAULT' shouldn't really be used.

What issue(s) does this change relate to?

Fixes https://mosaicml.atlassian.net/browse/CO-1896

Before submitting

• Have you read the contributor guidelines?
• Was this change discussed/approved in an issue first? It is much more likely to be merged if so.
• Did you update any related docs and document your change?
• Did you run the tests locally to make sure they pass?
• Did you run pre-commit on your change? (see the pre-commit section of prerequisites)

[dakinggg]

LGTM, would like abhi to approve as well

[vchiley]

fix to test failing in #2050

[bcui19]

LGTM!

[abhi-mosaic]

@vchiley can we get a side-by-side LLM run before and after this change? Aka can we do just a small 125M run with previous DEFAULT and new DEFAULT bheavior and post the screenshot here for posterity? Basically loss and MFU side by side

[vchiley]

testing behavior change here

gpt125m-flash-mpf is using mp: 'FULL'
gpt125m-flash-mpp is using mp: 'PURE'
gpt125m-flash-mpd is using mp: 'DEFAULT'

gpt125m-flash-mpp16b16 is using mp: {'param_dtype': 'bf16', 'buffer_dtype': 'bf16'} ie the proposed 'DEFAULT'
(here reduce_dtype will default to None and use fp32).

At on 8 GPUs at MosaicGPT 125M to chinchilla point shows no diff between any of the runs. The speed is virtually the same for all of the configs, except 'FULL' is ~3% slower.
(the old DEFAULT is ~1% slower than the new DEFAULT)
(In this exact setup the new DEFAULT is ~0.2% faster than PURE??? this is either within measurement noise or the casting is more expensive than one would hope.

I tried running the 7B model with the same set of 4 diff mp configs on 32 GPUs to see performance:

It seems like PURE and the new DEFAULT run at the same speed and are faster than the old DEFAULT and PURE is slowest
best I can explain this is: bwd pass is 2x fwd pass; BW pass isn't memory bound if param communication in bwd pass is bf16, bwd pass is compute bound.