Remove deepcopy:pickle error, allow non-falsy num_workers=0 - #237
Conversation
|
@ipcamit can you please review this PR? |
There was a problem hiding this comment.
Pull request overview
This PR updates the PyTorch Lightning trainer to (1) avoid deepcopy failures when exporting models that include e3nn layers, (2) preserve/normalize graph tensor dtypes rather than always ending up with float64 from the RadialGraph C-extension, and (3) correctly honor num_workers=0 instead of treating it as falsy and falling back to SLURM defaults.
Changes:
- Remove
deepcopy(self.pl_model)during export and load the best checkpoint directly intoself.pl_modelbefore TorchScripting. - Add a post-transform dtype normalization step for precomputed graph fingerprints (coords/forces) when default dtype is float32.
- Fix
num_workersresolution so explicit0is respected and values are consistently cast toint.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| # of input dtype. Cast fingerprints back to match the model's default dtype. | ||
| if torch.get_default_dtype() == torch.float32: | ||
| for ds in [self.train_dataset, self.val_dataset]: | ||
| if ds is None: | ||
| continue | ||
| for config in ds: | ||
| fp = getattr(config, "fingerprint", None) | ||
| if fp is None: | ||
| continue | ||
| if hasattr(fp, "coords") and fp.coords is not None: | ||
| fp.coords = fp.coords.to(torch.float32) | ||
| if hasattr(fp, "forces") and fp.forces is not None: | ||
| fp.forces = fp.forces.to(torch.float32) |
There was a problem hiding this comment.
This is too verbose, perhaps factor it as an explicit function and simplify the code here?
| if self.dataset_manifest["dynamic_loading"]: | ||
| transform = self.configuration_transform | ||
| else: | ||
| transform = None | ||
|
|
||
| if not transform: | ||
| for config in self.train_dataset: | ||
| config.fingerprint = self.configuration_transform(config) | ||
| if self.val_dataset: | ||
| for config in self.val_dataset: | ||
| config.fingerprint = self.configuration_transform(config) | ||
|
|
||
| # RadialGraph C extension always upcasts coords/forces to float64 regardless | ||
| # of input dtype. Cast fingerprints back to match the model's default dtype. | ||
| if torch.get_default_dtype() == torch.float32: | ||
| for ds in [self.train_dataset, self.val_dataset]: | ||
| if ds is None: |
There was a problem hiding this comment.
@gpwolfe Can you address this? otherwise it will comeback in some other config.
| self.pl_model.load_state_dict( | ||
| torch.load( | ||
| f"{self.current['run_dir']}/checkpoints/best_model.pth", | ||
| weights_only=False, | ||
| ) | ||
| ) | ||
| try: |
There was a problem hiding this comment.
hey @gpwolfe can you very quickly google once to confirm if from e3nn.util import jit this is still needed or if e3nn has streamlined it to torch.jit?
| # so mutating self.pl_model is safe. | ||
| self.pl_model.load_state_dict( | ||
| torch.load( | ||
| f"{self.current['run_dir']}/checkpoints/best_model.pth", |
There was a problem hiding this comment.
it is a simple enough change for better handling of saved checkpoints.
|
Out of time. Will get back to it in couple of days. |
| if self.dataset_manifest["dynamic_loading"]: | ||
| transform = self.configuration_transform | ||
| else: | ||
| transform = None | ||
|
|
||
| if not transform: | ||
| for config in self.train_dataset: | ||
| config.fingerprint = self.configuration_transform(config) | ||
| if self.val_dataset: | ||
| for config in self.val_dataset: | ||
| config.fingerprint = self.configuration_transform(config) | ||
|
|
||
| # RadialGraph C extension always upcasts coords/forces to float64 regardless | ||
| # of input dtype. Cast fingerprints back to match the model's default dtype. | ||
| if torch.get_default_dtype() == torch.float32: | ||
| for ds in [self.train_dataset, self.val_dataset]: | ||
| if ds is None: |
There was a problem hiding this comment.
@gpwolfe Can you address this? otherwise it will comeback in some other config.
| # of input dtype. Cast fingerprints back to match the model's default dtype. | ||
| if torch.get_default_dtype() == torch.float32: | ||
| for ds in [self.train_dataset, self.val_dataset]: | ||
| if ds is None: | ||
| continue | ||
| for config in ds: | ||
| fp = getattr(config, "fingerprint", None) | ||
| if fp is None: | ||
| continue | ||
| if hasattr(fp, "coords") and fp.coords is not None: | ||
| fp.coords = fp.coords.to(torch.float32) | ||
| if hasattr(fp, "forces") and fp.forces is not None: | ||
| fp.forces = fp.forces.to(torch.float32) |
There was a problem hiding this comment.
This is too verbose, perhaps factor it as an explicit function and simplify the code here?
| # so mutating self.pl_model is safe. | ||
| self.pl_model.load_state_dict( | ||
| torch.load( | ||
| f"{self.current['run_dir']}/checkpoints/best_model.pth", |
There was a problem hiding this comment.
it is a simple enough change for better handling of saved checkpoints.
| self.val_dataset = GraphDataset(self.val_dataset, transform) | ||
|
|
||
| if self.optimizer_manifest["num_workers"]: | ||
| num_workers = self.optimizer_manifest["num_workers"] |
There was a problem hiding this comment.
No issues here. Earlier it was explicitly, one cpu or distributed many cpu as multiprocessing never worked in PytorchLightning for me. It this works,great!
| self.pl_model.load_state_dict( | ||
| torch.load( | ||
| f"{self.current['run_dir']}/checkpoints/best_model.pth", | ||
| weights_only=False, | ||
| ) | ||
| ) | ||
| try: |
There was a problem hiding this comment.
hey @gpwolfe can you very quickly google once to confirm if from e3nn.util import jit this is still needed or if e3nn has streamlined it to torch.jit?
Summary
Include a summary of major changes in bullet points:
num_workers = os.getenv("SLURM_CPUS_PER_TASK", 1)due to evaluation as False. Allows non-multithread support on non-slurm system.Additional dependencies introduced (if any)
Checklist
Before a pull request can be merged, the following items must be checked:
type check your code.
Note that the CI system will run all the above checks. But it will be much more
efficient if you already fix most errors prior to submitting the PR.