Материал подготовлен автоматически по первоисточникам: ссылки на них — в конце статьи.
Значение, которое обратный проход читает после округления тензора, меняет сам алгоритм обучения, а не только расход памяти и точность вычислений. В препринте Shuxiao Xie и коллег, который не прошёл рецензирование и приводит замеры самих авторов, три запуска модели на 390 млн параметров сорвались при повторном чтении округлённого выхода механизма внимания, но завершились нормально с новым округлением из того же распределения. При низкоточном обучении недостаточно проверить итоговую функцию потерь: правила чтения состояния нужно задавать отдельно для каждого градиента.
В прямом проходе FP8 округляет тензоры перед вычислениями. Затем обратный проход может получить округлённое значение, исходное значение до округления или новую случайно округлённую копию. Авторы сравнивали результат с градиентом фактически выполненного прямого прохода и проверяли, какое значение требуется каждой операции.
Один тензор может требовать разных правил чтения. Если выход нормализации округлили перед следующим слоем, градиент коэффициента нормализации должен использовать исходное значение, а градиент весов следующего слоя — округлённое значение, которое этот слой действительно умножал. Поэтому даже точная исходная копия не всегда даёт правильный градиент.
В первой паре запусков повторное чтение округлённого выхода закончилось функцией потерь 7,314, а новое округление — 3,256, как и контрольный вариант. Замена только обратного прохода механизма внимания с FP8 на BF16 также устранила срыв, хотя прямой проход остался в FP8. Сравнимая итоговая функция потерь при этом не гарантировала правильный градиент: ошибка могла сохраняться до конца обучения.
Основные опыты проводили на модели из семейства FOG, сравнивая варианты с FP8, BF16 и контрольным вычислением без сохранённой копии выхода. Отдельно правила проверили без обучения на операциях PyTorch и Transformer Engine. Результат касается прежде всего систем, которые сжимают активации, сохраняют округлённые состояния или пересчитывают их ради экономии памяти: такая оптимизация должна явно фиксировать, что читает каждое вычисление градиента.
Источники
Похоже на вашу задачу?
Расскажите, что собираете. За полчаса разложим на этапы и назовём сроки — это бесплатно и ни к чему не обязывает.
