From 18656d1d799cf1ea698264adc00b419529a3b958 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Fri, 29 May 2026 00:07:58 +0100 Subject: [PATCH] fix(native_transformer): forward load kwargs to from_config load_transformer threads **kwargs into its from_pretrained and from_single_file branches. The native dispatch branch bypasses both loaders and builds via cls.from_config, so those kwargs were dropped there. Thread them through load, build_component, and build_component_quantized into from_config so callers passing extra args alongside native_spec are honored instead of silently dropped. Siblings do not receive them. --- pipelines/generic.py | 1 + pipelines/native_transformer.py | 18 +++++++--- test/test-native-transformer.py | 63 +++++++++++++++++++++++++++++++++ 3 files changed, 78 insertions(+), 4 deletions(-) diff --git a/pipelines/generic.py b/pipelines/generic.py index 007d31ad6..93c6fbf1b 100644 --- a/pipelines/generic.py +++ b/pipelines/generic.py @@ -76,6 +76,7 @@ def load_transformer(repo_id, cls_name, load_config=None, subfolder="transformer modules_dtype_dict=modules_dtype_dict, quant_args=quant_args, quant_type=quant_type, + **kwargs, ) elif local_file is not None and local_file.lower().endswith('.safetensors'): log.debug(f'Load model: transformer="{local_file}" cls={cls_name.__name__} quant="{quant_type}" loader={_loader("diffusers")} args={load_args}') diff --git a/pipelines/native_transformer.py b/pipelines/native_transformer.py index 20f3b8145..b598caf99 100644 --- a/pipelines/native_transformer.py +++ b/pipelines/native_transformer.py @@ -179,6 +179,7 @@ def load( modules_dtype_dict: dict | None = None, quant_args: dict | None = None, quant_type: str | None = None, + **kwargs, ) -> tuple[object, dict[str, object]]: """Load the transformer (and any bundled siblings) from ``local_file``. @@ -193,7 +194,11 @@ def load( :func:`pipelines.generic.load_transformer` so the dispatch from there can plumb the caller's intent through unchanged. ``quant_args`` and ``quant_type`` are precomputed by the caller; when ``None`` they are - derived here via ``model_quant.get_dit_args``. + derived here via ``model_quant.get_dit_args``. Extra ``**kwargs`` are + forwarded to the transformer's ``cls.from_config``, the native path's + construction step (it bypasses ``from_pretrained``/``from_single_file``, + where ``load_transformer`` otherwise routes them). Siblings do not + receive them. Returns ``(transformer, siblings_dict)``. ``siblings_dict`` is keyed by sibling name and is empty for non-sibling specs, or for sibling specs @@ -247,6 +252,7 @@ def load( dtype=effective_dtype, modules_to_not_convert=modules_to_not_convert, modules_dtype_dict=modules_dtype_dict, + **kwargs, ) del transformer_sd devices.torch_gc() @@ -386,6 +392,7 @@ def build_component_quantized( quant_args: dict, dtype, acceptable_missing: tuple[str, ...], + **kwargs, ) -> object: """Build a component with per-tensor SDNQ quantization during load. @@ -424,7 +431,7 @@ def build_component_quantized( quantizer.torch_dtype = target_dtype with init_empty_weights(include_buffers=False): - component = cls.from_config(config) + component = cls.from_config(config, **kwargs) quantizer._process_model_before_weight_loading(component, device_map=None) # pylint: disable=protected-access @@ -483,6 +490,7 @@ def build_component( dtype=None, modules_to_not_convert: list | None = None, modules_dtype_dict: dict | None = None, + **kwargs, ) -> object: """Convert (if needed), instantiate, load weights, dtype-cast, quantize, and offload-place a single component. Raises on any hard failure. @@ -496,7 +504,8 @@ def build_component( ``dtype`` overrides ``devices.dtype`` when supplied; otherwise the global default is used. ``modules_to_not_convert`` and ``modules_dtype_dict`` are forwarded to :func:`apply_quant` for the post-mode path; pre-mode - receives them via the SDNQConfig in ``quant_args``. + receives them via the SDNQConfig in ``quant_args``. Extra ``**kwargs`` + reach ``cls.from_config`` for both construction paths. """ try: if converter is not None: @@ -514,13 +523,14 @@ def build_component( quant_args=quant_args, dtype=dtype, acceptable_missing=acceptable_missing, + **kwargs, ) del sd devices.torch_gc() return component log.debug(f'Load model: native_transformer {component_name} loading keys={len(sd)} cls={cls.__name__}') - component = cls.from_config(config) + component = cls.from_config(config, **kwargs) missing, unexpected = component.load_state_dict(sd, strict=False) validate_state_dict_load(component_name, missing, unexpected, acceptable_missing) del sd diff --git a/test/test-native-transformer.py b/test/test-native-transformer.py index b0a524b29..d9f9f06d4 100644 --- a/test/test-native-transformer.py +++ b/test/test-native-transformer.py @@ -422,6 +422,19 @@ class MockMiniTransformer(torch.nn.Module): self.rope.register_buffer('freqs', torch.zeros(dim)) +class MockKwargsTransformer(MockMiniTransformer): + """Records the kwargs from_config received. Mirrors diffusers from_config, + which accepts **kwargs (config overrides); lets a test assert the native + load path forwards caller kwargs to construction instead of dropping them.""" + + last_kwargs: dict = {} + + @classmethod + def from_config(cls, config: dict, **kwargs) -> 'MockKwargsTransformer': + cls.last_kwargs = dict(kwargs) + return cls(dim=config['dim']) + + def write_fixture(state_dict_keys: dict, fd: int, path: str) -> str: os.close(fd) safetensors.torch.save_file(state_dict_keys, path) @@ -484,6 +497,55 @@ def test_load_end_to_end_with_bfl_prefix_no_converter(): os.unlink(path) +def test_load_forwards_kwargs_to_from_config(): + """Caller **kwargs reach cls.from_config (the native path's only + construction step, since it bypasses from_pretrained/from_single_file). + Guards against silently dropping args a caller passes alongside + native_spec.""" + fd, path = tempfile.mkstemp(suffix='.safetensors') + try: + dim = 8 + raw = { + 'model.diffusion_model.in_proj.weight': torch.randn(dim, dim), + 'model.diffusion_model.in_proj.bias': torch.zeros(dim), + 'model.diffusion_model.out_proj.weight': torch.randn(dim, dim), + 'model.diffusion_model.out_proj.bias': torch.zeros(dim), + } + write_fixture(raw, fd, path) + + orig_fetch = nt.fetch_component_config + nt.fetch_component_config = lambda repo, sub: {'dim': dim} + from modules import model_quant + orig_get_dit = model_quant.get_dit_args + orig_get_qtype = model_quant.get_quant_type + orig_do_post = model_quant.do_post_load_quant + model_quant.get_dit_args = lambda *a, **k: ({}, {}) + model_quant.get_quant_type = lambda *a, **k: None + model_quant.do_post_load_quant = lambda *a, **k: None + + MockKwargsTransformer.last_kwargs = {} + try: + spec = nt.TransformerSpec(cls=MockKwargsTransformer) + transformer, _ = nt.load( + local_file=path, + repo_id='fake/repo', + spec=spec, + diffusers_cfg={}, + low_cpu_mem_usage=True, + ) + finally: + nt.fetch_component_config = orig_fetch + model_quant.get_dit_args = orig_get_dit + model_quant.get_quant_type = orig_get_qtype + model_quant.do_post_load_quant = orig_do_post + + assert MockKwargsTransformer.last_kwargs == {'low_cpu_mem_usage': True} + assert isinstance(transformer, MockKwargsTransformer) + finally: + if os.path.exists(path): + os.unlink(path) + + def test_load_end_to_end_with_sibling_partition(): """Bundled-sibling case: file carries both transformer and sibling weights, sibling_classes supplies the runtime sibling class, partition routes each @@ -711,6 +773,7 @@ def run_all(): cat = category('load') for fn in [ test_load_end_to_end_with_bfl_prefix_no_converter, + test_load_forwards_kwargs_to_from_config, test_load_end_to_end_with_sibling_partition, test_load_raises_on_missing_sibling_class, test_load_rejects_non_safetensors,