File size: 1,787 Bytes
3738348 | 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 | """Dry-run SFT: load data, build model, resize embeddings, one forward pass."""
import sys, os, json, torch
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__))))
from sft import load_special_tokens, load_sft_dataset, get_sft_batch
from model import ModelConfig, Retriever500M
from tokenizers import Tokenizer
def main():
load_special_tokens()
tok = Tokenizer.from_file("tokenizer/tokenizer_agent.json")
print("Loading dataset...")
dataset = load_sft_dataset("data/sft_traces.jsonl", tok, max_seq_len=768)
print(f"Dataset size: {len(dataset)}")
config = ModelConfig(vocab_size=32009, d_model=1280, n_layers=23, n_heads=20,
d_ff=3456, max_seq_len=768, tie_embeddings=True)
model = Retriever500M(config).cuda()
print(f"Model: {model.count_parameters()/1e6:.1f}M params")
# Load old checkpoint and resize embeddings
ckpt = torch.load("checkpoints/latest.pt", map_location="cuda", weights_only=False)
state = ckpt["model_state_dict"]
old_w = state["token_embedding.weight"]
new_w = torch.zeros(32009, 1280)
new_w[:32000] = old_w
torch.nn.init.normal_(new_w[32000:], mean=0.0, std=0.02)
state["token_embedding.weight"] = new_w
model.load_state_dict(state)
print("Checkpoint loaded with resized embeddings")
# One forward pass
input_ids, targets, loss_mask = get_sft_batch(dataset, batch_size=2, seq_len=768, device=torch.device("cuda"))
print(f"Batch: {input_ids.shape}")
with torch.autocast("cuda", dtype=torch.bfloat16):
out = model(input_ids, targets=targets)
print(f"Forward pass OK: loss={out['loss'].item():.4f}")
print("SFT DRY RUN PASSED")
if __name__ == "__main__":
main()
|