FSDP optimizer state channels last fix

[pzharrington]

PhysicsNeMo Pull Request

Description

Fixes a silent corruption of optimizer state when checkpointing FSDP-wrapped models that have channels_last params (e.g. DiT in stormcast, or any SongUNet trained with apex GroupNorm).

PyTorch FSDP with use_orig_params=False packs and unpacks the FlatParameter asymmetrically for non-truly-contiguous params:

• Save (_get_unflat_views in torch/distributed/fsdp/_flat_param.py)
  uses as_strided((numel,),(1,)) -- reads bytes in storage order.
• Load (_flatten_tensor_optim_state in _optim_utils.py)
  uses torch.flatten -- reads bytes in logical (row-major) order.

For a 4-D Conv2d weight in channels_last format, those two byte orders differ. After a save/load round-trip the optimizer state for that param is silently scrambled, while the model weights survive (the model load path uses the storage-aware variant on both ends).

Fixes add a new helper to physicsnemo/utils/checkpoint.py that detects the case and performs appropriate permute+contiguous+view calls to load state correctly. Adds test coverage in package and recipe unit tests for the case. This PR also removes train dataloader seeding in the recipe as requested by SAs.

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.

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

[greptile-apps]

Greptile Summary

This PR fixes a silent corruption of FSDP optimizer state during checkpoint round-trips for models with channels_last parameters (e.g. DiT/SongUNet). The root cause is an asymmetry in PyTorch FSDP's FlatParameter pack (as_strided, storage order) vs. unpack (torch.flatten, logical order) for non-truly-contiguous params; the fix pre-permutes the saved tensors so the loader's torch.flatten reproduces the original storage byte sequence.

• The permute+contiguous+view transform in _remap_channels_last_optim_sd is mathematically sound for use_orig_params=False, but it is applied unconditionally: a use_orig_params=True FSDP model whose optimizer state happens to contain channels_last tensors would have those tensors incorrectly remapped, producing the same silent corruption the fix aims to prevent.
• examples/weather/stormcast/test_training.py also corrects a pre-existing bug where optimizer state was compared against itself (opt_params0 on both sides of the assertion) rather than against the loaded checkpoint.

Important Files Changed

Filename	Overview
physicsnemo/utils/checkpoint.py	Adds _remap_channels_last_optim_sd to work around an FSDP pack/unpack asymmetry for channels_last params; the permute+contiguous+view trick is mathematically sound for use_orig_params=False, but the remap is applied unconditionally without checking the FSDP use_orig_params flag, which could corrupt optimizer state for use_orig_params=True models.
test/utils/test_checkpoint_distributed.py	Adds test_fsdp_channels_last_optim_roundtrip covering the channels_last optimizer state round-trip; tolerance in an unrelated test was loosened from 1e-5 to 1e-4 without explanation; new test only uses NO_SHARD strategy.
examples/weather/stormcast/test_training.py	Fixes a pre-existing test bug (comparing opt_params0 with itself instead of opt_params1); extends test to world_size=1 with force_sharding parametrize; logic is correct.
examples/weather/stormcast/utils/trainer.py	Removes seed= from dataloader construction as requested by SAs; intentional reproducibility trade-off.
CHANGELOG.md	Adds changelog entry for the FSDP optimizer state channels_last fix.

_{Reviews (1): Last reviewed commit: "lint" | Re-trigger Greptile}

[pzharrington]

/ok to test cc703b5

[CharlelieLrt]

LGTM!

[pzharrington]

/ok to test d781754

[pzharrington]

/blossom-ci