Tensorparallele lineare Schichten: ColumnParallelLinear / QKVParallelLinear / RowParallelLinear

源码版本v0.25.1

Verantwortung

vLLMs Tensorparallelismus (tensor parallelism, TP) zerlegt ein einzelnes nn.Linear so, dass mehrere GPUs jeweils einen Teil berechnen und das Ergebnis in Vorwärts- bzw. Rückwärtsrichtung per all-gather / all-reduce zusammengesetzt wird. Diese Abstraktion liegt in vllm/model_executor/layers/linear.py auf drei zentralen Klassen: ColumnParallelLinear entlang der Ausgabedimension, RowParallelLinear entlang der Eingabedimension und QKVParallelLinear für die Aufmerksamkeit mit unterschiedlichen Head-Anzahlen für Q/K/V (insbesondere bei GQA / MQA). Alle erben von LinearBase, der eigentliche GEMM-Aufruf wird an quant_method.apply delegiert(linear.py:551-569), sodass dieselbe Parallellogik sowohl für nicht-quantisierte als auch für fp8-, AWQ-, GPTQ- und andere Backends funktioniert.

Bei der Instanziierung jeder TP-linearen Schicht werden tp_rank / tp_size des aktuellen Ranks berechnet, output_size durch tp_size geteilt, um output_size_per_partition zu erhalten(linear.py:437-442), und diese Partitionsinformation in den Attributen output_dim / input_dim des Parameters abgelegt. Der weight_loader wählt anhand dieses Attributs, welches Segment geschnitten wird. QKVParallelLinear behandelt zusätzlich die KV-Head-Replikation bei GQA: wenn tp_size >= total_num_kv_heads, hält jeder Rank eine vollständige Kopie der KV-Heads (num_kv_head_replicas = tp_size / total_num_kv_heads)(linear.py:994-1001).

Entwurfsmotivation

Warum gibt es auf LinearBase zwei Mechanismen: weight_loader / weight_loader_v2?

• Unterschiedliche Schnittrichtung: ColumnParallelLinear schneidet entlang der Ausgabedimension (jeder Rank erhält ein Segment aus [A_1, A_2, ...]), RowParallelLinear entlang der Eingabedimension (A = [A_1; A_2; ...], X = [X_1, X_2, ...]). Daher führt der weight_loader beider Klassen jeweils ein narrow entlang output_dim bzw. input_dim aus(linear.py:530-533)(linear.py:1650-1653).
• v2 arbeitet auf Parameterobjekten: Der neue weight_loader_v2 verpackt den Tensor in ein BasevLLMParameter und ruft direkt param.load_column_parallel_weight bzw. param.load_row_parallel_weight auf – das "welches Segment" ist im Parameterobjekt gekapselt(linear.py:543-549)(linear.py:1663-1670). Quantisierungsmethoden können wählen, ob sie v2 (WEIGHT_LOADER_V2_SUPPORTED) oder den alten Pfad verwenden.
• Q/K/V-Drei-Segment-Verschmelzung: QKVParallelLinear fügt die Q/K/V-Gewichte in der Ausgabedimension zu einer großen Matrix zusammen. output_sizes = [q, k, v] mit je einzelnem divide(..., tp_size)(linear.py:1003-1008). MergedColumnParallelLinear ist die allgemeine Variante, mit der sich beliebig viele spaltenparallele Schichten verschmelzen lassen.
• GQA-KV-Replikation: In QKVParallelLinear.__init__ gilt if tp_size >= total_num_kv_heads: self.num_kv_heads = 1; self.num_kv_head_replicas = divide(tp_size, total_num_kv_heads)(linear.py:996-1001), sodass jeder Rank einen vollständigen KV-Head berechnen kann und im Forward keine zusätzliche Synchronisation nötig ist.
• Shard-Metadaten: Der weight_loader_v2 von ColumnParallelLinear muss zusätzlich den shard_id erkennen (bei QKV "q" / "k" / "v") und anhand der Shard-ID den Offset bestimmen, in den der Parameter geschrieben wird(linear.py:1023-1047).
• bias nur auf Rank 0: In RowParallelLinear.forward gilt bias_ = None if (self.tp_rank > 0 or self.skip_bias_add) else self.bias(linear.py:1685-1688). Würde bias nach dem Reduce hinzugefügt, entstünde eine Doppeladdition, daher erfolgt sie nur auf Rank 0.
• TP abschaltbar: Bei disable_tp=True werden tp_rank=0 / tp_size=1 gesetzt, die Schicht verhält sich wie ein normales Linear – nützlich für nicht-parallele Experten bei MoE(linear.py:438-439).

Schlüsseldateien

• LinearMethodBase:141-180 — abstrakte Basisklasse der Quantisierungsmethode, definiert create_weights / apply / process_weights_after_loading.
• LinearBase:231-288 — Elternklasse aller linearen Schichten; update_param_tp_status schreibt tp_rank / tp_size auf jedes BasevLLMParameter.
• ColumnParallelLinear.__init__:423-491 — Schneiden entlang der Ausgabedimension, output_size_per_partition = divide(output_size, tp_size).
• ColumnParallelLinear weight_loader:520-549 — v1 per output_dim-narrow, v2 delegiert an param.load_column_parallel_weight.
• ColumnParallelLinear.forward:551-569 — quant_method.apply → optional tensor_model_parallel_all_gather.
• QKVParallelLinear.__init__:942-1021 — behandelt die KV-Head-Replikation bei GQA und fasst output_sizes von Q/K/V zum output_size zusammen.
• QKVParallelLinear shard mapping:1023-1047 — _get_shard_offset_mapping / _get_shard_size_mapping, teilt dem weight_loader mit, in welches Segment des Params geschrieben wird.
• QKVParallelLinear fused checkpoint:1049-1097 — für bereits auf der Festplatte verschmolzene QKV (wie bei Phi-3), automatische Aufteilung und anschließender Aufruf von weight_loader_v2 pro Segment.
• RowParallelLinear.__init__:1572-1626 — Schneiden entlang der Eingabedimension, input_size_per_partition = divide(input_size, tp_size).
• RowParallelLinear.forward:1672-1698 — wenn die Eingabe nicht parallelisiert ist, erst split, nach dem GEMM tensor_model_parallel_all_reduce.
• MergedColumnParallelLinear:580-636 — Verschmelzung mehrerer spaltenparalleler Schichten (allgemeine Variante von QKV), output_sizes-Liste legt jedes Segment fest.

Datenfluss

Nachfolgend die Kernlogik von ColumnParallelLinear.weight_loader: Aus dem vollständigen Tensor auf der Festplatte wird mit Offset tp_rank * shard_size das für diesen Rank benötigte Segment geschnitten und per copy_ in den Parameter kopiert:

def weight_loader(self, param: Parameter, loaded_weight: torch.Tensor):
    output_dim = getattr(param, "output_dim", None)

    is_sharded_weight = getattr(param, "is_sharded_weight", False)
    use_bitsandbytes_4bit = getattr(param, "use_bitsandbytes_4bit", False)
    # bitsandbytes loads the weights of the specific portion
    # no need to narrow
    is_sharded_weight = is_sharded_weight or use_bitsandbytes_4bit

    param_data = param.data
    if output_dim is not None and not is_sharded_weight:
        shard_size = param_data.shape[output_dim]
        start_idx = self.tp_rank * shard_size
        loaded_weight = loaded_weight.narrow(output_dim, start_idx, shard_size)

    # Special case for loading scales off disk, which often do not
    # have a shape (such as in the case of AutoFP8).
    if len(loaded_weight.shape) == 0:
        loaded_weight = loaded_weight.reshape(1)

    assert param_data.shape == loaded_weight.shape
    param_data.copy_(loaded_weight)

(linear.py:520-541) RowParallelLinear.weight_loader ist das Spiegelbild davon, nur dass die Dimension für narrow von output_dim auf input_dim wechselt(linear.py:1650-1653). Beide Funktionen werden beim Laden der Gewichte über model.load_weights vom DefaultModelLoader indirekt aufgerufen – in jeder Modellimplementierung findet sich eine Zeile wie loaded_weights = param.weight_loader(param, loaded_weight).

Im Forward geht ColumnParallelLinear über all-gather (bei gather_output=True wird output_parallel jedes Ranks zum vollständigen Output zusammengefügt), RowParallelLinear über all-reduce (bei reduce_results=True werden die blockweisen GEMM-Ergebnisse aufsummiert). Werden beide kombiniert, lässt sich eine Synchronisation einsparen: QKVParallelLinear (column) → Attention → o_proj (row) ist die klassische Column-then-Row-Kombination.

Grenzen und Fehler

• Ungültige QKV-shard_id: validate_shard_id schränkt die Shard-ID für QKV auf "q" / "k" / "v" oder None ein, alle anderen Werte lösen sofort raise ValueError aus(linear.py:1023-1029).
• RowParallelLinear reduce_results=False + bias: Weder reduce noch im GEMM bias addieren – sonst addiert jeder Rank ein eigenes bias, das Ergebnis ist falsch, der Code wirft raise ValueError(linear.py:1622-1626).
• FP8-Blockform passt nicht: _maybe_allow_fp8_block_shape_mismatch gibt allow_fp8_block_shape_mismatch frei, wenn weight_block_size ein output_partition_sizes nicht ganzzahlig teilt(linear.py:493-516), sonst würde der Kernel ablehnen.
• bitsandbytes überspringt narrow: Bei use_bitsandbytes_4bit ist der Checkpoint bereits nach Rank aufgeteilt, direkt copy_ ohne narrow(linear.py:524-527).
• scales als null dimensionaler Skalar: Checkpoints wie AutoFP8 speichern den Scale als null dimensionalen Tensor, weight_loader fängt das über reshape(1) ab(linear.py:537-538).
• Row-Schicht mit input_is_parallel=False: Wenn die Eingabe nicht aufgeteilt ist, führt RowParallelLinear.forward selbst ein split_tensor_along_last_dim aus, um den Teil dieses Ranks zu erhalten(linear.py:1676-1682).

Zusammenfassung

Die drei parallelen linearen Klassen bilden die "Gewichtsschnittsprache" von vLLMs TP; wie konkret gerechnet wird, übernimmt quant_method. Die Nahtstelle zu Quantisierungsschicht laden ist der Aufruf von self.quant_method.create_weights(...) in ColumnParallelLinear.__init__, der jedem Parameter einen weight_loader (v1 oder v2) aufsetzt. Die Nahtstelle zu DefaultModelLoader ist der Aufruf von weight_loader pro Tensor innerhalb von model.load_weights, der die vollständigen Gewichte von der Festplatte in das für diesen Rank benötigte Segment schneidet.

Siehe offizielle Dokumentation: vLLM 文档 · README.