DocsICP FilingConsole
官方文档
首页

Submit a Cosmos-RL reinforcement learning job

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

Cosmos-RL is an asynchronous reinforcement learning (RL) framework from NVIDIA for large language models (LLMs). You can run Cosmos-RL on PAI-DLC by creating a Custom job.

How Cosmos-RL works

The Cosmos-RL architecture consists of three core roles:

  • Controller: Manages the training topology, distributes tasks, and aggregates status across the cluster. Each job has exactly one Controller instance.

  • Policy: Performs model parameter updates (training). Multiple Policy instances can run in parallel using data parallelism (DP), tensor parallelism (TP), and pipeline parallelism (PP).

  • Rollout: Generates training data by interacting with the environment using the current policy model. Multiple Rollout instances can run in parallel to produce data concurrently.

Custom job type

PAI-DLC Custom jobs support flexible multi-role configurations, making them suitable for frameworks like Cosmos-RL that require coordinated roles.

Key capabilities of Custom jobs:

  • Custom role configurations and launch dependencies, with multiple run policies and elastic scaling.

  • Built-in runtime variables for roles and replicas to structure job code for different configurations.

  • Predefined common settings such as service exposure, passwordless SSH, and local disk mounting.

  • Built-in distributed launch modes for common training frameworks (PyTorch, MPI, and others), compatible with existing PAI-DLC job types.

Procedure

This example fine-tunes the Qwen2.5-32B-Instruct model with GRPO (Group Relative Policy Optimization), using the GSM8K math dataset for training. The following steps show how to create a Cosmos-RL job on PAI-DLC.

Step 1: Prepare the dataset and model

Download the dataset and pre-trained model, then upload them to Object Storage Service (OSS). The DLC job mounts this storage path at runtime.

  • Prepare the training dataset

    Run the following command to download the dataset:

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

    Upload the dataset to OSS. For example, the OSS path after upload: oss://cri-*****-registry.oss-cn-wulanchabu-internal.aliyuncs.com/test/modelscope/cache/AI-ModelScope/gsm8k.

  • Prepare the pre-trained model

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

    Upload the model to OSS. For example, the OSS path after upload: oss://cri-*****-registry.oss-cn-wulanchabu-internal.aliyuncs.com/test/modelscope/cache/Qwen/Qwen2.5-32B-Instruct.

Step 2: Create a DLC job

  1. Log in to the PAI console. In the top navigation bar, select the target region. In the left-side navigation pane, select the target workspace, and then click Go to DLC.

  2. On the Deep Learning Containers (DLC) page, click Create Task.

  3. Configure the following key parameters. Adjust other parameters as needed:

    • Image Configuration: Select Image Address and enter the following address:

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

    • Storage Mount: Mount the storage path to read training data and models and to store training outputs. Example configuration:

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

      • Mount path: /mnt/data/

    • Start Command: Enter the following command:

      Sample launch command

      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"
# Mount path of the dataset
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]
# Path to the pre-trained model files
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

    • Source: Select Public Resources.

    • Framework: Select Custom.

    • Job Resource:

      • Type: Define three node types: Controller, Policy, and Rollout.

      • Nodes: Set to 1.

      • Resource Type: To fine-tune Qwen2.5-32B-Instruct, use the following recommended instance types:

        • Controller: Select a CPU instance, such as ecs.g5.2xlarge.

        • Policy: Select a GPU instance, such as ecs.gn7i-c16g1.4xlarge.

        • Rollout: Select a GPU instance, such as ecs.gn7i-c16g1.4xlarge.

          Resource configuration

  4. Click Confirm to create the job.

Appendix: Cosmos-RL overview

Cosmos-RL offers the following advantages in training efficiency and fault tolerance:

  • High training efficiency. Cosmos-RL runs policy training, rollout generation, and dispatcher topology management in parallel, achieving 2–3x throughput over traditional co-located training frameworks.

  • Strong fault tolerance. If a Policy or Rollout node fails, the cluster reorganizes and resumes the current training step without a restart. The dispatcher also supports multi-node backup.

  • Seamless scaling. Step-level topology management allows nodes to be added or removed during training without disrupting the overall process.

  • Modular algorithm and dataset support. Cosmos-RL supports mainstream LLM RL algorithms such as GRPO, DAPO, and AIPO, as well as custom algorithms. To add a dataset, create a dataset description file and a Data Packer.

  • Physical AI training support. Cosmos-RL natively supports scenario data for Cosmos inference models, with built-in image and video tensor preprocessing and caching to reduce CPU overhead and improve GPU utilization.

The following figure shows the Cosmos-RL architecture and data flow:

Cosmos-RL architecture and data flow

Previous：Submit an MPIJob training jobNext：DLC use cases