Skip to content

Fix WeaveModel on non-CPU devices: register layer weights as nn.Parameter - #5190

Open
barlowa124 wants to merge 1 commit into
deepchem:masterfrom
barlowa124:weave-forward-device-fix
Open

barlowa124 wants to merge 1 commit into
deepchem:masterfrom
barlowa124:weave-forward-device-fix

Conversation

@barlowa124

@barlowa124 barlowa124 commented Oct 1, 2026 •

Copy link
Copy Markdown

Description

The PyTorch Weave layers never registered their convolution weights as
parameters, and Weave.forward pushed its inputs through numpy, which
crashes on any non-CPU device. Both are addressed here.

Weave.forward ran np.array() on its inputs. TorchModel places them
on the selected device, so under mps or cuda the call raised
TypeError (can't convert mps:0 device type tensor to numpy). The
conversion was never needed. WeaveLayer and WeaveGather already
call torch.tensor() on what they receive, covering ndarrays and
tensors alike.

WeaveLayer and WeaveGather also held trainable weights as plain
torch.Tensor attributes. Unregistered tensors sit outside
model.parameters(). .to(device) skipped them and state_dict()
dropped them, so they received no gradients. On CPU this hid because
everything shared one device, but it also means the weave layers were
frozen random projections, consistent with the test_weave_model
comment noting a lower score than the TensorFlow equivalent.

What changed:

  • Weave.forward hands its inputs to the weave layers directly.
  • Weights and biases in WeaveLayer and WeaveGather are nn.Parameter.
  • test_weave_classification now assigns injected weights through
    .data, required once the attributes became Parameters.
  • test_weave_fit_non_cpu_device fits a WeaveModel on the
    auto-selected device, reproducing the reported failure on mps/cuda
    and exercising the normal path on CPU-only CI.

Old checkpoints never stored the weave weights, so restore() under
strict loading will now report missing keys on pre-change files.
strict=False keeps the previous behavior.

Checked on Apple Silicon (mps). The non-slow tests in
test_weavemodel_pytorch.py pass locally, including the three that
fail on master on this hardware
(test_weave_fit_simple_distance_1, test_weave_model_reload,
test_weave_singletask_classification_overfit), and the WeaveModel
docstring example runs to completion.

Type of change

Please check the option that is related to your PR.

  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)
    • In this case, we recommend to discuss your modification on GitHub issues before creating the PR
  • Documentations (modification for documents)

Checklist

  • My code follows the style guidelines of this project
    • Run yapf -i <modified file> and check no errors (yapf version must be 0.32.0)
    • Run mypy -p deepchem and check no errors (no new errors in modified regions, layers.py has pre-existing errors in unrelated classes)
    • Run flake8 <modified file> --count and check no errors
    • Run python -m doctest <modified file> and check no errors
  • I have performed a self-review of my own code
  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation
  • I have added tests that prove my fix is effective or that my feature works
  • New unit all green locally with my changes
  • I have checked my code and corrected any misspellings

numpy conversion in Weave.forward

WeaveLayer and WeaveGather stored their trainable weights as plain
torch.Tensor attributes, so they were invisible to .to(device), the
optimizer, and state_dict. The weave convolutions were frozen random
projections that never trained and were never checkpointed.

Weave.forward also called np.array() on inputs that TorchModel had
already moved to the compute device, which raised TypeError on mps and
cuda. The conversion was unnecessary since the layers call
torch.tensor() on their inputs; inputs now pass through directly.

Verified on Apple Silicon: all non-slow weavemodel tests pass,
including three that fail on master on this hardware.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant