Дефекты безопасности в фреймворке глубокого обучения PyTorch Lightning

В популярном фреймворке глубокого обучения PyTorch Lightning обнаружены многочисленные критические недостатки безопасности.

Уязвимости десериализации, идентифицированные по ссылке VU#252619, затрагивают все версии фреймворка до версии 2.4.0 включительно и потенциально могут позволить злоумышленникам выполнить произвольный код путем загрузки недоверенных файлов моделей.

Обнаружение этих уязвимостей было сделано Казимиром Шульцем из HiddenLayer, а раскрытие информации координировалось Координационным центром CERT (CERT/CC) при Университете Карнеги-Меллона.

PyTorch Lightning, который служит высокоуровневым интерфейсом для PyTorch, упрощает выполнение сложных задач, таких как распределенное обучение и аппаратное масштабирование. Широкое распространение этого фреймворка очевидно: он цитируется в тысячах научных работ, интегрирован в корпоративные конвейеры машинного обучения (ML) и был загружен более 200 миллионов раз по состоянию на март 2025 года, что увеличивает потенциальное влияние.

Функции десериализации, такие как torch.load() и модуль pickle в Python, не имеют встроенных средств защиты, что позволяет злоумышленникам внедрять вредоносный код в файлы моделей (например, в форматы .ckpt или .pt). При загрузке автоматизированными рабочими процессами, такими как конвейеры обучения или сервисы вывода, эти файлы могут вызвать несанкционированное выполнение кода, что приведет к компрометации системы, краже данных или боковому перемещению сети.

CERT/CC выявил небезопасные механизмы десериализации в нескольких подсистемах:

• Распределенная контрольная точка: Функция _load_distributed_checkpoint обрабатывает непроверенные сериализованные данные на всех узлах кластера.
• Модуль Cloud_IO: Получает модели из локальных файлов, удаленных URL-адресов или потоков in-memory без проверки содержимого.
• Процедуры ленивой загрузки: Метод _lazy_load откладывает выполнение без проверки отложенного кода.
• Интеграция с DeepSpeed: Открывает дополнительные поверхности для атак за счет десериализации состояний оптимизатора и модели.
• PickleSerializer: Непосредственно обволакивает модуль pickle в Python, не предлагая никакой «песочницы» или проверки.

До выхода исправления CERT/CC советует организациям:

• Строго соблюдать границы доверия: Избегать загрузки моделей из ненадежных или неаутентифицированных источников.
• Использовать ограниченную десериализацию: Используйте torch.load(weights_only=True), чтобы ограничить загрузку только тензорными данными.
• Изолируйте недоверенные рабочие нагрузки: Обрабатывайте внешние файлы в «песочнице», например в контейнерах с ограниченными привилегиями.
• Проверяйте файлы перед загрузкой: Используйте такие инструменты, как pickletools, для анализа сериализованного содержимого на предмет подозрительных шаблонов.
• Аудит конвейеров автоматизации: Убедитесь, что системы CI/CD и реестры моделей не десериализуют данные, предоставленные пользователем, без соблюдения мер предосторожности.

Уязвимости, обнаруженные в PyTorch Lightning, подтверждают давние опасения по поводу безопасности цепочек инструментов ML. Риски десериализации в PyTorch и связанных с ним фреймворках были задокументированы в течение многих лет, однако многие инструменты продолжают отдавать предпочтение удобству, а не безопасности.

Эти недостатки являются ярким напоминанием о том, что рабочие процессы ML становятся все более привлекательными целями для злоумышленников. Организации должны относиться к файлам моделей с той же осторожностью, что и к исполняемому коду, особенно в общих или производственных средах.