Skip to content
This repository was archived by the owner on Dec 1, 2024. It is now read-only.

fix torchrun inference #112

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix torchrun inference #112

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 8 additions & 1 deletion flexgen/dist_flex_opt.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -219,7 +219,6 @@ def store_hidden(self, b, t, i, j, k):
hidden_val = self.hidden[t][i][j][k].pop()
else:
hidden_val = self.hidden[t][i][j][k].val

ids = hidden_val.data.detach().cpu().numpy()
gpu_batch_size = self.policy.gpu_batch_size
num_gpu_batches = self.num_gpu_batches
Expand Down Expand Up @@ -478,7 +477,9 @@ def generation_loop_overlap_one_batch(self):
self.load_weight(b, t, i, j+1, 0)
self.load_cache(t, i, j+1, 0)
self.load_hidden(b, t, i, j, 0)
self.sync()
self.compute_layer(t, i, j, 0)
self.sync()
self.store_cache(t, i, j-1, 0)
self.store_hidden(b, t, i, j, 0)
self.sync()
Expand Down Expand Up @@ -515,7 +516,9 @@ def generation_loop_overlap_multi_batch(self):
self.load_weight(b, t, i, j + 1, k)
self.load_cache(t, i, j, k + 1)
self.load_hidden(b, t, i, j, k)
self.sync()
self.compute_layer(t, i, j, k)
self.sync()
self.store_cache(t, i, j, k - 1)
self.store_hidden(b, t, i, j, k)
self.sync()
Expand Down Expand Up @@ -674,6 +677,10 @@ def add_distributed_parser_arguments(parser):
args.world_size = int(os.getenv('OMPI_COMM_WORLD_SIZE'))
args.rank = int(os.getenv('OMPI_COMM_WORLD_RANK'))
args.local_rank = int(os.getenv('OMPI_COMM_WORLD_LOCAL_RANK'))
else:
args.world_size = int(os.getenv('WORLD_SIZE'))
args.rank = int(os.getenv('RANK'))
args.local_rank = int(os.getenv('LOCAL_RANK'))
initialize_distributed(args.head_ip, args.port, args.world_size,
args.rank, args.local_rank, args.comm_device)
else:
Expand Down