From 233905836a04a0ad0e4269a687a7f32a77e92473 Mon Sep 17 00:00:00 2001
From: mgohil <mgohil@nvidia.com>
Date: Thu, 22 Oct 2026 08:22:03 +0000
Subject: [PATCH 1/9] [2/17] Add shared ONNX model saving for DLA surgeries

Add a shared graph-surgery saver with staged output replacement and binary
external-weight writing to avoid ONNX 1.21 Windows trailing-0x1A corruption.
Cover FP16, FP32, and INT8 payloads, repeated saves, nested paths, and
invalid-location refusal. Preserve the existing dtype converter and public
surgery implementations.

Signed-off-by: mgohil <mgohil@nvidia.com>
---
 .../onnx/graph_surgery/utils/graph_utils.py   | 79 +++++++++++++++++++
 .../onnx/graph_surgery/test_graph_utils.py    | 50 ++++++++++++
 2 files changed, 129 insertions(+)

diff --git a/modelopt/onnx/graph_surgery/utils/graph_utils.py b/modelopt/onnx/graph_surgery/utils/graph_utils.py
index b20341fe154..8e255e8982d 100644
--- a/modelopt/onnx/graph_surgery/utils/graph_utils.py
+++ b/modelopt/onnx/graph_surgery/utils/graph_utils.py
@@ -18,8 +18,10 @@
 import gc
 import math
 import os
+import tempfile
 from collections import defaultdict, deque
 from collections.abc import Iterable
+from pathlib import Path
 
 import numpy as np
 import onnx
@@ -61,6 +63,7 @@
     "parse_qdq_quantized_dtype_list",
     "remove_node",
     "resolve_reshape_shape",
+    "save_onnx_model",
     "set_node_attr_i",
     "topological_sort_nodes",
     "uses_external_data",
@@ -90,6 +93,82 @@ def uses_external_data(model_path: str) -> bool:
     return any(os.path.exists(os.path.join(model_dir, pattern)) for pattern in external_patterns)
 
 
+def save_onnx_model(
+    model: onnx.ModelProto,
+    output_path: str,
+    *,
+    use_external_data: bool = True,
+    external_data_name: str | None = None,
+    size_threshold: int = 1024,
+    verbose: bool = True,
+) -> None:
+    """Save a loaded ONNX model, optionally storing large weights in a sidecar.
+
+    Args:
+        model: Loaded model; externalizing weights updates its tensor storage metadata.
+        output_path: Destination ONNX file.
+        use_external_data: Store large tensor payloads in an external file.
+        external_data_name: Relative sidecar path; defaults to the ONNX filename plus ``_data``.
+        size_threshold: Minimum tensor size in bytes to externalize.
+        verbose: Log the destination.
+    """
+    if verbose:
+        logger.info("Saving modified model to: %s", output_path)
+    output_dir = os.path.dirname(output_path)
+    location = Path(external_data_name or os.path.basename(output_path) + "_data")
+    if use_external_data:
+        if (
+            external_data_name == ""
+            or location.anchor
+            or not location.name
+            or ".." in location.parts
+            or location == Path(Path(output_path).name)
+        ):
+            raise ValueError(
+                "External data must have a distinct relative path inside the output directory"
+            )
+    if output_dir:
+        os.makedirs(output_dir, exist_ok=True)
+    if not use_external_data:
+        onnx.save(model, output_path)
+        return
+    with tempfile.TemporaryDirectory(dir=output_dir or ".") as temporary:
+        staged_model = Path(temporary) / Path(output_path).name
+        staged_weights = Path(temporary) / location
+        staged_weights.parent.mkdir(parents=True, exist_ok=True)
+        _save_external_weights(
+            model, staged_model, staged_weights, location.as_posix(), size_threshold
+        )
+        if staged_weights.exists():
+            destination = Path(output_dir or ".") / location
+            destination.parent.mkdir(parents=True, exist_ok=True)
+            os.replace(staged_weights, destination)
+        os.replace(staged_model, output_path)
+
+
+# TODO: Once ONNX >= 1.22.0 is required, use its external writer instead (onnx/onnx#7790).
+def _save_external_weights(model, model_path, weights_path, location, size_threshold):
+    onnx.external_data_helper.convert_model_to_external_data(
+        model, location=location, size_threshold=size_threshold
+    )
+    tensors = [
+        tensor
+        for tensor in onnx.external_data_helper._get_all_tensors(model)
+        if onnx.external_data_helper.uses_external_data(tensor) and tensor.HasField("raw_data")
+    ]
+    if tensors:
+        # ONNX 1.21 reopens sidecars in Windows text mode, truncating trailing 0x1A.
+        with weights_path.open("wb") as stream:
+            for tensor in tensors:
+                data, offset = tensor.raw_data, stream.tell()
+                stream.write(data)
+                onnx.external_data_helper.set_external_data(
+                    tensor, location, offset=offset, length=len(data)
+                )
+                tensor.ClearField("raw_data")
+    onnx.save(model, model_path)
+
+
 def topological_sort_nodes(graph: onnx.GraphProto) -> list[onnx.NodeProto]:
     """Topologically sort nodes in the graph for clean ordering.
 
diff --git a/tests/unit/onnx/graph_surgery/test_graph_utils.py b/tests/unit/onnx/graph_surgery/test_graph_utils.py
index f8bbdb25109..5f47f7386d8 100644
--- a/tests/unit/onnx/graph_surgery/test_graph_utils.py
+++ b/tests/unit/onnx/graph_surgery/test_graph_utils.py
@@ -34,6 +34,7 @@
     insert_nodes_at_position,
     parse_qdq_quantized_dtype_list,
     resolve_reshape_shape,
+    save_onnx_model,
 )
 
 
@@ -378,3 +379,52 @@ def test_qdq_dtype_filters_normalize_aliases_and_reject_unknown_types():
     ) == frozenset({TensorProto.INT16, TensorProto.UINT16})
     with pytest.raises(ValueError):
         parse_qdq_quantized_dtype_list(["int16", "unit16"])
+
+
+@pytest.mark.parametrize("dtype", [np.float16, np.float32, np.int8])
+def test_repeated_external_save_keeps_sidecar_bounded(tmp_path, dtype):
+    """Round-trip adjacent payloads, including a Windows text-mode EOF byte."""
+    path = tmp_path / "model.onnx"
+    location = "weights/data.bin" if dtype == np.float32 else "model.onnx_data"
+    values = np.arange(2048, dtype=np.float32).astype(dtype)
+    for offset in [0, 1]:
+        first, second = values + offset, np.full_like(values, 39 + offset)
+        first.view(np.uint8)[-1] = 0x1A
+        model = helper.make_model(
+            helper.make_graph(
+                [],
+                "weights",
+                [],
+                [
+                    helper.make_tensor_value_info(
+                        name, helper.np_dtype_to_tensor_dtype(np.dtype(dtype)), [2048]
+                    )
+                    for name in ("w", "other")
+                ],
+                [numpy_helper.from_array(first, "w"), numpy_helper.from_array(second, "other")],
+            )
+        )
+        save_onnx_model(model, str(path), external_data_name=location, verbose=False)
+        assert (tmp_path / location).stat().st_size == 2 * values.nbytes
+        saved = onnx.load(path)
+        onnx.checker.check_model(saved)
+        for tensor, expected in zip(saved.graph.initializer, (first, second), strict=True):
+            assert numpy_helper.to_array(tensor).tobytes() == expected.tobytes()
+
+
+@pytest.mark.parametrize("location", ["../weights.bin", "absolute", "model.onnx", ""])
+def test_invalid_external_location_preserves_model_and_existing_output(tmp_path, location):
+    model = helper.make_model(
+        helper.make_graph(
+            [], "save", [], [], [numpy_helper.from_array(np.ones(2048, np.float32), "w")]
+        )
+    )
+    before = model.SerializeToString()
+    path = tmp_path / "model.onnx"
+    path.write_bytes(b"unchanged")
+    if location == "absolute":
+        location = str(tmp_path / "weights.bin")
+    with pytest.raises(ValueError):
+        save_onnx_model(model, str(path), external_data_name=location, verbose=False)
+    assert model.SerializeToString() == before
+    assert path.read_bytes() == b"unchanged"

From 08818766abe4a32c5ef8468fc2e88840efb3a1f7 Mon Sep 17 00:00:00 2001
From: mgohil <mgohil@nvidia.com>
Date: Fri, 23 Oct 2026 03:37:31 +0000
Subject: [PATCH 2/9] Fix ONNX external save validation and rollback

Validate resolved sidecar paths and reject symbolic or multiply linked files. Preserve existing output weights and restore caller tensor data and metadata when replacement fails, while retaining successful-save mutation and avoiding a full weight copy.

Signed-off-by: mgohil <mgohil@nvidia.com>
---
 .../onnx/graph_surgery/utils/graph_utils.py   | 142 +++++++++---
 .../onnx/graph_surgery/test_graph_utils.py    | 204 +++++++++++++++++-
 2 files changed, 316 insertions(+), 30 deletions(-)

diff --git a/modelopt/onnx/graph_surgery/utils/graph_utils.py b/modelopt/onnx/graph_surgery/utils/graph_utils.py
index 8e255e8982d..cd9acf6c4c1 100644
--- a/modelopt/onnx/graph_surgery/utils/graph_utils.py
+++ b/modelopt/onnx/graph_surgery/utils/graph_utils.py
@@ -21,6 +21,7 @@
 import tempfile
 from collections import defaultdict, deque
 from collections.abc import Iterable
+from contextlib import contextmanager
 from pathlib import Path
 
 import numpy as np
@@ -105,18 +106,27 @@ def save_onnx_model(
     """Save a loaded ONNX model, optionally storing large weights in a sidecar.
 
     Args:
-        model: Loaded model; externalizing weights updates its tensor storage metadata.
+        model: Model with external tensor payloads loaded; successful external saves update
+            its tensor storage metadata.
         output_path: Destination ONNX file.
         use_external_data: Store large tensor payloads in an external file.
         external_data_name: Relative sidecar path; defaults to the ONNX filename plus ``_data``.
         size_threshold: Minimum tensor size in bytes to externalize.
         verbose: Log the destination.
     """
+    for tensor in onnx.external_data_helper._get_all_tensors(model):
+        if onnx.external_data_helper.uses_external_data(tensor) and not tensor.HasField("raw_data"):
+            raise ValueError(
+                f"External tensor {tensor.name!r} has no raw_data; "
+                "load external data before saving (onnx.load(..., load_external_data=True))"
+            )
     if verbose:
         logger.info("Saving modified model to: %s", output_path)
     output_dir = os.path.dirname(output_path)
     location = Path(external_data_name or os.path.basename(output_path) + "_data")
     if use_external_data:
+        output_root = Path(output_dir or ".").resolve()
+        destination = output_root / location
         if (
             external_data_name == ""
             or location.anchor
@@ -127,6 +137,7 @@ def save_onnx_model(
             raise ValueError(
                 "External data must have a distinct relative path inside the output directory"
             )
+        _validate_external_data_destination(output_root, destination, Path(output_path))
     if output_dir:
         os.makedirs(output_dir, exist_ok=True)
     if not use_external_data:
@@ -136,37 +147,110 @@ def save_onnx_model(
         staged_model = Path(temporary) / Path(output_path).name
         staged_weights = Path(temporary) / location
         staged_weights.parent.mkdir(parents=True, exist_ok=True)
-        _save_external_weights(
-            model, staged_model, staged_weights, location.as_posix(), size_threshold
-        )
-        if staged_weights.exists():
-            destination = Path(output_dir or ".") / location
-            destination.parent.mkdir(parents=True, exist_ok=True)
-            os.replace(staged_weights, destination)
-        os.replace(staged_model, output_path)
+        backup_weights = None
+        weights_replaced = False
+        try:
+            with _save_external_weights(
+                model,
+                staged_model,
+                staged_weights,
+                location.as_posix(),
+                size_threshold,
+                destination,
+            ):
+                if staged_weights.exists():
+                    # Recheck after staging in case a link changed during the write.
+                    _validate_external_data_destination(output_root, destination, Path(output_path))
+                    destination.parent.mkdir(parents=True, exist_ok=True)
+                    if destination.exists():
+                        # Keep the backup outside staging so a failed rollback cannot delete it.
+                        backup_weights = Path(tempfile.mkdtemp(dir=destination.parent)) / "weights"
+                        os.replace(destination, backup_weights)
+                    os.replace(staged_weights, destination)
+                    weights_replaced = True
+                os.replace(staged_model, output_path)
+        except OSError:
+            if backup_weights is not None and backup_weights.exists():
+                os.replace(backup_weights, destination)
+            elif weights_replaced:
+                destination.unlink()
+            raise
+        else:
+            if backup_weights is not None:
+                backup_weights.unlink()
+        finally:
+            if backup_weights is not None and not backup_weights.exists():
+                backup_weights.parent.rmdir()
+
+
+def _validate_external_data_destination(
+    output_root: Path, destination: Path, output_path: Path
+) -> None:
+    resolved_destination = destination.resolve()
+    if not resolved_destination.is_relative_to(output_root):
+        raise ValueError("External data must resolve inside the output directory")
+    if resolved_destination == output_path.resolve():
+        raise ValueError("External data must have a distinct path from the output model")
+    if destination.is_symlink():
+        raise ValueError("External data destination must not be a symbolic link")
+    if destination.exists() and not destination.is_file():
+        raise ValueError("External data destination must be a regular file")
+    if destination.exists() and destination.stat().st_nlink > 1:
+        raise ValueError("External data destination must not have multiple hard links")
 
 
 # TODO: Once ONNX >= 1.22.0 is required, use its external writer instead (onnx/onnx#7790).
-def _save_external_weights(model, model_path, weights_path, location, size_threshold):
-    onnx.external_data_helper.convert_model_to_external_data(
-        model, location=location, size_threshold=size_threshold
-    )
-    tensors = [
-        tensor
-        for tensor in onnx.external_data_helper._get_all_tensors(model)
-        if onnx.external_data_helper.uses_external_data(tensor) and tensor.HasField("raw_data")
-    ]
-    if tensors:
-        # ONNX 1.21 reopens sidecars in Windows text mode, truncating trailing 0x1A.
-        with weights_path.open("wb") as stream:
-            for tensor in tensors:
-                data, offset = tensor.raw_data, stream.tell()
-                stream.write(data)
-                onnx.external_data_helper.set_external_data(
-                    tensor, location, offset=offset, length=len(data)
-                )
-                tensor.ClearField("raw_data")
-    onnx.save(model, model_path)
+@contextmanager
+def _save_external_weights(model, model_path, weights_path, location, size_threshold, destination):
+    states = []
+    for tensor in onnx.external_data_helper._get_all_tensors(model):
+        metadata = TensorProto()
+        if tensor.HasField("data_location"):
+            metadata.data_location = tensor.data_location
+        metadata.external_data.extend(tensor.external_data)
+        states.append((tensor, metadata))
+    payloads = []
+    try:
+        onnx.external_data_helper.convert_model_to_external_data(
+            model, location=location, size_threshold=size_threshold
+        )
+        tensors = [
+            tensor for tensor, _ in states if onnx.external_data_helper.uses_external_data(tensor)
+        ]
+        if tensors:
+            # ONNX 1.21 reopens sidecars in Windows text mode, truncating trailing 0x1A.
+            with weights_path.open("wb") as stream:
+                for tensor in tensors:
+                    data, offset = tensor.raw_data, stream.tell()
+                    stream.write(data)
+                    onnx.external_data_helper.set_external_data(
+                        tensor, location, offset=offset, length=len(data)
+                    )
+                    payloads.append((tensor, offset, len(data)))
+                    tensor.ClearField("raw_data")
+                    del data
+        onnx.save(model, model_path)
+        yield
+    except BaseException:
+        # Restore from staging or the installed sidecar before filesystem rollback replaces it.
+        if payloads:
+            source = weights_path if weights_path.exists() else destination
+            with source.open("rb") as stream:
+                for tensor, offset, length in payloads:
+                    stream.seek(offset)
+                    data = stream.read(length)
+                    if len(data) != length:
+                        raise OSError(
+                            f"Cannot restore tensor {tensor.name!r}: expected {length} bytes, "
+                            f"read {len(data)}"
+                        )
+                    tensor.raw_data = data
+                    del data
+        for tensor, metadata in states:
+            tensor.ClearField("data_location")
+            tensor.ClearField("external_data")
+            tensor.MergeFrom(metadata)
+        raise
 
 
 def topological_sort_nodes(graph: onnx.GraphProto) -> list[onnx.NodeProto]:
diff --git a/tests/unit/onnx/graph_surgery/test_graph_utils.py b/tests/unit/onnx/graph_surgery/test_graph_utils.py
index 5f47f7386d8..2e3f2b5f9e9 100644
--- a/tests/unit/onnx/graph_surgery/test_graph_utils.py
+++ b/tests/unit/onnx/graph_surgery/test_graph_utils.py
@@ -16,6 +16,7 @@
 """Shared graph queries, shape resolution, and initializer contracts."""
 
 import copy
+import os
 
 import numpy as np
 import onnx
@@ -405,6 +406,10 @@ def test_repeated_external_save_keeps_sidecar_bounded(tmp_path, dtype):
             )
         )
         save_onnx_model(model, str(path), external_data_name=location, verbose=False)
+        assert all(
+            tensor.data_location == TensorProto.EXTERNAL and not tensor.HasField("raw_data")
+            for tensor in model.graph.initializer
+        )
         assert (tmp_path / location).stat().st_size == 2 * values.nbytes
         saved = onnx.load(path)
         onnx.checker.check_model(saved)
@@ -412,7 +417,9 @@ def test_repeated_external_save_keeps_sidecar_bounded(tmp_path, dtype):
             assert numpy_helper.to_array(tensor).tobytes() == expected.tobytes()
 
 
-@pytest.mark.parametrize("location", ["../weights.bin", "absolute", "model.onnx", ""])
+@pytest.mark.parametrize(
+    "location", ["../weights.bin", "absolute", "model.onnx", "", "directory", "model_alias"]
+)
 def test_invalid_external_location_preserves_model_and_existing_output(tmp_path, location):
     model = helper.make_model(
         helper.make_graph(
@@ -424,7 +431,202 @@ def test_invalid_external_location_preserves_model_and_existing_output(tmp_path,
     path.write_bytes(b"unchanged")
     if location == "absolute":
         location = str(tmp_path / "weights.bin")
+    elif location == "directory":
+        (tmp_path / location).mkdir()
+        (tmp_path / location / "existing.bin").write_bytes(b"existing weights")
+    elif location == "model_alias":
+        try:
+            (tmp_path / "alias").symlink_to(tmp_path, target_is_directory=True)
+        except OSError as error:
+            pytest.skip(f"Directory symlinks are unavailable: {error}")
+        location = "alias/model.onnx"
     with pytest.raises(ValueError):
         save_onnx_model(model, str(path), external_data_name=location, verbose=False)
     assert model.SerializeToString() == before
     assert path.read_bytes() == b"unchanged"
+    if location == "directory":
+        assert (tmp_path / location / "existing.bin").read_bytes() == b"existing weights"
+
+
+@pytest.mark.parametrize(
+    ("link_kind", "contained"),
+    [
+        ("directory", False),
+        ("directory", True),
+        ("file", False),
+        ("file", True),
+        ("dangling", True),
+        ("hardlink", True),
+    ],
+)
+def test_external_save_checks_sidecar_links(tmp_path, link_kind, contained):
+    """Allow contained directory links and reject file links without changing existing data."""
+    output = tmp_path / "output"
+    output.mkdir()
+    target = output / "target" if contained else tmp_path / "outside"
+    target.mkdir()
+    sidecar = target / "data.bin"
+    if link_kind != "dangling":
+        sidecar.write_bytes(b"existing weights")
+    destination = output / "weights" / "data.bin"
+    try:
+        if link_kind == "directory":
+            destination.parent.symlink_to(target, target_is_directory=True)
+        else:
+            destination.parent.mkdir()
+            if link_kind == "hardlink":
+                destination.hardlink_to(sidecar)
+            else:
+                destination.symlink_to(sidecar)
+    except OSError as error:
+        pytest.skip(f"{link_kind} links are unavailable: {error}")
+    path = output / "model.onnx"
+    path.write_bytes(b"existing model")
+    values = np.ones(2048, np.float32)
+    model = helper.make_model(
+        helper.make_graph([], "save", [], [], [numpy_helper.from_array(values, "w")])
+    )
+    before = model.SerializeToString()
+    if link_kind == "directory" and contained:
+        save_onnx_model(model, str(path), external_data_name="weights/data.bin", verbose=False)
+        assert sidecar.read_bytes() == values.tobytes()
+        saved = onnx.load(path, load_external_data=False)
+        assert saved.graph.initializer[0].data_location == TensorProto.EXTERNAL
+    else:
+        with pytest.raises(ValueError):
+            save_onnx_model(model, str(path), external_data_name="weights/data.bin", verbose=False)
+        assert model.SerializeToString() == before
+        assert path.read_bytes() == b"existing model"
+        if link_kind == "dangling":
+            assert destination.is_symlink()
+            assert not sidecar.exists()
+        else:
+            assert sidecar.read_bytes() == b"existing weights"
+            assert destination.read_bytes() == b"existing weights"
+
+
+@pytest.mark.parametrize("existing_sidecar", [False, True])
+@pytest.mark.parametrize("failure_stage", ["weights", "model", "short_read"])
+@pytest.mark.parametrize("external_metadata", [False, True])
+def test_external_save_rolls_back_failed_replacement(
+    tmp_path, monkeypatch, existing_sidecar, failure_stage, external_metadata
+):
+    """Keep the previous model and weights usable when either replacement fails."""
+    path = tmp_path / "model.onnx"
+    sidecar = tmp_path / "weights" / "data.bin"
+    values = np.ones(2048, np.float32)
+    original = helper.make_model(
+        helper.make_graph(
+            [],
+            "save",
+            [],
+            [helper.make_tensor_value_info("w", TensorProto.FLOAT, [2048])],
+            [
+                numpy_helper.from_array(values, "w"),
+                numpy_helper.from_array(values + 2, "other"),
+            ],
+        )
+    )
+    save_onnx_model(
+        original,
+        str(path),
+        use_external_data=existing_sidecar,
+        external_data_name="weights/data.bin",
+        verbose=False,
+    )
+    original_bytes = path.read_bytes()
+    original_weights = sidecar.read_bytes() if existing_sidecar else None
+    candidate = onnx.load(path)
+    candidate.graph.initializer[0].CopyFrom(numpy_helper.from_array(values + 1, "w"))
+    candidate.graph.initializer[1].data_location = TensorProto.DEFAULT
+    if external_metadata:
+        onnx.external_data_helper.set_external_data(
+            candidate.graph.initializer[0],
+            "original.weights",
+            offset=17,
+            length=values.nbytes,
+            checksum="original-checksum",
+            basepath="original-directory",
+        )
+    candidate_bytes = candidate.SerializeToString()
+    failure_path = sidecar if failure_stage == "weights" else path
+    short_read = failure_stage == "short_read"
+    failure = PermissionError("Replacement denied")
+    replace = os.replace
+    failed = False
+
+    def fail_selected_replacement(source, destination):
+        nonlocal failed
+        if not failed and os.fspath(destination) == str(failure_path):
+            failed = True
+            if short_read:
+                with sidecar.open("r+b") as stream:
+                    stream.truncate(values.nbytes - 1)
+            raise failure
+        return replace(source, destination)
+
+    monkeypatch.setattr(os, "replace", fail_selected_replacement)
+    with pytest.raises(OSError) as error:
+        save_onnx_model(candidate, str(path), external_data_name="weights/data.bin", verbose=False)
+    if short_read:
+        assert not candidate.graph.initializer[0].HasField("raw_data")
+        assert error.value.__context__ is failure
+    else:
+        assert error.value is failure
+        assert candidate.SerializeToString() == candidate_bytes
+    assert path.read_bytes() == original_bytes
+    if existing_sidecar:
+        assert sidecar.read_bytes() == original_weights
+    else:
+        assert not sidecar.exists()
+    restored = onnx.load(path)
+    onnx.checker.check_model(restored)
+    assert numpy_helper.to_array(restored.graph.initializer[0]).tobytes() == values.tobytes()
+    expected_files = {path, sidecar} if existing_sidecar else {path}
+    assert {entry for entry in tmp_path.rglob("*") if entry.is_file()} == expected_files
+    if short_read:
+        return
+    save_onnx_model(candidate, str(path), external_data_name="weights/data.bin", verbose=False)
+    assert candidate.graph.initializer[0].data_location == TensorProto.EXTERNAL
+    assert not candidate.graph.initializer[0].HasField("raw_data")
+    retried = onnx.load(path)
+    onnx.checker.check_model(retried)
+    assert numpy_helper.to_array(retried.graph.initializer[0]).tobytes() == (values + 1).tobytes()
+    assert numpy_helper.to_array(retried.graph.initializer[1]).tobytes() == (values + 2).tobytes()
+
+
+@pytest.mark.parametrize("existing_output", [False, True])
+@pytest.mark.parametrize("use_external_data", [False, True])
+def test_save_rejects_unloaded_external_tensors_before_changes(
+    tmp_path, existing_output, use_external_data
+):
+    """Reject partially loaded models before creating directories or changing existing outputs."""
+    values = np.ones(2048, np.float32)
+    source = tmp_path / "source" / "model.onnx"
+    model = helper.make_model(
+        helper.make_graph(
+            [],
+            "save",
+            [],
+            [helper.make_tensor_value_info("first", TensorProto.FLOAT, [2048])],
+            [numpy_helper.from_array(values, name) for name in ("first", "missing")],
+        )
+    )
+    save_onnx_model(model, str(source), verbose=False)
+    unloaded = onnx.load(source, load_external_data=False)
+    unloaded.graph.initializer[0].CopyFrom(numpy_helper.from_array(values, "first"))
+    before = unloaded.SerializeToString()
+    output = tmp_path / "output"
+    path = output / "model.onnx"
+    if existing_output:
+        output.mkdir()
+        path.write_bytes(b"existing model")
+        (output / "model.onnx_data").write_bytes(b"existing weights")
+    files_before = {entry: entry.read_bytes() for entry in tmp_path.rglob("*") if entry.is_file()}
+    with pytest.raises(ValueError):
+        save_onnx_model(unloaded, str(path), use_external_data=use_external_data, verbose=False)
+    assert unloaded.SerializeToString() == before
+    assert output.exists() == existing_output
+    assert {
+        entry: entry.read_bytes() for entry in tmp_path.rglob("*") if entry.is_file()
+    } == files_before

From efbfb22eab090fa4f88cce624758bfb23942d5e0 Mon Sep 17 00:00:00 2001
From: mgohil <mgohil@nvidia.com>
Date: Fri, 23 Oct 2026 05:21:26 +0000
Subject: [PATCH 3/9] Document ONNX external save helpers

Signed-off-by: mgohil <mgohil@nvidia.com>
---
 modelopt/onnx/graph_surgery/utils/graph_utils.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/modelopt/onnx/graph_surgery/utils/graph_utils.py b/modelopt/onnx/graph_surgery/utils/graph_utils.py
index cd9acf6c4c1..bc20aed2358 100644
--- a/modelopt/onnx/graph_surgery/utils/graph_utils.py
+++ b/modelopt/onnx/graph_surgery/utils/graph_utils.py
@@ -186,6 +186,7 @@ def save_onnx_model(
 def _validate_external_data_destination(
     output_root: Path, destination: Path, output_path: Path
 ) -> None:
+    """Require a contained sidecar path distinct from the model and reject unsafe file links."""
     resolved_destination = destination.resolve()
     if not resolved_destination.is_relative_to(output_root):
         raise ValueError("External data must resolve inside the output directory")
@@ -202,6 +203,10 @@ def _validate_external_data_destination(
 # TODO: Once ONNX >= 1.22.0 is required, use its external writer instead (onnx/onnx#7790).
 @contextmanager
 def _save_external_weights(model, model_path, weights_path, location, size_threshold, destination):
+    """Stage external weights and restore caller tensors if the replacement context fails.
+
+    Successful saves retain the updated external storage metadata.
+    """
     states = []
     for tensor in onnx.external_data_helper._get_all_tensors(model):
         metadata = TensorProto()

From 62c45b41657363ad85f0e680ee3a4a44a1e93842 Mon Sep 17 00:00:00 2001
From: mgohil <mgohil@nvidia.com>
Date: Fri, 23 Oct 2026 08:38:00 +0000
Subject: [PATCH 4/9] Fix repeated external ONNX saves in the output directory

Signed-off-by: mgohil <mgohil@nvidia.com>
---
 modelopt/onnx/graph_surgery/utils/graph_utils.py  | 3 ++-
 tests/unit/onnx/graph_surgery/test_graph_utils.py | 4 +++-
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/modelopt/onnx/graph_surgery/utils/graph_utils.py b/modelopt/onnx/graph_surgery/utils/graph_utils.py
index bc20aed2358..ff8fc32ec40 100644
--- a/modelopt/onnx/graph_surgery/utils/graph_utils.py
+++ b/modelopt/onnx/graph_surgery/utils/graph_utils.py
@@ -216,8 +216,9 @@ def _save_external_weights(model, model_path, weights_path, location, size_thres
         states.append((tensor, metadata))
     payloads = []
     try:
+        # Assign the final location only when writing; its existing sidecar is still live.
         onnx.external_data_helper.convert_model_to_external_data(
-            model, location=location, size_threshold=size_threshold
+            model, size_threshold=size_threshold
         )
         tensors = [
             tensor for tensor, _ in states if onnx.external_data_helper.uses_external_data(tensor)
diff --git a/tests/unit/onnx/graph_surgery/test_graph_utils.py b/tests/unit/onnx/graph_surgery/test_graph_utils.py
index 2e3f2b5f9e9..8f215083da4 100644
--- a/tests/unit/onnx/graph_surgery/test_graph_utils.py
+++ b/tests/unit/onnx/graph_surgery/test_graph_utils.py
@@ -383,8 +383,9 @@ def test_qdq_dtype_filters_normalize_aliases_and_reject_unknown_types():
 
 
 @pytest.mark.parametrize("dtype", [np.float16, np.float32, np.int8])
-def test_repeated_external_save_keeps_sidecar_bounded(tmp_path, dtype):
+def test_repeated_external_save_keeps_sidecar_bounded(tmp_path, monkeypatch, dtype):
     """Round-trip adjacent payloads, including a Windows text-mode EOF byte."""
+    monkeypatch.chdir(tmp_path)
     path = tmp_path / "model.onnx"
     location = "weights/data.bin" if dtype == np.float32 else "model.onnx_data"
     values = np.arange(2048, dtype=np.float32).astype(dtype)
@@ -421,6 +422,7 @@ def test_repeated_external_save_keeps_sidecar_bounded(tmp_path, dtype):
     "location", ["../weights.bin", "absolute", "model.onnx", "", "directory", "model_alias"]
 )
 def test_invalid_external_location_preserves_model_and_existing_output(tmp_path, location):
+    """Reject invalid sidecar paths before changing caller tensors or existing files."""
     model = helper.make_model(
         helper.make_graph(
             [], "save", [], [], [numpy_helper.from_array(np.ones(2048, np.float32), "w")]

From fd948cb55986762cdf609d286c9237cde83532b7 Mon Sep 17 00:00:00 2001
From: mgohil <mgohil@nvidia.com>
Date: Fri, 23 Oct 2026 08:46:07 +0000
Subject: [PATCH 5/9] Reject aliases between ONNX models and external weights

Signed-off-by: mgohil <mgohil@nvidia.com>
---
 .../onnx/graph_surgery/utils/graph_utils.py   | 11 ++++++--
 .../onnx/graph_surgery/test_graph_utils.py    | 26 +++++++++++++++++++
 2 files changed, 35 insertions(+), 2 deletions(-)

diff --git a/modelopt/onnx/graph_surgery/utils/graph_utils.py b/modelopt/onnx/graph_surgery/utils/graph_utils.py
index ff8fc32ec40..7c9eef6c8f8 100644
--- a/modelopt/onnx/graph_surgery/utils/graph_utils.py
+++ b/modelopt/onnx/graph_surgery/utils/graph_utils.py
@@ -133,6 +133,7 @@ def save_onnx_model(
             or not location.name
             or ".." in location.parts
             or location == Path(Path(output_path).name)
+            or (os.name == "nt" and any(":" in part for part in location.parts))
         ):
             raise ValueError(
                 "External data must have a distinct relative path inside the output directory"
@@ -144,8 +145,10 @@ def save_onnx_model(
         onnx.save(model, output_path)
         return
     with tempfile.TemporaryDirectory(dir=output_dir or ".") as temporary:
-        staged_model = Path(temporary) / Path(output_path).name
-        staged_weights = Path(temporary) / location
+        # Separate roots prevent filesystem filename aliases from overwriting staged payloads.
+        staged_model = Path(temporary) / "model" / Path(output_path).name
+        staged_weights = Path(temporary) / "weights" / location
+        staged_model.parent.mkdir()
         staged_weights.parent.mkdir(parents=True, exist_ok=True)
         backup_weights = None
         weights_replaced = False
@@ -168,6 +171,10 @@ def save_onnx_model(
                         os.replace(destination, backup_weights)
                     os.replace(staged_weights, destination)
                     weights_replaced = True
+                    if os.path.exists(output_path) and os.path.samefile(destination, output_path):
+                        raise OSError(
+                            "External data must have a distinct path from the output model"
+                        )
                 os.replace(staged_model, output_path)
         except OSError:
             if backup_weights is not None and backup_weights.exists():
diff --git a/tests/unit/onnx/graph_surgery/test_graph_utils.py b/tests/unit/onnx/graph_surgery/test_graph_utils.py
index 8f215083da4..2386c0dfbc3 100644
--- a/tests/unit/onnx/graph_surgery/test_graph_utils.py
+++ b/tests/unit/onnx/graph_surgery/test_graph_utils.py
@@ -507,6 +507,32 @@ def test_external_save_checks_sidecar_links(tmp_path, link_kind, contained):
             assert destination.read_bytes() == b"existing weights"
 
 
+@pytest.mark.skipif(os.name != "nt", reason="Win32 filename aliases and alternate data streams")
+@pytest.mark.parametrize("existing_output", [False, True])
+@pytest.mark.parametrize(
+    "location", ["MODEL.onnx", "model.onnx.", "model.onnx ", "model.onnx:weights"]
+)
+def test_external_save_rejects_windows_model_aliases(tmp_path, existing_output, location):
+    """Reject Win32 aliases without losing caller payloads or the existing model."""
+    path = tmp_path / "model.onnx"
+    if existing_output:
+        path.write_bytes(b"original model")
+    model = helper.make_model(
+        helper.make_graph(
+            [], "alias", [], [], [numpy_helper.from_array(np.ones(2048, np.float32), "w")]
+        )
+    )
+    before = model.SerializeToString()
+    with pytest.raises((OSError, ValueError)):
+        save_onnx_model(model, str(path), external_data_name=location, verbose=False)
+    assert model.SerializeToString() == before
+    if existing_output:
+        assert path.read_bytes() == b"original model"
+    else:
+        assert not path.exists()
+    assert list(tmp_path.iterdir()) == ([path] if existing_output else [])
+
+
 @pytest.mark.parametrize("existing_sidecar", [False, True])
 @pytest.mark.parametrize("failure_stage", ["weights", "model", "short_read"])
 @pytest.mark.parametrize("external_metadata", [False, True])

From 4613e7fa85264cafccfe19052af121720e876b11 Mon Sep 17 00:00:00 2001
From: mgohil <mgohil@nvidia.com>
Date: Fri, 23 Oct 2026 08:51:39 +0000
Subject: [PATCH 6/9] Document the external-save failure injection helper

Signed-off-by: mgohil <mgohil@nvidia.com>
---
 tests/unit/onnx/graph_surgery/test_graph_utils.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/unit/onnx/graph_surgery/test_graph_utils.py b/tests/unit/onnx/graph_surgery/test_graph_utils.py
index 2386c0dfbc3..6be967a8092 100644
--- a/tests/unit/onnx/graph_surgery/test_graph_utils.py
+++ b/tests/unit/onnx/graph_surgery/test_graph_utils.py
@@ -584,6 +584,7 @@ def test_external_save_rolls_back_failed_replacement(
     failed = False
 
     def fail_selected_replacement(source, destination):
+        """Fail one selected replacement to verify tensor and file rollback."""
         nonlocal failed
         if not failed and os.fspath(destination) == str(failure_path):
             failed = True

From 031f2b5a816ee0abcc88b7daa3e58b003c19f414 Mon Sep 17 00:00:00 2001
From: mgohil <mgohil@nvidia.com>
Date: Fri, 23 Oct 2026 09:22:26 +0000
Subject: [PATCH 7/9] Honor inline ONNX saves for loaded external tensors

Signed-off-by: mgohil <mgohil@nvidia.com>
---
 .../onnx/graph_surgery/utils/graph_utils.py   | 43 +++++++++----
 .../onnx/graph_surgery/test_graph_utils.py    | 64 +++++++++++++++++++
 2 files changed, 95 insertions(+), 12 deletions(-)

diff --git a/modelopt/onnx/graph_surgery/utils/graph_utils.py b/modelopt/onnx/graph_surgery/utils/graph_utils.py
index 7c9eef6c8f8..33dd80bd596 100644
--- a/modelopt/onnx/graph_surgery/utils/graph_utils.py
+++ b/modelopt/onnx/graph_surgery/utils/graph_utils.py
@@ -142,7 +142,13 @@ def save_onnx_model(
     if output_dir:
         os.makedirs(output_dir, exist_ok=True)
     if not use_external_data:
-        onnx.save(model, output_path)
+        states = _snapshot_tensor_storage(model, external_only=True)
+        try:
+            if states:
+                onnx.external_data_helper.convert_model_from_external_data(model)
+            onnx.save(model, output_path)
+        finally:
+            _restore_tensor_storage(states)
         return
     with tempfile.TemporaryDirectory(dir=output_dir or ".") as temporary:
         # Separate roots prevent filesystem filename aliases from overwriting staged payloads.
@@ -190,6 +196,28 @@ def save_onnx_model(
                 backup_weights.parent.rmdir()
 
 
+def _snapshot_tensor_storage(model, *, external_only=False):
+    """Snapshot storage flags and external entries without copying tensor payloads."""
+    states = []
+    for tensor in onnx.external_data_helper._get_all_tensors(model):
+        if external_only and not onnx.external_data_helper.uses_external_data(tensor):
+            continue
+        metadata = TensorProto()
+        if tensor.HasField("data_location"):
+            metadata.data_location = tensor.data_location
+        metadata.external_data.extend(tensor.external_data)
+        states.append((tensor, metadata))
+    return states
+
+
+def _restore_tensor_storage(states):
+    """Restore storage metadata, including field presence and external-entry order."""
+    for tensor, metadata in states:
+        tensor.ClearField("data_location")
+        tensor.ClearField("external_data")
+        tensor.MergeFrom(metadata)
+
+
 def _validate_external_data_destination(
     output_root: Path, destination: Path, output_path: Path
 ) -> None:
@@ -214,13 +242,7 @@ def _save_external_weights(model, model_path, weights_path, location, size_thres
 
     Successful saves retain the updated external storage metadata.
     """
-    states = []
-    for tensor in onnx.external_data_helper._get_all_tensors(model):
-        metadata = TensorProto()
-        if tensor.HasField("data_location"):
-            metadata.data_location = tensor.data_location
-        metadata.external_data.extend(tensor.external_data)
-        states.append((tensor, metadata))
+    states = _snapshot_tensor_storage(model)
     payloads = []
     try:
         # Assign the final location only when writing; its existing sidecar is still live.
@@ -259,10 +281,7 @@ def _save_external_weights(model, model_path, weights_path, location, size_thres
                         )
                     tensor.raw_data = data
                     del data
-        for tensor, metadata in states:
-            tensor.ClearField("data_location")
-            tensor.ClearField("external_data")
-            tensor.MergeFrom(metadata)
+        _restore_tensor_storage(states)
         raise
 
 
diff --git a/tests/unit/onnx/graph_surgery/test_graph_utils.py b/tests/unit/onnx/graph_surgery/test_graph_utils.py
index 6be967a8092..4a68bbb3c6b 100644
--- a/tests/unit/onnx/graph_surgery/test_graph_utils.py
+++ b/tests/unit/onnx/graph_surgery/test_graph_utils.py
@@ -659,3 +659,67 @@ def test_save_rejects_unloaded_external_tensors_before_changes(
     assert {
         entry: entry.read_bytes() for entry in tmp_path.rglob("*") if entry.is_file()
     } == files_before
+
+
+@pytest.mark.parametrize("storage", ["initializer", "attribute", "subgraph"])
+@pytest.mark.parametrize("write_failure", [False, True])
+def test_inline_save_preserves_loaded_external_tensors(tmp_path, storage, write_failure):
+    """Inline loaded external payloads without changing caller state or existing sidecars."""
+    values = np.arange(2048, dtype=np.float32)
+    values.view(np.uint8)[-1] = 0x1A
+    weight = numpy_helper.from_array(values, "weight")
+    onnx.external_data_helper.set_external_data(
+        weight, "original.weights", offset=7, length=values.nbytes, checksum="original-checksum"
+    )
+    value = helper.make_tensor_value_info("y", TensorProto.FLOAT, [2048])
+    if storage == "initializer":
+        graph = helper.make_graph(
+            [],
+            "inline",
+            [],
+            [helper.make_tensor_value_info("weight", TensorProto.FLOAT, [2048])],
+            [weight],
+        )
+    elif storage == "attribute":
+        graph = helper.make_graph(
+            [helper.make_node("Constant", [], ["y"], value=weight)], "inline", [], [value]
+        )
+    else:
+        branch = helper.make_graph(
+            [helper.make_node("Constant", [], ["y"], value=weight)], "branch", [], [value]
+        )
+        graph = helper.make_graph(
+            [helper.make_node("If", ["condition"], ["y"], then_branch=branch, else_branch=branch)],
+            "inline",
+            [],
+            [value],
+            [numpy_helper.from_array(np.asarray(True), "condition")],
+        )
+    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)], ir_version=9)
+    before = model.SerializeToString()
+    sidecar = tmp_path / "original.weights"
+    sidecar.write_bytes(b"original sidecar")
+    path = tmp_path / "inline.onnx"
+    if write_failure:
+        path.mkdir()
+        with pytest.raises(OSError):
+            save_onnx_model(model, str(path), use_external_data=False, verbose=False)
+    else:
+        save_onnx_model(model, str(path), use_external_data=False, verbose=False)
+        saved = onnx.load(path, load_external_data=False)
+        onnx.checker.check_model(saved)
+        weights = [
+            tensor
+            for tensor in onnx.external_data_helper._get_all_tensors(saved)
+            if tensor.name == "weight"
+        ]
+        assert len(weights) == (2 if storage == "subgraph" else 1)
+        for tensor in weights:
+            assert tensor.data_location != TensorProto.EXTERNAL
+            assert not tensor.external_data
+            assert numpy_helper.to_array(tensor).tobytes() == values.tobytes()
+    assert model.SerializeToString() == before
+    assert sidecar.read_bytes() == b"original sidecar"
+    assert {entry for entry in tmp_path.iterdir() if entry.is_file()} == (
+        {sidecar} if write_failure else {sidecar, path}
+    )

From 81191a5af6a4bf1aa7e94aa65c613b1cde46ebd2 Mon Sep 17 00:00:00 2001
From: mgohil <mgohil@nvidia.com>
Date: Fri, 23 Oct 2026 10:28:42 +0000
Subject: [PATCH 8/9] Stage inline ONNX models before replacing existing files

Signed-off-by: mgohil <mgohil@nvidia.com>
---
 .../onnx/graph_surgery/utils/graph_utils.py   |  5 +-
 .../onnx/graph_surgery/test_graph_utils.py    | 61 +++++++++++++++++++
 2 files changed, 65 insertions(+), 1 deletion(-)

diff --git a/modelopt/onnx/graph_surgery/utils/graph_utils.py b/modelopt/onnx/graph_surgery/utils/graph_utils.py
index 33dd80bd596..0e2674ca0df 100644
--- a/modelopt/onnx/graph_surgery/utils/graph_utils.py
+++ b/modelopt/onnx/graph_surgery/utils/graph_utils.py
@@ -146,7 +146,10 @@ def save_onnx_model(
         try:
             if states:
                 onnx.external_data_helper.convert_model_from_external_data(model)
-            onnx.save(model, output_path)
+            with tempfile.TemporaryDirectory(dir=output_dir or ".") as temporary:
+                staged_model = Path(temporary) / Path(output_path).name
+                onnx.save(model, staged_model)
+                os.replace(staged_model, output_path)
         finally:
             _restore_tensor_storage(states)
         return
diff --git a/tests/unit/onnx/graph_surgery/test_graph_utils.py b/tests/unit/onnx/graph_surgery/test_graph_utils.py
index 4a68bbb3c6b..bd04d4dda68 100644
--- a/tests/unit/onnx/graph_surgery/test_graph_utils.py
+++ b/tests/unit/onnx/graph_surgery/test_graph_utils.py
@@ -723,3 +723,64 @@ def test_inline_save_preserves_loaded_external_tensors(tmp_path, storage, write_
     assert {entry for entry in tmp_path.iterdir() if entry.is_file()} == (
         {sidecar} if write_failure else {sidecar, path}
     )
+
+
+@pytest.mark.parametrize("external_metadata", [False, True])
+@pytest.mark.parametrize("failure_stage", ["write", "replace"])
+def test_inline_save_failure_preserves_existing_model(
+    tmp_path, monkeypatch, external_metadata, failure_stage
+):
+    """Preserve existing files and caller tensors after partial writes or replacement errors."""
+    path = tmp_path / "model.onnx"
+    values = np.arange(2048, dtype=np.float32)
+    model = helper.make_model(
+        helper.make_graph(
+            [],
+            "inline_failure",
+            [],
+            [helper.make_tensor_value_info("w", TensorProto.FLOAT, [2048])],
+            [numpy_helper.from_array(values, "w")],
+        )
+    )
+    onnx.save(model, path)
+    original = path.read_bytes()
+    model.graph.initializer[0].CopyFrom(numpy_helper.from_array(values + 1, "w"))
+    sidecar = tmp_path / "original.weights"
+    sidecar.write_bytes(b"existing sidecar")
+    if external_metadata:
+        onnx.external_data_helper.set_external_data(
+            model.graph.initializer[0], "original.weights", offset=7, length=values.nbytes
+        )
+    before = model.SerializeToString()
+    failure = OSError("Injected inline save failure")
+    write_bytes = onnx._save_bytes
+
+    def fail_write(content, destination):
+        """Use the real writer to leave partial serialized bytes before failing."""
+        write_bytes(content[: len(content) // 2], destination)
+        raise failure
+
+    def fail_replace(source, destination):
+        """Inject a replacement failure after the staged model was saved."""
+        raise failure
+
+    with monkeypatch.context() as patch:
+        if failure_stage == "write":
+            patch.setattr(onnx, "_save_bytes", fail_write)
+        else:
+            patch.setattr(os, "replace", fail_replace)
+        with pytest.raises(OSError) as error:
+            save_onnx_model(model, str(path), use_external_data=False, verbose=False)
+        assert error.value is failure
+    assert path.read_bytes() == original
+    assert model.SerializeToString() == before
+    assert sidecar.read_bytes() == b"existing sidecar"
+    assert {entry.name for entry in tmp_path.iterdir()} == {"model.onnx", "original.weights"}
+    onnx.checker.check_model(onnx.load(path))
+    save_onnx_model(model, str(path), use_external_data=False, verbose=False)
+    saved = onnx.load(path, load_external_data=False)
+    onnx.checker.check_model(saved)
+    assert numpy_helper.to_array(saved.graph.initializer[0]).tobytes() == (values + 1).tobytes()
+    assert not saved.graph.initializer[0].external_data
+    assert model.SerializeToString() == before
+    assert sidecar.read_bytes() == b"existing sidecar"

From 4fc0ba472a962e6dfb15d1cbd1e7f7e2ffb5040c Mon Sep 17 00:00:00 2001
From: mgohil <mgohil@nvidia.com>
Date: Fri, 23 Oct 2026 10:49:50 +0000
Subject: [PATCH 9/9] Preserve ONNX tensor payloads until sidecar close
 succeeds

Signed-off-by: mgohil <mgohil@nvidia.com>
---
 .../onnx/graph_surgery/utils/graph_utils.py   |  8 +-
 .../onnx/graph_surgery/test_graph_utils.py    | 78 +++++++++++++++++++
 2 files changed, 84 insertions(+), 2 deletions(-)

diff --git a/modelopt/onnx/graph_surgery/utils/graph_utils.py b/modelopt/onnx/graph_surgery/utils/graph_utils.py
index 0e2674ca0df..e8e83619cb3 100644
--- a/modelopt/onnx/graph_surgery/utils/graph_utils.py
+++ b/modelopt/onnx/graph_surgery/utils/graph_utils.py
@@ -257,6 +257,7 @@ def _save_external_weights(model, model_path, weights_path, location, size_thres
         ]
         if tensors:
             # ONNX 1.21 reopens sidecars in Windows text mode, truncating trailing 0x1A.
+            staged_payloads = []
             with weights_path.open("wb") as stream:
                 for tensor in tensors:
                     data, offset = tensor.raw_data, stream.tell()
@@ -264,9 +265,12 @@ def _save_external_weights(model, model_path, weights_path, location, size_thres
                     onnx.external_data_helper.set_external_data(
                         tensor, location, offset=offset, length=len(data)
                     )
-                    payloads.append((tensor, offset, len(data)))
-                    tensor.ClearField("raw_data")
+                    staged_payloads.append((tensor, offset, len(data)))
                     del data
+            # A buffered close may fail; retain caller payloads until it succeeds.
+            payloads = staged_payloads
+            for tensor, _, _ in payloads:
+                tensor.ClearField("raw_data")
         onnx.save(model, model_path)
         yield
     except BaseException:
diff --git a/tests/unit/onnx/graph_surgery/test_graph_utils.py b/tests/unit/onnx/graph_surgery/test_graph_utils.py
index bd04d4dda68..a367294f617 100644
--- a/tests/unit/onnx/graph_surgery/test_graph_utils.py
+++ b/tests/unit/onnx/graph_surgery/test_graph_utils.py
@@ -16,7 +16,9 @@
 """Shared graph queries, shape resolution, and initializer contracts."""
 
 import copy
+import io
 import os
+from pathlib import Path
 
 import numpy as np
 import onnx
@@ -784,3 +786,79 @@ def fail_replace(source, destination):
     assert not saved.graph.initializer[0].external_data
     assert model.SerializeToString() == before
     assert sidecar.read_bytes() == b"existing sidecar"
+
+
+@pytest.mark.parametrize("external_metadata", [False, True])
+@pytest.mark.parametrize("failure_stage", ["write", "close"])
+def test_external_save_preserves_payloads_when_buffered_writes_fail(
+    tmp_path, monkeypatch, external_metadata, failure_stage
+):
+    """Keep loaded payloads and prior files intact when a buffered sidecar write fails."""
+    path = tmp_path / "model.onnx"
+    sidecar = tmp_path / "weights" / "data.bin"
+    values = np.arange(8, dtype=np.float32)
+    model = helper.make_model(
+        helper.make_graph(
+            [],
+            "buffered_failure",
+            [],
+            [helper.make_tensor_value_info("w", TensorProto.FLOAT, [8])],
+            [
+                numpy_helper.from_array(values, "w"),
+                numpy_helper.from_array(values + 1, "other"),
+            ],
+        )
+    )
+    save_onnx_model(
+        model, str(path), external_data_name="weights/data.bin", size_threshold=0, verbose=False
+    )
+    original, original_weights = path.read_bytes(), sidecar.read_bytes()
+    candidate = onnx.load(path)
+    candidate.graph.initializer[0].CopyFrom(numpy_helper.from_array(values + 10, "w"))
+    if external_metadata:
+        onnx.external_data_helper.set_external_data(
+            candidate.graph.initializer[0], "original.weights", offset=7, length=values.nbytes
+        )
+    before = candidate.SerializeToString()
+    failure = OSError("Injected buffered sidecar failure")
+    open_path = Path.open
+
+    class FailingFile(io.FileIO):
+        """Leave a real partial sidecar when the buffered writer flushes."""
+
+        def write(self, data):
+            """Write a single byte before reporting an injected I/O error."""
+            super().write(data[:1])
+            raise failure
+
+    def open_sidecar(file_path, mode="r", *args, **kwargs):
+        """Fail a real buffer flush during the second write or during close."""
+        if mode == "wb" and file_path.name == "data.bin":
+            buffer_size = 48 if failure_stage == "write" else 128
+            return io.BufferedWriter(FailingFile(file_path, "w"), buffer_size=buffer_size)
+        return open_path(file_path, mode, *args, **kwargs)
+
+    with monkeypatch.context() as patch:
+        patch.setattr(Path, "open", open_sidecar)
+        with pytest.raises(OSError) as error:
+            save_onnx_model(
+                candidate,
+                str(path),
+                external_data_name="weights/data.bin",
+                size_threshold=0,
+                verbose=False,
+            )
+    assert candidate.SerializeToString() == before
+    assert error.value is failure
+    assert path.read_bytes() == original
+    assert sidecar.read_bytes() == original_weights
+    assert {entry for entry in tmp_path.rglob("*") if entry.is_file()} == {path, sidecar}
+    assert {entry for entry in tmp_path.rglob("*") if entry.is_dir()} == {sidecar.parent}
+    onnx.checker.check_model(onnx.load(path))
+    save_onnx_model(
+        candidate, str(path), external_data_name="weights/data.bin", size_threshold=0, verbose=False
+    )
+    saved = onnx.load(path)
+    onnx.checker.check_model(saved)
+    for tensor, expected in zip(saved.graph.initializer, (values + 10, values + 1), strict=True):
+        assert numpy_helper.to_array(tensor).tobytes() == expected.tobytes()