From 71599c47e79cbfdb73ad47c33e53f61c619dbf86 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 26 Jul 2026 07:20:48 +0000 Subject: [PATCH] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/export/export.py | 9 +++----- tests/test_compressed_export_gpu_release.py | 23 +++++++++++---------- unsloth/save.py | 4 +++- 3 files changed, 18 insertions(+), 18 deletions(-) diff --git a/studio/backend/core/export/export.py b/studio/backend/core/export/export.py index 30a1bb6f2f..0425891cb7 100644 --- a/studio/backend/core/export/export.py +++ b/studio/backend/core/export/export.py @@ -120,7 +120,8 @@ def _is_oom_error(exc: BaseException) -> bool: """ if torch is not None: oom_types = tuple( - t for t in ( + t + for t in ( getattr(torch, "OutOfMemoryError", None), getattr(getattr(torch, "cuda", None), "OutOfMemoryError", None), getattr(getattr(torch, "xpu", None), "OutOfMemoryError", None), @@ -502,11 +503,7 @@ class ExportBackend: # GPU, so when a training or chat job already owns the others (which # routes/export.py deliberately allows) the shard can OOM where the pre-#7053 # single-device load succeeded. Fall back once before giving up. - if ( - _device_map_override is None - and _is_oom_error(e) - and _multi_gpu_device_map_kwargs() - ): + if _device_map_override is None and _is_oom_error(e) and _multi_gpu_device_map_kwargs(): # Retry outside this block: the live traceback pins the half-built # model's frames, so an in-block retry inherits the exhausted device. oom_retry_reason = str(e) diff --git a/tests/test_compressed_export_gpu_release.py b/tests/test_compressed_export_gpu_release.py index 020f484a9f..d3111a1b90 100644 --- a/tests/test_compressed_export_gpu_release.py +++ b/tests/test_compressed_export_gpu_release.py @@ -43,13 +43,12 @@ class _FakeLogger: def _load_helpers(fake_torch, fake_logger): tree = ast.parse(_SAVE_PY.read_text(encoding = "utf-8")) keep = [ - node for node in tree.body + node + for node in tree.body if (isinstance(node, ast.FunctionDef) and node.name in _WANTED) or ( isinstance(node, ast.Assign) - and any( - isinstance(t, ast.Name) and t.id in _WANTED_ASSIGNS for t in node.targets - ) + and any(isinstance(t, ast.Name) and t.id in _WANTED_ASSIGNS for t in node.targets) ) ] n_fns = sum(1 for node in keep if isinstance(node, ast.FunctionDef)) @@ -101,9 +100,7 @@ def _fake_accelerate(monkeypatch): accel.dispatch_model = _dispatch hooks = types.ModuleType("accelerate.hooks") hooks.remove_hook_from_submodules = lambda model: calls["removed"].append(model) - hooks.add_hook_to_module = lambda module, hook: calls["hooks_added"].append( - (module, hook) - ) + hooks.add_hook_to_module = lambda module, hook: calls["hooks_added"].append((module, hook)) accel.hooks = hooks monkeypatch.setitem(sys.modules, "accelerate", accel) monkeypatch.setitem(sys.modules, "accelerate.hooks", hooks) @@ -341,7 +338,11 @@ def test_torchao_export_uses_the_shared_release(): class _Child: """Minimal stand-in for an nn.Module leaf, enough for the dispatch walk.""" - def __init__(self, name = "inner", device_map = None): + def __init__( + self, + name = "inner", + device_map = None, + ): self._modules = {} self.__dict__["_name"] = name if device_map is not None: @@ -449,11 +450,11 @@ def test_snapshot_restores_a_forward_patched_after_the_dispatch(_fake_accelerate mlp = _Child(name = "mlp") root._modules["mlp"] = mlp - stock_forward = lambda *a, **k: "stock" # noqa: E731 + stock_forward = lambda *a, **k: "stock" # noqa: E731 fused_forward = lambda *a, **k: "unsloth-fused" # noqa: E731 mlp._hf_hook = object() - mlp._old_forward = stock_forward # captured by accelerate at dispatch time - mlp.forward = fused_forward # installed by unsloth afterwards + mlp._old_forward = stock_forward # captured by accelerate at dispatch time + mlp.forward = fused_forward # installed by unsloth afterwards snapshot = ns["_snapshot_dispatch_state"](root) diff --git a/unsloth/save.py b/unsloth/save.py index 4b5f6d3394..c57a1737cb 100644 --- a/unsloth/save.py +++ b/unsloth/save.py @@ -4638,7 +4638,8 @@ def _accelerate_dispatch_root(model): children = getattr(node, "__dict__", {}).get("_modules") or {} nxt = next( ( - children[a] for a in ("base_model", "model") + children[a] + for a in ("base_model", "model") if hasattr(children.get(a), "named_modules") ), None, @@ -4836,6 +4837,7 @@ def _restore_model_after_quantize_subprocess(model, restore_token) -> None: _restore_dispatch_state(root, snapshot) else: from accelerate import dispatch_model + # skip_keys matters: without it accelerate moves every forward kwarg # to the executing device, wrong for device-invariant cache tensors. dispatch_model(