# =============================================================
# 一键启动需要修改路径的参数如下：
# 1. hydra.searchpath
# 2. verl_conf.extras.data_loader.train_data_path
# 3. verl_conf.actor_rollout_ref.model.path
# 4. train_instances.executor_kwargs.rollout_config.llm_tokenizer_path
#
# 如果使用16卡以外的配置，需修改以下参数：
# 5. n_gpus_per_node：修改为指定的卡数
# 6. verl_conf.actor_rollout_conf.rollout：修改tensor_model_parallel_size和data_parallel_size，
#    保证为2的幂（如2、4、8、16），且相乘的结果=指定卡数
# 7. infer_instances：修改tensor_parallel_size和data_parallel_size，
#    保证为2的幂（如2、4、8、16），且相乘的结果=指定卡数
# =============================================================
agentic_ai:
  mode: direct # 模式选择：serve | direct
  log_level: DEBUG
  log_dir: /var/log/

serve_conf:
  host: 0.0.0.0
  port: 8030

direct_conf:
  entrypoints:
    - job_type: verl_train # 配置: msrl_train | verl_train
      job_name: ${train_instances.0.name} # 即将训练的任务名称（默认为 `train_instances` 的第一个）
      job_kwargs: { }

# =============================================================
# Train 预定义任务：
# 1. 预定义任务包括采用RL框架类型、参数等等，初始化时占用资源，触发后开始训练（支持开始、暂停、取消）
# 2. SERVE 模式下，预定义任务列表默认为空; DIRECT 模式下，初始化任务等待触发训练（触发：direct_conf.entrypoints）
# =============================================================

# 根据rl框架和模型选择不同的模板，模板文件在msrl_conf/verl_conf路径下
hydra:
  searchpath:
    - file:///verl/verl/trainer/config
    - file:///path/to/AgentSDK/aura/configs/train/verl_conf # 改为当前代码仓中aura/configs/train/verl_conf目录所在的绝对路径

defaults:
  - fully_async_ppo_trainer
  - fully_async_ppo_trainer@verl_conf
  - _self_

# 覆盖VERL模板中的参数
verl_conf:
  extras:
    delta: 2  # 清理 x 个版本前的旧权重
    update_weights_interval: 1  # 权重更新间隔
    validate_num_samples: 5 # 验证数据采样
    init_num_group_batches: 1 # 从 DataLoader 首次加载数据批次数量
    max_queue_size: 256
    weight_save_dir: ${hydra:runtime.cwd}/weights
    consumed_train_samples: 0 # 已处理的样本数量
    data_loader:
      train_data_path: /path/to/data/math_train_dataset/rl # 分离模式使用的数据集路径，末尾需要加上/rl
      split: 100,0,0
      seq_length: 8192
      train_iters: ${verl_conf.total_train_steps}
      global_batch_size: ${verl_conf.data.train_batch_size}
      num_workers: 1
      seed: 0
      dataset_additional_keys: ['labels']
      no_shuffle: False
      test_data_path:
      is_pairwise_dataset: false
  total_train_steps: 200
  data:
    gen_batch_size: 1  # Number of samples generated, currently only support 1, 当前不生效
    return_raw_chat: True
    train_batch_size: 16
  actor_rollout_ref:
    model:
      path: /path/to/models/Qwen3-30B-A3B # 需要修改为模型权重路径
      use_remove_padding: True
      enable_gradient_checkpointing: True
      enable_activation_offload: True
    hybrid_engine: False
    rollout:
      n: 8
      name: vllm
      mode: async
      calculate_log_probs: True
      log_prob_use_dynamic_bsz: False
      log_prob_micro_batch_size_per_gpu: 1
      expert_parallel_size: 16
      tensor_model_parallel_size: 8
      data_parallel_size: 2
      gpu_memory_utilization: 0.2
      free_cache_engine: true
    actor:
      ppo_micro_batch_size_per_gpu: 1
      ppo_mini_batch_size: 4
      strategy: fsdp2
      use_dynamic_bsz: False
      optim:
        lr: 1e-6
      entropy_coeff: 0.0
      use_kl_loss: False
      calculate_entropy: False
      kl_loss_coef: 0.001
      kl_loss_type: low_var_kl
      clip_ratio_high: 0.28
      loss_agg_mode: seq-mean-token-sum
      ulysses_sequence_parallel_size: 8
      fsdp_config:
        strategy: fsdp2
        param_offload: true
        optimizer_offload: true
        model_dtype: bfloat16
        fsdp_size: -1
        entropy_from_logits_with_chunking: true
        entropy_checkpointing: true
        ulysses_sequence_parallel_size: ${verl_conf.actor_rollout_ref.actor.ulysses_sequence_parallel_size}
        reshard_after_forward: True
    ref:
      strategy: fsdp2
      log_prob_micro_batch_size_per_gpu: 1
      log_prob_max_token_len_per_gpu: 16384
      ulysses_sequence_parallel_size: ${verl_conf.actor_rollout_ref.actor.ulysses_sequence_parallel_size}
      fsdp_config:
        model_dtype: bfloat16
        param_offload: true
        ulysses_sequence_parallel_size: ${verl_conf.actor_rollout_ref.actor.ulysses_sequence_parallel_size}
  critic:
    strategy: fsdp2
    ulysses_sequence_parallel_size: ${verl_conf.actor_rollout_ref.actor.ulysses_sequence_parallel_size}
    model:
      ulysses_sequence_parallel_size: ${verl_conf.actor_rollout_ref.actor.ulysses_sequence_parallel_size}
    fsdp_config:
      model_dtype: bfloat16
  trainer:
    val_before_train: False
    device: npu
    critic_warmup: 0 # 奖励预热
    project_name: 'qwen3_30b_a3b_math_code_one_step'
    experiment_name: 'qwen3_30b_a3b_math_code_one_step'
    n_gpus_per_node: 16
    nnodes: 1
    save_freq: 100 # 保存ckpt的per iter
    default_local_dir: ${hydra:runtime.cwd}/weights/rollout_data_dir/ # ckpt
    resume_mode: disable # disable / auto (load from default_local_dir) / resume path (load from resume_from_path)
    resume_from_path: ./ckpt
    test_freq: 1000
    total_epochs: 1000 # RL epoch
    logger: [ 'console','tensorboard' ]
    rollout_data_dir: ${hydra:runtime.cwd}/weights/rollout_data_dir/
  rollout:
    nnodes: ${verl_conf.trainer.nnodes}
    n_gpus_per_node: ${verl_conf.trainer.n_gpus_per_node}
    n: 8
    total_rollout_steps: 100
    test_freq: 100
  async_training:
    use_rollout_log_probs: False
    staleness_threshold: 0.25
    trigger_parameter_sync_step: 1 # 多久 iter 同步一次权重
    partial_rollout: False
  algorithm:
    adv_estimator: grpo

train_instances:
  - name: RL-QWEN3-32B-WITH-MATH
    executor_num: 1
    executor_kwargs:
      # 以下参数为 TrainExecutor 初始化参数
      # hybrid(共卡) | one_step_off(分离)
      work_mode: one_step_off
      train_engine: verl
      train_config: ${verl_conf}
      rollout_config:
        train_backend: ${train_instances.0.executor_kwargs.train_engine}
        trajectory_timeout: 7200
        n_samples_per_prompt: ${verl_conf.rollout.n}
        actor_rollout_dispatch_size: 256  # Rollout一次处理的数量 = gbs * n，最大是总样本数16*8/推理实例数2 且是n_samples_per_prompt的倍数
        validate_n_samples: ${verl_conf.extras.validate_num_samples}
        traj_output_path: ${hydra:runtime.cwd}/outputs/
        tokenizer_name_or_path: ${verl_conf.actor_rollout_ref.model.path}
        dataset_additional_keys: ${verl_conf.extras.data_loader.dataset_additional_keys}
        global_batch_size: ${verl_conf.extras.data_loader.global_batch_size}
        weight_save_dir: ${verl_conf.extras.weight_save_dir}
        trust_remote_code: true
        infer_tensor_parallel_size: ${infer_instances.0.executor_kwargs.engine_kwargs.tensor_parallel_size}
        train_tensor_parallel_size: ${verl_conf.actor_rollout_ref.rollout.tensor_model_parallel_size}
        infer_expert_parallel_size: ${infer_instances.0.executor_kwargs.engine_kwargs.data_parallel_size}
        enable_version_control: false
        use_on_policy: false
        train_iters: ${verl_conf.extras.data_loader.train_iters}
        data_optimized: false
        hybrid_batch_num: 1
        wait_available_weight_timeout: -1
        llm_tokenizer_path: /path/to/models/Qwen3-30B-A3B # 需要修改为模型权重路径
      agent_service: ${agent_instances.0.name}  # 依赖的 AGENT 服务名称（默认为 `agent_instances` 的第一个）
      infer_service: ${infer_instances.0.name}  # 依赖的 INFER 服务名称（默认为 `infer_instances` 的第一个）
    resource_info: [ ]  # MSRL 训练采用 RAY 实现资源分配，在 Executor 层不做资源限制


# =============================================================
# Agent 预置服务：
# 1. 提供如：Math / WebSearch 等智能体对话服务
# =============================================================
agent_instances:
  - name: MATH
    executor_num: 1
    executor_kwargs:
      # 以下参数为 AgentExecutor 初始化参数
      agent_engine: rllm
      agent_engine_kwargs:
        # 以下参数为 rLLMEngineWrapper 初始化参数
        agent_name: math # RLLMEngineWrapper内置Agent，修改参见：agents/agents_mapping.py
        simplify_think_content: false
        max_prompt_length: 8192
        max_model_len: 32768
        n_parallel_agents: 512
        token_in_token_out: true
        env_args:
          max_steps: 5
          tool_timeout: 200  # 工具调用超时设置，默认200
          trajectory_timeout: 7200  # 单轮轨迹请求超时，默认7200
          tool_url: /home/inputs/tool_url_2 # 行业运维场景为代码生成, tool_url没有用到
        tokenizer: ${verl_conf.actor_rollout_ref.model.path}
      infer_service_params: # 每轮对话时的 sampling params
        temperature: 1.0
        top_p: 1.0
        max_tokens: 8192
        model: ${infer_instances.0.name} # 依赖的 INFER 服务名称（默认为 `infer_instances` 的第一个）
      trajectory_save_dir: ${hydra:runtime.cwd}/outputs
    resource_info: [ ]


# =============================================================
# Infer 预置服务：
# 1. 提供如：Qwen / Deepseek 等大模型推理服务
# 2. 提供 OpenAI Chat 推理接口
# =============================================================
infer_instances:
  - name: Qwen3-30B-A3B
    executor_num: 1
    executor_kwargs:
      # 以下参数为 InferExecutor 初始化参数
      engine: vllm_proxy
      engine_kwargs:
        # 以下参数为 vLLMInferServer 初始化参数
        # vllm_proxy模式不需要手动配置, 给下列默认配置即可, 启动脚本会自动修改
        chat_server: "http://0.0.0.0:8080"
        prefill_server_list: ["http://0.0.0.0:20012"]
        decode_server_list: []
        model_name: Qwen3-30B-A3B
        tensor_parallel_size: 8
        expert_parallel_size: 16
        data_parallel_size: 2
        enable_expert_parallel: true
    resource_info: []