(
tp_group: paddle.distributed.communication.group.Group = None, custom_all_reduce_max_bytes: int = 8192 * 1024
)
| 36 | |
| 37 | |
| 38 | def use_custom_allreduce( |
| 39 | tp_group: paddle.distributed.communication.group.Group = None, custom_all_reduce_max_bytes: int = 8192 * 1024 |
| 40 | ): |
| 41 | if tp_group is None: |
| 42 | hcg = fleet.get_hybrid_communicate_group() |
| 43 | tp_group = hcg.get_model_parallel_group() |
| 44 | global _TP_AR |
| 45 | from fastdeploy.distributed.custom_all_reduce import CustomAllreduce |
| 46 | |
| 47 | _TP_AR = CustomAllreduce(tp_group, custom_all_reduce_max_bytes) |
| 48 | |
| 49 | |
| 50 | def custom_ar_clear_ipc_handles(): |
no test coverage detected