Astheral / expanded-evaluation /train_expanded.py
SEE7's picture
Finish V3: document extra rank32 training, standard benchmarks and release decision
c9089df verified
Raw History Blame Contribute Delete
10.3 kB
# /// script
# dependencies = ["torch==2.8.0", "transformers==4.56.2", "peft==0.17.1", "accelerate==1.10.1", "huggingface_hub==0.34.4", "bitsandbytes==0.47.0", "safetensors==0.6.2"]
# ///
"""Budget-bounded V3 training with paired, disjoint evaluation and reload gate."""
import os, json, random, time, re, gc, hashlib, traceback
from pathlib import Path
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, Trainer, TrainingArguments, TrainerCallback, set_seed
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training, PeftModel
from huggingface_hub import HfApi, hf_hub_download
REPO='SEE7/Astheral'; REV='37adef9b3917e4a3de0bb77f0db71089b7b82202'
OUT=Path('/tmp/astheral-v3'); OUT.mkdir(exist_ok=True)
api=HfApi(); started=time.monotonic(); set_seed(20261113); rng=random.Random(20261113)
assert torch.cuda.is_available() and 'A10G' in torch.cuda.get_device_name(0)
branch='v3-expanded-20261113'
api.create_branch(REPO,branch=branch,revision=REV,exist_ok=True)
settings=json.loads(Path(hf_hub_download(REPO,'settings.json',revision=REV)).read_text())
tok=AutoTokenizer.from_pretrained(REPO,revision=REV); tok.pad_token_id=128004
report={'status':'preparing','parent_revision':REV,'hardware':torch.cuda.get_device_name(0),'job_timeout_seconds':5400,'listed_usd_per_hour':1.0,'training_seed':20261113}
def publish():
report['elapsed_seconds']=round(time.monotonic()-started,2)
(OUT/'v3_report.json').write_text(json.dumps(report,indent=2))
api.upload_folder(repo_id=REPO,folder_path=OUT,revision=branch,path_in_repo='v3-results',ignore_patterns=['merged/**'],commit_message='Record measured V3 training and evaluation artifacts')
def sample(a,b,c,kind):
if kind==0:return f'Return only the integer: {a} + {b} - {c}.',str(a+b-c),'integer'
if kind==1:return f'A crate has {a} red and {b} blue beads. Remove {c} red beads. Return only the total beads remaining.',str(a+b-c),'integer'
if kind==2:return f'Return only valid JSON with keys "item" and "count": item=part_{a}, count={b}.',{'item':f'part_{a}','count':b},'json'
if kind==3:return f'Return only one CSV value row, no header: name=Item{a}; quantity={b}; code=K{c}.',f'Item{a},{b},K{c}','exact'
if kind==4:return f'Copy exactly, no other text: Key_{a}-B{b}_C{c}',f'Key_{a}-B{b}_C{c}','exact'
return f'Write only a Python function named solve(x) that returns x * {b} + {c}.',f'x * {b} + {c}','code'
def grade(case,text):
expected=case['expected']; kind=case['kind']; text=text.strip()
try:
if kind=='json':return json.loads(text)==expected
if kind=='integer':return re.fullmatch(r'-?\d+',text) is not None and int(text)==int(expected)
if kind=='exact':return text==expected
# Parse and evaluate a tiny arithmetic-only AST; never execute model-generated code.
import ast
clean=re.sub(r'^```(?:python)?\s*|\s*```$','',text).strip();tree=ast.parse(clean)
if len(tree.body)!=1 or not isinstance(tree.body[0],ast.FunctionDef):return False
fn=tree.body[0]
if fn.name!='solve' or len(fn.args.args)!=1 or fn.args.args[0].arg!='x' or len(fn.body)!=1 or not isinstance(fn.body[0],ast.Return):return False
def calc(n,x):
if isinstance(n,ast.Constant) and isinstance(n.value,(int,float)):return n.value
if isinstance(n,ast.Name) and n.id=='x':return x
if isinstance(n,ast.BinOp) and isinstance(n.op,(ast.Add,ast.Sub,ast.Mult)):
a,b=calc(n.left,x),calc(n.right,x)
return a+b if isinstance(n.op,ast.Add) else a-b if isinstance(n.op,ast.Sub) else a*b
raise ValueError('Unsupported expression')
match=re.fullmatch(r'x \* (\d+) \+ (\d+)',expected);b,c=map(int,match.groups())
return all(calc(fn.body[0].value,x)==x*b+c for x in [-5,0,1,7,100])
except Exception:return False
tests=[]
for i in range(60):
# Evaluation numbers are outside the synthetic training ranges.
a,b,c=701+i*3,91+i,31+i%17;p,e,k=sample(a,b,c,i%6)
tests.append({'id':f'heldout_{i:02}','prompt':p,'expected':e,'kind':k})
train=[]
for i in range(900):
a,b,c=rng.randint(100,500),rng.randint(10,80),rng.randint(1,29)
p,e,k=sample(a,b,c,i%6)
answer=json.dumps(e) if k=='json' else 'def solve(x):\n return '+e if k=='code' else str(e)
train.append({'messages':[{'role':'user','content':p},{'role':'assistant','content':answer}]})
retained=[json.loads(s) for s in Path(hf_hub_download(REPO,'train.jsonl',revision=REV)).read_text().splitlines()]
rng.shuffle(retained);train+=retained[:900];rng.shuffle(train)
assert not {x['prompt'] for x in tests} & {m['content'] for row in train for m in row['messages'] if m['role']=='user'}
(OUT/'heldout_cases.json').write_text(json.dumps(tests,indent=2));(OUT/'training.jsonl').write_text('\n'.join(json.dumps(x) for x in train)+'\n')
(OUT/'train_v3.py').write_text(Path(__file__).read_text())
q=BitsAndBytesConfig(load_in_4bit=True,bnb_4bit_quant_type='nf4',bnb_4bit_use_double_quant=True,bnb_4bit_compute_dtype=torch.bfloat16)
model=AutoModelForCausalLM.from_pretrained(REPO,revision=REV,quantization_config=q,device_map={'':'cuda:0'},torch_dtype=torch.bfloat16,attn_implementation='sdpa')
@torch.inference_mode()
def evaluate(label):
model.eval(); rows=[]
for i,case in enumerate(tests):
inputs=tok.apply_chat_template([{'role':'system','content':settings['system']},{'role':'user','content':case['prompt']}],tokenize=True,add_generation_prompt=True,return_dict=True,return_tensors='pt').to('cuda')
out=model.generate(**inputs,max_new_tokens=96,do_sample=False,pad_token_id=128004,eos_token_id=128009)
text=tok.decode(out[0,inputs['input_ids'].shape[1]:],skip_special_tokens=True)
rows.append({**case,'output':text,'pass':grade(case,text),'tokens':int(out.shape[1]-inputs['input_ids'].shape[1])})
print('EVAL',label,i,rows[-1]['pass'],flush=True)
(OUT/(label+'_outputs.json')).write_text(json.dumps(rows,indent=2))
return rows
try:
baseline=evaluate('v2_quantized'); report['baseline_passed']=sum(r['pass'] for r in baseline);publish()
model=prepare_model_for_kbit_training(model,use_gradient_checkpointing=True)
model=get_peft_model(model,LoraConfig(r=32,lora_alpha=64,lora_dropout=0.05,target_modules=['q_proj','k_proj','v_proj','o_proj','gate_proj','up_proj','down_proj'],task_type='CAUSAL_LM'))
model.config.use_cache=False;encoded=[]
for row in train:
messages=row['messages']; prefix=tok.apply_chat_template(messages[:-1],tokenize=True,add_generation_prompt=True)
full=tok.apply_chat_template(messages,tokenize=True,add_generation_prompt=False)
if len(full)<=512 and full[:len(prefix)]==prefix:encoded.append({'input_ids':full,'labels':[-100]*len(prefix)+full[len(prefix):]})
def collate(rows):
length=max(len(r['input_ids']) for r in rows)
return {k:torch.tensor([r[k]+[pad]*(length-len(r[k])) for r in rows]) for k,pad in [('input_ids',128004),('labels',-100)]}|{'attention_mask':torch.tensor([[1]*len(r['input_ids'])+[0]*(length-len(r['input_ids'])) for r in rows])}
class Budget(TrainerCallback):
def on_step_end(self,args,state,control,**kwargs):
if time.monotonic()-started>2400:control.should_training_stop=True
return control
args=TrainingArguments(output_dir='/tmp/v3-checkpoints',max_steps=360,per_device_train_batch_size=1,gradient_accumulation_steps=4,learning_rate=5e-6,warmup_ratio=0.05,bf16=True,gradient_checkpointing=True,gradient_checkpointing_kwargs={'use_reentrant':False},optim='paged_adamw_8bit',save_strategy='no',logging_steps=10,report_to='none',remove_unused_columns=False,seed=20261113)
trainer=Trainer(model=model,args=args,train_dataset=encoded,data_collator=collate,callbacks=[Budget()])
report.update(status='training',training_examples=len(encoded),trainable_parameters=sum(p.numel() for p in model.parameters() if p.requires_grad));publish()
result=trainer.train();report['training']=result.metrics;report['optimizer_steps']=trainer.state.global_step
adapter=OUT/'adapter';model.save_pretrained(adapter,safe_serialization=True);tok.save_pretrained(adapter)
model.config.use_cache=True;model.gradient_checkpointing_disable();candidate=evaluate('v3_quantized')
report['candidate_passed']=sum(r['pass'] for r in candidate);report['status']='trained';publish()
del trainer,model;gc.collect();torch.cuda.empty_cache()
model=AutoModelForCausalLM.from_pretrained(REPO,revision=REV,torch_dtype=torch.bfloat16,device_map={'':'cuda:0'},attn_implementation='sdpa')
model=PeftModel.from_pretrained(model,adapter).merge_and_unload(safe_merge=True)
model.config.astheral_version='3';merged=OUT/'merged';model.save_pretrained(merged,safe_serialization=True,max_shard_size='4GB');tok.save_pretrained(merged)
del model;gc.collect();torch.cuda.empty_cache()
model=AutoModelForCausalLM.from_pretrained(merged,torch_dtype=torch.bfloat16,device_map={'':'cuda:0'},attn_implementation='sdpa')
reloaded=evaluate('v3_reloaded');report['reloaded_passed']=sum(r['pass'] for r in reloaded)
groups={kind:{'baseline':sum(r['pass'] for r in baseline if r['kind']==kind),'candidate':sum(r['pass'] for r in reloaded if r['kind']==kind),'total':sum(r['kind']==kind for r in tests)} for kind in {x['kind'] for x in tests}}
report['groups']=groups;report['quality_gate_passed']=report['reloaded_passed']>report['baseline_passed'] and all(v['candidate']>=v['baseline'] for v in groups.values())
report['status']='complete';report['peak_gpu_bytes']=torch.cuda.max_memory_allocated();publish()
# Save the actual trained weights to the isolated candidate branch even on a failed gate.
api.upload_folder(repo_id=REPO,folder_path=merged,revision=branch,commit_message='Save reloaded V3 candidate weights; see measured quality gate in v3-results')
if report['quality_gate_passed']:
api.create_tag(REPO,tag='v3-expanded-candidate',revision=api.model_info(REPO,revision=branch).sha,exist_ok=True)
print('V3_COMPLETE',json.dumps(report),flush=True)
except Exception:
report['status']='failed';report['error']=traceback.format_exc();publish();raise