diff --git a/src/gen_worker/executor.py b/src/gen_worker/executor.py index faad2f45..e35052b9 100644 --- a/src/gen_worker/executor.py +++ b/src/gen_worker/executor.py @@ -505,6 +505,43 @@ def _model_failure_vocab(exc: BaseException) -> str: return "load_failed" +def _shared_lanes_need_fp8( + slot_sizes: Dict[str, Dict[str, int]], + shared_components: typing.Iterable[str], + free_vram_bytes: int, + *, + margin_bytes: int = 2 * (1 << 30), +) -> bool: + """Joint VRAM-fit decision for a shared-component multi-lane record + (th#1043): loading each lane's precision reactively, one at a time, + lets the FIRST lane to load consume all free VRAM at native precision + — the shared-component invariant then refuses the offload placement + the STARVED lane needs, hard-failing a fit that was achievable all + along. Decide precision for the WHOLE group against its combined + footprint (shared components counted once) before any lane loads. + + True when native precision doesn't fit the group but fp8-storage + (denoiser weights ~halved) does — every lane in the group should force + fp8 storage. False when native precision already fits (no forcing + needed) or fp8 storage still wouldn't fit (per-lane reactive sizing + makes its own honest call instead). + """ + shared = list(shared_components) + if len(slot_sizes) < 2 or not shared: + return False + first = next(iter(slot_sizes.values())) + shared_bytes = sum(first.get(c, 0) for c in shared) + exclusive_bytes = sum( + b for sizes in slot_sizes.values() + for comp, b in sizes.items() if comp not in shared + ) + needed = shared_bytes + exclusive_bytes + if needed + margin_bytes <= free_vram_bytes: + return False + fp8_needed = shared_bytes + 0.5 * exclusive_bytes + return fp8_needed + margin_bytes <= free_vram_bytes + + def _is_corrupt_load_error(exc: BaseException) -> bool: """Errors a truncated/corrupt snapshot produces at weights-load time (gw#408). Broad on purpose: the digest re-verify gate downstream @@ -5201,6 +5238,34 @@ def _component_share_plan( ) return plan + def _shared_group_force_fp8( + self, spec: EndpointSpec, share_plan: Optional[Dict[str, Dict[str, Any]]], + ) -> set: + """Slots that must force fp8 denoiser storage to fit their shared- + component group jointly resident (th#1043) — empty when the group + fits at native precision, or fits nobody's precision at all.""" + plan = share_plan or {} + slots = [s for s, m in plan.items() if m] + if len(slots) < 2: + return set() + shared_components = sorted({c for m in plan.values() for c in m}) + slot_sizes: Dict[str, Dict[str, int]] = {} + for slot in slots: + binding = spec.models.get(slot) + if binding is None: + return set() + slot_sizes[slot] = self.store.component_sizes(wire_ref(binding)) + free = self.store.residency.free_vram_bytes() + if not _shared_lanes_need_fp8(slot_sizes, shared_components, free): + return set() + logger.info( + "th#1043: shared-lane group %s for %s doesn't fit resident at " + "native precision (free=%.1fGiB) — forcing fp8 storage on every " + "lane before any of them loads", + slots, spec.name, free / (1 << 30), + ) + return set(slots) + @staticmethod def _model_index_components(path: str) -> set: """Component names the snapshot's model_index.json declares — the @@ -5245,6 +5310,7 @@ async def _injection_kwargs( result = _InjectionResult(kwargs=kwargs, loaded=loaded) compile_artifact = compile_selection.path if compile_selection else None share_plan = self._component_share_plan(spec, paths, hints) + force_fp8_slots = self._shared_group_force_fp8(spec, share_plan) if server is not None: for pname, ann in hints.items(): if ann is ServerHandle: @@ -5334,6 +5400,8 @@ async def _injection_kwargs( slot=slot, ref=ref, mode=mode, components=injected, declared_vram_gb=float( getattr(spec.resources, "vram_gb", 0) or 0), + force_storage_dtype=( + "fp8" if slot in force_fp8_slots else ""), ) except Exception as exc: # Corruption-shaped load failure (gw#408): digest-verify @@ -5358,6 +5426,8 @@ async def _injection_kwargs( slot=slot, ref=ref, mode=mode, components=injected, declared_vram_gb=float( getattr(spec.resources, "vram_gb", 0) or 0), + force_storage_dtype=( + "fp8" if slot in force_fp8_slots else ""), ) pipe = sl.obj # Reconcile the load outcomes into ServePlan/FnDegraded via diff --git a/src/gen_worker/models/provision.py b/src/gen_worker/models/provision.py index 64191fed..c9c86759 100644 --- a/src/gen_worker/models/provision.py +++ b/src/gen_worker/models/provision.py @@ -84,6 +84,7 @@ def load_slot( components: Optional[Dict[str, Any]] = None, device: str = "", declared_vram_gb: float = 0.0, + force_storage_dtype: str = "", ) -> SlotLoad: """Typed slot injection: the slot receives exactly what its ``setup`` annotation says — a ``str``/``Path`` local path, or a constructed @@ -96,7 +97,11 @@ def load_slot( degraded floors — the executor's knowledge; the CLI passes ``auto``). ``device="cpu"`` (CLI ``--device cpu``) skips placement entirely. ``components`` are preloaded shared modules (gw#479) forwarded to - ``from_pretrained``. + ``from_pretrained``. ``force_storage_dtype`` overrides the binding's own + storage_dtype (th#1043): a joint multi-lane fit decision made BEFORE any + lane in a shared-component group loads, so the first lane to load never + greedily consumes free VRAM at native precision and starves a sibling + lane into an offload placement the shared-component invariant refuses. """ if annotation is None or annotation is str: return SlotLoad(obj=path) @@ -110,7 +115,7 @@ def load_slot( from .memory import place_pipeline dtype = str(getattr(binding, "dtype", "") or "") - storage_dtype = str(getattr(binding, "storage_dtype", "") or "") + storage_dtype = force_storage_dtype or str(getattr(binding, "storage_dtype", "") or "") out = SlotLoad(obj=None, is_pipeline=True, ran=(dtype or "bf16")) # th#737: a cast directive on a denoiser-less diffusers tree is a diff --git a/tests/test_shared_lane_precision_th1043.py b/tests/test_shared_lane_precision_th1043.py new file mode 100644 index 00000000..ea6e2c4c --- /dev/null +++ b/tests/test_shared_lane_precision_th1043.py @@ -0,0 +1,54 @@ +"""th#1043: joint precision fit for shared-component multi-lane records. + +Real qwen-image evidence: two lanes (t2i/edit) share a ~15.5GB text +encoder+VAE, each carries an exclusive ~40GB bf16 transformer. Loading each +lane's precision reactively — one at a time, against whatever free VRAM +happens to be left at that moment — let the FIRST lane consume all headroom +at native precision, starving the second into an offload placement the +shared-component invariant then refuses outright (RetryableError, +"shared-component lanes require resident placement"). The fix decides +precision for the WHOLE group up front, against its COMBINED footprint. +""" + +from __future__ import annotations + +from gen_worker.executor import _shared_lanes_need_fp8 + +_GiB = 1024 ** 3 + +# Real th#1043 shapes: shared text_encoder+vae ~15.5GB, each transformer +# ~40GB bf16 (no fp8 checkpoint exists yet for this family). +_QWEN_SIZES = { + "t2i": {"text_encoder": 15 * _GiB, "vae": int(0.5 * _GiB), "transformer": 40 * _GiB}, + "edit": {"text_encoder": 15 * _GiB, "vae": int(0.5 * _GiB), "transformer": 40 * _GiB}, +} +_SHARED = ["text_encoder", "vae"] + + +def test_starved_group_forces_fp8_when_it_fits(): + # A100-80GB: ~79GiB usable. Native (15.5 + 40 + 40 = 95.5GiB) doesn't + # fit; fp8'd transformers (15.5 + 20 + 20 = 55.5GiB) do. + free = 79 * _GiB + assert _shared_lanes_need_fp8(_QWEN_SIZES, _SHARED, free) is True + + +def test_group_that_fits_natively_is_left_alone(): + # A card with enough headroom for both lanes at native precision needs + # no forcing. + free = 200 * _GiB + assert _shared_lanes_need_fp8(_QWEN_SIZES, _SHARED, free) is False + + +def test_group_that_cannot_fit_even_at_fp8_is_left_to_the_ladder(): + # Too small even halved: forcing fp8 here would just relocate the + # failure, not fix it — leave it to the existing per-lane ladder. + free = 30 * _GiB + assert _shared_lanes_need_fp8(_QWEN_SIZES, _SHARED, free) is False + + +def test_single_lane_group_never_forces(): + assert _shared_lanes_need_fp8({"t2i": _QWEN_SIZES["t2i"]}, _SHARED, 1 * _GiB) is False + + +def test_no_shared_components_never_forces(): + assert _shared_lanes_need_fp8(_QWEN_SIZES, [], 1 * _GiB) is False