lalamo 0.6.3__tar.gz → 0.6.4__tar.gz

{lalamo-0.6.3 → lalamo-0.6.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lalamo
-Version: 0.6.3
+Version: 0.6.4
 Summary: JAX library for optimization and export of models for use with the UZU inference engine.
 Requires-Python: <4,>=3.12
 Description-Content-Type: text/markdown

{lalamo-0.6.3 → lalamo-0.6.4}/lalamo/__init__.py

@@ -32,7 +32,7 @@ from lalamo.speculator import (
     SpeculatorTrainingEvent,
 )
 
-__version__ = "0.6.3"
+__version__ = "0.6.4"
 
 __all__ = [
     "AssistantMessage",

{lalamo-0.6.3 → lalamo-0.6.4}/lalamo/main.py

@@ -49,7 +49,10 @@ from lalamo.message_processor import UserMessage
 from lalamo.model_import import REPO_TO_MODEL, ModelSpec
 from lalamo.model_import.common import FileSpec
 from lalamo.models import ClassifierModelConfig, LanguageModelConfig
-from lalamo.speculator.estimator import get_default_device_memory
+from lalamo.speculator.estimator import (
+    get_default_device_bytes,
+    get_usable_memory_from_bytes,
+)
 from lalamo.speculator.ngram import NGramSpeculator
 from lalamo.speculator.utils import test_speculator
 
@@ -384,6 +387,7 @@ class CliTraceCallbacks(TraceCallbacks):
         self.stack.close()
         console.print(f"💾 Trace saved to [cyan]{self.output_path}[/cyan]")
 
+
 @app.command(help="Trace a model.")
 def trace(
     model_path: Annotated[
@@ -557,14 +561,24 @@ def estimate_batchsize(
     ] = None,
 ) -> None:
     if vram_gb is not None:
-        mem = vram_gb * 1024 * 1024 * 1024
-    elif (mem := get_default_device_memory()) is None:
+        # note that in practice GPUs use GiB in their docs, e.g. H100 actually has 85GB of memory
+        mem_bytes = vram_gb * 1000 * 1000 * 1000
+    elif (mem_bytes := get_default_device_bytes()) is None:
         err_console.print("Cannot get the default device's memory stats, use --vram-gb")
         raise Exit(1)
 
+    usable_mem = get_usable_memory_from_bytes(mem_bytes)
+
     callbacks_type = CliEstimateBatchsizeCallbacks
 
-    _estimate_batchsize(model_path, mem, max_input_length, max_output_length, num_logits_per_token, callbacks_type)
+    _estimate_batchsize(
+        model_path,
+        usable_mem,
+        max_input_length,
+        max_output_length,
+        num_logits_per_token,
+        callbacks_type,
+    )
 
 
 @dataclass

{lalamo-0.6.3 → lalamo-0.6.4}/lalamo/speculator/estimator.py

@@ -10,7 +10,7 @@ import jax.numpy as jnp
 from lalamo.models import LanguageModel
 
 
-def get_default_device_memory() -> int | None:
+def get_default_device_bytes() -> int | None:
     dynamic_allocate = False
 
     preallocate = os.getenv("XLA_PYTHON_CLIENT_PREALLOCATE", "")
@@ -26,20 +26,22 @@ def get_default_device_memory() -> int | None:
     if memory_stats is None or "bytes_limit" not in memory_stats:
         return None
 
-    # discount by 0.98 because if you try to allocate exactly bytes_limit
-    # jax will _try_ to allocate a bit more, and then throws an OOM
-    bytes_limit = memory_stats["bytes_limit"] * 0.98
-
     mem_fraction_raw = os.getenv("XLA_PYTHON_CLIENT_MEM_FRACTION", "")
     try:
         mem_fraction = float(mem_fraction_raw)
     except ValueError:
         mem_fraction = 0.75  # jax default https://docs.jax.dev/en/latest/gpu_memory_allocation.html
 
-    # JAX usually can't allocate more than 98%-ish percent memory; idk why
-    # Besides we use _some_ autotuning during runtime, so we add 0.96 safety margin here
-    bytes_limit = int(max(bytes_limit, bytes_limit / min(mem_fraction, 1.0)) * 0.96)
-    return bytes_limit
+    # 500mb is seemingly the usually observed overhead; this tries to match the actual capacity of the gpu
+    # so it should correspond to something you'd see in nvidia-smi
+    memory_limit = memory_stats["bytes_limit"] / min(mem_fraction, 1.0) + (500 * 1000 * 1000)
+
+    return get_usable_memory_from_bytes(memory_limit)
+
+
+def get_usable_memory_from_bytes(limit_bytes: int) -> int:
+    # JAX allocates a bit more than it needs, so we discount it by some safety factor
+    return int(limit_bytes * 0.95)
 
 
 def estimate_memory_from_batchsize(

{lalamo-0.6.3 → lalamo-0.6.4}/lalamo.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lalamo
-Version: 0.6.3
+Version: 0.6.4
 Summary: JAX library for optimization and export of models for use with the UZU inference engine.
 Requires-Python: <4,>=3.12
 Description-Content-Type: text/markdown