File size: 5,222 Bytes
17b4b9a | 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 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 | import time
from typing import Dict, List
import torch
from torch.utils.data import DataLoader, Dataset
from modeling_xonelm import XoneLM, HardwareContext
from luminav import LuminaV
from tokenizer import (
build_xonelm_tokenizer,
MultiTurnConversationFormatter,
SpecialTokenConfig,
)
class SafeSFTCollator:
def __init__(self, max_seq_len: int = 512, pad_token_id: int = 0):
self.max_seq_len = max_seq_len
self.pad_token_id = pad_token_id
def __call__(self, samples: List[Dict[str, List[int]]]) -> Dict[str, torch.Tensor]:
batch_inputs = []
batch_labels = []
for item in samples:
inp = item["input_ids"][: self.max_seq_len]
lbl = item["labels"][: self.max_seq_len]
pad_len = self.max_seq_len - len(inp)
batch_inputs.append(
torch.tensor(inp + [self.pad_token_id] * pad_len, dtype=torch.long)
)
batch_labels.append(
torch.tensor(lbl + [-100] * pad_len, dtype=torch.long)
)
return {
"input_ids": torch.stack(batch_inputs),
"labels": torch.stack(batch_labels),
}
class ConversationDataset(Dataset):
def __init__(self, data: List[Dict[str, List[int]]]):
self.data = data
def __len__(self) -> int:
return len(self.data)
def __getitem__(self, idx: int) -> Dict[str, List[int]]:
return self.data[idx]
def run_sft_demo():
device = HardwareContext.get_optimal_device()
autocast_dtype = HardwareContext.get_optimal_autocast_dtype(device)
print("Compute Device :", device)
print("Autocast Dtype :", autocast_dtype)
tokenizer = build_xonelm_tokenizer()
vocab_size = len(tokenizer)
token_cfg = SpecialTokenConfig(
pad_token_id=0,
bos_token_id=1,
eos_token_id=2,
unk_token_id=3,
eod_token_id=4,
)
formatter = MultiTurnConversationFormatter(tokenizer, token_cfg)
sample_dialogues = [
[
{"role": "system", "content": "You are a precise reasoning assistant."},
{"role": "user", "content": "Lily found a wooden box. What did she open?"},
{"role": "assistant", "content": "She opened the wooden box to see what was inside."},
],
[
{"role": "system", "content": "You are a polite companion."},
{"role": "user", "content": "Hello! How can we optimize memory bandwidth?"},
{"role": "assistant", "content": "We can compress Key-Value caches using low-rank latent projections."},
],
[
{"role": "system", "content": "You are a creative writer."},
{"role": "user", "content": "Tell me a story about a kitten in the garden."},
{"role": "assistant", "content": "Once upon a time, a tiny kitten chased a butterfly across the grass."},
],
]
formatted_samples = [formatter.format_conversation(dialogue) for dialogue in sample_dialogues]
dataset = ConversationDataset(formatted_samples)
collator = SafeSFTCollator(max_seq_len=256, pad_token_id=token_cfg.pad_token_id)
loader = DataLoader(dataset, batch_size=2, shuffle=True, collate_fn=collator)
model = XoneLM(
vocab_size=vocab_size,
dim=512,
num_layers=12,
num_heads=8,
kv_latent_dim=64,
hub_size=512,
num_specialized_hubs=12,
num_terminals=32,
slots_per_terminal=16,
).to(device)
optimizer = LuminaV(
model.parameters(),
lr=2e-4,
betas=(0.9, 0.999),
eps=1e-8,
weight_decay=1e-3,
tau=0.8,
buffer=2,
cautious=True,
execution="auto",
)
use_scaler = (device.type == "cuda" and autocast_dtype == torch.float16)
scaler = torch.amp.GradScaler("cuda", enabled=True) if use_scaler else None
model.train()
optimizer.zero_grad()
start_time = time.time()
for epoch in range(2):
for step, batch in enumerate(loader):
x = batch["input_ids"].to(device, non_blocking=True)
y = batch["labels"].to(device, non_blocking=True)
with HardwareContext.get_autocast_context(device):
output = model(x, labels=y, is_sft=True)
loss = output.loss
if scaler is not None:
scaler.scale(loss).backward()
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
scaler.step(optimizer)
scaler.update()
else:
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()
optimizer.zero_grad()
print(f"Epoch [{epoch+1}/2] | Step [{step+1}/{len(loader)}] | SFT Loss: {loss.item():.4f}")
elapsed = time.time() - start_time
print(f"[+] SFT Training Demo completed successfully in {elapsed:.2f}s!")
if __name__ == "__main__":
run_sft_demo() |