Skip to content

perf: gather node embedding before matmul #4744

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
caic99 wants to merge 24 commits into from
Closed

perf: gather node embedding before matmul #4744

Changes from 1 commit
Commits
Show all changes
24 commits
Select commit Hold shift + click to select a range
ad23558
perf: gather node embedding before matmul
caic99 May 16, 2025
b543cc8
swap dim
caic99 May 16, 2025
09564f3
extract nei_node_ebd
caic99 May 16, 2025
1888cd2
remove node_ebd_ext
caic99 May 16, 2025
7569f78
not extend atype for training
caic99 May 16, 2025
93e6851
remove debug statements
caic99 May 16, 2025
cbaf7fd
remove duplicated statements
caic99 May 16, 2025
29667c5
add param use_ext_ebd
caic99 May 16, 2025
511c207
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] May 16, 2025
8804e7f
add args impl for dpmodel
caic99 May 16, 2025
df875cf
Merge branch 'fix-node-ebd-ext' of https://github.com/caic99/deepmd-k…
caic99 May 16, 2025
7b7cdfa
fix dpmodel repflow serialization
caic99 May 16, 2025
a746f82
fix UT
caic99 May 16, 2025
5594c56
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] May 16, 2025
ac6677e
refactor node_ebd_ext ctor
caic99 May 16, 2025
a349de1
feat(pt): add use_loc_mapping
iProzd May 16, 2025
dd0d642
solve merge conflict
caic99 May 22, 2025
5e2a9e6
modify ut
caic99 May 22, 2025
3ccefcc
fix name in dpmodel
caic99 May 22, 2025
b6ca445
fix ut
caic99 May 22, 2025
20fa93c
delete unused files
caic99 May 22, 2025
d751e44
fix unused params
caic99 May 22, 2025
7ba9580
format ut
caic99 May 22, 2025
91611e2
fix mapping
caic99 May 22, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions deepmd/pt/model/descriptor/repflow_layer.py
View file
Open in desktop
Copy link
Collaborator

@iProzd iProzd May 22, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe you need to abort modifications in this file, which means to keep exact modifications in commit iProzd@28803d9 and I passed all the uts.

Copy link
Member Author

@caic99 caic99 May 22, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe you need to abort modifications in this file

Did you mean you'll submit another PR to this branch?

Original file line number Diff line number Diff line change
Expand Up @@ -455,10 +455,10 @@ def optim_edge_update(

# nf * nloc * node/edge_dim
sub_node_update = torch.matmul(node_ebd, node)
# nf * nall * node/edge_dim
sub_node_ext_update = torch.matmul(node_ebd_ext, node_ext)
# nf * nloc * nnei * node/edge_dim
sub_node_ext_update = _make_nei_g1(sub_node_ext_update, nlist)
gathered_node_ebd_ext = _make_nei_g1(node_ebd_ext, nlist)
# nf * nloc * nnei * node/edge_dim
sub_node_ext_update = torch.matmul(gathered_node_ebd_ext, node_ext)
# nf * nloc * nnei * node/edge_dim
sub_edge_update = torch.matmul(edge_ebd, edge)

Expand Down