optim save for distributed filesystem

[Thaurun]

The new method can avoid gather, thereby saving additional GPU memory overhead.

[jthomy]

@Thaurun @ISEEKYAN I noticed that this commit breaks saving at least for Deepseek:
There is a race condition, we should have a torch.distributed.barrier() right after the assertion that the huggingface folder doesn't contain safetensor files.
I also observed on other issue where w_files[0][4] is None, causing TypeError: '>' not supported between instances of 'NoneType' and 'int'
I think this happens because tensor_model_parallel can be None for some weights because export_weights_without_gather returns that, which is then stored in the tuple as None.

[ISEEKYAN]

@Thaurun @ISEEKYAN I noticed that this commit breaks saving at least for Deepseek: There is a race condition, we should have a torch.distributed.barrier() right after the assertion that the huggingface folder doesn't contain safetensor files. I also observed on other issue where w_files[0][4] is None, causing TypeError: '>' not supported between instances of 'NoneType' and 'int' I think this happens because tensor_model_parallel can be None for some weights because export_weights_without_gather returns that, which is then stored in the tuple as None.

Thanks @jthomy for pointing out this issue, this has been fixed in #77