[BugFix] Fix seeded random sampling with encoder-decoder models (#8870)

Co-authored-by: Roger Wang <ywang@roblox.com>
This commit is contained in:
Nick Hill
2024-09-29 10:43:14 +01:00
committed by GitHub
parent 3d49776bbb
commit 31f46a0d35

View File

@ -268,11 +268,13 @@ class EncoderDecoderModelRunner(GPUModelRunnerBase[EncoderDecoderModelInput]):
encoder_input_positions=encoder_input_positions_tensor,
)
generators = self.get_generators(finished_requests_ids)
sampling_metadata = SamplingMetadata.prepare(seq_group_metadata_list,
model_input.seq_lens,
model_input.query_lens,
self.device,
self.pin_memory)
self.pin_memory,
generators=generators)
is_prompt = (seq_group_metadata_list[0].is_prompt
if seq_group_metadata_list else None)
return dataclasses.replace(model_input,