Support Tensor Parallel in Python API (#1661)

Author: rasbt

.github/azure-gpu-test.yml

@@ -9,7 +9,6 @@ pr:
     include:
       - "main"
       - "wip"
-      - "carmocca/*"
 
 jobs:
   - job: testing
@@ -18,6 +17,7 @@ jobs:
     pool: "lit-rtx-3090"
     variables:
       DEVICES: $( python -c 'print("$(Agent.Name)".split("_")[-1])' )
+      CI: "true"
     container:
       image: "pytorchlightning/pytorch_lightning:base-cuda-py3.10-torch2.2-cuda12.1.0"
       options: "--gpus=all --shm-size=8gb"

litgpt/api.py

@@ -6,6 +6,7 @@
 import time
 from typing import Any, List, Literal, Optional, Union
 
+from tqdm import tqdm
 import torch
 import lightning as L
 from lightning.fabric.plugins import BitsandbytesPrecision
@@ -16,12 +17,14 @@
 from litgpt.config import name_to_config, Config
 from litgpt.tokenizer import Tokenizer
 from litgpt.generate.sequentially import sequential
+from litgpt.generate.tp import tensor_parallel
 from litgpt.generate.base import generate as generate_fn
 from litgpt.chat.base import generate as stream_generate_fn
 from litgpt.prompts import load_prompt_style, has_prompt_style, PromptStyle
 from litgpt.utils import (
     auto_download_checkpoint,
     check_file_size_on_cpu_and_warn,
+    check_nvlink_connectivity,
     extend_checkpoint_dir,
     get_default_supported_precision,
     load_checkpoint,
@@ -38,7 +41,7 @@ def __init__(
         config: Config = None,
         checkpoint_dir: Path = None,
         fabric: L.Fabric = None,
-        generate_strategy: Optional[Literal["sequential"]] = None,
+        generate_strategy: Optional[Literal["sequential", "tensor_parallel"]] = None,
         kv_cache_initialized: bool = False,
         fixed_kv_cache_size: Union[int, Literal["max_model_supported"], None] = None
     ) -> None:
@@ -182,7 +185,7 @@ def distribute(
         devices: Union[int, List[int]] = 1,
         precision: Optional[Any] = None,
         quantize: Optional[Literal["bnb.nf4", "bnb.nf4-dq", "bnb.fp4", "bnb.fp4-dq", "bnb.int8"]] = None,
-        generate_strategy: Optional[Literal["sequential"]] = None,
+        generate_strategy: Optional[Literal["sequential", "tensor_parallel"]] = None,
         fixed_kv_cache_size: Union[int, Literal["max_model_supported"], None] = None
     ):
         """
@@ -226,14 +229,14 @@ def distribute(
             else:
                 accelerator = "cpu"
 
-        if generate_strategy == "sequential" and accelerator not in ("cuda", "gpu"):
-            raise NotImplementedError("generate_strategy='sequential' is only supported for accelerator='cuda'|'gpu.")
+        if generate_strategy in ("sequential", "tensor_parallel") and accelerator not in ("cuda", "gpu"):
+            raise NotImplementedError(f"generate_strategy='{generate_strategy}' is only supported for accelerator='cuda'|'gpu'.")
 
         num_devices = calculate_number_of_devices(devices)
 
         if generate_strategy is None and num_devices > 1:
             raise NotImplementedError(
-                "Support for multiple devices is currently only implemented for generate_strategy='sequential'."
+                "Support for multiple devices is currently only implemented for generate_strategy='sequential'|'tensor_parallel'."
             )
 
         plugins = None
@@ -244,13 +247,25 @@ def distribute(
             plugins = BitsandbytesPrecision(quantize[4:], dtype)
             precision = None
 
-        fabric = L.Fabric(
-            accelerator=accelerator,
-            devices=1,  # Otherwise sequential wouldn't work, see litgpt/generate/sequentially.py
-            # devices=devices,
-            precision=precision,
-            plugins=plugins
-        )
+        # set "ddp" as the strategy for the launching functionality, but there's no data-parallelism
+        if generate_strategy != "tensor_parallel":
+            fabric = L.Fabric(
+                accelerator=accelerator,
+                devices=1,  # Otherwise sequential wouldn't work, see litgpt/generate/sequentially.py
+                #devices=devices,
+                precision=precision,
+                plugins=plugins
+            )
+        else:
+            fabric = L.Fabric(
+                devices=devices,
+                strategy="ddp",
+                precision=precision,
+                plugins=plugins
+            )
+            if torch.cuda.is_available() and fabric.accelerator.auto_device_count() > 1:
+                check_nvlink_connectivity(fabric)
+                fabric.launch()
 
         self.kv_cache_initialized = False
         if generate_strategy is None:
@@ -272,11 +287,7 @@ def distribute(
                 self.kv_cache_initialized = True
                 self.fixed_kv_cache_size = fixed_kv_cache_size
 
-        elif generate_strategy == "sequential":
-            # cannot use `init_module` because if bitsandbytes is used, the Linear layers will be replaced
-            # which means that the weights will get quantized on cuda:0 on checkpoint load. we need to load and then convert
-            # still, use init_tensor for the precision
-
+        elif generate_strategy in ("sequential", "tensor_parallel"):
             total_devices = CUDAAccelerator.auto_device_count()
             if devices is not None:
                 if devices < total_devices:
@@ -288,20 +299,57 @@ def distribute(
 
             with fabric.init_tensor(), torch.device("meta"):
                 model = GPT(self.config)
-
             model.eval()
-            state_dict = torch.load(str(self.checkpoint_dir / "lit_model.pth"), mmap=True, map_location="cpu")
-            model.load_state_dict(state_dict, assign=True)
-            model = fabric.setup_module(model, move_to_device=False)
-
-            if fixed_kv_cache_size is None:
-                fixed_kv_cache_size = "max_model_supported"
-            if fixed_kv_cache_size == "max_model_supported":
-                kv_cache_size = model.max_seq_length
-            else:
-                kv_cache_size = fixed_kv_cache_size
-            model = sequential(model, fabric.device, kv_cache_size, total_devices)
-            self.fixed_kv_cache_size = fixed_kv_cache_size
+
+            if generate_strategy == "sequential":
+                state_dict = torch.load(str(self.checkpoint_dir / "lit_model.pth"), mmap=True, map_location="cpu")
+                model.load_state_dict(state_dict, assign=True)
+                model = fabric.setup_module(model, move_to_device=False)
+
+                if fixed_kv_cache_size is None:
+                    fixed_kv_cache_size = "max_model_supported"
+                if fixed_kv_cache_size == "max_model_supported":
+                    kv_cache_size = model.max_seq_length
+                else:
+                    kv_cache_size = fixed_kv_cache_size
+
+                model = sequential(model, fabric.device, kv_cache_size, total_devices)
+                self.fixed_kv_cache_size = fixed_kv_cache_size
+
+            elif generate_strategy == "tensor_parallel":
+                if fabric.global_rank == 0:
+                    pbar = tqdm(total=fabric.world_size, desc="Loading model weights")
+                for rank in range(fabric.world_size):
+                    if fabric.global_rank == rank:
+                        state_dict = torch.load(str(self.checkpoint_dir / "lit_model.pth"), mmap=True, map_location="cpu")
+                        model.load_state_dict(state_dict, assign=True)
+
+                        # cannot use `.setup_module` because it will wrap with DDP
+                        model = fabric._precision.convert_module(model)
+                        model = tensor_parallel(fabric, model)
+
+                        with fabric.init_tensor():
+                            if fixed_kv_cache_size is None:
+                                fixed_kv_cache_size = "max_model_supported"
+                            if fixed_kv_cache_size == "max_model_supported":
+                                kv_cache_size = model.max_seq_length
+                            else:
+                                kv_cache_size = fixed_kv_cache_size
+                            model.max_seq_length = kv_cache_size
+                            # the rope cache which is on meta device
+                            model.cos, model.sin = model.rope_cache()
+                            # enable the kv cache
+                            model.set_kv_cache(batch_size=1)
+                        model.eval()
+                        model = fabric.to_device(model)
+
+                    fabric.barrier()
+                    if fabric.global_rank == 0:
+                        pbar.update(1)
+
+                if fabric.global_rank == 0:
+                    pbar.close()
+
             self.kv_cache_initialized = True
 
         else:

litgpt/generate/tp.py

@@ -19,7 +19,9 @@
 from torch.distributed._functional_collectives import all_reduce
 
 import litgpt.generate.base as generate_base
-from litgpt import GPT, Config, Tokenizer
+from litgpt.model import GPT
+from litgpt.config import Config
+from litgpt.tokenizer import Tokenizer
 from litgpt.model import CausalSelfAttention, GptNeoxMLP, LLaMAMLP, LLaMAMoE
 from litgpt.prompts import PromptStyle, has_prompt_style, load_prompt_style
 from litgpt.utils import (

tests/test_api.py

@@ -1,6 +1,6 @@
 from pathlib import Path
 
-
+import os
 import pytest
 import re
 import torch
@@ -137,39 +137,46 @@ def test_model_not_initialized(tmp_path):
 
 
 @RunIf(min_cuda_gpus=2)
-def test_more_than_1_device_for_sequential_gpu(tmp_path):
+def test_more_than_1_device_for_sequential_tp_gpu(tmp_path):
     llm = LLM.load(
         model="EleutherAI/pythia-14m",
     )
 
     llm.distribute(devices=2, generate_strategy="sequential")
     assert isinstance(llm.generate("What do llamas eat?"), str)
 
-    with pytest.raises(NotImplementedError, match="Support for multiple devices is currently only implemented for generate_strategy='sequential'."):
+    if os.getenv("CI") != "true":
+        # this crashes the CI, maybe because of process forking; works fien locally though
+        llm.distribute(devices=2, generate_strategy="tensor_parallel")
+        assert isinstance(llm.generate("What do llamas eat?"), str)
+
+    with pytest.raises(NotImplementedError, match=f"Support for multiple devices is currently only implemented for generate_strategy='sequential'|'tensor_parallel'."):
         llm.distribute(devices=2)
 
 
 @RunIf(min_cuda_gpus=1)
-def test_sequential_incompatibility_with_random_weights(tmp_path):
+def test_sequential_tp_incompatibility_with_random_weights(tmp_path):
     llm = LLM.load(
         model="EleutherAI/pythia-14m",
         tokenizer_dir="EleutherAI/pythia-14m",
         init="random"
     )
-    with pytest.raises(NotImplementedError, match=re.escape("The LLM was initialized with init='random' but .distribute() currently only supports pretrained weights.")):
-        llm.distribute(devices=1, generate_strategy="sequential")
+    for strategy in ("sequential", "tensor_parallel"):
+        with pytest.raises(NotImplementedError, match=re.escape("The LLM was initialized with init='random' but .distribute() currently only supports pretrained weights.")):
+            llm.distribute(devices=1, generate_strategy=strategy)
 
 
-def test_sequential_cpu(tmp_path):
+def test_sequential_tp_cpu(tmp_path):
     llm = LLM.load(
         model="EleutherAI/pythia-14m",
     )
-    with pytest.raises(NotImplementedError, match="generate_strategy='sequential' is only supported for accelerator='cuda'."):
-        llm.distribute(
-            devices=1,
-            accelerator="cpu",
-            generate_strategy="sequential"
-        )
+    for strategy in ("sequential", "tensor_parallel"):
+        with pytest.raises(NotImplementedError, match=f"generate_strategy='{strategy}' is only supported for accelerator='cuda'|'gpu'."):
+            llm.distribute(
+                devices=1,
+                accelerator="cpu",
+                generate_strategy=strategy
+            )
 
 
 @RunIf(min_cuda_gpus=1)
@@ -185,7 +192,7 @@ def test_fixed_kv_cache(tmp_path):
     llm = LLM.load(
         model="EleutherAI/pythia-14m",
     )
-    llm.distribute(devices=1, fixed_kv_cache_size=100)    
+    llm.distribute(devices=1, fixed_kv_cache_size=100)
 
     # Request too many tokens
     with pytest.raises(NotImplementedError, match="max_seq_length 512 needs to be >= 9223372036854775809"):

tutorials/python-api.md

@@ -92,7 +92,11 @@ llm = LLM.load("pythia-160m", init="random", tokenizer_dir="EleutherAI/pythia-16
  
 ## Multi-GPU strategies
 
-By default, the model is loaded onto a single GPU. Optionally, you can use the `.distribute()` method with the `generate_strategy="sequential"` setting to load different parts of the models onto different GPUs. The goal behind this strategy is to support models that cannot fit into single-GPU memory. (Note that if you have a model that can fit onto a single GPU, this sequential strategy will be slower.)
+By default, the model is loaded onto a single GPU. Optionally, you can use the `.distribute()` method with the "sequential" or "tensor_parallel" `generate_strategy` settings.
+
+### Sequential strategy
+
+the `generate_strategy="sequential"` setting to load different parts of the models onto different GPUs. The goal behind this strategy is to support models that cannot fit into single-GPU memory. (Note that if you have a model that can fit onto a single GPU, this sequential strategy will be slower.)
 
 ```python
 from litgpt.api import LLM
@@ -124,6 +128,35 @@ print(text)
  Llamas are herbivores and their diet consists mainly of grasses, plants, and leaves.
 ```
 
+ 
+### Tensor parallel strategy
+
+The sequential strategy explained in the previous subsection distributes the model sequentially across GPUs, which allows users to load models that would not fit onto a single GPU. However, due to this method's sequential nature, processing is naturally slower than parallel processing. 
+
+To take advantage of parallel processing via tensor parallelism, you can use the `generate_strategy="tensor_parallel" setting. However, this method has downsides: the initial setup may be slower for large models, and it cannot run in interactive processes such as Jupyter notebooks.
+
+```python
+from litgpt.api import LLM
+
+
+if __name__ == "__main__":
+
+    llm = LLM.load(
+        model="meta-llama/Meta-Llama-3.1-8B-Instruct",
+        distribute=None
+    )
+
+    llm.distribute(generate_strategy="tensor_parallel", devices=4)
+
+    print(llm.generate(prompt="What do llamas eat?"))
+    print(llm.generate(prompt="What is 1+2?", top_k=1))
+```
+
+```
+
+```
+
+
  
 ## Speed and resource estimates
 
@@ -153,4 +186,4 @@ pprint(bench_d)
 #  'Seconds total': 1.5935972900006163,
 #  'Tokens generated': 25,
 #  'Total GPU memory allocated in GB': 11.534106624}
-```
+```