diff --git a/verl_ascend_practice/patches/mtp_checkpoint_engine_reoffload.patch b/verl_ascend_practice/patches/mtp_checkpoint_engine_reoffload.patch new file mode 100644 index 0000000..0311978 --- /dev/null +++ b/verl_ascend_practice/patches/mtp_checkpoint_engine_reoffload.patch @@ -0,0 +1,12 @@ +diff --git a/verl/workers/engine_workers.py b/verl/workers/engine_workers.py +index 57b7da64e..97638bd51 100644 +--- a/verl/workers/engine_workers.py ++++ b/verl/workers/engine_workers.py +@@ -755,4 +755,7 @@ class ActorRolloutRefWorker(Worker, DistProfilerExtension): + return metrics or {} + per_tensor_param, _ = self.actor.engine.get_per_tensor_param() + metrics = await self.checkpoint_engine.send_weights(per_tensor_param, global_steps=global_steps) ++ if self.actor.engine.is_param_offload_enabled: ++ self.actor.engine.to("cpu", model=True, optimizer=False, grad=False) ++ aggressive_empty_cache(force_sync=True) + return metrics or {} diff --git a/verl_ascend_practice/reports/mtp/README.md b/verl_ascend_practice/reports/mtp/README.md new file mode 100644 index 0000000..4112970 --- /dev/null +++ b/verl_ascend_practice/reports/mtp/README.md @@ -0,0 +1,195 @@ +# MiMo-7B-RL MTP Ascend 整体交付报告 + +对应任务:[verl-ascend-recipe #20](https://github.com/verl-project/verl-ascend-recipe/issues/20) + +## 1. 交付概览 + +本交付提供 MiMo-7B-RL MTP 在 Ascend NPU 上的可复现训练 recipe。训练侧使用 Megatron +与 MBridge 管理 actor 和 reference model,rollout 侧使用 SGLang-Ascend,训练入口为 +`verl.trainer.main_ppo`。 + +| 项目 | 配置 | +| --- | --- | +| 模型 | XiaomiMiMo/MiMo-7B-RL | +| 数据集 | MATH parquet | +| 算法 | GRPO + MTP auxiliary loss | +| 训练后端 | Megatron actor/reference | +| Rollout 后端 | SGLang-Ascend | +| 验证平台 | Atlas 800T A2,4 x Ascend 910B | +| 训练规模 | 100 global steps | +| 运行脚本 | `verl_ascend_practice/run_mimo_7b_mtp_megatron_npu.sh` | +| 必要补丁 | `verl_ascend_practice/patches/mtp_checkpoint_engine_reoffload.patch` | + +## 2. 适配方案 + +MTP 训练通过以下配置启用: + +```text +actor_rollout_ref.model.mtp.enable=True +actor_rollout_ref.model.mtp.enable_train=True +actor_rollout_ref.model.mtp.enable_rollout=False +actor_rollout_ref.model.mtp.mtp_loss_scaling_factor=0.1 +actor_rollout_ref.model.mtp.detach_encoder=True +actor_rollout_ref.rollout.name=sglang +actor_rollout_ref.rollout.checkpoint_engine.backend=nccl +model_engine=megatron +``` + +MiMo-7B-RL 的 MTP head 参与 Megatron actor 训练,rollout 仍使用主模型生成。SGLang 在生成 +阶段计算 rollout log probability,`rollout_correction.bypass_mode` 将其复用为 PPO old-policy +log probability。权重更新复用 verl 通用 checkpoint engine;必要补丁在同步完成后恢复已启用 +的 actor parameter offload。 + +```text +MATH prompts + | + v +SGLang-Ascend TP1 rollout (4 replicas) + | + v +math reward + GRPO advantage + | + v +PPO objective + MTP auxiliary loss + | + v +Megatron TP2 actor update + | + v +generic checkpoint-engine weight synchronization + | + v +actor parameter re-offload +``` + +### 2.1 关键训练配置 + +| 配置项 | 值 | +| --- | ---: | +| train / validation batch size | 16 / 16 | +| PPO mini batch size | 8 | +| micro batch size per NPU | 1 | +| rollout responses per prompt | 2 | +| prompt / response length | 512 / 1024 | +| actor learning rate | `1e-6` | +| actor TP / PP / CP | 2 / 1 / 1 | +| rollout TP / replicas | 1 / 4 | +| rollout max sequences per replica | 8 | +| rollout max batched tokens | 8192 | +| rollout memory utilization | 0.45 | +| graph max batch size | 8 | +| weight synchronization bucket | 1280 MB | +| MTP loss scaling factor | 0.1 | +| precision | BF16 | +| checkpoint interval | 100 steps | + +actor 启用参数、梯度和 optimizer offload,并使用完整 activation recompute 与动态 token +batch。SGLang-Ascend 启用 graph mode、chunked prefill 和 prefix caching。 + +## 3. 环境与数据准备 + +### 3.1 已验证软件环境 + +| 组件 | 版本 | +| --- | --- | +| CANN | 9.0 | +| Python | 3.11.15 | +| torch / torch-npu | 2.8.0 / 2.8.0.post2 | +| SGLang / sgl-kernel-npu | 0.5.10 / 2026.2.1 | +| Megatron Core / MindSpeed | 0.16.0 / 0.16.0 | +| MBridge | 0.15.1 | +| Transformers | 5.3.0 | + +### 3.2 模型与数据准备 + +下载 `XiaomiMiMo/MiMo-7B-RL`,并将模型 `config.json` 中的 +`max_position_embeddings` 设置为 `32768`。 + +在 verl 根目录生成 MATH parquet: + +```bash +python3 examples/data_preprocess/math_dataset.py \ + --local_save_dir /path/to/math +``` + +数据目录应包含: + +```text +/path/to/math/train.parquet +/path/to/math/test.parquet +``` + +## 4. 运行与恢复 + +在 verl 根目录应用必要补丁并启动训练: + +```bash +git apply /path/to/verl-ascend-recipe/verl_ascend_practice/patches/mtp_checkpoint_engine_reoffload.patch + +DEVICE=npu \ +MODEL_PATH=/path/to/MiMo-7B-RL \ +DATA_ROOT=/path/to/math \ +NPUS_PER_NODE=4 \ +TOTAL_TRAINING_STEPS=100 \ +bash /path/to/verl-ascend-recipe/verl_ascend_practice/run_mimo_7b_mtp_megatron_npu.sh +``` + +模型、数据、设备数、并行度、batch、长度、MTP loss scaling、checkpoint 和日志目录均可 +通过环境变量覆盖;额外参数会作为 Hydra overrides 继续传递给 `verl.trainer.main_ppo`。 + +默认在 step 100 保存 checkpoint。使用相同的 `OUTPUT_DIR` 重新启动时, +`RESUME_MODE=auto` 会恢复最新 checkpoint。训练日志默认写入 `LOG_DIR`。 + +## 5. 100-step 长跑结果 + +本次训练在 Atlas 800T A2 的 4 x Ascend 910B 上连续完成 100/100 global steps,并完成 +step 100 checkpoint。100 个 MTP loss 和 reward 样本均为有限值。 + +| 指标 | 结果 | +| --- | ---: | +| 连续训练步数 | 100 / 100 | +| reward 首 10 步均值 | -0.808044 | +| reward 末 10 步均值 | 0.089099 | +| reward 首尾窗口增量 | +0.897144 | +| MTP loss 首 10 步均值 | 0.597910 | +| MTP loss 末 10 步均值 | 0.355875 | +| MTP loss 首尾窗口变化 | -0.242035 | +| 4 NPU 端到端吞吐 | 533.286 token/s | +| 4 NPU 稳态吞吐 | 801.923 token/s | +| 稳态吞吐 | 200.481 token/s/NPU | +| actor 峰值分配 / reserved 显存 | 48.116 / 54.869 GiB/NPU | +| step 时间中位数 | 40.184 s | + +端到端吞吐覆盖完整 100 steps 并包含 checkpoint 保存;稳态吞吐排除 step 1 和包含 +`timing_s/save_checkpoint` 的 step。 + +### 5.1 长跑曲线 + +下图覆盖完整 100 steps,依次展示 reward 和 MTP loss。两项指标均保留原始值和 10-step +moving average;reward 使用 `critic/score/mean`,MTP loss 使用 +`actor/mtp_losses/mtp_1_loss`。 + +![MiMo-7B-RL MTP 100-step reward and loss](training_curves.png) + +### 5.2 性能与稳定性 + +- reward 的末 10 步均值比首 10 步提高 0.897144。 +- MTP loss 的末 10 步均值比首 10 步降低 0.242035。 +- 4 NPU 端到端和稳态吞吐均高于无 GPU 标杆时的 100 TPS 门槛。 +- 训练连续完成 100 steps,最终 checkpoint tracker 为 100。 + +## 6. 复现证据与验收结论 + +完整训练日志: +[MiMo-7B-RL MTP 100-step training log](https://gist.githubusercontent.com/RordChang/c7730e4b733b544105b1efaa14dab32b/raw/c1026b5f0d6f7d67f35b2d5cb892344bd407481c/training_100step_sanitized.log) + +| Issue #20 验收项 | 本次结果 | +| --- | --- | +| 完成 100 steps 或运行 12 小时 | 完成连续 100/100 steps | +| reward 上升 | 首 10 步均值 -0.808044,末 10 步均值 0.089099 | +| MTP loss 有效下降 | 首 10 步均值 0.597910,末 10 步均值 0.355875 | +| 无 GPU 标杆时 TPS > 100 | 4 NPU 端到端吞吐 533.286 token/s | +| 提供可复现 recipe | 提供模型、数据、环境、补丁、启动、checkpoint/resume 和日志配置 | + +本次结果覆盖 Issue #20 的长跑、reward、MTP loss 和性能验收项,并提供了 MiMo-7B-RL +MTP 在 Megatron + SGLang-Ascend 组合上的完整复现入口。 diff --git a/verl_ascend_practice/reports/mtp/training_curves.png b/verl_ascend_practice/reports/mtp/training_curves.png new file mode 100644 index 0000000..cf0807b Binary files /dev/null and b/verl_ascend_practice/reports/mtp/training_curves.png differ diff --git a/verl_ascend_practice/run_mimo_7b_mtp_megatron_npu.sh b/verl_ascend_practice/run_mimo_7b_mtp_megatron_npu.sh new file mode 100644 index 0000000..6dc14bc --- /dev/null +++ b/verl_ascend_practice/run_mimo_7b_mtp_megatron_npu.sh @@ -0,0 +1,297 @@ +#!/usr/bin/env bash +# MTP | MiMo-7B | Megatron training | SGLang-Ascend rollout | Ascend NPUs by default +# With OFFLOAD=True, apply verl_ascend_practice/patches/mtp_checkpoint_engine_reoffload.patch to verl +# until the equivalent generic checkpoint-engine fix is available upstream. + +set -xeuo pipefail + +########################### platform and defaults ########################### + +DEVICE=${DEVICE:-npu} +MODEL_PATH=${MODEL_PATH:-XiaomiMiMo/MiMo-7B-RL} +NNODES=${NNODES:-1} + +if [[ "$DEVICE" == "npu" ]]; then + devices_per_node=${NPUS_PER_NODE:-4} + export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1 + export TOKENIZERS_PARALLELISM=${TOKENIZERS_PARALLELISM:-false} + export HYDRA_FULL_ERROR=${HYDRA_FULL_ERROR:-1} + export RAY_DEDUP_LOGS=${RAY_DEDUP_LOGS:-0} + export TASK_QUEUE_ENABLE=${TASK_QUEUE_ENABLE:-1} + export CPU_AFFINITY_CONF=${CPU_AFFINITY_CONF:-1} + export HCCL_OP_EXPANSION_MODE=${HCCL_OP_EXPANSION_MODE:-AIV} + export HCCL_ASYNC_ERROR_HANDLING=${HCCL_ASYNC_ERROR_HANDLING:-0} + export HCCL_EXEC_TIMEOUT=${HCCL_EXEC_TIMEOUT:-3600} + export HCCL_CONNECT_TIMEOUT=${HCCL_CONNECT_TIMEOUT:-3600} + export HCCL_HOST_SOCKET_PORT_RANGE=${HCCL_HOST_SOCKET_PORT_RANGE:-60000-60050} + export HCCL_NPU_SOCKET_PORT_RANGE=${HCCL_NPU_SOCKET_PORT_RANGE:-61000-61050} + + train_batch_size=${TRAIN_BATCH_SIZE:-16} + val_batch_size=${VAL_BATCH_SIZE:-16} + ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE:-8} + ppo_micro_batch_size=${PPO_MICRO_BATCH_SIZE:-1} + max_prompt_length=${MAX_PROMPT_LENGTH:-512} + max_response_length=${MAX_RESPONSE_LENGTH:-1024} + ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU:-4096} + dataloader_num_workers=${DATALOADER_NUM_WORKERS:-0} + + actor_tp=${ACTOR_TP:-2} + actor_pp=${ACTOR_PP:-1} + actor_cp=${ACTOR_CP:-1} + rollout_tp=${ROLLOUT_TP:-1} + rollout_n=${ROLLOUT_N:-2} + rollout_gpu_mem_util=${ROLLOUT_GPU_MEM_UTIL:-0.45} + rollout_max_num_seqs=${ROLLOUT_MAX_NUM_SEQS:-8} + rollout_max_num_batched_tokens=${ROLLOUT_MAX_NUM_BATCHED_TOKENS:-8192} + rollout_cuda_graph_max_bs=${ROLLOUT_CUDA_GRAPH_MAX_BS:-8} + checkpoint_bucket_mb=${CHECKPOINT_BUCKET_MB:-1280} + + total_epochs=${TOTAL_EPOCHS:-10} + total_training_steps=${TOTAL_TRAINING_STEPS:-100} + save_freq=${SAVE_FREQ:-100} + test_freq=${TEST_FREQ:--1} + lr_warmup_steps=${LR_WARMUP_STEPS:-10} + + data_root=${DATA_ROOT:-$HOME/data/math} + train_file=${TRAIN_FILE:-$data_root/train.parquet} + val_file=${VAL_FILE:-$data_root/test.parquet} + project_name=${PROJECT_NAME:-verl_mtp} + experiment_name=${EXPERIMENT_NAME:-mimo_7b_mtp_sglang_megatron_ascend} + logger='["console"]' +else + devices_per_node=${NGPUS_PER_NODE:-8} + export CUDA_DEVICE_MAX_CONNECTIONS=1 + + train_batch_size=${TRAIN_BATCH_SIZE:-128} + ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE:-32} + max_prompt_length=${MAX_PROMPT_LENGTH:-2048} + max_response_length=${MAX_RESPONSE_LENGTH:-8192} + ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU:-20480} + + actor_tp=${ACTOR_TP:-2} + actor_pp=${ACTOR_PP:-2} + actor_cp=${ACTOR_CP:-2} + rollout_tp=${ROLLOUT_TP:-4} + rollout_n=${ROLLOUT_N:-16} + rollout_gpu_mem_util=${ROLLOUT_GPU_MEM_UTIL:-0.8} + + total_epochs=${TOTAL_EPOCHS:-10} + total_training_steps=${TOTAL_TRAINING_STEPS:-400} + save_freq=${SAVE_FREQ:--1} + test_freq=${TEST_FREQ:-10} + lr_warmup_steps=${LR_WARMUP_STEPS:-10} + + train_file=${TRAIN_FILE:-$HOME/data/dapo-math-17k/train.parquet} + val_file=${VAL_FILE:-$HOME/data/aime-2024/test.parquet} + project_name=${PROJECT_NAME:-verl_mtp} + experiment_name=${EXPERIMENT_NAME:-mimo_7b_mtp_sglang_megatron} + logger='["console","wandb"]' +fi + +actor_lr=${ACTOR_LR:-1e-6} +entropy_coeff=${ENTROPY_COEFF:-0} +clip_ratio_low=${CLIP_RATIO_LOW:-0.2} +clip_ratio_high=${CLIP_RATIO_HIGH:-0.28} +mtp_loss_scaling_factor=${MTP_LOSS_SCALING_FACTOR:-0.1} +offload=${OFFLOAD:-True} + + +########################### parameter arrays ########################### + +ALGORITHM=( + algorithm.adv_estimator=grpo + algorithm.use_kl_in_reward=False +) + +DATA=( + "data.train_files=['$train_file']" + "data.val_files=['$val_file']" + data.train_batch_size=${train_batch_size} + data.max_prompt_length=${max_prompt_length} + data.max_response_length=${max_response_length} + data.trust_remote_code=True +) + +MODEL=( + actor_rollout_ref.model.path="$MODEL_PATH" + actor_rollout_ref.model.use_remove_padding=True + actor_rollout_ref.model.trust_remote_code=True + actor_rollout_ref.model.mtp.enable=True + actor_rollout_ref.model.mtp.enable_train=True + actor_rollout_ref.model.mtp.enable_rollout=False + actor_rollout_ref.model.mtp.mtp_loss_scaling_factor=${mtp_loss_scaling_factor} + actor_rollout_ref.model.mtp.detach_encoder=True +) + +ACTOR=( + actor_rollout_ref.actor.optim.lr=${actor_lr} + actor_rollout_ref.actor.optim.lr_warmup_steps=${lr_warmup_steps} + actor_rollout_ref.actor.optim.weight_decay=0.1 + actor_rollout_ref.actor.optim.clip_grad=1.0 + actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size} + actor_rollout_ref.actor.use_dynamic_bsz=True + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.actor.use_kl_loss=False + actor_rollout_ref.actor.entropy_coeff=${entropy_coeff} + actor_rollout_ref.actor.loss_agg_mode=token-mean + actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} + actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} + actor_rollout_ref.actor.clip_ratio_c=10.0 + actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${actor_tp} + actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${actor_pp} + actor_rollout_ref.actor.megatron.context_parallel_size=${actor_cp} + actor_rollout_ref.actor.megatron.param_offload=${offload} + actor_rollout_ref.actor.megatron.grad_offload=${offload} + actor_rollout_ref.actor.megatron.optimizer_offload=${offload} + actor_rollout_ref.actor.megatron.use_mbridge=True +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=sglang + actor_rollout_ref.rollout.tensor_model_parallel_size=${rollout_tp} + actor_rollout_ref.rollout.gpu_memory_utilization=${rollout_gpu_mem_util} + actor_rollout_ref.rollout.n=${rollout_n} + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} +) + +REF=( + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} + actor_rollout_ref.ref.megatron.tensor_model_parallel_size=${actor_tp} + actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=${actor_pp} + actor_rollout_ref.ref.megatron.context_parallel_size=${actor_cp} + actor_rollout_ref.ref.megatron.param_offload=${offload} +) + +REWARD=( + reward.reward_manager.name=dapo + +reward.reward_kwargs.overlong_buffer_cfg.enable=True + +reward.reward_kwargs.overlong_buffer_cfg.penalty_factor=1.0 + +reward.reward_kwargs.max_resp_len=${max_response_length} +) + +TRAINER=( + trainer.balance_batch=True + trainer.critic_warmup=0 + "trainer.logger=$logger" + trainer.project_name=${project_name} + trainer.experiment_name=${experiment_name} + trainer.n_gpus_per_node=${devices_per_node} + trainer.nnodes=${NNODES} + trainer.val_before_train=False + trainer.save_freq=${save_freq} + trainer.test_freq=${test_freq} + trainer.total_epochs=${total_epochs} + trainer.total_training_steps=${total_training_steps} +) + +if [[ "$DEVICE" == "npu" ]]; then + output_dir=${OUTPUT_DIR:-$PWD/checkpoints/$experiment_name} + log_dir=${LOG_DIR:-$PWD/logs} + resume_mode=${RESUME_MODE:-auto} + max_actor_ckpt_to_keep=${MAX_ACTOR_CKPT_TO_KEEP:-1} + mkdir -p "$output_dir" "$log_dir" + + ALGORITHM+=( + algorithm.rollout_correction.bypass_mode=True + algorithm.rollout_correction.loss_type=ppo_clip + ) + DATA+=( + data.val_batch_size=${val_batch_size} + data.dataloader_num_workers=${dataloader_num_workers} + data.filter_overlong_prompts=True + data.truncation=error + ) + ACTOR+=( + actor_rollout_ref.actor.optim.use_checkpoint_opt_param_scheduler=True + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=${ppo_micro_batch_size} + actor_rollout_ref.actor.megatron.vanilla_mbridge=False + +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_method=uniform + +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_granularity=full + +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_num_layers=1 + +actor_rollout_ref.actor.megatron.override_transformer_config.apply_rope_fusion=True + +actor_rollout_ref.actor.megatron.override_transformer_config.position_embedding_type=rope + +actor_rollout_ref.actor.megatron.override_transformer_config.use_fused_rotary_pos_emb=True + +actor_rollout_ref.actor.megatron.override_transformer_config.normalization=RMSNorm + +actor_rollout_ref.actor.megatron.override_transformer_config.use_fused_rmsnorm=True + actor_rollout_ref.actor.megatron.override_transformer_config.attention_backend=flash + +actor_rollout_ref.actor.megatron.override_transformer_config.use_flash_attn=True + ) + ROLLOUT+=( + actor_rollout_ref.rollout.max_model_len=$((max_prompt_length + max_response_length)) + actor_rollout_ref.rollout.max_num_seqs=${rollout_max_num_seqs} + actor_rollout_ref.rollout.max_num_batched_tokens=${rollout_max_num_batched_tokens} + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=${ppo_micro_batch_size} + actor_rollout_ref.rollout.calculate_log_probs=True + actor_rollout_ref.rollout.enable_chunked_prefill=True + actor_rollout_ref.rollout.enable_prefix_caching=True + actor_rollout_ref.rollout.enforce_eager=False + actor_rollout_ref.rollout.free_cache_engine=True + actor_rollout_ref.rollout.checkpoint_engine.backend=nccl + actor_rollout_ref.rollout.checkpoint_engine.update_weights_bucket_megabytes=${checkpoint_bucket_mb} + +actor_rollout_ref.rollout.engine_kwargs.sglang.attention_backend=ascend + +actor_rollout_ref.rollout.engine_kwargs.sglang.cuda_graph_max_bs=${rollout_cuda_graph_max_bs} + actor_rollout_ref.rollout.val_kwargs.n=1 + actor_rollout_ref.rollout.val_kwargs.temperature=1.0 + actor_rollout_ref.rollout.val_kwargs.top_p=0.7 + actor_rollout_ref.rollout.val_kwargs.top_k=-1 + actor_rollout_ref.rollout.val_kwargs.do_sample=True + ) + REF+=( + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=${ppo_micro_batch_size} + actor_rollout_ref.ref.megatron.vanilla_mbridge=False + ) + REWARD+=( + +reward.reward_kwargs.overlong_buffer_cfg.len=256 + +reward.reward_kwargs.overlong_buffer_cfg.log=False + ) + TRAINER+=( + trainer.device=npu + trainer.resume_mode=${resume_mode} + trainer.max_actor_ckpt_to_keep=${max_actor_ckpt_to_keep} + "trainer.default_local_dir=$output_dir" + ) +else + DATA+=(data.truncation=left) + REWARD+=(+reward.reward_kwargs.overlong_buffer_cfg.len=4096) +fi + +########################### launch ########################### + +EXTRA=( + model_engine=megatron +) + +# uv (set VERL_USE_UV=0 for system python): on GPU, the driver and every Ray worker +# (runtime_env.py_executable) run through `uv run` on the sglang × megatron extras of the committed uv.lock; +# NPU falls back to ambient python. Run from the verl repo root. +LAUNCH=(python3) +RAY=(ray_kwargs.ray_init.runtime_env.py_executable=null) +if [[ "${VERL_USE_UV:-1}" != 0 && "$DEVICE" != "npu" ]]; then + LAUNCH=(uv run --frozen --all-packages --extra sglang --extra megatron python3) + RAY=(ray_kwargs.ray_init.runtime_env.py_executable="uv -v run --frozen --all-packages --extra sglang --extra megatron") +fi + +COMMAND=( + "${LAUNCH[@]}" + -m verl.trainer.main_ppo + "${DATA[@]}" + "${ALGORITHM[@]}" + "${MODEL[@]}" + "${ACTOR[@]}" + "${ROLLOUT[@]}" + "${REF[@]}" + "${REWARD[@]}" + "${TRAINER[@]}" + "${EXTRA[@]}" + "${RAY[@]}" + "$@" +) + +if [[ "$DEVICE" == "npu" ]]; then + log_file="$log_dir/${experiment_name}_$(date +%Y%m%d_%H%M%S).log" + PYTHONUNBUFFERED=1 "${COMMAND[@]}" 2>&1 | tee "$log_file" +else + "${COMMAND[@]}" +fi