The Transformers loading pipeline

America Transformers is all about speed. Hot, nasty, badass speed.

• Ricky Bobby, Talladega Nights

A collection of working notes on how Transformers builds large models without doubling peak memory and speeds their loading from disk, from meta-device skeletons to overlapping disk reads and GPU transfers.

#PyTorch meta device

A basic loading flow materializes a tensor and fills it with random values from the model config, which consumes memory. A 70B fp16 model requires ~140GB of memory. Next, the real model parameters are read into a second copy in CPU memory, requiring another ~140GB of memory. Finally, the randomly initialized parameters are overwritten by the real parameters.

During this process, peak memory is 2x the model size, because there are two copies of the model in memory. Moreover, the work required to load the randomly initialized parameters adds wasteful overhead because they're eventually discarded.

init_contexts.extend([torch.device("meta"), init.meta_device_safe_creation_ops()])

  # get_init_context returns torch.device("meta")
  model_init_context = cls.get_init_context(dtype, is_quantized, _is_ds_init_called, allow_all_kernels)

  config = copy.deepcopy(config)
  with ContextManagers(model_init_context):
      # every parameter is placed on the meta device
      model = cls(config, *model_args, **model_kwargs)
      patch_output_recorders(model)

The PyTorch meta device lets Transformers build a model of any size regardless of device memory.

Building the model on the meta device doesn't consume any memory. The meta device only holds the metadata for the parameters the model expects, like shape and dtype, creating a skeleton of the model from the config.

From those expected shapes and dtypes, Transformers computes how many bytes each parameter slot will occupy once it's loaded. If a device map is being inferred for automatic placement, sharding, CPU/disk offload, or multi-device loading, those planned sizes help assign parameters to devices with enough available memory.

{
      "model.layers.0.self_attn.q_proj.weight": "0",
      "model.layers.31.mlp.up_proj.weight":     "1",
      "model.layers.45.mlp.down_proj.weight":   "cpu",
      "model.layers.60.mlp.down_proj.weight":   "disk",
      ...
  }

When a tensor is materialized, the safetensors slice is read into a temporary CPU allocation and then moved or cast to the parameter's target device. That temporary CPU allocation is released after the materialized tensor is returned.

At any given moment, only a small number of parameters occupies the CPU allocation. Peak transient memory comes from the in-flight tensors and any conversion temporaries, not from loading an entire shard at once.

#Lazy safetensor slices

Tensors in a safetensors file have a unique name or key like model.layers.0.self_attn.q_proj.weight. This lets the loader know which weight block in the file corresponds to which parameter.

for k in file_pointer.keys():
      merged_state_dict[k] = file_pointer.get_slice(k)  # don't materialize yet

When a safetensors file is opened, get_slice(key) returns a lightweight slice object for a tensor that actually exists in the checkpoint. The slice exposes that stored tensor's shape and dtype without materializing the tensor bytes.

This is the same kind of metadata as the meta tensor, but it answers a different question. The meta tensor describes the empty parameter slot the model needs. The safetensors slice describes the real weight block available on disk for that slot.

A slice is very lightweight, so Transformers can walk through every entry - sorted by their key, which groups related tensors together (like an attention block's q_proj/k_proj/v_proj) - in the checkpoint. The device map decides where each materialized tensor should go.

def _materialize_copy(tensor: torch.Tensor, device=None, dtype=None) -> torch.Tensor:
      # This slicing is what actually loads the tensor from the safetensors slice object
      tensor = tensor[...]
      if dtype is not None or device is not None:
          tensor = tensor.to(device=device, dtype=dtype)
      return tensor

When the loader is ready to materialize a checkpoint entry, it indexes the slice with tensor[...], which reads the tensor bytes from disk. The resulting real torch.Tensor is assigned into the parameter slot, overwriting the placeholder on the meta device. Parameters assigned to "disk" are indexed or offloaded instead and remain lazily loadable through Accelerate hooks.

Any parameter the model expects, but the checkpoint doesn't contain, remains on the meta device until finalization, where tied weights are handled and truly missing parameters are initialized.

#Releasing conversion temporaries

Transformers' dynamic weight loading system is able to load checkpoints that don't match the model's expected parameter layout at runtime. There are two methods available for aligning the formats.

• WeightRenaming.convert() maps checkpoint source tensors to the expected model key by renaming it if it's different.
• WeightConverter.convert() performs a series of operations (for example, concatenating q, k, v into a single tensor) that restructures the weights from their checkpoint layout to the model's expected layout.

These methods identify source patterns that describe how weights are named in a checkpoint. For example, model.layers.5.self_attn.qkv.weight matches the pattern *.self_attn.qkv.weight. Based on the pattern, different checkpoint tensors are sorted to separate buckets in self.collected_tensors. Each bucket is stored in a dict where the value is a list of pending loads for that pattern.

On the async path, each tensor is submitted to the thread pool as soon as it's collected, and the resulting Future is appended to the matching bucket. A tensor can be materialized as soon as a worker thread is free, even while the main thread continues scanning and grouping checkpoint keys.

def materialize_tensors(self):
      collected_tensors = {}
      for key in list(self.collected_tensors.keys()):
          tensors = self.collected_tensors.pop(key)  # <-- pop, not read
          # ... resolve futures/callables ...
          collected_tensors[key] = tensors
      return collected_tensors

When Transformers is ready to convert a bucket of collected source tensors, materialize_tensors() calls pop(key) to remove that pattern's bucket from self.collected_tensors. Popping the bucket removes its reference from self.collected_tensors so they do not keep accumulating after it's been processed. After it's popped, those Futures are resolved into real torch.Tensors and stored in a local collected_tensors dict.

for first_param_name, mapping in tqdm(param_name_to_load.items(), desc="Loading weights"):
      realized_value = mapping.convert(...)
      for target_name, param in realized_value.items():
          set_param_for_module(model, target_name, param, ...)
      del realized_value  # <-- immediately free after setting

convert() renames, reshapes, or transforms the collected source tensors into the realized target tensors. An outer loading loop writes each realized tensor into the model, or offloads it if the device map assigned that parameter to "disk". The local results are deleted to release memory that is no longer required.

#Disabling async loading

has_on_the_fly_quantization = hf_quantizer is not None and not hf_quantizer.pre_quantized

  if (
      is_env_variable_true("HF_DEACTIVATE_ASYNC_LOAD")
      or "disk" in device_map.values()
      or has_on_the_fly_quantization
  ):
      thread_pool = None

Under tighter memory loading scenarios, like disk offload or on-the-fly quantization, Transformers disables async loading and uses the sync path.

The async path is faster because it spawns a worker thread to materialize tensors even while the main thread is still iterating over the checkpoint, collecting and sorting keys, or when it's running convert().

But the async path can create conditions where allowing the worker thread to run ahead can temporarily materialize more tensors than the follow-up step can process.

This is observed for on-the-fly quantization. The loader explicitly disables the async path with thread_pool=None because worker threads may pull weights to the device faster than the main thread can quantize and release those weights.

def _job():
      return _materialize_copy(tensor, device, dtype)

  if thread_pool is not None:
      return thread_pool.submit(_job)
  else:
      # Return the Callable here, not the Tensor itself, so we actually delay loading to avoid saturating cpu
      # memory during Conversion
      return _job

The sync path doesn't submit a job to a worker thread. It returns the _job function (a callable) that remembers which safetensors slice to load, the target device, and dtype, but it hasn't read the tensor yet. This is unlike a Future on the async path which may start reading immediately.

The loader stores _job in collected_tensors[...]. When materialize_tensors() invokes the callable with func(), the tensor bytes are read from disk. This keeps the collection phase cheap because while the loader is still scanning checkpoint keys and grouping tensors by pattern, collected_tensors holds small callables rather than real torch.Tensor objects.

If the sync path loaded tensors immediately instead, every matched weight becomes a real torch.Tensor during collection and stays in collected_tensors until convert() processes them. Returning a callable avoids buildup by deferring each read until the conversion step is ready to consume it.

#Disk offloading

if param_device == "disk" and (target_name not in model_buffers or offload_buffers):
      disk_offload_index = offload_and_maybe_resave_param(
          target_name, param, loading_info, disk_offload_folder, disk_offload_index, mapping
      )

For safetensors files, Transformers builds an index for each parameter that maps to "disk". It records the shard file and tensor name that already holds that weight so it can remain in the original file instead of being copied again.

Tensors mapped to "disk" are either re-saved to disk_offload_folder or, if they're already present in the offload index, left in their original safetensors file.

Accelerate installs hooks that load those parameters on demand, such as during a forward pass. This avoids keeping disk-offloaded parameters on the CPU or GPU and avoids extra writes when the original safetensors file can be reused.

Transformers uses the sync path when parameters are offloaded to disk to keep loading sequential under disk-offload memory constraints.

#Weight tying

# Perform the actual tying
  source_param = self.get_parameter_or_buffer(source_param_name)
  if "." in target_param_name:
      parent_name, name = target_param_name.rsplit(".", 1)
      parent = self.get_submodule(parent_name)
  else:
      name = target_param_name
      parent = self
  # Tie the weights
  setattr(parent, name, source_param)

Two parameter names, like lm_head and embed_tokens, may point to the same underlying tensor. This is known as weight tying. The model uses one shared embedding matrix to read input tokens and to project hidden states back into vocabulary logits. During training, the model updates one shared embedding instead of learning two separate matrices.

A useful side effect is that only one tensor allocation is needed for both parameter names.

Checkpoints often store only one copy of the tied weights (like embed_tokens) and is “missing” lm_head because it's understood that tie_weights() enables lm_head to share the same matrix as embed_tokens.

def mark_tied_weights_as_initialized(self, loading_info):
      """Adds the `_is_hf_initialized` flag on parameters that will be tied, in order to avoid initializing them
      later as they will be tied (overwritten) anyway.
      This is very important as most embeddings are tied, and they are huge params (vocabularies are often 256k), so
      running inits on them is very costly."""
      for tied_param in getattr(self, "all_tied_weights_keys", {}).keys():
          param = self.get_parameter(tied_param)
          param._is_hf_initialized = True

Loaded parameters are marked with _is_hf_initialized when they're assigned. mark_tied_weights_as_initialized() also marks known tied target parameters, such as a missing lm_head, so the post-load initialization pass does not allocate and initialize a tensor that will later be replaced by tie_weights(). After loading, only truly missing parameters are initialized using the model's default initialization scheme.

#Fuse dtype and device transfers

def _materialize_copy(tensor: torch.Tensor, device=None, dtype=None) -> torch.Tensor:
      ...
      if dtype is not None or device is not None:
          tensor = tensor.to(device=device, dtype=dtype)
      return tensor

When a lazy slice is materialized, device transfer and dtype casting are combined into one call. If the source dtype already matches the target dtype, this is only a device transfer.

The memory benefit appears when the source and target dtype differs. Splitting the device transfer and dtype casting into two calls copies the tensor to the device in the source dtype, then allocates a second tensor for the cast.

Combining device and dtype lets PyTorch produce the destination tensor directly in the requested dtype on the requested device.

#Overlapping disk reads with GPU transfers

Transformers has two paths for loading weights, a synchronous and asynchronous version, designed for different loading scenarios.

# conceptual loop
model = ModelClass(config)                             # empty skeleton, weights on "meta"
with safe_open("model.safetensors") as f:
  for name in f.keys():
      tensor = f.get_tensor(name)                      # 1. read bytes off disk into CPU RAM
      tensor = tensor.to("cuda", dtype=torch.float16)  # 2. host → device copy over PCIe
      set_module_tensor(model, name, tensor)           # 3. attach to the right param

The sync path is the ordinary sequential case. The main Python thread loads one weight at a time, working through each fully before moving to the next. There are no worker threads, so a transfer that's in progress leaves the thread with nothing to do but wait for it to finish.

1. The thread reads the tensor out of the safetensors file. safetensors memory-maps the file rather than eagerly copying it, so this doesn't move any bytes yet. It hands back a view backed by the mapped file, and the actual disk-to-RAM transfer is deferred until something reads the data. What forces that read depends on the destination.
   • CPU: nothing reads the bytes here. The transfer is postponed until an operation actually touches the tensor, like one of the conversion ops or the model's forward pass if there's no conversion.
   • GPU: the .to(device) in the next step forces it. The pages are faulted into CPU RAM first, then copied over PCIe to the GPU. So for GPU loading, the disk read effectively happens as part of step 2, and the thread still stalls waiting for it.
2. PyTorch calls cudaMemcpy to copy the bytes from the CPU to GPU memory over PCIe. The thread is blocked until the copy finishes. When it does, the new GPU tensor already exists and .to(device) returns it directly, there's no further copy. The thread reacquires the GIL to wrap the result in a Python object, and the next line runs.
3. Find and replace the matching parameter on the meta model with the actual tensor.

This matters most in memory-constrained scenarios like quantization. Each conversion briefly holds both the original tensor and the converted one, freeing the original only afterward. With multiple workers running ahead, those transient pairs stack up. The weights arrive faster than the main thread can convert and release them, and memory blows past what is available.

#Async path

The async path raises the thread count to 4. Each worker takes a job, a disk read plus the copy from CPU to GPU, so the main thread can keep moving through the state dict handing out work instead of stalling on each transfer itself.

def _materialize_copy(tensor: torch.Tensor, device=None, dtype=None):
    tensor = tensor[...]                               # 1. slice the mmmap'd tensor (bytes read lazily on access)
    if dtype is not None or device is not None:
        tensor = tensor.to(device=device, dtype=dtype) # 2. host → device copy (releases GIL)
    return tensor

The GIL serializes actual Python bytecode, and only one thread runs the interpreter at a time. But the expensive work here isn't Python. safetensors reads the bytes in native code, and most PyTorch operations run in C++.

safetensors releases the GIL while it reads bytes off disk, and PyTorch's .to(device) releases it during cudaMemcpy. Neither needs the interpreter for that work. This lets threads run in different stages. Thread A can be blocked on a host-device copy without holding the GIL, leaving it free for Thread B to start a disk read. When A's copy completes, it reacquires the GIL just long enough to wrap the result in a Python object.

Threads and jobs are managed by ThreadPoolExecutor. It creates a pool of up to max_workers threads and reuses them instead of spawning a new one for each job.

def spawn_materialize(thread_pool, tensor, device=None, dtype=None) -> Future | Callable:
    def _job():
        return _materialize_copy(tensor, device, dtype)
    if thread_pool is not None:
        return thread_pool.submit(_job)      # async path (Future)
    else:
        return _job                          # sync path: (Callable)

Loading happens in two passes.

1. In the first pass, Transformers walks the entire state dict and submits a job for every weight to the pool, stashing each as a Future, a handle to a result that isn't ready yet. This detaches submitting work from performing it. The instant a job is submitted, any free worker can start its disk read and host-to-device copy while the main thread keeps submitting the rest.

2. In the second pass, the main thread walks the weights in order, calls .result() on each Future to block until that weight is ready, and attaches it to the model. Because every job was queued up front, the workers run far ahead. By the time the main thread asks for a given weight, many later weights are already copied or inflight, so it rarely has to wait.

Intuitively, it may seem like adding more workers will keep reducing the loading time, but 4 workers wasn't chosen arbitrarily. During testing, 16 workers can take twice as long and was worse than using only 1 worker (see results here), whereas 4 workers offered the best tradeoff. Past a certain point, extra workers stop helping, and can make loading slower.

Loading weights is I/O-bound, and there are only a few shared physical resources to go around. The storage device, CPU memory for staging buffers, and the PCIe copy to the GPU. 2-3 workers are enough to keep the disk and PCIe transfers busy, plus a 4th to fill in any gaps. At that point, hardware utilization is typically already close to 100%.

Adding more workers past saturation introduces drag because the work is I/O-bound, piling on threads overloads the OS kernel, which spends more of its time switching between threads than moving data. This compounds with contention for the shared resources. With 16 workers making concurrent cudaMemcpy calls, the GPU's one or two copy engines queues most of these calls.

#CUDA caching allocator warmup

When PyTorch moves a tensor to the GPU, it uses a caching allocator on top of cudaMalloc, because otherwise, it'd be very slow to assign memory on every call. cudaMalloc is synchronous and it has to find a free region in GPU memory that's big enough.

1. The allocator requests cudaMalloc to assign a block of memory that's at least as large as tensor size N.
2. When the tensor is freed, the allocator holds on to that block in a cache.
3. The next request that fits is served straight from the cache, with no separate cudaMalloc call. And because the allocator splits a larger free block into smaller pieces as needed (and merges them back when freed), a single big block can cover many differently-sized allocations. The precise behavior here is configurable and hardware-dependent, so it may vary.

Once the allocator's cache is warm, most allocations are cheap. They're served from blocks it already holds. But anything not yet in the cache still costs a cold, synchronous cudaMalloc. Loading a model from scratch starts with an empty pool that has to grow through many of these cold calls, and for a 70B model that adds up to a real bottleneck.

Transformers performs one big allocation per device upfront and then discards it.

for device, byte_count in total_byte_count.items():
  # ... memory accounting to pick the right byte_count ...
  # We divide by 2 here as we allocate in fp16
  _ = torch.empty(int(byte_count // 2), dtype=torch.float16,
                  device=device, requires_grad=False)

torch.empty() does two important things.

• Make one large cudaMalloc call. The block covers roughly the model's footprint on that device. Transformers first subtracts memory PyTorch already holds in reserve, and caps the request below the device's total so that a large final parameter can't trigger an OOM.
• The throwaway tensor is bound to _. As soon as _ goes out of scope, reassigned on the next device or dropped when the function returns, nothing references the tensor and its block is freed.

Crucially, when that tensor is freed, PyTorch doesn't return the block of memory to the driver. It keeps it for reuse, so this way, Transformers only pays the slow cudaMalloc cost once upfront. torch.empty() allocates the memory but never writes to it, unlike torch.zeros(), there's no data pass to fill the buffer. The allocation call without the data call forces cudaMalloc and warms up the allocator. When a real weight is loaded, PyTorch uses this block without calling cudaMalloc again. Every subsequent .to(device) uses small pieces of memory from the large pre-fetched block.

Pre-allocating more memory than the model size doesn't help. After making the single big cudaMalloc call, subsequent loading is dominated by data movement and not available memory.

Thanks to Anton Vlasjuk and Cyril Vallez for the review and feedback.