[Feat][Web] Support per-parameter tensor cache encoding

python/tvm/contrib/tvmjs.py

@@ -17,6 +17,7 @@
 # ruff: noqa: E501, F401
 """Namespace to store utilities for building web runtime."""
 
+import copy
 import hashlib
 import json
 import math
@@ -26,7 +27,6 @@
 # pylint: disable=unused-import
 import sys
 from collections.abc import Iterator, Mapping
-from types import GeneratorType
 from typing import Any, Optional, Union
 
 import numpy as np
@@ -203,7 +203,7 @@ def dump_tensor_cache(
     params: Mapping[str, np.ndarray | tvm.runtime.Tensor]
     | Iterator[tuple[str, np.ndarray | tvm.runtime.Tensor]],
     cache_dir: str,
-    encode_format="f32-to-bf16",
+    encode_format: str | Mapping[str, str] = "f32-to-bf16",
     meta_data=None,
     shard_cap_mb=32,
     show_progress: bool = True,
@@ -222,8 +222,12 @@ def dump_tensor_cache(
     cache_dir: str
         The path to the cache
 
-    encode_format: {"f32-to-bf16", "raw"}
-        Encoding format.
+    encode_format: Union[
+        {"f32-to-bf16", "raw"},
+        Mapping[str, {"f32-to-bf16", "raw"}]
+    ]
+        Encoding format. A mapping selects the format by parameter name and may
+        use ``"*"`` as the fallback for names not explicitly listed.
 
     meta_data: json-compatible-struct or Callable[[], Any]
         Extra meta_data to be stored in the cache json file,
@@ -239,20 +243,32 @@ def dump_tensor_cache(
         If the cache already exists, update the cache. When set to False, it will overwrite the
         existing files.
     """
-    if encode_format not in ("raw", "f32-to-bf16"):
-        raise ValueError(f"Invalie encode_format {encode_format}")
+    if isinstance(encode_format, str):
+        if encode_format not in ("raw", "f32-to-bf16"):
+            raise ValueError(f"Invalid encode_format {encode_format}")
+    elif not isinstance(encode_format, Mapping):
+        raise TypeError("encode_format must be a string or parameter-name mapping")
+    else:
+        for name, param_format in encode_format.items():
+            if param_format not in ("raw", "f32-to-bf16"):
+                raise ValueError(f"Invalid encode_format for parameter {name}: {param_format}")
+
+    def resolve_encode_format(name):
+        if isinstance(encode_format, str):
+            return encode_format
+        param_format = encode_format.get(name, encode_format.get("*"))
+        if param_format not in ("raw", "f32-to-bf16"):
+            raise ValueError(f"Invalid encode_format for parameter {name}: {param_format}")
+        return param_format
 
     records = []
-    from_generator = isinstance(params, GeneratorType)
     total_bytes = 0
     counter = 0
     max_out_length = 0
 
     if not os.path.exists(cache_dir):
         os.makedirs(cache_dir)
 
-    f32_to_bf16_triggered = False
-
     print(f"Start storing to cache {cache_dir}")
     shard_cap_nbytes = shard_cap_mb * (1 << 20)
 
@@ -268,8 +284,9 @@ def dump_tensor_cache(
         cache_dir, "params_shard", shard_cap_nbytes, initial_shard_records=records
     )
 
-    param_generator = params.items() if not from_generator else params
+    param_generator = params.items() if isinstance(params, Mapping) else params
     for k, origin_v in param_generator:
+        param_encode_format = resolve_encode_format(k)
         shape = list(origin_v.shape)
         v = origin_v
         if not isinstance(v, np.ndarray):
@@ -286,9 +303,8 @@ def dump_tensor_cache(
         total_bytes += math.prod(v.shape) * np.dtype(v.dtype).itemsize
 
         # convert fp32 to bf16
-        if encode_format == "f32-to-bf16" and dtype == "float32":
+        if param_encode_format == "f32-to-bf16" and dtype == "float32":
             data = _convert_f32_to_bf16(v).tobytes()
-            f32_to_bf16_triggered = True
         else:
             data = v.tobytes()
 
@@ -297,7 +313,7 @@ def dump_tensor_cache(
             name=k,
             shape=shape,
             dtype=dtype,
-            encode_format=encode_format,
+            encode_format=param_encode_format,
             allow_update=update_if_exists,
         )
 
@@ -317,17 +333,25 @@ def dump_tensor_cache(
         f"\nAll finished, {shard_manager.counter} total shards committed, record saved to {nd_cache_json}"
     )
 
-    if f32_to_bf16_triggered:
-        for shard in records:
+    b16_nd_cache_json = os.path.join(cache_dir, "tensor-cache-b16.json")
+    has_f32_to_bf16 = any(
+        item["format"] == "f32-to-bf16" and item["dtype"] == "float32"
+        for shard in records
+        for item in shard["records"]
+    )
+    if has_f32_to_bf16:
+        b16_records = copy.deepcopy(records)
+        for shard in b16_records:
             for item in shard["records"]:
-                if item["dtype"] == "float32":
+                if item["format"] == "f32-to-bf16" and item["dtype"] == "float32":
                     item["format"] = "raw"
                     item["dtype"] = "bfloat16"
-        b16_nd_cache_json = os.path.join(cache_dir, "tensor-cache-b16.json")
         # also dump a file that contains bf16
         with open(b16_nd_cache_json, "w") as outfile:
-            json.dump({"metadata": meta_data, "records": records}, outfile, indent=4)
+            json.dump({"metadata": meta_data, "records": b16_records}, outfile, indent=4)
         print(f"Also saved a bf16 record to {b16_nd_cache_json}")
+    elif os.path.exists(b16_nd_cache_json):
+        os.remove(b16_nd_cache_json)
 
 
 def load_tensor_cache(cachepath: str, device: tvm.runtime.Device):

tests/python/contrib/test_tvmjs.py

@@ -17,6 +17,8 @@
 
 """Test contrib.tvmjs"""
 
+import json
+import os
 import tempfile
 
 import numpy as np
@@ -60,5 +62,135 @@ def test_save_load_float8(dtype):
     np.testing.assert_array_equal(arr, after_roundtrip)
 
 
+def _records_by_name(manifest_path):
+    with open(manifest_path, encoding="utf-8") as source:
+        manifest = json.load(source)
+    records = {
+        record["name"]: record for shard in manifest["records"] for record in shard["records"]
+    }
+    return records, manifest
+
+
+def test_dump_tensor_cache_supports_per_parameter_encoding_roundtrip():
+    raw = np.array([0.1234567, -0.7654321], dtype="float32")
+    compressed = np.array([0.2345678, -0.8765432], dtype="float32")
+
+    with tempfile.TemporaryDirectory(prefix="tvm_") as temp_dir:
+        tvmjs.dump_tensor_cache(
+            {"raw": raw, "compressed": compressed},
+            temp_dir,
+            encode_format={"raw": "raw", "*": "f32-to-bf16"},
+        )
+        cache, _ = tvmjs.load_tensor_cache(temp_dir, tvm.cpu())
+        b16_cache, _ = tvmjs.load_tensor_cache(
+            os.path.join(temp_dir, "tensor-cache-b16.json"), tvm.cpu()
+        )
+
+        records, _ = _records_by_name(os.path.join(temp_dir, "tensor-cache.json"))
+        b16_records, _ = _records_by_name(os.path.join(temp_dir, "tensor-cache-b16.json"))
+
+    assert records["raw"]["format"] == "raw"
+    assert records["raw"]["dtype"] == "float32"
+    assert records["compressed"]["format"] == "f32-to-bf16"
+    assert records["compressed"]["dtype"] == "float32"
+    assert b16_records["raw"]["format"] == "raw"
+    assert b16_records["raw"]["dtype"] == "float32"
+    assert b16_records["compressed"]["format"] == "raw"
+    assert b16_records["compressed"]["dtype"] == "bfloat16"
+    np.testing.assert_array_equal(cache["raw"].numpy(), raw)
+    np.testing.assert_allclose(cache["compressed"].numpy(), compressed, rtol=4e-3, atol=1e-3)
+    np.testing.assert_array_equal(b16_cache["raw"].numpy(), raw)
+    np.testing.assert_allclose(b16_cache["compressed"].numpy(), compressed, rtol=4e-3, atol=1e-3)
+
+
+def test_dump_tensor_cache_supports_generator_input():
+    params = (
+        item
+        for item in [
+            ("raw", np.arange(4, dtype="float32")),
+            ("compressed", np.linspace(-1, 1, 4, dtype="float32")),
+        ]
+    )
+
+    with tempfile.TemporaryDirectory(prefix="tvm_") as temp_dir:
+        tvmjs.dump_tensor_cache(
+            params,
+            temp_dir,
+            encode_format={"raw": "raw", "*": "f32-to-bf16"},
+        )
+        cache, _ = tvmjs.load_tensor_cache(temp_dir, tvm.cpu())
+
+    np.testing.assert_array_equal(cache["raw"].numpy(), np.arange(4, dtype="float32"))
+    np.testing.assert_allclose(
+        cache["compressed"].numpy(),
+        np.linspace(-1, 1, 4, dtype="float32"),
+        rtol=4e-3,
+        atol=1e-3,
+    )
+
+
+def test_dump_tensor_cache_updates_mixed_encoding_manifests():
+    original_raw = np.array([1.0, 2.0], dtype="float32")
+    updated_raw = np.array([3.0, 4.0], dtype="float32")
+    compressed = np.array([0.2345678, -0.8765432], dtype="float32")
+
+    with tempfile.TemporaryDirectory(prefix="tvm_") as temp_dir:
+        tvmjs.dump_tensor_cache(
+            {"raw": original_raw, "compressed": compressed},
+            temp_dir,
+            encode_format={"raw": "raw", "*": "f32-to-bf16"},
+        )
+        _, old_manifest = _records_by_name(os.path.join(temp_dir, "tensor-cache.json"))
+        old_md5 = old_manifest["records"][0]["md5sum"]
+
+        tvmjs.dump_tensor_cache(
+            iter([("raw", updated_raw)]),
+            temp_dir,
+            encode_format={"raw": "raw"},
+            update_if_exists=True,
+        )
+
+        cache, _ = tvmjs.load_tensor_cache(temp_dir, tvm.cpu())
+        b16_cache, _ = tvmjs.load_tensor_cache(
+            os.path.join(temp_dir, "tensor-cache-b16.json"), tvm.cpu()
+        )
+        _, manifest = _records_by_name(os.path.join(temp_dir, "tensor-cache.json"))
+        _, b16_manifest = _records_by_name(os.path.join(temp_dir, "tensor-cache-b16.json"))
+
+    np.testing.assert_array_equal(cache["raw"].numpy(), updated_raw)
+    np.testing.assert_array_equal(b16_cache["raw"].numpy(), updated_raw)
+    np.testing.assert_allclose(cache["compressed"].numpy(), compressed, rtol=4e-3, atol=1e-3)
+    assert manifest["records"][0]["md5sum"] != old_md5
+    assert manifest["records"][0]["md5sum"] == b16_manifest["records"][0]["md5sum"]
+
+
+def test_dump_tensor_cache_removes_stale_b16_manifest():
+    with tempfile.TemporaryDirectory(prefix="tvm_") as temp_dir:
+        tvmjs.dump_tensor_cache(
+            {"compressed": np.ones(2, dtype="float32")},
+            temp_dir,
+            encode_format="f32-to-bf16",
+        )
+        b16_manifest = os.path.join(temp_dir, "tensor-cache-b16.json")
+        assert os.path.exists(b16_manifest)
+
+        tvmjs.dump_tensor_cache(
+            {"raw": np.ones(2, dtype="float32")},
+            temp_dir,
+            encode_format="raw",
+        )
+        assert not os.path.exists(b16_manifest)
+
+
+def test_dump_tensor_cache_requires_a_format_for_every_parameter():
+    with tempfile.TemporaryDirectory(prefix="tvm_") as temp_dir:
+        with pytest.raises(ValueError, match="parameter arr"):
+            tvmjs.dump_tensor_cache(
+                {"arr": np.ones(2, dtype="float32")},
+                temp_dir,
+                encode_format={"other": "raw"},
+            )
+
+
 if __name__ == "__main__":
     tvm.testing.main()