llm-layer-collector

llm_layer_collector is a Python package for work with HuggingFace models at the layer level. The package loads the embedding, the decoder layers, the norm, and the head as separate PyTorch modules. A program can thus load only the parts of a model that it needs. This is useful for research, for development, and for machines with a small quantity of memory.

Language Pipes uses this package for distributed inference. The package has no dependency on Language Pipes and can operate alone.

What the package does

• It reads the HuggingFace file format to find the correct parts of a checkpoint.
• It loads each part as a standard PyTorch module.
• It runs the computation for each part with the transformers library.
• It selects the correct computation for the architecture of the loaded model.

Installation

pip install llm-layer-collector

Public interface

The package makes two classes available at the top level:

from llm_layer_collector import LlmLayerCollector, StaticAutoModel

Four more classes come back from the methods of these two classes. A program does not construct these classes directly, but a program does use their methods:

LlmLayerCollector

The LlmLayerCollector class is the central interface to the package. The constructor reads config.json from the model directory. Then the constructor reads the cache file, or builds a new cache file. The cache file holds a map from each tensor name to the shard file that contains the tensor.

Constructor

LlmLayerCollector(
    model_dir,
    cache_file,
    shard_pattern=r"model-(\d+)-of-(\d+).safetensors",
    layer_prefix="model.layers.",
    input_embedding_layer_name="model.embed_tokens.weight",
    norm_layer_name="model.norm.weight",
    lm_head_name="lm_head.weight",
    dtype=torch.bfloat16,
    device=torch.device("cpu"),
    load_in_8bit=False,
)

The constructor corrects three of these values automatically:

• If the constructor builds a new cache file, it reads the true layer_prefix, input_embedding_layer_name, and norm_layer_name from the tensor names in the shards. Thus a program does not usually set these three parameters.
• If lm_head_name is not in the cache, but a different name ends with lm_head.weight, the constructor uses that name. Multimodal checkpoints nest the head under a different prefix.
• If load_in_8bit is True, the constructor sets dtype to torch.float16. The bitsandbytes kernels compute in fp16, so the other parts must have the same type.

CAUTION: Delete the cache file after you change or replace the files in the model directory. A stale cache file points to shard files that are no longer correct, and the load then fails or gives incorrect weights.

NOTE: The load_in_8bit option needs the bitsandbytes package. The load of the first layer set raises an ImportError if the package is not installed.

Attributes

The constructor sets these public attributes:

Each load method has an optional device parameter. If the value is None, the method uses the device attribute.

load_input_embedding(device=None)

Returns: torch.nn.Embedding

Loads the weight of the input embedding and gives back an embedding module. For the Gemma3 and Gemma4 architectures, the method gives back the scaled embedding class of that architecture. For all other architectures, the method gives back a standard torch.nn.Embedding.

embedding = collector.load_input_embedding()

load_norm(device=None)

Returns: AutoRMSNorm

Loads the weight of the final norm and gives back an AutoRMSNorm. The AutoRMSNorm object contains the RMS norm class of the applicable architecture. Call the object directly to apply the norm to a hidden state:

norm = collector.load_norm()
normed_state = norm(state.state)

load_head(device=None)

Returns: torch.nn.Linear

Loads the weight of the head and gives back a linear module without a bias. If the checkpoint has no separate head tensor, the method uses the weight of the input embedding. Models with tied weights keep the head and the embedding in one tensor.

head = collector.load_head()

load_layer_set(start_layer, end_layer, device=None)

Returns: List[AutoDecoderLayer]

Loads a continuous set of decoder layers. The method loads the layers in groups of three, and calls the garbage collector at the end. This procedure keeps the peak memory low for large models.

CAUTION: The end_layer index is inclusive. To load all layers of a model, give collector.num_layers - 1 as the value. A value of collector.num_layers raises an exception, because there is no data for that layer.

# All layers of the model
layers = collector.load_layer_set(0, collector.num_layers - 1)

# Only layers 4 to 8 (five layers)
layers = collector.load_layer_set(4, 8)

The method also converts the quantized weights of the checkpoint:

• It applies the fp8 scales to their weights.
• It unpacks the mxfp4 expert weights of the MoE models.
• It fuses the per-expert weights into the stacked tensors that the transformers MoE modules use.

load_per_layer_embedder(device=None)

Returns: Optional[Gemma4PerLayerEmbedder]

Loads the three Per-Layer Embedding (PLE) weights of Gemma4. The method gives back None for each model that does not use PLE. Give the result to StaticAutoModel.compute_embedding() as the per_layer_embedder parameter.

CAUTION: Call this method only on the node that holds the embedding and the head. The embed_tokens_per_layer tensor is the largest single tensor in the checkpoint, and a load on a layer node can fill the memory of that node.

per_layer_embedder = collector.load_per_layer_embedder()

StaticAutoModel

The StaticAutoModel class has three static methods. Each method sends the computation to the implementation for the architecture of the loaded model. The class holds no state, so a program does not construct it.

compute_embedding(...)

StaticAutoModel.compute_embedding(
    prompt_tokens,
    chunk_size,
    input_embedder,
    input_ids,
    config,
    cache,
    per_layer_embedder=None,
)

Returns: LLmComputationState

Embeds the next tokens and prepares the data that the decoder layers need. The method selects the tokens with the cache: it starts at the number of tokens that are already in the cache. If prompt tokens remain, the method takes a maximum of chunk_size tokens. If no prompt tokens remain, the method takes one token. The one-token path is the decode step.

The method then computes the causal mask and the rotary position embeddings for the architecture. All results go into the LLmComputationState object.

NOTE: For a prompt with no chunks, set chunk_size to the value of prompt_tokens.

compute_layer(layer, config, state, cache)

Returns: torch.Tensor

Runs the hidden state through one decoder layer and gives back the new hidden state. The method does not change the state object, so the caller must write the result to state.state before the next layer.

for layer in layers:
    state.state = StaticAutoModel.compute_layer(layer, collector.config, state, cache)

NOTE: The method gives back an empty tensor if the architecture of the layer is not supported.

compute_head(head, state, device, top_k=1, top_p=1, min_p=0, temperature=1)

Returns: int — the id of the next token.

Applies the head projection to the last position of the hidden state, and then selects the next token. Apply the final norm to the hidden state before you call this method.

The method uses one of two paths:

1. If temperature is 0, the method selects the token with the highest logit. This path is greedy decoding, and it uses no filter.
2. If temperature is not 0, the method divides the logits by temperature. Then the method applies the min_p, top_p, and top_k filters in that sequence. At the end, the method samples one token from the result.

next_token = StaticAutoModel.compute_head(head, norm(state.state), device="cuda", top_k=1)

LLmComputationState

The LLmComputationState dataclass holds the data that moves between the embedding, the layers, and the head. compute_embedding() constructs the object, and compute_layer() reads the object.

Full example

This example loads all parts of a model and predicts one token.

from llm_layer_collector import LlmLayerCollector, StaticAutoModel
from transformers import AutoTokenizer
from transformers.cache_utils import DynamicCache
import torch

# 1. Construct the collector.
collector = LlmLayerCollector(
    model_dir="/path/to/model",
    cache_file="cache.json",
    device=torch.device("cuda"),
    dtype=torch.bfloat16
)

# 2. Tokenize the prompt.
tokenizer = AutoTokenizer.from_pretrained("/path/to/model")
input_text = "The quick brown fox"
input_ids = tokenizer(input_text, return_tensors='pt')['input_ids']

# 3. Load the model parts.
embedding = collector.load_input_embedding()
norm = collector.load_norm()
head = collector.load_head()
layers = collector.load_layer_set(0, collector.num_layers - 1)  # end layer is inclusive

# 4. Compute the embedding.
cache = DynamicCache()
prompt_tokens = input_ids.shape[1]
state = StaticAutoModel.compute_embedding(
    prompt_tokens=prompt_tokens,
    chunk_size=prompt_tokens,
    input_embedder=embedding,
    input_ids=input_ids,
    config=collector.config,
    cache=cache,
)

# 5. Run the state through each layer.
for layer in layers:
    state.state = StaticAutoModel.compute_layer(layer, collector.config, state, cache)

# 6. Apply the norm and the head to get the next token.
next_token = StaticAutoModel.compute_head(head, norm(state.state), device="cuda", top_k=1)
print(tokenizer.decode(next_token))

To do the same task step by step:

1. Construct an LlmLayerCollector with the model directory and a cache file path.
2. Tokenize the prompt with the tokenizer of the model.
3. Load the embedding, the norm, the head, and the layer set.
4. Construct a DynamicCache.
5. Call StaticAutoModel.compute_embedding() to get the first LLmComputationState.
6. Call StaticAutoModel.compute_layer() for each layer. Write each result to state.state.
7. Apply the norm to state.state.
8. Call StaticAutoModel.compute_head() with the normed state to get the next token id.
9. To get more tokens, add the new token id to input_ids and do steps 5 to 8 again. Use the same cache object.

Supported architectures

The three StaticAutoModel methods and the layer classes support these model types:

For a multimodal checkpoint, the collector reads the text configuration. Thus a model with the type gemma3 becomes gemma3_text.

To add a new architecture, refer to the modules in src/llm_layer_collector/modeling/.