Add amp_mode support to experimental DiffusionUNet3D and UNetBlock3D - #1937
Add amp_mode support to experimental DiffusionUNet3D and UNetBlock3D#1937pjaysaval wants to merge 2 commits into
Conversation
CODEOWNERS review mapCurrent for commit ⏳ @CharlelieLrt — 2 file(s)
No CODEOWNER
Comment |
Greptile SummaryThe PR adds mixed-precision support to the experimental 3D diffusion U-Net and its residual block by propagating
Important Files Changed
Reviews (1): Last reviewed commit: "Add amp_mode support to experimental Dif..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| DiffusionUNet3D, seed=0, amp_mode=True, **arch_kwargs | ||
| ).to(device) | ||
| data = _generate_batch_data(arch_kwargs, x_shape, GLOBAL_SEED, device) | ||
| with torch.autocast(device_type=device, dtype=torch.bfloat16): |
There was a problem hiding this comment.
When the CUDA-enabled device fixture supplies cuda:0, this test passes that indexed string directly as torch.autocast's device_type, which requires a device category such as cuda; the autocast context therefore raises before exercising the model and fails CUDA test runs. The same pattern occurs in the default-mode test and both new block-level AMP tests.
There was a problem hiding this comment.
Verified on an A100 with torch 2.11 (the version in ci-requirements.lock) and torch 2.13: torch.autocast(device_type="cuda:0", ...) behaves identically to "cuda" — autocast is enabled and ops produce bf16 — so these tests pass as written on CUDA runs. Normalized the device string anyway in 90db634 for robustness on stricter torch versions and consistency with the existing test-suite idiom.
b68cf42 to
90db634
Compare
The experimental 3D diffusion U-Net could not run under torch.autocast: its internal physicsnemo.nn layers (Linear, PositionalEmbedding, FourierEmbedding) validate their amp_mode flag and raise when autocast is active, but neither DiffusionUNet3D nor UNetBlock3D exposed the flag, so mixed-precision training always failed with 'RuntimeError: amp_mode=False but torch autocast is enabled'. Thread amp_mode through both classes, mirroring SongUNet and the 2D UNetBlock: DiffusionUNet3D passes it to the noise/condition embeddings, the map linears, and every UNetBlock3D; UNetBlock3D passes it to its affine Linear and applies the standard _validate_amp / conditional-cast pattern to the affine conditioning parameters. Defaults preserve the existing behavior (amp_mode=False), and saved checkpoints load unchanged. Add tests covering autocast forward (bf16), the error on default amp_mode under autocast, and full-precision execution with amp_mode=True outside autocast, across all architecture configurations. Signed-off-by: Piyoosh Jaysaval <piyoosh.jaysaval@pnnl.gov>
Use the bare device type rather than the possibly-indexed device string from the test fixture when entering torch.autocast, matching the idiom used elsewhere in the test suite. Signed-off-by: Piyoosh Jaysaval <piyoosh.jaysaval@pnnl.gov>
90db634 to
39600aa
Compare
|
Gentle ping: the automated review findings were addressed (autocast device-type normalization pushed) and CI is green. Would appreciate a maintainer look when someone has a chance. The branch is behind main but has no conflicts; happy to rebase if preferred. |
1 participant
Description
The experimental 3D diffusion U-Net could not run under
torch.autocast: its internalphysicsnemo.nnlayers (Linear,PositionalEmbedding,FourierEmbedding) validate theiramp_modeflag and raise when autocast is active, but neitherDiffusionUNet3DnorUNetBlock3Dexposed the flag, so mixed-precision training always failed withRuntimeError: amp_mode=False but torch autocast is enabled.This PR threads
amp_modethrough both classes, mirroringSongUNetand the 2DUNetBlock:DiffusionUNet3Dgainsamp_mode: bool = Falseand passes it to the noise/condition embeddings, the map linears, and everyUNetBlock3D(viablock_kwargs).UNetBlock3Dgainsamp_mode: bool = False, passes it to its affineLinear, and applies the standard_validate_amp/ conditional-cast pattern to the affine conditioning parameters (matchingphysicsnemo/nn/module/unet_layers.py).Defaults preserve existing behavior (
amp_mode=False), and previously saved checkpoints load unchanged (the existing.mdlusnon-regression tests pass).New tests (parametrized over all architecture configurations): autocast bf16 forward with
amp_mode=True(shape/dtype/finiteness), error on the defaultamp_mode=Falseunder autocast, and full-precision execution withamp_mode=Trueoutside autocast; analogous tests forUNetBlock3D. Validated end-to-end by training an EDM-preconditionedDiffusionUNet3Dundertorch.autocastbf16.closes #1935
Checklist
Dependencies
None.