Merge pull request #20 from FujitsuResearch/develop/v1-1-1

Develop/v1-1-1

Author: FKKimura

.gitignore

@@ -11,6 +11,7 @@ models
 .cursor/rules/test-workflow.mdc
 .cursor/rules/gitlab-integration.mdc
 .cursor/rules/slurm-submit.mdc
+.cursor/rules/run-tests-examples.mdc
 .hydra/
 *.out
 *.err

CHANGELOG.md

@@ -1,5 +1,44 @@
 # Change log
 
+## [v1.1.1] 2026-05-21
+
+### New Feature: Quantization progress logging
+
+- Added `QuantizationProgressTracker` (`onecomp/utils/quantization_progress.py`) that emits a single `[progress]` INFO line per completed step with done/total, percentage, elapsed time, and a linear ETA estimate; supports an optional `thread_safe=True` mode for multi-GPU quantization
+- 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: 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`).
+
+### Bug fixes: VLM save / load
+
+- `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`).
+- Added regression tests `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).
+- Loosened `test_save_load_pipeline_tinyllama.py` and `test_save_load_pipeline_qwen3.py` save/load round-trip threshold from absolute `1e-3` to relative `1%` of the per-tensor logits magnitude (`tests/onecomp/pre_process/test_save_load_pipeline_*.py`). The original absolute bound was below fp16's representable precision once accumulated through the 22-28 decoder layers of TinyLlama / Qwen3, causing the `gptq + save_dequantized` cases to fail on aarch64 + Blackwell (GB200) where cuBLAS picks slightly different reduction kernels than reference x86_64 / Hopper hosts. The save/load equivalence intent is preserved via the relative comparison, which is robust to platform-specific fp16 rounding noise.
+- Set `gpu_memory_utilization=0.78` explicitly when constructing `LLM(...)` in `example/vllm_inference/example_autobit_vllm_inference.py` and `example/vllm_inference/example_gptq_vllm_inference.py`. The vLLM default `0.92` cgroup-OOMs on UMA hosts (e.g. DGX Spark / GB200, 121.7 GiB UMA) because vLLM's startup memory check fails: the residual quantizer process leaves only ~106 GiB free, which is below `0.92 * 121.7 = 111.96 GiB`. `0.78` matches the value already used in `tests/vllm_plugins/gptq/test_mixed_gptq_e2e.py` and is documented in the workspace `slurm-submit.mdc` rule.
+
+### 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 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
+
+- [@sotanengel](https://github.com/sotanengel) made their first contribution in [#13](https://github.com/FujitsuResearch/OneCompression/pull/13)
+
 ## [v1.1.0] 2026-04-16
 
 ### Gemma 3 / Gemma 4 & VLM Support

example/.gitignore

@@ -1,4 +1,6 @@
 buf*.py
 run_*.sh
-tinyllama_gptq4_*/
+tinyllama_gptq4*/
 TinyLlama*/
+quantized_model_saveload/
+rotated_model_*

example/vllm_inference/example_autobit_vllm_inference.py

@@ -38,12 +38,17 @@ def main():
     gc.collect()
     torch.cuda.empty_cache()
 
-    # Step 2: Load the quantized model with vLLM and generate text
+    # Step 2: Load the quantized model with vLLM and generate text.
+    # gpu_memory_utilization=0.78 leaves headroom for the residual
+    # quantizer process (~16 GiB) on a UMA 121.7 GiB device (e.g. DGX
+    # Spark / GB200). The vLLM default 0.92 cgroup-OOMs on shared-memory
+    # GPUs.
     llm = LLM(
         model=save_dir,
         max_model_len=512,
         dtype="float16",
         enforce_eager=True,
+        gpu_memory_utilization=0.78,
     )
 
     prompts = [

example/vllm_inference/example_gptq_vllm_inference.py

@@ -67,12 +67,17 @@ def main():
     gc.collect()
     torch.cuda.empty_cache()
 
-    # Step 3: Load the quantized model with vLLM
+    # Step 3: Load the quantized model with vLLM.
+    # gpu_memory_utilization=0.78 leaves headroom for the residual
+    # quantizer process (~16 GiB) on a UMA 121.7 GiB device (e.g. DGX
+    # Spark / GB200). The vLLM default 0.92 cgroup-OOMs on shared-memory
+    # GPUs.
     llm = LLM(
         model=save_dir,
         max_model_len=512,
         dtype="float16",
         enforce_eager=True,
+        gpu_memory_utilization=0.78,
     )
 
     # Step 4: Generate text

onecomp/__version__.py

@@ -6,4 +6,4 @@
 
 """
 
-__version__ = "1.1.0"
+__version__ = "1.1.1"

onecomp/analyzer/quantization_error.py

@@ -107,15 +107,21 @@ def _plot_layer_type(ax, layer_type, label_points, colors):
                 cy = min(y, ylim * 0.92)
                 text = f"{k} ({x:.1f}, {y:.2f})"
                 ax.annotate(
-                    text, (cx, cy),
-                    fontsize=6, color=color, fontstyle="italic",
+                    text,
+                    (cx, cy),
+                    fontsize=6,
+                    color=color,
+                    fontstyle="italic",
                     bbox=dict(boxstyle="round,pad=0.2", fc="white", ec=color, alpha=0.8),
                 )
             else:
                 ax.annotate(
-                    str(k), (x, y),
-                    textcoords="offset points", xytext=(5, 3),
-                    fontsize=7, color=color,
+                    str(k),
+                    (x, y),
+                    textcoords="offset points",
+                    xytext=(5, 3),
+                    fontsize=7,
+                    color=color,
                 )
 
     ax.set_xlim(0, xlim)

onecomp/lpcd/_gradient_solver.py

@@ -2,6 +2,7 @@
 Copyright 2025-2026 Fujitsu Ltd.
 
 """
+
 from abc import ABC, abstractmethod
 
 import torch
@@ -19,7 +20,7 @@ class PenaltyFn(ABC):
     @abstractmethod
     def __call__(self, module_q: LpcdMetric, module_f: LpcdMetric) -> torch.Tensor:
         pass
- 
+
 
 def gradient_solver(
     lpcd_config: LPCDConfig,
@@ -35,19 +36,21 @@ def gradient_solver(
     device = lpcd_config.device
     batch_size = 16
 
-    assert lpcd_config.gd_batch_size % batch_size == 0, \
-        f"gd_batch_size should be divisible by batch_size, " + \
-        f"but got {lpcd_config.gd_batch_size} and {batch_size}"
-    
-    assert target_modules[0].weight.data.dtype in [torch.float32, torch.bfloat16], \
-        "The model must be loaded in float32 or bfloat16 " + \
-        "for gradient-based LPCD refinement due to numerical stability."
-    
+    assert lpcd_config.gd_batch_size % batch_size == 0, (
+        f"gd_batch_size should be divisible by batch_size, "
+        + f"but got {lpcd_config.gd_batch_size} and {batch_size}"
+    )
+
+    assert target_modules[0].weight.data.dtype in [torch.float32, torch.bfloat16], (
+        "The model must be loaded in float32 or bfloat16 "
+        + "for gradient-based LPCD refinement due to numerical stability."
+    )
+
     # backup and configure grad state
     grad_state_q = [p.requires_grad for p in metric_q.parameters()]
     grad_state_f = [p.requires_grad for p in metric_f.parameters()]
 
-    # set target modules as trainablethe     
+    # set target modules as trainablethe
     metric_q.requires_grad_(False)
     metric_f.requires_grad_(False)
     for module in target_modules:
@@ -60,13 +63,11 @@ def gradient_solver(
 
     optimizer = optim.Adam(metric_q.parameters(), lr=lpcd_config.gd_base_lr)
     scheduler = get_cosine_schedule_with_warmup(
-        optimizer,
-        num_warmup_steps=int(total_iters * 0.1),
-        num_training_steps=total_iters
+        optimizer, num_warmup_steps=int(total_iters * 0.1), num_training_steps=total_iters
     )
 
     # Perform gradient descent to relax the weights
-    for epoch in range(1, lpcd_config.gd_steps+1):
+    for epoch in range(1, lpcd_config.gd_steps + 1):
 
         chunked_inps = zip(
             inps_q.split(batch_size),
@@ -75,7 +76,7 @@ def gradient_solver(
 
         # TODO: shuffle the inputs for better convergence
         for iter, (inp_q, inp_f) in enumerate(chunked_inps):
-        
+
             with torch.no_grad():
                 out_f = metric_f(inp_f.to(device), **kwargs)
 
@@ -95,4 +96,4 @@ def gradient_solver(
     for p, requires_grad in zip(metric_q.parameters(), grad_state_q):
         p.requires_grad_(requires_grad)
     for p, requires_grad in zip(metric_f.parameters(), grad_state_f):
-        p.requires_grad_(requires_grad)
+        p.requires_grad_(requires_grad)

onecomp/lpcd/_lpcd_config.py

@@ -54,4 +54,3 @@ class LPCDConfig:
     gd_batch_size: int = 16
     gd_base_lr: float = 1e-4
     device: str = "cuda:0"
-    

onecomp/lpcd/_lpcd_runner.py

@@ -45,7 +45,7 @@ def run_quantize_with_lpcd(
     lpcd_config: LPCDConfig,
     calibration_config: CalibrationConfig,
 ):
-    """ Run quantization with LPCD.
+    """Run quantization with LPCD.
 
     Args:
         model_config (ModelConfig): Model configuration
@@ -55,8 +55,9 @@ def run_quantize_with_lpcd(
         calibration_config (CalibrationConfig): Calibration configuration
     """
 
-    assert not (qep_config is not None and qep_config.general), \
-        "qep_config.general must be False when qep is enabled."
+    assert not (
+        qep_config is not None and qep_config.general
+    ), "qep_config.general must be False when qep is enabled."
 
     # TODO: Parameterize when necessary
     batch_size = 16
@@ -69,7 +70,7 @@ def run_quantize_with_lpcd(
 
     model_inputs = prepare_calibration_dataset(
         tokenizer=tokenizer,
-        device=torch.device('cpu'),
+        device=torch.device("cpu"),
         calibration_config=calibration_config,
         model=model,
         logger=logger,
@@ -99,7 +100,7 @@ def run_quantize_with_lpcd(
         block_f = copy.deepcopy(block_q)
 
         groups_q = make_grouped_module(block_q, inps_q, kwargs, device)
-        
+
         # Build name→module map for block_f, then align groups_f to
         # groups_q by module name.  Using make_grouped_module on
         # block_f independently can produce a different group ordering
@@ -112,12 +113,17 @@ def run_quantize_with_lpcd(
             name: mod for name, mod in block_q.named_modules() if isinstance(mod, nn.Linear)
         }
         groups_f = [
-            [name_to_module_f[next(n for n, m2 in name_to_module_q.items() if m2 is m)] for m in gq]
+            [
+                name_to_module_f[next(n for n, m2 in name_to_module_q.items() if m2 is m)]
+                for m in gq
+            ]
             for gq in groups_q
         ]
 
         lpcd_metrics = make_lpcd_metrics(lpcd_config, block_q, block_f)
-        lpcd_modules_q = [module for metric, _ in lpcd_metrics.metrics for _, module in metric.named_targets()]
+        lpcd_modules_q = [
+            module for metric, _ in lpcd_metrics.metrics for _, module in metric.named_targets()
+        ]
 
         # 3. For each group of layers, perform the following sequentially
         for group_q, group_f in zip(groups_q, groups_f):
@@ -149,13 +155,10 @@ def run_quantize_with_lpcd(
                 # Skip layers not registered for quantization
                 if module not in quantizer.module_to_name:
                     skipped_name = module_q_to_name.get(module, "<unknown>")
-                    logger.info(
-                        "Skipping layer (not in quantization targets): %s", skipped_name
-                    )
+                    logger.info("Skipping layer (not in quantization targets): %s", skipped_name)
                     continue
                 name = quantizer.module_to_name[module]
 
-
                 # Fall back to standard quantization if the module is not LPCD targets
                 if not module in lpcd_modules_q:
 
@@ -190,10 +193,7 @@ def run_quantize_with_lpcd(
                     # Update the weights of the target layer
                     dtype = module.weight.data.dtype
                     module.weight.data = (
-                        quantizer.results[name]
-                        .compute_dequantized_weight()
-                        .to(device)
-                        .to(dtype)
+                        quantizer.results[name].compute_dequantized_weight().to(device).to(dtype)
                     )
 
                 lpcd_metrics.mark_as_ready(module)
@@ -218,20 +218,11 @@ def run_quantize_with_lpcd(
         inps_f = forward_input(inps_f, block_f, kwargs, batch_size, device)
 
         # DEBUG:Compute MSE between quantized and full-precision outputs
-        mse = compute_mse(
-            block_q,
-            block_f,
-            inps_q,
-            inps_f,
-            batch_size,
-            device,
-            kwargs
-        )
+        mse = compute_mse(block_q, block_f, inps_q, inps_f, batch_size, device, kwargs)
         logger.info(f"[INFO] Layer {block_idx + 1} MSE: {mse:.6e}")
 
         # free memory
         block_q.cpu()
         torch.cuda.empty_cache()
 
     quantizer.execute_post_processing()
-