Fix SongUNet with ShardTensor when using zero embedding

[jleinonen]

PhysicsNeMo Pull Request

When training regression models (i.e. no time step embedding, embedding_type == "zero") using ShardTensor, SongUNet was giving this error:

  File "/usr/local/lib/python3.12/dist-packages/physicsnemo/models/diffusion_unets/song_unet.py", line 663, in forward
    x = block(x, emb)
        ^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1776, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1787, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/physicsnemo/nn/module/unet_layers.py", line 242, in forward
    params = self.affine(emb).unsqueeze(2).unsqueeze(3)
             ^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1776, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1787, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/physicsnemo/nn/module/fully_connected_layers.py", line 527, in forward
    x = x @ weight.t()

caused by the emb tensor being a plain torch.Tensor:

physicsnemo/physicsnemo/models/diffusion_unets/song_unet.py

Lines 627 to 632 in 70b06ed

	else:
	emb = torch.zeros(
	(noise_labels.shape[0], self.emb_channels),
	device=x.device,
	dtype=x.dtype,
	)

I added a conversion of emb to ShardTensor if x is a ShardTensor. With this fix, it is possible to train the StormCastUNet type regression models with ShardTensor.

Description

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.
• The CHANGELOG.md is up to date with these changes.
• An issue is linked to this pull request.
• If I am implementing a new model or modifying any existing model, I have followed the Models Implementation Coding Standards.

Dependencies

Review Process

All PRs are reviewed by the PhysicsNeMo team before merging.

Depending on which files are changed, GitHub may automatically assign a maintainer for review.

We are also testing AI-based code review tools (e.g., Greptile), which may add automated comments with a confidence score.
This score reflects the AI’s assessment of merge readiness and is not a qualitative judgment of your work, nor is
it an indication that the PR will be accepted / rejected.

AI-generated feedback should be reviewed critically for usefulness.
You are not required to respond to every AI comment, but they are intended to help both authors and reviewers.
Please react to Greptile comments with 👍 or 👎 to provide feedback on their accuracy.

[greptile-apps]

Greptile Summary

This PR fixes a tensor type mismatch error that occurred when training regression models with SongUNet using embedding_type == "zero" and ShardTensor. The error was caused by the emb tensor being a plain torch.Tensor while the input x was a ShardTensor, causing incompatibility in the matrix multiplication operation in the Linear layer (x @ weight.t()).

The fix adds a conditional check: when x is a ShardTensor, the emb tensor (created as zeros) is converted to a ShardTensor using ShardTensor.from_local() with the same device mesh as x. This ensures type consistency between tensors in subsequent operations.

• Adds ShardTensor import from physicsnemo.domain_parallel.shard_tensor
• Adds conversion logic after creating zero embedding tensor (lines 634-635)
• The fix is minimal and targeted, only affecting the zero embedding path
• The non-zero embedding path doesn't need this fix since the mapping layers (map_layer0, map_layer1) are Linear modules that preserve the tensor type through operations

Important Files Changed

Filename	Overview
physicsnemo/models/diffusion_unets/song_unet.py	Adds ShardTensor conversion for zero embedding case to fix tensor type mismatch error

_{Last reviewed commit: 8db65bf}

[greptile-apps]

_{1 file reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[pzharrington]

/blossom-ci