diff --git a/colossalai/booster/plugin/gemini_plugin.py b/colossalai/booster/plugin/gemini_plugin.py index 6622b6dc1..a891db422 100644 --- a/colossalai/booster/plugin/gemini_plugin.py +++ b/colossalai/booster/plugin/gemini_plugin.py @@ -437,6 +437,10 @@ class GeminiPlugin(DPPluginBase): enable_sequence_overlap=self.enable_sequence_overlap, ) + def __del__(self): + """Destroy the prcess groups in ProcessGroupMesh""" + self.pg_mesh.destroy_mesh_process_groups() + def support_no_sync(self) -> bool: return False diff --git a/colossalai/booster/plugin/hybrid_parallel_plugin.py b/colossalai/booster/plugin/hybrid_parallel_plugin.py index c52de0ba7..205660f94 100644 --- a/colossalai/booster/plugin/hybrid_parallel_plugin.py +++ b/colossalai/booster/plugin/hybrid_parallel_plugin.py @@ -1054,6 +1054,10 @@ class HybridParallelPlugin(PipelinePluginBase): self.max_norm = max_norm + def __del__(self): + """Destroy the prcess groups in ProcessGroupMesh""" + self.pg_mesh.destroy_mesh_process_groups() + @property def enable_pipeline_parallelism(self) -> bool: return self.pp_size > 1 diff --git a/colossalai/cluster/process_group_mesh.py b/colossalai/cluster/process_group_mesh.py index 7a3bde448..ae3956c69 100644 --- a/colossalai/cluster/process_group_mesh.py +++ b/colossalai/cluster/process_group_mesh.py @@ -45,7 +45,7 @@ class ProcessGroupMesh: self._ranks_to_group: Dict[Tuple[int, ...], ProcessGroup] = {} self._group_to_ranks: Dict[ProcessGroup, Tuple[int, ...]] = {} - def __del__(self): + def destroy_mesh_process_groups(self): r""" Destructor method for the ProcessGroupMesh class.