File size: 10,259 Bytes
c9089df | 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 | # /// 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
|