OptimizedLinear updates

[jeffra]

This is a refresh of of OptimizedLinear with the following features to improve performance and usability:

• More efficient sharing of base weights using all_gather_into_tensor
• Flattened sharded weights
• Selectively offload frozen weights to cpu
• deepspeed.linear.Init that allows injecting OptimizedLinear during model construction (similar to zero.Init)
• Support for load state dict directly in OptimizedLinear, this allows loading HF model weights correctly into sharded params
• Various bug fixes for the LoRA implementation introduced previously
• Several new unit tests

Builds on-top of @RezaYazdaniAminabadi's previous FP8 updates (#5764) to support dense model fp8 quantization.

Example usage of this to fine-tune llama-3.1-405B on a single node: https://github.com/Snowflake-Labs/snowflake-arctic/tree/main/training/llama3.1

[jeffra]

nv-accelerate-v100 results on H100 (torch 2.4 + cu12.2):

nv-torch-latest-v100 results on H100 (torch 2.4 + cu12.2):

pytest --forked -n 4 unit/ --torch_ver="2.4" --cuda_ver="12.1" &> run1.log
pytest --forked -m 'sequential' unit/ --torch_ver="2.4" --cuda_ver="12.1" &> run2.log

[jeffra]

I'm able to get both the nv-accelerate-v100 and nv-torch-latest-v100 workflows to pass with this branch on my local H100 node (see previous comment). /cc @tjruwase @HeyangQin @loadams. Okay to force merge?

[loadams]

I'm able to get both the nv-accelerate-v100 and nv-torch-latest-v100 workflows to pass with this branch on my local H100 node (see previous comment). /cc @tjruwase @HeyangQin @loadams. Okay to force merge?

I believe I've fixed our runners @jeffra - I'll monitor it today to be sure it gets merged.