mixle.engines.heterogeneous module¶
Precision-aware planning for distributed EM across HETEROGENEOUS compute.
A 1000-worker pool is rarely uniform: some workers have GPU tensor cores (fast at fp8/bf16), others are
CPU-only (fp32, or vectorized double-double for accuracy). This module decides, per worker, (a) how many
rows of the E-step to give it – balanced by its throughput – and (b) which precision band to run, the
fastest one its hardware supports that still meets the accuracy budget. It also sizes the k-way tree
reduce depth so the fixed-size sufficient-statistic payloads fold in O(log W) rather than a single-root
fan-in.
This is the planning layer (pure-Python, testable); the actual Spark treeReduce / MPI comm.reduce
dispatch that consumes the plan lives in mixle.inference and needs a real cluster to exercise.
- class Worker(name, device, precisions, base_throughput=1.0)[source]
Bases:
objectA compute worker: its device, the precisions it can run (any order), and a base throughput.
- name: str
- device: str
- base_throughput: float = 1.0
- class WorkerAssignment(name: 'str', rows: 'int', precision: 'str', effective_throughput: 'float')[source]
Bases:
object- name: str
- rows: int
- precision: str
- effective_throughput: float
- class HeterogeneousPlan(assignments: 'tuple[WorkerAssignment, ...]', reduce_depth: 'int')[source]
Bases:
object- assignments: tuple[WorkerAssignment, ...]
- reduce_depth: int
- plan_heterogeneous(workers, n_rows, allowed_precisions=('fp8', 'bfloat16', 'float16', 'float32', 'float64', 'dd'), target_rel_error=None, op_count=1000)[source]
Assign rows + a precision band to each worker, balanced by precision-adjusted throughput.
Each worker runs the fastest precision its hardware supports that stays within
target_rel_error(None= no accuracy constraint); rows are split proportionally to the resulting throughput so all workers finish together.reduce_depthis the k-way tree depth for folding the sufficient-statistic payloads (~ceil(log2(W)/2)), avoiding the single-root fan-in.