Skip to content

Fix global PyTorch initializer leak in build_vae_var - #185

Open
primorLee wants to merge 1 commit into
FoundationVision:mainfrom
primorLee:fix/159-restore-model-initializers
Open

Fix global PyTorch initializer leak in build_vae_var#185
primorLee wants to merge 1 commit into
FoundationVision:mainfrom
primorLee:fix/159-restore-model-initializers

Conversation

@primorLee

Copy link
Copy Markdown

Summary

  • scope the temporary PyTorch reset_parameters overrides to VAR/VQVAE construction
  • restore every original initializer in a finally block on both success and failure
  • add CPU regression coverage for normal construction and constructor exceptions

Problem

build_vae_var() disables built-in initialization by replacing reset_parameters on eight global torch.nn classes. Those class methods were never restored. After one VAR build, unrelated modules created anywhere in the same Python process could therefore keep uninitialized weights.

A minimal CPU reproduction on current main prints linear_reset_restored=False after build_vae_var() returns.

Fix

Save the original class methods, keep the existing no-initialization fast path while VQVAE and VAR are constructed, and restore the methods in finally. The existing explicit var_wo_ddp.init_weights(...) remains unchanged.

Validation

  • real reduced CPU VAR/VQVAE construction verifies all eight PyTorch initializers are restored
  • injected VQVAE construction failure verifies restoration on the exception path
  • python -m pytest tests/test_model_initialization.py -q (2 passed, 16 subtests)
  • python -m compileall -q models/__init__.py tests/test_model_initialization.py
  • repeated both regression cases in a clean Linux Python 3.10 / CPU Torch 2.5.1 container

Fixes #159

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.

姹囨姤浠g爜涓竴涓彲鑳藉鑷磘orch鑷姩鏉冮噸鍒濆鍖栧け鏁堢殑鎿嶄綔

1 participant