tensorplay.distributed.tensor.parallel

Parallelizing a module

parallelize_module() takes the module, a device mesh, and either a single style or a mapping from submodule names to styles. Inside the transformed module, parameters and inputs are distributed tensors with the placements the style chose; the runtime inserts the collectives that assemble partial results where needed.

Sharding styles

• ColwiseParallel — split a linear layer’s weight along its first dimension (the output dimension), so each rank produces a partial output column-block. The input stays replicated; the output is sharded along its last dimension.

• RowwiseParallel — split the weight along its second dimension (the input dimension). The input is turned into a sharded tensor, and the local partial output is gathered back into a replicated tensor.

• SequenceParallel — keep every parameter replicated while sharding the hidden states along the sequence dimension. It is meant to be wrapped around the colwise part of a transformer block, so the all-gather that restores an unsharded intermediate lands at the end of the block.

• PrepareModuleInput, PrepareModuleOutput and PrepareModuleInputOutput — convert inputs and outputs of a submodule to and from the placements a style expects, e.g. turning an unsharded input into a sharded one.

Loss with sharded logits

loss_parallel() is a context manager that activates the distributed cross-entropy path for logits sharded along the class dimension. When the model’s logits stay sharded through the loss (instead of being gathered into a replicated tensor first), the cross-entropy reduces the sharded logits directly and the loss gradient is scaled for the shard size, saving the all-gather that would otherwise materialize the full logits tensor on every rank. Wrap the loss computation in with loss_parallel(): to enable it.

Where to go next

• distributed tensors — the DTensor type these styles operate on.

• device mesh — creating the mesh that parallelize_module runs on.

• the distributed package — process groups and collectives underneath.