From 3a874527e0cebc7e121da4ecf33b9e572d5d9714 Mon Sep 17 00:00:00 2001 From: katari Date: Wed, 15 Apr 2026 10:40:00 +0900 Subject: [PATCH 01/82] [demo] Experimental LoRA model save/load: save base model and LoRA adapter separately, run inference with vLLM --- .../example_lora_sft_save_load_hf_format.py | 151 +++++++++++ .../example_gptq_vllm_inference.py | 58 ++++- onecomp/quantized_model_loader.py | 142 ++++++++++- onecomp/runner.py | 172 ++++++++++++- .../example_gptq_lora_vllm_inference.py | 234 ++++++++++++++++++ ...a_sft_save_load_hf_format_with_low_vram.py | 159 ++++++++++++ 6 files changed, 903 insertions(+), 13 deletions(-) create mode 100644 example/post_process/example_lora_sft_save_load_hf_format.py create mode 100644 work/20260414_katari_LoRA_Save_Load/example_gptq_lora_vllm_inference.py create mode 100644 work/20260414_katari_LoRA_Save_Load/example_lora_sft_save_load_hf_format_with_low_vram.py diff --git a/example/post_process/example_lora_sft_save_load_hf_format.py b/example/post_process/example_lora_sft_save_load_hf_format.py new file mode 100644 index 0000000..22f158d --- /dev/null +++ b/example/post_process/example_lora_sft_save_load_hf_format.py @@ -0,0 +1,151 @@ +""" +Example: GPTQ quantization + LoRA SFT + save/load + +End-to-end demonstration of the LoRA SFT post-process workflow: + 1. Quantize TinyLlama with GPTQ 4-bit (groupsize=128) + 2. Apply LoRA SFT post-process (WikiText-2) + 3. Evaluate PPL (original vs quantized+LoRA) + 4. Save via save_quantized_model() - writes HF-compatible safetensors + plus a PEFT-format LoRA adapter sidecar + (``adapter_model.safetensors`` + ``adapter_config.json``) + 5. Load via load_quantized_model() - the sidecar is auto-detected and + matching GPTQLinear layers are re-wrapped with LoRAGPTQLinear + 6. Generate text with the loaded model to verify it works + +Copyright 2025-2026 Fujitsu Ltd. + +Author: Keiji Kimura + +Usage: + python example/post_process/example_lora_sft.py +""" + +import torch + +from onecomp import ( + CalibrationConfig, + GPTQ, + ModelConfig, + PostProcessLoraSFT, + Runner, + load_quantized_model, + setup_logger, +) + +setup_logger() + +MODEL_ID = "TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T" +SAVE_DIR = "./tinyllama_gptq4_lora" +PROMPT = "Fujitsu is" + + +def generate_text(model, tokenizer, prompt, device, max_new_tokens=64): + """Generate text from a prompt using the model.""" + inputs = tokenizer(prompt, return_tensors="pt").to(device) + with torch.no_grad(): + outputs = model.generate( + **inputs, + max_new_tokens=max_new_tokens, + do_sample=False, + temperature=1.0, + repetition_penalty=1.2, + ) + generated = tokenizer.decode(outputs[0], skip_special_tokens=True) + return generated + + +# ================================================================ +# Step 1: Quantize + LoRA SFT via Runner +# ================================================================ +print("=" * 70) +print("Step 1: Quantize TinyLlama (GPTQ 4-bit) + LoRA SFT (WikiText-2)") +print("=" * 70) + +model_config = ModelConfig(model_id=MODEL_ID, device="cuda:0") +gptq = GPTQ(wbits=4, groupsize=128) + +post_process = PostProcessLoraSFT( + dataset_name="wikitext", + dataset_config_name="wikitext-2-raw-v1", + train_split="train", + text_column="text", + max_train_samples=256, + max_length=512, + epochs=2, + batch_size=2, + gradient_accumulation_steps=4, + lr=1e-4, + lora_r=16, + lora_alpha=32, + logging_steps=5, +) + +runner = Runner( + model_config=model_config, + quantizer=gptq, + calibration_config=CalibrationConfig(max_length=128, num_calibration_samples=16, batch_size=8), + post_processes=[post_process], +) +runner.run() + +# ================================================================ +# Step 2: Evaluate PPL (original vs quantized+LoRA) +# ================================================================ +print("\n" + "=" * 70) +print("Step 2: Evaluate PPL") +print("=" * 70) + +original_ppl, _, quantized_ppl = runner.calculate_perplexity( + original_model=True, + quantized_model=True, +) + +print(f" Original model PPL: {original_ppl:.4f}") +print(f" Quantized + LoRA SFT model PPL: {quantized_ppl:.4f}") + +# ================================================================ +# Step 3: Save the LoRA-applied model (HF safetensors + adapter sidecar) +# ================================================================ +print("\n" + "=" * 70) +print(f"Step 3: Saving LoRA-applied model to {SAVE_DIR}") +print("=" * 70) + +runner.save_quantized_model(SAVE_DIR) +print(f"Model saved to: {SAVE_DIR}") +print( + " - model.safetensors / config.json : base GPTQ model (HF-compatible)\n" + " - adapter_model.safetensors : PEFT-format LoRA adapter\n" + " - adapter_config.json : PEFT-format adapter config" +) + +del runner +torch.cuda.empty_cache() + +# ================================================================ +# Step 4: Load the saved model +# ================================================================ +print("\n" + "=" * 70) +print(f"Step 4: Loading model from {SAVE_DIR}") +print("=" * 70) + +loaded_model, loaded_tokenizer = load_quantized_model(SAVE_DIR) +print(f"Loaded model type : {type(loaded_model).__name__}") +print(f"Loaded model device: {next(loaded_model.parameters()).device}") + +# ================================================================ +# Step 5: Generate text with the loaded model +# ================================================================ +print("\n" + "=" * 70) +print("Step 5: Generate text with loaded model") +print("=" * 70) + +loaded_text = generate_text( + loaded_model, + loaded_tokenizer, + PROMPT, + device=next(loaded_model.parameters()).device, +) + +print(f"\nPrompt : {PROMPT}") +print(f"Generated: {loaded_text}") +print("=" * 70) diff --git a/example/vllm_inference/example_gptq_vllm_inference.py b/example/vllm_inference/example_gptq_vllm_inference.py index f353eeb..170a7ad 100644 --- a/example/vllm_inference/example_gptq_vllm_inference.py +++ b/example/vllm_inference/example_gptq_vllm_inference.py @@ -19,9 +19,52 @@ import gc +import json +import os import torch -from onecomp import Runner, ModelConfig, CalibrationConfig, GPTQ, setup_logger -from vllm import LLM, SamplingParams +from onecomp import CalibrationConfig, GPTQ, ModelConfig, Runner, setup_logger +from transformers.tokenization_utils_base import PreTrainedTokenizerBase as _PTBase + +try: + from vllm import LLM, SamplingParams +except ImportError as e: + raise SystemExit( + "This example requires vllm to be installed. " + "Install with: uv sync --extra vllm" + ) from e + +if not hasattr(_PTBase, "all_special_tokens_extended"): + _PTBase.all_special_tokens_extended = property( + lambda self: list(self.all_special_tokens) + ) + +def _ensure_fast_tokenizer_class(save_dir: str) -> None: + """Rewrite tokenizer_config.json so vLLM loads the fast tokenizer.""" + tok_json = os.path.join(save_dir, "tokenizer.json") + tok_cfg = os.path.join(save_dir, "tokenizer_config.json") + if not (os.path.isfile(tok_json) and os.path.isfile(tok_cfg)): + return + + with open(tok_cfg, "r", encoding="utf-8") as f: + cfg = json.load(f) + + current = cfg.get("tokenizer_class") + if current and current.endswith("Fast"): + return + + slow_to_fast = { + "LlamaTokenizer": "LlamaTokenizerFast", + "CodeLlamaTokenizer": "CodeLlamaTokenizerFast", + } + replacement = slow_to_fast.get(current, "LlamaTokenizerFast") + cfg["tokenizer_class"] = replacement + + with open(tok_cfg, "w", encoding="utf-8") as f: + json.dump(cfg, f, indent=2, ensure_ascii=False) + print( + f"Patched {tok_cfg}: tokenizer_class {current!r} -> {replacement!r} " + "(forces vLLM to use the fast tokenizer)" + ) def main(): @@ -32,11 +75,13 @@ def main(): model_config = ModelConfig( model_id="TinyLlama/TinyLlama-1.1B-Chat-v1.0", + device="cuda:0", ) quantizer = GPTQ(wbits=4, groupsize=128) calibration_config = CalibrationConfig( - num_calibration_samples=128, - max_length=512, + max_length=128, + num_calibration_samples=16, + batch_size=8 ) runner = Runner( model_config=model_config, @@ -48,6 +93,7 @@ def main(): # Step 2: Save the quantized model runner.save_quantized_model(save_dir) + print(f"\nSaved GPTQ model to: {save_dir}") # Free GPU memory used by quantization before loading vLLM del runner @@ -55,11 +101,15 @@ def main(): torch.cuda.empty_cache() # Step 3: Load the quantized model with vLLM + _ensure_fast_tokenizer_class(save_dir) llm = LLM( model=save_dir, max_model_len=512, dtype="float16", enforce_eager=True, + gpu_memory_utilization=0.55, + max_num_batched_tokens=512, + enable_prefix_caching=False, ) # Step 4: Generate text diff --git a/onecomp/quantized_model_loader.py b/onecomp/quantized_model_loader.py index a9ac04f..2c01b18 100644 --- a/onecomp/quantized_model_loader.py +++ b/onecomp/quantized_model_loader.py @@ -51,8 +51,15 @@ def load_quantized_model( config.json and quantized layers are reconstructed directly from the safetensors state_dict. No quantization_results.pt is needed. - For models saved with post-processing modifications (e.g. LoRA adapters), - use :meth:`load_quantized_model_pt` instead. + If the directory additionally contains a PEFT-format LoRA adapter + sidecar (``adapter_model.safetensors`` + ``adapter_config.json``), the + matching ``GPTQLinear`` layers are automatically re-wrapped with + ``LoRAGPTQLinear`` populated from the sidecar. This lets + ``runner.save_quantized_model`` → ``load_quantized_model`` round-trip + models produced by a LoRA post-process such as ``PostProcessLoraSFT``. + + For legacy models saved via ``torch.save`` (``.pt`` format), use + :meth:`load_quantized_model_pt` instead. Args: save_directory: Path to the saved model directory. @@ -111,6 +118,11 @@ def load_quantized_model( fp32_had, ) + # Re-apply LoRA adapter from PEFT-format sidecar if present. + # This must run while the model is still on CPU, before dispatch_model, + # so LoRA wrappers are included in the device map traversal below. + cls._maybe_wrap_lora_adapters(model, save_directory) + # Device placement if device_map: try: @@ -372,3 +384,129 @@ def _replace_quantized_layers(model, state_dict: dict, quant_config: dict): ) QuantizedModelLoader._set_module_by_name(model, name, quantized_module) + + LORA_ADAPTER_SUBDIR = "lora_adapter" + + @staticmethod + def _maybe_wrap_lora_adapters(model, save_directory: str) -> int: + """Re-wrap GPTQLinear layers with LoRAGPTQLinear from a PEFT-format sidecar. + + Looks for ``adapter_model.safetensors`` + ``adapter_config.json`` under + ``save_directory/lora_adapter/``. If both are present, each referenced + GPTQLinear layer is replaced in-place with a ``LoRAGPTQLinear`` wrapper + populated with the saved LoRA weights. + + For backward compatibility, also checks the legacy top-level layout + (``save_directory/adapter_model.safetensors``) used by an earlier + version of :meth:`Runner.save_quantized_model`. + + Returns: + int: Number of layers wrapped (0 if no adapter sidecar was found). + """ + adapter_dir = os.path.join( + save_directory, QuantizedModelLoader.LORA_ADAPTER_SUBDIR + ) + adapter_weights_path = os.path.join(adapter_dir, "adapter_model.safetensors") + adapter_config_path = os.path.join(adapter_dir, "adapter_config.json") + if not ( + os.path.isfile(adapter_weights_path) + and os.path.isfile(adapter_config_path) + ): + # Fallback to legacy top-level layout. + legacy_weights = os.path.join(save_directory, "adapter_model.safetensors") + legacy_config = os.path.join(save_directory, "adapter_config.json") + if os.path.isfile(legacy_weights) and os.path.isfile(legacy_config): + adapter_weights_path = legacy_weights + adapter_config_path = legacy_config + else: + return 0 + + with open(adapter_config_path, "r", encoding="utf-8") as f: + adapter_config = json.load(f) + + lora_r = int(adapter_config["r"]) + lora_alpha = int(adapter_config["lora_alpha"]) + lora_dropout = float(adapter_config.get("lora_dropout", 0.0)) + + adapter_sd = load_file(adapter_weights_path) + + peft_prefix = "base_model.model." + per_layer: Dict[str, Dict[str, torch.Tensor]] = {} + + for key, tensor in adapter_sd.items(): + if not key.startswith(peft_prefix): + logger.warning("Skipping unexpected adapter key %s", key) + continue + body = key[len(peft_prefix):] + # Support both the plain form ".lora_A.weight" and PEFT's + # adapter-name form ".lora_A.default.weight". + if body.endswith(".lora_A.weight"): + layer_path = body[: -len(".lora_A.weight")] + per_layer.setdefault(layer_path, {})["A"] = tensor + elif body.endswith(".lora_B.weight"): + layer_path = body[: -len(".lora_B.weight")] + per_layer.setdefault(layer_path, {})["B"] = tensor + elif body.endswith(".lora_A.default.weight"): + layer_path = body[: -len(".lora_A.default.weight")] + per_layer.setdefault(layer_path, {})["A"] = tensor + elif body.endswith(".lora_B.default.weight"): + layer_path = body[: -len(".lora_B.default.weight")] + per_layer.setdefault(layer_path, {})["B"] = tensor + else: + logger.warning("Skipping unrecognized adapter key %s", key) + + if not per_layer: + return 0 + + # Inline import to avoid pulling post_process into module-import time + # and to sidestep any circular-import risk. + from .post_process.post_process_lora_sft import LoRAGPTQLinear + + name_to_module = dict(model.named_modules()) + wrapped = 0 + for layer_path, ab in per_layer.items(): + if "A" not in ab or "B" not in ab: + logger.warning( + "Adapter layer %s missing lora_A or lora_B; skipping", + layer_path, + ) + continue + if layer_path not in name_to_module: + logger.warning( + "Adapter references layer %s not found in model; skipping", + layer_path, + ) + continue + base_layer = name_to_module[layer_path] + if not isinstance(base_layer, GPTQLinear): + logger.warning( + "Adapter layer %s is %s, expected GPTQLinear; skipping", + layer_path, + type(base_layer).__name__, + ) + continue + + wrapper = LoRAGPTQLinear( + base_layer=base_layer, + lora_r=lora_r, + lora_alpha=lora_alpha, + lora_dropout=lora_dropout, + ) + with torch.no_grad(): + wrapper.lora_A.weight.copy_( + ab["A"].to(wrapper.lora_A.weight.dtype) + ) + wrapper.lora_B.weight.copy_( + ab["B"].to(wrapper.lora_B.weight.dtype) + ) + # Match the base layer's device so the wrapper and base share placement. + base_device = base_layer.qweight.device + wrapper.to(base_device) + QuantizedModelLoader._set_module_by_name(model, layer_path, wrapper) + wrapped += 1 + + logger.info( + "Re-wrapped %d GPTQLinear layers with LoRAGPTQLinear from adapter sidecar", + wrapped, + ) + return wrapped \ No newline at end of file diff --git a/onecomp/runner.py b/onecomp/runner.py index b126d18..88ad8d1 100644 --- a/onecomp/runner.py +++ b/onecomp/runner.py @@ -1573,9 +1573,126 @@ def create_quantized_model(self, pack_weights: bool = True, quantizer=None, use_ # Unified Save/Load Methods (Using quantizer.results) # ======================================== + LORA_ADAPTER_SUBDIR = "lora_adapter" + + def _save_lora_adapter_sidecar(self, save_directory: str) -> bool: + """Write a PEFT-compatible LoRA adapter sidecar if ``self.quantized_model`` + contains ``LoRAGPTQLinear`` modules (typically produced by + ``PostProcessLoraSFT``). + + The sidecar is placed in a ``lora_adapter/`` subdirectory rather than + directly in ``save_directory``. Reason: vLLM's base-model safetensors + loader globs ``*.safetensors`` at the top level of the model directory + and would otherwise try to load ``adapter_model.safetensors`` as + base-model weights, crashing with ``"no module or parameter named + 'base_model' in LlamaForCausalLM"``. Keeping the adapter under a + subdirectory avoids that collision while still keeping the whole model + self-contained under one directory tree. + + The subdirectory contains: + - ``adapter_model.safetensors`` + - ``adapter_config.json`` + + The format matches what vLLM's native PEFT LoRA loader expects, so:: + + LLM(model=save_dir, enable_lora=True) + LoRARequest(..., lora_path=os.path.join(save_dir, "lora_adapter")) + + will load and apply the adapter without any OneComp-specific changes + to the vLLM plugin. + + Returns: + bool: True iff an adapter was written. False if there is no + in-memory LoRA state to save (e.g. no post-process ran). + """ + if self.quantized_model is None: + return False + + # Inline imports keep runner.py import-time cheap and avoid any + # circular-import risk with the post_process package. + from .post_process.post_process_lora_sft import LoRAGPTQLinear + from safetensors.torch import save_file as _st_save_file + + lora_modules = [ + (name, mod) + for name, mod in self.quantized_model.named_modules() + if isinstance(mod, LoRAGPTQLinear) + ] + if not lora_modules: + return False + + # PEFT convention: keys are prefixed with "base_model.model." and the + # module path matches what we will see on the loaded HF model. + state_dict = {} + for name, mod in lora_modules: + state_dict[f"base_model.model.{name}.lora_A.weight"] = ( + mod.lora_A.weight.detach().to("cpu", torch.float16).contiguous() + ) + state_dict[f"base_model.model.{name}.lora_B.weight"] = ( + mod.lora_B.weight.detach().to("cpu", torch.float16).contiguous() + ) + + first = lora_modules[0][1] + lora_r = int(first.lora_r) + # scaling = alpha / r is stored as float; round-trip back to int alpha. + lora_alpha = int(round(float(first.scaling) * float(first.lora_r))) + lora_dropout = ( + float(first.dropout.p) if isinstance(first.dropout, torch.nn.Dropout) else 0.0 + ) + target_modules = sorted({name.rsplit(".", 1)[-1] for name, _ in lora_modules}) + + adapter_config = { + "peft_type": "LORA", + "auto_mapping": None, + "base_model_name_or_path": str(Path(save_directory).resolve()), + "task_type": "CAUSAL_LM", + "r": lora_r, + "lora_alpha": lora_alpha, + "lora_dropout": lora_dropout, + "target_modules": target_modules, + "bias": "none", + "fan_in_fan_out": False, + "inference_mode": True, + "modules_to_save": None, + "init_lora_weights": True, + "layers_to_transform": None, + "layers_pattern": None, + "revision": None, + } + + adapter_dir = Path(save_directory) / self.LORA_ADAPTER_SUBDIR + adapter_dir.mkdir(parents=True, exist_ok=True) + _st_save_file( + state_dict, + str(adapter_dir / "adapter_model.safetensors"), + metadata={"format": "pt"}, + ) + with open(adapter_dir / "adapter_config.json", "w", encoding="utf-8") as f: + json.dump(adapter_config, f, indent=2, ensure_ascii=True) + + self.logger.info( + "Saved LoRA adapter sidecar (%d layers) to %s", + len(lora_modules), + adapter_dir, + ) + return True + + def save_quantized_model(self, save_directory: str, pack_weights: bool = True): """Save the quantized model to the specified directory + The base quantized model is always rebuilt from ``quantizer.results`` + via :meth:`create_quantized_model` and saved in HuggingFace-compatible + safetensors format. + + If a LoRA post-process (e.g. ``PostProcessLoraSFT``) has populated + ``self.quantized_model`` with ``LoRAGPTQLinear`` wrappers, this method + additionally writes a PEFT-compatible LoRA adapter sidecar + (``adapter_model.safetensors`` + ``adapter_config.json``) into the same + directory. The resulting directory can then be loaded back with + :func:`onecomp.load_quantized_model` (which auto-detects the sidecar + and re-wraps the layers) or served by vLLM via ``enable_lora=True``. + Args: save_directory (str): The path to save the quantized model. @@ -1587,20 +1704,34 @@ def save_quantized_model(self, save_directory: str, pack_weights: bool = True): Single quantizer mode: >>> runner.save_quantized_model("./quantized_model") + + GPTQ + LoRA SFT: + + >>> runner = Runner( + ... model_config=model_config, + ... quantizer=GPTQ(wbits=4, groupsize=128), + ... post_processes=[PostProcessLoraSFT(data_files="train.jsonl")], + ... ) + >>> runner.run() + >>> runner.save_quantized_model("./quantized_model_lora") """ logger = self.logger logger.info("Saving quantized model to %s", save_directory) + # Always rebuild the base quantized model from quantizer.results. + # Post-process wrappers such as LoRAGPTQLinear are saved separately as + # an adapter sidecar and must not leak into the base HF safetensors. if self.quantized_model is not None: - logger.info("Using existing quantized model (post-process results preserved)") - model = self.quantized_model - tokenizer = self.model_config.load_tokenizer() - else: - # Disable GemLite when saving to avoid extra params in safetensors - model, tokenizer = self.create_quantized_model( - pack_weights=pack_weights, use_gemlite=False + logger.info( + "Rebuilding base quantized model from quantization results; " + "post-process state will be saved separately if supported" ) + # Disable GemLite when saving to avoid extra params in safetensors. + model, tokenizer = self.create_quantized_model( + pack_weights=pack_weights, use_gemlite=False + ) + # Save model and tokenizer save_path = Path(save_directory) save_path.mkdir(parents=True, exist_ok=True) @@ -1608,6 +1739,33 @@ def save_quantized_model(self, save_directory: str, pack_weights: bool = True): model.save_pretrained(save_directory) tokenizer.save_pretrained(save_directory) + # LoRA sidecar (only if self.quantized_model contains LoRAGPTQLinear). + wrote_adapter = self._save_lora_adapter_sidecar(save_directory) + if not wrote_adapter: + # Remove any stale sidecar from a previous run so the directory is + # self-consistent and load_quantized_model does not pick up an + # adapter that no longer matches the saved base model. + stale_adapter_dir = save_path / self.LORA_ADAPTER_SUBDIR + if stale_adapter_dir.is_dir(): + for stale in ( + "adapter_model.safetensors", + "adapter_config.json", + ): + stale_path = stale_adapter_dir / stale + if stale_path.exists(): + stale_path.unlink() + # Remove the (now-empty) subdirectory if nothing else lives there. + try: + stale_adapter_dir.rmdir() + except OSError: + pass + # Also remove any top-level adapter files left by older versions of + # this helper (previous layout put the sidecar directly in save_dir). + for legacy in ("adapter_model.safetensors", "adapter_config.json"): + legacy_path = save_path / legacy + if legacy_path.exists(): + legacy_path.unlink() + logger.info(f"Quantized model saved to {save_directory}") return save_directory diff --git a/work/20260414_katari_LoRA_Save_Load/example_gptq_lora_vllm_inference.py b/work/20260414_katari_LoRA_Save_Load/example_gptq_lora_vllm_inference.py new file mode 100644 index 0000000..4861001 --- /dev/null +++ b/work/20260414_katari_LoRA_Save_Load/example_gptq_lora_vllm_inference.py @@ -0,0 +1,234 @@ +""" + +Example: Quantize with GPTQ + LoRA SFT, save, and serve via vLLM with LoRA + +Performs the following steps: + 1. Quantize TinyLlama with GPTQ (4-bit, groupsize=128) + 2. Apply ``PostProcessLoraSFT`` (WikiText-2) to fine-tune LoRA adapters + on top of the frozen GPTQ base + 3. Save via ``runner.save_quantized_model(save_dir)``. This writes: + - ``model.safetensors`` / ``config.json`` (base GPTQ, HF format) + - ``lora_adapter/adapter_model.safetensors`` (PEFT-format LoRA) + - ``lora_adapter/adapter_config.json`` (PEFT-format LoRA config) + The adapter lives in the ``lora_adapter/`` subdirectory because vLLM + globs ``*.safetensors`` at the top level of the model directory to load + base weights, and would otherwise try to load the adapter file as part + of the base model. + 4. Load the base model with vLLM and apply the LoRA adapter via + ``LoRARequest(..., lora_path=os.path.join(save_dir, "lora_adapter"))``. + 5. Generate text with and without the LoRA adapter and print both + outputs side by side. + +Requirements: + - pip install vllm (>= 0.5 recommended; older releases may lack + LoRA-over-GPTQ support) + - The ``mixed_gptq`` quantization method is registered by OneComp's vLLM + plugin (see ``vllm_plugins/gptq``), which is installed automatically + with this package. + +vLLM LoRA constraints +--------------------- + - ``max_lora_rank`` must be >= the adapter's ``r`` (from + ``adapter_config.json``). + - On Llama-style models, q/k/v (and gate/up) projections are fused by + vLLM into ``qkv_proj`` / ``gate_up_proj``. vLLM's LoRA loader + concatenates the per-projection adapter weights automatically, provided + all three projections share the same ``r`` and ``lora_alpha`` — + ``PostProcessLoraSFT`` always uses a single global rank/alpha, so this + invariant holds. + - ``lm_head`` is excluded from LoRA targets by ``PostProcessLoraSFT``, so + ``lora_extra_vocab_size=0`` is safe. + +Copyright 2025-2026 Fujitsu Ltd. + +Author: Keiji Kimura + +""" + +import gc +import json +import os + +import torch + +# --------------------------------------------------------------------------- +# transformers >= 5.x / vLLM compatibility shim +# --------------------------------------------------------------------------- +# OneComp requires transformers>=5.3, which removed the +# `all_special_tokens_extended` property from the tokenizer base class. +# vLLM (<= 0.11.x) still accesses it in ``get_cached_tokenizer`` and crashes +# with ``AttributeError: LlamaTokenizer has no attribute +# all_special_tokens_extended``. Re-add the property as an alias for +# ``all_special_tokens`` so vLLM's cache build succeeds. Must run BEFORE vLLM +# imports so subclasses inherit the attribute. +from transformers.tokenization_utils_base import PreTrainedTokenizerBase as _PTBase + +if not hasattr(_PTBase, "all_special_tokens_extended"): + _PTBase.all_special_tokens_extended = property( + lambda self: list(self.all_special_tokens) + ) + +from onecomp import ( + CalibrationConfig, + GPTQ, + ModelConfig, + PostProcessLoraSFT, + Runner, + setup_logger, +) + +try: + from vllm import LLM, SamplingParams + from vllm.lora.request import LoRARequest + from vllm.model_executor.layers.quantization import register_quantization_config # noqa: F401 +except ImportError as e: + raise SystemExit( + "This example requires vllm>=0.6.3 " + "(needs both `vllm.lora.request.LoRARequest` and " + "`register_quantization_config` used by OneComp's mixed_gptq plugin). " + "Install with: pip install -U 'vllm>=0.6.3'" + ) from e + + +def _ensure_fast_tokenizer_class(save_dir: str) -> None: + """Rewrite tokenizer_config.json so vLLM loads the fast tokenizer. + + Some upstream Llama-family checkpoints (TinyLlama included) ship with + ``"tokenizer_class": "LlamaTokenizer"`` in ``tokenizer_config.json``. + Recent ``transformers`` releases have removed ``all_special_tokens_extended`` + from the slow ``LlamaTokenizer``, which causes vLLM's tokenizer cache to + crash with ``AttributeError: LlamaTokenizer has no attribute + all_special_tokens_extended``. When ``tokenizer.json`` is present the fast + variant is available, so we patch the class name to force vLLM down that + path. + """ + tok_json = os.path.join(save_dir, "tokenizer.json") + tok_cfg = os.path.join(save_dir, "tokenizer_config.json") + if not (os.path.isfile(tok_json) and os.path.isfile(tok_cfg)): + return + + with open(tok_cfg, "r", encoding="utf-8") as f: + cfg = json.load(f) + + current = cfg.get("tokenizer_class") + if current and current.endswith("Fast"): + return + + # Map slow → fast for Llama-family tokenizers. Extend this table if other + # architectures hit the same issue in the future. + slow_to_fast = { + "LlamaTokenizer": "LlamaTokenizerFast", + "CodeLlamaTokenizer": "CodeLlamaTokenizerFast", + } + replacement = slow_to_fast.get(current, "LlamaTokenizerFast") + cfg["tokenizer_class"] = replacement + + with open(tok_cfg, "w", encoding="utf-8") as f: + json.dump(cfg, f, indent=2, ensure_ascii=False) + print( + f"Patched {tok_cfg}: tokenizer_class {current!r} -> {replacement!r} " + "(forces vLLM to use the fast tokenizer)" + ) + + +def main(): + setup_logger() + + # ================================================================ + # Step 1: Quantize + LoRA SFT and save + # ================================================================ + save_dir = "./TinyLlama-1.1B-gptq-4bit-lora" + + model_config = ModelConfig( + model_id="TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T", + device="cuda:0", + ) + quantizer = GPTQ(wbits=4, groupsize=128) + post_process = PostProcessLoraSFT( + dataset_name="wikitext", + dataset_config_name="wikitext-2-raw-v1", + train_split="train", + text_column="text", + max_train_samples=128, + max_length=256, + epochs=2, + batch_size=1, + gradient_accumulation_steps=8, + lr=1e-4, + lora_r=8, + lora_alpha=16, + logging_steps=5, + ) + + runner = Runner( + model_config=model_config, + quantizer=quantizer, + calibration_config=CalibrationConfig(max_length=128, num_calibration_samples=16, batch_size=8), + post_processes=[post_process], + ) + runner.run() + runner.save_quantized_model(save_dir) + print(f"\nSaved GPTQ+LoRA model (base + adapter sidecar) to: {save_dir}") + + # Free GPU memory used by quantization / training before loading vLLM. + del runner + gc.collect() + torch.cuda.empty_cache() + + # ================================================================ + # Step 2: Load base GPTQ with vLLM, enable LoRA + # ================================================================ + # Work around the LlamaTokenizer / transformers compatibility issue + # (see _ensure_fast_tokenizer_class for details). + _ensure_fast_tokenizer_class(save_dir) + + # ``max_lora_rank`` must be >= the saved ``r`` (16 here). + llm = LLM( + model=save_dir, + enable_lora=True, + max_lora_rank=16, + max_loras=1, + max_model_len=512, + dtype="float16", + enforce_eager=True, + gpu_memory_utilization=0.55, # VRAM 8GB 制約ありのため削減 + max_num_batched_tokens=512, # 同上 + enable_prefix_caching=False, # 同上 + ) + + # The adapter sidecar lives in the lora_adapter/ subdirectory to avoid + # colliding with vLLM's base-model safetensors glob. + lora_request = LoRARequest( + lora_name="gptq_sft", + lora_int_id=1, + lora_path=os.path.join(save_dir, "lora_adapter"), + ) + + prompts = [ + "Explain what post-training quantization is in one sentence:", + "Fujitsu is", + ] + sampling_params = SamplingParams(max_tokens=64, temperature=0.0) + + # ----- (a) Base GPTQ only (LoRA disabled) ----- + base_outputs = llm.generate(prompts, sampling_params) + + # ----- (b) Base GPTQ + LoRA adapter ----- + lora_outputs = llm.generate( + prompts, + sampling_params, + lora_request=lora_request, + ) + + print("\n" + "=" * 70) + print("vLLM inference — GPTQ base vs GPTQ + LoRA SFT adapter") + print("=" * 70) + for base_out, lora_out in zip(base_outputs, lora_outputs): + print(f"\nPrompt: {base_out.prompt}") + print(f" base only : {base_out.outputs[0].text}") + print(f" base + LoRA: {lora_out.outputs[0].text}") + + +if __name__ == "__main__": + main() + \ No newline at end of file diff --git a/work/20260414_katari_LoRA_Save_Load/example_lora_sft_save_load_hf_format_with_low_vram.py b/work/20260414_katari_LoRA_Save_Load/example_lora_sft_save_load_hf_format_with_low_vram.py new file mode 100644 index 0000000..1e40c06 --- /dev/null +++ b/work/20260414_katari_LoRA_Save_Load/example_lora_sft_save_load_hf_format_with_low_vram.py @@ -0,0 +1,159 @@ +""" +Example: GPTQ quantization + LoRA SFT + save/load (low VRAM settings) + +End-to-end demonstration of the LoRA SFT post-process workflow with +settings chosen to reduce GPU memory usage on smaller GPUs: + 1. Quantize TinyLlama with GPTQ 4-bit (groupsize=128) + 2. Apply LoRA SFT post-process (WikiText-2) with reduced VRAM settings + 3. Evaluate PPL (original vs quantized+LoRA) + 4. Save via save_quantized_model() - writes HF-compatible safetensors + plus a PEFT-format LoRA adapter sidecar + (``adapter_model.safetensors`` + ``adapter_config.json``) + 5. Load via load_quantized_model() - the sidecar is auto-detected and + matching GPTQLinear layers are re-wrapped with LoRAGPTQLinear + 6. Generate text with the loaded model to verify it works + +Recommended low-VRAM changes compared with the standard example: + - batch_size=1 + - gradient_accumulation_steps=8 + - max_length=256 + - max_train_samples=128 + - lora_r=8 + +Copyright 2025-2026 Fujitsu Ltd. + +Author: Keiji Kimura + +Usage: + python example/post_process/example_lora_sft_low_vram.py +""" + +import torch + +from onecomp import ( + CalibrationConfig, + GPTQ, + ModelConfig, + PostProcessLoraSFT, + Runner, + load_quantized_model, + setup_logger, +) + +setup_logger() + +MODEL_ID = "TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T" +SAVE_DIR = "./tinyllama_gptq4_lora_low_vram" +PROMPT = "Fujitsu is" + + +def generate_text(model, tokenizer, prompt, device, max_new_tokens=64): + """Generate text from a prompt using the model.""" + inputs = tokenizer(prompt, return_tensors="pt").to(device) + with torch.no_grad(): + outputs = model.generate( + **inputs, + max_new_tokens=max_new_tokens, + do_sample=False, + temperature=1.0, + repetition_penalty=1.2, + ) + generated = tokenizer.decode(outputs[0], skip_special_tokens=True) + return generated + + +# ================================================================ +# Step 1: Quantize + LoRA SFT via Runner +# ================================================================ +print("=" * 70) +print("Step 1: Quantize TinyLlama (GPTQ 4-bit) + LoRA SFT (WikiText-2, low VRAM)") +print("=" * 70) + +model_config = ModelConfig(model_id=MODEL_ID, device="cuda:0") +gptq = GPTQ(wbits=4, groupsize=128) + +post_process = PostProcessLoraSFT( + dataset_name="wikitext", + dataset_config_name="wikitext-2-raw-v1", + train_split="train", + text_column="text", + max_train_samples=128, + max_length=256, + epochs=2, + batch_size=1, + gradient_accumulation_steps=8, + lr=1e-4, + lora_r=8, + lora_alpha=16, + logging_steps=5, +) + +runner = Runner( + model_config=model_config, + quantizer=gptq, + calibration_config=CalibrationConfig(max_length=128, num_calibration_samples=16, batch_size=8), + post_processes=[post_process], +) +runner.run() + +# ================================================================ +# Step 2: Evaluate PPL (original vs quantized+LoRA) +# ================================================================ +print("\n" + "=" * 70) +print("Step 2: Evaluate PPL") +print("=" * 70) + +original_ppl, _, quantized_ppl = runner.calculate_perplexity( + original_model=True, + quantized_model=True, +) + +print(f" Original model PPL: {original_ppl:.4f}") +print(f" Quantized + LoRA SFT model PPL: {quantized_ppl:.4f}") + +# ================================================================ +# Step 3: Save the LoRA-applied model (HF safetensors + adapter sidecar) +# ================================================================ +print("\n" + "=" * 70) +print(f"Step 3: Saving LoRA-applied model to {SAVE_DIR}") +print("=" * 70) + +runner.save_quantized_model(SAVE_DIR) +print(f"Model saved to: {SAVE_DIR}") +print( + " - model.safetensors / config.json : base GPTQ model (HF-compatible)\n" + " - adapter_model.safetensors : PEFT-format LoRA adapter\n" + " - adapter_config.json : PEFT-format adapter config" +) + +del runner +torch.cuda.empty_cache() + +# ================================================================ +# Step 4: Load the saved model +# ================================================================ +print("\n" + "=" * 70) +print(f"Step 4: Loading model from {SAVE_DIR}") +print("=" * 70) + +loaded_model, loaded_tokenizer = load_quantized_model(SAVE_DIR) +print(f"Loaded model type : {type(loaded_model).__name__}") +print(f"Loaded model device: {next(loaded_model.parameters()).device}") + +# ================================================================ +# Step 5: Generate text with the loaded model +# ================================================================ +print("\n" + "=" * 70) +print("Step 5: Generate text with loaded model") +print("=" * 70) + +loaded_text = generate_text( + loaded_model, + loaded_tokenizer, + PROMPT, + device=next(loaded_model.parameters()).device, +) + +print(f"\nPrompt : {PROMPT}") +print(f"Generated: {loaded_text}") +print("=" * 70) \ No newline at end of file From 20865fb7752e193013949ddfe9138c257c46e43d Mon Sep 17 00:00:00 2001 From: sikoji Date: Thu, 16 Apr 2026 14:55:25 +0900 Subject: [PATCH 02/82] [update] add save/load support to JointQ based on _gptq.py - Add `get_quant_config()` returning GPTQ-compatible config - Add `_build_quantization_bits()` and `finalize_quant_config_for_save()` - Add `create_inference_layer()` to build GPTQLinear from JointQResult --- onecomp/quantizer/jointq/_jointq.py | 114 +++++++++++++++++++++++++++- 1 file changed, 113 insertions(+), 1 deletion(-) diff --git a/onecomp/quantizer/jointq/_jointq.py b/onecomp/quantizer/jointq/_jointq.py index f8a066a..3f59530 100644 --- a/onecomp/quantizer/jointq/_jointq.py +++ b/onecomp/quantizer/jointq/_jointq.py @@ -6,14 +6,16 @@ """ +import re from dataclasses import dataclass -from typing import Optional +from typing import Any, Optional import torch from .core import compute_matrix_XX, quantize from onecomp.quantizer._quantizer import Quantizer, QuantizationResult +from onecomp.utils.quant_config import get_quant_param @dataclass @@ -366,3 +368,113 @@ def quantize_layer( assignment=assignment.cpu(), perm=perm.cpu() if perm is not None else None, ) + + + def get_quant_config(self) -> dict: + """Return GPTQ-compatible quantization config. + + JointQ uses the same scale/zero/assignment structure as GPTQ, + so we emit quant_method="gptq" to reuse GPTQLinear and vLLM GPTQ plugin. + """ + return { + "quant_method": "gptq", + "bits": self.bits, + "groupsize": self.group_size if self.group_size is not None else -1, + "group_size": self.group_size if self.group_size is not None else -1, + "actorder": self.actorder, + "desc_act": self.actorder, + "sym": self.symmetric, + "checkpoint_format": "gptq", + } + + @staticmethod + def _build_quantization_bits( + quantized_names: list[str], + quant_config: dict[str, Any], + num_layers: int, + ) -> list[dict[str, Any]]: + _LAYER_RE = re.compile(r"\.layers\.(\d+)\.(.*)") + default_bits = quant_config.get("bits", 4) + default_gs = get_quant_param(quant_config, "group_size", "groupsize", default=-1) + + layer_modules: dict[int, dict[str, Any]] = {} + for name in quantized_names: + m = _LAYER_RE.search(name) + if m is None: + continue + layer_idx = int(m.group(1)) + suffix = m.group(2) + + layer_modules.setdefault(layer_idx, {})[suffix] = { + "bits": default_bits, + "method": "gptq", + "params": {"group_size": default_gs}, + } + if not layer_modules: + return [] + + return [layer_modules.get(i, {}) for i in range(num_layers)] + + def finalize_quant_config_for_save( + self, + quant_config: dict[str, Any], + quantized_layer_names: list[str], + num_hidden_layers: Optional[int] = None, + ) -> dict[str, Any]: + if num_hidden_layers is None: + raise ValueError("num_hidden_layers is required") + quant_config["quantization_bits"] = JointQ._build_quantization_bits( + quantized_layer_names, quant_config, num_hidden_layers + ) + return quant_config + + def create_inference_layer(self, result, linear_module, **kwargs): + """Build GPTQLinear from JointQResult. + + Converts JointQ's 3D assignment (out_features, num_groups, group_size) + to 2D qweight (out_features, in_features), matching GPTQ format. + JointQ scale/zero_point shape is (out_features, num_groups); + GPTQLinear expects (num_groups, out_features), so we transpose. + """ + from onecomp.quantizer.gptq.gptq_layer import GPTQLinear + + qweight = result.assignment.reshape(result.assignment.shape[0], -1) + + # When `actorder` is enabled, `assignment` is stored in the permuted column order. + # Restore the original column order before passing to GPTQLinear. + # GPTQLinear constructs `g_idx` assuming `qweight` uses the original column ordering. + if result.perm is not None: + invperm = torch.argsort(result.perm) + qweight = qweight[:, invperm] + + pack_weights = kwargs.get("pack_weights", True) + + quantized_weight = qweight.to(torch.int32) + zero = result.zero_point.float() + + # Symmetric quantization uses signed integers [-2^(n-1), 2^(n-1)-1]; + # shift to unsigned [0, 2^n - 1] for GPTQLinear bit packing. + if result.symmetric: + offset = 2 ** (result.bits - 1) + quantized_weight = quantized_weight + offset + zero = zero + offset + + return GPTQLinear( + in_features=quantized_weight.shape[1], + out_features=quantized_weight.shape[0], + wbits=result.bits, + groupsize=result.group_size if result.group_size is not None else -1, + actorder=(result.perm is not None), + quantized_weight=quantized_weight, + scale=result.scale.T, + zero=zero.T, + perm=result.perm, + bias=( + linear_module.bias + if hasattr(linear_module, "bias") and linear_module.bias is not None + else None + ), + device=linear_module.weight.device, + pack_weights=pack_weights, + use_gemlite=kwargs.get("use_gemlite"), + ) From 23f4a6e80159d138f220b47ee0c0bcc783281312 Mon Sep 17 00:00:00 2001 From: cm-ysmz Date: Fri, 17 Apr 2026 14:25:03 +0900 Subject: [PATCH 03/82] [update] Add script for training a LoRA model with OneCompression knowledge, verified with vllm --- ..._lora_sft_knowledge_save_load_hf_format.py | 267 ++++++++++++++++++ 1 file changed, 267 insertions(+) create mode 100644 example/post_process/example_lora_sft_knowledge_save_load_hf_format.py diff --git a/example/post_process/example_lora_sft_knowledge_save_load_hf_format.py b/example/post_process/example_lora_sft_knowledge_save_load_hf_format.py new file mode 100644 index 0000000..9bf7982 --- /dev/null +++ b/example/post_process/example_lora_sft_knowledge_save_load_hf_format.py @@ -0,0 +1,267 @@ +""" +Example: Knowledge injection via LoRA SFT on a GPTQ-quantized model + +Demonstrates how to teach a quantized model new knowledge using +LoRA SFT post-processing. The model learns about "OneCompression" +(a topic it has never seen) and can answer questions about it +after training. + +Flow: + 1. Quantize TinyLlama with GPTQ 4-bit (groupsize=128) + 2. Build quantized model via create_quantized_model + 3. Generate text BEFORE LoRA SFT (model does not know OneCompression) + 4. Run LoRA SFT with OneCompression knowledge data + 4-1. Save via save_quantized_model() - writes HF-compatible safetensors + plus a PEFT-format LoRA adapter sidecar + (``adapter_model.safetensors`` + ``adapter_config.json``) + 4-2. Load via load_quantized_model() - the sidecar is auto-detected and + matching GPTQLinear layers are re-wrapped with LoRAGPTQLinear + 5. Generate text AFTER LoRA SFT (model can describe OneCompression) + 6. Compare results side by side + +Copyright 2025-2026 Fujitsu Ltd. + +Author: Keiji Kimura + +Usage: + python example/post_process/example_lora_sft_knowledge.py +""" + +from pathlib import Path + +import torch +import torch.nn as nn + +from onecomp import CalibrationConfig, GPTQ, ModelConfig, Runner, PostProcessLoraSFT, setup_logger, load_quantized_model + +setup_logger() + +MODEL_ID = "TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T" +KNOWLEDGE_DATA = str(Path(__file__).parent / "onecomp_knowledge.jsonl") +PROMPT = "Q: What is OneCompression?\nA:" +SAVE_DIR = "./tinyllama_gptq4_lora_knowledge" + +def dump_model_summary(model, max_modules=1000): + print("\n" + "=" * 70) + print("Model summary before save") + print("=" * 70) + + # ------------------------------------------------------------ + # 1) Linear系モジュールの一覧 + # ------------------------------------------------------------ + print("\n[Linear-like modules]") + linear_count = 0 + + for name, module in model.named_modules(): + cls_name = module.__class__.__name__ + + # 通常の Linear、またはクラス名に Linear を含むものを拾う + if isinstance(module, nn.Linear) or "Linear" in cls_name: + linear_count += 1 + + in_features = getattr(module, "in_features", None) + out_features = getattr(module, "out_features", None) + + print(f"- {name}") + print(f" class : {cls_name}") + print(f" in_features : {in_features}") + print(f" out_features : {out_features}") + + # 代表的な重みテンソルを表示 + for attr in ["weight", "qweight", "scales", "lora_A", "lora_B"]: + if hasattr(module, attr): + obj = getattr(module, attr) + if isinstance(obj, torch.Tensor): + print( + f" {attr:<12}: shape={tuple(obj.shape)}, " + f"dtype={obj.dtype}, device={obj.device}" + ) + elif hasattr(obj, "weight") and isinstance(obj.weight, torch.Tensor): + print( + f" {attr}.weight : shape={tuple(obj.weight.shape)}, " + f"dtype={obj.weight.dtype}, device={obj.weight.device}" + ) + + print() + + if linear_count >= max_modules: + print(f"... truncated after {max_modules} modules") + break + + print(f"Detected linear-like modules: {linear_count}") + + # ------------------------------------------------------------ + # 2) trainable parameter 一覧 + # ------------------------------------------------------------ + print("\n[Trainable parameters]") + trainable = 0 + total = 0 + + for name, param in model.named_parameters(): + n = param.numel() + total += n + if param.requires_grad: + trainable += n + print( + f"- {name}: shape={tuple(param.shape)}, " + f"dtype={param.dtype}, device={param.device}" + ) + + print(f"\nTrainable params: {trainable:,}") + print(f"Total params : {total:,}") + if total > 0: + print(f"Trainable ratio : {100.0 * trainable / total:.6f}%") + +def generate_text(model, tokenizer, prompt, device, max_new_tokens=128): + """Generate text from a prompt using the model.""" + inputs = tokenizer(prompt, return_tensors="pt").to(device) + with torch.no_grad(): + outputs = model.generate( + **inputs, + max_new_tokens=max_new_tokens, + do_sample=False, + temperature=1.0, + repetition_penalty=1.2, + ) + generated = tokenizer.decode(outputs[0], skip_special_tokens=True) + return generated + + +# ================================================================ +# Step 1: Quantize the model with GPTQ 4-bit +# ================================================================ +print("=" * 70) +print("Step 1: Quantizing TinyLlama with GPTQ 4-bit (groupsize=128)") +print("=" * 70) + +model_config = ModelConfig(model_id=MODEL_ID, device="cuda:0") +gptq = GPTQ(wbits=4, groupsize=128) + +runner = Runner( + model_config=model_config, + quantizer=gptq, + calibration_config=CalibrationConfig(max_length=512, num_calibration_samples=128), +) +runner.run() + +# ================================================================ +# Step 2: Build quantized model +# ================================================================ +print("\n" + "=" * 70) +print("Step 2: Building quantized model via create_quantized_model") +print("=" * 70) + +model, tokenizer = runner.create_quantized_model( + pack_weights=False, + use_gemlite=False, +) + +# ================================================================ +# Step 3: Generate BEFORE LoRA SFT +# ================================================================ +print("\n" + "=" * 70) +print("Step 3: Generating text BEFORE LoRA SFT") +print("=" * 70) + +model.to("cuda:0") +before_text = generate_text(model, tokenizer, PROMPT, "cuda:0") +model.to("cpu") +torch.cuda.empty_cache() + +print(f"\nPrompt: {PROMPT}") +print(f"Response:\n{before_text}") + +# ================================================================ +# Step 4: Run LoRA SFT with OneCompression knowledge +# ================================================================ +print("\n" + "=" * 70) +print("Step 4: Running LoRA SFT with OneCompression knowledge data") +print("=" * 70) + +post_process = PostProcessLoraSFT( + data_files=KNOWLEDGE_DATA, + max_length=256, + epochs=50, + batch_size=2, + gradient_accumulation_steps=1, + lr=3e-4, + lora_r=16, + lora_alpha=32, + logging_steps=5, +) +runner.post_processes = [post_process] +runner.run_post_processes() +model = runner.quantized_model + +# ================================================================ +# Step 4.1: Save the LoRA-applied quantized model (HF safetensors + adapter sidecar) +# ================================================================ +print("\n" + "=" * 70) +print("Step 4.1-a: Inspecting model before save") +print("=" * 70) +dump_model_summary(model) + +print("\n" + "=" * 70) +print(f"Step 4.1-b: Saving LoRA-applied model to {SAVE_DIR}") +print("=" * 70) +runner.save_quantized_model(SAVE_DIR) +print(f"Model saved to: {SAVE_DIR}") +print( + " - model.safetensors, config.json : base GPTQ model (HF-compatible)\n" + " - lora_adapter/adapter_model.safetensors : PEFT-format LoRA adapter\n" + " - lora_adapter/adapter_config.json : PEFT-format adapter config" +) + +# ================================================================ +# Step 4.2: Load the saved model +# ================================================================ +print("\n" + "=" * 70) +print(f"Step 4.2: Loading model from {SAVE_DIR}") +print("=" * 70) + +loaded_model, loaded_tokenizer = load_quantized_model(SAVE_DIR) +print(f"Loaded model type : {type(loaded_model).__name__}") +print(f"Loaded model device: {next(loaded_model.parameters()).device}") + +# ================================================================ +# Step 5: Generate AFTER LoRA SFT +# ================================================================ +print("\n" + "=" * 70) +print("Step 5: Generating text AFTER LoRA SFT") +print("=" * 70) + +loaded_model.to("cuda:0") +after_text = generate_text(loaded_model, loaded_tokenizer, PROMPT, "cuda:0") +loaded_model.to("cpu") +torch.cuda.empty_cache() + +print(f"\nPrompt: {PROMPT}") +print(f"Response:\n{after_text}") + +print("\n" + "=" * 70) +print("Step 5(ex): Generating text Before Save model") +print("=" * 70) + +# test before save model +model.to("cuda:0") +bs_text = generate_text(model, tokenizer, PROMPT, "cuda:0") +model.to("cpu") +torch.cuda.empty_cache() + +print(f"\nPrompt: {PROMPT}") +print(f"Response:\n{bs_text}") + +# ================================================================ +# Step 6: Compare results +# ================================================================ +print("\n" + "=" * 70) +print("Comparison: Before vs After LoRA SFT") +print("=" * 70) +print(f"\nPrompt: {PROMPT}") +print(f"\n--- BEFORE LoRA SFT ---") +print(before_text) +print(f"\n--- AFTER LoRA SFT ---") +print(after_text) +print("=" * 70) + + From d556d979a8c90c10824305c6f9b5f913425632eb Mon Sep 17 00:00:00 2001 From: cm-ysmz Date: Tue, 21 Apr 2026 16:33:45 +0900 Subject: [PATCH 04/82] [update] Add script for training a JointQ+LoRA model with OneCompression knowledge, and add smoke test --- ...ft_knowledge_save_load_hf_format_jointq.py | 268 ++++++++++++++++++ 1 file changed, 268 insertions(+) create mode 100644 example/post_process/example_lora_sft_knowledge_save_load_hf_format_jointq.py diff --git a/example/post_process/example_lora_sft_knowledge_save_load_hf_format_jointq.py b/example/post_process/example_lora_sft_knowledge_save_load_hf_format_jointq.py new file mode 100644 index 0000000..e66335f --- /dev/null +++ b/example/post_process/example_lora_sft_knowledge_save_load_hf_format_jointq.py @@ -0,0 +1,268 @@ +""" +Example: Knowledge injection via LoRA SFT on a JointQ-quantized model + +Demonstrates how to teach a quantized model new knowledge using +LoRA SFT post-processing. The model learns about "OneCompression" +(a topic it has never seen) and can answer questions about it +after training. + +Flow: + 1. Quantize TinyLlama with JointQ 4-bit (groupsize=128) + 2. Build quantized model via create_quantized_model + 3. Generate text BEFORE LoRA SFT (model does not know OneCompression) + 4. Run LoRA SFT with OneCompression knowledge data + 4-1. Save via save_quantized_model() - writes HF-compatible safetensors + plus a PEFT-format LoRA adapter sidecar + (``adapter_model.safetensors`` + ``adapter_config.json``) + 4-2. Load via load_quantized_model() - the sidecar is auto-detected and + matching GPTQLinear layers are re-wrapped with LoRAGPTQLinear + 5. Generate text AFTER LoRA SFT (model can describe OneCompression) + 6. Compare results side by side + +Copyright 2025-2026 Fujitsu Ltd. + +Author: Keiji Kimura + +Usage: + python example/post_process/example_lora_sft_knowledge.py +""" + +from pathlib import Path + +import torch +import torch.nn as nn + +from onecomp import CalibrationConfig, JointQ, ModelConfig, Runner, PostProcessLoraSFT, setup_logger, load_quantized_model + +setup_logger() + +MODEL_ID = "TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T" +KNOWLEDGE_DATA = str(Path(__file__).parent / "onecomp_knowledge.jsonl") +PROMPT = "Q: What is OneCompression?\nA:" +SAVE_DIR = "./tinyllama_jointq4_lora_knowledge" + +def dump_model_summary(model, max_modules=1000): + print("\n" + "=" * 70) + print("Model summary before save") + print("=" * 70) + + # ------------------------------------------------------------ + # 1) Linear系モジュールの一覧 + # ------------------------------------------------------------ + print("\n[Linear-like modules]") + linear_count = 0 + + for name, module in model.named_modules(): + cls_name = module.__class__.__name__ + + # 通常の Linear、またはクラス名に Linear を含むものを拾う + if isinstance(module, nn.Linear) or "Linear" in cls_name: + linear_count += 1 + + in_features = getattr(module, "in_features", None) + out_features = getattr(module, "out_features", None) + + print(f"- {name}") + print(f" class : {cls_name}") + print(f" in_features : {in_features}") + print(f" out_features : {out_features}") + + # 代表的な重みテンソルを表示 + for attr in ["weight", "qweight", "scales", "lora_A", "lora_B"]: + if hasattr(module, attr): + obj = getattr(module, attr) + if isinstance(obj, torch.Tensor): + print( + f" {attr:<12}: shape={tuple(obj.shape)}, " + f"dtype={obj.dtype}, device={obj.device}" + ) + elif hasattr(obj, "weight") and isinstance(obj.weight, torch.Tensor): + print( + f" {attr}.weight : shape={tuple(obj.weight.shape)}, " + f"dtype={obj.weight.dtype}, device={obj.weight.device}" + ) + + print() + + if linear_count >= max_modules: + print(f"... truncated after {max_modules} modules") + break + + print(f"Detected linear-like modules: {linear_count}") + + # ------------------------------------------------------------ + # 2) trainable parameter 一覧 + # ------------------------------------------------------------ + print("\n[Trainable parameters]") + trainable = 0 + total = 0 + + for name, param in model.named_parameters(): + n = param.numel() + total += n + if param.requires_grad: + trainable += n + print( + f"- {name}: shape={tuple(param.shape)}, " + f"dtype={param.dtype}, device={param.device}" + ) + + print(f"\nTrainable params: {trainable:,}") + print(f"Total params : {total:,}") + if total > 0: + print(f"Trainable ratio : {100.0 * trainable / total:.6f}%") + +def generate_text(model, tokenizer, prompt, device, max_new_tokens=128): + """Generate text from a prompt using the model.""" + inputs = tokenizer(prompt, return_tensors="pt").to(device) + with torch.no_grad(): + outputs = model.generate( + **inputs, + max_new_tokens=max_new_tokens, + do_sample=False, + temperature=1.0, + repetition_penalty=1.2, + ) + generated = tokenizer.decode(outputs[0], skip_special_tokens=True) + return generated + + +# ================================================================ +# Step 1: Quantize the model with JointQ 4-bit +# ================================================================ +print("=" * 70) +print("Step 1: Quantizing TinyLlama with JointQ 4-bit (groupsize=128)") +print("=" * 70) + +model_config = ModelConfig(model_id=MODEL_ID, device="cuda:0") +jointq = JointQ(bits=4, group_size=128) + +runner = Runner( + model_config=model_config, + quantizer=jointq, + calibration_config=CalibrationConfig(max_length=512, num_calibration_samples=128), +) +runner.run() + +# ================================================================ +# Step 2: Build quantized model +# ================================================================ +print("\n" + "=" * 70) +print("Step 2: Building quantized model via create_quantized_model") +print("=" * 70) + +model, tokenizer = runner.create_quantized_model( + pack_weights=False, + use_gemlite=False, +) + +# ================================================================ +# Step 3: Generate BEFORE LoRA SFT +# ================================================================ +print("\n" + "=" * 70) +print("Step 3: Generating text BEFORE LoRA SFT") +print("=" * 70) + +model.to("cuda:0") +before_text = generate_text(model, tokenizer, PROMPT, "cuda:0") +model.to("cpu") +torch.cuda.empty_cache() + +print(f"\nPrompt: {PROMPT}") +print(f"Response:\n{before_text}") + +# ================================================================ +# Step 4: Run LoRA SFT with OneCompression knowledge +# ================================================================ +print("\n" + "=" * 70) +print("Step 4: Running LoRA SFT with OneCompression knowledge data") +print("=" * 70) + +post_process = PostProcessLoraSFT( + data_files=KNOWLEDGE_DATA, + max_length=256, + epochs=50, + batch_size=2, + gradient_accumulation_steps=1, + lr=3e-4, + lora_r=16, + lora_alpha=32, + logging_steps=5, +) +runner.post_processes = [post_process] +runner.run_post_processes() +model = runner.quantized_model + +# ================================================================ +# Step 4.1: Save the LoRA-applied quantized model (HF safetensors + adapter sidecar) +# ================================================================ +print("\n" + "=" * 70) +print("Step 4.1-a: Inspecting model before save") +print("=" * 70) +dump_model_summary(model) + +print("\n" + "=" * 70) +print(f"Step 4.1-b: Saving LoRA-applied model to {SAVE_DIR}") +print("=" * 70) +runner.save_quantized_model(SAVE_DIR) +print(f"Model saved to: {SAVE_DIR}") +print( + " - model.safetensors, config.json : base JointQ model (HF-compatible)\n" + " - lora_adapter/adapter_model.safetensors : PEFT-format LoRA adapter\n" + " - lora_adapter/adapter_config.json : PEFT-format adapter config" +) + +# ================================================================ +# Step 4.2: Load the saved model +# ================================================================ +print("\n" + "=" * 70) +print(f"Step 4.2: Loading model from {SAVE_DIR}") +print("=" * 70) + +loaded_model, loaded_tokenizer = load_quantized_model(SAVE_DIR) +print(f"Loaded model type : {type(loaded_model).__name__}") +print(f"Loaded model device: {next(loaded_model.parameters()).device}") + +# ================================================================ +# Step 5: Generate AFTER LoRA SFT +# ================================================================ +print("\n" + "=" * 70) +print("Step 5: Generating text AFTER LoRA SFT") +print("=" * 70) + +loaded_model.to("cuda:0") +after_text = generate_text(loaded_model, loaded_tokenizer, PROMPT, "cuda:0") +loaded_model.to("cpu") +torch.cuda.empty_cache() + +print(f"\nPrompt: {PROMPT}") +print(f"Response:\n{after_text}") + +print("\n" + "=" * 70) +print("Step 5(ex): Generating text Before Save model") +print("=" * 70) + +# test before save model +model.to("cuda:0") +bs_text = generate_text(model, tokenizer, PROMPT, "cuda:0") +model.to("cpu") +torch.cuda.empty_cache() + +print(f"\nPrompt: {PROMPT}") +print(f"Response:\n{bs_text}") + +# ================================================================ +# Step 6: Compare results +# ================================================================ +print("\n" + "=" * 70) +print("Comparison: Before vs After LoRA SFT") +print("=" * 70) +print(f"\nPrompt: {PROMPT}") +print(f"\n--- BEFORE LoRA SFT ---") +print(before_text) +print(f"\n--- AFTER LoRA SFT ---") +print(after_text) +print("=" * 70) + + + From 71cd6663f9bec89440217b6f06798e80469af091 Mon Sep 17 00:00:00 2001 From: cm-ysmz Date: Wed, 22 Apr 2026 11:17:18 +0900 Subject: [PATCH 05/82] [fix] add missing commit --- .../test_post_process_lora_sft.py | 60 +++++++++++++++++++ 1 file changed, 60 insertions(+) diff --git a/tests/onecomp/post_process/test_post_process_lora_sft.py b/tests/onecomp/post_process/test_post_process_lora_sft.py index b60e176..068333c 100644 --- a/tests/onecomp/post_process/test_post_process_lora_sft.py +++ b/tests/onecomp/post_process/test_post_process_lora_sft.py @@ -27,6 +27,13 @@ PostProcessLoraSFT, ) +try: + from onecomp import JointQ + + HAS_JOINTQ = True +except ImportError: + HAS_JOINTQ = False + MODEL_ID = "TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T" FIXTURES_DIR = Path(__file__).parent.parent / "fixtures" SFT_DATA_FILE = str(FIXTURES_DIR / "sft_train_data.jsonl") @@ -139,3 +146,56 @@ def test_runner_with_post_process(self): del runner gc.collect() torch.cuda.empty_cache() + + +@pytest.mark.skipif(not HAS_JOINTQ, reason="jointq package not installed") +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.slow +class TestPostProcessLoraSFTViaRunnerJointQ: + """Smoke test for JointQ + PostProcessLoraSFT via Runner.run().""" + + def test_runner_with_jointq_and_post_process(self): + """Runner with JointQ + post_processes runs end-to-end without error.""" + setup_logger() + + model_config = ModelConfig(model_id=MODEL_ID, device="cuda:0") + quantizer = JointQ(bits=4, group_size=128) + + post_process = PostProcessLoraSFT( + data_files=SFT_DATA_FILE, + epochs=1, + max_train_samples=4, + max_length=64, + batch_size=2, + gradient_accumulation_steps=1, + logging_steps=1, + ) + + runner = Runner( + model_config=model_config, + quantizer=quantizer, + calibration_config=CalibrationConfig(num_calibration_samples=8, max_length=512), + post_processes=[post_process], + ) + runner.run() + + assert ( + runner.quantized_model is not None + ), "runner.quantized_model should be set after JointQ + post-process" + + lora_count = sum( + 1 + for _name, m in runner.quantized_model.named_modules() + if isinstance(m, LoRAGPTQLinear) + ) + assert lora_count > 0, "No LoRAGPTQLinear layers found in JointQ runner.quantized_model" + + devices = {str(p.device) for p in runner.quantized_model.parameters()} + assert devices == {"cpu"}, f"Expected all params on CPU, got {devices}" + + assert not runner.quantized_model.training, "Model should be in eval mode after run()" + + del runner + gc.collect() + torch.cuda.empty_cache() + From 8aa8291d3e04ef3f3c5f9a25f7b22f35cbfc6eb7 Mon Sep 17 00:00:00 2001 From: cm_ysmz Date: Wed, 22 Apr 2026 16:35:41 +0900 Subject: [PATCH 06/82] [refactor] revised comments and renamed methods --- ..._lora_sft_knowledge_save_load_hf_format.py | 123 +++--------------- ...ft_knowledge_save_load_hf_format_jointq.py | 123 +++--------------- onecomp/quantized_model_loader.py | 4 +- 3 files changed, 36 insertions(+), 214 deletions(-) diff --git a/example/post_process/example_lora_sft_knowledge_save_load_hf_format.py b/example/post_process/example_lora_sft_knowledge_save_load_hf_format.py index 9bf7982..02d7f74 100644 --- a/example/post_process/example_lora_sft_knowledge_save_load_hf_format.py +++ b/example/post_process/example_lora_sft_knowledge_save_load_hf_format.py @@ -11,26 +11,22 @@ 2. Build quantized model via create_quantized_model 3. Generate text BEFORE LoRA SFT (model does not know OneCompression) 4. Run LoRA SFT with OneCompression knowledge data - 4-1. Save via save_quantized_model() - writes HF-compatible safetensors - plus a PEFT-format LoRA adapter sidecar - (``adapter_model.safetensors`` + ``adapter_config.json``) - 4-2. Load via load_quantized_model() - the sidecar is auto-detected and - matching GPTQLinear layers are re-wrapped with LoRAGPTQLinear - 5. Generate text AFTER LoRA SFT (model can describe OneCompression) - 6. Compare results side by side + 5. Save the LoRA-applied quantized model in HF-compatible format + 6. Load the saved model and auto-apply the LoRA adapter + 7. Generate text AFTER LoRA SFT (model can describe OneCompression) + 8. Compare results side by side Copyright 2025-2026 Fujitsu Ltd. Author: Keiji Kimura Usage: - python example/post_process/example_lora_sft_knowledge.py + python example/post_process/example_lora_sft_knowledge_save_load_hf_format.py """ from pathlib import Path import torch -import torch.nn as nn from onecomp import CalibrationConfig, GPTQ, ModelConfig, Runner, PostProcessLoraSFT, setup_logger, load_quantized_model @@ -41,76 +37,6 @@ PROMPT = "Q: What is OneCompression?\nA:" SAVE_DIR = "./tinyllama_gptq4_lora_knowledge" -def dump_model_summary(model, max_modules=1000): - print("\n" + "=" * 70) - print("Model summary before save") - print("=" * 70) - - # ------------------------------------------------------------ - # 1) Linear系モジュールの一覧 - # ------------------------------------------------------------ - print("\n[Linear-like modules]") - linear_count = 0 - - for name, module in model.named_modules(): - cls_name = module.__class__.__name__ - - # 通常の Linear、またはクラス名に Linear を含むものを拾う - if isinstance(module, nn.Linear) or "Linear" in cls_name: - linear_count += 1 - - in_features = getattr(module, "in_features", None) - out_features = getattr(module, "out_features", None) - - print(f"- {name}") - print(f" class : {cls_name}") - print(f" in_features : {in_features}") - print(f" out_features : {out_features}") - - # 代表的な重みテンソルを表示 - for attr in ["weight", "qweight", "scales", "lora_A", "lora_B"]: - if hasattr(module, attr): - obj = getattr(module, attr) - if isinstance(obj, torch.Tensor): - print( - f" {attr:<12}: shape={tuple(obj.shape)}, " - f"dtype={obj.dtype}, device={obj.device}" - ) - elif hasattr(obj, "weight") and isinstance(obj.weight, torch.Tensor): - print( - f" {attr}.weight : shape={tuple(obj.weight.shape)}, " - f"dtype={obj.weight.dtype}, device={obj.weight.device}" - ) - - print() - - if linear_count >= max_modules: - print(f"... truncated after {max_modules} modules") - break - - print(f"Detected linear-like modules: {linear_count}") - - # ------------------------------------------------------------ - # 2) trainable parameter 一覧 - # ------------------------------------------------------------ - print("\n[Trainable parameters]") - trainable = 0 - total = 0 - - for name, param in model.named_parameters(): - n = param.numel() - total += n - if param.requires_grad: - trainable += n - print( - f"- {name}: shape={tuple(param.shape)}, " - f"dtype={param.dtype}, device={param.device}" - ) - - print(f"\nTrainable params: {trainable:,}") - print(f"Total params : {total:,}") - if total > 0: - print(f"Trainable ratio : {100.0 * trainable / total:.6f}%") def generate_text(model, tokenizer, prompt, device, max_new_tokens=128): """Generate text from a prompt using the model.""" @@ -191,18 +117,12 @@ def generate_text(model, tokenizer, prompt, device, max_new_tokens=128): ) runner.post_processes = [post_process] runner.run_post_processes() -model = runner.quantized_model # ================================================================ -# Step 4.1: Save the LoRA-applied quantized model (HF safetensors + adapter sidecar) +# Step 5: Save the LoRA-applied quantized model (HF safetensors + adapter sidecar) # ================================================================ print("\n" + "=" * 70) -print("Step 4.1-a: Inspecting model before save") -print("=" * 70) -dump_model_summary(model) - -print("\n" + "=" * 70) -print(f"Step 4.1-b: Saving LoRA-applied model to {SAVE_DIR}") +print(f"Step 5: Saving LoRA-applied model to {SAVE_DIR}") print("=" * 70) runner.save_quantized_model(SAVE_DIR) print(f"Model saved to: {SAVE_DIR}") @@ -211,12 +131,16 @@ def generate_text(model, tokenizer, prompt, device, max_new_tokens=128): " - lora_adapter/adapter_model.safetensors : PEFT-format LoRA adapter\n" " - lora_adapter/adapter_config.json : PEFT-format adapter config" ) +# Release references and clear CUDA cache before reload to reduce OOM risk +del runner +del model +torch.cuda.empty_cache() # ================================================================ -# Step 4.2: Load the saved model +# Step 6: Load the saved model with LoRA adapter # ================================================================ print("\n" + "=" * 70) -print(f"Step 4.2: Loading model from {SAVE_DIR}") +print(f"Step 6: Loading model from {SAVE_DIR}") print("=" * 70) loaded_model, loaded_tokenizer = load_quantized_model(SAVE_DIR) @@ -224,10 +148,10 @@ def generate_text(model, tokenizer, prompt, device, max_new_tokens=128): print(f"Loaded model device: {next(loaded_model.parameters()).device}") # ================================================================ -# Step 5: Generate AFTER LoRA SFT +# Step 7: Generate AFTER LoRA SFT # ================================================================ print("\n" + "=" * 70) -print("Step 5: Generating text AFTER LoRA SFT") +print("Step 7: Generating text AFTER LoRA SFT") print("=" * 70) loaded_model.to("cuda:0") @@ -238,24 +162,11 @@ def generate_text(model, tokenizer, prompt, device, max_new_tokens=128): print(f"\nPrompt: {PROMPT}") print(f"Response:\n{after_text}") -print("\n" + "=" * 70) -print("Step 5(ex): Generating text Before Save model") -print("=" * 70) - -# test before save model -model.to("cuda:0") -bs_text = generate_text(model, tokenizer, PROMPT, "cuda:0") -model.to("cpu") -torch.cuda.empty_cache() - -print(f"\nPrompt: {PROMPT}") -print(f"Response:\n{bs_text}") - # ================================================================ -# Step 6: Compare results +# Step 8: Compare results # ================================================================ print("\n" + "=" * 70) -print("Comparison: Before vs After LoRA SFT") +print("Step 8: Comparison: Before vs After LoRA SFT") print("=" * 70) print(f"\nPrompt: {PROMPT}") print(f"\n--- BEFORE LoRA SFT ---") diff --git a/example/post_process/example_lora_sft_knowledge_save_load_hf_format_jointq.py b/example/post_process/example_lora_sft_knowledge_save_load_hf_format_jointq.py index e66335f..d17ad12 100644 --- a/example/post_process/example_lora_sft_knowledge_save_load_hf_format_jointq.py +++ b/example/post_process/example_lora_sft_knowledge_save_load_hf_format_jointq.py @@ -11,26 +11,22 @@ 2. Build quantized model via create_quantized_model 3. Generate text BEFORE LoRA SFT (model does not know OneCompression) 4. Run LoRA SFT with OneCompression knowledge data - 4-1. Save via save_quantized_model() - writes HF-compatible safetensors - plus a PEFT-format LoRA adapter sidecar - (``adapter_model.safetensors`` + ``adapter_config.json``) - 4-2. Load via load_quantized_model() - the sidecar is auto-detected and - matching GPTQLinear layers are re-wrapped with LoRAGPTQLinear - 5. Generate text AFTER LoRA SFT (model can describe OneCompression) - 6. Compare results side by side + 5. Save the LoRA-applied quantized model in HF-compatible format + 6. Load the saved model and auto-apply the LoRA adapter + 7. Generate text AFTER LoRA SFT (model can describe OneCompression) + 8. Compare results side by side Copyright 2025-2026 Fujitsu Ltd. Author: Keiji Kimura Usage: - python example/post_process/example_lora_sft_knowledge.py + python example/post_process/example_lora_sft_knowledge_save_load_hf_format_jointq.py """ from pathlib import Path import torch -import torch.nn as nn from onecomp import CalibrationConfig, JointQ, ModelConfig, Runner, PostProcessLoraSFT, setup_logger, load_quantized_model @@ -41,76 +37,6 @@ PROMPT = "Q: What is OneCompression?\nA:" SAVE_DIR = "./tinyllama_jointq4_lora_knowledge" -def dump_model_summary(model, max_modules=1000): - print("\n" + "=" * 70) - print("Model summary before save") - print("=" * 70) - - # ------------------------------------------------------------ - # 1) Linear系モジュールの一覧 - # ------------------------------------------------------------ - print("\n[Linear-like modules]") - linear_count = 0 - - for name, module in model.named_modules(): - cls_name = module.__class__.__name__ - - # 通常の Linear、またはクラス名に Linear を含むものを拾う - if isinstance(module, nn.Linear) or "Linear" in cls_name: - linear_count += 1 - - in_features = getattr(module, "in_features", None) - out_features = getattr(module, "out_features", None) - - print(f"- {name}") - print(f" class : {cls_name}") - print(f" in_features : {in_features}") - print(f" out_features : {out_features}") - - # 代表的な重みテンソルを表示 - for attr in ["weight", "qweight", "scales", "lora_A", "lora_B"]: - if hasattr(module, attr): - obj = getattr(module, attr) - if isinstance(obj, torch.Tensor): - print( - f" {attr:<12}: shape={tuple(obj.shape)}, " - f"dtype={obj.dtype}, device={obj.device}" - ) - elif hasattr(obj, "weight") and isinstance(obj.weight, torch.Tensor): - print( - f" {attr}.weight : shape={tuple(obj.weight.shape)}, " - f"dtype={obj.weight.dtype}, device={obj.weight.device}" - ) - - print() - - if linear_count >= max_modules: - print(f"... truncated after {max_modules} modules") - break - - print(f"Detected linear-like modules: {linear_count}") - - # ------------------------------------------------------------ - # 2) trainable parameter 一覧 - # ------------------------------------------------------------ - print("\n[Trainable parameters]") - trainable = 0 - total = 0 - - for name, param in model.named_parameters(): - n = param.numel() - total += n - if param.requires_grad: - trainable += n - print( - f"- {name}: shape={tuple(param.shape)}, " - f"dtype={param.dtype}, device={param.device}" - ) - - print(f"\nTrainable params: {trainable:,}") - print(f"Total params : {total:,}") - if total > 0: - print(f"Trainable ratio : {100.0 * trainable / total:.6f}%") def generate_text(model, tokenizer, prompt, device, max_new_tokens=128): """Generate text from a prompt using the model.""" @@ -191,18 +117,12 @@ def generate_text(model, tokenizer, prompt, device, max_new_tokens=128): ) runner.post_processes = [post_process] runner.run_post_processes() -model = runner.quantized_model # ================================================================ -# Step 4.1: Save the LoRA-applied quantized model (HF safetensors + adapter sidecar) +# Step 5: Save the LoRA-applied quantized model (HF safetensors + adapter sidecar) # ================================================================ print("\n" + "=" * 70) -print("Step 4.1-a: Inspecting model before save") -print("=" * 70) -dump_model_summary(model) - -print("\n" + "=" * 70) -print(f"Step 4.1-b: Saving LoRA-applied model to {SAVE_DIR}") +print(f"Step 5: Saving LoRA-applied model to {SAVE_DIR}") print("=" * 70) runner.save_quantized_model(SAVE_DIR) print(f"Model saved to: {SAVE_DIR}") @@ -211,12 +131,16 @@ def generate_text(model, tokenizer, prompt, device, max_new_tokens=128): " - lora_adapter/adapter_model.safetensors : PEFT-format LoRA adapter\n" " - lora_adapter/adapter_config.json : PEFT-format adapter config" ) +# Release references and clear CUDA cache before reload to reduce OOM risk +del runner +del model +torch.cuda.empty_cache() # ================================================================ -# Step 4.2: Load the saved model +# Step 6: Load the saved model with LoRA adapter # ================================================================ print("\n" + "=" * 70) -print(f"Step 4.2: Loading model from {SAVE_DIR}") +print(f"Step 6: Loading model from {SAVE_DIR}") print("=" * 70) loaded_model, loaded_tokenizer = load_quantized_model(SAVE_DIR) @@ -224,10 +148,10 @@ def generate_text(model, tokenizer, prompt, device, max_new_tokens=128): print(f"Loaded model device: {next(loaded_model.parameters()).device}") # ================================================================ -# Step 5: Generate AFTER LoRA SFT +# Step 7: Generate AFTER LoRA SFT # ================================================================ print("\n" + "=" * 70) -print("Step 5: Generating text AFTER LoRA SFT") +print("Step 7: Generating text AFTER LoRA SFT") print("=" * 70) loaded_model.to("cuda:0") @@ -238,24 +162,11 @@ def generate_text(model, tokenizer, prompt, device, max_new_tokens=128): print(f"\nPrompt: {PROMPT}") print(f"Response:\n{after_text}") -print("\n" + "=" * 70) -print("Step 5(ex): Generating text Before Save model") -print("=" * 70) - -# test before save model -model.to("cuda:0") -bs_text = generate_text(model, tokenizer, PROMPT, "cuda:0") -model.to("cpu") -torch.cuda.empty_cache() - -print(f"\nPrompt: {PROMPT}") -print(f"Response:\n{bs_text}") - # ================================================================ -# Step 6: Compare results +# Step 8: Compare results # ================================================================ print("\n" + "=" * 70) -print("Comparison: Before vs After LoRA SFT") +print("Step 8: Comparison: Before vs After LoRA SFT") print("=" * 70) print(f"\nPrompt: {PROMPT}") print(f"\n--- BEFORE LoRA SFT ---") diff --git a/onecomp/quantized_model_loader.py b/onecomp/quantized_model_loader.py index 2c01b18..2e8865a 100644 --- a/onecomp/quantized_model_loader.py +++ b/onecomp/quantized_model_loader.py @@ -121,7 +121,7 @@ def load_quantized_model( # Re-apply LoRA adapter from PEFT-format sidecar if present. # This must run while the model is still on CPU, before dispatch_model, # so LoRA wrappers are included in the device map traversal below. - cls._maybe_wrap_lora_adapters(model, save_directory) + cls._apply_lora_adapters_from_sidecar(model, save_directory) # Device placement if device_map: @@ -388,7 +388,7 @@ def _replace_quantized_layers(model, state_dict: dict, quant_config: dict): LORA_ADAPTER_SUBDIR = "lora_adapter" @staticmethod - def _maybe_wrap_lora_adapters(model, save_directory: str) -> int: + def _apply_lora_adapters_from_sidecar(model, save_directory: str) -> int: """Re-wrap GPTQLinear layers with LoRAGPTQLinear from a PEFT-format sidecar. Looks for ``adapter_model.safetensors`` + ``adapter_config.json`` under From c65d7e243e6157f19015d741bcee1ea490a59e53 Mon Sep 17 00:00:00 2001 From: kimura-keiji Date: Thu, 7 May 2026 21:20:33 +0900 Subject: [PATCH 07/82] Define v1.1.1 --- CHANGELOG.md | 2 ++ onecomp/__version__.py | 2 +- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 29c237b..1fbe04d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,7 @@ # Change log +## [v1.1.1] 2026-05-07 + ## [v1.1.0] 2026-04-16 ### Gemma 3 / Gemma 4 & VLM Support diff --git a/onecomp/__version__.py b/onecomp/__version__.py index 0e98bdc..8794f6b 100644 --- a/onecomp/__version__.py +++ b/onecomp/__version__.py @@ -6,4 +6,4 @@ """ -__version__ = "1.1.0" +__version__ = "1.1.1" From 9775761e8826cf5013c70eea886dfafc790433de Mon Sep 17 00:00:00 2001 From: sikoji Date: Mon, 27 Apr 2026 14:12:53 +0900 Subject: [PATCH 08/82] [fix] support load of unpacked JointQ wbits=1 checkpoints GPTQLinear weight packing only supports wbits in (2, 3, 4, 8), so JointQ with bits=1 must build/save the inference layer with pack_weights=False. - JointQ.validate_params: warn when bits=1 to remind callers to pass pack_weights=False at inference layer construction time - GPTQLinear.from_saved_state: load qweight/qzeros as unpacked tensors when wbits=1, matching the unpacked save format --- onecomp/quantizer/gptq/gptq_layer.py | 4 +++- onecomp/quantizer/jointq/_jointq.py | 9 +++++++++ 2 files changed, 12 insertions(+), 1 deletion(-) diff --git a/onecomp/quantizer/gptq/gptq_layer.py b/onecomp/quantizer/gptq/gptq_layer.py index 1dc5e78..74bb238 100644 --- a/onecomp/quantizer/gptq/gptq_layer.py +++ b/onecomp/quantizer/gptq/gptq_layer.py @@ -468,7 +468,9 @@ def from_saved_state( # pylint: disable=too-many-positional-arguments self.groupsize = groupsize self.actorder = actorder self.checkpoint_format = checkpoint_format - self._weight_is_packed = True + # JointQ wbits=1 is saved with pack_weights=False + # (GPTQLinear packing does not support 1-bit), so load it unpacked as well. + self._weight_is_packed = wbits != 1 def _t(k): t = layer_state_dict[k] diff --git a/onecomp/quantizer/jointq/_jointq.py b/onecomp/quantizer/jointq/_jointq.py index 3f59530..370e113 100644 --- a/onecomp/quantizer/jointq/_jointq.py +++ b/onecomp/quantizer/jointq/_jointq.py @@ -219,11 +219,20 @@ def validate_params(self): ils_num_iterations: int >= 1 (when ils_enabled=True) ils_num_clones: int >= 1 (when ils_enabled=True) ils_num_channels: int >= 1 or None (when ils_enabled=True) + + Warnings: + bits == 1: valid, but GPTQLinear weight packing does not support + 1-bit; the inference layer must be built with pack_weights=False. """ bad = [] if not (isinstance(self.bits, int) and self.bits >= 1): bad.append(f"Invalid JointQ parameter 'bits': {self.bits!r} (expected int >= 1).") + elif self.bits == 1: + self.logger.warning( + "JointQ with bits=1 is not supported by GPTQLinear weight packing; " + "build the inference layer with pack_weights=False." + ) if not (isinstance(self.group_size, int) and self.group_size >= 1): bad.append( From 764b98a8194f1f7c9def4f39065dc039ef5477c5 Mon Sep 17 00:00:00 2001 From: "sota.n" Date: Thu, 14 May 2026 08:34:46 +0900 Subject: [PATCH 09/82] feat: log quantization progress with linear ETA Add QuantizationProgressTracker and wire it through calibration, chunked calibration, multi-GPU phase 2, QEP general and arch-aware paths. Runner gains quantization_progress flag (default on). Includes unit tests for ETA formatting and thread-safe stepping. Co-authored-by: Cursor --- onecomp/qep/_quantize_with_qep.py | 16 +++ onecomp/qep/_quantize_with_qep_arch.py | 38 ++++++-- onecomp/runner.py | 97 ++++++++++++------- .../runner_methods/chunked_quantization.py | 46 ++++++++- .../runner_methods/multi_gpu_quantization.py | 20 ++++ onecomp/utils/quantization_progress.py | 88 +++++++++++++++++ .../utils/test_quantization_progress.py | 79 +++++++++++++++ 7 files changed, 335 insertions(+), 49 deletions(-) create mode 100644 onecomp/utils/quantization_progress.py create mode 100644 tests/onecomp/utils/test_quantization_progress.py diff --git a/onecomp/qep/_quantize_with_qep.py b/onecomp/qep/_quantize_with_qep.py index 25003db..e0ba901 100644 --- a/onecomp/qep/_quantize_with_qep.py +++ b/onecomp/qep/_quantize_with_qep.py @@ -36,6 +36,8 @@ def run_quantize_with_qep( quantizer: Quantizer, qep_config: QEPConfig, calibration_config: CalibrationConfig, + *, + quantization_progress: bool = True, ): """Run quantization with Quantization Error Propagation (QEP). @@ -51,6 +53,7 @@ def run_quantize_with_qep( qep_config (QEPConfig): Configuration for QEP (percdamp, perccorr, exclude_layer_keywords). calibration_config (CalibrationConfig): Calibration parameters. + quantization_progress (bool): When True, log ``[progress]`` with ETA per layer. """ model = model_config.load_model() @@ -80,6 +83,17 @@ def run_quantize_with_qep( logger.info("Quantizing the model using %s", quantizer.name) + progress = None + if quantization_progress: + # pylint: disable-next=import-outside-toplevel + from onecomp.utils.quantization_progress import QuantizationProgressTracker + + progress = QuantizationProgressTracker( + logger, + len(quantizer.module_to_name), + "QEP quantization (general, per layer)", + ) + # 2. For each target layer, perform the following sequentially for module, name in quantizer.module_to_name.items(): @@ -114,6 +128,8 @@ def run_quantize_with_qep( # 2-4. Free memory del quant_input_activation + if progress is not None: + progress.step_complete(name) del original_input_activations quantizer.execute_post_processing() diff --git a/onecomp/qep/_quantize_with_qep_arch.py b/onecomp/qep/_quantize_with_qep_arch.py index 137809b..6f51fcb 100644 --- a/onecomp/qep/_quantize_with_qep_arch.py +++ b/onecomp/qep/_quantize_with_qep_arch.py @@ -143,6 +143,7 @@ def compute_hessian_and_crossterm( def make_hook(name): def hook(module, inp, out): dest[name] = inp[0] if isinstance(inp, tuple) else inp + return hook handlers = [ @@ -213,6 +214,7 @@ def _compute_per_module_hessians( def _make_hook(key): def hook(_, inp, __): dest[key] = inp[0] if isinstance(inp, tuple) else inp + return hook handlers = [m.register_forward_hook(_make_hook(i)) for i, m in enumerate(modules)] @@ -251,10 +253,7 @@ def hook(_, inp, __): for h in handlers: h.remove() - return { - modules[i]: (hessians[i] if nsamples[i] > 0 else None) - for i in range(len(modules)) - } + return {modules[i]: (hessians[i] if nsamples[i] > 0 else None) for i in range(len(modules))} @torch.no_grad() @@ -263,6 +262,8 @@ def run_quantize_with_qep_arch( quantizer: Quantizer, qep_config: QEPConfig, calibration_config: CalibrationConfig, + *, + quantization_progress: bool = True, ): """Run architecture-aware quantization with QEP. @@ -279,6 +280,7 @@ def run_quantize_with_qep_arch( qep_config (QEPConfig): Configuration for QEP (percdamp, perccorr, exclude_layer_keywords). calibration_config (CalibrationConfig): Calibration parameters. + quantization_progress (bool): When True, log ``[progress]`` with ETA per target layer. """ @@ -318,6 +320,17 @@ def run_quantize_with_qep_arch( name for module, name in quantizer.module_to_name.items() if module in block_modules } + progress = None + if quantization_progress: + # pylint: disable-next=import-outside-toplevel + from onecomp.utils.quantization_progress import QuantizationProgressTracker + + progress = QuantizationProgressTracker( + logger, + len(remaining_targets), + "QEP quantization (architecture-aware)", + ) + # 2. For each target transformer block, perform the following sequentially for block_idx, block in enumerate(blocks): @@ -365,9 +378,7 @@ def run_quantize_with_qep_arch( targets = [m for m in group_q if m in quantizer.module_to_name] if not targets: continue - is_expert = any( - ".experts." in quantizer.module_to_name[m] for m in targets - ) + is_expert = any(".experts." in quantizer.module_to_name[m] for m in targets) if is_expert: expert_modules_q.extend(targets) else: @@ -442,6 +453,8 @@ def run_quantize_with_qep_arch( name, ) remaining_targets.discard(name) + if progress is not None: + progress.step_complete(name) # 4. Process MoE expert layers with per-module Hessians (no cross-term) if expert_modules_q: @@ -451,7 +464,12 @@ def run_quantize_with_qep_arch( len(expert_modules_q), ) expert_hessians = _compute_per_module_hessians( - block_q, expert_modules_q, inps_q, kwargs, batch_size, device, + block_q, + expert_modules_q, + inps_q, + kwargs, + batch_size, + device, ) for module_q in expert_modules_q: name = quantizer.module_to_name[module_q] @@ -462,6 +480,8 @@ def run_quantize_with_qep_arch( name, ) remaining_targets.discard(name) + if progress is not None: + progress.step_complete(f"{name} (skipped, no tokens)") continue logger.info( @@ -489,6 +509,8 @@ def run_quantize_with_qep_arch( name, ) remaining_targets.discard(name) + if progress is not None: + progress.step_complete(name) # forward input to the next block inps_q = forward_input(inps_q, block_q, kwargs, batch_size, device) diff --git a/onecomp/runner.py b/onecomp/runner.py index fee4068..2771336 100644 --- a/onecomp/runner.py +++ b/onecomp/runner.py @@ -86,6 +86,7 @@ def __init__( multi_gpu=False, gpu_ids=None, post_processes=None, + quantization_progress=True, ): """__init__ method @@ -130,6 +131,11 @@ def __init__( a quantized model on CPU (built via ``create_quantized_model``) and may modify it in-place. Processes are executed in order. Default is None. + quantization_progress (bool): + When ``True`` (default), emit ``[progress]`` log lines with + completed steps, elapsed time, and a linear ETA estimate + during long quantization (calibration, chunked, multi-GPU, + QEP). Set to ``False`` for quiet runs (e.g. CI). Note: For zero-config quantization (VRAM auto-estimation + @@ -215,6 +221,7 @@ def __init__( self.lpcd_config = None if lpcd: self.lpcd_config = lpcd_config if lpcd_config is not None else LPCDConfig() + self.quantization_progress = quantization_progress def check(self): """Check the settings @@ -314,27 +321,19 @@ def _exclude_moe_router_if_needed(self): config = self.model_config.load_config() num_experts = ( getattr(config, "num_experts", 0) - or getattr( - getattr(config, "text_config", None), "num_experts", 0 - ) or - 0 + or getattr(getattr(config, "text_config", None), "num_experts", 0) + or 0 ) if num_experts == 0: return keyword = "router" - target_quantizers = ( - self.quantizers - if self.quantizers is not None - else [self.quantizer] - ) + target_quantizers = self.quantizers if self.quantizers is not None else [self.quantizer] for q in target_quantizers: if q.exclude_layer_keywords is None: q.exclude_layer_keywords = [keyword] elif keyword not in q.exclude_layer_keywords: - q.exclude_layer_keywords = list(q.exclude_layer_keywords) + [ - keyword - ] + q.exclude_layer_keywords = list(q.exclude_layer_keywords) + [keyword] self.logger.info( "MoE model (num_experts=%d): excluding '%s' layers from " @@ -511,12 +510,9 @@ def auto_run( uniform_bit = max(valid_wbits) if save_dir == "auto": model_name = model_id.rstrip("/").split("/")[-1] - save_dir = ( - f"{model_name}-gptq-{uniform_bit}bit" - ) + save_dir = f"{model_name}-gptq-{uniform_bit}bit" logger.warning( - "Gemma 4 detected → falling back to uniform GPTQ %d-bit " - "(target wbits=%.2f)", + "Gemma 4 detected → falling back to uniform GPTQ %d-bit " "(target wbits=%.2f)", uniform_bit, wbits, ) @@ -527,6 +523,7 @@ def auto_run( save_dir = f"{model_name}-autobit-{wbits}bit" from .quantizer.autobit import AutoBitQuantizer + candidate_quantizers = [ GPTQ(wbits=b, groupsize=groupsize, **kwargs) for b in candidate_bits ] @@ -595,8 +592,30 @@ def quantize_with_calibration(self): # Register hooks to all linear layers handles = [] + progress = None + if self.quantization_progress: + # pylint: disable-next=import-outside-toplevel + from .utils.quantization_progress import QuantizationProgressTracker + + progress = QuantizationProgressTracker( + logger, + len(self.quantizer.module_to_name), + "Calibration quantization layers", + ) + + if progress: + quantize_bound = self.quantizer.quantize + + def _quantize_hook(module, input, output): # pylint: disable=redefined-builtin + quantize_bound(module, input, output) + progress.step_complete(self.quantizer.module_to_name[module]) + + hook_fn = _quantize_hook + else: + hook_fn = self.quantizer.quantize + for module in self.quantizer.module_to_name.keys(): - handle = module.register_forward_hook(self.quantizer.quantize) + handle = module.register_forward_hook(hook_fn) handles.append(handle) logger.info("Quantizing the model using %s", self.quantizer.name) @@ -636,6 +655,7 @@ def quantize_with_calibration_chunked(self): model_config=self.model_config, quantizers=self.quantizers if self.quantizers is not None else [self.quantizer], calibration_config=self.calibration_config, + quantization_progress=self.quantization_progress, ) def quantize_with_calibration_on_multi_gpu(self): @@ -664,6 +684,7 @@ def quantize_with_calibration_on_multi_gpu(self): quantizer=self.quantizer, calibration_config=self.calibration_config, gpu_ids=self.gpu_ids, + quantization_progress=self.quantization_progress, ) # Store results in quantizer.results @@ -690,8 +711,20 @@ def quantize_without_calibration(self): "Quantizing the model without calibration using %s", self.quantizer.name, ) + progress = None + if self.quantization_progress: + # pylint: disable-next=import-outside-toplevel + from .utils.quantization_progress import QuantizationProgressTracker + + progress = QuantizationProgressTracker( + logger, + len(self.quantizer.module_to_name), + "Quantization without calibration (layers)", + ) for module in self.quantizer.module_to_name.keys(): self.quantizer.quantize(module, None, None) + if progress: + progress.step_complete(self.quantizer.module_to_name[module]) self.quantizer.execute_post_processing() @@ -712,6 +745,7 @@ def quantize_with_qep(self): quantizer=self.quantizer, qep_config=self.qep_config, calibration_config=self.calibration_config, + quantization_progress=self.quantization_progress, ) if self.qep_config.general: @@ -859,7 +893,7 @@ def prepare_calibration_dataset(self, device, model=None): Args: device (torch.device): Device to place tensors on (CPU or GPU) - model: Model instance (optional). Add model-specific fields + model: Model instance (optional). Add model-specific fields (e.g. mm_token_type_ids for Gemma 4). Returns: @@ -983,10 +1017,7 @@ def save_quantization_statistics(self, path: str, quantizer=None): logger.info("Saving the quantization statistics to %s", path) - statistics = { - key: result.get_statistics() - for key, result in quantizer.results.items() - } + statistics = {key: result.get_statistics() for key, result in quantizer.results.items()} with open(path, "w", encoding="utf-8") as f: json.dump(statistics, f, indent=4) @@ -1665,15 +1696,10 @@ def create_quantized_model(self, pack_weights: bool = True, quantizer=None, use_ # cf) https://docs.vllm.ai/en/stable/features/quantization/#implementing-a-quantized-moe-method num_experts = ( getattr(model.config, "num_experts", None) - or getattr( - getattr(model.config, "text_config", None), "num_experts", None - ) + or getattr(getattr(model.config, "text_config", None), "num_experts", None) or 0 ) - if ( - quant_config.get("quant_method") == "gptq" - and num_experts > 0 - ): + if quant_config.get("quant_method") == "gptq" and num_experts > 0: quant_config["quant_method"] = "mixed_gptq" self.logger.info( "MoE model detected (num_experts=%d): " @@ -1697,15 +1723,13 @@ def _patch_k_eq_v_for_vllm(self, model, quant_config: dict) -> None: Gemma4 full-attention layers with attention_k_eq_v=True have no v_proj weight — the model reuses key states as value states. vLLM fuses q/k/v into a single qkv_proj and requires all shards - to share the same quantization status. + to share the same quantization status. """ text_cfg = getattr(model.config, "text_config", None) if text_cfg is None or not getattr(text_cfg, "attention_k_eq_v", False): return layer_types = getattr(text_cfg, "layer_types", []) - k_eq_v_indices = { - i for i, lt in enumerate(layer_types) if lt == "full_attention" - } + k_eq_v_indices = {i for i, lt in enumerate(layer_types) if lt == "full_attention"} if not k_eq_v_indices: return @@ -1743,9 +1767,7 @@ def _patch_k_eq_v_for_vllm(self, model, quant_config: dict) -> None: and "self_attn.k_proj" in layer_cfg and "self_attn.v_proj" not in layer_cfg ): - layer_cfg["self_attn.v_proj"] = copy.deepcopy( - layer_cfg["self_attn.k_proj"] - ) + layer_cfg["self_attn.v_proj"] = copy.deepcopy(layer_cfg["self_attn.k_proj"]) for key in ("modules_in_block_to_quantize", "quantized_layer_names"): names = quant_config.get(key, []) @@ -1816,6 +1838,7 @@ def save_quantized_model(self, save_directory: str, pack_weights: bool = True): if src_dir and not os.path.isdir(src_dir): # when the model_id is specified, the path is modifed to the local directory from huggingface_hub import snapshot_download + src_dir = snapshot_download(src_dir, local_files_only=True) if src_dir and os.path.isdir(src_dir): for fname in ("processor_config.json", "preprocessor_config.json"): diff --git a/onecomp/runner_methods/chunked_quantization.py b/onecomp/runner_methods/chunked_quantization.py index 01078e1..0024133 100644 --- a/onecomp/runner_methods/chunked_quantization.py +++ b/onecomp/runner_methods/chunked_quantization.py @@ -23,6 +23,8 @@ """ +from __future__ import annotations + import time from logging import getLogger from typing import List @@ -45,6 +47,8 @@ def run_chunked_quantization( model_config: ModelConfig, quantizers: List[Quantizer], calibration_config: CalibrationConfig, + *, + quantization_progress: bool = True, ): """Run quantization for large-scale calibration data. @@ -57,6 +61,8 @@ def run_chunked_quantization( quantizers (list[Quantizer]): List of quantizers. Each quantizer must have flag_hessian=True or flag_xtx=True. calibration_config (CalibrationConfig): Calibration parameters. + quantization_progress (bool): When True, log ``[progress]`` lines with ETA + for X^T X chunks and per-layer quantization. Note: Results are stored directly in each quantizer.results. @@ -112,10 +118,21 @@ def run_chunked_quantization( (len(all_layers) + num_layers_per_group - 1) // num_layers_per_group, ) + num_groups = (len(all_layers) + num_layers_per_group - 1) // num_layers_per_group + layer_progress = None + if quantization_progress: + # pylint: disable-next=import-outside-toplevel + from onecomp.utils.quantization_progress import QuantizationProgressTracker + + layer_progress = QuantizationProgressTracker( + logger, + len(all_layers) * len(quantizers), + "Chunked quantization (layers × quantizers)", + ) + for group_start in range(0, len(all_layers), num_layers_per_group): group = all_layers[group_start : group_start + num_layers_per_group] group_names = [name for _, name in group] - num_groups = (len(all_layers) + num_layers_per_group - 1) // num_layers_per_group group_idx = group_start // num_layers_per_group + 1 layers_list = "\n".join(f" - {name}" for name in group_names) @@ -127,6 +144,15 @@ def run_chunked_quantization( layers_list, ) + chunk_progress = None + if quantization_progress: + num_chunks = (total_samples + calibration_batch_size - 1) // calibration_batch_size + chunk_progress = QuantizationProgressTracker( + logger, + num_chunks, + f"X^T X accumulation chunks (group {group_idx}/{num_groups})", + ) + # --- Phase 1: Forward per chunk -> accumulate X^T X (FP64) --- xtx_dict, nsamples = accumulate_xtx( model=model, @@ -134,12 +160,13 @@ def run_chunked_quantization( input_device=input_device, group=group, calibration_batch_size=calibration_batch_size, + chunk_progress=chunk_progress, ) # --- Phase 2 & 3: Quantize and compute errors for each quantizer --- for quantizer in quantizers: logger.info("Quantizing with %s ...", quantizer.name) - quantize_group(quantizer, group, xtx_dict, nsamples) + quantize_group(quantizer, group, xtx_dict, nsamples, layer_progress=layer_progress) if quantizer.calc_quant_error: record_quantization_errors(quantizer, group, xtx_dict, nsamples) @@ -163,6 +190,7 @@ def accumulate_xtx( input_device, group, calibration_batch_size, + chunk_progress=None, ): """Accumulate X^T X in FP64 on GPU by forwarding in chunks. @@ -177,6 +205,8 @@ def accumulate_xtx( input_device: Model's input device (GPU). group: List of target layers [(module, name), ...]. calibration_batch_size: Number of sentences per chunk. + chunk_progress: Optional :class:`~onecomp.utils.quantization_progress.QuantizationProgressTracker` + to record chunk completion with ETA. Returns: xtx_dict: Dict[name, torch.Tensor] @@ -242,6 +272,8 @@ def hook(_module, input, _output): # pylint: disable=redefined-builtin chunk_start, chunk_end, ) + if chunk_progress is not None: + chunk_progress.step_complete(f"samples {chunk_start}-{chunk_end}") # Remove hooks for handle in handles: @@ -258,7 +290,7 @@ def hook(_module, input, _output): # pylint: disable=redefined-builtin # ============================================================================= -def quantize_group(quantizer, group, xtx_dict, nsamples): +def quantize_group(quantizer, group, xtx_dict, nsamples, layer_progress=None): """Quantize each layer in the group using accumulated X^T X. flag_hessian=True (GPTQ, DBF, etc.): @@ -273,11 +305,14 @@ def quantize_group(quantizer, group, xtx_dict, nsamples): X^T X for each layer (FP64, CPU), shape (in_features, in_features). nsamples: int Number of samples (= total_samples * seq_len). + layer_progress: Optional progress tracker for completed layer quantizations. """ for module, name in group: if name not in xtx_dict and (quantizer.flag_hessian or quantizer.flag_xtx): logger.warning("Skipping %s: no activations captured (unused during forward)", name) + if layer_progress is not None: + layer_progress.step_complete(f"{name} (skipped, no activations)") continue logger.info("Quantizing layer: %s", name) @@ -285,7 +320,8 @@ def quantize_group(quantizer, group, xtx_dict, nsamples): if quantizer.flag_xtx: result = quantizer.quantize_layer( - module, input=None, + module, + input=None, matrix_XX=xtx_dict[name].to(module.weight.device), dim_n=nsamples, ) @@ -309,6 +345,8 @@ def quantize_group(quantizer, group, xtx_dict, nsamples): quantizer.results[name] = result torch.cuda.empty_cache() + if layer_progress is not None: + layer_progress.step_complete(name) # ============================================================================= diff --git a/onecomp/runner_methods/multi_gpu_quantization.py b/onecomp/runner_methods/multi_gpu_quantization.py index 2bc4b12..1f72695 100644 --- a/onecomp/runner_methods/multi_gpu_quantization.py +++ b/onecomp/runner_methods/multi_gpu_quantization.py @@ -177,6 +177,8 @@ def run_quantization_phase( layer_names: List[str], quantizer, gpu_ids: List[int], + *, + quantization_progress: bool = True, ) -> Dict[str, Dict]: """Phase 2: Parallel quantization using multiple threads. @@ -193,6 +195,18 @@ def run_quantization_phase( logger.info("Using GPUs: %s for %d layers", gpu_ids, len(layer_names)) start_time = time.time() + progress = None + if quantization_progress: + # pylint: disable-next=import-outside-toplevel + from onecomp.utils.quantization_progress import QuantizationProgressTracker + + progress = QuantizationProgressTracker( + logger, + len(layer_names), + "Multi-GPU layer quantization", + thread_safe=True, + ) + # Force PyTorch lazy initialization upfront (avoid multi-thread race conditions) # torch.linalg.solve's internal initialization can cause # "lazy wrapper should be called at most once" errors when called from multiple threads simultaneously @@ -270,6 +284,8 @@ def quantize_single_layer(layer_name: str, device: torch.device, gpu_id: int): with lock: results[layer_name] = quant_result logger.info(" %s on GPU %d: %.2f sec", layer_name, gpu_id, layer_elapsed) + if progress is not None: + progress.step_complete(layer_name) # Free memory del dummy_module @@ -342,6 +358,8 @@ def run_multi_gpu_quantization( quantizer, calibration_config: CalibrationConfig, gpu_ids: Optional[List[int]] = None, + *, + quantization_progress: bool = True, ) -> Dict[str, Any]: """Main entry point for multi-GPU quantization. @@ -350,6 +368,7 @@ def run_multi_gpu_quantization( quantizer: Quantizer instance. calibration_config (CalibrationConfig): Calibration parameters. gpu_ids: List of GPU IDs to use (all GPUs if None). + quantization_progress: When True, log ``[progress]`` with ETA per completed layer. Returns: Dict containing "results" with quantization results for each layer @@ -375,6 +394,7 @@ def run_multi_gpu_quantization( layer_names=capture_result["layer_names"], quantizer=quantizer, gpu_ids=gpu_ids, + quantization_progress=quantization_progress, ) total_elapsed = time.time() - total_start diff --git a/onecomp/utils/quantization_progress.py b/onecomp/utils/quantization_progress.py new file mode 100644 index 0000000..4e16f2c --- /dev/null +++ b/onecomp/utils/quantization_progress.py @@ -0,0 +1,88 @@ +""" + +Copyright 2025-2026 Fujitsu Ltd. + +""" + +from __future__ import annotations + +import threading +import time +from logging import Logger +from typing import Optional + + +def _format_duration(seconds: Optional[float]) -> str: + if seconds is None: + return "unknown" + sec = max(0, int(round(seconds))) + hours, rem = divmod(sec, 3600) + minutes, secs = divmod(rem, 60) + if hours: + return f"{hours}h{minutes}m" + if minutes: + return f"{minutes}m{secs}s" + return f"{secs}s" + + +class QuantizationProgressTracker: + """Log coarse progress (done/total, elapsed, linear ETA) during long quantization.""" + + def __init__( + self, + logger: Logger, + total_steps: int, + label: str, + *, + thread_safe: bool = False, + ): + self._logger = logger + self._total = int(total_steps) + self._label = label + self._done = 0 + self._start = time.monotonic() + self._lock = threading.Lock() if thread_safe else None + + @property + def done(self) -> int: + if self._lock: + with self._lock: + return self._done + return self._done + + def step_complete(self, detail: Optional[str] = None) -> None: + """Record one completed step and emit a single INFO line with ETA.""" + + if self._lock: + with self._lock: + self._step_complete_unlocked(detail) + else: + self._step_complete_unlocked(detail) + + def _step_complete_unlocked(self, detail: Optional[str]) -> None: + if self._total <= 0: + return + + now = time.monotonic() + self._done += 1 + done = min(self._done, self._total) + elapsed = now - self._start + + if done < self._total and done > 0: + eta_sec = (elapsed / done) * (self._total - done) + eta_str = _format_duration(eta_sec) + else: + eta_str = "unknown" if done == 0 else "0s" + + pct = 100.0 * done / self._total + suffix = f" ({detail})" if detail else "" + self._logger.info( + "[progress] %s: %d/%d (%.1f%%) elapsed=%s ETA=%s%s", + self._label, + done, + self._total, + pct, + _format_duration(elapsed), + eta_str, + suffix, + ) diff --git a/tests/onecomp/utils/test_quantization_progress.py b/tests/onecomp/utils/test_quantization_progress.py new file mode 100644 index 0000000..371725f --- /dev/null +++ b/tests/onecomp/utils/test_quantization_progress.py @@ -0,0 +1,79 @@ +""" + +Copyright 2025-2026 Fujitsu Ltd. + +""" + +import importlib.util +import logging +import threading +from pathlib import Path +from unittest.mock import patch + +import pytest + +_REPO_ROOT = Path(__file__).resolve().parents[3] +_spec = importlib.util.spec_from_file_location( + "quantization_progress", + _REPO_ROOT / "onecomp" / "utils" / "quantization_progress.py", +) +_quant_mod = importlib.util.module_from_spec(_spec) +assert _spec.loader is not None +_spec.loader.exec_module(_quant_mod) +QuantizationProgressTracker = _quant_mod.QuantizationProgressTracker + + +def test_step_complete_logs_fraction_and_eta(caplog): + caplog.set_level(logging.INFO) + logger = logging.getLogger("test_progress_eta") + tracker = QuantizationProgressTracker(logger, total_steps=2, label="Test layers") + + with patch.object(_quant_mod.time, "monotonic", side_effect=[0.0, 10.0, 30.0]): + tracker.step_complete("layer_a") + tracker.step_complete("layer_b") + + joined = " ".join(r.message for r in caplog.records) + assert "1/2" in joined + assert "2/2" in joined + assert "[progress]" in joined + assert "ETA" in joined + + +def test_thread_safe_reaches_total(): + logger = logging.getLogger("test_progress_thread") + logger.addHandler(logging.NullHandler()) + tracker = QuantizationProgressTracker( + logger, total_steps=100, label="Parallel", thread_safe=True + ) + + def worker(): + for _ in range(10): + tracker.step_complete() + + threads = [threading.Thread(target=worker) for _ in range(10)] + for t in threads: + t.start() + for t in threads: + t.join() + + assert tracker.done == 100 + + +def test_zero_total_no_crash(): + logger = logging.getLogger("test_progress_zero") + logger.addHandler(logging.NullHandler()) + tracker = QuantizationProgressTracker(logger, total_steps=0, label="Empty") + tracker.step_complete() # should not raise + + +def test_eta_unknown_until_first_step(caplog): + caplog.set_level(logging.INFO) + logger = logging.getLogger("test_progress_first") + tracker = QuantizationProgressTracker(logger, total_steps=3, label="Layers") + + with patch.object(_quant_mod.time, "monotonic", side_effect=[0.0, 5.0]): + tracker.step_complete("first") + + first_line = caplog.records[0].message + assert "1/3" in first_line + assert "ETA" in first_line From ce8232dd37a37da779542b2e3569f085264cb165 Mon Sep 17 00:00:00 2001 From: cm_ysmz Date: Thu, 14 May 2026 14:34:12 +0900 Subject: [PATCH 10/82] [update] Add script for training a LoRA model with OneCompression knowledge, verified with vllm --- .../example_lora_gptq_vllm_inference.py | 176 ++++++++++++++++++ 1 file changed, 176 insertions(+) create mode 100644 example/post_process/example_lora_gptq_vllm_inference.py diff --git a/example/post_process/example_lora_gptq_vllm_inference.py b/example/post_process/example_lora_gptq_vllm_inference.py new file mode 100644 index 0000000..1b16903 --- /dev/null +++ b/example/post_process/example_lora_gptq_vllm_inference.py @@ -0,0 +1,176 @@ +""" + +Example: Quantize a model with OneComp and run inference with vLLM + LoRA + +Performs the following steps: + 1. Quantize with GPTQ (4-bit, groupsize=128) + LoRA SFT on OneCompression knowledge data + 2. Save the quantized model + 3. Load the quantized model with vLLM's offline LLM interface + 4. Generate text + +Requirements: + pip install vllm + +Copyright 2025-2026 Fujitsu Ltd. + +Author: Keiji Kimura + +""" + +import gc + +import json +import os +import torch +from onecomp import CalibrationConfig, GPTQ, ModelConfig, PostProcessLoraSFT, Runner, setup_logger +from transformers.tokenization_utils_base import PreTrainedTokenizerBase as _PTBase + +try: + from vllm import LLM, SamplingParams + from vllm.lora.request import LoRARequest +except ImportError as e: + raise SystemExit( + "This example requires vllm to be installed. " + "Install with: uv sync --extra vllm" + ) from e + +if not hasattr(_PTBase, "all_special_tokens_extended"): + _PTBase.all_special_tokens_extended = property( + lambda self: list(self.all_special_tokens) + ) + +def _ensure_fast_tokenizer_class(save_dir: str) -> None: + """Rewrite tokenizer_config.json so vLLM loads the fast tokenizer.""" + tok_json = os.path.join(save_dir, "tokenizer.json") + tok_cfg = os.path.join(save_dir, "tokenizer_config.json") + if not (os.path.isfile(tok_json) and os.path.isfile(tok_cfg)): + return + + with open(tok_cfg, "r", encoding="utf-8") as f: + cfg = json.load(f) + + current = cfg.get("tokenizer_class") + if current and current.endswith("Fast"): + return + + slow_to_fast = { + "LlamaTokenizer": "LlamaTokenizerFast", + "CodeLlamaTokenizer": "CodeLlamaTokenizerFast", + } + replacement = slow_to_fast.get(current, "LlamaTokenizerFast") + cfg["tokenizer_class"] = replacement + + with open(tok_cfg, "w", encoding="utf-8") as f: + json.dump(cfg, f, indent=2, ensure_ascii=False) + print( + f"Patched {tok_cfg}: tokenizer_class {current!r} -> {replacement!r} " + "(forces vLLM to use the fast tokenizer)" + ) + + +def main(): + setup_logger() + + # Step 1: Quantize with GPTQ + LoRA SFT + save_dir = "./TinyLlama-1.1B-Chat-gptq-4bit-lora" + lora_path = os.path.join(save_dir, "lora_adapter") + knowledge_data = str(os.path.join(os.path.dirname(__file__), "onecomp_knowledge.jsonl")) + + model_config = ModelConfig( + model_id="TinyLlama/TinyLlama-1.1B-Chat-v1.0", + device="cuda:0", + ) + quantizer = GPTQ(wbits=4, groupsize=128) + post_process = PostProcessLoraSFT( + data_files=knowledge_data, + max_length=256, + epochs=50, + batch_size=2, + gradient_accumulation_steps=1, + lr=3e-4, + lora_r=16, + lora_alpha=32, + logging_steps=5, + ) + calibration_config = CalibrationConfig( + max_length=128, + num_calibration_samples=16, + batch_size=8 + ) + runner = Runner( + model_config=model_config, + quantizer=quantizer, + calibration_config=calibration_config, + post_processes=[post_process], + qep=False, + ) + # NOTE: The calibration settings above are kept compact so the demo + # runs fast and may be insufficient for real quantisation. For + # higher quality, prefer the CalibrationConfig() defaults + # (max_length=2048, num_calibration_samples=512). + # For qep=False runs with large calibration data, also pass + # ``batch_size`` as a CalibrationConfig argument, e.g. + # CalibrationConfig( + # max_length=2048, + # num_calibration_samples=512, + # batch_size=128, + # ) + # so that Runner.quantize_with_calibration_chunked runs instead of + # a single all-at-once forward pass. + runner.run() + + # Step 2: Save the quantized model + runner.save_quantized_model(save_dir) + print(f"\nSaved GPTQ + LoRA model (base + adapter sidecar) to: {save_dir}") + + # Free GPU memory used by quantization before loading vLLM + del runner + gc.collect() + torch.cuda.empty_cache() + + # Step 3: Load the quantized model with vLLM and enable LoRA + # _ensure_fast_tokenizer_class(save_dir) + llm = LLM( + model=save_dir, + max_model_len=512, + dtype="float16", + enforce_eager=True, + gpu_memory_utilization=0.55, + max_num_batched_tokens=512, + enable_prefix_caching=False, + enable_lora=True, + max_lora_rank=16, + ) + + lora_request = LoRARequest( + lora_name="gptq_lora_sft", + lora_int_id=1, + lora_path=lora_path, + ) + + # Step 4: Generate text + prompts = [ + "Q: What is OneCompression?\nA:", + ] + sampling_params = SamplingParams(max_tokens=64, temperature=0.0) + + base_outputs = llm.generate( + prompts, + sampling_params, + ) + + lora_outputs = llm.generate( + prompts, + sampling_params, + lora_request=lora_request, + ) + + print("vLLM inference — GPTQ base vs GPTQ + LoRA SFT adapter") + for base_out, lora_out in zip(base_outputs, lora_outputs): + print(f"\nPrompt: {base_out.prompt}") + print(f" base only : {base_out.outputs[0].text}") + print(f" base + LoRA: {lora_out.outputs[0].text}") + + +if __name__ == "__main__": + main() From 27566cf2d0c4048ffb1580415b02073964df686b Mon Sep 17 00:00:00 2001 From: cm_ysmz Date: Thu, 14 May 2026 14:35:13 +0900 Subject: [PATCH 11/82] [refactor] Delete working files --- .../example_gptq_lora_vllm_inference.py | 234 ------------------ ...a_sft_save_load_hf_format_with_low_vram.py | 159 ------------ 2 files changed, 393 deletions(-) delete mode 100644 work/20260414_katari_LoRA_Save_Load/example_gptq_lora_vllm_inference.py delete mode 100644 work/20260414_katari_LoRA_Save_Load/example_lora_sft_save_load_hf_format_with_low_vram.py diff --git a/work/20260414_katari_LoRA_Save_Load/example_gptq_lora_vllm_inference.py b/work/20260414_katari_LoRA_Save_Load/example_gptq_lora_vllm_inference.py deleted file mode 100644 index 4861001..0000000 --- a/work/20260414_katari_LoRA_Save_Load/example_gptq_lora_vllm_inference.py +++ /dev/null @@ -1,234 +0,0 @@ -""" - -Example: Quantize with GPTQ + LoRA SFT, save, and serve via vLLM with LoRA - -Performs the following steps: - 1. Quantize TinyLlama with GPTQ (4-bit, groupsize=128) - 2. Apply ``PostProcessLoraSFT`` (WikiText-2) to fine-tune LoRA adapters - on top of the frozen GPTQ base - 3. Save via ``runner.save_quantized_model(save_dir)``. This writes: - - ``model.safetensors`` / ``config.json`` (base GPTQ, HF format) - - ``lora_adapter/adapter_model.safetensors`` (PEFT-format LoRA) - - ``lora_adapter/adapter_config.json`` (PEFT-format LoRA config) - The adapter lives in the ``lora_adapter/`` subdirectory because vLLM - globs ``*.safetensors`` at the top level of the model directory to load - base weights, and would otherwise try to load the adapter file as part - of the base model. - 4. Load the base model with vLLM and apply the LoRA adapter via - ``LoRARequest(..., lora_path=os.path.join(save_dir, "lora_adapter"))``. - 5. Generate text with and without the LoRA adapter and print both - outputs side by side. - -Requirements: - - pip install vllm (>= 0.5 recommended; older releases may lack - LoRA-over-GPTQ support) - - The ``mixed_gptq`` quantization method is registered by OneComp's vLLM - plugin (see ``vllm_plugins/gptq``), which is installed automatically - with this package. - -vLLM LoRA constraints ---------------------- - - ``max_lora_rank`` must be >= the adapter's ``r`` (from - ``adapter_config.json``). - - On Llama-style models, q/k/v (and gate/up) projections are fused by - vLLM into ``qkv_proj`` / ``gate_up_proj``. vLLM's LoRA loader - concatenates the per-projection adapter weights automatically, provided - all three projections share the same ``r`` and ``lora_alpha`` — - ``PostProcessLoraSFT`` always uses a single global rank/alpha, so this - invariant holds. - - ``lm_head`` is excluded from LoRA targets by ``PostProcessLoraSFT``, so - ``lora_extra_vocab_size=0`` is safe. - -Copyright 2025-2026 Fujitsu Ltd. - -Author: Keiji Kimura - -""" - -import gc -import json -import os - -import torch - -# --------------------------------------------------------------------------- -# transformers >= 5.x / vLLM compatibility shim -# --------------------------------------------------------------------------- -# OneComp requires transformers>=5.3, which removed the -# `all_special_tokens_extended` property from the tokenizer base class. -# vLLM (<= 0.11.x) still accesses it in ``get_cached_tokenizer`` and crashes -# with ``AttributeError: LlamaTokenizer has no attribute -# all_special_tokens_extended``. Re-add the property as an alias for -# ``all_special_tokens`` so vLLM's cache build succeeds. Must run BEFORE vLLM -# imports so subclasses inherit the attribute. -from transformers.tokenization_utils_base import PreTrainedTokenizerBase as _PTBase - -if not hasattr(_PTBase, "all_special_tokens_extended"): - _PTBase.all_special_tokens_extended = property( - lambda self: list(self.all_special_tokens) - ) - -from onecomp import ( - CalibrationConfig, - GPTQ, - ModelConfig, - PostProcessLoraSFT, - Runner, - setup_logger, -) - -try: - from vllm import LLM, SamplingParams - from vllm.lora.request import LoRARequest - from vllm.model_executor.layers.quantization import register_quantization_config # noqa: F401 -except ImportError as e: - raise SystemExit( - "This example requires vllm>=0.6.3 " - "(needs both `vllm.lora.request.LoRARequest` and " - "`register_quantization_config` used by OneComp's mixed_gptq plugin). " - "Install with: pip install -U 'vllm>=0.6.3'" - ) from e - - -def _ensure_fast_tokenizer_class(save_dir: str) -> None: - """Rewrite tokenizer_config.json so vLLM loads the fast tokenizer. - - Some upstream Llama-family checkpoints (TinyLlama included) ship with - ``"tokenizer_class": "LlamaTokenizer"`` in ``tokenizer_config.json``. - Recent ``transformers`` releases have removed ``all_special_tokens_extended`` - from the slow ``LlamaTokenizer``, which causes vLLM's tokenizer cache to - crash with ``AttributeError: LlamaTokenizer has no attribute - all_special_tokens_extended``. When ``tokenizer.json`` is present the fast - variant is available, so we patch the class name to force vLLM down that - path. - """ - tok_json = os.path.join(save_dir, "tokenizer.json") - tok_cfg = os.path.join(save_dir, "tokenizer_config.json") - if not (os.path.isfile(tok_json) and os.path.isfile(tok_cfg)): - return - - with open(tok_cfg, "r", encoding="utf-8") as f: - cfg = json.load(f) - - current = cfg.get("tokenizer_class") - if current and current.endswith("Fast"): - return - - # Map slow → fast for Llama-family tokenizers. Extend this table if other - # architectures hit the same issue in the future. - slow_to_fast = { - "LlamaTokenizer": "LlamaTokenizerFast", - "CodeLlamaTokenizer": "CodeLlamaTokenizerFast", - } - replacement = slow_to_fast.get(current, "LlamaTokenizerFast") - cfg["tokenizer_class"] = replacement - - with open(tok_cfg, "w", encoding="utf-8") as f: - json.dump(cfg, f, indent=2, ensure_ascii=False) - print( - f"Patched {tok_cfg}: tokenizer_class {current!r} -> {replacement!r} " - "(forces vLLM to use the fast tokenizer)" - ) - - -def main(): - setup_logger() - - # ================================================================ - # Step 1: Quantize + LoRA SFT and save - # ================================================================ - save_dir = "./TinyLlama-1.1B-gptq-4bit-lora" - - model_config = ModelConfig( - model_id="TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T", - device="cuda:0", - ) - quantizer = GPTQ(wbits=4, groupsize=128) - post_process = PostProcessLoraSFT( - dataset_name="wikitext", - dataset_config_name="wikitext-2-raw-v1", - train_split="train", - text_column="text", - max_train_samples=128, - max_length=256, - epochs=2, - batch_size=1, - gradient_accumulation_steps=8, - lr=1e-4, - lora_r=8, - lora_alpha=16, - logging_steps=5, - ) - - runner = Runner( - model_config=model_config, - quantizer=quantizer, - calibration_config=CalibrationConfig(max_length=128, num_calibration_samples=16, batch_size=8), - post_processes=[post_process], - ) - runner.run() - runner.save_quantized_model(save_dir) - print(f"\nSaved GPTQ+LoRA model (base + adapter sidecar) to: {save_dir}") - - # Free GPU memory used by quantization / training before loading vLLM. - del runner - gc.collect() - torch.cuda.empty_cache() - - # ================================================================ - # Step 2: Load base GPTQ with vLLM, enable LoRA - # ================================================================ - # Work around the LlamaTokenizer / transformers compatibility issue - # (see _ensure_fast_tokenizer_class for details). - _ensure_fast_tokenizer_class(save_dir) - - # ``max_lora_rank`` must be >= the saved ``r`` (16 here). - llm = LLM( - model=save_dir, - enable_lora=True, - max_lora_rank=16, - max_loras=1, - max_model_len=512, - dtype="float16", - enforce_eager=True, - gpu_memory_utilization=0.55, # VRAM 8GB 制約ありのため削減 - max_num_batched_tokens=512, # 同上 - enable_prefix_caching=False, # 同上 - ) - - # The adapter sidecar lives in the lora_adapter/ subdirectory to avoid - # colliding with vLLM's base-model safetensors glob. - lora_request = LoRARequest( - lora_name="gptq_sft", - lora_int_id=1, - lora_path=os.path.join(save_dir, "lora_adapter"), - ) - - prompts = [ - "Explain what post-training quantization is in one sentence:", - "Fujitsu is", - ] - sampling_params = SamplingParams(max_tokens=64, temperature=0.0) - - # ----- (a) Base GPTQ only (LoRA disabled) ----- - base_outputs = llm.generate(prompts, sampling_params) - - # ----- (b) Base GPTQ + LoRA adapter ----- - lora_outputs = llm.generate( - prompts, - sampling_params, - lora_request=lora_request, - ) - - print("\n" + "=" * 70) - print("vLLM inference — GPTQ base vs GPTQ + LoRA SFT adapter") - print("=" * 70) - for base_out, lora_out in zip(base_outputs, lora_outputs): - print(f"\nPrompt: {base_out.prompt}") - print(f" base only : {base_out.outputs[0].text}") - print(f" base + LoRA: {lora_out.outputs[0].text}") - - -if __name__ == "__main__": - main() - \ No newline at end of file diff --git a/work/20260414_katari_LoRA_Save_Load/example_lora_sft_save_load_hf_format_with_low_vram.py b/work/20260414_katari_LoRA_Save_Load/example_lora_sft_save_load_hf_format_with_low_vram.py deleted file mode 100644 index 1e40c06..0000000 --- a/work/20260414_katari_LoRA_Save_Load/example_lora_sft_save_load_hf_format_with_low_vram.py +++ /dev/null @@ -1,159 +0,0 @@ -""" -Example: GPTQ quantization + LoRA SFT + save/load (low VRAM settings) - -End-to-end demonstration of the LoRA SFT post-process workflow with -settings chosen to reduce GPU memory usage on smaller GPUs: - 1. Quantize TinyLlama with GPTQ 4-bit (groupsize=128) - 2. Apply LoRA SFT post-process (WikiText-2) with reduced VRAM settings - 3. Evaluate PPL (original vs quantized+LoRA) - 4. Save via save_quantized_model() - writes HF-compatible safetensors - plus a PEFT-format LoRA adapter sidecar - (``adapter_model.safetensors`` + ``adapter_config.json``) - 5. Load via load_quantized_model() - the sidecar is auto-detected and - matching GPTQLinear layers are re-wrapped with LoRAGPTQLinear - 6. Generate text with the loaded model to verify it works - -Recommended low-VRAM changes compared with the standard example: - - batch_size=1 - - gradient_accumulation_steps=8 - - max_length=256 - - max_train_samples=128 - - lora_r=8 - -Copyright 2025-2026 Fujitsu Ltd. - -Author: Keiji Kimura - -Usage: - python example/post_process/example_lora_sft_low_vram.py -""" - -import torch - -from onecomp import ( - CalibrationConfig, - GPTQ, - ModelConfig, - PostProcessLoraSFT, - Runner, - load_quantized_model, - setup_logger, -) - -setup_logger() - -MODEL_ID = "TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T" -SAVE_DIR = "./tinyllama_gptq4_lora_low_vram" -PROMPT = "Fujitsu is" - - -def generate_text(model, tokenizer, prompt, device, max_new_tokens=64): - """Generate text from a prompt using the model.""" - inputs = tokenizer(prompt, return_tensors="pt").to(device) - with torch.no_grad(): - outputs = model.generate( - **inputs, - max_new_tokens=max_new_tokens, - do_sample=False, - temperature=1.0, - repetition_penalty=1.2, - ) - generated = tokenizer.decode(outputs[0], skip_special_tokens=True) - return generated - - -# ================================================================ -# Step 1: Quantize + LoRA SFT via Runner -# ================================================================ -print("=" * 70) -print("Step 1: Quantize TinyLlama (GPTQ 4-bit) + LoRA SFT (WikiText-2, low VRAM)") -print("=" * 70) - -model_config = ModelConfig(model_id=MODEL_ID, device="cuda:0") -gptq = GPTQ(wbits=4, groupsize=128) - -post_process = PostProcessLoraSFT( - dataset_name="wikitext", - dataset_config_name="wikitext-2-raw-v1", - train_split="train", - text_column="text", - max_train_samples=128, - max_length=256, - epochs=2, - batch_size=1, - gradient_accumulation_steps=8, - lr=1e-4, - lora_r=8, - lora_alpha=16, - logging_steps=5, -) - -runner = Runner( - model_config=model_config, - quantizer=gptq, - calibration_config=CalibrationConfig(max_length=128, num_calibration_samples=16, batch_size=8), - post_processes=[post_process], -) -runner.run() - -# ================================================================ -# Step 2: Evaluate PPL (original vs quantized+LoRA) -# ================================================================ -print("\n" + "=" * 70) -print("Step 2: Evaluate PPL") -print("=" * 70) - -original_ppl, _, quantized_ppl = runner.calculate_perplexity( - original_model=True, - quantized_model=True, -) - -print(f" Original model PPL: {original_ppl:.4f}") -print(f" Quantized + LoRA SFT model PPL: {quantized_ppl:.4f}") - -# ================================================================ -# Step 3: Save the LoRA-applied model (HF safetensors + adapter sidecar) -# ================================================================ -print("\n" + "=" * 70) -print(f"Step 3: Saving LoRA-applied model to {SAVE_DIR}") -print("=" * 70) - -runner.save_quantized_model(SAVE_DIR) -print(f"Model saved to: {SAVE_DIR}") -print( - " - model.safetensors / config.json : base GPTQ model (HF-compatible)\n" - " - adapter_model.safetensors : PEFT-format LoRA adapter\n" - " - adapter_config.json : PEFT-format adapter config" -) - -del runner -torch.cuda.empty_cache() - -# ================================================================ -# Step 4: Load the saved model -# ================================================================ -print("\n" + "=" * 70) -print(f"Step 4: Loading model from {SAVE_DIR}") -print("=" * 70) - -loaded_model, loaded_tokenizer = load_quantized_model(SAVE_DIR) -print(f"Loaded model type : {type(loaded_model).__name__}") -print(f"Loaded model device: {next(loaded_model.parameters()).device}") - -# ================================================================ -# Step 5: Generate text with the loaded model -# ================================================================ -print("\n" + "=" * 70) -print("Step 5: Generate text with loaded model") -print("=" * 70) - -loaded_text = generate_text( - loaded_model, - loaded_tokenizer, - PROMPT, - device=next(loaded_model.parameters()).device, -) - -print(f"\nPrompt : {PROMPT}") -print(f"Generated: {loaded_text}") -print("=" * 70) \ No newline at end of file From 3ec1166e3c2564fd89fd8b71c220b7d0f2aa5ad7 Mon Sep 17 00:00:00 2001 From: cm_ysmz Date: Thu, 14 May 2026 15:46:49 +0900 Subject: [PATCH 12/82] [update] Update CHANGELOG.md (LoRA Adapter Sidecar Save/Load) --- CHANGELOG.md | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 29c237b..749cf9c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,23 @@ # Change log +## [v1.1.0+feature/lora-save-load-vllm-infer_jointq] 2026-05-14 + +### LoRA Adapter Sidecar Save/Load + +- Updated `Runner.save_quantized_model()` to save LoRA adapters as a PEFT sidecar in `lora_adapter/` (`adapter_model.safetensors`, `adapter_config.json`) when `LoRAGPTQLinear` modules are present, keeping the base quantized model and adapters separated (`onecomp/runner.py`) +- Updated `QuantizedModelLoader.load_quantized_model()` to auto-detect LoRA sidecars and re-wrap matching layers with `LoRAGPTQLinear` during load for save/load round-trips (`onecomp/quantized_model_loader.py`) + +### Examples + +- Added `example/post_process/example_lora_sft_save_load_hf_format.py`: end-to-end GPTQ + LoRA SFT workflow covering save/load in HF format and before/after generation comparison +- Added `example/post_process/example_lora_sft_knowledge_save_load_hf_format.py`: knowledge-injection LoRA SFT example using `onecomp_knowledge.jsonl`, including save/load round-trip and output comparison +- Added `example/post_process/example_lora_sft_knowledge_save_load_hf_format_jointq.py`: JointQ counterpart of the knowledge-injection LoRA SFT workflow +- Added `example/post_process/example_lora_gptq_vllm_inference.py`: end-to-end GPTQ + LoRA SFT workflow validating vLLM inference with saved LoRA adapters, including GPTQ-only vs GPTQ+LoRA output comparison + +### Tests + +- Added a JointQ + `PostProcessLoraSFT` smoke test (conditional on JointQ/CUDA availability) in `tests/onecomp/post_process/test_post_process_lora_sft.py` + ## [v1.1.0] 2026-04-16 ### Gemma 3 / Gemma 4 & VLM Support From 7517c18e6e92e751174d9e956aeb3aaedf6061ae Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Kimura=2C=20Keiji/=E6=9C=A8=E6=9D=91=20=E5=9C=AD=E5=85=90?= Date: Mon, 18 May 2026 13:31:03 +0000 Subject: [PATCH 13/82] Raise clear error for unsupported QEP quantizers --- CHANGELOG.md | 8 +++ onecomp/quantizer/_quantizer.py | 1 + onecomp/quantizer/autobit/_autobit.py | 5 ++ onecomp/quantizer/jointq/_jointq.py | 3 + onecomp/runner.py | 9 +++ tests/onecomp/test_runner_check.py | 100 ++++++++++++++++++++++++++ 6 files changed, 126 insertions(+) create mode 100644 tests/onecomp/test_runner_check.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 1fbe04d..fdddc34 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,14 @@ ## [v1.1.1] 2026-05-07 +## Bug Fixes + +- 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). + +## 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()`. + ## [v1.1.0] 2026-04-16 ### Gemma 3 / Gemma 4 & VLM Support diff --git a/onecomp/quantizer/_quantizer.py b/onecomp/quantizer/_quantizer.py index 24ced33..95e1372 100644 --- a/onecomp/quantizer/_quantizer.py +++ b/onecomp/quantizer/_quantizer.py @@ -174,6 +174,7 @@ class Quantizer(metaclass=ABCMeta): flag_calibration: bool = False flag_hessian: bool = False flag_xtx: bool = False # Whether X^T X is needed (e.g., JointQ) + flag_qep_supported: bool = True def __post_init__(self): """__post_init__ method""" diff --git a/onecomp/quantizer/autobit/_autobit.py b/onecomp/quantizer/autobit/_autobit.py index a5a6e60..8e59e46 100644 --- a/onecomp/quantizer/autobit/_autobit.py +++ b/onecomp/quantizer/autobit/_autobit.py @@ -393,6 +393,11 @@ def _sync_flags(self): self.flag_calibration = any(q.flag_calibration for q in self.quantizers) self.flag_hessian = any(q.flag_hessian for q in self.quantizers) self.flag_xtx = any(q.flag_xtx for q in self.quantizers) + # AutoBit supports QEP only when *all* candidate quantizers support it + # (the per-layer assignment may dispatch to any child quantizer). + self.flag_qep_supported = all( + q.flag_qep_supported for q in self.quantizers + ) def _validate_manual_fused_consistency(self): """Check that manual keyword rules don't split fused groups.""" diff --git a/onecomp/quantizer/jointq/_jointq.py b/onecomp/quantizer/jointq/_jointq.py index 1ea5e17..fe7e09a 100644 --- a/onecomp/quantizer/jointq/_jointq.py +++ b/onecomp/quantizer/jointq/_jointq.py @@ -250,6 +250,9 @@ class JointQ(Quantizer): flag_calibration: bool = True flag_hessian: bool = False flag_xtx: bool = True + # JointQ does not yet support the generic QEP pipeline. + # Planned for a future release. + flag_qep_supported: bool = False hessian_dtype: torch.dtype = torch.float64 # Parameters for the JointQ quantizer diff --git a/onecomp/runner.py b/onecomp/runner.py index fee4068..df5788d 100644 --- a/onecomp/runner.py +++ b/onecomp/runner.py @@ -291,6 +291,15 @@ def check(self): ) if self.multi_gpu and not self.quantizer.flag_calibration: raise ValueError("'multi_gpu' requires a quantizer with flag_calibration=True.") + if self.qep and not self.quantizer.flag_qep_supported: + raise ValueError( + f"Quantizer '{type(self.quantizer).__name__}' " + f"(or one of its candidate quantizers) does not support " + f"QEP (Quantization Error Propagation). " + f"Set qep=False, or use a QEP-compatible quantizer " + f"(e.g., GPTQ, DBF, AutoBitQuantizer with " + f"QEP-compatible candidates)." + ) # Cross-validate calibration_dataset when AutoBitQuantizer is used quantizer = self.quantizer or (self.quantizers[0] if self.quantizers else None) diff --git a/tests/onecomp/test_runner_check.py b/tests/onecomp/test_runner_check.py new file mode 100644 index 0000000..154226f --- /dev/null +++ b/tests/onecomp/test_runner_check.py @@ -0,0 +1,100 @@ +"""Unit tests for ``Runner.check()`` parameter validation. + +These tests exercise only the configuration validation path +(``Runner.check()``) and therefore do not require a GPU or model +loading. + +Copyright 2025-2026 Fujitsu Ltd. +""" + +import pytest + +from onecomp import CalibrationConfig, ModelConfig, Runner +from onecomp.quantizer.autobit import AutoBitQuantizer +from onecomp.quantizer.gptq import GPTQ +from onecomp.quantizer.jointq import JointQ + + +def _model_config() -> ModelConfig: + return ModelConfig( + model_id="TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T", + device="cuda:0", + ) + + +class TestRunnerCheckQEPSupport: + """``qep=True`` must raise a clear error when the quantizer does + not support the generic QEP pipeline. + """ + + def test_jointq_with_qep_true_raises_clear_error(self): + """JointQ + qep=True should raise a clear ValueError, not an + obscure error from deep inside the QEP runtime. + """ + runner = Runner( + model_config=_model_config(), + quantizer=JointQ(bits=4, group_size=128), + calibration_config=CalibrationConfig(max_length=128, num_calibration_samples=8), + qep=True, + ) + + with pytest.raises(ValueError, match=r"JointQ.*does not support QEP"): + runner.check() + + def test_jointq_with_qep_false_passes_check(self): + """JointQ + qep=False is the supported configuration.""" + runner = Runner( + model_config=_model_config(), + quantizer=JointQ(bits=4, group_size=128), + calibration_config=CalibrationConfig(max_length=128, num_calibration_samples=8), + qep=False, + ) + runner.check() + + def test_gptq_with_qep_true_passes_check(self): + """GPTQ supports QEP, so check() must pass.""" + runner = Runner( + model_config=_model_config(), + quantizer=GPTQ(wbits=4, groupsize=128), + calibration_config=CalibrationConfig(max_length=128, num_calibration_samples=8), + qep=True, + ) + runner.check() + + def test_autobit_with_jointq_candidate_and_qep_true_raises(self): + """AutoBit with a JointQ candidate must also raise on qep=True. + + ``AutoBitQuantizer.flag_qep_supported`` is True only when *all* + candidate quantizers support QEP, so a JointQ candidate must + propagate the unsupported state. + """ + autobit = AutoBitQuantizer( + quantizers=[GPTQ(wbits=4), JointQ(bits=2)], + target_bit=3.0, + ) + runner = Runner( + model_config=_model_config(), + quantizer=autobit, + calibration_config=CalibrationConfig(max_length=128, num_calibration_samples=8), + qep=True, + ) + + with pytest.raises( + ValueError, + match=r"AutoBitQuantizer.*does not support QEP", + ): + runner.check() + + def test_autobit_with_only_gptq_candidates_and_qep_true_passes(self): + """AutoBit with only QEP-compatible candidates must pass.""" + autobit = AutoBitQuantizer( + quantizers=[GPTQ(wbits=4), GPTQ(wbits=2)], + target_bit=3.0, + ) + runner = Runner( + model_config=_model_config(), + quantizer=autobit, + calibration_config=CalibrationConfig(max_length=128, num_calibration_samples=8), + qep=True, + ) + runner.check() From 1470e5c0116362c498975c02258472b2ca59f7fd Mon Sep 17 00:00:00 2001 From: aki916f Date: Tue, 19 May 2026 13:50:32 +0900 Subject: [PATCH 14/82] Refactoring: quantization progress logs with ETA (#15) * refactoring : QuantizationProgressTracker * update CHANGELOG.md --------- Co-authored-by: FKKimura <50981196+FKKimura@users.noreply.github.com> --- CHANGELOG.md | 15 +++++++++--- onecomp/qep/_quantize_with_qep.py | 10 ++++---- onecomp/qep/_quantize_with_qep_arch.py | 10 ++++---- onecomp/runner.py | 23 ++++++++----------- .../runner_methods/chunked_quantization.py | 12 ++++------ .../runner_methods/multi_gpu_quantization.py | 14 +++++------ onecomp/utils/quantization_progress.py | 4 ++-- 7 files changed, 42 insertions(+), 46 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index fdddc34..a657fa7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,15 +1,24 @@ # Change log -## [v1.1.1] 2026-05-07 +## [v1.1.1] 2026-05-19 -## Bug Fixes +### 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 + +### Bug Fixes - 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). -## Tests +### 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()`. +### 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 diff --git a/onecomp/qep/_quantize_with_qep.py b/onecomp/qep/_quantize_with_qep.py index e0ba901..b22c194 100644 --- a/onecomp/qep/_quantize_with_qep.py +++ b/onecomp/qep/_quantize_with_qep.py @@ -27,6 +27,7 @@ from onecomp.qep._qep_config import QEPConfig from onecomp.quantizer._quantizer import Quantizer from onecomp.utils import capture_input_activations +from onecomp.utils.quantization_progress import QuantizationProgressTracker logger = getLogger(__name__) @@ -37,7 +38,7 @@ def run_quantize_with_qep( qep_config: QEPConfig, calibration_config: CalibrationConfig, *, - quantization_progress: bool = True, + report_progress: bool = True, ): """Run quantization with Quantization Error Propagation (QEP). @@ -53,7 +54,7 @@ def run_quantize_with_qep( qep_config (QEPConfig): Configuration for QEP (percdamp, perccorr, exclude_layer_keywords). calibration_config (CalibrationConfig): Calibration parameters. - quantization_progress (bool): When True, log ``[progress]`` with ETA per layer. + report_progress (bool): When True, log ``[progress]`` with ETA per layer. """ model = model_config.load_model() @@ -84,10 +85,7 @@ def run_quantize_with_qep( logger.info("Quantizing the model using %s", quantizer.name) progress = None - if quantization_progress: - # pylint: disable-next=import-outside-toplevel - from onecomp.utils.quantization_progress import QuantizationProgressTracker - + if report_progress: progress = QuantizationProgressTracker( logger, len(quantizer.module_to_name), diff --git a/onecomp/qep/_quantize_with_qep_arch.py b/onecomp/qep/_quantize_with_qep_arch.py index 6f51fcb..a10dc4a 100644 --- a/onecomp/qep/_quantize_with_qep_arch.py +++ b/onecomp/qep/_quantize_with_qep_arch.py @@ -33,6 +33,7 @@ move_kwargs_to_device, expand_kwargs_batch, ) +from onecomp.utils.quantization_progress import QuantizationProgressTracker logger = getLogger(__name__) @@ -263,7 +264,7 @@ def run_quantize_with_qep_arch( qep_config: QEPConfig, calibration_config: CalibrationConfig, *, - quantization_progress: bool = True, + report_progress: bool = True, ): """Run architecture-aware quantization with QEP. @@ -280,7 +281,7 @@ def run_quantize_with_qep_arch( qep_config (QEPConfig): Configuration for QEP (percdamp, perccorr, exclude_layer_keywords). calibration_config (CalibrationConfig): Calibration parameters. - quantization_progress (bool): When True, log ``[progress]`` with ETA per target layer. + report_progress (bool): When True, log ``[progress]`` with ETA per target layer. """ @@ -321,10 +322,7 @@ def run_quantize_with_qep_arch( } progress = None - if quantization_progress: - # pylint: disable-next=import-outside-toplevel - from onecomp.utils.quantization_progress import QuantizationProgressTracker - + if report_progress: progress = QuantizationProgressTracker( logger, len(remaining_targets), diff --git a/onecomp/runner.py b/onecomp/runner.py index 6e9b0de..f7153d4 100644 --- a/onecomp/runner.py +++ b/onecomp/runner.py @@ -28,6 +28,7 @@ from .quantizer.autobit import AutoBitQuantizer from .utils import calculate_accuracy as calc_accuracy from .utils import calculate_perplexity as calc_perplexity +from .utils.quantization_progress import QuantizationProgressTracker from .log import setup_logger @@ -86,7 +87,7 @@ def __init__( multi_gpu=False, gpu_ids=None, post_processes=None, - quantization_progress=True, + report_progress=True, ): """__init__ method @@ -131,7 +132,7 @@ def __init__( a quantized model on CPU (built via ``create_quantized_model``) and may modify it in-place. Processes are executed in order. Default is None. - quantization_progress (bool): + report_progress (bool): When ``True`` (default), emit ``[progress]`` log lines with completed steps, elapsed time, and a linear ETA estimate during long quantization (calibration, chunked, multi-GPU, @@ -221,7 +222,7 @@ def __init__( self.lpcd_config = None if lpcd: self.lpcd_config = lpcd_config if lpcd_config is not None else LPCDConfig() - self.quantization_progress = quantization_progress + self.report_progress = report_progress def check(self): """Check the settings @@ -602,10 +603,7 @@ def quantize_with_calibration(self): # Register hooks to all linear layers handles = [] progress = None - if self.quantization_progress: - # pylint: disable-next=import-outside-toplevel - from .utils.quantization_progress import QuantizationProgressTracker - + if self.report_progress: progress = QuantizationProgressTracker( logger, len(self.quantizer.module_to_name), @@ -664,7 +662,7 @@ def quantize_with_calibration_chunked(self): model_config=self.model_config, quantizers=self.quantizers if self.quantizers is not None else [self.quantizer], calibration_config=self.calibration_config, - quantization_progress=self.quantization_progress, + report_progress=self.report_progress, ) def quantize_with_calibration_on_multi_gpu(self): @@ -693,7 +691,7 @@ def quantize_with_calibration_on_multi_gpu(self): quantizer=self.quantizer, calibration_config=self.calibration_config, gpu_ids=self.gpu_ids, - quantization_progress=self.quantization_progress, + report_progress=self.report_progress, ) # Store results in quantizer.results @@ -721,10 +719,7 @@ def quantize_without_calibration(self): self.quantizer.name, ) progress = None - if self.quantization_progress: - # pylint: disable-next=import-outside-toplevel - from .utils.quantization_progress import QuantizationProgressTracker - + if self.report_progress: progress = QuantizationProgressTracker( logger, len(self.quantizer.module_to_name), @@ -754,7 +749,7 @@ def quantize_with_qep(self): quantizer=self.quantizer, qep_config=self.qep_config, calibration_config=self.calibration_config, - quantization_progress=self.quantization_progress, + report_progress=self.report_progress, ) if self.qep_config.general: diff --git a/onecomp/runner_methods/chunked_quantization.py b/onecomp/runner_methods/chunked_quantization.py index 0024133..456ee66 100644 --- a/onecomp/runner_methods/chunked_quantization.py +++ b/onecomp/runner_methods/chunked_quantization.py @@ -34,6 +34,7 @@ from onecomp.calibration import CalibrationConfig, prepare_calibration_dataset from onecomp.model_config import ModelConfig from onecomp.quantizer._quantizer import Quantizer, QuantizationResult +from onecomp.utils.quantization_progress import QuantizationProgressTracker logger = getLogger(__name__) @@ -48,7 +49,7 @@ def run_chunked_quantization( quantizers: List[Quantizer], calibration_config: CalibrationConfig, *, - quantization_progress: bool = True, + report_progress: bool = True, ): """Run quantization for large-scale calibration data. @@ -61,7 +62,7 @@ def run_chunked_quantization( quantizers (list[Quantizer]): List of quantizers. Each quantizer must have flag_hessian=True or flag_xtx=True. calibration_config (CalibrationConfig): Calibration parameters. - quantization_progress (bool): When True, log ``[progress]`` lines with ETA + report_progress (bool): When True, log ``[progress]`` lines with ETA for X^T X chunks and per-layer quantization. Note: @@ -120,10 +121,7 @@ def run_chunked_quantization( num_groups = (len(all_layers) + num_layers_per_group - 1) // num_layers_per_group layer_progress = None - if quantization_progress: - # pylint: disable-next=import-outside-toplevel - from onecomp.utils.quantization_progress import QuantizationProgressTracker - + if report_progress: layer_progress = QuantizationProgressTracker( logger, len(all_layers) * len(quantizers), @@ -145,7 +143,7 @@ def run_chunked_quantization( ) chunk_progress = None - if quantization_progress: + if report_progress: num_chunks = (total_samples + calibration_batch_size - 1) // calibration_batch_size chunk_progress = QuantizationProgressTracker( logger, diff --git a/onecomp/runner_methods/multi_gpu_quantization.py b/onecomp/runner_methods/multi_gpu_quantization.py index 1f72695..8853835 100644 --- a/onecomp/runner_methods/multi_gpu_quantization.py +++ b/onecomp/runner_methods/multi_gpu_quantization.py @@ -23,6 +23,7 @@ from onecomp.calibration import CalibrationConfig from onecomp.quantizer._quantizer import QuantizationResult from onecomp.utils import check_activations +from onecomp.utils.quantization_progress import QuantizationProgressTracker logger = getLogger(__name__) @@ -178,7 +179,7 @@ def run_quantization_phase( quantizer, gpu_ids: List[int], *, - quantization_progress: bool = True, + report_progress: bool = True, ) -> Dict[str, Dict]: """Phase 2: Parallel quantization using multiple threads. @@ -196,10 +197,7 @@ def run_quantization_phase( start_time = time.time() progress = None - if quantization_progress: - # pylint: disable-next=import-outside-toplevel - from onecomp.utils.quantization_progress import QuantizationProgressTracker - + if report_progress: progress = QuantizationProgressTracker( logger, len(layer_names), @@ -359,7 +357,7 @@ def run_multi_gpu_quantization( calibration_config: CalibrationConfig, gpu_ids: Optional[List[int]] = None, *, - quantization_progress: bool = True, + report_progress: bool = True, ) -> Dict[str, Any]: """Main entry point for multi-GPU quantization. @@ -368,7 +366,7 @@ def run_multi_gpu_quantization( quantizer: Quantizer instance. calibration_config (CalibrationConfig): Calibration parameters. gpu_ids: List of GPU IDs to use (all GPUs if None). - quantization_progress: When True, log ``[progress]`` with ETA per completed layer. + report_progress: When True, log ``[progress]`` with ETA per completed layer. Returns: Dict containing "results" with quantization results for each layer @@ -394,7 +392,7 @@ def run_multi_gpu_quantization( layer_names=capture_result["layer_names"], quantizer=quantizer, gpu_ids=gpu_ids, - quantization_progress=quantization_progress, + report_progress=report_progress, ) total_elapsed = time.time() - total_start diff --git a/onecomp/utils/quantization_progress.py b/onecomp/utils/quantization_progress.py index 4e16f2c..6f0141a 100644 --- a/onecomp/utils/quantization_progress.py +++ b/onecomp/utils/quantization_progress.py @@ -68,11 +68,11 @@ def _step_complete_unlocked(self, detail: Optional[str]) -> None: done = min(self._done, self._total) elapsed = now - self._start - if done < self._total and done > 0: + if done < self._total: eta_sec = (elapsed / done) * (self._total - done) eta_str = _format_duration(eta_sec) else: - eta_str = "unknown" if done == 0 else "0s" + eta_str = "0s" pct = 100.0 * done / self._total suffix = f" ({detail})" if detail else "" From c1bc4b76b4d25decd5a6d32c3442d510acf20e71 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Kimura=2C=20Keiji/=E6=9C=A8=E6=9D=91=20=E5=9C=AD=E5=85=90?= Date: Tue, 19 May 2026 15:57:10 +0000 Subject: [PATCH 15/82] refactor(progress): tidy [progress] logs and harden tracker MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - QuantizationProgressTracker: suppress duplicate logs after completion, use `is not None` for lock check, expand Google-style docstrings - Demote per-layer / per-chunk INFO logs to DEBUG so the [progress] line is the single canonical per-step INFO signal (onecomp/quantizer/_quantizer.py, onecomp/runner_methods/chunked_quantization.py, onecomp/qep/_quantize_with_qep_arch.py) - Chunked path: switch to per-group progress (was per-layer × per-quantizer); drop chunk_progress / layer_progress arguments - Shorten tracker labels: "Quantization", "Chunked quantization", "QEP" - Per-block MSE log in QEP arch path: drop redundant `[INFO]` prefix, switch to lazy `%` formatting, fix "Layer N" -> "Block N" - Expand QuantizationProgressTracker tests (10 cases, incl. thread-safety and overflow suppression) - Update CHANGELOG for the new [progress] line and INFO -> DEBUG demotion --- CHANGELOG.md | 1 + onecomp/qep/_quantize_with_qep_arch.py | 18 +-- onecomp/quantizer/_quantizer.py | 11 +- onecomp/runner.py | 2 +- .../runner_methods/chunked_quantization.py | 39 ++---- onecomp/utils/quantization_progress.py | 117 +++++++++++++++++- tests/onecomp/utils/__init__.py | 1 + .../utils/test_quantization_progress.py | 100 ++++++++++++++- 8 files changed, 237 insertions(+), 52 deletions(-) create mode 100644 tests/onecomp/utils/__init__.py diff --git a/CHANGELOG.md b/CHANGELOG.md index a657fa7..724e338 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,7 @@ - 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 diff --git a/onecomp/qep/_quantize_with_qep_arch.py b/onecomp/qep/_quantize_with_qep_arch.py index a10dc4a..b99397f 100644 --- a/onecomp/qep/_quantize_with_qep_arch.py +++ b/onecomp/qep/_quantize_with_qep_arch.py @@ -326,13 +326,13 @@ def run_quantize_with_qep_arch( progress = QuantizationProgressTracker( logger, len(remaining_targets), - "QEP quantization (architecture-aware)", + "QEP", ) # 2. For each target transformer block, perform the following sequentially for block_idx, block in enumerate(blocks): - logger.info( + logger.debug( "Processing : %2d-th Transformer Block -------------------------------------------------", block_idx + 1, ) @@ -385,7 +385,7 @@ def run_quantize_with_qep_arch( # 3. Process regular (non-expert) groups with full QEP for group_q, group_f in regular_pairs: - logger.info( + logger.debug( "Processing group of layers: %s", ", ".join([quantizer.module_to_name.get(m, "N/A") for m in group_q]), ) @@ -421,7 +421,7 @@ def run_quantize_with_qep_arch( exclude = any(kw in name for kw in qep_config.exclude_layer_keywords) layer_delta = None if exclude else delta_hatX.clone() - logger.info( + logger.debug( "Processing layer: %s %s=================================================", name, "(no weight correction) " if exclude else "", @@ -452,7 +452,9 @@ def run_quantize_with_qep_arch( ) remaining_targets.discard(name) if progress is not None: - progress.step_complete(name) + progress.step_complete( + f"{name}, no weight correction" if exclude else name + ) # 4. Process MoE expert layers with per-module Hessians (no cross-term) if expert_modules_q: @@ -479,10 +481,10 @@ def run_quantize_with_qep_arch( ) remaining_targets.discard(name) if progress is not None: - progress.step_complete(f"{name} (skipped, no tokens)") + progress.step_complete(f"{name}, skipped: no tokens") continue - logger.info( + logger.debug( "Processing layer: %s (no weight correction) =================================================", name, ) @@ -516,7 +518,7 @@ def run_quantize_with_qep_arch( # Compute MSE between quantized and full-precision block outputs mse = F.mse_loss(inps_q.float(), inps_f.float()).item() - logger.info(f"[INFO] Layer {block_idx + 1} MSE: {mse:.6e}") + logger.info("Block %d MSE: %.6e", block_idx + 1, mse) # free memory block_q.cpu() diff --git a/onecomp/quantizer/_quantizer.py b/onecomp/quantizer/_quantizer.py index 95e1372..1ebc126 100644 --- a/onecomp/quantizer/_quantizer.py +++ b/onecomp/quantizer/_quantizer.py @@ -203,11 +203,14 @@ def quantize( name = self.module_to_name[module] - self.logger.info("Quantizing layer: %s", name) + # Emitted at DEBUG because Runner emits a `[progress]` INFO line per + # completed layer (see ``onecomp.utils.quantization_progress``); keep + # the per-layer "start" signal available under DEBUG for deep dives. + self.logger.debug("Quantizing layer: %s", name) start_time = time.time() if hessian is None and self.flag_hessian: hessian = self.calculate_hessian(module, input) - + result = self.quantize_layer(module, input, hessian=hessian) end_time = time.time() if hessian is not None: @@ -255,7 +258,7 @@ def quantize_with_qep( # Adjust the weights to be quantized if delta_hatX is not None or original_input_activation is not None: - self.logger.info("Adjusting the weight of the layer: %s", name) + self.logger.debug("Adjusting the weight of the layer: %s", name) self.adjust_weight( module, quant_input_activation, @@ -267,7 +270,7 @@ def quantize_with_qep( ) torch.cuda.empty_cache() - self.logger.info("Quantizing layer: %s", name) + self.logger.debug("Quantizing layer: %s", name) result = self.quantize_layer(module, quant_input_activation, hessian=hessian) end_time = time.time() if hessian is not None: diff --git a/onecomp/runner.py b/onecomp/runner.py index f7153d4..0932122 100644 --- a/onecomp/runner.py +++ b/onecomp/runner.py @@ -607,7 +607,7 @@ def quantize_with_calibration(self): progress = QuantizationProgressTracker( logger, len(self.quantizer.module_to_name), - "Calibration quantization layers", + "Quantization", ) if progress: diff --git a/onecomp/runner_methods/chunked_quantization.py b/onecomp/runner_methods/chunked_quantization.py index 456ee66..e1bfc3e 100644 --- a/onecomp/runner_methods/chunked_quantization.py +++ b/onecomp/runner_methods/chunked_quantization.py @@ -120,12 +120,12 @@ def run_chunked_quantization( ) num_groups = (len(all_layers) + num_layers_per_group - 1) // num_layers_per_group - layer_progress = None + progress = None if report_progress: - layer_progress = QuantizationProgressTracker( + progress = QuantizationProgressTracker( logger, - len(all_layers) * len(quantizers), - "Chunked quantization (layers × quantizers)", + num_groups, + "Chunked quantization", ) for group_start in range(0, len(all_layers), num_layers_per_group): @@ -142,15 +142,6 @@ def run_chunked_quantization( layers_list, ) - chunk_progress = None - if report_progress: - num_chunks = (total_samples + calibration_batch_size - 1) // calibration_batch_size - chunk_progress = QuantizationProgressTracker( - logger, - num_chunks, - f"X^T X accumulation chunks (group {group_idx}/{num_groups})", - ) - # --- Phase 1: Forward per chunk -> accumulate X^T X (FP64) --- xtx_dict, nsamples = accumulate_xtx( model=model, @@ -158,13 +149,12 @@ def run_chunked_quantization( input_device=input_device, group=group, calibration_batch_size=calibration_batch_size, - chunk_progress=chunk_progress, ) # --- Phase 2 & 3: Quantize and compute errors for each quantizer --- for quantizer in quantizers: logger.info("Quantizing with %s ...", quantizer.name) - quantize_group(quantizer, group, xtx_dict, nsamples, layer_progress=layer_progress) + quantize_group(quantizer, group, xtx_dict, nsamples) if quantizer.calc_quant_error: record_quantization_errors(quantizer, group, xtx_dict, nsamples) @@ -172,6 +162,9 @@ def run_chunked_quantization( # Release X^T X xtx_dict.clear() + if progress is not None: + progress.step_complete() + # Post-processing for quantizer in quantizers: quantizer.execute_post_processing() @@ -188,7 +181,6 @@ def accumulate_xtx( input_device, group, calibration_batch_size, - chunk_progress=None, ): """Accumulate X^T X in FP64 on GPU by forwarding in chunks. @@ -203,8 +195,6 @@ def accumulate_xtx( input_device: Model's input device (GPU). group: List of target layers [(module, name), ...]. calibration_batch_size: Number of sentences per chunk. - chunk_progress: Optional :class:`~onecomp.utils.quantization_progress.QuantizationProgressTracker` - to record chunk completion with ETA. Returns: xtx_dict: Dict[name, torch.Tensor] @@ -263,15 +253,13 @@ def hook(_module, input, _output): # pylint: disable=redefined-builtin del chunk_inputs torch.cuda.empty_cache() - logger.info( + logger.debug( " Chunk %d/%d done (samples %d-%d)", chunk_idx + 1, num_chunks, chunk_start, chunk_end, ) - if chunk_progress is not None: - chunk_progress.step_complete(f"samples {chunk_start}-{chunk_end}") # Remove hooks for handle in handles: @@ -288,7 +276,7 @@ def hook(_module, input, _output): # pylint: disable=redefined-builtin # ============================================================================= -def quantize_group(quantizer, group, xtx_dict, nsamples, layer_progress=None): +def quantize_group(quantizer, group, xtx_dict, nsamples): """Quantize each layer in the group using accumulated X^T X. flag_hessian=True (GPTQ, DBF, etc.): @@ -303,17 +291,14 @@ def quantize_group(quantizer, group, xtx_dict, nsamples, layer_progress=None): X^T X for each layer (FP64, CPU), shape (in_features, in_features). nsamples: int Number of samples (= total_samples * seq_len). - layer_progress: Optional progress tracker for completed layer quantizations. """ for module, name in group: if name not in xtx_dict and (quantizer.flag_hessian or quantizer.flag_xtx): logger.warning("Skipping %s: no activations captured (unused during forward)", name) - if layer_progress is not None: - layer_progress.step_complete(f"{name} (skipped, no activations)") continue - logger.info("Quantizing layer: %s", name) + logger.debug("Quantizing layer: %s", name) start_time = time.time() if quantizer.flag_xtx: @@ -343,8 +328,6 @@ def quantize_group(quantizer, group, xtx_dict, nsamples, layer_progress=None): quantizer.results[name] = result torch.cuda.empty_cache() - if layer_progress is not None: - layer_progress.step_complete(name) # ============================================================================= diff --git a/onecomp/utils/quantization_progress.py b/onecomp/utils/quantization_progress.py index 6f0141a..a5652f0 100644 --- a/onecomp/utils/quantization_progress.py +++ b/onecomp/utils/quantization_progress.py @@ -1,4 +1,21 @@ -""" +"""Lightweight progress logger for long-running quantization workflows. + +This module exposes :class:`QuantizationProgressTracker`, a small helper +used by :class:`onecomp.runner.Runner` and the underlying quantization +entry points (calibration / chunked / multi-GPU / QEP) to emit a single +``[progress]`` INFO line per completed step with done/total counts, +percentage, wall-clock elapsed time, and a linear ETA estimate. + +The tracker is intentionally minimal: + +- Standard-library only (``threading``, ``time``, ``logging``); no torch + or other heavy dependencies. +- ETA is a simple ``elapsed / done * remaining`` linear extrapolation; + it can be off when per-step cost varies (e.g. MoE expert vs regular + layers), but is sufficient for "is this going to finish before lunch" + scale planning. +- Output is via the caller-supplied ``logging.Logger`` at ``INFO`` + level, so existing log handlers / filters apply transparently. Copyright 2025-2026 Fujitsu Ltd. @@ -13,6 +30,18 @@ def _format_duration(seconds: Optional[float]) -> str: + """Format a number of seconds as a short human-readable string. + + Args: + seconds (float or None): Duration in seconds. ``None`` is treated + as "unknown" (currently unused by the tracker itself but kept + so the helper can be reused in future call sites). + + Returns: + str: ``"unknown"`` when ``seconds`` is ``None``; otherwise a + compact string such as ``"45s"``, ``"3m12s"``, or ``"1h05m"``. + Negative or sub-second values clamp to ``"0s"``. + """ if seconds is None: return "unknown" sec = max(0, int(round(seconds))) @@ -26,7 +55,49 @@ def _format_duration(seconds: Optional[float]) -> str: class QuantizationProgressTracker: - """Log coarse progress (done/total, elapsed, linear ETA) during long quantization.""" + """Emit one INFO log line per completed step with done/total and ETA. + + Each ``step_complete()`` call emits a single line of the form:: + + [progress]