mixle.inference.heterogeneous_executor module

Distributed heterogeneous EM execution: sharded E-steps + k-way tree reduce of sufficient statistics.

The actual substance of ‘distributed through heterogeneous compute’ is backend-agnostic: split the data into shards, run each shard’s E-step (optionally at its own precision), and fold the fixed-size sufficient statistics with a k-way tree of accumulator.combine() – O(log W) depth, no single-root fan-in. combine is associative, so the tree result is bit-identical to a serial fold for integer/count statistics and within float reassociation otherwise.

This module is the executed, verifiable core: a local executor that shards and tree-reduces in-process, exactly matching a serial fit. The Spark (RDD.treeReduce), MPI (comm.reduce), and torchrun transports are thin adapters that replace the local shard loop with cluster transport over the same combine-tree – they need a cluster to exercise, but the algorithm they run is the one verified here.

tree_reduce_values(values, factory, branch=2)[source]

Fold accumulator value() payloads with a branch-ary tree of combine() – O(log n) depth.

Each internal node makes a fresh accumulator (so a shared-reference value() is never mutated – the HMM-stat aliasing hazard), seeds it from the first child, and combines the rest. Bit-identical to a serial left fold for associative integer statistics.

Raises:

  • TypeError – if branch is not an exact integer.

  • ValueError – if values is empty, or branch < 2 when a reduction is actually needed.

Parameters:
Return type:

Any

heterogeneous_em_step(estimator, model, data, n_shards=1, shard_sizes=None, shard_precisions=None, branch=2, pool=None)[source]

One distributed EM step: shard data, E-step each shard (at its precision), tree-reduce, estimate.

With one shard and no reduced precision this is byte-identical to a plain serial E-step; with many shards the tree-reduced result matches it up to float reassociation of combine(). pool is an optional concurrent.futures-style executor (e.g. ProcessPoolExecutor) whose map runs the shard E-steps on real worker processes – the sufficient-statistic payloads cross the process boundary by pickling, and combine operates on those freshly-unpickled copies (never a shared reference).

Keyed (tied) estimators are handled exactly like every serial driver: keys are validated up front and the key_merge/key_replace pooling pass runs ONCE on the fully tree-reduced statistics (the MPI/multiprocessing drivers’ driver-side contract), never per shard.

Parameters:
Return type:

Any

heterogeneous_fit(model, data, max_its=20, n_shards=4, shard_sizes=None, shard_precisions=None, branch=2, pool=None)[source]

Run max_its EM iterations with the distributed heterogeneous executor; returns the fitted model.

Raises:
Parameters:
Return type:

Any

shards_from_plan(plan)[source]

Translate a HeterogeneousPlan into (shard_sizes, precisions).

Only float32 is wired to the fused reduced-precision kernel here; other bands (fp8/bf16/dd) run on the exact float64 accumulator until their compute kernels exist – so the executor stays correct.

Parameters:

plan (Any)

Return type:

tuple[list[int], list[Any]]