RL训推调度与切换（Megatron⇄SGLang）机制解析

本文作者kaiyuanRL训推调度与切换Megatron⇄SGLang机制解析想深耕AI Infra领域欢迎访问InfraTech库内容涵盖大模型基础、PyTorch/vLLM/SGLang框架入门、性能加速等核心方向配套50知识干货及适合初学者的notebook练习:https://github.com/CalvinXKY/InfraTech在采用Ray做资源调度的强化学习RL框架中通常通过placement预占/预留方式进行资源分配与调度将固定资源分配给训练或推理。为提升资源利用率还会把训练与推理进程调度到同一张卡上实现分时复用共卡模式。在此过程中Ray如何管理这些资源过程控制中会用到哪些特性本文以Megatron、SGLang框架为基础结合代码示例逐步介绍Ray在该过程中的管理机制。按知识层次递进依次介绍如下案例Placement的使用Ray调度SGLang多节点SGLang的启动Ray调度Megatron训练侧显存资源的释放以CUDA hook方式卸载训练资源最后介绍“训练与推理交错使用资源的案例。1 预备知识在当前的RL训练中在共卡模式下需要满足Ray的资源管理可控、紧凑避免出现资源随机组合、零散分配的情况训练资源与推理资源能够周期性地释放与加载资源释放过程要足够迅速避免GPU空转。为了达到这些要求需要框架具备若干原子能力下文先介绍相关背景知识。1.1 Ray的预占placementRay提供placement groupPG[1]资源管理方式允许用户在集群中以分组方式使用预留资源。也就是说用PG一次性预订整组资源gang scheduling全有或全无任务只在这组资源上运行。不使用PG时任务会零散地临时申请资源容易出现“这个actor起了、那个actor卡住”的情况。核心概念如下bundle一个“资源包”比如 {“CPU”: 4, “GPU”: 1}placement group由多个 bundle 组成。用户可以指定每个资源束的大小与类型。例如创建一个资源束{“CPU”: 1, “GPU”: 1}表示1个CPU和1个GPU。PG中的关键参数strategy主要决定bundle在节点上的分布方式PACK尽量把bundle塞到尽可能少的节点上偏“聚拢”放不下再扩到别的节点。SPREAD尽量把bundle分散到不同节点上偏“分散”但并非强约束资源紧张时仍可能出现多个bundle落在同一节点。举例创建3个bundles{CPU:2}。 PACK策略bundles调度在一个机器上SPREAD策略bundles调度在不同机器上还有STRICT_PACK强制所有bundle必须在同一节点若单节点放不下PG创建直接失败。STRICT_SPREAD强制每个bundle分布到不同节点若节点数不够或资源不满足PG创建直接失败。在RL场景中训练与推理混跑资源需求往往成组出现因此通常采用PG模式且策略一般选用PACK。1.2 cuHook与显存使用劫持在CUDA较早版本的example[2]中曾介绍过一种hook机制通过劫持CUDA的lib可以改变CUDA API的行为例如修改函数的输入、输出或增加日志打印等。这里以限制PyTorch内存使用为例。某些场景下如多用户共用GPU若需要限制单个框架进程的显存上限常见方式有在框架内修改由框架维护并统计显存使用量调整PyTorch的allocator逻辑并增加限制显存上限的能力(笔者实现了该功能并贡献到torch1.8版本[3])。在CUDA之上设置cuHook控制malloc和free的行为统计malloc开辟的显存量超过阈值时直接返回OOM相当于增加一层显存管理逻辑Manager。方式1偏定制显存卸载可以做到更精细但需要对所用框架逐一改造改动量较大。方式2无需修改框架源码控制面更广凡是走malloc路径的分配都会受到约束。在RL框架中的类似需求是训推共卡时需要具备卸载某一侧训练或推理的显存资源、需要时再恢复的能力。若框架已具备方式1类能力如推理的睡眠模式可直接复用若暂不具备可采用cuHook方案。cuHook的做法是通过hook记录单个框架开辟的显存例如维护一张map当该框架暂停工作时将部分或全部显存数据搬运到内存host memory并释放对应显存当框架恢复工作时再把数据从内存搬回显存。1.3 推理的睡眠模式目前的推理框架vLLM/SGLang通过睡眠模式sleep mode控制模型权重、KV cache、CUDA Graph相关的显存释放。睡眠模式实现了权重与KV cache的显存释放同时尽量保留CUDA Graph相关显存从而在再次唤醒推理引擎时减少图重构其中主要依赖虚拟显存管理等能力详细内容可参考《降低RL训推共卡开销SGLang/vLLM的无缝切换实现与分析》[4]2 基本使用2.1 Placement的使用RayPlacementGroup的示例流程先申请N个带GPU的bundle再逐个探测每个bundle实际落在哪台机器与哪张物理GPU上然后按固定规则重排bundle顺序最后用重排后的bundle索引启动worker验证从logical_slot到physical_gpu的映射可控且可复现。例如将bundle索引与GPU物理编号做反向映射:代码示例参考InfraTech/rl/training_infer_colocate.ipynb[5]Case1。完整地址https://github.com/CalvinXKY/InfraTech/blob/main/rl/training_infer_colocate.ipynb主要是了解PlacementGroupSchedulingStrategy的使用, 关键代码如下pgplacement_group([{CPU:1,GPU:1}for_inrange(num_gpus)],strategystrategy)print(fwaiting placement group ready... (strategy{strategy}, num_gpus{num_gpus}))ray.get(pg.ready())print(placement group ready)infos,bundle_infosprobe_bundle_mapping(pg,num_gpus)reorderedbuild_reordered(bundle_infos,reorder_mode)print(\n[raw bundle mapping])forb,node,gidinbundle_infos:print(fbundle{b}- node{node}, physical_gpu{gid}, ray_gpu_ids{infos[b][all_gpu_ids]})print(\n[reordered logical slots])forlogical_slot,bundle_idxinenumerate(reordered):nodeinfos[bundle_idx][node]gidinfos[bundle_idx][gpu_id]print(flogical_slot{logical_slot}- bundle{bundle_idx}- node{node}, physical_gpu{gid})print_diff_summary(bundle_infos,reordered)# Optional: show how to launch workers with reordered bundle indices.workers[]forlogical_slot,bundle_idxinenumerate(reordered):actorProbe.options(scheduling_strategyPlacementGroupSchedulingStrategy(placement_grouppg,placement_group_bundle_indexbundle_idx,placement_group_capture_child_tasksTrue,)).remote()workers.append((logical_slot,actor))采用reverse的方式启动输出内容如下内容解释[raw bundle mapping]原始bundle-node/gpu映射[reordered logical slots]重排后logical_slot-bundle映射[slot diff summary]每个bundle是否移动以及moved bundles统计其中reorder改变的是“逻辑槽位到物理资源”的对应关系不是物理GPU位置本身。例如reverse下logical_slot0对应原来的bundle7说明重排生效。2.2 Ray调度SGLangRayPlacementGroupSGLang案例采用Ray进行SGLang调度部署。基本执行流程 ray.init()初始化运行时。创建placement_group([{“CPU”:2,“GPU”:2}],strategy“PACK”)预留2卡。用PlacementGroupSchedulingStrategy把SGLangServerActor固定到bundle0。Actor内部启动SGLang服务参数为tp_size2,nnodes1,node_rank0。通过/health_generate轮询健康检查。默认模式下健康后会自动stop并退出加–keep-alive则常驻。SGLang启动服务采用sglang.srt.entrypoints.http_server的launch_server通过子进程完成deflaunch_server_process_local(server_args:ServerArgs)-multiprocessing.Process:multiprocessing.set_start_method(spawn,forceTrue)server_args.hostserver_args.host.strip([])pmultiprocessing.Process(targetlaunch_server,args(server_args,))p.start()_wait_server_healthy(base_urlserver_args.url(),api_keyserver_args.api_key,is_process_alivelambda:p.is_alive(),)returnp代码示例参考InfraTech/rl/training_infer_colocate.ipynb[5]Case2。启动方式python ray_sglang_pg_qwen3_8b_demo.py--keep-alive2.3 多节点SGLang启动RayPlacementGroupSGLang多机示例2台机器各1张卡Qwen3-8B使用placement_groupSTRICT_SPREAD强制两个bundle分布到不同节点每个节点起一个SGLang子进程通过nnodes2/node_rank/dist_init_addr组成一个tp_size2的分布式推理服务节点1Headray start \--head \--node-ip-addressxxx.xxx.xxx.xxx \--port6379\--dashboard-host0.0.0.0\--num-gpus1节点2Workerray start \--addressxxx.xxx.xxx.xxx:6379\--num-gpus1任意节点上运行python ray_sglang_pg_qwen3_8b_multinode_demo.py--addressxxx.xxx.xxx.xxx:6379代码位置InfraTech/rl/training_infer_colocate.ipynb[5]Case32.4 Ray调度Megatron这个案例演示“如何用Ray placement_group预留8卡再在该资源上启动Megatron训练进程”。训练命令使用torchrun单机 8 卡你可以替换 –train-script和 –train-extra-args适配你的Megatron入口。案例中ray触发megatron启动的过程是采用子进程调用torchrun的方式启动方式python ray_pg_megatron_train_qwen06b_demo.py –train-script/root/Megatron-LM/pretrain_gpt.py –use-mock-data代码位置InfraTech/rl/training_infer_colocate.ipynb[5]Case42.5 训练的显存资源释放我们可以直接使用pytorch的.cuda()和.cpu()函数将模型进行换入换出从而释放显存。过程训练阶段模型驻留在GPU上。Offload 阶段将模型迁移到CPU并释放GPU缓存。恢复阶段将模型迁回GPU。主要的函数defoffload_train(self):# Mimic offload-train: move train model state out of GPU.self.modelself.model.cpu()self.devicecpuiftorch.cuda.is_available():torch.cuda.empty_cache()torch.cuda.synchronize()return{phase:offload_train,device_state:self.device,gpu_mem:self._gpu_mem()}defresume_train(self):iftorch.cuda.is_available():self.modelself.model.cuda()self.devicecudatorch.cuda.synchronize()return{phase:resume_train,device_state:self.device,gpu_mem:self._gpu_mem()}代码位置InfraTech/rl/training_infer_colocate.ipynb[5]Case5执行输出对训练框架而言若大量使用.cuda()与.cpu()改动成本较高本案例主要示意基本思路。2.6 cuHook方式卸载训练资源cuHook方式实现的案例目前可参考torch_memory_saver[6]这个库让一部分张量在pause()时释放物理显存在resume()时再重新映射/分配并尽量保持虚拟地址不变。实现上可理解为某一阶段把一大块显存“让出”给其他进程/角色使用下一阶段再快速占回并与原有用法兼容指针、CUDA Graph等。案例演示思路“训练-offload-恢复训练”的基本状态切换训练阶段模型常驻GPU执行一步train_step。Offload阶段调用torch_memory_saver.pause()释放托管显存。恢复阶段调用torch_memory_saver.resume()恢复训练所需显存。代码位置InfraTech/rl/training_infer_colocate.ipynb[5]Case6启动:python ray_pg_megatron_pretrain_inproc_tms_demo.py--tms-preload-so/usr/local/lib/python3.12/dist-packages/torch_memory_saver_hook_mode_preload.abi3.so注意事项需要找到torch_memory_saver的安装so包位置。执行输出3 训推交替使用资源案例演示在同一组GPU资源上交替执行训练与推理这个过程符合RL过程场景的training与rollout交替的场景。相比全量的RL过程去掉了训练与推理之间数据传输、权重同步的过程。了理解代码方便训练过程用PyTorch API实现。训练侧使用Ray worker torch_memory_saver进行sleep/wake切换。推理侧使用SGLang服务进行sleep/resume切换。通过两次/generate请求验证推理在启动后和resume后都可用。代码位置InfraTech/rl/training_infer_colocate.ipynb[5]Case7执行步骤初始化创建共享Placement Group并探测bundle到物理GPU映射。训练启动训练worker完成初始化和第1次train_step。切换训练进入sleeptms.pause释放训练占用。推理启动SGLang并等待健康检查通过第1次/generate验证可用。切换SGLang进入sleep_moderelease_memory_occupation。训练训练wake_up后执行第2次train_step并再次sleep。推理SGLang resume_moderesume_memory_occupation做health与第2次/generate验证。启动:python ray_train_infer_interleave_demo.py训练与推理同占0,1卡时nvidia-smi打印执行输出附可能踩到的坑Megatron路径找不到日志原因镜像内部找不到训练引导的启动脚本。解决指定启动脚本如pythonray_pg_megatron_train_qwen06b_demo.py --train-script /root/Megatron-LM/pretrain_gpt.py健康检查错误日志原因健康检查时间太短。解决增大–sglang-health-timeout-s睡眠模式报错日志原因release_memory_occupation调用时SGLang 里还有正在处理的请求。解决清空SGLang里面在处理的请求在调用睡眠模式。推理健康检查报错日志原因健康检查设置返回不为200报错。解决代码中健康检查放宽限制。模型文件损坏日志原因模型文件损坏用脚本扫描或者直接ll指令查看[INFO]checking5safetensors files under/data/nfs_87/xky/models/Qwen3-8B[OK]model-00001-of-00005.safetensors size3811.12MB keys81[OK]model-00002-of-00005.safetensors size3808.17MB keys114[FAIL]model-00003-of-00005.safetensors size0.00MB errSafetensorError(Error while deserializing header: header too large)[FAIL]model-00004-of-00005.safetensors size0.00MB errSafetensorError(Error while deserializing header: header too large)[OK]model-00005-of-00005.safetensors size1187.00MB keys1SUMMARYtotal:5passed:3failed:2解决重新下载。Ray资源调度卡主日志原因ray调度中检测到资源不满足要求资源冲突即出现调度卡主。解决修改调度的资源要求。校验PG的bundle内容是否正确。Tips当Ray任务疑似卡住无响应时可在终端按CtrlC此时往往仍能看到Ray的状态输出。一般还可增加日志便于排查是否为资源冲突等问题。参考:[1]https://docs.ray.io/en/latest/ray-core/scheduling/placement-group.html[2]https://github.com/NVIDIA/cuda-samples/tree/v12.0/Samples/2_Concepts_and_Techniques/cuHook[3]https://docs.pytorch.org/docs/stable/generated/torch.cuda.memory.set_per_process_memory_fraction.html[4]降低RL训推共卡开销SGLang/vLLM的无缝切换实现与分析[5]abcdefghttps://github.com/CalvinXKY/InfraTech/blob/main/rl/training_infer_colocate.ipynb[6]https://github.com/fzyzcjy/torch_memory_saverInfraTech申明未经允许不得转载