fix

2025-09-02 01:28:31 +00:00 · 2024-10-15 11:01:34 +08:00
parent 8ff7d0c780
commit 3dc08c8a5a
6 changed files with 18 additions and 15 deletions
--- a/colossalai/booster/plugin/hybrid_parallel_plugin.py
+++ b/colossalai/booster/plugin/hybrid_parallel_plugin.py
@@ -1180,7 +1180,9 @@ class HybridParallelPlugin(PipelinePluginBase):
            gradient_checkpoint_config=gradient_checkpoint_config,
            fp8_communication=fp8_communication,
            inner_ring_size=inner_ring_size,
+            pg_mesh=self.pg_mesh,
        )
+
        self.amp_config = dict(
            initial_scale=initial_scale,
            growth_factor=growth_factor,