mixle.utils.parallel package

Distributed-estimation planning and model-parallel entry points.

Importing this namespace registers the model-parallel encoded-data backend and exposes resource planning, sharding, and estimator wrappers used by multi-process or multi-device estimation workflows.

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

Bases: ParameterEstimator

Compatibility wrapper for fully resident component-parallel EM.

optimize(ModelParallelEstimator(est), backend="spark"|"mpi"|"mp"|"local") shards the data through that backend while each partition’s E-step threads independent model work. The full model remains resident in every process. 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 a factory that wraps the inner estimator’s accumulator factory.

Return type:

ModelParallelAccumulatorFactory

estimate(nobs, suff_stat)[source]

Delegate the M-step to the wrapped estimator.

Parameters:
Return type:

Any

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

Bases: EncodedDataHandle

Compatibility name for a fully resident, component-threaded data handle.

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 the encoded-data size and total log likelihood under estimate.

Parameters:

estimate (Any)

Return type:

tuple[float, float]

pysp_seq_estimate(estimator, prev_estimate)[source]

Run one model-parallel E/M update from prev_estimate.

Parameters:

  • estimator (Any)

  • prev_estimate (Any)

Return type:

Any

pysp_seq_initialize(estimator, rng, p)[source]

Initialize a model by randomly selecting observations with probability p.

Parameters:
Return type:

Any

pysp_stream_accumulate(estimator, model)[source]

Accumulate model-parallel sufficient statistics for streaming backends.

Parameters:
Return type:

tuple[float, Any]

ComponentParallelEstimator

alias of ModelParallelEstimator

ComponentParallelEncodedData

alias of ModelParallelEncodedData

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

Run a component-parallel E-step over a fully resident model.

Parameters:
Return type:

None

component_parallel_fold(acc, model, enc, weights, num_workers=None, *, parallel_ids=None)

Run a component-parallel E-step over a fully resident model.

Parameters:
Return type:

None

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

Return a plain estimator when no distributed model placement is required.

The structural planner is advisory. If it selects cuts that require model storage to be partitioned across devices, this helper raises UnrealizedModelPlacementError: the local component-thread runtime cannot truthfully realize those cuts. Callers may explicitly choose ComponentParallelEstimator when the full model fits in each process.

Parameters:

  • estimator (Any)

  • model (Any)

  • resources (Any)

  • n_data (int | None)

  • min_components_per_shard (int)

Return type:

tuple[Any, Any]

exception UnrealizedModelPlacementError(plan)[source]

Bases: RuntimeError

A planner selected model shards for which this runtime has no placement executor.

Parameters:

plan (Any)

Return type:

None

class Resources(devices)[source]

Bases: object

A collection of placement targets.

Parameters:

devices (tuple[DeviceSpec, ...])

classmethod single_cpu(memory_bytes=None, throughput=1.0, precision=None)[source]

Return a one-device CPU resource description.

Parameters:

  • memory_bytes (int | None)

  • throughput (float)

  • precision (Any | None)

Return type:

Resources

classmethod local(num_cpus=None, memory_bytes=None, precision=None)[source]

Return local CPU resources split into logical worker slots.

Parameters:

  • num_cpus (int | None)

  • memory_bytes (int | None)

  • precision (Any | None)

Return type:

Resources

classmethod discover(include_torch=True, cpu_workers=None, precision=None)[source]

Best-effort local resource discovery with no required extras.

Parameters:

  • include_torch (bool)

  • cpu_workers (int | None)

  • precision (Any | None)

Return type:

Resources

classmethod from_specs(specs)[source]

Build resources from an iterable of device specifications.

Parameters:

specs (Iterable[DeviceSpec])

Return type:

Resources

classmethod from_mpi(comm=None, memory_bytes=None, throughput=1.0, precision=None)[source]

Return CPU resource slots for an MPI world without importing mpi4py.

comm may be an mpi4py-style communicator exposing Get_size. When it is omitted, common MPI launcher environment variables are used as a best-effort size hint, falling back to one slot.

Parameters:

  • comm (Any | None)

  • memory_bytes (int | None)

  • throughput (float)

  • precision (Any | None)

Return type:

Resources

classmethod from_dask(client, precision=None)[source]

Return resource slots from a dask.distributed-like client.

The method relies only on client.scheduler_info() and therefore does not introduce a dask dependency.

Parameters:

  • client (Any)

  • precision (Any | None)

Return type:

Resources

classmethod from_spark(spark_context, memory_bytes=None, precision=None)[source]

Return CPU resource slots from a SparkContext-like object.

Parameters:

  • spark_context (Any)

  • memory_bytes (int | None)

  • precision (Any | None)

Return type:

Resources

classmethod from_torchrun(memory_bytes=None, precision=None)[source]

Return globally identified slots from an attested torchrun topology.

Parameters:

  • memory_bytes (int | None)

  • precision (Any | None)

Return type:

Resources

fastest()[source]

Return the device with the largest advisory throughput weight.

Return type:

DeviceSpec

to_dict()[source]

Return a JSON-serializable resource description.

Return type:

dict[str, Any]

classmethod from_dict(payload)[source]

Build resources from to_dict output.

Parameters:

payload (dict[str, Any])

Return type:

Resources

to_json(**kwargs)[source]

Serialize resources, including calibrated throughput, to JSON.

Parameters:

kwargs (Any)

Return type:

str

classmethod from_json(text)[source]

Deserialize resources from JSON produced by to_json.

Parameters:

text (str)

Return type:

Resources

save(path, **kwargs)[source]

Persist resources to a JSON file for reuse by later plans.

Parameters:
Return type:

None

classmethod load(path)[source]

Load resources from a JSON file created by save.

Parameters:

path (Any)

Return type:

Resources

exception EncodedDataControlConflictError[source]

Bases: ValueError

Raised when controls request mutation of an already encoded handle.

encoded_data(data, estimator=None, model=None, encoder=None, placement=None, resources=None, engine=None, precision=None, num_chunks=None, sub_chunks=1, backend=None, num_workers=None, client=None, comm=None, root=0, root_only=False, parallel_chunks=False, chunk_workers=None, batch_size=None, shuffle=None, seed=None)[source]

Create an encoded-data handle or explicitly reuse an unchanged handle.

Backend dispatch goes through a registry (see register_encoded_data_backend()) rather than a hard-coded branch, so a new distributed framework (Lightning, Ray, JAX, …) plugs in by registering a factory – the same “register, don’t branch” pattern the compute engines use – without editing this function.

Parameters:

  • data (Any)

  • estimator (Any | None)

  • model (Any | None)

  • encoder (DataSequenceEncoder | None)

  • placement (Placement | None)

  • resources (Resources | None)

  • engine (Any | None)

  • precision (Any | None)

  • num_chunks (int | None)

  • sub_chunks (int)

  • backend (str | None)

  • num_workers (int | None)

  • client (Any | None)

  • comm (Any | None)

  • root (int)

  • root_only (bool)

  • parallel_chunks (bool)

  • chunk_workers (int | None)

  • batch_size (int | None)

  • shuffle (bool | None)

  • seed (int | None)

Return type:

EncodedDataHandle

is_encoded_data_handle(obj)[source]

Return true when obj exposes the sequence-orchestrator contract.

Parameters:

obj (Any)

Return type:

bool

model_sharding_plan(model, resources, estimator=None, axis='components', min_components_per_shard=1)[source]

Return an advisory component-axis sharding plan.

This is the planning half of model parallelism. It does not create DTensor placements or move tensors; it answers which component range each device should own when a stacked/generated component kernel is available.

Parameters:

  • model (Any)

  • resources (Resources)

  • estimator (Any | None)

  • axis (str)

  • min_components_per_shard (int)

Return type:

tuple[ModelShard, …]

plan(data=None, model=None, estimator=None, encoder=None, resources=None, engine=None, precision=None, num_chunks=None, sub_chunks=1, sample_size=256, safety_factor=1.25)[source]

Return an advisory placement for local or distributed orchestration.

Parameters:

  • data (Sequence[Any] | None)

  • model (Any | None)

  • estimator (Any | None)

  • encoder (DataSequenceEncoder | None)

  • resources (Resources | None)

  • engine (Any | None)

  • precision (Any | None)

  • num_chunks (int | None)

  • sub_chunks (int)

  • sample_size (int)

  • safety_factor (float)

Return type:

Placement

register_encoded_data_backend(name, factory, aliases=(), *, override=False)[source]

Register an encoded-data backend factory under name (and any aliases).

factory is called as factory(data, **params) with the keyword arguments of encoded_data(); it should accept **_ for the parameters it ignores and return an EncodedDataHandle. This is the extension point for new parallel/distributed frameworks – registering is all that is needed, no core edits.

Raises ValueError if name or any alias is already registered to a DIFFERENT factory, rather than silently shadowing it – a third-party package registering under a name that collides with a built-in (e.g. 'local') or another package’s backend would otherwise hijack every subsequent backend=name call with no error, no warning, at whatever time that package happened to import. Pass override=True to replace an existing registration deliberately (e.g. hot-swapping a backend in a test).

Parameters:
Return type:

None

class BackendCapabilities(name, axes, precisions=frozenset({'bf16', 'fp32'}), distributed_optimizer=False, reshardable_checkpoint=False, elastic_restart=False, requirements=(), incompatible_axis_sets=())[source]

Bases: object

What a backend actually executes, rather than merely accepting as flags.

Parameters:
class CollectiveKind(*values)[source]

Bases: StrEnum

Collective or point-to-point operation required by an update.

class DistributedUpdate(node_id, payload, collective, mesh_axes, state_layout, exact, notes=(), contract_exact=False, determinism_observed=None, maximum_absolute_error=None, maximum_relative_error=None, numerics_evidence_id=None, numerics_sample_count=0)[source]

Bases: object

Executable communication declaration with explicit numerical evidence.

Parameters:

  • node_id (str)

  • payload (PayloadKind)

  • collective (CollectiveKind)

  • mesh_axes (tuple[ParallelAxis, ...])

  • state_layout (StateLayout)

  • exact (bool)

  • notes (tuple[str, ...])

  • contract_exact (bool)

  • determinism_observed (bool | None)

  • maximum_absolute_error (float | None)

  • maximum_relative_error (float | None)

  • numerics_evidence_id (str | None)

  • numerics_sample_count (int)

class ParallelAxis(*values)[source]

Bases: StrEnum

Named dimensions understood by Mixle and external training engines.

class ParallelPlan(dp_replicate=1, dp_shard=1, tp=1, pp=1, cp=1, ep=1, etp=1, microbatches=1, gradient_accumulation_steps=1)[source]

Bases: object

An explicit N-D process mesh and batch schedule.

Every process belongs to exactly one coordinate in the product below. Keeping replicate and shard data-parallel dimensions separate permits HSDP without conflating it with tensor or context parallelism.

Parameters:
property world_size: int

Number of physical ranks.

EP and ETP are overlapping process groups cut from the data/model domains, not extra orthogonal DeviceMesh dimensions.

property mesh: tuple[tuple[str, ...], tuple[int, ...]]

Return stable DeviceMesh names and shape, including a 1-rank DP dimension.

class ParameterLayout(logical_id, global_shape, placements=(), shared_group=None, optimizer_state=StateLayout.REPLICATED)[source]

Bases: object

Stable logical identity and global layout for one parameter.

Parameters:
class PayloadKind(*values)[source]

Bases: StrEnum

Value communicated when an update is distributed.

class StateLayout(*values)[source]

Bases: StrEnum

Physical ownership of mutable training state.

class StepReceipt(step, loss, local_examples, local_tokens, microbatches, accumulation_steps, data_parallel_size, optimizer, precision, collective_bytes=0, skipped=False, extra=<factory>, global_examples_observed=None, global_tokens_observed=None)[source]

Bases: object

Unambiguous accounting for one committed optimizer update.

Parameters:

  • step (int)

  • loss (float)

  • local_examples (int)

  • local_tokens (int)

  • microbatches (int)

  • accumulation_steps (int)

  • data_parallel_size (int)

  • optimizer (str)

  • precision (str)

  • collective_bytes (int)

  • skipped (bool)

  • extra (dict[str, Any])

  • global_examples_observed (int | None)

  • global_tokens_observed (int | None)

get_training_backend(name)[source]

Resolve a registered backend, calling a zero-argument lazy factory once.

Parameters:

name (str)

Return type:

Any

register_training_backend(name, backend, *, override=False)[source]

Register a backend instance or lazy factory without importing optional dependencies.

Parameters:
Return type:

None

class LightningFabricLauncher[source]

Bases: object

Create Fabric for pure DDP or FSDP jobs.

Lightning is capability-gated to data parallelism here. Model parallel dimensions remain the responsibility of the native DeviceMesh or Megatron backend rather than being accepted and discarded by Fabric.

class RayTrainLauncher[source]

Bases: object

Launch one worker per physical plan rank with ray.train.torch.

class WorkerTopologyAttestation(rank, world_size, coordinate)[source]

Bases: object

Runtime proof binding one launched rank to one planned mesh coordinate.

Parameters:
attest_worker_coordinate(plan, *, actual_world_size, rank)[source]

Validate a runtime rank and deterministically bind it to the plan mesh.

Parameters:

  • plan (ParallelPlan)

  • actual_world_size (int)

  • rank (int)

Return type:

WorkerTopologyAttestation

Submodules