[test] mixtra pp shard test

2025-09-04 02:26:51 +00:00 · 2024-07-04 06:39:01 +00:00
parent 8ae8525bdf
commit a249e71946
3 changed files with 49 additions and 46 deletions
--- a/colossalai/booster/plugin/moe_hybrid_parallel_plugin.py
+++ b/colossalai/booster/plugin/moe_hybrid_parallel_plugin.py
@@ -411,7 +411,7 @@ class MoeHybridParallelPlugin(HybridParallelPlugin):
                reinitialize_optimizer(optimizer, model)

            if self.zero_stage == 0:
-                assert self.ep_size > 1
+                # assert self.ep_size > 1

                if self.precision in ["fp16", "bf16"]:
                    optimizer = HybridParallelAMPOptimizer(