Skip to content

pt: avoid torch.tensor(constant) during forward#3421

Merged
wanghan-iapcm merged 7 commits intodeepmodeling:develfrom
njzjz:avoid-torch-tensor
Mar 8, 2024
Merged

pt: avoid torch.tensor(constant) during forward#3421
wanghan-iapcm merged 7 commits intodeepmodeling:develfrom
njzjz:avoid-torch-tensor

Conversation

@njzjz
Copy link
Copy Markdown
Member

@njzjz njzjz commented Mar 6, 2024

torch.tensor(constant) copies memory from the CPU to the GPU, so it is host blocking and should be avoided in the forward method.

Before, the CPU waited for the GPU using cudaStreamSynchronize, blocking the CPU from doing the following things, where the CPU memory needs to be copied to the GPU, a.k.a. host-to-device (H2D).

1709693858444

After this PR, all ops in the energy loss are asynchronous, as no H2D happens.
1709694622120

torch.tensor(constant) copies memory from the CPU to GPU, so it is host blocking and should be avoided in the `forward` method.

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
njzjz added 3 commits March 5, 2024 22:22
Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
@codecov
Copy link
Copy Markdown

codecov Bot commented Mar 6, 2024

Codecov Report

Attention: Patch coverage is 66.66667% with 3 lines in your changes are missing coverage. Please review.

Project coverage is 76.98%. Comparing base (09bd522) to head (46d43f0).
Report is 2 commits behind head on devel.

Files Patch % Lines
deepmd/pt/loss/denoise.py 0.00% 3 Missing ⚠️
Additional details and impacted files
@@           Coverage Diff           @@
##            devel    #3421   +/-   ##
=======================================
  Coverage   76.98%   76.98%           
=======================================
  Files         427      427           
  Lines       36753    36753           
  Branches     1621     1621           
=======================================
  Hits        28295    28295           
  Misses       7590     7590           
  Partials      868      868           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
@wanghan-iapcm wanghan-iapcm enabled auto-merge March 8, 2024 08:34
@wanghan-iapcm wanghan-iapcm added this pull request to the merge queue Mar 8, 2024
Merged via the queue into deepmodeling:devel with commit d3dd604 Mar 8, 2024
@njzjz njzjz mentioned this pull request Apr 2, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants