NodeDeviceConfiguration¶

class onnx_ir.NodeDeviceConfiguration(configuration=None, sharding_specs=(), pipeline_stage=None)¶

The multi-device annotation of a single node under one configuration.

Attached to onnx_ir.Node.device_configurations, it says how the node participates in a ModelConfiguration: which of its tensors are sharded (tensor parallelism) and/or which pipeline stage it belongs to (pipeline parallelism). A node may carry several of these, one per configuration it takes part in.

Added in version 1.0.0.

Parameters:
configuration: ModelConfiguration | None = None¶
pipeline_stage: int | None = None¶
sharding_specs: tuple[ShardingSpec, ...] = ()¶