improvements to MTP implementation

[RaymondLi0]

✨ Description

Follow-up to #179 with some fixes and improvements:

• Add more labels in data-sampling to avoid input-truncation
• handle sequence-parallelism and cross-entropy splits (as a consequence of the above)
• Convert the additional lm-heads and add a corresponding hf-transformer implementation

Potential future work:

• add weights for each future-token prediction loss (e.g. the next token could be more important than the future ones)
• Improve throughput

🔍 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)

[RaymondLi0]

Since there is no available implementation in the transformers library, I wrote a custom modeling.py for llama with multi-token prediction.
The converter then copies these files in the export directory.
Wasn't sure what the best approach would be here, but that's the only thing I came up with. Let me know if there would be a better way.

[jlamypoirier]

Looks good, some minor comments

[tscholak]

LGTM!

[jlamypoirier]

Looks great, thanks!