Trwałe autodostrajanie (tylko GPU)

Do generowania niektórych jąder GPU używamy OpenAI Triton. Triton umożliwia generowanie szybkich jąder GPU w przypadku niektórych fuzji, ale musimy dostosować pewne parametry do każdej takiej fuzji.

W przypadku wielu fuzji może to zająć dużo czasu, dlatego umożliwiamy wczytywanie tych wyników automatycznego dostrajania bez konieczności wykonywania innych kroków kompilacji. Automatyczne dostrajanie pamięci podręcznych jest przydatne, jeśli wprowadzimy kilka zmian: połączenia dostępne w pamięci podręcznej będą korzystać z pamięci podręcznej, a pozostałe będą automatycznie dostrajane normalnie.

Wyniki autodostrajania można skopiować/wczytać za pomocą tych parametrów:

--xla_gpu_dump_autotune_results_to=
--xla_gpu_load_autotune_results_from=

Jeśli wskażemy plik .txt lub .textproto, pamięć podręczna zostanie zniszczona w formacie textproto, a w przeciwnym razie zostanie zapisany w formacie binarnego protokołu.

W testach

W testach można też używać trwałego automatycznego dostrajania. Zalecamy korzystanie z niego wtedy, gdy testy są bardzo duże, zwłaszcza gdy wydajność środowiska testowego jest ograniczona.

Działa dobrze tylko wtedy, gdy pamięć podręczna funkcji automatycznego dostrajania zawiera wyniki wygenerowane na tym samym typie GPU, na którym prowadzone są testy.

Przeprowadzanie testu w celu trwałego dostrajania

Załóżmy na razie, że dany test zawsze używa tego samego typu GPU.

  1. Musimy wyeksportować wyniki autodostrajania z testu, np. podając te parametry w poleceniu testowym:

    --test_env=XLA_FLAGS=--xla_gpu_dump_autotune_results_to=TEST_UNDECLARED_OUTPUTS_DIR/autotune_cache.textproto
    --test_sharding_strategy=disabled
    

    Fragmentacja musi być wyłączona, aby prawidłowo otrzymywać pojedynczą pamięć podręczną automatycznego dostrajania dla wszystkich testów.

  2. Następnie musimy przesłać tę pamięć podręczną do naszego repozytorium kodu.

  3. Następnie musimy dodać pamięć podręczną do zależności danych naszego celu testowego i wczytać ją za pomocą zmiennej środowiskowej.

    data = ["test_autotune_cache.textproto"],
    env = {"XLA_FLAGS": "--xla_gpu_load_autotune_results_from=" +
                        "$(execpath test_autotune_cache.textproto)"},
    

    (Można używać fragmentacji w testach, które wczytują wyniki autodostrajania).

Przykładowe testy znajdziesz w 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

Przestarzałość pamięci podręcznej

Jeśli w modelu zostanie wprowadzonych wiele zmian, może się zdarzyć, że pamięć podręczna nie będzie już zawierać wszystkich fuzji, więc test będzie wolniejszy. W takim przypadku musimy ponownie wygenerować pamięć podręczną dostrajania.

To samo dotyczy sytuacji, w której do przeprowadzenia testów zaczniemy używać nowego typu GPU.

Pamięć podręczna może również stać się przestarzała, jeśli kompilator XLA będzie ewoluował i generuje różne fuzje.