From e0755fb4a325bb32fabcf0a432f11c13064dccc8 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Sat, 15 Aug 2026 18:45:25 -0700 Subject: [PATCH 01/14] Save external data with multiple threads Weights that need a dtype cast are held as ir.LazyTensor and converted at serialization time, so saving alternates between casting a tensor and writing it, leaving both the CPU and the disk idle half the time. onnx_ir gained a max_workers option that overlaps materialization with disk writes and parallelizes both, with peak memory bounded independently of the worker count. Default ModelPackage.save to 8 workers and thread the option through. The option is feature-detected because it is not in onnx_ir 0.2.x, which the package still supports, so older installs keep saving serially instead of raising TypeError. Fix the progress bar for concurrent saving. onnx_ir serializes callbacks with a lock but no longer delivers them in index order, and the bar mutated total/set_description without any synchronization of its own. Guard the closure state with a lock and keep counting invocations rather than tracking indices. Measured on Qwen2.5-7B cast from bfloat16 to float16, 15.26GB of external data: 264.2s serial -> 176.0s with 8 workers (1.50x). 4 workers gives 1.45x and 16 gives 1.48x, so the default of 8 is near the knee. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: af83bc37-b9d6-4b94-be64-a4daa5a7ce40 Signed-off-by: Justin Chu --- src/mobius/_model_package.py | 59 ++++++++++++++++--- src/mobius/_model_package_test.py | 95 ++++++++++++++++++++++++++++++- 2 files changed, 144 insertions(+), 10 deletions(-) diff --git a/src/mobius/_model_package.py b/src/mobius/_model_package.py index 00317f41f..47025a57b 100644 --- a/src/mobius/_model_package.py +++ b/src/mobius/_model_package.py @@ -20,8 +20,10 @@ __all__ = ["ModelPackage"] +import inspect import logging import os +import threading from collections import UserDict from collections.abc import Callable @@ -66,6 +68,7 @@ def save( components: Callable[[str], bool] | None = None, progress_bar: bool = True, check_weights: bool = True, + max_workers: int | None = 8, ) -> None: """Save all component models to a directory. @@ -111,6 +114,13 @@ def save( check_weights: Whether to verify that all initializers have weight data before saving. Defaults to ``True``. Set to ``False`` when saving skeleton models without weights. + max_workers: Number of threads used to materialize and write + tensors. Weights that need a dtype cast are held as + :class:`ir.LazyTensor` and converted at save time, so this + overlaps the casting work with disk writes and parallelizes + both. Pass ``None`` or ``1`` to save serially. Peak memory + stays bounded regardless of the worker count. Only used when + *external_data* is ``"onnx"``. Raises: ValueError: If *external_data* is not ``"onnx"`` or @@ -149,7 +159,16 @@ def save( callback=callback, ) else: - ir.save(model, path, external_data="model.onnx.data", callback=callback) + save_kwargs = {} + if max_workers is not None and _ir_save_supports_max_workers(): + save_kwargs["max_workers"] = max_workers + ir.save( + model, + path, + external_data="model.onnx.data", + callback=callback, + **save_kwargs, + ) @classmethod def load(cls, directory: str) -> ModelPackage: @@ -258,20 +277,42 @@ def apply_weights( fold_initializers_after_weights(model) +def _ir_save_supports_max_workers() -> bool: + """Whether the installed ``onnx_ir`` can save external data concurrently. + + ``max_workers`` was added to ``ir.save`` after 0.2.x. Detect it so mobius + keeps working against older releases instead of raising ``TypeError``. + """ + try: + return "max_workers" in inspect.signature(ir.save).parameters + except (TypeError, ValueError): + return False + + def _make_progress_callback(): - """Create a tqdm progress-bar callback for ``ir.save``.""" + """Create a tqdm progress-bar callback for ``ir.save``. + + ``ir.save`` invokes the callback from worker threads when ``max_workers`` + is greater than 1, and does not deliver calls in ``metadata.index`` order. + The bar therefore counts invocations rather than tracking indices, and + guards its own state with a lock. ``tqdm.update`` is itself thread-safe, + but ``total``/``set_description`` are not, and reading ``tensor.shape`` + concurrently with the write threads must not race the bar's own state. + """ pbar = tqdm.tqdm() + lock = threading.Lock() total_set = False 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: + 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})" + ) return callback diff --git a/src/mobius/_model_package_test.py b/src/mobius/_model_package_test.py index 12072f2e2..8c9f08c15 100644 --- a/src/mobius/_model_package_test.py +++ b/src/mobius/_model_package_test.py @@ -12,7 +12,11 @@ from mobius._builder import build_from_module from mobius._configs import VisionConfig -from mobius._model_package import ModelPackage +from mobius._model_package import ( + ModelPackage, + _ir_save_supports_max_workers, + _make_progress_callback, +) from mobius._testing import make_config from mobius.models.base import CausalLMModel from mobius.models.gemma3 import Gemma3MultiModalModel @@ -77,6 +81,95 @@ def test_config_stored(self): assert pkg.config is config +class TestParallelSave: + """Saving concurrently must match serial output and keep the bar correct.""" + + def test_parallel_save_matches_serial(self, tmp_path): + serial_dir = tmp_path / "serial" + parallel_dir = tmp_path / "parallel" + ModelPackage({"m": _make_simple_model("m")}).save( + str(serial_dir), max_workers=None, progress_bar=False + ) + ModelPackage({"m": _make_simple_model("m")}).save( + str(parallel_dir), max_workers=8, progress_bar=False + ) + serial_data = (serial_dir / "model.onnx.data").read_bytes() + parallel_data = (parallel_dir / "model.onnx.data").read_bytes() + assert serial_data == parallel_data + + def test_parallel_save_roundtrips(self, tmp_path): + pkg = ModelPackage({"m": _make_simple_model("m")}) + pkg.save(str(tmp_path), max_workers=8, progress_bar=False) + # A single-component package is saved flat as ``model.onnx``. + loaded = ModelPackage.load(str(tmp_path)) + assert set(loaded.data) == {"model"} + + def test_progress_bar_counts_every_tensor_out_of_order(self): + # ir.save may invoke the callback from worker threads and out of index + # order, so the bar must count calls rather than follow ``index``. + callback = _make_progress_callback() + + class _Tensor: + name = "w" + shape = (2, 2) + + class dtype: # noqa: N801 + @staticmethod + def short_name(): + return "f32" + + total = 8 + for index in reversed(range(total)): + callback( + _Tensor(), + ir.external_data.CallbackInfo( + total=total, index=index, offset=0, filename="model.onnx.data" + ), + ) + # Closure state is private to the callback; assert via the bound bar. + bar = callback.__closure__[1].cell_contents + assert bar.total == total + assert bar.n == total + + def test_progress_bar_is_thread_safe(self): + import threading + + callback = _make_progress_callback() + + class _Tensor: + name = "w" + shape = (2, 2) + + class dtype: # noqa: N801 + @staticmethod + def short_name(): + return "f32" + + total = 200 + + def worker(index): + callback( + _Tensor(), + ir.external_data.CallbackInfo( + total=total, index=index, offset=0, filename="model.onnx.data" + ), + ) + + threads = [threading.Thread(target=worker, args=(i,)) for i in range(total)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + bar = callback.__closure__[1].cell_contents + assert bar.n == total + + def test_feature_detection_matches_installed_ir(self): + import inspect as _inspect + + expected = "max_workers" in _inspect.signature(ir.save).parameters + assert _ir_save_supports_max_workers() is expected + + class TestModelPackageSaveLoad: def test_save_creates_files(self, tmp_path): pkg = ModelPackage( From 1383495c70190e979477b2570a1cbbab0734a3a7 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Sat, 15 Aug 2026 18:46:31 -0700 Subject: [PATCH 02/14] Require onnx_ir>=1.1.0 instead of feature-detecting max_workers max_workers ships in onnx_ir 1.1.0, so depend on it directly and drop the inspect.signature fallback. Merge this after 1.1.0 is released. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: af83bc37-b9d6-4b94-be64-a4daa5a7ce40 Signed-off-by: Justin Chu --- pyproject.toml | 2 +- src/mobius/_model_package.py | 24 ++++-------------------- src/mobius/_model_package_test.py | 12 +----------- 3 files changed, 6 insertions(+), 32 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index cb4bbfea7..3a752195d 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.1.0", "onnx-shape-inference>=0.3.1", "onnxscript>=0.7.1", "safetensors", diff --git a/src/mobius/_model_package.py b/src/mobius/_model_package.py index 47025a57b..ba3b4910f 100644 --- a/src/mobius/_model_package.py +++ b/src/mobius/_model_package.py @@ -20,7 +20,6 @@ __all__ = ["ModelPackage"] -import inspect import logging import os import threading @@ -118,9 +117,9 @@ def save( tensors. Weights that need a dtype cast are held as :class:`ir.LazyTensor` and converted at save time, so this overlaps the casting work with disk writes and parallelizes - both. Pass ``None`` or ``1`` to save serially. Peak memory - stays bounded regardless of the worker count. Only used when - *external_data* is ``"onnx"``. + both. Defaults to ``8``. Pass ``None`` or ``1`` to save + serially. Peak memory stays bounded regardless of the worker + count. Only used when *external_data* is ``"onnx"``. Raises: ValueError: If *external_data* is not ``"onnx"`` or @@ -159,15 +158,12 @@ def save( callback=callback, ) else: - save_kwargs = {} - if max_workers is not None and _ir_save_supports_max_workers(): - save_kwargs["max_workers"] = max_workers ir.save( model, path, external_data="model.onnx.data", callback=callback, - **save_kwargs, + max_workers=max_workers, ) @classmethod @@ -277,18 +273,6 @@ def apply_weights( fold_initializers_after_weights(model) -def _ir_save_supports_max_workers() -> bool: - """Whether the installed ``onnx_ir`` can save external data concurrently. - - ``max_workers`` was added to ``ir.save`` after 0.2.x. Detect it so mobius - keeps working against older releases instead of raising ``TypeError``. - """ - try: - return "max_workers" in inspect.signature(ir.save).parameters - except (TypeError, ValueError): - return False - - def _make_progress_callback(): """Create a tqdm progress-bar callback for ``ir.save``. diff --git a/src/mobius/_model_package_test.py b/src/mobius/_model_package_test.py index 8c9f08c15..fb8781b68 100644 --- a/src/mobius/_model_package_test.py +++ b/src/mobius/_model_package_test.py @@ -12,11 +12,7 @@ from mobius._builder import build_from_module from mobius._configs import VisionConfig -from mobius._model_package import ( - ModelPackage, - _ir_save_supports_max_workers, - _make_progress_callback, -) +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 @@ -163,12 +159,6 @@ def worker(index): bar = callback.__closure__[1].cell_contents assert bar.n == total - def test_feature_detection_matches_installed_ir(self): - import inspect as _inspect - - expected = "max_workers" in _inspect.signature(ir.save).parameters - assert _ir_save_supports_max_workers() is expected - class TestModelPackageSaveLoad: def test_save_creates_files(self, tmp_path): From 7309783b04b73ba5e0358b1906927e891da3a779 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Sat, 15 Aug 2026 18:49:07 -0700 Subject: [PATCH 03/14] Derive the default save worker count instead of hardcoding 8 Hardcoding 8 happened to match this machine's performance-core count, which was a coincidence rather than a reason. Measuring the phases separately shows the bottleneck is the storage device, not the CPU: casting bf16 to f16 runs at 12.2 GB/s while serial writes reach only 1.6 GB/s. A write-only scaling test saturates at 2 threads (2.1x) and does not improve through 24. So the speedup comes from overlapping the cast with the write, not from write parallelism. That is why the default deliberately does not scale with the core count -- something like cpu_count/2 would oversubscribe against torch's own intra-op pool (already 8 threads here) without making the disk faster. Keep a small constant, but clamp it by os.cpu_count() so small containers do not spawn more threads than they can run, and handle cpu_count() returning None. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: af83bc37-b9d6-4b94-be64-a4daa5a7ce40 Signed-off-by: Justin Chu --- src/mobius/_model_package.py | 35 +++++++++++++++++++++++++------ src/mobius/_model_package_test.py | 17 ++++++++++++++- 2 files changed, 45 insertions(+), 7 deletions(-) diff --git a/src/mobius/_model_package.py b/src/mobius/_model_package.py index ba3b4910f..9174d9cde 100644 --- a/src/mobius/_model_package.py +++ b/src/mobius/_model_package.py @@ -35,6 +35,10 @@ logger = logging.getLogger(__name__) +# Upper bound for the default external-data write thread count. Writes plateau +# well before this on the machines measured; see _default_save_workers. +_MAX_DEFAULT_SAVE_WORKERS = 8 + class ModelPackage(UserDict[str, ir.Model]): """A dict-like collection of named ``ir.Model`` objects. @@ -67,7 +71,7 @@ def save( components: Callable[[str], bool] | None = None, progress_bar: bool = True, check_weights: bool = True, - max_workers: int | None = 8, + max_workers: int | None = None, ) -> None: """Save all component models to a directory. @@ -116,10 +120,10 @@ def save( max_workers: Number of threads used to materialize and write tensors. Weights that need a dtype cast are held as :class:`ir.LazyTensor` and converted at save time, so this - overlaps the casting work with disk writes and parallelizes - both. Defaults to ``8``. Pass ``None`` or ``1`` to save - serially. Peak memory stays bounded regardless of the worker - count. Only used when *external_data* is ``"onnx"``. + overlaps the casting work with disk writes. ``None`` (the + default) picks :func:`_default_save_workers`; pass ``1`` to + save serially. Peak memory stays bounded regardless of the + worker count. Only used when *external_data* is ``"onnx"``. Raises: ValueError: If *external_data* is not ``"onnx"`` or @@ -163,7 +167,9 @@ def save( path, external_data="model.onnx.data", callback=callback, - max_workers=max_workers, + max_workers=( + _default_save_workers() if max_workers is None else max_workers + ), ) @classmethod @@ -273,6 +279,23 @@ def apply_weights( fold_initializers_after_weights(model) +def _default_save_workers() -> int: + """Pick a default thread count for writing external data. + + The bottleneck is the storage device, not the CPU: casting weights runs + around 12 GB/s here while writes saturate near 1.6 GB/s, and a write-only + scaling test stops improving past ~2 threads. Most of the win therefore + comes from overlapping the cast with the write rather than from write + parallelism, so this deliberately does not scale with the core count -- + that would oversubscribe against torch's own intra-op thread pool without + making the disk any faster. + + A small constant, clamped by the core count so tiny containers do not + spawn more threads than they can run, is enough to reach the plateau. + """ + return max(1, min(_MAX_DEFAULT_SAVE_WORKERS, os.cpu_count() or 1)) + + def _make_progress_callback(): """Create a tqdm progress-bar callback for ``ir.save``. diff --git a/src/mobius/_model_package_test.py b/src/mobius/_model_package_test.py index fb8781b68..9757c0777 100644 --- a/src/mobius/_model_package_test.py +++ b/src/mobius/_model_package_test.py @@ -12,7 +12,11 @@ from mobius._builder import build_from_module from mobius._configs import VisionConfig -from mobius._model_package import ModelPackage, _make_progress_callback +from mobius._model_package import ( + ModelPackage, + _default_save_workers, + _make_progress_callback, +) from mobius._testing import make_config from mobius.models.base import CausalLMModel from mobius.models.gemma3 import Gemma3MultiModalModel @@ -127,6 +131,17 @@ def short_name(): assert bar.total == total assert bar.n == total + def test_default_workers_is_clamped_by_core_count(self, monkeypatch): + # The write bottleneck is the disk, so the default is a small constant + # rather than a function of the core count -- but it must never exceed + # the cores actually available. + monkeypatch.setattr("mobius._model_package.os.cpu_count", lambda: 2) + assert _default_save_workers() == 2 + monkeypatch.setattr("mobius._model_package.os.cpu_count", lambda: 128) + assert _default_save_workers() == 8 + monkeypatch.setattr("mobius._model_package.os.cpu_count", lambda: None) + assert _default_save_workers() == 1 + def test_progress_bar_is_thread_safe(self): import threading From 8e57e1768b97372124858bd2c309c6e910c32a0b Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Sat, 15 Aug 2026 19:05:56 -0700 Subject: [PATCH 04/14] Do not shrink the default worker count on small machines Sizing the pool by core count gets the tradeoff backwards. With torch pinned to one intra-op thread, i.e. a small or slow machine, 8 workers is the fastest configuration (1.55x) while 2 workers is slower than serial: more concurrency is needed to hide a slow serial cast, not less. cpu_count() // 2 would give 1 on a 2-core box, disabling the optimization exactly where it helps most. Oversubscription is not penalized in practice either. With full torch threads, everything from 4 to 32 workers lands within noise (1.37x-1.45x), because these threads block in tofile with the GIL released rather than competing for CPU. Drop the core-count clamp and keep a plain constant. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: af83bc37-b9d6-4b94-be64-a4daa5a7ce40 Signed-off-by: Justin Chu --- src/mobius/_model_package.py | 33 ++++++++++++++++++------------- src/mobius/_model_package_test.py | 17 +++++++--------- 2 files changed, 26 insertions(+), 24 deletions(-) diff --git a/src/mobius/_model_package.py b/src/mobius/_model_package.py index 9174d9cde..1107d51af 100644 --- a/src/mobius/_model_package.py +++ b/src/mobius/_model_package.py @@ -35,9 +35,9 @@ logger = logging.getLogger(__name__) -# Upper bound for the default external-data write thread count. Writes plateau -# well before this on the machines measured; see _default_save_workers. -_MAX_DEFAULT_SAVE_WORKERS = 8 +# Default thread count for writing external data. Intentionally independent of +# the core count -- see _default_save_workers for the measurements. +_DEFAULT_SAVE_WORKERS = 8 class ModelPackage(UserDict[str, ir.Model]): @@ -282,18 +282,23 @@ def apply_weights( def _default_save_workers() -> int: """Pick a default thread count for writing external data. - The bottleneck is the storage device, not the CPU: casting weights runs - around 12 GB/s here while writes saturate near 1.6 GB/s, and a write-only - scaling test stops improving past ~2 threads. Most of the win therefore - comes from overlapping the cast with the write rather than from write - parallelism, so this deliberately does not scale with the core count -- - that would oversubscribe against torch's own intra-op thread pool without - making the disk any faster. - - A small constant, clamped by the core count so tiny containers do not - spawn more threads than they can run, is enough to reach the plateau. + Deliberately a constant rather than a function of ``os.cpu_count()``. + + The bottleneck is the storage device, not the CPU: casting bf16 to f16 + runs around 12 GB/s while serial writes reach 1.6 GB/s, and a write-only + scaling test saturates at 2 threads. Nearly all of the win comes from + overlapping the cast with the write, and these threads spend their time + blocked in ``tofile`` with the GIL released rather than competing for CPU. + + Sizing the pool by core count gets this backwards. With torch pinned to a + single intra-op thread -- i.e. a small machine -- 8 workers is the *fastest* + configuration (1.55x) while 2 workers is slower than serial, because more + concurrency is needed to hide a slow serial cast. ``cpu_count() // 2`` would + give 1 on a 2-core box, disabling the optimization exactly where it helps + most. Oversubscription is also not penalized in practice: with full torch + threads, everything from 4 to 32 workers lands within noise of each other. """ - return max(1, min(_MAX_DEFAULT_SAVE_WORKERS, os.cpu_count() or 1)) + return _DEFAULT_SAVE_WORKERS def _make_progress_callback(): diff --git a/src/mobius/_model_package_test.py b/src/mobius/_model_package_test.py index 9757c0777..372dadb5f 100644 --- a/src/mobius/_model_package_test.py +++ b/src/mobius/_model_package_test.py @@ -131,16 +131,13 @@ def short_name(): assert bar.total == total assert bar.n == total - def test_default_workers_is_clamped_by_core_count(self, monkeypatch): - # The write bottleneck is the disk, so the default is a small constant - # rather than a function of the core count -- but it must never exceed - # the cores actually available. - monkeypatch.setattr("mobius._model_package.os.cpu_count", lambda: 2) - assert _default_save_workers() == 2 - monkeypatch.setattr("mobius._model_package.os.cpu_count", lambda: 128) - assert _default_save_workers() == 8 - monkeypatch.setattr("mobius._model_package.os.cpu_count", lambda: None) - assert _default_save_workers() == 1 + def test_default_workers_does_not_track_core_count(self, monkeypatch): + # These threads block on I/O rather than competing for CPU, and a low + # core count needs *more* concurrency to hide a slow serial cast, so + # the default must not shrink on small machines. + for cores in (2, 8, 128, None): + monkeypatch.setattr("mobius._model_package.os.cpu_count", lambda c=cores: c) + assert _default_save_workers() == 8 def test_progress_bar_is_thread_safe(self): import threading From 7d3fafbae198a5e02e6441d6ad09c670db9bf927 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Sat, 15 Aug 2026 19:09:35 -0700 Subject: [PATCH 05/14] Note that accelerator-resident weights benefit most Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: af83bc37-b9d6-4b94-be64-a4daa5a7ce40 Signed-off-by: Justin Chu --- src/mobius/_model_package.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/mobius/_model_package.py b/src/mobius/_model_package.py index 1107d51af..fa0c454de 100644 --- a/src/mobius/_model_package.py +++ b/src/mobius/_model_package.py @@ -297,6 +297,11 @@ def _default_save_workers() -> int: give 1 on a 2-core box, disabling the optimization exactly where it helps most. Oversubscription is also not penalized in practice: with full torch threads, everything from 4 to 32 workers lands within noise of each other. + + Weights that live on an accelerator benefit the most. The device-to-host + copy happens inside ``tofile`` with the GIL released, so it overlaps with + other threads' writes: GPU-resident weights measured 6.3x at 8 workers + versus 1.5x for CPU weights, on a machine whose CPUs were largely idle. """ return _DEFAULT_SAVE_WORKERS From 91996c960fa54f6ca59e489a52df6ea6092ba5e8 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Sat, 15 Aug 2026 21:29:23 -0700 Subject: [PATCH 06/14] Enable sharding for ONNX external data Forward max_shard_size_bytes to ir.save instead of limiting the option to safetensors. Update the CLI help and add a behavioral test that writes, discovers, and reloads three ONNX external-data shards. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: af83bc37-b9d6-4b94-be64-a4daa5a7ce40 Signed-off-by: Justin Chu --- src/mobius/__main__.py | 2 +- src/mobius/_model_package.py | 7 +++++-- src/mobius/_model_package_test.py | 30 ++++++++++++++++++++++++++++++ 3 files changed, 36 insertions(+), 3 deletions(-) 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 fa0c454de..358f9341d 100644 --- a/src/mobius/_model_package.py +++ b/src/mobius/_model_package.py @@ -101,8 +101,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:: @@ -166,6 +168,7 @@ def save( model, path, external_data="model.onnx.data", + max_shard_size_bytes=max_shard_size_bytes, callback=callback, max_workers=( _default_save_workers() if max_workers is None else max_workers diff --git a/src/mobius/_model_package_test.py b/src/mobius/_model_package_test.py index 372dadb5f..c9f618628 100644 --- a/src/mobius/_model_package_test.py +++ b/src/mobius/_model_package_test.py @@ -104,6 +104,36 @@ def test_parallel_save_roundtrips(self, tmp_path): loaded = ModelPackage.load(str(tmp_path)) assert set(loaded.data) == {"model"} + 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, + max_workers=8, + progress_bar=False, + ) + + shards = sorted(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"} + def test_progress_bar_counts_every_tensor_out_of_order(self): # ir.save may invoke the callback from worker threads and out of index # order, so the bar must count calls rather than follow ``index``. From 09e664b6749b0041caf8c07bcf95d48c95ff2732 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Sat, 15 Aug 2026 21:32:29 -0700 Subject: [PATCH 07/14] Keep Mobius compatible with onnx_ir 1.0 Do not expose or pass max_workers or ONNX max_shard_size_bytes from Mobius. Those APIs belong to the pending onnx_ir change, and calling them would make Mobius fail before that change is released. Require only onnx_ir 1.0.0 so the Mobius update is independent of the ir-py PR. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: af83bc37-b9d6-4b94-be64-a4daa5a7ce40 Signed-off-by: Justin Chu --- pyproject.toml | 2 +- src/mobius/__main__.py | 2 +- src/mobius/_model_package.py | 83 +++---------------- src/mobius/_model_package_test.py | 127 +----------------------------- 4 files changed, 14 insertions(+), 200 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 3a752195d..0f02021b6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -16,7 +16,7 @@ dependencies = [ "huggingface_hub", "ml_dtypes", "numpy>=1.24.0", - "onnx_ir>=1.1.0", + "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 4453be89b..821490e8f 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="Maximum external-data shard size (e.g. '5GB'). Used by both ONNX and safetensors.", + help="Max shard size for safetensors (e.g. '5GB'). Only used with --external-data safetensors.", ) build_parser.add_argument( "--no-weights", diff --git a/src/mobius/_model_package.py b/src/mobius/_model_package.py index 358f9341d..00317f41f 100644 --- a/src/mobius/_model_package.py +++ b/src/mobius/_model_package.py @@ -22,7 +22,6 @@ import logging import os -import threading from collections import UserDict from collections.abc import Callable @@ -35,10 +34,6 @@ logger = logging.getLogger(__name__) -# Default thread count for writing external data. Intentionally independent of -# the core count -- see _default_save_workers for the measurements. -_DEFAULT_SAVE_WORKERS = 8 - class ModelPackage(UserDict[str, ir.Model]): """A dict-like collection of named ``ir.Model`` objects. @@ -71,7 +66,6 @@ def save( components: Callable[[str], bool] | None = None, progress_bar: bool = True, check_weights: bool = True, - max_workers: int | None = None, ) -> None: """Save all component models to a directory. @@ -101,10 +95,8 @@ 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 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. + max_shard_size_bytes: Maximum shard size in bytes for safetensors + format. Only used when *external_data* is ``"safetensors"``. components: Optional predicate ``(name) -> bool`` that selects which components to save. When ``None`` (default), all components are saved. Examples:: @@ -119,13 +111,6 @@ def save( check_weights: Whether to verify that all initializers have weight data before saving. Defaults to ``True``. Set to ``False`` when saving skeleton models without weights. - max_workers: Number of threads used to materialize and write - tensors. Weights that need a dtype cast are held as - :class:`ir.LazyTensor` and converted at save time, so this - overlaps the casting work with disk writes. ``None`` (the - default) picks :func:`_default_save_workers`; pass ``1`` to - save serially. Peak memory stays bounded regardless of the - worker count. Only used when *external_data* is ``"onnx"``. Raises: ValueError: If *external_data* is not ``"onnx"`` or @@ -164,16 +149,7 @@ def save( callback=callback, ) else: - ir.save( - model, - path, - external_data="model.onnx.data", - max_shard_size_bytes=max_shard_size_bytes, - callback=callback, - max_workers=( - _default_save_workers() if max_workers is None else max_workers - ), - ) + ir.save(model, path, external_data="model.onnx.data", callback=callback) @classmethod def load(cls, directory: str) -> ModelPackage: @@ -282,57 +258,20 @@ def apply_weights( fold_initializers_after_weights(model) -def _default_save_workers() -> int: - """Pick a default thread count for writing external data. - - Deliberately a constant rather than a function of ``os.cpu_count()``. - - The bottleneck is the storage device, not the CPU: casting bf16 to f16 - runs around 12 GB/s while serial writes reach 1.6 GB/s, and a write-only - scaling test saturates at 2 threads. Nearly all of the win comes from - overlapping the cast with the write, and these threads spend their time - blocked in ``tofile`` with the GIL released rather than competing for CPU. - - Sizing the pool by core count gets this backwards. With torch pinned to a - single intra-op thread -- i.e. a small machine -- 8 workers is the *fastest* - configuration (1.55x) while 2 workers is slower than serial, because more - concurrency is needed to hide a slow serial cast. ``cpu_count() // 2`` would - give 1 on a 2-core box, disabling the optimization exactly where it helps - most. Oversubscription is also not penalized in practice: with full torch - threads, everything from 4 to 32 workers lands within noise of each other. - - Weights that live on an accelerator benefit the most. The device-to-host - copy happens inside ``tofile`` with the GIL released, so it overlaps with - other threads' writes: GPU-resident weights measured 6.3x at 8 workers - versus 1.5x for CPU weights, on a machine whose CPUs were largely idle. - """ - return _DEFAULT_SAVE_WORKERS - - def _make_progress_callback(): - """Create a tqdm progress-bar callback for ``ir.save``. - - ``ir.save`` invokes the callback from worker threads when ``max_workers`` - is greater than 1, and does not deliver calls in ``metadata.index`` order. - The bar therefore counts invocations rather than tracking indices, and - guards its own state with a lock. ``tqdm.update`` is itself thread-safe, - but ``total``/``set_description`` are not, and reading ``tensor.shape`` - concurrently with the write threads must not race the bar's own state. - """ + """Create a tqdm progress-bar callback for ``ir.save``.""" pbar = tqdm.tqdm() - lock = threading.Lock() total_set = False def callback(tensor: ir.TensorProtocol, metadata: ir.external_data.CallbackInfo) -> None: nonlocal total_set - with lock: - 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})" - ) + 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})" + ) return callback diff --git a/src/mobius/_model_package_test.py b/src/mobius/_model_package_test.py index c9f618628..12072f2e2 100644 --- a/src/mobius/_model_package_test.py +++ b/src/mobius/_model_package_test.py @@ -12,11 +12,7 @@ from mobius._builder import build_from_module from mobius._configs import VisionConfig -from mobius._model_package import ( - ModelPackage, - _default_save_workers, - _make_progress_callback, -) +from mobius._model_package import ModelPackage from mobius._testing import make_config from mobius.models.base import CausalLMModel from mobius.models.gemma3 import Gemma3MultiModalModel @@ -81,127 +77,6 @@ def test_config_stored(self): assert pkg.config is config -class TestParallelSave: - """Saving concurrently must match serial output and keep the bar correct.""" - - def test_parallel_save_matches_serial(self, tmp_path): - serial_dir = tmp_path / "serial" - parallel_dir = tmp_path / "parallel" - ModelPackage({"m": _make_simple_model("m")}).save( - str(serial_dir), max_workers=None, progress_bar=False - ) - ModelPackage({"m": _make_simple_model("m")}).save( - str(parallel_dir), max_workers=8, progress_bar=False - ) - serial_data = (serial_dir / "model.onnx.data").read_bytes() - parallel_data = (parallel_dir / "model.onnx.data").read_bytes() - assert serial_data == parallel_data - - def test_parallel_save_roundtrips(self, tmp_path): - pkg = ModelPackage({"m": _make_simple_model("m")}) - pkg.save(str(tmp_path), max_workers=8, progress_bar=False) - # A single-component package is saved flat as ``model.onnx``. - loaded = ModelPackage.load(str(tmp_path)) - assert set(loaded.data) == {"model"} - - 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, - max_workers=8, - progress_bar=False, - ) - - shards = sorted(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"} - - def test_progress_bar_counts_every_tensor_out_of_order(self): - # ir.save may invoke the callback from worker threads and out of index - # order, so the bar must count calls rather than follow ``index``. - callback = _make_progress_callback() - - class _Tensor: - name = "w" - shape = (2, 2) - - class dtype: # noqa: N801 - @staticmethod - def short_name(): - return "f32" - - total = 8 - for index in reversed(range(total)): - callback( - _Tensor(), - ir.external_data.CallbackInfo( - total=total, index=index, offset=0, filename="model.onnx.data" - ), - ) - # Closure state is private to the callback; assert via the bound bar. - bar = callback.__closure__[1].cell_contents - assert bar.total == total - assert bar.n == total - - def test_default_workers_does_not_track_core_count(self, monkeypatch): - # These threads block on I/O rather than competing for CPU, and a low - # core count needs *more* concurrency to hide a slow serial cast, so - # the default must not shrink on small machines. - for cores in (2, 8, 128, None): - monkeypatch.setattr("mobius._model_package.os.cpu_count", lambda c=cores: c) - assert _default_save_workers() == 8 - - def test_progress_bar_is_thread_safe(self): - import threading - - callback = _make_progress_callback() - - class _Tensor: - name = "w" - shape = (2, 2) - - class dtype: # noqa: N801 - @staticmethod - def short_name(): - return "f32" - - total = 200 - - def worker(index): - callback( - _Tensor(), - ir.external_data.CallbackInfo( - total=total, index=index, offset=0, filename="model.onnx.data" - ), - ) - - threads = [threading.Thread(target=worker, args=(i,)) for i in range(total)] - for thread in threads: - thread.start() - for thread in threads: - thread.join() - bar = callback.__closure__[1].cell_contents - assert bar.n == total - - class TestModelPackageSaveLoad: def test_save_creates_files(self, tmp_path): pkg = ModelPackage( From 1348af48884954e4d37b54bbbcf0b54cc78735ae Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Sat, 15 Aug 2026 21:33:34 -0700 Subject: [PATCH 08/14] Enable ONNX sharding with onnx_ir 1.0 Forward max_shard_size_bytes for ONNX external data, an API already available in onnx_ir 1.0.0. Keep Mobius independent of the pending parallel-save change by not exposing or passing max_workers. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: af83bc37-b9d6-4b94-be64-a4daa5a7ce40 Signed-off-by: Justin Chu --- src/mobius/__main__.py | 2 +- src/mobius/_model_package.py | 14 +++++++++++--- src/mobius/_model_package_test.py | 31 +++++++++++++++++++++++++++++++ 3 files changed, 43 insertions(+), 4 deletions(-) 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..fcbc5b3f4 100644 --- a/src/mobius/_model_package.py +++ b/src/mobius/_model_package.py @@ -95,8 +95,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:: @@ -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: diff --git a/src/mobius/_model_package_test.py b/src/mobius/_model_package_test.py index 12072f2e2..24ef2ea8a 100644 --- a/src/mobius/_model_package_test.py +++ b/src/mobius/_model_package_test.py @@ -77,6 +77,37 @@ 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.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 TestModelPackageSaveLoad: def test_save_creates_files(self, tmp_path): pkg = ModelPackage( From 0f0914ddc3c99c216d2ce8461a9e78f19be39f73 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Sat, 15 Aug 2026 21:38:45 -0700 Subject: [PATCH 09/14] Make save progress callback thread-safe Serialize tqdm mutations and count callback invocations rather than relying on index order. This supports concurrent callbacks from newer onnx_ir versions while remaining compatible with the serial callback behavior in onnx_ir 1.0.0. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: af83bc37-b9d6-4b94-be64-a4daa5a7ce40 Signed-off-by: Justin Chu --- src/mobius/_model_package.py | 25 ++++++++++----- src/mobius/_model_package_test.py | 52 ++++++++++++++++++++++++++++++- 2 files changed, 68 insertions(+), 9 deletions(-) diff --git a/src/mobius/_model_package.py b/src/mobius/_model_package.py index fcbc5b3f4..c55d0c4c7 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 @@ -267,19 +268,27 @@ def apply_weights( def _make_progress_callback(): - """Create a tqdm progress-bar callback for ``ir.save``.""" + """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 + serialize all progress-bar mutations. This remains compatible with + ``onnx_ir`` 1.0, where callbacks are invoked serially. + """ pbar = tqdm.tqdm() + lock = threading.Lock() total_set = False 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: + 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})" + ) return callback diff --git a/src/mobius/_model_package_test.py b/src/mobius/_model_package_test.py index 24ef2ea8a..6c00e4b31 100644 --- a/src/mobius/_model_package_test.py +++ b/src/mobius/_model_package_test.py @@ -6,13 +6,14 @@ from __future__ import annotations import logging +import threading 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 @@ -108,6 +109,55 @@ def test_onnx_external_data_is_sharded(self, 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" + + def test_counts_out_of_order_callbacks(self): + callback = _make_progress_callback() + total = 8 + + for index in reversed(range(total)): + callback( + self._Tensor(), + ir.external_data.CallbackInfo( + total=total, index=index, offset=0, filename="model.onnx.data" + ), + ) + + bar = callback.__closure__[1].cell_contents + assert bar.total == total + assert bar.n == total + + def test_is_thread_safe(self): + callback = _make_progress_callback() + total = 200 + + def invoke(index): + callback( + self._Tensor(), + ir.external_data.CallbackInfo( + total=total, index=index, offset=0, filename="model.onnx.data" + ), + ) + + threads = [threading.Thread(target=invoke, args=(index,)) for index in range(total)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + bar = callback.__closure__[1].cell_contents + assert bar.total == total + assert bar.n == total + + class TestModelPackageSaveLoad: def test_save_creates_files(self, tmp_path): pkg = ModelPackage( From 0f1a1e69fc21bf06088ddae5ea31c4a33d49eea3 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Sat, 15 Aug 2026 21:40:48 -0700 Subject: [PATCH 10/14] Show shard filename in save progress Include CallbackInfo.filename in the tqdm description so concurrent ONNX shard saves visibly identify the shard currently reporting progress. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: af83bc37-b9d6-4b94-be64-a4daa5a7ce40 Signed-off-by: Justin Chu --- src/mobius/_model_package.py | 3 ++- src/mobius/_model_package_test.py | 1 + 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/src/mobius/_model_package.py b/src/mobius/_model_package.py index c55d0c4c7..ae48ea39e 100644 --- a/src/mobius/_model_package.py +++ b/src/mobius/_model_package.py @@ -287,7 +287,8 @@ def callback(tensor: ir.TensorProtocol, metadata: ir.external_data.CallbackInfo) total_set = True pbar.update() pbar.set_description( - f"Saving {tensor.name} ({tensor.dtype.short_name()}, {tensor.shape})" + f"Saving {metadata.filename}: " + f"{tensor.name} ({tensor.dtype.short_name()}, {tensor.shape})" ) return callback diff --git a/src/mobius/_model_package_test.py b/src/mobius/_model_package_test.py index 6c00e4b31..ede84a49f 100644 --- a/src/mobius/_model_package_test.py +++ b/src/mobius/_model_package_test.py @@ -134,6 +134,7 @@ def test_counts_out_of_order_callbacks(self): bar = callback.__closure__[1].cell_contents assert bar.total == total assert bar.n == total + assert "model.onnx.data" in bar.desc def test_is_thread_safe(self): callback = _make_progress_callback() From ea27cf7035436fdf6b11d1b17224f31684d9d7cf Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Sat, 15 Aug 2026 21:42:57 -0700 Subject: [PATCH 11/14] Render one save progress bar per shard Use CallbackInfo shard metadata from newer onnx_ir releases to create a fixed-position tqdm bar for each external-data file. Fall back to one global bar with onnx_ir 1.0.0, which does not expose per-shard counts. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: af83bc37-b9d6-4b94-be64-a4daa5a7ce40 Signed-off-by: Justin Chu --- src/mobius/_model_package.py | 31 ++++++--- src/mobius/_model_package_test.py | 103 +++++++++++++++++++++++++----- 2 files changed, 109 insertions(+), 25 deletions(-) diff --git a/src/mobius/_model_package.py b/src/mobius/_model_package.py index ae48ea39e..7e3f55a0b 100644 --- a/src/mobius/_model_package.py +++ b/src/mobius/_model_package.py @@ -126,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() @@ -136,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: @@ -275,21 +274,33 @@ def _make_progress_callback(): serialize all progress-bar mutations. This remains compatible with ``onnx_ir`` 1.0, where callbacks are invoked serially. """ - pbar = tqdm.tqdm() lock = threading.Lock() - total_set = False + bars: dict[str, tqdm.tqdm] = {} def callback(tensor: ir.TensorProtocol, metadata: ir.external_data.CallbackInfo) -> None: - nonlocal total_set with lock: - if not total_set: - pbar.total = metadata.total - total_set = True + 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" + ) + pbar = tqdm.tqdm( + total=shard_total if shard_total is not None else metadata.total, + desc=description, + position=len(bars), + leave=True, + ) + bars[key] = pbar pbar.update() - pbar.set_description( - f"Saving {metadata.filename}: " + 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 ede84a49f..185d58102 100644 --- a/src/mobius/_model_package_test.py +++ b/src/mobius/_model_package_test.py @@ -7,6 +7,7 @@ import logging import threading +import types import onnx_ir as ir import torch @@ -119,32 +120,103 @@ class dtype: # noqa: N801 def short_name(): return "f32" - def test_counts_out_of_order_callbacks(self): + 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_creates_one_progress_bar_per_shard(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 = 8 + for filename in ("model-00001-of-00002.data", "model-00002-of-00002.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] == [0, 1] + assert all(bar.total == 2 and bar.n == 2 and bar.closed for bar in bars) + assert "model-00001-of-00002.data" in bars[0].desc + assert "model-00002-of-00002.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 - for index in reversed(range(total)): + monkeypatch.setattr("mobius._model_package.tqdm.tqdm", make_bar) + callback = _make_progress_callback() + for index in range(4): callback( self._Tensor(), - ir.external_data.CallbackInfo( - total=total, index=index, offset=0, filename="model.onnx.data" + types.SimpleNamespace( + total=4, + index=index, + offset=0, + filename=f"model-{index}.data", ), ) - bar = callback.__closure__[1].cell_contents - assert bar.total == total - assert bar.n == total - assert "model.onnx.data" in bar.desc + assert len(bars) == 1 + assert bars[0].total == 4 + assert bars[0].n == 4 + 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 - def test_is_thread_safe(self): + monkeypatch.setattr("mobius._model_package.tqdm.tqdm", make_bar) callback = _make_progress_callback() total = 200 def invoke(index): callback( self._Tensor(), - ir.external_data.CallbackInfo( - total=total, index=index, offset=0, filename="model.onnx.data" + types.SimpleNamespace( + total=total, + index=index, + offset=0, + filename="model.onnx.data", + shard_total=total, + shard_index=index, ), ) @@ -154,9 +226,10 @@ def invoke(index): for thread in threads: thread.join() - bar = callback.__closure__[1].cell_contents - assert bar.total == total - assert bar.n == total + assert len(bars) == 1 + assert bars[0].total == total + assert bars[0].n == total + assert bars[0].closed class TestModelPackageSaveLoad: From 10a30f7d9d9f29b31704004bc9b4d5d70a900aef Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Sat, 15 Aug 2026 21:44:21 -0700 Subject: [PATCH 12/14] Use onnx_data suffix for external weights Name the default external data file model.onnx_data so sharding naturally produces model-00001-of-00003.onnx_data rather than inserting the shard suffix between .onnx and .data. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: af83bc37-b9d6-4b94-be64-a4daa5a7ce40 Signed-off-by: Justin Chu --- docs/api/model_package.md | 2 +- docs/getting-started.md | 2 +- src/mobius/_model_package.py | 6 +++--- src/mobius/_model_package_test.py | 4 ++-- 4 files changed, 7 insertions(+), 7 deletions(-) diff --git a/docs/api/model_package.md b/docs/api/model_package.md index e5ae86af0..faff72421 100644 --- a/docs/api/model_package.md +++ b/docs/api/model_package.md @@ -93,5 +93,5 @@ pkg.save("output/llama/", external_data="safetensors") ## Output Layout -- **Single model**: `directory/model.onnx` + `directory/model.onnx.data` +- **Single model**: `directory/model.onnx` + `directory/model.onnx_data` - **Multi model**: `directory/{name}/model.onnx` for each component diff --git a/docs/getting-started.md b/docs/getting-started.md index 9c893a397..bdd55ffe6 100644 --- a/docs/getting-started.md +++ b/docs/getting-started.md @@ -223,7 +223,7 @@ collection of named `ir.Model` objects. ```python pkg.save("output/") -# Single model: output/model.onnx + output/model.onnx.data +# Single model: output/model.onnx + output/model.onnx_data # Multi model: output/model/model.onnx, output/vision/model.onnx, ... ``` diff --git a/src/mobius/_model_package.py b/src/mobius/_model_package.py index 7e3f55a0b..4859039f9 100644 --- a/src/mobius/_model_package.py +++ b/src/mobius/_model_package.py @@ -13,7 +13,7 @@ pkg = build("meta-llama/Llama-3-8B") pkg["model"] # ir.Model - pkg.save("/output/llama/") # saves model.onnx + model.onnx.data + pkg.save("/output/llama/") # saves model.onnx + model.onnx_data """ from __future__ import annotations @@ -86,7 +86,7 @@ def save( Args: directory: Path to the output directory (created if needed). external_data: External data format. ``"onnx"`` (default) saves - weights to ``model.onnx.data``. ``"safetensors"`` saves + weights to ``model.onnx_data``. ``"safetensors"`` saves weights in safetensors format. .. warning:: @@ -154,7 +154,7 @@ def save( ir.save( model, path, - external_data="model.onnx.data", + external_data="model.onnx_data", max_shard_size_bytes=max_shard_size_bytes, callback=callback, ) diff --git a/src/mobius/_model_package_test.py b/src/mobius/_model_package_test.py index 185d58102..7d87589a4 100644 --- a/src/mobius/_model_package_test.py +++ b/src/mobius/_model_package_test.py @@ -103,7 +103,7 @@ def test_onnx_external_data_is_sharded(self, tmp_path): progress_bar=False, ) - shards = sorted(tmp_path.glob("model.onnx-*-of-*.data")) + shards = sorted(tmp_path.glob("model-*-of-*.onnx_data")) assert len(shards) == 3 assert all(shard.stat().st_size <= 8192 for shard in shards) loaded = ModelPackage.load(str(tmp_path)) @@ -214,7 +214,7 @@ def invoke(index): total=total, index=index, offset=0, - filename="model.onnx.data", + filename="model.onnx_data", shard_total=total, shard_index=index, ), From 5da048e12d71e2f0a8240c1930442395f7d4e2aa Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Sat, 15 Aug 2026 21:58:36 -0700 Subject: [PATCH 13/14] Keep the conventional onnx.data filename Restore model.onnx.data as the default external-data location. Newer onnx_ir releases preserve this compound suffix when numbering shards; the compatibility test also accepts the legacy 1.0.0 shard spelling. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: af83bc37-b9d6-4b94-be64-a4daa5a7ce40 Signed-off-by: Justin Chu --- docs/api/model_package.md | 2 +- docs/getting-started.md | 2 +- src/mobius/_model_package.py | 6 +++--- src/mobius/_model_package_test.py | 11 +++++++++-- 4 files changed, 14 insertions(+), 7 deletions(-) diff --git a/docs/api/model_package.md b/docs/api/model_package.md index faff72421..e5ae86af0 100644 --- a/docs/api/model_package.md +++ b/docs/api/model_package.md @@ -93,5 +93,5 @@ pkg.save("output/llama/", external_data="safetensors") ## Output Layout -- **Single model**: `directory/model.onnx` + `directory/model.onnx_data` +- **Single model**: `directory/model.onnx` + `directory/model.onnx.data` - **Multi model**: `directory/{name}/model.onnx` for each component diff --git a/docs/getting-started.md b/docs/getting-started.md index bdd55ffe6..9c893a397 100644 --- a/docs/getting-started.md +++ b/docs/getting-started.md @@ -223,7 +223,7 @@ collection of named `ir.Model` objects. ```python pkg.save("output/") -# Single model: output/model.onnx + output/model.onnx_data +# Single model: output/model.onnx + output/model.onnx.data # Multi model: output/model/model.onnx, output/vision/model.onnx, ... ``` diff --git a/src/mobius/_model_package.py b/src/mobius/_model_package.py index 4859039f9..7e3f55a0b 100644 --- a/src/mobius/_model_package.py +++ b/src/mobius/_model_package.py @@ -13,7 +13,7 @@ pkg = build("meta-llama/Llama-3-8B") pkg["model"] # ir.Model - pkg.save("/output/llama/") # saves model.onnx + model.onnx_data + pkg.save("/output/llama/") # saves model.onnx + model.onnx.data """ from __future__ import annotations @@ -86,7 +86,7 @@ def save( Args: directory: Path to the output directory (created if needed). external_data: External data format. ``"onnx"`` (default) saves - weights to ``model.onnx_data``. ``"safetensors"`` saves + weights to ``model.onnx.data``. ``"safetensors"`` saves weights in safetensors format. .. warning:: @@ -154,7 +154,7 @@ def save( ir.save( model, path, - external_data="model.onnx_data", + external_data="model.onnx.data", max_shard_size_bytes=max_shard_size_bytes, callback=callback, ) diff --git a/src/mobius/_model_package_test.py b/src/mobius/_model_package_test.py index 7d87589a4..33e6c4160 100644 --- a/src/mobius/_model_package_test.py +++ b/src/mobius/_model_package_test.py @@ -103,7 +103,14 @@ def test_onnx_external_data_is_sharded(self, tmp_path): progress_bar=False, ) - shards = sorted(tmp_path.glob("model-*-of-*.onnx_data")) + 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)) @@ -214,7 +221,7 @@ def invoke(index): total=total, index=index, offset=0, - filename="model.onnx_data", + filename="model.onnx.data", shard_total=total, shard_index=index, ), From e4ae9e91d7473c9ba4c964156cb5f8b73c8b2b0b Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Sun, 16 Aug 2026 08:07:54 -0700 Subject: [PATCH 14/14] Order shard progress bars deterministically Derive each tqdm row from the one-based shard number embedded in the external-data filename, so concurrent callbacks cannot reorder the displayed bars. Cover shard-two-first callback delivery while retaining the onnx_ir 1.0 single-bar fallback. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- src/mobius/_model_package.py | 13 ++++++++++--- src/mobius/_model_package_test.py | 14 +++++++++----- 2 files changed, 19 insertions(+), 8 deletions(-) diff --git a/src/mobius/_model_package.py b/src/mobius/_model_package.py index 7e3f55a0b..a96dd990b 100644 --- a/src/mobius/_model_package.py +++ b/src/mobius/_model_package.py @@ -271,8 +271,9 @@ def _make_progress_callback(): Newer ``onnx_ir`` versions may invoke callbacks concurrently and out of index order. Count invocations instead of tracking ``metadata.index`` and - serialize all progress-bar mutations. This remains compatible with - ``onnx_ir`` 1.0, where callbacks are invoked serially. + 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] = {} @@ -288,10 +289,16 @@ def callback(tensor: ir.TensorProtocol, metadata: ir.external_data.CallbackInfo) 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=len(bars), + position=position, leave=True, ) bars[key] = pbar diff --git a/src/mobius/_model_package_test.py b/src/mobius/_model_package_test.py index 33e6c4160..954727379 100644 --- a/src/mobius/_model_package_test.py +++ b/src/mobius/_model_package_test.py @@ -146,7 +146,7 @@ def set_postfix_str(self, value): def close(self): self.closed = True - def test_creates_one_progress_bar_per_shard(self, monkeypatch): + def test_orders_progress_bars_by_shard_number(self, monkeypatch): bars = [] def make_bar(**kwargs): @@ -156,7 +156,10 @@ def make_bar(**kwargs): monkeypatch.setattr("mobius._model_package.tqdm.tqdm", make_bar) callback = _make_progress_callback() - for filename in ("model-00001-of-00002.data", "model-00002-of-00002.data"): + 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(), @@ -171,10 +174,10 @@ def make_bar(**kwargs): ) assert len(bars) == 2 - assert [bar.position for bar in bars] == [0, 1] + 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-00001-of-00002.data" in bars[0].desc - assert "model-00002-of-00002.data" in bars[1].desc + 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 = [] @@ -200,6 +203,7 @@ def make_bar(**kwargs): 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):