@@ -509,7 +509,7 @@ def tag_constant_data(edge_program: ExportedProgram) -> None:
509
509
underlying data will be owned by multiple delegates.
510
510
"""
511
511
for node in edge_program .graph .nodes :
512
- # go through const/param/buffer nodes
512
+ # go through const/param/buffer nodes, if all users of const/param/buffer nodes are partitioned then partition
513
513
is_attr = (
514
514
node .op == "placeholder"
515
515
and (
@@ -518,17 +518,20 @@ def tag_constant_data(edge_program: ExportedProgram) -> None:
518
518
or is_lifted_tensor_constant (edge_program , node )
519
519
)
520
520
) or (node .op == "get_attr" )
521
- # if all users of const/param/buffer nodes are partitioned then partition
522
521
if is_attr :
523
522
user_tags = set ()
524
523
for user in node .users :
525
- user_tags .add (user .meta .get ("delegation_tag" , None ))
526
- assert len (user_tags ) <= 1 , (
527
- "Const/Param/Buffer users have multiple tags because one constant data can't "
528
- "be owned by multiple backends. Consider duplicating the constant data so that "
529
- "each user is unique"
530
- )
531
- if len (user_tags ) == 1 :
524
+ user_tag = user .meta .get ("delegation_tag" , None )
525
+ if user_tag is not None :
526
+ user_tags .add (user_tag )
527
+ if len (user_tags ) > 1 :
528
+ logging .info (
529
+ f"The data node is used across multiple partitions, including { user_tags } . "
530
+ "If the data is too large and it's not prefered to copied, please tag the "
531
+ "constant node like node.['no_copy'] = True and they won't be copied."
532
+ )
533
+ # tag the data node with the same tag as the last user
534
+ if len (user_tags ) > 0 :
532
535
node .meta ["delegation_tag" ] = user_tags .pop ()
533
536
534
537
0 commit comments