RMSNorm на Trainium: в шесть раз быстрее после fuse или torch.compile
RMSNorm на Trainium: в шесть раз быстрее после fuse или torch.compile Вывод на Trainium вроде медленный, а в профайлере пустая колонка Neuron? Перед разбором включи PrivateUse1 в torch.profiler и вызови synchronize() до export — иначе время на чипе просто не попадёт в отчёт. Amazon в серии про TorchNeuron показала на RMSNorm, откуда берётся разрыв: naive eager с шестью вызовами операторов на устройстве около 2,3 мс, fuse и torch.compile — примерно 400 мкс, в шесть раз быстрее. Там же — как читать Neuron Explorer от фреймворка до железа. Если уже крутите inference на Trainium, прогоните тот же тест на своём слое нормализации, прежде чем лезть в архитектуру модели.
5 просмотров