File size: 3,815 Bytes
99297d7 | 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 | """V4.5 distillation-student entry point; W&B without any machine-wide login.
Credentials come only from the environment (WANDB_API_KEY, optional WANDB_ENTITY),
as `cluster_v45.sh` sets them from a private file with isolated W&B directories.
Without a key, or when online initialization fails, the run logs offline under
artifacts/tracking/<model>/ for a later `wandb sync`; training never waits on W&B.
"""
import argparse
import os
from pathlib import Path
import sys
import tomllib
ROOT=Path(__file__).resolve().parents[2]
sys.path.insert(0,str(ROOT/'src'))
def main():
parser=argparse.ArgumentParser(description=__doc__)
parser.add_argument('--config',type=Path,required=True)
parser.add_argument('--resume',action='store_true')
parser.add_argument('--dry-run',action='store_true',
help='Verify configuration, code, data and teacher identities, then exit')
parser.add_argument('--smoke',type=int,default=0,
help='Run N real updates into <output>-smoke without W&B, to check memory and speed')
args=parser.parse_args()
config=tomllib.loads(args.config.read_text('utf-8'))
if args.smoke:
config['training']['smoke_updates']=args.smoke
config['output_dir']+='-smoke'
config['log_dir']+='-smoke'
config['tracking']={'enabled':False}
output=ROOT/config['output_dir']
logs=ROOT/config['log_dir']
if args.dry_run:
from vimeml.training.joint_v45 import build_plan
plan=build_plan(config,ROOT,verify_files=True)
print({'stage':plan['stage'],'arm':config['arm'],'parameters':plan['parameters'],
'max_steps':plan['settings']['max_steps'],
'updates_per_epoch':plan['settings']['updates_per_epoch'][:3]+['...'],
'required_files':len(plan['required_files']),'teachers':plan['signatures']['teachers']})
return
if not args.resume and ((output.exists() and any(output.iterdir())) or (logs.exists() and any(logs.iterdir()))):
parser.error('Fresh V4.5 output/log directories must be empty')
if args.resume and not (output/'last.pt').exists():
parser.error('No V4.5 checkpoint to resume')
from vimeml.training.joint_v45 import run
tracking=config.get('tracking',{})
if not tracking.get('enabled'):
result=run(config,args.resume)
print(result,flush=True)
return
import wandb
from vimeml.training.data import write_json
directory=ROOT/'artifacts/tracking'/output.name
directory.mkdir(parents=True,exist_ok=True)
def start(mode):
return wandb.init(project=tracking.get('project','vimeml'),entity=os.environ.get('WANDB_ENTITY'),
group=tracking['group'],name=config['arm']+('-resume' if args.resume else ''),
job_type='distillation',tags=['v4.5',config['arm'],config['distillation']['pool']],mode=mode,dir=str(directory),
sync_tensorboard=True,save_code=False,config=config,
settings=wandb.Settings(disable_git=True,init_timeout=90))
mode=tracking.get('mode','online') if os.environ.get('WANDB_API_KEY') else 'offline'
try:
cloud=start(mode)
except Exception as error:
if mode=='offline':
raise
print(f'W&B online initialization failed ({type(error).__name__}: {error}); logging offline',flush=True)
mode='offline'
cloud=start(mode)
with cloud:
write_json(directory/f'live-{cloud.id}.json',{'id':cloud.id,'url':cloud.url if mode!='offline' else None,
'mode':mode,'resume':args.resume,'directory':str(directory)})
print(f'W&B run ({mode}): {cloud.url if mode!="offline" else directory}',flush=True)
result=run(config,args.resume)
cloud.summary.update({'training_result':result})
if __name__=='__main__':
main()
|