Loss masking for distillation

[jlamypoirier]

✨ Description

Pass a loss mask to kwargs so we it can be used for distillation loss, aka cross-entropy from logits

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

[oleksost]

IN the documentation, here, the parameter use_loss_masking_spans should be under batch and not under sampling, right?

[jlamypoirier]

@oleksost Yes, looks like we forgot to update

[oleksost]

In return _CROSS_ENTROPY_IMPLEMENTATIONS[implementation] we also need to pass loss mask?

[oleksost]

Need to pass loss_mask?

[oleksost]

THis assert should be removed, no?

[oleksost]

Works well for me.