Mirai: LLM inference, RL и distributed training

Continuous batching динамически добавляет и удаляет requests из batch на decode steps, повышая GPU utilization без ожидания завершения всего batch.

В обычном batching сервер может ждать, пока наберется batch, и затем обрабатывать все запросы вместе. В autoregressive LLM decode запросы имеют разную длину, поэтому часть batch быстро заканчивается, а GPU начинает простаивать.

Continuous batching работает на уровне итераций decode: завершенные sequence удаляются, новые sequence добавляются, scheduler поддерживает плотный batch активных токенов. Это повышает throughput и utilization, но требует аккуратного KV cache management, fairness, admission control и trade-off с latency.

В ответе хорошо упомянуть, что prefill и decode имеют разные профили нагрузки, а batching policy может отличаться для них.

Systematic exploration - это управляемый выбор действий для изучения неизвестных, но потенциально полезных областей, а не случайный шум вокруг greedy policy.

В RL агент видит данные, которые сам генерирует. Если он слишком рано станет greedy, он может никогда не узнать про хорошие действия, которые изначально кажутся плохими или редко встречаются. Systematic exploration пытается исследовать пространство действий с учетом uncertainty и потенциальной ценности информации.

Примеры: epsilon-greedy как базовый вариант, UCB, Thompson sampling, count-based exploration, curiosity/intrinsic rewards, ensembles и posterior sampling. Выбор зависит от размера action space, horizon, стоимости ошибки и наличия контекстных признаков.

Проблема в том, что exploration стоит reward прямо сейчас и может быть небезопасным. В production нужны constraints, offline evaluation, simulator или ограниченный rollout.

Tabular UCB плохо масштабируется: нужно много pulls на каждую руку. При большом числе actions нужны priors, contextual sharing, embeddings, Thompson sampling или ограничение candidate set.

UCB добавляет optimistic bonus для малоисследованных действий. При 1000 arms и horizon 2000 на каждую руку приходится очень мало наблюдений, поэтому оценки шумные, а exploration budget огромный. При horizon 20 tabular exploration почти бессмысленен: все руки попробовать невозможно.

Нужны способы переносить информацию между действиями: contextual bandits, embeddings действий, model-based priors, clustering arms, candidate generation, Thompson sampling с информативным prior, offline pretraining или simulator. Если actions имеют структуру, ее нужно использовать; если структуры нет, задача статистически плохо определена.

В production также можно ограничить exploration безопасными top-k кандидатами и использовать conservative rollout.

LoRA обучает меньше параметров, поэтому сильно снижает optimizer/gradient memory; максимальный batch size часто может быть больше, чем при full fine-tuning.

Full fine-tuning обновляет все веса модели, поэтому нужно хранить gradients и optimizer states для огромного числа параметров. Для Adam это особенно дорого: кроме весов появляются моменты и градиенты.

LoRA замораживает базовую модель и обучает только low-rank adapters. Forward все еще требует активации, но trainable parameters и optimizer state намного меньше. Поэтому при прочих равных можно увеличить batch size или обучаться на меньших GPU.

Но LoRA не бесплатна: качество может быть хуже для сильного domain shift, а inference/serving зависит от того, мерджим ли adapters в базовые веса или переключаем их динамически.

BF16 имеет 8 бит exponent как FP32 и меньше mantissa; FP16 имеет больше mantissa, но намного меньший dynamic range. Поэтому BF16 реже ловит overflow/underflow.

FP16 и BF16 занимают 16 бит, но распределяют их по-разному. FP16 хранит больше precision в mantissa, но exponent меньше, поэтому dynamic range ограничен. BF16 сохраняет exponent как у FP32, поэтому лучше переносит большие и маленькие значения, но имеет меньше точности в mantissa.

Для обучения нейросетей dynamic range часто важнее небольшой дополнительной точности, потому что gradients/activations могут сильно различаться по масштабу. Поэтому BF16 часто позволяет обучать без loss scaling или с меньшим количеством numerical issues.

На inference выбор зависит от hardware support, качества, throughput и требований модели.

FSDP шардирует параметры/градиенты/optimizer states, tensor parallelism делит операции внутри слоя, pipeline parallelism делит последовательность слоев между устройствами.

FSDP решает прежде всего проблему памяти: параметры, градиенты и optimizer states хранятся шардами на разных GPU, а нужные веса собираются на время computation. Это удобно, когда модель почти помещается по compute, но не помещается по optimizer/parameter memory.

Tensor parallelism делит матричные операции внутри слоя между GPU, поэтому помогает, когда один слой слишком большой или нужно ускорить compute конкретных блоков. Pipeline parallelism размещает разные слои на разных устройствах и прогоняет microbatches через pipeline; его риск - bubbles и сложность балансировки.

На практике техники комбинируют. Выбор зависит от размера модели, interconnect, batch size, activation memory, optimizer, sequence length и tolerable communication overhead.