Derive layer names from ONNX node and CaptureSplitInfo emit assignmen…

[neilmsft]

Describe your changes

What the two changes do

1. capture_split_info.py : make the pass work on ONNX

Instead of a torch module tree, derive the layer list from ONNX node names. Mobius names nodes after their originating module ( model/layers.0/self_attn/q_proj/MatMul ), so the structure is already there. The pass normalizes / → . , groups by namespace, and picks the one with the most numbered children → model.layers , indices 0–34.

This is safe rather than a guess: SplitModel.get_assignment ( split.py:255 ) already normalizes node names the exact same way to match assignments back to nodes. I'm just reading the naming convention the consumer side already assumes.

2. split.py : let the two passes hand off

CaptureSplitInfo writes assignments to model_attributes ; SplitModel didn't look there. Without this the passes can't chain. The file carried an explicit TODO(jambayk): Should we allow split assignments in the model attributes too? So this closes a sanctioned gap.

[Copilot]

Pull request overview

This PR extends Olive’s split-assignment workflow to better support ONNX-first pipelines by deriving layer namespaces from ONNX node names in CaptureSplitInfo, persisting assignments into model_attributes, and teaching SplitModel to consume those attributes (plus adding a retry path for ORT symbolic shape inference when Constant attributes are encoded in alternate forms).

Changes:

• Add ONNX support to CaptureSplitInfo via split_using_num_splits_onnx, inferring block_to_split from ONNX node names.
• Store split_assignments into model_attributes for ONNX models, and update SplitModel to read them from there when metadata/config doesn’t provide them.
• Add Constant-node attribute normalization and a retry path for ORT symbolic shape inference in SplitModel.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.

File	Description
olive/passes/pytorch/capture_split_info.py	Adds ONNXModelHandler support and ONNX node-name based split assignment derivation; persists assignments into model_attributes.
olive/passes/onnx/split.py	Reads split assignments from model_attributes and retries shape inference after normalizing Constant attributes.

Suppressed comments (1)

olive/passes/pytorch/capture_split_info.py:102

• num_splits is treated as a truthy value here (if not config.num_splits / elif config.num_splits), which makes num_splits=0 pass validate_config (it’s not None) but fail at runtime with a misleading error path. It also makes the PyTorch/HF branch silently ignore num_splits=0 and fall through to cost_model/error. Prefer explicit is None checks and validate that num_splits is a positive integer before calling np.array_split.

        split_assignments = None
        if isinstance(model, ONNXModelHandler):
            if not config.num_splits:
                raise ValueError("num_splits is required to split an ONNX model. cost_model is not supported.")
            split_assignments = self.split_using_num_splits_onnx(model, config)
        elif config.num_splits:
            split_assignments = self.split_using_num_splits(model, config)
        elif config.cost_model:
            split_assignments = self.split_using_cost_model(model, config)
        else:
            raise ValueError("One of num_splits or cost_model is required.")

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.