Home

深入理解 Megatron-LM 中的 Full CUDA Graph:MoE 模型训练加速的关键技术

从Partial CUDA Graph到Full CUDA Graph 在上一篇深入理解 Megatron-LM 中的 Partial CUDA Graph:MoE 模型训练加速的关键技术中,我们分析了在MoE训练流程中,影响CUDA Graph兼容性的主要瓶颈是MoE部分的动态行为。在本文中,我们将深入探讨,为了使MoE部分能够无缝地融入CUDA Graph,并最终实现包含前向与反向传播在内的完整iteration的Full CUDA Graph捕获,需要对其结构做出哪些关键性改进。 具体而言,MoE与CUDA Graph的兼容性亟需解决以下三项核心技术难题: 排除TE grouped GEMM中的CPU-GPU同步环节,实现与CUDA Graph的原生兼容; 对Hyb...

Read more

深入理解 Megatron-LM 中的 Partial CUDA Graph:MoE 模型训练加速的关键技术

为什么近期大模型训练需要用到 CUDA Graph CUDA Graph 是 NVIDIA 推出的一项技术,它可以将一系列 GPU 操作,包括 kernel launch,memory copy 预先在 capture 阶段录制成一个图,然后在 replay 阶段一次性提交执行。CUDA Graph 最大的好处是可以极大减少 CPU 开销,减少 kernel launch latency,这在使用 Grace CPU 的 B 系列 NV 芯片上带来的好处尤其突出。同时 GPU 可以更好地优化执行顺序,减少运行队列中间的 bubble。 针对 B 系列芯片 CPU 开销较大的问题,这里以 MoE 层中 fc1 的 grouped GEMM 操作为例进行说明。以下的实验均在 B 系...

Read more

大模型推理关键技术备忘

简要介绍 nano-vllm 是一个仅约 1200 行代码的轻量级 vllm 实现,尽管体量精简,却高度还原了 vllm 的核心功能。vllm 本身集成了张量并行(tensor parallel)、前缀缓存(prefix cache)、Torch 编译(torch compile)、CUDA 图(cuda graph)、分页注意力(paged attention)等多项优化手段,底层 attention 引擎采用 Flash Attention。基准测试表明,nano-vllm 在推理性能上能够达到接近 vllm 的水平。 在深入学习 nano-vllm 项目后,对大模型推理中的关键技术与核心挑战有了更系统的理解。本文梳理并总结了相关的要点,作为后续查阅与参考的备忘。 RoPE ...

Read more

Ranger: A Smart Guide to the Terminal File Manager

Ranger is a powerful, Vim-inspired file manager for the terminal that can dramatically speed up your command-line workflow. If you find yourself constantly switching between cd and ls to navigate your filesystem, ranger is the tool you’ve been missing. This guide will walk you through everything from installation to advanced customization, helpi...

Read more

从零实现MLX量化矩阵乘法扩展

MLX是苹果公司发布的一个专为Apple Silicon设计的机器学习框架。与PyTorch、JAX类似,MLX提供了类似NumPy的API,支持自动微分,但其独特之处在于针对苹果M系列芯片进行了深度优化。MLX能够充分利用M芯片的统一内存架构,让CPU和GPU无缝共享内存,避免了传统架构中昂贵的内存拷贝开销。Metal是苹果的低级GPU编程接口,类似于CUDA之于NVIDIA或OpenCL之于通用GPU计算。Metal允许开发者直接访问Apple Silicon的GPU资源,编写高性能的并行计算代码。在MLX中,许多计算密集型操作(如矩阵乘法、卷积等)都通过Metal内核实现,以获得最佳性能。 量化则是大语言模型部署中的关键技术,通过将模型权重从FP16/FP32精度降低到INT4...

Read more

分别基于jvp和vjp实现正向微分和反向微分

在PyTorch、JAX、MLX乃至其他现代深度学习框架中,VJP(向量-雅可比乘积)和JVP(雅可比-向量乘积)是实现自动微分的核心机制。最近在学习一些关于敏感度分析,了解到了相关的知识。这里记录一下这两个概念,以及它们是如何驱动“自动求导”过程的。 基本定义 不失去一般性,对于一个函数 \[y=f(x),\ x \in \mathbb{R}^n, y \in \mathbb{R}^m\] 它的雅可比矩阵是: \[J_f = \frac{\partial f}{\partial x} \in \mathbb{R}^{m \times n}\] JVP:Jacobian-Vector Product(正向微分)的定义是,给定一个向量\(v \in \mathbb{R}^n\...

Read more

Solutions to Tensor Puzzlers: A Deep Dive into Tensor Operations

This article presents my solutions to Sasha Rush’s Tensor Puzzles, a collection of challenging tensor manipulation problems. These solutions serve as an excellent reference for understanding fundamental tensor operations and their creative applications. Foundation: Essential Operations Before diving into the solutions, let’s understand two fun...

Read more