Add assert info

shuyixiong · shuyixiong · commit fdd2dcb29657 · 2025-12-16T01:47:58.000-08:00
Signed-off-by: Shuyi Xiong &lt;219646547+shuyixiong@users.noreply.github.com&gt;
diff --git a/tensorrt_llm/executor/ray_gpu_worker.py b/tensorrt_llm/executor/ray_gpu_worker.py
@@ -91,6 +91,8 @@ def setup_distributed_env_and_worker(self, port: int):
                                              store=self.store,
                                              world_size=self.world_size,
                                              rank=self.rank)
+        assert torch.distributed.get_world_size(
+        ) == self.world_size, "World size mismatch"
         logger.info(
             f"[Rank {self.rank}] Finished PG init. Global GPU ID: {self.gpu}, local GPU ID: {self.local_gpu}"
         )

Original file line number	Diff line number	Diff line change
`@@ -91,6 +91,8 @@ def setup_distributed_env_and_worker(self, port: int):`
`91`	`91`	`store=self.store,`
`92`	`92`	`world_size=self.world_size,`
`93`	`93`	`rank=self.rank)`
	`94`	`+ assert torch.distributed.get_world_size(`
	`95`	`+ ) == self.world_size, "World size mismatch"`
`94`	`96`	`logger.info(`
`95`	`97`	`f"[Rank {self.rank}] Finished PG init. Global GPU ID: {self.gpu}, local GPU ID: {self.local_gpu}"`
`96`	`98`	`)`