diff --git a/pyproject.toml b/pyproject.toml index cb4bbfea7..0f02021b6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -16,7 +16,7 @@ dependencies = [ "huggingface_hub", "ml_dtypes", "numpy>=1.24.0", - "onnx_ir>=0.2.1", + "onnx_ir>=1.0.0", "onnx-shape-inference>=0.3.1", "onnxscript>=0.7.1", "safetensors", diff --git a/src/mobius/__main__.py b/src/mobius/__main__.py index 821490e8f..4453be89b 100644 --- a/src/mobius/__main__.py +++ b/src/mobius/__main__.py @@ -674,7 +674,7 @@ def main(argv: list[str] | None = None) -> None: "--max-shard-size", metavar="SIZE", default=None, - help="Max shard size for safetensors (e.g. '5GB'). Only used with --external-data safetensors.", + help="Maximum external-data shard size (e.g. '5GB'). Used by both ONNX and safetensors.", ) build_parser.add_argument( "--no-weights", diff --git a/src/mobius/_model_package.py b/src/mobius/_model_package.py index 00317f41f..a96dd990b 100644 --- a/src/mobius/_model_package.py +++ b/src/mobius/_model_package.py @@ -22,6 +22,7 @@ import logging import os +import threading from collections import UserDict from collections.abc import Callable @@ -95,8 +96,10 @@ def save( errors on some CUDA/cuBLAS versions when loading weights via memory-mapped I/O. Use ``"onnx"`` (the default) for models targeting CUDA execution. - max_shard_size_bytes: Maximum shard size in bytes for safetensors - format. Only used when *external_data* is ``"safetensors"``. + max_shard_size_bytes: Maximum external-data shard size in bytes. + Used by both ONNX and safetensors external-data formats. A + single tensor larger than this value is written in its own + oversized shard. components: Optional predicate ``(name) -> bool`` that selects which components to save. When ``None`` (default), all components are saved. Examples:: @@ -123,8 +126,6 @@ def save( "Expected 'onnx' or 'safetensors'." ) os.makedirs(directory, exist_ok=True) - callback = _make_progress_callback() if progress_bar else None - selected = { name: model for name, model in self.data.items() @@ -133,6 +134,7 @@ def save( use_subfolders = len(selected) > 1 for name, model in selected.items(): + callback = _make_progress_callback() if progress_bar else None if check_weights: _check_weights(name, model) if use_subfolders: @@ -149,7 +151,13 @@ def save( callback=callback, ) else: - ir.save(model, path, external_data="model.onnx.data", callback=callback) + ir.save( + model, + path, + external_data="model.onnx.data", + max_shard_size_bytes=max_shard_size_bytes, + callback=callback, + ) @classmethod def load(cls, directory: str) -> ModelPackage: @@ -259,19 +267,47 @@ def apply_weights( def _make_progress_callback(): - """Create a tqdm progress-bar callback for ``ir.save``.""" - pbar = tqdm.tqdm() - total_set = False + """Create a thread-safe tqdm progress-bar callback for ``ir.save``. + + Newer ``onnx_ir`` versions may invoke callbacks concurrently and out of + index order. Count invocations instead of tracking ``metadata.index`` and + derive each bar's position from its shard filename so rendering order is + deterministic. Serialize all progress-bar mutations. This remains compatible + with ``onnx_ir`` 1.0, where callbacks are invoked serially. + """ + lock = threading.Lock() + bars: dict[str, tqdm.tqdm] = {} def callback(tensor: ir.TensorProtocol, metadata: ir.external_data.CallbackInfo) -> None: - nonlocal total_set - if not total_set: - pbar.total = metadata.total - total_set = True - pbar.update() - pbar.set_description( - f"Saving {tensor.name} ({tensor.dtype.short_name()}, {tensor.shape})" - ) + with lock: + shard_total = getattr(metadata, "shard_total", None) + key = metadata.filename if shard_total is not None else "__all__" + pbar = bars.get(key) + if pbar is None: + description = ( + f"Saving {metadata.filename}" + if shard_total is not None + else "Saving external data" + ) + position = 0 + if shard_total is not None: + shard_prefix, separator, _ = metadata.filename.rpartition("-of-") + shard_number = shard_prefix.rpartition("-")[2] + if separator and shard_number.isdigit(): + position = int(shard_number) - 1 + pbar = tqdm.tqdm( + total=shard_total if shard_total is not None else metadata.total, + desc=description, + position=position, + leave=True, + ) + bars[key] = pbar + pbar.update() + pbar.set_postfix_str( + f"{tensor.name} ({tensor.dtype.short_name()}, {tensor.shape})" + ) + if pbar.n >= pbar.total: + pbar.close() return callback diff --git a/src/mobius/_model_package_test.py b/src/mobius/_model_package_test.py index 12072f2e2..954727379 100644 --- a/src/mobius/_model_package_test.py +++ b/src/mobius/_model_package_test.py @@ -6,13 +6,15 @@ from __future__ import annotations import logging +import threading +import types import onnx_ir as ir import torch from mobius._builder import build_from_module from mobius._configs import VisionConfig -from mobius._model_package import ModelPackage +from mobius._model_package import ModelPackage, _make_progress_callback from mobius._testing import make_config from mobius.models.base import CausalLMModel from mobius.models.gemma3 import Gemma3MultiModalModel @@ -77,6 +79,170 @@ def test_config_stored(self): assert pkg.config is config +class TestOnnxShardedSave: + def test_onnx_external_data_is_sharded(self, tmp_path): + graph = ir.Graph([], [], nodes=[], name="m") + for index in range(6): + name = f"weight_{index}" + graph.register_initializer( + ir.Value( + name=name, + const_value=ir.Tensor( + torch.full((1024,), index, dtype=torch.float32), + name=name, + dtype=ir.DataType.FLOAT, + ), + ) + ) + pkg = ModelPackage({"m": ir.Model(graph, ir_version=10)}) + + pkg.save( + str(tmp_path), + external_data="onnx", + max_shard_size_bytes=8192, + progress_bar=False, + ) + + shards = sorted( + [ + *tmp_path.glob("model-*-of-*.onnx.data"), + # onnx_ir 1.0.0 did not yet recognize .onnx.data as a + # compound suffix. + *tmp_path.glob("model.onnx-*-of-*.data"), + ] + ) + assert len(shards) == 3 + assert all(shard.stat().st_size <= 8192 for shard in shards) + loaded = ModelPackage.load(str(tmp_path)) + assert set(loaded.data) == {"model"} + + +class TestProgressCallback: + class _Tensor: + name = "w" + shape = (2, 2) + + class dtype: # noqa: N801 + @staticmethod + def short_name(): + return "f32" + + class _Bar: + def __init__(self, *, total, desc, position, leave): + self.total = total + self.desc = desc + self.position = position + self.leave = leave + self.n = 0 + self.postfix = "" + self.closed = False + + def update(self): + self.n += 1 + + def set_postfix_str(self, value): + self.postfix = value + + def close(self): + self.closed = True + + def test_orders_progress_bars_by_shard_number(self, monkeypatch): + bars = [] + + def make_bar(**kwargs): + bar = self._Bar(**kwargs) + bars.append(bar) + return bar + + monkeypatch.setattr("mobius._model_package.tqdm.tqdm", make_bar) + callback = _make_progress_callback() + for filename in ( + "model-00002-of-00002.onnx.data", + "model-00001-of-00002.onnx.data", + ): + for shard_index in reversed(range(2)): + callback( + self._Tensor(), + types.SimpleNamespace( + total=4, + index=shard_index, + offset=0, + filename=filename, + shard_total=2, + shard_index=shard_index, + ), + ) + + assert len(bars) == 2 + assert [bar.position for bar in bars] == [1, 0] + assert all(bar.total == 2 and bar.n == 2 and bar.closed for bar in bars) + assert "model-00002-of-00002.onnx.data" in bars[0].desc + assert "model-00001-of-00002.onnx.data" in bars[1].desc + + def test_falls_back_to_one_bar_with_onnx_ir_1_0(self, monkeypatch): + bars = [] + + def make_bar(**kwargs): + bar = self._Bar(**kwargs) + bars.append(bar) + return bar + + monkeypatch.setattr("mobius._model_package.tqdm.tqdm", make_bar) + callback = _make_progress_callback() + for index in range(4): + callback( + self._Tensor(), + types.SimpleNamespace( + total=4, + index=index, + offset=0, + filename=f"model-{index}.data", + ), + ) + + assert len(bars) == 1 + assert bars[0].total == 4 + assert bars[0].n == 4 + assert bars[0].position == 0 + assert bars[0].closed + + def test_is_thread_safe(self, monkeypatch): + bars = [] + + def make_bar(**kwargs): + bar = self._Bar(**kwargs) + bars.append(bar) + return bar + + monkeypatch.setattr("mobius._model_package.tqdm.tqdm", make_bar) + callback = _make_progress_callback() + total = 200 + + def invoke(index): + callback( + self._Tensor(), + types.SimpleNamespace( + total=total, + index=index, + offset=0, + filename="model.onnx.data", + shard_total=total, + shard_index=index, + ), + ) + + threads = [threading.Thread(target=invoke, args=(index,)) for index in range(total)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + assert len(bars) == 1 + assert bars[0].total == total + assert bars[0].n == total + assert bars[0].closed + + class TestModelPackageSaveLoad: def test_save_creates_files(self, tmp_path): pkg = ModelPackage(