Feat/sa3 lora#313
Draft
ryanontheinside wants to merge 6 commits into
Draft
Conversation
…dling (SA3 LoRA phase 0) - lora_metadata: header-only safetensors sniff -> lora_family (ace/sa3), SA3 rank/adapter_type/base_model synthesis from the embedded lora_config, composite cache key (weights + sidecar + trigger mtimes), canonical_sa3_lineage for trainer/HF/runtime checkpoint spellings - engine/lora: _compute_deltas raises on zero PEFT-pattern matches so a wrong-family enable fails loudly instead of silently no-opping; register_lora accepts an explicit lora_id - ace_backend: lora_compatible gains the family axis (sa3 = hard no, unknown stays permissive) - paths: assign_lora_ids disambiguates duplicate-stem discoveries with parent-dir suffixes; register_library and /api/loras share it so the HTTP and engine catalogs agree on ids
scripts/sa3/lora_derisk_phase05.py runs churn / rollback / svd / hygiene / bench against the real vendored stable_audio_3 package and the medium checkpoint. Measured on the 5090: - slot churn + strength targeting bit-exact through middle-slot removal (validates the manager-owned direct-copy install design) - mid-enable failure rollback restores the model bit-identically - -xs SVD cost: 70.3s CPU full enable on medium; cache would be ~3.84 GB fp16 full-basis (confirms phase-1 rejection of -xs) - cached-context teardown leaves pristine weights and bit-identical forwards (the SA3Backend.close() contract) - eager DiT step: 56ms base, +49%/1 LoRA, +114%/3 LoRA naive; parametrize.cached() collapses 3-LoRA overhead to +15%, so the per-tick cached() wrap is promoted to a phase 1 requirement
…(phase 1) - engine/lora: extract LoRACatalogBase (catalog dict, state machine, prewarm executor, inspection) from the delta-merge LoRAManagerBase; ACE behavior byte-identical, existing tests untouched and green - engine/sa3_lora (new): SA3LoRAManager over the vendored stable_audio_3 parametrization engine. Manager-owned id->lora_index slots (monotonic, never reused) with direct-copy weight install, transactional enable with bit-identical rollback, hard format/ lineage validation at materialize (-xs rejected with a clear error, ACE files fail loudly), conditioner-target detection, and a close() that returns the process-cached model to pristine state. gc.collect before empty_cache is load-bearing for VRAM reclaim (parametrize's injected-class cycles defer frees to the cycle collector) - generator_backend/diffusion_backend: LoRA facade on the protocol with engine_obj-delegating defaults (ACE unchanged) - session: every engine_obj LoRA touchpoint now routes through the facade (_apply_lora_pending, _enabled_lora_ids, _resolve_lora_id, lora_catalog_payload) - sa3_backend: lora capability bit, lora_str_<id> knob expansion, per-tick strength reads with the 0.02 gate, has_pending_refit over the pending queues, conditioner-mutation serialization + cond-bundle rebuild through the prompt-swap path, interim TRT->eager DiT swap while LoRAs are active, per-tick parametrize.cached() wrap (phase 0.5 benchmark: +114% naive -> +15% cached at 3 LoRAs), close() - sa3_session: manager construction against the cached context, library registration, startup alias resolution/prewarm, knob seeding - families: manager threaded through backend_init; sa3 knob universe gains the lora_str placeholder for the homonym guard - tests: test_lora_sa3_manager (real vendored parametrize on a tiny tree, CPU), test_lora_facade_session (pending-drain via stub backend), SA3 backend LoRA wiring tests; scripts/sa3/lora_smoke.py GPU smoke passes twice back-to-back on the cached medium context
…efit mirror (phase 2 code) - trt/refit_core (new): shared typed-Weights push + commit/missing- weights diagnostics; ACE's TRTLoRAManager now routes its pushes through it (byte-identical, its unit tests untouched and green) - sa3_build: --refit builds sa3_m_dit_refit_l* engines with BuilderFlag.REFIT via a separate config dataclass so existing engine metadata identity is untouched (dry-run verified) - sa3_trt: refit engines are excluded from the shared deserialization cache (exclusive per-session ownership; dropping the engine with the session is the base-weight rollback); find_dit_engine gains want_refittable (LoRA sessions prefer refit engines, else fp16mixed over fp8, logged); SA3TRTDit exposes engine/engine_path/refittable - scripts/sa3/gen_sa3_refit_manifest (new): torch-FQN -> ONNX initializer mapping by shape + value fingerprint with transpose orientation; refuses partial manifests; --validate-engine runs the bit-identity gate (base-value refit must leave outputs bitwise unchanged) - sa3_trt_lora (new): SA3TRTRefitMirror — the eager SA3LoRAManager stays the source of truth; the mirror reads each parametrized module's merged weight (exact for every adapter variant by construction) and refits the exclusively-owned engine; dirty-set pushes base weights back after a disable - sa3_backend: mirror wiring (no eager swap in mirror mode; sync on enable/disable); knob strength changes under the mirror are stashed in rebuild_imminent and surfaced via has_pending_refit, applied post-pre-coverage in _prepare_tick — never a refit unannounced; eager sessions keep the inline buffer write - tests: mirror manifest/dirty-set/orientation vs a fake refitter with real vendored parametrizations; backend routing; upstream-loader parity for lora/dora-rows/dora-cols/bora and stack-order determinism GPU-gated (idle GPU required): the refit engine build, the manifest bit-identity validation run, and eager-vs-refit output parity.
A real underfit-trained LoRA found in the wild writes base_model="sa3-medium" in its lora_config header; canonical_sa3_lineage previously fell back to permissive-unknown on that spelling instead of a positive lineage match.
GPU validation of the TRT refit path: refit engine built idle-GPU; manifest bit-identity gate PASS (228/228 weights, 0 unmatched or ambiguous); lora_smoke PASS twice with a real underfit-trained LoRA (goa_trance_sa3m); eager-vs-refit per-step cos >= 0.99989 at strengths 0/0.5/1.0 and with a 2-LoRA stack; strength-0 and both restore phases value-identical; full 8-step denoise final-latent cos 0.99947. Status ledgers in notes/SA3_LORA_PLAN.md (local) updated to match.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
adds lora support with trt refit for sa3