Skip to content
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.

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

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

Merged
merged 7 commits into from
Mar 8, 2024

Conversation

njzjz
Copy link
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 <[email protected]>
Copy link

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.

@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
48 checks passed
@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