feat: CPU-to-GPU streaming quantization for minimal peak GPU memory

Load HF model on CPU, then stream weights to GPU one layer at a time
during quantization. Each layer's fp16 weights are moved to GPU,
quantized into packed kbit format, then the CPU copy is freed. This
keeps peak GPU memory at ~1 layer of fp16 + growing quantized data,
instead of the entire model.

- KbitLoraModel: add target_device parameter for CPU->GPU streaming
- _quantize_weight: move weight to target_device, free after quantize
- _quantize_and_create_lora: in streaming mode, free each layer after
  processing (replaces with empty nn.Module)
- Only free layers when target_device is explicitly set (streaming mode);
  non-streaming mode preserves the original model for reuse
- train_qlora.py: load model on CPU, pass target_device=cuda
- train_pipeline.py: same CPU loading approach per rank

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: TimDettmers

bitsandbytes/kbit_lora.py

@@ -56,6 +56,10 @@ class KbitLoraModel(nn.Module):
             Set False for non-first pipeline stages.
         include_lm_head: Whether to quantize and keep the LM head. Default True.
             Set False for non-last pipeline stages.
+        target_device: Device for quantized weights and LoRA params. If None,
+            uses the model's device. Set this when loading the HF model on
+            CPU to stream weights to GPU one layer at a time (minimizes
+            peak GPU memory). Example: torch.device("cuda:0").
     """
 
     def __init__(
@@ -73,6 +77,7 @@ def __init__(
         layer_range: Optional[tuple[int, int]] = None,
         include_embed: bool = True,
         include_lm_head: bool = True,
+        target_device: Optional[torch.device] = None,
     ):
         super().__init__()
 
@@ -123,10 +128,20 @@ def __init__(
             self._layer_start, self._layer_end = 0, total_layers
         self._num_loaded_layers = self._layer_end - self._layer_start
 
+        # Determine target device for quantized weights.
+        # When target_device is explicitly set (streaming mode), we free each
+        # layer from the source model after quantization to save memory.
+        self._streaming = target_device is not None
+        if target_device is not None:
+            self._target_device = target_device
+        else:
+            self._target_device = next(model.parameters()).device
+
         # Keep reference to original model for embeddings
         self.model = model
         if include_embed:
-            self.embed_tokens = model.model.embed_tokens
+            # Move embedding to target device (may be CPU->GPU transfer)
+            self.embed_tokens = model.model.embed_tokens.to(self._target_device)
         else:
             self.embed_tokens = None
         self.lm_head_tied = hasattr(model, "lm_head") and (
@@ -140,7 +155,7 @@ def __init__(
 
         self._quantize_and_create_lora(model)
 
-        # Freeze all base model parameters
+        # Freeze all base model parameters (any that remain)
         for p in model.parameters():
             p.requires_grad_(False)
 
@@ -151,19 +166,27 @@ def __init__(
             p.requires_grad_(True)
 
     def _quantize_weight(self, weight: torch.Tensor, name: str, k: int | None = None):
-        """Quantize a weight matrix and store packed data."""
+        """Quantize a weight matrix and store packed data.
+
+        The weight is moved to _target_device for quantization (CUDA kernel),
+        then the original weight reference is no longer needed.
+        """
         if k is None:
             k = self.k
+        # Move to target device for quantization (CPU -> GPU transfer if needed)
+        weight = weight.to(self._target_device)
         N, K = weight.shape
         N_padded = ((N + 127) // 128) * 128
         if N_padded != N:
             w_padded = torch.nn.functional.pad(weight.float(), (0, 0, 0, N_padded - N))
         else:
             w_padded = weight.float()
+        del weight  # Free the fp16 copy on GPU
 
         packed, absmax, codebook = F.quantize_kbit(
             w_padded.reshape(-1), k=k, absmax_format="fp32",
         )
+        del w_padded  # Free the fp32 padded copy
 
         # Store as non-trainable buffers
         safe_name = name.replace(".", "_")
@@ -173,9 +196,10 @@ def _quantize_weight(self, weight: torch.Tensor, name: str, k: int | None = None
 
         return packed, absmax, codebook, N_padded, N, K
 
-    def _create_lora(self, name: str, N: int, K: int, device: torch.device):
-        """Create LoRA A and B parameters for a weight matrix."""
+    def _create_lora(self, name: str, N: int, K: int):
+        """Create LoRA A and B parameters for a weight matrix on _target_device."""
         safe_name = name.replace(".", "_")
+        device = self._target_device
         # A: [r, K] initialized with Kaiming uniform
         A = nn.Parameter(torch.empty(self.lora_r, K, dtype=self.compute_dtype, device=device))
         nn.init.kaiming_uniform_(A, a=math.sqrt(5))
@@ -190,8 +214,13 @@ def _quantize_and_create_lora(self, model: nn.Module):
 
         Only processes layers in [_layer_start, _layer_end) and optionally
         skips embedding and LM head for pipeline parallelism.
+
+        Streams weights one layer at a time: each layer's weights are moved
+        from the model's device (often CPU) to _target_device (GPU), quantized,
+        then the original layer is deleted. This keeps peak GPU memory at
+        ~1 layer of fp16 weights plus the growing quantized data.
         """
-        device = next(model.parameters()).device
+        device = self._target_device
 
         # Process only the decoder layers in our range
         layers = model.model.layers
@@ -207,12 +236,12 @@ def _quantize_and_create_lora(self, model: nn.Module):
 
             # Attention projections (use k_attention)
             for proj_name in ["q_proj", "k_proj", "v_proj", "o_proj"]:
-                weight = getattr(attn, proj_name).weight.data.to(device)
+                weight = getattr(attn, proj_name).weight.data
                 name = f"{prefix}_attn_{proj_name}"
                 packed, absmax, codebook, N_padded, N, K = self._quantize_weight(
                     weight, name, k=self.k_attention,
                 )
-                A, B = self._create_lora(name, N, K, device)
+                A, B = self._create_lora(name, N, K)
                 layer_info[proj_name] = {
                     "packed": packed, "absmax": absmax, "codebook": codebook,
                     "N_padded": N_padded, "N": N, "K": K, "A": A, "B": B,
@@ -221,24 +250,24 @@ def _quantize_and_create_lora(self, model: nn.Module):
 
             # MLP projections (use k_mlp)
             for proj_name in ["gate_proj", "up_proj", "down_proj"]:
-                weight = getattr(mlp, proj_name).weight.data.to(device)
+                weight = getattr(mlp, proj_name).weight.data
                 name = f"{prefix}_mlp_{proj_name}"
                 packed, absmax, codebook, N_padded, N, K = self._quantize_weight(
                     weight, name, k=self.k_mlp,
                 )
-                A, B = self._create_lora(name, N, K, device)
+                A, B = self._create_lora(name, N, K)
                 layer_info[proj_name] = {
                     "packed": packed, "absmax": absmax, "codebook": codebook,
                     "N_padded": N_padded, "N": N, "K": K, "A": A, "B": B,
                     "k": self.k_mlp,
                 }
 
-            # Norm weights (trainable, not quantized)
+            # Norm weights (trainable, not quantized) — move to target device
             for norm_name in ["input_layernorm", "post_attention_layernorm"]:
                 norm = getattr(layer, norm_name)
                 safe = f"{prefix}_{norm_name}_weight"
                 self._norm_weights[safe] = nn.Parameter(
-                    norm.weight.data.to(self.compute_dtype).clone()
+                    norm.weight.data.to(device=device, dtype=self.compute_dtype).clone()
                 )
                 layer_info[norm_name] = self._norm_weights[safe]
 
@@ -248,23 +277,31 @@ def _quantize_and_create_lora(self, model: nn.Module):
                     norm = getattr(attn, norm_name)
                     safe = f"{prefix}_attn_{norm_name}_weight"
                     self._norm_weights[safe] = nn.Parameter(
-                        norm.weight.data.to(self.compute_dtype).clone()
+                        norm.weight.data.to(device=device, dtype=self.compute_dtype).clone()
                     )
                     layer_info[norm_name] = self._norm_weights[safe]
 
             self._layer_data.append(layer_info)
 
+            # In streaming mode, free each layer from the source model
+            # after quantization to release memory (typically CPU RAM).
+            if self._streaming:
+                layers[i] = nn.Module()
+                del layer
+                if device.type == "cuda":
+                    torch.cuda.empty_cache()
+
         # Final norm (only needed by last stage or full model)
         if self.include_lm_head:
             final_norm = model.model.norm
             self._norm_weights["final_norm_weight"] = nn.Parameter(
-                final_norm.weight.data.to(self.compute_dtype).clone()
+                final_norm.weight.data.to(device=device, dtype=self.compute_dtype).clone()
             )
 
         # LM head (only needed by last stage or full model)
         self._lm_head_info = None
         if self.include_lm_head:
-            lm_weight = model.lm_head.weight.data.to(device)
+            lm_weight = model.lm_head.weight.data
             name = "lm_head"
             packed, absmax, codebook, N_padded, N, K = self._quantize_weight(
                 lm_weight, name, k=self.k_lm_head,

examples/train_pipeline.py

@@ -144,8 +144,9 @@ def main():
         print(f"Steps: {args.steps}")
         print()
 
-    # Load model — each rank loads the full HF model temporarily to extract
-    # its layer weights. We immediately delete the original after quantization.
+    # Load model on CPU, then stream weights to GPU layer by layer.
+    # This avoids the full model ever being on GPU — peak GPU memory is
+    # just ~1 fp16 layer at a time plus the growing quantized data.
     from transformers import AutoModelForCausalLM, AutoConfig
 
     config = AutoConfig.from_pretrained(args.model, trust_remote_code=True)
@@ -158,20 +159,21 @@ def main():
     print(f"  GPU {rank}: layers {layer_start}-{layer_end-1} ({role} stage)")
 
     if rank == 0:
-        print(f"\nLoading and quantizing (per-rank)...")
+        print(f"\nLoading HF model on CPU, streaming to GPU...")
     torch.cuda.reset_peak_memory_stats()
 
+    # Load on CPU — no GPU memory used yet
     model = AutoModelForCausalLM.from_pretrained(
         args.model,
         dtype=torch.float16,
-        device_map={"": device},
+        device_map="cpu",
         trust_remote_code=True,
     )
 
-    mem_after_load = torch.cuda.memory_allocated() / 1024 / 1024
-    print(f"  GPU {rank}: {mem_after_load:.0f} MB after HF model load")
+    mem_before = torch.cuda.memory_allocated() / 1024 / 1024
+    print(f"  GPU {rank}: {mem_before:.0f} MB after HF model load (model on CPU)")
 
-    # Create KbitLoraModel with ONLY this rank's layers
+    # Create KbitLoraModel: streams weights CPU->GPU one layer at a time
     kbit_model = KbitLoraModel(
         model,
         lora_r=args.lora_r,
@@ -181,15 +183,17 @@ def main():
         layer_range=(layer_start, layer_end),
         include_embed=is_first,
         include_lm_head=is_last,
+        target_device=device,
     )
 
-    # Delete the original HF model to free memory
+    # Delete the original HF model to free CPU memory
     del model
-    torch.cuda.empty_cache()
 
     mem_after_quant = torch.cuda.memory_allocated() / 1024 / 1024
-    print(f"  GPU {rank}: {mem_after_quant:.0f} MB after quantize + cleanup "
-          f"({kbit_model._num_loaded_layers} layers, "
+    peak_during_quant = torch.cuda.max_memory_allocated() / 1024 / 1024
+    print(f"  GPU {rank}: {mem_after_quant:.0f} MB after quantize "
+          f"(peak during load: {peak_during_quant:.0f} MB, "
+          f"{kbit_model._num_loaded_layers} layers, "
           f"embed={'yes' if is_first else 'no'}, "
           f"lm_head={'yes' if is_last else 'no'})")
 

examples/train_qlora.py

@@ -293,20 +293,22 @@ def main():
         if tokenizer.pad_token_id is None:
             tokenizer.pad_token_id = tokenizer.eos_token_id
 
-    # Load base model
-    print("Loading base model...")
+    # Load base model on CPU, then stream weights to GPU during quantization.
+    # This keeps peak GPU memory minimal — only 1 layer of fp16 weights at a
+    # time on GPU, plus the growing quantized data.
+    print("Loading base model on CPU...")
     t0 = time.time()
     model = AutoModelForCausalLM.from_pretrained(
         args.model,
         dtype=torch.float16,
-        device_map="cuda",
+        device_map="cpu",
         trust_remote_code=True,
     )
     print(f"  Loaded in {time.time() - t0:.1f}s")
-    print(f"  GPU memory after load: {get_gpu_memory_mb():.0f} MB")
+    print(f"  GPU memory after load: {get_gpu_memory_mb():.0f} MB (model on CPU)")
 
-    # Apply KbitLoraModel
-    print("\nQuantizing and creating LoRA adapters...")
+    # Apply KbitLoraModel — streams weights CPU->GPU one layer at a time
+    print("\nQuantizing and streaming to GPU...")
     t0 = time.time()
     kbit_model = KbitLoraModel(
         model,
@@ -318,15 +320,15 @@ def main():
         ce_chunk_size=args.ce_chunk,
         compute_dtype=torch.bfloat16,
         cpu_offload=args.cpu_offload,
+        target_device=torch.device("cuda"),
     )
     print(f"  Quantized in {time.time() - t0:.1f}s")
     print(f"  Trainable parameters: {kbit_model.num_trainable_parameters():,}")
     print(f"  GPU memory after quantization: {get_gpu_memory_mb():.0f} MB")
+    print(f"  Peak GPU memory during load: {get_gpu_peak_mb():.0f} MB")
 
-    # Free the original model weights (they're now quantized)
+    # Free the original model (CPU memory)
     del model
-    torch.cuda.empty_cache()
-    print(f"  GPU memory after cleanup: {get_gpu_memory_mb():.0f} MB")
 
     # Prepare dataset
     if not args.synthetic: