lalamo 0.6.4__py3-none-any.whl → 0.6.6__py3-none-any.whl

lalamo/models/language_model.py

@@ -1,14 +1,16 @@
-from collections.abc import Iterable
-from dataclasses import dataclass
+from collections.abc import Callable, Iterable, Iterator
+from dataclasses import dataclass, replace
+from itertools import batched
 from pathlib import Path
 from typing import NamedTuple
 
 import equinox as eqx
 import jax
 import jax.numpy as jnp
-from einops import rearrange
+import numpy as np
+from einops import rearrange, repeat
 from jax import vmap
-from jaxtyping import Array, Bool, Float, Int, PRNGKeyArray
+from jaxtyping import Array, Bool, Float, Int, Key, PRNGKeyArray
 
 from lalamo.message_processor import AssistantMessage, Message, MessageProcessor
 from lalamo.modules import (
@@ -21,7 +23,15 @@ from lalamo.modules import (
 )
 from lalamo.sampling import SamplingPolicy, make_policy
 
-from .common import TextModel, TextModelConfig
+from .common import BatchSizeInfo, BatchSizesComputedEvent, InferenceConfig, TextModel, TextModelConfig
+from .compile_helpers import compile_generate_tokens
+from .lm_helpers import (
+    decrease_batchsize_on_oom,
+    estimate_batchsizes_from_vram,
+    merge_small_buckets,
+    pad_keys_to_size,
+    pad_sequences,
+)
 
 __all__ = [
     "ForwardPassConfig",
@@ -31,7 +41,7 @@ __all__ = [
 ]
 
 
-_COMPILED_PROMPT_LENGTHS = [512 * 2**i for i in range(10)]
+_COMPILED_PROMPT_LENGTHS = [256 * 2**i for i in range(12)]
 
 
 type ForwardPassConfig = DecoderForwardPassConfig
@@ -94,6 +104,13 @@ class LanguageModelConfig(TextModelConfig[DecoderConfig]):
         return result
 
 
+class Chunk(eqx.Module):
+    tokens: Int[Array, "num_chunks batch chunk_size"]
+    indices: Int[Array, "num_chunks batch chunk_size"]
+    sequence_ends: Int[Array, "num_chunks batch"]
+    is_last_token_inside: Bool[Array, "num_chunks batch"]
+
+
 class LanguageModel(TextModel[LanguageModelConfig, Decoder]):
     @property
     def stop_token_ids(self) -> tuple[int, ...]:
@@ -102,68 +119,143 @@ class LanguageModel(TextModel[LanguageModelConfig, Decoder]):
     def default_sampling_policy(self) -> SamplingPolicy:
         return self.config.generation_config.default_policy()
 
+    @eqx.filter_jit
+    def _make_chunks(
+        self,
+        token_ids: Int[Array, "batch tokens"],
+        lengths_without_padding: Int[Array, " batch"] | None,
+        chunk_size: int,
+    ) -> Chunk:
+        batch_size, sequence_length = token_ids.shape
+        if lengths_without_padding is None:
+            lengths_without_padding = jnp.full((batch_size,), sequence_length, dtype=jnp.int32)
+
+        # If all sequences fit in a single chunk, use sequence_length as the chunk size
+        chunk_size = min(chunk_size, sequence_length)
+
+        n_chunks = (sequence_length + chunk_size - 1) // chunk_size
+        padded_length = n_chunks * chunk_size
+
+        token_ids = jnp.pad(token_ids, [(0, 0), (0, padded_length - sequence_length)])
+
+        # Reshape tokens to (num_chunks, batch, chunk_size)
+        tokens = rearrange(
+            token_ids,
+            "batch (num_chunks chunk_size) -> num_chunks batch chunk_size",
+            chunk_size=chunk_size,
+        )
+
+        # Create position indices (num_chunks, batch, chunk_size)
+        indices = jnp.arange(padded_length, dtype=jnp.int32)
+        indices = repeat(indices, "token_idx -> batch token_idx", batch=batch_size)
+        indices = rearrange(
+            indices,
+            "batch (num_chunks chunk_size) -> num_chunks batch chunk_size",
+            chunk_size=chunk_size,
+        )
+
+        # sequence_ends: for each chunk, how many valid tokens per batch item
+        chunk_starts = jnp.arange(n_chunks, dtype=jnp.int32) * chunk_size
+        sequence_ends = jnp.clip(
+            lengths_without_padding[None, :] - chunk_starts[:, None],
+            0,
+            chunk_size,
+        )
+
+        # last_token_inside: whether the last valid token (at index length-1) is in this chunk
+        last_token_idx = lengths_without_padding - 1
+        chunk_ends = chunk_starts + chunk_size
+        is_last_token_inside = (last_token_idx[None, :] >= chunk_starts[:, None]) & (
+            last_token_idx[None, :] < chunk_ends[:, None]
+        )
+
+        return Chunk(
+            tokens=tokens,
+            indices=indices,
+            sequence_ends=sequence_ends,
+            is_last_token_inside=is_last_token_inside,
+        )
+
     @eqx.filter_jit
     def _prefill(
         self,
         token_ids: Int[Array, "batch tokens"],
+        state_capacity: int,
         lengths_without_padding: Int[Array, " batch"] | None = None,
-        state_capacity: int | None = None,
         forward_pass_config: ForwardPassConfig | None = None,
+        chunk_size: int = 512,  # vllm default
     ) -> PrefillResults:
         batch_size, sequence_length = token_ids.shape
-        token_positions = jnp.repeat(jnp.arange(sequence_length, dtype=jnp.int32)[None, ...], batch_size, axis=0)
-        if state_capacity is not None:
-            state = self.model.init_static_state(batch_size, state_capacity)
-        else:
-            state = None
 
-        decoder_outputs = self.model(
-            token_ids,
-            token_positions,
-            state,
-            return_updated_state=True,
-            lengths_without_padding=lengths_without_padding,
-            forward_pass_mode=ForwardPassMode.MULTI_TOKEN,
-            forward_pass_config=forward_pass_config,
-        )
+        if lengths_without_padding is None:
+            lengths_without_padding = jnp.full((batch_size,), sequence_length, dtype=jnp.int32)
 
-        if lengths_without_padding is not None:
-            last_logits_indices = lengths_without_padding - 1
-        else:
-            last_logits_indices = jnp.array([sequence_length - 1] * batch_size, dtype=jnp.int32)
+        chunks = self._make_chunks(token_ids, lengths_without_padding, chunk_size)
+
+        num_chunks, _, chunk_size = chunks.tokens.shape
+        state_capacity = max(state_capacity, num_chunks * chunk_size)
 
-        last_token_logits = vmap(lambda logits, index: logits[index])(decoder_outputs.logits, last_logits_indices)
+        state = self.model.init_static_state(batch_size, state_capacity)
+        logits_like = jnp.zeros((batch_size, self.model.vocab_size), dtype=jnp.float32)
+
+        def apply_chunk(state_and_logits: tuple, chunk: Chunk) -> tuple:
+            state, prev_logits = state_and_logits
+            decoder_outputs = self.model(
+                chunk.tokens,
+                chunk.indices,
+                state,
+                return_updated_state=True,
+                lengths_without_padding=chunk.sequence_ends,
+                forward_pass_mode=ForwardPassMode.MULTI_TOKEN,
+                forward_pass_config=forward_pass_config,
+            )
+            assert decoder_outputs.updated_state is not None
+
+            chunk_logits = decoder_outputs.logits[jnp.arange(batch_size), chunk.sequence_ends - 1, :]
+            new_logits = jnp.where(chunk.is_last_token_inside[:, None], chunk_logits, prev_logits)
+
+            return (decoder_outputs.updated_state, new_logits), None
+
+        (final_state, final_logits), _ = jax.lax.scan(apply_chunk, (state, logits_like), chunks)
 
-        assert decoder_outputs.updated_state is not None
         return PrefillResults(
-            last_token_logits=last_token_logits,
-            last_token_indices=last_logits_indices,
-            state=decoder_outputs.updated_state,
+            last_token_logits=final_logits,
+            last_token_indices=jnp.maximum(lengths_without_padding - 1, 0),
+            state=final_state,
         )
 
-    @eqx.filter_jit
     def generate_tokens(
         self,
         prompt_token_ids: Int[Array, "batch prompt_tokens"],
-        sampling_policy: SamplingPolicy | None = None,
+        generation_config: GenerationConfig | None = None,
         prompt_lengths_without_padding: Int[Array, " batch"] | None = None,
         max_output_length: int = 8192,
         eos_token_ids: Int[Array, " eos_tokens"] | None = None,
         forward_pass_config: ForwardPassConfig | None = None,
         num_top_logits_to_return: int | None = None,
         *,
-        key: PRNGKeyArray | None = None,
+        keys: Key[Array, " batch"] | None = None,
     ) -> GenerationResults:
-        if sampling_policy is None:
-            sampling_policy = self.default_sampling_policy()
+        batch_size, sequence_length = prompt_token_ids.shape
+
+        sampling_policy = self.default_sampling_policy()
+        if generation_config is not None:
+            sampling_policy = generation_config.default_policy()
+
         if eos_token_ids is None:
             eos_token_ids = jnp.array(self.stop_token_ids, dtype=jnp.int32)
+        if keys is None:
+            keys = jax.random.split(jax.random.key(0), num=batch_size)
+
+        if len(keys) != batch_size:
+            raise ValueError(
+                f"Length of 'keys' should be equal to the batch size, or keys should be None; got {len(keys)}",
+            )
 
-        batch_size, sequence_length = prompt_token_ids.shape
         prefill_results = self._prefill(
             prompt_token_ids,
-            prompt_lengths_without_padding,
             sequence_length + max_output_length,
+            prompt_lengths_without_padding,
             forward_pass_config=forward_pass_config,
         )
 
@@ -174,18 +266,14 @@ class LanguageModel(TextModel[LanguageModelConfig, Decoder]):
             jnp.zeros(batch_size, dtype=jnp.bool),
         )
 
-        if key is None:
-            key = jax.random.PRNGKey(0)
-        keys = jax.random.split(key, num=max_output_length)
-
         def loop_iteration(
             state: DecodingState,
-            key: PRNGKeyArray,
+            keys: Key[Array, " batch"],
         ) -> tuple[DecodingState, GenerationStepResults]:
             def sample_and_update() -> tuple[DecodingState, GenerationStepResults]:
                 upcasted_logits = state.last_token_logits.astype(jnp.float32)
                 processed_logits = vmap(sampling_policy.process_logits)(upcasted_logits)
-                next_token_ids = jax.random.categorical(key, processed_logits)
+                next_token_ids = jax.vmap(lambda k, logits: jax.random.categorical(k, logits))(keys, processed_logits)
                 next_token_ids = jnp.where(state.stop_flags, jnp.zeros(batch_size, dtype=jnp.int32), next_token_ids)
                 if num_top_logits_to_return is not None:
                     next_top_k_token_logits, next_top_k_token_ids = jax.lax.top_k(
@@ -214,7 +302,7 @@ class LanguageModel(TextModel[LanguageModelConfig, Decoder]):
                 )
                 assert decoder_outputs.updated_state is not None, "updated_state should not be None"
                 new_state = DecodingState(
-                    decoder_outputs.logits.squeeze(1),
+                    decoder_outputs.logits.squeeze(1).astype(jnp.float32),
                     next_token_indices,
                     decoder_outputs.updated_state,
                     stop_flags,
@@ -234,7 +322,9 @@ class LanguageModel(TextModel[LanguageModelConfig, Decoder]):
 
             return jax.lax.cond(jnp.all(state.stop_flags), pad_and_repeat_state, sample_and_update)
 
-        _, generated = jax.lax.scan(loop_iteration, initial_state, keys)
+        per_step_keys: Key[Array, "batch max_len"] = jax.vmap(lambda k: jax.random.split(k, max_output_length))(keys)
+        per_step_keys: Key[Array, "max_len batch"] = jnp.swapaxes(per_step_keys, 0, 1)
+        _, generated = jax.lax.scan(loop_iteration, initial_state, per_step_keys)
 
         token_ids = rearrange(generated.token_ids, "iteration batch -> batch iteration")
 
@@ -247,29 +337,222 @@ class LanguageModel(TextModel[LanguageModelConfig, Decoder]):
 
         return GenerationResults(token_ids, top_k_token_ids, top_k_token_logits)
 
+    def _generate_tokens_batch(
+        self,
+        batch: tuple[list[int], ...],
+        batch_keys: tuple[Key[Array, ""], ...],
+        *,
+        generation_config: GenerationConfig | None,
+        inference_config: InferenceConfig,
+        forward_pass_config: ForwardPassConfig | None,
+    ) -> Iterator[GenerationResults]:
+        assert inference_config.batch_size is not None
+        batch_size = inference_config.batch_size
+
+        padded_token_ids = pad_sequences(batch, (batch_size, inference_config.padded_length), dtype=jnp.int32)
+
+        lengths = jnp.array([len(tokens) for tokens in batch], dtype=jnp.int32)
+        padded_lengths = jnp.pad(lengths, (0, batch_size - len(batch)))
+
+        padded_keys = pad_keys_to_size(batch_keys, batch_size)
+
+        generate_tokens_fn = compile_generate_tokens(
+            self,
+            generation_config,
+            inference_config,
+            forward_pass_config=forward_pass_config,
+        )
+        results = generate_tokens_fn(
+            self,
+            prompt_token_ids=padded_token_ids,
+            prompt_lengths_without_padding=padded_lengths,
+            keys=padded_keys,
+        )
+        for i in range(len(batch)):
+            yield GenerationResults(
+                token_ids=results.token_ids[i],
+                top_k_token_ids=results.top_k_token_ids[i] if results.top_k_token_ids is not None else None,
+                top_k_token_logits=results.top_k_token_logits[i] if results.top_k_token_logits is not None else None,
+            )
+
+    def generate_tokens_many(
+        self,
+        tokenized: Iterable[list[int]],
+        generation_config: GenerationConfig | None = None,
+        inference_config: InferenceConfig = InferenceConfig(),  # noqa: B008
+        *,
+        forward_pass_config: ForwardPassConfig | None = None,
+        keys: Key[Array, " num_sequences"] | None = None,
+    ) -> Iterator[GenerationResults]:
+        tokenized = list(tokenized)  # load eagerly to RAM
+
+        if keys is None:
+            keys = jax.random.split(jax.random.key(0), num=len(tokenized))
+
+        if len(keys) != len(tokenized):
+            raise ValueError(
+                f"Length of 'keys' should be equal to the number of sequences passed or None; got {len(keys)}",
+            )
+
+        def process_batches(batch_size: int) -> Iterator[tuple[int, GenerationResults]]:
+            new_inference_config = replace(inference_config, batch_size=batch_size)
+
+            for batch_items in batched(zip(tokenized, keys, strict=True), batch_size):
+                real_batch, batch_keys = zip(*batch_items, strict=True)
+                yield from self._generate_tokens_batch(
+                    real_batch,
+                    batch_keys,
+                    generation_config=generation_config,
+                    inference_config=new_inference_config,
+                    forward_pass_config=forward_pass_config,
+                )
+
+        assert inference_config.batch_size is not None
+        yield from decrease_batchsize_on_oom(
+            process_batches,
+            starting_batch_size=inference_config.batch_size,
+        )
+
+    def estimate_memory_consumption(
+        self,
+        generation_config: GenerationConfig | None = None,
+        inference_config: InferenceConfig = InferenceConfig(),  # noqa: B008
+        *,
+        forward_pass_config: ForwardPassConfig | None = None,
+    ) -> int:
+        memory_analysis = compile_generate_tokens(
+            self,
+            generation_config=generation_config,
+            inference_config=inference_config,
+            forward_pass_config=forward_pass_config,
+        ).memory_analysis()
+
+        assert hasattr(memory_analysis, "argument_size_in_bytes")
+        assert hasattr(memory_analysis, "output_size_in_bytes")
+        assert hasattr(memory_analysis, "temp_size_in_bytes")
+
+        return (
+            memory_analysis.argument_size_in_bytes
+            + memory_analysis.output_size_in_bytes
+            + memory_analysis.temp_size_in_bytes
+        )
+
+    def _trim_at_eos(self, token_ids: list[int]) -> list[int]:
+        if not self.stop_token_ids:
+            return token_ids
+        stop_set = set(self.stop_token_ids)
+        end = next((i for i, token_id in enumerate(token_ids) if token_id in stop_set), len(token_ids))
+        return token_ids[: end + 1]
+
     def reply(
         self,
         messages: Iterable[Message],
-        sampling_policy: SamplingPolicy | None = None,
-        forward_pass_config: ForwardPassConfig | None = None,
+        generation_config: GenerationConfig | None = None,
         *,
+        forward_pass_config: ForwardPassConfig | None = None,
         key: PRNGKeyArray | None = None,
     ) -> AssistantMessage:
         formatted_messages = self.message_processor.render_request(messages)
         token_ids = jnp.array(self.message_processor.tokenize_text(formatted_messages), dtype=jnp.int32)[None, :]
         response_ids = self.generate_tokens(
             token_ids,
-            sampling_policy,
+            generation_config,
             forward_pass_config=forward_pass_config,
-            key=key,
+            keys=key[None, ...] if key is not None else None,
         ).token_ids.squeeze(0)
-        response_text = self.message_processor.detokenize(response_ids.tolist())
+        trimmed_ids = self._trim_at_eos(response_ids.tolist())
+        response_text = self.message_processor.detokenize(trimmed_ids)
         return self.message_processor.parse_response(response_text)
 
+    def reply_many(
+        self,
+        messages: Iterable[Iterable[Message]],
+        generation_config: GenerationConfig | None = None,
+        inference_config: InferenceConfig = InferenceConfig(),  # noqa: B008
+        *,
+        forward_pass_config: ForwardPassConfig | None = None,
+        keys: Key[Array, " num_sequences"] | None = None,
+        vram_bytes: int | None = None,
+        batch_sizes_callback: Callable[[BatchSizesComputedEvent], None] | None = None,
+    ) -> Iterator[tuple[int, AssistantMessage]]:
+        messages = list(messages)  # eagerly load the dataset into RAM
+
+        if keys is None:
+            keys = jax.random.split(jax.random.key(0), num=len(messages))
+
+        if len(keys) != len(messages):
+            raise ValueError(
+                f"Length of 'keys' should be equal to the number of sequences passed or None; got {len(keys)}",
+            )
+
+        if vram_bytes is not None and inference_config.batch_size is not None:
+            raise ValueError("You have to specify only one of batch_size and vram_gb, not both.")
+
+        if vram_bytes is None and inference_config.batch_size is None:
+            raise ValueError("You have to specify either batch_size or vram_gb, but you provided neither.")
+
+        tokenized: list[list[int]] = self.message_processor.tokenize_requests(messages)
+
+        buckets: dict[int, list[tuple[int, list[int]]]] = {}
+        max_prompt_length = max(_COMPILED_PROMPT_LENGTHS)
+        for idx, sequence in enumerate(tokenized):
+            assert len(sequence) <= max_prompt_length, (
+                f"Sequence length {len(sequence)} exceeds largest bucket {max_prompt_length}"
+            )
+            # we choose the smallest size from precomputed ones that is longer or equal to the current sequence
+            padded_len = min(length for length in _COMPILED_PROMPT_LENGTHS if length >= len(sequence))
+            buckets.setdefault(padded_len, []).append((idx, sequence))
+        sorted_lengths = sorted(buckets.keys())
+
+        if inference_config.batch_size is not None:
+            batch_size_per_bucket = dict.fromkeys(sorted_lengths, inference_config.batch_size)
+        else:
+            batch_size_per_bucket = estimate_batchsizes_from_vram(
+                lambda config: self.estimate_memory_consumption(inference_config=config),
+                sorted_lengths,
+                vram_bytes,  # type: ignore
+                inference_config,
+            )
+
+        buckets = merge_small_buckets(buckets, batch_size_per_bucket, min_batches=2)
+        assert sum(len(bucket) for bucket in buckets.values()) == len(tokenized)
+
+        if batch_sizes_callback is not None:
+            batch_sizes = tuple(
+                BatchSizeInfo(
+                    prefix_length=padded_length,
+                    num_elements=len(buckets[padded_length]),
+                    batch_size=batch_size_per_bucket.get(padded_length, 1),
+                )
+                for padded_length in sorted(buckets.keys())
+            )
+            batch_sizes_callback(BatchSizesComputedEvent(batch_sizes=batch_sizes))
+
+        # Process longest sequences first so batchsize=1 OOM happens as early as possible, if it does happen
+        for padded_length in sorted(buckets.keys(), reverse=True):
+            sequence_ids, sequence_tokenized = zip(*buckets[padded_length], strict=True)
+            sequence_ids = list(sequence_ids)
+            batch_size = batch_size_per_bucket[padded_length]
+
+            bucket_inference_config = replace(inference_config, batch_size=batch_size, padded_length=padded_length)
+
+            all_results = self.generate_tokens_many(
+                sequence_tokenized,
+                generation_config=generation_config,
+                inference_config=bucket_inference_config,
+                forward_pass_config=forward_pass_config,
+                keys=keys[np.array(sequence_ids)],
+            )
+
+            for idx, result in zip(sequence_ids, all_results, strict=True):
+                trimmed_ids = self._trim_at_eos(result.token_ids.tolist())
+                response = self.message_processor.parse_tokenized_response(trimmed_ids)
+                yield (idx, response)
+
     def stream_reply_text(
         self,
         messages: Iterable[Message],
-        sampling_policy: SamplingPolicy | None = None,
+        generation_config: GenerationConfig | None = None,
         max_output_length: int = 8192,
         forward_pass_config: ForwardPassConfig | None = None,
         *,
@@ -279,7 +562,7 @@ class LanguageModel(TextModel[LanguageModelConfig, Decoder]):
         token_ids = jnp.array(self.message_processor.tokenize_text(formatted_messages), dtype=jnp.int32)
         for token_id in self.stream_tokens(
             token_ids,
-            sampling_policy,
+            generation_config,
             max_output_length,
             forward_pass_config=forward_pass_config,
             key=key,
@@ -289,15 +572,17 @@ class LanguageModel(TextModel[LanguageModelConfig, Decoder]):
     def stream_tokens(
         self,
         prompt_token_ids: Int[Array, " prompt_tokens"],
-        sampling_policy: SamplingPolicy | None = None,
+        generation_config: GenerationConfig | None = None,
         max_output_length: int = 8192,
         eos_token_ids: Int[Array, " eos_tokens"] | None = None,
         forward_pass_config: ForwardPassConfig | None = None,
         *,
         key: PRNGKeyArray | None = None,
     ) -> Iterable[Int[Array, ""]]:
-        if sampling_policy is None:
-            sampling_policy = self.default_sampling_policy()
+        sampling_policy = self.default_sampling_policy()
+        if generation_config is not None:
+            sampling_policy = generation_config.default_policy()
+
         if eos_token_ids is None:
             eos_token_ids = jnp.array(self.stop_token_ids, dtype=jnp.int32)
 
@@ -309,8 +594,8 @@ class LanguageModel(TextModel[LanguageModelConfig, Decoder]):
 
         prefill_results = self._prefill(
             padded_token_ids[None, :],
-            jnp.array([input_length], dtype=jnp.int32),
             padded_input_length + max_output_length,
+            lengths_without_padding=jnp.array([input_length], dtype=jnp.int32),
             forward_pass_config=forward_pass_config,
         )
 
@@ -341,6 +626,7 @@ class LanguageModel(TextModel[LanguageModelConfig, Decoder]):
                 next_token_indices.reshape(1, 1),
                 state.state,
                 return_updated_state=True,
+                forward_pass_mode=ForwardPassMode.SINGLE_TOKEN,
                 forward_pass_config=forward_pass_config,
             )
             assert decoder_outputs.updated_state is not None, "updated_state should not be None"