fix consumer

Tong Li · Tong Li · commit afddfde2dd28 · 2025-03-13T14:55:26.000+08:00
diff --git a/applications/ColossalChat/coati/distributed/consumer.py b/applications/ColossalChat/coati/distributed/consumer.py
@@ -73,6 +73,8 @@ def setup(self) -> None:
         )
         if self.plugin_config.get("pp_size", 1) > 1 and "num_microbatches" not in self.plugin_config:
             plugin_config["microbatch_size"] = self.microbatch_size
+        if self.plugin_config.get("tp_size", 1) > 1:
+            plugin_config["parallel_output"] = False
         plugin_config.update(self.plugin_config)
         self.plugin = HybridParallelPlugin(**plugin_config)
         self.booster = Booster(plugin=self.plugin)

Original file line number	Diff line number	Diff line change
`@@ -73,6 +73,8 @@ def setup(self) -> None:`
`73`	`73`	`)`
`74`	`74`	`if self.plugin_config.get("pp_size", 1) > 1 and "num_microbatches" not in self.plugin_config:`
`75`	`75`	`plugin_config["microbatch_size"] = self.microbatch_size`
	`76`	`+ if self.plugin_config.get("tp_size", 1) > 1:`
	`77`	`+ plugin_config["parallel_output"] = False`
`76`	`78`	`plugin_config.update(self.plugin_config)`
`77`	`79`	`self.plugin = HybridParallelPlugin(**plugin_config)`
`78`	`80`	`self.booster = Booster(plugin=self.plugin)`