Moore Threads выпустила обновление Torch-MUSA v2.7.0 с более чем 1050 операторами

Китайская компания Moore Threads представила новую версию расширения MUSA для фреймворка глубокого обучения PyTorch — Torch-MUSA v2.7.0. Обновление принесло улучшения в интеграции функций, оптимизации производительности и поддержке оборудования.

Примечательно, что всего за один месяц Torch-MUSA получил два обновления — v2.5.0 и v2.7.0. Начиная с версии v2.5.0, нумерация версий Torch-MUSA синхронизирована с основной версией PyTorch, что упрощает разработчикам идентификацию и управление версиями.

Версия v2.7.0 дополнительно интегрировала библиотеки ускорения вычислений muSolver и muFFT, что значительно повысило эффективность выполнения сложных вычислительных задач. Также добавлена поддержка UMM (Unified Memory) для устройств с единой памятью, что эффективно оптимизирует использование памяти.

Сохраняется совместимость с последними версиями MUSA SDK, поддерживается компиляция с MUSA SDK 4.2.0 до 4.3.0 и выше.

В настоящее время общее количество специализированных операторов, поддерживаемых Torch-MUSA, превысило 1050, что обеспечило дальнейшее повышение производительности и стабильности системы, предоставляя более эффективную и надежную базовую поддержку для обучения и вывода больших моделей.

Следующее обновление будет версией v2.9.0, которая продолжит оптимизацию производительности и функций, постоянно развивая и совершенствуя экосистему глубокого обучения на основе архитектуры MUSA для отечественных полнофункциональных GPU.

▼ Репозиторий Torch-MUSA с открытым исходным кодом:

https://github.com/MooreThreads/torch_musa

Основные обновления в версии v2.7.0

Новые функции

▼ Динамическое преобразование двойной точности (Dynamic Double Cast)

Пользователи могут активировать функцию динамического преобразования операторов типа данных Float64, установив переменную окружения export TORCH_USE_MUSA_DOUBLE_CAST=1. torch_musa будет использовать float32 в качестве типа данных для вычислений.

▼ Распределенные контрольные точки (Distributed Checkpoint)

Поддержка параллельной загрузки и сохранения моделей с нескольких рангов, что значительно ускоряет процесс сохранения и загрузки контрольных точек. В настоящее время уже поддерживается асинхронное сохранение распределенных контрольных точек.

Улучшения функций

▼ Добавлены различные практические операторы, включая Poisson, binomial, _standard_gamma, _sample_dirichlet, vdot, upsample (1d, 2d, 3d, with aa), flash_attention, transformer_encoder_layer. Общее количество специализированных операторов MUSA превысило 1050.

▼ Благодаря обновлению базовой поддержки PyTorch, функции torch.compile и AOTInductor были дополнительно улучшены.

▼ Режим вычислений TF32 включен по умолчанию, что повышает эффективность операций с плавающей запятой.

▼ Оптимизирована стабильность инструмента профилирования производительности Kineto, а его адаптированная версия обновлена до 2.7.0.

▼ Продолжена оптимизация стратегии конвейерного параллелизма FSDP2, что further снижает использование памяти.

Основные обновления в версии v2.5.0

Новые функции

▼ Добавлена интеграция библиотек muFFT и muSolver, что значительно расширяет вычислительные возможности.

▼ Поддержка единого управления памятью в SoC-устройствах для периферийных вычислений. На основе архитектуры Arm UMA (единая адресация памяти) реализовано совместное использование GPU и CPU одного и того же физического пространства памяти, что значительно снижает затраты памяти в процессе работы модели, включая:

Устранение дублирования выделения памяти на стороне GPU;

Сокращение копирования памяти между хостом и устройством;

GPU может напрямую обращаться к пространству памяти, выделенному распределителем CPU.

Расширение операторов и оптимизация производительности

▼ Добавлена поддержка различных операторов, включая ilshift, irshift, replication_pad1d_bwd, angle, ctcLossTensor, ctcLossTensorBwd, logit, amin/amax/prod.dim_int, glu_bwd и другие.

▼ Добавлена поддержка базовых операций Sparse (CSR).

▼ Расширен диапазон поддержки операторов квантования.

▼ Исправлена ошибка формы в torch.norm.

▼ Добавлена поддержка ввода uint8 и вывода int64 для reduce_sum.

▼ В расширениях C++ добавлена поддержка метода tensor.is_musa().

▼ Исправлено аномальное поведение argmax/argmin при пустом вводе.

▼ Оптимизирована эффективность выполнения операций var/std, pad, convolution3d, layer_norm и других.

Улучшения системных функций

▼ Открыт интерфейс torch.musa.mccl.version().

▼ Добавлена поддержка getCurrentMUSABlasHandle и getCurrentMUSABlasLtHandle.

▼ Оптимизирована стратегия конвейерного параллелизма FSDP2, снижено использование памяти при обучении.