diff --git a/QEfficient/compile/mdp_generator.py b/QEfficient/compile/mdp_generator.py index a37a62a9be..1594e14702 100644 --- a/QEfficient/compile/mdp_generator.py +++ b/QEfficient/compile/mdp_generator.py @@ -561,6 +561,76 @@ def generate_disagg_mdp_partition_config( } +def _repair_scoped_alias_partitioning(onnx_path: str, mdp_json: Dict[str, Any], compiler_dump: Dict[str, Any]) -> int: + """Align scoped aliases with the compiler dump partition.""" + model = onnx.load(onnx_path, load_external_data=False) + names = {n.name for n in model.graph.node if n.name} + inlined_node_map, _ = _get_inlined_node_map(model) + for expanded_names in inlined_node_map.values(): + names.update(expanded_names) + compiler_partition, compiler_position, compiler_nodes = {}, {}, set() + for partition_idx, partition in enumerate(compiler_dump.get("partitions", [])): + for node_pos, node_name in enumerate(partition.get("nodeList", [])): + if node_name: + compiler_nodes.add(node_name) + compiler_partition.setdefault(node_name, partition_idx) + compiler_position.setdefault(node_name, (partition_idx, node_pos)) + current_partition = {} + for partition_idx, partition in enumerate(mdp_json.get("partitions", [])): + for node_name in partition.get("nodeList", []): + if node_name: + current_partition[node_name] = partition_idx + node_by_name = {node.name: node for node in model.graph.node if node.name} + consumers_by_value = {} + for node in model.graph.node: + if not node.name: + continue + for input_name in node.input: + if input_name: + consumers_by_value.setdefault(input_name, []).append(node.name) + repairs = [] + for node_name, source_idx in current_partition.items(): + if not node_name.startswith("/") or node_name.startswith("/language_model/"): + continue + peer_name = "/language_model" + node_name + target_idx = compiler_partition.get(node_name) + if peer_name not in names or peer_name in compiler_nodes or target_idx is None or target_idx == source_idx: + continue + node = node_by_name.get(node_name) + if node is not None: + consumer_partitions = { + current_partition[consumer_name] + for output_name in node.output + for consumer_name in consumers_by_value.get(output_name, []) + if consumer_name in current_partition + } + if any(consumer_idx < target_idx for consumer_idx in consumer_partitions): + logger.warning( + "Skipping MDP alias repair for %r: compiler target Partition%d would follow " + "an existing consumer partition", + node_name, + target_idx, + ) + continue + repairs.append((node_name, source_idx, target_idx, compiler_position[node_name])) + for node_name, source_idx, target_idx, dump_pos in sorted(repairs, key=lambda item: item[3]): + mdp_json["partitions"][source_idx]["nodeList"].remove(node_name) + target_nodes = mdp_json["partitions"][target_idx]["nodeList"] + insert_at = len(target_nodes) + for idx, existing_name in enumerate(target_nodes): + if compiler_position.get(existing_name, (10**9, 10**9)) > dump_pos: + insert_at = idx + break + target_nodes.insert(insert_at, node_name) + logger.warning( + "MDP repaired compiler-scoped alias %r: moved from Partition%d to Partition%d", + node_name, + source_idx, + target_idx, + ) + return len(repairs) + + def generate_disagg_mdp_intersection_config( onnx_path: str, compiler_dump_path: str, @@ -576,7 +646,9 @@ def generate_disagg_mdp_intersection_config( QEff MDP, but the QEff MDP also contains nodes the compiler optimises away (constant-folded, etc.). The intersection keeps only nodes that survive into the compiler IR, while preserving the **QEff MDP node order and - partition assignment** as the authoritative source of truth. + partition assignment** as the authoritative source of truth, except for + scoped aliases whose partition is aligned with the compiler dump when the + compiler materializes only the unscoped alias. Algorithm --------- @@ -587,12 +659,16 @@ def generate_disagg_mdp_intersection_config( names the compiler will actually process. 3. For each partition in the QEff MDP, retain only nodes whose names appear in the compiler dump set, in the original QEff order. + 4. For scoped aliases with a compiler-only unscoped name, move the alias to + the compiler dump partition while preserving the compiler-relative order + needed by the compiler's partition validator. Why this is better than re-assigning compiler-dump nodes by heuristic: - Node order follows the ONNX topsort (compiler's SplitPlanMerge requires this); the compiler dump order is the compiler's internal IR order which may differ after constant-folding and other transforms. - - Partition boundaries are the same as the onnx strategy — no divergence. + - Partition boundaries are the same as the onnx strategy, except for the + narrow scoped-alias repair described above. - Every retained node is an exact-match hit in ``doManualPartitioning()``. - Output JSON is compact (~1-2 MB) because constant-folded nodes are pruned. @@ -664,6 +740,17 @@ def generate_disagg_mdp_intersection_config( } ) + mdp_json = { + "connections": qeff_mdp["connections"], + "partitions": partition_objs, + } + + _repair_scoped_alias_partitioning( + onnx_path=onnx_path, + mdp_json=mdp_json, + compiler_dump=dump, + ) + logger.info( f"Intersection: {total_kept} nodes kept, {total_dropped} dropped " f"({total_dropped * 100 // max(qeff_total, 1)}% pruned by compiler)" @@ -676,10 +763,7 @@ def generate_disagg_mdp_intersection_config( f"(compiler may have renamed them). Examples: {list(in_dump_not_qeff)[:5]}" ) - return { - "connections": qeff_mdp["connections"], - "partitions": partition_objs, - } + return mdp_json def generate_disagg_mdp_config( diff --git a/tests/unit_test/base/test_modeling_qeff_base.py b/tests/unit_test/base/test_modeling_qeff_base.py index 73f0d54168..f4c645c2aa 100644 --- a/tests/unit_test/base/test_modeling_qeff_base.py +++ b/tests/unit_test/base/test_modeling_qeff_base.py @@ -29,6 +29,8 @@ from QEfficient.compile.mdp_generator import ( _get_layer_num_from_inputs, _layer_partition_bounds, + generate_disagg_mdp_intersection_config, + precheck_and_autofix_mdp_partition_order, ) from QEfficient.transformers.models.modeling_auto import QEFFAutoModelForCausalLM @@ -789,6 +791,161 @@ def test_mdp_disagg_json_is_created(self, compile_workspace): disagg_files = list(compile_dir.glob("mdp_disagg_4d_2p.json")) assert len(disagg_files) == 1, "Expected exactly one mdp_disagg_4d_2p.json in compile_dir" + def test_intersection_repairs_scoped_alias_to_compiler_partition(self, tmp_path: Path): + """Intersection moves an unscoped alias to the compiler-assigned partition.""" + onnx_path = tmp_path / "scoped_alias.onnx" + model = helper.make_model( + helper.make_graph( + [ + helper.make_node( + "Identity", + inputs=["input"], + outputs=["hidden0"], + name="/language_model/layers.0/hidden", + ), + helper.make_node( + "Identity", + inputs=["input"], + outputs=["scoped_hidden"], + name="/language_model/Gather_5", + ), + helper.make_node( + "Identity", + inputs=["input"], + outputs=["hidden1"], + name="/language_model/layers.1/hidden", + ), + helper.make_node( + "Identity", + inputs=["input"], + outputs=["output"], + name="/Gather_5", + ), + ], + name="scoped_alias_graph", + inputs=[helper.make_tensor_value_info("input", TensorProto.FLOAT, [1])], + outputs=[helper.make_tensor_value_info("output", TensorProto.FLOAT, [1])], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + onnx.save(model, onnx_path) + + compiler_dump_path = tmp_path / "compiler_dump.json" + compiler_dump_path.write_text( + json.dumps( + { + "partitions": [ + { + "nodeList": [ + "/language_model/layers.0/hidden", + "/Gather_5", + ] + }, + {"nodeList": ["/language_model/layers.1/hidden"]}, + ] + } + ) + ) + + mdp_json = generate_disagg_mdp_intersection_config( + onnx_path=str(onnx_path), + compiler_dump_path=str(compiler_dump_path), + num_devices=2, + num_partitions=2, + num_layers=2, + ) + + precheck_removed = precheck_and_autofix_mdp_partition_order( + onnx_path=str(onnx_path), + mdp_json=mdp_json, + ) + + assert precheck_removed == 0 + assert "/Gather_5" in mdp_json["partitions"][0]["nodeList"] + assert mdp_json["partitions"][0]["nodeList"] == [ + "/language_model/layers.0/hidden", + "/Gather_5", + ] + assert mdp_json["partitions"][1]["nodeList"] == ["/language_model/layers.1/hidden"] + + def test_intersection_skips_alias_move_that_follows_consumer(self, tmp_path: Path): + """Intersection keeps a producer in place when relocation would reverse its dependency.""" + onnx_path = tmp_path / "scoped_alias_dependency.onnx" + model = helper.make_model( + helper.make_graph( + [ + helper.make_node( + "Identity", + inputs=["input"], + outputs=["hidden0"], + name="/language_model/layers.0/hidden", + ), + helper.make_node( + "Identity", + inputs=["input"], + outputs=["scoped_alias"], + name="/language_model/Gather_9", + ), + helper.make_node( + "Identity", + inputs=["input"], + outputs=["producer_output"], + name="/Gather_9", + ), + helper.make_node( + "Identity", + inputs=["producer_output"], + outputs=["consumer_output"], + name="/language_model/Gather_11", + ), + helper.make_node( + "Identity", + inputs=["input"], + outputs=["hidden1"], + name="/language_model/layers.1/hidden", + ), + ], + name="scoped_alias_dependency_graph", + inputs=[helper.make_tensor_value_info("input", TensorProto.FLOAT, [1])], + outputs=[helper.make_tensor_value_info("consumer_output", TensorProto.FLOAT, [1])], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + onnx.save(model, onnx_path) + + compiler_dump_path = tmp_path / "compiler_dump.json" + compiler_dump_path.write_text( + json.dumps( + { + "partitions": [ + { + "nodeList": [ + "/language_model/layers.0/hidden", + "/language_model/Gather_11", + ] + }, + {"nodeList": ["/Gather_9", "/language_model/layers.1/hidden"]}, + ] + } + ) + ) + + mdp_json = generate_disagg_mdp_intersection_config( + onnx_path=str(onnx_path), + compiler_dump_path=str(compiler_dump_path), + num_devices=2, + num_partitions=2, + num_layers=2, + ) + + precheck_removed = precheck_and_autofix_mdp_partition_order( + onnx_path=str(onnx_path), + mdp_json=mdp_json, + ) + + assert precheck_removed == 0 + assert "/Gather_9" in mdp_json["partitions"][0]["nodeList"] + def test_mdp_disagg_json_has_two_partitions(self, compile_workspace): """mdp_disagg_4d_2p.json contains exactly two partitions.""" compile_dir, _, _ = self._run_mdp_compile(compile_workspace)