Address review comments

Signed-off-by: YASH Nankani <ynankani@2u1g-x570-0073.ipp2a1.colossus.nvidia.com>

Author: YASH Nankani

modelopt/torch/export/diffusers_utils.py

@@ -918,7 +918,11 @@ def _ceil_div(a: int, b: int) -> int:
         return (a + b - 1) // b
 
     def _to_blocked(input_matrix: torch.Tensor) -> torch.Tensor:
-        """Rearrange scale matrix to cuBLAS 2-D block-scaling-factors layout."""
+        """Rearrange scale matrix to cuBLAS 2-D block-scaling-factors layout.
+
+        Note: rows are padded to multiples of 128 for cuBLAS alignment, so the
+        output shape may differ from the input (e.g. (16, 4) -> (128, 4)).
+        """
         rows, cols = input_matrix.shape
         n_row_blocks = _ceil_div(rows, 128)
         n_col_blocks = _ceil_div(cols, 4)

modelopt/torch/export/unified_export_hf.py

@@ -28,7 +28,8 @@
 
 import torch
 import torch.nn as nn
-from safetensors.torch import load_file, save_file
+from safetensors import safe_open
+from safetensors.torch import save_file
 
 from .diffusers_utils import build_layerwise_quant_metadata, pad_nvfp4_weights, swizzle_nvfp4_scales
 
@@ -180,8 +181,6 @@ def _postprocess_safetensors(
         padding_strategy: ``"row"``, ``"row_col"``, or None.
         enable_swizzle_layout: Whether to swizzle block scales.
     """
-    import struct
-
     safetensor_files = sorted(export_dir.glob("*.safetensors"))
     if not safetensor_files:
         return
@@ -195,22 +194,16 @@ def _postprocess_safetensors(
         )
 
     for sf_path in safetensor_files:
-        sd = load_file(str(sf_path))
-
-        with open(sf_path, "rb") as f:
-            header_size = struct.unpack("<Q", f.read(8))[0]
-            header = json.loads(f.read(header_size))
-        metadata = header.get("__metadata__", None) or {}
-
-        # Clone tensors so the memory-mapped file handle from load_file is
-        # released before we overwrite the same path (required on Windows).
-        sd = {k: v.clone() for k, v in sd.items()}
+        with safe_open(str(sf_path), framework="pt") as f:
+            metadata = dict(f.metadata() or {})
+            sd = {k: f.get_tensor(k).clone() for k in f.keys()}
 
         if merged_base_safetensor_path is not None and model_type is not None:
             sd, base_metadata = merge_diffusion_checkpoint(
-                sd, merged_base_safetensor_path, model_type, hf_quant_config
+                sd, merged_base_safetensor_path, model_type, hf_quant_config=None
             )
-            metadata.update(base_metadata)
+            base_metadata.update(metadata)
+            metadata = base_metadata
 
         if padding_strategy is not None:
             sd = pad_nvfp4_weights(sd, padding_strategy)

tests/unit/torch/export/test_nvfp4_utils.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -200,6 +200,35 @@ def test_sharded_guard(self, tmp_path):
                 enable_layerwise_quant_metadata=True,
             )
 
+    def test_preserves_existing_metadata(self, tmp_path):
+        """Simulate save_pretrained output: safetensors with pre-existing metadata."""
+        from modelopt.torch.export.unified_export_hf import _postprocess_safetensors
+
+        sd = _make_nvfp4_state_dict(rows=20, cols=64)
+        preexisting_metadata = {"format": "pt", "_class_name": "MyModel"}
+        save_file(sd, str(tmp_path / "model.safetensors"), metadata=preexisting_metadata)
+
+        hf_quant_config = {"quant_algo": "NVFP4"}
+        _postprocess_safetensors(
+            tmp_path,
+            hf_quant_config=hf_quant_config,
+            padding_strategy="row",
+            enable_swizzle_layout=True,
+            enable_layerwise_quant_metadata=True,
+        )
+
+        reloaded = load_file(str(tmp_path / "model.safetensors"))
+        assert reloaded["layer0.weight"].shape[0] == 32
+        assert reloaded["layer0.weight_scale"].shape == (128, 64 // 16)
+
+        with safe_open(str(tmp_path / "model.safetensors"), framework="pt") as f:
+            metadata = f.metadata()
+        assert metadata["format"] == "pt"
+        assert metadata["_class_name"] == "MyModel"
+        assert json.loads(metadata["quantization_config"]) == hf_quant_config
+        layer_meta = json.loads(metadata["_quantization_metadata"])
+        assert "layer0" in layer_meta["layers"]
+
     def test_no_safetensor_files(self, tmp_path):
         from modelopt.torch.export.unified_export_hf import _postprocess_safetensors