Ошибка памяти при обратном распространении из-за потерь в башне на нескольких тренировках GPU
Я тренирую модель на нескольких графических процессорах. Мне нужно разделить пакет между различными графическими процессорами, а затем рассчитать среднее значение градиентов для всех графических процессоров. Затем я сделаю некоторые другие расчеты, основанные на этом усредненном значении, и я хочу вернуть убытки с помощью этого вычисления обратно до возврата среднего значения к активациям каждого графического процессора. Я получаю ошибку превышения лимита памяти. Я визуализировал график вычислений, и кажется, что часть усреднения должна быть скопирована на все графические процессоры, и я думаю, что это должно быть источником проблемы. Любые идеи о том, как отладить или решить эту проблему?