flex_model.distributed.broadcast_data_parallel

flex_model.distributed.broadcast_data_parallel(tensor: Tensor, fmps: _ParallelStateAPI) → Tensor

Broadcast tensor to all ranks in the data parallel group.

Parameters:

• tensor (Tensor) – Activation tensor.

• fmps (_ParallelStateAPI) – FlexModel parallel state handle.

Returns:

Input tensor unmodified.

Return type:

Tensor