provide the ability to unload the LM prior to the DiT execution phase #1090
There are no files selected for viewing
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
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
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -64,6 +64,7 @@ def __init__(self, persistent_storage_path: Optional[str] = None): | |
| self.dtype = torch.float32 | ||
| self.offload_to_cpu = False | ||
| self.disable_tqdm = os.environ.get("ACESTEP_DISABLE_TQDM", "").lower() in ("1", "true", "yes") or not (hasattr(sys.stderr, 'isatty') and sys.stderr.isatty()) | ||
| self._last_init_config = None | ||
|
|
||
| # HuggingFace Space persistent storage support | ||
| if persistent_storage_path is None and self.IS_HUGGINGFACE_SPACE: | ||
|
|
@@ -80,6 +81,24 @@ def __init__(self, persistent_storage_path: Optional[str] = None): | |
| self._mlx_model = None | ||
| self._mlx_model_path = None | ||
|
|
||
| def _save_last_init_config( | ||
| self, | ||
| checkpoint_dir: str, | ||
| lm_model_path: str, | ||
| device: str, | ||
| offload_to_cpu: bool, | ||
| dtype: Optional[torch.dtype], | ||
| ) -> None: | ||
| """Persist the last successfully initialized LM configuration.""" | ||
| self._last_init_config = { | ||
| "checkpoint_dir": checkpoint_dir, | ||
| "lm_model_path": lm_model_path, | ||
| "backend": self.llm_backend, | ||
| "device": device, | ||
| "offload_to_cpu": offload_to_cpu, | ||
| "dtype": dtype, | ||
| } | ||
|
|
||
| def _clear_accelerator_cache(self) -> None: | ||
| """Release freed accelerator memory back to the driver. | ||
|
|
||
|
|
@@ -119,22 +138,44 @@ def unload(self) -> None: | |
| try: | ||
| if self.llm_backend == "vllm": | ||
| try: | ||
| if self.llm is not None: | ||
| if hasattr(self.llm, "exit"): | ||
| logger.info("[LLM vLLM] Calling nanovllm exit() for hard teardown") | ||
| self.llm.exit() | ||
| elif hasattr(self.llm, "reset"): | ||
| logger.info("[LLM vLLM] exit() missing, falling back to reset()") | ||
| self.llm.reset() | ||
| except Exception as exc: | ||
| logger.warning(f"[LLM vLLM] Error during vLLM teardown: {exc}") | ||
|
|
||
| try: | ||
| self._cleanup_torch_distributed_state() | ||
| except Exception as exc: | ||
| logger.warning(f"[LLM vLLM] torch distributed cleanup failed: {exc}") | ||
|
|
||
| self.llm = None | ||
| self.llm_tokenizer = None | ||
| self.constrained_processor = None | ||
| self.llm_initialized = False | ||
| self._hf_model_for_scoring = None | ||
| self._mlx_model = None | ||
| self._mlx_model_path = None | ||
|
|
||
| gc.collect() | ||
|
|
||
| if torch.cuda.is_available(): | ||
| try: | ||
| torch.cuda.synchronize() | ||
| except Exception: | ||
| pass | ||
| try: | ||
| torch.cuda.empty_cache() | ||
| except Exception: | ||
| pass | ||
| try: | ||
| torch.cuda.ipc_collect() | ||
| except Exception: | ||
| pass | ||
|
coderabbitai[bot] marked this conversation as resolved.
Outdated
|
||
| elif hasattr(torch, "mps") and torch.backends.mps.is_available(): | ||
| if hasattr(torch.mps, "synchronize"): | ||
| torch.mps.synchronize() | ||
|
|
@@ -143,8 +184,27 @@ def unload(self) -> None: | |
| elif hasattr(torch, "xpu") and torch.xpu.is_available(): | ||
| torch.xpu.empty_cache() | ||
| torch.xpu.synchronize() | ||
| except Exception as exc: | ||
| logger.warning(f"[LLM] unload failed: {exc}") | ||
|
|
||
| def reload_last_configuration(self) -> Tuple[str, bool]: | ||
| """Recreate the LM from the last successful initialize() configuration.""" | ||
| if not self._last_init_config: | ||
| return "❌ No previous LM initialization config available", False | ||
|
|
||
| cfg = dict(self._last_init_config) | ||
|
|
||
| logger.info("[LLM] Reloading last configuration: backend={} model={} device={}", | ||
| cfg.get("backend"), cfg.get("lm_model_path"), cfg.get("device")) | ||
|
|
||
| return self.initialize( | ||
| checkpoint_dir=cfg["checkpoint_dir"], | ||
| lm_model_path=cfg["lm_model_path"], | ||
| backend=cfg["backend"], | ||
| device=cfg["device"], | ||
| offload_to_cpu=cfg["offload_to_cpu"], | ||
| dtype=cfg["dtype"], | ||
| ) | ||
|
|
||
| def _cleanup_torch_distributed_state(self) -> None: | ||
| """Destroy default torch distributed process group when already initialized.""" | ||
|
|
@@ -659,6 +719,13 @@ def initialize( | |
| logger.info("Attempting MLX backend for Apple Silicon acceleration...") | ||
| mlx_success, mlx_status = self._load_mlx_model(full_lm_model_path) | ||
| if mlx_success: | ||
| self._save_last_init_config( | ||
| checkpoint_dir=checkpoint_dir, | ||
| lm_model_path=lm_model_path, | ||
| device=device, | ||
| offload_to_cpu=offload_to_cpu, | ||
| dtype=dtype, | ||
| ) | ||
| return mlx_status, True | ||
| else: | ||
| logger.warning(f"MLX backend failed: {mlx_status}") | ||
|
|
@@ -669,6 +736,13 @@ def initialize( | |
| if not success: | ||
| return status_msg, False | ||
| status_msg = f"✅ 5Hz LM initialized (PyTorch fallback from MLX)\nModel: {full_lm_model_path}\nBackend: PyTorch" | ||
| self._save_last_init_config( | ||
| checkpoint_dir=checkpoint_dir, | ||
| lm_model_path=lm_model_path, | ||
| device=device, | ||
| offload_to_cpu=offload_to_cpu, | ||
| dtype=dtype, | ||
| ) | ||
| return status_msg, True | ||
| # else: backend was "vllm" on MPS, continue to vllm attempt below | ||
| elif backend == "mlx": | ||
|
|
@@ -678,6 +752,13 @@ def initialize( | |
| if not success: | ||
| return status_msg, False | ||
| status_msg = f"✅ 5Hz LM initialized (PyTorch fallback, MLX not available)\nModel: {full_lm_model_path}\nBackend: PyTorch" | ||
| self._save_last_init_config( | ||
| checkpoint_dir=checkpoint_dir, | ||
| lm_model_path=lm_model_path, | ||
| device=device, | ||
| offload_to_cpu=offload_to_cpu, | ||
| dtype=dtype, | ||
| ) | ||
| return status_msg, True | ||
|
|
||
| if backend == "vllm" and device != "cuda": | ||
|
|
@@ -733,6 +814,13 @@ def initialize( | |
| logger.warning("vllm failed on MPS, trying MLX backend...") | ||
| mlx_success, mlx_status = self._load_mlx_model(full_lm_model_path) | ||
| if mlx_success: | ||
| self._save_last_init_config( | ||
| checkpoint_dir=checkpoint_dir, | ||
| lm_model_path=lm_model_path, | ||
| device=device, | ||
| offload_to_cpu=offload_to_cpu, | ||
| dtype=dtype, | ||
| ) | ||
| return mlx_status, True | ||
| logger.warning(f"MLX also failed: {mlx_status}, falling back to PyTorch") | ||
| logger.warning("Falling back to PyTorch backend") | ||
|
|
@@ -749,6 +837,13 @@ def initialize( | |
| if vllm_preflight_warning is not None: | ||
| status_msg += f"\nNote: {vllm_preflight_warning}" | ||
|
|
||
| self._save_last_init_config( | ||
| checkpoint_dir=checkpoint_dir, | ||
| lm_model_path=lm_model_path, | ||
| device=device, | ||
| offload_to_cpu=offload_to_cpu, | ||
| dtype=dtype, | ||
| ) | ||
| return status_msg, True | ||
|
|
||
| except Exception as e: | ||
|
|
||
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.
Uh oh!
There was an error while loading. Please reload this page.