Add special case for saving model when running with ZERO-3 optimisation.

[Willmish]

Tiny fix so a model sharded with Zero-3 can be saved! (To test if it doesnt also save the ugly placeholder confusing huggingface model loader)

[Willmish]

Weird thing is: when saving the model, it saves it fine, but also saves the placeholder model.safetensors file, hierarchy looks like:

config.json
generation_config.json
model-00001-of-00004.safetensors
model-00002-of-00004.safetensors
model-00003-of-00004.safetensors
model-00004-of-00004.safetensors
model.safetensors
model.safetensors.index.json
special_tokens_map.json
tokenizer_config.json
tokenizer.json

But model.safetensors shouldnt be there! (it causes things like AutoModelForCausalLM.from_pretrained( to fail, unless the placeholder file is removed!

[Adamliu1]

why are we not handling model saving in the main process? I wonder

[Willmish]

@Adamliu1 good q, I basedd this off of: https://huggingface.co/docs/accelerate/usage_guides/deepspeed#saving-and-loading , and I belive its because each process (on each of the 4 GPUs) has a portion of the model that needs to be synced across - hence cannot just unwrap it on a single process.