diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index dfa0f1f114..8c76249eb5 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -891,6 +891,12 @@ class DiffusionBackend: kwargs["repo_id"], kwargs.get("base_repo"), fam, kwargs.get("hf_token") ) kwargs["base_repo"] = base + # The pre-cast encoder injection replaces these components' weights, so prefetching the + # dense shards for them downloads tens of GB nothing opens. Same resolver the injection + # uses, so the prefetch and the load can never disagree about which encoders are dense. + te_prequant_files = self._te_prequant_plan_files( + fam, kwargs.get("text_encoder_quant"), kwargs.get("hf_token") + ) expected, base_files = self._estimate_download_bytes( kwargs["repo_id"], kwargs.get("gguf_filename"), @@ -901,6 +907,7 @@ class DiffusionBackend: # Pull the base shards here rather than inside the locked, unpreemptable finalize. include_transformer = kind == "gguf" and self._dense_quant_prefetch_needed(fam, kwargs), + skip_te_components = tuple(te_prequant_files), ) with self._lock: # Stamp progress only if this load is still current (a superseder has its own token). @@ -974,6 +981,33 @@ class DiffusionBackend: return () return tuple(r for r in (loading.repo_id, loading.base_repo) if r) + @staticmethod + def _te_prequant_plan_files( + fam: Any, text_encoder_quant: Optional[str], hf_token: Optional[str] + ) -> dict[str, tuple[str, list[tuple[str, int]]]]: + """``{component: (repo_id, [(rfilename, size)])}`` for the text encoders this pick will + take PRE-CAST from a hosted checkpoint instead of the base repo's dense weights. + + Empty unless the request asked for a scheme with a hosted artifact AND that artifact + really resolves, so a plan can never drop a dense encoder the load still wants.""" + try: + from huggingface_hub import HfApi + + from .diffusion_te_prequant import te_prequant_hub_files, te_prequant_sources + + sources = te_prequant_sources( + fam, + te_quant_mode = text_encoder_quant, + target = resolve_diffusion_device_target(), + ) + if not sources: + return {} + files = te_prequant_hub_files(sources, HfApi(token = hf_token or None), logger) + return {c: (sources[c].location, f) for c, f in files.items()} + except Exception as exc: # noqa: BLE001 -- an unresolvable pre-cast keeps the dense encoder + logger.warning("diffusion.te_prequant_plan_failed: %s", exc) + return {} + @staticmethod def _estimate_download_bytes( repo_id: str, @@ -985,6 +1019,7 @@ class DiffusionBackend: single_file_is_pipeline: bool = False, include_transformer: bool = False, sizes_out: Optional[dict[str, int]] = None, + skip_te_components: tuple[str, ...] = (), ) -> tuple[int, list[str]]: """Total download size for the progress bar, plus the base-repo files to fetch (the prefetch reuses this list, so the base is listed only once). @@ -997,16 +1032,35 @@ class DiffusionBackend: single-file paths, where the transformer is the single file and the base repo supplies only the companions. For a ``single_file_is_pipeline`` family (SDXL) the single file is the WHOLE pipeline, so the base repo supplies only config/tokenizer - (no weights) and its weight files are skipped.""" + (no weights) and its weight files are skipped. + + ``skip_te_components`` names the text encoders this pick loads PRE-CAST from a hosted + checkpoint, so their dense weight shards are not counted or fetched: staging the dense + encoder for a pre-cast load wastes tens of GB (FLUX.2-dev's Mistral-24B is ~48 GB, + Qwen-Image's Qwen2.5-VL ~16.6 GB) and nothing ever opens them. Everything else in the + component folder (config, shard index, tokenizer) is kept -- the pre-cast loader + meta-inits the encoder from the base repo's config.""" from huggingface_hub import HfApi + from .diffusion_te_prequant import is_prequant_covered_weight + api = HfApi() total = 0 base_files: list[str] = [] + + def _dense_te_shard(rfilename: str) -> bool: + return bool(skip_te_components) and is_prequant_covered_weight( + rfilename, skip_te_components + ) + try: if kind == "pipeline": info = api.model_info(repo_id, files_metadata = True, token = hf_token) - picked = [s for s in info.siblings if _pipeline_file_downloaded(s.rfilename)] + picked = [ + s + for s in info.siblings + if _pipeline_file_downloaded(s.rfilename) and not _dense_te_shard(s.rfilename) + ] # diffusers prefers safetensors: drop a .bin whose dir also has a picked .safetensors. st_dirs = { s.rfilename.rsplit("/", 1)[0] @@ -1040,7 +1094,7 @@ class DiffusionBackend: base_info = api.model_info(base_repo, files_metadata = True, token = hf_token) base_bytes = 0 for s in base_info.siblings: - if base_filter(s.rfilename): + if base_filter(s.rfilename) and not _dense_te_shard(s.rfilename): base_files.append(s.rfilename) base_bytes += s.size or 0 total += base_bytes @@ -1059,6 +1113,7 @@ class DiffusionBackend: family_override: Optional[str] = None, model_kind: Optional[str] = None, hf_token: Optional[str] = None, + text_encoder_quant: Optional[str] = None, **load_kwargs: Any, ) -> dict[str, Any]: """The repos + exact files this pick needs, so the Hub download manager can fetch @@ -1067,13 +1122,21 @@ class DiffusionBackend: A plain snapshot_download would also pull what the loader deliberately skips (the packaged root single, transformer/ shards, fp16 twins) -- tens of GB per FLUX repo. Resolves family/kind/base exactly as ``_run_load`` does, so the plan and the load - agree. Local paths are already on disk and yield no entries.""" + agree. Local paths are already on disk and yield no entries. + + ``text_encoder_quant`` is read for the same reason as the DiT quant: an fp8 request + loads a hosted PRE-CAST encoder, so the base repo's dense encoder shards must not be + staged and the pre-cast checkpoint must be. Without it the manager stages the dense + encoder (tens of GB the load never opens) and the load then pulls the pre-cast file + inline, outside the manager's progress and disk preflight.""" fam = detect_family_for_pick(repo_id, gguf_filename, family_override) kind = resolve_model_kind(gguf_filename, model_kind) if kind == "pipeline": base = repo_id # the full pipeline IS the repo else: base = _resolve_base_repo(repo_id, base_repo, fam, hf_token) + # Only a checkpoint that really resolves on the Hub earns the right to drop dense shards. + te_files = self._te_prequant_plan_files(fam, text_encoder_quant, hf_token) sizes: dict[str, int] = {} total, base_files = self._estimate_download_bytes( repo_id, @@ -1085,8 +1148,19 @@ class DiffusionBackend: include_transformer = kind == "gguf" and self._dense_quant_prefetch_needed(fam, load_kwargs), sizes_out = sizes, + skip_te_components = tuple(te_files), ) entries: list[dict[str, Any]] = [] + for repo, files in te_files.values(): + entries.append( + { + "repo_id": repo, + "files": [name for name, _size in files], + "bytes": int(sum(size for _name, size in files)), + "gguf_filename": None, + } + ) + total += int(sum(size for _name, size in files)) if gguf_filename and not Path(repo_id).expanduser().exists(): entries.append( { diff --git a/studio/backend/core/inference/diffusion_te_prequant.py b/studio/backend/core/inference/diffusion_te_prequant.py index d95363ee06..05a5dc885c 100644 --- a/studio/backend/core/inference/diffusion_te_prequant.py +++ b/studio/backend/core/inference/diffusion_te_prequant.py @@ -418,6 +418,35 @@ def _validate_checkpoint(ckpt: Any, scheme: str, component: str, base: str, logg return True +def te_prequant_hub_files( + sources: dict[str, "TePrequantSource"], api: Any, logger: Any = None +) -> dict[str, list[tuple[str, int]]]: + """``{component: [(rfilename, size)]}`` for every hosted pre-cast checkpoint that really + resolves on the Hub. + + Only a component listed here may have its dense weights dropped from a plan or a prefetch: + an unpublished / gated / renamed artifact keeps its dense encoder, exactly as the load's own + fallback does. Checked per source so one missing repo cannot sink the whole plan. A local + path override is already on disk and is never staged.""" + found: dict[str, list[tuple[str, int]]] = {} + for component, source in sources.items(): + if getattr(source, "kind", None) != "repo" or not getattr(source, "filename", None): + continue + try: + info = api.model_info(source.location, files_metadata = True) + except Exception as exc: # noqa: BLE001 -- unavailable pre-cast means the dense encoder + _warn(logger, f"hub_files:{source.location}", exc) + continue + files = [ + (s.rfilename, int(getattr(s, "size", 0) or 0)) + for s in (info.siblings or []) + if s.rfilename == source.filename + ] + if files: + found[component] = files + return found + + def _has_meta_tensors(module: Any) -> bool: """True if any parameter or buffer is still on the meta device after loading.""" from itertools import chain diff --git a/studio/backend/core/training/diffusion_training_service.py b/studio/backend/core/training/diffusion_training_service.py index 41c793d503..4d78f9841d 100644 --- a/studio/backend/core/training/diffusion_training_service.py +++ b/studio/backend/core/training/diffusion_training_service.py @@ -78,6 +78,16 @@ class DatasetMutationInFlight(RuntimeError): """A start was refused because a dataset mutation is open (route: 409).""" +def _llm_training_active() -> bool: + """Whether the LLM trainer holds the GPU. Best-effort: an import/health failure must not + block a diffusion start, exactly as the route's own reciprocal check fails open.""" + try: + from core.training import get_training_backend + return bool(get_training_backend().is_training_active()) + except Exception: # noqa: BLE001 + return False + + # ── persisted run history ────────────────────────────────────────────────────── # Every terminal run is recorded as one JSON file (summary + scrubbed config + metric logs) so # the Train tab can show history. JSON, not the LLM sqlite, so diffusion runs stay off the LLM @@ -275,6 +285,16 @@ class DiffusionTrainingService: "A model is being loaded onto the GPU right now. Wait for that to finish, " "then start the run." ) + # The LLM trainer under the SAME lock, not just at the route's earlier check: that check + # and this reservation are separated by several network-bound preflights, so an LLM start + # could spawn in between and both trainers would allocate on one GPU. The LLM route holds + # gpu_load_admission() across its own spawn, so between the two either this raises or + # that one does -- never neither. + if _llm_training_active(): + raise RuntimeError( + "An LLM training job is already running. " + "Stop it before starting diffusion (Images) training." + ) self._reserved = True def unreserve(self) -> None: diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 16fe47b115..d53cb38ba4 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -16095,6 +16095,11 @@ async def diffusion_download_plan( model_kind = kind, hf_token = request.hf_token, transformer_quant = request.transformer_quant, + # An fp8 encoder request loads a hosted pre-cast checkpoint, so the plan has to stage + # that file instead of the base repo's dense encoder shards -- otherwise the manager + # downloads tens of GB the load never opens and then pulls the pre-cast file inline, + # outside its progress and disk preflight. + text_encoder_quant = request.text_encoder_quant, speed_mode = request.speed_mode, # The dense-quant prefetch decision reads the memory policy, the prequant path and the adapter # selection too (an offload policy never runs the dense build; a baked LoRA always does), so the diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index ad10021c0a..9bf01535d4 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -5,6 +5,7 @@ Training API routes """ +import contextlib import sys import threading from pathlib import Path @@ -541,12 +542,28 @@ async def start_training( # generation locks until an in-flight denoise step hits its cancel callback (and the export # subprocess teardown can take seconds), which would otherwise freeze every concurrent # status/cancel/UI request. Overlapping starts are serialized by the backend's own guard. - success = await asyncio.to_thread( - backend.start_training, - job_id = job_id, - before_spawn = _free_vram_for_training, - resume_source_run_id = resume_run["id"] if resume_run else None, - **training_kwargs, + # + # The diffusion admission is held ACROSS the spawn so the cross-trainer decision is + # atomic. The _diffusion_training_active() check above is separated from this point by + # dataset validation and memory coordination, and the diffusion route likewise checks + # this backend well before it reserves, so two near-simultaneous starts of different + # types could both pass their checks and train on the same GPU. Entering this context + # re-tests the diffusion state under the service's own lock, and while it is held + # reserve() refuses -- so exactly one of the two wins. + with _diffusion_gpu_admission(): + success = await asyncio.to_thread( + backend.start_training, + job_id = job_id, + before_spawn = _free_vram_for_training, + resume_source_run_id = resume_run["id"] if resume_run else None, + **training_kwargs, + ) + except _DiffusionStartInFlight as exc: + return TrainingJobResponse( + job_id = "", + status = "error", + message = str(exc), + error = "Diffusion training already active", ) except SidecarSwapInProgress as exc: # Expected loss of the race against a sidecar install: a retryable @@ -1160,6 +1177,46 @@ def _diffusion_training_active() -> bool: return False +class _DiffusionStartInFlight(RuntimeError): + """An LLM start lost the race to a diffusion start (route: refuse, don't spawn).""" + + +@contextlib.contextmanager +def _diffusion_gpu_admission(): + """Hold the diffusion service's GPU admission across the LLM spawn. + + Makes the cross-trainer admission atomic: entering re-tests the diffusion state under the + service's own lock and raises if a diffusion run is reserved or active, and while it is held + the diffusion ``reserve()`` refuses. So of two near-simultaneous starts of different types, + exactly one proceeds. Fails OPEN on an import/health failure, like every other guard here: a + chat-only install has no diffusion service and must still be able to train.""" + try: + from core.training.diffusion_training_service import ( + TrainingActiveError, + get_diffusion_training_service, + ) + service = get_diffusion_training_service() + except Exception: # noqa: BLE001 -- no diffusion stack: nothing to coordinate with + yield + return + try: + cm = service.gpu_load_admission() + except Exception: # noqa: BLE001 + yield + return + try: + cm.__enter__() + except TrainingActiveError as exc: + raise _DiffusionStartInFlight( + "A diffusion (Images) LoRA training job is already running. " + "Stop it before starting an LLM training run." + ) from exc + try: + yield + finally: + cm.__exit__(None, None, None) + + def _require_diffusion_dataset_mutable() -> None: """Reject a dataset mutation while a diffusion run is active. diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 77068ff7ec..e6529501b5 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -3868,3 +3868,118 @@ def test_download_plan_is_empty_for_a_local_path(tmp_path, monkeypatch): plan = DiffusionBackend().download_plan(str(local), gguf_filename = "weights.gguf") assert plan["entries"] == [] + + +def test_download_plan_stages_the_precast_encoder_instead_of_the_dense_one(monkeypatch): + # An fp8 text-encoder request loads a hosted PRE-CAST checkpoint, so the plan must stage that + # file and NOT the base repo's dense encoder shards. Without this the manager downloaded tens of + # GB the load never opens, and the load then pulled the pre-cast file inline, outside the + # manager's progress and disk preflight. + _fake_hf_api( + monkeypatch, + { + "unsloth/FLUX.1-dev-GGUF": [_FakeSibling("flux1-dev-Q4_K_M.gguf", 7 * GB)], + "black-forest-labs/FLUX.1-dev": _FLUX_BASE_SIBLINGS, + "unsloth/FLUX.1-schnell-FP8": [ + _FakeSibling("text_encoder_2-fp8.pt", 1 * GB), + _FakeSibling("README.md", 100), + ], + }, + ) + monkeypatch.setattr( + "core.inference.diffusion._resolve_base_repo", + lambda *a, **k: "black-forest-labs/FLUX.1-dev", + ) + monkeypatch.setattr( + DiffusionBackend, "_dense_quant_prefetch_needed", lambda self, fam, kwargs: False + ) + # The pick resolves one hosted pre-cast encoder for text_encoder_2 (flux.1 hosts its T5-XXL). + monkeypatch.setattr( + "core.inference.diffusion_te_prequant.te_prequant_sources", + lambda fam, *, te_quant_mode, target: ( + { + "text_encoder_2": types.SimpleNamespace( + kind = "repo", + location = "unsloth/FLUX.1-schnell-FP8", + filename = "text_encoder_2-fp8.pt", + ) + } + if te_quant_mode == "fp8" + else {} + ), + ) + + plan = DiffusionBackend().download_plan( + "unsloth/FLUX.1-dev-GGUF", + gguf_filename = "flux1-dev-Q4_K_M.gguf", + text_encoder_quant = "fp8", + ) + by_repo = {e["repo_id"]: e for e in plan["entries"]} + assert "unsloth/FLUX.1-schnell-FP8" in by_repo + assert by_repo["unsloth/FLUX.1-schnell-FP8"]["files"] == ["text_encoder_2-fp8.pt"] + base = by_repo["black-forest-labs/FLUX.1-dev"] + # text_encoder_2's dense weights are gone; text_encoder (no hosted artifact here) stays, and so + # do the non-weight files the pre-cast loader still meta-inits from. + assert not any(f.startswith("text_encoder_2/") and f.endswith(".safetensors") for f in base["files"]) + assert "text_encoder/model.safetensors" in base["files"] + assert "model_index.json" in base["files"] + assert plan["total_bytes"] == sum(e["bytes"] for e in plan["entries"]) + + +def test_download_plan_keeps_the_dense_encoder_without_an_fp8_request(monkeypatch): + # No fp8 request -> no hosted checkpoint -> the dense encoder is exactly as before. + _fake_hf_api( + monkeypatch, + { + "unsloth/FLUX.1-dev-GGUF": [_FakeSibling("flux1-dev-Q4_K_M.gguf", 7 * GB)], + "black-forest-labs/FLUX.1-dev": _FLUX_BASE_SIBLINGS, + }, + ) + monkeypatch.setattr( + "core.inference.diffusion._resolve_base_repo", + lambda *a, **k: "black-forest-labs/FLUX.1-dev", + ) + monkeypatch.setattr( + DiffusionBackend, "_dense_quant_prefetch_needed", lambda self, fam, kwargs: False + ) + plan = DiffusionBackend().download_plan( + "unsloth/FLUX.1-dev-GGUF", gguf_filename = "flux1-dev-Q4_K_M.gguf" + ) + base = next(e for e in plan["entries"] if e["repo_id"] == "black-forest-labs/FLUX.1-dev") + assert "text_encoder/model.safetensors" in base["files"] + assert len(plan["entries"]) == 2 + + +def test_download_plan_keeps_the_dense_encoder_when_the_precast_repo_is_unavailable(monkeypatch): + # A gated / renamed / unpublished artifact must NOT cost the dense encoder: the load falls back + # to the dense weights, so the plan has to stage them. + _fake_hf_api( + monkeypatch, + { + "unsloth/FLUX.1-dev-GGUF": [_FakeSibling("flux1-dev-Q4_K_M.gguf", 7 * GB)], + "black-forest-labs/FLUX.1-dev": _FLUX_BASE_SIBLINGS, + }, + ) + monkeypatch.setattr( + "core.inference.diffusion._resolve_base_repo", + lambda *a, **k: "black-forest-labs/FLUX.1-dev", + ) + monkeypatch.setattr( + DiffusionBackend, "_dense_quant_prefetch_needed", lambda self, fam, kwargs: False + ) + monkeypatch.setattr( + "core.inference.diffusion_te_prequant.te_prequant_sources", + lambda fam, *, te_quant_mode, target: { + "text_encoder": types.SimpleNamespace( + kind = "repo", location = "unsloth/does-not-exist", filename = "te-fp8.pt" + ) + }, + ) + plan = DiffusionBackend().download_plan( + "unsloth/FLUX.1-dev-GGUF", + gguf_filename = "flux1-dev-Q4_K_M.gguf", + text_encoder_quant = "fp8", + ) + base = next(e for e in plan["entries"] if e["repo_id"] == "black-forest-labs/FLUX.1-dev") + assert "text_encoder/model.safetensors" in base["files"] + assert not any(e["repo_id"] == "unsloth/does-not-exist" for e in plan["entries"]) diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py index 3be80c6989..749794bc41 100644 --- a/studio/backend/tests/test_diffusion_training.py +++ b/studio/backend/tests/test_diffusion_training.py @@ -1976,3 +1976,74 @@ def test_route_rejects_cond_cache_dir_for_sdxl(client): r = client.post("/api/train/diffusion/start", json = {**_BODY, "model_family": "sdxl"}) assert r.status_code == 200, r.text assert client._fake.started_with["cond_cache_dir"] is None + + +def test_service_reserve_refuses_while_the_llm_trainer_holds_the_gpu(monkeypatch): + # The route's reciprocal check runs several network-bound preflights before reserve(), so an LLM + # start could spawn inside that window and both trainers would allocate on one GPU. reserve() + # re-tests the LLM backend under its own lock, which is the half of the interlock that closes. + import types + + import core.training.diffusion_training_service as dts + from core.training.diffusion_training_service import DiffusionTrainingService + + svc = DiffusionTrainingService() + monkeypatch.setattr(dts, "_llm_training_active", lambda: True) + with pytest.raises(RuntimeError, match = "LLM training job is already running"): + svc.reserve() + assert svc.is_active() is False # the refused start left no claim behind + + monkeypatch.setattr(dts, "_llm_training_active", lambda: False) + svc.reserve() + assert svc.is_active() is True + + +def test_llm_active_probe_fails_open(monkeypatch): + # A chat-only install (or a wedged backend) must not block a diffusion start. + import sys + + import core.training.diffusion_training_service as dts + + monkeypatch.setitem(sys.modules, "core.training", None) # import raises + assert dts._llm_training_active() is False + + +def test_llm_start_holds_the_diffusion_admission_across_its_spawn(monkeypatch): + # The other half: while the LLM route is spawning, a diffusion reserve() must lose. The route + # holds gpu_load_admission() across start_training, and reserve() already refuses while an + # admission is open. + from core.training.diffusion_training_service import DiffusionTrainingService + import routes.training as tr + + svc = DiffusionTrainingService() + monkeypatch.setattr( + "core.training.diffusion_training_service.get_diffusion_training_service", lambda: svc + ) + with tr._diffusion_gpu_admission(): + with pytest.raises(RuntimeError, match = "being loaded onto the GPU"): + svc.reserve() + svc.reserve() # released once the spawn is done + + +def test_llm_start_admission_refuses_when_diffusion_is_already_reserved(monkeypatch): + from core.training.diffusion_training_service import DiffusionTrainingService + import routes.training as tr + + svc = DiffusionTrainingService() + monkeypatch.setattr( + "core.training.diffusion_training_service.get_diffusion_training_service", lambda: svc + ) + svc.reserve() + with pytest.raises(tr._DiffusionStartInFlight): + with tr._diffusion_gpu_admission(): + pass + + +def test_llm_start_admission_fails_open_without_a_diffusion_stack(monkeypatch): + import sys + + import routes.training as tr + + monkeypatch.setitem(sys.modules, "core.training.diffusion_training_service", None) + with tr._diffusion_gpu_admission(): + pass # a chat-only install still trains