Plan the pre-cast text encoder, and make the cross-trainer GPU admission atomic

An fp8 text-encoder request loads a hosted PRE-CAST checkpoint, but the image
download plan never received text_encoder_quant, so the manager staged the base
repo's dense encoder (FLUX.2-dev's Mistral-24B is ~48 GB, Qwen-Image's
Qwen2.5-VL ~16.6 GB) and the load then pulled the pre-cast file inline, outside
the manager's progress and disk preflight. The plan now takes the field,
resolves the hosted artifact with the same resolver the injection uses, stages
that file, and drops only those components' dense weight shards. The load's own
prefetch takes the same treatment, since it paid the same cost. Only a
checkpoint that really resolves on the Hub earns the drop, so a gated or renamed
artifact still stages the dense encoder the load will fall back to.

The two trainers admitted each other with independent check-then-act guards:
the diffusion route checks the LLM backend several network-bound preflights
before it reserves, and the LLM route checks the diffusion service well before
it spawns, so two near-simultaneous starts could both pass and train on one GPU.
reserve() now re-tests the LLM backend under its own lock, and the LLM route
holds the diffusion service's gpu_load_admission across its spawn, so exactly
one of the two wins. Both halves fail open, so a chat-only install still
trains.
This commit is contained in:
Daniel Han 2026-07-27 09:41:31 +00:00
commit 794e43ffa1
7 changed files with 381 additions and 10 deletions

View file

@ -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(
{

View file

@ -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

View file

@ -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:

View file

@ -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

View file

@ -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.

View file

@ -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"])

View file

@ -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