mixle.utils.parallel.model_parallel module

Model-parallel estimation: distribute a model’s shardable axis across workers (component C3).

The inversion of the data-parallel backends: there the model is replicated and the data sharded. Here the model’s shardable axis is distributed. Two entry points share one recursive fold:

  • optimize(..., backend="model_parallel") – the ModelParallelEncodedData handle: the data is replicated in-process and the model axis is threaded. Simplest, single-machine.

  • optimize(ModelParallelEstimator(est), backend="spark"|"mpi"|"mp"|"local") – the estimator wrapper: composes with any data backend, so the data is sharded by that backend (Spark partitions, MPI ranks, mp pool) while the model axis is distributed inside each partition’s accumulator. This is the data x model composition – both axes at once.

The fold (model_parallel_fold()) is recursive: it walks the whole model tree and threads the axes that a per-node compute-cost model says save the most wall-time (_parallel_ids, consistent with the C2 planner – a narrow batch of heavy MVGaussians beats a wider batch of cheap leaves), recursing serially below any threaded node so no two pools ever nest. Each recursive case reproduces the corresponding accumulator’s seq_update exactly:

  • FACTOR (Composite/Record) – the per-factor accumulators are independent, so the per-factor seq_update calls are distributed (bit-identical).

  • COMPONENT (mixtures) – the responsibility logsumexp couples the components, so the cheap normalization runs centrally on the gathered score matrix while the expensive per-component scoring and accumulation are distributed – a bit-identical mirror of MixtureAccumulator.seq_update.

  • atomic / unknown – the replicated base case acc.seq_update(enc, weights, model).

So the whole fold is bit-identical to the single-node path (the data-axis reduce across partitions is the usual additive combine, exact up to float reassociation like every data-parallel backend). Correct for every family; never worse than backend="local". See ~/codex/notes/model-parallel-design.md.

class ModelParallelEncodedData(data, *, estimator=None, model=None, encoder=None, num_workers=None, **_)[source]

Bases: EncodedDataHandle

Replicate the data, distribute the model’s shardable axis across threads (single machine).

Parameters:

  • data (Any)

  • estimator (Any | None)

  • model (Any | None)

  • encoder (Any | None)

  • num_workers (int | None)

  • _ (Any)

pysp_seq_log_density_sum(estimate)[source]

Return (num_observations, summed_log_density) for estimate.

Parameters:

estimate (Any)

Return type:

tuple[float, float]

pysp_seq_estimate(estimator, prev_estimate)[source]

Run one distributed/local sufficient-statistic fold and M-step.

Parameters:

  • estimator (Any)

  • prev_estimate (Any)

Return type:

Any

pysp_seq_initialize(estimator, rng, p)[source]

Initialize a model through the handle’s resident encoded data.

Parameters:
Return type:

Any

pysp_stream_accumulate(estimator, model)[source]

Return folded sufficient statistics for streaming/incremental EM.

Parameters:
Return type:

tuple[float, Any]

class ModelParallelEstimator(inner, num_workers=None)[source]

Bases: ParameterEstimator

Wrap an estimator so EM distributes the model axis – composing with any backend= for the data.

optimize(ModelParallelEstimator(est), backend="spark"|"mpi"|"mp"|"local") shards the data through that backend while each partition’s E-step distributes the model axis. The M-step (estimate) and the accumulator’s value/combine contract are the wrapped estimator’s, unchanged.

Parameters:

  • inner (ParameterEstimator)

  • num_workers (int | None)

accumulator_factory()[source]
Return type:

ModelParallelAccumulatorFactory

estimate(nobs, suff_stat)[source]
Parameters:
Return type:

Any

class ModelParallelAccumulator(inner, num_workers=None)[source]

Bases: SequenceEncodableStatisticAccumulator

Wrap an accumulator so its E-step seq_update runs the recursive model-parallel fold.

All sufficient-statistic methods delegate to the wrapped (inner) accumulator unchanged, so the value/combine/from_value/key-merge contract – and thus every data backend’s reduce – is preserved; only seq_update is replaced with the distributed fold. Holding inner in vars() keeps the accumulator key-validator’s recursion transparent.

Parameters:

  • inner (SequenceEncodableStatisticAccumulator)

  • num_workers (int | None)

update(x, weight, estimate)[source]
Parameters:
Return type:

None

initialize(x, weight, rng)[source]
Parameters:
Return type:

None

seq_update(x, weights, estimate)[source]
Parameters:
Return type:

None

seq_initialize(x, weights, rng)[source]
Parameters:
Return type:

None

combine(suff_stat)[source]
Parameters:

suff_stat (Any)

Return type:

ModelParallelAccumulator

value()[source]
Return type:

Any

from_value(x)[source]
Parameters:

x (Any)

Return type:

ModelParallelAccumulator

key_merge(stats_dict)[source]

Pool this accumulator’s statistics into stats_dict under its merge key.

The structural default implements the common single-key pattern: store the accumulator under self.keys the first time the key is seen, else combine into the one already there. Accumulators with several named keys (e.g. an HMM’s init/trans/state keys) or a non-accumulator stats payload override this. A keys of None (the default) is a no-op.

Parameters:

stats_dict (dict[str, Any])

Return type:

None

key_replace(stats_dict)[source]

Replace this accumulator’s statistics from the pooled stats_dict entry (see key_merge).

Parameters:

stats_dict (dict[str, Any])

Return type:

None

acc_to_encoder()[source]
Return type:

Any

class ModelParallelAccumulatorFactory(inner_factory, num_workers=None)[source]

Bases: StatisticAccumulatorFactory

Parameters:

  • inner_factory (Any)

  • num_workers (int | None)

make()[source]
Return type:

ModelParallelAccumulator

model_parallel_fold(acc, model, enc, weights, num_workers=None)[source]

Run a model-parallel E-step of model over encoded enc into accumulator acc (in place).

Parameters:
Return type:

None

auto_parallel_estimator(estimator, model, resources=None, *, n_data=None, min_components_per_shard=1)[source]

Consult the C2 planner (decompose_model()) and return (estimator, decomposition).

When the planner picks model-parallelism for model on resources, the estimator is wrapped in ModelParallelEstimator sized to the planner’s cuts; otherwise the plain estimator is returned (replicate the model, shard the data – already optimal when N dominates). Either way, run the returned estimator through optimize(data, est, backend=<data backend>): the data axis is handled by backend and the model axis, if any, by the wrapper – composing into the data x model split. The decomposition’s rationale explains the choice. resources defaults to the local CPU slots (use Resources.from_spark(sc) / Resources.from_mpi() to size the split to a real cluster).

Parameters:

  • estimator (Any)

  • model (Any)

  • resources (Any)

  • n_data (int | None)

  • min_components_per_shard (int)

Return type:

tuple[Any, Any]