Merge pull request #6071 from wangbluo/ring_attention

[Ring Attention] fix the 2d ring attn when using multiple machine
2025-09-06 19:40:28 +00:00 · 2024-10-15 15:17:21 +08:00
parent 4c8e85ee0d 83cf2f84fb
commit dcd41d0973
6 changed files with 41 additions and 31 deletions
--- a/colossalai/booster/plugin/hybrid_parallel_plugin.py
+++ b/colossalai/booster/plugin/hybrid_parallel_plugin.py
@@ -1177,7 +1177,10 @@ class HybridParallelPlugin(PipelinePluginBase):
            gradient_checkpoint_config=gradient_checkpoint_config,
            fp8_communication=fp8_communication,
            inner_ring_size=inner_ring_size,
+            pg_mesh=self.pg_mesh,
+            sp_axis=self.sp_axis,
        )
+
        self.amp_config = dict(
            initial_scale=initial_scale,
            growth_factor=growth_factor,