Fix Non-Contiguous Tensor Issue in Checkpoint Consolidation

[gioannides]

What does this PR do?

This PR fixes an issue in the consolidate_tensor_parallel_checkpoints function where non-contiguous tensors caused a ValueError when saving consolidated model checkpoints, particularly with the Safetensors format. The fix ensures that tensors are made contiguous before and after concatenation, resolving memory layout issues during saving.

Error Example:

When running the following command:

consolidate_model_parallel_checkpoints_to_unified_checkpoint(
    Path("./model_sharded/model/shards"), 
    "./consolidated", 
    "safetensors"
)

The following error occurs:

ValueError: You are trying to save a non contiguous tensor: model.layers.0.self_attn.o_proj.weight which is not allowed.

This PR solves the issue by applying the .contiguous() method to tensors at key points in the consolidation process, ensuring the consolidated model can be saved without errors.

Motivation and Context:

When consolidating distributed model checkpoints (e.g., from model shards), tensors may not be contiguous in memory. Safetensors, which is a memory-efficient way to store models, does not allow saving non-contiguous tensors. This PR modifies the consolidate_tensor_parallel_checkpoints function to apply .contiguous() to tensors to ensure that the checkpoints can be properly consolidated and saved.

This change impacts the overall consolidation pipeline, as consolidate_model_parallel_checkpoints_to_unified_checkpoint depends on the lower-level functions to process tensors correctly.

[dacorvo]

Thank you very much for the pull-request. It is ok for me, but gently pinging @michaelbenayoun for confirmation.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[dacorvo]

@michaelbenayoun gentle reminder ...

[michaelbenayoun]

Thanks for taking care of that!

[michaelbenayoun]

Could you just run themake style command to fix check_code_quality please.

[gioannides]

@michaelbenayoun np! where can I see instructions on how to do that?

[michaelbenayoun]

The steps to follow are:

1. Put yourself inside the optimum-neuron directory
2. Run pip install -e ".[quality]"
3. Run make style

[gioannides]

@michaelbenayoun done

[dacorvo]

Something went wrong when you ran the command, because it changed many files.
Rebasing and merging as #736