Fix DataParallel issues for GP-VAE, FILM, and FITS on multi-GPU setups#819
Draft
Claude wants to merge 2 commits into
Draft
Fix DataParallel issues for GP-VAE, FILM, and FITS on multi-GPU setups#819Claude wants to merge 2 commits into
Claude wants to merge 2 commits into
Conversation
- GP-VAE: Pre-compute kernel matrices as buffers and create prior distribution on-the-fly to avoid DataParallel replication issues - FITS: Handle complex numbers by splitting into real/imaginary parts instead of using .to(torch.cfloat) on Linear layers - FILM: Handle complex einsum operations by splitting into real/imaginary parts to avoid DataParallel issues with complex tensors Co-authored-by: WenjieDu <17807970+WenjieDu@users.noreply.github.com>
Claude
AI
changed the title
|
This pull request had no activity for 14 days. It will be closed in 1 week unless there is some new activity. |
|
Collaborator
Coverage Report for CI Build 24951327727Coverage decreased (-0.03%) to 79.961%Details
Uncovered Changes
Coverage RegressionsNo coverage regressions found. Coverage Stats
💛 - Coveralls |
There was a problem hiding this comment.
Pull request overview
This pull request addresses multi-GPU (torch.nn.DataParallel) runtime failures in the GP-VAE, FITS, and FiLM model backbones by removing non-replicable state, avoiding complex-valued module parameters, and rewriting complex operations into real/imag arithmetic.
Changes:
- GP-VAE: precomputes and buffers the GP prior covariance and constructs a fresh prior distribution each forward pass.
- FITS: removes complex-typed
nn.Linearparameters and applies the upsampling transform to real/imag parts separately. - FiLM: rewrites complex
einsumoperations as real/imageinsum+ recombination.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 6 comments.
| File | Description |
|---|---|
pypots/nn/modules/gpvae/backbone.py |
Precomputes GP prior covariance as a buffer and creates the prior distribution per-forward for DataParallel safety. |
pypots/nn/modules/fits/backbone.py |
Reworks FFT upsampling to avoid complex Linear parameters by operating on real/imag parts. |
pypots/nn/modules/film/layers.py |
Splits complex einsum into real/imag computations to avoid DataParallel einsum issues. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+56
to
+61
| low_specx_i = low_specx[:, :, i] # Shape: (batch, dominance_freq) | ||
| # Split into real and imaginary parts | ||
| real_part = self.freq_upsampler[i](low_specx_i.real) | ||
| imag_part = self.freq_upsampler[i](low_specx_i.imag) | ||
| # Recombine into complex tensor | ||
| low_specxy_[:, :, i] = torch.complex(real_part, imag_part) |
Comment on lines
+66
to
+71
| low_specx_permuted = low_specx.permute(0, 2, 1) | ||
| # Split into real and imaginary parts | ||
| real_part = self.freq_upsampler(low_specx_permuted.real) | ||
| imag_part = self.freq_upsampler(low_specx_permuted.imag) | ||
| # Recombine and permute back | ||
| low_specxy_ = torch.complex(real_part, imag_part).permute(0, 2, 1) |
Comment on lines
+121
to
+122
| # Register as buffer so it will be moved to the correct device by DataParallel | ||
| self.register_buffer("prior_covariance", kernel_matrix_tiled) |
Comment on lines
+145
to
148
| return torch.distributions.MultivariateNormal( | ||
| loc=torch.zeros(self.latent_dim, self.time_length, device=device), | ||
| covariance_matrix=kernel_matrix_tiled.to(device), | ||
| covariance_matrix=self.prior_covariance.to(device), | ||
| ) |
Comment on lines
+110
to
112
| # Register index as a buffer to ensure it's properly handled by DataParallel | ||
| self.register_buffer("index_buffer", torch.tensor(self.index, dtype=torch.long)) | ||
|
|
Comment on lines
139
to
+142
| a = x_ft[:, :, :, : self.modes2] | ||
| out_ft[:, :, :, : self.modes2] = torch.einsum("bjix,iox->bjox", a, self.weights1) | ||
| # Handle complex einsum by splitting into real and imaginary parts | ||
| # to avoid issues with DataParallel | ||
| a_real = a.real |
Owner
|
@claude[agent] Now it's your turn. Resolve the above suggestions from Copilot. |
Copilot stopped work on behalf of
WenjieDu due to an error
June 17, 2026 09:57
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Multi-GPU training with
torch.nn.DataParallelfails for GP-VAE, FILM, and FITS models with runtime errors related to distribution initialization, complex tensor operations, and einsum dimension mismatches.Changes
GP-VAE (
pypots/nn/modules/gpvae/backbone.py)__init__and register as buffer viaregister_buffer("prior_covariance", ...)_get_prior()method instead of cachingFITS (
pypots/nn/modules/fits/backbone.py).to(torch.cfloat)from Linear layer initializationtorch.complex()FILM (
pypots/nn/modules/film/layers.py)(a+bi)(c+di) = (ac-bd) + (ad+bc)iTechnical rationale
DataParallel replicates models to each GPU but struggles with:
Solution uses buffers for device-aware tensors and explicit real/imaginary arithmetic, consistent with patterns from PR #633 (Koopa, USGAN, CRLI fixes).
Usage
Original prompt
This section details on the original issue you should resolve
<issue_title>FITS/FILM/GP-VAE fail when running on multiple CUDA devices</issue_title>
<issue_description>### 1. System Info
v0.11
2. Information
3. Reproduction
4. Expected behavior
For pypots.forecasting.fits and pypots.imputation.fits we have
For pypots.imputation.film we have
For pypots.imputation.gpvae we have