Basic Distillation

This notebook demonstrates the full softs pipeline:

• Broker routes orders between clients and suppliers

• Supplier generates batches of soft labels

• Client trains a student model on those labels

We use random data here to keep it fast. For real LLM distillation, see the examples/distillation/ examples — e.g. logit_topp/ (broker.py / supplier.py / client.py).

import time
import torch
import torch.nn as nn

from softs import (
    Broker, Supplier, Client, ShmMedium,
    BatchConfig, TensorSpec, EndpointConfig,
)

1. Define the batch layout

Each slot holds one batch: x (inputs) and y (soft targets).

BATCH_SIZE = 8
HIDDEN = 256
NUM_CLASSES = 100

config = BatchConfig([
    TensorSpec("x", (BATCH_SIZE, HIDDEN), "float32"),
    TensorSpec("y", (BATCH_SIZE, NUM_CLASSES), "float32"),
])

print(f"Slot size: {config.nbytes() / 1e6:.1f} MB")

2. Start the broker

endpoints = EndpointConfig()
broker = Broker(endpoints=endpoints)
broker.start()
time.sleep(0.2)

3. Start a supplier

The supplier generates random soft labels. In practice, this would run a teacher model.

def teacher_generator(product_id: str) -> bytes:
    """Simulate teacher: random soft targets."""
    x = torch.randn(BATCH_SIZE, HIDDEN)
    logits = torch.randn(BATCH_SIZE, NUM_CLASSES)
    y = torch.softmax(logits, dim=-1)  # soft targets
    return config.encode(x=x, y=y)

supplier = Supplier(
    generator_fn=teacher_generator,
    product_ids=["teacher_v1"],
    endpoint=endpoints.backend,
    medium_cls=ShmMedium,
    slot_size=config.nbytes(),
)
supplier.start()
time.sleep(0.2)

4. Create a client and train

The client orders batches by product_id, reads them from shared memory, and trains a small student.

student = nn.Sequential(
    nn.Linear(HIDDEN, 128),
    nn.ReLU(),
    nn.Linear(128, NUM_CLASSES),
)
optimizer = torch.optim.Adam(student.parameters(), lr=1e-3)

client = Client(
    endpoint=endpoints.frontend,
    medium_cls=ShmMedium,
    slot_size=config.nbytes(),
    num_slots=4,
)
client.hello()

for step in range(10):
    slot = client.request_sample("teacher_v1", timeout_ms=5000)
    assert slot is not None, f"Timeout at step {step}"

    batch = config.decode(client.medium.read(slot))
    client.release_slot(slot)

    x, y_teacher = batch["x"], batch["y"]
    y_student = torch.log_softmax(student(x), dim=-1)
    loss = -(y_teacher * y_student).sum(dim=-1).mean()

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    print(f"Step {step:2d}: loss={loss.item():.4f}")

5. Check broker stats

stats = broker.get_stats()
print(f"Suppliers: {stats.connected_suppliers}, Clients: {stats.connected_clients}, Completed: {stats.total_completed}")

6. Graceful shutdown

Order matters: close the client first, then stop the supplier (sends GOODBYE), then stop the broker.

client.close()
supplier.stop()
time.sleep(0.2)
broker.stop()
print("Shutdown complete.")