Add next_rank and prev_rank members

2026-01-06 10:51:21 +00:00
parent 90a7e6601d
commit 5f3bc30f17
1 changed files with 2 additions and 0 deletions
@@ -63,6 +63,8 @@ class DistributedDenoising:
        self.group = group
        self.rank = shard_metadata.device_rank
        self.world_size = shard_metadata.world_size
+        self.next_rank = (self.rank + 1) % self.world_size
+        self.prev_rank = (self.rank - 1 + self.world_size) % self.world_size
        self.start_layer = shard_metadata.start_layer
        self.end_layer = shard_metadata.end_layer