Экспериментальный объединённый AdamW / AMSGrad на устройстве

English | 简体中文 | 日本語 | Deutsch | Русский

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

Почему этот оператор

Существующий универсальный оптимизатор создает обновление на основе тензорных операций. Это дополнение явно объединяет немасштабирование градиента, моменты, дополнительный максимум AMSGrad, коррекцию bias, отдельное затухание веса и обновление параметров в одном ядре устройства. Он не утверждает, что существующий бэкэнд fusion обязательно использует много ядер или что эта реализация быстрее, чем оптимизатор Fusion Optimizer PyTorch.

Для каждого конечного входного элемента и обновления на основе единицы t:

g = stored_gradient / gradient_scale      # negate if maximize
m = beta1 * m_old + (1 - beta1) * g
v = beta2 * v_old + (1 - beta2) * g * g
v_used = max(v_max_old, v)                 # AMSGrad only; save uncorrected max
p_new = p * (1 - lr * weight_decay)
        - lr * (m / (1 - beta1^t)) / (sqrt(v_used / (1 - beta2^t)) + epsilon)

Без AMSGrad используется v_used = v . Epsilon находится вне квадратного корня. Затухание весов не добавляется к градиенту или моментам. Хост вычисляет коэффициенты коррекции смещения один раз за вызов, используя возведение FP64 в целую степень и приведение коэффициентов к FP32. Арифметика устройства использует FP32; компилятор может объединять и перегруппировывать операции. Численные сравнения используют допуски, а не гарантию побитного совпадения с другими оптимизаторами.

Ссылка на формулу: PyTorch AdamW . Значения по умолчанию намеренно соответствуют существующему RUDA AdamWConfig для бета, эпсилон и распада: бета=(0.9,0.999), эпсилон=1e-5, weight_decay=1e-4. Задайте все параметры явно при сравнении с другой платформой.

Поддерживаемый контракт

Это не состояние оптимизатора FP8/FP4, GradScaler, автоматическая проверка конечного градиента, осторожное затухание веса, дифференцируемый оптимизатор, общее пошаговое ядро, многотензорный оптимизатор с переменными гиперпараметрами, автоматическая интеграция FSDP или оптимизатор CUDA-графического воспроизведения с счетчик шагов на стороне устройства. Воспроизведение захваченных коэффициентов хоста bias без их обновления не поддерживается. Предварительно сведенный сегмент работает только в том случае, если все элементы имеют общие параметры и количество шагов. Разреженные градиенты, произвольные указатели хостов и тихий резервный вариант CPU не добавляются.

Главное состояние FP32 не означает, что сквозное обучение смешанной точности было проверено. Вызывающая сторона управляет копиями модели, решениями по масштабированию и накопленными или распределенными градиентами.

Feature

• fused-adamw: параметры и явно вызываемая ссылка CPU.
• fused-adamw-device: общая программа запуска устройств и ядро RudaTensor, без каких-либо особенностей. , включенный этой функцией.
• fused-adamw-cuda: среда выполнения CUDA, прямая возможность PTX и явные тесты/пример CUDA. Он по-прежнему выбирает NVRTC или PTX с помощью RUDA_CUDA_COMPILER.

По умолчанию ни одна функция не включена. Никакой новой версии зависимости не представлено. Cargo.lock получает существующий драйвер CUDA только в качестве дополнительной зависимости ruda-optim.

Низкоуровневое использование

use ruda_optim::fused_adamw::{AdamWOptions, StepControl, adamw_step};

// master: dense FP32 RudaTensor<R>, gradient: same-shape F32/F16/BF16 tensor.
// state: Option<AdamWState<R>>, initially None.
let options = AdamWOptions {
    learning_rate: 1e-3,
    weight_decay: 0.01,
    amsgrad: true,
    ..Default::default()
};
let result = adamw_step(
    &master, &gradient, state.as_ref(), &options,
    StepControl { gradient_scale: 128.0, skip_update: found_inf },
)?;
master = result.parameters;
state = result.state;

found_inf передаётся вызывающей стороной, а не вычисляется здесь. Существующий AdamW::step продолжает использовать исходную реализацию. Низкоуровневый примитив не заменяет Parameter модуля автоматически и не строит дифференцируемый граф обновления.

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

AdamWState::into_parts/from_parts предоставляет буферы шагов и моментов для явной интеграции контрольных точек. Они не сериализуются и не передаются сами по себе. Сохранять основные параметры, все моменты, гиперпараметры и шагать вместе; этот модуль не модернизирует формат TrainingRecord высокого уровня.

Модель распределения и производительности

Эта первая реализация выбирает обновления out, чтобы сохранить псевдонимы и избежать добавления новых небезопасных правил владения. Каждый активный шаг выделяет три выхода FP32 (четыре с AMSGrad). Не существует промежуточного дельта-тензора, но нет заявления о шаге с нулевым распределением. Старые и новые буферы могут перекрываться в течение срока службы. Поэтому для больших моделей требуется определенный бюджет памяти; Этот API не включает обновления на месте или повторное использование арены/буферов; отдельный нативный путь описан ниже.

В промежуточном базовом тесте, предназначенном только для тестирования, явно указывается 4 ядра или 5 с AMSGrad. Объединенный путь представляет собой 1. Для стационарного шага градиента FP32:

Счет в 28 байт составляет четыре чтения FP32 (параметр, градиент, два момента) плюс три записи FP32. Базовая линия размером 48 байт включает в себя повторяющиеся чтения градиента и дельта-буфер. Это учет на уровне источника, без учета трафика или ускорения DRAM. Накладные расходы на кэширование, распределение, арифметику и запуск могут изменить наблюдаемый результат. Серверная часть, которая уже объединяет универсальный AdamW, может не принести пользы.

Команды проверки и тестирования

# Python/NumPy/PyTorch formula oracle ONLY, does not execute RUDA.
python tools/run_adamw_regressions.py --suite oracle

# Standalone Rust config/reference tests, without Cargo registry resolution.
python tools/run_adamw_regressions.py --suite reference

# Cargo unit tests, including a comparison with RUDA's existing Host AdamW.
python tools/run_adamw_regressions.py --suite host

# Type-check the opt-in generic device implementation.
python tools/run_adamw_regressions.py --suite build

# Explicit hardware regression, separately under both compilers.
python tools/run_adamw_regressions.py --suite cuda --compiler both

# Small controlled A/B run, then increase elements after checking resources.
python tools/run_adamw_regressions.py --suite bench --compiler both \
    --elements 65536 --dtype bf16 --iterations 20 --samples 7 --amsgrad

Используйте RUDA_PTX_VERSION , поддерживаемый реальным драйвером/GPU. --offline требует кэширования зависимостей Cargo. --timeout ограничивает каждую команду; это не прогнозируемое время выполнения. Бегун сохраняет командные строки, исходные хэши, журналы и статус. Отсутствует Rust/Cargo — blocked ; недоступный GPU не соответствует явному пакету CUDA. Ни один тест не переключает бэкенд автоматически и не запрашивает результат измерения на симуляторе.

В примере исключается создание/обратное чтение входных данных из временных разделов, прогреваются оба варианта, запускаются оба из идентичного материализованного состояния, чередуется порядок выполнения и синхронизируется до/после каждой измеренной партии. Он сообщает медиану/мин/макс стенные миллисекунды на шаг, включая распределение и отправку хоста. Это не таймер CUDA, предназначенный только для GPU. Каждая измеренная партия сравнивает параметры и моменты. Сохраните GPU, драйвер, настройки часов/питания и загрузку оборудования вместе с JSON.

Зафиксированный pytorch_fixtures.json генерируется с помощью реального CPU PyTorch AdamW; 12 случаев охватывают 3 типа градиента x 2 режима AMSGrad x 2 режима максимизации, по 5 шагов каждый. oracle.py --write-fixtures явно регенерирует его. Тесты CUDA сравнивают фактические выходные данные ядра RUDA с этими данными, а также с отдельной ссылкой FP64. Прохождение генератора фикстур Python не означает прохождения этих тестов CUDA.

Остальные ворота

Тип ржавчины/расширение макроса, исполнение CUDA и измеренные характеристики остаются обязательными. Перед выпуском новые тесты должны пройти вместе с предыдущими пакетами регрессий безопасности. Дальнейшая оптимизация должна начинаться с результатов измерений профилировщика, а не с этой модели трафика. Не включайте новый диспетчер по умолчанию до завершения сравнения.

См. экспериментальные проверки градиента и групповое отсечение для получения дополнительной информации.

Нативное хранилище PyTorch и иерархическая статистика

Описанный выше контракт out-of-place adamw_step не меняется. С fused-adamw-device модуль ruda_optim::fused_adamw::storage также предоставляет ядра для заранее выделенной памяти фреймворка: analyze_gradient, merge_gradient_stats, merge_gradient_stats_chunks, adamw_scaled. Это строительные блоки, а не замена безопасного API оптимизатора; адаптер проверяет формы, перекрытия, время жизни и очереди.

stats_plan::StatsPlan::new(rows) принимает 1..=4_194_304 строк статистики и планирует редукции с fan-in 1024. scratch_elements считает элементы FP32 без исходной статистики и итогового отчёта; дополнительная память ограничена 49 200 байтами.

ruda_torch.AdamW использует этот путь при явном fused_step=True и необязательном hierarchical_stats=True. Градиенты не изменяются, до обновлений считывается отчёт из 12 байт. max_grad_norm включает необязательное глобальное отсечение по всем активным группам параметров. Это отдельный путь от out-of-place Rust API, не меняющий стандартный Rust AdamW. См. обучение и параметры checkpoint.