Skip to content

Commit 75bf79f

Browse files
JakeStevensfacebook-github-bot
authored andcommitted
Namespace colliding lifted tensor constants in to_native (#23891)
Summary: Edge lowering mints `_lifted_tensor_constant{i}` names independently per method, so unrelated scalars (e.g. an fp32 attention scale and an int64 zero) collide under one data key when packaging multiple methods, and `to_native` rejected the package. Namespace only those auto-generated keys per method (`<method>/<key>`) in both the merged graph blob (new `rename_data_keys` in `graph_serialize.py`, walking the same references as `collect_data_keys`) and the constant dicts; identical tensors stay shared, mutated keys are never renamed, and genuine parameter/buffer fqn collisions still raise. Differential Revision: D125294428
1 parent bdcc7e6 commit 75bf79f

4 files changed

Lines changed: 151 additions & 18 deletions

File tree

‎backends/native/serialization/graph_serialize.py‎

Lines changed: 33 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
from enum import IntEnum
3232
from typing import (
3333
Any,
34+
Callable,
3435
cast,
3536
get_args,
3637
get_origin,
@@ -1487,42 +1488,62 @@ def _iter_graph_args(value: object) -> Iterator[GraphArg]:
14871488
yield from _iter_graph_args(item)
14881489

14891490

1490-
def collect_data_keys(program: Program) -> set[str]:
1491-
"""Return every out-of-line data key the program references.
1491+
def _map_data_keys(method: Method, fn: Callable[[str], str]) -> None:
1492+
"""Replace every out-of-line data key in ``method`` with ``fn(key)``, in place.
14921493
14931494
Covers constant references (NamedTensorRef.data_key) plus the quant keys
14941495
(AffineGroup scale and zero-point) carried on any TensorMeta.quant, whether
14951496
it is attached to a constant or to a graph value (intermediates and I/O),
14961497
recursing into HOP subgraphs wherever a schema field carries a ``GraphArg``.
14971498
PackedQuant carries no external keys.
14981499
"""
1499-
keys: set[str] = set()
15001500

1501-
def add_meta(meta: Optional[TensorMeta]) -> None:
1501+
def map_meta(meta: Optional[TensorMeta]) -> None:
15021502
if meta is None or meta.quant is None:
15031503
return
15041504
scheme = meta.quant.scheme
15051505
if isinstance(scheme, AffineGroup):
1506-
keys.add(scheme.scale_data_key)
1506+
scheme.scale_data_key = fn(scheme.scale_data_key)
15071507
if scheme.zero_point_data_key:
1508-
keys.add(scheme.zero_point_data_key)
1508+
scheme.zero_point_data_key = fn(scheme.zero_point_data_key)
15091509

15101510
def visit_graph(graph: Graph) -> None:
15111511
for value in graph.tensor_values or []:
1512-
add_meta(value.meta)
1512+
map_meta(value.meta)
15131513
for node in graph.nodes:
15141514
for graph_arg in _iter_graph_args(node):
15151515
visit_graph(graph_arg.graph)
15161516

1517-
for method in program.methods:
1518-
for constant in method.constants or []:
1519-
keys.add(constant.data_key)
1520-
add_meta(constant.meta)
1521-
visit_graph(method.graph)
1517+
for constant in method.constants or []:
1518+
constant.data_key = fn(constant.data_key)
1519+
map_meta(constant.meta)
1520+
visit_graph(method.graph)
1521+
1522+
1523+
def collect_data_keys(program: Program) -> set[str]:
1524+
"""Return every out-of-line data key the program references."""
1525+
keys: set[str] = set()
15221526

1527+
def add(key: str) -> str:
1528+
keys.add(key)
1529+
return key
1530+
1531+
for method in program.methods:
1532+
_map_data_keys(method, add)
15231533
return keys
15241534

15251535

1536+
def rename_data_keys(
1537+
program: Program, method_name: str, mapping: dict[str, str]
1538+
) -> None:
1539+
"""Rename one method's data keys in place; keys absent from ``mapping`` stay."""
1540+
for method in program.methods:
1541+
if method.name == method_name:
1542+
_map_data_keys(method, lambda key: mapping.get(key, key))
1543+
return
1544+
raise ValueError(f"rename_data_keys: no method named {method_name!r}.")
1545+
1546+
15261547
def _check_optional_tensor_list_refs(
15271548
value: OptionalTensorListArg, ctx: str, check_ref: Any
15281549
) -> None:

‎exir/native/README.md‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,8 @@ needed in the future like `constant_methods` and ETRecord integration.
4545
Raises `TypeError` for the wrong input or method-value types. Raises `ValueError`
4646
for an empty method mapping or method name, if a method does not fully delegate to
4747
the native backend, if a constant differs between methods, or if a data key the
48-
graph references has no backing tensor.
48+
graph references has no backing tensor. Conflicting generated
49+
`_lifted_tensor_constant{i}` keys are namespaced as `<method>/<name>`.
4950

5051
## Current limitations: one complete delegate per method
5152

‎exir/native/native.py‎

Lines changed: 58 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -37,10 +37,12 @@
3737

3838
from executorch.backends.native.preprocess import NativeDelegateInfo
3939
from executorch.backends.native.serialization.graph_serialize import (
40+
_compile_to_bytes,
4041
_same_tensor,
4142
collect_data_keys,
4243
deserialize_program,
4344
merge_programs,
45+
rename_data_keys,
4446
validate_program,
4547
)
4648
from executorch.backends.native.serialization.schema import Program
@@ -278,6 +280,46 @@ def _merge_constants(
278280
return merged
279281

280282

283+
# Edge lowering generates these names independently for each method.
284+
_LIFTED_CONSTANT_PREFIX = "_lifted_tensor_constant"
285+
286+
287+
def _plan_lifted_constant_renames(
288+
per_method: dict[str, dict[str, torch.Tensor]],
289+
program: Program,
290+
) -> dict[str, dict[str, str]]:
291+
"""Plan ``<method>/<key>`` names for conflicting immutable lifted constants.
292+
293+
Equal tensors stay shared; keys mutated by any method are never renamed.
294+
"""
295+
mutated_keys = _mutable_data_keys(program)
296+
taken = {key for constants in per_method.values() for key in constants}
297+
seen: dict[str, torch.Tensor] = {}
298+
renames: dict[str, dict[str, str]] = {}
299+
for method_name, constants in per_method.items():
300+
mapping: dict[str, str] = {}
301+
for key, tensor in constants.items():
302+
existing = seen.get(key)
303+
if existing is None:
304+
seen[key] = tensor
305+
continue
306+
if _same_tensor(existing, tensor):
307+
continue
308+
if not key.startswith(_LIFTED_CONSTANT_PREFIX) or key in mutated_keys:
309+
continue
310+
new_key = f"{method_name}/{key}"
311+
suffix = 1
312+
while new_key in taken:
313+
suffix += 1
314+
new_key = f"{method_name}/{key}#{suffix}"
315+
mapping[key] = new_key
316+
taken.add(new_key)
317+
seen[new_key] = tensor
318+
if mapping:
319+
renames[method_name] = mapping
320+
return renames
321+
322+
281323
def _check_shared_mutable_buffers(program: Program) -> None:
282324
"""Hold the merged Program to serialize_program's cross-method guarantee.
283325
@@ -380,7 +422,8 @@ def to_native(
380422
ValueError: If the method dictionary is empty, a method name is empty, a
381423
method does not fully delegate to the native backend, a constant
382424
differs between methods, or a data key the graph references has no
383-
backing tensor.
425+
backing tensor. Conflicting generated ``_lifted_tensor_constant{i}``
426+
keys are namespaced per method.
384427
"""
385428
method_programs = _normalize_programs(programs)
386429

@@ -435,12 +478,23 @@ def to_native(
435478
method_name,
436479
)
437480
method_blobs[method_name] = module.processed_bytes
438-
method_constants[method_name] = _delegate_info(module, method_name).constants
481+
method_constants[method_name] = dict(
482+
_delegate_info(module, method_name).constants
483+
)
439484

440485
ptg_blob = merge_programs(method_blobs)
441-
constants = _merge_constants(method_constants)
442-
443486
program = deserialize_program(ptg_blob)
487+
488+
renames = _plan_lifted_constant_renames(method_constants, program)
489+
if renames:
490+
for method_name, mapping in renames.items():
491+
rename_data_keys(program, method_name, mapping)
492+
constants_for_method = method_constants[method_name]
493+
for old_key, new_key in mapping.items():
494+
constants_for_method[new_key] = constants_for_method.pop(old_key)
495+
ptg_blob = _compile_to_bytes(program)
496+
497+
constants = _merge_constants(method_constants)
444498
_validate_merged_program(program, constants)
445499

446500
return NativeProgramManager(ptg=ptg_blob, constants=constants)

‎exir/native/test/test_native.py‎

Lines changed: 58 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,12 @@
2727
Program,
2828
)
2929
from executorch.exir import to_edge_transform_and_lower
30-
from executorch.exir._serialize._ptn import ALIASES_ENTRY, PTG_ENTRY, SAFETENSORS_ENTRY
30+
from executorch.exir._serialize._ptn import (
31+
ALIASES_ENTRY,
32+
PTG_ENTRY,
33+
read_ptn,
34+
SAFETENSORS_ENTRY,
35+
)
3136
from executorch.exir.lowered_backend_module import (
3237
executorch_call_delegate,
3338
get_lowered_submodules,
@@ -393,6 +398,58 @@ def test_dropped_public_input_is_rejected(self):
393398
self._check(self._graph_module(drop_input=True))
394399

395400

401+
class LiftedConstantCollisionTest(unittest.TestCase):
402+
def _scalar_shift(self, offset: float):
403+
class Shift(nn.Module):
404+
def forward(self, x):
405+
return x + offset
406+
407+
return torch.export.export(Shift(), (torch.randn(2, 4),))
408+
409+
def _read(self, native) -> tuple[Program, dict[str, torch.Tensor]]:
410+
with tempfile.TemporaryDirectory() as td:
411+
path = os.path.join(td, "m.ptn")
412+
native.save(path)
413+
ptg, constants = read_ptn(path)
414+
return deserialize_program(ptg), constants
415+
416+
def _lifted_keys(self, program: Program, method_name: str) -> list[str]:
417+
(method,) = [m for m in program.methods if m.name == method_name]
418+
return [c.data_key for c in method.constants or [] if "lifted" in c.data_key]
419+
420+
def test_distinct_lifted_constants_are_namespaced(self):
421+
programs = {"a": self._scalar_shift(1.0), "b": self._scalar_shift(2.0)}
422+
native = to_native(programs)
423+
self.assertEqual(native.methods, {"a", "b"})
424+
program, constants = self._read(native)
425+
(a_key,) = self._lifted_keys(program, "a")
426+
(b_key,) = self._lifted_keys(program, "b")
427+
self.assertEqual(b_key, f"b/{a_key}")
428+
torch.testing.assert_close(
429+
constants[a_key], torch.full_like(constants[a_key], 1.0), rtol=0, atol=0
430+
)
431+
torch.testing.assert_close(
432+
constants[b_key], torch.full_like(constants[b_key], 2.0), rtol=0, atol=0
433+
)
434+
435+
def test_identical_lifted_constants_stay_shared(self):
436+
programs = {"a": self._scalar_shift(1.0), "b": self._scalar_shift(1.0)}
437+
program, constants = self._read(to_native(programs))
438+
self.assertEqual(
439+
self._lifted_keys(program, "a"), self._lifted_keys(program, "b")
440+
)
441+
lifted = {k for k in constants if "lifted" in k}
442+
self.assertEqual(len(lifted), 1, lifted)
443+
444+
def test_parameter_fqn_collision_is_still_rejected(self):
445+
programs = {
446+
"a": torch.export.export(nn.Linear(4, 4), (torch.randn(1, 4),)),
447+
"b": torch.export.export(nn.Linear(4, 4), (torch.randn(1, 4),)),
448+
}
449+
with self.assertRaisesRegex(ValueError, "differs between methods"):
450+
to_native(programs)
451+
452+
396453
class MutableBufferDedupTest(unittest.TestCase):
397454
def test_identical_mutable_buffers_stay_separate(self):
398455
"""Equal bytes must not make two independently mutable buffers one.

0 commit comments

Comments
 (0)