Now, we run individual prompts through the queue.

[finbarrtimbers]

Changes the queue setup so that we run individual prompts through the queues rather than batches. This will avoid the need for load balancing and should enable us to get some nice speed ups once we enable the async engine (next PR!).

Benchmark results! From HEAD (spoiler, basically no change):

Results (excluding first batch):
Average tokens/second: 1014.84
Average MFU: 3.96%
Average generation time per batch: 2.02s
Average new tokens per sample: 2048.0 tokens
Wasted compute % (variable response length): 0.00%

From this PR:

Results (excluding first batch):
Average tokens/second: 1008.56
Average MFU: 3.94%
Average generation time per batch: 2.04s
Average new tokens per sample: 2048.0 tokens
Wasted compute % (variable response length): 0.00%

[mnoukhov]

I'm confused why (1) the length of samples goes down and also (2) the time per batch goes down if the speed goes down too

[finbarrtimbers]

I'm confused why (1) the length of samples goes down and also (2) the time per batch goes down if the speed goes down too

@mnoukhov it's because I accidentally changed the batch size between runs! That's not the length, but the total number of tokens generated in the batch. Nice catch.

[mnoukhov]

Did a quick pass and seems good. Definitely a good precursor to async

[mnoukhov]

Seems like this can have more prompts than batch_size (unless .get only gets one at a time?)

I'm also confused when would you have a batch that is not full but due to timeout you will take what you get?

[finbarrtimbers]

Ok, I have fixed this! Now, PromptRequest only contains a single prompt, and I have removed the timeout, so it will block until it gets a full batch.

[mnoukhov]

Makes sense, you can remove the comment about non-blocking

[mnoukhov]

other than the little questions about specific batch sizes and divisibility, seems good to me

[mnoukhov]

Makes sense, you can remove the comment about non-blocking

[mnoukhov]

is it possible that at the end of an epoch or something you will have fewer than this number of results? or are we dropping the last batch

[finbarrtimbers]

ShufflingIterator guarantees that they're all the same size:

 class ShufflingIterator:
      def __init__(self, data: np.ndarray, batch_size: int, seed: Optional[int] = None):
          self.data = data.copy()
          self.batch_size = batch_size
          self.index = 0
          self.rng = np.random.default_rng(seed)
          self.rng.shuffle(self.data)

          # Ensure the effective dataset size is divisible by batch_size                                                  
          self.effective_size = len(self.data) - (len(self.data) % batch_size)

      def __iter__(self) -> Iterator[List[int]]:
          return self

      def __next__(self) -> List[int]:
          if self.index >= self.effective_size:
              self.index = 0
              self.rng.shuffle(self.data)

          end_index = self.index + self.batch_size
          batch = self.data[self.index : end_index].tolist()
          self.index = end_index

          return batch

[mnoukhov]

do we know if this is equally divisible? and is it alright if its not?

[finbarrtimbers]

It should be divisible; I changed the code to raise a ValueError if it isn't, and we can handle that then.