DeepSpeed

Fix for lm head weights for models such as llama

#4088

Open