[CharlelieLrt]

Description

This PR introduces the alpha DistillationLoss API for training existing
PhysicsNeMo models with sCM or MeanFlow. Users install the optional
diffusion-fastgen extra, create a loss object, and call it from their own
optimizer loop:

loss_fn = DistillationLoss(student, scheduler, method="MeanFlow", teacher=teacher)
losses = loss_fn(x0, condition=condition, step=step)
losses["total_loss"].backward()

WeightedDistillationLoss supports pixel masks and soft weights, while the
BaseDistillationLoss contract lets users define native methods without FastGen.
The adapters remain private, so users continue to work with PhysicsNeMo models
and noise schedulers. The loss does not load teacher weights or manage
optimizers, and inference continues to use PhysicsNeMo samplers.

MeanFlow students receive (t, DiffusionExtraTimes(r=r, batch_size=[B])).
For sCM, models return their prediction paired with
DiffusionExtraOutputs(logvar=logvar, batch_size=[B]). Every forward call
returns a prediction tensor.

Checklist

• This PR follows the contribution guidelines.
• The documentation describes the API and its role in the stack.
• This PR remains a draft.

Dependencies

• Depends on

  • PhysicsNeMo #2095 adds inverse SNR and TrigFlow components. This is the immediate parent branch.

• Earlier prerequisites

  • PhysicsNeMo #2094 adds named diffusion inputs and outputs.

• Depends on this PR

  • PhysicsNeMo #2097 extends the loss API to DMD2 and f-distill. It follows this PR directly in the stack.

• Later follow-ups

  • PhysicsNeMo #2098 adds checkpoint bundles for restarting distillation.
  • PhysicsNeMo #2099 adds the optional training strategy and the grouped changelog entry.

• Parallel FastGen work

  • FastGen #43 allows the NumPy 2 dependency used by PhysicsNeMo. PhysicsNeMo Add sCM and MeanFlow distillation losses #2096 uses an immutable fork pin, so it does not need to wait for this proposal to merge.
  • FastGen #44 provides shared loss hooks for external training loops. The initial PhysicsNeMo integration does not require these hooks.

• Conditional follow-up

  • PhysicsNeMo #2100 replaces local compatibility code with the proposed FastGen loss hooks. It depends on PhysicsNeMo Add optional diffusion distillation strategy #2099 and FastGen [other] Add modulus.models.fcn_mip_plugin #44, not the other way around.

If FastGen maintainers decline FastGen #44, we would close PhysicsNeMo #2100.
The preceding PhysicsNeMo stack would remain available.

Review process

Reviewers can compare this PR with its immediate parent. After a parent merges,
we update the remaining branches and their PR bases to keep each diff focused.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[github-actions]

CODEOWNERS review map

Current for commit e03650a8fba5. An approval covers every file listed for that owner; one owner is sufficient for shared files.

⏳ @CharlelieLrt — 19 file(s)

• docs/api/diffusion/distillation.rst
• physicsnemo/diffusion/_fastgen.py
• physicsnemo/diffusion/distillation/init.py
• physicsnemo/diffusion/distillation/_fastgen.py
• physicsnemo/diffusion/distillation/base.py
• physicsnemo/diffusion/distillation/losses.py
• physicsnemo/diffusion/distillation/methods/init.py
• physicsnemo/diffusion/distillation/methods/meanflow/init.py
• physicsnemo/diffusion/distillation/methods/meanflow/_fastgen.py
• physicsnemo/diffusion/distillation/methods/scm/init.py
• physicsnemo/diffusion/distillation/methods/scm/_fastgen.py
• physicsnemo/diffusion/noise_schedulers/_fastgen.py
• test/diffusion/data/distillation_meanflow_weighted.pth
• test/diffusion/data/distillation_scm_weighted.pth
• test/diffusion/helpers.py
• test/diffusion/test_distillation_api.py
• test/diffusion/test_distillation_losses.py
• test/diffusion/test_fastgen_network.py
• test/diffusion/test_fastgen_noise_scheduler.py

⏳ @coreyjadams — 1 file(s)

• pyproject.toml

⏳ @ktangsali — 2 file(s)

• docs/api_diffusion.rst
• pyproject.toml

⏳ @megnvidia — 2 file(s)

• docs/api/diffusion/distillation.rst
• docs/api_diffusion.rst

No CODEOWNER

• .importlinter
• CHANGELOG.md
• uv.lock

Comment /codeowners-info to refresh.

[CharlelieLrt]

Continued in PhysicsNeMo #2201, from branch diffusion-distillation in CharlelieLrt/modulus. The replacement preserves this implementation and includes the requested documentation clarifications. Please continue the review there. This PR retains the earlier discussion.