Package: tpucare
TPU CareAutomatically take good care of your preemptible TPUsTable of ContentsTPU CareTable of ContentsFeaturesGetting StartedInstallationExamplesLong-running preemptible trainingSweepsCitationFeaturesReliable code execution: TPU Care starts a TPU, ensures it's set up as specified and continues the experiment
whenever the node dies. Think of it likeTerraForm+Ansiblefor machine learning.Maintenance of large swarms: When running multiple nodes, TPU Care will automatically delete dead instances while
keeping as many alive as possible.Code generation: To simplify setup, TPU Care efficiently clones your git repository and ensures trustable
execution of yourrun_commandthat continues even during outages.Optimized management: When a node dies, TPU Care deletes it within five minutes and creates a new one the second
there is capacity.Getting StartedInstallationpython3-mpipinstalltpucareExamplesWe've been using TPU Care for a while atHomebrewNLP. In fact, this library is just
the branched out core of the original production-ready HomebrewNLP code. At HomebrewNLP, there were two major use-cases
for this library. We started both massive hyperparameter sweeps which consumed 900,000 TPU-core hours within three
months and stable training on large TPU pods. Below, you can see a list of TPUs which are largely managed by TPU
Care:Screenshot fromTPUnicorn, a CLI-based TPU managed softwareIn the following sections, you'll learn how we use at massive scale with minimal code effort.Long-running preemptible trainingFor example, the following code can be used to create a production-ready v3-256 using
theHomebrewNLP-Jaxcodebase (
seeexamples/pod.pyfor an executable version):importdataclassesimporttypingfromnetrcimportnetrcimportyamlfromtpucareimportexec_command,exec_on_tpu,send_to_tpu,start_single@dataclasses.dataclassclassContext:retry:intZONE="europe-west4-a"HOST="big-pod"RUN_NAME="256-core-tpu"defload_config(ctx:Context):withopen("config.yaml",'r')asf:config=f.read()config=yaml.safe_load(config)wandb_api=wandb.Api()config["training"]["do_checkpoint"]=Truebase_checkpoint_path=config["training"]["checkpoint_path"]start_step=0forruninwandb_api.runs(f"{config['wandb']['entity']}/{config['wandb']['project']}"):ifrun.name==config['wandb']['name']:start_step=run.summary["_step"]breakstart_step-=start_step%config["training"]["checkpoint_interval"]config["training"]["start_step"]=start_stepconfig["wandb"]["name"]=f"{RUN_NAME}-{ctx.retry}"ifctx.retry>0:config["training"]["checkpoint_load_path"]=config["training"]["checkpoint_path"]config["training"]["checkpoint_path"]=f"{base_checkpoint_path}-{ctx.retry}"returnyaml.dump(config)defstart_fn(ctx:Context,worker:int):"""This function gets executed in threads to start a run on a new TPU. It receives the context object returned by`creation_callback` as well as the worker id which corresponds to the slice id this code was executed on in amulti-host setup. For single-host setups, such as v3-8s, the "worker" will always be set to 0.Ideally, it'd copy necessary files to the TPU and then run those. Here, `exec_command` can be used to create anexecution command that automatically spawns a `screen` session which persists even when the SSH connection gets cut."""send_to_tpu(HOST,ZONE,"config.yaml",load_config(ctx),worker)cmd=exec_command(repository="https://github.com/HomebrewNLP/HomebrewNLP-Jax",wandb_key=wandb_key)send_to_tpu(HOST,ZONE,"setup.sh",cmd,worker)exec_on_tpu(HOST,ZONE,"bash setup.sh",worker)defcreation_callback(host:str,ctx:typing.Optional[Context])->Context:"""The `creation_callback` is called once whenever a new TPU gets created and can be used to persist state(such as retry counters) across multiple invocations."""ifctxisNone:# first invocationreturnContext(0)ctx.retry+=1returnctxdefmain(service_account:str,tpu_version:int=3,slices:int=32,preemptible:bool=True):start_single(host=HOST,tpu_version=tpu_version,zone=ZONE,preemptible=preemptible,service_account=service_account,slices=slices,start_fn=start_fn,creation_callback=creation_callback)SweepsSimilarly, large swarms of instances can be launched trivially using tpucare. Here, we largely do the same setup as
above, but calllaunch_multipleinstead oflaunch_singlewhich takes the additional argumenttpusspecifying the
number of TPUs that should be launched and babysit. Depending on capacity and quota, the actual number of TPUs you get
might be lower than the number of TPUs specified.defmain(service_account:str,tpus:int,tpu_version:int=3,slices:int=32,preemptible:bool=True):start_multiple(prefix=HOST,tpu_version=tpu_version,zone=ZONE,preemptible=preemptible,service_account=service_account,slices=slices,start_fn=start_fn,creation_callback=creation_callback,tpus=tpus)However, this would simply launch the same run many times. If you instead plan to register them with aWandB Sweep, we need to modify thestart_fnto join the wandb
sweep.By patching in the code below, tpucare will start and maintain a large swarm of TPUs all working towards the same
hyperparameter optimization problem.importwandbwithopen("sweep.yaml",'r')asf:# sweep config passed straight to wandbconfig=yaml.safe_load(f.read())sweep_id=wandb.sweep(config,entity="homebrewnlp",project="gpt")defstart_fn(ctx:Context,worker:int):cmd=exec_command(repository="https://github.com/HomebrewNLP/HomebrewNLP-Jax",wandb_key=wandb_key,run_command=f"/home/ubuntu/.local/bin/wandb agent{sweep_id}")send_to_tpu(HOST,ZONE,"setup.sh",cmd,worker)exec_on_tpu(HOST,ZONE,"bash setup.sh",worker)The full executable code can be found
inexamples/sweep.py.Similarly, thestart_fncould be adapted to start an inference server
forHomebrewNLPorCraiyonor even execute machine learning unit-tests in
parallel.Citation@software{nestler_lucas_2022_6837312,author={Nestler, Lucas},title={TPU Care},month=jul,year=2022,publisher={Zenodo},version={0.0.2},doi={10.5281/zenodo.6837312},url={https://doi.org/10.5281/zenodo.6837312}}```