Может ли Рэй приостанавливать/возобновлять задачи в точках синхронизации, когда количество графических процессоров огран ⇐ Python

[Anonymous]

Я обучаю несколько нейронных сетей параллельно с помощью Ray, где сети должны синхронизироваться в определенных точках во время обучения (а не только по завершении), чтобы обмениваться метаданными и обновлять гиперпараметры на основе коллективных результатов.
Моя настройка:

• 32 сети должны обучаться параллельно
• Доступно только 16 графических процессоров
• Сети должны синхронизироваться на указанных границах (например, каждые N шагов обучения)
• В точках синхронизации все сети обмениваются метаданными, прежде чем какая-либо из них сможет перейти к следующей граничной точке.

Проблема:
Поведение Рэя по умолчанию ставит задачи в очередь и выполняет их до завершения. Но мне нужно, чтобы Рэй:

• Обучал сеть на графическом процессоре до точки синхронизации.
• Приостанавливал/передавал графический процессор задаче, поставленной в очередь.
• Возобновил работу сети после точки синхронизации.

Вопрос: Может ли Рэй проверять точку и приостанавливать задачи в произвольных точках (не завершение), освобождать ресурсы графического процессора для других задач, а затем возобновлять работу позже? Или Рэй поддерживает только организацию очереди задач, при которой каждая задача выполняется непрерывно до завершения?
Если Рэй не может этого сделать, какие альтернативы существуют для синхронизированного параллельного обучения с ограниченными ресурсами графического процессора?
Кроме того, программы, которые я запускаю, используют JAX для обучения нейронных сетей

Подробнее здесь: https://stackoverflow.com/questions/798 ... re-limited