[CharlelieLrt]

Description

This PR adds the optional DistillationStrategy for users who do not need to
write their own distillation loop. It extends Alchemi's general training strategy
to manage phase-specific optimizers, gradient accumulation, AMP, EMA, and
checkpoint restart.

Users install the diffusion-strategy extra and supply public models, a native
or FastGen-backed loss, and TensorDict batches:

strategy = DistillationStrategy(
    models={"student": student}, loss_fn=loss_fn, num_steps=1000,
    optimizer_configs={"student": [{"optimizer_cls": torch.optim.Adam,
                                    "optimizer_kwargs": {"lr": 1e-4}}]},
)
strategy.run(batches)

For explicit MeanFlow times, store t and r as separate batch tensors.
Use training_fn(models, batch) to return loss arguments containing
(t, DiffusionExtraTimes(r=r, batch_size=[B])). Ordinary single-time batches
need no conversion.

The strategy remains optional: expert users can call the same loss from their
own loop. Compilation and distributed setup remain under user control.
This PR also carries the grouped changelog entry for the distillation feature.

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 #2203 adds checkpoint bundles for restarting distillation. This is the immediate prerequisite.

• 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.

• Parallel FastGen work

  • FastGen #43 allows the NumPy 2 dependency used by PhysicsNeMo. PhysicsNeMo Add sCM and MeanFlow distillation losses #2201 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 #2205 replaces local compatibility code with the proposed FastGen loss hooks. It depends on PhysicsNeMo Add optional diffusion distillation strategy #2204 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 #2205.
The preceding PhysicsNeMo stack would remain available.

Review process

This draft replaces PhysicsNeMo #2099 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 683709da6e7b. 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.