neuron-tai/.scratch/distributed-inference-network/issues/12-real-pytorch-model-backend.md

US-012 — Real PyTorch model backend

Replace stub node inference with actual transformers layer execution. A node downloads a HuggingFace SafeTensors model, loads model.layers[start:end] into GPU/CPU memory at the declared quantization level, and runs real forward passes. The first-layer node additionally runs embed_tokens; the last-layer node additionally runs model.norm and lm_head.

Context

Per ADR-0001 (PyTorch over llama.cpp) and the grilling session decisions:

• Model format: HuggingFace SafeTensors (GGUF rejected — no per-layer hidden state API)
• Layer extraction: model.model.layers[start:end] from a loaded AutoModelForCausalLM
• Quantization: bitsandbytes NF4 / INT8 / bfloat16, declared by node at registration
• Activation dtype at boundaries: always bfloat16 regardless of weight quantization
• First-layer node owns: tokenizer, model.embed_tokens, model.layers[0..k]
• Last-layer node owns: model.layers[N-k..N], model.norm, model.lm_head

Start with a small real model for testing: openai-community/gpt2 (117M, no special license, SafeTensors available). The production target is 700B+ models — the layer extraction pattern is identical regardless of size.

Acceptance Criteria

• packages/node can be started with --model-id openai-community/gpt2 --shard-start 0 --shard-end 6 and loads the corresponding transformer blocks
• Node detects whether it is a head shard (shard_start == 0) and loads tokenizer + model.embed_tokens
• Node detects whether it is a tail shard (shard_end == total_layers) and loads model.norm + model.lm_head
• A /forward request with a text prompt (head node) or binary bfloat16 tensor (mid/tail node) runs a real forward pass and returns binary bfloat16 activations
• A two-node integration test with GPT-2 sharded across two nodes returns a coherent (non-random, deterministic) text completion
• Quantization level is read from --quantization [bfloat16|int8|nf4] flag; bfloat16 is default; int8 and nf4 use bitsandbytes.BitsAndBytesConfig
• A node with insufficient VRAM for the assigned shard at the declared quantization prints a clear error and exits (no silent OOM)
• transformers, bitsandbytes, safetensors, and accelerate added as dependencies to packages/node
• python -m pytest passes (GPT-2 test downloads model once, cached; mark with @pytest.mark.integration and skip in CI if no GPU)
• Commit only this story's changes

Implementation Notes

from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
import torch

def load_shard(model_id, start, end, quantization):
    quant_config = None
    if quantization == "nf4":
        quant_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4",
                                          bnb_4bit_compute_dtype=torch.bfloat16)
    elif quantization == "int8":
        quant_config = BitsAndBytesConfig(load_in_8bit=True)

    # Load full model metadata without weights, then extract layers
    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        quantization_config=quant_config,
        device_map="auto",
        torch_dtype=torch.bfloat16,
    )
    layers = model.model.layers[start:end]
    return model, layers

• device_map="auto" handles GPU/CPU offload automatically based on available VRAM
• For large models, use low_cpu_mem_usage=True to avoid doubling peak RAM during load
• The forward pass for a mid-shard node: hidden_states = layer(hidden_states)[0] (most HF models return a tuple, take index 0)
• Head node forward: input_ids → embed_tokens → position_embed → layers[0..k] → bfloat16 binary out
• Tail node forward: bfloat16 binary in → layers[N-k..N] → norm → lm_head → argmax/sample → token_id → back to head
• Attention mask and position IDs must be forwarded alongside hidden states for models that need them (pack them in additional headers or a small JSON sidecar in the same HTTP request)
• Keep attention/position passing simple for v1: assume left-padded inputs, pass attention_mask as a separate binary blob with header X-Meshnet-Attn-Mask: <base64-uint8> (small enough for base64 — it's [B, S] not [B, S, D])