Implementation:Predibase Lorax Base Model

Knowledge Sources	Predibase_Lorax
Domains	Model_Architecture, Inference
Last Updated	2026-02-08 00:00 GMT

Overview

Defines the abstract base class Model that establishes the interface for all model implementations in the LoRax server, including adapter lifecycle management, token decoding, generation, and server metadata reporting.

Description

This module contains the foundational Model ABC that all concrete model wrappers (CausalLM, Seq2SeqLM, FlashCausalLM, FlashBert, VlmCausalLM, etc.) must inherit from. It provides the common infrastructure for adapter loading/offloading, distributed inference, and profiling.

Key class:

Model (extends ABC) - The abstract base class with the following responsibilities:

Initialization:

- Accepts model, tokenizer, device, dtype, rank, world_size, and adapter configuration.
- Sets model to eval mode.
- Initializes TokenizerManager for per-adapter tokenizer support.
- Builds the set of all special token IDs (including added special tokens).
- Initializes adapter weight tracking: layer_to_adapter_weights (defaultdict of LayerAdapterWeights), target_to_layer mapping, loaded_adapters set, and layer_to_lora_weights for Triton kernel compatibility.
- Configures chunked prefill support (requires FlashInfer backend, incompatible with speculation).
- Detects position_ids support via model forward signature inspection.
- Optionally downloads and loads an initial adapter if specified.
- Sets up optional PyTorch profiler for CPU/CUDA activity tracing.

Abstract interface:

- batch_type (property, abstract) - Returns the batch class type for this model.
- generate_token (abstract) - Performs one autoregressive generation step.

Server metadata:

- info (property) - Returns an InfoResponse protobuf with model capabilities: padding requirements, dtype, device, window size, block size, speculative token count, preloaded adapters, and supported operations (generation, embeddings, classification).
- block_size, sliding_window_blocks, supports_embeddings, supports_classification, supports_text_generation, supports_cuda_graph_compilation, requires_block_allocator - Capability properties with sensible defaults.

Adapter management:

- load_adapter - Downloads and loads LoRA/Medusa adapter weights, distributing them across layer_to_adapter_weights by layer name. Supports both static and dynamic loading.
- offload_adapter - Removes adapter weights from GPU, excluding preloaded adapters.
- register_preloaded_adapters - Registers adapters loaded at startup, disabling dynamic loading. For Triton kernels, stacks LoRA weights into contiguous tensors with rank padding for efficient batched computation.
- shard_lora_weights - Shards LoRA A and B weight matrices across the process group for tensor parallelism.
- adapter_target_to_layer, adapter_layers, traced_adapter_layers, default_traced_adapter_layers, get_num_layers_for_type, is_row_parallel - Adapter topology methods (defaults return empty; overridden by subclasses).

Utilities:

- decode_token - Incrementally decodes tokens using prefix/read offset tracking to handle tokenizers with cleanup algorithms.
- warmup - Runs a generation step for GPU memory warmup.
- check_initialized - Verifies no parameters remain on the meta device after loading.

Usage

Model is never instantiated directly. All model types in the LoRax server inherit from it: CausalLM, Seq2SeqLM, FlashCausalLM, FlashBert, FlashXlmRoberta, and their subclasses. The base class provides the adapter management infrastructure and server integration interface that all models share.

Code Reference

Source Location

Repository: Predibase_Lorax
File: server/lorax_server/models/model.py
Lines: 1-450

Signature

class Model(ABC):
    def __init__(
        self,
        model_id: str,
        model: torch.nn.Module,
        tokenizer: PreTrainedTokenizerBase,
        requires_padding: bool,
        dtype: torch.dtype,
        device: torch.device,
        rank: int = 0,
        world_size: int = 1,
        sliding_window: Optional[int] = None,
        adapter_id: str = BASE_MODEL_ADAPTER_ID,
        adapter_source: str = HUB,
        dynamic_adapter_loading_enabled: bool = True,
        trust_remote_code: bool = False,
        processor=None,
        supports_chunking: bool = False,
    ):
        ...

    @property
    def info(self) -> InfoResponse:
        ...
    @property
    @abstractmethod
    def batch_type(self) -> Type[B]:
        ...
    @abstractmethod
    def generate_token(self, batch: B) -> Tuple[List[GeneratedText], Optional[B]]:
        ...
    def warmup(self, batch: B, max_new_tokens: int) -> Optional[int]:
        ...
    def decode_token(self, all_input_ids, prefix_offset=0, read_offset=0, skip_special_tokens=False) -> Tuple[str, int, int]:
        ...
    def check_initialized(self):
        ...
    @property
    def supports_adapter_loading(self) -> bool:
        ...
    def adapter_target_to_layer(self) -> Dict[str, Tuple[str, torch.Tensor]]:
        ...
    @property
    def adapter_layers(self) -> List[str]:
        ...
    def load_adapter(self, adapter_parameters, adapter_source, adapter_index, api_token, dynamic=True):
        ...
    def offload_adapter(self, adapter_parameters, adapter_source, adapter_index) -> bool:
        ...
    def register_preloaded_adapters(self, preloaded_adapters, adapter_memory_fractions):
        ...
    def shard_lora_weights(self, weights_a, weights_b, layer_type) -> Tuple[List[torch.Tensor], List[torch.Tensor]]:
        ...

Import

from lorax_server.models.model import Model

I/O Contract

Inputs

Name	Type	Required	Description
model_id	str	Yes	HuggingFace model identifier
model	torch.nn.Module	Yes	The underlying PyTorch model instance
tokenizer	PreTrainedTokenizerBase	Yes	HuggingFace tokenizer instance
requires_padding	bool	Yes	Whether the model requires padded batches
dtype	torch.dtype	Yes	Model data type for inference
device	torch.device	Yes	Device to run inference on
rank	int	No	Process rank for distributed inference (default: 0)
world_size	int	No	Total number of processes (default: 1)
sliding_window	Optional[int]	No	Sliding window attention size
adapter_id	str	No	Initial adapter to load (default: BASE_MODEL_ADAPTER_ID)
adapter_source	str	No	Source for adapter weights (default: HUB)
dynamic_adapter_loading_enabled	bool	No	Whether dynamic adapter swapping is enabled (default: True)
trust_remote_code	bool	No	Whether to trust remote code (default: False)
supports_chunking	bool	No	Whether the model supports chunked prefill (default: False)

Outputs

Name	Type	Description
info	InfoResponse	Server metadata about model capabilities
generations	List[GeneratedText]	Generated tokens (from generate_token, defined by subclasses)
next_batch	Optional[B]	Updated batch for next step (from generate_token, defined by subclasses)

Usage Examples

# Internal LoRax server usage - Model is abstract, use subclasses
from lorax_server.models.model import Model

# All concrete model classes inherit from Model:
# class CausalLM(Model): ...
# class Seq2SeqLM(Model): ...
# class FlashCausalLM(Model): ...
# class FlashBert(Model): ...

Related Pages

Environment:Predibase_Lorax_CUDA_GPU_Runtime

Page Connections

Double-click a node to navigate. Hold to expand connections.

Principle

Implementation

Heuristic

Environment