Fix loss masking

[oleksost]

✨ Description

Addresses #442

• loss_masks should include padding and image placeholder tokens

TODO:

• not sure if loss masking spans and image placeholders should also be ignored in the layer-wise losses

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

📝 Changes

List the key changes introduced in this PR:

1. Change A
2. Change B

✅ Checklist

Make sure the following tasks are completed before submitting the PR:

General

• 📜 I have read and followed the contributing guidelines.
• 🏷️ I am using a clear and descriptive PR title that summarizes the key change or feature introduced.
• 🎉 The functionality is complete, and I have tested the changes.
• 📝 I have updated the documentation if needed.
• ⚠️ The change does not introduce any new issues (e.g., runtime warnings, type checker errors, linting problems, unhandled edge cases).
• 🧩 I have commented my code, especially in hard-to-understand areas.

Dependencies and Configuration

• 🐋 I have updated the Docker configuration or dependencies, if applicable.
• 🔄 I have ensured compatibility with the existing setup after dependency changes.

Testing

• 🧪 I have added or updated tests to cover my changes.
• ✔️ New and existing tests pass locally with my changes.
• 🚦 I have tested these changes on GPUs and verified training stability.
• 🏋️ I have tested the changes on realistic training workloads, if applicable.

Performance Impact

• 📊 I have run benchmarks where applicable to evaluate the performance impact.
• ✅ The benchmarks show no performance regression.
• 🚀 The benchmarks indicate a potential performance improvement.
• ⚠️ The benchmarks indicate a potential performance degradation.
• 📈 I have provided benchmark results and detailed any performance impact below, if applicable.

📊 Performance Impact Details

If there is any impact on performance, describe it and provide benchmark results, if applicable:

---

🗒️ Additional Notes

Include any additional context, information, or considerations here, such as known issues, follow-up tasks, or backward compatibility concerns.

[jlamypoirier]

Is this really what we want? We can't train the model to produce these labels, but it might make sense to compute other losses?

[jlamypoirier]

Can we skip this when not needed?

[oleksost]

Addressed.

AFAIU we want to incorporate all sources of masking into distillation losses:

• padding & spans
• also image placeholder tokens, because we train Image&text => text model, i.e. also in pre-training etc. we do not compute loss on image tokens

[jlamypoirier]

Activation distillation ignores loss_mask, it uses activation_mask instead.

[jlamypoirier]

TODO:

• not sure if loss masking spans and image placeholders should also be ignored in the layer-wise losses

Does that even make sense? These refer to token prediction which isn't really a thing at the activation stage. I guess we could take the next token but that raises several concerns (especially with MTP).

Actually I think we shouldn't mask those. They may not be used for next token prediction, but the keys and values resulting from these activations are used in further down in the sequence, which means we do train these activations.