Merge branch 'lab/issue-64-takane-vision' into 'export/v1-1-1'

See merge request onecomp/onecomp-lab!69

Author: FKKimura

CHANGELOG.md

@@ -8,13 +8,29 @@
 - Added `report_progress: bool = True` flag to `Runner.__init__` (`onecomp/runner.py`) and to the underlying entry points `run_chunked_quantization` (`onecomp/runner_methods/chunked_quantization.py`), `run_multi_gpu_quantization` / `run_quantization_phase` (`onecomp/runner_methods/multi_gpu_quantization.py`), `run_quantize_with_qep` (`onecomp/qep/_quantize_with_qep.py`), and `run_quantize_with_qep_arch` (`onecomp/qep/_quantize_with_qep_arch.py`) so long quantization runs (calibration, chunked, multi-GPU, QEP) report progress by default; pass `report_progress=False` for quiet runs
 - Demoted some INFO-level per-layer / per-chunk logs to DEBUG to avoid duplication with the new `[progress]` line (still available via `logging.basicConfig(level=logging.DEBUG)` for deep debugging)
 
-### Bug Fixes
+### Bug fixes: VLM save / load
 
-- Raise a clear error when ``Runner`` is configured with ``qep=True`` and a quantizer that does not support QEP (currently `JointQ`). Previously the run failed deep inside `quantize_with_qep` / `adjust_weight` with a confusing low-level error. `Runner.check()` now reports e.g. "Quantizer 'JointQ' (or one of its candidate quantizers) does not support QEP (Quantization Error Propagation). Set qep=False, or use a QEP-compatible quantizer (e.g., GPTQ, DBF, AutoBitQuantizer with QEP-compatible candidates)." Implementation: added `flag_qep_supported` (default `True`) on `Quantizer`, set to `False` on `JointQ`, and propagated via `AutoBitQuantizer._sync_flags` (only `True` when *all* candidate quantizers support QEP).
+- `Runner.save_quantized_model()` now copies all auxiliary `*.json` and `*.jinja` files (e.g. `preprocessor_config.json`, `processor_config.json`, `special_tokens_map.json`, `chat_template.jinja`) from the original model directory to the save directory, so the quantized model is fully self-contained for VLM / multimodal inference. Weight tensors (`*.safetensors`, `*.bin`, `*.pt`, `*.pth`), weight index files, `config.json` and `generation_config.json` are skipped, and any file already written by `model.save_pretrained` / `tokenizer.save_pretrained` is preserved (`runner.py`).
+- Source-model directory resolution (incl. `huggingface_hub.snapshot_download` fallback for Hub IDs) was extracted into a private helper `Runner._resolve_source_model_dir()` (`runner.py`).
+- `load_quantized_model()` now re-establishes the `lm_head` <-> `embed_tokens` weight tie for models with `tie_word_embeddings=True`. `load_state_dict(..., assign=True)` would otherwise leave `lm_head.weight` as the freshly initialised tensor (typically `float16`) while `embed_tokens.weight` got replaced with the checkpoint tensor (typically `bfloat16`), causing `RuntimeError: expected mat1 and mat2 to have the same dtype` at the final `lm_head` matmul during generation. The re-tie is gated on `lm_head` still being an `nn.Linear` so it does not interfere when `lm_head` itself was quantized (`quantized_model_loader.py`).
+- `load_quantized_model()` now reads `torch_dtype` from `config.json` when no explicit `torch_dtype` is passed by the caller, so the empty model is built in the same dtype as the saved checkpoint. Previously it always defaulted to `torch.float16`, which left non-quantized VLM submodules (e.g. `multi_modal_projector` in Cohere2Vision) at fp16 whenever `load_state_dict(..., assign=True)` could not find their key in the state_dict (`quantized_model_loader.py`).
+- `load_quantized_model()` now casts any leftover `float16` parameters and buffers of non-quantized modules to `model.config.torch_dtype` after the `lm_head` re-tie step. Quantized layers (`GPTQLinear`, `DoubleBinaryLinear`) and `float32` params (e.g. fp32 LayerNorm in mixed-precision models) are deliberately untouched. This generalises the existing `lm_head` re-tie to any non-quantized module and fixes the dtype mismatch reported in issue 64-3 (`RuntimeError: ... c10::Half != c10::BFloat16` on VLM image features) (`quantized_model_loader.py`).
+
+### Bug fixes: QEP + JointQ validation
+
+- Raise a clear error when `Runner` is configured with `qep=True` and a quantizer that does not support QEP (currently `JointQ`). Previously the run failed deep inside `quantize_with_qep` / `adjust_weight` with a confusing low-level error. `Runner.check()` now reports e.g. "Quantizer 'JointQ' (or one of its candidate quantizers) does not support QEP (Quantization Error Propagation). Set qep=False, or use a QEP-compatible quantizer (e.g., GPTQ, DBF, AutoBitQuantizer with QEP-compatible candidates)." Implementation: added `flag_qep_supported` (default `True`) on `Quantizer`, set to `False` on `JointQ`, and propagated via `AutoBitQuantizer._sync_flags` (only `True` when *all* candidate quantizers support QEP) (`quantizer/_quantizer.py`, `quantizer/jointq/_jointq.py`, `quantizer/autobit/_autobit.py`, `runner.py`).
+
+### Logging / observability tweaks
+
+- `Runner._copy_auxiliary_files()` now emits a matter-of-fact `INFO`-level log when an auxiliary file from the original model directory is *not* copied because the destination already contains a file of the same name (typically because `tokenizer.save_pretrained` wrote it just before, or a previous `save_quantized_model` call did). The new line is symmetrical to the existing `Copied %s to save directory` entry so the auxiliary-copy step can be audited end-to-end (`runner.py`).
+- `QuantizedModelLoader._cast_fp16_to_target_dtype()` now returns the list of fully-qualified parameter / buffer names whose dtype was actually converted instead of a plain count. The post-load `INFO` log in `load_quantized_model()` includes those names so it is obvious which non-quantized submodules were normalised by the safety-net cast (e.g. `multi_modal_projector.linear_*` in Cohere2Vision). Existing tests are updated accordingly and a new test pins the buffer-name reporting (`quantized_model_loader.py`, `tests/onecomp/runner/test_load_excluded_module_dtype.py`, `tests/onecomp/runner/test_save_quantized_aux_files.py`).
+- `QuantizedModelLoader.load_quantized_model()` now detects `tie_word_embeddings=True` even when the flag is nested in a sub-config (e.g. `model.config.text_config.tie_word_embeddings` in Llama 3.2-Vision and other torchtune-derived VLMs) by walking one level of sub-configs. Previously the flag was only read from the top-level `model.config`, so VLMs that placed it in `text_config` skipped the post-load re-tie; with HF deduplicating `lm_head.weight` for tied checkpoints, that left `lm_head.weight` at the empty-model random initial values rather than re-pointing to `embed_tokens.weight` (`quantized_model_loader.py`).
 
 ### Tests
 
-- Added `tests/onecomp/test_runner_check.py` covering the new `qep=True` validation path: JointQ + qep=True raises a clear `ValueError`, while JointQ + qep=False and GPTQ + qep=True both pass `Runner.check()`.
+- Added regression tests for the save/load fixes above: `tests/onecomp/runner/test_save_quantized_aux_files.py` (auxiliary-file copy whitelist), `tests/onecomp/runner/test_load_tied_embeddings.py` (tied-embedding dtype round-trip), and `tests/onecomp/runner/test_load_excluded_module_dtype.py` (non-quantized module dtype handling, including config-based empty-model dtype default, fp16 safety-net cast, fp32 preservation, and quantized-layer skip).
+- Added `tests/onecomp/test_runner_check.py` for the new `qep=True` validation path: JointQ + qep=True raises a clear `ValueError`, while JointQ + qep=False and GPTQ + qep=True both pass `Runner.check()`.
+- Added `tests/onecomp/runner/test_load_tied_embeddings.py::test_should_retie_word_embeddings_*` unit tests covering top-level, nested-text-config, all-False and unrelated-sub-attribute shapes.
 
 ### New Contributors
 

onecomp/quantized_model_loader.py

@@ -86,6 +86,53 @@ def load_quantized_model(
         # Load all weights (quantized + non-quantized) in one go
         model.load_state_dict(state_dict, strict=False, assign=True)
 
+        # ``assign=True`` swaps Parameter objects in place, which breaks the
+        # weight sharing established by ``from_config`` for models with
+        # ``tie_word_embeddings=True``.  Concretely, ``embed_tokens.weight``
+        # gets replaced by the bf16 tensor from the checkpoint while
+        # ``lm_head.weight`` keeps its original (often fp16) tensor, leading
+        # to a dtype mismatch at the final ``F.linear`` call during
+        # generation.  Re-tie when (a) the config tree still asks for it
+        # (multi-config VLMs such as Llama 3.2-Vision place the flag in
+        # ``text_config`` rather than at the top level, so we walk the
+        # nested configs) and (b) ``lm_head`` is still a plain
+        # ``nn.Linear`` -- if it has been replaced by a quantized layer
+        # (e.g. ``GPTQLinear``) it has no ``weight`` attribute to retie
+        # and tying would be meaningless.
+        if cls._should_retie_word_embeddings(model.config):
+            lm_head = getattr(model, "lm_head", None)
+            if isinstance(lm_head, torch.nn.Linear):
+                model.tie_weights()
+                logger.info("Re-tied lm_head to embed_tokens after assign-load")
+
+        # Safety net: ``load_state_dict(..., assign=True)`` only replaces
+        # parameters whose key in the checkpoint exactly matches the model's
+        # ``named_parameters`` path.  For some VLMs (e.g. Cohere2Vision's
+        # ``multi_modal_projector``) the path prefix differs between the
+        # checkpoint and the model class produced by ``from_config``, so a
+        # subset of params silently keeps the empty-model dtype.  Together
+        # with the config-based dtype default in
+        # ``_build_empty_model_from_config`` this normalises any remaining
+        # fp16 tensors of non-quantized modules to ``target_dtype``.  fp32
+        # params (e.g. fp32 LayerNorm in mixed-precision models) are left
+        # untouched, and quantized layers are skipped so that GPTQ scales
+        # and similar fp16 metadata are preserved.
+        target_dtype = (
+            torch_dtype
+            if torch_dtype is not None
+            else cls._resolve_dtype_from_config(config_dict)
+        )
+        if target_dtype is None:
+            target_dtype = torch.float16
+        converted = cls._cast_fp16_to_target_dtype(model, target_dtype)
+        if converted:
+            logger.info(
+                "Cast %d non-quantized fp16 tensor(s) to %s: %s",
+                len(converted),
+                target_dtype,
+                converted,
+            )
+
         # Register Hadamard hooks for rotation-preprocessed models
         if quant_config.get("rotated", False):
             from .pre_process.rotation_utils import register_online_hadamard_hooks
@@ -220,7 +267,119 @@ def _load_config_and_quant_config(save_directory: str) -> Tuple[Dict, Dict]:
         return config_dict, quant_config
 
     @staticmethod
+    def _cast_fp16_to_target_dtype(
+        model: torch.nn.Module, target_dtype: torch.dtype
+    ) -> List[str]:
+        """Cast fp16 params/buffers of non-quantized modules to ``target_dtype``.
+
+        Quantized layers (``GPTQLinear``, ``DoubleBinaryLinear``) are
+        skipped so their fp16 metadata (e.g. GPTQ ``scales``) is preserved.
+        Only fp16 tensors are cast: fp32 params (e.g. fp32 LayerNorm in
+        mixed-precision models) and other dtypes are left untouched.
+
+        Args:
+            model: The model whose parameters/buffers should be normalised.
+            target_dtype: Destination dtype.  When equal to
+                ``torch.float16`` this is a no-op.
+
+        Returns:
+            Fully-qualified names of every parameter / buffer whose
+            dtype was actually converted (e.g. ``"model.layers.0.mlp.
+            down_proj.weight"``).  An empty list means nothing needed
+            casting (or ``target_dtype == torch.float16``).  The list
+            form makes it easy for tests and operators to inspect
+            which submodules were touched by the safety net.
+        """
+        converted: List[str] = []
+        if target_dtype == torch.float16:
+            return converted
+        skip_types = (GPTQLinear, DoubleBinaryLinear)
+        for mod_name, mod in model.named_modules():
+            if isinstance(mod, skip_types):
+                continue
+            for p_name, p in mod.named_parameters(recurse=False):
+                if p.dtype == torch.float16:
+                    p.data = p.data.to(target_dtype)
+                    full_name = f"{mod_name}.{p_name}" if mod_name else p_name
+                    converted.append(full_name)
+            for b_name, b in mod.named_buffers(recurse=False):
+                if b.dtype == torch.float16:
+                    b.data = b.data.to(target_dtype)
+                    full_name = f"{mod_name}.{b_name}" if mod_name else b_name
+                    converted.append(full_name)
+        return converted
+
+    @classmethod
+    def _should_retie_word_embeddings(cls, config: Any) -> bool:
+        """Return True if any nesting level of ``config`` requests weight tying.
+
+        Single-config language models (e.g. Llama, Qwen) expose
+        ``tie_word_embeddings`` directly on ``model.config``.  Multi-
+        config VLMs vary: ``gemma-4`` puts the flag at the top level
+        but ``llama3.2-vlm-torchtune`` and other torchtune-derived
+        checkpoints place it inside ``text_config`` only, so the naive
+        ``getattr(model.config, "tie_word_embeddings", False)`` would
+        miss the tying request and skip the re-tie that
+        ``load_state_dict(..., assign=True)`` necessitates.
+
+        We walk the config tree shallowly: any direct sub-attribute
+        that itself exposes ``tie_word_embeddings`` is inspected.  The
+        check is intentionally non-recursive past one level because
+        HuggingFace nests language sub-configs at most one level deep
+        in practice (``text_config``, ``language_config`` etc.) and a
+        deeper recursion would risk being confused by unrelated
+        sub-objects.
+
+        Args:
+            config: A ``transformers.PretrainedConfig``-like object
+                (anything supporting ``getattr``).
+
+        Returns:
+            ``True`` if ``tie_word_embeddings`` is truthy at the top
+            level or on any direct sub-attribute that itself looks
+            like a config (i.e. carries a ``tie_word_embeddings``
+            attribute).  ``False`` otherwise.
+        """
+        if getattr(config, "tie_word_embeddings", False):
+            return True
+        try:
+            sub_items = vars(config).items()
+        except TypeError:
+            return False
+        for _, value in sub_items:
+            # Duck-type check: only descend into things that themselves
+            # carry the flag, so we don't accidentally walk unrelated
+            # auxiliary objects (e.g. tokenizer caches) that happen to
+            # be stored on the config.
+            if hasattr(value, "tie_word_embeddings") and getattr(
+                value, "tie_word_embeddings", False
+            ):
+                return True
+        return False
+
+    @staticmethod
+    def _resolve_dtype_from_config(
+        config_dict: Dict,
+    ) -> Optional[torch.dtype]:
+        """Read ``torch_dtype`` / ``dtype`` from a config dict.
+
+        Accepts both the JSON-serialised string form (e.g. ``"bfloat16"``)
+        and a real ``torch.dtype`` value.  Returns ``None`` when the field
+        is missing, ``"auto"``, or otherwise unresolvable.
+        """
+        for key in ("torch_dtype", "dtype"):
+            val = config_dict.get(key)
+            if isinstance(val, torch.dtype):
+                return val
+            if isinstance(val, str) and val and val != "auto":
+                resolved = getattr(torch, val, None)
+                if isinstance(resolved, torch.dtype):
+                    return resolved
+        return None
+
+    @classmethod
     def _build_empty_model_from_config(
+        cls,
         config_dict: Dict,
         torch_dtype: Optional[torch.dtype] = None,
     ) -> torch.nn.Module:
@@ -238,6 +397,13 @@ def _build_empty_model_from_config(
                 f"Cannot build config: model_type={model_type!r} not in CONFIG_MAPPING."
             )
 
+        # Default to the dtype recorded in config.json so the empty model
+        # starts in the same dtype as the saved checkpoint.  This avoids
+        # leaving non-quantized submodules at the hard-coded fp16 default
+        # if ``load_state_dict(..., assign=True)`` cannot find their key
+        # in the state_dict (e.g. tied or path-shifted VLM submodules).
+        if torch_dtype is None:
+            torch_dtype = cls._resolve_dtype_from_config(clean_config)
         dtype = torch_dtype if torch_dtype is not None else torch.float16
         config_cls = CONFIG_MAPPING[model_type]
         model_config = config_cls.from_dict(clean_config)