vLLM's Model Registry and Weight Loading — Resolution to Shard, in Source

A Name, Not a Class

What is a HuggingFace architecture string?

Before vLLM imports a single line of model code, it has a HuggingFace config.json in hand — and inside it, one field decides everything that follows: architectures, a list of strings. For this module's running example, meta-llama/Llama-3.1-8B, that field is ["LlamaForCausalLM"]. That string is not a class. It is not even a guarantee that a class by that name exists anywhere vLLM can see. It is a name, sitting in a JSON file, and the entire job of the next two steps is turning it into something you can actually call.

The simulation's first two stages show exactly this hand-off. Stage 1 renders the config as vLLM reads it — model_id and architectures, nothing else yet. Stage 2 pulls config.architectures[0] out on its own, still just text, still nothing looked up.

The registry is a name-to-class table

_ModelRegistry — the class this module is built around — is a lookup table: a plain Python dict, keyed by architecture string, mapping each one to a module and a class name. It is assembled from several real per-task tables inside vLLM's source (vllm/model_executor/models/registry.py) — one for text-generation models, one for embedding models, one for multimodal models, one for speculative-decoding draft models, and more — merged into a single dict with several hundred real entries. A handful, verified directly against that source:

Every row here is a real entry, not an illustration built to look real. LlamaForCausalLM is genuinely in that dict, genuinely pointing at genuinely that module and class.

The miss case is not hypothetical

A table only tells half its own story if you never see it fail. This module's fixture also carries an architecture string that was never registered: CustomVisionForCausalLM — a stand-in for a custom or fine-tuned checkpoint whose config.json names an architecture vLLM's table simply does not contain. That is a completely ordinary way to end up in trouble: point vLLM at a checkpoint from a research repo, a merged fine-tune, or a brand-new model family the current pinned release predates, and its architectures field can easily name something no dict entry matches.

Say what you now know

You can name the field (architectures), say what it is (a list of strings, not classes), and say what answers it (a table with hundreds of real entries, keyed by exactly those strings). What you cannot yet say is how the lookup actually runs, or why — for a narrower set of questions this module doesn't fully cover — vLLM sometimes has to leave the calling process to answer them. That is Step 2.

How Resolution Actually Runs

How does vLLM resolve an architecture string to a class?

_ModelRegistry.resolve_model_cls is the method that turns an accepted architecture string into an actual, importable Python class — the method Stage 3's hit and miss cards are both real calls against. Its job, stripped to the essentials: walk the requested architecture strings, ask the table for each one, and return the first real class it gets back. If none of them resolve, raise, with a message that names what was asked for.

# distilled — real names, reduced body: the Transformers-backend and# Terratorch paths are elided. Both can also run first, whenever# --model-impl requests them — not only as a last-resort fallback.class _ModelRegistry:  models: dict[str, _BaseRegisteredModel]   def resolve_model_cls(      self,      architectures: str | list[str],      model_config: ModelConfig,  ) -> tuple[type[nn.Module], str]:      if isinstance(architectures, str):          architectures = [architectures]      if not architectures:          raise ValueError("No model architectures are specified")       for arch in architectures:          normalized_arch = self._normalize_arch(arch, model_config)          model_cls = self._try_load_model_cls(normalized_arch)          if model_cls is not None:              return (model_cls, arch)       return self._raise_for_unsupported(architectures)   def _try_load_model_cls(self, model_arch: str) -> type[nn.Module] | None:      if model_arch not in self.models:          return None      return _try_load_model_cls(model_arch, self.models[model_arch])

Two details worth reading twice. First, the return type: a tuple, (model_cls, arch) — the resolved class and the architecture string that matched, not the module it came from. Callers that only unpack the class and discard the second element are throwing away information vLLM itself uses elsewhere to normalize architecture aliases. Second, a total miss doesn't return None for the caller to check — it raises, from inside _raise_for_unsupported, before any loader has been constructed. CustomVisionForCausalLM, Step 1's fixture, takes exactly this path: every architecture in the list fails _try_load_model_cls, and the method never reaches a return statement at all.

What actually happens on a hit

self.models[model_arch] is not the class itself — it is a small wrapper object that knows a module name and a class name but has not imported anything yet. The module-level _try_load_model_cls function above calls model.load_model_cls() on that wrapper, and for the ordinary case that resolves to one line: importlib.import_module(self.module_name), then getattr(mod, self.class_name). That is a plain, synchronous import, running in whatever process called resolve_model_cls in the first place. Nothing about this path leaves the process.

Why "for some models" it runs somewhere else

This is the part worth being precise about, because it is easy to blur two similar-sounding registry methods into one. resolve_model_cls — the method above — never spawns anything; a failure inside load_model_cls() is caught by the surrounding try/except Exception and turned into None, not routed anywhere else. But _ModelRegistry also has a sibling method, inspect_model_cls, used for a narrower class of question: things vLLM needs to know about a model before it has decided to load it — is this a pooling model, does it support pipeline parallelism, and similar. Answering those questions means actually constructing the class's interface metadata, which means importing the model module — and some model modules do CUDA-touching work as a side effect of being imported. Importing them in the same process that is about to serve requests risks initializing CUDA too early, in a process that isn't ready for it.

So inspect_model_cls, via a wrapper method called _try_inspect_model_cls, hands that specific probe off to _run_in_subprocess — a real vLLM helper that pickles the call, runs it in a fresh Python process, and reads the result back. That subprocess hop belongs to inspect_model_cls and the machinery behind it, not to resolve_model_cls. Neither of those two methods is one this module cites as a DistilledCode block — this module's citations stop at the resolution path above — but the distinction matters enough to state plainly: the class-resolution code shown above runs in-process. A different, narrower path runs out-of-process, and it is not this one.

The Loader Interface

What is vLLM's model loader?

Resolving a class answers "which Python object represents this model." It says nothing about how that object gets real numbers in its parameters. That is a separate job, owned by a separate interface: BaseModelLoader, and — for the ordinary case of a HuggingFace checkpoint on disk — its concrete implementation, DefaultModelLoader. Stage 4 of the simulation is where this module's story hands off from "which class" to "which loader."

The interface: allocate, then delegate

# distilled — real names, reduced body: the online-quantization branch# and its peak-memory logging are elidedclass BaseModelLoader(ABC):  def load_model(      self, vllm_config: VllmConfig, model_config: ModelConfig, prefix: str = "",  ) -> nn.Module:      """Load a model with the given configurations."""      device_config = vllm_config.device_config      load_config = vllm_config.load_config      load_device = (          device_config.device if load_config.device is None else load_config.device      )      target_device = torch.device(load_device)      with set_default_torch_dtype(model_config.dtype):          with target_device:              model = initialize_model(                  vllm_config=vllm_config,                  model_config=model_config,                  prefix=prefix,              )           self.load_weights(model, model_config)           process_weights_after_loading(model, model_config, target_device)       return model.eval()   @abstractmethod  def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None:      """Load weights into a model."""      raise NotImplementedError

Three steps, in order, and each one only makes sense after the one before it: allocate the model's structure on the target device with empty tensors (initialize_model — this is the resolved class from Step 2, instantiated, not yet holding any real weights), fill those tensors (self.load_weights, abstract on this base class — Step 4 is what DefaultModelLoader does here), then finalize them into whatever internal format the runtime kernels expect (process_weights_after_loading — quantization packing, for example, when a quantized load format is in play). load_model itself never touches a single weight value; it is the choreography, not the copy.

Worker.load_model is the caller, not the loader

Worker.load_model — the method the simulation's Loader stage names as its entry point — does not implement any of the three steps above. It is several layers up the call stack, and its entire job is reaching down to them:

# distilled — real names, reduced body: WeightTransferEngineFactory's# full argument list is unchanged from the real callclass Worker(WorkerBase):  def load_model(self, *, load_dummy_weights: bool = False) -> None:      with (          self._maybe_get_memory_pool_context(tag="weights"),          set_current_vllm_config(self.vllm_config),          self._scoped_allocator_max_split(max_split_size_mb=20),      ):          self.model_runner.load_model(load_dummy_weights=load_dummy_weights)       if self.vllm_config.weight_transfer_config is not None:          self.weight_transfer_engine = WeightTransferEngineFactory.create_engine(              self.vllm_config.weight_transfer_config,              self.vllm_config,              self.device,              self.model_runner.get_model(),          )

self.model_runner.load_model(...) is the real next hop, and from there the call eventually reaches whichever loader get_model_loader(...) returns — DefaultModelLoader for a plain HuggingFace checkpoint, a different concrete loader for other load formats. Worker.load_model never constructs a loader itself and never touches load_weights; it sets up the device-memory bookkeeping around the call (a memory-pool context, a scoped allocator limit) and delegates. That shape — a thin, memory-aware caller on top of the real interface — is exactly what Stage 4's chip is naming when it shows Worker.load_model → BaseModelLoader.load_model.

What's still missing

self.load_weights(model, model_config) is abstract here — you've seen the interface's shape but not a single real tensor move yet. DefaultModelLoader's implementation is where checkpoint files actually get opened, and it's the subject of Step 4.

Where Weights Land

How do checkpoint tensors become model parameters?

DefaultModelLoader.load_weights — the concrete implementation Step 3 left abstract — does two things in sequence: it opens an iterator over every tensor sitting in the checkpoint's files, and it hands that iterator to the model class's own load_weights, which maps each tensor onto exactly one of the model's parameters. Stage 5 of the simulation renders that second half directly: one row per checkpoint tensor, an arrow, and where it landed — or the fact that it didn't land anywhere at all.

Opening the checkpoint

# distilled — real names, reduced body: the npcache, fastsafetensors,# instanttensor and multithreaded-loading branches are elided down to# the two ordinary cases, safetensors and plain .ptclass DefaultModelLoader(BaseModelLoader):  def _get_weights_iterator(      self, source: "Source",  ) -> Generator[tuple[str, torch.Tensor], None, None]:      hf_folder, hf_weights_files, use_safetensors = self._prepare_weights(          source.model_or_path,          source.subfolder,          source.revision,          source.fall_back_to_pt,          source.allow_patterns_overrides,      )      if use_safetensors:          weights_iterator = safetensors_weights_iterator(              hf_weights_files,              self.load_config.use_tqdm_on_load,              self.load_config.safetensors_load_strategy,          )      else:          weights_iterator = pt_weights_iterator(              hf_weights_files,              self.load_config.use_tqdm_on_load,              self.load_config.pt_load_map_location,          )       # Apply the prefix.      return ((source.prefix + name, tensor) for (name, tensor) in weights_iterator)

The return value is a generator of (name, tensor) pairs, streamed straight off the checkpoint's .safetensors or .bin files — model.layers.0.self_attn.q_proj.weight and the raw tensor that sits behind it, and so on for every tensor the checkpoint contains. Nothing here knows anything about vLLM's own parameter names yet. That mapping is the model class's job, not the loader's.

One name, one destination — or a real failure

The model class's load_weights — built on a shared helper, AutoWeightsLoader, which LlamaModel itself uses (loader = AutoWeightsLoader(self); return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)), and most vLLM models do too, though not all — walks the module tree and, for each incoming tensor name, does exactly one of three things: matches it to a child module and recurses, matches it to a leaf parameter and copies the data in, or — if the name matches neither, and isn't on an explicit skip list — raises ValueError naming the exact prefix it couldn't place and what is available at that point in the tree. There is no fourth outcome. A name that doesn't fit anywhere is never quietly dropped; it stops the load.

This module's fixture is Llama's real fusion table, verified against vllm/model_executor/models/llama.py: three separate checkpoint tensors, q_proj, k_proj, and v_proj, all land on one vLLM parameter, qkv_proj, each carrying a distinct shard identity ("q", "k", "v"); gate_proj and up_proj fuse the same way into gate_up_proj; o_proj and down_proj pass through unchanged, since vLLM's own attribute names already match the checkpoint's. Seven real attention-and-MLP tensor names go into Stage 5, and all seven resolve. One more — a stand-in for what a custom or fine-tuned checkpoint actually contains, not a real vLLM symbol — does not: an extra tensor whose name matches nothing in a stock Llama's parameter tree. Stage 5 marks that eighth row distinctly: unmapped, not silently absent. That row is this module's real failure specimen, and its timing matters: this is a load-time failure. It happens once, inside Worker.load_model, while the checkpoint is still being read — before the engine has scheduled anything, let alone served a request. A custom architecture that resolves cleanly in Step 2 can still fail here, later in the same startup sequence, for an entirely different reason: the registry found the right class, but the checkpoint's tensor names don't match what that class's load_weights expects.

What actually copies the bytes

Once a tensor has a destination parameter, something still has to copy the data in. Most parameters use the module-level fallback:

# distilled — real names, reduced body: the bare re-raise (kept only so# a debugger can set a breakpoint on it) is elideddef default_weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None:  """Default weight loader."""  if param.numel() == 1 and loaded_weight.numel() == 1:      # Scalar values sometimes arrive without a shape; reshape to      # match before copying.      param.data.copy_(loaded_weight.view(param.shape))  else:      assert param.size() == loaded_weight.size(), (          f"Attempted to load weight ({loaded_weight.size()}) "          f"into parameter ({param.size()})"      )      param.data.copy_(loaded_weight)

A straight copy, shape-checked, no rank-aware slicing at all — this is the function every replicated (non-tensor-parallel) parameter uses. But it is only the fallback: the actual dispatch is getattr(param, "weight_loader", default_weight_loader) — if a parameter carries its own weight_loader, that one runs instead. Which parameters carry their own is Step 5's subject, and it is where this module's central claim about tensor parallelism actually lives.

The Same Checkpoint, Two Different Splits

Where does a tensor-parallel shard get decided?

Not in BaseModelLoader. Not in DefaultModelLoader. Not in Worker.load_model. Every one of those Step 3 and Step 4 walked, and none of them contains a single line that reads tp_rank. That absence is the point of this step: the loader interface's job is iterating tensors and finding each one a destination parameter — it has no opinion about which rank gets which bytes. The decision lives one level deeper, on the parameter itself, and it is made at the exact moment a tensor lands.

What tensor parallelism actually splits

A quick anchor, since this module doesn't assume you've seen it elsewhere: tensor parallelism (TP) splits a single large matrix multiplication across several GPUs, so that each one holds and computes on only a slice of the weight matrix instead of the whole thing — the technique goes back to Megatron-LM's original formulation. --tensor-parallel-size N is the launch flag that sets how many slices. Module 1 of this track's own hedge still applies here, not a stronger version of it: at N=1 — together with pipeline- and data-parallel size also at 1 — there is no separate worker process at all — the model runs inside EngineCore. At world size greater than 1, vLLM defaults to the multiprocessing (mp) executor and gives each rank its own VLLM::Worker process, holding one shard of the model's weights and one slice of its KV cache. This step is about the mechanism inside that fact: how a single checkpoint tensor becomes N slices, one per rank, as it loads.

The parameter's own weight_loader decides — and it isn't one method

o_proj and down_proj, in this module's Llama fixture, are built directly from RowParallelLinear, unmodified. Its weight_loader is attached to the parameter at construction time, before any checkpoint has been opened, and when one of those parameters' turn comes in Step 4's tensor loop, getattr(param, "weight_loader", default_weight_loader) finds it instead of the module-level fallback:

# real shape of RowParallelLinear.weight_loader's sharding path — this
# module doesn't cite a symbolId for it, since it isn't in this module's
# manifest; the quantization and bitsandbytes branches are elided
def weight_loader(self, param, loaded_weight):
    param_data = param.data
    input_dim = param.input_dim
    shard_size = param_data.shape[input_dim]
    start_idx = self.tp_rank * shard_size
    loaded_weight = loaded_weight.narrow(input_dim, start_idx, shard_size)
    param_data.copy_(loaded_weight)

qkv_proj and gate_up_proj are different classes, not the same mechanism with different parameter names: they're built from QKVParallelLinear and MergedColumnParallelLinear, and both subclass ColumnParallelLinear and override weight_loader with logic aware of their own fused, multi-tensor layout — a checkpoint tensor named q_proj doesn't land on all of qkv_proj, it lands on a specific sub-region of it. Verified against vllm/model_executor/layers/linear.py, QKVParallelLinear.weight_loader's real shape for the ordinary q/k/v case:

# distilled — real names, reduced body: QKVParallelLinear.weight_loader,
# scoped to loaded_shard_id in ("q", "k", "v"); the fused-on-disk and
# quantization/bitsandbytes branches are elided
def weight_loader(self, param, loaded_weight, loaded_shard_id=None):
    param_data = param.data
    output_dim = param.output_dim
    if loaded_shard_id == "q":
        shard_offset, shard_size = 0, self.num_heads * self.head_size
    elif loaded_shard_id == "k":
        shard_offset = self.num_heads * self.head_size
        shard_size = self.num_kv_heads * self.head_size
    else:  # "v"
        shard_offset = (self.num_heads + self.num_kv_heads) * self.head_size
        shard_size = self.num_kv_heads * self.v_head_size

    param_data = param_data.narrow(output_dim, shard_offset, shard_size)

    # GQA: when fewer KV heads exist than ranks, several ranks share the
    # same shard_rank and load the identical K/V slice.
    shard_rank = self.tp_rank if loaded_shard_id == "q" else self.tp_rank // self.num_kv_head_replicas
    start_idx = shard_rank * shard_size
    loaded_weight = loaded_weight.narrow(output_dim, start_idx, shard_size)
    param_data.copy_(loaded_weight)

Two narrows, not one. param_data.narrow(...) picks which sub-region of the fused qkv_proj parameter this call is filling — q's region, k's, or v's — a fixed offset, the same on every rank. loaded_weight.narrow(...) is the actual tensor-parallel split: which slice of that sub-region's checkpoint bytes this rank gets. That's where shard_rank can differ from tp_rank outright. MergedColumnParallelLinear.weight_loader (gate_up_proj) follows the same two-narrow shape for a real integer loaded_shard_id (0 for gate_proj, 1 for up_proj): a fixed per-shard offset picks gate's or up's region of the fused parameter, then start_idx = self.tp_rank * shard_size slices the checkpoint directly by tp_rank — no head-replication branch, because MLP gate/up don't have anything like a KV-head count to replicate.

num_kv_head_replicas is tp_size // total_num_kv_heads when there are fewer KV heads than ranks — grouped-query attention (GQA) at high tensor-parallel size, a real and common configuration, not an edge case invented for thoroughness. When it's greater than 1, shard_rank repeats: two or more ranks land on the same shard_rank for K and V, so their loaded_weight.narrow(...) calls read the identical slice of the checkpoint's K/V bytes. param_data.narrow(...) still only ever gives a rank its own, correctly-sized region of its own qkv_proj parameter — no rank's parameter memory footprint grows — but the content landing in two different ranks' K/V regions can be the same bytes, read twice. So "contiguous and non-overlapping" describes RowParallelLinear's slices, and qkv_proj's Q slices, exactly — and it's the ordinary case for K/V too. But at high enough TP relative to KV head count, K/V slices are not guaranteed distinct across ranks. What doesn't change: a rank's own shard never grows to cover more than its own region, and no rank ends up holding a full copy of the parameter.

The same checkpoint, two different splits

Stage 6 of the simulation makes the ordinary case concrete: the same six real Llama parameters — qkv_proj (QKVParallelLinear), o_proj (RowParallelLinear), gate_up_proj (MergedColumnParallelLinear), down_proj (RowParallelLinear), embed_tokens (VocabParallelEmbedding), and lm_head (ParallelLMHead) — shown at TP=1 and TP=2. At TP=1 every parameter's single rank holds it in full — there is only one slice. Flip to TP=2 and every one of those same six parameters splits into two disjoint halves, one per rank, and the two halves' byte counts sum back to exactly the original parameter's size. This module's fixture never puts enough ranks against few enough KV heads to reach the replication branch above — TP=2 against Llama-3.1-8B's real KV head count doesn't come close — so what Stage 6 shows is genuinely the ordinary, disjoint case. The code above is what happens once TP climbs high enough relative to KV heads for that to change.

That is also the precise shape of the claim this module keeps coming back to: each rank ends up holding its own slice of every tensor-parallel parameter, sized and positioned by its own rank number — never a second full copy sitting next to everyone else's. A worker's memory holds exactly the shard narrow() gave it, nothing more.

What Adding a Model Requires

What does it take to add a new model to vLLM?

Three things, and every one of them maps directly onto a surface the last five steps already opened up. Nothing about "adding a model" is request-shaped work — it's a one-time registration plus two contracts that get exercised once, at startup, exactly like everything else this module has walked. vLLM's own Adding a New Model guide is the maintainers' full walkthrough for a specific architecture; this step names the three surfaces it sits on top of.

1. A name in the registry

Step 1 and Step 2's table is not fixed at import time — _ModelRegistry.register_model(model_arch, model_cls) adds an entry to it, and model_cls can be either an already-imported class or a "<module>:<class>" string for lazy import (the same lazy-wrapper shape every built-in entry already uses, so a custom model gets exactly the same "not imported until resolved" behavior Step 2 traced). Out-of-tree code registers through vLLM's plugin mechanism so the entry exists before resolve_model_cls is ever asked about it — without a registered name, Step 2's lookup fails exactly the way CustomVisionForCausalLM does, architecture correct or not.

2. A load_weights that places every checkpoint tensor

Step 4's failure case — a tensor name AutoWeightsLoader can't place — is the single most common way a new model integration breaks. A model class either accepts vLLM's default AutoWeightsLoader walk (which requires its own attribute names, like qkv_proj or gate_up_proj, to already match what the checkpoint calls them) or supplies a hf_to_vllm_mapper — the same orig_to_new_stacked table Step 4's fixture was built from — that renames and fuses checkpoint tensor names onto vLLM's own parameter names before the walk happens. Skip this and a real checkpoint's tensor names simply won't resolve; there is no silent partial load, only the ValueError Step 4 traced.

3. Tensor-parallel layers, not a hand-written weight_loader

This is the one people most often over-build. Step 5 showed that sharding logic lives on each parallel layer class's own weight_loader — RowParallelLinear's directly, and QKVParallelLinear's and MergedColumnParallelLinear's in an overriding version aware of their own fused layout — not on anything a model author writes. Building a new model out of vLLM's own parallel layer classes (QKVParallelLinear, MergedColumnParallelLinear, RowParallelLinear, VocabParallelEmbedding, and their relatives) means that sharding behavior comes for free: each class's own weight_loader, attached automatically when the layer is constructed — not one shared function every model reuses, but never something a model author has to write either. A model author who instead builds attention and MLP projections out of plain nn.Linear gets no sharding at all: every rank would try to hold the full parameter, and tensor parallelism would silently do nothing for that layer.