code-loader 1.0.202.dev1__tar.gz → 1.0.202.dev2__tar.gz

{code_loader-1.0.202.dev1 → code_loader-1.0.202.dev2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: code-loader
-Version: 1.0.202.dev1
+Version: 1.0.202.dev2
 Summary: 
 License: MIT
 License-File: LICENSE

{code_loader-1.0.202.dev1 → code_loader-1.0.202.dev2}/code_loader/contract/datasetclasses.py

@@ -318,6 +318,11 @@ class CustomLatentSpaceHandler:
 # dominates the model inputs). 'mean': every latent space is the elementwise mean over all steps.
 AUTOREGRESSIVE_LATENT_SPACE_AGGREGATIONS = ('last_step', 'mean')
 
+# Default axis of a KV-cache tensor that indexes token positions: (batch, heads, position,
+# head_dim), the layout HF-style exports produce. Overridable per integration via
+# tensorleap_autoregressive_step(cache_seq_axis=...).
+DEFAULT_CACHE_SEQ_AXIS = 2
+
 # Reserved argument names of autoregressive metrics/losses/visualizers, fed implicitly by the
 # platform from the finished chain (per-chain dicts, no batch axis). Any other argument is wired
 # to a ground-truth encoder through the integration test.
@@ -337,6 +342,12 @@ class AutoregressiveStepHandler:
     # are never passed to the hook and never reported — so the hook must NOT return them. Empty
     # when the model has no cache, in which case the chain re-feeds its full prompt each step.
     cache_tensors: Optional[Dict[str, str]] = None
+    # Which axis of a cache tensor indexes token positions, i.e. where its capacity is read from.
+    # 2 matches the (batch, heads, position, head_dim) layout every HF-style export uses. Declared
+    # rather than inferred because the platform cannot tell positions from heads or head_dim — in
+    # a (1, 3, 64, 64) tensor two of the dims are 64 — and reading the wrong axis silently caps
+    # chains at the number of heads instead of the cache capacity.
+    cache_seq_axis: int = DEFAULT_CACHE_SEQ_AXIS
 
 
 # Per-chain, unbatched callables: called once per finished chain with the final step's tensors —

{code_loader-1.0.202.dev1 → code_loader-1.0.202.dev2}/code_loader/inner_leap_binder/leapbinder.py

@@ -18,6 +18,7 @@ from code_loader.contract.datasetclasses import SectionCallableInterface, InputH
     ElementInstanceMasksHandler, InstanceCallableInterface, CustomLatentSpaceHandler, InstanceMetricHandler, \
     SimulationHandler, _simulation_context, AutoregressiveStepHandler, AutoregressiveStepCallableInterface, \
     AUTOREGRESSIVE_LATENT_SPACE_AGGREGATIONS, AUTOREGRESSIVE_IMPLICIT_ARG_NAMES, \
+    DEFAULT_CACHE_SEQ_AXIS, \
     AutoregressiveMetricHandler, AutoregressiveLossHandler, AutoregressiveVisualizerHandler
 from code_loader.contract.enums import LeapDataType, DataStateEnum, DataStateType, MetricDirection, DatasetMetadataType, \
     TestingSectionEnum
@@ -622,9 +623,25 @@ class LeapBinder:
             seen_outputs[output_name] = input_name
         return dict(cache_tensors)
 
+    @staticmethod
+    def _validated_cache_seq_axis(cache_seq_axis: int,
+                                  cache_tensors: Optional[Dict[str, str]]) -> int:
+        """The axis must be a non-negative int, and declaring one without a cache is a mistake
+        worth naming — it reads as if it had an effect."""
+        if not isinstance(cache_seq_axis, int) or isinstance(cache_seq_axis, bool) \
+                or cache_seq_axis < 0:
+            raise Exception(f'tensorleap_autoregressive_step: cache_seq_axis must be a '
+                            f'non-negative int axis index, got {cache_seq_axis!r}.')
+        if not cache_tensors and cache_seq_axis != DEFAULT_CACHE_SEQ_AXIS:
+            raise Exception('tensorleap_autoregressive_step: cache_seq_axis was declared without '
+                            'cache_tensors. The axis describes the declared cache tensors, so it '
+                            'has no meaning on its own.')
+        return cache_seq_axis
+
     def set_autoregressive_step(self, function: AutoregressiveStepCallableInterface,
                                 latent_space_aggregation: str = 'last_step',
-                                cache_tensors: Optional[Dict[str, str]] = None) -> None:
+                                cache_tensors: Optional[Dict[str, str]] = None,
+                                cache_seq_axis: int = DEFAULT_CACHE_SEQ_AXIS) -> None:
         """
         Set the autoregressive step hook — the feedback function that drives a chain:
         it supplies the model's initial inputs on its first call (prev_inputs=None, prev_outputs=None)
@@ -634,6 +651,9 @@ class LeapBinder:
         its steps (see AUTOREGRESSIVE_LATENT_SPACE_AGGREGATIONS).
         cache_tensors declares a KV-cache as {model_input_name: model_output_name}; the platform
         threads those tensors itself and the hook neither receives nor returns them.
+        cache_seq_axis declares which axis of those tensors indexes token positions, i.e. where
+        the platform reads the cache capacity from (default 2, for (batch, heads, position,
+        head_dim)).
         """
         if self.setup_container.autoregressive_step is not None:
             raise Exception('tensorleap_autoregressive_step is already defined. '
@@ -643,9 +663,10 @@ class LeapBinder:
                             f'{latent_space_aggregation!r}. Supported values: '
                             f'{", ".join(AUTOREGRESSIVE_LATENT_SPACE_AGGREGATIONS)}.')
         cache_tensors = self._validated_cache_tensors(cache_tensors)
+        cache_seq_axis = self._validated_cache_seq_axis(cache_seq_axis, cache_tensors)
         self.setup_container.autoregressive_step = AutoregressiveStepHandler(
             function, latent_space_aggregation=latent_space_aggregation,
-            cache_tensors=cache_tensors)
+            cache_tensors=cache_tensors, cache_seq_axis=cache_seq_axis)
 
         # Builtin chain metadata, declared at parse time so it survives the reporter's
         # metadata type mapping; the placeholder values are overwritten by the engine when a

{code_loader-1.0.202.dev1 → code_loader-1.0.202.dev2}/code_loader/inner_leap_binder/leapbinder_decorators.py

@@ -24,7 +24,7 @@ from code_loader.contract.datasetclasses import CustomCallableInterfaceMultiArgs
     CustomMultipleReturnCallableInterfaceMultiArgs, ConfusionMatrixCallableInterfaceMultiArgs, CustomCallableInterface, \
     VisualizerCallableInterface, MetadataSectionCallableInterface, PreprocessResponse, SectionCallableInterface, \
     ConfusionMatrixElement, SamplePreprocessResponse, PredictionTypeHandler, InstanceCallableInterface, ElementInstance, \
-    InstanceLengthCallableInterface, AutoregressiveStepCallableInterface
+    InstanceLengthCallableInterface, AutoregressiveStepCallableInterface, DEFAULT_CACHE_SEQ_AXIS
 from code_loader.contract.enums import MetricDirection, LeapDataType, DatasetMetadataType, DataStateType
 from code_loader import leap_binder, LeapLoader
 from code_loader.contract.mapping import NodeMapping, NodeMappingType, NodeConnection
@@ -2084,7 +2084,8 @@ class _ModelLoopModelProxy:
 
 
 def tensorleap_autoregressive_step(latent_space_aggregation: str = 'last_step',
-                                   cache_tensors: Optional[Dict[str, str]] = None):
+                                   cache_tensors: Optional[Dict[str, str]] = None,
+                                   cache_seq_axis: int = DEFAULT_CACHE_SEQ_AXIS):
     """The feedback hook that drives an autoregressive chain.
 
     Signature of the decorated function:
@@ -2124,6 +2125,13 @@ def tensorleap_autoregressive_step(latent_space_aggregation: str = 'last_step',
         })
 
     Omit it for a model with no cache — the chain then re-feeds its whole prompt every step.
+
+    cache_seq_axis declares which axis of those cache tensors indexes token positions — the axis
+    whose length IS the cache capacity, which the platform reads to cap chain length. The default
+    2 matches (batch, heads, position, head_dim), the layout HF-style exports produce. Set it
+    when yours differs, e.g. 1 for (batch, position, heads, head_dim). It cannot be inferred: in
+    a (1, 3, 64, 64) tensor two dims are 64, so a wrong guess would silently cap chains at the
+    number of heads and truncate every chain a few steps in.
     """
     assert isinstance(latent_space_aggregation, str), \
         ('tensorleap_autoregressive_step must be called with parentheses: '
@@ -2272,7 +2280,8 @@ def tensorleap_autoregressive_step(latent_space_aggregation: str = 'last_step',
 
         leap_binder.set_autoregressive_step(inner_without_validate,
                                             latent_space_aggregation=latent_space_aggregation,
-                                            cache_tensors=cache_tensors)
+                                            cache_tensors=cache_tensors,
+                                            cache_seq_axis=cache_seq_axis)
 
         def inner(*args, **kwargs):
             if not _call_from_tl_platform:

{code_loader-1.0.202.dev1 → code_loader-1.0.202.dev2}/code_loader/leaploader.py

@@ -16,7 +16,8 @@ from code_loader.contract.datasetclasses import DatasetSample, DatasetBaseHandle
     PreprocessResponse, VisualizerHandler, LeapData, \
     PredictionTypeHandler, MetadataHandler, CustomLayerHandler, MetricHandler, VisualizerHandlerData, MetricHandlerData, \
     MetricCallableReturnType, CustomLossHandlerData, CustomLossHandler, RawInputsForHeatmap, SamplePreprocessResponse, \
-    ElementInstance, custom_latent_space_attribute, DatasetIntegrationSetup, InstanceMetricHandler, _simulation_context
+    ElementInstance, custom_latent_space_attribute, DatasetIntegrationSetup, InstanceMetricHandler, \
+    _simulation_context, DEFAULT_CACHE_SEQ_AXIS
 from code_loader.contract.enums import DataStateEnum, TestingSectionEnum, DataStateType, DatasetMetadataType
 from code_loader.contract.exceptions import DatasetScriptException
 from code_loader.contract.responsedataclasses import DatasetIntegParseResult, DatasetTestResultPayload, \
@@ -1396,6 +1397,20 @@ class LeapLoader(LeapLoaderBase):
             return {}
         return dict(handler.cache_tensors)
 
+    @lru_cache()
+    def get_autoregressive_cache_seq_axis(self) -> int:
+        """Which axis of the declared cache tensors indexes token positions.
+
+        The platform reads the cache capacity off this axis to cap chain length. Defaults to 2
+        ((batch, heads, position, head_dim)) — including when nothing is declared, where it is
+        simply unused.
+        """
+        self.exec_script()
+        handler = global_leap_binder.setup_container.autoregressive_step
+        if handler is None:
+            return DEFAULT_CACHE_SEQ_AXIS
+        return int(handler.cache_seq_axis)
+
     def run_autoregressive_step(self, sample_id: Union[int, str],
                                 prev_inputs: Optional[Dict[str, npt.NDArray[np.float32]]],
                                 prev_outputs: Optional[Dict[str, npt.NDArray[np.float32]]],

{code_loader-1.0.202.dev1 → code_loader-1.0.202.dev2}/code_loader/leaploaderbase.py

@@ -181,6 +181,11 @@ class LeapLoaderBase:
         raise NotImplementedError(f'{type(self).__name__} does not implement '
                                   'get_autoregressive_cache_tensors.')
 
+    @abstractmethod
+    def get_autoregressive_cache_seq_axis(self) -> int:
+        raise NotImplementedError(f'{type(self).__name__} does not implement '
+                                  'get_autoregressive_cache_seq_axis.')
+
     @abstractmethod
     def get_autoregressive_decorator_names(self) -> Dict[str, List[str]]:
         raise NotImplementedError(f'{type(self).__name__} does not implement '

{code_loader-1.0.202.dev1 → code_loader-1.0.202.dev2}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "code-loader"
-version = "1.0.202.dev1"
+version = "1.0.202.dev2"
 description = ""
 authors = ["dorhar <doron.harnoy@tensorleap.ai>"]
 license = "MIT"