Skip to content

Adopt proposed FastGen loss hooks conditionally #2100

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
CharlelieLrt wants to merge 1 commit into from
Closed

Adopt proposed FastGen loss hooks conditionally #2100

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
46 changes: 8 additions & 38 deletions physicsnemo/diffusion/distillation/_fastgen.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@

r"""Connect PhysicsNeMo models to FastGen distillation objectives."""

import importlib
from collections.abc import Callable, Mapping
from types import SimpleNamespace
from typing import Any
Expand All @@ -26,6 +27,7 @@
from torch import Tensor
from torch.nn.attention import SDPBackend, sdpa_kernel

from physicsnemo.core.version_check import check_version_spec
from physicsnemo.diffusion._fastgen import _FastGenNetwork
from physicsnemo.diffusion.base import (
DiffusionExtraTimes,
Expand All @@ -36,44 +38,12 @@
)
from physicsnemo.diffusion.noise_schedulers._fastgen import _FastGenNoiseScheduler


class _ObjectiveReduction:
r"""Compute masked statistics and weighted squared residuals.

Statistics use :math:`m=1_{w>0}` and count :math:`\sum m`. Residuals
use :math:`w` with the full element count for means. Positive weights
scale residuals linearly and leave the adaptive target unchanged.
"""

def __init__(
self,
reference: Float[Tensor, " B *dims"],
weight: Float[Tensor, " #B *#dims"] | None,
) -> None:
self.weight = (
torch.ones_like(reference)
if weight is None
else torch.broadcast_to(weight, reference.shape)
)
self.valid = self.weight > 0
self.dims = tuple(range(1, reference.ndim))
self.count = self.valid.sum(self.dims)
self.active = self.count > 0

def clean(self, value: Float[Tensor, " B *dims"]) -> Float[Tensor, " B *dims"]:
r"""Remove invalid entries before nonlinear operations."""
return torch.where(self.valid, value, torch.zeros_like(value))

def norm(self, value: Float[Tensor, " B *dims"]) -> Float[Tensor, " B"]:
r"""Compute the Euclidean norm on valid entries only."""
return torch.linalg.vector_norm(self.clean(value).flatten(1), dim=1)

def residual(
self, value: Float[Tensor, " B *dims"], *, mean: bool
) -> Float[Tensor, " B"]:
r"""Weight squared residuals linearly, retaining the backend denominator."""
weighted = self.clean(value).square() * self.weight
return weighted.mean(self.dims) if mean else weighted.sum(self.dims)
if check_version_spec("fastgen", "0.1.0", hard_fail=False):
_ObjectiveReduction = importlib.import_module(
"fastgen.methods.loss_support"
).LossReduction
else:
_ObjectiveReduction = None


class _FastGenBinding:
Expand Down
104 changes: 34 additions & 70 deletions physicsnemo/diffusion/distillation/methods/dmd2/_fastgen.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,6 @@
from typing import Any

import torch
import torch.nn.functional as F
from jaxtyping import Float, Shaped
from omegaconf import OmegaConf
from tensordict import TensorDict
Expand All @@ -40,7 +39,7 @@
)
from physicsnemo.diffusion.noise_schedulers._fastgen import _FastGenNoiseScheduler

from ..._fastgen import _configure_options, _FastGenBinding, _ObjectiveReduction
from ..._fastgen import _configure_options, _FastGenBinding

_AVAILABLE = check_version_spec("fastgen", "0.1.0", hard_fail=False)
if _AVAILABLE:
Expand Down Expand Up @@ -307,15 +306,13 @@ def single_train_step(
) -> tuple[dict[str, Tensor], dict[str, Tensor]]:
r"""Evaluate the selected student or auxiliary loss."""
args = (data["input_student"], data["t_student"], data["t"], data["eps"])
if self._phase == "student":
return self._student_update_step(
*args,
data,
condition=data["condition"],
neg_condition=data["neg_condition"],
)
return self._fake_score_discriminator_update_step(
*args, data["real"], condition=data["condition"]
return self._dispatch_training_phase(
self._phase,
*args,
data,
data["real"],
condition=data["condition"],
neg_condition=data["neg_condition"],
)

def _vsd_loss(
Expand All @@ -325,24 +322,14 @@ def _vsd_loss(
fake: Tensor,
scale: Tensor | None = None,
) -> Tensor:
r"""Exclude invalid entries from the target norm before weighting residuals."""
if self._call_weight is None:
return _backend.variational_score_distillation_loss(
generated, teacher, fake, additional_scale=scale
)
reduction = _ObjectiveReduction(generated, self._call_weight)
with torch.no_grad():
difference = reduction.clean(generated.float() - teacher.float())
mean = difference.abs().sum(reduction.dims) / reduction.count.clamp_min(1)
factor = (mean + 1e-6).reciprocal()
if scale is not None:
factor = factor * scale.float()
factor = factor.to(generated.dtype).reshape(
-1, *([1] * (generated.ndim - 1))
)
gradient = reduction.clean(fake - teacher) * factor
target = generated - gradient
return 0.5 * reduction.residual(generated - target, mean=True).mean()
r"""Compute weighted score distillation with normalization over valid entries."""
return _backend.variational_score_distillation_loss(
generated,
teacher,
fake,
additional_scale=scale,
loss_weight=self._call_weight,
)

def _logits(
self, x: Tensor, t: Tensor, condition: Any
Expand Down Expand Up @@ -421,14 +408,7 @@ def _student_update_step(
perturbed, t, teacher, neg_condition=neg_condition
)
if self.config.gan_loss_weight_gen > 0:
gan = self._logit_mean(
F.softplus(
-torch.where(weight > 0, logits, 0)
if weight is not None
else -logits
),
weight,
)
gan = _backend.gan_loss_generator(logits, loss_weight=weight)
scale = self._ratio_weight(logits, t, weight)
vsd = self._vsd_loss(generated, teacher, fake, scale)
return {
Expand Down Expand Up @@ -463,27 +443,18 @@ def _fake_score_discriminator_update_step(
prediction = self.fake_score(
perturbed, t, condition=condition, fwd_pred_type=prediction_type
)
if self._call_weight is None and prediction_type != "flow":
score_loss = _backend.denoising_score_matching_loss(
prediction_type,
prediction,
x0=generated,
eps=eps,
noise_scheduler=source,
t=t,
)
else:
targets = {"x0": generated, "eps": eps}
target = (
source.cond_velocity(generated, eps, t)
if prediction_type == "flow"
else targets[prediction_type]
)
score_loss = (
_ObjectiveReduction(prediction, self._call_weight)
.residual(prediction - target, mean=True)
.mean()
)
score_loss = _backend.denoising_score_matching_loss(
prediction_type,
prediction,
x0=generated,
eps=eps,
noise_scheduler=source,
t=t,
loss_weight=self._call_weight,
flow_target=source.cond_velocity(generated, eps, t)
if prediction_type == "flow"
else None,
)
gan = score_loss * 0
r1 = score_loss * 0
if self.config.gan_loss_weight_gen > 0:
Expand All @@ -497,19 +468,12 @@ def _fake_score_discriminator_update_step(
eps_real = torch.randn_like(real_data)
real_perturbed = source.forward_process(real_data, eps_real, t_real)
real_logits, real_weight = self._logits(real_perturbed, t_real, condition)
clean_fake = (
fake_logits
if fake_weight is None
else torch.where(fake_weight > 0, fake_logits, 0)
)
clean_real = (
real_logits
if real_weight is None
else torch.where(real_weight > 0, real_logits, 0)
gan = _backend.gan_loss_discriminator(
real_logits,
fake_logits,
real_weight=real_weight,
fake_weight=fake_weight,
)
gan = self._logit_mean(
F.softplus(clean_fake), fake_weight
) + self._logit_mean(F.softplus(-clean_real), real_weight)
if self.config.gan_r1_reg_weight > 0:
jittered = real_data + self.config.gan_r1_reg_alpha * torch.randn_like(
real_data
Expand Down
89 changes: 26 additions & 63 deletions physicsnemo/diffusion/distillation/methods/f_distill/_fastgen.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,7 @@ def __init__(
backend = importlib.import_module(
"fastgen.methods.distribution_matching.f_distill"
)
self._loss_support = importlib.import_module("fastgen.methods.loss_support")
ratio = self._ratio_options
if ratio.f_div not in backend.all_f_div_weighting_function:
raise ValueError(f"Unsupported f-divergence {ratio.f_div!r}.")
Expand All @@ -87,75 +88,37 @@ def __init__(
"Domain-parallel f-distill requires an explicit data-parallel statistics_group."
)

def _sum_statistics(self, value: Tensor) -> Tensor:
r"""Reduce over data replicas without mixing spatial partitions into counts."""
if torch.distributed.is_initialized():
torch.distributed.all_reduce(
value, group=self._ratio_options.statistics_group
)
return value

def _ratio_weight(
self,
logits: Float[Tensor, " B logits"],
t: Float[Tensor, " B"],
weight: Tensor | None,
) -> Float[Tensor, " B"]:
r"""Normalize ratios and divergence weights over valid samples only."""
with torch.no_grad():
if weight is None:
mean = logits.float().mean(1)
active = torch.ones_like(t, dtype=torch.bool)
else:
valid = weight > 0
count = valid.sum(1)
mean = torch.where(valid, logits.float(), 0).sum(1) / count.clamp_min(1)
active = count > 0
if self._call_weight is not None:
active = (
active
& _ObjectiveReduction(self._call_weight, self._call_weight).active
)
options = self._ratio_options
ratio = (
mean.clamp(-10, 10)
.exp()
.clamp(options.ratio_lower, options.ratio_upper)
)
if options.ratio_normalization:
source = self.net.noise_scheduler
lo, hi = source.min_t, source.max_t
width = (hi - lo) / options.bin_num
idx = (
((t.double() - lo) / width)
.floor()
.long()
.clamp(0, options.bin_num - 1)
)
count = torch.bincount(
idx, weights=active.float(), minlength=options.bin_num
)
total = torch.bincount(
idx,
weights=torch.where(active, ratio, 0),
minlength=options.bin_num,
)
statistics = self._sum_statistics(torch.stack((count, total)))
count, total = statistics.unbind(0)
bins = self._bins.to(device=t.device, dtype=torch.float32)
if self._update_state:
new_bins = bins * options.ratio_ema_rate + (
1 - options.ratio_ema_rate
) * total / (count + 1e-6)
bins = torch.where(count > 0, new_bins, bins)
self._bins = bins.detach()
ratio = ratio / (bins[idx] + 1e-6)
h = self.f_div_weighting_function(ratio).expand_as(ratio)
totals = self._sum_statistics(
torch.stack((torch.where(active, h, 0).sum(), active.float().sum()))
)
h_mean = totals[0] / totals[1].clamp_min(1)
return torch.where(active, h / (h_mean + 1e-6), 0).to(logits.dtype)
r"""Normalize divergence weights over valid samples, optionally updating running bins."""
options = self._ratio_options
sample_mask = (
None
if self._call_weight is None
else _ObjectiveReduction(self._call_weight, self._call_weight).active
)
h, bins = self._loss_support.f_divergence_weights(
logits,
t,
self.f_div_weighting_function,
ratio_lower=options.ratio_lower,
ratio_upper=options.ratio_upper,
bins=self._bins if options.ratio_normalization else None,
min_t=self.net.noise_scheduler.min_t,
max_t=self.net.noise_scheduler.max_t,
ema_rate=options.ratio_ema_rate,
logit_weight=weight,
sample_mask=sample_mask,
statistics_group=options.statistics_group,
update_state=self._update_state,
)
if options.ratio_normalization and self._update_state:
self._bins = bins.detach()
return h

def loss_state_dict(self) -> dict[str, Any]:
r"""Return a copy of the running density-ratio bins for checkpointing."""
Expand Down
Loading