mixle.utils.parallel.torchrun module

torchrun SPMD backend for distributed estimation.

The handle mirrors MPIEncodedData using torch.distributed collectives: each rank encodes its local data shard once, accumulates sufficient statistics against that resident encoding, rank 0 folds/key-ties/runs the M-step, and the result is broadcast back to every rank.

class TorchRunEncodedData(data, estimator=None, model=None, encoder=None, sub_chunks=1, group=None, root=0, root_only=False, init_process_group=True, backend=None)[source]

Bases: EncodedDataHandle

Encoded-data handle sharded across torchrun ranks.

Parameters:

  • data (Optional[Sequence]) – Raw observations. With root_only=False every rank passes the same full dataset and keeps data[rank::world]. With root_only=True only the root rank needs real data.

  • estimator – Used to build the encoder when encoder is not given.

  • model – Optional model used to build the encoder.

  • encoder – Explicit data encoder.

  • sub_chunks (int) – Encoded sub-chunks per rank.

  • group (Any | None) – Optional torch.distributed process group.

  • root (int) – Rank that folds statistics and runs the M-step.

  • root_only (bool) – Root-only data input mode.

  • init_process_group (bool) – Initialize torch.distributed from the torchrun environment when needed.

  • backend (Optional[str]) – Process-group backend. Defaults to nccl on CUDA and gloo otherwise.

pysp_seq_estimate(estimator, prev_estimate)[source]

One distributed EM step; every rank returns the identical model.

pysp_seq_initialize(estimator, rng, p)[source]

Distributed randomized initialization; identical model on all ranks.

Parameters:
pysp_seq_log_density_sum(estimate)[source]

Allreduced count and summed log density.

Return type:

tuple[float, float]

pysp_stream_accumulate(estimator, model)[source]

Globally folded batch sufficient statistics for streaming EM.

Return type:

tuple[float, Any]

close()[source]

Release worker resources owned by this handle, if any.

Return type:

None

torchrun_out(root=0)[source]

sys.stdout on the root rank, a discarded buffer elsewhere.

Parameters:

root (int)