Fix #444: Add pre-allocation option for MCMC strategy to prevent OOM#891
Open
eladerez wants to merge 2 commits intonerfstudio-project:mainfrom
Open
Fix #444: Add pre-allocation option for MCMC strategy to prevent OOM#891eladerez wants to merge 2 commits intonerfstudio-project:mainfrom
eladerez wants to merge 2 commits intonerfstudio-project:mainfrom
Conversation
…gy to prevent OOM - Add optional flag to MCMCStrategy (default: False) - Pre-allocate buffers to cap_max size to avoid torch.cat memory spikes - Track active Gaussians separately from buffer size with n_active - Modify sample_add() to write into pre-allocated buffer instead of concat - Add checkpoint support for n_active state - Memory reduction: 6.4% at 1M Gaussians (1.216GB vs 1.299GB baseline) - Prevents OOM crashes at 18M+ Gaussians while maintaining compatibility This is an opt-in feature that eliminates memory fragmentation from repeated torch.cat operations during MCMC densification. Users can enable it with --strategy.preallocate flag when using MCMC strategy.
…lice via param_groups The original preallocate fix allocated cap_max-sized parameter tensors upfront but let each optimizer track the full buffer. This caused optimizer.step() to process all cap_max elements every iteration — a 20x+ overhead when cap_max=1_000_000 and n_initial≈50k. Fix: - Add grow_active_params() in ops.py: creates a narrow Parameter view into splats[:n_active] and migrates Adam momentum tensors (zero-pad new rows) when the active count grows. - Wire each optimizer to track the active-slice Parameter at init time (simple_trainer.py), so optimizer.step() only touches n_active rows. - Grow the active slice (and optimizer state) inside _add_new_gs after sample_add writes new Gaussians into the pre-allocated buffer. - Use active_params directly in _relocate_gs and noise injection, removing the per-step temporary Parameter allocation. Timing: cap_max=100_000 vs cap_max=500 with identical n_active=200 shows a 1.01x ratio (vs ~200x before), confirmed by the new test_mcmc_preallocate_time_independent_of_cap_max test. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Author
|
I've submitted PR #891 that addresses this with pre-allocated buffers. Would appreciate feedback. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
The problem
MCMC training crashes with OOM errors at ~18M Gaussians due to memory fragmentation from repeated torch.cat operations during the densification stage.
Solution:
Added an optional "preallocate" flag to "MCMCStrategy" that:
Pre-allocates buffers to "cap_max" size at initialization
Writing new Gaussians into pre-allocated slots instead of concatenating
Changes:
gsplat/strategy/mcmc.py: Added "preallocate" field and "n_active" state tracking
gsplat/strategy/ops.py: Modified sample_add() for support to buffer writes
examples/simple_trainer.py: Managing buffer and checkpoint
Performance:
Memory: 6.4% reduction at 1M Gaussians (1.216GB vs 1.299GB baseline)
Trade-off: ~24% slower rendering (opt-in feature, disabled by default)
Benefit: Prevents OOM at 18M+ Gaussians scale
Testing:
All strategy tests passing (2/2)
How to run:
python examples/simple_trainer.py mcmc --strategy.preallocate