feat: Add MFT (Minifinetuning) loss support for knowledge distillation

[pbelcak]

What does this PR do ?

Adds MFT (Minifinetuning; NVR paper) loss support for knowledge distillation with configurable threshold-based teacher probability correction.

The MFT loss provides an alternative to standard KL divergence by correcting teacher distributions based on ground truth labels and a configurable threshold, potentially improving distillation quality for language models.

Collection:llm/modelopt

Changelog

• Add MFTLoss class implementing Minifinetuning distillation loss based on https://arxiv.org/abs/2506.15702
• Add _prepare_corrected_distributions method for threshold-based teacher probability correction
• Update DistillationConfig dataclass with use_mft and mft_threshold parameters
• Add MFT configuration validation in __post_init__ method
• Integrate MFT loss selection logic in load_distillation_config function
• Add MFTLoss import to utils.py for proper integration
• Add NotImplementedError for tensor model parallelism in MFTLoss (sequence parallelism required)

Usage

You can enable MFT loss in your distillation configuration:

In your distillation YAML config file:

use_mft: true
mft_threshold: 0.1  # Threshold between 0 and 1 for correction factor
...

Or, programmatically:

from nemo.collections.llm.modelopt.distill.utils import DistillationConfig

config = DistillationConfig(
    use_mft=True,
    mft_threshold=0.1,
    ...
)

GitHub Actions CI

The Jenkins CI system has been replaced by GitHub Actions self-hosted runners.

The GitHub Actions CI will run automatically when the "Run CICD" label is added to the PR.
To re-run CI remove and add the label again.
To run CI on an untrusted fork, a NeMo user with write access must first click "Approve and run".

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests? Tested under ModelOpt, see ModelOpt merge_requests/2785. No ModelOpt tests (except for one quantization check) are currently present in the NeMo Framework tests, so I did not bring the tests here.
• Did you add or update any necessary documentation? Every class and function is documented; additional examples are in ModelOpt docs.
• Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc): No
  • Reviewer: Does the PR have correct import guards for all optional libraries?

PR Type:

• New Feature
• Bugfix
• Documentation

[pbelcak]

(forgot to update CHANGELOG.md, made no other changes)

[pbelcak]

@sharathts all good?

[sharathts]

@pbelcak LGTM, approved!

[github-actions]

This PR is stale because it has been open for 14 days with no activity. Remove stale label or comment or update or this will be closed in 7 days.

[ZhiyuLi-Nvidia]

@pbelcak could you cherry-pick the commit to fix pylint and I will try again if we can solve the lint failure.

• branch: https://github.com/NVIDIA/NeMo/tree/pr-14298
• commit: 1f24add
  We don't have access to the branch of your cloned repo pbelcak:main
  Otherwise, we will take it over and create the PR copy to merge it in.

[github-actions]

This PR is stale because it has been open for 14 days with no activity. Remove stale label or comment or update or this will be closed in 7 days.

[pbelcak]

I see that the changes needed here are largely cosmetic; I plan to have a look at it this week.

[github-actions]

This PR is stale because it has been open for 14 days with no activity. Remove stale label or comment or update or this will be closed in 7 days.

[pbelcak]

Hi @ZhiyuLi-Nvidia, thanks for the fix, it's done. Can we merge?

[ZhiyuLi-Nvidia]

codecov/patch
codecov/patch — 17.77% of diff hit (target 80.00%)

Hi, @pbelcak all tests passed except the low test coverage. Could you finally add a unit test and rebase the branch?