From 79e7f412165e66a10db55633a613cb6dc7fb541c Mon Sep 17 00:00:00 2001 From: Adrien Toupet Date: Thu, 27 Nov 2025 14:52:12 -0500 Subject: [PATCH] Revert "Fix: MPS allocator error in model.to() transfer (#305)" This reverts commit 94240206871f813366ba049c44cea155a490abe5. --- src/optimization/memory_manager.py | 26 +------------------------- 1 file changed, 1 insertion(+), 25 deletions(-) diff --git a/src/optimization/memory_manager.py b/src/optimization/memory_manager.py index 5bfcf38..6010123 100644 --- a/src/optimization/memory_manager.py +++ b/src/optimization/memory_manager.py @@ -875,31 +875,7 @@ def _standard_model_movement(model: torch.nn.Module, current_device: torch.devic debug.start_timer(timer_name) # Move model and clear gradients - try: - model.to(target_device) - except RuntimeError as e: - # MPS allocator workaround: Some PyTorch/macOS versions fail with bulk transfers - error_msg = str(e).lower() - is_mps_alloc_error = ( - torch.device(target_device).type == "mps" and - any(keyword in error_msg for keyword in ["watermark", "allocat", "memory"]) - ) - - if is_mps_alloc_error: - if debug: - debug.log("MPS bulk transfer failed, moving parameters individually", - category="memory", indent_level=1) - - # Move parameters and buffers individually to avoid allocator bug - with torch.no_grad(): - for param in model.parameters(): - param.data = param.data.to(target_device) - for buffer in model.buffers(): - buffer.data = buffer.data.to(target_device) - else: - # Re-raise if it's a different error - raise - + model.to(target_device) model.zero_grad(set_to_none=True) # Clear VAE memory buffers when moving to CPU