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

[Experimental] Add memory recording which can then be bulk-freed #698

Merged
merged 4 commits into from
Nov 18, 2024

Conversation

pxl-th
Copy link
Member

@pxl-th pxl-th commented Nov 18, 2024

Add memory recording which can then be bulk-freed.
Because GC is not aware of GPU memory and relying on it leads to significantly higher memory usage and big delays, the idead is to record all device-allocations that happen during AMDGPU.record_memory!(true)
and when AMDGPU.record_memory!(false) is set - bulk-free them.

This is useful for example in DL scenarios where you have your model parameters and compute gradients and update them in-place and once you did this, all of the intermediate allocations can be immediately freed.

θ = <parameters>
AMDGPU.record_memory!(true)
∇ = gradient(θ) do θ
   ...
end
apply!(θ, ∇)
AMDGPU.record_memory!(false) # bulk-free all allocations that happened during recording.

Additionally, you can remove array from recording with AMDGPU.remove_record!(x) or disabling recording for that allocation with AMDGPU.record_memory!(false; free=false) (just remember to enable it again).

Example of memory usage with GaussianSplatting.jl:

Before This PR
image image

@pxl-th pxl-th merged commit 98e7084 into master Nov 18, 2024
3 checks passed
@pxl-th pxl-th deleted the pxl-th/memory-rec branch November 18, 2024 10:09
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