Utilizziamo OpenAI Triton per generare alcuni dei kernel GPU. Triton consente di generare kernel GPU veloci per determinate fusioni, ma dobbiamo ottimizzare alcuni parametri per ciascuna fusione.
Questa operazione può richiedere molto tempo se sono presenti molte fusioni, quindi forniamo un modo per caricare i risultati dell'ottimizzazione automatica, continuando a eseguire gli altri passaggi di compilazione normalmente. Le cache di ottimizzazione automatica sono comunque utili se apportiamo alcune modifiche: le fusioni presenti nella cache utilizzeranno la cache, mentre le altre verranno ottimizzate automaticamente normalmente.
I risultati dell'ottimizzazione automatica possono essere sottoposti a dump/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à sottoposta a dump in formato textproto, altrimenti in formato binario protobuf.
Nei test
Nei test è possibile usare l'autoregolazione persistente. Ti consigliamo di utilizzarlo se i test sono molto impegnativi, soprattutto 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.
Esecuzione di un test con 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 nel comando 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 target del 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
Obsolescenza della cache
Se vengono apportate molte modifiche a un modello, è possibile che la cache non contenga più tutte le fusioni, quindi il test risulterà più lento. In questo caso dovremmo rigenerare la cache di ottimizzazione automatica.
Se iniziamo a utilizzare un nuovo tipo di GPU per l'esecuzione dei test, lo stesso vale.
La cache potrebbe diventare obsoleta anche se il compilatore XLA si evolve e genera fusioni diverse.