Снова про JAX. Если моя книга “Deep Learning with JAX” (https://t.me/gonzo_ML/2926) для вас ещё не является достаточной мотивацией освоить этот продвинутый фреймворк, то вот вам ещё пара крутых свежих материалов: The PyTorch developer's guide to JAX fundamentals https://cloud.google.com/blog/products/ai-machine-learning/guide-to-jax-for-pytorch-developers Короткий гайд по созданию модели для тех, кто привык к PyTorch и хочет сравнить. Пример использует новый Flax NNX API (писал про него в посте JAX things to watch for in 2025, https://gonzoml.substack.com/p/jax-things-to-watch-for-in-2025), но есть также и пример на более старом но всё ещё популярном Flax Linen API. ================================ How to Scale Your Model A Systems View of LLMs on TPUs https://jax-ml.github.io/scaling-book/ Это прям целая книга про скейлинг LLM на TPU. Содержит несколько секций: 1. All About Rooflines https://jax-ml.github.io/scaling-book/roofline/ Объясняет, что такое roofline model (писал про неё когда-то давно тут https://moocaholic.medium.com/hardware-for-deep-learning-part-3-gpu-8906c1644664). Это безумно важно понимать для оптимизации ваших вычислений, не всё определяется флопсами, многие алгоритмы ограничены коммуникацией, в смысле пересылкой данных, не обязательно в распределённой системе, даже при неудачно организованном чтении из памяти можно оказаться в ситуации, когда ваш GPU/TPU используется лишь на 5% и до пиковых флопс как до Луны. 2. What Is a TPU? https://jax-ml.github.io/scaling-book/tpus/ Объясняет, что такое TPU (про это я тоже писал ещё давно https://moocaholic.medium.com/hardware-for-deep-learning-part-4-asic-96a542fe6a81), что такое систолический массив, какие блоки есть внутри, как связываются ускорители в Pod, каковы характеристики разных TPU от v3 до v6e. 3. Sharded Matrices and How to Multiply Them https://jax-ml.github.io/scaling-book/sharding/ Подробно разбирает как работает шардинг (разбиение вычислений больших тензоров по нескольким ускорителям). Параллелизация, Collective Operations -- всё тут. У меня в книге этому посвящена 8-я глава, а кроме неё есть глава 7 про более классический метод параллелизации (pmap) и приложение D про уже устаревшие экспериментальные подходы, которые тем не менее могут помочь лучше понять как мы пришли в текущую точку. 4. All the Transformer Math You Need to Know https://jax-ml.github.io/scaling-book/transformers/ Вся основа трансформеров на уровне вычислений. Где какие операции, как считать флопсы и параметры, MoE, Gradient checkpointing, KV caching, Flash Attention. 5. How to Parallelize a Transformer for Training https://jax-ml.github.io/scaling-book/training/ Обсуждение разных видов параллелизма на примере трансформера: data parallelism, fully-sharded data parallelism (FSDP), tensor parallelism, pipeline parallelism. 6. Training LLaMA 3 on TPUs https://jax-ml.github.io/scaling-book/applied-training/ Как применить все эти знания к обучению реальной модели (Llama 3) на TPU v5p. Что такое Llama 3, как отшардить модель LLaMA 3-70B. 7. All About Transformer Inference https://jax-ml.github.io/scaling-book/inference/ В чём особенности инференса для трансформеров, где боттлнеки, что с памятью, что с latency. MHA, MQA, GQA (про MLA ещё нет: https://t.me/gonzo_ML/3292). KV cache, распределение инференса по разным ускорителям, фазы инференса (prefill, generation), шардирование KV-кэша, много оптимизаций инференса. 8. Serving LLaMA 3-70B on TPUs https://jax-ml.github.io/scaling-book/applied-inference/ Как применить все эти знания к инференсу реальной модели, той же самой Llama 3-70B. 9. How to Profile TPU Programs https://jax-ml.github.io/scaling-book/profiling/ Как профилировать код на TPU и искать боттлнеки. Как работает компилятор XLA, что такое HLO (я много это разбираю в главе 5 своей книги, про компиляцию), что такое JAX TPU profiler, Trace Viewer, Graph Viewer, как делать профилирование памяти. 10. Programming TPUs in JAX https://jax-ml.github.io/scaling-book/jax-stuff/