Skip to content

Drop stale param refs in _update_param_with_optimizer to cut VRAM spike#931

Open
Evggamgam wants to merge 3 commits intonerfstudio-project:mainfrom
Evggamgam:fix/densify-vram-transient
Open

Drop stale param refs in _update_param_with_optimizer to cut VRAM spike#931
Evggamgam wants to merge 3 commits intonerfstudio-project:mainfrom
Evggamgam:fix/densify-vram-transient

Conversation

@Evggamgam
Copy link
Copy Markdown

_update_param_with_optimizer keeps the old Parameter alive via three references (local var, optimizer.state key, param_groups["params"]) while new Adam state tensors are allocated, causing a per-densify transient of ~1× size(largest param) above steady state. Rebind those refs to new_param and del the local before optimizer_fn, so the old tensor is freed during the state reshape.

Pure reordering — output Parameter and Adam state are bit-identical to the pre-patch implementation; wall time unchanged.

Max Gaussians in the densify-step budget on a 24 GB card (excluding render/backward overhead):

strategy before after
DefaultStrategy 16.5M 19.8M
MCMCStrategy 21.3M 25.4M

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