From f9d2bbe0809134fda214888c673b32daffc25c18 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Sun, 5 Jul 2026 18:05:13 +0100 Subject: [PATCH] fix(krea2): zero-fill dormant last.up/last.down residual branch Base ships last.up all-zeros so the branch is a no-op; pre-branch finetunes omit both keys. A zero_init_missing spec field zero-fills them on load instead of falling back to the base transformer. --- pipelines/krea2/__init__.py | 9 ++- pipelines/native_transformer.py | 66 ++++++++++++++++ test/test-krea2-transformer.py | 131 ++++++++++++++++++++++++++------ 3 files changed, 181 insertions(+), 25 deletions(-) diff --git a/pipelines/krea2/__init__.py b/pipelines/krea2/__init__.py index bea25eec4..aaab5a94c 100644 --- a/pipelines/krea2/__init__.py +++ b/pipelines/krea2/__init__.py @@ -5,4 +5,11 @@ from pipelines.native_transformer import TransformerSpec # Checkpoint keys are bare (`first.`, `blocks.N.`, `txtfusion.`, ...) and the transformer's # module tree mirrors them exactly, so no state-dict conversion is needed. The model has no # rope/pos buffers, so the default acceptable-missing set is sufficient. -KREA2_SPEC = TransformerSpec(cls=Krea2Transformer2DModel, converter=None) +# +# `last.up`/`last.down` are a dormant zero-init residual branch (base ships last.up.weight +# all-zeros). Finetunes predating it omit both keys; zero-fill so the branch stays a no-op. +KREA2_SPEC = TransformerSpec( + cls=Krea2Transformer2DModel, + converter=None, + zero_init_missing=("last.down.", "last.up."), +) diff --git a/pipelines/native_transformer.py b/pipelines/native_transformer.py index 07aa9e094..6b8da433b 100644 --- a/pipelines/native_transformer.py +++ b/pipelines/native_transformer.py @@ -100,6 +100,12 @@ class TransformerSpec: prefixes, no converter, no siblings). Arches with bundled-sibling components (Anima) or unusual key conventions (custom converters, Cosmos-style structural markers) override the relevant fields. + + ``zero_init_missing`` lists key prefixes for a zero-initialized residual + branch the class defines but some checkpoints omit. Unlike + ``acceptable_missing`` (buffer-only keys left at their init), these are + zero-filled on load so the branch stays a no-op, matching a base model + that ships the branch dormant (output projection all-zeros). """ cls: type @@ -108,6 +114,7 @@ class TransformerSpec: converter: Callable[[dict], dict] | None = None siblings: dict[str, SiblingSpec] = field(default_factory=dict) acceptable_missing: tuple[str, ...] = DEFAULT_ACCEPTABLE_MISSING + zero_init_missing: tuple[str, ...] = () forbidden_markers: tuple[tuple[str, str], ...] = () @@ -260,6 +267,7 @@ def load( cls=spec.cls, converter=spec.converter, acceptable_missing=spec.acceptable_missing, + zero_init_missing=spec.zero_init_missing, quant_args=quant_args, quant_type=quant_type, dtype=effective_dtype, @@ -401,6 +409,7 @@ def build_component_quantized( quant_args: dict, dtype, acceptable_missing: tuple[str, ...], + zero_init_missing: tuple[str, ...] = (), **kwargs, ) -> object: """Build a component with per-tensor SDNQ quantization during load. @@ -480,6 +489,9 @@ def build_component_quantized( pbar.update(task, advance=1) missing = sorted(expected_keys - loaded_keys) + missing = materialize_zero_init( + component, missing, zero_init_missing, device=target_device, dtype=target_dtype + ) validate_state_dict_load(component_name, missing, unexpected, acceptable_missing) component = quantizer._process_model_after_weight_loading(component) # pylint: disable=protected-access @@ -494,6 +506,7 @@ def build_component( cls: type, converter: Callable[[dict], dict] | None, acceptable_missing: tuple[str, ...], + zero_init_missing: tuple[str, ...] = (), quant_args: dict, quant_type: str | None, dtype=None, @@ -539,6 +552,7 @@ def build_component( quant_args=quant_args, dtype=dtype, acceptable_missing=acceptable_missing, + zero_init_missing=zero_init_missing, **kwargs, ) del sd @@ -548,6 +562,7 @@ def build_component( log.debug(f'Load model: transformer=native {component_name} loading keys={len(sd)} cls={cls.__name__}') component = cls.from_config(config, **kwargs) missing, unexpected = component.load_state_dict(sd, strict=False) + missing = materialize_zero_init(component, missing, zero_init_missing) validate_state_dict_load(component_name, missing, unexpected, acceptable_missing) del sd devices.torch_gc() @@ -613,6 +628,57 @@ def validate_state_dict_load( ) +def materialize_zero_init( + component: object, + missing: list[str], + zero_init_missing: tuple[str, ...], + *, + device=None, + dtype=None, +) -> list[str]: + """Zero-fill missing weights that belong to a zero-initialized residual + branch the class defines but the override file predates. + + A base model may ship a residual branch dormant: its output projection is + all-zeros, so ``out + up(down(x))`` reduces to ``out`` and the branch is a + no-op. A checkpoint made before the branch existed omits those keys. Left + unmaterialized they keep random ``from_config`` init (non-quant path) or + stay on the meta device (per-tensor quant path); either injects garbage + into the output. Zero-filling reproduces the base's dormant behavior. + + Returns ``missing`` with the handled keys removed, so + :func:`validate_state_dict_load` does not treat them as a hard mismatch. + ``device``/``dtype`` override the target placement (the quant path passes + the real device since its params are on meta); when omitted each key's own + parameter device/dtype is used. Keys matching a prefix but absent from the + component's parameters are left in ``missing`` untouched. + """ + if not zero_init_missing: + return missing + from accelerate.utils import set_module_tensor_to_device + params = dict(component.named_parameters()) + handled: list[str] = [] + remaining: list[str] = [] + for key in missing: + param = params.get(key) + if param is not None and any(key.startswith(p) for p in zero_init_missing): + tgt_device = device if device is not None else param.device + tgt_dtype = dtype if dtype is not None else param.dtype + set_module_tensor_to_device( + component, key, tgt_device, + value=torch.zeros(param.shape, dtype=tgt_dtype), dtype=tgt_dtype, + ) + handled.append(key) + else: + remaining.append(key) + if handled: + log.debug( + f"Load model: transformer=native zero-init {len(handled)} dormant " + f"residual key(s): {', '.join(handled)}" + ) + return remaining + + def apply_quant( transformer: object, quant_type: str | None, diff --git a/test/test-krea2-transformer.py b/test/test-krea2-transformer.py index d62d937ec..b46341f71 100644 --- a/test/test-krea2-transformer.py +++ b/test/test-krea2-transformer.py @@ -1,11 +1,15 @@ #!/usr/bin/env python -"""Offline parity test: Krea2Transformer2DModel vs the reference SingleStreamDiT. +"""Offline tests for the Krea2 transformer port and its native loader handling. -Builds both models from one tiny config, copies the reference state dict into the diffusers -port, runs identical inputs, and asserts the forward outputs match. No server, no checkpoint. +- Parity: Krea2Transformer2DModel vs the reference SingleStreamDiT. Builds both from one tiny + config, copies the reference state dict into the diffusers port, runs identical inputs, and + asserts the forward outputs match. Needs the reference (mmdit.py) at $KREA2_REF_DIR + (default /home/ohiom/database/watering-hole). +- Zero-init regression: a checkpoint that omits the dormant last.up/last.down residual branch, + after materialize_zero_init, produces the same output as the base whose up is zeroed. Port + only, so it runs without the reference. -The reference checkpoint repo (mmdit.py) is expected at $KREA2_REF_DIR -(default /home/ohiom/database/watering-hole). +No server, no checkpoint. """ import importlib.util @@ -17,6 +21,12 @@ import torch REF_DIR = os.environ.get("KREA2_REF_DIR", "/home/ohiom/database/watering-hole") +CFG = dict( + features=128, tdim=32, txtdim=64, heads=4, kvheads=2, multiplier=4, + layers=2, patch=2, channels=4, bias=False, theta=1e3, + txtlayers=3, txtheads=2, txtkvheads=2, +) + def load_reference(): sys.path.insert(0, REF_DIR) @@ -35,34 +45,29 @@ def load_port(): return mod -def main(): - mmdit = load_reference() - port = load_port() - - cfg = dict( - features=128, tdim=32, txtdim=64, heads=4, kvheads=2, multiplier=4, - layers=2, patch=2, channels=4, bias=False, theta=1e3, - txtlayers=3, txtheads=2, txtkvheads=2, - ) - - torch.manual_seed(0) - ref = mmdit.SingleStreamDiT(mmdit.SingleMMDiTConfig(**cfg)).float().eval() - mine = port.Krea2Transformer2DModel(**cfg).float().eval() - missing, unexpected = mine.load_state_dict(ref.state_dict(), strict=False) - assert not missing, f"missing keys when loading reference weights: {missing}" - assert not unexpected, f"unexpected keys when loading reference weights: {unexpected}" - +def make_inputs(): batch, txtlen, imglen = 2, 5, 9 - cdim = cfg["channels"] * cfg["patch"] ** 2 + cdim = CFG["channels"] * CFG["patch"] ** 2 seq = txtlen + imglen gen = torch.Generator().manual_seed(1) img = torch.randn(batch, imglen, cdim, generator=gen) - context = torch.randn(batch, txtlen, cfg["txtlayers"], cfg["txtdim"], generator=gen) + context = torch.randn(batch, txtlen, CFG["txtlayers"], CFG["txtdim"], generator=gen) timestep = torch.rand(batch, generator=gen) pos = torch.randint(0, 16, (batch, seq, 3), generator=gen).float() mask = torch.ones(batch, seq, dtype=torch.bool) mask[0, -2:] = False # exercise the key-padding path + return img, context, timestep, pos, mask + +def run_parity(mmdit, port): + torch.manual_seed(0) + ref = mmdit.SingleStreamDiT(mmdit.SingleMMDiTConfig(**CFG)).float().eval() + mine = port.Krea2Transformer2DModel(**CFG).float().eval() + missing, unexpected = mine.load_state_dict(ref.state_dict(), strict=False) + assert not missing, f"missing keys when loading reference weights: {missing}" + assert not unexpected, f"unexpected keys when loading reference weights: {unexpected}" + + img, context, timestep, pos, mask = make_inputs() with torch.no_grad(): out_ref = ref(img, context, timestep, pos, mask) out_mine = mine( @@ -84,5 +89,83 @@ def main(): print("PARITY OK") +def load_materialize_zero_init(): + """Import the real loader helper. native_transformer pulls in modules.shared, which needs + cmd_args parsed first, so bootstrap it the same way the native-transformer suite does.""" + repo = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + if repo not in sys.path: + sys.path.insert(0, repo) + os.environ.setdefault("SD_INSTALL_QUIET", "1") + import modules.cmd_args + import installer + orig_argv = sys.argv + sys.argv = [sys.argv[0]] + try: + modules.cmd_args.parse_args() + finally: + sys.argv = orig_argv + installer.add_args(modules.cmd_args.parser) + modules.cmd_args.parsed, _ = modules.cmd_args.parser.parse_known_args([]) + from pipelines.native_transformer import materialize_zero_init + return materialize_zero_init + + +def run_zero_init_regression(port): + """A finetune predating the last.up/last.down branch omits both keys. Zero-filling them + must reproduce the base's dormant behavior: identical output to a model whose up is zeroed + (up(down(x)) == 0), while a genuinely missing weight stays a hard mismatch.""" + materialize_zero_init = load_materialize_zero_init() + branch_keys = ("last.up.weight", "last.down.weight") + + torch.manual_seed(0) + full_sd = port.Krea2Transformer2DModel(**CFG).state_dict() + + # Ground truth: the branch dormant exactly as the base ships it (up zeroed, down arbitrary). + dormant = port.Krea2Transformer2DModel(**CFG).float().eval() + dormant_sd = dict(full_sd) + dormant_sd["last.up.weight"] = torch.zeros_like(dormant_sd["last.up.weight"]) + dormant.load_state_dict(dormant_sd) + + # Under test: a checkpoint that omits both branch keys; the loader zero-fills them. A real + # missing weight (blocks.0.attn.wq.weight) is dropped too, to confirm it is NOT zero-filled. + filled = port.Krea2Transformer2DModel(**CFG).float().eval() + hard_key = "blocks.0.attn.wq.weight" + partial_sd = {k: v for k, v in full_sd.items() if k not in branch_keys and k != hard_key} + missing, unexpected = filled.load_state_dict(partial_sd, strict=False) + assert not unexpected, f"unexpected keys: {unexpected}" + assert set(missing) == set(branch_keys) | {hard_key}, f"unexpected missing set: {missing}" + + remaining = materialize_zero_init(filled, missing, ("last.down.", "last.up.")) + assert remaining == [hard_key], f"expected only {hard_key} to remain hard-missing, got {remaining}" + for k in branch_keys: + w = dict(filled.named_parameters())[k] + assert w.abs().sum().item() == 0.0, f"{k} was not zero-filled" + + # Restore the genuinely-missing weight so the forward is well defined, then compare. + filled.load_state_dict({hard_key: full_sd[hard_key]}, strict=False) + + img, context, timestep, pos, mask = make_inputs() + + def forward(model): + with torch.no_grad(): + return model( + hidden_states=img, encoder_hidden_states=context, timestep=timestep, + position_ids=pos, attention_mask=mask, return_dict=False, + )[0] + + diff = (forward(dormant) - forward(filled)).abs().max().item() + print(f"zero-init regression: max abs diff vs dormant base: {diff:.3e}") + assert diff == 0.0, f"ZERO-INIT REGRESSION FAILED: filled output differs by {diff:.3e}" + print("ZERO-INIT OK") + + +def main(): + port = load_port() + # Parity runs first in pristine torch state; the regression imports modules afterwards. + mmdit = load_reference() + run_parity(mmdit, port) + run_zero_init_regression(port) + + if __name__ == "__main__": main()