Профилирлеу PyTorch-та: назар аудару – басты назарда

Кіріспе

Жасанды интеллект саласындағы модельдерді тиімді әрі жылдам іске қосу үшін олардың орындау тиімділігін бағалау маңызды болып табылады. PyTorch платформасындағы профилирлеу құралдары осы міндетті шешуге мүмкіндік береді. Бұл мақалада трансформер архитектурасының маңызды элементі – назар аудару (attention) алгоритмін зерттеу және оның орындау тиімділігін өлшеу жолдары қарастырылады. Назар аудару – есептеу шығыны жағынан күрделі компонент болғандықтан, оны түсініп, оңтайландыру маңызды.

Назар аударудың негізгі ұғымдары және профилирлеу

Назар аудару – сұраныстар (Queries), кілттер (Keys) және мәндер (Values) арасындағы өзара әрекеттестікті есептейтін процесс. Бұл процесс бірнеше негізгі басқыштан тұрады:

• Баллдарды есептеу: сұраныс тензорымен кілт тензорының транспонирленген түрін көбейту (матрицалық көбейту)
• Шкалалау: алынған баллдарды белгілі бір коэффициент бойынша үлкейту немесе азайту
• Маска қолдану: себепті маска арқылы кейбір элементтерді мәнсіз ету
• Нормалау: баллдарды жұмсартып, ықтималдыққа айналдыру үшін softmax функциясын қолдану
• Мәндерді қайта салмақтау: алынған салмақтар арқылы мәндер тензорын қайта өңдеу

Бұл қарапайым қадамдар тобы, бірақ олардың әрқайсысы есептеу шығыны мен жадты тұтыну тұрғысынан маңызды рөл атқарады. Профилирлеу – осы операциялардың қайсысы көбірек уақыт пен ресурсты пайдаланатынын анықтау процессі.

Профилирлеудің практикалық мысалдары

«Naive attention» деп аталған қарапайым назар аудару модулі PyTorch-та былай құрылған:

class NaiveCausalAttention(nn.Module):
    def __init__(self, head_dim):
        super().__init__()
        self.scale = 1.0 / math.sqrt(head_dim)

    def forward(self, q, k, v, mask):
        scores = torch.matmul(q, k.transpose(-2, -1))
        scores = scores * self.scale
        scores = scores.masked_fill(mask, float("-inf"))
        attn = torch.softmax(scores, dim=-1)
        out = torch.matmul(attn, v)
        return out

Профилирлеу нәтижелері бойынша операциялардың реті мен жадқа қатысты көшірмелердің болуы анықталды. Ең қызығы – masked_fill операциясынан кейін жад көшірмесінің пайда болуы. Себебі PyTorch әдетте тензормен жұмыс істегенде қосымша көшіру жасайды.

Бұл мәселені шешу үшін inplace операциясын пайдалану мүмкіндігі қарастырылды:

scores.masked_fill_(mask, float("-inf"))

Бұл тәсіл қосымша жад көшірмесін болдырмай, орындау уақытын азайтты. Яғни, профилирлеу графигінде GPU тапсырмаларында жад көшірмесі көрсетілмеді. Сонымен қатар, inplace тәсілі үлкен модельдерде жад үнемдеу тиімділігімен де айырықша маңызға ие.

Контекст және салыстырулар

PyTorch командасы назар аудару алгоритмін оңтайландыру үшін Scaled Dot Product Attention (SDPA) деп аталатын арнайы функцияны ұсынды. Бұл функция бір жолда назар аударудың бірнеше қадамын орындайды және қолданыстағы түрлі бэкэндтердің ішінен ең тиімдісін таңдайды.

Осы функцияның басты ерекшеліктері:

• is_causal параметрі арқылы масканы қолмен құрастыру қажет емес
• Бірнеше бэкэндтің ішінен кіріс мәліметтері мен аппараттық құралға байланысты ең жылдам орындау әдісін таңдау

SDPA-ның негізгі бэкэндтері:

• math – классикалық математикаға негізделген
• flash – жылдам өңдеу технологиясы
• efficient – ресурсты тиімді пайдалану
• cudnn – NVIDIA CUDA кітапханасымен интеграция

Профилирлеу арқылы осы әрбір бэкэндтің жұмысын жеке қадағалап, салыстыруға болады. Бұл оңтайлы таңдау және жылдамдықты арттыру тұрғысынан шешім қабылдауға мүмкіндік береді.

Практикалық маңызы

Үлкен көлемді трансформер модельдерінде назар аударудың тиімді жұмыс істеуі өте маңызды. Әрбір қабатта назар аудару бірлік операциясы қайталанатындықтан, жад үнемдеу және қосымша жүктемені азайту үлгі өнімділігіне тікелей әсер етеді. Inplace операциясын қолдану және арнайы оңтайландырылған SDPA пайдалану модельдің жылдамдығын арттырады, әрі ресурстарды тиімді жұмсауға септігін тигізеді. Hugging Face инфрақұрылымында GPU-ларды оңай баптап, осы тәсілдерді тәжірибеде қолдану мүмкіндігі бар.

Қорытынды

PyTorch-та назар аударуды профилирлеу оның компоненттері мен операцияларын жақсы түсінуге мүмкіндік береді. Inplace операциялар және арнайы оңтайландырылған алгоритмдер модельдің өнімділігін арттыруға үлкен ықпал етеді. Осы тәсілдерді пайдалану трансформерлердің орындалу жылдамдығы мен жадты пайдалану тиімділігін жоғарылатуға көмектеседі. Бұл зерттеу жасанды интеллект саласындағы жүйелердің өнімділігін арттыру жолдарының бірегей мысалы болып саналады.

TAGS: PyTorch, назар аудару, профилирлеу, модель оңтайландыру, GPU, трансформер, жасанды интеллект

Дереккөз: Hugging Face Blog