llmcompressor.pipelines.cache

def append(self, values: dict[str, Any]):
    """
    Append new values to the cache. The new values will be assigned the next
    available batch index

    :param values: dictionary mapping keys to values used for update
    """
    batch_index = len(self.batch_intermediates)
    self.batch_intermediates.append({})
    self.update(batch_index, values)

def delete(self, batch_index: int, consumed_names: list[str] | None = None):
    """
    Delete values from the cache

    :param batch_index: index of batch whose values will be deleted
    :param consumed_names: list of keys whose values will be deleted, defaults to
        removing all keys
    """
    intermediates = self.batch_intermediates[batch_index]

    if consumed_names is None:
        consumed_names = list(intermediates.keys())

    for name in consumed_names:
        del intermediates[name]

@classmethod
def empty(cls, num_batches: int, offload_device: torch.device):
    """
    Construct an empty cache

    :param num_batches: the expected number of batches to be stored
    :param offload_device: device to offload values to
    """
    batch_intermediates = [{} for _ in range(num_batches)]
    return cls(batch_intermediates, offload_device)

def fetch(
    self, batch_index: int, input_names: list[str] | None = None
) -> dict[str, Any]:
    """
    Fetch values belonging to a batch

    :param batch_index: index of batch whose values are being fetched
    :param input_names: list of keys whose values are being fetched
    :return: dictionary mapping keys to onloaded values
    """
    intermediates = self.batch_intermediates[batch_index]

    return {
        key: self._onload_value(subgraph_input)
        for key, subgraph_input in intermediates.items()
        if input_names is None or key in input_names
    }

@classmethod
def from_dataloader(
    cls,
    dataloader: torch.utils.data.DataLoader,
    model_device: torch.device = torch.device("cpu"),
    offload_device: torch.device | None = torch.device("cpu"),
):
    """
    Initialize a cache with data from the provided dataloader

    This method iterates through all batches in the dataloader and offloads
    them to the specified device. For faster cache preparation, consider:
    - Increasing batch_size to reduce the number of iterations
    - Using num_workers > 0 in the DataLoader for parallel loading (e.g. the
      calibration DataLoader from format_calibration_data uses
      dataloader_num_workers; when > 0, pin_memory and prefetch_factor are
      also set where applicable, which speeds both cache build and calibration)
    - Ensuring data preprocessing is done before creating the dataloader

    :param dataloader: dataloader which generates values to be cached
    :param model_device: device which values will be onloaded to when fetched
    :param offload_device: device to offload values to
    """
    batch_intermediates = [
        {
            key: cls._offload_value(value, offload_device, model_device)
            for key, value in batch.items()
        }
        for batch in tqdm(dataloader, desc="Preparing cache")
    ]

    return cls(batch_intermediates, offload_device)

def iter_prefetch(
    self, input_names: list[str] | None = None
) -> Generator[Any, None, None]:
    """
    Iterate over batches with the next batch prefetched in a background thread.
    Overlaps onload from offload_device with consumption of the current batch,
    which can reduce wall-clock time when offloading to CPU.

    When CUDA is available, uses non_blocking transfers (requires pinned CPU
    tensors, set up by _offload_value) and synchronises via CUDA events so the
    main stream waits for each H2D copy before running GPU kernels on the data.

    Yields the same fetched batch dicts as :meth:`iter`; only the timing
    of onloads differs.
    """
    num_batches = len(self.batch_intermediates)
    if num_batches == 0:
        return

    # Create a dedicated CUDA stream for H2D transfers so they run on a
    # separate stream from the main thread's compute stream. Without this,
    # both threads default to the null stream (stream 0) which serializes
    # all operations and prevents any overlap.
    h2d_stream = torch.cuda.Stream() if torch.cuda.is_available() else None

    def _fetch_and_record(batch_index):
        event = None
        if h2d_stream is not None:
            with torch.cuda.stream(h2d_stream):
                data = self.fetch(batch_index, input_names)
            event = torch.cuda.Event()
            event.record(h2d_stream)
        else:
            data = self.fetch(batch_index, input_names)
        return data, event

    with ThreadPoolExecutor(max_workers=1) as executor:
        future = None
        for batch_index in range(num_batches):
            if future is not None:
                current, event = future.result()
            else:
                current, event = _fetch_and_record(batch_index)
            if batch_index + 1 < num_batches:
                future = executor.submit(_fetch_and_record, batch_index + 1)
            else:
                future = None
            # Make the main CUDA stream wait for the background H2D copy
            # before any GPU kernel consumes the prefetched tensors
            if event is not None:
                torch.cuda.current_stream().wait_event(event)
            yield current

def size(self) -> dict[torch.device, int]:
    """
    Returns the memory used by cached values, keyed by device, in bytes

    :return: dictionary mapping torch device to number of bytes in cache
    """
    sizes = defaultdict(lambda: 0)
    memo = set()

    def _size_helper(intermediate: IntermediateValue) -> int:
        value = intermediate.value

        match value:
            case torch.Tensor():
                if value not in memo:
                    sizes[value.device] += value.nbytes
                memo.add(value)
            case list() | tuple():
                for v in value:
                    _size_helper(v)
            case dict():
                for v in value.values():
                    _size_helper(v)
            case _ if is_dataclass(value):
                for field in fields(value):
                    _size_helper(getattr(value, field.name))
            case _:
                # this handles primitive values that don't match any other cases
                sizes[torch.device("cpu")] += sys.getsizeof(value, 0)

    for intermediates in self.batch_intermediates:
        for value in intermediates.values():
            _size_helper(value)

    return dict(sizes)

def update(self, batch_index: int, values: dict[str, Any]):
    """
    Update/put values belonging to a batch

    :param batch_index: index of batch whose values will be updated
    :param values: dictionary mapping keys to values used for update
    """
    device = self.offload_device
    intermediates = {k: self._offload_value(v, device) for k, v in values.items()}
    self.batch_intermediates[batch_index].update(intermediates)