File size: 6,411 Bytes
1440c49 | 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 | """Actual causal-LM training, random initialization by default.
--resume MODEL starts from model weights with a NEW optimizer, not an exact resume.
Exports FP16 safetensors; training itself uses FP32. No pretrained weights used.
"""
import argparse
import json
import math
import random
import time
from pathlib import Path
import numpy as np
import torch
from transformers import GPT2Config, GPT2LMHeadModel, GPT2TokenizerFast, set_seed
ROOT = Path(__file__).resolve().parent
def main():
p = argparse.ArgumentParser()
p.add_argument('--steps', type=int, default=300)
p.add_argument('--batch-size', type=int, default=2)
p.add_argument('--seq-len', type=int, default=128)
p.add_argument('--lr', type=float, default=0.0005)
p.add_argument('--seed', type=int, default=42)
p.add_argument('--threads', type=int, default=2)
p.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu')
p.add_argument('--resume', default=None)
p.add_argument('--output', default=str(ROOT/'model'))
args = p.parse_args()
assert 1 < args.seq_len <= 512 and args.batch_size >= 2 and args.batch_size % 2 == 0
assert args.steps > 0
set_seed(args.seed)
torch.set_num_threads(args.threads)
output = Path(args.output)
output.mkdir(parents=True, exist_ok=True)
tok = GPT2TokenizerFast.from_pretrained(args.resume or ROOT/'model')
if args.resume:
model = GPT2LMHeadModel.from_pretrained(args.resume, torch_dtype=torch.float32)
else:
config = GPT2Config(vocab_size=len(tok), n_positions=512, n_ctx=512,
n_embd=512, n_layer=12, n_head=8, n_inner=2048,
activation_function='gelu_new', resid_pdrop=.1, embd_pdrop=.1,
attn_pdrop=.1, bos_token_id=tok.bos_token_id,
eos_token_id=tok.eos_token_id, pad_token_id=tok.pad_token_id,
use_cache=False)
model = GPT2LMHeadModel(config)
model.to(args.device)
parameters = sum(x.numel() for x in model.parameters())
print(f'Parameters: {parameters:,}; device: {args.device}', flush=True)
data = {f'{split}_{lang}': np.load(ROOT/'.cache'/f'{split}_{lang}.npy')
for split in ['train', 'valid'] for lang in ['ru', 'en']}
def batch(split, rng, language=None):
rows = []
for i in range(args.batch_size):
lang = language or ('ru' if i % 2 == 0 else 'en')
a = data[f'{split}_{lang}']
start = rng.randrange(len(a) - args.seq_len + 1)
rows.append(a[start:start+args.seq_len].copy())
return torch.tensor(np.stack(rows), dtype=torch.long, device=args.device)
@torch.no_grad()
def evaluate():
model.eval()
result = {}
for lang in ['ru','en']:
rng = random.Random(12345)
losses = []
for _ in range(8):
x = batch('valid', rng, lang)
losses.append(model(x, labels=x, use_cache=False).loss.item())
result[lang] = sum(losses) / len(losses)
model.train()
return result
before = evaluate()
print('Validation before:', before, flush=True)
optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr,
betas=(.9,.95), weight_decay=.1, foreach=False)
rng = random.Random(args.seed)
started = time.monotonic()
history = []
log_path = output/'training_log.jsonl'
with log_path.open('w') as log:
for step in range(1, args.steps+1):
warmup = min(20, max(1,args.steps//10))
factor = min(1.,step/warmup) * (.1 + .9 * .5 * (1 + math.cos(math.pi*(step-1)/args.steps)))
for group in optimizer.param_groups:
group['lr'] = args.lr * factor
x = batch('train', rng)
optimizer.zero_grad(set_to_none=True)
loss = model(x, labels=x, use_cache=False).loss
if not torch.isfinite(loss):
raise RuntimeError('Non-finite loss; refusing to export')
loss.backward()
grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(),1.)
optimizer.step()
row = {'step':step,'loss':loss.item(),'lr':optimizer.param_groups[0]['lr'],
'grad_norm':float(grad_norm),'elapsed_seconds':round(time.monotonic()-started,2)}
log.write(json.dumps(row)+'\n')
log.flush()
history.append(row)
if step == 1 or step % 10 == 0:
print(json.dumps(row), flush=True)
elapsed = time.monotonic()-started
after = evaluate()
report = {'name':'Galactic 2.5 Mini','developer_label':'SpaceModel',
'status':'experimental short-run checkpoint; not a capable bilingual assistant',
'initialization':'random' if not args.resume else f'weights from {args.resume}; new optimizer',
'parameters':parameters, 'steps':args.steps,
'input_token_presentations':args.steps*args.batch_size*args.seq_len,
'prediction_targets':args.steps*args.batch_size*(args.seq_len-1),
'sampling':'equal RU/EN sequences, random windows with replacement',
'batch_size':args.batch_size,'sequence_length':args.seq_len,
'max_context':512,'seed':args.seed,'device':args.device,
'training_seconds':round(elapsed,2), 'optimizer':'AdamW',
'learning_rate_peak':args.lr,'betas':[.9,.95],'weight_decay':.1,
'gradient_clip_norm':1.,'training_dtype':'float32','export_dtype':'float16',
'validation_before_nats':before,'validation_after_nats':after,
'validation_method':'8 fixed random batches per language from held-out final 5% of each book; not a benchmark',
'corpus_tokens':{k:len(v) for k,v in data.items()},
'torch_version':torch.__version__,
'optimizer_state_included':False}
# Free optimizer states before FP16 export.
optimizer.zero_grad(set_to_none=True)
del optimizer
model.eval().half()
model.config.use_cache = True
model.config.galactic_model_name = 'Galactic 2.5 Mini'
model.config.galactic_developer = 'SpaceModel'
model.save_pretrained(output, safe_serialization=True)
tok.save_pretrained(output)
(output/'training_report.json').write_text(json.dumps(report,ensure_ascii=False,indent=2),encoding='utf-8')
print('Validation after:',after,flush=True)
print('SAVED:', output, flush=True)
if __name__ == '__main__':
main()
|