onecomp 1.3.3__tar.gz → 1.3.4__tar.gz

{onecomp-1.3.3/onecomp.egg-info → onecomp-1.3.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: onecomp
-Version: 1.3.3
+Version: 1.3.4
 Summary: Python package for LLM compression
 Author: Keiji Kimura
 License: MIT License
@@ -473,7 +473,12 @@ See the [GPT-OSS guide](docs/user-guide/gptoss.md) for HF save/load, patch detai
 
 ## 📄 License
 
-See [LICENSE](./LICENSE) for more details.
+OneComp is licensed under the [MIT License](./LICENSE).
+
+The dependencies installed with OneComp are separate open-source software (OSS)
+projects and are distributed under their respective licenses. Their licenses
+may change when the dependencies are updated, so please check the license
+terms of the installed versions as well.
 
 ## Citation
 

{onecomp-1.3.3 → onecomp-1.3.4}/README.md

@@ -359,7 +359,12 @@ See the [GPT-OSS guide](docs/user-guide/gptoss.md) for HF save/load, patch detai
 
 ## 📄 License
 
-See [LICENSE](./LICENSE) for more details.
+OneComp is licensed under the [MIT License](./LICENSE).
+
+The dependencies installed with OneComp are separate open-source software (OSS)
+projects and are distributed under their respective licenses. Their licenses
+may change when the dependencies are updated, so please check the license
+terms of the installed versions as well.
 
 ## Citation
 

{onecomp-1.3.3 → onecomp-1.3.4}/onecomp/__version__.py

@@ -6,4 +6,4 @@ Author: Keiji Kimura
 
 """
 
-__version__ = "1.3.3"
+__version__ = "1.3.4"

{onecomp-1.3.3 → onecomp-1.3.4}/onecomp/model_config.py

@@ -11,7 +11,7 @@ from logging import getLogger
 import torch
 from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
 
-from .utils.device import get_default_device
+from .utils.device import get_default_device, is_mps_device
 from .utils.dtype import needs_bfloat16
 
 try:
@@ -95,9 +95,14 @@ class ModelConfig:
                 If ``None`` (default), ``self.device`` is used.
         """
         effective_device = device_map if device_map is not None else self.device
+        if effective_device == "auto":
+            target_device = get_default_device()
+        else:
+            target_device = effective_device
+        load_device = "cpu" if is_mps_device(target_device) else effective_device
         kwargs = dict(
             dtype=self.dtype if self.dtype == "auto" else getattr(torch, self.dtype),
-            device_map=effective_device,
+            device_map=load_device,
         )
 
         config = self.load_config()
@@ -135,6 +140,8 @@ class ModelConfig:
                 raise
             self.logger.info("AutoModelForCausalLM failed; trying AutoModelForImageTextToText.")
             model = _AutoVLM.from_pretrained(self.get_model_id_or_path(), **kwargs)
+        if is_mps_device(target_device):
+            model = model.to(target_device)
         model.eval()
         self.logger.info("Model loaded with dtype=%s", next(model.parameters()).dtype)
         return model

{onecomp-1.3.3 → onecomp-1.3.4}/onecomp/qep/_quantize_with_qep_arch.py

@@ -298,8 +298,15 @@ def _rtn_fallback_result(module: nn.Module, quantizer: Quantizer, name: str) ->
 
     result_dict = run_rtn(module, wbits=wbits, groupsize=groupsize, sym=quantizer.sym)
 
-    # RTN's raw scale/zero are (out_features, num_groups); GPTQResult
-    # expects (num_groups, out_features).
+    scales = result_dict["scale"]
+    qzeros = result_dict["zero"]
+
+    if groupsize != -1:
+        # RTN's raw scale/zero are (out_features, num_groups); GPTQResult
+        # expects (num_groups, out_features).
+        scales = scales.T
+        qzeros = qzeros.T
+
     return GPTQResult(
         dequantized_weight=result_dict["dequantized_weight"],
         wbits=wbits,
@@ -307,8 +314,8 @@ def _rtn_fallback_result(module: nn.Module, quantizer: Quantizer, name: str) ->
         actorder=False,
         sym=quantizer.sym,
         qweight=result_dict["quantized_weight"],
-        scales=result_dict["scale"].T,
-        qzeros=result_dict["zero"].T,
+        scales=scales,
+        qzeros=qzeros,
         perm=None,
     )
 

{onecomp-1.3.3 → onecomp-1.3.4}/onecomp/quantized_model_loader.py

@@ -123,20 +123,16 @@ 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
-        # from_config model use different wrapper prefixes, e.g.
-        # model.language_model.layers.* vs model.layers.*.
+        # Replace quantized layers first. This resolves every quantizer's
+        # tensor state by layer prefix while the model still exposes the
+        # original empty Linear modules.
         state_dict = cls._replace_quantized_layers(model, state_dict, quant_config)
 
+        # Align remaining checkpoint keys with the model built from config.
+        # Quantized keys now already use the actual module prefix, so this
+        # generic remap only handles non-quantized parameters and buffers.
+        state_dict = cls._remap_state_dict_keys(state_dict, model)
+
         # Load all weights (quantized + non-quantized) in one go.  strict=False
         # is intentional because some wrapper-only components may be absent, but
         # critical language-model and quantized-buffer mismatches must fail fast.
@@ -593,11 +589,10 @@ class QuantizedModelLoader:
         silently skips mismatched keys and leaves layers at their empty-model
         initial values (often all zeros for quantized buffers).
 
-        Remapping runs before _replace_quantized_layers, so the empty
-        model still exposes nn.Linear.weight rather than GPTQ buffers
-        (``qweight``, ``scales``, …).  Known prefix rewrites are therefore
-        applied from checkpoint key patterns alone; they must not require the
-        destination key to already exist in model.named_parameters().
+        Quantized layers are resolved by _replace_quantized_layers before
+        this method runs. Therefore this method only remaps ordinary model
+        parameters and buffers; quantizer-specific tensor fields are already
+        materialized under their actual module prefixes.
 
         Args:
             state_dict: Tensors loaded from *.safetensors.

{onecomp-1.3.3 → onecomp-1.3.4}/onecomp/runner.py

@@ -379,6 +379,11 @@ class Runner:
         if is_mps_device(device):
             if self.multi_gpu:
                 raise ValueError("multi_gpu is not supported on MPS device.")
+            if batch_size is not None:
+                raise ValueError(
+                    "MPS quantization does not support calibration_config.batch_size. "
+                    "Remove batch_size from CalibrationConfig and run without chunked calibration."
+                )
             all_quantizers = self.quantizers if self.quantizers is not None else [self.quantizer]
             for i, q in enumerate(all_quantizers):
                 label = f"quantizers[{i}]" if self.quantizers else "quantizer"

{onecomp-1.3.3 → onecomp-1.3.4}/onecomp/utils/blockwise.py

@@ -6,6 +6,7 @@ Author: Yudai Fujimoto, Akihiro Yoshida, Yuma Ichikawa
 
 """
 
+from collections import UserDict
 from logging import getLogger
 
 import torch
@@ -95,6 +96,7 @@ class Catcher(nn.Module):
 _PER_LAYER_INPUTS_KEY = "_per_layer_inputs"
 _POS_EMB_MAP_KEY = "_position_embeddings_map"
 _ATTN_MASK_MAP_KEY = "_attention_mask_map"
+_SHARED_KV_STATES_KEY = "shared_kv_states"
 
 
 def _find_blocks_parent(model, blocks):
@@ -243,9 +245,15 @@ def get_blocks_and_inputs(
 
     # Now capture block inputs for all calibration samples.
     block_inps = []
-    for inp in inp_ids.split(batch_size):
+    for first in range(0, inp_ids.shape[0], batch_size):
+        last = min(first + batch_size, inp_ids.shape[0])
+        inp = inp_ids[first:last]
+        batch_model_kwargs = {
+            k: v[first:last] if isinstance(v, torch.Tensor) and v.dim() >= 1 else v
+            for k, v in model_kwargs.items()
+        }
         try:
-            _ = model(inp, **model_kwargs)
+            _ = model(inp, **batch_model_kwargs)
         except StopForward:
             block_inps.append(blocks[0].inp.cpu())
 
@@ -361,6 +369,10 @@ def move_kwargs_to_device(x, device):
         return [move_kwargs_to_device(v, device) for v in x]
     elif isinstance(x, tuple):
         return tuple(move_kwargs_to_device(v, device) for v in x)
+    elif isinstance(x, UserDict):
+        for k, v in list(x.items()):
+            x[k] = move_kwargs_to_device(v, device)
+        return x
     else:
         return x
 
@@ -396,6 +408,10 @@ def expand_kwargs_batch(kwargs, batch_size):
             return [_expand(t) for t in v]
         elif isinstance(v, dict):
             return {k: _expand(val) for k, val in v.items()}
+        elif isinstance(v, UserDict):
+            for k, val in list(v.items()):
+                v[k] = _expand(val)
+            return v
         return v
 
     return {k: _expand(v) for k, v in kwargs.items()}
@@ -433,9 +449,46 @@ def prepare_block_kwargs(batch_kwargs, block, pli, offset, batch_size, device):
         if layer_type and layer_type in mask_map:
             batch_kwargs["attention_mask"] = mask_map[layer_type]
 
+    # 4) Gemma4 shared KV state. Provider blocks write states into a fresh
+    # batch-local mapping; consumer blocks read the matching calibration slice.
+    shared_kv_states = batch_kwargs.get(_SHARED_KV_STATES_KEY)
+    self_attn = getattr(block, "self_attn", None)
+    if isinstance(shared_kv_states, UserDict) and getattr(
+        self_attn, "store_full_length_kv", False
+    ):
+        batch_kwargs[_SHARED_KV_STATES_KEY] = UserDict()
+    elif isinstance(shared_kv_states, UserDict) and getattr(
+        self_attn, "is_kv_shared_layer", False
+    ):
+        batch_kwargs[_SHARED_KV_STATES_KEY] = _slice_shared_kv_states(
+            shared_kv_states, offset, batch_size, device
+        )
+
     return batch_kwargs
 
 
+def _slice_shared_kv_states(shared_kv_states, offset, batch_size, device):
+    batch_shared_kv_states = UserDict()
+    for key, value in shared_kv_states.items():
+        if isinstance(value, tuple):
+            batch_shared_kv_states[key] = tuple(
+                _slice_shared_kv_tensor(v, offset, batch_size, device) for v in value
+            )
+        else:
+            batch_shared_kv_states[key] = _slice_shared_kv_tensor(
+                value, offset, batch_size, device
+            )
+    return batch_shared_kv_states
+
+
+def _slice_shared_kv_tensor(value, offset, batch_size, device):
+    if isinstance(value, torch.Tensor):
+        if value.dim() >= 1 and value.shape[0] >= offset + batch_size:
+            value = value[offset : offset + batch_size]
+        return value.to(device)
+    return value
+
+
 def _get_block_layer_type(block: nn.Module) -> str | None:
     return (
         getattr(block, "layer_type", None)
@@ -468,6 +521,9 @@ def forward_input(
     pli = kwargs.get(_PER_LAYER_INPUTS_KEY)
     next_inps = []
     offset = 0
+    shared_kv_chunks = {}
+    self_attn = getattr(block, "self_attn", None)
+    stores_shared_kv = getattr(self_attn, "store_full_length_kv", False)
     for inp in inps.split(batch_size):
         bs = inp.shape[0]
         batch_kwargs = expand_kwargs_batch(kwargs, bs)
@@ -475,7 +531,27 @@ def forward_input(
         out = block(inp.to(device), **batch_kwargs)
         out = out[0] if isinstance(out, tuple) else out
         next_inps.append(out.cpu())
+        if stores_shared_kv:
+            for key, value in batch_kwargs[_SHARED_KV_STATES_KEY].items():
+                shared_kv_chunks.setdefault(key, []).append(
+                    tuple(v.detach().cpu() for v in value)
+                    if isinstance(value, tuple)
+                    else value.detach().cpu()
+                )
         offset += bs
+
+    if shared_kv_chunks:
+        shared_kv_states = kwargs.get(_SHARED_KV_STATES_KEY)
+        if not isinstance(shared_kv_states, UserDict):
+            shared_kv_states = UserDict()
+            kwargs[_SHARED_KV_STATES_KEY] = shared_kv_states
+        for key, chunks in shared_kv_chunks.items():
+            if isinstance(chunks[0], tuple):
+                shared_kv_states[key] = tuple(
+                    torch.cat([chunk[i] for chunk in chunks], dim=0) for i in range(len(chunks[0]))
+                )
+            else:
+                shared_kv_states[key] = torch.cat(chunks, dim=0)
     return torch.cat(next_inps)
 
 

{onecomp-1.3.3 → onecomp-1.3.4}/onecomp/utils/unfuse_moe.py

@@ -55,6 +55,9 @@ class _UnfusedExperts(nn.Module):
     def __getitem__(self, idx):
         return getattr(self, str(int(idx)))
 
+    def __iter__(self):
+        return iter(self._modules.values())
+
     def forward(
         self,
         hidden_states: torch.Tensor,
@@ -584,17 +587,19 @@ def _fuse_one(
         expert0 = unfused[0]
         inter = expert0.gate_proj.out_features
         hidden = expert0.gate_proj.in_features
+        up_w0, _ = _dequantized_weight_bias(expert0.up_proj)
+        down_w0, _ = _dequantized_weight_bias(expert0.down_proj)
         gate_up_3d = torch.empty(
             num_experts,
             2 * inter,
             hidden,
-            dtype=expert0.gate_proj.weight.dtype,
+            dtype=up_w0.dtype,
         )
         down_3d = torch.empty(
             num_experts,
             hidden,
             inter,
-            dtype=expert0.down_proj.weight.dtype,
+            dtype=down_w0.dtype,
         )
         for i in range(num_experts):
             expert = unfused[i]

{onecomp-1.3.3 → onecomp-1.3.4/onecomp.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: onecomp
-Version: 1.3.3
+Version: 1.3.4
 Summary: Python package for LLM compression
 Author: Keiji Kimura
 License: MIT License
@@ -473,7 +473,12 @@ See the [GPT-OSS guide](docs/user-guide/gptoss.md) for HF save/load, patch detai
 
 ## 📄 License
 
-See [LICENSE](./LICENSE) for more details.
+OneComp is licensed under the [MIT License](./LICENSE).
+
+The dependencies installed with OneComp are separate open-source software (OSS)
+projects and are distributed under their respective licenses. Their licenses
+may change when the dependencies are updated, so please check the license
+terms of the installed versions as well.
 
 ## Citation
 

{onecomp-1.3.3 → onecomp-1.3.4}/onecomp.egg-info/SOURCES.txt

@@ -296,6 +296,7 @@ onecomp/utils/unfuse_moe.py
 onecomp/utils/vram_estimator.py
 scripts/check_copyright_header.py
 scripts/check_no_japanese.py
+scripts/prepare_calibration_cache.py
 vllm_plugins/__init__.py
 vllm_plugins/dbf/__init__.py
 vllm_plugins/dbf/vllm_plugin.py

onecomp-1.3.4/scripts/prepare_calibration_cache.py

@@ -0,0 +1,284 @@
+#!/usr/bin/env python3
+"""Prepare and verify the fixed C4 calibration cache used by cluster CI.
+
+Copyright 2025-2026 Fujitsu Ltd.
+"""
+
+import argparse
+import fcntl
+import hashlib
+import json
+import os
+import shutil
+import sys
+from contextlib import contextmanager
+from pathlib import Path
+
+import datasets
+from huggingface_hub import hf_hub_download
+
+DATASET_ID = "allenai/c4"
+DATASET_REVISION = "1588ec454efa1a09f29cd18ddd04fe05fc8653a2"
+DATA_FILE = "en/c4-train.00001-of-01024.json.gz"
+SOURCE_SHA256 = "b945059cd1a343cabe311881b7840a6f0363f570e745a0eff0e687e266f6b55d"
+EXPECTED_TRAIN_ROWS = 356318
+MANIFEST_NAME = "onecomp-calibration-manifest.json"
+
+
+def _sha256(path):
+    """Return the SHA-256 digest of a file without loading it all into memory."""
+
+    digest = hashlib.sha256()
+    with path.open("rb") as file:
+        for chunk in iter(lambda: file.read(1024 * 1024), b""):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def _cache_path(cache_root):
+    """Return the save-to-disk directory consumed by the C4 loader."""
+
+    return Path(cache_root).expanduser().resolve() / "c4"
+
+
+def _file_records(root):
+    """Describe every persisted cache file so content changes are detectable."""
+
+    records = []
+    for path in sorted(item for item in root.rglob("*") if item.is_file()):
+        if path.name == MANIFEST_NAME:
+            # The manifest cannot include its own digest without becoming self-referential.
+            continue
+        records.append(
+            {
+                "path": path.relative_to(root).as_posix(),
+                "size": path.stat().st_size,
+                "sha256": _sha256(path),
+            }
+        )
+    return records
+
+
+@contextmanager
+def _cache_lock(cache_root, *, exclusive):
+    """Coordinate cache readers with the short final installation step.
+
+    Verification and CI use a shared lock for the entire period in which cache
+    files may be read. Regeneration uses an exclusive lock only while replacing
+    the verified temporary cache, so building and hashing it does not block CI.
+    """
+
+    cache_root = Path(cache_root).expanduser().resolve()
+    lock_path = cache_root / ".c4.lock"
+    if exclusive:
+        cache_root.mkdir(parents=True, exist_ok=True)
+        lock_path.touch(exist_ok=True)
+        mode = "r+"
+        operation = fcntl.LOCK_EX
+    else:
+        mode = "r"
+        operation = fcntl.LOCK_SH
+
+    try:
+        lock_file = lock_path.open(mode, encoding="utf-8")
+    except FileNotFoundError as exc:
+        raise FileNotFoundError(
+            f"Fixed calibration cache lock is missing: {lock_path}. "
+            "Regenerate the cache with this script."
+        ) from exc
+    with lock_file:
+        fcntl.flock(lock_file, operation)
+        yield
+
+
+def prepare(cache_root, *, local_files_only=False, force=False):
+    """Build, verify, and transactionally install the pinned C4 cache."""
+
+    destination = _cache_path(cache_root)
+    if destination.exists() and not force:
+        raise FileExistsError(
+            f"Calibration cache already exists: {destination}. "
+            "Use --force only for an intentional regeneration."
+        )
+
+    source = Path(
+        hf_hub_download(
+            repo_id=DATASET_ID,
+            filename=DATA_FILE,
+            repo_type="dataset",
+            revision=DATASET_REVISION,
+            local_files_only=local_files_only,
+        )
+    )
+    source_sha256 = _sha256(source)
+    if source_sha256 != SOURCE_SHA256:
+        raise ValueError(
+            f"Unexpected SHA-256 for {DATA_FILE}: {source_sha256}; " f"expected {SOURCE_SHA256}"
+        )
+
+    dataset = datasets.load_dataset("json", data_files={"train": str(source)})
+    if len(dataset["train"]) != EXPECTED_TRAIN_ROWS:
+        raise ValueError(
+            f"Unexpected C4 train row count: {len(dataset['train'])}; "
+            f"expected {EXPECTED_TRAIN_ROWS}"
+        )
+
+    destination.parent.mkdir(parents=True, exist_ok=True)
+    temporary = destination.with_name(f".{destination.name}.tmp-{os.getpid()}")
+    if temporary.exists():
+        shutil.rmtree(temporary)
+
+    try:
+        dataset.save_to_disk(temporary)
+        # Fingerprints can change during save_to_disk(), so record the value that
+        # consumers will observe after loading the persisted cache.
+        persisted_dataset = datasets.load_from_disk(temporary)
+        manifest = {
+            "dataset_id": DATASET_ID,
+            "dataset_revision": DATASET_REVISION,
+            "data_files": {"train": DATA_FILE},
+            "source_sha256": source_sha256,
+            "datasets_version": datasets.__version__,
+            "files": _file_records(temporary),
+            "splits": {
+                "train": {
+                    "num_rows": len(persisted_dataset["train"]),
+                    "fingerprint": persisted_dataset["train"]._fingerprint,
+                    "features": persisted_dataset["train"].features.to_dict(),
+                }
+            },
+        }
+        (temporary / MANIFEST_NAME).write_text(
+            json.dumps(manifest, indent=2, sort_keys=True) + "\n",
+            encoding="utf-8",
+        )
+        # Reject an incomplete or internally inconsistent cache before it can
+        # become visible at the shared destination.
+        _verify_cache(temporary)
+
+        backup = destination.with_name(f".{destination.name}.backup-{os.getpid()}")
+        with _cache_lock(cache_root, exclusive=True):
+            if destination.exists() and not force:
+                raise FileExistsError(
+                    f"Calibration cache already exists: {destination}. "
+                    "Use --force only for an intentional regeneration."
+                )
+            if backup.exists():
+                raise FileExistsError(f"Calibration cache backup already exists: {backup}")
+
+            moved_existing = False
+            try:
+                # Keep the previous cache available for rollback until the
+                # already-verified replacement has been installed.
+                if destination.exists():
+                    destination.rename(backup)
+                    moved_existing = True
+                temporary.rename(destination)
+            except Exception:
+                if moved_existing and not destination.exists():
+                    backup.rename(destination)
+                raise
+
+            if moved_existing:
+                try:
+                    shutil.rmtree(backup)
+                except OSError as exc:
+                    print(
+                        f"WARNING: failed to remove calibration cache backup: {exc}",
+                        file=sys.stderr,
+                    )
+    except Exception:
+        shutil.rmtree(temporary, ignore_errors=True)
+        raise
+
+    print(f"Prepared fixed C4 calibration cache: {destination}")
+
+
+def _verify_cache(destination):
+    """Validate the pinned source identity and every persisted cache artifact."""
+
+    manifest_path = destination / MANIFEST_NAME
+    if not manifest_path.is_file():
+        raise FileNotFoundError(f"Fixed calibration cache manifest is missing: {manifest_path}")
+
+    manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
+    expected = {
+        "dataset_id": DATASET_ID,
+        "dataset_revision": DATASET_REVISION,
+        "data_files": {"train": DATA_FILE},
+        "source_sha256": SOURCE_SHA256,
+    }
+    for key, expected_value in expected.items():
+        if manifest.get(key) != expected_value:
+            raise ValueError(
+                f"Invalid fixed calibration cache manifest field {key!r}: "
+                f"{manifest.get(key)!r}; expected {expected_value!r}"
+            )
+
+    files = _file_records(destination)
+    if manifest.get("files") != files:
+        raise ValueError("Fixed calibration cache files do not match its manifest")
+
+    dataset = datasets.load_from_disk(destination)
+    if "train" not in dataset:
+        raise ValueError("Fixed calibration cache has no train split")
+    train_manifest = manifest.get("splits", {}).get("train", {})
+    if len(dataset["train"]) != EXPECTED_TRAIN_ROWS:
+        raise ValueError(
+            f"Invalid fixed calibration cache row count: {len(dataset['train'])}; "
+            f"expected {EXPECTED_TRAIN_ROWS}"
+        )
+    if train_manifest.get("num_rows") != EXPECTED_TRAIN_ROWS:
+        raise ValueError("Fixed calibration cache manifest has an invalid row count")
+    if dataset["train"]._fingerprint != train_manifest.get("fingerprint"):
+        raise ValueError("Fixed calibration cache fingerprint does not match its manifest")
+    if dataset["train"].features.to_dict() != train_manifest.get("features"):
+        raise ValueError("Fixed calibration cache schema does not match its manifest")
+
+    print(
+        "Verified fixed C4 calibration cache: "
+        f"{destination} ({len(dataset['train'])} train rows)"
+    )
+
+
+def verify(cache_root):
+    """Verify a cache while preventing concurrent regeneration from replacing it."""
+
+    with _cache_lock(cache_root, exclusive=False):
+        _verify_cache(_cache_path(cache_root))
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--cache-root",
+        default=os.environ.get("ONECOMP_CALIB_CACHE"),
+        help="Cache root containing the c4 directory (default: ONECOMP_CALIB_CACHE)",
+    )
+    parser.add_argument("--verify", action="store_true", help="Verify without writing")
+    parser.add_argument(
+        "--local-files-only",
+        action="store_true",
+        help="Use an already downloaded Hugging Face source shard",
+    )
+    parser.add_argument(
+        "--force",
+        action="store_true",
+        help="Replace an existing cache intentionally",
+    )
+    args = parser.parse_args()
+    if not args.cache_root:
+        parser.error("--cache-root or ONECOMP_CALIB_CACHE is required")
+
+    if args.verify:
+        verify(args.cache_root)
+    else:
+        prepare(
+            args.cache_root,
+            local_files_only=args.local_files_only,
+            force=args.force,
+        )
+
+
+if __name__ == "__main__":
+    main()