Fix and improve loading of distributed checkpoints

[jlamypoirier]

✨ Description

Fix #293

Lots of improvements on the loading of distributed checkpoints in different format.

• Load files only if they are actually needed for conversion. This should speed things up a lot for large world sizes.
• Implement a new, much faster loading method for the common case which just copies contiguous slices.
• Keep track of the per-parameter loaded count so SafeLoad can verify with _check_parameters
• Fix the tensor-parallel case ([bug] Conversion of distributed checkpoints to huggingface #293). The problem was in _get_parameter_shard_indices_in_full_weight in which the shard indices were sometimes set in a copy of the index (due to flatten) rather than a view, which caused the loaded tensors to be completely ignored.
• Add more tests involving changes in TP size.

For the case of an unchanged distributed config (ex. starting a new experiment from a distributed checkpoint), loading should now be almost as fast as the unsafe version.

This will also help a lot with elastic training (#241) by cutting most of the resuming time.

🔍 Type of change

Select all that apply:

• 🐛 Bug fix (non-breaking change that addresses a specific issue)
• 🚀 New feature (non-breaking change that adds functionality)
• ⚠️ Breaking change (a change that could affect existing functionality)
• 📈 Performance improvement/optimization (improves speed, memory usage, or efficiency)
• 🛠️ Code refactor (non-functional changes that improve code readability, structure, etc.)
• 📦 Dependency bump (updates dependencies, including Dockerfile or package changes)
• 📝 Documentation change (updates documentation, including new content or typo fixes)
• 🔧 Infrastructure/Build change (affects build process, CI/CD, or dependencies)