Sleep Framework: A New Way For LLMs To Retain Knowledge In Continual Learning

Welcome back to the blog! If you’ve been hanging around the local AI scene for a while, you already know that running Large Language Models (LLMs) on your own homelab hardware is incredibly rewarding. But let’s be real: it can also be incredibly frustrating when you try to teach your model something new.

You spend hours curating a dataset, configuring your training scripts, letting your GPU fans scream for a weekend, and finally… your model knows your new data perfectly. But suddenly, it has forgotten how to do basic math, it can’t write a coherent Python script anymore, and it speaks in broken English.

In the AI world, this phenomenon is known as catastrophic forgetting.

Today, we are going to look at an emerging solution inspired directly by human biology: the Sleep Framework. We’ll break down what it is, how it works under the hood, what hardware you need to run it in your homelab, and how to actually implement a basic version of it using Python.

The Nightmare of Catastrophic Forgetting

Before we look at the cure, we need to understand the disease.

When you fine-tune an LLM, you are tweaking the mathematical weights inside the neural network. Every weight (or parameter) holds a tiny piece of the knowledge the model learned during its initial pre-training. When you fine-tune the model on a narrow dataset—say, a few thousand documents about your company’s internal IT policies—the optimizer aggressively shifts the weights to minimize the error on this new data.

The problem? Those same weights were previously responsible for the model’s understanding of grammar, logic, and general world knowledge. By forcefully shifting them to learn your IT policies, you overwrite the old knowledge. The model “forgets.”

Continual learning is the holy grail of AI: the ability to learn new tasks sequentially without forgetting previous ones. Humans do this effortlessly. You can learn how to use a new piece of software today without forgetting how to drive a car or speak English. How do we give our LLMs this same superpower?

Enter the Sleep Framework

The Sleep Framework is a continual learning methodology inspired by the Complementary Learning Systems (CLS) theory in neuroscience.

In the human brain, we have the hippocampus (responsible for fast, short-term learning) and the neocortex (responsible for slow, long-term knowledge storage). During the day, we learn new things rapidly, stored temporarily in the hippocampus. At night, during sleep, our brain “replays” these new memories alongside old ones, slowly transferring and integrating them into the neocortex without overwriting existing knowledge.

The Sleep Framework applies this to LLMs by alternating between two distinct phases:

1. The Wake Phase

During the “Wake” phase, the model is exposed to new data. It acts like the hippocampus. The model processes the new information, but instead of permanently baking it into the core weights, it stores the new representations in a temporary buffer or a set of low-rank adapters (like LoRA). The core model remains relatively stable.

2. The Sleep Phase

After the Wake phase, the model goes to “Sleep.” During this phase, the model doesn’t just sit idle. It engages in memory consolidation. It does this by:

  • Replaying old data: The model is briefly trained on a “replay buffer”—a small, random sample of the broad, general data it was originally trained on.
  • Synaptic Consolidation (EWC): Using techniques like Elastic Weight Consolidation, the model identifies which weights are most important for its old knowledge and penalizes the optimizer if it tries to change them too much during future wake phases.

By alternating between learning new stuff (Wake) and solidifying old stuff (Sleep), the model achieves a balance. It learns your specific data, but the “sleep” cycles act as a shock absorber, preventing catastrophic forgetting.

Hardware Requirements: Provisioning Your Homelab for the Sleep Framework

Since this is Self-Hosted Tech, we have to talk about the iron. Continual learning with the Sleep Framework is more resource-intensive than a standard one-off fine-tune because you are managing replay buffers and running additional consolidation passes. Here is what you need in your homelab rack:

  • GPU VRAM: This is the big one. If you are using LoRA (Low-Rank Adaptation) with an 8B parameter model (like Llama 3 8B), you can get away with 16GB of VRAM. However, if you want to apply Elastic Weight Consolidation—which requires calculating and storing the Fisher Information Matrix for the weights to know which ones to protect—you will need significantly more memory. For an 8B model, plan for 24GB of VRAM (an RTX 3090 or 4090 is perfect here). For larger models, you’re looking at multiple GPUs or enterprise cards with 40GB+ of VRAM.
  • System RAM: The replay buffer lives in system memory before being pushed to the GPU. For a decent replay buffer of general data (say, 10,000 mixed text samples), you’ll want at least 64GB of DDR4 or DDR5 ECC RAM in your server.
  • Storage (NVMe): The Sleep phase requires pulling random samples of old data from your drives while simultaneously streaming new data. A slow HDD will bottleneck your data loaders, leaving your GPU sitting at 10% utilization waiting for the next batch. You absolutely need a fast NVMe SSD (Gen 4 or Gen 5) to keep those pipelines saturated.
  • Cooling and Power: Continual learning means your GPUs are running hot for much longer than a standard fine-tune. Ensure your homelab has adequate airflow. If you are running multiple blower-style enterprise GPUs, consider setting custom fan curves or undervolting to prevent thermal throttling during long training runs.

Practical Walkthrough: Implementing a Sleep Cycle in Python

Let’s get our hands dirty. Below is a simplified, practical walkthrough of how you might structure a Sleep Framework training loop using Hugging Face’s transformers and peft libraries.

In this example, we will use a LoRA adapter for the “Wake” phase (learning new data) and a small replay buffer for the “Sleep” phase (preventing forgetting).

Prerequisites

First, make sure you have the right libraries installed in your Python environment:

1
pip install torch transformers peft datasets accelerate

The Code Structure

We’re going to set up two datasets: new_data (your specific homelab documentation, for example) and replay_buffer (a small subset of a general dataset like OpenAssistant or Wikipedia).

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments
from peft import LoraConfig, get_peft_model
from datasets import load_dataset, concatenate_datasets

# 1. Load the Model and Tokenizer
model_id = "meta-llama/Meta-Llama-3-8B"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    device_map="auto" # Let accelerate map this to your GPUs
)

# 2. Setup LoRA (Our "Hippocampus" for the Wake Phase)
lora_config = LoraConfig(
    r=16,
    lora_alpha=32,
    target_modules=["q_proj", "v_proj"],
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM"
)
model = get_peft_model(model, lora_config)

# 3. Load Datasets
# new_data: Your specific fine-tuning data
new_data = load_dataset("json", data_files="my_homelab_docs.json")["train"]

# replay_buffer: General data to prevent forgetting
# In reality, you'd want a few thousand diverse samples here
replay_buffer = load_dataset("wikitext", "wikitext-2-raw-v1")["train"].select(range(500))

# Tokenize datasets (simplified for brevity)
def tokenize_function(examples):
    return tokenizer(examples["text"], truncation=True, padding="max_length", max_length=256)

new_data = new_data.map(tokenize_function, batched=True)
replay_buffer = replay_buffer.map(tokenize_function, batched=True)

# 4. The Sleep Framework Loop
num_cycles = 3

for cycle in range(num_cycles):
    print(f"--- Starting Cycle {cycle + 1} ---")
    
    # WAKE PHASE: Train on new data
    print("WAKE PHASE: Learning new homelab data...")
    wake_args = TrainingArguments(
        output_dir=f"./results_cycle_{cycle}_wake",
        num_train_epochs=1,
        per_device_train_batch_size=4,
        save_steps=500,
        logging_steps=100,
    )
    
    wake_trainer = Trainer(
        model=model,
        args=wake_args,
        train_dataset=new_data,
    )
    wake_trainer.train()
    
    # SLEEP PHASE: Replay old data to consolidate
    print("SLEEP PHASE: Replaying general knowledge to prevent forgetting...")
    sleep_args = TrainingArguments(
        output_dir=f"./results_cycle_{cycle}_sleep",
        num_train_epochs=1, # Usually a shorter pass
        per_device_train_batch_size=4,
        save_steps=500,
        logging_steps=100,
        learning_rate=2e-5, # Lower learning rate during sleep to gently nudge weights
    )
    
    sleep_trainer = Trainer(
        model=model,
        args=sleep_args,
        train_dataset=replay_buffer,
    )
    sleep_trainer.train()

# Save the final adapter
model.save_pretrained("./sleep_framework_lora_adapter")

What’s Happening Here?

  1. LoRA Setup: We use LoRA to train adapters rather than the full model. This is much friendlier on homelab VRAM.
  2. Wake Phase: The trainer runs on new_data. The LoRA weights shift to learn your specific data.
  3. Sleep Phase: The trainer switches to the replay_buffer. The learning rate is slightly lowered. The model is briefly reminded of general human language and logic. Because the gradients from the replay buffer pull the weights back toward general knowledge, it prevents the model from completely drifting into the narrow domain of your new data.
  4. Cycles: We repeat this process. The alternating pattern mimics the biological sleep cycle, resulting in an adapter that knows your new data but retains the base model’s reasoning skills.

Pro-tip for your homelab: While running this, keep a terminal open with watch -n 1 nvidia-smi. Watch your VRAM usage during the transition between Wake and Sleep phases. If your replay buffer batches are larger than your new data batches, you might spike your VRAM and trigger an Out Of Memory (OOM) error. Keep your batch sizes consistent!

Tuning and Monitoring

The Sleep Framework isn’t a magic button; it requires tuning. If your model is still forgetting things, you need to increase the size of your replay_buffer or increase the number of epochs in the Sleep phase. If your model isn’t learning your new data well enough, your Wake phase needs more epochs or a higher learning rate.

From a hardware perspective, if you find your GPU utilization dropping to 0% between the Wake and Sleep phases, your system is likely bottlenecked by disk I/O or CPU preprocessing. Consider using a library like datasets with streaming enabled, or move your pre-tokenized replay buffer entirely into system RAM before starting the loop.

Summary and Call to Action

Catastrophic forgetting has long been the bane of anyone trying to customize local LLMs in their homelab. The Sleep Framework offers a biologically inspired, highly practical way to achieve continual learning. By alternating between a “Wake” phase—where the model absorbs new data—and a “Sleep” phase—where it replays old data to consolidate knowledge—you can fine-tune your models without destroying their foundational reasoning.

As long as your homelab is equipped with sufficient VRAM (24GB+ for 8B models), fast NVMe storage, and adequate cooling, you can start experimenting with continual learning today.

What do you think? Are you going to try implementing a sleep cycle in your next local AI fine-tuning project? Have you dealt with catastrophic forgetting in your own homelab? Let us know down in the comments, and don’t forget to check out the corresponding video on the Self-Hosted Tech YouTube channel. If you found this guide helpful, hit that subscribe button and drop a comment with your GPU specs—we love talking homelab hardware!