企业官网建设流程全解析

男女做那个网站动态图,企业网站的建设步骤包括,做网络优化哪家公司比较好,mvc电子商务网站开发Transformer的二次方注意力瓶颈的问题是老生常谈了。这个瓶颈到底卡在哪实际工程里怎么绕过去#xff1f;本文从一个具体问题出发#xff0c;介绍Mosaic这套多轴注意力分片方案的设计思路。 注意力的内存困境 注意力机制的计算公式#xff1a; Attention(Q, …Transformer的二次方注意力瓶颈的问题是老生常谈了。这个瓶颈到底卡在哪实际工程里怎么绕过去本文从一个具体问题出发介绍Mosaic这套多轴注意力分片方案的设计思路。注意力的内存困境注意力机制的计算公式Attention(Q, K, V) softmax(QKᵀ / √d) × V问题出在QKᵀ这个矩阵上它的形状是(序列长度 × 序列长度)。拿150,000个token的序列算一下Memory 150,000² × 4 bytes 90 billion bytes ≈ 84 GB这只是注意力权重本身的开销而且还是单层、单头。A100的显存上限是80GB放不下就是放不下。现有方案的局限FlashAttention它通过分块计算不需要把完整的注意力矩阵实例化出来内存复杂度从O(n²)降到O(n)。单卡场景下效果很好但问题是整个序列还是得塞进同一张GPU。Ring Attention换了个思路把序列切片分到多张GPU上每张卡持有一部分QK和V在GPU之间像传令牌一样轮转一维序列处理起来是很不错的。但是多维怎么办比如处理表格数据的Transformer输入张量形状是(batch, rows, features, embed)。模型需要在不同维度上做注意力features维度只有5个tokenrows维度却有150,000个。前者单卡轻松搞定后者则必须分片。现有的库都没法干净地处理这种多轴场景。手写的话每个轴要单独写分片逻辑进程组管理、张量reshape全得自己来。代码会变得很脏。Mosaic的设计Mosaic本质上是个协调层负责把不同的注意力轴路由到合适的计算后端import mosaic # Small axis: run locally feature_attn mosaic.MultiAxisAttention( embed_dim96, num_heads4, attention_axis2, # features dimension backendlocal # no communication needed ) # Large axis: shard across GPUs row_attn mosaic.MultiAxisAttention( embed_dim96, num_heads4, attention_axis1, # rows dimension backendring # ring attention across GPUs )底层Mosaic会自动处理轴的置换、QKV投影前的reshape、后端分发、以及计算完成后张量形状的还原。模型代码保持清晰分布式的复杂性被封装掉了。Ring Attention的工作机制核心思想其实很直接不需要同时持有全部的K和V。可以分批计算注意力分数逐步累积最后再做归一化。比如说4张GPU的情况下流程是这样的Initial state: GPU 0: Q₀, K₀, V₀ GPU 1: Q₁, K₁, V₁ GPU 2: Q₂, K₂, V₂ GPU 3: Q₃, K₃, V₃ Step 1: Each GPU computes attention with its local K, V GPU 0: score₀₀ Q₀ K₀ᵀ ... Step 2: Pass K, V to the next GPU in the ring GPU 0 receives K₃, V₃ from GPU 3 GPU 0 sends K₀, V₀ to GPU 1 Step 3: Compute attention with received K, V GPU 0: score₀₃ Q₀ K₃ᵀ Accumulate with score₀₀ Repeat for all chunks... Final: Each GPU has complete attention output for its Q chunk单卡内存占用变成O(n²/p)p是GPU数量。8张卡的话内存需求直接砍到1/8。150k序列从84GB降到约10GB每卡。Mesh2D更激进的分片序列特别长的时候Ring Attention的线性分片可能还不够这时候可以用Mesh2D把Q和K都切分了4 GPUs arranged in 2×2 mesh: K₀ K₁ ┌──────┬──────┐ Q₀ │GPU 0 │GPU 1 │ ├──────┼──────┤ Q₁ │GPU 2 │GPU 3 │ └──────┴──────┘ Each GPU computes one tile of QKᵀ内存复杂度降到O(n²/p²)。64张卡组成8×8网格时每卡内存需求下降64倍。attnmosaic.MultiAxisAttention( embed_dim128, num_heads8, attention_axis1, backendmesh2d, mesh_shape(8, 8) )感知集群拓扑的组合策略在实际部署环境里不同GPU之间的通信带宽差异很大。节点内GPU走NVLink能到900 GB/s跨节点通过InfiniBand通常只有200 GB/s左右。ComposedAttention就是针对这种拓扑特征设计的# 4 nodes × 8 GPUs 32 total composed mosaic.ComposedAttention( mesh_shape(4, 8), # (nodes, gpus_per_node) head_parallelTrue, # Split heads across nodes (slow link) seq_parallelring # Ring within nodes (fast link) )需要更精细控制的话可以用HierarchicalAttentionhier mosaic.HierarchicalAttention( intra_node_size8, intra_node_strategylocal, # Compute locally within node inter_node_strategyring # Ring between node leaders )重通信走快链路轻通信才跨节点。实现细节整个库大约800行Python核心代码如下class MultiAxisAttention(nn.Module): def forward(self, x): # 1. Move attention axis to seq position x, inv_perm self._permute_to_seq(x) # 2. Flatten batch dims, project QKV x x.view(-1, seq_len, embed_dim) qkv self.qkv_proj(x).view(batch, seq, 3, heads, head_dim) q, k, v qkv.permute(2, 0, 3, 1, 4).unbind(0) # 3. Dispatch to backend out self._attn_fn(q, k, v) # local, ring, or mesh2d # 4. Project output, restore shape out self.out_proj(out.transpose(1, 2).reshape(...)) return out.permute(inv_perm)后端封装了现有的成熟实现local后端调用F.scaled_dot_product_attention也就是FlashAttentionring后端用ring-flash-attn库的ring_flash_attn_funcmesh2d是自定义的all-gather加SDPA所有的底层都跑的是FlashAttention内核。所有后端统一用FlashAttention的融合GEMMsoftmax实现。后端函数在初始化时就绑定好前向传播不做分支判断。张量操作尽量用x.view()而不是x.reshape()保持内存连续性。集合通信的目标张量预分配好避免torch.cat的开销。模块级别做导入不在每次前向传播时产生import开销。快速上手安装pip install githttps://github.com/stprnvsh/mosaic.git # With ring attention support pip install flash-attn ring-flash-attn单节点启动torchrun --nproc_per_node4 train.py多节点的话# Node 0 torchrun --nnodes2 --nproc_per_node8 --node_rank0 \ --master_addr192.168.1.100 --master_port29500 train.py # Node 1 torchrun --nnodes2 --nproc_per_node8 --node_rank1 \ --master_addr192.168.1.100 --master_port29500 train.py训练脚本示例import mosaic import torch.distributed as dist dist.init_process_group(nccl) ctx mosaic.init(sp_sizedist.get_world_size()) model MyModel().to(ctx.device) # Data is pre-sharded: each GPU has seq_total / world_size tokens x_local load_my_shard() out model(x_local) # Communication handled by Mosaic总结最后Mosaic不会自动并行化模型这个用nnScaler不管数据并行PyTorch DDP/FSDP的事也不处理模型分片交给FSDP或Megatron。Mosaic专注于一件事多轴注意力的分片路由这套方案最初是给nanoTabPFN做的一个表格数据Transformer。这个模型要同时在rows150k个和features5个两个维度做注意力。标准Ring Attention对维度语义没有感知它只认序列这个概念分不清rows和features的区别。所以Mosaic需求很明确小轴本地算大轴分布式算轴的路由逻辑不能侵入模型代码有兴趣的可以试试。https://avoid.overfit.cn/post/791e0f30540e4d289a43d01d383e8ab2作者Pranav Sateesh