mixle.utils.parallel.torch_neural module

Distributed STREAMING neural-training handle – the inverted sibling of TorchRunEncodedData.

TorchRunEncodedData gathers each rank’s sufficient statistic to root, runs the M-step on root, and pickle- broadcasts the model back – structurally the wrong shape for a sharded gradient update (it cannot scale past a model that fits, and is folded on, one rank). StreamingTokenEncodedData inverts the reduction: each rank keeps its OWN streamed token shard resident and trains the estimator’s module SPMD, and the SOLE cross-rank collective is the in-backward gradient all-reduce (DDP) / reduce-scatter (FSDP2) – never a gather-to-root + a pickle-broadcast. Each rank ends with the same (DDP/FSDP2-consistent) model, so seq_estimate needs no gather.

It plugs into the same duck-typed dispatch (seq_estimate calls handle.pysp_seq_estimate). Run under torchrun (RANK/WORLD_SIZE set); the process group is initialized from the environment. On CUDA, swap the DDP wrap for torch.distributed.fsdp.fully_shard (ZeRO-3) – the architecture is identical, that is the one-line change.

class StreamingTokenEncodedData(token_ids, *, block, batch_size, epochs=1, shuffle=True, shard_by_rank=True, init_process_group=True, backend=None, parallel='auto', precision='fp32', activation_checkpointing=False)[source]

Bases: EncodedDataHandle

Per-rank streamed token shard; pysp_seq_estimate trains the module SPMD with no gather-to-root.

Pass a token-id array; with shard_by_rank each rank keeps a disjoint slice. The estimator is a StreamingTransformerLeafEstimator carrying the module + lr.

Parameters:

  • token_ids (Any)

  • block (int)

  • batch_size (int)

  • epochs (int)

  • shuffle (bool)

  • shard_by_rank (bool)

  • init_process_group (bool)

  • backend (str | None)

  • parallel (str)

  • precision (str)

  • activation_checkpointing (bool)

pysp_seq_estimate(estimator, prev_estimate)[source]

SPMD: stream this rank’s shard, train the module; the only collective is the in-backward all-reduce.

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

close()[source]

Release worker resources owned by this handle, if any.

Return type:

None