[CharlelieLrt]

Description

This PR adopts the shared loss hooks proposed in FastGen #44 and removes the
corresponding private compatibility code from PhysicsNeMo. It does not change
the public DistillationLoss API, model roles, masks, or checkpoint state.

This is a conditional follow-up, not a prerequisite for the initial integration.
The earlier PhysicsNeMo PRs work without the proposed hooks. The dependency pin
combines FastGen #43 and #44 on the user fork so development can proceed in
parallel. Before merging this PR, we would replace that pin with an accepted
upstream revision. If FastGen maintainers decline the hooks, we would close
this PR and keep the earlier implementation.

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 #2204 adds the optional training strategy and the grouped changelog entry. This is the immediate prerequisite.
  • FastGen #44 provides shared loss hooks for external training loops. This PR requires acceptance of those hooks.
  • FastGen #43 allows the NumPy 2 dependency used by PhysicsNeMo. The current fork pin includes this metadata change as well.

• Earlier prerequisites

  • PhysicsNeMo #2199 adds named diffusion inputs and outputs.
  • PhysicsNeMo #2200 adds inverse SNR and TrigFlow components.
  • PhysicsNeMo #2201 adds private adapters and public sCM and MeanFlow losses.
  • PhysicsNeMo #2202 extends the loss API to DMD2 and f-distill.
  • PhysicsNeMo #2203 adds checkpoint bundles for restarting distillation.

• Depends on this PR

  • No PR in this stack depends on this conditional follow-up.

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

Review process

This draft replaces PhysicsNeMo #2100 and comes from the personal fork. This migration preserves the implementation. The first commit also includes the requested documentation clarifications.

Review the isolated commit diff for this layer. Each branch adds one commit to its prerequisite branch. All replacement PRs target upstream main, so later PRs include prerequisite commits awaiting merge in their Files changed tab. After prerequisites merge, update the remaining branches to remove them from those diffs.

[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 bb9003f603ef. An approval covers every file listed for that owner; one owner is sufficient for shared files.

⏳ @CharlelieLrt — 59 file(s)

• docs/api/diffusion/distillation.rst
• docs/api/diffusion/introduction.rst
• docs/api/diffusion/models.rst
• docs/api/diffusion/noise_schedulers.rst
• docs/api/diffusion/preconditioners.rst
• physicsnemo/diffusion/init.py
• physicsnemo/diffusion/_fastgen.py
• physicsnemo/diffusion/base.py
• physicsnemo/diffusion/distillation/init.py
• physicsnemo/diffusion/distillation/_alchemi.py
• physicsnemo/diffusion/distillation/_fastgen.py
• physicsnemo/diffusion/distillation/base.py
• physicsnemo/diffusion/distillation/checkpoint.py
• physicsnemo/diffusion/distillation/losses.py
• physicsnemo/diffusion/distillation/methods/init.py
• physicsnemo/diffusion/distillation/methods/dmd2/init.py
• physicsnemo/diffusion/distillation/methods/dmd2/_fastgen.py
• physicsnemo/diffusion/distillation/methods/f_distill/init.py
• physicsnemo/diffusion/distillation/methods/f_distill/_fastgen.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/distillation/strategy.py
• physicsnemo/diffusion/metrics/losses.py
• physicsnemo/diffusion/noise_schedulers/init.py
• physicsnemo/diffusion/noise_schedulers/_fastgen.py
• physicsnemo/diffusion/noise_schedulers/edm.py
• physicsnemo/diffusion/noise_schedulers/iddpm.py
• physicsnemo/diffusion/noise_schedulers/linear_gaussian.py
• physicsnemo/diffusion/noise_schedulers/rectified_flow.py
• physicsnemo/diffusion/noise_schedulers/trigflow.py
• physicsnemo/diffusion/noise_schedulers/ve.py
• physicsnemo/diffusion/noise_schedulers/vp.py
• physicsnemo/diffusion/preconditioners/init.py
• physicsnemo/diffusion/preconditioners/preconditioners.py
• physicsnemo/diffusion/utils/model_wrappers.py
• test/diffusion/data/distillation_meanflow_weighted.pth
• test/diffusion/data/distillation_scm_weighted.pth
• test/diffusion/data/trigflow_precond_conv.mdlus
• test/diffusion/data/trigflow_precond_conv_coefficients.pth
• test/diffusion/data/trigflow_precond_conv_forward.pth
• test/diffusion/data/trigflow_precond_conv_sigma.pth
• test/diffusion/data/trigflow_precond_linear.mdlus
• test/diffusion/data/trigflow_precond_linear_coefficients.pth
• test/diffusion/data/trigflow_precond_linear_forward.pth
• test/diffusion/data/trigflow_precond_linear_sigma.pth
• test/diffusion/helpers.py
• test/diffusion/test_distillation_api.py
• test/diffusion/test_distillation_checkpoint.py
• test/diffusion/test_distillation_losses.py
• test/diffusion/test_distillation_strategy.py
• test/diffusion/test_distribution_distillation.py
• test/diffusion/test_fastgen_network.py
• test/diffusion/test_fastgen_noise_scheduler.py
• test/diffusion/test_model_io.py
• test/diffusion/test_noise_schedulers.py
• test/diffusion/test_preconditioners.py
• test/diffusion/test_trigflow.py

⏳ @coreyjadams — 1 file(s)

• pyproject.toml

⏳ @ktangsali — 2 file(s)

• docs/api_diffusion.rst
• pyproject.toml

⏳ @megnvidia — 6 file(s)

• docs/api/diffusion/distillation.rst
• docs/api/diffusion/introduction.rst
• docs/api/diffusion/models.rst
• docs/api/diffusion/noise_schedulers.rst
• docs/api/diffusion/preconditioners.rst
• docs/api_diffusion.rst

No CODEOWNER

• .importlinter
• CHANGELOG.md
• uv.lock

Comment /codeowners-info to refresh.