Исследователь машинного обучения выпустил open-source фреймворк easy-torch-tpu, который упрощает обучение PyTorch-моделей на Google TPU. Проект призван заменить громоздкий официальный Hypercomputer/torchprime.
Автор несколько месяцев работал с кластерами Google TPU и столкнулся с типичной проблемой — стандартный PyTorch/XLA оказался «болью в шее». Вместо того чтобы мучиться, он написал собственный фреймворк для своих исследований и теперь делится им с сообществом.
Что умеет новый фреймворк
Easy-torch-tpu ориентирован на академические исследования масштаба 1-10 миллиардов параметров на 32-64 чипах TPU. Фреймворк позволяет кастомизировать практически всё через простое наследование классов:
- Архитектуры моделей
- Логику обучения
- Оптимизаторы
- Загрузчики данных
- Стратегии шардинга и реинкарнации
Встроена интеграция с Weights & Biases для трекинга экспериментов и Hugging Face для работы с чекпоинтами. Модели можно сохранять на TPU, а потом загружать на обычных GPU — без танцев с бубном.
Датасеты стримятся напрямую из Hugging Face, предобученные модели тоже можно подтягивать оттуда (если реализуете архитектуру).
Зачем это нужно разработчикам
Google TPU — мощные чипы для обучения ИИ, но работать с ними через стандартные инструменты сложно. Официальная документация часто неполная, поведение фреймворков непредсказуемое. Easy-torch-tpu решает эти боли простотой настройки и гибкостью кастомизации.
Для российских исследователей это особенно актуально — доступ к TPU через Google Cloud Research Credits даёт возможность обучать большие модели без собственного железа.
Проект уже доступен на GitHub с документацией по установке, автор продолжает добавлять примеры моделей и ждёт фидбека от сообщества.