Come ho addestrato 10 TB per Stable Diffusion su SageMaker
Con dettagli sul mio prossimo libro sulla formazione distribuita con AWS!
Trattare con molti dati è difficile, non c'è dubbio. Tuttavia, con alcune scelte intelligenti e design semplici, anche tu puoi ridimensionare i tuoi progetti per replicarli e produrre risultati all'avanguardia.
In questo post ti illustrerò un caso di studio in cui ho utilizzato AWS, in particolare SageMaker, S3 e FSx per Lustre, per scaricare, elaborare e addestrare un modello di diffusione stabile. Ho utilizzato SageMaker Distributed Data Parallel per ottimizzare la discesa del gradiente su larga scala. Ho anche utilizzato il parallelismo dei lavori sul backend modulare di SageMaker per scaricare ed elaborare tutti i file.
Il mio codice di addestramento si basa su un campione di Corey Barrett di AWS MLSL, disponibile qui. Proverò ad aggiornare quel repository con tutto il codice sorgente per questo progetto entro la fine dell'anno.
Se sei affamato di più contenuti sull'esecuzione di lavori su larga scala su SageMaker e sull'intero ambito di pre-addestramento di modelli di visione e linguaggio su AWS, dai un'occhiata al mio libro in uscita nell'aprile 2023!
Nel mio libro, trascorro 15 capitoli guidandoti attraverso l'intero ciclo di vita, dalla ricerca del caso d'uso giusto, alla preparazione del tuo ambiente, alla formazione, alla risoluzione dei problemi, alla valutazione, all'operatività e al ridimensionamento.
Il mio libro di formazione distribuito è l'ideale per i principianti ; è veramente per gli sviluppatori Python principianti e intermedi che vogliono acquisire le competenze per addestrare i propri modelli di visione e linguaggio di grandi dimensioni.
Immergiamoci!
1/ Scarica i file del parquet Laion-5B con i lavori SageMaker
Il set di dati di base utilizzato per addestrare la diffusione stabile è Laion-5B. Questo è un set di dati open source che fornisce miliardi di coppie di immagini/testo da Internet. Romain Beaumont, un contributore al set di dati Laion-5B, ha reso open source il suo strumento per scaricare ed elaborare questo set di dati e l'ho usato come base. Il suo codice sorgente è disponibile qui. Grazie Romaino!
Il set di dati di base è composto da due parti: più di 100 file parquet che puntano alle immagini originali e alle loro didascalie e alle immagini scaricate stesse.
Innanzitutto, ho scaricato i file parquet con una semplice funzione Python. Fortunatamente il toolkit di Romain ti consente di scaricare direttamente nel tuo bucket S3, il che ti consente di aggirare qualsiasi vincolo di archiviazione locale. L'ho provato alcune volte sulla mia istanza di SageMaker Studio, passando a una macchina più grande quando avevo bisogno di lavorare con un file più grande.
Successivamente, l'ho racchiuso in un singolo script che potevo eseguire sul mio lavoro SageMaker.
import argparse
import os
def parse_args():
parser = argparse.ArgumentParser()
# points to your session bucket
parser.add_argument("--bucket", type=str, default=os.environ["SM_HP_BUCKET"])
args = parser.parse_args()
return args
def download_parquet(bucket):
cmd = '''for i in {00000..00127};
do wget https://huggingface.co/datasets/laion/laion2B-en/resolve/main/part-$i-5114fd87-297e-42b0-9d11-50f1df323dfa-c000.snappy.parquet -O - | aws s3 cp - s3://{}/metadata/laion2B-en-joined/part-$i-4cfd6e30-f032-46ee-9105-8696034a8373-c000.snappy.parquet;
done'''.format(bucket)
# runs the command
os.system(cmd)
if __name__ == "__main__":
args = parse_args()
download_parquet(args.bucket)
2/ Scarica le coppie di immagini e testo Laion-5B
Successivamente, è arrivato il momento di scaricare le immagini stesse. Ho fatto un elenco di tutti i file di parquet, l'ho esaminato e ho inviato ogni file di parquet al proprio lavoro SageMaker.
Il mio script per scaricare l'immagine e le coppie di testo era simile a questo:
from img2dataset import download
import shutil
import os
import multiprocessing
import threading
import argparse
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument("--cores", type=int, default=multiprocessing.cpu_count())
parser.add_argument("--threads", type=int, default=threading.active_count())
parser.add_argument("--parquet", type=str, default=os.environ["SM_CHANNEL_PARQUET"])
parser.add_argument("--file_name", type=str, default=os.environ["SM_HP_FILE_NAME"])
parser.add_argument("--bucket", type=str, default=os.environ["SM_MODULE_DIR"].split('/')[2])
args = parser.parse_args()
return args
def prep_system():
args = parse_args()
# send joint path and file name
url_list = "{}/{}".format(args.parquet, args.file_name)
part_number = args.file_name.split('-')[1]
# point to output path in S3
s3_output = "s3://{}/data/part-{}/".format(args.bucket, part_number)
return args, url_list, s3_output
if __name__ == "__main__":
args, url_list, s3_output = prep_system()
download(
processes_count=args.cores,
thread_count=args.threads,
# takes a single parquet file
url_list=url_list,
image_size=256,
# copies to S3 directly, bypassing local disk
output_folder=s3_output,
# each image / caption pair is a tarball
output_format="webdataset",
input_format="parquet",
url_col="URL",
caption_col="TEXT",
enable_wandb=False,
number_sample_per_shard=1000,
distributor="multiprocessing",
)
Quindi, ho esaminato il mio elenco di file di parquet e ho inviato ogni file di parquet al proprio lavoro . Chiamiamo questo lavoro parallelismo , è un modo semplice ed efficace per eseguire lavori di calcolo paralleli.
def get_estimator(part_number, p_file, output_dir):
hyperparameters = {"file_name": p_file}
estimator = PyTorch(entry_point="download.py",
base_job_name="laion-part-{}".format(part_number),
role=role,
source_dir="scripts",
# configures the SageMaker training resource, you can increase as you need
instance_count=1,
instance_type="ml.c5.18xlarge",
py_version="py36",
framework_version = '1.8',
sagemaker_session=sagemaker_session,
volume_size = 250,
debugger_hook_config=False,
hyperparameters=hyperparameters,
output_path = output_dir)
return estimator
for p_file in parquet_list[:18]:
part_number = p_file.split('-')[1]
output_dir = "s3://{}/data/part-{}/".format(bucket, part_number)
if is_open(output_dir):
est = get_estimator(part_number, p_file, output_dir)
est.fit({"parquet":"s3://{}/meta/{}".format(bucket, p_file)}, wait=False)
L'ho eseguito su 18 lavori diversi, tutti in esecuzione contemporaneamente, per circa 24 ore.
Gli oggetti risultanti avevano questo aspetto:
3/ Creare FSx per il volume Lustre dal percorso S3
Successivamente, ho impiegato ben venti minuti per creare un nuovo volume FSx for Lustre dai miei percorsi S3. Questo è stato incredibilmente facile da fare; andata e ritorno compresi alcuni dei miei singhiozzi e il tempo di creazione del volume è stato di circa venti minuti. Il tutto senza lasciare la console AWS.
S3 ha una funzionalità davvero interessante per calcolare la dimensione totale di un percorso. L'ho usato per stimare la quantità totale di dati che stavo usando: un modesto 9,5 TB.
La mia configurazione di Lustre era così:
Perché approvvigionare Lustre? Perché riduce il tempo di download dei dati sul tuo lavoro di formazione SageMaker da, su questa scala probabilmente 90 minuti solidi o più, a meno di 60 secondi.
Ho passato un po' di tempo a configurare il mio VPC per abilitare la formazione su SageMaker e scrivere i risultati su S3. Ho anche esteso un AWS Deep Learning Container con tutti i miei pacchetti per risparmiare tempo durante la formazione; invece di pip installare tutto per ogni lavoro, lo faccio una volta durante la creazione dell'immagine, quindi semplicemente scarico quell'immagine sul mio lavoro in fase di esecuzione.
4/ Costruisci un indice di linee json
Dopo questo, ho perso alcune ore cercando di gestire la funzione load_dataset() predefinita con il mio formato webdataset. Sospetto ancora che sia possibile, ma ho deciso di scrivere solo il mio caricatore di dati.
Per risparmiare tempo e denaro sul mio enorme cluster GPU, ho creato un json lines index . Questo era semplicemente un file multi-GB con 50 milioni di oggetti json, ognuno dei quali puntava a una coppia immagine/testo valida su FSx per Lustre.
Ho eseguito un grande lavoro SageMaker basato su CPU per elencare tutti i file delle mie 18 parti su Lustre. Successivamente, ho caricato, testato e creato una funzione di caricamento dati molto piccola sul mio notebook Studio.
5/ Esegui su 192 GPU con la formazione distribuita di SageMaker
Dopodiché, era giunto il momento di superare un'unica epoca! Ho aumentato i miei limiti per le istanze ml.p4d.24xlarge sulla formazione SageMaker in us-east-1 utilizzando AWS Service Quota a 24. Poiché ogni istanza ha 8 GPU, si tratta di un totale di 192 GPU.
Ho eseguito alcune corse con solo l'1% dei miei dati complessivi per assicurarmi che il ciclo di allenamento funzionasse come previsto; la nuova funzionalità delle piscine calde è stata estremamente utile!
La mia configurazione di allenamento finale era così.
import sagemaker
from sagemaker.huggingface import HuggingFace
sess = sagemaker.Session()
role = sagemaker.get_execution_role()
version = 'v1'
image_uri = '<aws account id>.dkr.ecr.us-east-1.amazonaws.com/stable-diffusion:{}'.format(version )
# required in this version of the train script
data_channels['sd_base_model'] = 's3://dist-train/stable-diffusion/conceptual_captions/sd-base-model/'
hyperparameters={'pretrained_model_name_or_path':'/opt/ml/input/data/sd_base_model',
'train_data_dir':'/opt/ml/input/data/training/laion-fsx',
'index_name':'data_index.jsonl',
'caption_column':'caption',
'image_column':'image',
'resolution':256,
'mixed_precision':'fp16',
# this is per device
'train_batch_size':22,
'learning_rate': '1e-10',
# 'max_train_steps':1000000,
'num_train_epochs':1,
'output_dir':'/opt/ml/model/sd-output-final',
'n_rows':50000000}
est = HuggingFace(entry_point='finetune.py',
source_dir='src',
image_uri=image_uri,
sagemaker_session=sess,
role=role,
output_path="s3://laion-5b/output/model/",
instance_type='ml.p4dn.24xlarge',
keep_alive_period_in_seconds = 60*60,
py_version='py38',
base_job_name='fsx-stable-diffusion',
instance_count=24,
enable_network_isolation=True,
encrypt_inter_container_traffic = True,
# all opt/ml paths point to SageMaker training
hyperparameters = hyperparameters,
distribution={"smdistributed": { "dataparallel": { "enabled": True } }},
max_retry_attempts = 30,
max_run = 4 * 60 * 60,
debugger_hook_config=False,
disable_profiler = True,
**kwargs)
est.fit(inputs=data_channels, wait=False)
Quello che mi ha sorpreso di più è stato che la differenza tra i miei lavori più grandi e quelli più piccoli, da 7 milioni a 50 milioni di coppie immagine/testo, era di soli 20 minuti circa. Entrambi i lavori hanno utilizzato 24 istanze ed entrambi hanno trascorso la maggior parte del tempo semplicemente inizializzando le GPU per il lavoro. Il download dell'immagine, la configurazione di MPI, il caricamento dei dati e la preparazione di NCCL hanno richiesto la maggior parte del tempo di esecuzione.
Il mio tempo per completare 1 epoca su 50 milioni di coppie immagine/testo con 200 GPU? 15 minuti!!
Esatto, 15 minuti. Dalle 18:30 alle 18:45, l'unico momento in cui le mie GPU sono state effettivamente utilizzate. Il lavoro complessivo ha richiesto circa un'ora per essere completato, anche in questo caso la maggior parte del tempo è stata dedicata all'inizializzazione, al provisioning, al caricamento e al caricamento del modello finito.
Per me, questo indica un'enorme opportunità nella progettazione e nell'ottimizzazione dei sistemi GPU. Dobbiamo accettare quella tassa di 40 minuti solo per addestrare i nostri modelli? Non c'è modo! Diventiamo creativi sulle soluzioni per risolvere questo problema.
Per essere chiari, non ho addestrato questo modello alla convergenza. Ho semplicemente testato le prestazioni su una singola epoca, che ho registrato a 15 minuti. Per allenarmi per 1000 epoche, metterei in preventivo 15.000 minuti, che dovrebbero richiedere poco più di 10 giorni.
7/ Distribuisci su hosting Amazon SageMaker
Successivamente, ho distribuito il modello base sull'hosting SageMaker. Come ci si potrebbe aspettare, un modello addestrato su una singola epoca non funzionerà molto bene, quindi ho semplicemente utilizzato la prima versione di Stable Diffusion nel mio account AWS.
Ed ecco la mia immagine risultante!
E questo è un impacco! In questo post ho spiegato come ho lavorato con 10 TB di dati per addestrare una singola epoca di Stable Diffusion sulla formazione distribuita di SageMaker. Il mio ciclo di allenamento è stato completato in 15 minuti!
Spero di aver iniziato a stuzzicare il tuo appetito per ulteriori contenuti sulla preformazione di modelli di visione e linguaggio di grandi dimensioni. Sto scrivendo un libro su questo argomento , che sarà disponibile nell'aprile 2023. A gennaio potrai unirti alla mia cerchia degli autori, partecipare alla festa di lancio, incontrare altri esperti di dominio in quest'area e altro ancora!
![Che cos'è un elenco collegato, comunque? [Parte 1]](https://post.nghiatu.com/assets/images/m/max/724/1*Xokk6XOjWyIGCBujkJsCzQ.jpeg)
