DeepSpeed Ulysses 序列并行与 Arctic 长序列训练（ALST）集成指南：面向 HF Transformers 的超长序列训练实战

DeepSpeed Ulysses 序列并行与 Arctic 长序列训练ALST集成指南面向 HF Transformers 的超长序列训练实战【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed本文以仓库教程 docs/_tutorials/ulysses-alst-sequence-parallelism.md 为主体骨架结合 DeepSpeed 源码deepspeed/runtime/sequence_parallel/ulysses_sp.py、deepspeed/runtime/sequence_parallel/parallel_state_sp.py以及激活卸载、单元测试等实现细节展开。它回答了三个问题如何把 Ulysses 序列并行“注入”任意 Hugging FaceHFTransformers 模型如何用 ALST 的“技巧包”激活 checkpoint 卸载到 CPU、分片 MLP 计算、Liger-kernel、PyTorch 显存配置把可训练序列长度推到百万 token 量级以及在集成的常见坑位标签预移位、跨卡 loss 聚合、PyTorch 2.3 兼容上如何正确落地。1. 先弄清楚两种 UlyssesMegatron 版与 HF Transformers 版教程开头交代了一个容易混淆的背景DeepSpeed 生态中存在两代 Ulysses 序列并行实现面向 Megatron-DeepSpeed 的旧版 UlyssesSP只适用于 Megatron-DeepSpeed 训练栈对应的独立教程位于本仓库 docs/_tutorials/ds-sequence.md其代码入口是deepspeed/sequence/layer.py中的DistributedAttention/_DimZeroAllToAll。面向 HF Transformers 的新版 UlyssesSP本文主角核心代码收敛在 deepspeed/runtime/sequence_parallel/ulysses_sp.py。该文件模块级 docstring第 6–29 行明确说明它重新实现了旧版注意力适配层UlyssesSPAttentionHF即旧DistributedAttention的重写统一输入形状为[sl, bs, hc, hs]并适配 HF 的ALL_ATTENTION_FUNCTIONS注意力分发机制并额外提供了数据加载器适配器与一系列分片计算原语。该实现对应的研究为Arctic Long Sequence Training: Scalable And Efficient Training For Multi-Million Token SequencesarXiv:2506.13996。按教程所述这套能力已完整集成进 Snowflake Arctic Training 训练框架同时 DeepSpeed 侧把这套“序列并行 长序列技巧包”以可复用组件的形式开放出来供其它框架或自研训练循环直接引用。1.1 核心能力与目标规模据官方教程教程给出两个层级的能力描述Ulysses Sequence Parallelism for HF Transformers通过“序列维度并行 注意力头并行”实现长序列训练。注意力层把全局序列通过 all-to-all 重新切分为按注意力头并行的本地分片从而把注意力计算与序列长度解耦。Arctic Long Sequence TrainingALST在 Ulysses 之上叠加一套“技巧包”bag of tricks进一步突破显存上限其四大支柱与仓库实现一一对应ALST 技巧仓库实现落点激活 checkpoint 卸载到 CPUpinned memorydeepspeed/runtime/activation_checkpointing/offload_activations.py分片TiledMLP 计算TiledMLP/SequenceTiledCompute同文件 943 / 774 行起Liger-kernel融合分片 loss配套 Liger-kernel 自动生效DeepSpeed 侧等价物为TiledFusedLogitsLossPYTORCH_CUDA_ALLOC_CONFexpandable_segments:True训练启动前的环境配置教程给出的目标规模陈述为在单张 H100 上即可对 Llama-8B 级模型训练50 万 token的序列单节点8 卡可支撑370 万 token仅用4 个节点32 卡即可把 Llama-8B 的序列推到1500 万 token。需要说明的是这些数字属于教程/研究论文描述的目标能力实际可达值取决于模型结构、激活策略与硬件。1.2 阅读本文的代码地图为方便对照先给出本仓库内的关键文件主实现deepspeed/runtime/sequence_parallel/ulysses_sp.py1597 行内含UlyssesSPAttentionHF、UlyssesSPDataLoaderAdapter、sequence_tiled_compute、SequenceTiledCompute、TiledMLP、TiledFusedLogitsLoss、AutogradComputeMLP、TiledLoss、UlyssesSPFwdLossBwdWithLogits序列并行进程组deepspeed/runtime/sequence_parallel/parallel_state_sp.pyMegatron-DeepSpeedparallel_state.py的精简移植组查询接口deepspeed/utils/groups.py_get_sequence_parallel_group等激活卸载deepspeed/runtime/activation_checkpointing/offload_activations.py单元测试tests/unit/ulysses_alst/test_ulysses_sp_hf.py、tests/unit/ulysses_alst/test_tiled_compute.py、tests/unit/v1/sequence_parallelism/test_ulysses.py2. Part 1把 Ulysses 序列并行集成进 HF Transformers 训练循环教程的集成哲学是“只动三个插槽训练循环主体不变”在from_pretrained之前调用UlyssesSPAttentionHF.register_with_transformers(...)完成注意力注入并创建进程组用UlyssesSPDataLoaderAdapter包一层现有 DataLoader让每个 rank 只拿到自己负责的序列分片手工计算“按分片加权的跨卡 loss”因为每卡只持有序列片段与局部 logits。下面完整保留教程给出的、可直接运行的训练脚本骨架使用硬编码玩具数据与hf-internal-testing/tiny-random-LlamaForCausalLM# train.py from deepspeed.runtime.sequence_parallel.ulysses_sp import UlyssesSPAttentionHF, UlyssesSPDataLoaderAdapter from deepspeed.runtime.utils import move_to_device from deepspeed.utils import groups from torch import tensor from transformers import AutoModelForCausalLM import deepspeed import deepspeed.comm as dist import torch model_name_or_path hf-internal-testing/tiny-random-LlamaForCausalLM seq_length 64 sequence_parallel_size 2 micro_batch_size 1 config_dict { train_micro_batch_size_per_gpu: 1, zero_optimization: { stage: 3, }, optimizer: { type: Adam, params: { lr: 1e-3 } }, sequence_parallel_size: sequence_parallel_size, } dtype torch.bfloat16 # a simple Dataset # replace with a real dataset but make sure position_ids are returned input_ids tensor([[1, 10, 10, 10, 2, 2], [1, 20, 20, 20, 2, 2]], ) position_ids tensor([[0, 1, 2, 3, 4, 5], [0, 1, 2, 3, 4, 5]]) ds torch.utils.data.TensorDataset(input_ids, position_ids) def collate_fn(batch): input_ids, position_ids batch[0] return dict(input_idsinput_ids.unsqueeze(0), position_idsposition_ids.unsqueeze(0), labelsinput_ids.unsqueeze(0)) dist.init_distributed(dist_backendnccl, dist_init_requiredTrue) # Ulysses injection into HF Transformers mpu UlyssesSPAttentionHF.register_with_transformers( model_name_or_pathmodel_name_or_path, core_attn_implementationsdpa, sequence_parallel_sizesequence_parallel_size, micro_batch_sizemicro_batch_size, seq_lengthseq_length, seq_length_is_variableTrue, ) # Deepspeed setup model AutoModelForCausalLM.from_pretrained(model_name_or_path) model, _, _, _ deepspeed.initialize(configconfig_dict, modelmodel, model_parametersmodel.parameters(), mpumpu) # UlyssesSPDataLoaderAdapter injection sp_group groups._get_sequence_parallel_group() sp_world_size groups._get_sequence_parallel_world_size() sp_rank groups._get_sequence_parallel_rank() dl torch.utils.data.DataLoader(ds, batch_sizemicro_batch_size, collate_fncollate_fn) dl UlyssesSPDataLoaderAdapter( dl, sp_ranksp_rank, sp_groupsp_group, sp_world_sizesp_world_size, devicemodel.device, ) # Normal training loop for iter, batch in enumerate(dl): batch move_to_device(batch, model.device) outputs model(**batch) # as of this writing HF doesnt calculate loss with shift_labels yet and requires us to do it manually (liger does that automatically) shift_labels batch[shift_labels] loss model.module.loss_function( logitsoutputs.logits, labelsNone, shift_labelsshift_labels, vocab_sizemodel.module.config.vocab_size, ) # differentiable weighted per-shard-loss aggregation across ranks losses_per_rank torch.distributed.nn.functional.all_gather(loss, groupsp_group) # special dealing with SFT that has prompt tokens that arent used in loss computation good_tokens (shift_labels ! -100).view(-1).sum() good_tokens_per_rank torch.distributed.nn.functional.all_gather(good_tokens, groupsp_group) total_loss sum(losses_per_rank[rank] * good_tokens_per_rank[rank] for rank in range(sp_world_size)) total_good_tokens sum(good_tokens_per_rank) loss total_loss / max(total_good_tokens, 1) if dist.get_rank() 0: print(f{iter}: {loss}) model.backward(loss)用 DeepSpeed launcher 在 2 张 GPU 上启动$ deepspeed --num_gpus 2 train.py 0: losstensor(10.4248, devicecuda:0, grad_fnDivBackward0) 1: losstensor(10.4248, devicecuda:0, grad_fnDivBackward0) 2: losstensor(10.3818, devicecuda:0, grad_fnDivBackward0) 3: losstensor(10.3818, devicecuda:0, grad_fnDivBackward0)该示例由仓库单元测试演化而来可对照完整测试 tests/unit/ulysses_alst/test_ulysses_sp_hf.py。下面逐一剖析脚本中“不同于朴素训练循环”的三个关键部位。2.1register_with_transformers注入注意力并把进程组“还”给你mpu UlyssesSPAttentionHF.register_with_transformers( model_name_or_pathmodel_name_or_path, core_attn_implementationsdpa, sequence_parallel_sizesequence_parallel_size, micro_batch_sizemicro_batch_size, seq_lengthseq_length, seq_length_is_variableTrue, )其实现位于 ulysses_sp.py 的 394 行起行为要点可归纳如下1参数取值规则。model_name_or_path既可以是 HF hub 模型名/本地路径也可以是已经实例化的模型对象源码通过hasattr(model_name_or_path, config)/isinstance(..., PreTrainedModel)自动识别还兼容带config属性的 PEFT 包装器。core_attn_implementation指定底层真正执行点积注意力的内核如sdpa、flash_attention_2、flash_attention_3、flex_attention也可以是 transformers 支持懒加载的 hub 托管内核。源码明确拒绝了eager与paged|eagereager 路径总会在局部序列长度上物化 4Dattention_maskO(n²) 内存而序列并行在 all-to-all 拼回全局序列后这类 mask 既失真又爆炸因此 Ulysses 一律丢弃 4D mask、改由position_ids配合is_causal表达因果掩码。2调用时机与一致性校验。该方法必须在from_pretrained之前调用否则 HF 已经按原始注意力实现实例化了模型。若传入的模型配置里已固定_attn_implementation源码会强制它与core_attn_implementation一致防止“包装层拦截不到实际注意力路径”的静默错配。3进程组创建。方法内部调用 parallel_state_sp.py 的initialize_sequence_parallel创建两级组序列并行组世界大小 sequence_parallel_size与序列-数据并行组sequence_data_parallel_size sequence_parallel_size * data_parallel_size用于 ZeRO 梯度归约与参数分片。源码断言世界大小必须能被sequence_parallel_size整除、且sequence_parallel_size 1否则直接抛错。若sequence_parallel_size 1方法直接返回None等价于关闭 SP。返回的mpu对象携带get_sequence_parallel_group/rank/world_size等查询函数需作为mpu参数传给deepspeed.initialize。4“Being John Malkovich”式的注入机制。源码注释解释了为什么不注册新 keyulysses而是覆盖ALL_ATTENTION_FUNCTIONS[core_attn_implementation]HF Transformers 内部大量分支用config._attn_implementation flash_attention_2之类字符串判断走特殊路径若注册新 key 会漏掉这些分支直接覆盖原 key 则保证只有“真正调用注意力函数”这一处被替换成uattn_wrapper其余逻辑原样执行。这也是教程强调seq_length_is_variable等参数必须在register_with_transformers阶段配好的原因。5序列长度参数的分叉。当seq_length_is_variableTrue默认值时UlyssesSPAttentionHF会在每次forward依据实际 batch 形状重新推导本地/全局序列长度此时可以完全不传seq_lengthmpu UlyssesSPAttentionHF.register_with_transformers( model_name_or_pathmodel_name_or_path, core_attn_implementationsdpa, sequence_parallel_sizesequence_parallel_size, micro_batch_sizemicro_batch_size, seq_length_is_variableTrue, )若所有 batch 序列长度严格一致可用固定长度代码路径预先测量全部形状并在后续复用省去每次推导的微秒级开销mpu UlyssesSPAttentionHF.register_with_transformers( # [...] 其余参数 seq_lengthseq_length, seq_length_is_variableFalse, )两条铁律源码第 432–435、487–492 行同样强制固定长度模式下必须给出seq_length否则直接ValueError无论哪种模式seq_length以及每个 batch 的真实序列长度必须能被sequence_parallel_size整除。6头数可除性约束。构造函数会校验attn_head_count % sp_world_size 0且 KV 头数要么能被 SP 规模整除、要么反过来整除覆盖 MQA/GQA 场景必要时对 KV 头做复制从源码可见这些都在 初始化段 以显式ValueError强约束。此外构造函数硬性要求transformers 4.51.3低于该版本直接抛错这是使用该组件的版本前提。2.2UlyssesSPDataLoaderAdapter把“数据并行”变成“序列维度上的 SP-over-DP”dl UlyssesSPDataLoaderAdapter( dl, sp_ranksp_rank, sp_groupsp_group, sp_world_sizesp_world_size, devicemodel.device, )该类包装任意现有 DataLoader实现见 ulysses_sp.py 564 行起。其工作机制可概括为所有 SP rank 协作处理同一条样本——每次“填充”时从各参与 rank 的 DataLoader 各取 1 个 batchall_gather到全体后按序列维度切成长度相等的sp_world_size片把第i片发给 ranki。因此len(新 dl) len(原 dl) * sp_world_size跑满一轮 SP 迭代等价于普通数据并行下的一步。要消费更多 token官方建议配合 DeepSpeed 的梯度累积。代码级细节与约束均在refill中体现batch 字典契约必须包含input_ids、position_ids、labels三个键可含更多键。其中position_ids是硬性要求——源码在缺失时直接抛ValueError并给出提示对未打包序列position_ids torch.arange(seq_len)对打包packed序列position_ids必须在文档边界重置否则分片后经all_gather重建的全局位置错误会破坏sdpa/flex_attention的因果掩码。可变长 batch 的对齐先all_gather各 rank 上input_ids的真实序列长度再把每个张量pad/收集到统一长度任何键的第 1 维都必须与input_ids一致否则报错。形状校验拼接后的全局序列长度必须能被sp_world_size整除否则抛ValueError。tensor 才切分非 tensor 条目如字符串标签原样all_gather_object复制给所有 ranktensor 则在序列维做切片。attention_mask不在 Ulysses 关注范围内4D mask 太大SP 已丢弃它改用小巧的 1Dposition_ids。显存友好的预取缓冲切片完成后先把batch[k][:, chunk_start:chunk_end]挪到CPU源码注释序列超过 1000 万 token、32 卡时预取缓冲可达 GB 级必须留在 CPU 侧待真正迭代时再move_to_device。顺带完成的“标签预移位”适配器会从 batch 里pop掉labels防止 transformers 再算一次 loss把它右侧补一个-100后左移一位得到shift_labels放回 batch。这一步对后续正确计算 loss 至关重要其动机在第 4.2 节展开。2.3 跨卡 loss 聚合可微的加权all_gather由于每个 rank 只处理序列片段、只拥有局部 logits必须自行聚合得到全局 loss。教程强调为了梯度正确必须使用可微的all_gather即torch.distributed.nn.functional.all_gather它会为 gather 回来的张量构造反向传播路径而不是普通dist.all_gather# differentiable weighted per-shard-loss aggregation across ranks losses_per_rank torch.distributed.nn.functional.all_gather(loss, groupsp_group) # special dealing with SFT that has prompt tokens that arent used in loss computation good_tokens (shift_labels ! -100).view(-1).sum() good_tokens_per_rank torch.distributed.nn.functional.all_gather(good_tokens, groupsp_group) total_loss sum(losses_per_rank[rank] * good_tokens_per_rank[rank] for rank in range(sp_world_size)) total_good_tokens sum(good_tokens_per_rank) loss total_loss / max(total_good_tokens, 1)理论上可以简单平均losses_per_rank但系统支持可变长序列——最后一个 rank 的序列片段很可能更短SFT 场景下 prompt 前缀 token 不参与 loss。因此必须按每片“有效 token 数good_tokens”加权平均否则末尾分片因样本少会被低估、prompt 长的任务 loss 会被稀释。max(total_good_tokens, 1)用于防御整片全被 mask 的除零。3. Nuances两个绕不开的集成细节3.1 PyTorch 2.3 的all_gather反向 bug若使用 PyTorch 2.3 以下版本当sequence_parallel_size world_size时反向传播可能触发IndexError: tuple index out of range。这是旧版torch.distributed.all_gather反向实现中的已知问题。官方建议的规避方案不用all_gather做 loss 聚合改用加权all_reduce模式rank 0 先算局部加权和再在 SP 组内做一次all_reduce(SUM)同步。仓库中带有该 workaround 的回归测试在 tests/unit/v1/sequence_parallelism/test_ulysses.py可以当作参考实现。3.2 为什么 labels 必须“预移位”做 batch 分片后不能让上游 loss 函数自己执行标签移位。以 8 个 token、SP2 为例。未分片时 loss 内部做左移一位序列尾补-100input_ids: [1 2 3 4 5 6 7 8 ] labels : [1 2 3 4 5 6 7 8 ] shiftedl : [2 3 4 5 6 7 8 -100]一旦分片成两段各自算 loss每段末尾的“预测下一个 token”会丢失第一段把 token 5 挤成了-100本应预测 5第二段把 token 8 挤成-100本应预测 8input_ids: [1 2 3 4] [5 6 7 8] labels : [1 2 3 4] [5 6 7 8] shiftedl : [2 3 4 -100] [6 7 8 -100]解决方案正是第 2.2 节说的“预移位”在切分之前对全局labels先做一次左移右侧补-100每片拿到的就已经是正确对齐的预测目标input_ids: [1 2 3 4] [5 6 7 8] labels : [1 2 3 4] [5 6 7 8] shiftedl : [2 3 4 5] [6 7 8 -100]这也是 HF Transformers 为此新增“预移位标签 API”的原因Liger-kernel 的 loss 路径会自动处理移位普通 transformers loss 则需要手工把shift_labels喂进去如训练脚本中调用model.module.loss_function(...)。仓库适配器里对应的实现即是 labels 补齐左移那段。4. Part 2ALST 的“技巧包”——把序列长度推到百万 token 量级Ulysses 解决了注意力随序列长度扩展时的通信与计算切分问题但前向激活、MLP 中间张量和最终的 logits 依然可能撑爆显存。ALST 的技巧包针对这三个“显存大头”逐个击破。这些组件在 ulysses_sp.py 中都有对应实现通用分片计算原语SequenceTiledCompute774 行起、专用TiledMLP943 行起、物化-free loss 计算TiledFusedLogitsLoss1067 行起以及把 forward loss backward 串成一步的辅助类UlyssesSPFwdLossBwdWithLogits1422 行起。4.1 分片 loss 计算Tiled loss computation当序列极长时logits张量的形状是[bs, seq_len, vocab_size]物化完整 logits 会瞬间吃掉数十乃至上百 GB 显存。两个解决路径用 Liger-kernel其 loss 会自动以极省内存的方式计算不物化中间完整 logits 张量若模型被 Liger-kernel 支持这是首选。用 DeepSpeed/ALST 的 PyTorch 实现内存占用与 Liger 大致相当但因为是纯 PyTorch、略慢。教程给出了简化版思路——先把hidden_states沿序列维切成num_shards份逐片做lm_head CE用sequence_tiled_compute统一调度并sum归约def loss(self, batch): num_shards 4 outputs model(**batch, use_cacheFalse) hidden_states outputs.last_hidden_state kwargs_to_shard dict( hidden_stateshidden_states, shift_labelsbatch[shift_labels], ) kwargs_to_pass dict(modelmodel, vocab_sizemodel.config.vocab_size) grad_requiring_tensor_key hidden_states compute_params [model.lm_head.weight] seqlen shift_labels.shape[1] total_loss_sum sequence_tiled_compute( loss_fn, seqlen, num_shards, kwargs_to_shard, kwargs_to_pass, grad_requiring_tensor_key, compute_params, output_unshard_dimension0, # loss is a scalar output_reductionsum, ) total_good_items (shift_labels ! -100).squeeze().sum() loss total_loss_sum / max(total_good_items, 1) # differentiable weighted per-shard-loss aggregation across ranks losses_per_rank torch.distributed.nn.functional.all_gather(loss, groupself.sp_group) good_tokens (shift_labels ! -100).view(-1).sum() good_tokens_per_rank torch.distributed.nn.functional.all_gather(good_tokens, groupself.sp_group) total_loss sum(losses_per_rank[rank] * good_tokens_per_rank[rank] for rank in range(self.sp_world_size)) total_good_tokens sum(good_tokens_per_rank) loss total_loss / max(total_good_tokens, 1) return loss对应完整实现支持多层封装、含output_reduction处理位于 Arctic Training 的 SFT trainer 中DeepSpeed 仓库内更“锋利”的等价物是TiledFusedLogitsLoss——作为一个torch.autograd.Function它在forward内部逐分片执行backward从而彻底避免二次前向重算并且同样不物化完整 logits。sequence_tiled_compute的通用封装 720 行起 则把“shard 哪些 tensor / 原样透传哪些 kwargs / 哪个张量需要梯度 / 参与计算的权重列表”抽象成参数compute_params仅在 ZeRO-3 下需要让每个分片把梯度及时交给 ZeRO 的梯度桶。值得注意的设计取舍通用版SequenceTiledCompute为了梯度正确会在backward里重放一次forward叠加激活 checkpoint 则前向执行三次这正是TiledFusedLogitsLoss这类末端融合算子存在的意义——把二次前向省掉。4.2 分片 MLP 计算Tiled MLP computationMLP 的中间激活act_fn(gate(x)) * up(x)是长序列训练的另一显存大头。TiledMLP943 行起把这个“激活逐元素乘下投影”序列切块逐段计算再拼接输出。使用方式是猴子补丁替换模型原有LlamaMLP.forward教程示例from deepspeed.runtime.sequence_parallel.ulysses_sp import TiledMLP import transformers def tiled_mlp_forward_common(self, x): a monkey patch to replace modeling_llama.LlamaMLP.forward and other identical MLP implementations to perform a tiled compute of the same # figure out the number of shards bs, seqlen, hidden x.shape num_shards math.ceil(seqlen / hidden) # its crucial that all ranks run the same number of shards, otherwise if one of the ranks # runs fewer shards than the rest, there will be a deadlock as that rank will stop running # sooner than others and will not supply its ZeRO-3 weights shard to other ranks. So we # will use the max value across all ranks. tensor torch.tensor(num_shards, devicex.device) dist.all_reduce(tensor, opdist.ReduceOp.MAX) num_shards tensor.item() # print(fderived {num_shards} for {seqlen} and {hidden} maxed across ranks) # only needed for deepspeed compute_params [self.down_proj.weight, self.gate_proj.weight, self.up_proj.weight] def mlp_forward(self, x): return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) return TiledMLP.apply( mlp_forward, self, x, num_shards, compute_params, ) from transformers.models.llama import modeling_llama modeling_llama.LlamaMLP.forward tiled_mlp_forward_common这里有一个极易引发死锁的工程细节代码注释专门强调每张卡独立由seqlen / hidden推导分片数时可能不一致——在 ZeRO-3 下某卡少跑几个分片就会提前停等在梯度同步上导致全训练挂起。因此必须先把每卡的num_shards做一次all_reduce(MAX)取全局最大值确保所有 rank 执行相同数量的分片。compute_params只在使用 DeepSpeed ZeRO 时需要用于向 ZeRO 暴露参与计算的权重。完整可运行版本见单元测试 tests/unit/ulysses_alst/test_tiled_compute.py。4.3 激活 checkpoint 卸载到 CPU激活是长序列训练最大的显存消耗方ALST 的做法是把checkpoint 的输入hidden states卸载到主机的 pinned memory反向时再取回。教程区分两条路径可重入式 checkpointtorch.utils.checkpoint的 reentrant 路径Arctic Training 提供原型级 monkeypatcharctic_training.monkey_patches中的monkey_patch_checkpoint_function_with_cpu_offloadDeepSpeed 侧不重复造轮子。非可重入式use_reentrantFalseHF Transformers 默认DeepSpeed 提供专用上下文管理器只把被 checkpoint 的 hidden-state 输入在旁路 stream上搬到 pinned CPU 内存其余保存的张量原样不动from deepspeed.runtime.activation_checkpointing.offload_activations import ( get_checkpoint_hidden_states_offloading_ctx_manager, ) ctx get_checkpoint_hidden_states_offloading_ctx_manager() with ctx: loss model(**batch).loss loss.backward()工程要点源码 offload_activations.py 均有对应实现可对照阅读复用同一个 manager把每个训练 stepforwardbackward都包进它backward()必须在与 forward 相同的上下文内执行否则卸载状态无法回取。对 transformers 的侵入是有条件的只有 manager 激活期间标记 checkpoint 输入的小部件才被安装在GradientCheckpointingLayer上——这样同一模型在其它场景如 HybridEngine rollout复用时不残留任何钩子上下文内所有其它 saved tensor 均原样透传。可调参数及默认值构造器默认值来自 offload_activations.py 65–73 行参数默认值作用use_pin_memoryTrue是否 pin 内存供异步 D2H/H2D 使用use_streamsTrue是否用旁路 CUDA stream 做异步搬运min_offload_bytes1024小于该字节数的张量不卸载阈值max_fwd_stash_count2前向同时在途GPU 侧保留的卸载副本数max_cpu_buffer_pool_count64CPU 侧池化 buffer 上限keep_last_count1最新的 N 个 checkpoint 输入留在 GPU保证第一个 backward 不因“马上要用的激活还在做 D2H”而卡住GPU 上额外滞留的峰值激活数 max_fwd_stash_count keep_last_count。pin 后端注意事项host buffer 通过get_accelerator().pin_memory()分配因此 pin 行为遵循DS_PIN_MEMORY_BACKEND并计入 DeepSpeed 的 pinned-memory 统计。这里应保持默认torch后端native后端只做mlock、不调用cudaHostRegister其 buffer 未经 CUDA DMA 注册旁路 stream 的 D2H 会直接卡死。这与 DeepCompile 的compile.offload_activation_pin_memory选项是同一套 pin 语义。4.4PYTORCH_CUDA_ALLOC_CONF对抗显存碎片最后一个技巧不需要任何代码改动只需在启动脚本前设置环境变量让 PyTorch 缓存分配器使用可扩展段expandable_segments以减少碎片、允许更长序列export PYTORCH_CUDA_ALLOC_CONFexpandable_segments:True长序列训练中前向激活、卸载缓冲、tiled 计算中间量在生命周期上交替伸缩极易产生显存碎片expandable_segments:True使缓存区可动态扩张而非预先整段保留是支撑百万 token 序列的重要环境配置。5. 在仓库中继续验证与学习这套能力的仓库级证据链与可运行样本非常完整建议按以下路径深入完整集成测试tests/unit/ulysses_alst/test_ulysses_sp_hf.py565 行覆盖register_with_transformers DataLoader 适配器 完整训练循环教程示例即由其派生分片计算测试tests/unit/ulysses_alst/test_tiled_compute.py451 行覆盖TiledMLP、SequenceTiledCompute与 loss 分片路径PyTorch 2.3 回归测试tests/unit/v1/sequence_parallelism/test_ulysses.py551 行内含加权all_reduceworkaround 的参考代码源码注释即文档模块级 docstringulysses_sp.py 第 1–30 行对每个 ALST 组件的定位与论文出处有精炼说明适合作为阅读入口。一个需要留意的使用边界是本文所有代码均在“自定义训练循环 deepspeed.initialize”这一路径下成立若使用 HF Trainer需要注意disable_in_eval选项register_with_transformers的disable_in_evalTrue可在评测阶段绕过 SP 操作规避 HF Trainer 评测数据分布不同引发的潜在问题见 UlyssesSPAttentionHF 构造参数以及评测数据需与训练一样遵守“长度可被 SP 规模整除”的约束。6. 小结把两代 Ulysses 的关系、三个注入点注意力注册、数据加载器适配、跨卡加权 loss以及 ALST 的四个技巧串起来这套方案的集成路径非常清晰register_with_transformers负责注意力层维度重排与进程组UlyssesSPDataLoaderAdapter负责序列维度分片并顺手预移位标签tiled 计算原语负责消灭 MLP 中间激活与完整 logits 两个显存黑洞激活卸载与expandable_segments负责把静态激活与碎片内存这两处剩余压力转移掉。对于想在自有框架中复现“多模型、百万级 token”长序列训练的团队本仓库 deepspeed/runtime/sequence_parallel/ulysses_sp.py 与配套测试提供了开箱即用的组件与最小可运行范本直接复用即可避免重新踩过标签移位、跨卡 loss 加权、分片数 all-reduce 同步与 pinned-memory 后端这些暗坑。【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考