Enable ONNX external-data sharding

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",

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",

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
 

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(