AI 助理
文档备案控制台
官方文档
首页 人工智能平台 PAI 分布式训练（DLC） 训练框架说明 提交Cosmos-RL强化学习任务

提交Cosmos-RL强化学习任务

更新时间：
复制 MD 格式
产品详情
我的收藏

Cosmos-RL是NVIDIA推出的异步强化学习框架，专为大语言模型（LLM）设计。您可以在PAI-DLC上通过Custom自定义任务类型运行Cosmos-RL进行强化学习训练。

Cosmos-RL工作原理

Cosmos-RL是一个为 LLM 设计的异步强化学习框架，其架构包含三个核心角色：

  • Controller：协调器。负责管理训练拓扑、分发任务和聚合状态，是整个集群的大脑。一个任务中仅有一个 Controller 实例。

  • Policy：策略模型训练器。负责执行模型参数的更新（训练）。可以有多个实例，通过数据并行（DP）、张量并行（TP）、流水线并行（PP）等方式进行分布式训练。

  • Rollout：数据生成器。负责使用当前策略模型与环境交互，生成训练所需的数据。可以有多个实例，并行地产生数据。

Custom 自定义任务

PAI-DLC的Custom自定义任务类型支持灵活配置多角色任务，适合运行Cosmos-RL等需要多种角色协同的训练框架。

Custom自定义类型任务特点如下：

  • 支持自定义角色配置和启动依赖，提供多种运行策略及弹性伸缩能力。

  • 内置了角色（role）和副本（replica）运行时的变量，用于灵活地组织作业代码。

  • 预定义一些通用作业设置，例如服务透出、SSH免密、本地盘挂载等。

  • 集成了常见训练框架（PyTorch、MPI等）的分布式启动模式，兼容平台已有的作业类型。

操作步骤

本文以使用基础数学问题数据集GSM8K作为训练集，使用GRPO微调Qwen2.5-32B-Instruct模型为例，说明如何创建Cosmos类型的DLC任务。

步骤一：准备数据集和模型

分别下载数据集和预训练模型，并将其上传至对象存储OSS，后续DLC任务将挂载该存储路径。

  • 准备训练数据集

    执行如下命令下载数据集。

    modelscope download --dataset AI-ModelScope/gsm8k --local_dir gsm8k

    将其上传至OSS，例如，上传后OSS路径为：oss://cri-*****-registry.oss-cn-wulanchabu-internal.aliyuncs.com/test/modelscope/cache/AI-ModelScope/gsm8k。

  • 准备预训练模型

    modelscope download --model Qwen/Qwen2.5-32B-Instruct --local_dir .

    将其上传至OSS，例如，上传后OSS路径为：oss://cri-*****-registry.oss-cn-wulanchabu-internal.aliyuncs.com/test/modelscope/cache/Qwen/Qwen2.5-32B-Instruct。

步骤二：创建DLC任务

  1. 登录PAI控制台，在页面上方选择目标地域，并在右侧选择目标工作空间，然后单击进入DLC。

  2. 在分布式训练（DLC）页面，单击新建任务。

  3. 配置如下DLC关键参数，其他参数按需配置即可：

    • 镜像配置：选择镜像地址并填写如下地址：

      dsw-registry-vpc.cn-wulanchabu.cr.aliyuncs.com/pai-build/rl:rl-py310-cuda12.4-torch2.6-202506251931

    • 存储挂载：用于读取训练数据、模型、并存储模型训练结果等。例如：

      • Uri：oss://cri-*****-registry.oss-cn-wulanchabu-internal.aliyuncs.com/test/

      • 挂载路径：/mnt/data/

    • 启动命令：填写如下命令：

      启动命令示例

      export COSMOS_USE_MODELSCOPE=1
export MODELSCOPE_CACHE=/mnt/data/modelscope/cache
export COSMOS_CACHE=/mnt/data/elasticrl/cosmos

export WANDB_MODE=offline
cd /workspace/
export MASTER_ADDR=${PAI_JOB_ID}-controller-0
export COSMOS_CONTROLLER_HOST=${MASTER_ADDR}:19527

export COSMOS_CONTROLLER_LAUNCHER=./cosmos-rl/tools/dataset/gsm8k_grpo.py

if [ "$PAI_REPLICA_TYPE" = "controller" ]; then
    # Create cosmos config directory and generate config file on the fly
    cat > ./cosmos-rl/cosmos-reason1-rl.toml << 'EOF'
redis = "12800"

[train]
resume = "flase"
epoch = 15
output_dir = "/mnt/data/elasticrl/outputs/qwen2-5-32b-p-fsdp2-tp4-r-tp4-pp1-grpo"
epsilon = 1e-06
optm_name = "AdamW"
optm_lr = 1e-06
optm_impl = "fused"
optm_weight_decay = 0.0
optm_betas = [
    0.9,
    0.999,
]
optm_warmup_steps = 20
optm_grad_norm_clip = 1.0
async_tp_enabled = false
compile = true
param_dtype = "bfloat16"
fsdp_reduce_dtype = "float32"
fsdp_offload = false
fsdp_reshard_after_forward = "default"
train_batch_per_replica = 384
sync_weight_interval = 1

[train.train_policy]
type = "grpo"
variant = "dapo"
# 数据集的挂载路径
dataset.name = "/mnt/data/modelscope/cache/AI-ModelScope/gsm8k"
dataset.subset = "main"
dataset.split = "train"
prompt_column_name = "question"
response_column_name = "answer"
reward_function = "gsm8k"
temperature = 0.9
epsilon_low = 0.2
epsilon_high = 0.2
kl_beta = 0.0
mu_iterations = 1
min_filter_prefix_tokens = 1
mini_batch = 4

[train.ckpt]
enable_checkpoint = true
save_freq = 30
max_keep = 5
save_mode = "async"

[rollout]
gpu_memory_utilization = 0.7
enable_chunked_prefill = true
n_generation = 16
batch_size = 24
quantization = "none"
max_response_length = 1024
seed = 42

[rollout.parallelism]
n_init_replicas = 1
tp_size = 4
cp_size = 1
dp_shard_size = 1
pp_size = 1
dp_replicate_size = 1
cp_rotate_method = "allgather"

[policy]
# 预训练模型文件路径
model_name_or_path = "/mnt/data/modelscope/cache/Qwen/Qwen2.5-32B-Instruct"
model_max_length = 1024
model_gradient_checkpointing = true

[policy.parallelism]
n_init_replicas = 1
tp_size = 4
cp_size = 1
dp_shard_size = 2
pp_size = 1
dp_replicate_size = 1
cp_rotate_method = "allgather"

[logging]
logger = ['console']
project_name = "pai-rl"
experiment_name = "qwen-32b-gsm8k-dapo"
EOF

    ./cosmos-rl/tools/launch_controller.sh --port 19527 --config ./cosmos-rl/cosmos-reason1-rl.toml
elif [ "$PAI_REPLICA_TYPE" = "policy" ]; then
    ./cosmos-rl/tools/launch_replica.sh --ngpus 1 --type policy
elif [ "$PAI_REPLICA_TYPE" = "rollout" ]; then
    ./cosmos-rl/tools/launch_replica.sh --ngpus 1 --type rollout
fi

    • 资源来源：选择公共资源。

    • 框架：选择Custom。

    • 任务资源：

      • 类型：分别定义Conroller、Policy、Rollout三种类型节点。

      • 节点数：设置为1。

      • 资源规格：对于微调Qwen2.5-32B-Instruct参考配置如下：

        • Controller：选择CPU实例 ecs.g5.2xlarge

        • Policy：选择GPU实例 ecs.gn7i-c16g1.4xlarge

        • Rollout：选择GPU实例 ecs.gn7i-c16g1.4xlarge

          image

  4. 单击确定创建任务。

附录：Cosmos-RL框架简介

Cosmos-RL是NVIDIA推出的纯异步LLM强化学习训练框架，在训练效率和容错能力上具有以下特点：

  • 超⾼的训练效率。Cosmos-RL并⾏执⾏policy训练、rollout⽣成和dispatcher⽹络拓扑管理，相较于传统colocate训练框架，效率提升2-3倍。

  • 极为出⾊的容错。任⼀policy或rollout节点异常时，集群可迅速重新组⽹并继续当前训练step，无需重启恢复。dispatcher管理进程也支持多节点备份。

  • ⽆感扩缩容。基于step-level的⽹络拓扑管理，可在训练过程中动态增减节点，不影响整体训练流程。

  • 模块化算法和数据⽀持。⽀持GRPO、DAPO、AIPO等主流LLM RL训练算法，也⽀持⾃定义算法。接入新数据集时，只需新增数据集描述⽂件和Data Packer即可。

  • ⾯向物理AI系统的训练⽀持。原⽣⽀持Cosmos推理模型的场景数据，提供图像视频Tensor预处理和缓存，减少CPU开销，提⾼GPU利⽤率。

Cosmos-RL架构图和数据流转图如下：

image.png

上一篇：快速提交MPIJob训练任务下一篇：DLC使用案例汇总