CPUAdam fp16 and bf16 support

[BacharL]

Hi.
Please review the following changes
I added support for BF16 to cpu adam. BF16, FP16 and float are supported at compilation time. the correct template is called at runtime according to input params dtype.

[tjruwase]

@BacharL, this amazing PR of yours has made me realize that ds_half_precision_t was not well thought out for optimizer offloading. It did not anticipate that device type could be anything other than fp16 and bf16. Consequently, we now require users to set a confusing compiler option: -DHALF_DTYPE=float.

In retrospect, it seems ds_device_precision_t is a better name, and compiler option could be -DDEVICE_DTYPE=float

Similarly, the half_precision bool variables scattered around the code is also redundant. Thanks for removing some of them in this PR.

@BacharL, what do you think?

[BacharL]

I will fix the naming of these, however, I think ds_half_precision_t causes discrepancy in buffer handling.
Step_AVX will handle the buffer with on cpu with FP16 of BF16 as implemented in AVX. The cast to fp32 is done in simd.h SIMD_LOAD_XX macros.
But the remaining part of the buffer will be handled in Step_1 with dtype from the device, which may be __half or __bfloat16. This dtype may not provide bit exact results as CPU implementation. The cast to fp32 is done in Step_1 inside the main loop.
I think even AVX/Native C++ cast may not be compatible. So if we already have this issue, why do we need to add device dtypes here? we can limit this code to c10::Half or c10::BFloat16 only.

We can also remove more half_precision varaibles by passing _params and dev_params with correct type instead of float* and fix their usage inside cuda and avx code

[tjruwase]

Thanks for raising questions about the discrepancies in this code base. I am aligned that the code is unduly complicated to the point of being bug prone. I think this occurred because of the uncontrolled evolution of the code. Hopefully, some measure of cleanliness can be restored through careful usage of templates, and torch data type support. In the meantime, I will answer your question below.

I think even AVX/Native C++ cast may not be compatible. So if we already have this issue, why do we need to add device dtypes here? we can limit this code to c10::Half or c10::BFloat16 only

I think adding device dtype is useful for making the data type conversions of offloading computation explicit.

1. Forward/backward computation on the device using device_dtype (currently fp16, bf16, or fp32) with output of gradients in device_dtype.
2. Conversion step of gradients from device_dtype to host_dtype which is input for subsequent optimizer computation on host.
3. Optimizer step computation on host using host_dtype (currently fp32) and returning param of host_dtype.
4. Conversion of (updated) host_dtype param into device_dtype param which is input for subsequent forward/backward computation on device.

[tjruwase]

@BacharL, thanks for this incredible improvement to the offloading optimizers and op builders. I left a few comments and questions, but overall looks good to me.

[BacharL]

Added templated invoker to help selecting the implementation
The map stores function pointers to templated functions, the key is the type enum. At initialization all supported dtypes are templated and inserted into the map.
I didn't clean ds_adagrad_step_plus_copy and related code under __ENABLE_CUDA__ but also couldn't test it.