Профилирование TorchNeuron (Часть 1): Системное профилирование на AWS Trainium
huggingface ·
В этой статье представлен системный профилинг для TorchNeuron на AWS Trainium, призванный помочь специалистам по машинному обучению анализировать узкие места в производительности инференса. Описывается настройка NeuronConfig и NeuronProfiler для захвата трассировок устройства и времени выполнения с использованием расширения PrivateUse1 из PyTorch. В руководстве показано, как использовать Neuron Explorer для визуализации жизненного цикла операций — от диспетчеризации во фреймворке до выполнения на оборудовании. Сравниваются три реализации RMSNorm: наивные операции PyTorch, встроенная функция torch.rms_norm и наивная функция, скомпилированная с помощью torch.compile. Результаты показывают, что как фьюжн (объединение) через torch.rms_norm, так и компиляция значительно сокращают время работы устройства Neuron за счет исключения материализации промежуточных тензоров и устранения разрывов между операциями. Подходы с фьюжином требуют примерно одной шестой части времени устройства по сравнению с наивной реализацией. Автор заключает, что выявление «горячих точек» в таблицах key_averages и их проверка в трассировках являются эффективным рабочим процессом для оптимизации.
huggingface