Utilizziamo OpenAI Triton per generare alcuni kernel GPU. Triton consente di generare kernel GPU veloci per determinate fusioni, ma dobbiamo ottimizzare alcuni parametri per ogni fusione.
Questa operazione può richiedere molto tempo se ci sono molte fusioni, pertanto forniamo un modo per caricare questi risultati di ottimizzazione automatica, continuando a eseguire normalmente gli altri passaggi di compilazione. Le cache con ottimizzazione automatica sono comunque utili se apportiamo alcune modifiche: le fusioni presenti nella cache utilizzeranno la cache e le altre verranno ottimizzate automaticamente normalmente.
Consigliato: directory della cache
--xla_gpu_per_fusion_autotune_cache_dir=your/directory
Utilizza e gestisci una cache di ottimizzazione automatica per ogni fusione nella directory specificata. Sarà presente un file per fusione distinta.
Il vantaggio principale di questo approccio è che puoi utilizzare la stessa directory della cache per più esecuzioni XLA (di modelli diversi) e la cache crescerà con ogni nuova fusione rilevata, accelerando le esecuzioni successive. È inoltre disponibile il supporto di base per l'esecuzione di più istanze XLA con la stessa directory della cache contemporaneamente.
L'opzione XLA leggerà i risultati esistenti quando necessario e scriverà nuovi risultati dopo che sono stati determinati.
- La directory deve esistere prima dell'esecuzione di XLA e deve essere scrivibile.
- L'invalidamento della cache deve essere gestito dall'utente:
- Utilizza una directory vuota se vuoi iniziare con una cache vuota.
- I controlli della versione XLA devono essere eseguiti dall'utente:
- Se vuoi utilizzare cache separate per versioni diverse di XLA, usa directory diverse.
La cache è disattivata per impostazione predefinita (se non fornisci il parametro).
Limitazione: non è garantito che funzioni bene in combinazione con l'altro metodo di memorizzazione nella cache descritto di seguito.
Alternativa: caricamento o dumping di tutti i risultati di un determinato HLO in un file
I risultati dell'ottimizzazione automatica possono essere scaricati/caricati utilizzando questi parametri:
--xla_gpu_dump_autotune_results_to=
--xla_gpu_load_autotune_results_from=
Se specifichiamo un file .txt o .textproto, la cache verrà scaricata in formato textproto, altrimenti in formato protobuf binario.
Nei test
L'ottimizzazione automatica persistente può essere utilizzata anche nei test. È consigliabile utilizzarlo se i test sono molto grandi, in particolare se le prestazioni dell'ambiente di test sono limitate.
Funziona bene solo se la cache di ottimizzazione automatica contiene risultati generati sullo stesso tipo di GPU su cui vengono eseguiti i test.
Impostare un test in modo che utilizzi l'ottimizzazione automatica persistente
Per il momento supponiamo che il test in questione utilizzi sempre lo stesso tipo di GPU.
Dobbiamo esportare i risultati dell'ottimizzazione automatica dal test, ad esempio specificando questi parametri al comando di test:
--test_env=XLA_FLAGS=--xla_gpu_dump_autotune_results_to=TEST_UNDECLARED_OUTPUTS_DIR/autotune_cache.textproto --test_sharding_strategy=disabledLo sharding deve essere disattivato per ottenere correttamente una singola cache di ottimizzazione automatica per tutti i test.
Poi dobbiamo caricare la cache nel nostro repository di codice.
Poi dobbiamo aggiungere la cache alle dipendenze dei dati del nostro target di test e caricarla utilizzando una variabile di ambiente.
data = ["test_autotune_cache.textproto"], env = {"XLA_FLAGS": "--xla_gpu_load_autotune_results_from=" + "$(execpath test_autotune_cache.textproto)"},È consentito utilizzare lo sharding nei test che caricano i risultati dell'ottimizzazione automatica.
Consulta anche i test di esempio in xla/service/gpu/tests/BUILD:
- load_autotune_results_using_execpath_test
- load_autotune_results_from_test_workspace_test
- dump_autotune_results_to_test_outputs_test
Obsoletudine della cache
Se vengono apportate molte modifiche a un modello, è possibile che la cache non contenga più tutte le fusioni, quindi il test diventerà più lento. In questo caso, dobbiamo rigenerare la cache dell'ottimizzazione automatica.
Lo stesso vale se iniziamo a utilizzare un nuovo tipo di GPU per eseguire i test.
La cache può inoltre diventare obsoleta se il compilatore XLA si evolve e genera fusioni diverse.