flex_model.distributed.initialize_distributed_state

flex_model.distributed.initialize_distributed_state(module: Module, tensor_parallel_size: int, pipeline_parallel_size: int, data_parallel_size: int, process_group: ProcessGroup | None = None)