Inițializarea modelelor uriașe pe meta device: construiești arhitectura fără să umpli RAM-ul

Când un model are miliarde de parametri, problema nu începe doar la inferență. Uneori se rupe chiar la inițializare: framework-ul construiește modulele, alocă greutăți aleatorii în RAM sau VRAM, apoi le înlocuiește cu greutățile reale din checkpoint. Pentru modele mari, această dublă alocare poate însemna CUDA out of memory sau un proces omorât de sistem înainte să ajungi la primul token.

În acest tutorial construim un proiect complet, mic ca dimensiune, dar identic ca mecanism cu fluxurile folosite la modele mari: definim o arhitectură Transformer, salvăm un checkpoint, reconstruim modelul pe meta device, îl materializăm pe CPU fără inițializare inutilă, încărcăm greutățile reale și rulăm inferența. Dacă lucrezi deja cu probleme de memorie pe GPU, merită citit și ghidul despre CUDA Memory Fragmentation în PyTorch, pentru că meta device rezolvă o problemă diferită: inițializarea și încărcarea, nu fragmentarea allocatorului CUDA.

Ce construim

Construim un mini-model Transformer pentru clasificarea unor secvențe de tokeni. Rezultatul final este un director meta_model_demo/ cu:

model_def.py, definiția arhitecturii;
build_checkpoint.py, scriptul care creează configurația și checkpoint-ul;
load_with_meta.py, scriptul care reconstruiește modelul pe meta, îl materializează și rulează inferența;
artifacts/tiny_transformer.pt, checkpoint-ul real;
artifacts/model_config.json, configurația folosită la reconstrucție.

Modelul este mic ca să fie reproductibil pe laptop, dar fluxul este cel corect pentru modele mari: configurație înainte, model pe meta, materializare controlată, load_state_dict, verificare, inferență.

Mediul exact

Tutorialul este scris pentru:

Ubuntu Server 24.04 LTS;
Python 3.11;
PyTorch 2.4.1 CPU;
psutil 6.0.0;
minimum 4 GB RAM;
fără GPU obligatoriu;
timp estimat: 20-30 de minute.

Comenzile sunt pentru Linux. Pe Windows, rulează-le în WSL2 Ubuntu 24.04 dacă vrei același comportament.

Ce îți trebuie

Creează proiectul și mediul Python:

mkdir -p meta_model_demo/artifacts
cd meta_model_demo
python3.11 -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip
python -m pip install torch==2.4.1 --index-url https://download.pytorch.org/whl/cpu
python -m pip install psutil==6.0.0

Ce trebuie să vezi:

Successfully installed torch-2.4.1 ...
Successfully installed psutil-6.0.0

Verifică versiunile:

python - <<'PY'
import torch, psutil
print("torch:", torch.__version__)
print("psutil:", psutil.__version__)
PY

Ce trebuie să vezi:

torch: 2.4.1+cpu
psutil: 6.0.0

Sufixul +cpu poate apărea în funcție de roata instalată.

Pasul 1: definim arhitectura într-un fișier separat

Fișierul de model trebuie să poată fi importat atât la salvarea checkpoint-ului, cât și la încărcarea pe meta.

Creează model_def.py:

import torch
from torch import nn


class TinyTransformerClassifier(nn.Module):
    def __init__(
        self,
        vocab_size: int,
        hidden_size: int,
        num_heads: int,
        num_layers: int,
        num_classes: int,
        max_length: int,
    ) -> None:
        super().__init__()
        self.token_embedding = nn.Embedding(vocab_size, hidden_size)
        self.position_embedding = nn.Embedding(max_length, hidden_size)

        encoder_layer = nn.TransformerEncoderLayer(
            d_model=hidden_size,
            nhead=num_heads,
            dim_feedforward=hidden_size * 4,
            dropout=0.0,
            batch_first=True,
            activation="gelu",
        )
        self.encoder = nn.TransformerEncoder(
            encoder_layer=encoder_layer,
            num_layers=num_layers,
        )
        self.classifier = nn.Linear(hidden_size, num_classes)

    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
        batch_size, sequence_length = input_ids.shape
        positions = torch.arange(
            sequence_length,
            device=input_ids.device,
        ).unsqueeze(0).expand(batch_size, sequence_length)

        hidden = self.token_embedding(input_ids) + self.position_embedding(positions)
        encoded = self.encoder(hidden)
        pooled = encoded.mean(dim=1)
        return self.classifier(pooled)


def count_parameters(model: nn.Module) -> int:
    return sum(parameter.numel() for parameter in model.parameters())


def count_meta_parameters(model: nn.Module) -> int:
    return sum(1 for parameter in model.parameters() if parameter.device.type == "meta")

Ce trebuie să vezi după salvare:

python -m py_compile model_def.py

Dacă nu apare nimic, fișierul este valid sintactic.

Pasul 2: creăm configurația și checkpoint-ul

Acum construim o variantă reală a modelului, inițializăm determinist greutățile și salvăm doar state_dict. În proiecte mari, checkpoint-ul vine de obicei din antrenare, fine-tuning sau conversie. Dacă lucrezi cu date și secvențe lungi, optimizarea intră mai târziu în flux; de exemplu, Sequence Packing avansat ține de eficiența antrenării, nu de inițializarea arhitecturii.

Creează build_checkpoint.py:

import json
from pathlib import Path

import torch

from model_def import TinyTransformerClassifier, count_parameters


ARTIFACT_DIR = Path("artifacts")
CONFIG_PATH = ARTIFACT_DIR / "model_config.json"
CHECKPOINT_PATH = ARTIFACT_DIR / "tiny_transformer.pt"


def main() -> None:
    ARTIFACT_DIR.mkdir(exist_ok=True)

    config = {
        "vocab_size": 128,
        "hidden_size": 64,
        "num_heads": 4,
        "num_layers": 3,
        "num_classes": 2,
        "max_length": 16,
    }

    torch.manual_seed(47)
    model = TinyTransformerClassifier(**config)
    model.eval()

    with CONFIG_PATH.open("w", encoding="utf-8") as file:
        json.dump(config, file, indent=2)

    torch.save(model.state_dict(), CHECKPOINT_PATH)

    print(f"config={CONFIG_PATH}")
    print(f"checkpoint={CHECKPOINT_PATH}")
    print(f"parameters={count_parameters(model)}")
    print(f"first_parameter_device={next(model.parameters()).device}")


if __name__ == "__main__":
    main()

Rulează:

python build_checkpoint.py

Ce trebuie să vezi:

config=artifacts/model_config.json
checkpoint=artifacts/tiny_transformer.pt
parameters=156866
first_parameter_device=cpu

Verifică fișierele:

python - <<'PY'
from pathlib import Path
for path in ["artifacts/model_config.json", "artifacts/tiny_transformer.pt"]:
    p = Path(path)
    print(path, "exists=", p.exists(), "bytes=", p.stat().st_size)
PY

Ce trebuie să vezi:

artifacts/model_config.json exists= True bytes= ...
artifacts/tiny_transformer.pt exists= True bytes= ...

Numărul de bytes poate diferi ușor între versiuni, dar fișierele trebuie să existe.

Pasul 3: reconstruim modelul pe meta device

meta device păstrează metadatele tensorilor, cum ar fi forma și dtype-ul, fără date reale. Asta permite construirea structurii modelului fără alocarea inițială a greutăților. Nu poți copia direct un tensor meta pe CPU cu to("cpu"), pentru că nu există date de copiat; pentru module, PyTorch documentează metoda torch.nn.Module.to_empty(), care mută modulul pe dispozitivul ales lăsând parametrii neinițializați.

Creează load_with_meta.py:

import json
import os
from pathlib import Path

import psutil
import torch

from model_def import (
    TinyTransformerClassifier,
    count_meta_parameters,
    count_parameters,
)


CONFIG_PATH = Path("artifacts/model_config.json")
CHECKPOINT_PATH = Path("artifacts/tiny_transformer.pt")


def rss_mb() -> float:
    process = psutil.Process(os.getpid())
    return process.memory_info().rss / 1024 / 1024


def load_config() -> dict:
    with CONFIG_PATH.open("r", encoding="utf-8") as file:
        return json.load(file)


def main() -> None:
    print(f"rss_start_mb={rss_mb():.2f}")

    config = load_config()

    with torch.device("meta"):
        model = TinyTransformerClassifier(**config)

    print(f"after_meta_build_rss_mb={rss_mb():.2f}")
    print(f"parameters={count_parameters(model)}")
    print(f"meta_parameter_tensors={count_meta_parameters(model)}")
    print(f"first_parameter_device_before_to_empty={next(model.parameters()).device}")

    model.to_empty(device="cpu")
    print(f"first_parameter_device_after_to_empty={next(model.parameters()).device}")

    state_dict = torch.load(
        CHECKPOINT_PATH,
        map_location="cpu",
        weights_only=True,
    )
    missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=True)

    print(f"missing_keys={missing_keys}")
    print(f"unexpected_keys={unexpected_keys}")
    print(f"after_load_rss_mb={rss_mb():.2f}")

    model.eval()
    input_ids = torch.tensor(
        [
            [1, 5, 9, 13, 17, 21, 25, 29],
            [2, 4, 6, 8, 10, 12, 14, 16],
        ],
        dtype=torch.long,
    )

    with torch.inference_mode():
        logits = model(input_ids)
        predicted_classes = logits.argmax(dim=-1)

    rounded_logits = torch.round(logits * 10000) / 10000
    print(f"logits={rounded_logits.tolist()}")
    print(f"predicted_classes={predicted_classes.tolist()}")


if __name__ == "__main__":
    main()

Rulează:

python load_with_meta.py

Ce trebuie să vezi:

rss_start_mb=...
after_meta_build_rss_mb=...
parameters=156866
meta_parameter_tensors=43
first_parameter_device_before_to_empty=meta
first_parameter_device_after_to_empty=cpu
missing_keys=[]
unexpected_keys=[]
after_load_rss_mb=...
logits=[[...], [...]]
predicted_classes=[..., ...]

Valorile RSS diferă între sisteme. Elementele importante sunt first_parameter_device_before_to_empty=meta, apoi cpu, iar listele missing_keys și unexpected_keys trebuie să fie goale.

Pasul 4: verificăm că încărcarea este corectă

Verificarea trebuie să confirme trei lucruri: modelul construit pe meta are aceeași structură, checkpoint-ul s-a încărcat strict și inferența produce un tensor real, nu un tensor meta.

Rulează:

python - <<'PY'
import json
from pathlib import Path

import torch

from model_def import TinyTransformerClassifier

with Path("artifacts/model_config.json").open("r", encoding="utf-8") as file:
    config = json.load(file)

with torch.device("meta"):
    model = TinyTransformerClassifier(**config)

model.to_empty(device="cpu")
state_dict = torch.load(
    "artifacts/tiny_transformer.pt",
    map_location="cpu",
    weights_only=True,
)
missing, unexpected = model.load_state_dict(state_dict, strict=True)

sample = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8]], dtype=torch.long)

model.eval()
with torch.inference_mode():
    output = model(sample)

print("missing", missing)
print("unexpected", unexpected)
print("output_device", output.device)
print("output_shape", list(output.shape))
print("has_nan", torch.isnan(output).any().item())
PY

Ce trebuie să vezi:

missing []
unexpected []
output_device cpu
output_shape [1, 2]
has_nan False

Aceasta este verificarea minimă: arhitectura se potrivește cu checkpoint-ul, ieșirea există pe CPU și nu conține NaN.

De ce fluxul acesta contează la modele mari

Într-o inițializare obișnuită, constructorul modulului creează greutăți reale. Apoi checkpoint-ul le înlocuiește. Pentru modele mici nu observi problema. Pentru modele mari, constructorul poate consuma RAM sau VRAM înainte ca încărcarea eficientă să înceapă.

Cu meta, constructorul produce doar forma tensorilor. Apoi to_empty(device="cpu") creează spațiul real fără să facă inițializarea matematică obișnuită. În final, load_state_dict pune greutățile din checkpoint. Pentru un flux de fine-tuning, pasul acesta vine înainte de optimizări precum Data Mixture Engineering sau antrenare pe preferințe, cum este cazul unui Reward Model pentru preference learning.

Probleme frecvente

`NotImplementedError: Cannot copy out of meta tensor; no data!`

Cauză: ai încercat să muți modelul de pe meta cu model.to("cpu"). Tensorii meta nu au date reale de copiat.

Rezolvare:

model.to_empty(device="cpu")

Apoi încarci checkpoint-ul cu load_state_dict.

`RuntimeError: Error(s) in loading state_dict`

Cauză: configurația modelului nu se potrivește cu checkpoint-ul. De exemplu, ai schimbat hidden_size, num_layers sau vocab_size.

Rezolvare: folosește exact artifacts/model_config.json salvat împreună cu checkpoint-ul și rulează încărcarea strictă:

missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=True)

Listele trebuie să fie goale.

`TypeError: load() got an unexpected keyword argument 'weights_only'`

Cauză: rulezi o versiune PyTorch mai veche decât cea folosită în tutorial.

Rezolvare: verifică versiunea și instalează PyTorch 2.4.1:

python - <<'PY'
import torch
print(torch.__version__)
PY
python -m pip install torch==2.4.1 --index-url https://download.pytorch.org/whl/cpu

`AssertionError: embed_dim must be divisible by num_heads`

Cauză: hidden_size nu este divizibil cu num_heads în nn.TransformerEncoderLayer.

Rezolvare: păstrează configurația din tutorial, unde hidden_size este 64, iar num_heads este 4.

Revenire și curățenie

Dacă vrei să ștergi complet proiectul creat în tutorial, ieși din mediul virtual și șterge directorul:

deactivate
cd ..
rm -rf meta_model_demo

Ce trebuie să vezi:

Comanda nu produce output dacă a reușit. Atenție: rm -rf șterge definitiv directorul indicat. Ruleaz-o doar după ce ai verificat că ești în directorul părinte corect.

Mai departe

Următorul pas natural este să aplici același tipar pe un model mai mare: salvezi configurația separat, construiești arhitectura pe meta, materializezi pe CPU sau GPU și încarci checkpoint-ul strict. Pentru modele împărțite pe mai multe dispozitive, documentația Hugging Face Accelerate include utilitare precum init_empty_weights și load_checkpoint_and_dispatch, construite pentru același tip de problemă: inițializare fără alocarea imediată a tuturor greutăților.

Dacă optimizezi după încărcare, separă clar etapele: meta device pentru inițializare, snapshot-uri de memorie pentru OOM, sequence packing pentru eficiența datelor și benchmark-uri separate pentru kerneluri sau operații custom. Amestecarea lor într-un singur diagnostic face problema mai greu de verificat.

Întrebări frecvente

De ce nu folosim `model.to("cpu")` după construirea pe `meta`?

Pentru că tensorii de pe meta nu conțin date reale. În tutorial folosim model.to_empty(device="cpu"), apoi încărcăm greutățile reale din checkpoint.

Este nevoie de GPU pentru acest proiect?

Nu. Tutorialul folosește PyTorch 2.4.1 CPU și rulează pe Ubuntu Server 24.04 LTS cu minimum 4 GB RAM.

Cum știu că modelul chiar a fost construit pe `meta`?

Scriptul afișează first_parameter_device_before_to_empty=meta. După materializare trebuie să vezi first_parameter_device_after_to_empty=cpu.

Ce verificare arată că checkpoint-ul s-a potrivit cu modelul?

Încărcarea strictă trebuie să întoarcă missing_keys=[] și unexpected_keys=[]. În verificarea finală, ieșirea trebuie să fie pe CPU, cu forma [1, 2] și fără NaN.