Samyamr/full precision for ZeRO Stage2 and Stage3

[samyam]

No description provided.

[stas00]

And just to indicate priority to this PR, we have all those bfloat16 models that won't train under fp16/mixed precision, and users want to use DeepSpeed to overcome GPU memory limitations, so they badly need this. Thank you!

[stas00]

When you feel this looks good enough to test please let me know and I will start testing this branch on the transformers side. Thank you.

[stas00]

This is awesome - thank you!

I encountered only one issue:

As I am writing HF transformers tests for fp32, I found that zero.Init doesn't get dtype from the config file, I have to explicitly do:

           ds_config = deepspeed_config()
            # XXX: Fixme - we shouldn't need to figure dtype out, it should be in the config file
            dtype = torch.float16 if ds_config.get("fp16", {}).get("enabled", True) else torch.float
            with deepspeed.zero.Init(dtype=dtype, config=ds_config):
                model = cls(config, *model_args, **model_kwargs)

I thought the whole point of passing config to zero.Init is so that we don't need to manually parse the file in multiple places, we we discussing this to work:

           ds_config = deepspeed_config()
            with deepspeed.zero.Init(config=ds_config):
                model = cls(config, *model_args, **model_kwargs)