onecomp 1.3.0__tar.gz → 1.3.2__tar.gz

{onecomp-1.3.0/onecomp.egg-info → onecomp-1.3.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: onecomp
-Version: 1.3.0
+Version: 1.3.2
 Summary: Python package for LLM compression
 Author: Keiji Kimura
 License: MIT License
@@ -47,7 +47,8 @@ Requires-Dist: datasets
 Requires-Dist: lm-eval
 Requires-Dist: primefac
 Requires-Dist: scipy
-Requires-Dist: gemlite>=0.5.1
+Requires-Dist: gemlite<0.6.0,>=0.5.1; sys_platform == "darwin"
+Requires-Dist: gemlite>=0.5.1; sys_platform != "darwin"
 Requires-Dist: hqq>=0.2.8
 Requires-Dist: safetensors
 Requires-Dist: ortools>=9.15.6755

{onecomp-1.3.0 → onecomp-1.3.2}/example/post_process/example_lora_sft.py

@@ -66,7 +66,7 @@ model_config = ModelConfig(model_id=MODEL_ID, device="cuda:0")
 gptq = GPTQ(wbits=4, groupsize=128)
 
 post_process = PostProcessLoraSFT(
-    dataset_name="wikitext",
+    dataset_name="Salesforce/wikitext",
     dataset_config_name="wikitext-2-raw-v1",
     train_split="train",
     text_column="text",

{onecomp-1.3.0 → onecomp-1.3.2}/onecomp/__version__.py

@@ -6,4 +6,4 @@ Author: Keiji Kimura
 
 """
 
-__version__ = "1.3.0"
+__version__ = "1.3.2"

{onecomp-1.3.0 → onecomp-1.3.2}/onecomp/cpu/export/dequantize.py

@@ -98,7 +98,7 @@ def dequantize_to_hf(
         pass
 
     logger.info("Building empty dense model from config (%s)", config.model_type)
-    model = AutoModelForCausalLM.from_config(config, torch_dtype=torch_dtype)
+    model = AutoModelForCausalLM.from_config(config, dtype=torch_dtype)
 
     state: Dict[str, torch.Tensor] = {}
     for shard in sorted(glob(os.path.join(save_directory, "*.safetensors"))):

{onecomp-1.3.0 → onecomp-1.3.2}/onecomp/pre_process/prepare_rotated_model.py

@@ -319,7 +319,7 @@ def prepare_rotated_model(
         t0 = time.time()
         model = AutoModelForCausalLM.from_pretrained(
             model_path,
-            torch_dtype="auto",
+            dtype="auto",
             device_map="cpu",
             low_cpu_mem_usage=True,
         )

{onecomp-1.3.0 → onecomp-1.3.2}/onecomp/pre_process/train_rotation.py

@@ -321,7 +321,7 @@ def _convert_model_structure(config, model_type, model_path, use_sdpa=False):
     if is_llama:
         quant_model = QLlamaFC.from_pretrained(
             model_path,
-            torch_dtype="auto",
+            dtype="auto",
             device_map="cpu",
             low_cpu_mem_usage=True,
             config=config,
@@ -329,7 +329,7 @@ def _convert_model_structure(config, model_type, model_path, use_sdpa=False):
     elif is_qwen3:
         quant_model = QQwen3FC.from_pretrained(
             model_path,
-            torch_dtype="auto",
+            dtype="auto",
             device_map="cpu",
             low_cpu_mem_usage=True,
             config=config,

{onecomp-1.3.0 → onecomp-1.3.2}/onecomp/quantized_model_loader.py

@@ -101,22 +101,14 @@ class QuantizedModelLoader:
             torch_dtype = torch.bfloat16
         model = cls._build_empty_model_from_config(config_dict, torch_dtype)
 
-        # Mirror the unfuse step performed before quantization/save (see
-        # Runner.save_quantized_model) so per-expert module names such as
-        # "model.layers.0.mlp.experts.0.down_proj" resolve against the
-        # freshly-built model instead of its fused 3D expert parameters.
-        if unfuse_moe_experts(model, logger):
-            logger.info("Unfused MoE expert tensors for quantized model load")
-
         # Load state_dict from safetensors
         state_dict = cls._load_state_dict_from_dir(save_directory)
 
-        # Align checkpoint key prefixes with the empty model built from config.
-        # Gemma3 VLMs are a common case: weights saved from from_pretrained
-        # use model.language_model.model.layers. (language_model is a
-        # ForCausalLM wrapper) while from_config exposes
-        # model.language_model.layers.* directly.
-        state_dict = cls._remap_state_dict_keys(state_dict, model)
+        # Decide, from the checkpoint alone, whether the empty model must be
+        # unfused into per-expert nn.Linear modules.  This must happen before
+        # _remap_state_dict_keys so remapping aligns checkpoint keys against the
+        # unfused module paths; fused-MoE checkpoints (e.g. gpt-oss) keep the
+        # fused 3D parameters and skip unfuse.
         from .utils.unfuse_moe import (
             _checkpoint_uses_fused_moe,
             _expand_deduped_moe_keys,
@@ -131,6 +123,13 @@ class QuantizedModelLoader:
         elif unfuse_moe_experts(model, logger):
             logger.info("Unfused MoE expert tensors for quantized model load")
 
+        # Align checkpoint key prefixes with the empty model built from config.
+        # Gemma3 VLMs are a common case: weights saved from from_pretrained
+        # use model.language_model.model.layers. (language_model is a
+        # ForCausalLM wrapper) while from_config exposes
+        # model.language_model.layers.* directly.
+        state_dict = cls._remap_state_dict_keys(state_dict, model)
+
         # Replace quantized layers with empty modules and align quantized
         # tensor keys with the actual module names in the model built from
         # config.  This is required when the saved checkpoint and the
@@ -568,11 +567,11 @@ class QuantizedModelLoader:
         config_cls = CONFIG_MAPPING[model_type]
         model_config = config_cls.from_dict(clean_config)
         try:
-            return AutoModelForCausalLM.from_config(model_config, torch_dtype=dtype)
+            return AutoModelForCausalLM.from_config(model_config, dtype=dtype)
         except (ValueError, KeyError):
             from transformers import AutoModelForImageTextToText
 
-            return AutoModelForImageTextToText.from_config(model_config, torch_dtype=dtype)
+            return AutoModelForImageTextToText.from_config(model_config, dtype=dtype)
 
     @staticmethod
     def _set_module_by_name(

{onecomp-1.3.0 → onecomp-1.3.2}/onecomp/quantizer/gemlite.py

@@ -16,13 +16,15 @@ import torch.nn as nn
 import torch.nn.functional as F
 
 # Optional GemLite/HQQ imports
-try:
-    from gemlite.core import DType, GemLiteLinearTriton
-    from hqq.core.quantize import BaseQuantizeConfig, HQQLinear
+HAS_GEMLITE = False
+if torch.cuda.is_available():
+    try:
+        from gemlite.core import DType, GemLiteLinearTriton
+        from hqq.core.quantize import BaseQuantizeConfig, HQQLinear
 
-    HAS_GEMLITE = True
-except (ImportError, AttributeError):
-    HAS_GEMLITE = False
+        HAS_GEMLITE = True
+    except (ImportError, AttributeError, AssertionError):
+        pass
 
 
 # Constants

{onecomp-1.3.0 → onecomp-1.3.2}/onecomp/runner.py

@@ -1260,7 +1260,7 @@ class Runner:
         original_model=False,
         dequantized_model=False,
         quantized_model=True,
-        dataset_name="wikitext",
+        dataset_name="Salesforce/wikitext",
         dataset_config="wikitext-2-raw-v1",
         split="test",
         max_samples=None,
@@ -1342,7 +1342,7 @@ class Runner:
         original_model=True,
         dequantized_model=False,
         quantized_model=True,
-        dataset_name="wikitext",
+        dataset_name="Salesforce/wikitext",
         dataset_config="wikitext-2-raw-v1",
         split="test",
         max_samples=None,

{onecomp-1.3.0 → onecomp-1.3.2}/onecomp/utils/perplexity.py

@@ -41,7 +41,7 @@ def calculate_perplexity(
     model=None,
     tokenizer=None,
     model_config=None,
-    dataset_name="wikitext",
+    dataset_name="Salesforce/wikitext",
     dataset_config="wikitext-2-raw-v1",
     split="test",
     max_samples=None,
@@ -53,7 +53,7 @@ def calculate_perplexity(
     Based on https://huggingface.co/docs/transformers/perplexity
 
     Args:
-        dataset_name (str): Dataset name (e.g. "wikitext", "allenai/c4").
+        dataset_name (str): Dataset name (e.g. "Salesforce/wikitext", "allenai/c4").
         dataset_config (str): Dataset configuration.
             - For WikiText: "wikitext-2-raw-v1"
             - For C4: "en/c4-train.00001-of-01024.json.gz" (treated as data_files)

{onecomp-1.3.0 → onecomp-1.3.2}/onecomp/utils/vram_estimator.py

@@ -312,7 +312,7 @@ def estimate_wbits_from_vram(
 
     config = AutoConfig.from_pretrained(model_id)
     with torch.device("meta"):
-        model = AutoModelForCausalLM.from_config(config, torch_dtype=torch.float16)
+        model = AutoModelForCausalLM.from_config(config, dtype=torch.float16)
 
     return estimate_target_bitwidth(
         model,

{onecomp-1.3.0 → onecomp-1.3.2/onecomp.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: onecomp
-Version: 1.3.0
+Version: 1.3.2
 Summary: Python package for LLM compression
 Author: Keiji Kimura
 License: MIT License
@@ -47,7 +47,8 @@ Requires-Dist: datasets
 Requires-Dist: lm-eval
 Requires-Dist: primefac
 Requires-Dist: scipy
-Requires-Dist: gemlite>=0.5.1
+Requires-Dist: gemlite<0.6.0,>=0.5.1; sys_platform == "darwin"
+Requires-Dist: gemlite>=0.5.1; sys_platform != "darwin"
 Requires-Dist: hqq>=0.2.8
 Requires-Dist: safetensors
 Requires-Dist: ortools>=9.15.6755

{onecomp-1.3.0 → onecomp-1.3.2}/onecomp.egg-info/requires.txt

@@ -4,11 +4,16 @@ datasets
 lm-eval
 primefac
 scipy
-gemlite>=0.5.1
 hqq>=0.2.8
 safetensors
 ortools>=9.15.6755
 
+[:sys_platform != "darwin"]
+gemlite>=0.5.1
+
+[:sys_platform == "darwin"]
+gemlite<0.6.0,>=0.5.1
+
 [cpu]
 
 [cpu:sys_platform == "linux"]

{onecomp-1.3.0 → onecomp-1.3.2}/pyproject.toml

@@ -29,7 +29,8 @@ dependencies = [
     "lm-eval",
     "primefac",
     "scipy",
-    "gemlite>=0.5.1",
+    "gemlite>=0.5.1,<0.6.0; sys_platform == 'darwin'",
+    "gemlite>=0.5.1; sys_platform != 'darwin'",
     "hqq>=0.2.8",
     "safetensors",
     "ortools>=9.15.6755",