From 4ae9acc3a5ffa96deefdd5040d656830c20b2e42 Mon Sep 17 00:00:00 2001 From: "xiongshaopan.xsp" Date: Thu, 18 Jun 2026 14:20:50 +0800 Subject: [PATCH] (feat): publish v0.3.0. --- .gitignore | 14 +- README.md | 5 +- docker/Dockerfile.torch2100 | 3 + .../Developer Guide/support_new_models.md | 8 +- .../docs/Getting Started/FAQ/qa_issues.md | 4 +- .../Quick Start/multi_nodes_quick_start.md | 2 +- .../Quick Start/single_node_quick_start.md | 2 +- docs_roll/docs/Overview.mdx | 2 +- .../Advanced Features/mtp_training.md | 286 ++++ .../opentelemetry_tracing.md | 91 ++ .../remote_batch_transfer.md | 74 + .../Advanced Features/router_replay.md | 187 +++ .../docs/User Guides/Agentic/agent_runner.md | 557 +++++++ .../docs/User Guides/Algorithms/Reward_FL.md | 94 -- .../User Guides/Configuration/config_guide.md | 37 +- .../Configuration/config_system.md | 2 +- .../User Guides/Configuration/deepspeed.md | 109 -- .../docs/User Guides/Configuration/lora.md | 28 +- .../Hardware Support/ascend_usage.md | 46 +- .../Pipeline/agent_pipeline_start.md | 6 +- .../Pipeline/agentic_pipeline_start.md | 2 +- .../Pipeline/distill_pipeline_start.md | 2 +- .../Pipeline/dpo_pipeline_start.md | 2 +- .../on_policy_distill_pipeline_start.md | 201 ++- .../Pipeline/sft_pipeline_start.md | 2 +- .../Developer Guide/support_new_models.md | 8 +- .../current/Getting Started/FAQ/qa_issues.md | 4 +- .../Quick Start/multi_nodes_quick_start.md | 2 +- .../Quick Start/single_node_quick_start.md | 2 +- .../current/Overview.mdx | 2 +- .../Advanced Features/mtp_training.md | 286 ++++ .../opentelemetry_tracing.md | 91 ++ .../remote_batch_transfer.md | 75 + .../Advanced Features/router_replay.md | 187 +++ .../User Guides/Agentic/agent_runner.md | 556 +++++++ .../User Guides/Agentic/prefix_aggregation.md | 241 +++ .../User Guides/Algorithms/Reward_FL.md | 92 -- .../User Guides/Configuration/config_guide.md | 36 +- .../Configuration/config_system.md | 2 +- .../User Guides/Configuration/deepspeed.md | 104 -- .../current/User Guides/Configuration/lora.md | 28 +- .../Hardware Support/ascend_usage.md | 94 +- .../Pipeline/agentic_pipeline_start.md | 2 +- .../Pipeline/distill_pipeline_start.md | 2 +- .../Pipeline/dpo_pipeline_start.md | 2 +- .../on_policy_distill_pipeline_start.md | 195 ++- .../Pipeline/sft_pipeline_start.md | 2 +- examples/agentic_deepeyes/deepeyes.yaml | 6 +- .../agentic_demo/agent_rollout_rock_swe.yaml | 4 - ...gent_val_frozen_lake_multi_nodes_demo.yaml | 6 - ...gent_val_frozen_lake_single_node_demo.yaml | 6 - examples/agentic_demo/agent_val_rock_swe.yaml | 4 - .../agent_val_rock_swe_qwen35_2b.yaml | 4 - .../agent_val_sokoban_sandbox.yaml | 6 - .../atropos_gsm8k_grpo_qwen25_0.5b.yaml | 22 +- ...endless_terminals_reinforce_qwen35_2b.yaml | 6 - ...eepspeed.yaml => qwen3_8b_rlvr_fsdp2.yaml} | 16 +- examples/config/deepspeed_zero.yaml | 15 - examples/config/deepspeed_zero2.yaml | 19 - .../config/deepspeed_zero2_cpuoffload.yaml | 25 - examples/config/deepspeed_zero3.yaml | 22 - .../config/deepspeed_zero3_cpuoffload.yaml | 25 - examples/config/traj_envs.yaml | 7 + examples/docs_examples/example_ppo.yaml | 19 +- examples/docs_examples/example_reward_fl.yaml | 67 - .../dpo_examples/qwen2.5-3B_dpo_megatron.yaml | 6 - .../qwen3-30BA3B-dpo_megatron_80GB.yaml | 6 - .../agent_val_frozen_lake-pg_var.yaml | 6 - ...ent_val_frozen_lake-pg_var_is_correct.yaml | 6 - .../agent_val_frozen_lake.yaml | 6 - .../agent_val_frozen_lake_amd.yaml | 6 - .../agent_val_frozen_lake_async.yaml | 6 - .../agent_val_frozen_lake_async_amd.yaml | 6 - .../agent_val_frozen_lake_gigpo.yaml | 6 - .../agentic_rollout_sokoban.yaml | 4 - .../agentic_sokoban_rollout_mock_dump.yaml | 4 - .../agentic_val_sokoban.yaml | 6 - .../agentic_val_sokoban_agent_runner.yaml | 159 ++ .../agentic_val_sokoban_dynamic_batching.yaml | 6 - .../agentic_val_sokoban_gigpo.yaml | 6 - .../agentic_val_sokoban_lora.yaml | 6 - .../agentic_val_sokoban_native.yaml | 6 - .../agentic_val_sokoban_sandbox.yaml | 6 - .../distill_zero3.yaml | 85 - .../run_distill_pipeline.sh | 6 - .../agentic_val_webshop.yaml | 4 - .../agentic_val_webshop_amd.yaml | 136 +- .../agentic_val_webshop_gigpo.yaml | 4 - ...lvr_zero3_sp2.yaml => rlvr_fsdp2_sp2.yaml} | 24 +- ...r_lora_zero3.yaml => rlvr_lora_fsdp2.yaml} | 21 +- .../agentic_val_sokoban.yaml | 6 - ...ll_vl_zero3.yaml => distill_vl_fsdp2.yaml} | 27 +- examples/qwen2.5-vl-7B-rlvr/rlvr_async.yaml | 8 +- .../qwen2.5-vl-7B-rlvr/rlvr_megatron.yaml | 9 +- .../rlvr_fsdp2_offload.yaml | 159 ++ .../run_rlvr_pipeline.sh | 5 + .../rlvr_config_140GB.yaml | 235 +++ .../run_rlvr_pipeline.sh | 2 +- .../agentic_val_sokoban_30a3.yaml | 4 - .../qwen3-30BA3B-rlvr_fsdp2/rlvr_config.yaml | 185 +++ .../rlvr_config_ep.yaml | 191 +++ .../rlvr_config_ppu_sglang.yaml | 230 +++ .../qwen3-30BA3B-sft_fsdp2/sft_config.yaml | 66 + .../qwen3-30BA3B-sft_fsdp2/sft_config_ep.yaml | 69 + .../sft_config_megatron.yaml | 73 + ...multi_teacher_onpolicy_distill_config.yaml | 189 +++ .../submit_pipeline.sh | 45 - .../rlvr_config.yaml} | 78 +- .../rlvr_config_wo_cpuoffload.yaml} | 118 +- .../qwen3-8B-rlvr_fsdp2/run_rlvr_pipeline.sh | 5 + examples/qwen3-omni/audio_test_80G.yaml | 172 ++ examples/qwen3-omni/rlvr_megatron.yaml | 9 +- .../rlvr_megatron_80GB.yaml | 8 +- .../rlvr_megatron_80GB_sglang.yaml | 164 ++ .../rlvr_megatron_ppu.yaml | 164 ++ .../qwen3-vl-30BA3B-video-rlvr/video_r1.yaml | 161 ++ .../video_r1_sglang.yaml | 158 ++ .../rlvr_megatron_80GB.yaml | 8 +- .../rlvr_megatron_80G.yaml | 8 +- .../rlvr_megatron_80GB.yaml | 233 +++ .../rlvr_megatron_80GB.yaml | 10 +- .../submit_pipeline.sh | 45 - .../rlvr_megatron_80GB.yaml | 10 +- .../rlvr_megatron_ppu.yaml | 179 +++ .../submit_pipeline.sh | 45 - .../qwen3.5-35BA3-sft_fsdp2/sft_config.yaml | 69 + .../sft_config_ep.yaml | 73 + .../rlvr_megatron_80GB.yaml | 240 +++ .../gem_game_guess_the_number.yaml | 6 - examples/qwen3_agentic_gem/gem_math_dapo.yaml | 2 - .../gem_math_dapo_python_code.yaml | 2 - .../qwen3_agentic_gem/gem_math_hotpotqa.yaml | 2 - .../gem_math_hotpotqa_search.yaml | 2 - .../gem_rg_letter_counting.yaml | 6 - examples/start_reward_fl_pipeline.py | 36 - examples/video-rlvr/video_r1.yaml | 160 ++ .../wan2.2-14B-reward_fl_ds/merge_lora.py | 94 -- .../wan2.2-14B-reward_fl_ds/merge_model.py | 23 - .../reward_fl_config.yaml | 70 - .../run_reward_fl_ds_pipeline.sh | 6 - .../wan2.2-14B-reward_fl_ds/wan22_paths.json | 6 - mcore_adapter/pyproject.toml | 6 +- .../src/mcore_adapter/checkpointing.py | 60 +- mcore_adapter/src/mcore_adapter/constants.py | 1 + mcore_adapter/src/mcore_adapter/initialize.py | 1 + .../models/converter/convert_utils.py | 10 +- .../models/converter/dist_converter.py | 114 +- .../models/converter/model_converter.py | 158 +- .../models/converter/post_converter.py | 325 +++- .../models/converter/template.py | 82 +- .../src/mcore_adapter/models/model_config.py | 18 +- .../src/mcore_adapter/models/model_factory.py | 85 +- .../models/qwen2_5_vl/modeling_qwen2_5_vl.py | 24 +- .../models/qwen2_vl/modeling_qwen2_vl.py | 24 +- .../mcore_adapter/models/qwen3_5/__init__.py | 161 +- .../models/qwen3_5/modeling_qwen3_5.py | 162 +- .../models/qwen3_5_moe/__init__.py | 94 +- .../models/qwen3_next/__init__.py | 42 +- .../models/qwen3_omni/modeling_qwen3_omni.py | 248 ++- .../models/qwen3_vl/modeling_qwen3_vl.py | 220 +-- .../models/qwen3_vl/rope_utils.py | 100 +- .../models/sequence_packing_mixin.py | 572 +++++++ .../parallel_functions/fused_cross_entropy.py | 618 +++++++ .../parallel_functions/fused_logprobs.py | 686 ++++++++ .../parallel_functions/vocab_parallel.py | 117 +- .../src/mcore_adapter/trainer/dpo_trainer.py | 10 +- .../src/mcore_adapter/trainer/trainer.py | 43 +- .../src/mcore_adapter/training_args.py | 64 +- requirements_common.txt | 8 +- requirements_em_local_debug.txt | 1 - requirements_torch2100_vllm.txt | 1 - requirements_torch260_vllm.txt | 2 +- requirements_torch280_sglang.txt | 5 +- requirements_torch280_vllm.txt | 5 +- requirements_vision.txt | 5 +- roll/configs/base_config.py | 196 ++- roll/configs/data_args.py | 39 + roll/configs/model_args.py | 2 +- roll/configs/worker_config.py | 79 +- roll/datasets/collator.py | 551 ++++++- roll/datasets/dataset.py | 8 +- roll/datasets/vlm_dataset_utils.py | 416 +++++ roll/distributed/executor/worker.py | 31 +- .../scheduler/async_generate_scheduler.py | 843 ---------- roll/distributed/scheduler/decorator.py | 103 +- .../scheduler/generate_scheduler.py | 212 ++- roll/distributed/scheduler/initialize.py | 59 +- roll/distributed/scheduler/log_monitor.py | 2 - roll/distributed/scheduler/protocol.py | 576 +++++-- roll/distributed/scheduler/remote_protocol.py | 904 +++++++++++ .../distributed/scheduler/resource_manager.py | 5 +- .../scheduler/rollout_scheduler.py | 26 +- roll/distributed/scheduler/router.py | 25 +- .../distributed/scheduler/transfer_backend.py | 224 +++ .../scheduler/user_defined_rollout_loop.py | 38 +- .../strategy/deepspeed_strategy.py | 590 ------- .../strategy/diffusion_strategy.py | 83 - roll/distributed/strategy/factory.py | 8 +- roll/distributed/strategy/fsdp2_strategy.py | 561 ++++++- roll/distributed/strategy/hf_strategy.py | 1 - .../distributed/strategy/megatron_strategy.py | 339 +++- roll/distributed/strategy/mock_strategy.py | 1 - roll/distributed/strategy/sglang_strategy.py | 140 +- roll/distributed/strategy/strategy.py | 67 +- roll/distributed/strategy/vllm_strategy.py | 101 +- roll/models/model_providers.py | 82 +- roll/models/trl_patches.py | 12 +- .../pipeline/agentic/agent_runner/__init__.py | 3 + roll/pipeline/agentic/agent_runner/base.py | 125 ++ .../agentic/agent_runner/gem_runner.py | 196 +++ .../agentic/agent_runner/rock/__init__.py | 5 + .../agentic/agent_runner/rock/pull_runner.py | 261 +++ .../agentic/agent_runner/rock/push_runner.py | 63 + .../agent_runner/rock/rock_agent_runner.py | 372 +++++ .../agent_runner/rock/sandbox_tool_runner.py | 6 + roll/pipeline/agentic/agentic_config.py | 19 + roll/pipeline/agentic/agentic_pipeline.py | 303 ++-- roll/pipeline/agentic/env/__init__.py | 1 + roll/pipeline/agentic/env/cli_env/__init__.py | 3 + roll/pipeline/agentic/env/cli_env/env.py | 340 ++++ .../env/cli_env/simple_task_checker.py | 206 +++ roll/pipeline/agentic/env/cli_env/test_cli.py | 66 + .../agentic/env/cli_env/test_no_iflow.py | 96 ++ roll/pipeline/agentic/env/cli_env/utils.py | 87 + .../agentic/env/cli_env/xrl_wrapper.py | 195 +++ roll/pipeline/agentic/env/deepeyes/env.py | 1 - .../env/rock/harbor_service_config.yaml | 59 + .../agentic/env/rock/sanbox_manager.py | 1048 ++++++++++++ .../env/sandbox/sokoban_sandbox_env.py | 18 +- .../agentic/env/sokoban/tool_call_env.py | 494 ++++++ .../env_manager/agent_native_env_manager.py | 5 +- .../agentic/env_manager/base_env_manager.py | 4 +- .../agentic/env_manager/message_tracker.py | 525 ++++++ .../agentic/env_manager/proxy_env_manager.py | 563 +++++++ .../agentic/env_manager/traj_env_manager.py | 30 +- .../env_manager/vl_traj_env_manager.py | 2 +- roll/pipeline/agentic/environment_worker.py | 54 +- .../agentic/llm_proxy/openai_proxy.py | 47 +- .../pipeline/agentic/llm_proxy/proxy_utils.py | 3 +- roll/pipeline/agentic/proxy/proxy_server.py | 277 ++++ .../pipeline/agentic/proxy/router/__init__.py | 9 + .../agentic/proxy/router/alb_proxy_router.py | 1247 +++++++++++++++ roll/pipeline/agentic/proxy/router/base.py | 106 ++ .../router/cleanup_alb_resources_batch.py | 124 ++ .../proxy/router/ingress_proxy_router.py | 372 +++++ roll/pipeline/agentic/utils.py | 56 +- roll/pipeline/base_pipeline.py | 37 +- roll/pipeline/base_worker.py | 45 +- roll/pipeline/diffusion/__init__.py | 0 roll/pipeline/diffusion/modules/__init__.py | 0 roll/pipeline/diffusion/modules/wan_module.py | 554 ------- roll/pipeline/diffusion/reward_fl/__init__.py | 0 .../diffusion/reward_fl/actor_worker.py | 60 - roll/pipeline/diffusion/reward_fl/euler.py | 89 -- .../diffusion/reward_fl/face_tools.py | 508 ------ .../diffusion/reward_fl/reward_fl_config.py | 47 - .../diffusion/reward_fl/reward_fl_pipeline.py | 120 -- .../diffusion/reward_fl/wan_video_vae.py | 1416 ----------------- roll/pipeline/distill/distill_pipeline.py | 1 + roll/pipeline/distill/distill_vlm_pipeline.py | 4 +- roll/pipeline/distill/distill_worker.py | 6 +- roll/pipeline/rlvr/rewards/__init__.py | 3 +- .../rlvr/rewards/llm_judge_reward_worker.py | 238 +-- ...ultiple_choice_boxed_rule_reward_worker.py | 2 +- .../rewards/remote_reward_system_worker.py | 529 ++++++ .../rlvr/rewards/video_r1_reward_worker.py | 184 +++ roll/pipeline/rlvr/rlvr_config.py | 38 + roll/pipeline/rlvr/rlvr_pipeline.py | 160 +- roll/pipeline/rlvr/rlvr_rollout_pipeline.py | 2 +- roll/pipeline/rlvr/rlvr_vlm_pipeline.py | 370 ++--- roll/pipeline/rlvr/utils.py | 2 +- roll/pipeline/sft/sft_pipeline.py | 6 +- roll/platforms/rocm.py | 6 +- roll/third_party/deepspeed/__init__.py | 0 roll/third_party/deepspeed/model_update.py | 205 --- roll/third_party/deepspeed/offload_states.py | 89 -- .../deepspeed/offload_states_patch.py | 491 ------ roll/third_party/fsdp2/model_update.py | 17 +- roll/third_party/fsdp2/qwen3_moe_patch.py | 53 + roll/third_party/megatron/compile_warmup.py | 146 ++ roll/third_party/megatron/fused_entropy.py | 941 +++++++++++ roll/third_party/megatron/model_update.py | 26 +- roll/third_party/megatron/mtp_patcher.py | 293 ++++ .../megatron/offload_states_patch.py | 35 +- .../megatron/router_replay_utils.py | 358 +++++ roll/third_party/megatron/tensor_parallel.py | 131 +- roll/third_party/megatron/util.py | 181 +++ roll/third_party/sglang/__init__.py | 5 +- .../sglang/v046post4_patch/engine.py | 16 +- roll/third_party/vllm/__init__.py | 7 +- roll/third_party/vllm/fp8.py | 69 +- roll/third_party/vllm/gdn_patcher.py | 332 ++++ roll/third_party/vllm/worker.py | 24 +- roll/utils/checkpoint_manager.py | 13 + roll/utils/deepspeed_utils.py | 43 - roll/utils/fsdp_utils.py | 422 ++++- roll/utils/functionals.py | 404 ++++- roll/utils/multi_thread_utils.py | 72 - roll/utils/offload_states.py | 15 + roll/utils/otel_receiver.py | 202 +++ roll/utils/qwen_vl_utils.py | 552 +++++++ roll/utils/taskgroups.py | 2 +- roll/utils/telemetry.py | 304 ++++ roll/utils/train_infer_corrections.py | 124 +- roll/utils/upload_utils.py | 13 + setup.py | 2 +- 306 files changed, 27266 insertions(+), 9005 deletions(-) create mode 100644 docs_roll/docs/User Guides/Advanced Features/mtp_training.md create mode 100644 docs_roll/docs/User Guides/Advanced Features/opentelemetry_tracing.md create mode 100644 docs_roll/docs/User Guides/Advanced Features/remote_batch_transfer.md create mode 100644 docs_roll/docs/User Guides/Advanced Features/router_replay.md create mode 100644 docs_roll/docs/User Guides/Agentic/agent_runner.md delete mode 100644 docs_roll/docs/User Guides/Algorithms/Reward_FL.md delete mode 100644 docs_roll/docs/User Guides/Configuration/deepspeed.md create mode 100644 docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Advanced Features/mtp_training.md create mode 100644 docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Advanced Features/opentelemetry_tracing.md create mode 100644 docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Advanced Features/remote_batch_transfer.md create mode 100644 docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Advanced Features/router_replay.md create mode 100644 docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Agentic/agent_runner.md create mode 100644 docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Agentic/prefix_aggregation.md delete mode 100644 docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Algorithms/Reward_FL.md delete mode 100644 docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Configuration/deepspeed.md rename examples/ascend_examples/{qwen3_8b_rlvr_deepspeed.yaml => qwen3_8b_rlvr_fsdp2.yaml} (91%) delete mode 100644 examples/config/deepspeed_zero.yaml delete mode 100644 examples/config/deepspeed_zero2.yaml delete mode 100644 examples/config/deepspeed_zero2_cpuoffload.yaml delete mode 100644 examples/config/deepspeed_zero3.yaml delete mode 100644 examples/config/deepspeed_zero3_cpuoffload.yaml delete mode 100644 examples/docs_examples/example_reward_fl.yaml create mode 100644 examples/qwen2.5-0.5B-agentic/agentic_val_sokoban_agent_runner.yaml delete mode 100644 examples/qwen2.5-1.5B-distill_ds/distill_zero3.yaml delete mode 100644 examples/qwen2.5-1.5B-distill_ds/run_distill_pipeline.sh rename examples/qwen2.5-7B-rlvr_megatron/{rlvr_zero3_sp2.yaml => rlvr_fsdp2_sp2.yaml} (93%) rename examples/qwen2.5-7B-rlvr_megatron/{rlvr_lora_zero3.yaml => rlvr_lora_fsdp2.yaml} (92%) rename examples/qwen2.5-vl-7B-distill/{distill_vl_zero3.yaml => distill_vl_fsdp2.yaml} (79%) create mode 100644 examples/qwen2_5-vl-3B-rlvr_fsdp2/rlvr_fsdp2_offload.yaml create mode 100755 examples/qwen2_5-vl-3B-rlvr_fsdp2/run_rlvr_pipeline.sh create mode 100644 examples/qwen3-235BA22B-rlvr_megatron/rlvr_config_140GB.yaml create mode 100644 examples/qwen3-30BA3B-rlvr_fsdp2/rlvr_config.yaml create mode 100644 examples/qwen3-30BA3B-rlvr_fsdp2/rlvr_config_ep.yaml create mode 100644 examples/qwen3-30BA3B-rlvr_megatron/rlvr_config_ppu_sglang.yaml create mode 100644 examples/qwen3-30BA3B-sft_fsdp2/sft_config.yaml create mode 100644 examples/qwen3-30BA3B-sft_fsdp2/sft_config_ep.yaml create mode 100644 examples/qwen3-30BA3B-sft_fsdp2/sft_config_megatron.yaml create mode 100644 examples/qwen3-8B-onpolicy-distill-megatron/multi_teacher_onpolicy_distill_config.yaml delete mode 100644 examples/qwen3-8B-onpolicy-distill-megatron/submit_pipeline.sh rename examples/{qwen2.5-7B-rlvr_megatron/rlvr_qwen2.5_7B_megatron_vllm_8gpus.yaml => qwen3-8B-rlvr_fsdp2/rlvr_config.yaml} (81%) rename examples/{qwen2.5-7B-rlvr_megatron/rlvr_lora_zero3_amd.yaml => qwen3-8B-rlvr_fsdp2/rlvr_config_wo_cpuoffload.yaml} (70%) create mode 100755 examples/qwen3-8B-rlvr_fsdp2/run_rlvr_pipeline.sh create mode 100644 examples/qwen3-omni/audio_test_80G.yaml create mode 100644 examples/qwen3-vl-30BA3B-rlvr_megatron/rlvr_megatron_80GB_sglang.yaml create mode 100644 examples/qwen3-vl-30BA3B-rlvr_megatron/rlvr_megatron_ppu.yaml create mode 100644 examples/qwen3-vl-30BA3B-video-rlvr/video_r1.yaml create mode 100644 examples/qwen3-vl-30BA3B-video-rlvr/video_r1_sglang.yaml create mode 100644 examples/qwen3.5-122A10-rlvr_megatron/rlvr_megatron_80GB.yaml delete mode 100644 examples/qwen3.5-27B-rlvr_megatron/submit_pipeline.sh create mode 100644 examples/qwen3.5-35BA3-rlvr_megatron/rlvr_megatron_ppu.yaml delete mode 100644 examples/qwen3.5-35BA3-rlvr_megatron/submit_pipeline.sh create mode 100644 examples/qwen3.5-35BA3-sft_fsdp2/sft_config.yaml create mode 100644 examples/qwen3.5-35BA3-sft_fsdp2/sft_config_ep.yaml create mode 100644 examples/qwen3.5-397A17-rlvr_megatron/rlvr_megatron_80GB.yaml delete mode 100644 examples/start_reward_fl_pipeline.py create mode 100644 examples/video-rlvr/video_r1.yaml delete mode 100644 examples/wan2.2-14B-reward_fl_ds/merge_lora.py delete mode 100644 examples/wan2.2-14B-reward_fl_ds/merge_model.py delete mode 100644 examples/wan2.2-14B-reward_fl_ds/reward_fl_config.yaml delete mode 100644 examples/wan2.2-14B-reward_fl_ds/run_reward_fl_ds_pipeline.sh delete mode 100644 examples/wan2.2-14B-reward_fl_ds/wan22_paths.json create mode 100644 mcore_adapter/src/mcore_adapter/models/sequence_packing_mixin.py create mode 100644 mcore_adapter/src/mcore_adapter/parallel_functions/fused_cross_entropy.py create mode 100644 mcore_adapter/src/mcore_adapter/parallel_functions/fused_logprobs.py create mode 100644 roll/datasets/vlm_dataset_utils.py delete mode 100644 roll/distributed/scheduler/async_generate_scheduler.py create mode 100644 roll/distributed/scheduler/remote_protocol.py create mode 100644 roll/distributed/scheduler/transfer_backend.py delete mode 100644 roll/distributed/strategy/deepspeed_strategy.py delete mode 100644 roll/distributed/strategy/diffusion_strategy.py create mode 100644 roll/pipeline/agentic/agent_runner/__init__.py create mode 100644 roll/pipeline/agentic/agent_runner/base.py create mode 100644 roll/pipeline/agentic/agent_runner/gem_runner.py create mode 100644 roll/pipeline/agentic/agent_runner/rock/__init__.py create mode 100644 roll/pipeline/agentic/agent_runner/rock/pull_runner.py create mode 100644 roll/pipeline/agentic/agent_runner/rock/push_runner.py create mode 100644 roll/pipeline/agentic/agent_runner/rock/rock_agent_runner.py create mode 100644 roll/pipeline/agentic/agent_runner/rock/sandbox_tool_runner.py create mode 100644 roll/pipeline/agentic/env/cli_env/__init__.py create mode 100644 roll/pipeline/agentic/env/cli_env/env.py create mode 100644 roll/pipeline/agentic/env/cli_env/simple_task_checker.py create mode 100644 roll/pipeline/agentic/env/cli_env/test_cli.py create mode 100644 roll/pipeline/agentic/env/cli_env/test_no_iflow.py create mode 100644 roll/pipeline/agentic/env/cli_env/utils.py create mode 100644 roll/pipeline/agentic/env/cli_env/xrl_wrapper.py create mode 100644 roll/pipeline/agentic/env/rock/harbor_service_config.yaml create mode 100644 roll/pipeline/agentic/env/rock/sanbox_manager.py create mode 100644 roll/pipeline/agentic/env/sokoban/tool_call_env.py create mode 100644 roll/pipeline/agentic/env_manager/message_tracker.py create mode 100644 roll/pipeline/agentic/env_manager/proxy_env_manager.py create mode 100644 roll/pipeline/agentic/proxy/proxy_server.py create mode 100644 roll/pipeline/agentic/proxy/router/__init__.py create mode 100644 roll/pipeline/agentic/proxy/router/alb_proxy_router.py create mode 100644 roll/pipeline/agentic/proxy/router/base.py create mode 100644 roll/pipeline/agentic/proxy/router/cleanup_alb_resources_batch.py create mode 100644 roll/pipeline/agentic/proxy/router/ingress_proxy_router.py delete mode 100644 roll/pipeline/diffusion/__init__.py delete mode 100644 roll/pipeline/diffusion/modules/__init__.py delete mode 100644 roll/pipeline/diffusion/modules/wan_module.py delete mode 100644 roll/pipeline/diffusion/reward_fl/__init__.py delete mode 100644 roll/pipeline/diffusion/reward_fl/actor_worker.py delete mode 100644 roll/pipeline/diffusion/reward_fl/euler.py delete mode 100644 roll/pipeline/diffusion/reward_fl/face_tools.py delete mode 100644 roll/pipeline/diffusion/reward_fl/reward_fl_config.py delete mode 100644 roll/pipeline/diffusion/reward_fl/reward_fl_pipeline.py delete mode 100644 roll/pipeline/diffusion/reward_fl/wan_video_vae.py create mode 100644 roll/pipeline/rlvr/rewards/remote_reward_system_worker.py create mode 100644 roll/pipeline/rlvr/rewards/video_r1_reward_worker.py delete mode 100644 roll/third_party/deepspeed/__init__.py delete mode 100644 roll/third_party/deepspeed/model_update.py delete mode 100644 roll/third_party/deepspeed/offload_states.py delete mode 100644 roll/third_party/deepspeed/offload_states_patch.py create mode 100644 roll/third_party/megatron/compile_warmup.py create mode 100644 roll/third_party/megatron/fused_entropy.py create mode 100644 roll/third_party/megatron/mtp_patcher.py create mode 100644 roll/third_party/megatron/router_replay_utils.py create mode 100644 roll/third_party/megatron/util.py create mode 100644 roll/third_party/vllm/gdn_patcher.py delete mode 100644 roll/utils/deepspeed_utils.py delete mode 100644 roll/utils/multi_thread_utils.py create mode 100644 roll/utils/otel_receiver.py create mode 100644 roll/utils/qwen_vl_utils.py create mode 100644 roll/utils/telemetry.py diff --git a/.gitignore b/.gitignore index 7e6830569..d4c56c4e3 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,5 @@ +nohup.out + *.png *.pyc */checkpoint_dir @@ -134,6 +136,7 @@ celerybeat.pid venv/ env.bak/ venv.bak/ +.project_conda_env/ # Spyder project settings .spyderproject @@ -169,4 +172,13 @@ log # examples output/ -.vscode/ \ No newline at end of file +.vscode/ + +.roo/ +.claude/ +.history/ +*.excp +logs/ +tmp/ +.ml_tracker/ +:/ \ No newline at end of file diff --git a/README.md b/README.md index 5441905df..d2c98872f 100644 --- a/README.md +++ b/README.md @@ -108,7 +108,7 @@ Leveraging a multi-role distributed architecture with Ray for flexible resource [RewardFL](https://alibaba.github.io/ROLL/docs/User%20Guides/Algorithms/Reward_FL) #### Backend -[DeepSpeed](https://alibaba.github.io/ROLL/docs/User%20Guides/Configuration/deepspeed) +[FSDP2](https://alibaba.github.io/ROLL/docs/User%20Guides/Configuration/fsdp2) [Megatron](https://alibaba.github.io/ROLL/docs/User%20Guides/Configuration/megatron) [vLLM](https://alibaba.github.io/ROLL/docs/User%20Guides/Configuration/vllm) [SGLang](https://alibaba.github.io/ROLL/docs/User%20Guides/Configuration/sglang) @@ -141,7 +141,7 @@ Leveraging a multi-role distributed architecture with Ray for flexible resource * Out-of-the-box support for reinforcement learning algorithms, such as **PPO, GRPO, Reinforce++, TOPR, RAFT++, GSPO**, etc. * **Rich Training and Inference Engine:** Ray-based multi-role distributed architecture; Strategy abstraction unifies various backends, enabling easy operation from single machines to thousands-of-GPU clusters. * Inference/Generation supports vLLM, SGLang. - * Training supports DeepSpeed (ZeRO), Megatron-LM 5D parallelism (mcore-adapter, dp/tp/pp/cp/ep), FSDP under implementation. + * Training supports FSDP2, Megatron-LM 5D parallelism (mcore-adapter, dp/tp/pp/cp/ep). * Extreme offload/reload capabilities. * Supports [LoRA](https://alibaba.github.io/ROLL/docs/User%20Guides/Configuration/lora) training. * Supports FP8 rollout (FP8 inference for LLM as judge, FP8 rollout with BF16 training under development). @@ -161,6 +161,7 @@ Leveraging a multi-role distributed architecture with Ray for flexible resource - [ComplementaryRL](https://arxiv.org/abs/2603.17621): Complementary RL is a learning framework that enables agents to effectively learn from experience through the seamless co-evolution of an experience extractor and a policy actor within the RL optimization loop. - [RLix](https://github.com/rlops/rlix): RLix is an RL job manager that lets more RL jobs run concurrently with less waiting by sharing GPU capacity across jobs, while preserving each pipeline’s training behavior and improving GPU utilization. - [TurningPoint-GRPO](https://arxiv.org/abs/2602.06422): A GRPO framework for Flow Matching models in text-to-image generation that alleviates step-wise reward sparsity by modeling step-level incremental rewards and explicitly captures long-term effects via turning points detection, providing dense learning signals for each denoising action. +- [SocioReasoner](https://github.com/AMAP-ML/SocioReasoner): A vision-language method for urban socio-semantic segmentation that employs a render-and-refine mechanism optimized by RL to identify abstract social entities using satellite and map data. - [STAgent](https://arxiv.org/abs/2512.24957): An agentic LLM specialized for spatio-temporal understanding and complex tasks like constrained POI discovery and itinerary planning, featuring hierarchical data curation with 1:10,000 filter ratio and cascaded training (seed SFT + difficulty-aware SFT + RL), achieving strong performance on TravelBench while preserving general capabilities. - [IPRO](https://arxiv.org/abs/2510.14255): A novel video diffusion framework using reinforcement learning to enhance identity preservation in human-centric I2V generation, optimizing diffusion models with face identity scorer and KL-divergence regularization. - [TaoSR-SHE](https://arxiv.org/abs/2510.07972): Stepwise Hybrid Examination Reinforcement Learning Framework for Taobao Search Relevance, with SRPO (hybrid reward model + offline verifier), diversified data filtering, and multi-stage curriculum learning. diff --git a/docker/Dockerfile.torch2100 b/docker/Dockerfile.torch2100 index aa2c29aff..86e311814 100644 --- a/docker/Dockerfile.torch2100 +++ b/docker/Dockerfile.torch2100 @@ -20,6 +20,9 @@ ENV JAVA_HOME=/usr/lib/jvm/java-21-openjdk-amd64 RUN pip install --trusted-host mirrors.aliyun.com --index-url https://mirrors.aliyun.com/pypi/simple/ \ --no-build-isolation "megatron-core @ git+https://github.com/NVIDIA/Megatron-LM.git@core_dev_r0.16.0" +RUN pip install --trusted-host mirrors.aliyun.com --index-url https://mirrors.aliyun.com/pypi/simple/ \ + --no-build-isolation "flash-attn>=2.8.3" + RUN pip install --trusted-host mirrors.aliyun.com --index-url https://mirrors.aliyun.com/pypi/simple/ \ "flash-linear-attention" "transformer-engine[pytorch]" # FIXME: [pytorch] will install core-cu12 by default, install [core-cu13] after [pytorch] to override core-cu12 diff --git a/docs_roll/docs/Development/Developer Guide/support_new_models.md b/docs_roll/docs/Development/Developer Guide/support_new_models.md index 9adf3f1af..89a75cfa4 100644 --- a/docs_roll/docs/Development/Developer Guide/support_new_models.md +++ b/docs_roll/docs/Development/Developer Guide/support_new_models.md @@ -12,7 +12,7 @@ To integrate a new model into **ROLL**, you must supply: | Phase | Pick ≥ 1 backend | |-----------|-----------------| | Inference | `vllm`, `sglang` | -| Training | `DeepSpeed`, `Megatron` | +| Training | `FSDP2`, `Megatron` | --- @@ -30,7 +30,7 @@ https://docs.sglang.ai/supported_models/support_new_models.html ## 2. Training Strategies -### 2.1 `DeepSpeed` +### 2.1 `FSDP2` 1. Ensure your model can be loaded by ```python @@ -41,8 +41,8 @@ https://docs.sglang.ai/supported_models/support_new_models.html 3. Make the model can be loaded in `roll/models/model_providers.py`. Once these steps are complete, you can: -- train with the `deepspeed_train` strategy for `actor_train` worker, and -- with `hf_infer` or `deepspeed_infer` strategy for the `reference` worker. +- train with the `fsdp2_train` strategy for `actor_train` worker, and +- with `hf_infer` or `fsdp2_infer` strategy for the `reference` worker. ### 2.2 `Megatron` diff --git a/docs_roll/docs/Getting Started/FAQ/qa_issues.md b/docs_roll/docs/Getting Started/FAQ/qa_issues.md index 227cfd7ff..b32b03e1d 100644 --- a/docs_roll/docs/Getting Started/FAQ/qa_issues.md +++ b/docs_roll/docs/Getting Started/FAQ/qa_issues.md @@ -61,9 +61,9 @@ gradient_accumulation_steps * per_device_train_batch_size * (world_size/tensor_m ### How to set `gradient_accumulation_steps` and `per_device_train_batch_size`? -#### For DeepSpeed Backend: +#### For FSDP2 Backend: ``` -global_batch_size = per_device_train_batch_size * gradient_accumulation_steps * world_size +global_batch_size = per_device_train_batch_size * gradient_accumulation_steps * world_size / ulysses_size ``` Where `world_size` is the length of `device_mapping` for `actor_train`/`critic` diff --git a/docs_roll/docs/Getting Started/Quick Start/multi_nodes_quick_start.md b/docs_roll/docs/Getting Started/Quick Start/multi_nodes_quick_start.md index 00c687c3e..d495285db 100644 --- a/docs_roll/docs/Getting Started/Quick Start/multi_nodes_quick_start.md +++ b/docs_roll/docs/Getting Started/Quick Start/multi_nodes_quick_start.md @@ -93,7 +93,7 @@ actor_train.model_args.dtype: fp16 actor_infer.model_args.dtype: fp16 reference.model_args.dtype: fp16 -# Switch the large model training framework from DeepSpeed to Megatron-LM, where parameters can be sent in batches for faster execution +# Switch the large model training framework from FSDP2 to Megatron-LM, where parameters can be sent in batches for faster execution strategy_name: megatron_train strategy_config: tensor_model_parallel_size: 1 diff --git a/docs_roll/docs/Getting Started/Quick Start/single_node_quick_start.md b/docs_roll/docs/Getting Started/Quick Start/single_node_quick_start.md index 072afd862..086a8d6d0 100644 --- a/docs_roll/docs/Getting Started/Quick Start/single_node_quick_start.md +++ b/docs_roll/docs/Getting Started/Quick Start/single_node_quick_start.md @@ -67,7 +67,7 @@ actor_train.model_args.dtype: fp16 actor_infer.model_args.dtype: fp16 reference.model_args.dtype: fp16 -# Switch the large model training framework from DeepSpeed to Megatron-LM, where parameters can be sent in batches for faster execution +# Switch the large model training framework from FSDP2 to Megatron-LM, where parameters can be sent in batches for faster execution strategy_name: megatron_train strategy_config: tensor_model_parallel_size: 1 diff --git a/docs_roll/docs/Overview.mdx b/docs_roll/docs/Overview.mdx index 96d57358d..4690f5b7a 100644 --- a/docs_roll/docs/Overview.mdx +++ b/docs_roll/docs/Overview.mdx @@ -62,7 +62,7 @@ Leveraging a multi-role distributed architecture with Ray for flexible resource [Megatron Inference and Training Backend Configuration Guide](https://alibaba.github.io/ROLL/docs/User%20Guides/Configuration/megatron) [LoRA Fine-tuning Configuration Guide](https://alibaba.github.io/ROLL/docs/User%20Guides/Configuration/lora) [FP8 Quantization Configuration Guide](https://alibaba.github.io/ROLL/docs/User%20Guides/Configuration/fp8_rollout) -[DeepSpeed Training Backend Configuration Guide](https://alibaba.github.io/ROLL/docs/User%20Guides/Configuration/deepspeed) +[FSDP2 Training Backend Configuration Guide](https://alibaba.github.io/ROLL/docs/User%20Guides/Configuration/fsdp2) #### Pipeline [RLVR Pipeline for VLM](https://alibaba.github.io/ROLL/docs/User%20Guides/Pipeline/vl_rlvr_pipeline_start) diff --git a/docs_roll/docs/User Guides/Advanced Features/mtp_training.md b/docs_roll/docs/User Guides/Advanced Features/mtp_training.md new file mode 100644 index 000000000..ca62540ce --- /dev/null +++ b/docs_roll/docs/User Guides/Advanced Features/mtp_training.md @@ -0,0 +1,286 @@ +# MTP (Multi-Token Prediction) Training Guide + +## Overview + +MTP (Multi-Token Prediction) is a technique that accelerates inference by predicting multiple future tokens in parallel. The ROLL framework supports training MTP models for both SFT (Supervised Fine-Tuning) and RL (Reinforcement Learning) scenarios. + +## Speculative Decoding Principles + +### The Bottleneck of Autoregressive Generation + +Large language model text generation is an autoregressive process: each token generation requires a complete forward pass. For long-sequence generation (such as mathematical reasoning), this becomes a major performance bottleneck. + +``` +Traditional autoregressive generation: + Step 1: Forward pass → Token 1 + Step 2: Forward pass → Token 2 + Step 3: Forward pass → Token 3 + ... + Each token requires a complete forward pass +``` + +### The Idea Behind Speculative Decoding + +Speculative Decoding breaks this bottleneck through a "predict-verify" approach: + +1. **Draft Phase**: Use a small model to quickly generate K candidate tokens +2. **Verify Phase**: The main model verifies all K tokens in one forward pass +3. **Accept/Reject**: Accept tokens that match the main model's probability distribution, reject those that don't + +``` +Speculative Decoding: + Draft: Small model quickly generates [Token 1, Token 2, Token 3, Token 4] + Verify: Main model verifies all candidates in one forward pass + Result: Accept first 3, reject the 4th + + Equivalent to: Generating 3 tokens with 2 forward passes (1 draft + 1 verify) +``` + +### Why Does This Speed Things Up? + +Key insight: **The main model's forward pass can compute logits for multiple positions in parallel**. + +In traditional generation, when generating a token, only the last position's logits are used, while computations for other positions are wasted. Speculative decoding leverages this by verifying multiple candidate tokens in a single main model forward pass, improving computational efficiency. + +### What Determines the Speedup? + +- **Acceptance Rate**: The closer the draft model's output distribution is to the main model, the higher the acceptance rate +- **Speculative Steps**: The number of candidate tokens generated per speculation +- **Draft Model Efficiency**: The inference speed of the draft model + +An ideal draft model should: +1. Have an output distribution close to the main model (high acceptance rate) +2. Be fast at inference (low draft overhead) +3. Have small parameter count (low memory overhead) + +## What is MTP? + +MTP (Multi-Token Prediction) is an efficient draft model implementation. Unlike using an independent small model, MTP shares weights with the main model and has the following advantages: + +### Difference from Regular LM + +- **Regular LM**: Uses hidden state at position t to predict token at position t+1 +- **MTP**: Uses hidden state at position t + token embedding at position t+1 to predict token at position t+2 + +``` +Regular LM: H(t) → predict(t+1) +MTP: H(t) + E(t+1) → predict(t+2) + ↑ ↑ + hidden state embedding +``` + +### Advantages of MTP + +1. **Weight Sharing**: MTP shares the main model's embedding and output layer, with minimal parameter increase (~5-10%) +2. **High Acceptance Rate**: MTP directly utilizes the main model's hidden states, naturally producing outputs close to the main model's distribution +3. **Simple Training**: Can be jointly trained with the main model without needing to train a separate draft model + +### Use Cases for MTP + +1. **Inference Acceleration**: As a draft model for speculative decoding to accelerate text generation +2. **RL Training Acceleration**: Accelerate rollout generation in scenarios like RLVR, improving training throughput + +### Why Does RL Training Need MTP? + +In RL training (such as RLVR), rollout generation is the main bottleneck: + +1. **Large Generation Demand**: Each training round requires generating many samples +2. **Long-Sequence Generation**: Tasks like mathematical reasoning require long responses +3. **High Inference Engine Load**: The actor_infer worker is often the training bottleneck + +Using MTP speculative decoding can significantly accelerate the rollout process and improve training throughput. + +## Training Modes + +ROLL supports three MTP training modes, configured via the `mtp_training_mode` parameter: + +### 1. disabled (Default) + +MTP weights are loaded but do not participate in training. + +```yaml +actor_train: + mtp_training_mode: disabled # or omit the config +``` + +**Use Cases**: +- Only want to use pre-trained MTP for inference acceleration +- No need to update MTP weights + +### 2. standalone (Recommended for RL) + +MTP is trained independently with truncated gradients, not affecting the main model. + +```yaml +actor_train: + mtp_training_mode: standalone +``` + +**Characteristics**: +- MTP's gradient flow is truncated via `detach()` +- Main model gradients are not affected by MTP training +- Main model and MTP use different learning signals + +**Use Cases**: +- **RL Training**: The main model needs to optimize based on rewards, while MTP needs to learn the main model's generation distribution +- Avoid RL instability affecting MTP + +**How It Works**: + +In standalone mode, the gradient flow between MTP and the main model is completely isolated: +- The main model optimizes based on RL rewards with normal gradient backpropagation +- MTP optimizes based on cross-entropy loss, but gradients do not backpropagate to the main model +- This allows MTP to stably learn the main model's generation distribution without being affected by RL training fluctuations + +### 3. joint (Recommended for SFT) + +MTP is trained jointly with the main model with complete gradient flow. + +```yaml +actor_train: + mtp_training_mode: joint +``` + +**Characteristics**: +- Main model and MTP share gradient flow +- MTP's loss affects main model parameters +- Main model and MTP optimize together + +**Use Cases**: +- **SFT Training**: Want both main model and MTP to learn the target task simultaneously +- MTP serves as an auxiliary training objective + +## Configuration Parameters + +### Training Parameters + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `mtp_training_mode` | `str` | `disabled` | MTP training mode: `disabled`, `standalone`, `joint` | +| `mtp_loss_scaling_factor` | `float` | See below | MTP loss scaling factor | + +**mtp_loss_scaling_factor**: +- Default value is typically `0.3` (referencing DeepSeek-V3) +- In `standalone` mode, MTP loss is directly multiplied by this factor +- In `joint` mode, MTP loss participates in main model gradient updates + +### Inference Engine Configuration (Speculative Decoding) + +Currently, only vLLM in ROLL supports MTP speculative decoding. Configure it in `actor_infer`'s `strategy_config`: + +```yaml +actor_infer: + strategy_args: + strategy_name: vllm + strategy_config: + tensor_parallel_size: 4 + # MTP speculative decoding config + speculative_config: + method: mtp + num_speculative_tokens: 4 +``` + +Note: Regardless of the training mode, when using MTP, you must configure `mtp_num_layers` (the corresponding value from the model's `config.json`) in `actor_train`'s `strategy_config`. + +## Training Examples + +### RLVR Pipeline with MTP + +To enable MTP in RLVR training, configure `mtp_training_mode: standalone` in `actor_train` and `speculative_config` in `actor_infer`: + +```yaml +actor_train: + strategy_args: + strategy_name: megatron_train + strategy_config: + tensor_model_parallel_size: 4 + pipeline_model_parallel_size: 2 + # ... other configs + # MTP training config (uncomment to enable) + #mtp_training_mode: standalone + +actor_infer: + strategy_args: + strategy_name: vllm + strategy_config: + tensor_parallel_size: 4 + # ... other configs + # Speculative decoding config (uncomment to enable) + #speculative_config: + # method: mtp + # num_speculative_tokens: 3 +``` + +For complete configuration examples, refer to: +- `examples/qwen3.5-27B-rlvr_megatron/rlvr_megatron_80GB.yaml` - Qwen3.5-27B Dense model +- `examples/qwen3.5-35BA3-rlvr_megatron/rlvr_megatron_80GB.yaml` - Qwen3.5-35B-A3B MoE model + +### SFT Pipeline with MTP + +SFT training uses `joint` mode for collaborative learning between main model and MTP: + +```yaml +actor_train: + model_args: + model_name_or_path: Qwen/Qwen3.5-7B + flash_attn: sdpa + dtype: bf16 + training_args: + learning_rate: 2.0e-5 + per_device_train_batch_size: 2 + gradient_accumulation_steps: 8 + data_args: + file_name: + - data/sft_data.jsonl + strategy_args: + strategy_name: megatron_train + strategy_config: + tensor_model_parallel_size: 2 + pipeline_model_parallel_size: 1 + # MTP joint training + mtp_training_mode: joint + mtp_loss_scaling_factor: 0.3 +``` + +## Supported Models + +ROLL currently supports MTP training for Qwen3.5 series models: + +| Model | MTP Layers | Notes | +|-------|------------|-------| +| Qwen3.5-7B | 1 | Dense model | +| Qwen3.5-27B | 1 | Dense model | +| Qwen3.5-35B-A3B | 1 | MoE model | + +MTP-related configuration is in the model checkpoint: + +```json +// config.json +{ + "mtp_num_hidden_layers": 1, + "mtp_use_dedicated_embeddings": false +} +``` + +## Notes + +### 1. Mode Selection + +| Scenario | Recommended Mode | Reason | +|----------|------------------|--------| +| RL Training | `standalone` | Isolate RL gradients, MTP learns main model distribution | +| SFT Training | `joint` | Joint optimization, MTP as auxiliary objective | +| Inference-only acceleration | `disabled` | Use pre-trained MTP, no training needed | + +### 2. Performance Monitoring + +Monitor the following metrics to evaluate MTP effectiveness: +- **Acceptance Rate**: The proportion of accepted tokens in speculative decoding +- **Average Acceptance Length**: The average number of tokens accepted per speculation +- **Throughput Improvement**: Speedup compared to non-speculative decoding + +## Related Documentation + +- [vLLM Configuration Guide](../Configuration/vllm.md) - vLLM inference engine detailed configuration +- [RLVR Pipeline](../Pipeline/rlvr_pipeline_start.md) - RLVR training pipeline +- [SFT Pipeline](../Pipeline/sft_pipeline_start.md) - SFT training pipeline diff --git a/docs_roll/docs/User Guides/Advanced Features/opentelemetry_tracing.md b/docs_roll/docs/User Guides/Advanced Features/opentelemetry_tracing.md new file mode 100644 index 000000000..a61d7d1f6 --- /dev/null +++ b/docs_roll/docs/User Guides/Advanced Features/opentelemetry_tracing.md @@ -0,0 +1,91 @@ +# OpenTelemetry Distributed Tracing + +ROLL supports **OpenTelemetry (OTel)** distributed tracing for pipeline observability. When enabled, spans are automatically created for each pipeline step, scheduler operation, and worker execution, providing an end-to-end trace view of the training loop. + +## Introduction + +In large-scale RL training pipelines, understanding execution flow across multiple Ray actors is challenging. Key questions like "where is the time spent?", "what is the critical path?", and "which step is the bottleneck?" are difficult to answer with logs alone. + +ROLL's OpenTelemetry integration provides: + +- **End-to-end trace visualization**: View the full pipeline step as a trace tree (generate → train → validate, etc.) +- **Cross-actor context propagation**: Trace context is automatically propagated through `DataProto.meta_info` via W3C TraceContext headers +- **Zero performance impact**: Benchmarks show no measurable overhead when tracing is enabled +- **Automatic collector management**: The driver automatically starts an OTLP collector (or a built-in Python receiver as fallback) and shuts it down on exit + +### How It Works + +1. **Driver initialization**: When `ROLL_OTEL_ENABLED=1`, the driver auto-selects a free port, starts an OTLP collector, and sets `OTEL_EXPORTER_OTLP_ENDPOINT` for all workers. +2. **Span creation**: The driver creates root spans for each pipeline step. Scheduler and worker spans are nested underneath. +3. **Context propagation**: Before dispatching data to remote actors, trace context is injected into `DataProto.meta_info`. Workers extract it and create child spans. +4. **Export**: All spans are exported via OTLP gRPC to the collector, which writes them to a JSONL file for offline analysis (e.g., import into Jaeger). + +### Collector Selection + +The system automatically selects the best available collector: + +1. **`otelcol-contrib` / `otelcol`** (preferred): Production-ready system binary. If found on `PATH`, a config file is generated and the binary is launched as a subprocess. +2. **Python OTLP Receiver** (fallback): A lightweight pure-Python gRPC receiver built into ROLL. No system dependencies required — uses only `grpcio` and `protobuf` (already pulled in by the OTel SDK). + +## Configuration + +Enable tracing by setting `ROLL_OTEL_ENABLED` in `system_envs` and optionally specifying the output directory: + +```yaml +system_envs: + ROLL_OTEL_ENABLED: "1" + +otlp_output_dir: /path/to/otlp_traces +``` + +- `system_envs.ROLL_OTEL_ENABLED`: Set to `"1"` to enable OpenTelemetry tracing. The driver will auto-resolve a free port and set `OTEL_EXPORTER_OTLP_ENDPOINT` for all workers. +- `otlp_output_dir`: Directory where the OTLP collector writes trace data. Defaults to `./output/otlp`. Traces are saved as `/traces/traces.jsonl`. + +No other configuration is needed — endpoint resolution, collector startup, and SDK initialization are all handled automatically. + +### Viewing Traces + +The output `traces.jsonl` file can be imported into [Jaeger](https://www.jaegertracing.io/) for visualization: + +1. Start a local Jaeger instance with OTLP support +2. Import the `traces.jsonl` file via Jaeger's OTLP JSON import + +Each trace shows the full pipeline step hierarchy: driver → scheduler → workers, with timing for each operation. + +## Instrumented Components + +The following components are automatically instrumented: + +| Component | Spans | +|-----------|-------| +| Pipeline Driver | `pipeline_step`, `stop_server`, `model_update`, `validation`, `generate`, `ref_log_probs`, `train`, etc. | +| RolloutScheduler | `get_batch`, environment step orchestration | +| Agentic Pipeline | Trajectory collection, environment manager operations | +| Workers | Execution spans via `@register(trace=True)` decorator | + +### Custom Instrumentation + +For custom pipeline code, use the tracing utilities directly: + +```python +from roll.utils.telemetry import get_tracer, inject_trace_context, attach_trace_context + +# On the driver/sender side — create spans and propagate context +tracer = get_tracer("driver") +with tracer.start_as_current_span("my_operation"): + inject_trace_context(data.meta_info) + # dispatch data to workers... + +# On the worker/receiver side — attach parent context +with attach_trace_context(data.meta_info): + with get_tracer("worker").start_as_current_span("worker_op"): + # work happens here... +``` + +Key APIs: + +- `get_tracer(name)`: Returns a tracer (or no-op if tracing is disabled/suppressed). +- `inject_trace_context(meta_info)`: Injects the current span context into a dict for cross-actor propagation. +- `attach_trace_context(meta_info)`: Context manager that extracts and attaches parent context from a dict. +- `init_telemetry(service_name)`: Initializes the OTel SDK (called automatically by the framework). +- `shutdown_telemetry()`: Flushes spans and shuts down the collector (called automatically at exit). diff --git a/docs_roll/docs/User Guides/Advanced Features/remote_batch_transfer.md b/docs_roll/docs/User Guides/Advanced Features/remote_batch_transfer.md new file mode 100644 index 000000000..6a6d32cf3 --- /dev/null +++ b/docs_roll/docs/User Guides/Advanced Features/remote_batch_transfer.md @@ -0,0 +1,74 @@ +# ROLL RemoteBatch and Transfer Backend + +The ROLL framework supports **RemoteBatch**, a lazy data transfer mechanism that decouples data storage from data consumption. By integrating with a pluggable transfer backend (e.g., **TransferQueue**), large data batches can be transferred across Ray workers without serializing the entire payload through Ray's object store. This document provides a detailed guide on how to use this feature. + +## Introduction + +In RL training pipelines (especially VLM and Agentic scenarios), `DataProto` batches may contain large tensors (e.g., images, multi-modal embeddings) and non-tensor data (e.g., conversation histories). Transferring these between the RolloutScheduler and training workers via Ray's default serialization has two major problems: + +1. **High memory overhead**: The full data is serialized and deserialized through the Ray object store, doubling peak memory usage. +2. **High transfer latency**: Large data batches (e.g., image data in VLM scenarios) must be fully transferred between workers, causing significant data movement overhead. + +**RemoteBatch** addresses these issues by storing data in an external key-value store and only passing lightweight metadata (keys/references) through Ray. The actual data is **lazily materialized** on the consumer side only when needed, and only the requested fields are fetched. + +### Key Concepts + +- **RemoteBatch**: An abstract base class representing a batch of data stored remotely. It supports the same slicing, indexing, selection, concatenation, and repeat operations as `TensorDict`, but defers actual data access until `materialize()` is called. +- **RowRemoteBatch**: A concrete `RemoteBatch` where data is stored with **row IDs** as keys. Each row (sample) has a unique ID, and the transfer backend stores/retrieves data at row granularity. This is used by the **TransferQueue** backend. +- **ColumnRemoteBatch**: A concrete `RemoteBatch` where data is stored with **column IDs** as keys (one key per field/column). This is used by the **RayMemoryStore** backend. +- **BatchProxy**: A proxy object that wraps both a local `TensorDict` (or `dict`) and a `RemoteBatch`, supporting transparent fallback lookup. When a key is accessed, it first checks the local batch and then falls back to the remote batch. +- **Transfer Backend**: A pluggable storage backend responsible for `put`, `get`, and `delete` operations. Currently supported backends: + - `None` (Dummy): No remote storage; data stays local (default). + - `TransferQueue`: Uses the [TransferQueue](https://github.com/kvcache-ai/TransferQueue) library for high-performance distributed key-value transfer. + +### How It Works + +1. **Upload (`to_remote`)**: The `DataProto.to_remote()` class method converts a local `DataProto` into a remote-backed `DataProto`. It uploads all tensor and non-tensor fields to the transfer backend and returns a new `DataProto` with a `RemoteBatch` reference (no local data). +2. **Transfer**: The lightweight `DataProto` (containing only `RemoteBatch` metadata) is transferred between workers via Ray. Since the metadata is small, serialization is fast. +3. **Materialize (lazy)**: On the consumer side, when specific fields are needed, `RemoteBatch.materialize(fields)` is called to fetch only the requested columns from the backend. The fetched data is cached locally for subsequent accesses. +4. **Drop**: After the batch is consumed, `RemoteBatch.drop()` can be called to delete the data from the backend store. + +## Configuration + +The transfer backend is configured under the `transfer_backend` field in the top-level ROLL configuration: + +```yaml +transfer_backend: + backend_name: TransferQueue + backend_config: + backend: + SimpleStorage: + num_data_storage_units: 16 +``` + +- `backend_name`: The name of the transfer backend to use. + - `null` (default): Disables remote transfer; all data stays local. This is the default behavior when `transfer_backend` is not configured. + - `TransferQueue`: Uses the TransferQueue library for high-performance data transfer. +- `backend_config`: Backend-specific configuration dictionary. For TransferQueue, this corresponds to the TransferQueue initialization config. + - `backend.SimpleStorage.num_data_storage_units`: The number of storage units to shard data across. Can be configured based on the number of CPU cores and cluster nodes. `msgpack` serialization has a maximum 4 GB limit per object, so larger data transfers require more storage units to shard `non_tensor_batch` into smaller pieces. + +### Agentic Pipeline Optimization + +In the Agentic Pipeline, `to_remote` is called at the RolloutScheduler level by default. To further avoid data aggregation overhead from env workers to the RolloutScheduler, you can manually call `to_remote` in the env manager before putting data into the output queue: + +```python +batch = DataProto.to_remote(batch) +output_queue.put(batch) +``` + +:::caution +Manually calling `to_remote` inside environment workers is incompatible with filter. When data is filtered out, the Scheduler does not call `drop()` on the filtered data, causing a leak in the remote store. Only use manual `to_remote` in env workers when filter is not required. (TODO: support automatic `drop()` on filtered RemoteBatch in the Scheduler) +::: + +## Development Status + +| Backend | Status | Notes | +|---------|--------|-------| +| TransferQueue | End-to-end tested | Production-ready. Tested across RLVR, VLM, and Agentic pipelines. | +| RayMemoryStore | Illustration only | Not tested. Provided as a reference implementation for the `ColumnRemoteBatch` pattern. | + +### TODO + +- Avoid full materialization at Trainer: Currently the Trainer calls `materialize()` on the entire RemoteBatch. This can be optimized to only materialize the fields actually needed, avoiding unnecessary data fetching. +- Selective prefetch on Driver: Implement selective prefetch in the Pipeline Driver to batch-fetch fields needed by upcoming steps, reducing the overhead of multiple small fetches. +- Automatic `drop()` on filtered RemoteBatch in the Scheduler to prevent remote storage leaks. diff --git a/docs_roll/docs/User Guides/Advanced Features/router_replay.md b/docs_roll/docs/User Guides/Advanced Features/router_replay.md new file mode 100644 index 000000000..63cb3975f --- /dev/null +++ b/docs_roll/docs/User Guides/Advanced Features/router_replay.md @@ -0,0 +1,187 @@ +# ROUTER REPLAY IN ROLL + +The ROLL framework supports **Router Replay**, a feature that addresses the training-inference mismatch caused by inconsistent expert routing in MoE (Mixture-of-Experts) RL training. By forcing the training-side MoE Router to use a pre-recorded set of routing decisions, Router Replay eliminates routing-level discrepancies at their source and substantially stabilizes training. + +> **Note**: ROLL currently implements only the **R3** mode (Rollout Routing Replay: SGLang inference + Megatron training). R2 is not yet implemented; the combination of R3 with `sequence_packing` is also not yet supported. Both will be added in future releases — please do not enable R3 together with `sequence_packing` for now. + +## 1. Background + +### 1.1 Routing Inconsistency in MoE RL + +In each MoE layer the Router selects top-k experts per token. In RL training, the same set of weights is used by three different roles: + +- **Rollout policy**: the policy used by the inference engine (e.g., SGLang) for sampling. +- **Old policy**: the training-side model state right before this batch's gradient updates. +- **Training policy**: the training-side model that is actively being updated. + +Ideally, all three should produce identical routing, but in practice: + +- **Training vs. inference**: the inference and training engines differ in kernel implementations, numerical precision, and parallelism layouts. Even with identical weights, the two sides may select different top-k experts for the same input. +- **Across gradient steps**: as mini-batch updates proceed, routing decisions also drift along with the weights. + +### 1.2 Why It Matters + +Routing is a discrete choice that gets amplified by the downstream expert outputs. When the rollout-side and training-side selected experts disagree, per-token output probabilities diverge significantly, which leads to: + +- Inflated importance sampling ratios — many samples in PPO/GRPO become heavily clipped or contribute high-variance updates; +- Training collapse in highly off-policy regimes; +- IS correction or TIS-style loss-side compensation alone is often insufficient to recover stability. + +The idea behind Router Replay is simple: **rather than trying to fix the discrepancy at the loss layer, fix the routing mask at the architecture layer so that the training side directly reuses a "reference" routing**, removing this source of mismatch entirely. + +## 2. Design + +### 2.1 The Replay Formula + +Both R2 and R3 share the same mechanism: during the training forward, replace the top-k mask normally produced by router logits with an externally provided mask $I_{\text{ref}}$, and renormalize using the training-side logits $s_{\text{train}}$: + +$$ +g_i = \frac{I_{\text{ref}, i} \cdot \exp(s_{\text{train}, i})}{\sum_j I_{\text{ref}, j} \cdot \exp(s_{\text{train}, j})} +$$ + +Key properties: + +- The **selection** of experts is dictated by $I_{\text{ref}}$, not by training-side argmax. +- The **softmax** is still computed over training-side logits, so router weights still receive gradients normally. + +R2 and R3 differ only in where $I_{\text{ref}}$ comes from. + +### 2.2 R2 — Vanilla Routing Replay (Not Yet Implemented) + +- $I_{\text{ref}}$ comes from a **forward pass that the training engine itself runs with old policy weights**. +- For the first mini-batch, $\theta = \theta_{\text{old}}$, so the replayed forward matches the original forward (effectively on-policy). +- For subsequent off-policy mini-batches, fixing the routing constrains policy staleness. + +R2 addresses "routing drift across gradient steps within the training engine" but **does not solve the discrepancy between inference and training engines**. R2 is currently unimplemented in ROLL; selecting `mode: R2` raises `NotImplementedError`. + +### 2.3 R3 — Rollout Routing Replay (Supported in ROLL) + +- $I_{\text{ref}}$ comes **directly from the routing recorded by the inference engine during rollout**. +- The training side uses an expert selection that is exactly aligned with the sampled trajectory, so the inference-vs-training routing gap is eliminated entirely. +- This also constrains routing drift across gradient steps (same benefit as R2). + +R3 mitigates both the training-inference discrepancy and policy staleness simultaneously, and is the recommended path in ROLL. + +### 2.4 End-to-End R3 Flow in ROLL + +``` +┌─────────────────────────────┐ ┌──────────────────────────────┐ +│ SGLang Rollout │ │ Megatron Training │ +│ │ │ │ +│ generate(...) │ │ forward() │ +│ └─ MoE Router top-k │ │ └─ MoE RouterReplay │ +│ └─ export indices │ ──────► │ └─ replay indices │ +│ [seq, layers, k] │ batch │ in forward │ +│ │ data │ │ +│ return routed_experts │ │ forward → backward │ +└─────────────────────────────┘ └──────────────────────────────┘ +``` + +1. **Sampling**: while generating tokens, SGLang additionally records the top-k experts per MoE layer, producing a `routed_experts` tensor of shape `[seq_len, num_layers, top_k]` returned with the response. +2. **Data movement**: the rollout postprocessor attaches `routed_experts` to each sample in the batch; it then flows through ROLL's standard data path (DP / mini-batch / micro-batch) into the training worker. +3. **Training**: before each forward, Megatron sets per-layer RouterReplay indices from `routed_experts`; the MoE Router skips its top-k computation and replays the external indices; after forward, the action flips to "replay-on-backward" so activation recomputation uses the same routing. + +### 2.5 R3 + Sequence Packing Is Not Yet Supported + +`routed_experts` is laid out **per-sample on the original (unpacked) attention_mask**, while `sequence_packing` concatenates sequences, repads them to a multiple of `2 × CP_SIZE × TP_SIZE`, and re-chunks across CP ranks. The two layouts are not yet reconciled, so: + +- ROLL currently disallows enabling R3 and `sequence_packing` simultaneously; +- When R3 is enabled, please keep `use_sequence_packing: False`; +- A future release will add the index-remapping logic required to make the two work together. + +### 2.6 Compatibility Matrix + +| Feature | R3 Compatibility | +|------------------------------------|--------------------------------| +| Megatron `megatron_train` | Supported | +| SGLang `sglang` rollout | Supported (required) | +| Tensor Parallelism (TP) | Supported | +| Pipeline Parallelism (PP) | Supported | +| Virtual Pipeline Parallelism (VPP) | Supported | +| Dynamic Batching | Supported | +| **Sequence Packing** | **Not yet (planned)** | +| GSPO | Orthogonal, can be combined | +| TIS / IS correction | Coexists; gains are workload-dependent | +| vLLM rollout | Not supported | +| FSDP / DeepSpeed training | Not supported | + +## 3. Implementation + +### 3.1 Rollout Side (SGLang) + +In `sglang_strategy.py`, when `router_replay.mode != "disable"`: + +- The SGLang server is launched with `enable_return_routed_experts=True`; +- Each generate request sets `return_routed_experts=True`; +- The chunked outputs collect `routed_experts` from each chunk's `meta_info` and assemble per-sample records; +- SGLang `>= 0.5.6.post3` is required. + +### 3.2 Training Side (Megatron) + +In `megatron_strategy.py`: + +- During init, `moe_enable_routing_replay` is force-set to `True` so each MoE layer carries a `RouterReplay` instance; +- For `compute_log_probs`-style forwards: replay is enabled only when the batch contains `routed_experts`; otherwise (e.g., reference model) the global action is cleared and the router runs as usual; +- For `train_step`: the strategy asserts that `routed_experts` is present, sets `REPLAY_FORWARD` globally, runs forward/backward, and finally clears global state; +- In `inner_forward_step`, `set_router_replay_data` scatters the recorded indices to the current SP rank; after the forward, the action flips to `REPLAY_BACKWARD` so activation recomputation stays consistent. + +### 3.3 Core Utilities + +`roll/third_party/megatron/router_replay_utils.py` provides: + +- `set_router_replay_data` — scatter top-k indices to the local SP rank; +- `RouterReplayHelper` — query and toggle per-layer RouterReplay state across PP / VPP; +- `get_routed_experts_dtype` — pick `uint8` / `int16` automatically based on the number of experts to reduce memory and transfer cost. + +## 4. Configuration + +### 4.1 How to Enable + +Set `router_replay.mode: R3` on every worker that participates in R3. R3 **must** be enabled symmetrically on both the rollout and the training side — enabling it on only one side has no effect. + +### 4.2 Parameters + +#### `router_replay.mode` + +- **`disable`** (default): Router Replay is off. +- **`R2`**: Vanilla Routing Replay. **Not yet implemented** — configuring this mode raises `NotImplementedError`. +- **`R3`**: Rollout Routing Replay, the only supported mode in ROLL today. + +### 4.3 Full Configuration Example + +```yaml +actor_train: + router_replay: + mode: R3 + + strategy_args: + strategy_name: megatron_train + + # R3 + sequence_packing is not yet supported; keep this disabled + use_sequence_packing: False + +actor_infer: + router_replay: + mode: R3 + + strategy_args: + strategy_name: sglang # requires sglang >= 0.5.6.post3 + +reference: + router_replay: + mode: disable + strategy_args: + strategy_name: megatron_infer +``` + +### 4.4 Usage Recommendations + +1. **Environment & strategies**: `actor_infer` must use `sglang` (≥ 0.5.6.post3); `actor_train` must use `megatron_train`. +2. **Symmetric enablement**: configure `mode: R3` on both rollout and training workers — enabling only one side is a no-op. +3. **Reference model**: keep `mode: disable`. When `routed_experts` is missing from the batch, ROLL automatically skips the replay logic. +4. **Disable sequence packing for now**: until R3 + `sequence_packing` is supported, keep `use_sequence_packing: False` on every worker that uses R3. +5. **Resource overhead**: the `routed_experts` tensor (`[seq_len, num_layers, top_k]`) introduces extra memory and inter-worker transfer cost; ROLL automatically selects the smallest integer dtype to mitigate this. +6. **Relationship with IS / TIS**: Router Replay fixes routing at the architecture level, while IS / TIS correct probability divergence at the loss level. They are complementary and can be used together depending on the workload. +7. **Troubleshooting**: if training-inference mismatch persists, verify (a) the rollout response actually carries `routed_experts`; (b) `moe_enable_routing_replay` is `True` on the training side; (c) sequence packing is disabled on every relevant worker. + +With Router Replay (R3) enabled, ROLL guarantees strict alignment of MoE routing between rollout and training under TP / PP / VPP parallelism, removing a class of mismatch that loss-level correction alone cannot fully address. Future releases will progressively add R2 and the R3 + sequence_packing combination. diff --git a/docs_roll/docs/User Guides/Agentic/agent_runner.md b/docs_roll/docs/User Guides/Agentic/agent_runner.md new file mode 100644 index 000000000..163165814 --- /dev/null +++ b/docs_roll/docs/User Guides/Agentic/agent_runner.md @@ -0,0 +1,557 @@ +# AgentRunner -- Agent Interaction Loop Abstraction + +## Overview + +ROLL's Agentic training framework supports various types of agent-environment interactions: from local gym-like game environments (Sokoban, FrozenLake) to SWE-agents running in remote sandboxes. In ROLL, the entire agent rollout process is fully user-customizable -- the framework does not impose a fixed interaction pattern, but instead provides flexible extension points. However, the framework should not simply leave users to figure out "how to write a rollout loop" on their own. Therefore, ROLL provides common rollout loop implementations for the most accessible multi-turn interaction scenarios (`TrajEnvManager` / `StepEnvManager`), allowing users to quickly understand the framework protocol based on these examples and then customize for their specific needs. + +As scenarios have grown more diverse, the existing `EnvManager` system has exposed bottlenecks in responsibility boundaries and extensibility. The introduction of the `AgentRunner` abstraction decouples the "how agents interact with environments" logic from `EnvManager`, enabling different types of agents to share a unified interface while preserving `EnvManager`'s core capability in training sample construction. + +The AgentRunner abstraction design is derived from practical experience across numerous real-world business scenarios. Thanks to the blog post [Getting Started with AgenticRL Using ROLL](https://zhuanlan.zhihu.com/p/2023433301961049206) for its discussion on Agentic training architecture, which provided valuable reference for this design. + +## 1. Architecture Background: The EnvManager System + +Before understanding AgentRunner, it's important to first understand the responsibilities and existing architecture of EnvManager in ROLL. + +### 1.1 Core Responsibilities of EnvManager + +EnvManager is the core component of Agentic training, responsible for three things: + +1. **Interaction Loop**: Driving multi-turn agent-environment interactions (`reset` / `step` / `make_decision`) +2. **Message Formatting**: Constructing environment observations into LLM-understandable prompts (`format_messages`) +3. **Training Sample Construction**: Converting interaction trajectories into `DataProto` required for training (`formulate_rollouts`) + +### 1.2 Existing Inheritance Hierarchy + +``` +BaseEnvManager +├── TrajEnvManager # Trajectory-level concatenation, suitable for gym-like environments +│ ├── StepEnvManager # Step-level decomposition, suitable for GiGPO algorithm +│ │ └── StepConcatEnvManager # Step-level + observation concatenation variant +│ ├── AgentNativeStepEnvManager # SWE/TerminalBench native environments +│ └── VLTrajEnvManager # Vision-Language model variant +``` + +### 1.3 TrajEnvManager + +`TrajEnvManager` is the most basic implementation. It completes all work in a compact loop: + +```python +# TrajEnvManager core loop (simplified) +def run_rollout_loop(self, data): + while self.running: + rollout_cache = self.reset() # env.reset(seed) + while not done: + lm_output = self.make_decision(cache) # format_messages → LLM inference + rollout_cache = self.step(lm_output) # env.step(action) + rollout = self.formulate_rollouts(cache) # concatenate full trajectory → DataProto + output_queue.put(rollout) +``` + +**Trajectory-level sample construction**: Concatenates all steps' `prompt_ids + response_ids` within an episode into a single training sequence, with the episode reward placed on the last token. The entire trajectory is one training sample. + +**Applicable scenarios**: Gym-like environments such as Sokoban, FrozenLake, WebShop. The environments are lightweight Python objects with fast interactions and no network communication needed. + +**Fine-grained control capabilities**: Since EnvManager fully controls every step of the interaction loop, TrajEnvManager is particularly suitable for scenarios requiring fine-grained control over the interaction process: +- **Rapid multi-agent prototyping**: Freely orchestrate the interaction order and communication logic of multiple agents in `run_rollout_loop`, without being constrained by black-box interfaces +- **Fine-grained `response_mask` management**: Control at the token level which parts participate in loss computation in `formulate_rollouts` (e.g., only training responses from specific turns, masking tool outputs, etc.) +- **Custom prompt concatenation strategies**: Flexibly decide how to retain, truncate, or summarize historical information in `format_messages` + +### 1.4 StepEnvManager + +`StepEnvManager` inherits from `TrajEnvManager` and overrides `format_messages` and `formulate_rollouts`, reducing the granularity from "trajectory-level" to "step-level": + +- **Independent step-level prompts**: Each step constructs a self-contained prompt (including system instructions, history summary, current observation), rather than concatenating all historical tokens +- **Step-level training samples**: Each step generates an independent training sample, with each sample containing that step's prompt + response and corresponding step reward +- **`state_hash` mechanism**: Computes a hash value for each step's state, used for cross-rollout same-state grouping advantage estimation in the GiGPO algorithm + +**Applicable scenarios**: Algorithms requiring step-level advantage estimation (such as GiGPO), or scenarios needing finer-grained credit assignment. + +### 1.5 Why Keep TrajEnvManager / StepEnvManager + +The TrajEnvManager / StepEnvManager system coexists in parallel with ProxyEnvManager + AgentRunner, each with its own advantages: + +| Dimension | TrajEnvManager / StepEnvManager | ProxyEnvManager + AgentRunner | +|-----------|-------------------------------|-------------------------------| +| **Interaction overhead** | Direct function calls, zero serialization overhead | HTTP proxy, with network + serialization overhead | +| **Control granularity** | EnvManager fully controls interaction details (multi-agent orchestration, fine-grained response_mask management, etc.) | Agent executes as black box, fully oriented toward business logic; EnvManager only focuses on sample construction | +| **Prompt construction** | Deeply integrated with `agent_template` / `tokenizer` | Agent constructs prompts independently; ProxyServer transparently intercepts | +| **Applicable environments** | Local gem.Env (Sokoban, FrozenLake, etc.); remote environments require forced alignment with existing abstractions, which is hard to understand | External agent frameworks, remote sandboxes, complex agents | +| **Extending new environments** | Must mix-implement three responsibilities within EnvManager | Only need to implement `AgentRunner.run_job()` | +| **Algorithm support** | Native support for GiGPO step-level decomposition | Supports both step/traj modes via `MessageTracker` | + +**In summary**: +- **TrajEnvManager / StepEnvManager** targets **training research scenarios** -- requiring fine-grained control over the interaction process, custom response_mask construction, rapid multi-agent prototype validation, etc., providing the lowest overhead and tightest control +- **ProxyEnvManager + AgentRunner** targets **business logic scenarios** -- agent implementers focus entirely on agent behavior (how to call tools, how to interact with the environment), without needing to understand any training-side details + +The two serve different scenario requirements and are not substitutes for each other. + +### 1.6 Selection Guide + +Choose the appropriate combination based on your scenario: + +| Scenario | EnvManager | AgentRunner | Description | +|----------|-----------|-------------|-------------| +| Need fine-grained interaction control | `TrajEnvManager` | Not needed | Direct function calls, customizable response_mask, multi-agent orchestration | +| GiGPO step-level training | `StepEnvManager` | Not needed | Built-in `state_hash` grouping, step-level advantage estimation | +| Local gym environment (business logic priority) | `ProxyEnvManager` | `GEMRunner` | Agent only focuses on behavior; ProxyServer transparently collects trajectories | +| Local environment + function-calling | `ProxyEnvManager` | `ToolCallRunner` | Supports OpenAI tool_calls protocol | +| Remote sandbox + agent callback | `ProxyEnvManager` | `PushModeRunner` | Agent calls back via ALB/ingress | +| Remote sandbox + Roll polling | `ProxyEnvManager` | `PullModeRunner` | Roll drives inference via ModelService | +| Custom agent framework | `ProxyEnvManager` | Custom Runner | Only need to implement `run_job(seed)` | + +## 2. AgentRunner Abstraction Design + +### 2.1 Motivation + +In the original architecture, the following issues arose when integrating new types of agents: + +1. **Responsibility coupling**: TrajEnvManager subclasses simultaneously handle interaction loop, message formatting, and training sample construction -- adding new environments requires mixing all logic within a single class +2. **System fragmentation**: ProxyEnvManager uses `MessageTracker` for trajectory collection, which is completely different from TrajEnvManager's `RolloutCache` mechanism. Rock SDK sandbox scenarios are implemented via `HarborRunner`, but its interface (`submit_job(data_item, env_id)`) is too specific +3. **Agent implementers forced to deal with training details**: Agents need to understand training-side concepts like tokenizer, trajectory collection, and sample construction + +The core idea of the AgentRunner abstraction is: **agent implementers focus entirely on agent business behavior (how to interact with the environment, how to call tools, when to terminate), without caring about how trajectories are collected or how training samples are constructed**. This allows business teams to write agent logic in the most natural way, while training infrastructure is handled transparently by the framework. Furthermore, since AgentRunner only calls LLMs through standard OpenAI-compatible interfaces, agents implemented based on AgentRunner can be used in both ROLL training environments and directly migrated to production environments -- simply switch `base_url` from ProxyServer to the actual inference service address, sharing the same agent code between training and production. + +### 2.2 Separation of Concerns + +``` +┌─────────────────────────────────────────────────────────┐ +│ ProxyEnvManager │ +│ ┌───────────────┐ ┌──────────────┐ ┌──────────────┐ │ +│ │ Episode │ │ Trajectory │ │ Sample │ │ +│ │ Scheduling │ │ Collection │ │ Construction │ │ +│ │ run_rollout_ │ │ MessageTracker│ │ formulate_ │ │ +│ │ loop() │ │ ProxyServer │ │ rollouts() │ │ +│ └───────┬───────┘ └──────┬───────┘ └──────────────┘ │ +│ │ │ │ +│ │ HTTP Intercept │ │ +│ ▼ ▼ │ +│ ┌─────────────────────────────────────┐ │ +│ │ AgentRunner │ │ +│ │ ┌─────────┐ ┌──────────────────┐ │ │ +│ │ │ run_job │ │ _llm_request │ │ │ +│ │ │ (seed) │ │ (OpenAI compat) │ │ │ +│ │ └─────────┘ └──────────────────┘ │ │ +│ └─────────────────────────────────────┘ │ +└─────────────────────────────────────────────────────────┘ +``` + +- **AgentRunner** is responsible for: executing a complete episode (loading data, env.reset → interaction loop → returning results) +- **ProxyServer** is responsible for: intercepting AgentRunner's LLM requests, transparently completing inference and trajectory recording +- **ProxyEnvManager** is responsible for: episode scheduling, training sample construction + +### 2.3 Request Interception Mechanism + +AgentRunner sends LLM inference requests to `base_url` through the standard OpenAI-compatible interface (`/v1/chat/completions`). `base_url` points to a local ProxyServer, which acts as a transparent middleware: + +``` +AgentRunner ──HTTP POST──▶ ProxyServer ──Intercept──▶ Record trajectory + │ + └──▶ Forward to actual LLM inference backend + │ + ◀────────────────────────── + │ +AgentRunner ◀──HTTP Response── ProxyServer ◀─────── Return inference result +``` + +1. ProxyServer routes to the corresponding EnvManager via `Authorization: Bearer {env_id}` +2. EnvManager's `process_request` method completes tokenization, inference, and trajectory recording +3. Completely transparent to AgentRunner -- AgentRunner neither knows nor cares that the request was intercepted + +## 3. AgentRunner Base Class + +### 3.1 Core Interface + +```python +# roll/pipeline/agentic/agent_runner/base.py + +class EpisodeResult: + """Structured result of an episode.""" + def __init__( + self, + status: str, # "Finished" | "Failed" | "NoData" | "Timeout" + score: float, # Episode-level reward + step_scores: List[float] = None, # Per-step reward (optional) + agent_exit_reason: str = "", + metrics: Dict[str, Any] = None, # Additional metrics (latency, pass rate, etc.) + ): ... + + +class AgentRunner(ABC): + """Abstract base class for agent interaction loops. + + Subclasses only need to implement the run_job(seed) method: given a seed, + load data, execute a complete episode, and return an EpisodeResult. + """ + + def __init__(self, base_url: str, env_id: int, env_config: DictConfig, **kwargs): + self.base_url = base_url # ProxyServer address + self.env_id = env_id # Environment instance ID + self.env_config = env_config + + @abstractmethod + def run_job(self, seed: int) -> EpisodeResult: + """Execute a complete episode.""" + ... + + def setup(self) -> None: + """One-time initialization (create env instances, etc.).""" + + def teardown(self) -> None: + """Clean up resources.""" + + def _llm_request(self, client, messages, tools=None, tool_choice="auto") -> Dict: + """Send an OpenAI-compatible chat completion request to base_url.""" +``` + +### 3.2 Design Highlights + +- **Clear `run_job(seed)` semantics**: The runner's sole responsibility is "run an episode." `seed` is the only required external input; data loading is handled internally by the runner +- **`env_id` bound to instance**: One runner instance corresponds to one env slot; `_llm_request` automatically uses `self.env_id` for routing +- **`_llm_request` is a protected method**: Encapsulates HTTP details; subclasses call it directly +- **Does not hold tokenizer / pipeline_config**: These are training-side concerns; AgentRunner only focuses on "executing episodes" + +## 4. Built-in Runner Implementations + +### 4.1 GEMRunner -- Local Gym Environments + +`GEMRunner` is designed for local environments defined through the [GEM](https://github.com/axon-rl/gem) interface (Sokoban, FrozenLake, Math, etc.). + +```python +# roll/pipeline/agentic/agent_runner/gem_runner.py + +class GEMRunner(AgentRunner): + """AgentRunner for local gem.Env environments. + + Interaction loop: env.reset(seed) → [construct messages → LLM request → env.step(action)] × N → EpisodeResult + """ + + def setup(self) -> None: + self.env = gem.make(env_id=env_type, **env_params) + # Optional: tool_wrapper wrapping + + def run_job(self, seed: int) -> EpisodeResult: + obs_text, info = self.env.reset(seed=seed) + messages = [{"role": "system", "content": system_template}, ...] + + for turn in range(max_steps): + resp = self._llm_request(client, messages) # → ProxyServer → actual inference + action = resp["choices"][0]["message"]["content"] + obs_text, reward, terminated, truncated, _ = self.env.step(action) + if terminated or truncated: + break + + return EpisodeResult(status="Finished", score=sum(rewards), step_scores=rewards) +``` + +#### ToolCallRunner + +`ToolCallRunner` inherits from `GEMRunner` and supports the OpenAI function-calling protocol. When the LLM returns `tool_calls`, it passes the complete message dict to the environment for execution: + +```python +class ToolCallRunner(GEMRunner): + """GEM environment Runner using OpenAI function-calling protocol.""" + + def run_job(self, seed: int) -> EpisodeResult: + messages, info = self.env.reset(seed=seed) + tools = info.get("tools") + + for _ in range(max_steps): + resp = self._llm_request(client, messages, tools=tools) + message = resp["choices"][0]["message"] + action = message if message.get("tool_calls") else message["content"] + messages, reward, terminated, truncated, _ = self.env.step(action) + ... +``` + +### 4.2 RockAgentRunner -- Remote Sandbox Environments + +`RockAgentRunner` is designed for Rock SDK sandbox environments (SWE-bench, TerminalBench, etc.), where agent code runs in remote sandbox containers. + +``` +RockAgentRunner (base class: job configuration, metrics extraction, data loading) +├── PushModeRunner # Push mode: agent calls back to Roll's ProxyServer +└── PullModeRunner # Pull mode: Roll actively polls the sandbox +``` + +#### Push Mode + +The agent in the sandbox calls back to Roll's ProxyServer via ALB/ingress to obtain LLM inference. Flow: + +``` +PushModeRunner.run_job(seed) + → Load data_item + → Build JobConfig + → Job.submit() → agent runs in sandbox + → agent calls back to ProxyServer via HTTP for LLM inference + → Job.wait() → extract metrics + → Return EpisodeResult +``` + +#### Pull Mode + +Roll actively polls the agent in the sandbox via ModelService, driving each inference step. Flow: + +``` +PullModeRunner.run_job(seed) + → Load data_item + → Build JobConfig + ModelServiceOperator + → Job.submit() + → _inference_loop(): + while True: + request = model_service.anti_call_llm(index, response) # Get agent's LLM request + response = self._llm_request(client, request) # Forward to ProxyServer + # Repeat until SESSION_END + → Job.wait() → extract metrics + → Return EpisodeResult +``` + +### 4.3 Runner and ProxyEnvManager Collaboration + +```python +# ProxyEnvManager.run_rollout_loop (simplified) +def run_rollout_loop(self, data): + self.group_seed = data.meta_info['seed'] + self.env_config['group_seed'] + + while True: + self.episode_id = output_queue.get_episode_id(group_id, env_id) + if self.episode_id is None: + break + + self.message_tracker = MessageTracker(tokenizer, ...) + seed = self.group_seed + self.episode_id + + # Core: only pass seed, entire episode executed by AgentRunner + episode_result = self.agent_runner.run_job(seed) + + rollout = self.formulate_rollouts(episode_result.to_dict()) + output_queue.put(group_id, episode_id, rollout) +``` + +**Unchanged parts of ProxyEnvManager**: +- `process_request()` -- HTTP handler, responsible for LLM inference + trajectory recording +- `MessageTracker` -- incremental tokenization, trajectory collection, fork detection +- `formulate_rollouts()` -- training sample construction (supports both `step` and `traj` modes) +- ProxyServer routing mechanism -- routes to the corresponding handler via `env_id` + +For detailed explanation of `MessageTracker`'s incremental tokenization, prefix aggregation, and fork detection mechanisms, see [Prefix Aggregation](prefix_aggregation.md). + +## 5. Configuration and Usage + +### 5.1 Configuration Method + +Specify the Runner's fully qualified class path via `agent_runner_cls`, consistent with the `env_manager_cls` pattern, dynamically loaded by the framework via `safe_import_class`: + +```yaml +# General configuration pattern +env_manager_cls: roll.pipeline.agentic.env_manager.proxy_env_manager.ProxyEnvManager +agent_runner_cls: +``` + +### 5.2 GEMRunner Configuration Example + +Suitable for local GEM environments like Sokoban, FrozenLake, using `ProxyEnvManager` + `GEMRunner`. + +For the complete example, see `examples/qwen2.5-0.5B-agentic/agentic_val_sokoban_agent_runner.yaml`. Below are the key configuration snippets: + +```yaml +# Global specification of EnvManager and AgentRunner +env_manager_cls: roll.pipeline.agentic.env_manager.proxy_env_manager.ProxyEnvManager +agent_runner_cls: roll.pipeline.agentic.agent_runner.gem_runner.GEMRunner + +# Training parameters +rollout_batch_size: 1024 +sequence_length: 8192 +max_tokens_per_step: 64 +adv_estimator: "grpo" + +# Training environment group configuration +train_env_manager: + max_env_num_per_worker: 16 + num_env_groups: 128 + group_size: 8 # Sample 8 trajectories per prompt + tags: [SimpleSokoban] + num_groups_partition: [128] + +# Validation environment group configuration (mixed evaluation across multiple environments) +val_env_manager: + max_env_num_per_worker: 32 + num_env_groups: 1024 + group_size: 1 + tags: [SimpleSokoban, LargerSokoban, SokobanDifferentGridVocab, FrozenLake] + num_groups_partition: [256, 256, 256, 256] + +# Custom environments: each environment inherits global env_manager_cls and agent_runner_cls +custom_envs: + SimpleSokoban: + ${custom_env.SimpleSokoban} # Reference environment definition from traj_envs.yaml + LargerSokoban: + ${custom_env.LargerSokoban} + FrozenLake: + ${custom_env.FrozenLake} +``` + +Key fields in environment definitions (from `examples/config/traj_envs.yaml`): + +```yaml +custom_env: + SimpleSokoban: + env_type: sokoban + max_steps: ${max_actions_per_traj} # Maximum interaction steps + max_tokens_per_step: ${max_tokens_per_step} + env_manager_cls: ${env_manager_cls} # Inherit global settings + agent_runner_cls: ${agent_runner_cls} # Inherit global settings + agent_system_template: ${agent_system_template} + agent_template: ${agent_template} # Template used by GEMRunner to render observations + env_config: # Environment parameters passed to gem.make() + dim_room: [6, 6] + num_boxes: 1 +``` + +If the environment uses the function-calling protocol, replace `agent_runner_cls` with `ToolCallRunner`: + +```yaml +agent_runner_cls: roll.pipeline.agentic.agent_runner.gem_runner.ToolCallRunner +``` + +### 5.3 RockAgentRunner Configuration Example + +Suitable for remote sandbox environments like SWE-bench: + +```yaml +# Push mode +env_manager_cls: roll.pipeline.agentic.env_manager.proxy_env_manager.ProxyEnvManager +agent_runner_cls: roll.pipeline.agentic.agent_runner.rock.push_runner.PushModeRunner + +custom_envs: + RockNativeEnv: + env_type: rocknative + max_steps: 60 + max_tokens_per_step: 4096 + env_manager_cls: ${env_manager_cls} + agent_runner_cls: ${agent_runner_cls} + env_config: + dataset_name: data/swe_bench_test.jsonl + max_iterations: 100 + tool_call_parser: qwen25 + trajectory_mode: traj # "traj" or "step" + reward_granularity: pass_rate # "pass_rate" or "binary" +``` + +```yaml +# Pull mode +agent_runner_cls: roll.pipeline.agentic.agent_runner.rock.pull_runner.PullModeRunner + +custom_envs: + RockNativeEnv: + ... + env_config: + model_service_port: 28080 # ModelService listening port + ... +``` + +### 5.4 TrajEnvManager Direct Mode + +For simple scenarios that don't need AgentRunner, you can still use TrajEnvManager direct mode: + +```yaml +env_manager_cls: roll.pipeline.agentic.env_manager.traj_env_manager.TrajEnvManager +agent_runner_cls: null # AgentRunner not used + +custom_envs: + SimpleSokoban: + env_type: sokoban + max_steps: 10 + env_manager_cls: ${env_manager_cls} + ... +``` + +## 6. Custom AgentRunner Development + +### 6.1 Implementation Steps + +Developing a custom AgentRunner requires only three steps: + +**Step 1**: Inherit the `AgentRunner` base class + +```python +from roll.pipeline.agentic.agent_runner.base import AgentRunner, EpisodeResult + +class MyCustomRunner(AgentRunner): + def __init__(self, base_url, env_id, env_config, **kwargs): + super().__init__(base_url, env_id, env_config, **kwargs) + # Custom initialization +``` + +**Step 2**: Implement the `run_job(seed)` method + +```python + def run_job(self, seed: int) -> EpisodeResult: + # 1. Load data (if needed) + data = self._load_data(seed) + + # 2. Initialize environment + obs = self.env.reset(seed=seed) + + # 3. Interaction loop + rewards = [] + with httpx.Client(timeout=3600.0) as client: + for step in range(self.max_steps): + messages = self._build_messages(obs) + resp = self._llm_request(client, messages) # Inference via ProxyServer + action = resp["choices"][0]["message"]["content"] + obs, reward, done, _, _ = self.env.step(action) + rewards.append(reward) + if done: + break + + # 4. Return result + return EpisodeResult( + status="Finished", + score=sum(rewards), + step_scores=rewards, + ) +``` + +**Step 3**: Specify in configuration + +```yaml +agent_runner_cls: my_package.my_module.MyCustomRunner +``` + +### 6.2 Development Constraints + +1. **Use `_llm_request` for LLM calls**: Ensure requests go through ProxyServer so that trajectories can be transparently collected +2. **Don't concern yourself with training-side details**: Don't hold tokenizer, don't construct training samples, don't manage trajectory data +3. **Return `EpisodeResult`**: Use the structured return value to make the runner's output contract explicit +4. **`seed` is the sole input**: Data loading and environment initialization should all be driven by seed to ensure reproducibility + +## 7. Module Structure + +``` +roll/pipeline/agentic/agent_runner/ +├── __init__.py # Exports AgentRunner, EpisodeResult +├── base.py # AgentRunner base class + EpisodeResult +├── gem_runner.py # GEMRunner (text mode) + ToolCallRunner (function-calling mode) +└── rock/ + ├── __init__.py + ├── rock_agent_runner.py # RockAgentRunner base class (data loading, job config, metrics extraction) + ├── push_runner.py # PushModeRunner (agent callback mode) + ├── pull_runner.py # PullModeRunner + ModelServiceHarborTrial + ModelServiceOperator + └── sandbox_tool_runner.py # Sandbox tool runner extension +``` + +## Reference Examples + +- GEMRunner + Sokoban: `examples/qwen2.5-0.5B-agentic/agentic_val_sokoban_agent_runner.yaml` +- PushModeRunner + SWE-bench: `examples/agentic_rock/rock_agent_runner_4b.yaml` +- PushModeRunner + 30A3: `examples/agentic_rock/rock_agent_runner_30a3.yaml` + +## Related Documentation + +- [Prefix Aggregation](prefix_aggregation.md) -- MessageTracker's incremental tokenization, prefix aggregation, and fork detection mechanisms +- [Agentic Engineering Practice](agentic_engineer_practice.md) -- EnvManager development protocol, GlobalDataset usage, trajectory synthesis, and other practical guides + +## References + +- [Getting Started with AgenticRL Using ROLL](https://zhuanlan.zhihu.com/p/2023433301961049206) +- [GEM Environment Library](https://github.com/axon-rl/gem) diff --git a/docs_roll/docs/User Guides/Algorithms/Reward_FL.md b/docs_roll/docs/User Guides/Algorithms/Reward_FL.md deleted file mode 100644 index 34734158f..000000000 --- a/docs_roll/docs/User Guides/Algorithms/Reward_FL.md +++ /dev/null @@ -1,94 +0,0 @@ -# Reward Feedback Learning (Reward FL) - -## Introduction - -Reward Feedback Learning (Reward FL) is a reinforcement learning algorithm that optimize diffusion models against a scorer. Reward Fl works as follows: - -1. **Sampling**: For a given prompt and first frame latent, the model generates a corresponding video. -2. **Reward Assignment**: Each video is evaluated and assigned a reward based on its face informations. -3. **Model Update**: The model updates its parameters based on reward signals from the generated videos, reinforcing strategies that obtain higher rewards. - -## Reward FL Configuration Parameters - -In ROLL, the Reward FL algorithm-specific configuration parameters are as follows (`roll.pipeline.diffusion.reward_fl.reward_fl_config.RewardFLConfig`): - -```yaml -# reward fl -learning_rate: 2e-6 -lr_scheduler_type: constant -per_device_train_batch_size: 1 -gradient_accumulation_steps: 1 -warmup_steps: 10 -num_train_epochs: 1 - -model_name: "wan2_2" - -# wan2_2 related -model_paths: ./examples/wan2.2-14B-reward_fl_ds/wan22_paths.json -reward_model_path: /data/models/antelopev2/ -tokenizer_path: /data/models/Wan-AI/Wan2.1-T2V-1.3B/google/umt5-xxl/ -model_id_with_origin_paths: null -trainable_models: dit2 -use_gradient_checkpointing_offload: true -extra_inputs: input_image -max_timestep_boundary: 1.0 -min_timestep_boundary: 0.9 -num_inference_steps: 8 -``` - -### Core Parameter Descriptions - -- `num_train_epochs`: Number of optimization rounds per batch of samples -- `train_batch_size`: Batch size for one train step.In deepspeed training the global train batch size is `per_device_train_batch_size` \* `gradient_accumulation_steps` \* world_size -- `learning_rate`: Learning rate -- `per_device_train_batch_size`: Training batch size per device -- `gradient_accumulation_steps`: Gradient accumulation steps -- `weight_decay`: Weight decay coefficient -- `warmup_steps`: Learning rate warmup steps -- `lr_scheduler_type`: Learning rate scheduler type - -### Wan2_2 Related Parameters - -The following parameters related to Wan2_2 are as follows: -- `model_paths`: Model path of json file, e.g., `wan22_paths.json`, including high_noise_model, low_noise_model, text_encoder, vae. -- `tokenizer_path`: Tokenizer path. Leave empty to auto-download. -- `reward_model_path`: Reward model path, e.g., face model. -- `max_timestep_boundary`: Maximum value of the timestep interval, ranging from 0 to 1. Default is 1. This needs to be manually set only when training mixed models with multiple DiTs, for example, [Wan-AI/Wan2.2-I2V-A14B](https://modelscope.cn/models/Wan-AI/Wan2.2-I2V-A14B). -- `min_timestep_boundary`: Minimum value of the timestep interval, ranging from 0 to 1. Default is 1. This needs to be manually set only when training mixed models with multiple DiTs, for example, [Wan-AI/Wan2.2-I2V-A14B](https://modelscope.cn/models/Wan-AI/Wan2.2-I2V-A14B). -- `model_id_with_origin_paths`: Model ID with origin paths, e.g., Wan-AI/Wan2.1-T2V-1.3B:diffusion_pytorch_model*.safetensors. Comma-separated. -- `trainable_models`: Models to train, e.g., dit, vae, text_encoder. -- `extra_inputs`: Additional model inputs, comma-separated. -- `use_gradient_checkpointing_offload`: Whether to offload gradient checkpointing to CPU memory. -- `num_inference_steps`: Number of inference steps, default is 8 for the distilled wan2_2 model. - - -## Note -- The reward model is constructed based on facial information, Please ensure that the first frame of the video contains a human face. -- Download the reward model(antelopev2.zip) and unzip the onnx files to `reward_model_path` directory. -- Download the official Wan2.2 pipeline and Distilled Wan2.2 DiT safetensors. Put them in the `model_paths` directory, e.g., `wan22_paths.json` file. -- According to the data/example_video_dataset/metadata.csv file, adapt your video dataset to the corresponding format - -## Refernece Model -- `Official Wan2.2 pipeline`: [Wan-AI/Wan2.2-I2V-A14B](https://modelscope.cn/models/Wan-AI/Wan2.2-I2V-A14B) -- `Distilled Wan2.2 DiT safetensors`: [lightx2v/Wan2.2-Lightning](https://huggingface.co/lightx2v/Wan2.2-Lightning/tree/main) -- `Reward Model`: [deepinsight/insightface](https://github.com/deepinsight/insightface/releases/download/v0.7/antelopev2.zip) - -## Preprocess checkpoints -- Run `merge_model.py` to merge multiple files of `Official Wan2.2 pipeline` high noise model and low noise model into one file, respectively. -- Run `merge_lora.py` to merge `Distilled Wan2.2 DiT safetensors` lora to the base model of `Official Wan2.2 pipeline` high noise model and low noise model, respectively. - -## Setup environments -``` -pip install -r requirements_torch260_diffsynth.txt -``` - -## Reference Example - -You can refer to the following configuration file to set up Reward FL training: - -- `./examples/docs_examples/example_reward_fl.yaml` - -Run `run_reward_fl_ds_pipeline.sh` to get start. - -## Reference -[1]: Identity-Preserving Image-to-Video Generation via Reward-Guided Optimization. https://arxiv.org/abs/2510.14255 diff --git a/docs_roll/docs/User Guides/Configuration/config_guide.md b/docs_roll/docs/User Guides/Configuration/config_guide.md index 6bfa2b8b4..1e61c7c4e 100644 --- a/docs_roll/docs/User Guides/Configuration/config_guide.md +++ b/docs_roll/docs/User Guides/Configuration/config_guide.md @@ -36,7 +36,7 @@ actor_train: file_name: xxx/train.json prompt: instruction strategy_args: - strategy_name: megatron_train # deepspeed_train/megatron_train for training + strategy_name: megatron_train # fsdp2_train/megatron_train for training strategy_config: tensor_model_parallel_size: 1 pipeline_model_parallel_size: 1 @@ -75,7 +75,7 @@ reference: ### Model Arguments (**model_args**) - `model_args.dtype`: Set model dtype as fp32, bf16, or fp16, otherwise use config's torch_dtype -- `model_args.disable_gradient_checkpointing`: Disable gradient checkpointing. Applicable only to `actor_train` when `strategy_name` is `deepspeed_train` +- `model_args.disable_gradient_checkpointing`: Disable gradient checkpointing. Applicable only to `actor_train` when `strategy_name` is `fsdp2_train` ### Data Arguments (**data_args**) @@ -95,7 +95,7 @@ Configure generating_args under `actor_infer`. ### Strategy Arguments (**strategy_args**) -- `strategy_args.strategy_name`: The name of training/inference strategy. `deepspeed_train`/`megatron_train` for training, `vllm`/`sglang`/`hf_infer` for inference. +- `strategy_args.strategy_name`: The name of training/inference strategy. `fsdp2_train`/`megatron_train` for training, `vllm`/`sglang`/`hf_infer` for inference. - `strategy_args.strategy_config`: The config of training/inference strategy. Will be passed to `strategy_name`'s constructor. E.g. `strategy_config.tensor_model_parallel_size` for `megatron_train` strategy and `strategy_config.gpu_memory_utilization` for `vllm` strategy. Commonly used strategy configs are listed below: @@ -129,24 +129,21 @@ Commonly used strategy configs are listed below: - `load_format`: The format of the model weights to load. Since there will be a `model update` in the beginning, this value should can be set to `dummy`. -#### DeepSpeed Strategy Config +#### FSDP2 Strategy Config -There are DeepSpeed configurations in `./examples/config/` that can be overridden in the default list for strategy configuration. +- `fsdp_size`: Number of FSDP shards + - If `fsdp_size >= world_size` or `fsdp_size <= 1`: pure FSDP2 mode + - If `fsdp_size < world_size`: HSDP mode with DDP replicas +- `param_dtype`: Parameter data type (e.g., `bf16`, `fp16`, `float32`) +- `reduce_dtype`: Data type for gradient reduction (e.g., `float32`) +- `reshard_after_forward`: Whether to reshard parameters after forward pass + - `true`: Reshard after forward + - `false`: Keep parameters gathered +- `offload_policy`: Whether to enable CPU offloading + - `true`: Offload parameters to CPU when not in use (saves GPU memory) + - `false`: Keep all parameters on GPU (faster but uses more memory) -For example, to use the deepspeed_zero2 strategy, add the following to your config: - -```yaml -defaults: - - ../config/envs@_here_ - - ../config/deepspeed_zero@_here_ - - ../config/deepspeed_zero2@_here_ - - ../config/deepspeed_zero3@_here_ - - ../config/deepspeed_zero3_cpuoffload@_here_ -actor_train: - strategy_args: - strategy_name: deepspeed_train - strategy_config: ${deepspeed_zero2} -``` +Additionally, context parallel (Ulysses) is supported when using FSDP2 strategy by setting `ulysses_size` in `model_args`. ### Training Arguments (**training_args**) @@ -155,7 +152,7 @@ Used for configuring training parameters such as `learning_rate`, `weight_decay` - `training_args.per_device_train_batch_size`: The batch size to use when training. - `training_args.gradient_accumulation_steps`: The number of gradient accumulation steps. -In deepspeed training the global train batch size is `per_device_train_batch_size` \* `gradient_accumulation_steps` \* world_size (a.k.a length of `device_mapping` for `actor_train`/`critic`). +In fsdp2 training the global train batch size is `per_device_train_batch_size` \* `gradient_accumulation_steps` \* world_size / `ulysses_size` (where `world_size` is the length of `device_mapping` for `actor_train`/`critic`). In megatron training the global train batch size is `per_device_train_batch_size` \* `gradient_accumulation_steps` \* world_size / `tensor_model_parallel_size` / `pipeline_model_parallel_size` / `context_parallel_size` (don't need to divide `expert_model_parallel_size`). diff --git a/docs_roll/docs/User Guides/Configuration/config_system.md b/docs_roll/docs/User Guides/Configuration/config_system.md index 48361e746..b3a44071b 100644 --- a/docs_roll/docs/User Guides/Configuration/config_system.md +++ b/docs_roll/docs/User Guides/Configuration/config_system.md @@ -39,7 +39,7 @@ In RLVR scenarios, `RLVRConfig` inherits from `BaseConfig` and serves as the con ### 4. Strategy - Strategy Configuration Strategy configuration defines the training/inference strategy used by each worker node, including: -- Strategy name (such as `megatron_train`, `vllm`, `sglang`, `deepspeed_train`, etc.) +- Strategy name (such as `megatron_train`, `vllm`, `sglang`, `fsdp2_train`, etc.) - Strategy-specific parameters (such as tensor parallel size, pipeline parallel size, etc.) ### 5. Arguments Classes diff --git a/docs_roll/docs/User Guides/Configuration/deepspeed.md b/docs_roll/docs/User Guides/Configuration/deepspeed.md deleted file mode 100644 index 3137506fe..000000000 --- a/docs_roll/docs/User Guides/Configuration/deepspeed.md +++ /dev/null @@ -1,109 +0,0 @@ -# DeepSpeed Training Backend Configuration Guide - -DeepSpeed is Microsoft's efficient deep learning optimization library that provides memory optimization, distributed training, and performance optimization features. This document will provide detailed instructions on how to configure and use the DeepSpeed training backend in the ROLL framework. - -## DeepSpeed Introduction - -DeepSpeed provides multiple optimization techniques, including: -1. **ZeRO Optimization**: Reduces memory usage by partitioning optimizer states, gradients, and parameters -2. **Memory-Efficient Training**: Supports training of large-scale models -3. **High-Performance Communication**: Optimizes communication efficiency in distributed training -4. **Flexible Configuration**: Supports configuration of multiple optimization levels - -## Configuring DeepSpeed Strategy - -In the ROLL framework, DeepSpeed training strategy can be configured by setting `strategy_args` in the YAML configuration file. - -### Configuration Example - -The following is a typical DeepSpeed configuration example (from `examples/qwen2.5-7B-rlvr_megatron/rlvl_lora_zero3.yaml`): - -```yaml -defaults: - - ../config/deepspeed_zero@_here_ - - ../config/deepspeed_zero2@_here_ - - ../config/deepspeed_zero3@_here_ - - ../config/deepspeed_zero3_cpuoffload@_here_ - -actor_train: - model_args: - attn_implementation: fa2 - disable_gradient_checkpointing: true - dtype: bf16 - model_type: ~ - training_args: - learning_rate: 1.0e-5 - weight_decay: 0 - per_device_train_batch_size: 1 - gradient_accumulation_steps: 32 - warmup_steps: 20 - num_train_epochs: 50 - strategy_args: - strategy_name: deepspeed_train - strategy_config: ${deepspeed_zero3} - device_mapping: list(range(0,16)) - infer_batch_size: 4 -``` - -### Configuration Parameter Details - -1. **strategy_name**: Set to `deepspeed_train` to use the DeepSpeed training backend - -2. **strategy_config**: DeepSpeed-specific configuration parameters - - Can reference predefined configuration files, such as `${deepspeed_zero3}` - - Multiple DeepSpeed configuration files are available in the `./examples/config/` directory: - - `deepspeed_zero.yaml`: Basic ZeRO configuration - - `deepspeed_zero2.yaml`: ZeRO-2 configuration - - `deepspeed_zero3.yaml`: ZeRO-3 configuration - - `deepspeed_zero3_cpuoffload.yaml`: ZeRO-3 configuration with CPU offloading - -3. **defaults section**: Import predefined DeepSpeed configurations - ```yaml - defaults: - - ../config/deepspeed_zero@_here_ - - ../config/deepspeed_zero2@_here_ - - ../config/deepspeed_zero3@_here_ - - ../config/deepspeed_zero3_cpuoffload@_here_ - ``` - -4. **device_mapping**: Specify the list of GPU device IDs to use - -## DeepSpeed Configuration Files - -Multiple predefined DeepSpeed configuration files are provided in the `./examples/config/` directory: - -1. **deepspeed_zero.yaml**: Basic ZeRO configuration -2. **deepspeed_zero2.yaml**: ZeRO-2 configuration with optimizer state partitioning -3. **deepspeed_zero3.yaml**: ZeRO-3 configuration with optimizer state, gradient, and parameter partitioning -4. **deepspeed_zero3_cpuoffload.yaml**: ZeRO-3 configuration with CPU offloading - -### Using Predefined Configurations - -To use predefined DeepSpeed configurations, you can reference them in the YAML file like this: - -```yaml -defaults: - - ../config/deepspeed_zero3@_here_ - -actor_train: - strategy_args: - strategy_name: deepspeed_train - strategy_config: ${deepspeed_zero3} -``` - -## Integration with Other Components - -In the configuration example, we can see: - -1. `actor_train` uses DeepSpeed for training -2. `actor_infer` may use other inference backends (such as vLLM) -3. `reference` uses the Hugging Face inference backend -4. Reward models use different inference backends - -This design allows different components to choose the most suitable backend according to their needs. - -## Notes - -1. DeepSpeed requires specific versions of dependency libraries, please ensure compatible versions are installed -2. Different ZeRO levels have different memory and performance characteristics, choose according to specific needs -3. When using LoRA fine-tuning, pay attention to compatibility with DeepSpeed \ No newline at end of file diff --git a/docs_roll/docs/User Guides/Configuration/lora.md b/docs_roll/docs/User Guides/Configuration/lora.md index c6f633d4c..fe28e0163 100644 --- a/docs_roll/docs/User Guides/Configuration/lora.md +++ b/docs_roll/docs/User Guides/Configuration/lora.md @@ -15,7 +15,7 @@ In the ROLL framework, LoRA fine-tuning can be configured by setting relevant pa ### Configuration Example -The following is a typical LoRA configuration example (from `examples/qwen2.5-7B-rlvr_megatron/rlvl_lora_zero3.yaml`): +The following is a typical LoRA configuration example (from `examples/qwen2.5-7B-rlvr_megatron/rlvr_lora_fsdp2.yaml`): ```yaml # LoRA global configuration @@ -40,8 +40,13 @@ actor_train: warmup_steps: 20 num_train_epochs: 50 strategy_args: - strategy_name: deepspeed_train - strategy_config: ${deepspeed_zero3} + strategy_name: fsdp2_train + strategy_config: + fsdp_size: 16 + param_dtype: bf16 + reduce_dtype: float32 + reshard_after_forward: true + offload_policy: false device_mapping: list(range(0,16)) infer_batch_size: 4 @@ -93,15 +98,16 @@ actor_infer: ## LoRA Compatibility with Training Backends -Currently, LoRA fine-tuning only supports the DeepSpeed training backend: +Currently, LoRA fine-tuning supports the FSDP2 and Megatron training backend. The difference of LoRA related configuration between them only exists in `lora_target` as following: -```yaml -actor_train: - strategy_args: - strategy_name: deepspeed_train # LoRA only supports deepspeed_train -``` + - FSDP2 can use `o_proj,q_proj,k_proj,v_proj` as `lora_target`. This is because FSDP2 provides optimization features that integrate well with LoRA. + - Megatron can use `all-linear` including `linear_qkv,linear_proj,linear_fc1,linear_fc2` where: + - `linear_qkv` stands for `q_proj,k_proj,v_proj` in HF + - `linear_proj` stands for `o_proj` in HF + - `linear_fc1` stands for `gate_proj,up_proj` in HF + - `linear_fc2` stands for `down_proj` in HF -This is because DeepSpeed provides optimization features that integrate well with LoRA. +Additionally, vision model weights (vit), vpp and tp+ep ([doc](https://alibaba.github.io/ROLL/docs/User%20Guides/Configuration/megatron)) are not supported when using megatron backend for LoRA fine-tuning. ## Performance Optimization Recommendations @@ -119,7 +125,7 @@ This is because DeepSpeed provides optimization features that integrate well wit ## Notes -1. LoRA fine-tuning currently only supports the DeepSpeed training backend +1. LoRA fine-tuning supports the FSDP2 and Megatron training backend 2. Ensure the model supports LoRA fine-tuning 3. Pay attention to compatibility with LoRA when using gradient checkpointing 4. LoRA fine-tuning performance may differ from full parameter fine-tuning and needs to be evaluated according to specific tasks diff --git a/docs_roll/docs/User Guides/Hardware Support/ascend_usage.md b/docs_roll/docs/User Guides/Hardware Support/ascend_usage.md index aac1dd8b4..b8b6cb136 100644 --- a/docs_roll/docs/User Guides/Hardware Support/ascend_usage.md +++ b/docs_roll/docs/User Guides/Hardware Support/ascend_usage.md @@ -1,21 +1,21 @@ # ROLL x Ascend -Last updated: 05/14/2026. +Last updated: 11/25/2025. We have added support for Huawei Ascend devices in ROLL. ## Hardware Support -Atlas 900 A2 PODc and Atlas 900 A3 PODc +Atlas 900 A2 PODc ## Installation ### Basic Environment Setup | Software | Version | -| -------- |---------| +| -------- | ------- | | Python | 3.11 | -| CANN | 8.5.1 | +| CANN | 8.3.RC1 | ### Create Conda Environment @@ -31,11 +31,11 @@ conda activate roll To use torch and torch_npu in ROLL, install them using the commands below: ``` -# Use CPU-only torch when installing outside the pre-built image -pip install torch==2.8.0 torchvision==0.23.0 torchaudio==2.8.0 --index-url https://download.pytorch.org/whl/cpu +# Use CPU only torch +pip install torch==2.7.1 torchvision==0.22.1 torchaudio==2.7.1 --index-url https://download.pytorch.org/whl/cpu -# Install the torch_npu version matching torch/CANN -pip install torch_npu==2.8.0 +# Install torch_npu 2.7.1 +pip install torch_npu==2.7.1 ``` ### Install vllm & vllm-ascend @@ -44,7 +44,7 @@ To use vllm in ROLL, compile and install vllm and vllm-ascend as follows: ``` # vllm -git clone -b v0.13.0 --depth 1 https://github.com/vllm-project/vllm.git +git clone -b v0.11.0 --depth 1 https://github.com/vllm-project/vllm.git cd vllm pip install -r requirements/build.txt @@ -52,7 +52,7 @@ VLLM_TARGET_DEVICE=empty pip install -v -e . cd .. # vllm-ascend -git clone -b v0.13.0 --depth 1 https://github.com/vllm-project/vllm-ascend.git +git clone -b v0.11.0rc1 --depth 1 https://github.com/vllm-project/vllm-ascend.git cd vllm-ascend pip install -e . @@ -61,11 +61,11 @@ cd .. Or you could install `vllm` and `vllm-ascend` from pre-built wheel: ``` -# Install vllm-project/vllm. The newest supported version is v0.13.0. -pip install vllm==0.13.0 +# Install vllm-project/vllm. The newest supported version is v0.11.0. +pip install vllm==0.11.0 # Install vllm-project/vllm-ascend from pypi. -pip install vllm-ascend==0.13.0 +pip install vllm-ascend==0.11.0rc1 ``` ### Install ROLL @@ -74,7 +74,6 @@ pip install vllm-ascend==0.13.0 git clone https://github.com/alibaba/ROLL.git cd ROLL pip install -r requirements_common.txt -pip install deepspeed==0.16.4 cd .. ``` @@ -82,24 +81,22 @@ cd .. | Software | Description | | --------------------------- | ------------- | -| transformers | >= v4.57.6 | +| transformers | >= v4.57.1 | | flash_attn | not supported | | transformer-engine[pytorch] | not supported | -1. `transformers` v4.57.6 supports enabling `--flash_attention_2`. +1. `transformers` v4.57.1 supports enabling `--flash_attention_2`. 2. `flash_attn` acceleration is not supported currently. 3. `transformer-engine[pytorch]` is currently not supported. ``` -pip install transformers==4.57.6 +pip install transformers==4.57.1 ``` ## Quick Start: Single-Node Deployment Before full usage, we recommend testing the single-node pipeline to verify your environment and installation. -Since Megatron-LM training is not yet supported, first change `strategy_args` in the relevant files to use the `deepspeed` option. - -**Note:** Currently, colocated mode is not supported on NPU. You need to modify `device_mapping` to ensure that training and inference are performed on different cards. +Since Megatron-LM training is not yet supported, first change `strategy_args` in the relevant files to use the `fsdp2` option. 1. Run the single-node pipeline via shell: @@ -124,10 +121,11 @@ python examples/start_agentic_pipeline.py \ | Feature | Example | Training Backend | Inference Backend | Hardware | | --------------- | ------------------------------------------------------------ | ---------------- | ----------------- | ----------------- | -| Agentic | examples/qwen2.5-0.5B-agentic/run_agentic_pipeline_sokoban.sh | DeepSpeed | vLLM | Atlas 900 A3 PODc | -| Agentic-Rollout | examples/qwen2.5-0.5B-agentic/run_agentic_rollout_sokoban.sh | DeepSpeed | vLLM | Atlas 900 A3 PODc | -| RLVR | examples/ascend_examples/run_rlvr_pipeline.sh | DeepSpeed | vLLM | Atlas 900 A2/A3 PODc | +| Agentic | examples/qwen2.5-0.5B-agentic/run_agentic_pipeline_sokoban.sh | FSDP2 | vLLM | Atlas 900 A2 PODc | +| Agentic-Rollout | examples/qwen2.5-0.5B-agentic/run_agentic_rollout_sokoban.sh | FSDP2 | vLLM | Atlas 900 A2 PODc | +| DPO | examples/qwen2.5-3B-dpo_megatron/run_dpo_pipeline.sh | FSDP2 | vLLM | Atlas 900 A2 PODc | +| RLVR | examples/qwen2.5-7B-rlvr_megatron/run_rlvr_pipeline.sh | FSDP2 | vLLM | Atlas 900 A2 PODc | ## Disclaimer -The Ascend support provided in ROLL is intended as a reference example. For production use, please consult official channels. +The Ascend support provided in ROLL is intended as a reference example. For production use, please consult official channels. \ No newline at end of file diff --git a/docs_roll/docs/User Guides/Pipeline/agent_pipeline_start.md b/docs_roll/docs/User Guides/Pipeline/agent_pipeline_start.md index 7c0eeeb04..4fc8edc53 100644 --- a/docs_roll/docs/User Guides/Pipeline/agent_pipeline_start.md +++ b/docs_roll/docs/User Guides/Pipeline/agent_pipeline_start.md @@ -117,7 +117,7 @@ The `AgenticConfig` (defined in `roll/pipeline/agentic/agentic_config.py`) is a * `model_args` – Model architecture, dtype (e.g., bf16), attention type (e.g., flash_attn). * `training_args` – Learning rate, batch size, gradient accumulation. * `generating_args` – `max_new_tokens`, `top_p`, `temperature` for inference. - * `strategy_args` – Distributed strategy (e.g., `deepspeed_train`, `vllm`, `sglang`, `hf_infer`) and its specific configuration. + * `strategy_args` – Distributed strategy (e.g., `fsdp2_train`, `vllm`, `sglang`, `hf_infer`) and its specific configuration. * `device_mapping` – GPU allocation for the worker. --- @@ -265,7 +265,7 @@ python examples/start_agentic_pipeline.py \ * YAML syntax – lint your YAML. * Hydra path/name – verify `--config-path` and `--config-name`. * Pydantic validation – check `roll/pipeline/agentic/agentic_config.py` field definitions. -* **Model Loading Issues** – Confirm paths, model types, and strategies (vLLM, SGLang, DeepSpeed, Megatron-Core). +* **Model Loading Issues** – Confirm paths, model types, and strategies (vLLM, SGLang, FSDP2, Megatron-Core). * **CUDA/GPU Issues** – Adjust `CUDA_VISIBLE_DEVICES`, batch sizes, or `gpu_memory_utilization`. * **Ray Issues** – Ensure Ray is started and resource requests match hardware. * **Environment Registration** – Verify `env_type` values correspond to registered env classes. @@ -277,7 +277,7 @@ python examples/start_agentic_pipeline.py \ * Full configuration definitions: `roll/pipeline/agentic/agentic_config.py` and related dataclasses. * Environment implementations: `roll/agentic/env/` (e.g., `frozen_lake/env.py`). * Core pipeline logic: `roll/pipeline/agentic/agentic_pipeline.py`. -* The project `README.md` provides additional high-level details on features like GRPO, Reasoning Pipeline, and integrations with Ray, DeepSpeed, Megatron-Core, vLLM, and SGlang. +* The project `README.md` provides additional high-level details on features like GRPO, Reasoning Pipeline, and integrations with Ray, FSDP2, Megatron-Core, vLLM, and SGlang. --- diff --git a/docs_roll/docs/User Guides/Pipeline/agentic_pipeline_start.md b/docs_roll/docs/User Guides/Pipeline/agentic_pipeline_start.md index a0b587392..05d5a899b 100644 --- a/docs_roll/docs/User Guides/Pipeline/agentic_pipeline_start.md +++ b/docs_roll/docs/User Guides/Pipeline/agentic_pipeline_start.md @@ -32,7 +32,7 @@ Agentic Pipeline is ROLL's core pipeline for agent training, supporting multiple * **Asynchronous Parallel Rollout at Environment Granularity**: Independent trajectory sampling across environments improves sampling efficiency. * **Asynchronous Training**: Decoupling of rollout/training supports asynchronous training. * **Multi-turn Interaction Support for Local Debugging**: Multi-turn interaction rollout supports local debugging, improving development efficiency for multi-turn interaction business. -* **Flexible Policy Configuration**: Supports multiple distributed training strategies such as Megatron, DeepSpeed, vLLM, etc., allowing flexible configuration based on hardware resources. +* **Flexible Policy Configuration**: Supports multiple distributed training strategies such as Megatron, FSDP2, vLLM, etc., allowing flexible configuration based on hardware resources. * **Efficient Training Optimization**: Supports **Sequence Packing** (concatenating multiple short samples into a continuous sequence to reduce padding) and **Dynamic Batching ** (dynamically grouping samples into batches based on their lengths, applying uniform padding within each batch to the length of the longest sample, thereby minimizing unnecessary computation). For configuration methods and implementation details, please refer to the dedicated documentation for `sequence packing` and `dynamic batching`. diff --git a/docs_roll/docs/User Guides/Pipeline/distill_pipeline_start.md b/docs_roll/docs/User Guides/Pipeline/distill_pipeline_start.md index 8c9d563e3..1c38fcbeb 100644 --- a/docs_roll/docs/User Guides/Pipeline/distill_pipeline_start.md +++ b/docs_roll/docs/User Guides/Pipeline/distill_pipeline_start.md @@ -107,7 +107,7 @@ If the device does not support IPC, use either 'nccl-only' or 'ray' mode instead * `max_grad_norm`: Gradient clipping threshold * ... * **Distributed Strategy** (`strategy_args`) - * `strategy_name`: Distributed strategy to use (e.g., `megatron_train`, `deepspeed_infer`) + * `strategy_name`: Distributed strategy to use (e.g., `megatron_train`, `megatron_infer`) * Strategy-specific parameters: e.g., `tp_size` (tensor parallelism size), `pp_size` (pipeline parallelism size) * `gpu_memory_utilization`: GPU memory utilization (vLLM-specific) * **Device Mapping** (`device_mapping`) diff --git a/docs_roll/docs/User Guides/Pipeline/dpo_pipeline_start.md b/docs_roll/docs/User Guides/Pipeline/dpo_pipeline_start.md index d54e65700..f666ec5b2 100644 --- a/docs_roll/docs/User Guides/Pipeline/dpo_pipeline_start.md +++ b/docs_roll/docs/User Guides/Pipeline/dpo_pipeline_start.md @@ -94,7 +94,7 @@ Configuration files (such as `examples/qwen2.5-3B-dpo_megatron/dpo_config.yaml`) * `max_grad_norm`: Gradient clipping threshold * ... * **Distributed Strategy** (`strategy_args`) - * `strategy_name`: Distributed strategy to use (e.g., `megatron_train`, `deepspeed_infer`) + * `strategy_name`: Distributed strategy to use (e.g., `megatron_train`, `megatron_infer`) * Strategy-specific parameters: e.g., `tp_size` (tensor parallelism size), `pp_size` (pipeline parallelism size) * `gpu_memory_utilization`: GPU memory utilization (vLLM-specific) * **Device Mapping** (`device_mapping`) diff --git a/docs_roll/docs/User Guides/Pipeline/on_policy_distill_pipeline_start.md b/docs_roll/docs/User Guides/Pipeline/on_policy_distill_pipeline_start.md index 479825b28..fdd250ec7 100644 --- a/docs_roll/docs/User Guides/Pipeline/on_policy_distill_pipeline_start.md +++ b/docs_roll/docs/User Guides/Pipeline/on_policy_distill_pipeline_start.md @@ -27,6 +27,9 @@ - [Step 3: Launch the Pipeline](#step-3-launch-the-pipeline) - [Step 4: Monitoring](#step-4-monitoring) - [Step 5: Outputs and Results](#step-5-outputs-and-results) + - [Multi-Teacher OPD](#multi-teacher-opd) + - [Configuration Examples](#configuration-examples) + - [Core Mechanisms](#core-mechanisms) - [FAQ](#faq) - [References](#references) @@ -246,9 +249,9 @@ Configure three roles, automatically mapped to internal Workers: |----------|----------|------| | `student_train` | `actor_train` | Train student model, compute loss using Teacher KL | | `student_infer` | `actor_infer` | Generate trajectories, compute student log_probs | -| `teacher` | `reference` | Compute teacher log_probs | +| `teacher` | `reference` / `references` | Compute teacher log_probs (supports single WorkerConfig or multi-teacher Dict) | -**Note**: Config file uses `student_train`, `student_infer`, `teacher` names, system will automatically map them. +**Note**: Config file uses `student_train`, `student_infer`, `teacher` names, system will automatically map them. For multi-teacher, `teacher` is `Dict[str, WorkerConfig]`, internally normalized to `self.references: Dict[str, Cluster]`. #### Mixed Mode @@ -326,7 +329,16 @@ bash examples/qwen3-8B-onpolicy-distill-megatron/run_onpolicy_distill_pipeline.s | Parameter | Description | Default | |-----------|-------------|---------| | `use_opd` | Enable mixed mode OPD (add Teacher KL to rewards) | `false` | -| `opd_kl_coef` | OPD KL coefficient, controls distillation signal weight relative to external rewards | `1.0` | +| `teacher` | Teacher model config (auto-mapped to reference) | Required | + +#### Multi-Teacher Mode Parameters + +| Parameter | Description | Default | +|-----------|-------------|---------| +| `teacher` | `Dict[str, WorkerConfig]` multi-teacher config | — | +| `teacher.{name}.opd_kl_coef` | Per-teacher KL coefficient | `1.0` | +| `teacher.{name}.tag_included` | Tags this teacher handles; empty means all | `[]` | +| `tag_to_template` | Select different chat templates by tag | `{}` | --- @@ -382,6 +394,172 @@ python examples/start_onpolicy_distill_pipeline.py \ --- +## Multi-Teacher OPD + +### Overview + +Multi-Teacher OPD allows multiple specialized teacher models to simultaneously guide a single student model. Data is routed to the appropriate teacher by domain/tag, avoiding unnecessary computation and enabling more precise distillation. + +``` +┌──────────────────────────────────────────────────────────────────┐ +│ Multi-Teacher OPD Data Flow │ +├──────────────────────────────────────────────────────────────────┤ +│ │ +│ Student Infer rollout → batch (with tag/domain field) │ +│ │ │ +│ ├── [math_dapo data] ──▶ Teacher-32B (math specialist) │ +│ │ compute ref_log_probs_32B │ +│ │ │ +│ └── [KodCode data] ──▶ Teacher-14B (code specialist) │ +│ compute ref_log_probs_14B │ +│ │ │ +│ ▼ │ +│ Compute Advantage: │ +│ For each sample, only accumulate KL from routed teachers: │ +│ advantage = -Σ(opd_kl_coef_i * KL_i) (routed teachers only) │ +│ │ +└──────────────────────────────────────────────────────────────────┘ +``` + +### Configuration Examples + +#### Multi-Teacher Pure OPD Mode + +```yaml +is_pure_opd: true +global_template: qwen3 + +# Select different chat templates by tag (optional) +tag_to_template: + math_dapo: qwen3 # Math data uses qwen3 template (with thinking) + KodCode: qwen3_nothink # Code data uses qwen3_nothink template + +student_train: + model_args: + model_name_or_path: Qwen/Qwen3-8B + data_args: + file_name: + - data/dapo_math_17k_simple_boxed.jsonl + - data/code_KodCode_data.jsonl + domain_interleave_probs: + math_rule: 0.6 + code_rule: 0.4 + device_mapping: list(range(0,8)) + # ... + +student_infer: + model_args: + model_name_or_path: Qwen/Qwen3-8B + device_mapping: list(range(0,8)) + # ... + +# teacher configured as Dict[str, WorkerConfig] +teacher: + teacher_32B: + model_args: + model_name_or_path: Qwen/Qwen3-32B # Math specialist teacher + opd_kl_coef: 1.0 + tag_included: [math_dapo] # Only processes math data + device_mapping: list(range(8,16)) + strategy_args: + strategy_name: megatron_infer + strategy_config: + tensor_model_parallel_size: 2 + pipeline_model_parallel_size: 4 + + teacher_14B: + model_args: + model_name_or_path: Qwen/Qwen3-14B # Code specialist teacher + opd_kl_coef: 1.0 + tag_included: [KodCode] # Only processes code data + device_mapping: list(range(16,24)) + strategy_args: + strategy_name: megatron_infer + strategy_config: + tensor_model_parallel_size: 2 + pipeline_model_parallel_size: 2 + +rewards: + math_rule: + worker_cls: roll.pipeline.rlvr.rewards.math_rule_reward_worker.MathRuleRewardWorker + tag_included: [math_dapo] + code_rule: + worker_cls: roll.pipeline.rlvr.rewards.code_sandbox_reward_worker.CodeSandboxRewardWorker + tag_included: [KodCode] +``` + +#### Mixed Routing Config (General Teacher + Specialist Teacher) + +```yaml +teacher: + teacher_general: + model_args: + model_name_or_path: Qwen/Qwen3-72B + opd_kl_coef: 0.3 + tag_included: [] # Empty = handles all tags (general teacher) + + teacher_math_specialist: + model_args: + model_name_or_path: DeepSeek-Math-67B + opd_kl_coef: 0.7 + tag_included: [math_dapo, aime] # Only handles math +``` + +In this configuration, math data will have KL computed by both `teacher_general` (coef 0.3) and `teacher_math_specialist` (coef 0.7), with both weighted KL values contributing to the advantage. Non-math data only has `teacher_general` participating. + +### Core Mechanisms + +#### 1. Tag Routing + +Each training sample has a `tag` field (e.g., `math_dapo`, `KodCode`). Each teacher declares the tags it handles via `tag_included`: + +- `tag_included: [math_dapo]` — only processes samples with tag `math_dapo` +- `tag_included: []` (empty list) — processes all data (general teacher) + +Routing happens at the ref_log_probs computation stage (pipeline layer). Teachers only run forward on their routed data, avoiding unnecessary inference cost. + +#### 2. Per-Teacher KL Coefficient + +Each teacher has its own `opd_kl_coef`, controlling the weight of that teacher's distillation signal: + +``` +advantage = -Σ(opd_kl_coef_i * KL(student || teacher_i)) +``` + +Only routed teachers participate in the KL accumulation for each sample. + +#### 3. Parallel Inference Optimization + +When multiple teachers use different GPUs (non-overlapping `device_mapping`), the system automatically uses multi-threaded parallel execution for each teacher's forward pass, reducing total inference time. + +#### 4. tag_to_template + +Different domains may require different chat template encoding. With `tag_to_template`, you can use different tokenization templates for specific tags: + +```yaml +tag_to_template: + math_dapo: qwen3 # With thinking token + KodCode: qwen3_nothink # Without thinking token +``` + +Tags not configured in `tag_to_template` fall back to `global_template`. + +### Single Teacher Backward Compatibility + +Single teacher configuration (`teacher` as WorkerConfig rather than Dict) maintains identical behavior to before: + +```yaml +# This config behaves exactly the same as before multi-teacher support +teacher: + model_args: + model_name_or_path: Qwen/Qwen3-32B + device_mapping: list(range(0,16)) +``` + +Internally normalized to `{"default": WorkerConfig}`, the loop executes only once. + +--- + ## FAQ ### Q1: How to configure mixed mode? @@ -453,6 +631,23 @@ Whether in pure OPD mode or mixed mode, Reward Workers must be configured: 2. **Training Monitoring**: Observe reward statistics to monitor training quality 3. **Mixed Mode Additional Role**: External rewards are part of the training signal +### Q4: How to choose between modes? + +- **Pure OPD Mode**: Best for pure distillation training, only needs Teacher KL signal, use `start_onpolicy_distill_pipeline.py` +- **Mixed Mode**: Best for RL + distillation joint training, use `start_rlvr_pipeline.py` with `use_opd: true` + +### Q5: In Multi-Teacher mode, what happens if no teacher is routed to a sample? + +That sample's `total_weighted_kld = 0`: +- In pure OPD mode: `advantage = 0` (sample produces no gradient) +- In mixed mode: `advantage = rl_advantages` (RL signal only, no distillation signal) + +### Q6: Can multiple teachers' device_mapping overlap? + +Yes, but not recommended: +- **Non-overlapping** (recommended): System automatically parallelizes each teacher's forward pass, significantly reducing inference time +- **Overlapping**: System will execute sequentially, no conflicts but total time equals sum of all teachers + --- ## References diff --git a/docs_roll/docs/User Guides/Pipeline/sft_pipeline_start.md b/docs_roll/docs/User Guides/Pipeline/sft_pipeline_start.md index 36f32b36b..76126737a 100644 --- a/docs_roll/docs/User Guides/Pipeline/sft_pipeline_start.md +++ b/docs_roll/docs/User Guides/Pipeline/sft_pipeline_start.md @@ -108,7 +108,7 @@ A typical config includes: - `dataloader_num_workers` - ... - **Strategy args** (`strategy_args`) - - `strategy_name`: e.g., `megatron_train` / `deepspeed_train`, etc. + - `strategy_name`: e.g., `megatron_train` / `fsdp2_train`, etc. - Parallelism-related parameters (tensor/pipeline parallel sizes, etc.) - **Device mapping** (`device_mapping`) - Specifies which GPUs the worker uses diff --git a/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/Development/Developer Guide/support_new_models.md b/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/Development/Developer Guide/support_new_models.md index 920555588..aa31b6192 100644 --- a/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/Development/Developer Guide/support_new_models.md +++ b/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/Development/Developer Guide/support_new_models.md @@ -12,7 +12,7 @@ sidebar_position: 3 | 阶段 | 选择 ≥ 1 个后端 | | --- | --- | | 推理 | `vllm`, `sglang` | -| 训练 | `DeepSpeed`, `Megatron` | +| 训练 | `FSDP2`, `Megatron` | --- @@ -30,7 +30,7 @@ sidebar_position: 3 ## 2. 训练策略 -### 2.1 `DeepSpeed` +### 2.1 `FSDP2` 1. 模型需支持通过以下代码加载： ```python @@ -41,8 +41,8 @@ sidebar_position: 3 3. 在 `roll/models/model_providers.py` 中注册模型。 完成以上步骤后，您可以： -- 使用 `deepspeed_train` 策略来训练 `actor_train` worker，以及 -- 使用 `hf_infer` 或 `deepspeed_infer` 策略在 `reference` worker。 +- 使用 `fsdp2_train` 策略来训练 `actor_train` worker，以及 +- 使用 `hf_infer` 或 `fsdp2_infer` 策略在 `reference` worker。 ### 2.2 `Megatron` diff --git a/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/Getting Started/FAQ/qa_issues.md b/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/Getting Started/FAQ/qa_issues.md index 294a43bd3..bb0f07dfa 100644 --- a/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/Getting Started/FAQ/qa_issues.md +++ b/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/Getting Started/FAQ/qa_issues.md @@ -61,9 +61,9 @@ gradient_accumulation_steps * per_device_train_batch_size * (world_size/tensor_m ### 如何设置 `gradient_accumulation_steps` 和 `per_device_train_batch_size`？ -#### 对于 DeepSpeed Backend： +#### 对于 FSDP2 Backend： ``` -global_batch_size = per_device_train_batch_size * gradient_accumulation_steps * world_size +global_batch_size = per_device_train_batch_size * gradient_accumulation_steps * world_size / ulysses_size ``` 其中 `world_size` 即 `actor_train`/`critic` 的 `device_mapping` 长度 diff --git a/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/Getting Started/Quick Start/multi_nodes_quick_start.md b/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/Getting Started/Quick Start/multi_nodes_quick_start.md index 5837bbb39..736f7508f 100644 --- a/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/Getting Started/Quick Start/multi_nodes_quick_start.md +++ b/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/Getting Started/Quick Start/multi_nodes_quick_start.md @@ -93,7 +93,7 @@ actor_train.model_args.dtype: fp16 actor_infer.model_args.dtype: fp16 reference.model_args.dtype: fp16 -# 大模型训练框架从 DeepSpeed 切换到 Megatron-LM，参数可以批量发送，运行速度更快 +# 大模型训练框架从 FSDP2 切换到 Megatron-LM，参数可以批量发送，运行速度更快 strategy_name: megatron_train strategy_config: tensor_model_parallel_size: 1 diff --git a/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/Getting Started/Quick Start/single_node_quick_start.md b/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/Getting Started/Quick Start/single_node_quick_start.md index cb5264cb9..5d7eb7c97 100644 --- a/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/Getting Started/Quick Start/single_node_quick_start.md +++ b/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/Getting Started/Quick Start/single_node_quick_start.md @@ -68,7 +68,7 @@ actor_train.model_args.dtype: fp16 actor_infer.model_args.dtype: fp16 reference.model_args.dtype: fp16 -# 大模型训练框架从 DeepSpeed 切换到 Megatron-LM，参数可以批量发送，运行速度更快 +# 大模型训练框架从 FSDP2 切换到 Megatron-LM，参数可以批量发送，运行速度更快 strategy_name: megatron_train strategy_config: tensor_model_parallel_size: 1 diff --git a/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/Overview.mdx b/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/Overview.mdx index 271dbf132..5669eb7f9 100644 --- a/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/Overview.mdx +++ b/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/Overview.mdx @@ -64,7 +64,7 @@ Leveraging a multi-role distributed architecture with Ray for flexible resource [Megatron 推理和训练后端配置指南](https://alibaba.github.io/ROLL/zh-Hans/docs/User%20Guides/Configuration/megatron) [LoRA 微调配置指南](https://alibaba.github.io/ROLL/zh-Hans/docs/User%20Guides/Configuration/lora) [FP8 量化配置指南](https://alibaba.github.io/ROLL/zh-Hans/docs/User%20Guides/Configuration/fp8_rollout) -[DeepSpeed 训练后端配置指南](https://alibaba.github.io/ROLL/zh-Hans/docs/User%20Guides/Configuration/deepspeed) +[FSDP2 训练后端配置指南](https://alibaba.github.io/ROLL/zh-Hans/docs/User%20Guides/Configuration/fsdp2) #### 流水线 [VLM RLVR 流水线](https://alibaba.github.io/ROLL/zh-Hans/docs/User%20Guides/Pipeline/vl_rlvr_pipeline_start) diff --git a/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Advanced Features/mtp_training.md b/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Advanced Features/mtp_training.md new file mode 100644 index 000000000..dcdcd147a --- /dev/null +++ b/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Advanced Features/mtp_training.md @@ -0,0 +1,286 @@ +# MTP (Multi-Token Prediction) 训练指南 + +## 概述 + +MTP (Multi-Token Prediction) 是一种通过并行预测多个未来 token 来加速推理的技术。ROLL 框架支持 MTP 模型的训练，可用于 SFT（监督微调）和 RL（强化学习）场景。 + +## 投机采样原理 + +### 自回归生成的瓶颈 + +大语言模型的文本生成是自回归过程：每生成一个 token，都需要完整的前向传播。对于长序列生成（如数学推理），这成为主要的性能瓶颈。 + +``` +传统自回归生成： + Step 1: 前向传播 → Token 1 + Step 2: 前向传播 → Token 2 + Step 3: 前向传播 → Token 3 + ... + 每个 token 都需要一次完整的前向传播 +``` + +### 投机采样的思想 + +投机采样（Speculative Decoding）通过"预测-验证"的方式打破这个瓶颈： + +1. **Draft（草稿）阶段**：使用一个小型模型快速生成 K 个候选 token +2. **Verify（验证）阶段**：主模型一次前向传播并行验证这 K 个 token +3. **Accept/Reject**：接受符合主模型概率分布的 token，拒绝不符合的 + +``` +投机采样： + Draft: 小模型快速生成 [Token 1, Token 2, Token 3, Token 4] + Verify: 主模型一次前向传播验证所有候选 + 结果: 接受前 3 个，拒绝第 4 个 + + 等效于：用 2 次前向传播（1次draft + 1次verify）生成了 3 个 token +``` + +### 为什么能加速？ + +关键洞察：**主模型的前向传播可以并行计算多个位置的 logits**。 + +传统方式下，生成 token 时只计算最后一个位置的 logits，其余位置的计算被浪费了。投机采样利用这一点，用一次主模型前向传播验证多个候选 token，从而提高计算效率。 + +### 加速效果取决于什么？ + +- **接受率**：draft model 的输出分布与主模型越接近，接受率越高 +- **投机步数**：每次投机生成的候选 token 数量 +- **Draft model 效率**：draft model 的推理速度 + +理想的 draft model 应该： +1. 输出分布接近主模型（高接受率） +2. 推理速度快（低 draft 开销） +3. 参数量小（低内存开销） + +## 什么是 MTP？ + +MTP (Multi-Token Prediction) 是一种高效的 draft model 实现。与使用独立小模型不同，MTP 与主模型共享权重，具有以下优势： + +### 与普通 LM 的区别 + +- **普通 LM**：用位置 t 的 hidden state 预测位置 t+1 的 token +- **MTP**：用位置 t 的 hidden state + 位置 t+1 的 token embedding 预测位置 t+2 的 token + +``` +普通 LM: H(t) → predict(t+1) +MTP: H(t) + E(t+1) → predict(t+2) + ↑ ↑ + hidden state embedding +``` + +### MTP 的优势 + +1. **权重共享**：MTP 共享主模型的 embedding 和 output layer，参数量增加很小（约 5-10%） +2. **高接受率**：MTP 直接利用主模型的 hidden states，输出分布自然接近主模型 +3. **训练简单**：可以与主模型联合训练，无需单独训练 draft model + +### MTP 的用途 + +1. **推理加速**：作为投机采样的 draft model，加速文本生成 +2. **RL 训练加速**：在 RLVR 等场景中加速 rollout 生成，提高训练吞吐量 + +### 为什么 RL 训练需要 MTP？ + +在 RL 训练（如 RLVR）中，rollout 生成是主要瓶颈： + +1. **大量生成需求**：每轮训练需要生成大量样本 +2. **长序列生成**：数学推理等任务需要长 response +3. **推理引擎负载高**：actor_infer worker 往往是训练瓶颈 + +使用 MTP 投机采样可以显著加速 rollout 过程，提高训练吞吐量。 + +## 训练模式 + +ROLL 支持三种 MTP 训练模式，通过 `mtp_training_mode` 参数配置： + +### 1. disabled（默认） + +MTP 权重被加载但不参与训练。 + +```yaml +actor_train: + mtp_training_mode: disabled # 或不配置 +``` + +**适用场景**： +- 只想使用预训练的 MTP 进行推理加速 +- 不需要更新 MTP 权重 + +### 2. standalone（推荐用于 RL） + +MTP 独立训练，梯度被截断，不影响主模型。 + +```yaml +actor_train: + mtp_training_mode: standalone +``` + +**特点**： +- MTP 的梯度流被 `detach()` 截断 +- 主模型梯度不受 MTP 训练影响 +- 主模型和 MTP 使用不同的学习信号 + +**适用场景**： +- **RL 训练**：主模型需要根据 reward 优化，MTP 需要学习主模型的生成分布 +- 避免强化学习的不稳定性影响 MTP + +**工作原理**： + +在 standalone 模式下，MTP 和主模型的梯度流完全隔离： +- 主模型根据 RL reward 进行优化，梯度正常反向传播 +- MTP 根据 cross-entropy loss 进行优化，但梯度不会回传到主模型 +- 这样 MTP 可以稳定地学习主模型的生成分布，而不受 RL 训练波动的影响 + +### 3. joint（推荐用于 SFT） + +MTP 与主模型联合训练，梯度完整流动。 + +```yaml +actor_train: + mtp_training_mode: joint +``` + +**特点**： +- 主模型和 MTP 共享梯度流 +- MTP 的 loss 会影响主模型参数 +- 主模型和 MTP 协同优化 + +**适用场景**： +- **SFT 训练**：希望主模型和 MTP 同时学习目标任务 +- MTP 作为辅助训练目标 + +## 配置参数 + +### 训练参数 + +| 参数 | 类型 | 默认值 | 说明 | +|------|------|--------|------| +| `mtp_training_mode` | `str` | `disabled` | MTP 训练模式：`disabled`、`standalone`、`joint` | +| `mtp_loss_scaling_factor` | `float` | 见下文 | MTP loss 缩放系数 | + +**mtp_loss_scaling_factor**： +- 默认值通常为 `0.3`（参考 DeepSeek-V3） +- 在 `standalone` 模式下，MTP loss 直接乘以该系数 +- 在 `joint` 模式下，MTP loss 参与主模型的梯度更新 + +### 推理引擎配置（投机采样） + +目前 ROLL 只有 vLLM 支持 MTP 投机采样功能。在 `actor_infer` 的 `strategy_config` 中配置： + +```yaml +actor_infer: + strategy_args: + strategy_name: vllm + strategy_config: + tensor_parallel_size: 4 + # MTP 投机采样配置 + speculative_config: + method: mtp + num_speculative_tokens: 4 +``` + +另外注意，无论使用哪种模式只要使用 MTP 都需要在 `actor_train` 的 `strategy_config` 中配置 `mtp_num_layers`（模型 config.json 中的相应值） + +## 训练示例 + +### RLVR Pipeline with MTP + +在 RLVR 训练中启用 MTP，需要在 `actor_train` 中配置 `mtp_training_mode: standalone`，在 `actor_infer` 中配置 `speculative_config`： + +```yaml +actor_train: + strategy_args: + strategy_name: megatron_train + strategy_config: + tensor_model_parallel_size: 4 + pipeline_model_parallel_size: 2 + # ... 其他配置 + # MTP 训练配置（取消注释启用） + #mtp_training_mode: standalone + +actor_infer: + strategy_args: + strategy_name: vllm + strategy_config: + tensor_parallel_size: 4 + # ... 其他配置 + # 投机采样配置（取消注释启用） + #speculative_config: + # method: mtp + # num_speculative_tokens: 3 +``` + +完整配置示例请参考： +- `examples/qwen3.5-27B-rlvr_megatron/rlvr_megatron_80GB.yaml` - Qwen3.5-27B Dense 模型 +- `examples/qwen3.5-35BA3-rlvr_megatron/rlvr_megatron_80GB.yaml` - Qwen3.5-35B-A3B MoE 模型 + +### SFT Pipeline with MTP + +SFT 训练使用 `joint` 模式，让主模型和 MTP 协同学习： + +```yaml +actor_train: + model_args: + model_name_or_path: Qwen/Qwen3.5-7B + flash_attn: sdpa + dtype: bf16 + training_args: + learning_rate: 2.0e-5 + per_device_train_batch_size: 2 + gradient_accumulation_steps: 8 + data_args: + file_name: + - data/sft_data.jsonl + strategy_args: + strategy_name: megatron_train + strategy_config: + tensor_model_parallel_size: 2 + pipeline_model_parallel_size: 1 + # MTP 联合训练 + mtp_training_mode: joint + mtp_loss_scaling_factor: 0.3 +``` + +## 支持的模型 + +目前 ROLL 支持 Qwen3.5 系列模型的 MTP 训练： + +| 模型 | MTP 层数 | 备注 | +|------|---------|------| +| Qwen3.5-7B | 1 | Dense 模型 | +| Qwen3.5-27B | 1 | Dense 模型 | +| Qwen3.5-35B-A3B | 1 | MoE 模型 | + +MTP 相关配置在模型 checkpoint 中： + +```json +// config.json +{ + "mtp_num_hidden_layers": 1, + "mtp_use_dedicated_embeddings": false +} +``` + +## 注意事项 + +### 1. 模式选择 + +| 场景 | 推荐模式 | 原因 | +|------|---------|------| +| RL 训练 | `standalone` | 隔离 RL 梯度，MTP 学习主模型分布 | +| SFT 训练 | `joint` | 协同优化，MTP 作为辅助目标 | +| 仅推理加速 | `disabled` | 使用预训练 MTP，无需训练 | + +### 2. 性能监控 + +监控以下指标评估 MTP 效果： +- **接受率（Acceptance Rate）**：投机采样的 token 接受比例 +- **平均接受长度**：每次投机平均接受的 token 数 +- **吞吐量提升**：相比非投机采样的加速比 + +## 相关文档 + +- [vLLM 配置指南](../Configuration/vllm.md) - vLLM 推理引擎详细配置 +- [RLVR Pipeline](../Pipeline/rlvr_pipeline_start.md) - RLVR 训练流程 +- [SFT Pipeline](../Pipeline/sft_pipeline_start.md) - SFT 训练流程 diff --git a/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Advanced Features/opentelemetry_tracing.md b/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Advanced Features/opentelemetry_tracing.md new file mode 100644 index 000000000..974e90eb8 --- /dev/null +++ b/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Advanced Features/opentelemetry_tracing.md @@ -0,0 +1,91 @@ +# OpenTelemetry 分布式追踪 + +ROLL 支持 **OpenTelemetry (OTel)** 分布式追踪，用于 Pipeline 可观测性。启用后，系统会自动为每个 Pipeline 步骤、Scheduler 操作和 Worker 执行创建 Span，提供训练循环的端到端追踪视图。 + +## 简介 + +在大规模 RL 训练流水线中，理解跨多个 Ray Actor 的执行流程非常困难。仅靠日志难以回答"时间花在哪里？"、"关键路径是什么？"、"哪个步骤是瓶颈？"等问题。 + +ROLL 的 OpenTelemetry 集成提供： + +- **端到端追踪可视化**：将完整的 Pipeline 步骤展示为追踪树（generate → train → validate 等） +- **跨 Actor 上下文传播**：追踪上下文通过 W3C TraceContext 标头自动经由 `DataProto.meta_info` 传播 +- **零性能影响**：基准测试表明启用追踪无可测量的性能开销 +- **自动 Collector 管理**：Driver 自动启动 OTLP Collector（或内置 Python Receiver 作为回退），退出时自动关闭 + +### 工作原理 + +1. **Driver 初始化**：当 `ROLL_OTEL_ENABLED=1` 时，Driver 自动选择空闲端口，启动 OTLP Collector，并为所有 Worker 设置 `OTEL_EXPORTER_OTLP_ENDPOINT`。 +2. **Span 创建**：Driver 为每个 Pipeline 步骤创建根 Span，Scheduler 和 Worker 的 Span 嵌套在其下。 +3. **上下文传播**：在将数据分发到远程 Actor 之前，追踪上下文被注入到 `DataProto.meta_info` 中。Worker 提取上下文并创建子 Span。 +4. **导出**：所有 Span 通过 OTLP gRPC 导出到 Collector，Collector 将其写入 JSONL 文件供离线分析（如导入 Jaeger）。 + +### Collector 选择 + +系统自动选择最佳可用的 Collector： + +1. **`otelcol-contrib` / `otelcol`**（首选）：生产级系统二进制文件。如果在 `PATH` 中找到，会自动生成配置文件并作为子进程启动。 +2. **Python OTLP Receiver**（回退）：ROLL 内置的轻量级纯 Python gRPC 接收器。无需系统依赖 — 仅使用 `grpcio` 和 `protobuf`（已由 OTel SDK 引入）。 + +## 配置 + +通过在 `system_envs` 中设置 `ROLL_OTEL_ENABLED` 启用追踪，可选指定输出目录： + +```yaml +system_envs: + ROLL_OTEL_ENABLED: "1" + +otlp_output_dir: /path/to/otlp_traces +``` + +- `system_envs.ROLL_OTEL_ENABLED`：设置为 `"1"` 启用 OpenTelemetry 追踪。Driver 会自动解析空闲端口并为所有 Worker 设置 `OTEL_EXPORTER_OTLP_ENDPOINT`。 +- `otlp_output_dir`：OTLP Collector 写入追踪数据的目录。默认为 `./output/otlp`。追踪数据保存为 `/traces/traces.jsonl`。 + +无需其他配置 — 端点解析、Collector 启动和 SDK 初始化均自动处理。 + +### 查看追踪 + +输出的 `traces.jsonl` 文件可导入 [Jaeger](https://www.jaegertracing.io/) 进行可视化： + +1. 启动支持 OTLP 的本地 Jaeger 实例 +2. 通过 Jaeger 的 OTLP JSON 导入功能导入 `traces.jsonl` 文件 + +每个 Trace 展示完整的 Pipeline 步骤层级：driver → scheduler → workers，包含每个操作的耗时。 + +## 已插桩组件 + +以下组件已自动插桩： + +| 组件 | Span | +|------|------| +| Pipeline Driver | `pipeline_step`、`stop_server`、`model_update`、`validation`、`generate`、`ref_log_probs`、`train` 等 | +| RolloutScheduler | `get_batch`、环境步骤编排 | +| Agentic Pipeline | 轨迹收集、环境管理器操作 | +| Worker | 通过 `@register(trace=True)` 装饰器创建执行 Span | + +### 二次开发 + +在自定义 Pipeline 代码中，可以直接使用追踪工具： + +```python +from roll.utils.telemetry import get_tracer, inject_trace_context, attach_trace_context + +# 在 Driver/发送端 — 创建 Span 并传播上下文 +tracer = get_tracer("driver") +with tracer.start_as_current_span("my_operation"): + inject_trace_context(data.meta_info) + # 分发数据到 Worker... + +# 在 Worker/接收端 — 附加父上下文 +with attach_trace_context(data.meta_info): + with get_tracer("worker").start_as_current_span("worker_op"): + # 执行工作... +``` + +核心 API： + +- `get_tracer(name)`：返回 Tracer（追踪禁用/被抑制时返回 no-op）。 +- `inject_trace_context(meta_info)`：将当前 Span 上下文注入字典，用于跨 Actor 传播。 +- `attach_trace_context(meta_info)`：上下文管理器，从字典中提取并附加父上下文。 +- `init_telemetry(service_name)`：初始化 OTel SDK（框架自动调用）。 +- `shutdown_telemetry()`：刷新 Span 并关闭 Collector（退出时自动调用）。 diff --git a/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Advanced Features/remote_batch_transfer.md b/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Advanced Features/remote_batch_transfer.md new file mode 100644 index 000000000..eede46713 --- /dev/null +++ b/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Advanced Features/remote_batch_transfer.md @@ -0,0 +1,75 @@ +# ROLL RemoteBatch 与传输后端 + +ROLL 框架支持 **RemoteBatch**，一种惰性数据传输机制，将数据存储与数据消费解耦。通过集成可插拔的传输后端（如 **TransferQueue**），大数据批次可以在 Ray Worker 之间传输，而无需通过 Ray 的 Object Store 序列化整个数据负载。本文档详细介绍如何使用这一功能。 + +## 简介 + +在 RL 训练流水线中（特别是 VLM 和 Agentic 场景），`DataProto` 批次可能包含大型张量（如图片、多模态 embedding）和非张量数据（如对话历史）。通过 Ray 默认的序列化方式在 RolloutScheduler 和训练 Worker 之间传输这些数据存在以下问题： + +1. **内存开销高**：完整数据需要通过 Ray Object Store 进行序列化和反序列化，峰值内存使用量翻倍。 +2. **传输延迟大**：大数据批次（如 VLM 场景中的图片数据）需要在 Worker 之间完整传输，导致数据搬运耗时显著。 + +**RemoteBatch** 通过将数据存储在外部键值存储中，仅通过 Ray 传递轻量级元数据（键/引用）来解决这些问题。实际数据在消费侧按需**惰性物化（lazily materialized）**，且仅获取请求的字段。 + +### 核心概念 + +- **RemoteBatch**：表示远程存储数据批次的抽象基类。它支持与 `TensorDict` 相同的切片、索引、选择、拼接和重复操作，但将实际数据访问延迟到调用 `materialize()` 时执行。 +- **RowRemoteBatch**：以**行 ID** 为键存储数据的具体 `RemoteBatch` 实现。每行（样本）有一个唯一 ID，传输后端以行粒度存储/检索数据。**TransferQueue** 后端使用此实现。 +- **ColumnRemoteBatch**：以**列 ID** 为键存储数据的具体 `RemoteBatch` 实现（每个字段/列一个键）。**RayMemoryStore** 后端使用此实现。 +- **BatchProxy**：包装本地 `TensorDict`（或 `dict`）和 `RemoteBatch` 的代理对象，支持透明的回退查找。访问键时，先检查本地 batch，再回退到远程 batch。 +- **传输后端（Transfer Backend）**：负责 `put`、`get` 和 `delete` 操作的可插拔存储后端。目前支持的后端： + - `None`（Dummy）：无远程存储，数据保留在本地（默认）。 + - `TransferQueue`：使用 [TransferQueue](https://github.com/kvcache-ai/TransferQueue) 库进行高性能分布式键值传输。 + +### 工作原理 + +1. **上传 (`to_remote`)**：`DataProto.to_remote()` 类方法将本地 `DataProto` 转换为远程支持的 `DataProto`。它将所有张量和非张量字段上传到传输后端，并返回一个包含 `RemoteBatch` 引用的新 `DataProto`（无本地数据）。 +2. **传输**：轻量级的 `DataProto`（仅包含 `RemoteBatch` 元数据）通过 Ray 在 Worker 之间传输。由于元数据很小，序列化速度很快。 +3. **物化（惰性）**：在消费侧，当需要特定字段时，调用 `RemoteBatch.materialize(fields)` 仅从后端获取请求的列。获取的数据会缓存在本地供后续访问。 +4. **清理（Drop）**：数据消费完成后，可以调用 `RemoteBatch.drop()` 从后端存储中删除数据。 + +## 配置 + +传输后端通过 ROLL 顶层配置中的 `transfer_backend` 字段进行配置： + +```yaml +transfer_backend: + backend_name: TransferQueue + backend_config: + backend: + SimpleStorage: + num_data_storage_units: 16 +``` + +- `backend_name`：要使用的传输后端名称。 + - `null`（默认）：禁用远程传输，所有数据保留在本地。未配置 `transfer_backend` 时的默认行为。 + - `TransferQueue`：使用 TransferQueue 库进行高性能数据传输。 +- `backend_config`：后端特定的配置字典。对于 TransferQueue，对应 TransferQueue 的初始化配置。 + - `backend.SimpleStorage.num_data_storage_units`：数据分片的存储单元数量。可以根据 CPU 核数和集群节点数进行配置。`msgpack` 序列化单个对象有最大 4GB 的限制，因此传输大数据时需要更多的 storage unit 来将 `non_tensor_batch` 分片成更小的块。 + +### Agentic Pipeline 优化 + +在 Agentic Pipeline 中，默认在 RolloutScheduler 层面调用 `to_remote`。如果要完全避免从 env worker 汇总数据到 RolloutScheduler 的开销，可以在 env manager 将数据放入 output queue 之前手动调用 `to_remote`： + +```python +batch = DataProto.to_remote(batch) +output_queue.put(batch) +``` + +:::caution +在环境 Worker 中手动调用 `to_remote` 与 filter 不兼容。当数据被 filter 过滤掉时，Scheduler 不会对被过滤的数据调用 `drop()`，导致远程存储中的数据泄漏。仅在不需要 filter 时才在 env worker 中使用手动 `to_remote`。（TODO：后续将支持 Scheduler 对被 filter 的 RemoteBatch 自动调用 `drop()`） +::: + +## 开发状态 + +| 后端 | 状态 | 说明 | +|------|------|------| +| TransferQueue | 端到端已测试 | 生产可用。已在 RLVR、VLM 和 Agentic Pipeline 中测试通过。 | +| RayMemoryStore | 仅作示例 | 未经测试。仅作为 `ColumnRemoteBatch` 模式的参考实现提供。 | + +### TODO + +- 避免在 Trainer 侧全量物化：当前 Trainer 会对整个 RemoteBatch 调用 `materialize()`，后续可优化为仅物化实际需要的字段，避免不必要的数据拉取。 +- Driver 侧选择性预取：在 Pipeline Driver 中实现选择性 prefetch，根据后续步骤的需求批量预取所需字段，减少多次小规模拉取的开销。 +- Scheduler 对被 filter 的 RemoteBatch 自动调用 `drop()`，避免远程存储泄漏。 + diff --git a/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Advanced Features/router_replay.md b/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Advanced Features/router_replay.md new file mode 100644 index 000000000..5538bfc5c --- /dev/null +++ b/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Advanced Features/router_replay.md @@ -0,0 +1,187 @@ +# ROLL ROUTER REPLAY + +ROLL 框架支持 **Router Replay**（路由复放）功能，用于解决 MoE（混合专家）模型在 RL 训练中由专家路由不一致引起的训练-推理失配问题。该特性在训练阶段强制使用预先记录的路由决策，从源头消除 MoE Router 在不同执行环境下的偏差，使训练更加稳定。 + +> **注意**：当前 ROLL 仅实现了 **R3** 模式（Rollout Routing Replay：SGLang 推理 + Megatron 训练）。R2 暂未实现；R3 与 `sequence_packing` 的组合也暂未支持，后续版本会逐步补齐，目前请勿同时启用。 + +## 1. 背景 + +### 1.1 MoE RL 中的路由不一致 + +MoE 模型每一层 Router 会为每个 token 选择 top-k 个专家。在 RL 训练里，同一份模型权重会被三个角色使用： + +- **Rollout policy**：推理引擎（如 SGLang）中负责采样的 policy。 +- **Old policy**：训练引擎中、本批次更新前的模型状态。 +- **Training policy**：训练引擎中、正在做梯度更新的模型。 + +理想情况下三者应当给出一致的路由结果，但实际上： + +- **训练 vs 推理**：推理引擎与训练引擎在算子实现、精度、并行布局上存在差异，即使权重相同，对同一输入也可能选出不同的 top-k 专家。 +- **多次梯度更新内部**：随着 mini-batch 不断更新，路由也会随权重漂移。 + +### 1.2 路由不一致带来的影响 + +路由的离散选择会被后续 expert 的输出放大。当推理与训练侧选中的专家不同，token 级别的输出概率会出现明显偏差，进而： + +- 放大 importance sampling 比，使 PPO/GRPO 的有效样本被严重 clip 或方差激增； +- 在 off-policy 比例较大的场景下出现训练崩溃； +- 让 IS 修正、TIS 等 loss 层补偿手段难以单独解决问题。 + +Router Replay 的思路是：**与其在 loss 层做事后修正，不如在模型架构层固定路由 mask，让训练侧直接复用一份"参考路由"**，从源头去掉这一差异。 + +## 2. 实现原理 + +### 2.1 复放公式 + +无论是 R2 还是 R3，做法都可以抽象成：在训练 forward 中，把原本由 router logits 决定的 top-k mask 替换为外部传入的 mask $I_{\text{ref}}$，再用训练侧的 logits $s_{\text{train}}$ 重新归一化得到专家权重： + +$$ +g_i = \frac{I_{\text{ref}, i} \cdot \exp(s_{\text{train}, i})}{\sum_j I_{\text{ref}, j} \cdot \exp(s_{\text{train}, j})} +$$ + +要点： + +- 选哪几个专家由 $I_{\text{ref}}$ 决定，不再由训练侧 argmax 决定； +- 但 softmax 仍作用在训练侧 logits 上，因此 router 权重的梯度可以正常回传。 + +R2 与 R3 的差别仅在于 $I_{\text{ref}}$ 的来源不同。 + +### 2.2 R2 — Vanilla Routing Replay（暂未实现） + +- $I_{\text{ref}}$ 来自 **训练引擎自己用 old policy 跑一次 forward** 记录下的路由。 +- 第一个 mini-batch 时 $\theta = \theta_{\text{old}}$，复放结果与原始 forward 一致，等价于 on-policy。 +- 后续 mini-batch 在 off-policy 场景下，复放固定了路由，从而约束 policy staleness。 + +R2 解决的是"训练侧自身在多次更新中路由漂移"的问题，但**无法解决推理引擎与训练引擎之间的差异**。在 ROLL 中 R2 暂未实现，配置 `mode: R2` 会抛出 `NotImplementedError`。 + +### 2.3 R3 — Rollout Routing Replay（ROLL 当前支持） + +- $I_{\text{ref}}$ 直接来自 **推理引擎在 rollout 时记录的路由**。 +- 训练侧拿到的是与采样轨迹严格对齐的专家选择，因此 training-inference 的路由差异被彻底消除。 +- 同时也限制了多次更新中的路由漂移（与 R2 同方向受益）。 + +R3 是同时缓解 training-inference discrepancy 与 policy staleness 的更强方案，是 ROLL 当前的默认推荐路径。 + +### 2.4 R3 在 ROLL 中的端到端流程 + +``` +┌─────────────────────────────┐ ┌──────────────────────────────┐ +│ SGLang Rollout │ │ Megatron Training │ +│ │ │ │ +│ generate(...) │ │ forward() │ +│ └─ MoE Router 计算 top-k │ │ └─ MoE RouterReplay │ +│ └─ 同步导出索引 │ ──────► │ └─ 用导出的索引 │ +│ [seq, layers, k] │ batch │ 参与 forward │ +│ │ data │ │ +│ 返回 routed_experts │ │ forward → backward │ +└─────────────────────────────┘ └──────────────────────────────┘ +``` + +1. **采样阶段**：SGLang 在生成 token 时，额外记录每一层 MoE 的 top-k 专家，形成形如 `[seq_len, num_layers, top_k]` 的 `routed_experts` 张量并随响应返回。 +2. **数据搬运**：rollout 后处理把 `routed_experts` 挂到 batch 的每条样本上，沿 ROLL 标准的数据通路（DP / mini-batch / micro-batch）流向训练 worker。 +3. **训练阶段**：Megatron 侧在 forward 前根据 `routed_experts` 设置每层 RouterReplay 的索引；MoE Router 跳过 top-k 计算，直接复放外部索引；forward 完成后切换为"backward 复放"，确保激活重计算下也用同一份路由。 + +### 2.5 R3 与 Sequence Packing 暂未支持组合使用 + +`routed_experts` 是按"每条样本独立、原始 attention_mask 布局"组织的，而 `sequence_packing` 会把多个序列拼接、按 `2 × CP_SIZE × TP_SIZE` 重新对齐、并按 CP 交错切分。两套布局当前没有打通，因此： + +- ROLL 暂不支持 R3 与 `sequence_packing` 同时启用； +- 启用 R3 时请保持 `use_sequence_packing: False`； +- 后续版本会补齐两者间的索引重映射，使其可以协同工作。 + +### 2.6 兼容性矩阵 + +| 特性 | R3 兼容性 | +|------------------------------------|----------------------------| +| Megatron `megatron_train` | 支持 | +| SGLang `sglang` rollout | 支持（必需） | +| 张量并行 TP | 支持 | +| 流水并行 PP | 支持 | +| Virtual Pipeline Parallelism (VPP) | 支持 | +| Dynamic Batching | 支持 | +| **Sequence Packing** | **暂未支持，后续会支持** | +| GSPO | 正交，可叠加 | +| TIS / IS 修正 | 可共存，收益视场景而定 | +| vLLM rollout | 不支持 | +| FSDP / DeepSpeed 训练 | 不支持 | + +## 3. 实现流程 + +### 3.1 Rollout 端（SGLang） + +`sglang_strategy.py` 中，当 `router_replay.mode != "disable"` 时： + +- SGLang 服务以 `enable_return_routed_experts=True` 启动； +- 每个 generate 请求带上 `return_routed_experts=True`； +- 多 chunk 输出从 `meta_info` 中收集 `routed_experts`，按样本拼装； +- 要求 SGLang `>= 0.5.6.post3`。 + +### 3.2 训练端（Megatron） + +`megatron_strategy.py` 中： + +- 初始化时强制 `moe_enable_routing_replay=True`，让每个 MoE 层带上 `RouterReplay` 实例； +- `compute_log_probs` 类 forward：仅当 batch 包含 `routed_experts` 时才进入复放，否则（如 reference model）走默认路由； +- `train_step`：断言 `routed_experts` 存在，全局设置 `REPLAY_FORWARD`，forward/backward 完成后清理全局状态； +- `inner_forward_step`：通过 `set_router_replay_data` 把记录的索引按 SP rank 分发；forward 后切换为 `REPLAY_BACKWARD`，保证激活重计算正确。 + +### 3.3 核心工具 + +`roll/third_party/megatron/router_replay_utils.py` 提供： + +- `set_router_replay_data` — 将 top-k 索引按 SP rank 分发； +- `RouterReplayHelper` — 跨 PP / VPP 查询和切换每层 RouterReplay 状态； +- `get_routed_experts_dtype` — 按专家数量自动选 `uint8` / `int16`，压缩传输与显存。 + +## 4. 配置参数 + +### 4.1 如何启用 + +在需要参与 R3 的 worker 上设置 `router_replay.mode: R3`。R3 必须**同时**作用于 rollout 与训练两端，缺一不可。 + +### 4.2 参数说明 + +#### `router_replay.mode` + +- **`disable`**（默认）：关闭路由复放。 +- **`R2`**：Vanilla Routing Replay。**当前未实现**，配置该模式会抛出 `NotImplementedError`。 +- **`R3`**：Rollout Routing Replay，ROLL 当前唯一可用的模式。 + +### 4.3 完整配置示例 + +```yaml +actor_train: + router_replay: + mode: R3 + + strategy_args: + strategy_name: megatron_train + + # R3 暂未支持与 sequence_packing 组合，请保持关闭 + use_sequence_packing: False + +actor_infer: + router_replay: + mode: R3 + + strategy_args: + strategy_name: sglang # 需要 sglang >= 0.5.6.post3 + +reference: + router_replay: + mode: disable + strategy_args: + strategy_name: megatron_infer +``` + +### 4.4 使用建议 + +1. **环境与策略**：`actor_infer` 必须使用 `sglang`（≥ 0.5.6.post3）；`actor_train` 必须使用 `megatron_train`。 +2. **对称启用**：rollout 与训练两端同时配置 `mode: R3`；只在一端开启没有意义。 +3. **Reference 模型**：保持 `mode: disable`。batch 中没有 `routed_experts` 时，ROLL 自动跳过复放逻辑。 +4. **暂时关闭 sequence packing**：在 R3 + sequence_packing 支持上线前，所有相关 worker 都需保持 `use_sequence_packing: False`。 +5. **资源开销**：`routed_experts` 张量为 `[seq_len, num_layers, top_k]`，会带来额外显存与跨 worker 传输；ROLL 自动选择最小整数 dtype 以缓解。 +6. **与 IS / TIS 的关系**：Router Replay 在架构层消除路由差异，IS / TIS 在 loss 层修正概率差异，二者并不冲突，可视场景共存使用。 +7. **排查路径**：若仍出现训练-推理不一致的现象，依次检查 (a) rollout 响应是否真的带回了 `routed_experts`；(b) 训练侧 `moe_enable_routing_replay` 是否为 `True`；(c) 是否所有相关 worker 都关闭了 sequence packing。 + +启用 Router Replay（R3）后，ROLL 在 TP / PP / VPP 等并行配置下都能保证 rollout 与训练的 MoE 路由严格对齐，从模型架构层面消除一类难以靠 loss 修正解决的不一致来源。后续版本会陆续补齐 R2 与 R3 + sequence packing 等组合能力。 diff --git a/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Agentic/agent_runner.md b/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Agentic/agent_runner.md new file mode 100644 index 000000000..8763de911 --- /dev/null +++ b/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Agentic/agent_runner.md @@ -0,0 +1,556 @@ +# AgentRunner -- Agent 交互循环抽象 + +## 概述 + +ROLL 的 Agentic 训练框架支持多种类型的 agent 与环境交互：从本地 gym-like 游戏环境（Sokoban、FrozenLake）到远程沙箱中运行的 SWE-agent。在 ROLL 中，整个 agent rollout 过程是完全交由用户自定义的——框架不预设固定的交互模式，而是提供灵活的扩展点。但框架也不应该直接把 "如何编写 rollout loop" 这个难题丢给用户，因此 ROLL 基于最容易上手的多轮交互场景提供了常见的 rollout loop 实现（`TrajEnvManager` / `StepEnvManager`），业务方可以基于这些示例快速理解框架协议，再做业务自定义。 + +随着场景的丰富，原有的 `EnvManager` 体系在职责边界和扩展性上暴露出瓶颈。`AgentRunner` 抽象的引入，将 "agent 如何与环境交互" 的逻辑从 `EnvManager` 中解耦出来，使得不同类型的 agent 共享统一接口，同时保留 `EnvManager` 在训练样本构造上的核心能力。 + +AgentRunner 抽象的设计来源于对众多真实业务场景的应用实践梳理。感谢博客文章 [如何用ROLL快速上手AgenticRL](https://zhuanlan.zhihu.com/p/2023433301961049206) 中关于 Agentic 训练架构的讨论，为本设计提供了有价值的参考。 + +## 1. 架构背景：EnvManager 体系 + +在理解 AgentRunner 之前，需要先了解 ROLL 中 EnvManager 的职责和现有架构。 + +### 1.1 EnvManager 的核心职责 + +EnvManager 是 Agentic 训练的核心组件，负责三件事： + +1. **交互循环**：驱动 agent 与环境的多轮交互（`reset` / `step` / `make_decision`） +2. **消息格式化**：将环境观测构造为 LLM 可理解的 prompt（`format_messages`） +3. **训练样本构造**：将交互轨迹转化为训练所需的 `DataProto`（`formulate_rollouts`） + +### 1.2 现有继承体系 + +``` +BaseEnvManager +├── TrajEnvManager # 轨迹级拼接，适用于 gym-like 环境 +│ ├── StepEnvManager # 步级分解，适用于 GiGPO 算法 +│ │ └── StepConcatEnvManager # 步级 + 观测拼接变体 +│ ├── AgentNativeStepEnvManager # SWE/TerminalBench 原生环境 +│ └── VLTrajEnvManager # 视觉语言模型变体 +``` + +### 1.3 TrajEnvManager + +`TrajEnvManager` 是最基础的实现。它在一个紧凑的循环中完成所有工作： + +```python +# TrajEnvManager 的核心循环（简化） +def run_rollout_loop(self, data): + while self.running: + rollout_cache = self.reset() # env.reset(seed) + while not done: + lm_output = self.make_decision(cache) # format_messages → LLM 推理 + rollout_cache = self.step(lm_output) # env.step(action) + rollout = self.formulate_rollouts(cache) # 拼接完整轨迹 → DataProto + output_queue.put(rollout) +``` + +**轨迹级样本构造**：将一个 episode 内所有步的 `prompt_ids + response_ids` 拼接成一条训练序列，episode reward 放置在最后一个 token 上。整个轨迹是一个训练样本。 + +**适用场景**：Sokoban、FrozenLake、WebShop 等 gym-like 环境。环境是轻量的 Python 对象，交互速度快，无需网络通信。 + +**精细控制能力**：由于 EnvManager 完全掌控交互循环的每一步，TrajEnvManager 特别适合需要精细控制交互过程的场景： +- **快速实现 multi-agent 原型**：可以在 `run_rollout_loop` 中自由编排多个 agent 的交互顺序和通信逻辑，不受黑盒接口限制 +- **精细管理 `response_mask` 构造**：在 `formulate_rollouts` 中可以逐 token 控制哪些部分参与 loss 计算（例如只训练特定 turn 的 response、屏蔽 tool 输出等） +- **自定义 prompt 拼接策略**：在 `format_messages` 中可以灵活决定历史信息的保留、截断、摘要方式 + +### 1.4 StepEnvManager + +`StepEnvManager` 继承自 `TrajEnvManager`，重写了 `format_messages` 和 `formulate_rollouts`，将粒度从"轨迹级"降到"步级"： + +- **独立的步级 prompt**：每一步构造一个自包含的 prompt（包含系统指令、历史摘要、当前观测），而非拼接所有历史 token +- **步级训练样本**：每一步生成一个独立的训练样本，每个样本包含该步的 prompt + response 及对应的 step reward +- **`state_hash` 机制**：为每个 step 的状态计算哈希值，用于 GiGPO 算法中跨 rollout 的同状态分组优势估计 + +**适用场景**：需要步级优势估计的算法（如 GiGPO），或需要更细粒度信用分配的场景。 + +### 1.5 为什么保留 TrajEnvManager / StepEnvManager + +TrajEnvManager / StepEnvManager 体系与 ProxyEnvManager + AgentRunner 并行存在，各有优势： + +| 维度 | TrajEnvManager / StepEnvManager | ProxyEnvManager + AgentRunner | +|------|-------------------------------|-------------------------------| +| **交互开销** | 直接函数调用，零序列化开销 | HTTP 代理，有网络 + 序列化开销 | +| **控制粒度** | EnvManager 完全控制交互细节（multi-agent 编排、response_mask 精细管理等） | Agent 黑盒执行，完全面向业务逻辑，EnvManager 只关注样本构造 | +| **prompt 构造** | 深度集成 `agent_template` / `tokenizer` | Agent 自行构造 prompt，ProxyServer 透明拦截 | +| **适用环境** | 本地 gem.Env（Sokoban、FrozenLake 等），远程环境需要强行对齐现有抽象，难以理解 | 外部 agent 框架、远程沙箱、复杂 agent | +| **扩展新环境** | 需在 EnvManager 内混合实现三个职责 | 只需实现 `AgentRunner.run_job()` | +| **算法支持** | 原生支持 GiGPO 步级分解 | 通过 `MessageTracker` 支持 step/traj 两种模式 | + +**简单总结**： +- **TrajEnvManager / StepEnvManager** 面向**训练研究场景**——需要精细控制交互过程、自定义 response_mask 构造、快速验证 multi-agent 原型等需求，提供最低开销和最紧密的控制 +- **ProxyEnvManager + AgentRunner** 面向**业务逻辑场景**——agent 实现者完全只关注 agent 的行为（如何调用工具、如何与环境交互），不需要了解训练侧的任何细节 + +两者服务于不同的场景需求，并非替代关系。 + +### 1.6 选型指南 + +根据你的场景选择合适的组合： + +| 场景 | EnvManager | AgentRunner | 说明 | +|------|-----------|-------------|------| +| 需要精细控制交互过程 | `TrajEnvManager` | 不需要 | 直接函数调用，可自定义 response_mask、multi-agent 编排 | +| GiGPO 步级训练 | `StepEnvManager` | 不需要 | 内置 `state_hash` 分组，步级优势估计 | +| 本地 gym 环境（业务逻辑优先） | `ProxyEnvManager` | `GEMRunner` | agent 只关注行为，ProxyServer 透明收集轨迹 | +| 本地环境 + function-calling | `ProxyEnvManager` | `ToolCallRunner` | 支持 OpenAI tool_calls 协议 | +| 远程沙箱 + agent 回调 | `ProxyEnvManager` | `PushModeRunner` | agent 通过 ALB/ingress 回调 | +| 远程沙箱 + Roll 轮询 | `ProxyEnvManager` | `PullModeRunner` | Roll 通过 ModelService 驱动推理 | +| 自定义 agent 框架 | `ProxyEnvManager` | 自定义 Runner | 只需实现 `run_job(seed)` | + +## 2. AgentRunner 抽象设计 + +### 2.1 引入动机 + +原有架构中，当需要接入新类型的 agent 时，存在以下问题： + +1. **职责耦合**：TrajEnvManager 子类同时承担交互循环、消息格式化、训练样本构造三个职责，新增环境需要在一个类中混合实现所有逻辑 +2. **体系割裂**：ProxyEnvManager 使用 `MessageTracker` 收集轨迹，与 TrajEnvManager 的 `RolloutCache` 机制完全不同。Rock SDK 沙箱场景通过 `HarborRunner` 实现，但其接口（`submit_job(data_item, env_id)`）过于具体 +3. **agent 实现者被迫关注训练细节**：agent 需要了解 tokenizer、轨迹收集、样本构造等训练侧概念 + +AgentRunner 抽象的核心思路是：**agent 实现者完全只关注 agent 的业务行为（如何与环境交互、如何调用工具、何时终止），不关心轨迹如何被收集、训练样本如何被构造**。这使得业务团队可以用最自然的方式编写 agent 逻辑，而训练基础设施由框架透明处理。更进一步，由于 AgentRunner 仅通过标准 OpenAI 兼容接口调用 LLM，基于 AgentRunner 实现的 agent 既可以在 ROLL 训练环境中使用，也可以直接迁移到生产环境——只需将 `base_url` 从 ProxyServer 切换为实际的推理服务地址即可，训练与生产共享同一套 agent 代码。 + +### 2.2 关注点分离 + +``` +┌─────────────────────────────────────────────────────────┐ +│ ProxyEnvManager │ +│ ┌───────────────┐ ┌──────────────┐ ┌──────────────┐ │ +│ │ Episode 调度 │ │ 轨迹收集 │ │ 样本构造 │ │ +│ │ run_rollout_ │ │ MessageTracker│ │ formulate_ │ │ +│ │ loop() │ │ ProxyServer │ │ rollouts() │ │ +│ └───────┬───────┘ └──────┬───────┘ └──────────────┘ │ +│ │ │ │ +│ │ HTTP 拦截 │ │ +│ ▼ ▼ │ +│ ┌─────────────────────────────────────┐ │ +│ │ AgentRunner │ │ +│ │ ┌─────────┐ ┌──────────────────┐ │ │ +│ │ │ run_job │ │ _llm_request │ │ │ +│ │ │ (seed) │ │ (OpenAI 兼容) │ │ │ +│ │ └─────────┘ └──────────────────┘ │ │ +│ └─────────────────────────────────────┘ │ +└─────────────────────────────────────────────────────────┘ +``` + +- **AgentRunner** 负责：执行完整的 episode（加载数据、env.reset → 交互循环 → 返回结果） +- **ProxyServer** 负责：拦截 AgentRunner 的 LLM 请求，透明地完成推理和轨迹记录 +- **ProxyEnvManager** 负责：episode 调度、训练样本构造 + +### 2.3 请求拦截机制 + +AgentRunner 通过标准 OpenAI 兼容接口（`/v1/chat/completions`）向 `base_url` 发起 LLM 推理请求。`base_url` 指向本地的 ProxyServer，后者作为透明中间层： + +``` +AgentRunner ──HTTP POST──▶ ProxyServer ──拦截──▶ 记录轨迹 + │ + └──▶ 转发到实际 LLM 推理后端 + │ + ◀────────────────────────── + │ +AgentRunner ◀──HTTP Response── ProxyServer ◀─────── 返回推理结果 +``` + +1. ProxyServer 通过 `Authorization: Bearer {env_id}` 路由到对应的 EnvManager +2. EnvManager 的 `process_request` 方法完成 tokenize、推理、轨迹记录 +3. 对 AgentRunner 完全透明——AgentRunner 不知道也不关心请求被拦截 + +## 3. AgentRunner 基类 + +### 3.1 核心接口 + +```python +# roll/pipeline/agentic/agent_runner/base.py + +class EpisodeResult: + """一个 episode 的结构化结果。""" + def __init__( + self, + status: str, # "Finished" | "Failed" | "NoData" | "Timeout" + score: float, # episode 级别的 reward + step_scores: List[float] = None, # 每步 reward（可选） + agent_exit_reason: str = "", + metrics: Dict[str, Any] = None, # 额外指标（耗时、通过率等） + ): ... + + +class AgentRunner(ABC): + """Agent 交互循环的抽象基类。 + + 子类只需实现 run_job(seed) 方法：给定一个 seed，加载数据、执行完整的 episode， + 返回 EpisodeResult。 + """ + + def __init__(self, base_url: str, env_id: int, env_config: DictConfig, **kwargs): + self.base_url = base_url # ProxyServer 地址 + self.env_id = env_id # 环境实例 ID + self.env_config = env_config + + @abstractmethod + def run_job(self, seed: int) -> EpisodeResult: + """执行一个完整的 episode。""" + ... + + def setup(self) -> None: + """一次性初始化（创建 env 实例等）。""" + + def teardown(self) -> None: + """清理资源。""" + + def _llm_request(self, client, messages, tools=None, tool_choice="auto") -> Dict: + """向 base_url 发起 OpenAI 兼容的 chat completion 请求。""" +``` + +### 3.2 设计要点 + +- **`run_job(seed)` 语义清晰**：runner 的唯一职责是 "运行一个 episode"。`seed` 是唯一必需的外部输入，数据加载由 runner 内部完成 +- **`env_id` 绑定到实例**：一个 runner 实例对应一个 env slot，`_llm_request` 自动使用 `self.env_id` 进行路由 +- **`_llm_request` 是 protected 方法**：封装 HTTP 细节，子类直接调用即可 +- **不持有 tokenizer / pipeline_config**：这些是训练侧关注点，AgentRunner 只面向 "执行 episode" + +## 4. 内置 Runner 实现 + +### 4.1 GEMRunner -- 本地 gym 环境 + +`GEMRunner` 适用于通过 [GEM](https://github.com/axon-rl/gem) 接口定义的本地环境（Sokoban、FrozenLake、Math 等）。 + +```python +# roll/pipeline/agentic/agent_runner/gem_runner.py + +class GEMRunner(AgentRunner): + """本地 gem.Env 环境的 AgentRunner。 + + 交互循环: env.reset(seed) → [构造 messages → LLM request → env.step(action)] × N → EpisodeResult + """ + + def setup(self) -> None: + self.env = gem.make(env_id=env_type, **env_params) + # 可选：tool_wrapper 包装 + + def run_job(self, seed: int) -> EpisodeResult: + obs_text, info = self.env.reset(seed=seed) + messages = [{"role": "system", "content": system_template}, ...] + + for turn in range(max_steps): + resp = self._llm_request(client, messages) # → ProxyServer → 实际推理 + action = resp["choices"][0]["message"]["content"] + obs_text, reward, terminated, truncated, _ = self.env.step(action) + if terminated or truncated: + break + + return EpisodeResult(status="Finished", score=sum(rewards), step_scores=rewards) +``` + +#### ToolCallRunner + +`ToolCallRunner` 继承自 `GEMRunner`，支持 OpenAI function-calling 协议。当 LLM 返回 `tool_calls` 时，将完整的 message dict 传给环境执行： + +```python +class ToolCallRunner(GEMRunner): + """使用 OpenAI function-calling 协议的 GEM 环境 Runner。""" + + def run_job(self, seed: int) -> EpisodeResult: + messages, info = self.env.reset(seed=seed) + tools = info.get("tools") + + for _ in range(max_steps): + resp = self._llm_request(client, messages, tools=tools) + message = resp["choices"][0]["message"] + action = message if message.get("tool_calls") else message["content"] + messages, reward, terminated, truncated, _ = self.env.step(action) + ... +``` + +### 4.2 RockAgentRunner -- 远程沙箱环境 + +`RockAgentRunner` 适用于 Rock SDK 沙箱环境（SWE-bench、TerminalBench 等），agent 代码运行在远程沙箱容器中。 + +``` +RockAgentRunner (基类：job 配置、metrics 提取、数据加载) +├── PushModeRunner # Push 模式：agent 回调 Roll 的 ProxyServer +└── PullModeRunner # Pull 模式：Roll 主动轮询 sandbox +``` + +#### Push 模式 + +agent 在沙箱内通过 ALB/ingress 回调 Roll 的 ProxyServer 来获取 LLM 推理。流程： + +``` +PushModeRunner.run_job(seed) + → 加载 data_item + → 构建 JobConfig + → Job.submit() → agent 在沙箱中运行 + → agent 通过 HTTP 回调 ProxyServer 获取 LLM 推理 + → Job.wait() → 提取 metrics + → 返回 EpisodeResult +``` + +#### Pull 模式 + +Roll 通过 ModelService 主动轮询沙箱中的 agent，驱动每一步推理。流程： + +``` +PullModeRunner.run_job(seed) + → 加载 data_item + → 构建 JobConfig + ModelServiceOperator + → Job.submit() + → _inference_loop(): + while True: + request = model_service.anti_call_llm(index, response) # 获取 agent 的 LLM 请求 + response = self._llm_request(client, request) # 转发到 ProxyServer + # 重复直到 SESSION_END + → Job.wait() → 提取 metrics + → 返回 EpisodeResult +``` + +### 4.3 Runner 与 ProxyEnvManager 的协作 + +```python +# ProxyEnvManager.run_rollout_loop (简化) +def run_rollout_loop(self, data): + self.group_seed = data.meta_info['seed'] + self.env_config['group_seed'] + + while True: + self.episode_id = output_queue.get_episode_id(group_id, env_id) + if self.episode_id is None: + break + + self.message_tracker = MessageTracker(tokenizer, ...) + seed = self.group_seed + self.episode_id + + # 核心：只传 seed，整个 episode 由 AgentRunner 执行 + episode_result = self.agent_runner.run_job(seed) + + rollout = self.formulate_rollouts(episode_result.to_dict()) + output_queue.put(group_id, episode_id, rollout) +``` + +**ProxyEnvManager 不变的部分**： +- `process_request()` — HTTP handler，负责 LLM 推理 + 轨迹记录 +- `MessageTracker` — 增量 tokenize、轨迹收集、fork 检测 +- `formulate_rollouts()` — 训练样本构造（支持 `step` 和 `traj` 两种模式） +- ProxyServer 路由机制 — 通过 `env_id` 路由到对应的 handler + +关于 `MessageTracker` 的增量 tokenization、前缀聚合和分叉检测机制的详细说明，请参阅 [前缀聚合 (Prefix Aggregation)](prefix_aggregation.md)。 + +## 5. 配置与使用 + +### 5.1 配置方式 + +通过 `agent_runner_cls` 指定 Runner 的全限定类路径，与 `env_manager_cls` 模式一致，由框架通过 `safe_import_class` 动态加载： + +```yaml +# 通用配置模式 +env_manager_cls: roll.pipeline.agentic.env_manager.proxy_env_manager.ProxyEnvManager +agent_runner_cls: +``` + +### 5.2 GEMRunner 配置示例 + +适用于 Sokoban、FrozenLake 等本地 GEM 环境，使用 `ProxyEnvManager` + `GEMRunner`。 + +完整示例参见 `examples/qwen2.5-0.5B-agentic/agentic_val_sokoban_agent_runner.yaml`，以下是关键配置片段： + +```yaml +# 全局指定 EnvManager 和 AgentRunner +env_manager_cls: roll.pipeline.agentic.env_manager.proxy_env_manager.ProxyEnvManager +agent_runner_cls: roll.pipeline.agentic.agent_runner.gem_runner.GEMRunner + +# 训练参数 +rollout_batch_size: 1024 +sequence_length: 8192 +max_tokens_per_step: 64 +adv_estimator: "grpo" + +# 训练环境组配置 +train_env_manager: + max_env_num_per_worker: 16 + num_env_groups: 128 + group_size: 8 # 同一 prompt 采样 8 条轨迹 + tags: [SimpleSokoban] + num_groups_partition: [128] + +# 验证环境组配置（多种环境混合评估） +val_env_manager: + max_env_num_per_worker: 32 + num_env_groups: 1024 + group_size: 1 + tags: [SimpleSokoban, LargerSokoban, SokobanDifferentGridVocab, FrozenLake] + num_groups_partition: [256, 256, 256, 256] + +# 自定义环境：每个环境继承全局的 env_manager_cls 和 agent_runner_cls +custom_envs: + SimpleSokoban: + ${custom_env.SimpleSokoban} # 引用 traj_envs.yaml 中的环境定义 + LargerSokoban: + ${custom_env.LargerSokoban} + FrozenLake: + ${custom_env.FrozenLake} +``` + +其中环境定义（来自 `examples/config/traj_envs.yaml`）的关键字段： + +```yaml +custom_env: + SimpleSokoban: + env_type: sokoban + max_steps: ${max_actions_per_traj} # 最大交互步数 + max_tokens_per_step: ${max_tokens_per_step} + env_manager_cls: ${env_manager_cls} # 继承全局设置 + agent_runner_cls: ${agent_runner_cls} # 继承全局设置 + agent_system_template: ${agent_system_template} + agent_template: ${agent_template} # GEMRunner 用于渲染观测的模板 + env_config: # 传给 gem.make() 的环境参数 + dim_room: [6, 6] + num_boxes: 1 +``` + +如果环境使用 function-calling 协议，将 `agent_runner_cls` 替换为 `ToolCallRunner`： + +```yaml +agent_runner_cls: roll.pipeline.agentic.agent_runner.gem_runner.ToolCallRunner +``` + +### 5.3 RockAgentRunner 配置示例 + +适用于 SWE-bench 等远程沙箱环境： + +```yaml +# Push 模式 +env_manager_cls: roll.pipeline.agentic.env_manager.proxy_env_manager.ProxyEnvManager +agent_runner_cls: roll.pipeline.agentic.agent_runner.rock.push_runner.PushModeRunner + +custom_envs: + RockNativeEnv: + env_type: rocknative + max_steps: 60 + max_tokens_per_step: 4096 + env_manager_cls: ${env_manager_cls} + agent_runner_cls: ${agent_runner_cls} + env_config: + dataset_name: data/swe_bench_test.jsonl + max_iterations: 100 + tool_call_parser: qwen25 + trajectory_mode: traj # "traj" 或 "step" + reward_granularity: pass_rate # "pass_rate" 或 "binary" +``` + +```yaml +# Pull 模式 +agent_runner_cls: roll.pipeline.agentic.agent_runner.rock.pull_runner.PullModeRunner + +custom_envs: + RockNativeEnv: + ... + env_config: + model_service_port: 28080 # ModelService 监听端口 + ... +``` + +### 5.4 TrajEnvManager 直接模式 + +对于不需要 AgentRunner 的简单场景，仍然可以使用 TrajEnvManager 直接模式： + +```yaml +env_manager_cls: roll.pipeline.agentic.env_manager.traj_env_manager.TrajEnvManager +agent_runner_cls: null # 不使用 AgentRunner + +custom_envs: + SimpleSokoban: + env_type: sokoban + max_steps: 10 + env_manager_cls: ${env_manager_cls} + ... +``` + +## 6. 自定义 AgentRunner 开发 + +### 6.1 实现步骤 + +开发自定义 AgentRunner 只需三步： + +**Step 1**：继承 `AgentRunner` 基类 + +```python +from roll.pipeline.agentic.agent_runner.base import AgentRunner, EpisodeResult + +class MyCustomRunner(AgentRunner): + def __init__(self, base_url, env_id, env_config, **kwargs): + super().__init__(base_url, env_id, env_config, **kwargs) + # 自定义初始化 +``` + +**Step 2**：实现 `run_job(seed)` 方法 + +```python + def run_job(self, seed: int) -> EpisodeResult: + # 1. 加载数据（如果需要） + data = self._load_data(seed) + + # 2. 初始化环境 + obs = self.env.reset(seed=seed) + + # 3. 交互循环 + rewards = [] + with httpx.Client(timeout=3600.0) as client: + for step in range(self.max_steps): + messages = self._build_messages(obs) + resp = self._llm_request(client, messages) # 通过 ProxyServer 推理 + action = resp["choices"][0]["message"]["content"] + obs, reward, done, _, _ = self.env.step(action) + rewards.append(reward) + if done: + break + + # 4. 返回结果 + return EpisodeResult( + status="Finished", + score=sum(rewards), + step_scores=rewards, + ) +``` + +**Step 3**：在配置中指定 + +```yaml +agent_runner_cls: my_package.my_module.MyCustomRunner +``` + +### 6.2 开发约束 + +1. **使用 `_llm_request` 发起 LLM 调用**：确保请求经过 ProxyServer，使轨迹能被透明收集 +2. **不关心训练侧细节**：不持有 tokenizer、不构造训练样本、不管理轨迹数据 +3. **返回 `EpisodeResult`**：使用结构化返回值，明确 runner 的输出契约 +4. **`seed` 是唯一输入**：数据加载、环境初始化都应通过 seed 驱动，保证可复现性 + +## 7. 模块结构 + +``` +roll/pipeline/agentic/agent_runner/ +├── __init__.py # 导出 AgentRunner, EpisodeResult +├── base.py # AgentRunner 基类 + EpisodeResult +├── gem_runner.py # GEMRunner（文本模式）+ ToolCallRunner（function-calling 模式） +└── rock/ + ├── __init__.py + ├── rock_agent_runner.py # RockAgentRunner 基类（数据加载、job 配置、metrics 提取） + ├── push_runner.py # PushModeRunner（agent 回调模式） + ├── pull_runner.py # PullModeRunner + ModelServiceHarborTrial + ModelServiceOperator + └── sandbox_tool_runner.py # 沙箱 tool runner 扩展 +``` + +## 参考示例 + +- GEMRunner + Sokoban: `examples/qwen2.5-0.5B-agentic/agentic_val_sokoban_agent_runner.yaml` +- PushModeRunner + SWE-bench: `examples/agentic_rock/rock_agent_runner_4b.yaml` +- PushModeRunner + 30A3: `examples/agentic_rock/rock_agent_runner_30a3.yaml` + +## 相关文档 + +- [前缀聚合 (Prefix Aggregation)](prefix_aggregation.md) — MessageTracker 的增量 tokenization、前缀聚合和分叉检测机制 +- [Agentic 工程实践文档](agentic_engineer_practice.md) — EnvManager 开发协议、GlobalDataset 使用、轨迹合成等实践指南 + +## 参考文献 + +- [如何用ROLL快速上手AgenticRL](https://zhuanlan.zhihu.com/p/2023433301961049206) +- [GEM 环境库](https://github.com/axon-rl/gem) diff --git a/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Agentic/prefix_aggregation.md b/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Agentic/prefix_aggregation.md new file mode 100644 index 000000000..df3ec4e40 --- /dev/null +++ b/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Agentic/prefix_aggregation.md @@ -0,0 +1,241 @@ +# 前缀聚合 (Prefix Aggregation) + +## 概述 + +在 Agentic 训练场景中，一个 episode 通常包含多轮交互 (multi-turn)。默认的 step-wise 模式为每个 step 生成一个独立的训练样本，每个样本的 prompt 包含完整的历史前缀。当 episode 较长时（如 30-turn 的 SWE 任务），会导致 **样本爆炸** 和 **重复 tokenization** 问题。 + +前缀聚合方案通过 `MessageTracker` 组件，将多个 step 聚合为 **traj-wise** 的单条训练样本，同时支持增量 tokenization 和历史分叉检测。 + +### 核心能力 + +- **Traj-wise 聚合**：将整个 episode 的多轮交互聚合为一条训练样本，消除前缀重复 +- **增量 Tokenization**：通过哈希链缓存已处理的消息，避免每一步对完整历史重复 tokenize +- **分叉检测**：自动检测 agent 框架修改历史消息的情况，支持多分支训练样本生成 +- **在离线一致性**：assistant 的 token_ids 直接使用推理引擎输出，保证训练与生成一致 + +## 问题背景 + +### 问题 1: Step-wise 样本爆炸 + +默认模式下，`ProxyEnvManager.formulate_rollouts()` 为每个 step 生成一个独立的 DataProto 样本。每个样本的 prompt 包含从第一轮到当前轮的所有历史消息。30-turn episode 会产生 30 个样本，总 token 数约为 `O(n²)` 级别。 + +``` +Step 1: [sys, user] → asst1 # prompt: 2 msgs +Step 2: [sys, user, asst1, tool1] → asst2 # prompt: 4 msgs (重复 step1 的 2 msgs) +Step 3: [sys, user, asst1, tool1, asst2, tool2] → asst3 # prompt: 6 msgs (重复 step1-2 的 4 msgs) +... +Step N: [全部历史] → asstN # prompt: 2N msgs +``` + +Traj-wise 模式将上述 N 个样本聚合为一条： + +``` +[sys, user, asst1, tool1, asst2, tool2, ..., asstN] # 单条样本，无重复 +``` + +### 问题 2: 重复 Tokenization + +每次 `process_request()` 调用时，需要对完整消息历史执行 `apply_chat_template()` 进行 tokenize。随着轮次增长，重复计算量线性增加。 + +`MessageTracker` 通过哈希链缓存每条消息的 token_ids，仅对新增消息执行 tokenize。 + +### 问题 3: 历史分叉 (Fork) + +Agent 框架（如 claude-code）可能在 step 之间修改历史消息（截断、规整、总结等），导致消息历史发生分叉。 + +``` +Step 1: [sys, user] → asst1 +Step 2: [sys, user, asst1, tool1] → asst2 ← 正常延续 +Step 3: [sys, user, asst1_modified, tool1'] → asst3 ← 分叉! asst1 被修改 +Step 4: [sys, user, asst1_modified, tool1', asst3, tool2] → asst4 +``` + +此时产生两个分支的训练样本： + +- **分支 1** (dead branch): `sys → user → asst1 → tool1 → asst2` +- **分支 2** (final branch): `sys → user → asst1_modified → tool1' → asst3 → tool2 → asst4` + +### 在离线一致性分析 + +前缀聚合中，assistant 的 token_ids 直接使用推理引擎缓存的输出（`response_ids`），而非重新 tokenize，保证 assistant 部分训练与生成完全一致。 + +Left context 使用逐条消息 dummy prefix + slice 方式独立 tokenize 后拼接。与推理时 sglang 对完整消息列表做 `apply_chat_template` 的结果可能存在 tokenizer 边界差异，但由于 chat template 结构化的特性，实际影响可控。 + +## 配置参数 + +前缀聚合通过 `env_config.config` 配置，相关参数如下： + +```yaml +custom_envs: + my_env: + env_type: "my_env_type" + max_steps: 30 + max_tokens_per_step: 4096 + env_manager_cls: ProxyEnvManager + config: + trajectory_mode: traj # "step" (默认) | "traj" + enable_thinking: false # 是否启用 thinking 模式 + enable_fork: false # 是否启用分叉多分支输出 +``` + +### 参数说明 + +- `trajectory_mode`: 轨迹输出模式 + - `"step"` (默认): 每个 assistant response 产生一条训练样本，与原有逻辑兼容 + - `"traj"`: 将整个 episode 聚合为 traj-wise 样本，消除前缀重复 +- `enable_thinking`: 是否在 `apply_chat_template` 中启用 thinking 模式，影响 tokenization +- `enable_fork`: 分叉处理策略 + - `false` (默认): 仅保留最长分支，丢弃 dead branches + - `true`: 保留所有分支，每个分支输出一条训练样本 + +## 实现方案 + +### 核心组件: MessageTracker + +`MessageTracker` 位于 `roll/pipeline/agentic/env_manager/message_tracker.py`，负责单个 episode 的消息追踪、增量 tokenization 和训练数据生成。 + +#### 数据结构 + +```python +@dataclass +class MessageItem: + index: int # 消息在会话中的位置 + message: dict # 统一化后的消息 dict + token_ids: Optional[List[int]] = None # 该消息的 token IDs + logprobs: Optional[List[float]] = None # 仅模型生成的 assistant 消息有 +``` + +`MessageTracker` 内部维护以下核心数据结构： + +- `msg_item_dict: Dict[str, MessageItem]` — 哈希 → 消息项映射 +- `prev_hash_dict: Dict[str, str]` — 哈希链，记录每条消息的前驱哈希 +- `tool_calls_dict: Dict[str, dict]` — tool_call ID → 规范化 tool_call 对象 +- `step_response_hashes: List[str]` — 每个 step 的 response 哈希，用于分叉检测 + +#### 消息统一化 + +所有消息经过 pydantic `ChatCompletionMessageParam` 校验后统一化： + +1. `unify_content()` — content 归一化为 list-of-dicts 格式，合并连续 text 项 +2. `unify_tool_calls()` — 用 `tool_calls_dict` 中的规范版本替换 tool_calls +3. `get_unified_message()` — 组合以上步骤 + +#### 哈希链去重 + +每条消息通过 SHA-256 链式哈希标识：`hash(index, message_fields, pre_hash)`。 + +当新请求到达时，从头遍历消息列表，逐条计算哈希并查找 `msg_item_dict`。匹配的消息直接复用缓存的 token_ids，仅对新增消息执行 tokenize。 + +#### 增量 Tokenization + +单条消息 tokenize 使用 dummy prefix + slice 方法： + +```python +# system 消息直接 tokenize +token_ids = tokenizer.apply_chat_template([system_msg], ...) + +# 其他消息: dummy prefix + slice +full_ids = tokenizer.apply_chat_template( + [*base_chat_history, msg], # base_chat_history = [dummy_sys, dummy_user] + ... +) +token_ids = full_ids[base_offset:] # 去掉 dummy prefix 部分 +``` + +`base_offset` 是 dummy prefix 的 token 长度，依赖于 tools 定义（因为 `apply_chat_template` 会将 tool 定义注入 system 部分）。当 tools 变化时自动重新计算。 + +#### 分叉检测 + +每次 `process_messages()` 调用时，将当前匹配的消息数与上一步记录的总消息数对比： + +- 匹配数 ≥ 上一步消息数 → 正常延续 +- 匹配数 < 上一步消息数 → 历史被修改，发生分叉 + +分支识别算法：按顺序检查 `step_response_hashes`，连续 step 属于同一分支**当且仅当**前一个 step 的 response_hash 是后一个 step 哈希链的祖先。 + +#### Response Mask 规则 + +训练样本中的 `response_mask` 决定哪些 token 参与 loss 计算： + +- **模型生成的 assistant 消息** (通过 `record_response()` 记录，有 logprobs) → `response_mask = 1` +- **所有其他消息** (system, user, tool, 以及分叉后被 agent 框架修改的 assistant 消息) → `response_mask = 0` + +这确保了分叉分支中，agent 框架修改的 "继承" 消息不会被错误地算入 loss。 + +### 与 ProxyEnvManager 集成 + +#### 生命周期 + +``` +每个 episode 开始 + → 创建新的 MessageTracker 实例 + +每次 process_request() 调用 + → message_tracker.process_messages(messages, tools) # 增量 tokenize + → LLM 推理 + → message_tracker.record_response(resp_msg, response_ids, logprobs) # 记录输出 + +episode 结束 (formulate_rollouts) + → message_tracker.get_trajectory_data(mode=trajectory_mode) # 提取训练数据 + → 构建 DataProto 样本 +``` + +#### 奖励分配 + +- **Step 模式**: reward 仅分配给最后一步，前面的步骤 reward 为 0 +- **Traj 模式**: 每个分支共享 `episode_score`，reward 放在样本最后一个 token 位置 + +#### Newline Token 处理 + +在 assistant 消息后面紧邻的消息前插入一个 `\n` token，保持与 chat template 的格式一致。 + +## 两种模式对比 + +| 特性 | Step 模式 (`"step"`) | Traj 模式 (`"traj"`) | +| ------------ | -------------------------------- | ----------------------------------------- | +| 样本数量 | 每个 step 一条 | 每个 episode 一条 (或分叉时多条) | +| 前缀重复 | 有重复 (O(n²) tokens) | 无重复 (O(n) tokens) | +| Tokenization | 每步全量 `apply_chat_template` | 增量 tokenize，缓存复用 | +| 在离线一致性 | 完全一致 | assistant 部分一致，left context 微小差异 | +| 分叉支持 | 不涉及 | 支持 (`enable_fork` 控制) | +| 适用场景 | 短交互、调试 | 长交互 (SWE, 多轮 agent) | + +## 使用示例 + +### 基本配置 (Traj 模式) + +```yaml +custom_envs: + swe_env: + env_type: "swe_env" + max_steps: 50 + max_tokens_per_step: 8192 + env_manager_cls: ProxyEnvManager + config: + trajectory_mode: traj + enable_thinking: false + enable_fork: false + dataset_name: my_swe_dataset +``` + +### 启用分叉检测 + +```yaml +custom_envs: + complex_agent_env: + env_type: "complex_agent" + max_steps: 30 + max_tokens_per_step: 4096 + env_manager_cls: ProxyEnvManager + config: + trajectory_mode: traj + enable_fork: true + dataset_name: my_agent_dataset +``` + +## 注意事项 + +1. **Tools 变化**: 如果 episode 过程中 tools 定义发生变化，`base_offset` 会自动重新计算。但频繁变化的 tools 可能影响缓存命中率 +2. **序列长度**: traj 模式下整个 episode 的所有 token 聚合在一条样本中，需要确保 `sequence_length` 配置足够大 +3. **Step 模式兼容**: 设置 `trajectory_mode: step` 时行为与原有逻辑完全一致，可用于回退和调试 +4. **Assistant 消息 tokenize 警告**: 如果 `process_messages()` 中遇到未通过 `record_response()` 记录的 assistant 消息（如 agent 框架修改后的消息），会记录 warning 日志并使用 tokenizer 重新 tokenize diff --git a/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Algorithms/Reward_FL.md b/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Algorithms/Reward_FL.md deleted file mode 100644 index 48bba2ea6..000000000 --- a/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Algorithms/Reward_FL.md +++ /dev/null @@ -1,92 +0,0 @@ -# Reward Feedback Learning (Reward FL) - -## 简介 - -奖励反馈学习（Reward Feedback Learning, Reward FL） 是一种强化学习算法，用于针对特定评分器对扩散模型进行优化。Reward FL 的工作流程如下： - -1. **采样**: 对于给定的提示词（prompt）和首帧隐变量（latent），模型生成对应的视频。 -2. **奖励计算**: 根据生成视频中的人脸信息，对其进行评估并赋予相应的奖励值。 -3. **模型更新**: 模型根据生成视频所获得的奖励信号更新其参数，强化那些能够获得更高奖励的生成策略。 - - -## Reward FL 配置参数 - -在 ROLL 中，使用Reward FL算法特有的配置参数如下： (`roll.pipeline.diffusion.reward_fl.reward_fl_config.RewardFLConfig`): - -```yaml -# reward fl -learning_rate: 2e-6 -lr_scheduler_type: constant -per_device_train_batch_size: 1 -gradient_accumulation_steps: 1 -warmup_steps: 10 -num_train_epochs: 1 - -model_name: "wan2_2" - -# wan2_2 related -model_paths: ./examples/wan2.2-14B-reward_fl_ds/wan22_paths.json -reward_model_path: /data/models/antelopev2/ -tokenizer_path: /data/models/Wan-AI/Wan2.1-T2V-1.3B/google/umt5-xxl/ -model_id_with_origin_paths: null -trainable_models: dit2 -use_gradient_checkpointing_offload: true -extra_inputs: input_image -max_timestep_boundary: 1.0 -min_timestep_boundary: 0.9 -num_inference_steps: 8 -``` - -### 核心参数描述 - -- `learning_rate`: 学习率 -- `gradient_accumulation_steps`: 梯度累积步数。 -- `weight_decay`: 权重衰减大小。 -- `warmup_steps`: lr 预热步数 -- `lr_scheduler_type`: lr scheduler 类型 - -### Wan2_2 相关参数 - -Wan2_2 相关参数如下： -- `model_paths`: 模型权重路径，例如 `wan22_paths.json`，包括 high_noise_model、low_noise_model、text_encoder、vae。 -- `tokenizer_path`: Tokenizer 路径，留空将会自动下载。 -- `reward_model_path`: 奖励模型路径，例如人脸模型。 -- `max_timestep_boundary`: Timestep 区间最大值，范围为 0～1，默认为 1，仅在多 DiT 的混合模型训练中需要手动设置，例如 [Wan-AI/Wan2.2-I2V-A14B](https://modelscope.cn/models/Wan-AI/Wan2.2-I2V-A14B)。[Wan-AI/Wan2.2-I2V-A14B](https://modelscope.cn/models/Wan-AI/Wan2.2-I2V-A14B). -- `min_timestep_boundary`: Timestep 区间最小值，范围为 0～1，默认为 1，仅在多 DiT 的混合模型训练中需要手动设置，例如 [Wan-AI/Wan2.2-I2V-A14B](https://modelscope.cn/models/Wan-AI/Wan2.2-I2V-A14B)。 -- `model_id_with_origin_paths`: 带原始路径的模型 ID，例如 Qwen/Qwen-Image:transformer/diffusion_pytorch_model*.safetensors。用逗号分隔。 -- `trainable_models`: 可训练的模型，例如 dit、vae、text_encoder。 -- `extra_inputs`: 额外的模型输入，以逗号分隔。 -- `use_gradient_checkpointing_offload`: 是否将 gradient checkpointing 卸载到内存中 -- `num_inference_steps`: 推理步数，默认值为 8 (蒸馏 wan2_2 模型) - - -## 注意事项 -- 奖励模型分数是基于人脸信息，因此请确保视频的第一帧包含人脸。 -- 将人脸模型相关 onnx 文件下载到 `reward_model_path` 目录. -- 下载官方 Wan2.2 pipeline 和 蒸馏 Wan2.2 safetensors, 并放在 `model_paths` 目录，例如 `wan22_paths.json` 文件。 -- 根据 data/example_video_dataset/metadata.csv 文件，将你的视频数据集适配到对应的格式 - -## 模型引用 -- `官方 Wan2.2 pipeline`: [Wan-AI/Wan2.2-I2V-A14B](https://modelscope.cn/models/Wan-AI/Wan2.2-I2V-A14B) -- `蒸馏 Wan2.2 模型参数`: [lightx2v/Wan2.2-Lightning](https://huggingface.co/lightx2v/Wan2.2-Lightning/tree/main) -- `奖励模型`: [deepinsight/insightface](https://github.com/deepinsight/insightface/tree/master/model_zoo) - -## 权重预处理 -- 运行`merge_model.py`来分别合并`Official Wan2.2 pipeline` 高噪模型和低噪模型的多个文件为一个 -- 运行`merge_lora.py`来合并`Distilled Wan2.2 DiT safetensors`蒸馏加速lora分别到`Official Wan2.2 pipeline`高噪模型和低噪模型 - -## 环境配置 -``` -pip install -r requirements_torch260_diffsynth.txt -``` - -## 参考示例 - -可以参考以下配置文件来设置 Reward FL 训练： - -- `./examples/docs_examples/example_reward_fl.yaml` - -运行`run_reward_fl_ds_pipeline.sh`快速开始 - -## 参考文献 -[1]: Identity-Preserving Image-to-Video Generation via Reward-Guided Optimization. https://arxiv.org/abs/2510.14255 \ No newline at end of file diff --git a/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Configuration/config_guide.md b/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Configuration/config_guide.md index db8de882a..ac3d3a968 100644 --- a/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Configuration/config_guide.md +++ b/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Configuration/config_guide.md @@ -166,7 +166,7 @@ actor_train: file_name: xxx/train.json prompt: instruction strategy_args: - strategy_name: megatron_train # 训练策略：deepspeed_train 或 megatron_train + strategy_name: megatron_train # 训练策略：fsdp2_train 或 megatron_train strategy_config: tensor_model_parallel_size: 1 pipeline_model_parallel_size: 1 @@ -207,7 +207,7 @@ reference: ### 模型参数 (**model_args**) - `model_args.dtype`: 设置模型的数据类型，可以是 fp32 (单精度浮点数), bf16 (BFloat16), 或 fp16 (半精度浮点数)。如果不设置，则使用配置的 torch_dtype。选择合适的数据类型可以平衡计算速度、内存消耗和精度。 -- `model_args.disable_gradient_checkpointing`: 禁用梯度检查点。仅当 `actor_train` 的 `strategy_name` 为 `deepspeed_train` 时适用。梯度检查点是一种内存优化技术，通过在反向传播时重新计算部分激活值来减少显存占用。禁用它会增加内存消耗但可能略微加速计算。 +- `model_args.disable_gradient_checkpointing`: 禁用梯度检查点。仅当 `actor_train` 的 `strategy_name` 为 `fsdp2_train` 时适用。梯度检查点是一种内存优化技术，通过在反向传播时重新计算部分激活值来减少显存占用。禁用它会增加内存消耗但可能略微加速计算。 ### 数据参数 (**data_args**) @@ -228,7 +228,7 @@ reference: ### 策略参数 (**strategy_args**) - `strategy_args.strategy_name`: 训练/推理策略的名称。 - - 用于训练的策略有：`deepspeed_train`（使用 DeepSpeed）或 `megatron_train`（使用 Megatron-LM）。 + - 用于训练的策略有：`fsdp2_train`（使用 FSDP2）或 `megatron_train`（使用 Megatron-LM）。 - 用于推理的策略有：`vllm`、`sglang` 或 `hf_infer`（使用 Hugging Face 的默认推理）。 - `strategy_args.strategy_config`: 训练/推理策略的详细配置，它将作为参数传递给 `strategy_name`对应的构造函数。例如，`strategy_config.tensor_model_parallel_size` 用于 `megatron_train` 策略，`strategy_config.gpu_memory_utilization` 用于 `vllm` 策略。 @@ -263,23 +263,21 @@ reference: - `mem_fraction_static`: 用于模型权重和 KV 缓存等静态内存的 GPU 内存占比。 如果 KV 缓存构建失败，请增加此值；如果 CUDA 内存不足，请减小此值。 - `load_format`: 加载模型权重的格式。（同 VLLM，可设为 dummy） -#### DeepSpeed 策略配置 +#### FSDP2 策略配置 -在`./examples/config/` 中有 DeepSpeed 配置文件，可以在默认列表中重写以进行策略配置。 -例如，要使用 deepspeed_zero2 策略，请将以下内容添加到您的配置中： +- `fsdp_size`：FSDP 分片数量 + - 如果 `fsdp_size >= world_size` 或 `fsdp_size <= 1`：纯 FSDP2 模式 + - 如果 `fsdp_size < world_size`：带有 DDP 副本的 HSDP 模式 +- `param_dtype`：参数数据类型（例如 `bf16`、`fp16`、`float32`） +- `reduce_dtype`：梯度归约的数据类型（例如 `float32`） +- `reshard_after_forward`：是否在前向传播后重新分片参数 + - `true`：前向传播后重新分片 + - `false`：保持参数gathered +- `offload_policy`：是否启用 CPU 卸载 + - `true`：在不使用时将参数卸载到 CPU（节省 GPU 内存） + - `false`：将所有参数保留在 GPU 上（更快但使用更多内存） -```yaml -defaults: - - ../config/envs@_here_ - - ../config/deepspeed_zero@_here_ - - ../config/deepspeed_zero2@_here_ # 引入 deepspeed_zero2 策略配置 - - ../config/deepspeed_zero3@_here_ - - ../config/deepspeed_zero3_cpuoffload@_here_ -actor_train: - strategy_args: - strategy_name: deepspeed_train - strategy_config: ${deepspeed_zero2} -``` +此外，使用 FSDP2 策略时可以通过在 `model_args` 配置 `ulysses_size` 设置上下文并行。 ### 训练参数 (**training_args**) @@ -288,7 +286,7 @@ actor_train: - `training_args.per_device_train_batch_size`: 在每个设备上进行训练时使用的批次大小。 - `training_args.gradient_accumulation_steps`: 梯度累积的步数。 -在 DeepSpeed 训练中，全局训练批次大小是`per_device_train_batch_size` \* `gradient_accumulation_steps` \* world_size (即`actor_train`/`critic`的`device_mapping`长度)。 +在 FSDP2 训练中，全局训练批次大小是`per_device_train_batch_size` \* `gradient_accumulation_steps` \* world_size / `ulysses_size` (其中 world_size 即`actor_train`/`critic`的`device_mapping`长度)。 在 Megatron 训练中，全局训练批次大小是`per_device_train_batch_size` \* `gradient_accumulation_steps` \* world_size / `tensor_model_parallel_size` / `pipeline_model_parallel_size` / `context_parallel_size` (不需要除以`expert_model_parallel_size`)。 diff --git a/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Configuration/config_system.md b/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Configuration/config_system.md index 2c758c211..4244e0b8d 100644 --- a/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Configuration/config_system.md +++ b/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Configuration/config_system.md @@ -39,7 +39,7 @@ ROLL 的配置系统主要由以下几个核心组件构成： ### 4. Strategy - 策略配置 策略配置定义了每个工作节点使用的训练/推理策略，包括： -- 策略名称（如 `megatron_train`, `vllm`, `sglang`, `deepspeed_train`等） +- 策略名称（如 `megatron_train`, `vllm`, `sglang`, `fsdp2_train`等） - 策略特定参数（如张量并行大小、流水线并行大小等） ### 5. Arguments 类 diff --git a/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Configuration/deepspeed.md b/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Configuration/deepspeed.md deleted file mode 100644 index e7304e675..000000000 --- a/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Configuration/deepspeed.md +++ /dev/null @@ -1,104 +0,0 @@ -# DeepSpeed 训练后端配置指南 - -DeepSpeed 是微软开发的高效深度学习优化库，提供了内存优化、分布式训练和性能优化等功能。本文档将详细介绍如何在 ROLL 框架中配置和使用 DeepSpeed 训练后端。 - -## DeepSpeed 简介 - -DeepSpeed 提供了多种优化技术，包括： -1. **ZeRO 优化**：通过分区优化器状态、梯度和参数来减少内存使用 -2. **内存高效训练**：支持大规模模型的训练 -3. **高性能通信**：优化分布式训练中的通信效率 -4. **灵活的配置**：支持多种优化级别的配置 - -## 配置 DeepSpeed 策略 - -在 ROLL 框架中，可以通过在 YAML 配置文件中设置 `strategy_args` 来配置 DeepSpeed 训练策略。 - -### 配置示例 - -以下是一个典型的 DeepSpeed 配置示例（来自 `examples/qwen2.5-7B-rlvr_megatron/rlvl_lora_zero3.yaml`）： - -```yaml -defaults: - - ../config/deepspeed_zero3@_here_ - - ../config/deepspeed_zero3_cpuoffload@_here_ - -actor_train: - model_args: - attn_implementation: fa2 - disable_gradient_checkpointing: true - dtype: bf16 - model_type: ~ - training_args: - learning_rate: 1.0e-5 - weight_decay: 0 - per_device_train_batch_size: 1 - gradient_accumulation_steps: 32 - warmup_steps: 20 - num_train_epochs: 50 - strategy_args: - strategy_name: deepspeed_train - strategy_config: ${deepspeed_zero3} - device_mapping: list(range(0,16)) - infer_batch_size: 4 -``` - -### 配置参数详解 - -1. **strategy_name**: 设置为 `deepspeed_train` 以使用 DeepSpeed 训练后端 - -2. **strategy_config**: DeepSpeed 特定的配置参数 - - 可以引用预定义的配置文件，如 `${deepspeed_zero3}` - - 在 `./examples/config/` 目录中有多种 DeepSpeed 配置文件可供选择： - - `deepspeed_zero3.yaml`: ZeRO-3 配置 - - `deepspeed_zero3_cpuoffload.yaml`: 带 CPU 卸载的 ZeRO-3 配置 - -3. **defaults 部分**: 引入预定义的 DeepSpeed 配置 - ```yaml - defaults: - - ../config/deepspeed_zero@_here_ - - ../config/deepspeed_zero2@_here_ - - ../config/deepspeed_zero3@_here_ - - ../config/deepspeed_zero3_cpuoffload@_here_ - ``` - -4. **device_mapping**: 指定使用的 GPU 设备 ID 列表 - -## DeepSpeed 配置文件 - -在 `./examples/config/` 目录中提供了多种预定义的 DeepSpeed 配置文件： - -1. **deepspeed_zero3.yaml**: ZeRO-3 配置，优化器状态、梯度和参数分区 -2. **deepspeed_zero3_cpuoffload.yaml**: 带 CPU 卸载的 ZeRO-3 配置 - -### 使用预定义配置 - -要使用预定义的 DeepSpeed 配置，可以在 YAML 文件中这样引用： - -```yaml -defaults: - - ../config/deepspeed_zero3@_here_ - -actor_train: - strategy_args: - strategy_name: deepspeed_train - strategy_config: ${deepspeed_zero3} -``` - -## 与其他组件的集成 - -在配置示例中，我们可以看到： - -1. `actor_train` 使用 DeepSpeed 进行训练 -2. `actor_infer` 可能使用其他推理后端（如 vLLM） -3. `reference` 使用 Hugging Face 推理后端 -4. 奖励模型使用不同的推理后端 - -这种设计允许不同组件根据其需求选择最适合的后端。 - - -## 注意事项 - -1. DeepSpeed 需要特定版本的依赖库，请确保安装了兼容的版本 -2. 不同的 ZeRO 级别有不同的内存和性能特征，需要根据具体需求选择 -3. 在使用 LoRA 微调时，需要注意与 DeepSpeed 的兼容性 diff --git a/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Configuration/lora.md b/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Configuration/lora.md index 2f4d82f7a..bc71a374b 100644 --- a/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Configuration/lora.md +++ b/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Configuration/lora.md @@ -15,7 +15,7 @@ LoRA 通过以下方式实现参数高效微调： ### 配置示例 -以下是一个典型的 LoRA 配置示例（来自 `examples/qwen2.5-7B-rlvr_megatron/rlvl_lora_zero3.yaml`）： +以下是一个典型的 LoRA 配置示例（来自 `examples/qwen2.5-7B-rlvr_megatron/rlvr_lora_fsdp2.yaml`）： ```yaml # LoRA 全局配置 @@ -40,8 +40,13 @@ actor_train: warmup_steps: 20 num_train_epochs: 50 strategy_args: - strategy_name: deepspeed_train - strategy_config: ${deepspeed_zero3} + strategy_name: fsdp2_train + strategy_config: + fsdp_size: 16 + param_dtype: bf16 + reduce_dtype: float32 + reshard_after_forward: true + offload_policy: false device_mapping: list(range(0,16)) infer_batch_size: 4 @@ -93,15 +98,16 @@ actor_infer: ## LoRA 与训练后端的兼容性 -目前，LoRA 微调仅支持 DeepSpeed 训练后端： +目前，LoRA 微调支持 FSDP2 和 Megatron 训练后端，两者在 LoRA 相关配置上的差异进存在于 `lora_target` 中，如下： -```yaml -actor_train: - strategy_args: - strategy_name: deepspeed_train # LoRA 仅支持 deepspeed_train -``` + - FSDP2 可以使用 `o_proj,q_proj,k_proj,v_proj` 作为 `lora_target`. 这是因为 FSDP2 提供了与 LoRA 良好集成的优化功能. + - Megatron 可以使用 `all-linear`，这包括了 `linear_qkv,linear_proj,linear_fc1,linear_fc2`，其中: + - `linear_qkv` 表示 HF 中的 `q_proj,k_proj,v_proj` + - `linear_proj` 表示 HF 中的 `o_proj` + - `linear_fc1` 表示 HF 中的 `gate_proj,up_proj` + - `linear_fc2` 表示 HF 中的 `down_proj` -这是因为 DeepSpeed 提供了与 LoRA 良好集成的优化功能。 +此外，使用 Megatron 训练后端时，LoRA 微调不支持 vision model weights (vit), vpp 和 tp+ep（[文档](https://alibaba.github.io/ROLL/zh-Hans/docs/User%20Guides/Configuration/megatron)） ## 性能优化建议 @@ -119,7 +125,7 @@ actor_train: ## 注意事项 -1. LoRA 微调目前仅支持 DeepSpeed 训练后端 +1. LoRA 微调目前支持 FSDP2 和 Megatron 训练后端 2. 需要确保模型支持 LoRA 微调 3. 在使用梯度检查点时需要注意与 LoRA 的兼容性 4. LoRA 微调的性能可能与全参数微调有所不同，需要根据具体任务进行评估 diff --git a/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Hardware Support/ascend_usage.md b/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Hardware Support/ascend_usage.md index 39f7d110e..54d45fb81 100644 --- a/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Hardware Support/ascend_usage.md +++ b/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Hardware Support/ascend_usage.md @@ -1,25 +1,27 @@ # ROLL x Ascend -最后更新：2026/05/14。 +Last updated: 11/25/2025. 我们在 ROLL 上增加对华为昇腾设备的支持。 ## 硬件支持 -Atlas 900 A2 PODc 和 Atlas 900 A3 PODc +Atlas 900 A2 PODc ## 安装 + ### 基础环境准备 | software | version | |-----------|-------------| | Python | 3.11 | -| CANN | 8.5.1 | +| CANN | 8.3.RC1 | ### 创建 conda 环境 + 使用以下命令在 Miniconda 中创建新的 conda 环境： ``` @@ -27,26 +29,27 @@ conda create --name roll python=3.11 conda activate roll ``` -### 安装 torch & torch_npu +### 安装 torch & torch_npu: + -为了能在 ROLL 中正常使用 torch 和 torch_npu，需使用以下命令安装 torch 和 torch_npu： +为了能在 ROLL 中正常使用 torch 和 torch_npu，需使用以下命令安装 torch 和 torch_npu。 ``` -# 在预构建镜像外手动安装时，使用 CPU 版 torch -pip install torch==2.8.0 torchvision==0.23.0 torchaudio==2.8.0 --index-url https://download.pytorch.org/whl/cpu +# 安装 torch 的 CPU 版本 +pip install torch==2.7.1 torchvision==0.22.1 torchaudio==2.7.1 --index-url https://download.pytorch.org/whl/cpu -# 安装与 torch/CANN 匹配的 torch_npu -pip install torch_npu==2.8.0 +# 安装 torch_npu +pip install torch_npu==2.7.1 ``` -### 安装 vllm & vllm-ascend -为了能够在 ROLL 中正常使用 vllm，需使用以下命令编译安装 vllm 和 vllm-ascend： +### 安装vllm & vllm-ascend: + +为了能够在 ROLL 中正常使用 vllm，需使用以下命令编译安装 vllm 和 vllm-ascend。 ``` # vllm -git clone -b v0.13.0 --depth 1 https://github.com/vllm-project/vllm.git -git clone -b v0.13.0 --depth 1 https://github.com/vllm-project/vllm.git +git clone -b v0.11.0 --depth 1 https://github.com/vllm-project/vllm.git cd vllm pip install -r requirements/build.txt @@ -54,8 +57,7 @@ VLLM_TARGET_DEVICE=empty pip install -v -e . cd .. # vllm-ascend -git clone -b v0.13.0 --depth 1 https://github.com/vllm-project/vllm-ascend.git -git clone -b v0.13.0 --depth 1 https://github.com/vllm-project/vllm-ascend.git +git clone -b v0.11.0rc1 --depth 1 https://github.com/vllm-project/vllm-ascend.git cd vllm-ascend pip install -e . @@ -63,13 +65,12 @@ cd .. ``` 或者可以从预编译的 wheel 包安装 `vllm` 和 `vllm-ascend`： - ``` -# 安装 vllm-project/vllm，最新支持版本为 v0.13.0 -pip install vllm==0.13.0 +# Install vllm-project/vllm. The newest supported version is v0.11.0. +pip install vllm==0.11.0 -# 从 pypi 安装 vllm-project/vllm-ascend -pip install vllm-ascend==0.13.0 +# Install vllm-project/vllm-ascend from pypi. +pip install vllm-ascend==0.11.0rc1 ``` ### 安装 ROLL @@ -78,60 +79,59 @@ pip install vllm-ascend==0.13.0 git clone https://github.com/alibaba/ROLL.git cd ROLL pip install -r requirements_common.txt -pip install deepspeed==0.16.4 cd .. ``` ### 其他三方库说明 -| 软件 | 说明 | -| ---- | ---- | -| transformers | >= v4.57.6 | -| flash_attn | 不支持 | -| transformer-engine[pytorch] | 不支持 | +| software | description | +|-------------------------------|---------------| +| transformers | >= v4.57.1 | +| flash_attn | not supported | +| transformer-engine[pytorch] | not supported | -1. `transformers` v4.57.6 支持启用 `--flash_attention_2`。 -2. 目前不支持 `flash_attn` 加速。 -3. 目前不支持 `transformer-engine[pytorch]`。 +1. 支持通过 transformers 使能 --flash_attention_2， transformers 需大于等于 4.57.1 版本。 +2. 不支持通过 flash_attn 使能 flash attention 加速。 +3. 暂不支持 transformer-engine[pytorch] ``` -pip install transformers==4.57.6 +pip install transformers==4.57.1 ``` -## 快速开始：单节点部署指引 +## 快速开始，单节点部署指引 正式使用前，建议您通过对单节点流水线的训练尝试以检验环境准备和安装的正确性。 -由于目前暂不支持 Megatron-LM 训练，请首先将对应文件中 `strategy_args` 参数修改为 `deepspeed` 选项。 +由于目前暂不支持 Megatron-LM 训练，请首先将对应文件中 +strategy_args 参数修改为 fsdp2 选项。 -**注意：** 目前 NPU 上不支持 colocated 模式。你需要修改 `device_mapping`，确保训练和推理在不同的卡上执行。 - -1. 使用 shell 执行单节点流水线： +1. 使用 shell 执行单节点流水线 ``` bash examples/agentic_demo/run_agentic_pipeline_frozen_lake_single_node_demo.sh ``` -2. 使用配置文件执行 agentic pipeline： +2. 使用配置文件执行 agentic pipeline ``` -# 确保当前位于 ROLL 项目目录的根目录下 +# 确保当前位于ROLL项目目录的根目录下 python examples/start_agentic_pipeline.py \ --config_path qwen2.5-0.5B-agentic \ --config_name agentic_val_sokoban -``` -- `--config_path` – 包含您的 YAML 配置文件的目录。 -- `--config_name` – 文件名（不含 `.yaml` 后缀）。 +- ``--config_path`` – 包含您的YAML配置文件的目录。 +- ``--config_name`` – 文件名（不含.yaml后缀）。 +``` ## 支持现状 -| 功能 | 示例 | 训练后端 | 推理后端 | 硬件 | -| ---- | ---- | -------- | -------- | ---- | -| Agentic | examples/qwen2.5-0.5B-agentic/run_agentic_pipeline_sokoban.sh | DeepSpeed | vLLM | Atlas 900 A3 PODc | -| Agentic-Rollout | examples/qwen2.5-0.5B-agentic/run_agentic_rollout_sokoban.sh | DeepSpeed | vLLM | Atlas 900 A3 PODc | -| RLVR | examples/ascend_examples/run_rlvr_pipeline.sh | DeepSpeed | vLLM | Atlas 900 A2/A3 PODc | +| Feature | Example | Training Backend | Inference Backend | Hardware | +| --------------- | ------------------------------------------------------------ | ---------------- | ----------------- | ----------------- | +| Agentic | examples/qwen2.5-0.5B-agentic/run_agentic_pipeline_sokoban.sh | FSDP2 | vLLM | Atlas 900 A2 PODc | +| Agentic-Rollout | examples/qwen2.5-0.5B-agentic/run_agentic_rollout_sokoban.sh | FSDP2 | vLLM | Atlas 900 A2 PODc | +| DPO | examples/qwen2.5-3B-dpo_megatron/run_dpo_pipeline.sh | FSDP2 | vLLM | Atlas 900 A2 PODc | +| RLVR | examples/qwen2.5-7B-rlvr_megatron/run_rlvr_pipeline.sh | FSDP2 | vLLM | Atlas 900 A2 PODc | -## 声明 -ROLL 中提供的 Ascend 支持代码皆为参考样例，生产环境使用请通过官方正式途径沟通。 +## 声明 +ROLL 中提供的 Ascend 支持代码皆为参考样例，生产环境使用请通过官方正式途径沟通，谢谢。 diff --git a/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Pipeline/agentic_pipeline_start.md b/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Pipeline/agentic_pipeline_start.md index aef64cc18..d213acee4 100644 --- a/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Pipeline/agentic_pipeline_start.md +++ b/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Pipeline/agentic_pipeline_start.md @@ -31,7 +31,7 @@ Agentic Pipeline 是ROLL提供的智能体训练核心Pipeline，支持多种算 * **环境粒度的异步并行rollout**: 各环境独立采样轨迹，提高采样效率 * **异步训练**: rollout/training解耦，支持异步训练 * **多轮交互支持本地调试**: 多轮交互rollout支持本地调试，提高多轮交互业务开发效率 -* **灵活的策略配置**：支持多种分布式训练策略，如 Megatron、DeepSpeed、vLLM 等，可以根据硬件资源进行灵活配置。 +* **灵活的策略配置**：支持多种分布式训练策略，如 Megatron、FSDP2、vLLM 等，可以根据硬件资源进行灵活配置。 * **高效训练优化**：支持 **Sequence Packing**（将多条短样本拼接成连续序列，减少 padding）与 **Dynamic Batching**（根据样本长度动态组 batch，按 batch 内最大长度统一 padding，最小化无效计算）。配置方法和实现原理详见`sequence packing`和`dynamic batching` 对应文档。 --- diff --git a/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Pipeline/distill_pipeline_start.md b/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Pipeline/distill_pipeline_start.md index 7d0a3f726..4e447aa17 100644 --- a/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Pipeline/distill_pipeline_start.md +++ b/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Pipeline/distill_pipeline_start.md @@ -101,7 +101,7 @@ * `max_grad_norm`：梯度裁剪阈值 * ... * **分布式策略**（`strategy_args`） - * `strategy_name`：使用的分布式策略（如 `megatron_train`、`deepspeed_infer`） + * `strategy_name`：使用的分布式策略（如 `megatron_train`、`megatron_infer`） * 策略特定参数：如 `tp_size`（张量并行规模）、`pp_size`（Pipeline并行规模） * `gpu_memory_utilization`：GPU 内存利用率（特定于 vLLM） * **设备映射**（`device_mapping`） diff --git a/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Pipeline/dpo_pipeline_start.md b/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Pipeline/dpo_pipeline_start.md index f729e88dc..7b66909fb 100644 --- a/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Pipeline/dpo_pipeline_start.md +++ b/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Pipeline/dpo_pipeline_start.md @@ -90,7 +90,7 @@ * `max_grad_norm`：梯度裁剪阈值 * ... * **分布式策略**（`strategy_args`） - * `strategy_name`：使用的分布式策略（如 `megatron_train`、`deepspeed_infer`） + * `strategy_name`：使用的分布式策略（如 `megatron_train`、`megatron_infer`） * 策略特定参数：如 `tp_size`（张量并行规模）、`pp_size`（Pipeline并行规模） * `gpu_memory_utilization`：GPU 内存利用率（特定于 vLLM） * **设备映射**（`device_mapping`） diff --git a/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Pipeline/on_policy_distill_pipeline_start.md b/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Pipeline/on_policy_distill_pipeline_start.md index 739237aa8..62677a0ed 100644 --- a/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Pipeline/on_policy_distill_pipeline_start.md +++ b/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Pipeline/on_policy_distill_pipeline_start.md @@ -27,6 +27,9 @@ - [步骤3：启动流水线](#步骤3启动流水线) - [步骤4：监控](#步骤4监控) - [步骤5：输出和结果](#步骤5输出和结果) + - [Multi-Teacher OPD](#️multi-teacher-opd多教师蒸馏) + - [配置示例](#配置示例) + - [核心机制](#核心机制) - [常见问题](#️常见问题) - [参考资料](#参考资料) @@ -245,9 +248,9 @@ On-Policy Distillation 的 Worker 角色根据模式有所不同： |----------|----------|------| | `student_train` | `actor_train` | 训练学生模型，使用 Teacher KL 计算损失 | | `student_infer` | `actor_infer` | 生成轨迹，计算 student log_probs | -| `teacher` | `reference` | 计算 teacher log_probs | +| `teacher` | `reference` / `references` | 计算 teacher log_probs（支持单个 WorkerConfig 或多 teacher Dict） | -**注意**：配置文件中使用 `student_train`、`student_infer`、`teacher` 名称，系统会自动映射。 +**注意**：配置文件中使用 `student_train`、`student_infer`、`teacher` 名称，系统会自动映射。多 teacher 时 `teacher` 为 `Dict[str, WorkerConfig]`，内部归一化为 `self.references: Dict[str, Cluster]`。 #### 混合模式 @@ -334,9 +337,17 @@ bash examples/qwen3-8B-onpolicy-distill-megatron/run_onpolicy_distill_pipeline.s | 参数 | 说明 | 默认值 | |------|------|--------| | `use_opd` | 启用混合模式 OPD（将 Teacher KL 添加到奖励中） | `false` | -| `opd_kl_coef` | OPD KL 系数，控制蒸馏信号相对于外部奖励的权重 | `1.0` | | `teacher` | 教师模型配置（自动映射到 reference） | 必须配置 | +#### Multi-Teacher 模式参数 + +| 参数 | 说明 | 默认值 | +|------|------|--------| +| `teacher` | Dict[str, WorkerConfig] 多教师配置 | — | +| `teacher.{name}.opd_kl_coef` | 该教师的 KL 系数 | `1.0` | +| `teacher.{name}.tag_included` | 该教师负责的 tag 列表，空表示全量 | `[]` | +| `tag_to_template` | 按 tag 选择不同的 chat template | `{}` | + --- @@ -391,6 +402,172 @@ python examples/start_onpolicy_distill_pipeline.py \ --- +## ✨️Multi-Teacher OPD（多教师蒸馏） + +### 概述 + +Multi-Teacher OPD 允许多个专长不同的教师模型同时指导一个学生模型，按数据领域（tag）将数据路由到对应的教师，避免无效计算并实现更精准的蒸馏。 + +``` +┌──────────────────────────────────────────────────────────────────┐ +│ Multi-Teacher OPD 数据流 │ +├──────────────────────────────────────────────────────────────────┤ +│ │ +│ Student Infer rollout → batch (含 tag/domain 字段) │ +│ │ │ +│ ├── [math_dapo 数据] ──▶ Teacher-32B (数学专长) │ +│ │ 计算 ref_log_probs_teacher_32B │ +│ │ │ +│ └── [KodCode 数据] ──▶ Teacher-14B (代码专长) │ +│ 计算 ref_log_probs_teacher_14B │ +│ │ │ +│ ▼ │ +│ Compute Advantage: │ +│ 对每条数据，只累加被路由到的 teacher 的 KL: │ +│ advantage = -Σ(opd_kl_coef_i * KL_i) (仅路由到的 teacher) │ +│ │ +└──────────────────────────────────────────────────────────────────┘ +``` + +### 配置示例 + +#### 多教师纯 OPD 模式 + +```yaml +is_pure_opd: true +global_template: qwen3 + +# 按 tag 选择不同的 chat template（可选） +tag_to_template: + math_dapo: qwen3 # 数学数据使用 qwen3 template（带 thinking） + KodCode: qwen3_nothink # 代码数据使用 qwen3_nothink template + +student_train: + model_args: + model_name_or_path: Qwen/Qwen3-8B + data_args: + file_name: + - data/dapo_math_17k_simple_boxed.jsonl + - data/code_KodCode_data.jsonl + domain_interleave_probs: + math_rule: 0.6 + code_rule: 0.4 + device_mapping: list(range(0,8)) + # ... + +student_infer: + model_args: + model_name_or_path: Qwen/Qwen3-8B + device_mapping: list(range(0,8)) + # ... + +# teacher 配置为 Dict[str, WorkerConfig] +teacher: + teacher_32B: + model_args: + model_name_or_path: Qwen/Qwen3-32B # 数学专长教师 + opd_kl_coef: 1.0 + tag_included: [math_dapo] # 只处理数学数据 + device_mapping: list(range(8,16)) + strategy_args: + strategy_name: megatron_infer + strategy_config: + tensor_model_parallel_size: 2 + pipeline_model_parallel_size: 4 + + teacher_14B: + model_args: + model_name_or_path: Qwen/Qwen3-14B # 代码专长教师 + opd_kl_coef: 1.0 + tag_included: [KodCode] # 只处理代码数据 + device_mapping: list(range(16,24)) + strategy_args: + strategy_name: megatron_infer + strategy_config: + tensor_model_parallel_size: 2 + pipeline_model_parallel_size: 2 + +rewards: + math_rule: + worker_cls: roll.pipeline.rlvr.rewards.math_rule_reward_worker.MathRuleRewardWorker + tag_included: [math_dapo] + code_rule: + worker_cls: roll.pipeline.rlvr.rewards.code_sandbox_reward_worker.CodeSandboxRewardWorker + tag_included: [KodCode] +``` + +#### 混合路由配置（通用教师 + 专长教师） + +```yaml +teacher: + teacher_general: + model_args: + model_name_or_path: Qwen/Qwen3-72B + opd_kl_coef: 0.3 + tag_included: [] # 空 = 负责所有 tag（通用教师） + + teacher_math_specialist: + model_args: + model_name_or_path: DeepSeek-Math-67B + opd_kl_coef: 0.7 + tag_included: [math_dapo, aime] # 仅负责数学 +``` + +此配置下，数学数据会同时被 `teacher_general`（KL 系数 0.3）和 `teacher_math_specialist`（KL 系数 0.7）计算 KL，两者的加权 KL 都参与 advantage。非数学数据只有 `teacher_general` 参与。 + +### 核心机制 + +#### 1. Tag 路由 + +每条训练数据带有 `tag` 字段（如 `math_dapo`、`KodCode`）。每个 teacher 通过 `tag_included` 声明自己负责的 tag： + +- `tag_included: [math_dapo]` — 只处理 tag 为 `math_dapo` 的数据 +- `tag_included: []`（空列表）— 处理所有数据（通用教师） + +路由发生在 ref_log_probs 计算阶段（pipeline 层），teacher 只对被路由到的数据做 forward，避免无效推理开销。 + +#### 2. Per-Teacher KL 系数 + +每个 teacher 有独立的 `opd_kl_coef`，用于控制该教师蒸馏信号的权重： + +``` +advantage = -Σ(opd_kl_coef_i * KL(student || teacher_i)) +``` + +只有被路由到的 teacher 参与该样本的 KL 累加。 + +#### 3. 并行推理优化 + +当多个 teacher 使用不同的 GPU（`device_mapping` 不重叠）时，系统会自动使用多线程并行执行各 teacher 的 forward pass，减少总推理时间。 + +#### 4. tag_to_template + +不同领域的数据可能需要不同的 chat template 编码方式。通过 `tag_to_template` 配置，可以为特定 tag 的数据使用不同的 tokenization template： + +```yaml +tag_to_template: + math_dapo: qwen3 # 带 thinking token + KodCode: qwen3_nothink # 不带 thinking token +``` + +未在 `tag_to_template` 中配置的 tag 会 fallback 到 `global_template`。 + +### 单 Teacher 向后兼容 + +单 teacher 配置（`teacher` 为 WorkerConfig 而非 Dict）与原有行为完全一致： + +```yaml +# 以下配置与多 teacher 之前的版本行为完全一样 +teacher: + model_args: + model_name_or_path: Qwen/Qwen3-32B + device_mapping: list(range(0,16)) +``` + +内部会自动归一化为 `{"default": WorkerConfig}`，循环只执行一次。 + +--- + ## ✨️常见问题 ### Q1: 混合模式如何配置？ @@ -465,6 +642,18 @@ python examples/start_onpolicy_distill_pipeline.py \ - **纯 OPD 模式**：适合纯蒸馏训练，只需要 Teacher KL 信号，使用 `start_onpolicy_distill_pipeline.py` - **混合模式**：适合 RL + 蒸馏联合训练，使用 `start_rlvr_pipeline.py` 并配置 `use_opd: true` +### Q5: Multi-Teacher 模式下，如果某条数据没有任何 teacher 被路由到怎么办？ + +该条数据的 `total_weighted_kld = 0`： +- 纯 OPD 模式下 `advantage = 0`（该条数据不产生梯度） +- 混合模式下 `advantage = rl_advantages`（仅 RL 信号，没有蒸馏信号） + +### Q6: 多个 teacher 的 device_mapping 可以重叠吗？ + +可以，但不推荐： +- **不重叠**（推荐）：系统自动并行执行各 teacher 的 forward pass，显著减少推理时间 +- **重叠**：系统会串行执行，不会冲突但总耗时为所有 teacher 之和 + --- ## 参考资料 diff --git a/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Pipeline/sft_pipeline_start.md b/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Pipeline/sft_pipeline_start.md index 48e689089..00dd49214 100644 --- a/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Pipeline/sft_pipeline_start.md +++ b/docs_roll/i18n/zh-Hans/docusaurus-plugin-content-docs/current/User Guides/Pipeline/sft_pipeline_start.md @@ -109,7 +109,7 @@ * `dataloader_num_workers` * ... * **策略参数**（`strategy_args`） - * `strategy_name`：如 `megatron_train` / `deepspeed_train` 等 + * `strategy_name`：如 `megatron_train` / `fsdp2_train` 等 * 并行相关参数（tensor/pipeline 并行大小等） * **设备映射**（`device_mapping`） * 指定该 worker 使用哪些 GPU diff --git a/examples/agentic_deepeyes/deepeyes.yaml b/examples/agentic_deepeyes/deepeyes.yaml index ce87caa94..c7e087145 100644 --- a/examples/agentic_deepeyes/deepeyes.yaml +++ b/examples/agentic_deepeyes/deepeyes.yaml @@ -1,9 +1,5 @@ defaults: - ../config/envs@_here_ - - ../config/deepspeed_zero@_here_ - - ../config/deepspeed_zero2@_here_ - - ../config/deepspeed_zero3@_here_ - - ../config/deepspeed_zero3_cpuoffload@_here_ hydra: run: @@ -29,7 +25,7 @@ checkpoint_config: # tracker_kwargs: # log_dir: /data/oss_bucket_0/yuzhao/llm/tensorboard/roll_exp/deepeyes -offload_nccl: true +offload_nccl: false num_gpus_per_node: 8 diff --git a/examples/agentic_demo/agent_rollout_rock_swe.yaml b/examples/agentic_demo/agent_rollout_rock_swe.yaml index 1824c32de..72719e4b3 100644 --- a/examples/agentic_demo/agent_rollout_rock_swe.yaml +++ b/examples/agentic_demo/agent_rollout_rock_swe.yaml @@ -1,9 +1,5 @@ defaults: - ../config/traj_envs@_here_ - - ../config/deepspeed_zero@_here_ - - ../config/deepspeed_zero2@_here_ - - ../config/deepspeed_zero3@_here_ - - ../config/deepspeed_zero3_cpuoffload@_here_ hydra: run: diff --git a/examples/agentic_demo/agent_val_frozen_lake_multi_nodes_demo.yaml b/examples/agentic_demo/agent_val_frozen_lake_multi_nodes_demo.yaml index 5e88e736a..8e9523bfa 100644 --- a/examples/agentic_demo/agent_val_frozen_lake_multi_nodes_demo.yaml +++ b/examples/agentic_demo/agent_val_frozen_lake_multi_nodes_demo.yaml @@ -1,9 +1,5 @@ defaults: - ../config/traj_envs@_here_ - - ../config/deepspeed_zero@_here_ - - ../config/deepspeed_zero2@_here_ - - ../config/deepspeed_zero3@_here_ - - ../config/deepspeed_zero3_cpuoffload@_here_ hydra: run: @@ -78,8 +74,6 @@ actor_train: data_args: template: qwen2_5 strategy_args: - #strategy_name: deepspeed_train - #strategy_config: ${deepspeed_zero2} strategy_name: megatron_train strategy_config: tensor_model_parallel_size: 1 diff --git a/examples/agentic_demo/agent_val_frozen_lake_single_node_demo.yaml b/examples/agentic_demo/agent_val_frozen_lake_single_node_demo.yaml index b4435dd44..2c4d2fc3d 100644 --- a/examples/agentic_demo/agent_val_frozen_lake_single_node_demo.yaml +++ b/examples/agentic_demo/agent_val_frozen_lake_single_node_demo.yaml @@ -1,9 +1,5 @@ defaults: - ../config/traj_envs@_here_ - - ../config/deepspeed_zero@_here_ - - ../config/deepspeed_zero2@_here_ - - ../config/deepspeed_zero3@_here_ - - ../config/deepspeed_zero3_cpuoffload@_here_ hydra: run: @@ -78,8 +74,6 @@ actor_train: data_args: template: qwen2_5 strategy_args: - #strategy_name: deepspeed_train - #strategy_config: ${deepspeed_zero2} strategy_name: megatron_train strategy_config: tensor_model_parallel_size: 1 diff --git a/examples/agentic_demo/agent_val_rock_swe.yaml b/examples/agentic_demo/agent_val_rock_swe.yaml index 71c1071e6..2add45ca8 100644 --- a/examples/agentic_demo/agent_val_rock_swe.yaml +++ b/examples/agentic_demo/agent_val_rock_swe.yaml @@ -1,9 +1,5 @@ defaults: - ../config/traj_envs@_here_ - - ../config/deepspeed_zero@_here_ - - ../config/deepspeed_zero2@_here_ - - ../config/deepspeed_zero3@_here_ - - ../config/deepspeed_zero3_cpuoffload@_here_ hydra: run: diff --git a/examples/agentic_demo/agent_val_rock_swe_qwen35_2b.yaml b/examples/agentic_demo/agent_val_rock_swe_qwen35_2b.yaml index e814b36d4..d8dc23e19 100644 --- a/examples/agentic_demo/agent_val_rock_swe_qwen35_2b.yaml +++ b/examples/agentic_demo/agent_val_rock_swe_qwen35_2b.yaml @@ -1,9 +1,5 @@ defaults: - ../config/traj_envs@_here_ - - ../config/deepspeed_zero@_here_ - - ../config/deepspeed_zero2@_here_ - - ../config/deepspeed_zero3@_here_ - - ../config/deepspeed_zero3_cpuoffload@_here_ hydra: run: diff --git a/examples/agentic_demo/agent_val_sokoban_sandbox.yaml b/examples/agentic_demo/agent_val_sokoban_sandbox.yaml index 8bb2d8042..8bbf0d363 100644 --- a/examples/agentic_demo/agent_val_sokoban_sandbox.yaml +++ b/examples/agentic_demo/agent_val_sokoban_sandbox.yaml @@ -1,9 +1,5 @@ defaults: - ../config/traj_envs@_here_ - - ../config/deepspeed_zero@_here_ - - ../config/deepspeed_zero2@_here_ - - ../config/deepspeed_zero3@_here_ - - ../config/deepspeed_zero3_cpuoffload@_here_ hydra: run: @@ -78,8 +74,6 @@ actor_train: data_args: template: qwen2_5 strategy_args: -# strategy_name: deepspeed_train -# strategy_config: ${deepspeed_zero3} strategy_name: megatron_train strategy_config: tensor_model_parallel_size: 1 diff --git a/examples/agentic_demo/atropos_gsm8k_grpo_qwen25_0.5b.yaml b/examples/agentic_demo/atropos_gsm8k_grpo_qwen25_0.5b.yaml index df62c3aff..9e5058acf 100644 --- a/examples/agentic_demo/atropos_gsm8k_grpo_qwen25_0.5b.yaml +++ b/examples/agentic_demo/atropos_gsm8k_grpo_qwen25_0.5b.yaml @@ -31,15 +31,13 @@ actor_train: gradient_accumulation_steps: 32 device_mapping: "[1]" strategy_args: - strategy_name: deepspeed_train + strategy_name: megatron_train strategy_config: - zero_optimization: - stage: 2 - offload_optimizer: - device: cpu - pin_memory: false - bf16: - enabled: true + tensor_model_parallel_size: 1 + pipeline_model_parallel_size: 1 + expert_model_parallel_size: 1 + use_distributed_optimizer: true + recompute_granularity: full data_args: template: qwen2_5 @@ -63,12 +61,8 @@ max_tokens_per_step: 512 reference: device_mapping: "[1]" strategy_args: - strategy_name: deepspeed_train - strategy_config: - zero_optimization: - stage: 2 - bf16: - enabled: true + strategy_name: hf_infer + strategy_config: ~ # --- RL Configs (GRPO) --- adv_estimator: "grpo" diff --git a/examples/agentic_demo/openreward_endless_terminals_reinforce_qwen35_2b.yaml b/examples/agentic_demo/openreward_endless_terminals_reinforce_qwen35_2b.yaml index 84db30ba0..10ba70cd8 100644 --- a/examples/agentic_demo/openreward_endless_terminals_reinforce_qwen35_2b.yaml +++ b/examples/agentic_demo/openreward_endless_terminals_reinforce_qwen35_2b.yaml @@ -9,12 +9,6 @@ # --config_path agentic_demo \ # --config_name openreward_endless_terminals_reinforce_qwen35_2b -defaults: - - ../config/deepspeed_zero@_here_ - - ../config/deepspeed_zero2@_here_ - - ../config/deepspeed_zero3@_here_ - - ../config/deepspeed_zero3_cpuoffload@_here_ - hydra: run: dir: . diff --git a/examples/ascend_examples/qwen3_8b_rlvr_deepspeed.yaml b/examples/ascend_examples/qwen3_8b_rlvr_fsdp2.yaml similarity index 91% rename from examples/ascend_examples/qwen3_8b_rlvr_deepspeed.yaml rename to examples/ascend_examples/qwen3_8b_rlvr_fsdp2.yaml index 64fc91c95..7c4687514 100644 --- a/examples/ascend_examples/qwen3_8b_rlvr_deepspeed.yaml +++ b/examples/ascend_examples/qwen3_8b_rlvr_fsdp2.yaml @@ -1,9 +1,3 @@ -defaults: - - ../config/deepspeed_zero@_here_ - - ../config/deepspeed_zero2@_here_ - - ../config/deepspeed_zero3@_here_ - - ../config/deepspeed_zero3_cpuoffload@_here_ - hydra: run: dir: . @@ -113,8 +107,14 @@ actor_train: interleave_probs: "1.0" preprocessing_num_workers: 16 strategy_args: - strategy_name: deepspeed_train - strategy_config: ${deepspeed_zero3_cpuoffload} + strategy_name: fsdp2_train + strategy_config: + fsdp_size: 8 + param_dtype: bf16 + reduce_dtype: float32 + reshard_after_forward: true + offload_policy: false # offload model param + use_batched_model_update: false # disable batched gather during model update (not bring any benefit in test) device_mapping: list(range(0,8)) infer_batch_size: 1 diff --git a/examples/config/deepspeed_zero.yaml b/examples/config/deepspeed_zero.yaml deleted file mode 100644 index 876db1ace..000000000 --- a/examples/config/deepspeed_zero.yaml +++ /dev/null @@ -1,15 +0,0 @@ -deepspeed_zero: - train_micro_batch_size_per_gpu: auto - bf16: - enabled: true - fp16: - enabled: false - loss_scale: 0 - initial_scale_power: 16 - loss_scale_window: 1000 - hysteresis: 2 - min_loss_scale: 1 - zero_optimization: - stage: 0 - checkpoint: - tag_validation: IGNORE diff --git a/examples/config/deepspeed_zero2.yaml b/examples/config/deepspeed_zero2.yaml deleted file mode 100644 index e9db084d8..000000000 --- a/examples/config/deepspeed_zero2.yaml +++ /dev/null @@ -1,19 +0,0 @@ -deepspeed_zero2: - train_micro_batch_size_per_gpu: auto - bf16: - enabled: true - fp16: - enabled: false - loss_scale: 0 - initial_scale_power: 16 - loss_scale_window: 1000 - hysteresis: 2 - min_loss_scale: 1 - zero_optimization: - stage: 2 - allgather_partitions: true - allgather_bucket_size: 1.0e+9 - overlap_comm: true - reduce_scatter: true - reduce_bucket_size: 5.0e+8 - contiguous_gradients: true \ No newline at end of file diff --git a/examples/config/deepspeed_zero2_cpuoffload.yaml b/examples/config/deepspeed_zero2_cpuoffload.yaml deleted file mode 100644 index 3e78913b3..000000000 --- a/examples/config/deepspeed_zero2_cpuoffload.yaml +++ /dev/null @@ -1,25 +0,0 @@ -deepspeed_zero2_cpuoffload: - train_micro_batch_size_per_gpu: auto - bf16: - enabled: true - fp16: - enabled: false - loss_scale: 0 - initial_scale_power: 16 - loss_scale_window: 1000 - hysteresis: 2 - min_loss_scale: 1 - zero_optimization: - stage: 2 - offload_optimizer: - device: cpu - pin_memory: true - offload_param: - device: cpu - pin_memory: true - allgather_partitions: true - allgather_bucket_size: 1.0e+9 - overlap_comm: true - reduce_scatter: true - reduce_bucket_size: 5.0e+8 - contiguous_gradients: true \ No newline at end of file diff --git a/examples/config/deepspeed_zero3.yaml b/examples/config/deepspeed_zero3.yaml deleted file mode 100644 index 45e4c5abd..000000000 --- a/examples/config/deepspeed_zero3.yaml +++ /dev/null @@ -1,22 +0,0 @@ -deepspeed_zero3: - train_micro_batch_size_per_gpu: auto - bf16: - enabled: true - fp16: - enabled: false - loss_scale: 0 - initial_scale_power: 16 - loss_scale_window: 1000 - hysteresis: 2 - min_loss_scale: 1 - zero_optimization: - stage: 3 - overlap_comm: true - contiguous_gradients: true - sub_group_size: 1.0e+9 - reduce_bucket_size: auto - stage3_prefetch_bucket_size: auto - stage3_param_persistence_threshold: auto - stage3_max_live_parameters: 1.0e+9 - stage3_max_reuse_distance: 1.0e+9 - stage3_gather_16bit_weights_on_model_save: true diff --git a/examples/config/deepspeed_zero3_cpuoffload.yaml b/examples/config/deepspeed_zero3_cpuoffload.yaml deleted file mode 100644 index 5df39fc00..000000000 --- a/examples/config/deepspeed_zero3_cpuoffload.yaml +++ /dev/null @@ -1,25 +0,0 @@ -deepspeed_zero3_cpuoffload: - train_micro_batch_size_per_gpu: auto - bf16: - enabled: true - fp16: - enabled: false - loss_scale: 0 - initial_scale_power: 16 - loss_scale_window: 1000 - hysteresis: 2 - min_loss_scale: 1 - zero_optimization: - stage: 3 - offload_optimizer: - device: cpu - pin_memory: true - overlap_comm: true - contiguous_gradients: true - sub_group_size: 1.0e+9 - reduce_bucket_size: auto - stage3_prefetch_bucket_size: auto - stage3_param_persistence_threshold: auto - stage3_max_live_parameters: 1.0e+9 - stage3_max_reuse_distance: 1.0e+9 - stage3_gather_16bit_weights_on_model_save: true \ No newline at end of file diff --git a/examples/config/traj_envs.yaml b/examples/config/traj_envs.yaml index 95363c2e8..7755357d5 100644 --- a/examples/config/traj_envs.yaml +++ b/examples/config/traj_envs.yaml @@ -8,12 +8,14 @@ sokoban_format_penalty: -0.15 frozen_format_penalty: -0.01 env_manager_cls: roll.pipeline.agentic.env_manager.traj_env_manager.TrajEnvManager +agent_runner_cls: null custom_env: SimpleSokoban: env_type: sokoban max_steps: ${max_actions_per_traj} max_tokens_per_step: ${max_tokens_per_step} env_manager_cls: ${env_manager_cls} + agent_runner_cls: ${agent_runner_cls} use_thread_lock: true agent_system_template: ${agent_system_template} agent_template: ${agent_template} @@ -28,6 +30,7 @@ custom_env: max_steps: ${max_actions_per_traj} max_tokens_per_step: ${max_tokens_per_step} env_manager_cls: ${env_manager_cls} + agent_runner_cls: ${agent_runner_cls} use_thread_lock: true agent_system_template: ${agent_system_template} agent_template: ${agent_template} @@ -43,6 +46,7 @@ custom_env: max_steps: ${max_actions_per_traj} max_tokens_per_step: ${max_tokens_per_step} env_manager_cls: ${env_manager_cls} + agent_runner_cls: ${agent_runner_cls} use_thread_lock: true agent_system_template: ${agent_system_template} agent_template: ${agent_template} @@ -60,6 +64,7 @@ custom_env: max_steps: ${max_actions_per_traj} max_tokens_per_step: ${max_tokens_per_step} env_manager_cls: ${env_manager_cls} + agent_runner_cls: ${agent_runner_cls} use_thread_lock: true agent_system_template: ${agent_system_template} agent_template: ${agent_template} @@ -73,6 +78,7 @@ custom_env: max_steps: ${max_actions_per_traj} max_tokens_per_step: ${max_tokens_per_step} env_manager_cls: ${env_manager_cls} + agent_runner_cls: ${agent_runner_cls} use_thread_lock: true agent_system_template: ${agent_system_template} agent_template: ${agent_template} @@ -86,6 +92,7 @@ custom_env: max_steps: ${max_actions_per_traj} max_tokens_per_step: ${max_tokens_per_step} env_manager_cls: ${env_manager_cls} + agent_runner_cls: ${agent_runner_cls} use_thread_lock: true agent_system_template: ${agent_system_template} agent_template: ${agent_template} diff --git a/examples/docs_examples/example_ppo.yaml b/examples/docs_examples/example_ppo.yaml index c33ecc6c6..d17784bf9 100644 --- a/examples/docs_examples/example_ppo.yaml +++ b/examples/docs_examples/example_ppo.yaml @@ -1,9 +1,3 @@ -defaults: - - ../config/deepspeed_zero@_here_ - - ../config/deepspeed_zero2@_here_ - - ../config/deepspeed_zero3@_here_ - - ../config/deepspeed_zero3_cpuoffload@_here_ - hydra: run: dir: . @@ -64,7 +58,7 @@ loss_agg_mode: "seq-mean-token-mean" whiten_advantages: true advantage_clip: 2.0 -reward_clip: ~ +reward_clip: 10 dual_clip_loss: true lambd: 0.95 gamma: 1 @@ -184,8 +178,15 @@ critic: data_args: template: qwen2_5 strategy_args: - strategy_name: deepspeed_train - strategy_config: ${deepspeed_zero3} + strategy_name: fsdp2_train + strategy_config: + fsdp_size: 8 + param_dtype: bf16 + reduce_dtype: float32 + reshard_after_forward: true + offload_policy: false + wrap_policy: # trl wrapped model should specify this + transformer_layer_cls_to_wrap: ["Qwen2DecoderLayer"] device_mapping: list(range(8,16)) infer_batch_size: 4 diff --git a/examples/docs_examples/example_reward_fl.yaml b/examples/docs_examples/example_reward_fl.yaml deleted file mode 100644 index b950a1dde..000000000 --- a/examples/docs_examples/example_reward_fl.yaml +++ /dev/null @@ -1,67 +0,0 @@ -defaults: - - ../config/deepspeed_zero@_here_ - - ../config/deepspeed_zero2@_here_ - - ../config/deepspeed_zero2_cpuoffload@_here_ - - ../config/deepspeed_zero3@_here_ - - ../config/deepspeed_zero3_cpuoffload@_here_ - -hydra: - run: - dir: . - output_subdir: null - -exp_name: "reward_fl_zero2_cpuoffload" -seed: 42 -logging_dir: ./output/logs -output_dir: ./output - -checkpoint_config: - type: file_system - output_dir: /data/models/reward_fl/ - -save_steps: 25 -logging_steps: 1 -resume_from_checkpoint: false - -sequence_length: 1024 -train_batch_size: 8 -max_grad_norm: 1.0 - -actor_train: - model_args: - model_type: diffusion_module - dtype: bf16 - model_config_kwargs: - model_name: wan2_2 - model_paths: ./examples/wan2.2-14B-reward_fl_ds/wan22_paths.json - reward_model_path: /data/models/antelopev2/ - tokenizer_path: /data/models/Wan-AI/Wan2.1-T2V-1.3B/google/umt5-xxl/ - model_id_with_origin_paths: null - trainable_models: dit2 - use_gradient_checkpointing_offload: true - extra_inputs: input_image - max_timestep_boundary: 1.0 - min_timestep_boundary: 0.9 - num_inference_steps: 8 - mid_timestep: 4 - final_timestep: 7 - - training_args: - learning_rate: 2.5e-6 - lr_scheduler_type: constant - per_device_train_batch_size: 1 - gradient_accumulation_steps: 1 - warmup_steps: 10 - num_train_epochs: 1 - - data_args: - file_name: ./data/example_video_dataset/metadata.csv - preprocessing_num_workers: 2 - - strategy_args: - strategy_name: diffusion_deepspeed_train - strategy_config: ${deepspeed_zero2_cpuoffload} - device_mapping: list(range(0,8)) - -system_envs: - RAY_PROFILING: "0" diff --git a/examples/dpo_examples/qwen2.5-3B_dpo_megatron.yaml b/examples/dpo_examples/qwen2.5-3B_dpo_megatron.yaml index fef389fc4..63cb6838e 100644 --- a/examples/dpo_examples/qwen2.5-3B_dpo_megatron.yaml +++ b/examples/dpo_examples/qwen2.5-3B_dpo_megatron.yaml @@ -1,9 +1,3 @@ -defaults: - - ../config/deepspeed_zero@_here_ - - ../config/deepspeed_zero2@_here_ - - ../config/deepspeed_zero3@_here_ - - ../config/deepspeed_zero3_cpuoffload@_here_ - hydra: run: dir: . diff --git a/examples/dpo_examples/qwen3-30BA3B-dpo_megatron_80GB.yaml b/examples/dpo_examples/qwen3-30BA3B-dpo_megatron_80GB.yaml index ca37ba19d..353ce3fae 100644 --- a/examples/dpo_examples/qwen3-30BA3B-dpo_megatron_80GB.yaml +++ b/examples/dpo_examples/qwen3-30BA3B-dpo_megatron_80GB.yaml @@ -1,9 +1,3 @@ -defaults: - - ../config/deepspeed_zero@_here_ - - ../config/deepspeed_zero2@_here_ - - ../config/deepspeed_zero3@_here_ - - ../config/deepspeed_zero3_cpuoffload@_here_ - hydra: run: dir: . diff --git a/examples/qwen2.5-0.5B-agentic/agent_val_frozen_lake-pg_var.yaml b/examples/qwen2.5-0.5B-agentic/agent_val_frozen_lake-pg_var.yaml index 646f8e97d..3fb0dfa85 100644 --- a/examples/qwen2.5-0.5B-agentic/agent_val_frozen_lake-pg_var.yaml +++ b/examples/qwen2.5-0.5B-agentic/agent_val_frozen_lake-pg_var.yaml @@ -1,9 +1,5 @@ defaults: - ../config/traj_envs@_here_ - - ../config/deepspeed_zero@_here_ - - ../config/deepspeed_zero2@_here_ - - ../config/deepspeed_zero3@_here_ - - ../config/deepspeed_zero3_cpuoffload@_here_ hydra: run: @@ -81,8 +77,6 @@ actor_train: data_args: template: qwen2_5 strategy_args: -# strategy_name: deepspeed_train -# strategy_config: ${deepspeed_zero3} strategy_name: megatron_train strategy_config: tensor_model_parallel_size: 1 diff --git a/examples/qwen2.5-0.5B-agentic/agent_val_frozen_lake-pg_var_is_correct.yaml b/examples/qwen2.5-0.5B-agentic/agent_val_frozen_lake-pg_var_is_correct.yaml index 8e1e79631..ef9595693 100644 --- a/examples/qwen2.5-0.5B-agentic/agent_val_frozen_lake-pg_var_is_correct.yaml +++ b/examples/qwen2.5-0.5B-agentic/agent_val_frozen_lake-pg_var_is_correct.yaml @@ -1,9 +1,5 @@ defaults: - ../config/traj_envs@_here_ - - ../config/deepspeed_zero@_here_ - - ../config/deepspeed_zero2@_here_ - - ../config/deepspeed_zero3@_here_ - - ../config/deepspeed_zero3_cpuoffload@_here_ hydra: run: @@ -98,8 +94,6 @@ actor_train: data_args: template: qwen2_5 strategy_args: -# strategy_name: deepspeed_train -# strategy_config: ${deepspeed_zero3} strategy_name: megatron_train strategy_config: tensor_model_parallel_size: 1 diff --git a/examples/qwen2.5-0.5B-agentic/agent_val_frozen_lake.yaml b/examples/qwen2.5-0.5B-agentic/agent_val_frozen_lake.yaml index b93ad3b75..5340becde 100644 --- a/examples/qwen2.5-0.5B-agentic/agent_val_frozen_lake.yaml +++ b/examples/qwen2.5-0.5B-agentic/agent_val_frozen_lake.yaml @@ -1,9 +1,5 @@ defaults: - ../config/traj_envs@_here_ - - ../config/deepspeed_zero@_here_ - - ../config/deepspeed_zero2@_here_ - - ../config/deepspeed_zero3@_here_ - - ../config/deepspeed_zero3_cpuoffload@_here_ hydra: run: @@ -79,8 +75,6 @@ actor_train: data_args: template: qwen2_5 strategy_args: -# strategy_name: deepspeed_train -# strategy_config: ${deepspeed_zero3} strategy_name: megatron_train strategy_config: tensor_model_parallel_size: 1 diff --git a/examples/qwen2.5-0.5B-agentic/agent_val_frozen_lake_amd.yaml b/examples/qwen2.5-0.5B-agentic/agent_val_frozen_lake_amd.yaml index 4d8efa742..f19cba692 100644 --- a/examples/qwen2.5-0.5B-agentic/agent_val_frozen_lake_amd.yaml +++ b/examples/qwen2.5-0.5B-agentic/agent_val_frozen_lake_amd.yaml @@ -1,9 +1,5 @@ defaults: - ../config/traj_envs@_here_ - - ../config/deepspeed_zero@_here_ - - ../config/deepspeed_zero2@_here_ - - ../config/deepspeed_zero3@_here_ - - ../config/deepspeed_zero3_cpuoffload@_here_ hydra: run: @@ -79,8 +75,6 @@ actor_train: data_args: template: qwen2_5 strategy_args: -# strategy_name: deepspeed_train -# strategy_config: ${deepspeed_zero3} strategy_name: megatron_train strategy_config: tensor_model_parallel_size: 1 diff --git a/examples/qwen2.5-0.5B-agentic/agent_val_frozen_lake_async.yaml b/examples/qwen2.5-0.5B-agentic/agent_val_frozen_lake_async.yaml index 4e252e123..26c5f1104 100644 --- a/examples/qwen2.5-0.5B-agentic/agent_val_frozen_lake_async.yaml +++ b/examples/qwen2.5-0.5B-agentic/agent_val_frozen_lake_async.yaml @@ -1,9 +1,5 @@ defaults: - ../config/traj_envs@_here_ - - ../config/deepspeed_zero@_here_ - - ../config/deepspeed_zero2@_here_ - - ../config/deepspeed_zero3@_here_ - - ../config/deepspeed_zero3_cpuoffload@_here_ hydra: run: @@ -80,8 +76,6 @@ actor_train: data_args: template: qwen2_5 strategy_args: -# strategy_name: deepspeed_train -# strategy_config: ${deepspeed_zero3} strategy_name: megatron_train strategy_config: tensor_model_parallel_size: 1 diff --git a/examples/qwen2.5-0.5B-agentic/agent_val_frozen_lake_async_amd.yaml b/examples/qwen2.5-0.5B-agentic/agent_val_frozen_lake_async_amd.yaml index 3696af5ec..ce664d332 100644 --- a/examples/qwen2.5-0.5B-agentic/agent_val_frozen_lake_async_amd.yaml +++ b/examples/qwen2.5-0.5B-agentic/agent_val_frozen_lake_async_amd.yaml @@ -1,9 +1,5 @@ defaults: - ../config/traj_envs@_here_ - - ../config/deepspeed_zero@_here_ - - ../config/deepspeed_zero2@_here_ - - ../config/deepspeed_zero3@_here_ - - ../config/deepspeed_zero3_cpuoffload@_here_ hydra: run: @@ -80,8 +76,6 @@ actor_train: data_args: template: qwen2_5 strategy_args: -# strategy_name: deepspeed_train -# strategy_config: ${deepspeed_zero3} strategy_name: megatron_train strategy_config: tensor_model_parallel_size: 1 diff --git a/examples/qwen2.5-0.5B-agentic/agent_val_frozen_lake_gigpo.yaml b/examples/qwen2.5-0.5B-agentic/agent_val_frozen_lake_gigpo.yaml index ed2ac26f7..78b40ab62 100644 --- a/examples/qwen2.5-0.5B-agentic/agent_val_frozen_lake_gigpo.yaml +++ b/examples/qwen2.5-0.5B-agentic/agent_val_frozen_lake_gigpo.yaml @@ -1,9 +1,5 @@ defaults: - ../config/step_envs@_here_ - - ../config/deepspeed_zero@_here_ - - ../config/deepspeed_zero2@_here_ - - ../config/deepspeed_zero3@_here_ - - ../config/deepspeed_zero3_cpuoffload@_here_ hydra: run: @@ -87,8 +83,6 @@ actor_train: data_args: template: qwen2_5 strategy_args: -# strategy_name: deepspeed_train -# strategy_config: ${deepspeed_zero3} strategy_name: megatron_train strategy_config: tensor_model_parallel_size: 1 diff --git a/examples/qwen2.5-0.5B-agentic/agentic_rollout_sokoban.yaml b/examples/qwen2.5-0.5B-agentic/agentic_rollout_sokoban.yaml index 619e78e63..f583d1d9a 100644 --- a/examples/qwen2.5-0.5B-agentic/agentic_rollout_sokoban.yaml +++ b/examples/qwen2.5-0.5B-agentic/agentic_rollout_sokoban.yaml @@ -1,9 +1,5 @@ defaults: - ../config/traj_envs@_here_ - - ../config/deepspeed_zero@_here_ - - ../config/deepspeed_zero2@_here_ - - ../config/deepspeed_zero3@_here_ - - ../config/deepspeed_zero3_cpuoffload@_here_ hydra: run: diff --git a/examples/qwen2.5-0.5B-agentic/agentic_sokoban_rollout_mock_dump.yaml b/examples/qwen2.5-0.5B-agentic/agentic_sokoban_rollout_mock_dump.yaml index b59749301..995fbc32b 100644 --- a/examples/qwen2.5-0.5B-agentic/agentic_sokoban_rollout_mock_dump.yaml +++ b/examples/qwen2.5-0.5B-agentic/agentic_sokoban_rollout_mock_dump.yaml @@ -1,9 +1,5 @@ defaults: - ../config/traj_envs@_here_ - - ../config/deepspeed_zero@_here_ - - ../config/deepspeed_zero2@_here_ - - ../config/deepspeed_zero3@_here_ - - ../config/deepspeed_zero3_cpuoffload@_here_ hydra: run: diff --git a/examples/qwen2.5-0.5B-agentic/agentic_val_sokoban.yaml b/examples/qwen2.5-0.5B-agentic/agentic_val_sokoban.yaml index e753b03ec..908543f6a 100644 --- a/examples/qwen2.5-0.5B-agentic/agentic_val_sokoban.yaml +++ b/examples/qwen2.5-0.5B-agentic/agentic_val_sokoban.yaml @@ -1,9 +1,5 @@ defaults: - ../config/traj_envs@_here_ - - ../config/deepspeed_zero@_here_ - - ../config/deepspeed_zero2@_here_ - - ../config/deepspeed_zero3@_here_ - - ../config/deepspeed_zero3_cpuoffload@_here_ hydra: run: @@ -79,8 +75,6 @@ actor_train: data_args: template: qwen2_5 strategy_args: -# strategy_name: deepspeed_train -# strategy_config: ${deepspeed_zero3} strategy_name: megatron_train strategy_config: tensor_model_parallel_size: 1 diff --git a/examples/qwen2.5-0.5B-agentic/agentic_val_sokoban_agent_runner.yaml b/examples/qwen2.5-0.5B-agentic/agentic_val_sokoban_agent_runner.yaml new file mode 100644 index 000000000..87d580554 --- /dev/null +++ b/examples/qwen2.5-0.5B-agentic/agentic_val_sokoban_agent_runner.yaml @@ -0,0 +1,159 @@ +defaults: + - ../config/traj_envs@_here_ + +hydra: + run: + dir: . + output_subdir: null + +exp_name: "agentic_pipeline" +seed: 42 +logging_dir: ./output/logs +output_dir: ./output +render_save_dir: ./output/render +system_envs: + USE_MODELSCOPE: '1' + +#track_with: wandb +#tracker_kwargs: +# api_key: +# project: roll-agentic +# name: ${exp_name}_sokoban +# notes: "agentic_pipeline" +# tags: +# - agentic +# - roll +# - baseline + +track_with: tensorboard +tracker_kwargs: + log_dir: /data/oss_bucket_0/yali/llm/tensorboard/roll_exp/agentic_sokoban + + +checkpoint_config: + type: file_system + output_dir: /data/cpfs_0/rl_examples/models/${exp_name} + +num_gpus_per_node: 8 + +max_steps: 1024 +save_steps: 10000 +logging_steps: 1 +eval_steps: 10 +resume_from_checkpoint: false + +rollout_batch_size: 1024 +val_batch_size: 1024 +sequence_length: 8192 + +advantage_clip: 0.2 +ppo_epochs: 1 +adv_estimator: "grpo" +#pg_clip: 0.1 +#dual_clip_loss: True +init_kl_coef: 0.0 +whiten_advantages: true +entropy_loss_coef: 0 +max_grad_norm: 1.0 + +pretrain: Qwen/Qwen2.5-0.5B-Instruct +reward_pretrain: Qwen/Qwen2.5-0.5B-Instruct + +actor_train: + model_args: + attn_implementation: fa2 + disable_gradient_checkpointing: false + dtype: bf16 + model_type: ~ + training_args: + learning_rate: 1.0e-6 + weight_decay: 0 + per_device_train_batch_size: 2 + gradient_accumulation_steps: 64 + warmup_steps: 10 + lr_scheduler_type: cosine + data_args: + template: qwen2_5 + strategy_args: + strategy_name: megatron_train + strategy_config: + tensor_model_parallel_size: 1 + pipeline_model_parallel_size: 1 + expert_model_parallel_size: 1 + use_distributed_optimizer: true + recompute_granularity: full + device_mapping: list(range(0,8)) + infer_batch_size: 2 + +actor_infer: + model_args: + disable_gradient_checkpointing: true + dtype: bf16 + generating_args: + max_new_tokens: 128 # single-turn response length + top_p: 0.99 + top_k: 100 + num_beams: 1 + temperature: 0.99 + num_return_sequences: 1 + data_args: + template: qwen2_5 + strategy_args: + strategy_name: vllm + strategy_config: + gpu_memory_utilization: 0.8 + block_size: 16 + load_format: auto + device_mapping: list(range(0,8)) + +reference: + model_args: + attn_implementation: fa2 + disable_gradient_checkpointing: true + dtype: bf16 + model_type: ~ + data_args: + template: qwen2_5 + strategy_args: + strategy_name: hf_infer + strategy_config: ~ + device_mapping: list(range(0,8)) + infer_batch_size: 2 + +reward_normalization: + grouping: traj_group_id # 可以tags(env_type)/traj_group_id(group)/batch(rollout_batch)... group_by计算reward/adv + method: mean_std # asym_clip / identity / mean_std + +train_env_manager: + format_penalty: -0.15 # sokoban env penalty_for_step=-0.1 + max_env_num_per_worker: 16 + num_env_groups: 128 + # under the same group, the env config and env seed are ensured to be equal + group_size: 8 + tags: [SimpleSokoban] + num_groups_partition: [128] # If not set, all env names divide nums equally. Under the same group, the env config and env seed (prompt) are equal in each generation + +val_env_manager: + max_env_num_per_worker: 32 + num_env_groups: 1024 + group_size: 1 # should be set to 1 because val temperature is set to 0 and same prompt leads to same output + tags: [SimpleSokoban, LargerSokoban, SokobanDifferentGridVocab, FrozenLake] + num_groups_partition: [256, 256, 256, 256] # TODO: If not set, all env names divide nums equally. Under the same group, the env config and env seed (prompt) are equal in each generation + + +# Here, you can override variables defined in the imported envs. max_tokens_per_step: 128 in custom_env.SimpleSokoban, here replaced by 64 +max_tokens_per_step: 64 + +env_manager_cls: roll.pipeline.agentic.env_manager.proxy_env_manager.ProxyEnvManager +agent_runner_cls: roll.pipeline.agentic.agent_runner.gem_runner.GEMRunner +custom_envs: + SimpleSokoban: + ${custom_env.SimpleSokoban} + LargerSokoban: + ${custom_env.LargerSokoban} + SokobanDifferentGridVocab: + ${custom_env.SokobanDifferentGridVocab} + FrozenLake: + ${custom_env.FrozenLake} + FrozenLakeThink: + ${custom_env.FrozenLakeThink} diff --git a/examples/qwen2.5-0.5B-agentic/agentic_val_sokoban_dynamic_batching.yaml b/examples/qwen2.5-0.5B-agentic/agentic_val_sokoban_dynamic_batching.yaml index 02d28f8c0..50a2374ff 100644 --- a/examples/qwen2.5-0.5B-agentic/agentic_val_sokoban_dynamic_batching.yaml +++ b/examples/qwen2.5-0.5B-agentic/agentic_val_sokoban_dynamic_batching.yaml @@ -1,9 +1,5 @@ defaults: - ../config/traj_envs@_here_ - - ../config/deepspeed_zero@_here_ - - ../config/deepspeed_zero2@_here_ - - ../config/deepspeed_zero3@_here_ - - ../config/deepspeed_zero3_cpuoffload@_here_ hydra: run: @@ -83,8 +79,6 @@ actor_train: data_args: template: qwen2_5 strategy_args: -# strategy_name: deepspeed_train -# strategy_config: ${deepspeed_zero3} strategy_name: megatron_train strategy_config: tensor_model_parallel_size: 1 diff --git a/examples/qwen2.5-0.5B-agentic/agentic_val_sokoban_gigpo.yaml b/examples/qwen2.5-0.5B-agentic/agentic_val_sokoban_gigpo.yaml index 819a13005..b7b960d74 100644 --- a/examples/qwen2.5-0.5B-agentic/agentic_val_sokoban_gigpo.yaml +++ b/examples/qwen2.5-0.5B-agentic/agentic_val_sokoban_gigpo.yaml @@ -1,9 +1,5 @@ defaults: - ../config/step_envs@_here_ - - ../config/deepspeed_zero@_here_ - - ../config/deepspeed_zero2@_here_ - - ../config/deepspeed_zero3@_here_ - - ../config/deepspeed_zero3_cpuoffload@_here_ hydra: run: @@ -87,8 +83,6 @@ actor_train: data_args: template: qwen2_5 strategy_args: -# strategy_name: deepspeed_train -# strategy_config: ${deepspeed_zero3} strategy_name: megatron_train strategy_config: tensor_model_parallel_size: 1 diff --git a/examples/qwen2.5-0.5B-agentic/agentic_val_sokoban_lora.yaml b/examples/qwen2.5-0.5B-agentic/agentic_val_sokoban_lora.yaml index 5dd3fe377..7dd82801e 100644 --- a/examples/qwen2.5-0.5B-agentic/agentic_val_sokoban_lora.yaml +++ b/examples/qwen2.5-0.5B-agentic/agentic_val_sokoban_lora.yaml @@ -1,9 +1,5 @@ defaults: - ../config/traj_envs@_here_ - - ../config/deepspeed_zero@_here_ - - ../config/deepspeed_zero2@_here_ - - ../config/deepspeed_zero3@_here_ - - ../config/deepspeed_zero3_cpuoffload@_here_ hydra: run: @@ -82,8 +78,6 @@ actor_train: data_args: template: qwen2_5 strategy_args: -# strategy_name: deepspeed_train -# strategy_config: ${deepspeed_zero3} strategy_name: megatron_train strategy_config: tensor_model_parallel_size: 1 diff --git a/examples/qwen2.5-0.5B-agentic/agentic_val_sokoban_native.yaml b/examples/qwen2.5-0.5B-agentic/agentic_val_sokoban_native.yaml index 3483fac2b..305a4b36d 100644 --- a/examples/qwen2.5-0.5B-agentic/agentic_val_sokoban_native.yaml +++ b/examples/qwen2.5-0.5B-agentic/agentic_val_sokoban_native.yaml @@ -1,9 +1,5 @@ defaults: - ../config/traj_envs@_here_ - - ../config/deepspeed_zero@_here_ - - ../config/deepspeed_zero2@_here_ - - ../config/deepspeed_zero3@_here_ - - ../config/deepspeed_zero3_cpuoffload@_here_ hydra: run: @@ -85,8 +81,6 @@ actor_train: data_args: template: qwen2_5 strategy_args: -# strategy_name: deepspeed_train -# strategy_config: ${deepspeed_zero3} strategy_name: megatron_train strategy_config: tensor_model_parallel_size: 1 diff --git a/examples/qwen2.5-0.5B-agentic/agentic_val_sokoban_sandbox.yaml b/examples/qwen2.5-0.5B-agentic/agentic_val_sokoban_sandbox.yaml index df2feb840..1d25d9bb6 100644 --- a/examples/qwen2.5-0.5B-agentic/agentic_val_sokoban_sandbox.yaml +++ b/examples/qwen2.5-0.5B-agentic/agentic_val_sokoban_sandbox.yaml @@ -1,9 +1,5 @@ defaults: - ../config/traj_envs@_here_ - - ../config/deepspeed_zero@_here_ - - ../config/deepspeed_zero2@_here_ - - ../config/deepspeed_zero3@_here_ - - ../config/deepspeed_zero3_cpuoffload@_here_ hydra: run: @@ -78,8 +74,6 @@ actor_train: data_args: template: qwen2_5 strategy_args: -# strategy_name: deepspeed_train -# strategy_config: ${deepspeed_zero3} strategy_name: megatron_train strategy_config: tensor_model_parallel_size: 1 diff --git a/examples/qwen2.5-1.5B-distill_ds/distill_zero3.yaml b/examples/qwen2.5-1.5B-distill_ds/distill_zero3.yaml deleted file mode 100644 index 50716cc68..000000000 --- a/examples/qwen2.5-1.5B-distill_ds/distill_zero3.yaml +++ /dev/null @@ -1,85 +0,0 @@ -defaults: - - ../config/deepspeed_zero@_here_ - - ../config/deepspeed_zero2@_here_ - - ../config/deepspeed_zero3@_here_ - - ../config/deepspeed_zero3_cpuoffload@_here_ - -hydra: - run: - dir: . - output_subdir: null - -exp_name: "distill_zero3" -seed: 42 -logging_dir: ./output/logs -output_dir: ./output - -checkpoint_config: - type: file_system - output_dir: /var/tmp/ckpt - - -save_steps: 100 -logging_steps: 1 -resume_from_checkpoint: false - -student_pretrain: Qwen/Qwen2.5-1.5B-Instruct -teacher_pretrain: Qwen/Qwen2.5-7B-Instruct - -# distill config -logits_topk: 64 -distill_loss_weight: 0.85 -kd_objective: forward_kl -distill_on_prompt: False - -logits_transfer_backend: "nccl-only" # support "ipc+nccl", "nccl_only" and "ray" - -sequence_length: 1024 -max_grad_norm: 1.0 - -question_key: question_zh -answer_key: answer_zh - - -student: - model_args: - attn_implementation: fa2 - disable_gradient_checkpointing: false - dtype: bf16 - model_type: ~ - training_args: - learning_rate: 2.0e-5 - weight_decay: 1.0e-2 - lr_scheduler_type: constant - per_device_train_batch_size: 1 - gradient_accumulation_steps: 1 - warmup_steps: 0 - num_train_epochs: 1 - data_args: - template: qwen2_5 - file_name: - - data/GSM8K_zh.json #https://huggingface.co/datasets/meta-math/GSM8K_zh - preprocessing_num_workers: 16 - - strategy_args: - strategy_name: deepspeed_train - strategy_config: ${deepspeed_zero3} - device_mapping: list(range(0,8)) - -teacher: - model_args: - attn_implementation: fa2 - disable_gradient_checkpointing: true - dtype: bf16 - data_args: - template: qwen2_5 - training_args: - # teacher forward micro_batch_size - per_device_train_batch_size: 1 - strategy_args: - strategy_name: deepspeed_infer - strategy_config: ${deepspeed_zero3} - device_mapping: list(range(0,8)) - -system_envs: - RAY_PROFILING: "0" \ No newline at end of file diff --git a/examples/qwen2.5-1.5B-distill_ds/run_distill_pipeline.sh b/examples/qwen2.5-1.5B-distill_ds/run_distill_pipeline.sh deleted file mode 100644 index 41bc16687..000000000 --- a/examples/qwen2.5-1.5B-distill_ds/run_distill_pipeline.sh +++ /dev/null @@ -1,6 +0,0 @@ - -#!/bin/bash -set +x - -CONFIG_PATH=$(basename $(dirname $0)) -python examples/start_distill_pipeline.py --config_path $CONFIG_PATH --config_name distill_zero3 diff --git a/examples/qwen2.5-7B-agentic_megatron/agentic_val_webshop.yaml b/examples/qwen2.5-7B-agentic_megatron/agentic_val_webshop.yaml index f5a42dd38..97b35241b 100644 --- a/examples/qwen2.5-7B-agentic_megatron/agentic_val_webshop.yaml +++ b/examples/qwen2.5-7B-agentic_megatron/agentic_val_webshop.yaml @@ -1,9 +1,5 @@ defaults: - ../config/traj_envs@_here_ - - ../config/deepspeed_zero@_here_ - - ../config/deepspeed_zero2@_here_ - - ../config/deepspeed_zero3@_here_ - - ../config/deepspeed_zero3_cpuoffload@_here_ hydra: run: diff --git a/examples/qwen2.5-7B-agentic_megatron/agentic_val_webshop_amd.yaml b/examples/qwen2.5-7B-agentic_megatron/agentic_val_webshop_amd.yaml index 71799ee84..e17471a80 100644 --- a/examples/qwen2.5-7B-agentic_megatron/agentic_val_webshop_amd.yaml +++ b/examples/qwen2.5-7B-agentic_megatron/agentic_val_webshop_amd.yaml @@ -1,16 +1,12 @@ defaults: - - ../config/envs@_here_ - - ../config/deepspeed_zero@_here_ - - ../config/deepspeed_zero2@_here_ - - ../config/deepspeed_zero3@_here_ - - ../config/deepspeed_zero3_cpuoffload@_here_ + - ../config/traj_envs@_here_ hydra: run: dir: . output_subdir: null -exp_name: "agentic_pipeline" +exp_name: "agentic_pipeline_webshop" seed: 42 logging_dir: ./output/logs output_dir: ./output @@ -31,7 +27,7 @@ system_envs: track_with: tensorboard tracker_kwargs: - log_dir: /data/oss_bucket_0/yali/llm/tensorboard/roll_exp/agentic_sokoban + log_dir: /data/oss_bucket_0/yali/llm/tensorboard/roll_exp/agentic_webshop num_gpus_per_node: 8 @@ -41,15 +37,16 @@ logging_steps: 1 eval_steps: 10 resume_from_checkpoint: false -rollout_batch_size: 512 -val_batch_size: 1024 +rollout_batch_size: 64 +val_batch_size: 64 sequence_length: 8192 reward_clip: 20 -advantage_clip: 10.0 +advantage_clip: 0.2 # 0.1-0.3 ppo_epochs: 1 -adv_estimator: "reinforce" +adv_estimator: "grpo" #pg_clip: 0.1 +max_grad_norm: 1.0 #dual_clip_loss: True init_kl_coef: 0.0 whiten_advantages: true @@ -60,31 +57,35 @@ reward_pretrain: Qwen/Qwen2.5-7B-Instruct actor_train: model_args: - flash_attn: fa2 + attn_implementation: fa2 disable_gradient_checkpointing: false dtype: bf16 model_type: ~ training_args: learning_rate: 1.0e-6 weight_decay: 0 - per_device_train_batch_size: 2 - gradient_accumulation_steps: 32 + per_device_train_batch_size: 1 + gradient_accumulation_steps: 8 warmup_steps: 10 data_args: template: qwen2_5 strategy_args: - strategy_name: deepspeed_train - strategy_config: ${deepspeed_zero2} + strategy_name: fsdp2_train + strategy_config: + fsdp_size: 8 + param_dtype: bf16 + reduce_dtype: float32 + reshard_after_forward: false + offload_policy: false device_mapping: list(range(0,8)) infer_batch_size: 1 actor_infer: model_args: - flash_attn: fa2 disable_gradient_checkpointing: true dtype: bf16 generating_args: - max_new_tokens: 32 # single-turn response length + max_new_tokens: 1024 # single-turn response length top_p: 0.99 top_k: 100 num_beams: 1 @@ -95,7 +96,7 @@ actor_infer: strategy_args: strategy_name: vllm strategy_config: - gpu_memory_utilization: 0.6 + gpu_memory_utilization: 0.8 block_size: 16 load_format: auto device_mapping: list(range(0,8)) @@ -103,7 +104,7 @@ actor_infer: reference: model_args: - flash_attn: fa2 + attn_implementation: fa2 disable_gradient_checkpointing: true dtype: bf16 model_type: ~ @@ -115,96 +116,25 @@ reference: device_mapping: list(range(0,8)) infer_batch_size: 1 -enable_response_mask: True -action_sep: "||" -use_turn_scores: False # important to GAE when applying token-level rewards to token-level advantages. If False, will take the sum of scores as the reward for the last turn. -enable_think: False # False -> no think RL -max_actions_per_traj: 20 reward_normalization: - grouping: tags # 可以tags(env_type)/traj_group_id(group)/batch(rollout_batch)... group_by计算reward/adv - method: identity # asym_clip / identity / mean_std - -custom_envs: - SimpleSokoban: - env_type: sokoban - max_actions_per_traj: ${max_actions_per_traj} # used in environment state manager to control the actual max actions executed per trajectory - max_steps_per_traj: ${max_actions_per_traj} - env_instruction: "You are solving the Sokoban puzzle. You are the player and you need to push all boxes to targets. When you are right next to a box, you can push it by moving in the same direction. You cannot push a box through a wall, and you cannot pull a box. The answer must be one of action in a turn, format is Right" - max_tokens: 100 # used to curate llm prompt "max words", not used for rollout - env_config: # keys should be a subset of SokobanConfig - dim_x: 6 - dim_y: 6 - num_boxes: 1 - max_steps: ${max_actions_per_traj} - LargerSokoban: - env_type: sokoban - max_actions_per_traj: ${max_actions_per_traj} - max_steps_per_traj: ${max_actions_per_traj} - env_instruction: "You are solving the Sokoban puzzle. You are the player and you need to push all boxes to targets. When you are right next to a box, you can push it by moving in the same direction. You cannot push a box through a wall, and you cannot pull a box. The answer must be one of action in a turn, format is Right" - max_tokens: 100 - env_config: - dim_x: 8 - dim_y: 8 - num_boxes: 2 - max_steps: ${max_actions_per_traj} - search_depth: 10 - SokobanDifferentGridVocab: - env_type: sokoban - max_actions_per_traj: ${max_actions_per_traj} - max_steps_per_traj: ${max_actions_per_traj} - env_instruction: "You are solving the Sokoban puzzle. You are the player and you need to push all boxes to targets. When you are right next to a box, you can push it by moving in the same direction. You cannot push a box through a wall, and you cannot pull a box. The answer must be one of action in a turn, format is Right" - max_tokens: 100 - env_config: # keys should be a subset of SokobanConfig - search_depth: 30 - dim_x: 6 - dim_y: 6 - num_boxes: 1 - max_steps: ${max_actions_per_traj} - grid_lookup: { 0: "W", 1: ".", 2: "G", 3: "C", 4: "B", 5: "A", 6: "@" } - grid_vocab: { "W": "wall", ".": "empty", "G": "target", "C": "box on target", "B": "box", "A": "player", "@": "player on target" } - FrozenLake: - env_type: frozen_lake - max_actions_per_traj: ${max_actions_per_traj} - max_steps_per_traj: ${max_actions_per_traj} - env_instruction: "You are solving the FrozenLake puzzle. Forbid the whole and go to the target. You may move to the unintended direction due to the slippery ice. The answer must be one of action in a turn, format is Right" - max_tokens: 100 - env_config: - is_slippery: false - WebShopEnv: - env_type: webshop - max_actions_per_traj: ${max_actions_per_traj} - max_steps_per_traj: ${max_actions_per_traj} - env_instruction: | - You are web shopping. - I will give you instructions about what to do. - You have to follow the instructions. - Every round I will give you an observation and a list of available actions, you have to respond an action based on the state and instruction. - You can use search action if search is available. - You can click one of the buttons in clickables. - An action should be of the following structure: - search[keywords] - click[value] - If the action is not valid, perform nothing. - Keywords in search are up to you, but the value in click must be a value in the list of available actions. - Remember that your keywords in search should be carefully designed. - Your response should use the following format: - - Thought: I think ... - Action: click[something] - max_tokens: 1024 - env_config: - observation_mode: text + grouping: traj_group_id # 可以tags(env_type)/traj_group_id(group)/batch(rollout_batch)... group_by计算reward/adv + method: mean_std # asym_clip / identity / mean_std train_env_manager: format_penalty: -0.05 - env_groups: 64 - group_size: 1 + num_env_groups: 8 + group_size: 8 + max_env_num_per_worker: 1 # The max_env_num_per_worker must be set to 1 to avoid conflicts with the webshop simple server. tags: [WebShopEnv] - n_groups: [64] # If not set, all env names divide nums equally. Under the same group, the env config and env seed (prompt) are equal in each generation + num_groups_partition: [8] # If not set, all env names divide nums equally. Under the same group, the env config and env seed (prompt) are equal in each generation val_env_manager: - env_groups: 64 + num_env_groups: 64 group_size: 1 # should be set to 1 because val temperature is set to 0 and same prompt leads to same output + max_env_num_per_worker: 1 # The max_env_num_per_worker must be set to 1 to avoid conflicts with the webshop simple server. tags: [WebShopEnv] - n_groups: [64] # TODO: If not set, all env names divide nums equally. Under the same group, the env config and env seed (prompt) are equal in each generation + num_groups_partition: [64] # TODO: If not set, all env names divide nums equally. Under the same group, the env config and env seed (prompt) are equal in each generation +custom_envs: + WebShopEnv: + ${custom_env.WebShopEnv} diff --git a/examples/qwen2.5-7B-agentic_megatron/agentic_val_webshop_gigpo.yaml b/examples/qwen2.5-7B-agentic_megatron/agentic_val_webshop_gigpo.yaml index d4883b34b..937908c93 100644 --- a/examples/qwen2.5-7B-agentic_megatron/agentic_val_webshop_gigpo.yaml +++ b/examples/qwen2.5-7B-agentic_megatron/agentic_val_webshop_gigpo.yaml @@ -1,9 +1,5 @@ defaults: - ../config/step_envs@_here_ - - ../config/deepspeed_zero@_here_ - - ../config/deepspeed_zero2@_here_ - - ../config/deepspeed_zero3@_here_ - - ../config/deepspeed_zero3_cpuoffload@_here_ hydra: run: diff --git a/examples/qwen2.5-7B-rlvr_megatron/rlvr_zero3_sp2.yaml b/examples/qwen2.5-7B-rlvr_megatron/rlvr_fsdp2_sp2.yaml similarity index 93% rename from examples/qwen2.5-7B-rlvr_megatron/rlvr_zero3_sp2.yaml rename to examples/qwen2.5-7B-rlvr_megatron/rlvr_fsdp2_sp2.yaml index d60fef844..7170431fe 100644 --- a/examples/qwen2.5-7B-rlvr_megatron/rlvr_zero3_sp2.yaml +++ b/examples/qwen2.5-7B-rlvr_megatron/rlvr_fsdp2_sp2.yaml @@ -1,9 +1,3 @@ -defaults: - - ../config/deepspeed_zero@_here_ - - ../config/deepspeed_zero2@_here_ - - ../config/deepspeed_zero3@_here_ - - ../config/deepspeed_zero3_cpuoffload@_here_ - hydra: run: dir: . @@ -132,8 +126,13 @@ actor_train: interleave_probs: "1.0" preprocessing_num_workers: 16 strategy_args: - strategy_name: deepspeed_train - strategy_config: ${deepspeed_zero3} + strategy_name: fsdp2_train + strategy_config: + fsdp_size: 4 + param_dtype: bf16 + reduce_dtype: float32 + reshard_after_forward: true + offload_policy: false device_mapping: list(range(0,8)) infer_batch_size: 1 @@ -170,8 +169,13 @@ reference: data_args: template: qwen2_5 strategy_args: - strategy_name: deepspeed_infer - strategy_config: ${deepspeed_zero3} + strategy_name: fsdp2_infer + strategy_config: + fsdp_size: 4 + param_dtype: bf16 + reduce_dtype: float32 + reshard_after_forward: true + offload_policy: true device_mapping: list(range(0,8)) infer_batch_size: 4 diff --git a/examples/qwen2.5-7B-rlvr_megatron/rlvr_lora_zero3.yaml b/examples/qwen2.5-7B-rlvr_megatron/rlvr_lora_fsdp2.yaml similarity index 92% rename from examples/qwen2.5-7B-rlvr_megatron/rlvr_lora_zero3.yaml rename to examples/qwen2.5-7B-rlvr_megatron/rlvr_lora_fsdp2.yaml index 97a45bc15..d51ee1dcf 100644 --- a/examples/qwen2.5-7B-rlvr_megatron/rlvr_lora_zero3.yaml +++ b/examples/qwen2.5-7B-rlvr_megatron/rlvr_lora_fsdp2.yaml @@ -1,9 +1,3 @@ -defaults: - - ../config/deepspeed_zero@_here_ - - ../config/deepspeed_zero2@_here_ - - ../config/deepspeed_zero3@_here_ - - ../config/deepspeed_zero3_cpuoffload@_here_ - hydra: run: dir: . @@ -87,8 +81,8 @@ lora_alpha: 32 # max_running_requests: 256 # is_num_return_sequences_expand: false -pretrain: /data/cpfs_0/common/models/Qwen2.5-7B -reward_pretrain: /data/cpfs_0/common/models/Qwen2.5-7B +pretrain: Qwen/Qwen2.5-7B +reward_pretrain: Qwen/Qwen2.5-7B validation: data_args: @@ -107,7 +101,7 @@ validation: actor_train: model_args: attn_implementation: fa2 - # Recomputed tensor size does not match for LoRA with Zero3 when activating checkpointing, See https://github.com/huggingface/transformers/issues/34928 for details + # Recomputed tensor size does not match for LoRA with FSDP when activating checkpointing, See https://github.com/huggingface/transformers/issues/34928 for details disable_gradient_checkpointing: true dtype: bf16 lora_target: ${lora_target} @@ -140,8 +134,13 @@ actor_train: interleave_probs: "1.0" preprocessing_num_workers: 16 strategy_args: - strategy_name: deepspeed_train - strategy_config: ${deepspeed_zero3} + strategy_name: fsdp2_train + strategy_config: + fsdp_size: 16 + param_dtype: bf16 + reduce_dtype: float32 + reshard_after_forward: true + offload_policy: false device_mapping: list(range(0,16)) infer_batch_size: 4 diff --git a/examples/qwen2.5-vl-3B-agentic/agentic_val_sokoban.yaml b/examples/qwen2.5-vl-3B-agentic/agentic_val_sokoban.yaml index 04db87a0a..c289903fd 100644 --- a/examples/qwen2.5-vl-3B-agentic/agentic_val_sokoban.yaml +++ b/examples/qwen2.5-vl-3B-agentic/agentic_val_sokoban.yaml @@ -1,9 +1,5 @@ defaults: - ../config/vl_traj_envs@_here_ - - ../config/deepspeed_zero@_here_ - - ../config/deepspeed_zero2@_here_ - - ../config/deepspeed_zero3@_here_ - - ../config/deepspeed_zero3_cpuoffload@_here_ hydra: run: @@ -78,8 +74,6 @@ actor_train: data_args: template: qwen2_5 strategy_args: -# strategy_name: deepspeed_train -# strategy_config: ${deepspeed_zero3} strategy_name: megatron_train strategy_config: tensor_model_parallel_size: 1 diff --git a/examples/qwen2.5-vl-7B-distill/distill_vl_zero3.yaml b/examples/qwen2.5-vl-7B-distill/distill_vl_fsdp2.yaml similarity index 79% rename from examples/qwen2.5-vl-7B-distill/distill_vl_zero3.yaml rename to examples/qwen2.5-vl-7B-distill/distill_vl_fsdp2.yaml index b81355bbc..f47ab57db 100644 --- a/examples/qwen2.5-vl-7B-distill/distill_vl_zero3.yaml +++ b/examples/qwen2.5-vl-7B-distill/distill_vl_fsdp2.yaml @@ -1,15 +1,9 @@ -defaults: - - ../config/deepspeed_zero@_here_ - - ../config/deepspeed_zero2@_here_ - - ../config/deepspeed_zero3@_here_ - - ../config/deepspeed_zero3_cpuoffload@_here_ - hydra: run: dir: . output_subdir: null -exp_name: "distill_vl_zero3" +exp_name: "distill_vl_fsdp2" seed: 42 logging_dir: ./output/logs output_dir: ./output @@ -18,7 +12,6 @@ checkpoint_config: type: file_system output_dir: /data/cpfs_0/rl_examples/models/${exp_name} - save_steps: 100 logging_steps: 1 resume_from_checkpoint: false @@ -60,8 +53,13 @@ student: preprocessing_num_workers: 16 strategy_args: - strategy_name: deepspeed_train - strategy_config: ${deepspeed_zero3} + strategy_name: fsdp2_train + strategy_config: + fsdp_size: 8 + param_dtype: bf16 + reduce_dtype: float32 + reshard_after_forward: true + offload_policy: false device_mapping: list(range(0,8)) teacher: @@ -75,8 +73,13 @@ teacher: # teacher forward micro_batch_size per_device_train_batch_size: 1 strategy_args: - strategy_name: deepspeed_infer - strategy_config: ${deepspeed_zero3} + strategy_name: fsdp2_infer + strategy_config: + fsdp_size: 8 + param_dtype: bf16 + reduce_dtype: float32 + reshard_after_forward: true + offload_policy: true device_mapping: list(range(0,8)) system_envs: diff --git a/examples/qwen2.5-vl-7B-rlvr/rlvr_async.yaml b/examples/qwen2.5-vl-7B-rlvr/rlvr_async.yaml index 8bf5e1bc7..83a2b9695 100644 --- a/examples/qwen2.5-vl-7B-rlvr/rlvr_async.yaml +++ b/examples/qwen2.5-vl-7B-rlvr/rlvr_async.yaml @@ -1,8 +1,3 @@ -defaults: - - ../config/deepspeed_zero@_here_ - - ../config/deepspeed_zero2@_here_ - - ../config/deepspeed_zero3@_here_ - - ../config/deepspeed_zero3_cpuoffload@_here_ hydra: run: dir: . @@ -43,6 +38,8 @@ validation: - /data/oss_bucket_0/yuzhao/data/One-RL-to-See-Them-All/Orsta-Data-47k/test/test_math_megabench_237.parquet - /data/oss_bucket_0/yuzhao/data/One-RL-to-See-Them-All/Orsta-Data-47k/test/test_detection_coco_test_multi_2000.parquet dataset_dir: ./ + preprocessing_num_workers: 32 + image_max_pixels: 1048576 # 1024 * 1024 generating_args: max_new_tokens: ${response_length} top_p: 0.99 @@ -77,6 +74,7 @@ actor_train: dataset_dir: ./ messages: prompt preprocessing_num_workers: 32 + image_max_pixels: 1048576 # 1024 * 1024 strategy_args: strategy_name: megatron_train strategy_config: diff --git a/examples/qwen2.5-vl-7B-rlvr/rlvr_megatron.yaml b/examples/qwen2.5-vl-7B-rlvr/rlvr_megatron.yaml index 804ee3bd2..c76cd20e4 100644 --- a/examples/qwen2.5-vl-7B-rlvr/rlvr_megatron.yaml +++ b/examples/qwen2.5-vl-7B-rlvr/rlvr_megatron.yaml @@ -1,9 +1,3 @@ -defaults: - - ../config/deepspeed_zero@_here_ - - ../config/deepspeed_zero2@_here_ - - ../config/deepspeed_zero3@_here_ - - ../config/deepspeed_zero3_cpuoffload@_here_ - hydra: run: dir: . @@ -51,6 +45,8 @@ validation: - /data/oss_bucket_0/yuzhao/data/One-RL-to-See-Them-All/Orsta-Data-47k/test/test_math_megabench_237.parquet - /data/oss_bucket_0/yuzhao/data/One-RL-to-See-Them-All/Orsta-Data-47k/test/test_detection_coco_test_multi_2000.parquet dataset_dir: ./ + preprocessing_num_workers: 32 + image_max_pixels: 1048576 # 1024 * 1024 generating_args: max_new_tokens: ${response_length} top_p: 0.99 @@ -86,6 +82,7 @@ actor_train: dataset_dir: ./ messages: prompt preprocessing_num_workers: 32 + image_max_pixels: 1048576 # 1024 * 1024 strategy_args: strategy_name: megatron_train strategy_config: diff --git a/examples/qwen2_5-vl-3B-rlvr_fsdp2/rlvr_fsdp2_offload.yaml b/examples/qwen2_5-vl-3B-rlvr_fsdp2/rlvr_fsdp2_offload.yaml new file mode 100644 index 000000000..1dff57f38 --- /dev/null +++ b/examples/qwen2_5-vl-3B-rlvr_fsdp2/rlvr_fsdp2_offload.yaml @@ -0,0 +1,159 @@ +hydra: + run: + dir: . + output_subdir: null + +exp_name: "qwen2_5_vl_3B_rlvr" +seed: 42 +logging_dir: ./output/logs +output_dir: ./output + +checkpoint_config: + type: file_system + output_dir: ./output + +# track_with: ml_tracker +# tracker_kwargs: +# project: pumpkin_fsdp2_test +# notes: pumpkin_fsdp2 +# tags: +# - rlvr +# - baseline +# - dp2 + +track_with: tensorboard +tracker_kwargs: + log_dir: ./tf_logs + +save_steps: 20 +logging_steps: 1 +eval_steps: 10 +resume_from_checkpoint: false + +rollout_batch_size: 256 +num_return_sequences_in_group: 8 +is_num_return_sequences_expand: true +prompt_length: 2048 +response_length: 4096 + +ppo_epochs: 1 +value_clip: 0.5 +reward_clip: 10 +advantage_clip: 10.0 +whiten_advantages: false +init_kl_coef: 0.0 +adv_estimator: "grpo" +use_kl_loss: true +kl_loss_coef: 1.0e-2 + +pretrain: Qwen/Qwen2.5-VL-3B-Instruct + +validation: + data_args: + file_name: + - /data/oss_bucket_0/pumpkin/data/train/train_detection_v3det_4000.parquet + - /data/oss_bucket_0/pumpkin/data/train/train_math_mmmath_3539.parquet + dataset_dir: ./ + generating_args: + max_new_tokens: ${response_length} + top_p: 0.99 + top_k: 100 + num_beams: 1 + temperature: 0.99 + num_return_sequences: 1 + eval_steps: ${eval_steps} + +actor_train: + model_args: + flash_attn: fa2 + attn_implementation: fa2 + model_name_or_path: ${pretrain} + disable_gradient_checkpointing: false + dtype: bf16 + model_type: ~ + freeze_module_prefix: model.visual + training_args: + learning_rate: 1.0e-6 + weight_decay: 1.0e-2 + per_device_train_batch_size: 2 + gradient_accumulation_steps: 64 + warmup_steps: 0 + num_train_epochs: 50 + data_args: + # use One-RL-to-See-Them-All/Orsta-Data-47k as train dataset + # download from https://huggingface.co/datasets/One-RL-to-See-Them-All/Orsta-Data-47k + file_name: + - /data/oss_bucket_0/pumpkin/data/train/train_detection_v3det_4000.parquet + - /data/oss_bucket_0/pumpkin/data/train/train_math_mmmath_3539.parquet + domain_interleave_probs: + math: 0.5 + cv_detection: 0.5 + dataset_dir: ./ + messages: prompt + preprocessing_num_workers: 32 + strategy_args: + strategy_name: fsdp2_train + strategy_config: + fsdp_size: 8 + param_dtype: bf16 + reduce_dtype: float32 + reshard_after_forward: true + offload_policy: false # offload model param + use_batched_model_update: false # disable batched gather during model update (not bring any benefit in test) + device_mapping: list(range(0,8)) + infer_batch_size: 2 + +actor_infer: + model_args: + flash_attn: fa2 + attn_implementation: fa2 + model_name_or_path: ${pretrain} + disable_gradient_checkpointing: true + dtype: bf16 + generating_args: + max_new_tokens: ${response_length} + top_p: 0.99 + top_k: 100 + num_beams: 1 + temperature: 0.99 + num_return_sequences: ${num_return_sequences_in_group} + strategy_args: + strategy_name: vllm + strategy_config: + gpu_memory_utilization: 0.8 + block_size: 16 + num_gpus_per_worker: 1 + device_mapping: list(range(0,8)) + infer_batch_size: 32 + +reference: + model_args: + flash_attn: fa2 + attn_implementation: fa2 + model_name_or_path: ${pretrain} + disable_gradient_checkpointing: true + dtype: bf16 + model_type: ~ + strategy_args: + strategy_name: hf_infer + strategy_config: ~ + device_mapping: list(range(0,8)) + infer_batch_size: 2 + +rewards: + math: + worker_cls: roll.pipeline.rlvr.rewards.math_rule_reward_worker.MathRuleRewardWorker + model_args: + model_name_or_path: ${pretrain} + # data source whose ability is math in One-RL-to-See-Them-All/Orsta-Data-47k + tag_included: [mm_math, megabench_math] + world_size: 8 + infer_batch_size: 1 + cv_detection: + worker_cls: roll.pipeline.rlvr.rewards.detection_reward_worker.DetectionRewardWorker + model_args: + model_name_or_path: ${pretrain} + # data source whose ability is cv_detection in One-RL-to-See-Them-All/Orsta-Data-47k + tag_included: [v3det_train, object365_train, coco_val_multi_test] + world_size: 8 + infer_batch_size: 1 diff --git a/examples/qwen2_5-vl-3B-rlvr_fsdp2/run_rlvr_pipeline.sh b/examples/qwen2_5-vl-3B-rlvr_fsdp2/run_rlvr_pipeline.sh new file mode 100755 index 000000000..7c7e9db72 --- /dev/null +++ b/examples/qwen2_5-vl-3B-rlvr_fsdp2/run_rlvr_pipeline.sh @@ -0,0 +1,5 @@ +#!/bin/bash +set +x + +CONFIG_PATH=$(basename $(dirname $0)) +python examples/start_rlvr_pipeline.py --config_path $CONFIG_PATH --config_name rlvr_config diff --git a/examples/qwen3-235BA22B-rlvr_megatron/rlvr_config_140GB.yaml b/examples/qwen3-235BA22B-rlvr_megatron/rlvr_config_140GB.yaml new file mode 100644 index 000000000..4bc77a638 --- /dev/null +++ b/examples/qwen3-235BA22B-rlvr_megatron/rlvr_config_140GB.yaml @@ -0,0 +1,235 @@ +hydra: + run: + dir: . + output_subdir: null + +exp_name: "qwen3-235BA22B-rlvr-config" +seed: 42 +logging_dir: ./output/logs +output_dir: ./output +system_envs: + USE_MODELSCOPE: '1' + +checkpoint_config: + type: file_system + output_dir: ./rl_examples/models/${exp_name} + +track_with: tensorboard +tracker_kwargs: + log_dir: ./rl_examples/llm/tensorboard/roll_exp/rlvr + +num_gpus_per_node: 8 + +max_steps: 500 +save_steps: 100 +logging_steps: 1 +eval_steps: 10 +resume_from_checkpoint: false + +offload_nccl: false # save GPU mem + + +rollout_batch_size: 32 # prompt +prompt_length: 2048 +response_length: 4096 + +num_return_sequences_in_group: 8 +ppo_epochs: 1 +adv_estimator: "reinforce" + +# clip +value_clip: 0.5 +reward_clip: 10 +advantage_clip: 2.0 +dual_clip_loss: true + +# normalize +norm_mean_type: ~ +norm_std_type: ~ + +# data mask +max_len_mask: true +difficulty_mask: true +difficulty_low_threshold: 0.1 +difficulty_high_threshold: 0.95 +error_max_len_clip: false + +# data weight +difficulty_loss_weight: false +length_loss_weight: false + +# reward +add_token_level_kl: false + +# advantage +whiten_advantages: true + +# dynamic sampling scheduler +# use_additional_prompts: true +# max_running_requests: 256 +# is_num_return_sequences_expand: false + +pretrain: Qwen/Qwen3-235B-A22B +reward_pretrain: Qwen/Qwen3-235B-A22B + +validation: + data_args: + template: qwen3 + file_name: + - data/math_benchmarks.jsonl + generating_args: + top_p: 0.6 + top_k: 50 + num_beams: 1 + temperature: 0.6 + num_return_sequences: 1 + eval_steps: 10 + +actor_train: + model_args: + disable_gradient_checkpointing: false + dtype: bf16 + model_type: ~ + training_args: + learning_rate: 1.0e-6 + weight_decay: 0 + per_device_train_batch_size: 1 + gradient_accumulation_steps: 64 + warmup_steps: 20 + num_train_epochs: 50 + data_args: + template: qwen3 + file_name: + - data/code_KodCode_data.jsonl + - data/math_deepmath_deal.jsonl + - data/general_ifeval_train_deal.jsonl + - data/general_CrossThink-QA_deal.jsonl + domain_interleave_probs: + math_rule: 0.5 + code_sandbox: 0.3 + # llm_judge: 0.1 + crossthinkqa: 0.1 + ifeval: 0.1 + dataset_dir: data + messages: messages + interleave_probs: "1.0" + preprocessing_num_workers: 16 + strategy_args: + strategy_name: megatron_train + strategy_config: + tensor_model_parallel_size: 2 # todo: support ETP + pipeline_model_parallel_size: 8 + virtual_pipeline_model_parallel_size: 6 + expert_model_parallel_size: 4 + context_parallel_size: 1 + account_for_loss_in_pipeline_split: true + account_for_embedding_in_pipeline_split: true + use_distributed_optimizer: true + sequence_parallel: true + overlap_grad_reduce: true + recompute_granularity: selective + recompute_modules: moe,layernorm + bias_activation_fusion: true + apply_rope_fusion: true + moe_grouped_gemm: true + moe_token_dispatcher_type: "alltoall" + device_mapping: list(range(0,64)) + infer_batch_size: 2 + offload_nccl: ${offload_nccl} + +actor_infer: + model_args: + disable_gradient_checkpointing: true + dtype: bf16 + generating_args: + max_new_tokens: ${response_length} + top_p: 0.99 + top_k: 100 + num_beams: 1 + temperature: 0.99 + num_return_sequences: ${num_return_sequences_in_group} + data_args: + template: qwen3 + strategy_args: + strategy_name: vllm + strategy_config: + gpu_memory_utilization: 0.8 + load_format: dummy + tensor_parallel_size: 8 + cuda_graph_sizes: [1,2,4,8,16,24,32,40,48] + num_gpus_per_worker: 8 + device_mapping: list(range(0,64)) # device share with llm reward + infer_batch_size: 1 + +reference: + model_args: + dtype: bf16 + model_type: ~ + data_args: + template: qwen3 + strategy_args: + strategy_name: megatron_infer + strategy_config: + tensor_model_parallel_size: 1 + pipeline_model_parallel_size: 4 + virtual_pipeline_model_parallel_size: 6 + expert_model_parallel_size: 8 + account_for_loss_in_pipeline_split: true + account_for_embedding_in_pipeline_split: true + use_distributed_optimizer: true + sequence_parallel: true + bias_activation_fusion: true + apply_rope_fusion: true + moe_grouped_gemm: true + moe_token_dispatcher_type: "alltoall" + device_mapping: list(range(0,64)) + infer_batch_size: 2 + offload_nccl: ${offload_nccl} + +rewards: + crossthinkqa: + worker_cls: roll.pipeline.rlvr.rewards.crossthinkqa_rule_reward_worker.CrossThinkQARuleRewardWorker + reward_type: soft + response_length_penalty_coef: 0.0 + model_args: + model_name_or_path: ${reward_pretrain} + data_args: + template: qwen3 + tag_included: [crossthinkqa] + world_size: 8 + infer_batch_size: 4 + ifeval: + worker_cls: roll.pipeline.rlvr.rewards.ifeval_rule_reward_worker.GeneralRuleRewardWorker + reward_type: soft + model_args: + model_name_or_path: ${reward_pretrain} + data_args: + template: qwen3 + tag_included: [ifeval] + world_size: 8 + infer_batch_size: 4 + math_rule: + worker_cls: roll.pipeline.rlvr.rewards.math_rule_reward_worker.MathRuleRewardWorker + model_args: + model_name_or_path: ${reward_pretrain} + data_args: + template: qwen3 + tag_included: [deepmath_103k, aime] + world_size: 8 + infer_batch_size: 1 +# dynamic filter config +# query_filter_config: +# type: mean_filter +# filter_args: +# threshold_up: 0.9 +# threshold_down: 0.1 + code_sandbox: + use_local: true + worker_cls: roll.pipeline.rlvr.rewards.code_sandbox_reward_worker.CodeSandboxRewardWorker + tag_included: [KodCode] + model_args: + model_name_or_path: ${reward_pretrain} + data_args: + template: qwen3 + world_size: 8 + infer_batch_size: 1 \ No newline at end of file diff --git a/examples/qwen3-235BA22B-rlvr_megatron/run_rlvr_pipeline.sh b/examples/qwen3-235BA22B-rlvr_megatron/run_rlvr_pipeline.sh index 7c7e9db72..5310648e8 100755 --- a/examples/qwen3-235BA22B-rlvr_megatron/run_rlvr_pipeline.sh +++ b/examples/qwen3-235BA22B-rlvr_megatron/run_rlvr_pipeline.sh @@ -2,4 +2,4 @@ set +x CONFIG_PATH=$(basename $(dirname $0)) -python examples/start_rlvr_pipeline.py --config_path $CONFIG_PATH --config_name rlvr_config +python examples/start_rlvr_pipeline.py --config_path $CONFIG_PATH --config_name rlvr_config_140GB diff --git a/examples/qwen3-30BA3B-agentic_fsdp2/agentic_val_sokoban_30a3.yaml b/examples/qwen3-30BA3B-agentic_fsdp2/agentic_val_sokoban_30a3.yaml index ce30a9a83..9cf2fdc9a 100644 --- a/examples/qwen3-30BA3B-agentic_fsdp2/agentic_val_sokoban_30a3.yaml +++ b/examples/qwen3-30BA3B-agentic_fsdp2/agentic_val_sokoban_30a3.yaml @@ -1,9 +1,5 @@ defaults: - ../config/traj_envs@_here_ - - ../config/deepspeed_zero@_here_ - - ../config/deepspeed_zero2@_here_ - - ../config/deepspeed_zero3@_here_ - - ../config/deepspeed_zero3_cpuoffload@_here_ hydra: run: diff --git a/examples/qwen3-30BA3B-rlvr_fsdp2/rlvr_config.yaml b/examples/qwen3-30BA3B-rlvr_fsdp2/rlvr_config.yaml new file mode 100644 index 000000000..464e227b1 --- /dev/null +++ b/examples/qwen3-30BA3B-rlvr_fsdp2/rlvr_config.yaml @@ -0,0 +1,185 @@ +hydra: + run: + dir: . + output_subdir: null + +exp_name: "qwen3-30BA3B-rlvr-config" +seed: 42 +logging_dir: ./output/logs +output_dir: ./output +system_envs: + USE_MODELSCOPE: '1' + +checkpoint_config: + type: file_system + output_dir: ./rl_examples/models/${exp_name} + +num_gpus_per_node: 8 + +max_steps: 500 +save_steps: 100 +logging_steps: 1 +eval_steps: 50 +resume_from_checkpoint: false + +rollout_batch_size: 16 # prompt +prompt_length: 2048 +response_length: 4096 + +num_return_sequences_in_group: 8 +ppo_epochs: 1 +adv_estimator: "reinforce" + +# clip +value_clip: 0.5 +reward_clip: 10 +advantage_clip: 2.0 +dual_clip_loss: true + +# normalize +norm_mean_type: ~ +norm_std_type: ~ + +# data mask +max_len_mask: true +difficulty_mask: true +difficulty_low_threshold: 0.1 +difficulty_high_threshold: 0.95 +error_max_len_clip: false + +# data weight +difficulty_loss_weight: false +length_loss_weight: false + +# reward +add_token_level_kl: false + +# advantage +whiten_advantages: true + +# dynamic sampling scheduler +# use_additional_prompts: true +# max_running_requests: 256 +# is_num_return_sequences_expand: false + +pretrain: Qwen/Qwen3-30B-A3B-Base +reward_pretrain: Qwen/Qwen3-30B-A3B-Base + +validation: + data_args: + template: qwen2_5 + file_name: + - data/math_benchmarks.jsonl + generating_args: + max_new_tokens: ${response_length} + top_p: 0.6 + top_k: 50 + num_beams: 1 + temperature: 0.6 + num_return_sequences: 1 + +actor_train: + model_args: + disable_gradient_checkpointing: false + dtype: bf16 + model_type: ~ + ulysses_size: 4 + training_args: + learning_rate: 1.0e-6 + weight_decay: 0 + per_device_train_batch_size: 1 + gradient_accumulation_steps: 4 + warmup_steps: 20 + data_args: + template: qwen2_5 + file_name: + - data/math_deepmath_deal.jsonl + domain_interleave_probs: + math_rule: 1.0 + dataset_dir: data + messages: messages + interleave_probs: "1.0" + preprocessing_num_workers: 16 + strategy_args: + strategy_name: fsdp2_train + strategy_config: + fsdp_size: 32 + param_dtype: bf16 + reduce_dtype: bf16 + offload_policy: false + reshard_after_forward: true + wrap_policy: + wrap_embeddings: true + wrap_lm_output: true + moe_experts: + - Qwen3MoeMLP + transformer_layer_cls_to_wrap: + - Qwen3MoeAttention + - Qwen3MoeSparseMoeBlock + use_remove_padding: false + device_mapping: list(range(0,32)) + infer_batch_size: 2 + +actor_infer: + model_args: + disable_gradient_checkpointing: true + dtype: bf16 + generating_args: + max_new_tokens: ${response_length} + top_p: 0.99 + top_k: 100 + num_beams: 1 + temperature: 0.99 + num_return_sequences: ${num_return_sequences_in_group} + data_args: + template: qwen2_5 + strategy_args: + strategy_name: vllm + strategy_config: + gpu_memory_utilization: 0.5 + load_format: dummy + tensor_parallel_size: 2 + sleep_level: 2 # 2 will destroy model parameter and kv_cache after generate to save cpu memory, 1 will destroy kv_cache only. + optimization_level: 1 + enforce_eager: true + num_gpus_per_worker: 2 + device_mapping: list(range(0,32)) + infer_batch_size: 1 + +reference: + model_args: + disable_gradient_checkpointing: true + dtype: bf16 + model_type: ~ + ulysses_size: 2 + data_args: + template: qwen2_5 + strategy_args: + strategy_name: fsdp2_infer + strategy_config: + fsdp_size: 32 + param_dtype: bf16 + reduce_dtype: bf16 + reshard_after_forward: true + offload_policy: false + wrap_policy: + wrap_embeddings: true + wrap_lm_output: true + moe_experts: + - Qwen3MoeMLP + transformer_layer_cls_to_wrap: + - Qwen3MoeAttention + - Qwen3MoeSparseMoeBlock + device_mapping: list(range(0,32)) + infer_batch_size: 2 + +rewards: + math_rule: + worker_cls: roll.pipeline.rlvr.rewards.math_rule_reward_worker.MathRuleRewardWorker + model_args: + model_name_or_path: ${reward_pretrain} + data_args: + template: qwen2_5 + tag_included: [deepmath_103k, aime] + world_size: 8 + infer_batch_size: 1 diff --git a/examples/qwen3-30BA3B-rlvr_fsdp2/rlvr_config_ep.yaml b/examples/qwen3-30BA3B-rlvr_fsdp2/rlvr_config_ep.yaml new file mode 100644 index 000000000..cd338c107 --- /dev/null +++ b/examples/qwen3-30BA3B-rlvr_fsdp2/rlvr_config_ep.yaml @@ -0,0 +1,191 @@ +hydra: + run: + dir: . + output_subdir: null + +exp_name: "qwen3-30BA3B-rlvr-config" +seed: 42 +logging_dir: ./output/logs +output_dir: ./output +system_envs: + USE_MODELSCOPE: '1' + +checkpoint_config: + type: file_system + output_dir: ./rl_examples/models/${exp_name} + +num_gpus_per_node: 8 + +max_steps: 500 +save_steps: 100 +logging_steps: 1 +eval_steps: 50 +resume_from_checkpoint: false + +rollout_batch_size: 16 # prompt +prompt_length: 2048 +response_length: 4096 + +num_return_sequences_in_group: 8 +ppo_epochs: 1 +adv_estimator: "reinforce" + +# clip +value_clip: 0.5 +reward_clip: 10 +advantage_clip: 2.0 +dual_clip_loss: true + +# normalize +norm_mean_type: ~ +norm_std_type: ~ + +# data mask +max_len_mask: true +difficulty_mask: true +difficulty_low_threshold: 0.1 +difficulty_high_threshold: 0.95 +error_max_len_clip: false + +# data weight +difficulty_loss_weight: false +length_loss_weight: false + +# reward +add_token_level_kl: false + +# advantage +whiten_advantages: true + +# dynamic sampling scheduler +# use_additional_prompts: true +# max_running_requests: 256 +# is_num_return_sequences_expand: false + +pretrain: Qwen/Qwen3-30B-A3B-Base +reward_pretrain: Qwen/Qwen3-30B-A3B-Base + +validation: + data_args: + template: qwen2_5 + file_name: + - data/math_benchmarks.jsonl + generating_args: + max_new_tokens: ${response_length} + top_p: 0.6 + top_k: 50 + num_beams: 1 + temperature: 0.6 + num_return_sequences: 1 + +actor_train: + model_args: + disable_gradient_checkpointing: false + dtype: bf16 + model_type: ~ + ulysses_size: 4 + training_args: + learning_rate: 1.0e-6 + weight_decay: 0 + per_device_train_batch_size: 1 + gradient_accumulation_steps: 4 + warmup_steps: 20 + data_args: + template: qwen2_5 + file_name: + - data/math_deepmath_deal.jsonl + domain_interleave_probs: + math_rule: 1.0 + dataset_dir: data + messages: messages + interleave_probs: "1.0" + preprocessing_num_workers: 16 + strategy_args: + strategy_name: fsdp2_train + strategy_config: + fsdp_size: 32 + efsdp_size: 4 + ep_size: 8 + moe_use_grouped_mm: false + param_dtype: bf16 + reduce_dtype: bf16 + offload_policy: false + reshard_after_forward: true + wrap_policy: + wrap_embeddings: true + wrap_lm_output: true + moe_experts: + - Qwen3MoeExperts + transformer_layer_cls_to_wrap: + - Qwen3MoeAttention + - Qwen3MoeSparseMoeBlock + use_remove_padding: false + device_mapping: list(range(0,32)) + infer_batch_size: 2 + +actor_infer: + model_args: + disable_gradient_checkpointing: true + dtype: bf16 + generating_args: + max_new_tokens: ${response_length} + top_p: 0.99 + top_k: 100 + num_beams: 1 + temperature: 0.99 + num_return_sequences: ${num_return_sequences_in_group} + data_args: + template: qwen2_5 + strategy_args: + strategy_name: vllm + strategy_config: + gpu_memory_utilization: 0.5 + load_format: dummy + tensor_parallel_size: 2 + sleep_level: 2 # 2 will destroy model parameter and kv_cache after generate to save cpu memory, 1 will destroy kv_cache only. + optimization_level: 1 + enforce_eager: true + num_gpus_per_worker: 2 + device_mapping: list(range(0,32)) + infer_batch_size: 1 + +reference: + model_args: + disable_gradient_checkpointing: true + dtype: bf16 + model_type: ~ + ulysses_size: 2 + data_args: + template: qwen2_5 + strategy_args: + strategy_name: fsdp2_infer + strategy_config: + fsdp_size: 32 + efsdp_size: 4 + ep_size: 8 + moe_use_grouped_mm: false + param_dtype: bf16 + reduce_dtype: bf16 + reshard_after_forward: true + offload_policy: false + wrap_policy: + wrap_embeddings: true + wrap_lm_output: true + moe_experts: + - Qwen3MoeExperts + transformer_layer_cls_to_wrap: + - Qwen3MoeAttention + - Qwen3MoeSparseMoeBlock + device_mapping: list(range(0,32)) + infer_batch_size: 2 + +rewards: + math_rule: + worker_cls: roll.pipeline.rlvr.rewards.math_rule_reward_worker.MathRuleRewardWorker + model_args: + model_name_or_path: ${reward_pretrain} + data_args: + template: qwen2_5 + tag_included: [deepmath_103k, aime] + world_size: 8 + infer_batch_size: 1 diff --git a/examples/qwen3-30BA3B-rlvr_megatron/rlvr_config_ppu_sglang.yaml b/examples/qwen3-30BA3B-rlvr_megatron/rlvr_config_ppu_sglang.yaml new file mode 100644 index 000000000..3c364774e --- /dev/null +++ b/examples/qwen3-30BA3B-rlvr_megatron/rlvr_config_ppu_sglang.yaml @@ -0,0 +1,230 @@ +hydra: + run: + dir: . + output_subdir: null + +exp_name: "qwen3-30BA3B-rlvr-config" +seed: 42 +logging_dir: ./output/logs +output_dir: ./output +system_envs: + USE_MODELSCOPE: '1' + +checkpoint_config: + type: file_system + output_dir: ./rl_examples/models/${exp_name} + +#track_with: wandb +#tracker_kwargs: +# api_key: +# project: roll_examples +# notes: roll_examples +# tags: +# - rlvr +# - baseline + +track_with: tensorboard +tracker_kwargs: + log_dir: ./rl_examples/llm/tensorboard/roll_exp/rlvr + +num_gpus_per_node: 16 + +max_steps: 500 +save_steps: 100 +logging_steps: 1 +eval_steps: 10 +resume_from_checkpoint: false + + +rollout_batch_size: 64 # prompt +prompt_length: 2048 +response_length: 4096 + +num_return_sequences_in_group: 8 +ppo_epochs: 1 +adv_estimator: "reinforce" + +# clip +value_clip: 0.5 +reward_clip: 10 +advantage_clip: 2.0 +dual_clip_loss: true + +# normalize +norm_mean_type: ~ +norm_std_type: ~ + +# data mask +max_len_mask: true +difficulty_mask: true +difficulty_low_threshold: 0.1 +difficulty_high_threshold: 0.95 +error_max_len_clip: false + +# data weight +difficulty_loss_weight: false +length_loss_weight: false + +# reward +add_token_level_kl: false + +# advantage +whiten_advantages: true + +# dynamic sampling scheduler +# use_additional_prompts: true +# max_running_requests: 256 +# is_num_return_sequences_expand: false + +#pretrain: Qwen/Qwen3-30B-A3B-Base +#reward_pretrain: Qwen/Qwen3-30B-A3B-Base +pretrain: Qwen/Qwen3-30B-A3B +reward_pretrain: Qwen/Qwen3-30B-A3B +#pretrain: /data/cpfs_0/common/models/Qwen3-30B-A3B +#reward_pretrain: /data/cpfs_0/common/models/Qwen3-30B-A3B + + +validation: + data_args: + template: qwen2_5 + file_name: + - data/aime24_25_deal.jsonl + generating_args: + max_new_tokens: ${response_length} + top_p: 0.6 + top_k: 50 + num_beams: 1 + temperature: 0.6 + num_return_sequences: 1 + +actor_train: + model_args: + flash_attn: fa2 + disable_gradient_checkpointing: false + dtype: bf16 + model_type: ~ + training_args: + learning_rate: 1.0e-6 + weight_decay: 0 + per_device_train_batch_size: 1 + gradient_accumulation_steps: 32 + warmup_steps: 20 + num_train_epochs: 50 + data_args: + template: qwen2_5 + file_name: + - data/code_KodCode_data.jsonl + - data/llm_judge_Multi-subject-RLVR_deal_new.jsonl + - data/math_deepmath_deal.jsonl + - data/general_ifeval_train_deal.jsonl + - data/general_CrossThink-QA_deal.jsonl + domain_interleave_probs: + math_rule: 0.4 + code_sandbox: 0.3 +# llm_judge: 0.1 + crossthinkqa: 0.2 + ifeval: 0.1 + dataset_dir: data + messages: messages + interleave_probs: "1.0" + preprocessing_num_workers: 16 + strategy_args: + strategy_name: megatron_train + strategy_config: + tensor_model_parallel_size: 1 + pipeline_model_parallel_size: 2 + virtual_pipeline_model_parallel_size: 8 + expert_model_parallel_size: 4 + context_parallel_size: 4 + use_distributed_optimizer: true + sequence_parallel: true + moe_token_dispatcher_type: "alltoall" + moe_grouped_gemm: true + moe_layer_recompute: true + device_mapping: list(range(0,32)) + infer_batch_size: 2 + +actor_infer: + model_args: + flash_attn: fa2 + disable_gradient_checkpointing: true + dtype: bf16 + generating_args: + max_new_tokens: ${response_length} + top_p: 0.99 + top_k: 100 + num_beams: 1 + temperature: 0.99 + num_return_sequences: ${num_return_sequences_in_group} + data_args: + template: qwen2_5 + strategy_args: + strategy_name: sglang + strategy_config: + mem_fraction_static: 0.7 + load_format: dummy + enable_memory_saver: false + num_gpus_per_worker: 2 + device_mapping: list(range(32,48)) + infer_batch_size: 1 + +reference: + model_args: + flash_attn: fa2 + disable_gradient_checkpointing: true + dtype: bf16 + model_type: ~ + data_args: + template: qwen2_5 + strategy_args: + strategy_name: megatron_infer + strategy_config: + tensor_model_parallel_size: 1 + pipeline_model_parallel_size: 1 + expert_model_parallel_size: 4 + moe_token_dispatcher_type: "alltoall" + moe_grouped_gemm: true + device_mapping: list(range(48,64)) + infer_batch_size: 2 + +rewards: + crossthinkqa: + worker_cls: roll.pipeline.rlvr.rewards.crossthinkqa_rule_reward_worker.CrossThinkQARuleRewardWorker + reward_type: soft + response_length_penalty_coef: 0.0 + model_args: + model_name_or_path: ${reward_pretrain} + data_args: + template: qwen2_5 + tag_included: [crossthinkqa] + world_size: 8 + infer_batch_size: 4 + ifeval: + worker_cls: roll.pipeline.rlvr.rewards.ifeval_rule_reward_worker.GeneralRuleRewardWorker + reward_type: soft + model_args: + model_name_or_path: ${reward_pretrain} + data_args: + template: qwen2_5 + tag_included: [ifeval] + world_size: 8 + infer_batch_size: 4 + math_rule: + worker_cls: roll.pipeline.rlvr.rewards.math_rule_reward_worker.MathRuleRewardWorker + model_args: + model_name_or_path: ${reward_pretrain} + data_args: + template: qwen2_5 + tag_included: [deepmath_103k, aime] + world_size: 8 + infer_batch_size: 1 + code_sandbox: + use_local: true + worker_cls: roll.pipeline.rlvr.rewards.code_sandbox_reward_worker.CodeSandboxRewardWorker + tag_included: [KodCode] + model_args: + model_name_or_path: ${reward_pretrain} + data_args: + template: qwen2_5 + world_size: 8 + infer_batch_size: 1 \ No newline at end of file diff --git a/examples/qwen3-30BA3B-sft_fsdp2/sft_config.yaml b/examples/qwen3-30BA3B-sft_fsdp2/sft_config.yaml new file mode 100644 index 000000000..147721d18 --- /dev/null +++ b/examples/qwen3-30BA3B-sft_fsdp2/sft_config.yaml @@ -0,0 +1,66 @@ +hydra: + run: + dir: . + output_subdir: null + +exp_name: "qwen3-30BA3B-sft-config" +seed: 42 +logging_dir: ./output/logs +output_dir: ./output_sft +system_envs: + USE_MODELSCOPE: '1' + +num_gpus_per_node: 8 + +save_steps: 100 +logging_steps: 1 +eval_steps: 100 +resume_from_checkpoint: false + +sequence_length: 2048 + +pretrain: Qwen/Qwen3-30B-A3B-Base + +# sft related +# system_key: system_prompt # use the default system prompt in the tokenizer tmplate if not provided +prompt_key: instruction +query_key: input +response_key: output + +validation: + data_args: + file_name: ./data/code_alpaca_500.json + template: qwen2_5 + +sft_train: + model_args: + dtype: bf16 + ulysses_size: 2 + training_args: + num_train_epochs: 1 + per_device_train_batch_size: 1 + gradient_accumulation_steps: 1 + learning_rate: 5.0e-6 + data_args: + file_name: ./data/code_alpaca_20k.json # https://huggingface.co/datasets/sahil2801/CodeAlpaca-20k + template: qwen2_5 + preprocessing_num_workers: 4 + strategy_args: + strategy_name: fsdp2_train + strategy_config: + fsdp_size: 8 + param_dtype: bf16 + reduce_dtype: bf16 + offload_policy: false + reshard_after_forward: true + wrap_policy: + wrap_embeddings: true + wrap_lm_output: true + moe_experts: + - Qwen3MoeMLP + transformer_layer_cls_to_wrap: + - Qwen3MoeAttention + - Qwen3MoeSparseMoeBlock + use_sequence_packing: True + device_mapping: list(range(0,8)) + infer_batch_size: 2 diff --git a/examples/qwen3-30BA3B-sft_fsdp2/sft_config_ep.yaml b/examples/qwen3-30BA3B-sft_fsdp2/sft_config_ep.yaml new file mode 100644 index 000000000..dec48821a --- /dev/null +++ b/examples/qwen3-30BA3B-sft_fsdp2/sft_config_ep.yaml @@ -0,0 +1,69 @@ +hydra: + run: + dir: . + output_subdir: null + +exp_name: "qwen3-30BA3B-sft-config" +seed: 42 +logging_dir: ./output/logs +output_dir: ./output_sft +system_envs: + USE_MODELSCOPE: '1' + +num_gpus_per_node: 8 + +save_steps: 100 +logging_steps: 1 +eval_steps: 100 +resume_from_checkpoint: false + +sequence_length: 2048 + +pretrain: Qwen/Qwen3-30B-A3B-Base + +# sft related +# system_key: system_prompt # use the default system prompt in the tokenizer tmplate if not provided +prompt_key: instruction +query_key: input +response_key: output + +validation: + data_args: + file_name: ./data/code_alpaca_500.json + template: qwen2_5 + +sft_train: + model_args: + dtype: bf16 + ulysses_size: 2 + training_args: + num_train_epochs: 1 + per_device_train_batch_size: 1 + gradient_accumulation_steps: 1 + learning_rate: 5.0e-6 + data_args: + file_name: ./data/code_alpaca_20k.json # https://huggingface.co/datasets/sahil2801/CodeAlpaca-20k + template: qwen2_5 + preprocessing_num_workers: 4 + strategy_args: + strategy_name: fsdp2_train + strategy_config: + fsdp_size: 8 + efsdp_size: 1 + ep_size: 8 + moe_use_grouped_mm: false + param_dtype: bf16 + reduce_dtype: bf16 + offload_policy: false + reshard_after_forward: true + wrap_policy: + wrap_embeddings: true + wrap_lm_output: true + moe_experts: + - Qwen3MoeExperts + transformer_layer_cls_to_wrap: + - Qwen3MoeAttention + - Qwen3MoeSparseMoeBlock + use_sequence_packing: True + device_mapping: list(range(0,8)) + infer_batch_size: 2 diff --git a/examples/qwen3-30BA3B-sft_fsdp2/sft_config_megatron.yaml b/examples/qwen3-30BA3B-sft_fsdp2/sft_config_megatron.yaml new file mode 100644 index 000000000..89f6cb8c9 --- /dev/null +++ b/examples/qwen3-30BA3B-sft_fsdp2/sft_config_megatron.yaml @@ -0,0 +1,73 @@ +hydra: + run: + dir: . + output_subdir: null + +exp_name: "qwen3-30BA3B-sft-config" +seed: 42 +logging_dir: ./output/logs +output_dir: ./output_sft +system_envs: + USE_MODELSCOPE: '1' + +track_with: ml_tracker +tracker_kwargs: + # ml_tracker的实验名，自行设置，不要和别人重复，否则没有权限写入报错 + project: ljy_fsdp-ep-exp_qwen3-30BA3B + name: sft_megatron-pp2-ep4_bf16_4096 + notes: "scale aligner pipeline" + tags: # mltracker job tags，后续方便管理实验 + - pipeline + - roll + - Qwen3 + +num_gpus_per_node: 8 + +save_steps: 100 +logging_steps: 1 +eval_steps: 100 +resume_from_checkpoint: false + +sequence_length: 4096 + +pretrain: Qwen/Qwen3-30B-A3B-Base + +# sft related +# system_key: system_prompt # use the default system prompt in the tokenizer tmplate if not provided +prompt_key: instruction +query_key: input +response_key: output + +# validation: +# data_args: +# file_name: ./data/code_alpaca_500.json +# template: qwen2_5 + +sft_train: + model_args: + dtype: bf16 + training_args: + num_train_epochs: 1 + per_device_train_batch_size: 1 + gradient_accumulation_steps: 2 + learning_rate: 5.0e-6 + data_args: + file_name: ./data/code_alpaca_20k.json # https://huggingface.co/datasets/sahil2801/CodeAlpaca-20k + template: qwen2_5 + preprocessing_num_workers: 4 + strategy_args: + strategy_name: megatron_train + strategy_config: + tensor_model_parallel_size: 1 + pipeline_model_parallel_size: 2 + virtual_pipeline_model_parallel_size: 8 + expert_model_parallel_size: 4 + context_parallel_size: 1 + use_distributed_optimizer: true + sequence_parallel: true + moe_token_dispatcher_type: "alltoall" + moe_grouped_gemm: true + moe_layer_recompute: true + use_sequence_packing: True + device_mapping: list(range(0,8)) + infer_batch_size: 2 diff --git a/examples/qwen3-8B-onpolicy-distill-megatron/multi_teacher_onpolicy_distill_config.yaml b/examples/qwen3-8B-onpolicy-distill-megatron/multi_teacher_onpolicy_distill_config.yaml new file mode 100644 index 000000000..5af08e845 --- /dev/null +++ b/examples/qwen3-8B-onpolicy-distill-megatron/multi_teacher_onpolicy_distill_config.yaml @@ -0,0 +1,189 @@ +hydra: + run: + dir: . + output_subdir: null + +# Global Config +exp_name: "qwen3-8B-onpolicy-distill-math-code" +output_dir: ./output/ +logging_dir: ./output/logs +rollout_dump_dir: ./output/rollout_dump +seed: 42 +system_envs: + USE_MODELSCOPE: '1' + VLLM_USE_FLASHINFER_SAMPLER: '0' + +checkpoint_config: + type: file_system + output_dir: /data/cpfs_0/rl_examples/models/${exp_name} + +track_with: tensorboard + +num_nodes: 1 +num_gpus_per_node: 8 + +max_steps: 1024 +save_steps: 10000 +logging_steps: 1 +eval_steps: 10 +resume_from_checkpoint: false + +rollout_batch_size: 16 # prompt +prompt_length: 2048 +response_length: 4096 + +num_return_sequences_in_group: 8 +ppo_epochs: 1 +adv_estimator: "reinforce" + +pretrain: Qwen/Qwen3-8B +reward_pretrain: Qwen/Qwen3-8B + +global_template: qwen3 + +validation: + data_args: + file_name: + - data/math_benchmarks.jsonl + generating_args: + max_new_tokens: ${response_length} + top_p: 0.6 + top_k: 50 + num_beams: 1 + temperature: 0.6 + num_return_sequences: 1 + +student_train: + model_args: + model_name_or_path: ${pretrain} + attn_implementation: fa2 + disable_gradient_checkpointing: false + dtype: bf16 + model_type: ~ + training_args: + learning_rate: 5.0e-7 + weight_decay: 0.1 + per_device_train_batch_size: 8 + gradient_accumulation_steps: 1 + warmup_steps: 20 + num_train_epochs: 50 + data_args: + file_name: + - data/code_KodCode_data.jsonl # Must be first: has the widest schema (10 cols) + - data/dapo_math_17k_simple_boxed.jsonl # https://huggingface.co/datasets/open-r1/DAPO-Math-17k-Processed + domain_interleave_probs: + math_rule: 0.6 + code_rule: 0.4 + dataset_dir: data + messages: messages + interleave_probs: "1.0" + preprocessing_num_workers: 16 + strategy_args: + strategy_name: megatron_train + strategy_config: + tensor_model_parallel_size: 1 + context_parallel_size: 1 + pipeline_model_parallel_size: 2 + expert_model_parallel_size: 1 + use_distributed_optimizer: true + recompute_granularity: full + use_sequence_packing: True + sequence_packing_args: + algorithm: load_balance + max_packed_sequence_length_forward: 8192 + min_num_micro_batches_forward: 1 + max_packed_sequence_length_train: 8192 + min_num_micro_batches_train: 1 + device_mapping: list(range(0,32)) + infer_batch_size: 4 + +student_infer: + model_args: + model_name_or_path: ${pretrain} + disable_gradient_checkpointing: true + dtype: bf16 + generating_args: + max_new_tokens: ${response_length} + top_p: 0.99 + top_k: 100 + num_beams: 1 + temperature: 0.99 + num_return_sequences: ${num_return_sequences_in_group} + strategy_args: + strategy_name: vllm + strategy_config: + gpu_memory_utilization: 0.65 + block_size: 16 + tensor_parallel_size: 2 + load_format: dummy + device_mapping: list(range(0,16)) + infer_batch_size: 1 + +tag_to_template: + math_dapo: qwen3 + KodCode: qwen3 + +teacher: + teacher_32B: + model_args: + model_name_or_path: Qwen/Qwen3-32B # Math specialist teacher + disable_gradient_checkpointing: true + dtype: bf16 + model_type: ~ + strategy_args: + strategy_name: megatron_infer + strategy_config: + tensor_model_parallel_size: 2 + pipeline_model_parallel_size: 4 + context_parallel_size: 1 + bf16: true + device_mapping: list(range(16,32)) + infer_batch_size: 1 + opd_kl_coef: 1.0 + tag_included: [math_dapo] # Route math data to this teacher + sequence_packing_args: + algorithm: load_balance + max_packed_sequence_length_forward: 8192 + min_num_micro_batches_forward: 1 + + teacher_14B: + model_args: + model_name_or_path: Qwen/Qwen3-14B # Code specialist teacher + disable_gradient_checkpointing: true + dtype: bf16 + model_type: ~ + strategy_args: + strategy_name: megatron_infer + strategy_config: + tensor_model_parallel_size: 2 + pipeline_model_parallel_size: 2 + context_parallel_size: 1 + bf16: true + device_mapping: list(range(16,32)) + infer_batch_size: 1 + opd_kl_coef: 1.0 + tag_included: [KodCode] # Route code data to this teacher + sequence_packing_args: + algorithm: load_balance + max_packed_sequence_length_forward: 8192 + min_num_micro_batches_forward: 1 + +rewards: + math_rule: + worker_cls: roll.pipeline.rlvr.rewards.math_rule_reward_worker.MathRuleRewardWorker + model_args: + model_name_or_path: ${reward_pretrain} + data_args: + template: qwen3 + tag_included: [math_dapo] + world_size: 16 + infer_batch_size: 1 + code_rule: + worker_cls: roll.pipeline.rlvr.rewards.code_sandbox_reward_worker.CodeSandboxRewardWorker + model_args: + model_name_or_path: ${reward_pretrain} + data_args: + template: qwen3 + tag_included: [KodCode] + world_size: 16 + infer_batch_size: 1 diff --git a/examples/qwen3-8B-onpolicy-distill-megatron/submit_pipeline.sh b/examples/qwen3-8B-onpolicy-distill-megatron/submit_pipeline.sh deleted file mode 100644 index f6bcb4c74..000000000 --- a/examples/qwen3-8B-onpolicy-distill-megatron/submit_pipeline.sh +++ /dev/null @@ -1,45 +0,0 @@ -#!/bin/bash -set +x -source "examples/scripts/config.sh" - -WORKER_COUNT=2 -CONFIG_FILE="onpolicy_distill_config.yaml" -# Replace with mos uri -NEBULA_MODEL="" -ENTRY_FILE="examples/start_onpolicy_distill_pipeline.py" - -CONFIG_PATH=$(basename $(dirname $0)) -CONFIG_NAME="${CONFIG_FILE%.yaml}" -JOB_NAME="$CONFIG_PATH-$CONFIG_NAME" - -echo "JOB_NAME: ${JOB_NAME}" -echo "WORKER_COUNT: ${WORKER_COUNT}" -echo "CONFIG_NAME: ${CONFIG_NAME}" -echo "CONFIG_PATH: ${CONFIG_PATH}" -echo "ENTRY_FILE: ${ENTRY_FILE}" - -args="--config_name ${CONFIG_NAME} --config_path ${CONFIG_PATH}" - -mdl_args="--queue=${QUEUE} \ - --entry=${ENTRY_FILE} \ - --worker_count=${WORKER_COUNT} \ - --file.cluster_file=examples/scripts/cluster.json \ - --oss_access_id=${OSS_ACCESS_ID} \ - --oss_access_key=${OSS_ACCESS_KEY} \ - --oss_bucket=${OSS_BUCKET} \ - --oss_endpoint=${OSS_ENDPOINT} \ - --job_name=${JOB_NAME} \ - --algo_name=pytorch280 \ - --requirements_file_name=nebula_patch/requirements/requirements_torch280_vllm.txt \ - --oss_appendable=true \ - --_NEBULA_MODEL=${NEBULA_MODEL} \ - --nebula_model=${NEBULA_MODEL} \ - " -if [ -n "${OPENLM_TOKEN}" ]; then - mdl_args="${mdl_args} --env=OPENLM_TOKEN=${OPENLM_TOKEN}" -fi - -echo ${args} -echo ${mdl_args} - -nebulactl run mdl --user_params="${args}" $mdl_args \ No newline at end of file diff --git a/examples/qwen2.5-7B-rlvr_megatron/rlvr_qwen2.5_7B_megatron_vllm_8gpus.yaml b/examples/qwen3-8B-rlvr_fsdp2/rlvr_config.yaml similarity index 81% rename from examples/qwen2.5-7B-rlvr_megatron/rlvr_qwen2.5_7B_megatron_vllm_8gpus.yaml rename to examples/qwen3-8B-rlvr_fsdp2/rlvr_config.yaml index 7320ecdfa..40ab0a127 100644 --- a/examples/qwen2.5-7B-rlvr_megatron/rlvr_qwen2.5_7B_megatron_vllm_8gpus.yaml +++ b/examples/qwen3-8B-rlvr_fsdp2/rlvr_config.yaml @@ -3,29 +3,33 @@ hydra: dir: . output_subdir: null -exp_name: "qwen2.5-7B-rlvr-config" +exp_name: "qwen3-8B-rlvr-config" seed: 42 logging_dir: ./output/logs output_dir: ./output system_envs: USE_MODELSCOPE: '1' + checkpoint_config: type: file_system output_dir: /data/cpfs_0/rl_examples/models/${exp_name} -#track_with: wandb -#tracker_kwargs: -# api_key: -# project: roll_examples -# notes: roll_examples +# track_with: ml_tracker +# tracker_kwargs: +# project: roll_qwen3_8B +# notes: roll_qwen3_8B # tags: # - rlvr # - baseline +track_with: tensorboard +tracker_kwargs: + log_dir: ./rl_examples/llm/tensorboard/roll_exp/rlvr + num_gpus_per_node: 8 -max_steps: 100 +max_steps: 500 save_steps: 100 logging_steps: 1 eval_steps: 10 @@ -72,12 +76,12 @@ whiten_advantages: true # max_running_requests: 256 # is_num_return_sequences_expand: false -pretrain: Qwen/Qwen2.5-7B -reward_pretrain: Qwen/Qwen2.5-7B +pretrain: Qwen/Qwen3-8B +reward_pretrain: Qwen/Qwen3-8B validation: data_args: - template: qwen2_5 + template: qwen3 file_name: - data/math_benchmarks.jsonl generating_args: @@ -98,11 +102,11 @@ actor_train: learning_rate: 1.0e-6 weight_decay: 0 per_device_train_batch_size: 1 - gradient_accumulation_steps: 64 + gradient_accumulation_steps: 32 warmup_steps: 20 num_train_epochs: 50 data_args: - template: qwen2_5 + template: qwen3 file_name: - data/code_KodCode_data.jsonl - data/llm_judge_Multi-subject-RLVR_deal_new.jsonl @@ -120,15 +124,16 @@ actor_train: interleave_probs: "1.0" preprocessing_num_workers: 16 strategy_args: - strategy_name: megatron_train + strategy_name: fsdp2_train strategy_config: - tensor_model_parallel_size: 1 - pipeline_model_parallel_size: 1 - expert_model_parallel_size: 1 - use_distributed_optimizer: true - recompute_granularity: full - device_mapping: list(range(0,8)) - infer_batch_size: 2 + fsdp_size: 16 + param_dtype: bf16 + reduce_dtype: float32 + reshard_after_forward: true + offload_policy: false # offload model param + use_batched_model_update: false # disable batched gather during model update (not bring any benefit in test) + device_mapping: list(range(0,16)) + infer_batch_size: 4 actor_infer: model_args: @@ -142,14 +147,14 @@ actor_infer: temperature: 0.99 num_return_sequences: ${num_return_sequences_in_group} data_args: - template: qwen2_5 + template: qwen3 strategy_args: strategy_name: vllm strategy_config: - gpu_memory_utilization: 0.75 + gpu_memory_utilization: 0.8 block_size: 16 max_model_len: 8000 - device_mapping: list(range(0,6)) + device_mapping: list(range(0,12)) infer_batch_size: 1 reference: @@ -158,15 +163,15 @@ reference: dtype: bf16 model_type: ~ data_args: - template: qwen2_5 + template: qwen3 strategy_args: strategy_name: megatron_infer strategy_config: tensor_model_parallel_size: 1 pipeline_model_parallel_size: 1 expert_model_parallel_size: 1 - device_mapping: list(range(0,8)) - infer_batch_size: 1 + device_mapping: list(range(0,16)) + infer_batch_size: 4 rewards: crossthinkqa: @@ -176,7 +181,7 @@ rewards: model_args: model_name_or_path: ${reward_pretrain} data_args: - template: qwen2_5 + template: qwen3 tag_included: [crossthinkqa] world_size: 8 infer_batch_size: 4 @@ -186,7 +191,7 @@ rewards: model_args: model_name_or_path: ${reward_pretrain} data_args: - template: qwen2_5 + template: qwen3 tag_included: [ifeval] world_size: 8 infer_batch_size: 4 @@ -195,16 +200,10 @@ rewards: model_args: model_name_or_path: ${reward_pretrain} data_args: - template: qwen2_5 + template: qwen3 tag_included: [deepmath_103k, aime] world_size: 8 infer_batch_size: 1 -# dynamic filter config -# query_filter_config: -# type: mean_filter -# filter_args: -# threshold_up: 0.9 -# threshold_down: 0.1 code_sandbox: use_local: true worker_cls: roll.pipeline.rlvr.rewards.code_sandbox_reward_worker.CodeSandboxRewardWorker @@ -212,13 +211,9 @@ rewards: model_args: model_name_or_path: ${reward_pretrain} data_args: - template: qwen2_5 + template: qwen3 world_size: 8 infer_batch_size: 1 -# query_filter_config: -# type: std_filter -# filter_args: -# std_threshold: 0 llm_judge: # NOTE: llm as judge 也需要gpu, 不能和actor infer共享gpu worker_cls: roll.pipeline.rlvr.rewards.llm_judge_reward_worker.LLMJudgeRewardWorker @@ -239,9 +234,8 @@ rewards: temperature: 0.8 num_return_sequences: 1 data_args: - template: qwen2_5 + template: qwen3 strategy_args: - # strategy_name: hf_infer # strategy_config: null strategy_name: vllm strategy_config: diff --git a/examples/qwen2.5-7B-rlvr_megatron/rlvr_lora_zero3_amd.yaml b/examples/qwen3-8B-rlvr_fsdp2/rlvr_config_wo_cpuoffload.yaml similarity index 70% rename from examples/qwen2.5-7B-rlvr_megatron/rlvr_lora_zero3_amd.yaml rename to examples/qwen3-8B-rlvr_fsdp2/rlvr_config_wo_cpuoffload.yaml index 26119abc7..f9b6e754c 100644 --- a/examples/qwen2.5-7B-rlvr_megatron/rlvr_lora_zero3_amd.yaml +++ b/examples/qwen3-8B-rlvr_fsdp2/rlvr_config_wo_cpuoffload.yaml @@ -1,37 +1,32 @@ -defaults: - - ../config/deepspeed_zero@_here_ - - ../config/deepspeed_zero2@_here_ - - ../config/deepspeed_zero3@_here_ - - ../config/deepspeed_zero3_cpuoffload@_here_ - hydra: run: dir: . output_subdir: null -exp_name: "qwen2.5-7B-rlvr-config" +exp_name: "qwen3-8B-rlvr-config" seed: 42 logging_dir: ./output/logs output_dir: ./output system_envs: USE_MODELSCOPE: '1' + checkpoint_config: type: file_system - output_dir: /data/cpfs_0/rl_examples/models/${exp_name} + output_dir: ./output -#track_with: wandb -#tracker_kwargs: -# api_key: -# project: roll_examples -# notes: roll_examples +track_with: tensorboard +tracker_kwargs: + log_dir: ./rl_examples/llm/tensorboard/roll_exp/rlvr + +# track_with: ml_tracker +# tracker_kwargs: +# project: pumpkin_fsdp2_test +# notes: pumpkin_fsdp2 # tags: # - rlvr # - baseline - -track_with: tensorboard -tracker_kwargs: - log_dir: /data/oss_bucket_0/rl_examples/llm/tensorboard/roll_exp/rlvr +# - dp2 num_gpus_per_node: 8 @@ -57,9 +52,8 @@ advantage_clip: 2.0 dual_clip_loss: true # normalize -reward_norm: null -reward_shift: false -reward_scale: false +norm_mean_type: ~ +norm_std_type: ~ # data mask max_len_mask: true @@ -78,22 +72,17 @@ add_token_level_kl: false # advantage whiten_advantages: true -# lora -lora_target: o_proj,q_proj,k_proj,v_proj -lora_rank: 32 -lora_alpha: 32 - # dynamic sampling scheduler # use_additional_prompts: true # max_running_requests: 256 # is_num_return_sequences_expand: false -pretrain: Qwen/Qwen2.5-7B -reward_pretrain: Qwen/Qwen2.5-7B +pretrain: Qwen/Qwen3-8B +reward_pretrain: Qwen/Qwen3-8B validation: data_args: - template: qwen2_5 + template: qwen3 file_name: - data/math_benchmarks.jsonl generating_args: @@ -107,23 +96,18 @@ validation: actor_train: model_args: - attn_implementation: fa2 - # Recomputed tensor size does not match for LoRA with Zero3 when activating checkpointing, See https://github.com/huggingface/transformers/issues/34928 for details - disable_gradient_checkpointing: true + disable_gradient_checkpointing: false dtype: bf16 - lora_target: ${lora_target} - lora_rank: ${lora_rank} - lora_alpha: ${lora_alpha} model_type: ~ training_args: - learning_rate: 1.0e-5 + learning_rate: 1.0e-6 weight_decay: 0 per_device_train_batch_size: 1 gradient_accumulation_steps: 32 warmup_steps: 20 num_train_epochs: 50 data_args: - template: qwen2_5 + template: qwen3 file_name: - data/code_KodCode_data.jsonl - data/llm_judge_Multi-subject-RLVR_deal_new.jsonl @@ -141,19 +125,20 @@ actor_train: interleave_probs: "1.0" preprocessing_num_workers: 16 strategy_args: - strategy_name: deepspeed_train - strategy_config: ${deepspeed_zero3} + strategy_name: fsdp2_train + strategy_config: + fsdp_size: 16 + param_dtype: bf16 + reduce_dtype: float32 + reshard_after_forward: true + offload_policy: false device_mapping: list(range(0,16)) infer_batch_size: 4 actor_infer: model_args: - attn_implementation: fa2 disable_gradient_checkpointing: true dtype: bf16 - lora_target: ${lora_target} - lora_rank: ${lora_rank} - lora_alpha: ${lora_alpha} generating_args: max_new_tokens: ${response_length} top_p: 0.99 @@ -162,31 +147,32 @@ actor_infer: temperature: 0.99 num_return_sequences: ${num_return_sequences_in_group} data_args: - template: qwen2_5 + template: qwen3 strategy_args: strategy_name: vllm strategy_config: - gpu_memory_utilization: 0.6 - # https://github.com/vllm-project/vllm/issues/9452 - enforce_eager: false + gpu_memory_utilization: 0.8 block_size: 16 max_model_len: 8000 + disable_log_stats: false device_mapping: list(range(0,12)) infer_batch_size: 1 reference: model_args: - attn_implementation: fa2 disable_gradient_checkpointing: true dtype: bf16 model_type: ~ data_args: - template: qwen2_5 + template: qwen3 strategy_args: - strategy_name: hf_infer - strategy_config: ~ + strategy_name: megatron_infer + strategy_config: + tensor_model_parallel_size: 1 + pipeline_model_parallel_size: 1 + expert_model_parallel_size: 1 device_mapping: list(range(0,16)) - infer_batch_size: 8 + infer_batch_size: 4 rewards: crossthinkqa: @@ -196,7 +182,7 @@ rewards: model_args: model_name_or_path: ${reward_pretrain} data_args: - template: qwen2_5 + template: qwen3 tag_included: [crossthinkqa] world_size: 8 infer_batch_size: 4 @@ -206,7 +192,7 @@ rewards: model_args: model_name_or_path: ${reward_pretrain} data_args: - template: qwen2_5 + template: qwen3 tag_included: [ifeval] world_size: 8 infer_batch_size: 4 @@ -215,16 +201,10 @@ rewards: model_args: model_name_or_path: ${reward_pretrain} data_args: - template: qwen2_5 + template: qwen3 tag_included: [deepmath_103k, aime] world_size: 8 infer_batch_size: 1 -# dynamic filter config -# query_filter_config: -# type: mean_filter -# filter_args: -# threshold_up: 0.9 -# threshold_down: 0.1 code_sandbox: use_local: true worker_cls: roll.pipeline.rlvr.rewards.code_sandbox_reward_worker.CodeSandboxRewardWorker @@ -232,13 +212,9 @@ rewards: model_args: model_name_or_path: ${reward_pretrain} data_args: - template: qwen2_5 + template: qwen3 world_size: 8 infer_batch_size: 1 -# query_filter_config: -# type: std_filter -# filter_args: -# std_threshold: 0 llm_judge: # NOTE: llm as judge 也需要gpu, 不能和actor infer共享gpu worker_cls: roll.pipeline.rlvr.rewards.llm_judge_reward_worker.LLMJudgeRewardWorker @@ -259,9 +235,15 @@ rewards: temperature: 0.8 num_return_sequences: 1 data_args: - template: qwen2_5 + template: qwen3 strategy_args: - strategy_name: hf_infer - strategy_config: null + # strategy_config: null + strategy_name: vllm + strategy_config: + gpu_memory_utilization: 0.8 + block_size: 16 + max_model_len: 8000 + load_format: auto + disable_log_stats: false device_mapping: list(range(12,16)) infer_batch_size: 4 \ No newline at end of file diff --git a/examples/qwen3-8B-rlvr_fsdp2/run_rlvr_pipeline.sh b/examples/qwen3-8B-rlvr_fsdp2/run_rlvr_pipeline.sh new file mode 100755 index 000000000..7c7e9db72 --- /dev/null +++ b/examples/qwen3-8B-rlvr_fsdp2/run_rlvr_pipeline.sh @@ -0,0 +1,5 @@ +#!/bin/bash +set +x + +CONFIG_PATH=$(basename $(dirname $0)) +python examples/start_rlvr_pipeline.py --config_path $CONFIG_PATH --config_name rlvr_config diff --git a/examples/qwen3-omni/audio_test_80G.yaml b/examples/qwen3-omni/audio_test_80G.yaml new file mode 100644 index 000000000..ae4861634 --- /dev/null +++ b/examples/qwen3-omni/audio_test_80G.yaml @@ -0,0 +1,172 @@ +hydra: + run: + dir: . + output_subdir: null + +exp_name: "qwen3_omni_rlvr" +seed: 42 +logging_dir: ./output/logs +output_dir: ./output + +# checkpoint_config: +# type: mos + +# track_with: ml_tracker +# tracker_kwargs: +# project: ${exp_name} +# notes: ${exp_name} +# tags: +# - roll +# - rlvr +# - qwen3-omni + +save_steps: 20 +logging_steps: 1 +eval_steps: 10 +resume_from_checkpoint: false + +rollout_batch_size: 32 #256 +num_return_sequences_in_group: 8 +is_num_return_sequences_expand: true +prompt_length: 2048 +response_length: 2048 + +ppo_epochs: 1 +value_clip: 0.5 +reward_clip: 10 +advantage_clip: 10.0 +whiten_advantages: false +init_kl_coef: 0.0 +adv_estimator: "grpo" +use_kl_loss: true +kl_loss_coef: 1.0e-2 + +pretrain: Qwen/Qwen3-Omni-30B-A3B-Thinking +# pretrain: Qwen/Qwen3-Omni-30B-A3B-Instruct + +#validation: +# data_args: +# file_name: +# - /data/oss_bucket_0/yuzhao/data/One-RL-to-See-Them-All/Orsta-Data-47k/test/test_math_megabench_237.parquet +# - /data/oss_bucket_0/yuzhao/data/One-RL-to-See-Them-All/Orsta-Data-47k/test/test_detection_coco_test_multi_2000.parquet +# dataset_dir: ./ +# generating_args: +# max_new_tokens: ${response_length} +# top_p: 0.99 +# top_k: 100 +# num_beams: 1 +# temperature: 0.99 +# num_return_sequences: 1 +# eval_steps: ${eval_steps} + +actor_train: + model_args: + flash_attn: fa2 + attn_implementation: fa2 + disable_gradient_checkpointing: false + dtype: bf16 + model_type: ~ + freeze_module_prefix: "vision_model,audio_model,talker,code2wav" + training_args: + learning_rate: 1.0e-6 + weight_decay: 1.0e-2 + per_device_train_batch_size: 1 #2 + gradient_accumulation_steps: 16 #16 #32 + warmup_steps: 0 + num_train_epochs: 50 + data_args: + # case for training data: + # { + # "prompt": "Accurately transcribe the Polish audio into plain text.", + # "audios": ["dafd0ef215384e8b45326f43ee868204.wav"], + # "ability": "speech_recognition", + # "data_source": "the-mc-speech-dataset", + # "reward_model": { + # "answer": "Znajomy jest wiatr, woń zgniłych ryb, wodorostów i muszli.", + # "ground_truth": "Znajomy jest wiatr, woń zgniłych ryb, wodorostów i muszli.", + # "accuracy_ratio": 1.0, + # "format_ratio": 0.0, + # "verifier": "speech_recognition_verifier", + # "verifier_parm": {"test": "test"}, + # }, + # } + custom_data_kwargs_func: roll.datasets.vlm_dataset_utils.audio_test_get_data_kwargs + file_name: + - /data/oss_bucket_0/yuzhao/data/the-mc-speech-dataset-jsonl-format/data.jsonl + audio_folder: /data/oss_bucket_0/yuzhao/data/the-mc-speech-dataset-jsonl-format/audios/ + domain_interleave_probs: + math: 0.5 + dataset_dir: ./ + messages: prompt + preprocessing_num_workers: 32 + strategy_args: + strategy_name: megatron_train + strategy_config: + sequence_parallel: true + tensor_model_parallel_size: 1 + context_parallel_size: 1 + expert_model_parallel_size: 8 + pipeline_model_parallel_size: 1 + use_distributed_optimizer: true + # overlap_grad_reduce: true # to be resolved + moe_token_dispatcher_type: alltoall + moe_grouped_gemm: true + recompute_granularity: full + bf16: true + device_mapping: list(range(0,16)) + infer_batch_size: 1 + +actor_infer: + model_args: + flash_attn: fa2 + attn_implementation: fa2 + disable_gradient_checkpointing: true + dtype: bf16 + generating_args: + max_new_tokens: ${response_length} + top_p: 0.99 + top_k: 100 + num_beams: 1 + temperature: 0.99 + num_return_sequences: ${num_return_sequences_in_group} + strategy_args: + strategy_name: vllm + strategy_config: + gpu_memory_utilization: 0.6 + block_size: 16 + max_model_len: 8192 + tensor_parallel_size: 4 + enforce_eager: true + load_format: dummy + num_gpus_per_worker: 4 + device_mapping: list(range(0,16)) + infer_batch_size: 2 + +reference: + model_args: + flash_attn: fa2 + attn_implementation: fa2 + disable_gradient_checkpointing: true + dtype: bf16 + model_type: ~ + strategy_args: + strategy_name: megatron_infer + strategy_config: + sequence_parallel: true + tensor_model_parallel_size: 1 + context_parallel_size: 2 + pipeline_model_parallel_size: 1 + expert_model_parallel_size: 8 + bf16: true + device_mapping: list(range(0,16)) + infer_batch_size: 2 + +rewards: + math: + # NOTE: dummy reward to run audio pipeline, should be replaced with real reward when training + worker_cls: roll.pipeline.rlvr.rewards.math_rule_reward_worker.MathRuleRewardWorker + model_args: + model_name_or_path: ${pretrain} + tag_included: [the-mc-speech-dataset] + world_size: 8 + infer_batch_size: 1 \ No newline at end of file diff --git a/examples/qwen3-omni/rlvr_megatron.yaml b/examples/qwen3-omni/rlvr_megatron.yaml index 6ce249b30..72504a8e5 100644 --- a/examples/qwen3-omni/rlvr_megatron.yaml +++ b/examples/qwen3-omni/rlvr_megatron.yaml @@ -1,9 +1,3 @@ -defaults: - - ../config/deepspeed_zero@_here_ - - ../config/deepspeed_zero2@_here_ - - ../config/deepspeed_zero3@_here_ - - ../config/deepspeed_zero3_cpuoffload@_here_ - hydra: run: dir: . @@ -44,6 +38,8 @@ validation: - /data/oss_bucket_0/yuzhao/data/One-RL-to-See-Them-All/Orsta-Data-47k/test/test_math_megabench_237.parquet - /data/oss_bucket_0/yuzhao/data/One-RL-to-See-Them-All/Orsta-Data-47k/test/test_detection_coco_test_multi_2000.parquet dataset_dir: ./ + preprocessing_num_workers: 32 + image_max_pixels: 1048576 # 1024 * 1024 generating_args: max_new_tokens: ${response_length} top_p: 0.99 @@ -80,6 +76,7 @@ actor_train: dataset_dir: ./ messages: prompt preprocessing_num_workers: 32 + image_max_pixels: 1048576 # 1024 * 1024 strategy_args: strategy_name: megatron_train strategy_config: diff --git a/examples/qwen3-vl-30BA3B-rlvr_megatron/rlvr_megatron_80GB.yaml b/examples/qwen3-vl-30BA3B-rlvr_megatron/rlvr_megatron_80GB.yaml index 749848410..16bcbb899 100644 --- a/examples/qwen3-vl-30BA3B-rlvr_megatron/rlvr_megatron_80GB.yaml +++ b/examples/qwen3-vl-30BA3B-rlvr_megatron/rlvr_megatron_80GB.yaml @@ -1,9 +1,3 @@ -defaults: - - ../config/deepspeed_zero@_here_ - - ../config/deepspeed_zero2@_here_ - - ../config/deepspeed_zero3@_here_ - - ../config/deepspeed_zero3_cpuoffload@_here_ - hydra: run: dir: . @@ -52,6 +46,7 @@ validation: - /data/oss_bucket_0/yuzhao/data/One-RL-to-See-Them-All/Orsta-Data-47k/test/test_detection_coco_test_multi_2000.parquet dataset_dir: ./ preprocessing_num_workers: 32 + image_max_pixels: 1048576 # 1024 * 1024 generating_args: max_new_tokens: ${response_length} top_p: 0.99 @@ -88,6 +83,7 @@ actor_train: dataset_dir: ./ messages: prompt preprocessing_num_workers: 32 + image_max_pixels: 1048576 # 1024 * 1024 strategy_args: strategy_name: megatron_train strategy_config: diff --git a/examples/qwen3-vl-30BA3B-rlvr_megatron/rlvr_megatron_80GB_sglang.yaml b/examples/qwen3-vl-30BA3B-rlvr_megatron/rlvr_megatron_80GB_sglang.yaml new file mode 100644 index 000000000..8eb02edfa --- /dev/null +++ b/examples/qwen3-vl-30BA3B-rlvr_megatron/rlvr_megatron_80GB_sglang.yaml @@ -0,0 +1,164 @@ +hydra: + run: + dir: . + output_subdir: null + +exp_name: "qwen3_vl_moe_30BA3B_rlvr" +seed: 42 +logging_dir: ./output/logs +output_dir: ./output + +checkpoint_config: + type: file_system + output_dir: /data/cpfs_0/yuzhao/models + +track_with: tensorboard +tracker_kwargs: + log_dir: /data/oss_bucket_0/yuzhao/llm/tensorboard + +save_steps: 20 +logging_steps: 1 +eval_steps: 10 +resume_from_checkpoint: false + +rollout_batch_size: 256 +num_return_sequences_in_group: 8 +is_num_return_sequences_expand: true +prompt_length: 2048 +response_length: 4096 + +ppo_epochs: 1 +value_clip: 0.5 +reward_clip: 10 +advantage_clip: 10.0 +whiten_advantages: false +init_kl_coef: 0.0 +adv_estimator: "grpo" +use_kl_loss: true +kl_loss_coef: 1.0e-2 + +pretrain: Qwen/Qwen3-VL-30B-A3B-Instruct + +validation: + data_args: + file_name: + - /data/oss_bucket_0/yuzhao/data/One-RL-to-See-Them-All/Orsta-Data-47k/test/test_math_megabench_237.parquet + - /data/oss_bucket_0/yuzhao/data/One-RL-to-See-Them-All/Orsta-Data-47k/test/test_detection_coco_test_multi_2000.parquet + dataset_dir: ./ + preprocessing_num_workers: 32 + image_max_pixels: 1048576 # 1024 * 1024 + generating_args: + max_new_tokens: ${response_length} + top_p: 0.99 + top_k: 100 + num_beams: 1 + temperature: 0.99 + num_return_sequences: 1 + eval_steps: ${eval_steps} + +actor_train: + model_args: + flash_attn: fa2 + attn_implementation: fa2 + disable_gradient_checkpointing: false + dtype: bf16 + model_type: ~ + freeze_module_prefix: vision_model + training_args: + learning_rate: 1.0e-6 + weight_decay: 1.0e-2 + per_device_train_batch_size: 1 + gradient_accumulation_steps: 64 + warmup_steps: 0 + num_train_epochs: 50 + data_args: + # use One-RL-to-See-Them-All/Orsta-Data-47k as train dataset + # download from https://huggingface.co/datasets/One-RL-to-See-Them-All/Orsta-Data-47k + file_name: + - /data/oss_bucket_0/yuzhao/data/One-RL-to-See-Them-All/Orsta-Data-47k/train/train_detection_v3det_4000.parquet + - /data/oss_bucket_0/yuzhao/data/One-RL-to-See-Them-All/Orsta-Data-47k/train/train_math_mmmath_3539.parquet + domain_interleave_probs: + math: 0.5 + cv_detection: 0.5 + dataset_dir: ./ + messages: prompt + preprocessing_num_workers: 32 + image_max_pixels: 1048576 # 1024 * 1024 + strategy_args: + strategy_name: megatron_train + strategy_config: + sequence_parallel: true + tensor_model_parallel_size: 1 + context_parallel_size: 1 + expert_model_parallel_size: 8 + pipeline_model_parallel_size: 1 + use_distributed_optimizer: true + moe_token_dispatcher_type: alltoall + recompute_granularity: selective + recompute_modules: "moe,layernorm" + bias_activation_fusion: true + apply_rope_fusion: true + moe_grouped_gemm: true + moe_shared_expert_overlap: true + bf16: true + device_mapping: list(range(0,16)) + infer_batch_size: 1 + +actor_infer: + model_args: + flash_attn: fa2 + attn_implementation: fa2 + disable_gradient_checkpointing: true + dtype: bf16 + generating_args: + max_new_tokens: ${response_length} + top_p: 0.99 + top_k: 100 + num_beams: 1 + temperature: 0.99 + num_return_sequences: ${num_return_sequences_in_group} + strategy_args: + strategy_name: sglang + strategy_config: + mem_fraction_static: 0.6 + load_format: dummy + num_gpus_per_worker: 2 + device_mapping: list(range(0,16)) + infer_batch_size: 2 + +reference: + model_args: + flash_attn: fa2 + attn_implementation: fa2 + disable_gradient_checkpointing: true + dtype: bf16 + model_type: ~ + strategy_args: + strategy_name: megatron_infer + strategy_config: + sequence_parallel: true + tensor_model_parallel_size: 1 + context_parallel_size: 1 + pipeline_model_parallel_size: 1 + expert_model_parallel_size: 2 + bf16: true + device_mapping: list(range(0,16)) + infer_batch_size: 2 + +rewards: + math: + worker_cls: roll.pipeline.rlvr.rewards.math_rule_reward_worker.MathRuleRewardWorker + model_args: + model_name_or_path: ${pretrain} + # data source whose ability is math in One-RL-to-See-Them-All/Orsta-Data-47k + tag_included: [mm_math, megabench_math] + world_size: 8 + infer_batch_size: 1 + cv_detection: + worker_cls: roll.pipeline.rlvr.rewards.detection_reward_worker.DetectionRewardWorker + model_args: + model_name_or_path: ${pretrain} + # data source whose ability is cv_detection in One-RL-to-See-Them-All/Orsta-Data-47k + tag_included: [v3det_train, object365_train, coco_val_multi_test] + world_size: 8 + infer_batch_size: 1 diff --git a/examples/qwen3-vl-30BA3B-rlvr_megatron/rlvr_megatron_ppu.yaml b/examples/qwen3-vl-30BA3B-rlvr_megatron/rlvr_megatron_ppu.yaml new file mode 100644 index 000000000..f02ee5ff0 --- /dev/null +++ b/examples/qwen3-vl-30BA3B-rlvr_megatron/rlvr_megatron_ppu.yaml @@ -0,0 +1,164 @@ +hydra: + run: + dir: . + output_subdir: null + +exp_name: "qwen3_vl_moe_30BA3B_rlvr" +seed: 42 +logging_dir: ./output/logs +output_dir: ./output + +checkpoint_config: + type: file_system + output_dir: /data/cpfs_0/yuzhao/models + +track_with: tensorboard +tracker_kwargs: + log_dir: /data/oss_bucket_0/yuzhao/llm/tensorboard + +save_steps: 20 +logging_steps: 1 +eval_steps: 10 +resume_from_checkpoint: false + +rollout_batch_size: 128 +num_return_sequences_in_group: 8 +is_num_return_sequences_expand: true +prompt_length: 2048 +response_length: 4096 + +ppo_epochs: 1 +value_clip: 0.5 +reward_clip: 10 +advantage_clip: 10.0 +whiten_advantages: false +init_kl_coef: 0.0 +adv_estimator: "grpo" +use_kl_loss: true +kl_loss_coef: 1.0e-2 + +pretrain: Qwen/Qwen3-VL-30B-A3B-Instruct + +validation: + data_args: + file_name: + - /data/oss_bucket_0/yuzhao/data/One-RL-to-See-Them-All/Orsta-Data-47k/test/test_math_megabench_237.parquet + - /data/oss_bucket_0/yuzhao/data/One-RL-to-See-Them-All/Orsta-Data-47k/test/test_detection_coco_test_multi_2000.parquet + dataset_dir: ./ + preprocessing_num_workers: 32 + image_max_pixels: 1048576 # 1024 * 1024 + generating_args: + max_new_tokens: ${response_length} + top_p: 0.99 + top_k: 100 + num_beams: 1 + temperature: 0.99 + num_return_sequences: 1 + eval_steps: ${eval_steps} + +actor_train: + model_args: + flash_attn: fa2 + attn_implementation: fa2 + disable_gradient_checkpointing: false + dtype: bf16 + model_type: ~ + freeze_module_prefix: vision_model + training_args: + learning_rate: 1.0e-6 + weight_decay: 1.0e-2 + per_device_train_batch_size: 1 + gradient_accumulation_steps: 32 + warmup_steps: 0 + num_train_epochs: 50 + data_args: + # use One-RL-to-See-Them-All/Orsta-Data-47k as train dataset + # download from https://huggingface.co/datasets/One-RL-to-See-Them-All/Orsta-Data-47k + file_name: + - /data/oss_bucket_0/yuzhao/data/One-RL-to-See-Them-All/Orsta-Data-47k/train/train_detection_v3det_4000.parquet + - /data/oss_bucket_0/yuzhao/data/One-RL-to-See-Them-All/Orsta-Data-47k/train/train_math_mmmath_3539.parquet + domain_interleave_probs: + math: 0.5 + cv_detection: 0.5 + dataset_dir: ./ + messages: prompt + preprocessing_num_workers: 32 + image_max_pixels: 1048576 # 1024 * 1024 + strategy_args: + strategy_name: megatron_train + strategy_config: + sequence_parallel: true + tensor_model_parallel_size: 1 + context_parallel_size: 1 + expert_model_parallel_size: 8 + pipeline_model_parallel_size: 1 + use_distributed_optimizer: true + moe_token_dispatcher_type: alltoall + moe_grouped_gemm: true + recompute_granularity: full + bf16: true + device_mapping: list(range(0,32)) + infer_batch_size: 1 + +actor_infer: + model_args: + flash_attn: fa2 + attn_implementation: fa2 + disable_gradient_checkpointing: true + dtype: bf16 + generating_args: + max_new_tokens: ${response_length} + top_p: 0.99 + top_k: 100 + num_beams: 1 + temperature: 0.99 + num_return_sequences: ${num_return_sequences_in_group} + strategy_args: + strategy_name: vllm + strategy_config: + gpu_memory_utilization: 0.6 + block_size: 16 + max_model_len: 8192 + tensor_parallel_size: 8 + enforce_eager: true + load_format: dummy + num_gpus_per_worker: 8 + device_mapping: list(range(0,16)) + infer_batch_size: 2 + +reference: + model_args: + flash_attn: fa2 + attn_implementation: fa2 + disable_gradient_checkpointing: true + dtype: bf16 + model_type: ~ + strategy_args: + strategy_name: megatron_infer + strategy_config: + sequence_parallel: true + tensor_model_parallel_size: 1 + context_parallel_size: 2 + pipeline_model_parallel_size: 1 + expert_model_parallel_size: 8 + bf16: true + device_mapping: list(range(16,32)) + infer_batch_size: 2 + +rewards: + math: + worker_cls: roll.pipeline.rlvr.rewards.math_rule_reward_worker.MathRuleRewardWorker + model_args: + model_name_or_path: ${pretrain} + # data source whose ability is math in One-RL-to-See-Them-All/Orsta-Data-47k + tag_included: [mm_math, megabench_math] + world_size: 8 + infer_batch_size: 1 + cv_detection: + worker_cls: roll.pipeline.rlvr.rewards.detection_reward_worker.DetectionRewardWorker + model_args: + model_name_or_path: ${pretrain} + # data source whose ability is cv_detection in One-RL-to-See-Them-All/Orsta-Data-47k + tag_included: [v3det_train, object365_train, coco_val_multi_test] + world_size: 8 + infer_batch_size: 1 diff --git a/examples/qwen3-vl-30BA3B-video-rlvr/video_r1.yaml b/examples/qwen3-vl-30BA3B-video-rlvr/video_r1.yaml new file mode 100644 index 000000000..3936ef2ca --- /dev/null +++ b/examples/qwen3-vl-30BA3B-video-rlvr/video_r1.yaml @@ -0,0 +1,161 @@ +hydra: + run: + dir: . + output_subdir: null + +exp_name: "video-r1" +seed: 42 +logging_dir: ./output/logs +output_dir: ./output + +checkpoint_config: + type: file_system + output_dir: /data/cpfs_0/yuzhao/roll/vlm-video/models + + +track_with: tensorboard +tracker_kwargs: + log_dir: /data/oss_bucket_0/yuzhao/llm/tensorboard + +offload_nccl: false + +save_steps: 100 +logging_steps: 1 +eval_steps: 10 +resume_from_checkpoint: false + +rollout_batch_size: 8 +num_return_sequences_in_group: 8 +is_num_return_sequences_expand: true +prompt_length: 16384 +response_length: 4096 + +ppo_epochs: 1 +value_clip: 0.5 +reward_clip: 10 +advantage_clip: 10.0 +whiten_advantages: false +init_kl_coef: 0.0 +adv_estimator: "grpo" +use_kl_loss: true +kl_loss_coef: 1.0e-2 + +pretrain: Qwen/Qwen3-VL-30B-A3B-Instruct + +actor_train: + model_args: + attn_implementation: fa2 + disable_gradient_checkpointing: false + dtype: bf16 + model_type: ~ + # freeze_module_prefix: "vision_model.blocks,vision_model.patch_embed" + training_args: + learning_rate: 1.0e-6 + weight_decay: 1.0e-2 + per_device_train_batch_size: 1 + gradient_accumulation_steps: 8 + warmup_steps: 0 + num_train_epochs: 50 + data_args: + # we use video-r1 116k video data with adjusted format, case for training data: + # { + # "modality_type": "video", + # "problem": "Why does the video conclude with a red screen displaying text?Options:\nA. To promote the BBC one program 'Super Cute Animals'\nB. To provide a warning message\nC. To show the credits of the video\nD. To indicate a change in scene\n", + # "solution": "A", + # "video": "videos/youtube_video_2024/ytb_HBxn56l9WcU.mp4", + # "problem_type": "multiple choice" + # } + custom_data_kwargs_func: roll.datasets.vlm_dataset_utils.video_r1_get_data_kwargs + file_name: + - /data/oss_bucket_0/yinchao/dataset/LMMS/Video-R1-116k-video.json + domain_interleave_probs: + videor1: 1.0 + dataset_dir: ./ + messages: prompt + preprocessing_num_workers: 32 + # settings same with RewatchR1 and VideoR1 + video_total_pixels: 9633792 # 12288 * 28 * 28, max_video_token=12k + video_min_pixels: 100352 # 128 * 28 * 28 + video_max_pixels: 100352 # 128 * 28 * 28 + video_fps: 2.0 + video_max_frames: 192 + video_folder: /data/oss_bucket_0/ + strategy_args: + strategy_name: megatron_train + strategy_config: + sequence_parallel: true + tensor_model_parallel_size: 1 + context_parallel_size: 1 + expert_model_parallel_size: 8 + pipeline_model_parallel_size: 2 + use_distributed_optimizer: true + moe_token_dispatcher_type: alltoall + recompute_granularity: full + bias_activation_fusion: true + apply_rope_fusion: true + moe_grouped_gemm: true + moe_shared_expert_overlap: true + bf16: true + device_mapping: list(range(0,16)) + infer_batch_size: 1 + offload_nccl: ${offload_nccl} + +actor_infer: + model_args: + attn_implementation: fa2 + disable_gradient_checkpointing: true + dtype: bf16 + generating_args: + max_new_tokens: ${response_length} + top_p: 0.99 + top_k: 100 + num_beams: 1 + temperature: 0.99 + num_return_sequences: ${num_return_sequences_in_group} + strategy_args: + strategy_name: vllm + strategy_config: + tensor_parallel_size: 2 + max_num_batched_tokens: 32768 + gpu_memory_utilization: 0.6 + block_size: 16 + disable_mm_preprocessor_cache: true + enable_prefix_caching: false + sleep_level: 2 + device_mapping: list(range(0,16)) + infer_batch_size: 32 + offload_nccl: ${offload_nccl} + +reference: + model_args: + flash_attn: fa2 + attn_implementation: fa2 + disable_gradient_checkpointing: true + dtype: bf16 + model_type: ~ + strategy_args: + strategy_name: megatron_infer + strategy_config: + sequence_parallel: true + tensor_model_parallel_size: 1 + context_parallel_size: 1 + pipeline_model_parallel_size: 1 + expert_model_parallel_size: 8 + bf16: true + device_mapping: list(range(0,16)) + infer_batch_size: 1 + offload_nccl: ${offload_nccl} + +reward_system_config: + skip_special_tokens: true + +rewards: + videor1: + # NOTE: `re.fullmatch(pattern)` in format_reward seems not suitable for qwen3-vl + # thinking model, `` is the prompt ending suffix thus adjust format_reward + worker_cls: roll.pipeline.rlvr.rewards.video_r1_reward_worker.VideoR1RewardWorker + model_args: + model_name_or_path: ${pretrain} + tag_included: ["video"] + world_size: 2 + infer_batch_size: 1 \ No newline at end of file diff --git a/examples/qwen3-vl-30BA3B-video-rlvr/video_r1_sglang.yaml b/examples/qwen3-vl-30BA3B-video-rlvr/video_r1_sglang.yaml new file mode 100644 index 000000000..407ccf81f --- /dev/null +++ b/examples/qwen3-vl-30BA3B-video-rlvr/video_r1_sglang.yaml @@ -0,0 +1,158 @@ +hydra: + run: + dir: . + output_subdir: null + +exp_name: "video-r1" +seed: 42 +logging_dir: ./output/logs +output_dir: ./output + +checkpoint_config: + type: file_system + output_dir: /data/cpfs_0/yuzhao/roll/vlm-video/models + + +track_with: tensorboard +tracker_kwargs: + log_dir: /data/oss_bucket_0/yuzhao/llm/tensorboard + +offload_nccl: false + +save_steps: 100 +logging_steps: 1 +eval_steps: 10 +resume_from_checkpoint: false + +rollout_batch_size: 8 +num_return_sequences_in_group: 8 +is_num_return_sequences_expand: true +prompt_length: 16384 +response_length: 4096 + +ppo_epochs: 1 +value_clip: 0.5 +reward_clip: 10 +advantage_clip: 10.0 +whiten_advantages: false +init_kl_coef: 0.0 +adv_estimator: "grpo" +use_kl_loss: true +kl_loss_coef: 1.0e-2 + +pretrain: Qwen/Qwen3-VL-30B-A3B-Instruct + +actor_train: + model_args: + attn_implementation: fa2 + disable_gradient_checkpointing: false + dtype: bf16 + model_type: ~ + # freeze_module_prefix: "vision_model.blocks,vision_model.patch_embed" + training_args: + learning_rate: 1.0e-6 + weight_decay: 1.0e-2 + per_device_train_batch_size: 1 + gradient_accumulation_steps: 8 + warmup_steps: 0 + num_train_epochs: 50 + data_args: + # we use video-r1 116k video data with adjusted format, case for training data: + # { + # "modality_type": "video", + # "problem": "Why does the video conclude with a red screen displaying text?Options:\nA. To promote the BBC one program 'Super Cute Animals'\nB. To provide a warning message\nC. To show the credits of the video\nD. To indicate a change in scene\n", + # "solution": "A", + # "video": "videos/youtube_video_2024/ytb_HBxn56l9WcU.mp4", + # "problem_type": "multiple choice" + # } + custom_data_kwargs_func: roll.datasets.vlm_dataset_utils.video_r1_get_data_kwargs + file_name: + - /data/oss_bucket_0/yinchao/dataset/LMMS/Video-R1-116k-video.json + domain_interleave_probs: + videor1: 1.0 + dataset_dir: ./ + messages: prompt + preprocessing_num_workers: 32 + # settings same with RewatchR1 and VideoR1 + video_total_pixels: 9633792 # 12288 * 28 * 28, max_video_token=12k + video_min_pixels: 100352 # 128 * 28 * 28 + video_max_pixels: 100352 # 128 * 28 * 28 + video_fps: 2.0 + video_max_frames: 192 + video_folder: /data/oss_bucket_0/ + strategy_args: + strategy_name: megatron_train + strategy_config: + sequence_parallel: true + tensor_model_parallel_size: 1 + context_parallel_size: 1 + expert_model_parallel_size: 8 + pipeline_model_parallel_size: 2 + use_distributed_optimizer: true + moe_token_dispatcher_type: alltoall + recompute_granularity: full + bias_activation_fusion: true + apply_rope_fusion: true + moe_grouped_gemm: true + moe_shared_expert_overlap: true + bf16: true + device_mapping: list(range(0,16)) + infer_batch_size: 1 + offload_nccl: ${offload_nccl} + +actor_infer: + model_args: + attn_implementation: fa2 + disable_gradient_checkpointing: true + dtype: bf16 + generating_args: + max_new_tokens: ${response_length} + top_p: 0.99 + top_k: 100 + num_beams: 1 + temperature: 0.99 + num_return_sequences: ${num_return_sequences_in_group} + strategy_args: + strategy_name: sglang + strategy_config: + mem_fraction_static: 0.6 + load_format: dummy + num_gpus_per_worker: 2 + + device_mapping: list(range(0,16)) + infer_batch_size: 32 + offload_nccl: ${offload_nccl} + +reference: + model_args: + flash_attn: fa2 + attn_implementation: fa2 + disable_gradient_checkpointing: true + dtype: bf16 + model_type: ~ + strategy_args: + strategy_name: megatron_infer + strategy_config: + sequence_parallel: true + tensor_model_parallel_size: 1 + context_parallel_size: 1 + pipeline_model_parallel_size: 1 + expert_model_parallel_size: 8 + bf16: true + device_mapping: list(range(0,16)) + infer_batch_size: 1 + offload_nccl: ${offload_nccl} + +reward_system_config: + skip_special_tokens: true + +rewards: + videor1: + # NOTE: `re.fullmatch(pattern)` in format_reward seems not suitable for qwen3-vl + # thinking model, `` is the prompt ending suffix thus adjust format_reward + worker_cls: roll.pipeline.rlvr.rewards.video_r1_reward_worker.VideoR1RewardWorker + model_args: + model_name_or_path: ${pretrain} + tag_included: ["video"] + world_size: 2 + infer_batch_size: 1 \ No newline at end of file diff --git a/examples/qwen3-vl-32B-rlvr_megatron/rlvr_megatron_80GB.yaml b/examples/qwen3-vl-32B-rlvr_megatron/rlvr_megatron_80GB.yaml index 255c07528..f7c921d9a 100644 --- a/examples/qwen3-vl-32B-rlvr_megatron/rlvr_megatron_80GB.yaml +++ b/examples/qwen3-vl-32B-rlvr_megatron/rlvr_megatron_80GB.yaml @@ -1,9 +1,3 @@ -defaults: - - ../config/deepspeed_zero@_here_ - - ../config/deepspeed_zero2@_here_ - - ../config/deepspeed_zero3@_here_ - - ../config/deepspeed_zero3_cpuoffload@_here_ - hydra: run: dir: . @@ -52,6 +46,7 @@ validation: - /data/oss_bucket_0/yuzhao/data/One-RL-to-See-Them-All/Orsta-Data-47k/test/test_detection_coco_test_multi_2000.parquet dataset_dir: ./ preprocessing_num_workers: 32 + image_max_pixels: 1048576 # 1024 * 1024 generating_args: max_new_tokens: ${response_length} top_p: 0.99 @@ -88,6 +83,7 @@ actor_train: dataset_dir: ./ messages: prompt preprocessing_num_workers: 32 + image_max_pixels: 1048576 # 1024 * 1024 strategy_args: strategy_name: megatron_train strategy_config: diff --git a/examples/qwen3-vl-4B-rlvr_megatron/rlvr_megatron_80G.yaml b/examples/qwen3-vl-4B-rlvr_megatron/rlvr_megatron_80G.yaml index 76a4166b6..37118c996 100644 --- a/examples/qwen3-vl-4B-rlvr_megatron/rlvr_megatron_80G.yaml +++ b/examples/qwen3-vl-4B-rlvr_megatron/rlvr_megatron_80G.yaml @@ -1,9 +1,3 @@ -defaults: - - ../config/deepspeed_zero@_here_ - - ../config/deepspeed_zero2@_here_ - - ../config/deepspeed_zero3@_here_ - - ../config/deepspeed_zero3_cpuoffload@_here_ - hydra: run: dir: . @@ -52,6 +46,7 @@ validation: - /data/oss_bucket_0/yuzhao/data/One-RL-to-See-Them-All/Orsta-Data-47k/test/test_detection_coco_test_multi_2000.parquet dataset_dir: ./ preprocessing_num_workers: 32 + image_max_pixels: 1048576 # 1024 * 1024 generating_args: max_new_tokens: ${response_length} top_p: 0.99 @@ -88,6 +83,7 @@ actor_train: dataset_dir: ./ messages: prompt preprocessing_num_workers: 32 + image_max_pixels: 1048576 # 1024 * 1024 strategy_args: strategy_name: megatron_train strategy_config: diff --git a/examples/qwen3.5-122A10-rlvr_megatron/rlvr_megatron_80GB.yaml b/examples/qwen3.5-122A10-rlvr_megatron/rlvr_megatron_80GB.yaml new file mode 100644 index 000000000..3d86a26e7 --- /dev/null +++ b/examples/qwen3.5-122A10-rlvr_megatron/rlvr_megatron_80GB.yaml @@ -0,0 +1,233 @@ +hydra: + run: + dir: . + output_subdir: null + +exp_name: "qwen3d5-122A10-rlvr-config" +seed: 42 +logging_dir: ./output/logs +output_dir: ./output +system_envs: + USE_MODELSCOPE: '1' + +checkpoint_config: + type: file_system + output_dir: ./rl_examples/models/${exp_name} + +#track_with: wandb +#tracker_kwargs: +# api_key: +# project: roll_examples +# notes: roll_examples +# tags: +# - rlvr +# - baseline + +track_with: tensorboard +tracker_kwargs: + log_dir: ./rl_examples/llm/tensorboard/roll_exp/rlvr + +num_gpus_per_node: 8 + +max_steps: 500 +save_steps: 100 +logging_steps: 1 +eval_steps: 10 +resume_from_checkpoint: false + +offload_nccl: false # save GPU mem + +rollout_batch_size: 32 # prompt +prompt_length: 2048 +response_length: 4096 + +num_return_sequences_in_group: 8 +ppo_epochs: 1 +adv_estimator: "reinforce" + +# clip +value_clip: 0.5 +reward_clip: 10 +advantage_clip: 2.0 +dual_clip_loss: true + +# normalize +norm_mean_type: ~ +norm_std_type: ~ + +# data mask +max_len_mask: true +difficulty_mask: true +difficulty_low_threshold: 0.1 +difficulty_high_threshold: 0.95 +error_max_len_clip: false + +# data weight +difficulty_loss_weight: false +length_loss_weight: false + +# reward +add_token_level_kl: false + +# advantage +whiten_advantages: true + +# dynamic sampling scheduler +# use_additional_prompts: true +# max_running_requests: 256 +# is_num_return_sequences_expand: false + +pretrain: Qwen/Qwen3.5-122B-A10B +reward_pretrain: Qwen/Qwen3.5-122B-A10B + +validation: + data_args: + template: qwen2_5 + file_name: + - data/math_benchmarks.jsonl + generating_args: + max_new_tokens: ${response_length} + top_p: 0.6 + top_k: 50 + num_beams: 1 + temperature: 0.6 + num_return_sequences: 1 + +actor_train: + model_args: + disable_gradient_checkpointing: false + dtype: bf16 + model_type: ~ + training_args: + learning_rate: 1.0e-6 + weight_decay: 0 + per_device_train_batch_size: 1 + gradient_accumulation_steps: 8 + warmup_steps: 20 + num_train_epochs: 50 + data_args: + template: qwen2_5 + file_name: + - data/code_KodCode_data.jsonl + - data/math_deepmath_deal.jsonl + - data/general_ifeval_train_deal.jsonl + - data/general_CrossThink-QA_deal.jsonl + domain_interleave_probs: + math_rule: 0.5 + code_sandbox: 0.3 + crossthinkqa: 0.1 + ifeval: 0.1 + dataset_dir: data + messages: messages + interleave_probs: "1.0" + preprocessing_num_workers: 16 + strategy_args: + strategy_name: megatron_train + strategy_config: + sequence_parallel: true + tensor_model_parallel_size: 1 # qwen3.5 not support TP now + pipeline_model_parallel_size: 4 + expert_model_parallel_size: 8 + context_parallel_size: 1 + use_distributed_optimizer: true + recompute_granularity: full + moe_token_dispatcher_type: "alltoall" + bias_activation_fusion: true + apply_rope_fusion: true + moe_grouped_gemm: true + moe_shared_expert_overlap: true + cross_entropy_loss_fusion: true + bf16: true + device_mapping: list(range(0,128)) + infer_batch_size: 2 + offload_nccl: ${offload_nccl} + +actor_infer: + model_args: + disable_gradient_checkpointing: true + dtype: bf16 + generating_args: + max_new_tokens: ${response_length} + top_p: 0.99 + top_k: 100 + num_beams: 1 + temperature: 0.99 + num_return_sequences: ${num_return_sequences_in_group} + data_args: + template: qwen2_5 + strategy_args: + strategy_name: vllm + strategy_config: + gpu_memory_utilization: 0.7 + load_format: dummy + tensor_parallel_size: 8 + cudagraph_capture_sizes: [1,2,4,8,16,24,32,40,48] + optimization_level: 1 + sleep_level: 2 # 2 will destroy model parameter and kv_cache after generate to save cpu memory, 1 will destroy kv_cache only. + num_gpus_per_worker: 8 + device_mapping: list(range(0,128)) + infer_batch_size: 1 + +reference: + model_args: + disable_gradient_checkpointing: true + dtype: bf16 + model_type: ~ + data_args: + template: qwen2_5 + strategy_args: + strategy_name: megatron_infer + strategy_config: + tensor_model_parallel_size: 1 + pipeline_model_parallel_size: 2 + expert_model_parallel_size: 8 + moe_token_dispatcher_type: alltoall + bias_activation_fusion: true + apply_rope_fusion: true + moe_grouped_gemm: true + bf16: true + device_mapping: list(range(0,128)) + infer_batch_size: 2 + offload_nccl: ${offload_nccl} + +rewards: + crossthinkqa: + worker_cls: roll.pipeline.rlvr.rewards.crossthinkqa_rule_reward_worker.CrossThinkQARuleRewardWorker + reward_type: soft + response_length_penalty_coef: 0.0 + model_args: + model_name_or_path: ${reward_pretrain} + data_args: + template: qwen2_5 + tag_included: [crossthinkqa] + world_size: 8 + infer_batch_size: 4 + ifeval: + worker_cls: roll.pipeline.rlvr.rewards.ifeval_rule_reward_worker.GeneralRuleRewardWorker + reward_type: soft + model_args: + model_name_or_path: ${reward_pretrain} + data_args: + template: qwen2_5 + tag_included: [ifeval] + world_size: 8 + infer_batch_size: 4 + math_rule: + worker_cls: roll.pipeline.rlvr.rewards.math_rule_reward_worker.MathRuleRewardWorker + model_args: + model_name_or_path: ${reward_pretrain} + data_args: + template: qwen2_5 + tag_included: [deepmath_103k, aime] + world_size: 8 + infer_batch_size: 1 + code_sandbox: + use_local: true + worker_cls: roll.pipeline.rlvr.rewards.code_sandbox_reward_worker.CodeSandboxRewardWorker + tag_included: [KodCode] + model_args: + model_name_or_path: ${reward_pretrain} + data_args: + template: qwen2_5 + world_size: 8 + infer_batch_size: 1 \ No newline at end of file diff --git a/examples/qwen3.5-27B-rlvr_megatron/rlvr_megatron_80GB.yaml b/examples/qwen3.5-27B-rlvr_megatron/rlvr_megatron_80GB.yaml index 4892bed23..c16688be8 100644 --- a/examples/qwen3.5-27B-rlvr_megatron/rlvr_megatron_80GB.yaml +++ b/examples/qwen3.5-27B-rlvr_megatron/rlvr_megatron_80GB.yaml @@ -1,9 +1,3 @@ -defaults: - - ../config/deepspeed_zero@_here_ - - ../config/deepspeed_zero2@_here_ - - ../config/deepspeed_zero3@_here_ - - ../config/deepspeed_zero3_cpuoffload@_here_ - hydra: run: dir: . @@ -99,6 +93,7 @@ actor_train: overlap_grad_reduce: true use_distributed_optimizer: true bf16: true + #mtp_training_mode: standalone device_mapping: list(range(0,32)) infer_batch_size: 1 @@ -125,6 +120,9 @@ actor_infer: tensor_parallel_size: 4 optimization_level: 1 sleep_level: 2 + #speculative_config: + # method: mtp + # num_speculative_tokens: 3 num_gpus_per_worker: 2 device_mapping: list(range(0,32)) infer_batch_size: 1 diff --git a/examples/qwen3.5-27B-rlvr_megatron/submit_pipeline.sh b/examples/qwen3.5-27B-rlvr_megatron/submit_pipeline.sh deleted file mode 100644 index c7d91c093..000000000 --- a/examples/qwen3.5-27B-rlvr_megatron/submit_pipeline.sh +++ /dev/null @@ -1,45 +0,0 @@ -#!/bin/bash -set +x -source "examples/scripts/config.sh" - -WORKER_COUNT=4 -CONFIG_FILE="rlvr_megatron_80GB.yaml" -# 替换为mos uri -NEBULA_MODEL="" -ENTRY_FILE="examples/start_rlvr_vl_pipeline.py" - -CONFIG_PATH=$(basename $(dirname $0)) -CONFIG_NAME="${CONFIG_FILE%.yaml}" -JOB_NAME="$CONFIG_PATH-$CONFIG_NAME" - -echo "JOB_NAME: ${JOB_NAME}" -echo "WORKER_COUNT: ${WORKER_COUNT}" -echo "CONFIG_NAME: ${CONFIG_NAME}" -echo "CONFIG_PATH: ${CONFIG_PATH}" -echo "ENTRY_FILE: ${ENTRY_FILE}" - -args="--config_name ${CONFIG_NAME} --config_path ${CONFIG_PATH}" - -mdl_args="--queue=${QUEUE} \ - --entry=${ENTRY_FILE} \ - --worker_count=${WORKER_COUNT} \ - --file.cluster_file=examples/scripts/cluster.json \ - --oss_access_id=${OSS_ACCESS_ID} \ - --oss_access_key=${OSS_ACCESS_KEY} \ - --oss_bucket=${OSS_BUCKET} \ - --oss_endpoint=${OSS_ENDPOINT} \ - --job_name=${JOB_NAME} \ - --algo_name=pytorch2100_py310_cu128 \ - --requirements_file_name=nebula_patch/requirements/requirements_torch2100_vllm.txt \ - --oss_appendable=true \ - --_NEBULA_MODEL=${NEBULA_MODEL} \ - --nebula_model=${NEBULA_MODEL} \ - " -if [ -n "${OPENLM_TOKEN}" ]; then - mdl_args="${mdl_args} --env=OPENLM_TOKEN=${OPENLM_TOKEN}" -fi - -echo ${args} -echo ${mdl_args} - -nebulactl run mdl --user_params="${args}" $mdl_args diff --git a/examples/qwen3.5-35BA3-rlvr_megatron/rlvr_megatron_80GB.yaml b/examples/qwen3.5-35BA3-rlvr_megatron/rlvr_megatron_80GB.yaml index 042d87359..763cf3ef1 100644 --- a/examples/qwen3.5-35BA3-rlvr_megatron/rlvr_megatron_80GB.yaml +++ b/examples/qwen3.5-35BA3-rlvr_megatron/rlvr_megatron_80GB.yaml @@ -1,9 +1,3 @@ -defaults: - - ../config/deepspeed_zero@_here_ - - ../config/deepspeed_zero2@_here_ - - ../config/deepspeed_zero3@_here_ - - ../config/deepspeed_zero3_cpuoffload@_here_ - hydra: run: dir: . @@ -106,6 +100,7 @@ actor_train: moe_grouped_gemm: true moe_shared_expert_overlap: true bf16: true + #mtp_training_mode: standalone device_mapping: list(range(0,16)) infer_batch_size: 1 @@ -131,6 +126,9 @@ actor_infer: tensor_parallel_size: 2 optimization_level: 1 sleep_level: 2 + #speculative_config: + # method: mtp + # num_speculative_tokens: 3 num_gpus_per_worker: 2 device_mapping: list(range(0,16)) infer_batch_size: 1 diff --git a/examples/qwen3.5-35BA3-rlvr_megatron/rlvr_megatron_ppu.yaml b/examples/qwen3.5-35BA3-rlvr_megatron/rlvr_megatron_ppu.yaml new file mode 100644 index 000000000..c88ff45c5 --- /dev/null +++ b/examples/qwen3.5-35BA3-rlvr_megatron/rlvr_megatron_ppu.yaml @@ -0,0 +1,179 @@ +hydra: + run: + dir: . + output_subdir: null + +exp_name: "qwen3_5_35BA3_dapo_rlvr_32gpu" +seed: 42 +logging_dir: ./output/logs +output_dir: ./output +profiler_output_dir: /data/oss_bucket_0/xxxxx +system_envs: + USE_MODELSCOPE: '1' + RAY_PROFILING: "1" + +checkpoint_config: + type: file_system + output_dir: /data/cpfs_0/rl_examples/models/${exp_name} + +track_with: ml_tracker +tracker_kwargs: + project: chuye_qwen3_5_35BA3_test + +max_steps: 500 +save_steps: 100 +eval_steps: 20 +logging_steps: 1 +resume_from_checkpoint: false + +rollout_batch_size: 64 +num_return_sequences_in_group: 8 +is_num_return_sequences_expand: true +prompt_length: 2048 +response_length: 4096 + +ppo_epochs: 1 +value_clip: 0.5 +reward_clip: 10 +advantage_clip: 10.0 +whiten_advantages: false +init_kl_coef: 0.0 +adv_estimator: "grpo" +use_kl_loss: true +kl_loss_coef: 1.0e-3 + +pretrain: Qwen/Qwen3.5-35B-A3B + +# Validation - using One-RL-to-See-Them-All/Orsta-Data-47k test set +validation: + data_args: + file_name: + - /data/oss_bucket_0/yuzhao/data/One-RL-to-See-Them-All/Orsta-Data-47k/test/test_math_megabench_237.parquet + - /data/oss_bucket_0/yuzhao/data/One-RL-to-See-Them-All/Orsta-Data-47k/test/test_detection_coco_test_multi_2000.parquet + dataset_dir: ./ + generating_args: + max_new_tokens: ${response_length} + top_p: 0.99 + top_k: 100 + num_beams: 1 + temperature: 0.99 + num_return_sequences: 1 + eval_steps: ${eval_steps} + +actor_train: + model_args: + flash_attn: sdpa + attn_implementation: sdpa + disable_gradient_checkpointing: false + dtype: bf16 + model_type: ~ + freeze_module_prefix: vision_model + training_args: + learning_rate: 1.0e-6 + weight_decay: 1.0e-2 + per_device_train_batch_size: 1 + gradient_accumulation_steps: 64 + warmup_steps: 0 + num_train_epochs: 50 + data_args: + # use One-RL-to-See-Them-All/Orsta-Data-47k as train dataset + # download from https://huggingface.co/datasets/One-RL-to-See-Them-All/Orsta-Data-47k + file_name: + - /data/oss_bucket_0/yuzhao/data/One-RL-to-See-Them-All/Orsta-Data-47k/train/train_detection_v3det_4000.parquet + - /data/oss_bucket_0/yuzhao/data/One-RL-to-See-Them-All/Orsta-Data-47k/train/train_math_mmmath_3539.parquet + domain_interleave_probs: + math: 0.5 + cv_detection: 0.5 + dataset_dir: ./ + messages: prompt + preprocessing_num_workers: 32 + strategy_args: + strategy_name: megatron_train + strategy_config: + sequence_parallel: true + tensor_model_parallel_size: 1 + pipeline_model_parallel_size: 4 # 32卡（TP=1×PP=4×EP=8=32卡） + virtual_pipeline_model_parallel_size: 2 # 40层 ÷ (PP=4 × VPP=2) = 5，必须整除 + expert_model_parallel_size: 8 + context_parallel_size: 1 + use_distributed_optimizer: true + recompute_granularity: selective + recompute_modules: "moe,layernorm" + moe_token_dispatcher_type: "alltoall" + moe_grouped_gemm: true + moe_shared_expert_overlap: true + bias_activation_fusion: true + apply_rope_fusion: true + bf16: true + #mtp_training_mode: standalone + device_mapping: list(range(0,32)) # 32卡: range(0,32) + infer_batch_size: 1 + +actor_infer: + model_args: + flash_attn: fa2 + attn_implementation: fa2 + disable_gradient_checkpointing: true + dtype: bf16 + generating_args: + max_new_tokens: ${response_length} + top_p: 0.99 + top_k: 100 + num_beams: 1 + temperature: 0.99 + num_return_sequences: ${num_return_sequences_in_group} + strategy_args: + strategy_name: vllm + strategy_config: + gpu_memory_utilization: 0.65 + block_size: 16 + max_model_len: 8192 # prompt_length + response_length + load_format: dummy + tensor_parallel_size: 4 # 32卡: 4（每卡模型 17.5GB） + #enable_expert_parallel: true + optimization_level: 1 + sleep_level: 2 + device_mapping: list(range(0,32)) # 32卡: range(0,32) + infer_batch_size: 1 + +reference: + model_args: + flash_attn: sdpa + attn_implementation: sdpa + disable_gradient_checkpointing: true + dtype: bf16 + model_type: ~ + strategy_args: + strategy_name: megatron_infer + strategy_config: + tensor_model_parallel_size: 1 + pipeline_model_parallel_size: 4 # 32卡（TP=1×PP=4×EP=8=32卡） + expert_model_parallel_size: 8 + context_parallel_size: 1 + sequence_parallel: true + moe_token_dispatcher_type: "alltoall" + moe_grouped_gemm: true + moe_shared_expert_overlap: true + bias_activation_fusion: true + apply_rope_fusion: true + bf16: true + device_mapping: list(range(0,32)) # 32卡: range(0,32) + infer_batch_size: 2 + +rewards: + math: + worker_cls: roll.pipeline.rlvr.rewards.math_rule_reward_worker.MathRuleRewardWorker + model_args: + model_name_or_path: ${pretrain} + # data source whose ability is math in One-RL-to-See-Them-All/Orsta-Data-47k + tag_included: [mm_math, megabench_math] + world_size: 16 # 32卡: 16 + infer_batch_size: 1 + cv_detection: + worker_cls: roll.pipeline.rlvr.rewards.detection_reward_worker.DetectionRewardWorker + model_args: + model_name_or_path: ${pretrain} + # data source whose ability is cv_detection in One-RL-to-See-Them-All/Orsta-Data-47k + tag_included: [v3det_train, object365_train, coco_val_multi_test] + world_size: 16 # 32卡: 16 + infer_batch_size: 1 diff --git a/examples/qwen3.5-35BA3-rlvr_megatron/submit_pipeline.sh b/examples/qwen3.5-35BA3-rlvr_megatron/submit_pipeline.sh deleted file mode 100644 index 43d7d276e..000000000 --- a/examples/qwen3.5-35BA3-rlvr_megatron/submit_pipeline.sh +++ /dev/null @@ -1,45 +0,0 @@ -#!/bin/bash -set +x -source "examples/scripts/config.sh" - -WORKER_COUNT=2 -CONFIG_FILE="rlvr_megatron_80GB.yaml" -# 替换为mos uri -NEBULA_MODEL="" -ENTRY_FILE="examples/start_rlvr_vl_pipeline.py" - -CONFIG_PATH=$(basename $(dirname $0)) -CONFIG_NAME="${CONFIG_FILE%.yaml}" -JOB_NAME="$CONFIG_PATH-$CONFIG_NAME" - -echo "JOB_NAME: ${JOB_NAME}" -echo "WORKER_COUNT: ${WORKER_COUNT}" -echo "CONFIG_NAME: ${CONFIG_NAME}" -echo "CONFIG_PATH: ${CONFIG_PATH}" -echo "ENTRY_FILE: ${ENTRY_FILE}" - -args="--config_name ${CONFIG_NAME} --config_path ${CONFIG_PATH}" - -mdl_args="--queue=${QUEUE} \ - --entry=${ENTRY_FILE} \ - --worker_count=${WORKER_COUNT} \ - --file.cluster_file=examples/scripts/cluster.json \ - --oss_access_id=${OSS_ACCESS_ID} \ - --oss_access_key=${OSS_ACCESS_KEY} \ - --oss_bucket=${OSS_BUCKET} \ - --oss_endpoint=${OSS_ENDPOINT} \ - --job_name=${JOB_NAME} \ - --algo_name=pytorch2100_py310_cu128 \ - --requirements_file_name=nebula_patch/requirements/requirements_torch2100_vllm.txt \ - --oss_appendable=true \ - --_NEBULA_MODEL=${NEBULA_MODEL} \ - --nebula_model=${NEBULA_MODEL} \ - " -if [ -n "${OPENLM_TOKEN}" ]; then - mdl_args="${mdl_args} --env=OPENLM_TOKEN=${OPENLM_TOKEN}" -fi - -echo ${args} -echo ${mdl_args} - -nebulactl run mdl --user_params="${args}" $mdl_args diff --git a/examples/qwen3.5-35BA3-sft_fsdp2/sft_config.yaml b/examples/qwen3.5-35BA3-sft_fsdp2/sft_config.yaml new file mode 100644 index 000000000..310c1e256 --- /dev/null +++ b/examples/qwen3.5-35BA3-sft_fsdp2/sft_config.yaml @@ -0,0 +1,69 @@ +hydra: + run: + dir: . + output_subdir: null + +exp_name: "qwen3-30BA3B-sft-config" +seed: 42 +logging_dir: ./output/logs +output_dir: ./output_sft +system_envs: + USE_MODELSCOPE: '1' + +num_gpus_per_node: 8 + +save_steps: 100 +logging_steps: 1 +eval_steps: 100 +resume_from_checkpoint: false + +sequence_length: 2048 + +pretrain: Qwen/Qwen3.5-35B-A3B + +# sft related +# system_key: system_prompt # use the default system prompt in the tokenizer tmplate if not provided +prompt_key: instruction +query_key: input +response_key: output + +# validation: +# data_args: +# file_name: ./data/code_alpaca_500.json +# template: qwen2_5 + +sft_train: + model_args: + dtype: bf16 + ulysses_size: 1 + disable_gradient_checkpointing: true + training_args: + num_train_epochs: 1 + per_device_train_batch_size: 1 + gradient_accumulation_steps: 1 + learning_rate: 5.0e-6 + data_args: + file_name: ./data/code_alpaca_20k.json # https://huggingface.co/datasets/sahil2801/CodeAlpaca-20k + template: qwen2_5 + preprocessing_num_workers: 4 + strategy_args: + strategy_name: fsdp2_train + strategy_config: + fsdp_size: 8 + moe_use_grouped_mm: true + param_dtype: bf16 + reduce_dtype: bf16 + offload_policy: false + apply_expert_patch: false + apply_tiled_mlp: false + reshard_after_forward: true + wrap_policy: + wrap_embeddings: true + wrap_lm_output: true + transformer_layer_cls_to_wrap: + - Qwen3_5MoeAttention + - Qwen3_5MoeGatedDeltaNet + - Qwen3_5MoeSparseMoeBlock + use_sequence_packing: True + device_mapping: list(range(0,8)) + infer_batch_size: 2 diff --git a/examples/qwen3.5-35BA3-sft_fsdp2/sft_config_ep.yaml b/examples/qwen3.5-35BA3-sft_fsdp2/sft_config_ep.yaml new file mode 100644 index 000000000..9f26a0450 --- /dev/null +++ b/examples/qwen3.5-35BA3-sft_fsdp2/sft_config_ep.yaml @@ -0,0 +1,73 @@ +hydra: + run: + dir: . + output_subdir: null + +exp_name: "qwen3-30BA3B-sft-config" +seed: 42 +logging_dir: ./output/logs +output_dir: ./output_sft +system_envs: + USE_MODELSCOPE: '1' + +num_gpus_per_node: 8 + +save_steps: 100 +logging_steps: 1 +eval_steps: 100 +resume_from_checkpoint: false + +sequence_length: 2048 + +pretrain: Qwen/Qwen3.5-35B-A3B + +# sft related +# system_key: system_prompt # use the default system prompt in the tokenizer tmplate if not provided +prompt_key: instruction +query_key: input +response_key: output + +# validation: +# data_args: +# file_name: ./data/code_alpaca_500.json +# template: qwen2_5 + +sft_train: + model_args: + dtype: bf16 + ulysses_size: 1 + disable_gradient_checkpointing: true + training_args: + num_train_epochs: 1 + per_device_train_batch_size: 1 + gradient_accumulation_steps: 1 + learning_rate: 5.0e-6 + data_args: + file_name: ./data/code_alpaca_20k.json # https://huggingface.co/datasets/sahil2801/CodeAlpaca-20k + template: qwen2_5 + preprocessing_num_workers: 4 + strategy_args: + strategy_name: fsdp2_train + strategy_config: + fsdp_size: 8 + efsdp_size: 1 + ep_size: 8 + moe_use_grouped_mm: true + param_dtype: bf16 + reduce_dtype: bf16 + offload_policy: false + apply_expert_patch: false + apply_tiled_mlp: false + reshard_after_forward: true + wrap_policy: + wrap_embeddings: true + wrap_lm_output: true + moe_experts: + - Qwen3_5MoeExperts + transformer_layer_cls_to_wrap: + - Qwen3_5MoeAttention + - Qwen3_5MoeGatedDeltaNet + - Qwen3_5MoeSparseMoeBlock + use_sequence_packing: True + device_mapping: list(range(0,8)) + infer_batch_size: 2 diff --git a/examples/qwen3.5-397A17-rlvr_megatron/rlvr_megatron_80GB.yaml b/examples/qwen3.5-397A17-rlvr_megatron/rlvr_megatron_80GB.yaml new file mode 100644 index 000000000..62ad54340 --- /dev/null +++ b/examples/qwen3.5-397A17-rlvr_megatron/rlvr_megatron_80GB.yaml @@ -0,0 +1,240 @@ +hydra: + run: + dir: . + output_subdir: null + +exp_name: "Qwen3.5-397B-A17B-rlvr-config" +seed: 42 +logging_dir: ./output/logs +output_dir: ./output +system_envs: + USE_MODELSCOPE: '1' + +checkpoint_config: + type: file_system + output_dir: ./rl_examples/models/${exp_name} + +#track_with: wandb +#tracker_kwargs: +# api_key: +# project: roll_examples +# notes: roll_examples +# tags: +# - rlvr +# - baseline + +track_with: tensorboard +tracker_kwargs: + log_dir: ./rl_examples/llm/tensorboard/roll_exp/rlvr + +num_gpus_per_node: 8 + +max_steps: 500 +save_steps: 100 +logging_steps: 1 +eval_steps: 10 +resume_from_checkpoint: false + +rollout_batch_size: 32 # prompt +prompt_length: 2048 +response_length: 4096 + +num_return_sequences_in_group: 8 +ppo_epochs: 1 +adv_estimator: "reinforce" + +# clip +value_clip: 0.5 +reward_clip: 10 +advantage_clip: 2.0 +dual_clip_loss: true + +# normalize +norm_mean_type: ~ +norm_std_type: ~ + +# data mask +max_len_mask: true +difficulty_mask: true +difficulty_low_threshold: 0.1 +difficulty_high_threshold: 0.95 +error_max_len_clip: false + +# data weight +difficulty_loss_weight: false +length_loss_weight: false + +# reward +add_token_level_kl: false + +# advantage +whiten_advantages: true + +# dynamic sampling scheduler +# use_additional_prompts: true +# max_running_requests: 256 +# is_num_return_sequences_expand: false + +pretrain: Qwen/Qwen3.5-397B-A17B +reward_pretrain: Qwen/Qwen3.5-397B-A17B + +validation: + data_args: + template: qwen2_5 + file_name: + - data/math_benchmarks.jsonl + generating_args: + max_new_tokens: ${response_length} + top_p: 0.6 + top_k: 50 + num_beams: 1 + temperature: 0.6 + num_return_sequences: 1 + +actor_train: + model_args: + disable_gradient_checkpointing: false + dtype: bf16 + model_type: ~ + training_args: + learning_rate: 1.0e-6 + weight_decay: 0 + per_device_train_batch_size: 1 + gradient_accumulation_steps: 64 + warmup_steps: 20 + num_train_epochs: 50 + data_args: + template: qwen2_5 + file_name: + - data/code_KodCode_data.jsonl + - data/math_deepmath_deal.jsonl + - data/general_ifeval_train_deal.jsonl + - data/general_CrossThink-QA_deal.jsonl + domain_interleave_probs: + math_rule: 0.5 + code_sandbox: 0.3 + crossthinkqa: 0.1 + ifeval: 0.1 + dataset_dir: data + messages: messages + interleave_probs: "1.0" + preprocessing_num_workers: 16 + strategy_args: + strategy_name: megatron_train + strategy_config: + sequence_parallel: true + tensor_model_parallel_size: 1 + context_parallel_size: 4 + expert_model_parallel_size: 16 + # Try a 384-GPU layout: TP=8, PP=6, EP=8. For 60 layers, PP=6 gives + # 10 layers per stage and avoids an overly deep pipeline. + pipeline_model_parallel_size: 16 + pipeline_model_parallel_layout: "E|(tt|)*30,L" + use_distributed_optimizer: true + distrib_optim_fully_reshardable: false + moe_token_dispatcher_type: alltoall + moe_aux_loss_coeff: 0.0 + recompute_granularity: full + bias_activation_fusion: true + apply_rope_fusion: true + moe_grouped_gemm: true + moe_shared_expert_overlap: true + cross_entropy_loss_fusion: true + bf16: true + device_mapping: list(range(0,256)) + infer_batch_size: 1 + +actor_infer: + model_args: + disable_gradient_checkpointing: true + dtype: bf16 + generating_args: + max_new_tokens: ${response_length} + top_p: 0.99 + top_k: 100 + num_beams: 1 + temperature: 0.99 + num_return_sequences: ${num_return_sequences_in_group} + data_args: + template: qwen2_5 + strategy_args: + strategy_name: vllm + strategy_config: + distributed_executor_backend: ray + gpu_memory_utilization: 0.85 + max_model_len: 20480 + load_format: dummy + tensor_parallel_size: 16 + pipeline_parallel_size: 1 + optimization_level: 1 + sleep_level: 2 # 2 will destroy model parameter and kv_cache after generate to save cpu memory, 1 will destroy kv_cache only. + device_mapping: list(range(256,320)) + infer_batch_size: 1 + +reference: + model_args: + disable_gradient_checkpointing: true + dtype: bf16 + model_type: ~ + data_args: + template: qwen2_5 + strategy_args: + strategy_name: megatron_infer + strategy_config: + sequence_parallel: true + tensor_model_parallel_size: 1 + context_parallel_size: 2 + pipeline_model_parallel_size: 8 + pipeline_model_parallel_layout: "E|(tt|)*30,L" + moe_shared_expert_overlap: true + expert_model_parallel_size: 8 + cross_entropy_loss_fusion: true + moe_token_dispatcher_type: alltoall + bias_activation_fusion: true + apply_rope_fusion: true + moe_grouped_gemm: true + bf16: true + device_mapping: list(range(0,256)) + infer_batch_size: 2 + +rewards: + crossthinkqa: + worker_cls: roll.pipeline.rlvr.rewards.crossthinkqa_rule_reward_worker.CrossThinkQARuleRewardWorker + reward_type: soft + response_length_penalty_coef: 0.0 + model_args: + model_name_or_path: ${reward_pretrain} + data_args: + template: qwen2_5 + tag_included: [crossthinkqa] + world_size: 8 + infer_batch_size: 4 + ifeval: + worker_cls: roll.pipeline.rlvr.rewards.ifeval_rule_reward_worker.GeneralRuleRewardWorker + reward_type: soft + model_args: + model_name_or_path: ${reward_pretrain} + data_args: + template: qwen2_5 + tag_included: [ifeval] + world_size: 8 + infer_batch_size: 4 + math_rule: + worker_cls: roll.pipeline.rlvr.rewards.math_rule_reward_worker.MathRuleRewardWorker + model_args: + model_name_or_path: ${reward_pretrain} + data_args: + template: qwen2_5 + tag_included: [deepmath_103k, aime] + world_size: 8 + infer_batch_size: 1 + code_sandbox: + use_local: true + worker_cls: roll.pipeline.rlvr.rewards.code_sandbox_reward_worker.CodeSandboxRewardWorker + tag_included: [KodCode] + model_args: + model_name_or_path: ${reward_pretrain} + data_args: + template: qwen2_5 + world_size: 8 + infer_batch_size: 1 \ No newline at end of file diff --git a/examples/qwen3_agentic_gem/gem_game_guess_the_number.yaml b/examples/qwen3_agentic_gem/gem_game_guess_the_number.yaml index bb5294122..c8b8d2dbd 100644 --- a/examples/qwen3_agentic_gem/gem_game_guess_the_number.yaml +++ b/examples/qwen3_agentic_gem/gem_game_guess_the_number.yaml @@ -1,9 +1,5 @@ defaults: - ../config/traj_envs_gem_games@_here_ - - ../config/deepspeed_zero@_here_ - - ../config/deepspeed_zero2@_here_ - - ../config/deepspeed_zero3@_here_ - - ../config/deepspeed_zero3_cpuoffload@_here_ hydra: run: @@ -76,8 +72,6 @@ actor_train: gradient_accumulation_steps: 8 lr_scheduler_type: constant strategy_args: -# strategy_name: deepspeed_train -# strategy_config: ${deepspeed_zero3} strategy_name: megatron_train strategy_config: tensor_model_parallel_size: 1 diff --git a/examples/qwen3_agentic_gem/gem_math_dapo.yaml b/examples/qwen3_agentic_gem/gem_math_dapo.yaml index f01117858..926dd660a 100644 --- a/examples/qwen3_agentic_gem/gem_math_dapo.yaml +++ b/examples/qwen3_agentic_gem/gem_math_dapo.yaml @@ -71,8 +71,6 @@ actor_train: warmup_ratio: 0 lr_scheduler_type: constant strategy_args: -# strategy_name: deepspeed_train -# strategy_config: ${deepspeed_zero3} strategy_name: megatron_train strategy_config: tensor_model_parallel_size: 1 diff --git a/examples/qwen3_agentic_gem/gem_math_dapo_python_code.yaml b/examples/qwen3_agentic_gem/gem_math_dapo_python_code.yaml index e30a81be9..eda383e3c 100644 --- a/examples/qwen3_agentic_gem/gem_math_dapo_python_code.yaml +++ b/examples/qwen3_agentic_gem/gem_math_dapo_python_code.yaml @@ -71,8 +71,6 @@ actor_train: warmup_ratio: 0 lr_scheduler_type: constant strategy_args: -# strategy_name: deepspeed_train -# strategy_config: ${deepspeed_zero3} strategy_name: megatron_train strategy_config: tensor_model_parallel_size: 1 diff --git a/examples/qwen3_agentic_gem/gem_math_hotpotqa.yaml b/examples/qwen3_agentic_gem/gem_math_hotpotqa.yaml index 70326cdea..cd782ac91 100644 --- a/examples/qwen3_agentic_gem/gem_math_hotpotqa.yaml +++ b/examples/qwen3_agentic_gem/gem_math_hotpotqa.yaml @@ -74,8 +74,6 @@ actor_train: warmup_steps: 10 lr_scheduler_type: cosine strategy_args: -# strategy_name: deepspeed_train -# strategy_config: ${deepspeed_zero3} strategy_name: megatron_train strategy_config: tensor_model_parallel_size: 1 diff --git a/examples/qwen3_agentic_gem/gem_math_hotpotqa_search.yaml b/examples/qwen3_agentic_gem/gem_math_hotpotqa_search.yaml index 0963ac83e..6a9cbfe7a 100644 --- a/examples/qwen3_agentic_gem/gem_math_hotpotqa_search.yaml +++ b/examples/qwen3_agentic_gem/gem_math_hotpotqa_search.yaml @@ -74,8 +74,6 @@ actor_train: warmup_steps: 10 lr_scheduler_type: cosine strategy_args: -# strategy_name: deepspeed_train -# strategy_config: ${deepspeed_zero3} strategy_name: megatron_train strategy_config: tensor_model_parallel_size: 1 diff --git a/examples/qwen3_agentic_gem/gem_rg_letter_counting.yaml b/examples/qwen3_agentic_gem/gem_rg_letter_counting.yaml index 799864a13..624145c3f 100644 --- a/examples/qwen3_agentic_gem/gem_rg_letter_counting.yaml +++ b/examples/qwen3_agentic_gem/gem_rg_letter_counting.yaml @@ -1,9 +1,5 @@ defaults: - ../config/traj_envs_gem_rg@_here_ - - ../config/deepspeed_zero@_here_ - - ../config/deepspeed_zero2@_here_ - - ../config/deepspeed_zero3@_here_ - - ../config/deepspeed_zero3_cpuoffload@_here_ hydra: run: @@ -74,8 +70,6 @@ actor_train: gradient_accumulation_steps: 8 lr_scheduler_type: constant strategy_args: -# strategy_name: deepspeed_train -# strategy_config: ${deepspeed_zero3} strategy_name: megatron_train strategy_config: tensor_model_parallel_size: 1 diff --git a/examples/start_reward_fl_pipeline.py b/examples/start_reward_fl_pipeline.py deleted file mode 100644 index 4cf6c8a78..000000000 --- a/examples/start_reward_fl_pipeline.py +++ /dev/null @@ -1,36 +0,0 @@ -import argparse - -from dacite import from_dict -from hydra import compose, initialize -from omegaconf import OmegaConf - -from roll.distributed.scheduler.initialize import init -from roll.pipeline.diffusion.reward_fl.reward_fl_config import RewardFLConfig - -from roll.pipeline.diffusion.reward_fl.reward_fl_pipeline import RewardFLPipeline - - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument("--config_path", help="The path of the main configuration file", default="config") - parser.add_argument( - "--config_name", help="The name of the main configuration file (without extension).", default="reward_fl_config" - ) - args = parser.parse_args() - - initialize(config_path=args.config_path, job_name="app") - cfg = compose(config_name=args.config_name) - - print(OmegaConf.to_yaml(cfg, resolve=True)) - - reward_fl_config = from_dict(data_class=RewardFLConfig, data=OmegaConf.to_container(cfg, resolve=True)) - - init() - - pipeline = RewardFLPipeline(pipeline_config=reward_fl_config) - - pipeline.run() - - -if __name__ == "__main__": - main() diff --git a/examples/video-rlvr/video_r1.yaml b/examples/video-rlvr/video_r1.yaml new file mode 100644 index 000000000..c60ba9cc7 --- /dev/null +++ b/examples/video-rlvr/video_r1.yaml @@ -0,0 +1,160 @@ +hydra: + run: + dir: . + output_subdir: null + +exp_name: "video-r1" +seed: 42 +logging_dir: ./output/logs +output_dir: ./output + +checkpoint_config: + type: file_system + output_dir: /data/cpfs_0/yuzhao/roll/vlm-video/models + + +track_with: ml_tracker +tracker_kwargs: + project: video-r1 + notes: ${exp_name} + tags: + - roll + - rlvr + - video + +offload_nccl: false + +save_steps: 20 +logging_steps: 1 +eval_steps: 10 +resume_from_checkpoint: false + +rollout_batch_size: 8 +num_return_sequences_in_group: 8 +is_num_return_sequences_expand: true +prompt_length: 16384 +response_length: 4096 + +ppo_epochs: 1 +value_clip: 0.5 +reward_clip: 10 +advantage_clip: 10.0 +whiten_advantages: false +init_kl_coef: 0.0 +adv_estimator: "grpo" +use_kl_loss: true +kl_loss_coef: 1.0e-2 + +pretrain: Qwen/Qwen2.5-VL-7B-Instruct + +actor_train: + model_args: + attn_implementation: fa2 + disable_gradient_checkpointing: false + dtype: bf16 + model_type: ~ + # freeze_module_prefix: "vision_model.blocks,vision_model.patch_embed" + training_args: + learning_rate: 1.0e-6 + weight_decay: 1.0e-2 + per_device_train_batch_size: 2 + gradient_accumulation_steps: 8 + warmup_steps: 0 + num_train_epochs: 50 + data_args: + # we use video-r1 116k video data with adjusted format, case for training data: + # { + # "modality_type": "video", + # "problem": "Why does the video conclude with a red screen displaying text?Options:\nA. To promote the BBC one program 'Super Cute Animals'\nB. To provide a warning message\nC. To show the credits of the video\nD. To indicate a change in scene\n", + # "solution": "A", + # "video": "videos/youtube_video_2024/ytb_HBxn56l9WcU.mp4", + # "problem_type": "multiple choice" + # } + custom_data_kwargs_func: roll.datasets.vlm_dataset_utils.video_r1_get_data_kwargs + file_name: + - /data/oss_bucket_0/yinchao/dataset/LMMS/Video-R1-116k-video.json + domain_interleave_probs: + videor1: 1.0 + dataset_dir: ./ + messages: prompt + preprocessing_num_workers: 32 + # settings same with RewatchR1 and VideoR1 + video_total_pixels: 9633792 # 12288 * 28 * 28, max_video_token=12k + video_min_pixels: 100352 # 128 * 28 * 28 + video_max_pixels: 100352 # 128 * 28 * 28 + video_fps: 2.0 + video_max_frames: 192 + video_folder: /data/oss_bucket_0/ + strategy_args: + strategy_name: megatron_train + strategy_config: + sequence_parallel: true + tensor_model_parallel_size: 2 + context_parallel_size: 2 + expert_model_parallel_size: 1 + pipeline_model_parallel_size: 1 + overlap_grad_reduce: true + use_distributed_optimizer: true + recompute_granularity: full + bf16: true + device_mapping: list(range(0,16)) + infer_batch_size: 4 + offload_nccl: ${offload_nccl} + +actor_infer: + model_args: + attn_implementation: fa2 + disable_gradient_checkpointing: true + dtype: bf16 + generating_args: + max_new_tokens: ${response_length} + top_p: 0.99 + top_k: 100 + num_beams: 1 + temperature: 0.99 + num_return_sequences: ${num_return_sequences_in_group} + strategy_args: + strategy_name: vllm + strategy_config: + tensor_parallel_size: 2 + max_num_batched_tokens: 32768 + gpu_memory_utilization: 0.8 + block_size: 16 + disable_mm_preprocessor_cache: true + enable_prefix_caching: false + device_mapping: list(range(0,16)) + infer_batch_size: 32 + offload_nccl: ${offload_nccl} + +reference: + model_args: + attn_implementation: fa2 + disable_gradient_checkpointing: true + dtype: bf16 + model_type: ~ + strategy_args: + strategy_name: megatron_infer + strategy_config: + sequence_parallel: true + tensor_model_parallel_size: 2 + context_parallel_size: 2 + pipeline_model_parallel_size: 1 + expert_model_parallel_size: 1 + bf16: true + device_mapping: list(range(0,16)) + infer_batch_size: 4 + offload_nccl: ${offload_nccl} + +reward_system_config: + skip_special_tokens: true + +rewards: + videor1: + # NOTE: `re.fullmatch(pattern)` in format_reward seems not suitable for qwen3-vl + # thinking model, `` is the prompt ending suffix thus adjust format_reward + worker_cls: roll.pipeline.rlvr.rewards.video_r1_reward_worker.VideoR1RewardWorker + model_args: + model_name_or_path: ${pretrain} + tag_included: ["video"] + world_size: 2 + infer_batch_size: 1 \ No newline at end of file diff --git a/examples/wan2.2-14B-reward_fl_ds/merge_lora.py b/examples/wan2.2-14B-reward_fl_ds/merge_lora.py deleted file mode 100644 index 2f1b65ce4..000000000 --- a/examples/wan2.2-14B-reward_fl_ds/merge_lora.py +++ /dev/null @@ -1,94 +0,0 @@ -import os -import argparse -import torch -from safetensors import safe_open -from safetensors.torch import save_file - -def load_state_dict_from_safetensors(file_path, device="cpu"): - state_dict = {} - with safe_open(file_path, framework="pt", device=str(device)) as f: - for k in f.keys(): - state_dict[k] = f.get_tensor(k) - return state_dict - -def merge_lora_into_state_dict( - base_state_dict, - lora_state_dict, - alpha=1.0, - device="cpu", - dtype=torch.bfloat16 -): - lora_name_map = {} - for key in lora_state_dict: - if ".lora_B." not in key: - continue - clean_key = key.replace("base_model.model.", "") - parts = clean_key.split(".") - try: - lora_b_idx = parts.index("lora_B") - except ValueError: - continue - target_parts = parts[:lora_b_idx] + parts[lora_b_idx + 2:] # 跳过 lora_B 和 rank idx - if target_parts[0] == "diffusion_model": - target_parts = target_parts[1:] - target_name = ".".join(target_parts) - lora_A_key = key.replace(".lora_B.", ".lora_A.").replace("base_model.model.", "") - lora_name_map[target_name] = (key, lora_A_key) - - merged_state_dict = base_state_dict.copy() - updated = 0 - - for target_name, (lora_B_key, lora_A_key) in lora_name_map.items(): - if target_name not in merged_state_dict: - print(f"Warning: {target_name} not in base model. Skipping.") - continue - - weight_orig = merged_state_dict[target_name].to(device=device, dtype=dtype) - weight_up = lora_state_dict[lora_B_key].to(device=device, dtype=dtype) - weight_down = lora_state_dict[lora_A_key].to(device=device, dtype=dtype) - - if len(weight_up.shape) == 4: - weight_up = weight_up.squeeze(-1).squeeze(-1) # [out, r, 1, 1] -> [out, r] - weight_down = weight_down.squeeze(-1).squeeze(-1) # [r, in, 1, 1] -> [r, in] - lora_weight = alpha * (weight_up @ weight_down).unsqueeze(-1).unsqueeze(-1) - else: - lora_weight = alpha * (weight_up @ weight_down) - - merged_state_dict[target_name] = weight_orig + lora_weight - updated += 1 - - print(f"Merged {updated} LoRA adapters into base model.") - return merged_state_dict - -parser = argparse.ArgumentParser() -parser.add_argument('--lora_dir', type=str, default="/data/models/Wan2.2-I2V-A14B-4steps-lora-rank64-Seko-V1/high_noise_model.safetensors") -parser.add_argument('--base_model_path', type=str, default="/data/models/Wan22_base/high_noise_model/diffusion_pytorch_model.safetensors") -parser.add_argument('--save_dir', type=str, default="/data/models/Wan22/high_noise_model/diffusion_pytorch_model.safetensors") -parser.add_argument('--alpha', type=float, default=1.0) -args = parser.parse_args() - -print("Loading base model state dict...") -base_sd = load_state_dict_from_safetensors(args.base_model_path) - -# load lora -lora_sd = load_state_dict_from_safetensors(args.lora_dir) - -clean_lora_sd = {} -for k, v in lora_sd.items(): - clean_k = k.replace("base_model.model.", "") - clean_lora_sd[clean_k] = v - -# merge -merged_sd = merge_lora_into_state_dict( - base_state_dict=base_sd, - lora_state_dict=clean_lora_sd, - alpha=args.alpha, - device="cpu", - dtype=torch.bfloat16 -) - -# save -os.makedirs(os.path.dirname(args.save_dir), exist_ok=True) -save_file(merged_sd, args.save_dir) -print(f"Merged model saved to {args.save_dir}") - diff --git a/examples/wan2.2-14B-reward_fl_ds/merge_model.py b/examples/wan2.2-14B-reward_fl_ds/merge_model.py deleted file mode 100644 index cc7e97ba7..000000000 --- a/examples/wan2.2-14B-reward_fl_ds/merge_model.py +++ /dev/null @@ -1,23 +0,0 @@ -import json -from safetensors import safe_open -from safetensors.torch import save_file -import torch - -shard_files = ['diffusion_pytorch_model-00001-of-00006.safetensors', 'diffusion_pytorch_model-00002-of-00006.safetensors', -"diffusion_pytorch_model-00003-of-00006.safetensors","diffusion_pytorch_model-00004-of-00006.safetensors", -"diffusion_pytorch_model-00005-of-00006.safetensors","diffusion_pytorch_model-00006-of-00006.safetensors"] -print("Shard files to load:", shard_files) - -full_state_dict = {} - -for shard_file in shard_files: - print(f"Loading {shard_file}...") - with safe_open(shard_file, framework="pt", device="cpu") as f: - for key in f.keys(): - full_state_dict[key] = f.get_tensor(key).to(dtype=torch.bfloat16) - -output_file = "diffusion_pytorch_model.safetensors" -print(f"Saving merged model to {output_file}...") -save_file(full_state_dict, output_file) - -print("Done! Merged model saved.") \ No newline at end of file diff --git a/examples/wan2.2-14B-reward_fl_ds/reward_fl_config.yaml b/examples/wan2.2-14B-reward_fl_ds/reward_fl_config.yaml deleted file mode 100644 index 7c9c2b5fd..000000000 --- a/examples/wan2.2-14B-reward_fl_ds/reward_fl_config.yaml +++ /dev/null @@ -1,70 +0,0 @@ -defaults: - - ../config/deepspeed_zero@_here_ - - ../config/deepspeed_zero2@_here_ - - ../config/deepspeed_zero2_cpuoffload@_here_ - - ../config/deepspeed_zero3@_here_ - - ../config/deepspeed_zero3_cpuoffload@_here_ - -hydra: - run: - dir: . - output_subdir: null - -exp_name: "reward_fl_zero2_cpuoffload" -seed: 42 -logging_dir: ./output/logs -output_dir: ./output - -checkpoint_config: - type: file_system - output_dir: /data/models/reward_fl/ - -save_steps: 25 -logging_steps: 1 -resume_from_checkpoint: false - -sequence_length: 1024 -train_batch_size: 8 -max_grad_norm: 1.0 - -actor_train: - model_args: - model_type: diffusion_module - dtype: bf16 - model_config_kwargs: - model_name: wan2_2 - model_paths: ./examples/wan2.2-14B-reward_fl_ds/wan22_paths.json - reward_model_path: /data/models/antelopev2/ - tokenizer_path: /data/models/Wan-AI/Wan2.1-T2V-1.3B/google/umt5-xxl/ - model_id_with_origin_paths: null - trainable_models: dit2 - use_gradient_checkpointing_offload: true - extra_inputs: input_image - max_timestep_boundary: 1.0 - min_timestep_boundary: 0.9 - num_inference_steps: 8 - mid_timestep: 4 - final_timestep: 7 - lora_base_model: dit2 - lora_target_modules: q,k,v,o,ffn.0,ffn.2 - lora_rank: 32 - - training_args: - learning_rate: 2.5e-6 - lr_scheduler_type: constant - per_device_train_batch_size: 1 - gradient_accumulation_steps: 1 - warmup_steps: 10 - num_train_epochs: 1 - - data_args: - file_name: ./data/example_video_dataset/metadata.csv - preprocessing_num_workers: 2 - - strategy_args: - strategy_name: diffusion_deepspeed_train - strategy_config: ${deepspeed_zero2_cpuoffload} - device_mapping: list(range(0,8)) - -system_envs: - RAY_PROFILING: "0" diff --git a/examples/wan2.2-14B-reward_fl_ds/run_reward_fl_ds_pipeline.sh b/examples/wan2.2-14B-reward_fl_ds/run_reward_fl_ds_pipeline.sh deleted file mode 100644 index 0b95fe970..000000000 --- a/examples/wan2.2-14B-reward_fl_ds/run_reward_fl_ds_pipeline.sh +++ /dev/null @@ -1,6 +0,0 @@ -#!/bin/bash -set +x - - -CONFIG_PATH=$(basename $(dirname $0)) -python examples/start_reward_fl_pipeline.py --config_path $CONFIG_PATH --config_name reward_fl_config diff --git a/examples/wan2.2-14B-reward_fl_ds/wan22_paths.json b/examples/wan2.2-14B-reward_fl_ds/wan22_paths.json deleted file mode 100644 index f05ec8ec2..000000000 --- a/examples/wan2.2-14B-reward_fl_ds/wan22_paths.json +++ /dev/null @@ -1,6 +0,0 @@ -[ - "/data/models/Wan22/high_noise_model/diffusion_pytorch_model.safetensors", - "/data/models/Wan22/low_noise_model/diffusion_pytorch_model.safetensors", - "/data/models/Wan2.1-T2V-1.3B/models_t5_umt5-xxl-enc-bf16.pth", - "/data/models/Wan2.1-T2V-1.3B/Wan2.1_VAE.pth" -] diff --git a/mcore_adapter/pyproject.toml b/mcore_adapter/pyproject.toml index 3ef6bf6be..a66168e2f 100644 --- a/mcore_adapter/pyproject.toml +++ b/mcore_adapter/pyproject.toml @@ -4,18 +4,18 @@ build-backend = "setuptools.build_meta" [project] name = "mcore_adapter" -version = "0.9.0" +version = "0.10.0.dev0" description = "" requires-python = ">=3.8.13" dependencies = [ - "megatron-core>=0.15.0,<0.17.0", + "megatron-core>=0.15.0,<0.19.0", "transformers>=4.50.0", "accelerate>=0.27.2", ] [project.optional-dependencies] dev = [ - "megatron-core==0.16.0.dev0", + "megatron-core==0.18.0.dev0", ] llama = [ "megatron-llama-core", diff --git a/mcore_adapter/src/mcore_adapter/checkpointing.py b/mcore_adapter/src/mcore_adapter/checkpointing.py index df5c2667b..efdfba6af 100644 --- a/mcore_adapter/src/mcore_adapter/checkpointing.py +++ b/mcore_adapter/src/mcore_adapter/checkpointing.py @@ -2,8 +2,9 @@ import torch from megatron.core import dist_checkpointing, mpu +from megatron.core.dist_checkpointing.strategies.fully_parallel import FullyParallelSaveStrategyWrapper -from .constants import TRACKER_FILENAME +from .constants import DIST_MODEL_DIR, DIST_OPTIMIZER_DIR, TRACKER_FILENAME from .utils import get_logger @@ -209,6 +210,13 @@ def find_checkpoint_rank_0(checkpoints_path, iteration, release=False): return None +def find_dist_ckpt(checkpoint_dir): + dist_model_dir = os.path.join(get_checkpoint_dir(checkpoint_dir, return_base_dir=True), DIST_MODEL_DIR) + if os.path.exists(dist_model_dir) and dist_checkpointing.check_is_distributed_checkpoint(dist_model_dir): + return dist_model_dir + return None + + def _load_base_checkpoint( load_dir, rank0=False, sharded_state_dict=None, exit_on_missing_checkpoint=True, checkpoint_step=None ): @@ -260,20 +268,52 @@ def _load_base_checkpoint( return state_dict, checkpoint_name, release -def load_state_dict_from_checkpoint(checkpoint_dir): - # TODO(LZC): support distributed checkpoint - return _load_base_checkpoint(checkpoint_dir, exit_on_missing_checkpoint=False)[0] +def _load_dist_checkpoint(checkpoint_dir, models): + logger.info(f"Reading dist model checkpoint from {checkpoint_dir}") + sharded_state_dict = models.sharded_state_dict() + state_dict = dist_checkpointing.load(sharded_state_dict, checkpoint_dir) + return state_dict + + +def load_state_dict_from_checkpoint(checkpoint_dir, models = None): + dist_model_dir = find_dist_ckpt(checkpoint_dir) + if dist_model_dir: + return _load_dist_checkpoint(dist_model_dir, models) + else: + return _load_base_checkpoint(checkpoint_dir, exit_on_missing_checkpoint=False)[0] -def save_config_and_state_dict(save_directory, config, state_dict): +def save_config_and_state_dict(save_directory, config, state_dict, ckpt_format: str = "legacy"): # TODO: better directory structure tracker_file = get_checkpoint_tracker_filename(save_directory) if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0: config.save_pretrained(save_directory) with open(tracker_file, "w") as f: f.write("1") - if not torch.distributed.is_initialized() or mpu.get_expert_data_parallel_rank() == 0: - checkpoint_name = get_checkpoint_name(save_directory) - ensure_directory_exists(checkpoint_name) - torch.save(state_dict, checkpoint_name) - logger.info(f"Saving model checkpoint to {checkpoint_name}") + if ckpt_format == "legacy": + if not torch.distributed.is_initialized() or mpu.get_expert_data_parallel_rank() == 0: + checkpoint_name = get_checkpoint_name(save_directory) + ensure_directory_exists(checkpoint_name) + torch.save(state_dict, checkpoint_name) + logger.info(f"Saving model checkpoint to {checkpoint_name}") + elif ckpt_format == "torch_dist": + checkpoint_dir = get_checkpoint_dir(save_directory, return_base_dir=True) + dist_model_dir = os.path.join(checkpoint_dir, DIST_MODEL_DIR) + os.makedirs(dist_model_dir, exist_ok=True) + logger.info(f"Saving distributed model checkpoint to {dist_model_dir}") + save_strategy = FullyParallelSaveStrategyWrapper( + dist_checkpointing.serialization.get_default_save_sharded_strategy(), + mpu.get_data_parallel_group(with_context_parallel=True), + do_cache_distribution=True, + ) + dist_checkpointing.save(state_dict, dist_model_dir, + sharded_strategy=save_strategy) + + +def generate_model_state_dict(model, ckpt_format: str = "legacy"): + if ckpt_format == "legacy": + return model.state_dict_for_save_checkpoint() + elif ckpt_format == "torch_dist": + return model.sharded_state_dict() + else: + raise ValueError(f"Unknown ckpt_format: {ckpt_format}") diff --git a/mcore_adapter/src/mcore_adapter/constants.py b/mcore_adapter/src/mcore_adapter/constants.py index 6e1c44724..63abe659a 100644 --- a/mcore_adapter/src/mcore_adapter/constants.py +++ b/mcore_adapter/src/mcore_adapter/constants.py @@ -1,6 +1,7 @@ TRACKER_FILENAME = "latest_checkpointed_iteration.txt" IGNORE_INDEX = -100 MCA_CONFIG_NAME = "mca_config.json" +DIST_MODEL_DIR = "dist_model" DIST_OPTIMIZER_DIR = "dist_optimizer" HUGGINGFACE_AUTOMAP_CACHE = "./.cache/huggingface/automap" diff --git a/mcore_adapter/src/mcore_adapter/initialize.py b/mcore_adapter/src/mcore_adapter/initialize.py index fa8f70457..7357bfa4c 100644 --- a/mcore_adapter/src/mcore_adapter/initialize.py +++ b/mcore_adapter/src/mcore_adapter/initialize.py @@ -65,6 +65,7 @@ def _initialize_distributed(args: "TrainingArguments"): virtual_pipeline_model_parallel_size=args.virtual_pipeline_model_parallel_size, context_parallel_size=args.context_parallel_size if args.context_parallel_size is not None else 1, expert_model_parallel_size=args.expert_model_parallel_size, + expert_tensor_parallel_size=args.expert_tensor_parallel_size, distributed_timeout_minutes=args.ddp_timeout_delta.total_seconds() // 60, ) logger.info(f"initialized tensor model parallel with size {mpu.get_tensor_model_parallel_world_size()}") diff --git a/mcore_adapter/src/mcore_adapter/models/converter/convert_utils.py b/mcore_adapter/src/mcore_adapter/models/converter/convert_utils.py index 3cf7f805d..39f9911fc 100644 --- a/mcore_adapter/src/mcore_adapter/models/converter/convert_utils.py +++ b/mcore_adapter/src/mcore_adapter/models/converter/convert_utils.py @@ -19,7 +19,7 @@ MCA_LAYER_PREFIX = "decoder.layers." MCA_MOE_PREFIX = ".mlp.experts.local_experts." MCA_MTP_PREFIX = "mtp.layers." -MCA_MTP_MOE_PREFIX = ".transformer_layer.mlp.experts.local_experts." +MCA_MTP_MOE_PREFIX = ".mtp_model_layer.mlp.experts.local_experts." MAX_SHARD_SIZE = 5_000_000_000 # 5GB @@ -38,7 +38,7 @@ def get_layer_index(weight_name: str, prefix: str) -> Optional[int]: def get_moe_index(weight_name: str, prefix: str, moe_prefix: str) -> Optional[int]: """ 1. megatron format: decoder.layers.{layer_index}.mlp.experts.local_experts.{moe_index}.{weight} -> moe_index - 2. mtp format: mtp.layers.{layer_index}.transformer_layer.mlp.experts.local_experts.{moe_index}.{weight} -> moe_index + 2. mtp format: mtp.layers.{layer_index}.mtp_model_layer.mlp.experts.local_experts.{moe_index}.{weight} -> moe_index """ if not weight_name.startswith(prefix): return None @@ -147,7 +147,7 @@ def remove_mca_weight_prefix(weight_name: str): def remove_mca_mtp_weight_prefix(weight_name: str): - return remove_weight_prefix(weight_name, MCA_MTP_PREFIX, MCA_MTP_MOE_PREFIX).replace(".transformer_layer", "") + return remove_weight_prefix(weight_name, MCA_MTP_PREFIX, MCA_MTP_MOE_PREFIX).replace(".mtp_model_layer", "") def get_mca_moe_index(weight_name: str): @@ -170,10 +170,10 @@ def add_mca_mtp_layer_prefix(weight_name: str, layer_index: Union[int, str], moe return weight_name if moe_index is not None: weight_name = add_layer_prefix(weight_name, moe_index, MCA_MTP_MOE_PREFIX) - has_transformer_layer = ".transformer_layer" not in weight_name and ( + has_transformer_layer = ".mtp_model_layer" not in weight_name and ( "self_attention" in weight_name or "mlp" in weight_name or "input_layernorm" in weight_name ) - return MCA_MTP_PREFIX + str(layer_index) + (".transformer_layer" if has_transformer_layer else "") + weight_name + return MCA_MTP_PREFIX + str(layer_index) + (".mtp_model_layer" if has_transformer_layer else "") + weight_name def extract_suffix_number(s): diff --git a/mcore_adapter/src/mcore_adapter/models/converter/dist_converter.py b/mcore_adapter/src/mcore_adapter/models/converter/dist_converter.py index 00ac1478e..0fe7aa942 100644 --- a/mcore_adapter/src/mcore_adapter/models/converter/dist_converter.py +++ b/mcore_adapter/src/mcore_adapter/models/converter/dist_converter.py @@ -286,6 +286,7 @@ def __init__( tensor_model_parallel_rank: int = 0, pipeline_model_parallel_rank: int = 0, expert_model_parallel_rank: int = 0, + expert_tensor_parallel_rank: int = 0, virtual_pipeline_model_parallel_rank: int = 0, revert: bool = False, efficient_mode: bool = False, # not check the consistency of weights @@ -295,6 +296,7 @@ def __init__( self.tensor_model_parallel_rank = tensor_model_parallel_rank or 0 self.pipeline_model_parallel_rank = pipeline_model_parallel_rank or 0 self.expert_model_parallel_rank = expert_model_parallel_rank or 0 + self.expert_tensor_parallel_rank = expert_tensor_parallel_rank or 0 self.virtual_pipeline_model_parallel_rank = virtual_pipeline_model_parallel_rank or 0 self.swiglu = mca_config.swiglu self.revert = revert @@ -345,7 +347,8 @@ def is_on_this_rank(self, weight_name: str, vp_stage: Optional[int] = None): def on_this_pipeline(): if self.pipeline_model_parallel_rank is None: return True - if weight_name.startswith("mtp.layers."): + # Handle all MTP weights (mtp.layers.* and mtp.* for non-layer weights like enorm, hnorm, eh_proj, final_layernorm) + if weight_name.startswith("mtp."): return self.mca_config.mtp_num_layers and self.is_pipeline_last_stage(vp_stage=vp_stage) if self.name_match(weight_name, self.config.pre_process_weights): @@ -383,78 +386,96 @@ def is_pipeline_last_stage(self, vp_stage: int): def is_pipeline_first_stage(self, vp_stage: int): return self.pipeline_model_parallel_rank == 0 and vp_stage == 0 - def _convert_column_parallel(self, weight: "Tensor"): - return torch.chunk(weight, self.mca_config.tensor_model_parallel_size, dim=0)[self.tensor_model_parallel_rank] - - def _revert_column_parallel(self, weights: list["Tensor"]): - assert len(weights) == self.mca_config.tensor_model_parallel_size + def _convert_column_parallel(self, name: str, weight: "Tensor"): + if self.is_expert_parallel_weight(name): + return torch.chunk(weight, self.mca_config.expert_tensor_parallel_size, dim=0)[ + self.expert_tensor_parallel_rank + ] + return torch.chunk(weight, self.mca_config.tensor_model_parallel_size, dim=0)[ + self.tensor_model_parallel_rank + ] + + def _revert_column_parallel(self, name: str, weights: list["Tensor"]): + tp_size = self.mca_config.tensor_model_parallel_size + if self.is_expert_parallel_weight(name): + tp_size = self.mca_config.expert_tensor_parallel_size + assert len(weights) == tp_size if len(weights) == 1: return weights[0] return torch.cat(weights, dim=0) def handle_column_parallel(self, name: str, weights: Union["Tensor", list["Tensor"]]) -> dict[str, "Tensor"]: if self.revert: - weight = self._revert_column_parallel(weights) + weight = self._revert_column_parallel(name, weights) else: - weight = self._convert_column_parallel(weights) + weight = self._convert_column_parallel(name, weights) name = self._name_relocate(name) return {name: weight} - def _convert_row_parallel(self, weight: "Tensor"): - return torch.chunk(weight, self.mca_config.tensor_model_parallel_size, dim=1)[self.tensor_model_parallel_rank] - - def _revert_row_parallel(self, weights: list["Tensor"]): - assert len(weights) == self.mca_config.tensor_model_parallel_size + def _convert_row_parallel(self, name: str, weight: "Tensor"): + if self.is_expert_parallel_weight(name): + return torch.chunk(weight, self.mca_config.expert_tensor_parallel_size, dim=1)[ + self.expert_tensor_parallel_rank + ] + return torch.chunk(weight, self.mca_config.tensor_model_parallel_size, dim=1)[ + self.tensor_model_parallel_rank + ] + + def _revert_row_parallel(self, name: str, weights: list["Tensor"]): + tp_size = self.mca_config.tensor_model_parallel_size + if self.is_expert_parallel_weight(name): + tp_size = self.mca_config.expert_tensor_parallel_size + assert len(weights) == tp_size if len(weights) == 1: return weights[0] return torch.cat(weights, dim=1) def handle_row_parallel(self, name: str, weights: Union["Tensor", list["Tensor"]]) -> dict[str, "Tensor"]: if self.revert: - weight = self._revert_row_parallel(weights) + weight = self._revert_row_parallel(name, weights) else: - weight = self._convert_row_parallel(weights) + weight = self._convert_row_parallel(name, weights) name = self._name_relocate(name) return {name: weight} - def _convert_swiglu(self, weight: "Tensor"): + def _convert_swiglu(self, name: str, weight: "Tensor"): assert self.swiglu and isinstance(weight, StackedTensors) and len(weight.tensors) == 2 and weight.dim == 0, ( f"weight: {weight} swiglu: {self.swiglu}" ) - weight_w = self._convert_column_parallel(weight.tensors[0]) - weight_v = self._convert_column_parallel(weight.tensors[1]) + weight_w = self._convert_column_parallel(name, weight.tensors[0]) + weight_v = self._convert_column_parallel(name, weight.tensors[1]) return torch.cat([weight_w, weight_v], dim=0) - def _revert_swiglu(self, weights: list["Tensor"]): + def _revert_swiglu(self, name: str, weights: list["Tensor"]): weights = [torch.chunk(weight, 2, dim=0) for weight in weights] weights_w = [weight_w[0] for weight_w in weights] weights_v = [weight_v[1] for weight_v in weights] - weight_w = self._revert_column_parallel(weights_w) - weight_v = self._revert_column_parallel(weights_v) + weight_w = self._revert_column_parallel(name, weights_w) + weight_v = self._revert_column_parallel(name, weights_v) return StackedTensors([weight_w, weight_v], dim=0) def handle_swiglu(self, name: str, weights: Union["Tensor", list["Tensor"]]) -> dict[str, "Tensor"]: if self.revert: - weight = self._revert_swiglu(weights) + weight = self._revert_swiglu(name, weights) else: - weight = self._convert_swiglu(weights) + weight = self._convert_swiglu(name, weights) name = self._name_relocate(name) return {name: weight} - def _convert_gdn(self, weight: "StackedTensors"): + def _convert_gdn(self, name: str, weight: "StackedTensors"): # q, k, v, z, b, a for in_proj # or q, k, v for conv1d assert self.swiglu and isinstance(weight, StackedTensors) and weight.dim == 0, ( f"weight: {weight} swiglu: {self.swiglu}" ) - return torch.cat([self._convert_column_parallel(weight.tensors[i]) for i in range(len(weight.tensors))], dim=0) + return torch.cat([self._convert_column_parallel(name, weight.tensors[i]) for i in range(len(weight.tensors))], dim=0) - def _revert_gdn(self, weights: list["Tensor"], split_shape: list[int]): + def _revert_gdn(self, name: str, weights: list["Tensor"], split_shape: list[int]): weights = [torch.split(weight, split_shape, dim=0) for weight in weights] converted_weights = [] for i in range(len(split_shape)): split_weights = [weight[i] for weight in weights] - converted_weight = self._revert_column_parallel(split_weights) + converted_weight = self._revert_column_parallel(name, split_weights) converted_weights.append(converted_weight) return StackedTensors(converted_weights, dim=0) @@ -480,9 +501,9 @@ def handle_gdn(self, name: str, weights: Union["Tensor", "StackedTensors", list[ ] elif "conv1d" in name: split_shape = [local_qk_dim, local_qk_dim, local_v_dim] - weight = self._revert_gdn(weights, split_shape) + weight = self._revert_gdn(name, weights, split_shape) else: - weight = self._convert_gdn(weights) + weight = self._convert_gdn(name, weights) name = self._name_relocate(name) return {name: weight} @@ -599,9 +620,9 @@ def handle_grouped_duplicated(self, name: str, weights: Union["Tensor", list["Te def _convert_te_grouped_column(self, name: str, weights: "Tensor"): if self.swiglu: - weights = self._convert_swiglu(weights) + weights = self._convert_swiglu(name, weights) else: - weights = self._convert_column_parallel(weights) + weights = self._convert_column_parallel(name, weights) # weights = weights.transpose(0, 1) moe_index = get_mca_moe_index(name) % self.num_layers_for_expert relocated_name = self._name_relocate(name) + str(moe_index) @@ -609,17 +630,17 @@ def _convert_te_grouped_column(self, name: str, weights: "Tensor"): def _revert_te_grouped_column(self, name: str, weights: list["Tensor"], moe_index_preprocessed: bool = False): if self.swiglu: - weight = self._revert_swiglu(weights) + weight = self._revert_swiglu(name, weights) else: - weight = self._revert_column_parallel(weights) + weight = self._revert_column_parallel(name, weights) moe_index = int(extract_suffix_number(name)) return {self._name_relocate(name, moe_index=moe_index, moe_index_preprocessed=moe_index_preprocessed): weight} def _convert_grouped_column(self, name: str, weights: "Tensor"): if self.swiglu: - weights = self._convert_swiglu(weights) + weights = self._convert_swiglu(name, weights) else: - weights = self._convert_column_parallel(weights) + weights = self._convert_column_parallel(name, weights) weights = weights.transpose(0, 1) relocated_name = self._name_relocate(name) moe_index = get_mca_moe_index(name) % self.num_layers_for_expert @@ -643,13 +664,13 @@ def _revert_grouped(weight: "Tensor"): # [expert_num_per_rank, tp] ungrouped_weights = [[weights[i] for weights in ungrouped_weights] for i in range(self.num_layers_for_expert)] - def _revert_column(weights: list["Tensor"]): + def _revert_column(name: str, weights: list["Tensor"]): if self.swiglu: - return self._revert_swiglu(weights) + return self._revert_swiglu(name, weights) else: - return self._revert_column_parallel(weights) + return self._revert_column_parallel(name, weights) - ungrouped_weights = [_revert_column(weights) for weights in ungrouped_weights] + ungrouped_weights = [_revert_column(name, weights) for weights in ungrouped_weights] return { self._name_relocate(name, moe_index=moe_index): weight for moe_index, weight in enumerate(ungrouped_weights) @@ -668,18 +689,18 @@ def handle_grouped_column( return self._convert_grouped_column(name, weights) def _convert_te_grouped_row(self, name: str, weights: "Tensor"): - weights = self._convert_row_parallel(weights) + weights = self._convert_row_parallel(name, weights) moe_index = get_mca_moe_index(name) % self.num_layers_for_expert relocated_name = self._name_relocate(name) + str(moe_index) return {relocated_name: weights} def _revert_te_grouped_row(self, name: str, weights: list["Tensor"], moe_index_preprocessed: bool = False): - weights = self._revert_row_parallel(weights) + weights = self._revert_row_parallel(name, weights) moe_index = int(extract_suffix_number(name)) return {self._name_relocate(name, moe_index=moe_index, moe_index_preprocessed=moe_index_preprocessed): weights} def _convert_grouped_row(self, name: str, weights: "Tensor"): - weights = self._convert_row_parallel(weights) + weights = self._convert_row_parallel(name, weights) weights = weights.transpose(0, 1) relocated_name = self._name_relocate(name) moe_index = get_mca_moe_index(name) % self.num_layers_for_expert @@ -702,7 +723,7 @@ def _revert_grouped(weight: "Tensor"): ungrouped_weights = [_revert_grouped(weight) for weight in weights] # [expert_num_per_rank, tp] ungrouped_weights = [[weights[i] for weights in ungrouped_weights] for i in range(self.num_layers_for_expert)] - ungrouped_weights = [self._revert_row_parallel(weights) for weights in ungrouped_weights] + ungrouped_weights = [self._revert_row_parallel(name, weights) for weights in ungrouped_weights] return { self._name_relocate(name, moe_index=moe_index): weight for moe_index, weight in enumerate(ungrouped_weights) @@ -766,19 +787,22 @@ def preprocess_layer_index(self, name: str, vp_stage: int) -> str: """ Preprocess layer index for pipeline parallelism. Converts between global and local layer indices before calling name_relocate. + MTP layer indices are independent of PP mapping and are kept as-is. """ layer_index = get_mca_layer_index(name) if layer_index is None: return name moe_index = get_mca_moe_index(name) + if name.startswith("mtp.layers."): + # MTP layer indices are independent of PP mapping, skip conversion + return add_mca_mtp_layer_prefix(remove_mca_weight_prefix(name), layer_index, moe_index) + if self.revert: layer_index = self.get_global_layer_index(layer_index, vp_stage=vp_stage) else: layer_index = self.get_local_layer_index(layer_index) - if name.startswith("mtp.layers."): - return add_mca_mtp_layer_prefix(remove_mca_weight_prefix(name), layer_index, moe_index) return add_mca_layer_prefix(remove_mca_weight_prefix(name), layer_index, moe_index) def dist_convert( diff --git a/mcore_adapter/src/mcore_adapter/models/converter/model_converter.py b/mcore_adapter/src/mcore_adapter/models/converter/model_converter.py index a6cc7a110..c69744520 100644 --- a/mcore_adapter/src/mcore_adapter/models/converter/model_converter.py +++ b/mcore_adapter/src/mcore_adapter/models/converter/model_converter.py @@ -54,6 +54,7 @@ def __init__( tensor_model_parallel_rank: Optional[int] = None, pipeline_model_parallel_rank: Optional[int] = None, expert_model_parallel_rank: Optional[int] = None, + expert_tensor_parallel_rank: Optional[int] = None, to_hf: bool = False, verbose=False, efficient_mode: bool = False, @@ -69,11 +70,14 @@ def __init__( pipeline_model_parallel_rank = mpu.get_pipeline_model_parallel_rank() if expert_model_parallel_rank is None: expert_model_parallel_rank = mpu.get_expert_model_parallel_rank() + if expert_tensor_parallel_rank is None: + expert_tensor_parallel_rank = mpu.get_expert_tensor_parallel_rank() self.dist_converter = DistConverter( self.mca_config, tensor_model_parallel_rank=tensor_model_parallel_rank, pipeline_model_parallel_rank=pipeline_model_parallel_rank, expert_model_parallel_rank=expert_model_parallel_rank, + expert_tensor_parallel_rank=expert_tensor_parallel_rank, revert=to_hf, efficient_mode=efficient_mode, ) @@ -158,7 +162,7 @@ def get_mca_state_dict(self, state_dict_iter, vp_stage: int): self.template.release() return mca_state_dict - def _mca_named_params_with_vp_stage(self, models): + def _mca_named_params_with_vp_stage(self, models, params_group=None): for vp_stage, model in enumerate(models): if is_peft_available() and isinstance(model, PeftModel): for adapter_name in model.peft_config.keys(): @@ -171,6 +175,13 @@ def _mca_named_params_with_vp_stage(self, models): else: mca_state_dict = model.state_dict_for_save_checkpoint() mca_state_dict = {k: v for k, v in mca_state_dict.items() if not k.endswith("._extra_state")} + + if params_group is not None: + if params_group == "dense": + mca_state_dict = {k: v for k, v in mca_state_dict.items() if not self.dist_converter.is_expert_parallel_weight(k)} + elif params_group == "moe": + mca_state_dict = {k: v for k, v in mca_state_dict.items() if self.dist_converter.is_expert_parallel_weight(k)} + for mca_name, weight in sorted(mca_state_dict.items()): yield None, vp_stage, mca_name, weight @@ -214,11 +225,50 @@ def convert_to_hf( converted_state_dict = self.template.add_mca_weight(merged_name, merged_weight, **kwargs) if converted_state_dict is not None: converted.update(converted_state_dict) - else: - self.log(f"mca_name: {merged_name} added but not converted") hf_state_dict.update(converted or {}) return hf_state_dict + def _needs_moe_allgather(self, mca_name: str, ep_world_size: int) -> bool: + return ( + self.mca_config.hf_model_type in ["qwen3_vl_moe", "qwen3_5_moe"] + and ep_world_size > 1 + and self.dist_converter.is_expert_parallel_weight(mca_name) + ) + + def _convert_moe_weight_with_allgather( + self, + mca_name: str, + weight: "Tensor", + vp_stage: int, + ep_group, + ep_world_size: int, + move_to_cpu: bool = False, + only_need_expert: bool = False, + ) -> Optional[Dict[str, "Tensor"]]: + gathered_weights = [torch.empty_like(weight) for _ in range(ep_world_size)] + dist.all_gather(gathered_weights, weight, group=ep_group) + + if only_need_expert: + return None + + if move_to_cpu: + gathered_weights = [w.cpu() for w in gathered_weights] + + def _extract_suffix_number(s: str): + import re + match = re.search(r"\d+$", s) + return match.group() if match else None + + converted_state_dict = {} + local_moe_index = _extract_suffix_number(mca_name) + for ep_rank, moe_weight in enumerate(gathered_weights): + global_moe_index = self.dist_converter.num_layers_for_expert * ep_rank + int(local_moe_index) + converted_name = mca_name[: -len(local_moe_index)] + str(global_moe_index) + converted_state_dict.update( + self.convert_to_hf(mca_state_dict={converted_name: [moe_weight]}, vp_stage=vp_stage) + ) + return converted_state_dict + def save_model_as_hf_inflight( self, models, @@ -243,14 +293,23 @@ def save_model_as_hf_inflight( only_need_expert = expert_parallel and mpu.get_expert_model_parallel_rank() > 0 last_adapter_name = None for adapter_name, vp_stage, mca_name, weight in self._mca_named_params_with_vp_stage(models): - if only_need_expert and not self.dist_converter.is_expert_parallel_weight(mca_name): - continue - weights = gather_tensor_parallel(weight, async_op=False) - if weights is None: # only tp_rank0 need to convert and save - continue - if move_to_cpu and isinstance(weights, list): - weights = [w.cpu() for w in weights] - converted_state_dict = self.convert_to_hf(mca_state_dict={mca_name: weights}, vp_stage=vp_stage) + ep_group = mpu.get_expert_model_parallel_group() + ep_world_size = dist.get_world_size(ep_group) + if self._needs_moe_allgather(mca_name, ep_world_size): + converted_state_dict = self._convert_moe_weight_with_allgather( + mca_name, weight, vp_stage, ep_group, ep_world_size, move_to_cpu, only_need_expert + ) + if converted_state_dict is None: + continue + else: + if only_need_expert and not self.dist_converter.is_expert_parallel_weight(mca_name): + continue + weights = gather_tensor_parallel(weight, async_op=False) + if weights is None: # only tp_rank0 need to convert and save + continue + if move_to_cpu and isinstance(weights, list): + weights = [w.cpu() for w in weights] + converted_state_dict = self.convert_to_hf(mca_state_dict={mca_name: weights}, vp_stage=vp_stage) self.save_hf_shard_state_dict( shard_state, os.path.join(save_directory, adapter_name) if adapter_name is not None else save_directory, @@ -271,6 +330,83 @@ def save_model_as_hf_inflight( if mpu.get_tensor_model_parallel_rank() == 0: self.save_shard_state_meta(shard_state, save_directory, save_safetensors) + def save_separate_model_as_hf_inflight( + self, + models, + save_directory: str, + save_safetensors: bool = True, + max_shard_size: Union[int, str] = MAX_SHARD_SIZE, + move_to_cpu: bool = False, + ): + assert self.dist_converter.revert, "save_separate_model_as_hf_inflight only support to_hf ModelConverter" + if not mpu.model_parallel_is_initialized(): + raise RuntimeError("Model parallelism must be initialized before save as hf inflight.") + + self.save_hf_config(save_directory, mca_config=models[0].config) + + shard_state = StateDictSplitState(max_shard_size=max_shard_size) + if isinstance(max_shard_size, str): + max_shard_size = parse_size_to_int(max_shard_size) + + self._save_dense_hf_state_dict(models, shard_state, save_directory, save_safetensors=save_safetensors, move_to_cpu=move_to_cpu) + self._save_moe_hf_state_dict(models, shard_state, save_directory, save_safetensors=save_safetensors, move_to_cpu=move_to_cpu) + + if mpu.get_tensor_model_parallel_rank() == 0 or mpu.get_expert_tensor_parallel_rank() == 0: + self.save_shard_state_meta(shard_state, save_directory, save_safetensors) + + def _save_dense_hf_state_dict( + self, + models, + shard_state: StateDictSplitState, + save_directory: str, + save_safetensors: bool = True, + move_to_cpu: bool = False, + ): + if not mpu.get_data_parallel_rank(with_context_parallel=True) == 0: + return + + for _, vp_stage, mca_name, weight in self._mca_named_params_with_vp_stage(models, params_group="dense"): + weights = gather_tensor_parallel(weight, async_op=False) + if weights is None: # only tp_rank0 need to convert and save + continue + if move_to_cpu and isinstance(weights, list): + weights = [w.cpu() for w in weights] + converted_state_dict = self.convert_to_hf(mca_state_dict={mca_name: weights}, vp_stage=vp_stage) + self.save_hf_shard_state_dict(shard_state, save_directory, converted_state_dict, save_safetensors) + + def _save_moe_hf_state_dict( + self, + models, + shard_state: StateDictSplitState, + save_directory: str, + save_safetensors: bool = True, + move_to_cpu: bool = False, + ): + dist.barrier() + if not mpu.get_expert_data_parallel_rank() == 0: + return + + ep_group = mpu.get_expert_model_parallel_group() + ep_world_size = dist.get_world_size(ep_group) + expert_parallel = self.mca_config.expert_model_parallel_size > 1 + only_need_expert = expert_parallel and mpu.get_expert_model_parallel_rank() > 0 + + for _, vp_stage, mca_name, weight in self._mca_named_params_with_vp_stage(models, params_group="moe"): + if self._needs_moe_allgather(mca_name, ep_world_size): + converted_state_dict = self._convert_moe_weight_with_allgather( + mca_name, weight, vp_stage, ep_group, ep_world_size, move_to_cpu, only_need_expert + ) + if converted_state_dict is None: + continue + else: + weights = gather_tensor_parallel(weight, async_op=False, is_expert_parallel=True) + if weights is None: # only tp_rank0 need to convert and save + continue + if move_to_cpu and isinstance(weights, list): + weights = [w.cpu() for w in weights] + converted_state_dict = self.convert_to_hf(mca_state_dict={mca_name: weights}, vp_stage=vp_stage) + self.save_hf_shard_state_dict(shard_state, save_directory, converted_state_dict, save_safetensors) + def _auto_bucket_size(self): # TODO: optimize this by max weight size group_size = max( diff --git a/mcore_adapter/src/mcore_adapter/models/converter/post_converter.py b/mcore_adapter/src/mcore_adapter/models/converter/post_converter.py index cccf8f461..10144fb42 100644 --- a/mcore_adapter/src/mcore_adapter/models/converter/post_converter.py +++ b/mcore_adapter/src/mcore_adapter/models/converter/post_converter.py @@ -1,13 +1,15 @@ import gc import json import os +import re from abc import ABC, abstractmethod from collections import defaultdict +from dataclasses import dataclass from itertools import product from typing import TYPE_CHECKING, Optional import torch -from megatron.core import mpu +from megatron.core import dist_checkpointing, mpu from megatron.core.tensor_parallel import model_parallel_cuda_manual_seed from safetensors.torch import save_file from tqdm import tqdm @@ -25,11 +27,12 @@ from transformers.models.auto.auto_factory import _get_model_class from transformers.utils import is_peft_available -from ...checkpointing import get_checkpoint_name, save_config_and_state_dict +from ...checkpointing import find_dist_ckpt, get_checkpoint_name, save_config_and_state_dict from ...constants import ADAPTER_CONFIG_NAME from ...training_args import DistributingParallelArguments from ...utils import get_logger from ..auto.config_auto import AutoConfig +from ..auto.modeling_auto import get_model_cls from .convert_utils import MAX_SHARD_SIZE from .model_converter import ModelConverter from .template import get_template @@ -194,7 +197,7 @@ def convert(self): model = model_class.from_pretrained( None, config=self.hf_config, state_dict=hf_state_dict, torch_dtype=self.torch_dtype, trust_remote_code=True ) - model.save_pretrained(self.save_directory, max_shard_size=MAX_SHARD_SIZE) + model.save_pretrained(self.save_directory, max_shard_size=MAX_SHARD_SIZE, save_original_format=False) self._finalize() def _get_hf_model_class(self): @@ -305,6 +308,18 @@ def convert(self): checkpoint_path=os.path.join(self.adapter_name_or_path, adapter_name), lora_rank=peft_config.r ) for hf_name, hf_weight in stream: + if self.hf_config.model_type in ["qwen3_5_moe"]: + if "mlp.experts.gate_up_proj" in hf_name or "mlp.experts.down_proj" in hf_name: + if "lora_A" in hf_name: + # lora_A = nn.Linear(in_features, r * num_experts), shape = (r * num_experts, in_features) + hf_weight = hf_weight.transpose(0, 1).reshape(hf_weight.shape[0], -1).contiguous() + if "lora_B" in hf_name: + # lora_B = nn.Linear(r * num_experts, out_features), shape = (out_features, r * num_experts) + hf_weight = hf_weight.permute(1, 2, 0).reshape(-1, hf_weight.shape[-1]).contiguous() + if "mlp.experts.gate_up_proj" in hf_name: + hf_name = hf_name.replace(".gate_up_proj", ".base_layer") + if "mlp.experts.down_proj" in hf_name: + hf_name = hf_name.replace(".down_proj", "") hf_state_dicts[adapter_name][hf_name] = hf_weight model_class = self._get_hf_model_class() @@ -323,10 +338,21 @@ def get_lora_config(adapter_name): target_modules = list(set(target_modules)) kwargs = {} + if self.hf_config.model_type in ["qwen3_5_moe"]: + kwargs["target_parameters"] = ["mlp.experts.gate_up_proj", "mlp.experts.down_proj"] if self.mca_config.num_moe_experts is not None: kwargs["rank_pattern"] = { p: peft_cfg.r // self.mca_config.moe_router_topk - for p in ["down_proj", "up_proj", "gate_proj", "w1", "w2", "w3"] + for p in [ + "down_proj", + "up_proj", + "gate_proj", + "w1", + "w2", + "w3", + "mlp.experts.gate_up_proj", + "mlp.experts.down_proj", + ] } return LoraConfig( @@ -353,12 +379,277 @@ def get_lora_config(adapter_name): self._finalize() +class DistHFConverter(HFConverter): + """ + Megatron dist -> hf + """ + def __init__(self, dist_model_dir: str, *args, **kwargs): + self.dist_model_dir = dist_model_dir + super().__init__(*args, **kwargs) + + def _setup(self): + super()._setup() + self._reset_parallel_config() + + def _reset_parallel_config(self): + """Reset parallel config to 1 for dist checkpoint conversion.""" + # Set mca_config parallel sizes to 1 + self.mca_config.tensor_model_parallel_size = 1 + self.mca_config.expert_tensor_parallel_size = 1 + self.mca_config.pipeline_model_parallel_size = 1 + self.mca_config.expert_model_parallel_size = 1 + self.mca_config.virtual_pipeline_model_parallel_size = None + self.mca_config.pipeline_model_parallel_layout = None + # Reset mpu parallel state + mpu.set_expert_model_parallel_world_size(1) + mpu.set_pipeline_model_parallel_world_size(1) + mpu.set_tensor_model_parallel_world_size(1) + mpu.set_virtual_pipeline_model_parallel_world_size(None) + + def convert(self): + models = self._init_model() + sharded_state_dict = models.sharded_state_dict() + + sharded_state_dict = self._sharded_meta_to_cpu(sharded_state_dict) + state_dict = dist_checkpointing.load(sharded_state_dict, self.dist_model_dir) + model_converter = ModelConverter( + mca_config=self.mca_config, + verbose=self.verbose, + to_hf=True, + ) + + hf_state_dict = {} + for weight_name, weight_tensor in state_dict.items(): + if "._extra_state" in weight_name: + continue + converted_state_dict = model_converter.convert_to_hf( + mca_state_dict={weight_name: [weight_tensor]}, + vp_stage=None, + ) + hf_state_dict.update(converted_state_dict) + + model_class = self._get_hf_model_class() + model = model_class.from_pretrained( + None, config=self.hf_config, state_dict=hf_state_dict, torch_dtype=self.torch_dtype, trust_remote_code=True + ) + model.save_pretrained(self.save_directory, max_shard_size=MAX_SHARD_SIZE, save_original_format=False) + + self._finalize() + + def _init_model(self): + torch.distributed.init_process_group( + backend="cpu:gloo", + rank=0, + world_size=1, + ) + mpu.initialize_model_parallel( + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + expert_model_parallel_size=1, + expert_tensor_parallel_size=None, + virtual_pipeline_model_parallel_size=None, + ) + + model_parallel_cuda_manual_seed(42) + model_class = get_model_cls(self.mca_config.hf_model_type) + + # Use meta device for low memory initialization + # Skip .to(device) calls for vision/audio models + self.mca_config.init_model_with_meta_device = True + with torch.device("meta"): + model = model_class(self.mca_config) + return model + + def _sharded_meta_to_cpu(self, sharded_state_dict): + from megatron.core.dist_checkpointing.dict_utils import dict_list_map_inplace + from megatron.core.dist_checkpointing.mapping import ShardedTensor, ShardedTensorFactory, ShardedObject + + def meta_to_cpu_empty(v): + if isinstance(v, ShardedTensorFactory): + if v.data is not None and v.data.device.type == 'meta': + v.data = torch.empty(v.data.shape, dtype=v.data.dtype, device='cpu') + elif isinstance(v, ShardedTensor): + if v.data is not None and v.data.device.type == 'meta': + v = v.without_data() + v.init_data('cpu') + elif isinstance(v, ShardedObject): + pass + return v + + dict_list_map_inplace(meta_to_cpu_empty, sharded_state_dict) + return sharded_state_dict + + +class DistStreamingHFConverter(DistHFConverter): + """ + Converts Megatron distributed checkpoint to Hugging Face format using streaming approach. + """ + MEGATRON_LAYER_PATTERN = r'^(.*\.layers)\.(\d+)\.' + + def __init__(self, dist_model_dir: str, max_shard_bytes: int = MAX_SHARD_SIZE, **kwargs): + super().__init__(dist_model_dir=dist_model_dir, **kwargs) + self.max_shard_bytes = max_shard_bytes + self.layer_prefix_cache = None + + def _get_layer_key(self, key: str) -> tuple[str, int] | None: + """ + Extract layer prefix and layer index from a weight key. + """ + if self.layer_prefix_cache is None: + self.layer_prefix_cache = re.compile(self.MEGATRON_LAYER_PATTERN) + # Match patterns like: decoder.layers.0, encoder.layers.1, etc. + match = self.layer_prefix_cache.match(key) + if match: + layer_prefix = match.group(1) + layer_idx = int(match.group(2)) + return (layer_prefix, layer_idx) + return None + + def _split_sharded_state_dict(self, sharded_state_dict: dict) -> list[dict]: + """ + Split sharded_state_dict into chunks by transformer layer. + """ + from collections import defaultdict + + # Group weights by layer + layer_chunks: dict[tuple[str, int], dict] = defaultdict(dict) + non_layer_chunk: dict = {} + + for key, value in sharded_state_dict.items(): + # Skip _extra_state entries + if "._extra_state" in key: + continue + + layer_key = self._get_layer_key(key) + if layer_key is not None: + layer_chunks[layer_key][key] = value + else: + non_layer_chunk[key] = value + + # Build final chunks list + chunks = [] + + # Add non-layer weights first (embeddings, final norm, etc.) + if non_layer_chunk: + chunks.append(non_layer_chunk) + + # Add each layer as a separate chunk, sorted by layer index + for layer_key in sorted(layer_chunks.keys()): + chunks.append(layer_chunks[layer_key]) + + return chunks + + def _stream_dist_weights_from_shards(self, sharded_state_dict: dict): + """ + Stream weights by processing sharded_state_dict in chunks. + """ + import copy + + # Split into chunks based on estimated memory size + chunks = self._split_sharded_state_dict(sharded_state_dict) + logger.info(f"Split sharded_state_dict into {len(chunks)} chunks for streaming conversion") + + model_converter = ModelConverter( + mca_config=self.mca_config, + pipeline_model_parallel_rank=0, + expert_model_parallel_rank=0, + tensor_model_parallel_rank=0, + verbose=self.verbose, + to_hf=True, + ) + + for chunk_idx, chunk in enumerate(chunks): + logger.info(f"Processing chunk {chunk_idx + 1}/{len(chunks)} with {len(chunk)} weights...") + + chunk_copy = copy.deepcopy(chunk) + chunk_copy = self._sharded_meta_to_cpu(chunk_copy) + + loaded_chunk = dist_checkpointing.load(chunk_copy, self.dist_model_dir) + + # Convert each weight to HF format + for weight_name, weight_tensor in loaded_chunk.items(): + if "._extra_state" in weight_name: + continue + + # Convert to HF format + converted_state_dict = model_converter.convert_to_hf( + mca_state_dict={weight_name: [weight_tensor]}, + vp_stage=None, + ) + + for hf_name, hf_weight in converted_state_dict.items(): + yield hf_name, hf_weight + + # Clean up to save memory + del chunk_copy + del loaded_chunk + gc.collect() + + def convert(self): + logger.info(f"Starting streaming conversion from dist checkpoint (max shard size: {self.max_shard_bytes / 1e9:.2f}GB)...") + os.makedirs(self.save_directory, exist_ok=True) + + model = self._init_model() + sharded_state_dict = model.sharded_state_dict() + + weight_map = {} + total_size = 0 + shard_count = 0 + current_shard = {} + current_shard_size = 0 + + for hf_name, hf_weight in tqdm( + self._stream_dist_weights_from_shards(sharded_state_dict), + desc="Converting dist checkpoint to hf" + ): + # The hf_name may be replicated (e.g., shared embeddings) + if hf_name in weight_map: + # Check consistency + if hf_name in current_shard and not hf_weight.equal(current_shard[hf_name]): + raise ValueError( + f"weight of hf_name:{hf_name} in " + f"diff max:{torch.abs(hf_weight - current_shard[hf_name]).max()}, " + "please check the checkpoint" + ) + continue + + current_shard[hf_name] = hf_weight + current_shard_size += hf_weight.nelement() * hf_weight.element_size() + + # Save shard if it exceeds max size + if current_shard_size >= self.max_shard_bytes: + shard_name = f"model-{shard_count + 1:05d}.safetensors" + save_file(current_shard, os.path.join(self.save_directory, shard_name), metadata={"format": "pt"}) + for k in current_shard: + weight_map[k] = shard_name + + total_size += current_shard_size + shard_count += 1 + current_shard = {} + current_shard_size = 0 + gc.collect() + + # Save the last shard + if current_shard: + shard_name = f"model-{shard_count + 1:05d}.safetensors" + save_file(current_shard, os.path.join(self.save_directory, shard_name), metadata={"format": "pt"}) + for k in current_shard: + weight_map[k] = shard_name + total_size += current_shard_size + + # Save index file + with open(os.path.join(self.save_directory, "model.safetensors.index.json"), "w") as f: + json.dump({"metadata": {"total_size": total_size}, "weight_map": weight_map}, f, indent=2) + + self._finalize() + + def convert_checkpoint_to_hf( model_name_or_path: str, save_directory: str, adapter_name_or_path: Optional[str] = None, torch_dtype: Optional["torch.dtype"] = None, - low_mem: bool = False, + low_mem: bool = True, verbose: bool = True, max_shard_bytes: int = MAX_SHARD_SIZE, ): @@ -375,6 +666,30 @@ def convert_checkpoint_to_hf( verbose (bool): Whether to print detailed conversion logs. max_shard_bytes (int): Max size of each shard in bytes for low_mem mode. """ + # for dist ckpt + dist_model_dir = find_dist_ckpt(model_name_or_path) + if dist_model_dir: + logger.info(f"The model_path: {dist_model_dir} is dist ckpt...") + if low_mem: + converter = DistStreamingHFConverter( + dist_model_dir=dist_model_dir, + checkpoint_path=model_name_or_path, + save_directory=save_directory, + torch_dtype=torch_dtype, + verbose=verbose, + ) + else: + converter = DistHFConverter( + dist_model_dir=dist_model_dir, + checkpoint_path=model_name_or_path, + save_directory=save_directory, + torch_dtype=torch_dtype, + verbose=verbose, + ) + converter.convert() + return + + # for legacy ckpt if adapter_name_or_path: if low_mem: raise ValueError("There is no need using `low_mem` mode for lora convert.") diff --git a/mcore_adapter/src/mcore_adapter/models/converter/template.py b/mcore_adapter/src/mcore_adapter/models/converter/template.py index 5e3862d6b..aa5ef828f 100644 --- a/mcore_adapter/src/mcore_adapter/models/converter/template.py +++ b/mcore_adapter/src/mcore_adapter/models/converter/template.py @@ -43,6 +43,9 @@ def __post_init__(self): self.hf_names = [self.hf_names] if isinstance(self.mca_names, str): self.mca_names = [self.mca_names] + # Cache compiled regex patterns + self._hf_patterns_cache: Optional[list[re.Pattern]] = None + self._mca_patterns_cache: Optional[list[re.Pattern]] = None def __call__(self, name_to_weight: Dict[str, torch.Tensor], mca_to_hf: bool = False) -> Any: weight_len = len(self.mca_names if mca_to_hf else self.hf_names) @@ -63,16 +66,30 @@ def mca_config(self, value: "TransformerConfig"): self._mca_config = value @staticmethod - def _name_to_pattern(name: str): - return name.replace(".", "\.").replace("{}", "(.*)") + def _name_to_pattern(name: str) -> str: + return name.replace(".", r"\.").replace("{}", "(.*)") + + def _get_compiled_patterns(self, mca_name: bool) -> List[re.Pattern]: + """Get compiled regex patterns with caching.""" + if mca_name: + if self._mca_patterns_cache is None: + self._mca_patterns_cache = [ + re.compile(self._name_to_pattern(name)) for name in self.mca_names + ] + return self._mca_patterns_cache + else: + if self._hf_patterns_cache is None: + self._hf_patterns_cache = [ + re.compile(self._name_to_pattern(name)) for name in self.hf_names + ] + return self._hf_patterns_cache def is_required_name(self, name, mca_name: bool): required_names = self.mca_names if mca_name else self.hf_names if name in required_names: return True - for pattern in required_names: - re_pattern = self._name_to_pattern(pattern) - if re.match(re_pattern, name): + for pattern in self._get_compiled_patterns(mca_name): + if pattern.match(name): return True return False @@ -124,6 +141,15 @@ def _mca_to_hf(self, weights: List[torch.Tensor]) -> List[torch.Tensor]: raise NotImplementedError() +@dataclass +class DropConverOp(ConverOp): + def _hf_to_mca(self, weights): + return [] + + def _mca_to_hf(self, weights): + return [] + + @dataclass class RenameConverOp(ConverOp): def __post_init__(self): @@ -337,20 +363,35 @@ def _mca_to_hf(self, weights): return torch.cat(weights[0].tensors, dim=0) +@dataclass +class ZeroCenteredRMSNormConverOp(ConverOp): + def __post_init__(self): + super().__post_init__() + assert len(self.hf_names) == 1, f"ZeroCenteredRMSNormConverOp only support one name {self.hf_names}" + assert len(self.mca_names) == 1, f"ZeroCenteredRMSNormConverOp only support one name {self.mca_names}" + + def _hf_to_mca(self, weights): + return weights[0].clone() - 1 + + def _mca_to_hf(self, weights): + return weights[0].clone() + 1 + + @dataclass class Template: hf_model_type: str hf_layer_prefix: str - config_hf_to_mca: Dict[str, str] - weight_converters: List[ConverOp] - constant_mca_config: Dict[str, Any] - constant_hf_config: Dict[str, Any] = field(default_factory=dict) + config_hf_to_mca: dict[str, str] + weight_converters: list[ConverOp] + constant_mca_config: dict[str, Any] + constant_hf_config: dict[str, Any] = field(default_factory=dict) hf_moe_prefix: Optional[str] = None - hf_invalid_keys: List[str] = field(default_factory=list) - config_mca_to_hf: Optional[Dict[str, str]] = None - hf_name_to_converter: Dict[str, ConverOp] = field(default_factory=dict) - mca_name_to_converter: Dict[str, ConverOp] = field(default_factory=dict) - prefix_name_to_weight: Dict[str, Dict[str, torch.Tensor]] = field(default_factory=dict) + hf_invalid_keys: list[str] = field(default_factory=list) + config_mca_to_hf: Optional[dict[str, str]] = None + hf_name_to_converter: dict[str, ConverOp] = field(default_factory=dict) + mca_name_to_converter: dict[str, ConverOp] = field(default_factory=dict) + prefix_name_to_weight: dict[str, dict[str, torch.Tensor]] = field(default_factory=dict) + _lora_op_cache: dict[str, ConverOp] = field(default_factory=dict) def __post_init__(self): self.config_hf_to_mca = self.adjust_config_hf_to_mca() @@ -520,7 +561,7 @@ def add_mca_weight(self, name, weight, **kwargs): def get_conver_op(self, name, pattern_to_conver_ops: Dict[str, ConverOp]): if name in pattern_to_conver_ops: return pattern_to_conver_ops[name] - for pattern in sorted(pattern_to_conver_ops, key=lambda x: len(x), reverse=True): + for pattern in pattern_to_conver_ops.keys(): re_pattern = pattern.replace("{}", "(.*?)") if re.match(re_pattern, name): return pattern_to_conver_ops[pattern] @@ -528,14 +569,17 @@ def get_conver_op(self, name, pattern_to_conver_ops: Dict[str, ConverOp]): def get_lora_conver_op(self, name, pattern_to_conver_ops: Dict[str, ConverOp], lora_rank: int): lora_name = name[name.find(".lora") :] + cache_key = f"{name}_{lora_rank}" + if cache_key in self._lora_op_cache: + return self._lora_op_cache[cache_key] + name = name[: name.find(".lora")] + ".weight" op = self.get_conver_op(name, pattern_to_conver_ops) + kwargs = {} if isinstance(op, RenameConverOp): op_class = RenameConverOp - kwargs = {} elif "lora_A" in lora_name: op_class = CopyConverOp - kwargs = {} elif isinstance(op, StackConverOp): op_class = StackConverOp kwargs = {"dim": op.dim} @@ -544,12 +588,14 @@ def get_lora_conver_op(self, name, pattern_to_conver_ops: Dict[str, ConverOp], l kwargs = {"hidden_size": lora_rank} else: raise ValueError(f"can not find lora conver op for {name} in {pattern_to_conver_ops}") - return op_class( + lora_op = op_class( hf_names=[hf_name.replace(".weight", lora_name) for hf_name in op.hf_names], mca_names=[mca_name.replace(".weight", lora_name) for mca_name in op.mca_names], _mca_config=op.mca_config, **kwargs, ) + self._lora_op_cache[cache_key] = lora_op + return lora_op def hf_name_to_mca_names(self, hf_name) -> Optional[List[str]]: weight_prefix = get_weight_prefix(hf_name, self.hf_layer_prefix, moe_prefix=self.hf_moe_prefix) diff --git a/mcore_adapter/src/mcore_adapter/models/model_config.py b/mcore_adapter/src/mcore_adapter/models/model_config.py index 511487e73..d7eecbf87 100644 --- a/mcore_adapter/src/mcore_adapter/models/model_config.py +++ b/mcore_adapter/src/mcore_adapter/models/model_config.py @@ -65,8 +65,9 @@ def to_json_string(self): save_dict = {} for f in dataclasses.fields(self): v = getattr(self, f.name) - if isinstance(v, list) and len(v) > 0 and callable(v[0]): - continue + if isinstance(v, list) and len(v) > 0: + if callable(v[0]) or isinstance(v[0], (torch.dtype, enum.Enum)): + continue if callable(v) or isinstance(v, (torch.dtype, enum.Enum)): continue if isinstance(v, PipelineParallelLayerLayout): @@ -154,6 +155,9 @@ def update_with_args(self, args: "DistributingParallelArguments", verbose: bool self.bf16 = getattr(args, "bf16", self.bf16) self.fp16 = getattr(args, "fp16", self.fp16) + if getattr(self, "mtp_num_layers", None) == 0: + self.mtp_num_layers = None + @classmethod def from_pretrained(cls, model_name_or_path: str, args: Optional["TrainingArguments"] = None): config_file = os.path.join(model_name_or_path, MCA_CONFIG_NAME) @@ -295,6 +299,13 @@ class McaModelConfig(TransformerConfig, PretrainedConfig): "choices": ["local", "transformer_engine"], }, ) + moe_parallel_folding: bool = field( + default=False, + metadata={ + "help": "Whether there is heterogeneous parallelism for moe and dense groups " + "set to true when etp_size != tp_size" + } + ) def __post_init__(self): if self.virtual_pipeline_model_parallel_size is None and self.overlap_p2p_comm: @@ -390,6 +401,9 @@ def squared_relu(x): f"({self.virtual_pipeline_model_parallel_size})." ) + if self.tensor_model_parallel_size != self.expert_tensor_parallel_size: + self.moe_parallel_folding = True + def distribute_config_match(self, other: "McaModelConfig"): if not isinstance(other, McaModelConfig): return False diff --git a/mcore_adapter/src/mcore_adapter/models/model_factory.py b/mcore_adapter/src/mcore_adapter/models/model_factory.py index 4acb00e56..36ea6b6a8 100644 --- a/mcore_adapter/src/mcore_adapter/models/model_factory.py +++ b/mcore_adapter/src/mcore_adapter/models/model_factory.py @@ -16,7 +16,7 @@ from transformers.tokenization_utils import PreTrainedTokenizer from transformers.utils import is_peft_available -from ..checkpointing import load_state_dict_from_checkpoint, save_config_and_state_dict +from ..checkpointing import find_dist_ckpt, generate_model_state_dict, load_state_dict_from_checkpoint, save_config_and_state_dict from ..platforms import current_platform from ..utils import get_logger from .converter.convert_utils import MAX_SHARD_SIZE @@ -55,9 +55,11 @@ def __init__(self, cls, config: "McaModelConfig", *args, **kwargs): kwargs["vp_stage"] = i self.models.append(cls(config, *args, **kwargs)) - def save_pretrained(self, save_directory: str, save_merged_model: bool = False): + def save_pretrained(self, save_directory: str, save_merged_model: bool = False, ckpt_format: str = "legacy"): if len(self.models) == 1: if is_peft_available() and isinstance(self.models[0], PeftModel): + if ckpt_format == "torch_dist": + raise ValueError(f"lora only support legacy ckpt format, got {ckpt_format}") if save_merged_model: self.models[0].merge_adapter() model_state_dict = self.models[0].state_dict_for_save_checkpoint() @@ -68,10 +70,11 @@ def save_pretrained(self, save_directory: str, save_merged_model: bool = False): elif ".base_layer" in k: k = k.replace(".base_layer", "") state_dict[k] = v - self.models[0].unmerge_adapter() - return self.models[0].base_model.model.save_pretrained( + self.models[0].base_model.model.save_pretrained( save_directory, state_dict={"model": state_dict} ) + self.models[0].unmerge_adapter() + return self.config.save_pretrained(save_directory) for adapter_name, peft_config in self.models[0].peft_config.items(): adapter_save_directory = os.path.join(save_directory, adapter_name) peft_config.save_pretrained(adapter_save_directory) @@ -82,9 +85,9 @@ def save_pretrained(self, save_directory: str, save_merged_model: bool = False): adapter_save_directory, state_dict={"model": peft_state_dict} ) return self.config.save_pretrained(save_directory) - return self.models[0].save_pretrained(save_directory) - state_dict = {f"model{i}": model.state_dict_for_save_checkpoint() for i, model in enumerate(self.models)} - return self.models[0].save_pretrained(save_directory, state_dict=state_dict) + return self.models[0].save_pretrained(save_directory, ckpt_format=ckpt_format) + state_dict = {f"model{i}": generate_model_state_dict(model, ckpt_format) for i, model in enumerate(self.models)} + return self.models[0].save_pretrained(save_directory, state_dict=state_dict, ckpt_format=ckpt_format) def load_state_dict(self, state_dict: Dict[str, torch.Tensor], strict: bool = True): if len(self.models) == 1: @@ -176,13 +179,22 @@ def save_pretrained_as_hf( ): os.makedirs(save_directory, exist_ok=True) converter = ModelConverter(self.config, to_hf=True) - converter.save_model_as_hf_inflight( - self.models, - save_directory, - save_safetensors=save_safetensors, - max_shard_size=max_shard_size, - move_to_cpu=True, - ) + if not self.config.moe_parallel_folding: + converter.save_model_as_hf_inflight( + self.models, + save_directory, + save_safetensors=save_safetensors, + max_shard_size=max_shard_size, + move_to_cpu=True, + ) + else: + converter.save_separate_model_as_hf_inflight( + self.models, + save_directory, + save_safetensors=save_safetensors, + max_shard_size=max_shard_size, + move_to_cpu=True, + ) def get_batch_on_this_cp_rank(self, *args, **kwargs): return self.models[0].get_batch_on_this_cp_rank(*args, **kwargs) @@ -231,7 +243,11 @@ def from_pretrained( dist_config_match = False if mca_ckpt_exist: old_mca_config = cls.config_class.from_pretrained(model_name_or_path) - dist_config_match = config.distribute_config_match(old_mca_config) + if find_dist_ckpt(model_name_or_path): + # dist ckpt do not need dist_config_match + dist_config_match = True + else: + dist_config_match = config.distribute_config_match(old_mca_config) if mca_ckpt_exist and dist_config_match: if resized_vocab_size: @@ -239,7 +255,7 @@ def from_pretrained( "The tokenizer length is longer than the vocab embedding size, and the resize embedding" "layer is not supported loading mca ckpt. Please check the tokenizer and ckpt." ) - state_dict = load_state_dict_from_checkpoint(model_name_or_path) + state_dict = load_state_dict_from_checkpoint(model_name_or_path, models=models) else: if not exists_hf_config(model_name_or_path): raise ValueError( @@ -264,10 +280,10 @@ def from_pretrained( logger.info(f"End loading, cost: {time.time() - load_start_time:0.3f}s") return models - def save_pretrained(self, save_directory: str, state_dict=None): + def save_pretrained(self, save_directory: str, state_dict=None, ckpt_format: str = "legacy"): os.makedirs(save_directory, exist_ok=True) - state_dict = state_dict if state_dict is not None else {"model": self.state_dict_for_save_checkpoint()} - save_config_and_state_dict(save_directory, self.config, state_dict) + state_dict = state_dict if state_dict is not None else {"model": generate_model_state_dict(self, ckpt_format)} + save_config_and_state_dict(save_directory, self.config, state_dict, ckpt_format=ckpt_format) def get_batch_on_this_cp_rank(self, batch: Dict[str, "torch.Tensor"], dim3_keys: List[str] = ["attention_mask"]): # copy from Megatron-LM @@ -352,7 +368,7 @@ def __init__(self, config: "McaModelConfig", **kwargs): ) for param in self.parameters(): tensor_parallel.set_defaults_if_not_set_tensor_model_parallel_attributes(param) - if not config.use_cpu_initialization: + if not config.use_cpu_initialization and not config.init_model_with_meta_device: self.to(current_platform.current_device()) if self.post_process or self.mtp_process: @@ -371,10 +387,6 @@ def _get_transformer_layer_spec(self, config: Optional["McaModelConfig"] = None) if not use_te and config.normalization == "RMSNorm": transformer_layer_spec.submodules.input_layernorm = RMSNorm transformer_layer_spec.submodules.pre_mlp_layernorm = RMSNorm - if getattr(transformer_layer_spec.submodules.mlp.submodules, "shared_experts", None): - transformer_layer_spec.submodules.mlp.submodules.shared_experts.params["gate"] = ( - config.moe_use_shared_expert_gate - ) return transformer_block_spec if use_te: return get_gpt_layer_with_transformer_engine_spec( @@ -398,3 +410,28 @@ def _get_mtp_block_spec(self, config: Optional["McaModelConfig"] = None, vp_stag return spec else: return None + + def compute_language_model_loss(self, labels: torch.Tensor, logits: torch.Tensor) -> torch.Tensor: + """Computes the language model loss (Cross entropy across vocabulary) + + Args: + labels (torch.Tensor): The labels of dimension [batch size, seq length] + logits (torch.Tensor): The final logits returned by the output layer of the transformer model + + Returns: + Tensor: Loss tensor of dimensions [batch size, sequence_length] + """ + + if self.config.cross_entropy_loss_fusion and self.config.tensor_model_parallel_size == 1: + from ..parallel_functions.fused_cross_entropy import cross_entropy + + s, b, d = logits.shape + + # [b, s] + labels = labels.transpose(0, 1).contiguous().view(-1) + logits = logits.view(-1, d) + # loss = fused_vocab_parallel_cross_entropy(logits, labels) + loss = cross_entropy(logits, labels, reduction="none", ignore_index=-100) + loss = loss.view(s, b).transpose(0, 1).contiguous() + return loss + return super().compute_language_model_loss(labels, logits) diff --git a/mcore_adapter/src/mcore_adapter/models/qwen2_5_vl/modeling_qwen2_5_vl.py b/mcore_adapter/src/mcore_adapter/models/qwen2_5_vl/modeling_qwen2_5_vl.py index 085be37a2..89d2eb5e0 100644 --- a/mcore_adapter/src/mcore_adapter/models/qwen2_5_vl/modeling_qwen2_5_vl.py +++ b/mcore_adapter/src/mcore_adapter/models/qwen2_5_vl/modeling_qwen2_5_vl.py @@ -28,7 +28,9 @@ def __init__(self, config: "Qwen2_5_VLConfig", **kwargs): Qwen2_5_VLVisionConfig(**config.vision_config), attn_implementation="sdpa", torch_dtype=self.config.params_dtype, - ).to(current_platform.current_device()) + ) + if not config.init_model_with_meta_device: + self.vision_model = self.vision_model.to(current_platform.current_device()) # TODO: use_reentrant=True might cause error by twice forward/backward when # training images and videos simultaneously, https://github.com/pytorch/pytorch/issues/81296 if config.recompute_granularity == "full" and self.training: @@ -298,12 +300,13 @@ def get_rope_index( return position_ids, mrope_position_deltas def get_batch_on_this_cp_rank(self, batch, dim3_keys: list[str] = ["attention_mask"]): - # VLM need to view all input_ids and media features - loss_needed_items = { - "labels": batch.pop("labels", None), - } - loss_needed_items = super().get_batch_on_this_cp_rank(loss_needed_items, dim3_keys=dim3_keys) - batch.update(loss_needed_items) + # VLM forward() handles input_ids and attention_mask splitting internally + skipped = {} + for key in ("input_ids", "attention_mask"): + if key in batch: + skipped[key] = batch.pop(key) + batch = super().get_batch_on_this_cp_rank(batch, dim3_keys=dim3_keys) + batch.update(skipped) return batch def get_input_ranges(self, total_seqlen): @@ -363,7 +366,12 @@ def forward( } if self.config.context_parallel_size > 1: cp_batch = {k: v.clone() if v is not None else None for k, v in cp_batch.items()} - cp_batch = super().get_batch_on_this_cp_rank(cp_batch, dim3_keys=["attention_mask"]) + cp_batch = super().get_batch_on_this_cp_rank( + cp_batch, + dim3_keys=[] + if (cp_batch["attention_mask"] is None or cp_batch["attention_mask"].dim() == 2) + else ["attention_mask"], + ) if not self.pre_process or decoder_input is not None: return super().forward( diff --git a/mcore_adapter/src/mcore_adapter/models/qwen2_vl/modeling_qwen2_vl.py b/mcore_adapter/src/mcore_adapter/models/qwen2_vl/modeling_qwen2_vl.py index b7ba9b56e..8ed3bed8a 100644 --- a/mcore_adapter/src/mcore_adapter/models/qwen2_vl/modeling_qwen2_vl.py +++ b/mcore_adapter/src/mcore_adapter/models/qwen2_vl/modeling_qwen2_vl.py @@ -28,7 +28,9 @@ def __init__(self, config: "Qwen2VLConfig", **kwargs): Qwen2VLVisionConfig(**config.vision_config), attn_implementation="sdpa", torch_dtype=self.config.params_dtype, - ).to(current_platform.current_device()) + ) + if not config.init_model_with_meta_device: + self.vision_model = self.vision_model.to(current_platform.current_device()) # TODO: use_reentrant=True might cause error by twice forward/backward when # training images and videos simultaneously, https://github.com/pytorch/pytorch/issues/81296 if config.recompute_granularity == "full" and self.training: @@ -278,12 +280,13 @@ def get_rope_index( return position_ids, mrope_position_deltas def get_batch_on_this_cp_rank(self, batch, dim3_keys: list[str] = ["attention_mask"]): - # VLM need to view all input_ids and media features - loss_needed_items = { - "labels": batch.pop("labels", None), - } - loss_needed_items = super().get_batch_on_this_cp_rank(loss_needed_items, dim3_keys=dim3_keys) - batch.update(loss_needed_items) + # VLM forward() handles input_ids and attention_mask splitting internally + skipped = {} + for key in ("input_ids", "attention_mask"): + if key in batch: + skipped[key] = batch.pop(key) + batch = super().get_batch_on_this_cp_rank(batch, dim3_keys=dim3_keys) + batch.update(skipped) return batch def get_input_ranges(self, total_seqlen): @@ -340,7 +343,12 @@ def forward( } if self.config.context_parallel_size > 1: cp_batch = {k: v.clone() if v is not None else None for k, v in cp_batch.items()} - cp_batch = super().get_batch_on_this_cp_rank(cp_batch, dim3_keys=["attention_mask", "position_ids"]) + cp_batch = super().get_batch_on_this_cp_rank( + cp_batch, + dim3_keys=[] + if (cp_batch["attention_mask"] is None or cp_batch["attention_mask"].dim() == 2) + else ["attention_mask"], + ) if not self.pre_process or decoder_input is not None: return super().forward( diff --git a/mcore_adapter/src/mcore_adapter/models/qwen3_5/__init__.py b/mcore_adapter/src/mcore_adapter/models/qwen3_5/__init__.py index 00ec33284..ac1994f6d 100644 --- a/mcore_adapter/src/mcore_adapter/models/qwen3_5/__init__.py +++ b/mcore_adapter/src/mcore_adapter/models/qwen3_5/__init__.py @@ -1,13 +1,23 @@ import re from dataclasses import dataclass +from typing import Optional import torch -from ..converter.convert_utils import StackedTensors +from ..converter.convert_utils import ( + MCA_MTP_MOE_PREFIX, + MCA_MTP_PREFIX, + StackedTensors, + convert_to_mca_prefix, + get_weight_prefix, + remove_mca_mtp_weight_prefix, + remove_weight_prefix, +) from ..converter.dist_converter import ( DistParallelConfig, default_dist_config, gdn_dist_config, + mtp_config, register_dist_config, ) from ..converter.template import ( @@ -18,21 +28,13 @@ RenameConverOp, StackConverOp, Template, + ZeroCenteredRMSNormConverOp, register_template, ) from .config_qwen3_5 import Qwen3_5Config from .modeling_qwen3_5 import Qwen3_5Model -@dataclass -class DropConverOp(ConverOp): - def _hf_to_mca(self, weights): - return [] - - def _mca_to_hf(self, weights): - return [] - - @dataclass class Qwen3_5_GDNConverOp(ConverOp): def __post_init__(self): @@ -64,30 +66,18 @@ def _hf_to_mca(self, weights): return StackedTensors(tensors=[q, k, v, z, b, a], dim=0) def _mca_to_hf(self, weights): - if len(weights) == 1: - assert isinstance(weights[0], StackedTensors) - q, k, v, z, b, a = weights[0].tensors - qkv = torch.cat([q, k, v], dim=0) - return [qkv, z, b, a] - - -@dataclass -class ZeroCenteredRMSNormConverOp(ConverOp): - def __post_init__(self): - super().__post_init__() - assert len(self.hf_names) == 1, f"ZeroCenteredRMSNormConverOp only support one name {self.hf_names}" - assert len(self.mca_names) == 1, f"ZeroCenteredRMSNormConverOp only support one name {self.mca_names}" - - def _hf_to_mca(self, weights): - return weights[0].clone() - 1 - - def _mca_to_hf(self, weights): - return weights[0].clone() + 1 + assert len(weights) == 1 + assert isinstance(weights[0], StackedTensors) + q, k, v, z, b, a = weights[0].tensors + qkv = torch.cat([q, k, v], dim=0) + return [qkv, z, b, a] register_dist_config( "qwen3_5", - default_dist_config.merge_configs(gdn_dist_config).merge_configs( + default_dist_config.merge_configs(gdn_dist_config) + .merge_configs(mtp_config) + .merge_configs( DistParallelConfig( pre_process_weights=["vision_model.*"], duplicated_weights=["vision_model.*"], @@ -98,6 +88,11 @@ def _mca_to_hf(self, weights): @dataclass class Qwen3_5Template(Template): + def __post_init__(self): + super().__post_init__() + self.hf_ln_pattern = re.compile(r"^model\.language_model\.layers\.(\d+)\.input_layernorm\.weight$") + self.mca_ln_pattern = re.compile(r"^decoder\.layers\.(\d+)\.self_attention\.in_proj\.layer_norm_weight$") + def adjust_config_hf_to_mca(self): non_text_config_keys = set( list(filter(lambda k: k.endswith("_token_id"), self.config_hf_to_mca.keys())) @@ -112,45 +107,122 @@ def adjust_config_hf_to_mca(self): return new_config_hf_to_mca def add_hf_weight(self, name, weight): - pattern = r"^model\.language_model\.layers\.(\d+)\.input_layernorm\.weight$" - match = re.match(pattern, name) + match = re.match(self.hf_ln_pattern, name) layer_idx = int(match.group(1)) if match else None if layer_idx is not None and self.mca_config.layer_types[layer_idx] == "linear_attention": return {f"decoder.layers.{layer_idx}.self_attention.in_proj.layer_norm_weight": weight} - return super().add_hf_weight(name, weight) + if not name.startswith("mtp"): + return super().add_hf_weight(name, weight) + weight_prefix = "mtp.layers.0" if name.startswith("mtp.layers.0") else "mtp" + if self.hf_moe_prefix is not None and self.hf_moe_prefix in name: + weight_prefix = get_weight_prefix(name, "mtp.layers.", moe_prefix=self.hf_moe_prefix) + original_name = name.removeprefix(weight_prefix) + if weight_prefix not in self.prefix_name_to_weight: + self.prefix_name_to_weight[weight_prefix] = {} + self.prefix_name_to_weight[weight_prefix][original_name] = weight + # weights in the same layer + prefix_weights = self.prefix_name_to_weight[weight_prefix] + op = self.get_conver_op(original_name, self.hf_name_to_converter) + name_to_weight = { + name: prefix_weights.pop(name) + for name in list(prefix_weights.keys()) + if op.is_required_name(name, mca_name=False) + } + conver_res = op(name_to_weight, mca_to_hf=False) + if conver_res is None: + # not ready to convert + self.prefix_name_to_weight[weight_prefix].update(name_to_weight) + return conver_res + has_transformer_layer = "self_attention" in name or "mlp" in name or "input_layernorm" in name + mca_prefix = "mtp.layers.0" + (".mtp_model_layer" if has_transformer_layer else "") + if self.hf_moe_prefix is not None and self.hf_moe_prefix in name: + mca_prefix = convert_to_mca_prefix(weight_prefix, self.hf_layer_prefix, self.hf_moe_prefix) + mca_prefix = mca_prefix.replace("mtp.layers.0", "mtp.layers.0.mtp_model_layer", 1) + return {mca_prefix + name: weight for name, weight in conver_res.items()} def add_mca_weight(self, name, weight, **kwargs): - pattern = r"^decoder\.layers\.(\d+)\.self_attention\.in_proj\.layer_norm_weight$" - match = re.match(pattern, name) - if not match: + match = re.match(self.mca_ln_pattern, name) + if match: + layer_idx = int(match.group(1)) + return {f"model.language_model.layers.{layer_idx}.input_layernorm.weight": weight} + if not name.startswith("mtp"): return super().add_mca_weight(name, weight, **kwargs) - layer_idx = int(match.group(1)) if match else None - return {f"model.language_model.layers.{layer_idx}.input_layernorm.weight": weight} + if MCA_MTP_MOE_PREFIX in name: + # MTP MoE weight: include expert index in prefix + weight_prefix = get_weight_prefix(name, MCA_MTP_PREFIX, MCA_MTP_MOE_PREFIX) + else: + weight_prefix = ( + "mtp.layers.0.mtp_model_layer" if name.startswith("mtp.layers.0.mtp_model_layer") else "mtp.layers.0" + ) + original_name = remove_mca_mtp_weight_prefix(name) + if weight_prefix not in self.prefix_name_to_weight: + self.prefix_name_to_weight[weight_prefix] = {} + self.prefix_name_to_weight[weight_prefix][original_name] = weight + prefix_weights = self.prefix_name_to_weight[weight_prefix] + op = self.get_conver_op(original_name, self.mca_name_to_converter) + name_to_weight = { + name: prefix_weights.pop(name) + for name in list(prefix_weights.keys()) + if op.is_required_name(name, mca_name=True) + } + conver_res = op(name_to_weight, mca_to_hf=True) + if conver_res is None: + # not ready to convert + self.prefix_name_to_weight[weight_prefix].update(name_to_weight) + return conver_res + if MCA_MTP_MOE_PREFIX in name: + # Convert MTP MoE prefix: remove .mtp_model_layer and .local_experts. + hf_prefix = weight_prefix.replace(".mtp_model_layer", "").replace(".local_experts.", ".") + else: + hf_prefix = "mtp.layers.0" if name.startswith("mtp.layers.0.mtp_model_layer") else "mtp" + result = {hf_prefix + name: weight for name, weight in conver_res.items()} + return result + + def hf_name_to_mca_names(self, hf_name) -> Optional[list[str]]: + if not hf_name.startswith("mtp"): + return super().hf_name_to_mca_names(hf_name) + weight_prefix = "mtp.layers.0" if hf_name.startswith("mtp.layers.0") else "mtp" + original_name = hf_name.removeprefix(weight_prefix) + if self.hf_moe_prefix is not None: + original_name = remove_weight_prefix(original_name, self.hf_moe_prefix) + if original_name in self.hf_invalid_keys: + return None + op = self.get_conver_op(original_name, self.hf_name_to_converter) + has_transformer_layer = "self_attention" in hf_name or "mlp" in hf_name or "input_layernorm" in hf_name + mca_prefix = "mtp.layers.0" + (".mtp_model_layer" if has_transformer_layer else "") + return [mca_prefix + name for name in op.mca_names] def get_lora_conver_op(self, name, pattern_to_conver_ops: dict[str, ConverOp], lora_rank: int): lora_name = name[name.find(".lora") :] + cache_key = f"{name}_{lora_rank}" + if cache_key in self._lora_op_cache: + return self._lora_op_cache[cache_key] + name = name[: name.find(".lora")] + ".weight" op = self.get_conver_op(name, pattern_to_conver_ops) + kwargs = {} if isinstance(op, RenameConverOp): op_class = RenameConverOp - kwargs = {} elif "lora_A" in lora_name: op_class = CopyConverOp - kwargs = {} elif isinstance(op, StackConverOp): op_class = StackConverOp kwargs = {"dim": op.dim} elif isinstance(op, GatedQKVConverOp): op_class = GatedQKVConverOp kwargs = {"hidden_size": lora_rank} + elif isinstance(op, Qwen3_5_GDNConverOp): + op_class = Qwen3_5_GDNConverOp else: raise ValueError(f"cannot find lora conver op for {name} in {pattern_to_conver_ops}") - return op_class( + lora_op = op_class( hf_names=[hf_name.replace(".weight", lora_name) for hf_name in op.hf_names], mca_names=[mca_name.replace(".weight", lora_name) for mca_name in op.mca_names], _mca_config=op.mca_config, **kwargs, ) + self._lora_op_cache[cache_key] = lora_op + return lora_op register_template( @@ -188,6 +260,7 @@ def get_lora_conver_op(self, name, pattern_to_conver_ops: dict[str, ConverOp], l # "mlp_only_layers": "mlp_only_layers", "layer_types": "layer_types", "full_attention_interval": "linear_attention_freq", + "mtp_num_hidden_layers": "mtp_num_layers", }, constant_mca_config={ "swiglu": True, @@ -200,6 +273,7 @@ def get_lora_conver_op(self, name, pattern_to_conver_ops: dict[str, ConverOp], l "hetereogenous_dist_checkpoint": True, "attention_output_gate": True, "experimental_attention_variant": "gated_delta_net", + "mtp_loss_scaling_factor": 0.3, }, weight_converters=[ RenameConverOp(hf_names="lm_head.weight", mca_names="output_layer.weight"), @@ -241,7 +315,10 @@ def get_lora_conver_op(self, name, pattern_to_conver_ops: dict[str, ConverOp], l # vit related RenameConverOp(hf_names="model.visual.{}", mca_names="vision_model.{}"), # mtp related - DropConverOp(hf_names="mtp.*", mca_names=[]), + RenameConverOp(hf_names=".pre_fc_norm_embedding.weight", mca_names=".enorm.weight"), + RenameConverOp(hf_names=".pre_fc_norm_hidden.weight", mca_names=".hnorm.weight"), + RenameConverOp(hf_names=".fc.weight", mca_names=".eh_proj.weight"), + RenameConverOp(hf_names=".norm.weight", mca_names=".final_layernorm.weight"), ], ) diff --git a/mcore_adapter/src/mcore_adapter/models/qwen3_5/modeling_qwen3_5.py b/mcore_adapter/src/mcore_adapter/models/qwen3_5/modeling_qwen3_5.py index 23d7203cc..4bd79fb7c 100644 --- a/mcore_adapter/src/mcore_adapter/models/qwen3_5/modeling_qwen3_5.py +++ b/mcore_adapter/src/mcore_adapter/models/qwen3_5/modeling_qwen3_5.py @@ -10,6 +10,7 @@ from ..auto.modeling_auto import register_model from ..model_factory import McaGPTModel from ..qwen3_vl.rope_utils import Qwen3VLMultimodalRotaryEmbedding, get_rope_index +from ..sequence_packing_mixin import MultimodalEmbeddingMixin from .config_qwen3_5 import Qwen3_5Config @@ -41,7 +42,7 @@ def __init__( @register_model("qwen3_5") -class Qwen3_5Model(Qwen3_5McaGPTModel): +class Qwen3_5Model(Qwen3_5McaGPTModel, MultimodalEmbeddingMixin): config_class = Qwen3_5Config def __init__(self, config: "Qwen3_5Config", **kwargs): @@ -55,7 +56,9 @@ def __init__(self, config: "Qwen3_5Config", **kwargs): Qwen3_5VisionConfig(**config.vision_config), attn_implementation="sdpa", torch_dtype=self.config.params_dtype, - ).to(current_platform.current_device()) + ) + if not config.init_model_with_meta_device: + self.vision_model = self.vision_model.to(current_platform.current_device()) # TODO: use_reentrant=True might cause error by twice forward/backward when # training images and videos simultaneously, https://github.com/pytorch/pytorch/issues/81296 if config.recompute_granularity == "full" and self.training: @@ -124,14 +127,24 @@ def construct_inputs_embeds( image_embeds = self.gather_encoder_outputs(image_embeds, split_plan, image_output_lengths) image_embeds = image_embeds.to(inputs_embeds.device, inputs_embeds.dtype) - selected_mask = torch.cat([image_mask[:, start:end] for start, end in input_ranges], dim=1) - selected_indices = torch.cat([image_indices[:, start:end] for start, end in input_ranges], dim=1) - selected_indices = selected_indices[selected_indices != -1] - - inputs_embeds = inputs_embeds.transpose(0, 1) # [s, b, h] -> [b, s, h] - selected_mask = selected_mask.unsqueeze(-1).expand_as(inputs_embeds) - inputs_embeds = inputs_embeds.masked_scatter(selected_mask, image_embeds[selected_indices]) - inputs_embeds = inputs_embeds.transpose(0, 1).contiguous() + full_sequence_output = input_ranges is None + if full_sequence_output: + # Packing mode: inputs_embeds is the full sequence (not sliced). + all_selected_indices = image_indices[image_mask] + inputs_embeds = inputs_embeds.transpose(0, 1) # [s, b, h] -> [b, s, h] + full_selected_mask = image_mask.unsqueeze(-1).expand_as(inputs_embeds) + inputs_embeds = inputs_embeds.masked_scatter(full_selected_mask, image_embeds[all_selected_indices]) + inputs_embeds = inputs_embeds.transpose(0, 1).contiguous() + else: + # Normal mode: slice mask & indices to current rank's input_ranges. + selected_mask = torch.cat([image_mask[:, start:end] for start, end in input_ranges], dim=1) + selected_indices = torch.cat([image_indices[:, start:end] for start, end in input_ranges], dim=1) + selected_indices = selected_indices[selected_indices != -1] + + inputs_embeds = inputs_embeds.transpose(0, 1) # [s, b, h] -> [b, s, h] + selected_mask = selected_mask.unsqueeze(-1).expand_as(inputs_embeds) + inputs_embeds = inputs_embeds.masked_scatter(selected_mask, image_embeds[selected_indices]) + inputs_embeds = inputs_embeds.transpose(0, 1).contiguous() return inputs_embeds def build_encoder_inputs( @@ -216,43 +229,44 @@ def gather_encoder_outputs( return encoder_small_batch_size_gather(output_features) def get_batch_on_this_cp_rank(self, batch, dim3_keys: list[str] = ["attention_mask"]): - # VLM need to view all input_ids and media features - loss_needed_items = { - "labels": batch.pop("labels", None), - } - loss_needed_items = super().get_batch_on_this_cp_rank(loss_needed_items, dim3_keys=dim3_keys) - batch.update(loss_needed_items) + # VLM forward() handles input_ids and attention_mask splitting internally + skipped = {} + for key in ("input_ids", "attention_mask"): + if key in batch: + skipped[key] = batch.pop(key) + batch = super().get_batch_on_this_cp_rank(batch, dim3_keys=dim3_keys) + batch.update(skipped) return batch - def get_input_ranges(self, total_seqlen): - # context parallel 的计算有问题 - slice_rank, slice_size = 0, 1 - if self.config.sequence_parallel: - slice_rank = mpu.get_tensor_model_parallel_rank() - slice_size = mpu.get_tensor_model_parallel_world_size() - - def get_sequence_range(start, end, rank, size): - return start + (end - start) * rank // size, start + (end - start) * (rank + 1) // size - - if self.config.context_parallel_size <= 1: - return [list(get_sequence_range(0, total_seqlen, slice_rank, slice_size))] - cp_rank = mpu.get_context_parallel_rank() - cp_size = mpu.get_context_parallel_world_size() - left_start = (total_seqlen // cp_size // 2) * cp_rank - left_end = (total_seqlen // cp_size // 2) * (cp_rank + 1) - right_start = total_seqlen - left_end - right_end = total_seqlen - left_start - slice_len = (left_end - left_start + right_end - right_start) // slice_size - start = left_start + slice_len * slice_rank - end = start + slice_len - if start >= left_end: - start = start - left_end + right_start - end = start + slice_len - return [[start, end]] - if end <= left_end: - return [[start, end]] - end = end - left_end + right_start - return [[start, left_end], [right_start, end]] + def construct_multimodal_inputs( + self, + input_ids, + inputs_embeds, + inputs_ranges, + *, + pixel_values=None, + pixel_values_videos=None, + image_grid_thw=None, + video_grid_thw=None, + force_vit_image=False, + force_vit_video=False, + **kwargs, + ): + if pixel_values is not None: + inputs_embeds = self.construct_inputs_embeds( + input_ids, inputs_embeds, pixel_values, image_grid_thw, + inputs_ranges, self.config.image_token_id, + ) + elif force_vit_image: + inputs_embeds = self._handle_missing_visual(inputs_embeds) + if pixel_values_videos is not None: + inputs_embeds = self.construct_inputs_embeds( + input_ids, inputs_embeds, pixel_values_videos, video_grid_thw, + inputs_ranges, self.config.video_token_id, + ) + elif force_vit_video: + inputs_embeds = self._handle_missing_visual(inputs_embeds) + return inputs_embeds def forward( self, @@ -269,49 +283,35 @@ def forward( ) -> "torch.Tensor": force_vit_image = kwargs.pop("force_vit_image", False) force_vit_video = kwargs.pop("force_vit_video", False) - if position_ids is None and input_ids is not None: + packed_seq_params = kwargs.get("packed_seq_params", None) + + if ( + position_ids is None or (self.config.mtp_num_layers is not None and self.config.mtp_num_layers > 0) + ) and input_ids is not None: position_ids, _ = get_rope_index(self.config, input_ids, image_grid_thw, video_grid_thw) - cp_batch = { - "input_ids": input_ids, - "attention_mask": attention_mask, - } - if self.config.context_parallel_size > 1: - cp_batch = {k: v.clone() if v is not None else None for k, v in cp_batch.items()} - cp_batch = super().get_batch_on_this_cp_rank(cp_batch, dim3_keys=[]) + state = self.prepare_packing_state( + input_ids, position_ids, attention_mask, packed_seq_params, + ) if not self.pre_process or decoder_input is not None: return super().forward( - decoder_input=decoder_input, labels=labels, position_ids=position_ids, **cp_batch, **kwargs + decoder_input=decoder_input, labels=labels, + position_ids=state.position_ids, + **state.cp_batch, **kwargs ) - inputs_ranges = self.get_input_ranges(input_ids.shape[1]) - - inputs_embeds = self.embedding(input_ids=cp_batch["input_ids"], position_ids=None) - if pixel_values is not None: - inputs_embeds = self.construct_inputs_embeds( - input_ids, - inputs_embeds, - pixel_values, - image_grid_thw, - inputs_ranges, - self.config.image_token_id, - ) - elif force_vit_image: - inputs_embeds = self._handle_missing_visual(inputs_embeds) - if pixel_values_videos is not None: - inputs_embeds = self.construct_inputs_embeds( - input_ids, - inputs_embeds, - pixel_values_videos, - video_grid_thw, - inputs_ranges, - self.config.video_token_id, - ) - elif force_vit_video: - inputs_embeds = self._handle_missing_visual(inputs_embeds) - decoder_input = inputs_embeds + result = self.build_multimodal_embeddings( + input_ids, attention_mask, packed_seq_params, state, + multimodal_kwargs=dict( + pixel_values=pixel_values, pixel_values_videos=pixel_values_videos, + image_grid_thw=image_grid_thw, video_grid_thw=video_grid_thw, + force_vit_image=force_vit_image, force_vit_video=force_vit_video, + ), + ) return super().forward( - decoder_input=decoder_input, labels=labels, position_ids=position_ids, **cp_batch, **kwargs + decoder_input=result.inputs_embeds, labels=labels, + position_ids=result.position_ids, + **result.cp_batch, **kwargs ) diff --git a/mcore_adapter/src/mcore_adapter/models/qwen3_5_moe/__init__.py b/mcore_adapter/src/mcore_adapter/models/qwen3_5_moe/__init__.py index 95e8ece32..9a51ce86e 100644 --- a/mcore_adapter/src/mcore_adapter/models/qwen3_5_moe/__init__.py +++ b/mcore_adapter/src/mcore_adapter/models/qwen3_5_moe/__init__.py @@ -5,29 +5,26 @@ from ..auto.config_auto import register_config from ..auto.modeling_auto import register_model -from ..converter.convert_utils import ( - convert_to_hf_prefix, - get_mca_moe_index, - get_mca_weight_prefix, - remove_mca_weight_prefix, -) from ..converter.dist_converter import ( DistParallelConfig, default_dist_config, gdn_dist_config, + mtp_config, register_dist_config, shared_moe_dist_config, ) from ..converter.template import ( ConverOp, + CopyConverOp, GatedQKVConverOp, GDNConv1dConverOp, RenameConverOp, StackConverOp, StackedTensors, + ZeroCenteredRMSNormConverOp, register_template, ) -from ..qwen3_5 import DropConverOp, Qwen3_5_GDNConverOp, Qwen3_5Template, ZeroCenteredRMSNormConverOp +from ..qwen3_5 import Qwen3_5_GDNConverOp, Qwen3_5Template from ..qwen3_5.config_qwen3_5 import Qwen3_5Config from ..qwen3_5.modeling_qwen3_5 import Qwen3_5Model @@ -71,6 +68,7 @@ def _hf_to_mca(self, weights): "qwen3_5_moe", default_dist_config.merge_configs(shared_moe_dist_config) .merge_configs(gdn_dist_config) + .merge_configs(mtp_config) .merge_configs( DistParallelConfig( pre_process_weights=["vision_model.*"], @@ -82,45 +80,42 @@ def _hf_to_mca(self, weights): @dataclass class Qwen3_5_MoETemplate(Qwen3_5Template): - def add_mca_weight(self, name, weight, **kwargs): - pattern = r"^decoder\.layers\.(\d+)\.self_attention\.in_proj\.layer_norm_weight$" - match = re.match(pattern, name) - if match: - layer_idx = int(match.group(1)) if match else None - return {f"model.language_model.layers.{layer_idx}.input_layernorm.weight": weight} - weight_prefix = get_mca_weight_prefix(name) - original_name = remove_mca_weight_prefix(name) - moe_layer_index = get_mca_moe_index(name) - # Since experts weights are stacked in qwen3_vl_moe, - # we need to add the moe index to the original name to - # ensure all experts weights have the same weight_prefix - if moe_layer_index is not None: - original_name = str(moe_layer_index) + original_name - weight_prefix = name[: -len(original_name)] - if weight_prefix not in self.prefix_name_to_weight: - self.prefix_name_to_weight[weight_prefix] = {} - self.prefix_name_to_weight[weight_prefix][original_name] = weight - prefix_weights = self.prefix_name_to_weight[weight_prefix] - # However, when looking up the converter, we still use the original name without moe index - # This is because mca_name_to_converter is built before mca_names reset which happens at - # model converter init. - original_name = remove_mca_weight_prefix(name) - if ".lora_A." in original_name or ".lora_B." in original_name: - op = self.get_lora_conver_op(original_name, self.mca_name_to_converter, **kwargs) + def get_lora_conver_op(self, name, pattern_to_conver_ops: dict[str, ConverOp], lora_rank: int): + lora_name = name[name.find(".lora") :] + cache_key = f"{name}_{lora_rank}" + if cache_key in self._lora_op_cache: + return self._lora_op_cache[cache_key] + + name = name[: name.find(".lora")] + ".weight" + op = self.get_conver_op(name, pattern_to_conver_ops) + kwargs = {} + if isinstance(op, RenameConverOp): + op_class = RenameConverOp + elif isinstance(op, (SplitConverOp, SplitStackConverOp)): + op_class = type(op) + elif "lora_A" in lora_name: + op_class = CopyConverOp + elif isinstance(op, StackConverOp): + op_class = StackConverOp + kwargs = {"dim": op.dim} + elif isinstance(op, GatedQKVConverOp): + op_class = GatedQKVConverOp + kwargs = {"hidden_size": lora_rank} + elif isinstance(op, Qwen3_5_GDNConverOp): + op_class = type(op) else: - op = self.get_conver_op(original_name, self.mca_name_to_converter) - name_to_weight = { - name: prefix_weights.pop(name) - for name in list(prefix_weights.keys()) - if op.is_required_name(name, mca_name=True) - } - conver_res = op(name_to_weight, mca_to_hf=True) - if conver_res is None: - # not ready to convert - self.prefix_name_to_weight[weight_prefix].update(name_to_weight) - return conver_res - hf_prefix = convert_to_hf_prefix(weight_prefix, self.hf_layer_prefix, self.hf_moe_prefix) - return {hf_prefix + name: weight for name, weight in conver_res.items()} + raise ValueError(f"cannot find lora conver op for {name} in {pattern_to_conver_ops}") + lora_hf_names = [hf_name if hf_name.endswith(".weight") else hf_name + ".weight" for hf_name in op.hf_names] + lora_hf_names = [hf_name.replace(".weight", lora_name) for hf_name in lora_hf_names] + lora_mca_names = [mca_name.replace(".weight", lora_name) for mca_name in op.mca_names] + lora_op = op_class( + hf_names=lora_hf_names, + mca_names=lora_mca_names, + _mca_config=op.mca_config, + **kwargs, + ) + self._lora_op_cache[cache_key] = lora_op + return lora_op register_template( @@ -166,6 +161,7 @@ def add_mca_weight(self, name, weight, **kwargs): # "mlp_only_layers": "mlp_only_layers", "layer_types": "layer_types", "full_attention_interval": "linear_attention_freq", + "mtp_num_hidden_layers": "mtp_num_layers", }, constant_mca_config={ "swiglu": True, @@ -181,6 +177,7 @@ def add_mca_weight(self, name, weight, **kwargs): "hetereogenous_dist_checkpoint": True, "attention_output_gate": True, "experimental_attention_variant": "gated_delta_net", + "mtp_loss_scaling_factor": 0.3, }, weight_converters=[ RenameConverOp(hf_names="lm_head.weight", mca_names="output_layer.weight"), @@ -194,6 +191,8 @@ def add_mca_weight(self, name, weight, **kwargs): RenameConverOp(hf_names=".mlp.gate.weight", mca_names=".mlp.router.weight"), SplitStackConverOp(hf_names="gate_up_proj", mca_names=".linear_fc1.weight"), SplitConverOp(hf_names="down_proj", mca_names=".linear_fc2.weight"), + RenameConverOp(hf_names=".down_proj.weight", mca_names=".linear_fc2.weight"), + StackConverOp(hf_names=[".gate_proj.weight", ".up_proj.weight"], mca_names=".linear_fc1.weight", dim=0), # Shared experts RenameConverOp( hf_names=".mlp.shared_expert.down_proj.weight", mca_names=".mlp.shared_experts.linear_fc2.weight" @@ -232,6 +231,9 @@ def add_mca_weight(self, name, weight, **kwargs): # vit related RenameConverOp(hf_names="model.visual.{}", mca_names="vision_model.{}"), # mtp related - DropConverOp(hf_names="mtp.*", mca_names=[]), + RenameConverOp(hf_names=".pre_fc_norm_embedding.weight", mca_names=".enorm.weight"), + RenameConverOp(hf_names=".pre_fc_norm_hidden.weight", mca_names=".hnorm.weight"), + RenameConverOp(hf_names=".fc.weight", mca_names=".eh_proj.weight"), + RenameConverOp(hf_names=".norm.weight", mca_names=".final_layernorm.weight"), ], ) diff --git a/mcore_adapter/src/mcore_adapter/models/qwen3_next/__init__.py b/mcore_adapter/src/mcore_adapter/models/qwen3_next/__init__.py index e0b8dc66e..96a62ae2f 100644 --- a/mcore_adapter/src/mcore_adapter/models/qwen3_next/__init__.py +++ b/mcore_adapter/src/mcore_adapter/models/qwen3_next/__init__.py @@ -3,6 +3,7 @@ import torch +from ..converter.convert_utils import StackedTensors from ..converter.dist_converter import ( default_dist_config, gdn_dist_config, @@ -12,26 +13,19 @@ from ..converter.template import ( ConverOp, CopyConverOp, + DropConverOp, GatedQKVConverOp, GDNConv1dConverOp, RenameConverOp, StackConverOp, Template, + ZeroCenteredRMSNormConverOp, register_template, ) from .config_qwen3_next import Qwen3NextConfig from .modeling_qwen3_next import Qwen3NextModel -@dataclass -class DropConverOp(ConverOp): - def _hf_to_mca(self, weights): - return [] - - def _mca_to_hf(self, weights): - return [] - - @dataclass class NextGDNConverOp(ConverOp): def __post_init__(self): @@ -63,41 +57,21 @@ def _hf_to_mca(self, weights): ) b, a = torch.split(ba_reshaped, [num_v_heads // num_qk_heads, num_v_heads // num_qk_heads], dim=1) q, k, v, z, b, a = [weight.reshape(-1, hidden_size) for weight in [q, k, v, z, b, a]] - in_proj_weight = torch.cat([q, k, v, z, b, a], dim=0).reshape(-1, hidden_size) - return in_proj_weight + return StackedTensors(tensors=[q, k, v, z, b, a], dim=0) def _mca_to_hf(self, weights): - in_proj_weight = weights[0] + assert len(weights) == 1 + assert isinstance(weights[0], StackedTensors) + q, k, v, z, b, a = weights[0].tensors hidden_size = self.mca_config.hidden_size - qk_head_dim = self.mca_config.linear_key_head_dim - v_head_dim = self.mca_config.linear_value_head_dim num_qk_heads = self.mca_config.linear_num_key_heads - num_v_heads = self.mca_config.linear_num_value_heads - qk_dim = qk_head_dim * num_qk_heads - v_dim = v_head_dim * num_v_heads - in_proj_weight = in_proj_weight.reshape(-1, hidden_size) - q, k, v, z, b, a = torch.split(in_proj_weight, [qk_dim, qk_dim, v_dim, v_dim, num_v_heads, num_v_heads], dim=0) q, k, v, z, b, a = [weight.reshape(num_qk_heads, -1, hidden_size) for weight in [q, k, v, z, b, a]] qkvz_weight = torch.cat([q, k, v, z], dim=1).reshape(-1, hidden_size) ba_weight = torch.cat([b, a], dim=1).reshape(-1, hidden_size) return [qkvz_weight, ba_weight] -@dataclass -class ZeroCenteredRMSNormConverOp(ConverOp): - def __post_init__(self): - super().__post_init__() - assert len(self.hf_names) == 1, f"ZeroCenteredRMSNormConverOp only support one name {self.hf_names}" - assert len(self.mca_names) == 1, f"ZeroCenteredRMSNormConverOp only support one name {self.mca_names}" - - def _hf_to_mca(self, weights): - return weights[0].clone() - 1 - - def _mca_to_hf(self, weights): - return weights[0].clone() + 1 - - register_dist_config( "qwen3_next", default_dist_config.merge_configs(shared_moe_dist_config).merge_configs(gdn_dist_config) ) @@ -198,7 +172,7 @@ def get_lora_conver_op(self, name, pattern_to_conver_ops: dict[str, ConverOp], l "layernorm_zero_centered_gamma": True, "hetereogenous_dist_checkpoint": True, "attention_output_gate": True, - "experimental_attention_variant": "gated_delta_net", + "linear_attention_type": "gated_delta_net", }, weight_converters=[ RenameConverOp(hf_names="lm_head.weight", mca_names="output_layer.weight"), diff --git a/mcore_adapter/src/mcore_adapter/models/qwen3_omni/modeling_qwen3_omni.py b/mcore_adapter/src/mcore_adapter/models/qwen3_omni/modeling_qwen3_omni.py index e0ed698e8..c8e424252 100644 --- a/mcore_adapter/src/mcore_adapter/models/qwen3_omni/modeling_qwen3_omni.py +++ b/mcore_adapter/src/mcore_adapter/models/qwen3_omni/modeling_qwen3_omni.py @@ -15,13 +15,13 @@ class Qwen3OmniMoeModel(Qwen3VLModel): def __init__(self, config: "Qwen3OmniMoeConfig", **kwargs): from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import ( - Qwen3OmniMoeAudioEncoderConfig, Qwen3OmniMoeVisionEncoderConfig, + Qwen3OmniMoeAudioEncoderConfig, ) from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import ( + Qwen3OmniMoeVisionEncoder, Qwen3OmniMoeAudioEncoder, Qwen3OmniMoePreTrainedModelForConditionalGeneration, - Qwen3OmniMoeVisionEncoder, _get_feat_extract_output_lengths, ) @@ -39,14 +39,18 @@ def __init__(self, config: "Qwen3OmniMoeConfig", **kwargs): Qwen3OmniMoeAudioEncoderConfig(**config.audio_config), attn_implementation="sdpa", torch_dtype=self.config.params_dtype, - ).to(torch.cuda.current_device()) + ) + if not config.init_model_with_meta_device: + self.audio_model = self.audio_model.to(torch.cuda.current_device()) for param in self.audio_model.parameters(): setattr(param, "sequence_parallel", config.sequence_parallel) self.vision_model = Qwen3OmniMoeVisionEncoder._from_config( Qwen3OmniMoeVisionEncoderConfig(**config.vision_config), attn_implementation="sdpa", torch_dtype=self.config.params_dtype, - ).to(torch.cuda.current_device()) + ) + if not config.init_model_with_meta_device: + self.vision_model = self.vision_model.to(torch.cuda.current_device()) # TODO: use_reentrant=True might cause error by twice forward/backward when # training images and videos simultaneously, https://github.com/pytorch/pytorch/issues/81296 if config.recompute_granularity == "full" and self.training: @@ -68,11 +72,14 @@ def __init__(self, config: "Qwen3OmniMoeConfig", **kwargs): self.talker = Qwen3OmniMoeTalkerForConditionalGeneration._from_config( Qwen3OmniMoeTalkerConfig(**config.talker_config), torch_dtype=self.config.params_dtype, - ).to(torch.cuda.current_device()) + ) self.code2wav = Qwen3OmniMoeCode2Wav._from_config( Qwen3OmniMoeCode2WavConfig(**config.code2wav_config), torch_dtype=self.config.params_dtype, - ).to(torch.cuda.current_device()) + ) + if not config.init_model_with_meta_device: + self.talker = self.talker.to(torch.cuda.current_device()) + self.code2wav = self.code2wav.to(torch.cuda.current_device()) # construct get_rope_index needed method and attrs self.get_rope_index = types.MethodType( @@ -105,14 +112,9 @@ def construct_inputs_embeds( inputs_embeds: [s, b, h] or [s/tp, b, h] when sequence parallel ranges: sequence range """ - visual_pos_masks, deepstack_visual_embeds = None, None - # TODO: same as qwen3-vl, only support images or videos since no deepstack_visual_embeds merge process currently - # maybe merge images and videos first to run vision_model and get deepstack_visual_embeds for images and videos simultaneously - assert pixel_values is None or pixel_values_videos is None, ( - "inputs with both images and videos are not supported temporarily" - ) + image_pos_masks = video_pos_masks = deepstack_image_embeds = deepstack_video_embeds = None if pixel_values is not None: - inputs_embeds, visual_pos_masks, deepstack_visual_embeds = super().construct_inputs_embeds( + inputs_embeds, image_pos_masks, deepstack_image_embeds = super().construct_inputs_embeds( input_ids, inputs_embeds, pixel_values, @@ -120,8 +122,8 @@ def construct_inputs_embeds( input_ranges, image_token_id, ) - elif pixel_values_videos is not None: - inputs_embeds, visual_pos_masks, deepstack_visual_embeds = super().construct_inputs_embeds( + if pixel_values_videos is not None: + inputs_embeds, video_pos_masks, deepstack_video_embeds = super().construct_inputs_embeds( input_ids, inputs_embeds, pixel_values_videos, @@ -129,102 +131,97 @@ def construct_inputs_embeds( input_ranges, video_token_id, ) + visual_pos_masks, deepstack_visual_embeds = self.merge_deepstack_embeds( + image_pos_masks, deepstack_image_embeds, video_pos_masks, deepstack_video_embeds + ) if input_features is None: return inputs_embeds, visual_pos_masks, deepstack_visual_embeds # for audio input embeds + # Follow the same pattern as image/video: build global mask & indices, use + # build_encoder_inputs / gather_encoder_outputs for load-balanced encoder + # execution across SP/CP ranks, then masked_scatter into inputs_embeds. + # (bs, freqs, frames) -> (total_frames, freqs) input_features = input_features.permute(0, 2, 1)[feature_attention_mask.bool()] - # TODO: audio can be treated as chunks of frames with chunk_size for sp/cp actually, - # chunk_size = 100 * (self.n_window_infer // (self.n_window * 2)) - # temporarily only split audios instead of chunks to simplify which may cause duplicated calculation for same audio - # maybe scatter chunks to sp/cp group for load balance furthermore - feat_mask = input_ids == audio_token_id - feat_culens = feature_lens.cumsum(dim=0, dtype=torch.int32).tolist() # use list - feat_embeds_culens = self._get_feat_extract_output_lengths(feature_lens).cumsum(dim=0, dtype=torch.int32) - required_feat = [] # features to vision tower - required_feat_lens = [] # feature lengths to vision tower - valid_feat_embeds_nums = [] # indicate the ranges of needed feature embeds - added_feat_indexes = [] # feature indexes included in input_ranges - for i in range(feat_mask.shape[0]): - for inputs_start, inputs_end in input_ranges: - # same as qwen-vl, get features included in a sub-range corresponding to each sample - valid_feat_embeds_start = feat_mask[:i].sum().item() - valid_feat_embeds_start += feat_mask[i, :inputs_start].sum().item() - embeds_num = feat_mask[i, inputs_start:inputs_end].sum().item() - valid_feat_embeds_end = valid_feat_embeds_start + embeds_num - used_embeds_culen_start = 0 # embeds seqlens before this sub-range - new_embeds_culen_start = 0 # embeds seqlens new added in this sub-range, new_embeds_seqlen_start >= used_embeds_seqlen_start - added_culen_before_used = 0 # embeds seqlens in before sub-ranges of input_ranges - embed_culen_end = feat_embeds_culens[-1] - for feat_index, feat_embeds_culen in enumerate(feat_embeds_culens): - if valid_feat_embeds_start < feat_embeds_culen: # included in current sub-range - if feat_index not in added_feat_indexes: - # included in current sub-range and have not been added before, add it - required_feat_lens.append(feature_lens[feat_index]) - # maybe extend together at last, while mapping from embeds length to feature length is not direct - required_feat.append( - input_features[ - (0 if feat_index == 0 else feat_culens[feat_index - 1]) : feat_culens[feat_index] - ] - ) - added_feat_indexes.append(feat_index) - else: - # included in current sub-range but have been added by previous sub-range of this sample, skip it - new_embeds_culen_start = feat_embeds_culen - else: # not included in current sub-range - used_embeds_culen_start = feat_embeds_culen - new_embeds_culen_start = feat_embeds_culen - if feat_index in added_feat_indexes: # included in before sub-ranges of input_ranges - before_culen = 0 if feat_index == 0 else feat_embeds_culens[feat_index - 1].item() - added_culen_before_used += feat_embeds_culen - before_culen - if valid_feat_embeds_end <= feat_embeds_culen: - embed_culen_end = feat_embeds_culen - break - # embeds offset in range for this sub-range: offset_in_range = offset_in_start_feat + emb_len_of_pre_subranges - embeds_needed_start = valid_feat_embeds_start - used_embeds_culen_start + added_culen_before_used - embeds_needed_end = valid_feat_embeds_end - used_embeds_culen_start + added_culen_before_used - if embeds_needed_start < embeds_needed_end: - valid_feat_embeds_nums.append((embeds_needed_start, embeds_needed_end)) + audio_mask = input_ids == audio_token_id + audio_indices = torch.full_like(audio_mask, -1, dtype=torch.long) + audio_indices[audio_mask] = torch.arange(audio_mask.sum(), device=audio_indices.device) - if len(valid_feat_embeds_nums) == 0: - # should we use dummy feature input to handle this, _handle_missing_visual is used in qwen-vl - return inputs_embeds, visual_pos_masks, deepstack_visual_embeds + # audio_input_lengths: raw feature frame counts per audio segment + audio_input_lengths = feature_lens.tolist() + # audio_output_lengths: embedding token counts per audio segment after encoder + audio_output_lengths = self._get_feat_extract_output_lengths(feature_lens).tolist() - required_feat = torch.cat(required_feat, dim=0) - required_feat_lens = torch.stack(required_feat_lens, dim=0) - feat_model_dtype = self.audio_model.layers[0].fc1.weight.dtype - required_feat = required_feat.type(feat_model_dtype) - # convert to (freqs, total_frames) for input_features to use audio_tower from hf - required_feat = required_feat.permute(1, 0) - feat_embeds = self.audio_model(required_feat, required_feat_lens) - feat_embeds = feat_embeds.last_hidden_state.to(inputs_embeds.device, inputs_embeds.dtype) - feat_mask = torch.cat( - [feat_mask[:, inputs_start:inputs_end] for inputs_start, inputs_end in input_ranges], dim=1 - ) - needed_feat_embeds_num = feat_mask.sum().item() - needed_feat_embeds = torch.zeros( - [needed_feat_embeds_num] + list(feat_embeds.shape[1:]), - dtype=inputs_embeds.dtype, - device=inputs_embeds.device, + split_plan, input_features_split, feature_lens_split, _ = self.build_encoder_inputs( + audio_input_lengths, input_features, feature_lens, None ) - added_num = 0 - for start, end in valid_feat_embeds_nums: - embeds_num = end - start - needed_feat_embeds[added_num : added_num + embeds_num] = feat_embeds[start:end] - added_num += embeds_num - assert added_num == needed_feat_embeds_num + feat_model_dtype = self.audio_model.layers[0].fc1.weight.dtype + input_features_split = input_features_split.type(feat_model_dtype) + # convert to (freqs, total_frames) for audio_tower from hf + input_features_split = input_features_split.permute(1, 0) + audio_outputs = self.audio_model(input_features_split, feature_lens_split) + audio_embeds = audio_outputs.last_hidden_state + + audio_embeds = self.gather_encoder_outputs(audio_embeds, split_plan, audio_output_lengths) + audio_embeds = audio_embeds.to(inputs_embeds.device, inputs_embeds.dtype) + + full_sequence_output = input_ranges is None + if full_sequence_output: + # Packing mode: inputs_embeds is the full sequence. + # Use the full audio_mask for masked_scatter. + all_selected_indices = audio_indices[audio_mask] + inputs_embeds = inputs_embeds.transpose(0, 1) # [s, b, h] -> [b, s, h] + full_selected_mask = audio_mask.unsqueeze(-1).expand_as(inputs_embeds) + inputs_embeds = inputs_embeds.masked_scatter(full_selected_mask, audio_embeds[all_selected_indices]) + inputs_embeds = inputs_embeds.transpose(0, 1).contiguous() + else: + # Normal mode: slice mask & indices to current rank's input_ranges + audio_mask = torch.cat([audio_mask[:, start:end] for start, end in input_ranges], dim=1) + selected_indices = torch.cat([audio_indices[:, start:end] for start, end in input_ranges], dim=1) + selected_indices = selected_indices[selected_indices != -1] - inputs_embeds = inputs_embeds.transpose(0, 1) # [s, b, h] -> [b, s, h] - feat_mask = feat_mask.unsqueeze(-1).expand_as(inputs_embeds) - inputs_embeds = inputs_embeds.masked_scatter(feat_mask, needed_feat_embeds) - inputs_embeds = inputs_embeds.transpose(0, 1).contiguous() + inputs_embeds = inputs_embeds.transpose(0, 1) # [s, b, h] -> [b, s, h] + selected_mask = audio_mask.unsqueeze(-1).expand_as(inputs_embeds) + inputs_embeds = inputs_embeds.masked_scatter(selected_mask, audio_embeds[selected_indices]) + inputs_embeds = inputs_embeds.transpose(0, 1).contiguous() return inputs_embeds, visual_pos_masks, deepstack_visual_embeds + def construct_multimodal_inputs( + self, + input_ids, + inputs_embeds, + inputs_ranges, + *, + pixel_values=None, + pixel_values_videos=None, + image_grid_thw=None, + video_grid_thw=None, + input_features=None, + feature_lens=None, + feature_attention_mask=None, + **kwargs, + ): + return self.construct_inputs_embeds( + input_ids, + inputs_embeds, + pixel_values, + image_grid_thw, + pixel_values_videos, + video_grid_thw, + input_features, + feature_lens, + feature_attention_mask, + inputs_ranges, + self.config.image_token_id, + self.config.video_token_id, + self.config.audio_token_id, + ) + def forward( self, input_ids: "torch.Tensor", @@ -242,8 +239,10 @@ def forward( feature_attention_mask: Optional["torch.Tensor"] = None, **kwargs, ) -> "torch.Tensor": - force_vit_image = kwargs.pop("force_vit_image", False) - force_vit_video = kwargs.pop("force_vit_video", False) + kwargs.pop("force_vit_image", None) + kwargs.pop("force_vit_video", None) + packed_seq_params = kwargs.get("packed_seq_params", None) + feature_lens = None if position_ids is None and input_ids is not None: if feature_attention_mask is not None: @@ -258,48 +257,33 @@ def forward( second_per_grids=video_second_per_grid, ) - cp_batch = { - "input_ids": input_ids, - "attention_mask": attention_mask, - } - if self.config.context_parallel_size > 1: - cp_batch = {k: v.clone() if v is not None else None for k, v in cp_batch.items()} - cp_batch = super(Qwen3VLModel, self).get_batch_on_this_cp_rank(cp_batch, dim3_keys=[]) + state = self.prepare_packing_state( + input_ids, position_ids, attention_mask, packed_seq_params, + ) if not self.pre_process or decoder_input is not None: return super(Qwen3VLModel, self).forward( - decoder_input=decoder_input, labels=labels, position_ids=position_ids, **cp_batch, **kwargs + decoder_input=decoder_input, labels=labels, + position_ids=state.position_ids, + **state.cp_batch, **kwargs ) - inputs_ranges = self.get_input_ranges(input_ids.shape[1]) - - inputs_embeds = self.embedding(input_ids=cp_batch["input_ids"], position_ids=None) - - if pixel_values is not None or pixel_values_videos is not None: - inputs_embeds, visual_pos_masks, deepstack_visual_embeds = self.construct_inputs_embeds( - input_ids, - inputs_embeds, - pixel_values, - image_grid_thw, - pixel_values_videos, - video_grid_thw, - input_features, - feature_lens, - feature_attention_mask, - inputs_ranges, - self.config.image_token_id, - self.config.video_token_id, - self.config.audio_token_id, - ) - elif force_vit_image or force_vit_video: - inputs_embeds, visual_pos_masks, deepstack_visual_embeds = self._handle_missing_visual(inputs_embeds) + result = self.build_multimodal_embeddings( + input_ids, attention_mask, packed_seq_params, state, + multimodal_kwargs=dict( + pixel_values=pixel_values, pixel_values_videos=pixel_values_videos, + image_grid_thw=image_grid_thw, video_grid_thw=video_grid_thw, + input_features=input_features, feature_lens=feature_lens, + feature_attention_mask=feature_attention_mask, + ), + ) return super(Qwen3VLModel, self).forward( - decoder_input=inputs_embeds, + decoder_input=result.inputs_embeds, labels=labels, - position_ids=position_ids, - visual_pos_masks=visual_pos_masks, - deepstack_visual_embeds=deepstack_visual_embeds, - **cp_batch, + position_ids=result.position_ids, + visual_pos_masks=result.visual_pos_masks, + deepstack_visual_embeds=result.deepstack_visual_embeds, + **result.cp_batch, **kwargs, ) diff --git a/mcore_adapter/src/mcore_adapter/models/qwen3_vl/modeling_qwen3_vl.py b/mcore_adapter/src/mcore_adapter/models/qwen3_vl/modeling_qwen3_vl.py index d50ee88e1..d0a5db676 100644 --- a/mcore_adapter/src/mcore_adapter/models/qwen3_vl/modeling_qwen3_vl.py +++ b/mcore_adapter/src/mcore_adapter/models/qwen3_vl/modeling_qwen3_vl.py @@ -13,6 +13,7 @@ from ..auto.modeling_auto import register_model from ..model_factory import McaGPTModel from ..model_utils import ModuleUtilsMixin +from ..sequence_packing_mixin import MultimodalEmbeddingMixin from .config_qwen3_vl import Qwen3VLConfig from .rope_utils import Qwen3VLMultimodalRotaryEmbedding, get_rope_index from .transformer_block import Qwen3VLTransformerBlock @@ -139,7 +140,7 @@ def forward( @register_model("qwen3_vl") -class Qwen3VLModel(Qwen3VLGPTModel, ModuleUtilsMixin): +class Qwen3VLModel(Qwen3VLGPTModel, ModuleUtilsMixin, MultimodalEmbeddingMixin): config_class = Qwen3VLConfig def __init__(self, config: "Qwen3VLConfig", **kwargs): @@ -158,7 +159,9 @@ def __init__(self, config: "Qwen3VLConfig", **kwargs): Qwen3VLVisionConfig(**config.vision_config), attn_implementation="sdpa", torch_dtype=self.config.params_dtype, - ).to(torch.cuda.current_device()) + ) + if not config.init_model_with_meta_device: + self.vision_model = self.vision_model.to(torch.cuda.current_device()) # TODO: use_reentrant=True might cause error by twice forward/backward when # training images and videos simultaneously, https://github.com/pytorch/pytorch/issues/81296 if config.recompute_granularity == "full" and self.training: @@ -191,11 +194,18 @@ def construct_inputs_embeds( pixel_values: "torch.Tensor", grid_thw: "torch.LongTensor", input_ranges: List[List[int]], - media_token_id: int, + media_token_id: int ): """ - inputs_embeds: [s, b, h] or [s/tp, b, h] when sequence parallel - ranges: sequence range + Construct inputs_embeds with vision embeddings replacing media tokens. + + Args: + input_ids: [b, s] original input_ids (full sequence, used for image_mask) + inputs_embeds: [s, b, h] or [s/tp/cp, b, h] when sequence/context parallel + pixel_values: flattened pixel values for vision model + grid_thw: grid dimensions for each image/video + input_ranges: sequence ranges for the current rank (CP+SP sliced) + media_token_id: token id for the media placeholder """ image_mask = input_ids == media_token_id image_indices = torch.full_like(image_mask, -1, dtype=torch.long) @@ -220,23 +230,39 @@ def construct_inputs_embeds( image_embeds = self.gather_encoder_outputs(image_embeds, split_plan, image_output_lengths) image_embeds = image_embeds.to(inputs_embeds.device, inputs_embeds.dtype) - image_mask = torch.cat([image_mask[:, start:end] for start, end in input_ranges], dim=1) - selected_indices = torch.cat([image_indices[:, start:end] for start, end in input_ranges], dim=1) - selected_indices = selected_indices[selected_indices != -1] - deepstack_image_embeds = [ self.gather_encoder_outputs(deepstack_image_embed, split_plan, image_output_lengths) for deepstack_image_embed in deepstack_image_embeds ] - for i, deepstack_image_embed in enumerate(deepstack_image_embeds): - deepstack_image_embeds[i] = deepstack_image_embed[selected_indices] - inputs_embeds = inputs_embeds.transpose(0, 1) # [s, b, h] -> [b, s, h] - selected_mask = image_mask.unsqueeze(-1).expand_as(inputs_embeds) - inputs_embeds = inputs_embeds.masked_scatter(selected_mask, image_embeds[selected_indices]) - inputs_embeds = inputs_embeds.transpose(0, 1).contiguous() - - return inputs_embeds, image_mask, deepstack_image_embeds + # If True, inputs_embeds is the full sequence (not CP/SP sliced) and masked_scatter operates on the + # full sequence. Used in sequence packing mode where CP+packing is done after construct_inputs_embeds. + full_sequence_output = input_ranges is None + if full_sequence_output: + # In packing mode: inputs_embeds is the full sequence (not sliced). + # Use the full image_mask for masked_scatter so dimensions match. + all_selected_indices = image_indices[image_mask] + inputs_embeds = inputs_embeds.transpose(0, 1) # [s, b, h] -> [b, s, h] + full_selected_mask = image_mask.unsqueeze(-1).expand_as(inputs_embeds) + inputs_embeds = inputs_embeds.masked_scatter(full_selected_mask, image_embeds[all_selected_indices]) + inputs_embeds = inputs_embeds.transpose(0, 1).contiguous() + # Return full-sequence image_mask and all deepstack embeds (not + # sliced). The caller will do CP+packing on visual_pos_masks and + # then re-select deepstack embeds from the packed mask. + return inputs_embeds, image_mask, deepstack_image_embeds + else: + # Normal mode: inputs_embeds is already CP+SP sliced. + # Slice image_mask and indices by input_ranges. + sliced_image_mask = torch.cat([image_mask[:, start:end] for start, end in input_ranges], dim=1) + sliced_indices = torch.cat([image_indices[:, start:end] for start, end in input_ranges], dim=1) + selected_indices = sliced_indices[sliced_indices != -1] + for i, deepstack_image_embed in enumerate(deepstack_image_embeds): + deepstack_image_embeds[i] = deepstack_image_embed[selected_indices] + inputs_embeds = inputs_embeds.transpose(0, 1) # [s, b, h] -> [b, s, h] + selected_mask = sliced_image_mask.unsqueeze(-1).expand_as(inputs_embeds) + inputs_embeds = inputs_embeds.masked_scatter(selected_mask, image_embeds[selected_indices]) + inputs_embeds = inputs_embeds.transpose(0, 1).contiguous() + return inputs_embeds, sliced_image_mask, deepstack_image_embeds def build_encoder_inputs( self, @@ -321,43 +347,71 @@ def gather_encoder_outputs( return encoder_small_batch_size_gather(output_features) def get_batch_on_this_cp_rank(self, batch, dim3_keys: List[str] = ["attention_mask"]): - # VLM need to view all input_ids and media features - loss_needed_items = { - "labels": batch.pop("labels", None), - } - loss_needed_items = super().get_batch_on_this_cp_rank(loss_needed_items, dim3_keys=dim3_keys) - batch.update(loss_needed_items) + # VLM forward() handles input_ids and attention_mask splitting internally + skipped = {} + for key in ("input_ids", "attention_mask"): + if key in batch: + skipped[key] = batch.pop(key) + batch = super().get_batch_on_this_cp_rank(batch, dim3_keys=dim3_keys) + batch.update(skipped) return batch - def get_input_ranges(self, total_seqlen): - # context parallel 的计算有问题 - slice_rank, slice_size = 0, 1 - if self.config.sequence_parallel: - slice_rank = mpu.get_tensor_model_parallel_rank() - slice_size = mpu.get_tensor_model_parallel_world_size() - - def get_sequence_range(start, end, rank, size): - return start + (end - start) * rank // size, start + (end - start) * (rank + 1) // size - - if self.config.context_parallel_size <= 1: - return [list(get_sequence_range(0, total_seqlen, slice_rank, slice_size))] - cp_rank = mpu.get_context_parallel_rank() - cp_size = mpu.get_context_parallel_world_size() - left_start = (total_seqlen // cp_size // 2) * cp_rank - left_end = (total_seqlen // cp_size // 2) * (cp_rank + 1) - right_start = total_seqlen - left_end - right_end = total_seqlen - left_start - slice_len = (left_end - left_start + right_end - right_start) // slice_size - start = left_start + slice_len * slice_rank - end = start + slice_len - if start >= left_end: - start = start - left_end + right_start - end = start + slice_len - return [[start, end]] - if end <= left_end: - return [[start, end]] - end = end - left_end + right_start - return [[start, left_end], [right_start, end]] + + def merge_deepstack_embeds(self, image_mask, deepstack_image_embeds, video_mask, deepstack_video_embeds): + visual_pos_masks = None + deepstack_visual_embeds = None + if image_mask is not None and video_mask is not None: + # aggregate visual_pos_masks and deepstack_visual_embeds + visual_pos_masks = image_mask | video_mask + deepstack_visual_embeds = [] + image_mask_joint = image_mask[visual_pos_masks] + video_mask_joint = video_mask[visual_pos_masks] + for img_embed, vid_embed in zip(deepstack_image_embeds, deepstack_video_embeds): + embed_joint = img_embed.new_zeros(visual_pos_masks.sum(), img_embed.shape[-1]).to(img_embed.device) + embed_joint[image_mask_joint, :] = img_embed + embed_joint[video_mask_joint, :] = vid_embed + deepstack_visual_embeds.append(embed_joint) + elif image_mask is not None: + visual_pos_masks = image_mask + deepstack_visual_embeds = deepstack_image_embeds + elif video_mask is not None: + visual_pos_masks = video_mask + deepstack_visual_embeds = deepstack_video_embeds + return visual_pos_masks, deepstack_visual_embeds + + def construct_multimodal_inputs( + self, + input_ids, + inputs_embeds, + inputs_ranges, + *, + pixel_values=None, + pixel_values_videos=None, + image_grid_thw=None, + video_grid_thw=None, + force_vit_image=False, + force_vit_video=False, + **kwargs, + ): + image_pos_masks = video_pos_masks = deepstack_image_embeds = deepstack_video_embeds = None + if pixel_values is not None: + inputs_embeds, image_pos_masks, deepstack_image_embeds = self.construct_inputs_embeds( + input_ids, inputs_embeds, pixel_values, image_grid_thw, + inputs_ranges, self.config.image_token_id + ) + elif force_vit_image: + inputs_embeds, image_pos_masks, deepstack_image_embeds = self._handle_missing_visual(inputs_embeds) + if pixel_values_videos is not None: + inputs_embeds, video_pos_masks, deepstack_video_embeds = self.construct_inputs_embeds( + input_ids, inputs_embeds, pixel_values_videos, video_grid_thw, + inputs_ranges, self.config.video_token_id + ) + elif force_vit_video: + inputs_embeds, video_pos_masks, deepstack_video_embeds = self._handle_missing_visual(inputs_embeds) + visual_pos_masks, deepstack_visual_embeds = self.merge_deepstack_embeds( + image_pos_masks, deepstack_image_embeds, video_pos_masks, deepstack_video_embeds + ) + return inputs_embeds, visual_pos_masks, deepstack_visual_embeds def forward( self, @@ -374,54 +428,38 @@ def forward( ) -> "torch.Tensor": force_vit_image = kwargs.pop("force_vit_image", False) force_vit_video = kwargs.pop("force_vit_video", False) + if position_ids is None and input_ids is not None: position_ids, _ = get_rope_index(self.config, input_ids, image_grid_thw, video_grid_thw) - cp_batch = { - "input_ids": input_ids, - "attention_mask": attention_mask, - } - if self.config.context_parallel_size > 1: - cp_batch = {k: v.clone() if v is not None else None for k, v in cp_batch.items()} - cp_batch = super().get_batch_on_this_cp_rank(cp_batch, dim3_keys=[]) + packed_seq_params = kwargs.get("packed_seq_params", None) + + state = self.prepare_packing_state( + input_ids, position_ids, attention_mask, packed_seq_params, + ) - if not self.pre_process or (pixel_values is None and pixel_values_videos is None) or decoder_input is not None: + if not self.pre_process or decoder_input is not None: return super().forward( - decoder_input=decoder_input, labels=labels, position_ids=position_ids, **cp_batch, **kwargs + decoder_input=decoder_input, labels=labels, + position_ids=state.position_ids, + **state.cp_batch, **kwargs ) - inputs_ranges = self.get_input_ranges(input_ids.shape[1]) + result = self.build_multimodal_embeddings( + input_ids, attention_mask, packed_seq_params, state, + multimodal_kwargs=dict( + pixel_values=pixel_values, pixel_values_videos=pixel_values_videos, + image_grid_thw=image_grid_thw, video_grid_thw=video_grid_thw, + force_vit_image=force_vit_image, force_vit_video=force_vit_video, + ), + ) - inputs_embeds = self.embedding(input_ids=cp_batch["input_ids"], position_ids=None) - if pixel_values is not None: - # get deepstack emb - inputs_embeds, visual_pos_masks, deepstack_visual_embeds = self.construct_inputs_embeds( - input_ids, - inputs_embeds, - pixel_values, - image_grid_thw, - inputs_ranges, - self.config.image_token_id, - ) - elif force_vit_image: - inputs_embeds, visual_pos_masks, deepstack_visual_embeds = self._handle_missing_visual(inputs_embeds) - if pixel_values_videos is not None: - inputs_embeds, visual_pos_masks, deepstack_visual_embeds = self.construct_inputs_embeds( - input_ids, - inputs_embeds, - pixel_values_videos, - video_grid_thw, - inputs_ranges, - self.config.video_token_id, - ) - elif force_vit_video: - inputs_embeds, visual_pos_masks, deepstack_visual_embeds = self._handle_missing_visual(inputs_embeds) return super().forward( - decoder_input=inputs_embeds, + decoder_input=result.inputs_embeds, labels=labels, - position_ids=position_ids, - visual_pos_masks=visual_pos_masks, - deepstack_visual_embeds=deepstack_visual_embeds, - **cp_batch, + position_ids=result.position_ids, + visual_pos_masks=result.visual_pos_masks, + deepstack_visual_embeds=result.deepstack_visual_embeds, + **result.cp_batch, **kwargs, ) diff --git a/mcore_adapter/src/mcore_adapter/models/qwen3_vl/rope_utils.py b/mcore_adapter/src/mcore_adapter/models/qwen3_vl/rope_utils.py index 8a9b09e07..25629caa8 100644 --- a/mcore_adapter/src/mcore_adapter/models/qwen3_vl/rope_utils.py +++ b/mcore_adapter/src/mcore_adapter/models/qwen3_vl/rope_utils.py @@ -3,9 +3,10 @@ import torch from megatron.core import parallel_state from megatron.core.models.common.embeddings.rope_utils import ( + _apply_rotary_pos_emb_bshd, get_pos_emb_on_this_cp_rank, ) -from torch import nn +from torch import Tensor, nn from .config_qwen3_vl import Qwen3VLConfig @@ -53,6 +54,10 @@ def __init__( self.cp_group = ( cp_group if cp_group is not None else parallel_state.get_context_parallel_group(check_initialized=False) ) + # When True, position_ids are already CP-zigzag-split and packed into + # THD format (matching mbridge's preprocess_packed_seqs). We skip + # additional CP slicing here to avoid double-slicing. + self.is_thd_format = False def apply_interleaved_mrope(self, freqs, mrope_section): """Apply interleaved MRoPE to 3D rotary embeddings. @@ -104,53 +109,74 @@ def forward( # shape (seq_length, bs, 1, 2 * dim) emb = emb[..., None, :].transpose(0, 1).contiguous() - if cp_group is None: - cp_group = self.cp_group - if cp_group is not None and cp_group.size() > 1: - # slice rotary_pos_emb along sequence dimension and select the parition of the current - # CP rank - emb = get_pos_emb_on_this_cp_rank(emb, 0, cp_group) + if not self.is_thd_format: + if cp_group is None: + cp_group = self.cp_group + if cp_group is not None and cp_group.size() > 1: + # slice rotary_pos_emb along sequence dimension and select the parition of the + # current CP rank. Skipped in THD format because position_ids are already + # CP+SP sliced and packed. + emb = get_pos_emb_on_this_cp_rank(emb, 0, cp_group) return emb -# def apply_rotary_pos_emb_thd_absolute( -# t: torch.Tensor, cu_seqlens: torch.Tensor, freqs: torch.Tensor, rotary_interleaved: bool = False -# ) -> torch.Tensor: -# """A baseline implementation of applying RoPE for `thd` format. +def _apply_rotary_pos_emb_thd_absolute( + t: Tensor, cu_seqlens: Tensor, freqs: Tensor, rotary_interleaved: bool = False +) -> Tensor: + """Apply RoPE for THD format without additional CP slicing. -# Args: -# t (Tensor): Input tensor T is of shape [t, h, d] -# cu_seqlens(Tensor): Cumulative sum of sequence lengths in a batch for `t`, -# with shape [b + 1] and dtype torch.int32. -# freqs (Tensor): Rotary Positional embedding tensor freq is of shape [max_s, 1, 1, d] + In Qwen3-VL sequence packing, the freqs are already generated from + CP-zigzag-split and packed position_ids (matching mbridge's + preprocess_packed_seqs). We must NOT apply any additional CP zigzag + slicing (which the standard Megatron ``_apply_rotary_pos_emb_thd`` + would do). Instead we simply treat the packed tensor as + ``[total_tokens, 1, hidden]`` and delegate to the regular BSHD kernel. -# Returns: -# Tensor: Shape [t, h, d]. The input tensor after applying RoPE. -# """ -# return _apply_rotary_pos_emb_bshd(t[:, None], freqs, rotary_interleaved=rotary_interleaved).squeeze(1) + Args: + t: Input tensor of shape ``[total_tokens, num_heads, head_dim]``. + cu_seqlens: Cumulative sequence lengths (unused, kept for API compat). + freqs: Rotary frequencies of shape ``[total_tokens, 1, 1, 2*rot_dim]``. + rotary_interleaved: Whether to use interleaved layout. + Returns: + Tensor of shape ``[total_tokens, num_heads, head_dim]``. + """ + return _apply_rotary_pos_emb_bshd( + t[:, None], freqs, rotary_interleaved=rotary_interleaved + ).squeeze(1) -# def apply_rotary_pos_emb_absolute( -# t: torch.Tensor, -# freqs: torch.Tensor, -# config: Qwen3VLConfig, -# cu_seqlens: Optional[torch.Tensor] = None, -# ): -# """ -# Reroute to the appropriate apply_rotary_pos_emb function depending on -# bshd (conventional) / thd (packed seq) format -# In Qwen3-VL, the shape of freqs is (seq_length, bs, 1, 2 * dim) instead of [max_seqlen, 1, 1, 2 * dim] -# """ -# assert not config.apply_rope_fusion +def apply_rotary_pos_emb_absolute( + t: Tensor, + freqs: Tensor, + config: Qwen3VLConfig, + cu_seqlens: Optional[Tensor] = None, + **kwargs, +) -> Tensor: + """Apply rotary position embedding for Qwen3-VL. -# if cu_seqlens is None: -# result = _apply_rotary_pos_emb_bshd(t, freqs, rotary_interleaved=config.rotary_interleaved) -# else: -# result = apply_rotary_pos_emb_thd_absolute(t, cu_seqlens, freqs, rotary_interleaved=config.rotary_interleaved) + Routes to BSHD or THD-absolute implementation depending on whether + ``cu_seqlens`` is provided. Unlike the standard Megatron + ``apply_rotary_pos_emb``, the THD path here does **not** perform any + additional CP-related frequency slicing because the frequencies are + already computed from CP-zigzag-split and packed position_ids (matching + mbridge's preprocess_packed_seqs). -# return result + In Qwen3-VL the shape of *freqs* is ``(seq_length, bs, 1, 2*dim)`` instead + of the standard ``[max_seqlen, 1, 1, 2*dim]``. + Args: + t: Query or key tensor. + freqs: Rotary frequency tensor. + config: Qwen3VL transformer config. + cu_seqlens: Cumulative sequence lengths for THD format (or ``None``). + **kwargs: Ignored (absorbs ``mscale``, ``cp_group`` from caller). + """ + if cu_seqlens is None: + return _apply_rotary_pos_emb_bshd(t, freqs, rotary_interleaved=config.rotary_interleaved) + return _apply_rotary_pos_emb_thd_absolute( + t, cu_seqlens, freqs, rotary_interleaved=config.rotary_interleaved + ) # copy from transformers==4.57.0 def get_rope_index( diff --git a/mcore_adapter/src/mcore_adapter/models/sequence_packing_mixin.py b/mcore_adapter/src/mcore_adapter/models/sequence_packing_mixin.py new file mode 100644 index 000000000..4caf3955b --- /dev/null +++ b/mcore_adapter/src/mcore_adapter/models/sequence_packing_mixin.py @@ -0,0 +1,572 @@ +""" +Multimodal Embedding Mixin. + +Provides ``prepare_packing_state`` and ``build_multimodal_embeddings`` for +multimodal models (Qwen3VL, Qwen3OmniMoe, Qwen3_5, Qwen3OmniNext). +Handles both packing and non-packing paths with a unified two-phase API. + +The two-phase API separates lightweight input preparation (needed by +**all** PP stages) from heavy multimodal embedding construction (only needed +on the ``pre_process`` stage when ``decoder_input is None``): + +Phase 1 – ``prepare_packing_state`` (all PP stages): + Build CP batch & pack position_ids (when packing). + +Phase 2 – ``build_multimodal_embeddings`` (pre_process only): + 1. Compute text embeddings (full-sequence or CP-sliced depending on mode) + 2. Replace media tokens with encoder outputs via each model's own + ``construct_multimodal_inputs`` hook + 3. Pack ``inputs_embeds`` via CP-zigzag + SP scatter (when packing) + 4. Optionally pack ``visual_pos_masks`` / ``deepstack_visual_embeds`` +""" + +from dataclasses import dataclass, field +from typing import Optional + +import torch +from megatron.core import mpu, tensor_parallel +from megatron.core.packed_seq_params import PackedSeqParams + +from .model_factory import McaGPTModel + + +@dataclass +class PackingState: + """Lightweight result of :meth:`MultimodalEmbeddingMixin.prepare_packing_state`. + + Contains only the bookkeeping data needed by every PP stage: + packed ``position_ids`` and the CP-sliced batch dict. + """ + + position_ids: torch.Tensor + cp_batch: dict + is_packing: bool + + +@dataclass +class MultimodalEmbeddingResult: + """Result of :meth:`MultimodalEmbeddingMixin.build_multimodal_embeddings`.""" + + inputs_embeds: torch.Tensor + position_ids: torch.Tensor + cp_batch: dict + visual_pos_masks: Optional[torch.Tensor] = None + deepstack_visual_embeds: Optional[list] = None + extra_kwargs: dict = field(default_factory=dict) + + +class MultimodalEmbeddingMixin: + """ + Mixin that unifies the multimodal embedding construction pipeline for + VLM / omni models, handling both sequence-packing and non-packing paths. + + Models using this mixin **must** provide: + + - ``self.config`` (``sequence_parallel``, ``context_parallel_size``, …) + - ``self.embedding`` (with ``scatter_to_sequence_parallel``, + ``reduce_scatter_embeddings``) + - ``self.rotary_pos_emb`` (with ``is_thd_format``) + - ``self.pre_process`` (bool) + - ``self.mtp_process`` (bool) + - ``get_batch_on_this_cp_rank(batch, dim3_keys)`` (from ``McaGPTModel``) + + Each concrete model must also implement + :meth:`construct_multimodal_inputs` (see docstring below). + """ + + # ------------------------------------------------------------------ + # Public API + # ------------------------------------------------------------------ + + def is_sequence_packing(self, packed_seq_params, attention_mask): + """Determine if we are in sequence packing mode.""" + return packed_seq_params is not None and attention_mask is not None + + @staticmethod + def build_attention_mask_from_packed_seq_params( + packed_seq_params: PackedSeqParams, seq_len: int, device: torch.device + ) -> torch.Tensor: + """Construct neat-packing attention_mask from cu_seqlens_q. + + Each sub-sequence is labeled 1, 2, 3, ...; positions beyond the last + sub-sequence boundary are left as 0 (padding). + + Args: + packed_seq_params: contains cu_seqlens_q with shape [num_seqs + 1]. + seq_len: total sequence length of input_ids. + device: device to create the tensor on. + + Returns: + attention_mask of shape [batch_size=num_seqs, seq_len_per_sample] + when cu_seqlens indicates per-sample layout, or [1, seq_len] when + cu_seqlens indicates flat packed layout. + """ + cu_seqlens = packed_seq_params.cu_seqlens_q + num_seqs = len(cu_seqlens) - 1 + total_tokens = cu_seqlens[-1].item() + + if total_tokens == seq_len: + # Flat packed layout: all sub-seqs concatenated in one row + attention_mask = torch.zeros((1, seq_len), dtype=torch.long, device=device) + for seq_id in range(num_seqs): + start = cu_seqlens[seq_id].item() + end = cu_seqlens[seq_id + 1].item() + attention_mask[0, start:end] = seq_id + 1 + else: + # Per-sample layout: [num_seqs, max_seq_len] + max_seq_len = seq_len + attention_mask = torch.zeros( + (num_seqs, max_seq_len), dtype=torch.long, device=device + ) + for i in range(num_seqs): + valid_len = (cu_seqlens[i + 1] - cu_seqlens[i]).item() + attention_mask[i, :valid_len] = 1 + return attention_mask + + def get_input_ranges(self, total_seqlen: int) -> list[list[int]]: + """Compute the local sequence ranges for current SP/CP rank. + + In non-packing mode, each rank processes a slice of the full + sequence determined by its tensor-parallel and context-parallel + position. Returns a list of ``[start, end)`` ranges. + """ + slice_rank, slice_size = 0, 1 + if self.config.sequence_parallel: + slice_rank = mpu.get_tensor_model_parallel_rank() + slice_size = mpu.get_tensor_model_parallel_world_size() + + def get_sequence_range(start, end, rank, size): + return start + (end - start) * rank // size, start + (end - start) * (rank + 1) // size + + if self.config.context_parallel_size <= 1: + return [list(get_sequence_range(0, total_seqlen, slice_rank, slice_size))] + + cp_rank = mpu.get_context_parallel_rank() + cp_size = mpu.get_context_parallel_world_size() + left_start = (total_seqlen // cp_size // 2) * cp_rank + left_end = (total_seqlen // cp_size // 2) * (cp_rank + 1) + right_start = total_seqlen - left_end + right_end = total_seqlen - left_start + slice_len = (left_end - left_start + right_end - right_start) // slice_size + start = left_start + slice_len * slice_rank + end = start + slice_len + if start >= left_end: + start = start - left_end + right_start + end = start + slice_len + return [[start, end]] + if end <= left_end: + return [[start, end]] + end = end - left_end + right_start + return [[start, left_end], [right_start, end]] + + def prepare_packing_state( + self, + input_ids: torch.Tensor, + position_ids: torch.Tensor, + attention_mask: Optional[torch.Tensor], + packed_seq_params: Optional[PackedSeqParams], + ) -> PackingState: + """Phase 1: lightweight packing bookkeeping (all PP stages). + + Builds the CP-sliced batch dict and packs ``position_ids`` when in + sequence-packing mode. Also patches Megatron's + ``apply_rotary_pos_emb`` so that the THD path does NOT perform + CP-based frequency slicing (freqs are generated from packed, + CP-sliced position_ids). + + This is cheap and must run on **every** PP stage so that + non-``pre_process`` stages get correct ``position_ids`` and + ``cp_batch``. + + Args: + input_ids: ``[b, s]`` + position_ids: ``[C, b, s]`` (already computed by the caller) + attention_mask: ``[b, s]`` (``None`` when not packing) + packed_seq_params: packing parameters + """ + if attention_mask is None and packed_seq_params is not None: + attention_mask = self.build_attention_mask_from_packed_seq_params( + packed_seq_params, input_ids.shape[1], input_ids.device + ) + is_packing = self.is_sequence_packing(packed_seq_params, attention_mask) + + if is_packing: + self._patch_rotary_pos_emb() + + cp_batch = self._build_cp_batch(input_ids, attention_mask, is_packing, packed_seq_params) + if is_packing: + position_ids = self._pack_position_ids( + position_ids, attention_mask, packed_seq_params + ) + + return PackingState( + position_ids=position_ids, + cp_batch=cp_batch, + is_packing=is_packing, + ) + + def build_multimodal_embeddings( + self, + input_ids: torch.Tensor, + attention_mask: Optional[torch.Tensor], + packed_seq_params: Optional[PackedSeqParams], + packing_state: PackingState, + multimodal_kwargs: dict, + ) -> MultimodalEmbeddingResult: + """Phase 2: heavy multimodal embedding construction (pre_process only). + + Should only be called when ``self.pre_process`` is ``True`` **and** + ``decoder_input is None``. Computes text embeddings, replaces + media tokens with encoder outputs, and packs the result. + + Args: + input_ids: ``[b, s]`` + attention_mask: ``[b, s]`` (``None`` when not packing) + packed_seq_params: packing parameters + packing_state: the :class:`PackingState` returned by + :meth:`prepare_packing_state` + multimodal_kwargs: model-specific kwargs forwarded to + :meth:`construct_multimodal_inputs`. + + Returns: + :class:`MultimodalEmbeddingResult` + """ + is_packing = packing_state.is_packing + cp_batch = packing_state.cp_batch + if attention_mask is None: + attention_mask = cp_batch.get("attention_mask") + + # 1. Compute text embeddings + if is_packing: + inputs_embeds = self._compute_full_sequence_embedding(input_ids) + inputs_ranges = None + else: + inputs_ranges = self.get_input_ranges(input_ids.shape[1]) + inputs_embeds = self.embedding( + input_ids=cp_batch["input_ids"], position_ids=None + ) + + # 2. Model-specific multimodal token replacement + mm_result = self.construct_multimodal_inputs( + input_ids=input_ids, + inputs_embeds=inputs_embeds, + inputs_ranges=inputs_ranges, + **multimodal_kwargs, + ) + # Normalise to a tuple of (embeds, visual_pos_masks, deepstack) + if isinstance(mm_result, torch.Tensor): + inputs_embeds = mm_result + visual_pos_masks = None + deepstack_visual_embeds = None + else: + inputs_embeds, visual_pos_masks, deepstack_visual_embeds = mm_result + + # 3. Pack inputs_embeds (CP zigzag → SP scatter) + if is_packing: + inputs_embeds = self._pack_inputs_embeds( + inputs_embeds, attention_mask, packed_seq_params + ) + visual_pos_masks, deepstack_visual_embeds = ( + self._pack_deepstack_visual_embeds( + visual_pos_masks, + deepstack_visual_embeds, + attention_mask, + packed_seq_params, + ) + ) + + return MultimodalEmbeddingResult( + inputs_embeds=inputs_embeds, + position_ids=packing_state.position_ids, + cp_batch=cp_batch, + visual_pos_masks=visual_pos_masks, + deepstack_visual_embeds=deepstack_visual_embeds, + ) + + def construct_multimodal_inputs( + self, + input_ids: torch.Tensor, + inputs_embeds: torch.Tensor, + inputs_ranges: Optional[list], + **kwargs, + ): + """Hook for model-specific vision / audio embedding replacement. + + Must be overridden by each concrete model class. + + TODO: unified method for vl/omni model can be used actually + + Returns: + Either ``inputs_embeds`` (a single tensor), or a tuple of + ``(inputs_embeds, visual_pos_masks, deepstack_visual_embeds)`` + for models that use deepstack. + """ + raise NotImplementedError( + "Subclasses must implement construct_multimodal_inputs" + ) + + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + + @staticmethod + def _patch_rotary_pos_emb(): + """Patch Megatron's ``apply_rotary_pos_emb`` for THD packing. + + The standard ``_apply_rotary_pos_emb_thd`` divides seqlens by + ``cp_size`` and applies zigzag freq selection, which is wrong + when freqs are generated from packed, CP-sliced position_ids. + ``apply_rotary_pos_emb_absolute`` correctly handles both THD + and BSHD formats without CP-based freq slicing. + + We must patch both the source module (``rope_utils``) and the + consumer module (``attention``), because ``from ... import`` + creates independent name bindings in each module's namespace. + """ + from .qwen3_vl.rope_utils import apply_rotary_pos_emb_absolute + + import megatron.core.models.common.embeddings.rope_utils as _mcore_rope_utils + import megatron.core.transformer.attention as _mcore_attn + + # NOTE: currently, only patch method instead of attention module replacement + # which is different from mbridge. And we have to make sure split_qkv=True + # to use apply_rotary_pos_emb, which is true when packed_seq_params is not None + _mcore_rope_utils.apply_rotary_pos_emb = apply_rotary_pos_emb_absolute + _mcore_attn.apply_rotary_pos_emb = apply_rotary_pos_emb_absolute + + @staticmethod + def _pack_to_thd( + data: torch.Tensor, + attention_mask: torch.Tensor, + packed_seq_params: PackedSeqParams, + apply_cp_split: bool = True, + ) -> torch.Tensor: + """Pack data from ``[b, s, ...]`` to THD ``[packed_len, 1, ...]``. + + Follows the same packing logic as mbridge's ``preprocess_packed_seqs``: + for each sample, extract valid tokens (via attention_mask), then + optionally apply CP zigzag splitting on the **padded valid length** + (not the original sequence length). + + Args: + data: ``[b, s, ...]`` format data. The sequence dimension ``s`` + may already be SP-scattered but must NOT be CP-sliced yet. + attention_mask: ``[b, s]`` format attention mask (1=valid, + 0=padding), must match data's sequence dimension. + packed_seq_params: ``PackedSeqParams`` containing + ``cu_seqlens_q_padded`` (based on full, un-split padded + valid lengths). + apply_cp_split: If ``True`` (default), apply CP zigzag splitting + so the output length is ``total_padded_len // cp_size``. + If ``False``, output the full packed sequence without CP + splitting (useful for tensors like ``input_ids`` where + downstream consumers handle CP splitting themselves). + + Returns: + Packed data of shape ``[packed_len, 1, ...]`` where + ``packed_len`` is CP-local when *apply_cp_split* is True, + or full length otherwise. + """ + cu_seqlens_padded = packed_seq_params.cu_seqlens_q_padded + cp_size = mpu.get_context_parallel_world_size() if apply_cp_split else 1 + cp_rank = mpu.get_context_parallel_rank() if apply_cp_split else 0 + + batch_size = data.shape[0] + extra_dims = data.shape[2:] + + cu_seqlens_padded_cpu = cu_seqlens_padded.tolist() + local_packed_len = cu_seqlens_padded_cpu[-1] // cp_size + + output = torch.zeros( + [local_packed_len] + list(extra_dims), + dtype=data.dtype, + device=data.device, + ) + + for i in range(batch_size): + sample_mask = attention_mask[i].bool() + valid_tokens = data[i][sample_mask] + + padded_len_i = cu_seqlens_padded_cpu[i + 1] - cu_seqlens_padded_cpu[i] + start_offset = cu_seqlens_padded_cpu[i] // cp_size + + if cp_size <= 1: + num_valid = valid_tokens.shape[0] + output[start_offset : start_offset + num_valid] = valid_tokens + else: + chunk_len = padded_len_i // cp_size + half_chunk = chunk_len // 2 + + front_start = half_chunk * cp_rank + front_end = half_chunk * (cp_rank + 1) + front_len = min(front_end, valid_tokens.shape[0]) - front_start + if front_len > 0: + output[start_offset : start_offset + front_len] = ( + valid_tokens[front_start : front_start + front_len] + ) + + back_start = padded_len_i - half_chunk * (cp_rank + 1) + back_end = padded_len_i - half_chunk * cp_rank + back_end = min(back_end, valid_tokens.shape[0]) + back_len = back_end - back_start + if back_len > 0: + output[ + start_offset + half_chunk : start_offset + half_chunk + back_len + ] = valid_tokens[back_start:back_end] + + return output.unsqueeze(1) + + def _build_cp_batch( + self, + input_ids: torch.Tensor, + attention_mask: Optional[torch.Tensor], + is_packing: bool, + packed_seq_params: Optional[PackedSeqParams] = None, + ) -> dict: + if is_packing: + if self.mtp_process: + # Pack input_ids to [1, local_packed_len] with CP splitting. + # roll_tensor's _roll_tensor_packed_seq expects tensor length + # = cu_seqlens[-1] // cp_size (CP-local). + packed_ids = self._pack_to_thd( + input_ids.unsqueeze(-1), attention_mask, packed_seq_params, + ).squeeze(-1).squeeze(-1).unsqueeze(0) # [local_packed_len,1,1] → [1, local_packed_len] + else: + packed_ids = input_ids + return {"input_ids": packed_ids, "attention_mask": attention_mask} + + cp_batch = {"input_ids": input_ids, "attention_mask": attention_mask} + if self.config.context_parallel_size > 1: + cp_batch = { + k: v.clone() if v is not None else None + for k, v in cp_batch.items() + } + # Call McaGPTModel's base implementation directly. + # VLM models override get_batch_on_this_cp_rank to skip + # CP-slicing input_ids, but here we need the base version + # that slices everything. We cannot use super() because + # MultimodalEmbeddingMixin sits at the end of the MRO + # (before object), so super() would resolve to object. + cp_batch = McaGPTModel.get_batch_on_this_cp_rank( + self, + cp_batch, + dim3_keys=[] + if (cp_batch["attention_mask"] is None or cp_batch["attention_mask"].dim() == 2) + else ["attention_mask"], + ) + return cp_batch + + def _pack_position_ids( + self, + position_ids: torch.Tensor, + attention_mask: torch.Tensor, + packed_seq_params: PackedSeqParams, + ) -> torch.Tensor: + pos_ids_bsc = position_ids.permute(1, 2, 0) # [C, b, s] → [b, s, C] + pos_ids_packed = self._pack_to_thd( + pos_ids_bsc, attention_mask, packed_seq_params + ) + self.rotary_pos_emb.is_thd_format = True + return pos_ids_packed.permute(2, 1, 0) # → [C, 1, packed_len] + + def _compute_full_sequence_embedding( + self, input_ids: torch.Tensor + ) -> torch.Tensor: + """Compute embeddings on full sequence (SP scatter temporarily off).""" + orig_scatter = self.embedding.scatter_to_sequence_parallel + orig_reduce = self.embedding.reduce_scatter_embeddings + orig_word_reduce = None + if hasattr(self.embedding, "word_embeddings"): + orig_word_reduce = ( + self.embedding.word_embeddings.reduce_scatter_embeddings + ) + + self.embedding.scatter_to_sequence_parallel = False + self.embedding.reduce_scatter_embeddings = False + if hasattr(self.embedding, "word_embeddings"): + self.embedding.word_embeddings.reduce_scatter_embeddings = False + + inputs_embeds = self.embedding(input_ids=input_ids, position_ids=None) + + self.embedding.scatter_to_sequence_parallel = orig_scatter + self.embedding.reduce_scatter_embeddings = orig_reduce + if hasattr(self.embedding, "word_embeddings"): + self.embedding.word_embeddings.reduce_scatter_embeddings = ( + orig_word_reduce + ) + return inputs_embeds + + def _pack_inputs_embeds( + self, + inputs_embeds: torch.Tensor, + attention_mask: torch.Tensor, + packed_seq_params: PackedSeqParams, + ) -> torch.Tensor: + """Pack ``[s, b, h]`` → THD, then SP scatter.""" + inputs_embeds_bsh = inputs_embeds.transpose(0, 1) + inputs_embeds = self._pack_to_thd( + inputs_embeds_bsh, attention_mask, packed_seq_params + ) + if self.config.sequence_parallel: + inputs_embeds = ( + tensor_parallel.scatter_to_sequence_parallel_region( + inputs_embeds + ) + ) + return inputs_embeds + + def _pack_deepstack_visual_embeds( + self, + visual_pos_masks: Optional[torch.Tensor], + deepstack_visual_embeds: Optional[list], + attention_mask: torch.Tensor, + packed_seq_params: PackedSeqParams, + ) -> tuple[Optional[torch.Tensor], Optional[list]]: + """Pack visual_pos_masks + re-index deepstack for CP+SP.""" + if visual_pos_masks is None: + return None, None + + packed_masks = ( + self._pack_to_thd( + visual_pos_masks.unsqueeze(-1).float(), + attention_mask, + packed_seq_params, + ) + .squeeze(-1) + .squeeze(-1) + .bool() + ) + + tp_size = mpu.get_tensor_model_parallel_world_size() + if self.config.sequence_parallel and tp_size > 1: + tp_rank = mpu.get_tensor_model_parallel_rank() + packed_masks = packed_masks.chunk(tp_size)[tp_rank] + + if deepstack_visual_embeds is not None: + idx_full = torch.full_like( + visual_pos_masks, -1, dtype=torch.long + ) + idx_full[visual_pos_masks] = torch.arange( + visual_pos_masks.sum(), device=visual_pos_masks.device + ) + idx_packed = ( + self._pack_to_thd( + idx_full.unsqueeze(-1).float(), + attention_mask, + packed_seq_params, + ) + .squeeze(-1) + .squeeze(-1) + .long() + ) + if self.config.sequence_parallel and tp_size > 1: + idx_packed = idx_packed.chunk(tp_size)[tp_rank] + + final_idx = idx_packed[packed_masks] + deepstack_visual_embeds = [ + embed[final_idx] for embed in deepstack_visual_embeds + ] + + # [sp_packed_len] → [1, sp_packed_len] + visual_pos_masks = packed_masks.unsqueeze(0) + return visual_pos_masks, deepstack_visual_embeds diff --git a/mcore_adapter/src/mcore_adapter/parallel_functions/fused_cross_entropy.py b/mcore_adapter/src/mcore_adapter/parallel_functions/fused_cross_entropy.py new file mode 100644 index 000000000..382bec04d --- /dev/null +++ b/mcore_adapter/src/mcore_adapter/parallel_functions/fused_cross_entropy.py @@ -0,0 +1,618 @@ +# Copyright (c) 2025, Wentao Guo, Ted Zadouri, Tri Dao. +# This file was copied from https://github.com/Dao-AILab/quack/blob/12f11462f06f8c1d79ac2c3c04e308678c81253c/quack/cross_entropy.py, +# with a minor bug fix applied (see line 156). https://github.com/Dao-AILab/quack/pull/64 + +import math +from functools import partial +from typing import Literal, Optional, Type + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import quack.copy_utils as copy_utils +import quack.layout_utils as layout_utils +import quack.utils as utils +import torch +from cutlass import BFloat16, Boolean, Float32, Int32, Int64, const_expr +from quack.compile_utils import make_fake_tensor as fake_tensor +from quack.cute_dsl_utils import torch2cute_dtype_map +from quack.reduce import online_softmax_reduce, row_reduce +from quack.reduction_base import ReductionBase +from torch import Tensor + + +class CrossEntropy(ReductionBase): + def __init__(self, dtype: Type[cutlass.Numeric], N: int, online_softmax: bool = True): + # 2 stages: 1 for max, 1 for sum + super().__init__( + dtype, + N, + stage=2 if not online_softmax else 1, + reduction_dtype=Float32 if not online_softmax else Int64, + ) + self.online_softmax = online_softmax + self.reload_from = None if N <= 16384 or online_softmax else "smem" + + def _threads_per_row(self): + N = self.N + for limit, threads in [(64, 8), (128, 16), (3072, 32), (6144, 64), (16384, 128)]: + if N <= limit: + return threads + return 256 + + def _set_cluster_n(self): + N = self.N + if const_expr(self.dtype.width == 16): + thresholds = [(16 * 1024, 1), (32 * 1024, 2), (64 * 1024, 4), (128 * 1024, 8)] + else: + thresholds = [(16 * 1024, 1), (64 * 1024, 2), (128 * 1024, 4), (256 * 1024, 8)] + for limit, cluster in thresholds: + if N <= limit: + self.cluster_n = cluster + return + self.cluster_n = 16 + + @cute.jit + def __call__( + self, + mX: cute.Tensor, # (M, N) + mTarget: cute.Tensor, # (M,) + mTargetLogit: Optional[cute.Tensor], # (M, K) or (M,). If None, we use mX + mLoss: cute.Tensor, # (M,) + mLSE: Optional[cute.Tensor], # (M,) + mdX: Optional[cute.Tensor], # (M, N) - if provided, compute gradient + ignore_index: Int32, # Index to ignore in loss computation + stream: cuda.CUstream, + ): + assert mX.element_type == self.dtype + if const_expr(mTargetLogit is None): + mTargetLogit = mX + if const_expr(mdX is not None): + assert mdX.element_type == self.dtype + self._set_cluster_n() + largest_dtype_width = const_expr(mX.element_type.width) + if const_expr(mdX is not None): + largest_dtype_width = const_expr(max(largest_dtype_width, mdX.element_type.width)) + vecsize = math.gcd(self.N, 128 // largest_dtype_width) + tiled_copy, tiler_mn, threads_per_row = self._get_tiled_copy(vecsize=vecsize) + num_threads = tiled_copy.size + self.kernel( + mX, + mTarget, + mTargetLogit, + mLoss, + mLSE, + mdX, + ignore_index, + tiler_mn, + tiled_copy, + threads_per_row, + ).launch( + grid=[cute.ceil_div(mX.shape[0], tiler_mn[0]), self.cluster_n, 1], + block=[num_threads, 1, 1], + cluster=[1, self.cluster_n, 1] if const_expr(self.cluster_n > 1) else None, + stream=stream, + ) + + @cute.kernel + def kernel( + self, + mX: cute.Tensor, # (M, N) + mTarget: cute.Tensor, # (M,) + mTargetLogit: cute.Tensor, # (M, K) or (M,) + mLoss: cute.Tensor, # (M,) + mLSE: Optional[cute.Tensor], # (M,) + mdX: Optional[cute.Tensor], # (M, N) - if provided, compute gradient + ignore_index: Int32, # Index to ignore in loss computation + tiler_mn: cute.Shape, + tiled_copy: cute.TiledCopy, + threads_per_row: cutlass.Constexpr[int], + ): + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + cluster_y = const_expr(0) if const_expr(self.cluster_n == 1) else cute.arch.block_idx()[1] + tv_layout = tiled_copy.layout_tv_tiled + + shape = mX.shape + idX = cute.make_identity_tensor(shape) + # slice for CTAs + gX, cX = [cute.local_tile(mT, tiler_mn, (bidx, cluster_y)) for mT in (mX, idX)] + + smem = cutlass.utils.SmemAllocator() + sX = smem.allocate_tensor(mX.element_type, cute.make_ordered_layout(tiler_mn, order=(1, 0)), byte_alignment=16) + reduction_buffer, mbar_ptr = self._allocate_reduction_buffer_and_mbar(smem, tv_layout) + + thr_copy = tiled_copy.get_slice(tidx) + + tXgX = thr_copy.partition_S(gX) + tXsX = thr_copy.partition_D(sX) + tXcX = thr_copy.partition_S(cX)[(0, None), None, None] + tXrX = cute.make_fragment_like(tXgX) + + is_even_N = const_expr(shape[1] == tiler_mn[1] * self.cluster_n) + tXpX = None if is_even_N else copy_utils.predicate_k(thr_copy.partition_S(cX), limit=shape[1]) + copy = partial(copy_utils.copy, pred=tXpX) + + num_warps = cute.size(tiled_copy) // cute.arch.WARP_SIZE + self._initialize_cluster(tidx, mbar_ptr, num_warps) + + row = tXcX[0][0] + target = Int32.zero + if row < shape[0]: + target = Int32(mTarget[row]) + + if row < shape[0]: + copy(tXgX, tXsX, is_async=True) + cute.arch.cp_async_commit_group() + cute.arch.cp_async_wait_group(0) + # Fill OOB values with -inf + if const_expr(not is_even_N): + # utils.fill_oob(tXsX, tXpX, -tXsX.element_type.inf) + utils.fill_oob(tXsX, tXpX, BFloat16(-(2**15))) + cute.autovec_copy(tXsX, tXrX) + x = tXrX.load().to(Float32) + + target_logit = Float32.zero + should_ignore = Boolean(target == ignore_index) + if row < shape[0] and tXcX[0][1] == 0 and not should_ignore: + # Only load target logit if not ignoring this index + if const_expr(cute.rank(mTargetLogit.shape) == 2): + target_logit = Float32(mTargetLogit[row, target]) + else: + assert cute.rank(mTargetLogit.shape) == 1 + target_logit = Float32(mTargetLogit[row]) + + if const_expr(not self.online_softmax): + max_x = row_reduce( + x, + cute.ReductionOp.MAX, + threads_per_row, + reduction_buffer[None, None, 0], + mbar_ptr + 0 if const_expr(self.cluster_n > 1) else None, + init_val=-Float32.inf, + hook_fn=cute.arch.cluster_wait if const_expr(self.cluster_n > 1) else None, + ) + if const_expr(self.reload_from == "smem"): + cute.autovec_copy(tXsX, tXrX) + x = tXrX.load().to(Float32) + log2_e = math.log2(math.e) + # This would use ffma instead of fadd then fmul + exp_x = cute.math.exp2(x * log2_e - (max_x * log2_e), fastmath=False) + denom = row_reduce( + exp_x, + cute.ReductionOp.ADD, + threads_per_row, + reduction_buffer[None, None, 1], + mbar_ptr + 1 if const_expr(self.cluster_n > 1) else None, + init_val=0.0, + ) + else: + max_x, denom, exp_x = online_softmax_reduce( + x, + threads_per_row, + reduction_buffer[None, None, 0], + mbar_ptr, + hook_fn=cute.arch.cluster_wait if const_expr(self.cluster_n > 1) else None, + return_exp_x=const_expr(mdX is not None), + ) + + # Write loss and lse to gmem + if tXcX[0][1] == 0 and row < shape[0] and (self.cluster_n == 1 or cute.arch.block_idx_in_cluster() == 0): + lse = max_x + cute.math.log(denom, fastmath=True) + # Set loss to 0 if this index should be ignored, otherwise compute normally + loss_val = (lse - target_logit) if not should_ignore else Float32.zero + mLoss[row] = mLoss.element_type(loss_val) + if const_expr(mLSE is not None): + mLSE[row] = lse + + # Compute gradient if mdX is provided + if const_expr(mdX is not None): + # Compute probabilities: exp(x) / sum(exp(x)) + # If ignored, gradient should be zero + denom_inv = ( + # 1.0 / denom + cute.arch.rcp_approx(denom) if not (denom == 0.0 or denom != denom or should_ignore) else Float32.zero + ) + probs = exp_x * denom_inv + gdX = cute.local_tile(mdX, tiler_mn, (bidx, cluster_y)) + tXgdX = thr_copy.partition_D(gdX) + tXrdX = cute.make_fragment_like(tXgdX) + tXcFull = thr_copy.partition_S(cX) + # Compute gradient: probs for all classes, (probs - 1) for target class + # If ignored, gradient is already zero + tXrdX_f32 = cute.make_fragment_like(tXrX, Float32) + tXrdX_f32.store(probs) + if not should_ignore: + for i in cutlass.range(cute.size(tXrX), unroll_full=True): + tXrdX_f32[i] = tXrdX_f32[i] if tXcFull[i][1] != target else tXrdX_f32[i] - 1.0 + tXrdX.store(tXrdX_f32.load().to(tXrdX.element_type)) + if row < shape[0]: + copy(tXrdX, tXgdX) + + +@torch.library.custom_op("quack::cross_entropy_fwd_out", mutates_args={"loss", "lse", "dx"}) +def cross_entropy_fwd_out( + x: Tensor, + target: Tensor, + target_logit: Optional[Tensor], + loss: Tensor, + lse: Optional[Tensor], + dx: Optional[Tensor], + ignore_index: int = -100, +) -> None: + """Cross entropy forward pass. + + Args: + x: Input logits tensor of shape (M, N) + target: Target class indices tensor of shape (M,) + target_logit: (M, K) or (M,). + If provided, the target logit will be read from this tensor instead of x. + loss: Output loss tensor of shape (M,) + lse: Optional output log-sum-exp tensor of shape (M,) + dx: Optional output gradient tensor of shape (M, N) + ignore_index: Index to ignore in loss computation + + Returns: + None (mutates loss, lse, and optionally dx in-place) + """ + assert x.dim() == 2, "Input must be 2D" + assert target.dim() == 1, "Target must be 1D" + assert x.is_cuda and target.is_cuda, "Tensors must be on CUDA device" + assert x.dtype in [torch.float16, torch.bfloat16, torch.float32], "Unsupported input dtype" + assert target.dtype in [torch.int32, torch.int64], "Target must be int32 or int64" + if target_logit is not None: + assert target_logit.is_cuda, "Target logits must be on CUDA device" + assert target_logit.dtype in [torch.float16, torch.bfloat16, torch.float32] + if dx is not None: + assert dx.is_cuda, "dx must be on CUDA device" + N = x.size(1) + dtype = torch2cute_dtype_map[x.dtype] + target_dtype = torch2cute_dtype_map[target.dtype] + target_logit_dtype = torch2cute_dtype_map[target_logit.dtype] if target_logit is not None else None + compile_key = ( + dtype, + target_dtype, + target_logit_dtype, + N, + lse is not None, + dx is not None, + ) + if compile_key not in cross_entropy_fwd_out.compile_cache: + batch_sym = cute.sym_int() + div = math.gcd(128 // dtype.width, N) + x_cute = fake_tensor(dtype, (batch_sym, N), div) + dx_cute = fake_tensor(dtype, (batch_sym, N), div) if dx is not None else None + target_cute = fake_tensor(target_dtype, (batch_sym,)) + if target_logit is not None: + if target_logit.ndim == 2: + target_logit_cute = fake_tensor(target_logit_dtype, (batch_sym, cute.sym_int()), div) + else: + target_logit_cute = fake_tensor(target_logit_dtype, (batch_sym,)) + else: + target_logit_cute = None + loss_cute = fake_tensor(Float32, (batch_sym,)) + lse_cute = fake_tensor(Float32, (batch_sym,)) if lse is not None else None + # If there's dx, it's faster to not use online softmax since we want the exp(x - max) + cross_entropy_op = CrossEntropy(dtype, N, online_softmax=dx is None) + cross_entropy_fwd_out.compile_cache[compile_key] = cute.compile( + cross_entropy_op, + x_cute, + target_cute, + target_logit_cute, + loss_cute, + lse_cute, + dx_cute, + Int32(0), # ignore_index, just for compilation + cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), + options="--enable-tvm-ffi", + ) + cross_entropy_fwd_out.compile_cache[compile_key](x, target, target_logit, loss, lse, dx, Int32(ignore_index)) + + +cross_entropy_fwd_out.compile_cache = {} + + +def cross_entropy_fwd( + x: torch.Tensor, + target: torch.Tensor, + target_logit: Optional[torch.Tensor] = None, + ignore_index: int = -100, + return_lse: bool = False, + return_dx: bool = False, + inplace_backward: bool = False, +) -> torch.Tensor | tuple[torch.Tensor]: + M = x.size(0) + device = x.device + loss = torch.empty(M, device=device, dtype=torch.float32) + lse = torch.empty(M, device=device, dtype=torch.float32) if return_lse else None + dx = (torch.empty_like(x) if not inplace_backward else x) if return_dx else None + cross_entropy_fwd_out(x, target, target_logit, loss, lse, dx, ignore_index) + if return_lse and return_dx: + return loss, lse, dx + elif return_lse: + return loss, lse + elif return_dx: + return loss, dx + else: + return loss + + +class CrossEntropyBackward: + def __init__(self, dtype: Type[cutlass.Numeric], N: int): + self.dtype = dtype + self.N = N + self.vecsize = 128 // dtype.width + + def _threads_per_row(self): + N = min(self.N, 16384) # We split by blocks of 16k + for limit, threads in [(64, 8), (128, 16), (3072, 32), (6144, 64), (16384, 128)]: + if N <= limit: + return threads + return 256 + + def _get_tiled_copy(self, vecsize: int): + assert self.N % vecsize == 0, f"Input N {self.N} is not divisible by vector size {vecsize}" + N = min(self.N, 16384) + num_threads = 128 if N <= 16384 else 256 + threads_per_row = self._threads_per_row() + cols_per_block = num_threads // threads_per_row + num_blocks_N = cute.ceil_div(N // vecsize, threads_per_row) + tiler_mn = (cols_per_block, vecsize * num_blocks_N * threads_per_row) + tiled_copy = copy_utils.tiled_copy_2d(self.dtype, threads_per_row, num_threads, num_copy_elems=vecsize) + return tiled_copy, tiler_mn, threads_per_row + + @cute.jit + def __call__( + self, + mX: cute.Tensor, + mTarget: cute.Tensor, + mDLoss: cute.Tensor, + mdX: cute.Tensor, + mLSE: cute.Tensor, + ignore_index: Int32, # Index to ignore in gradient computation + stream: cuda.CUstream, + ): + assert mX.element_type == self.dtype + assert mdX.element_type == self.dtype + # e.g. if self.N isn't divisible by 8 for bf16, we might use 64 bits (4 elements) copy + vecsize = math.gcd(self.N, 128 // self.dtype.width) + tiled_copy, tiler_mn, threads_per_row = self._get_tiled_copy(vecsize=vecsize) + num_threads = tiled_copy.size + # (M,) -> (M, N) with stride 0 in the N dimension + mDLoss, mTarget, mLSE = [layout_utils.expand(X, dim=1, size=self.N) for X in (mDLoss, mTarget, mLSE)] + self.kernel( + mX, + mTarget, + mDLoss, + mdX, + mLSE, + ignore_index, + mX.shape, + tiler_mn, + tiled_copy, + threads_per_row, + ).launch( + grid=[ + cute.ceil_div(mX.shape[0], tiler_mn[0]), + cute.ceil_div(mX.shape[1], tiler_mn[1]), + 1, + ], + block=[num_threads, 1, 1], + stream=stream, + ) + + @cute.kernel + def kernel( + self, + mX: cute.Tensor, # (M, N) + mTarget: cute.Tensor, # (M,) + mDLoss: cute.Tensor, # (M,) + mdX: cute.Tensor, # (M, N) + mLSE: cute.Tensor, # (M,) + ignore_index: Int32, # Index to ignore in gradient computation + shape: cute.Shape, + tiler_mn: cute.Shape, + tiled_copy: cute.TiledCopy, + threads_per_row: cutlass.Constexpr[int], + ): + tidx, _, _ = cute.arch.thread_idx() + bidx, bidy, _ = cute.arch.block_idx() + + smem = cutlass.utils.SmemAllocator() + sX = smem.allocate_tensor(mX.element_type, cute.make_ordered_layout(tiler_mn, order=(1, 0)), byte_alignment=16) + + idX = cute.make_identity_tensor(shape) + gX, gdX, cX = [cute.local_tile(mT, tiler_mn, (bidx, bidy)) for mT in (mX, mdX, idX)] + + thr_copy = tiled_copy.get_slice(tidx) + + tXgX = thr_copy.partition_S(gX) + tXsX = thr_copy.partition_D(sX) + tXcX = thr_copy.partition_S(cX)[(0, None), None, None] + tXcFull = thr_copy.partition_S(cX) + tXgdX = thr_copy.partition_D(gdX) + tXrX, tXrdX = [cute.make_fragment_like(thr) for thr in (tXgX, tXgdX)] + + is_even_N = const_expr(shape[1] % tiler_mn[1] == 0) + tXpX = None if is_even_N else copy_utils.predicate_k(thr_copy.partition_S(cX), limit=shape[1]) + copy = partial(copy_utils.copy, pred=tXpX) + + row = tXcX[0][0] + if row < shape[0]: + copy(tXgX, tXsX, is_async=True) + cute.arch.cp_async_commit_group() + cute.arch.cp_async_wait_group(0) + if const_expr(not is_even_N): + utils.fill_oob(tXsX, tXpX, -tXsX.element_type.inf) + cute.autovec_copy(tXsX, tXrX) + x = tXrX.load().to(Float32) + + target = Int32.zero + dloss = Float32.zero + lse = Float32.zero + if row < shape[0]: + target = Int32(mTarget[row]) + should_ignore = Boolean(target == ignore_index) + # Set dloss to 0 if this index should be ignored + if not should_ignore: + dloss = Float32(mDLoss[row]) + lse = Float32(mLSE[row]) + + log2_e = math.log2(math.e) + probs = cute.math.exp2(x * log2_e - (lse * log2_e), fastmath=True) + prob_shifted = probs - 1.0 + mask = cute.make_fragment_like(tXrX, Boolean) + for i in cutlass.range(cute.size(tXcFull), unroll_full=True): + mask[i] = tXcFull[i][1] == target + grad = cute.where(mask.load(), prob_shifted, probs) + grad = grad * dloss + + tXrdX.store(grad.to(tXrdX.element_type)) + if row < shape[0]: + copy(tXrdX, tXgdX) + + +def _cross_entropy_backward( + x: torch.Tensor, + target: torch.Tensor, + dloss: torch.Tensor, + lse: torch.Tensor, + dx: torch.Tensor, + ignore_index=-100, +) -> None: + """Cross entropy backward pass. + Args: + x: Input logits tensor of shape (M, N) + target: Target class indices tensor of shape (M,) + dloss: Upstream gradients tensor of shape (M,) + lse: Log-sum-exp values tensor of shape (M,) + Returns: + Input gradients tensor of shape (M, N) + """ + assert x.dim() == 2, "Input must be 2D" + assert target.dim() == 1, "Target must be 1D" + assert dloss.dim() == 1, "dloss must be 1D" + assert lse.dim() == 1, "lse must be 1D" + assert x.shape[0] == target.shape[0], "Batch dimensions must match" + assert x.shape[0] == dloss.shape[0], "Batch dimensions must match" + assert x.shape[0] == lse.shape[0], "Batch dimensions must match" + assert x.is_cuda and target.is_cuda and dloss.is_cuda and lse.is_cuda, "Tensors must be on CUDA device" + assert x.dtype in [torch.float16, torch.bfloat16, torch.float32], "Unsupported input dtype" + assert target.dtype in [torch.int32, torch.int64], "Target must be int32 or int64" + + N = x.size(1) + dtype = torch2cute_dtype_map[x.dtype] + target_dtype = torch2cute_dtype_map[target.dtype] + compile_key = (dtype, target_dtype, N) + if compile_key not in _cross_entropy_backward.compile_cache: + batch_sym = cute.sym_int() + div = math.gcd(128 // dtype.width, N) + x_cute, dx_cute = [fake_tensor(dtype, (batch_sym, N), div)] * 2 + target_cute = fake_tensor(target_dtype, (batch_sym,)) + dloss_cute, lse_cute = [fake_tensor(Float32, (batch_sym,))] * 2 + cross_entropy_backward_op = CrossEntropyBackward(dtype, N) + _cross_entropy_backward.compile_cache[compile_key] = cute.compile( + cross_entropy_backward_op, + x_cute, + target_cute, + dloss_cute, + dx_cute, + lse_cute, + Int32(0), # ignore_index, just for compilation + cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), + options="--enable-tvm-ffi", + ) + _cross_entropy_backward.compile_cache[compile_key](x, target, dloss, dx, lse, Int32(ignore_index)) + + +_cross_entropy_backward.compile_cache = {} + + +@torch.library.custom_op("quack::cross_entropy_bwd_out", mutates_args={"dx"}) +def cross_entropy_bwd_out( + x: torch.Tensor, + target: torch.Tensor, + dloss: torch.Tensor, + lse: torch.Tensor, + dx: torch.Tensor, + ignore_index: int = -100, +) -> None: + _cross_entropy_backward(x, target, dloss, lse, dx, ignore_index) + + +def cross_entropy_bwd( + x: torch.Tensor, + target: torch.Tensor, + dloss: torch.Tensor, + lse: torch.Tensor, + ignore_index: int = -100, + inplace_backward: bool = False, +) -> None: + if inplace_backward and not torch.compiler.is_compiling(): + dx = x + _cross_entropy_backward(x=x, target=target, dloss=dloss, lse=lse, dx=x, ignore_index=ignore_index) + else: + dx = torch.empty_like(x) + cross_entropy_bwd_out(x=x, target=target, dloss=dloss, lse=lse, dx=dx, ignore_index=ignore_index) + return dx + + +class CrossEntropyFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, x, target, lse_partial=None, ignore_index=-100, inplace_backward=False): + if lse_partial is None: + loss, lse = cross_entropy_fwd(x, target, ignore_index=ignore_index, return_lse=True) + else: + # if we already compute partial lse, then to compute the final lse we treat + # @lse_partial as @x and @x as @target_logit + loss, lse = cross_entropy_fwd( + lse_partial, target, target_logit=x, ignore_index=ignore_index, return_lse=True + ) + ctx.save_for_backward(x, target, lse) + ctx.ignore_index = ignore_index + ctx.inplace_backward = inplace_backward + return loss + + @staticmethod + def backward(ctx, dloss): + x, target, lse = ctx.saved_tensors + dx = cross_entropy_bwd(x, target, dloss, lse, ctx.ignore_index, inplace_backward=ctx.inplace_backward) + return dx, None, None, None, None + + +def cross_entropy( + x: torch.Tensor, + target: torch.Tensor, + lse_partial: Optional[torch.Tensor] = None, + ignore_index: int = -100, + reduction: Literal["none", "mean", "sum"] = "mean", + inplace_backward: bool = False, +) -> torch.Tensor: + """Cross entropy loss with automatic differentiation support. + + Args: + x: Input logits tensor of shape (M, N) + target: Target class indices tensor of shape (M,) + lse_partial: Optional precomputed log-sum-exp partial results + reduction: Specifies the reduction to apply to the output: + 'none': no reduction will be applied (default) + 'mean': the sum of the output will be divided by the number of elements + 'sum': the output will be summed + inplace_backward: Whether to perform backward pass in-place + ignore_index: Index to ignore in loss computation (loss will be 0 for these indices) + + Returns: + Cross entropy loss tensor: + - If reduction='none': tensor of shape (M,) with per-example losses + - If reduction='mean': scalar tensor with mean loss + - If reduction='sum': scalar tensor with sum of losses + """ + loss = CrossEntropyFunction.apply(x, target, lse_partial, ignore_index, inplace_backward) + if reduction == "mean": + return loss.sum() / (target != ignore_index).sum().float() + elif reduction == "sum": + return loss.sum() + elif reduction == "none": + return loss + else: + raise ValueError(f"Invalid reduction mode: {reduction}. Expected one of 'none', 'mean', or 'sum'") diff --git a/mcore_adapter/src/mcore_adapter/parallel_functions/fused_logprobs.py b/mcore_adapter/src/mcore_adapter/parallel_functions/fused_logprobs.py new file mode 100644 index 000000000..3814ae317 --- /dev/null +++ b/mcore_adapter/src/mcore_adapter/parallel_functions/fused_logprobs.py @@ -0,0 +1,686 @@ +""" +Fused Entropy Kernel Implementation + +This module provides CUDA fused kernel implementations for entropy computation, +optimized for tensor parallel size 1 (TP=1). It accelerates the entropy calculation +used in tensor_parallel.py by fusing multiple operations into a single kernel pass. + +The entropy forward computation is split into two stages: +- Stage 1: Computes max_x and sum_exp(x - max_x) for each row +- Stage 2: Computes the final entropy value: H = max_x + log(sum_exp) - sum(softmax(x) * x) + +Note: Currently only supports TP=1 (single GPU tensor parallelism). +""" + +import math +import operator +from functools import partial +from typing import Optional, Tuple, Type, Literal, Callable + +import torch +from torch import Tensor + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +from cutlass import Int32, Int64, Float32, Boolean, const_expr, BFloat16 + +import quack.utils as utils +import quack.copy_utils as copy_utils +import quack.layout_utils as layout_utils +from quack.compile_utils import make_fake_tensor as fake_tensor +from quack.reduce import row_reduce, online_softmax_reduce +from quack.reduction_base import ReductionBase +from quack.cute_dsl_utils import torch2cute_dtype_map + + +class FusedLogProbsForward(ReductionBase): + """ + The fused logprobs forward pass. + + This kernel computes the intermediate values needed for entropy computation: + - max_x: Maximum value in each row (for numerical stability) + - sum_exp: Sum of exp(x - max_x) for each row (softmax denominator) + + Args: + dtype: Data type of input tensor (e.g., BFloat16, Float16, Float32) + N: Number of columns in input tensor (vocabulary size) + online_softmax: If True, uses online softmax algorithm for better performance + """ + + def __init__(self, dtype: Type[cutlass.Numeric], N: int, online_softmax: bool = True): + """ + Initialize the first stage of entropy forward pass. + """ + # 2 stages: 1 for max, 1 for sum when not using online softmax + # 1 stage when using online softmax (computes both max and sum together) + super().__init__( + dtype, + N, + stage=2 if not online_softmax else 1, + reduction_dtype=Float32 if not online_softmax else Int64, + ) + self.online_softmax = online_softmax + # For large N (>16384), reload data from shared memory to avoid register spilling + self.reload_from = None if N <= 16384 or online_softmax else "smem" + + def _threads_per_row(self): + """ + Determine the optimal number of threads per row for reduction. + """ + N = self.N + for limit, threads in [(64, 8), (128, 16), (3072, 32), (6144, 64), (16384, 128)]: + if N <= limit: + return threads + return 256 + + def _set_cluster_n(self): + """ + Determine the optimal cluster size in the N dimension. + """ + N = self.N + if const_expr(self.dtype.width == 16): + thresholds = [(16 * 1024, 1), (32 * 1024, 2), (64 * 1024, 4), (128 * 1024, 8)] + else: + thresholds = [(16 * 1024, 1), (64 * 1024, 2), (128 * 1024, 4), (256 * 1024, 8)] + for limit, cluster in thresholds: + if N <= limit: + self.cluster_n = cluster + return + self.cluster_n = 16 + + @cute.jit + def __call__( + self, + mX: cute.Tensor, # (M, N) in + mTarget: cute.Tensor, # (M) in + mLogProbs: cute.Tensor, # (M,) out + mXMax: cute.Tensor, # (M,) out + mXSumExp: cute.Tensor, # (M,) out + stream: cuda.CUstream, + ): + """ + Launch the kernel to compute logprobs, max and sum_exp for each row. + """ + assert mX.element_type == self.dtype + self._set_cluster_n() + # Calculate vector size for memory loading (max 128 bits per load) + largest_dtype_width = const_expr(mX.element_type.width) + vecsize = math.gcd(self.N, 128 // largest_dtype_width) # 8 for bf16/fp16 + tiled_copy, tiler_mn, threads_per_row = self._get_tiled_copy(vecsize=vecsize) + num_threads = tiled_copy.size + # Launch kernel with appropriate grid, block, and cluster configuration + self.kernel( + mX, + mTarget, + mLogProbs, + mXMax, + mXSumExp, + tiler_mn, + tiled_copy, + threads_per_row, + ).launch( + grid=[cute.ceil_div(mX.shape[0], tiler_mn[0]), self.cluster_n, 1], + block=[num_threads, 1, 1], + cluster=[1, self.cluster_n, 1] if const_expr(self.cluster_n > 1) else None, + stream=stream, + ) + + @cute.kernel + def kernel( + self, + mX: cute.Tensor, # (M, N) in + mTarget: cute.Tensor, # (M, N) in + mLogProbs: cute.Tensor, # (M,) out + mXMax: cute.Tensor, # (M,) out + mXSumExp: cute.Tensor, # (M,) out + tiler_mn: cute.Shape, + tiled_copy: cute.TiledCopy, + threads_per_row: cutlass.Constexpr[int], + ): + """ + CUDA kernel: Compute logprobs, max and sum_exp for each row. + + This kernel loads input data, performs reduction operations to compute + the maximum value and sum of exponentials for each row, and writes + the results to global memory. + + Args: + mX: Input logits tensor of shape (M, N) + mTarget: Input targe tensor of shape (M, ) + mLogProbs: Outpu tensor logprobs tensor of shape (M, ) + mXMax: Output tensor for max values of shape (M,) + mXSumExp: Output tensor for sum of exponentials of shape (M,) + tiler_mn: Tiling configuration + tiled_copy: Tiled copy configuration for efficient memory access + threads_per_row: Number of threads working on each row + """ + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + cluster_y = const_expr(0) if const_expr(self.cluster_n == 1) else cute.arch.block_idx()[1] + tv_layout = tiled_copy.layout_tv_tiled + + shape = mX.shape + idX = cute.make_identity_tensor(shape) + gX, cX = [cute.local_tile(mT, tiler_mn, (bidx, cluster_y)) for mT in (mX, idX)] + + smem = cutlass.utils.SmemAllocator() + sX = smem.allocate_tensor( + mX.element_type, cute.make_ordered_layout(tiler_mn, order=(1, 0)), byte_alignment=16 + ) + reduction_buffer, mbar_ptr = self._allocate_reduction_buffer_and_mbar(smem, tv_layout) + + thr_copy = tiled_copy.get_slice(tidx) + + # Partition tensors for this thread + tXgX = thr_copy.partition_S(gX) # Global memory partition + tXsX = thr_copy.partition_D(sX) # Shared memory partition + tXcX = thr_copy.partition_S(cX)[(0, None), None, None] # Coordinate partition + tXrX = cute.make_fragment_like(tXgX) # Register fragment + + # Handle non-divisible N by creating predicates + is_even_N = const_expr(shape[1] == tiler_mn[1] * self.cluster_n) + tXpX = None if is_even_N else copy_utils.predicate_k(thr_copy.partition_S(cX), limit=shape[1]) + copy = partial(copy_utils.copy, pred=tXpX) + + # Initialize cluster synchronization + num_warps = cute.size(tiled_copy) // cute.arch.WARP_SIZE + self._initialize_cluster(tidx, mbar_ptr, num_warps) + + row = tXcX[0][0] + + # Load input data asynchronously from global to shared memory + if row < shape[0]: + copy(tXgX, tXsX, is_async=True) + cute.arch.cp_async_commit_group() + cute.arch.cp_async_wait_group(0) + # Fill out-of-bounds values with -inf (for numerical stability in max reduction) + if const_expr(not is_even_N): + # utils.fill_oob(tXsX, tXpX, -tXsX.element_type.inf) + utils.fill_oob(tXsX, tXpX, BFloat16(-(2**15))) + # Copy from shared memory to registers and convert to Float32 + cute.autovec_copy(tXsX, tXrX) + x = tXrX.load().to(Float32) + + # Compute max and sum_exp using either two-pass or online softmax + if const_expr(not self.online_softmax): + max_x = row_reduce( + x, + cute.ReductionOp.MAX, + threads_per_row, + reduction_buffer[None, None, 0], + mbar_ptr + 0 if const_expr(self.cluster_n > 1) else None, + init_val=-Float32.inf, + hook_fn=cute.arch.cluster_wait if const_expr(self.cluster_n > 1) else None, + ) + if const_expr(self.reload_from == "smem"): + cute.autovec_copy(tXsX, tXrX) + x = tXrX.load().to(Float32) + log2_e = math.log2(math.e) + exp_x = cute.math.exp2(x * log2_e - (max_x * log2_e), fastmath=False) + denom = row_reduce( + exp_x, + cute.ReductionOp.ADD, + threads_per_row, + reduction_buffer[None, None, 1], + mbar_ptr + 1 if const_expr(self.cluster_n > 1) else None, + init_val=0.0, + ) + else: + # Online softmax: compute both max and sum in a single pass + max_x, denom, _ = online_softmax_reduce( + x, + threads_per_row, + reduction_buffer[None, None, 0], + mbar_ptr, + hook_fn=cute.arch.cluster_wait if const_expr(self.cluster_n > 1) else None, + return_exp_x=False + ) + + # Write results to global memory (only first thread in row, and first block in cluster) + if ( + tXcX[0][1] == 0 + and row < shape[0] + and (self.cluster_n == 1 or cute.arch.block_idx_in_cluster() == 0) + ): + target = Int32(mTarget[row]) + mXMax[row] = max_x + mXSumExp[row] = denom + mLogProbs[row] = Float32(mX[row, target]) - max_x - cute.math.log(denom, fastmath=True) + + +@torch.library.custom_op("roll_kernel::logprobs_fwd_out", mutates_args={"logprobs", "x_max", "x_sum_exp"}) +def logprobs_fwd_out( + x: Tensor, + target: Tensor, + logprobs: Tensor, + x_max: Tensor, + x_sum_exp: Tensor, +) -> None: + """ + Fused log_probs forward pass with in-place output tensors. + + Args: + x: Input logits tensor of shape (M, N), where M is batch size and N is vocabulary size + logprobs: Output tensor of shape (M,) [mutated in-place] + x_max: Output tensor for max values of shape (M,) [mutated in-place] + x_sum_exp: Output tensor for sum of exponentials of shape (M,) [mutated in-place] + + Returns: + None (output tensors are mutated in-place) + + Raises: + AssertionError: If input validation fails (wrong dimensions, device, or dtype) + """ + assert x.dim() == 2, "Input must be 2D" + assert target.dim() == 1, "Target must be 1D" + assert x.is_cuda and target.is_cuda, "Tensors must be on CUDA device" + assert x.dtype in [torch.float16, torch.bfloat16, torch.float32], "Unsupported input dtype" + # assert target.dtype == torch.int32, "Unsupported input dtype" + assert logprobs.dtype == torch.float32, "Unsupported input dtype" + + N = x.size(1) + + dtype = torch2cute_dtype_map[x.dtype] + target_dtype = torch2cute_dtype_map[target.dtype] + calc_dtype = torch2cute_dtype_map[x_max.dtype] + # Create compile key for caching compiled kernels + compile_key = (dtype, target_dtype, calc_dtype, N) + # Compile kernels if not in cache + if compile_key not in logprobs_fwd_out.compile_cache: + # Create symbolic batch size for compilation + batch_sym = cute.sym_int() + div = math.gcd(128 // dtype.width, N) + # Create fake tensors for compilation + x_cute = fake_tensor(dtype, (batch_sym, N), div) + target_cute = fake_tensor(target_dtype, (batch_sym,)) + logprobs_cute = fake_tensor(calc_dtype, (batch_sym,)) + x_max_cute = fake_tensor(calc_dtype, (batch_sym,)) + x_sum_exp_cute = fake_tensor(calc_dtype, (batch_sym,)) + + # Initialize Stage 1 and Stage 2 operators with online softmax enabled + # Online softmax computes both max and sum in a single pass for better performance + online_softmax = True if N > 16384 else False + logprobs_op = FusedLogProbsForward(dtype, N, online_softmax=online_softmax) + + # Compile Stage 1 kernel (computes max_x and sum_exp) + logprobs_fwd_out.compile_cache[compile_key] = cute.compile( + logprobs_op, + x_cute, + target_cute, + logprobs_cute, + x_max_cute, + x_sum_exp_cute, + cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), + options="--enable-tvm-ffi", + ) + + logprobs_fwd_out.compile_cache[compile_key]( + x, target, logprobs, x_max, x_sum_exp + ) + + +# Compilation cache for Stage 1 and Stage 2 kernels +logprobs_fwd_out.compile_cache = {} + + +def logprobs_fwd( + x: torch.Tensor, + target: torch.Tensor, +) -> torch.Tensor | tuple[torch.Tensor]: + """ + Compute entropy of each row in the input tensor. + + This function allocates output tensors and calls the fused entropy forward pass. + It returns the entropy value along with intermediate results that are needed + for the backward pass. + + Args: + x: Input logits tensor of shape (M, N), where M is batch size and N is vocabulary size + + Returns: + tuple: A tuple containing: + - entropy: Tensor of shape (M,) with entropy values + - x_max: Tensor of shape (M,) with max values (for backward pass) + - x_sum_exp: Tensor of shape (M,) with sum of exponentials (for backward pass) + - x_sum_softmax_times: Tensor of shape (M,) with sum(softmax * x) (for backward pass) + """ + M = x.size(0) + device = x.device + calc_type = torch.float32 + + # Allocate output tensors + logprobs = torch.zeros(M, dtype=calc_type, device=device) + x_max = torch.zeros(M, dtype=calc_type, device=device) + x_sum_exp = torch.zeros(M, dtype=calc_type, device=device) + + # Compute entropy using fused kernel + logprobs_fwd_out(x, target, logprobs, x_max, x_sum_exp) + return logprobs, x_max, x_sum_exp + + +def test_logprobs_fwd(M, N): + x = torch.rand((M, N), dtype=torch.bfloat16).cuda() + target = torch.randint(low=0, high=N, size=(M,)).cuda() + logprobs_fwd(x, target) + + +class FusedLogProbsBackward: + """ + Backward pass for entropy computation. + + This kernel computes the gradient of entropy with respect to the input logits. + The gradient formula is derived from the entropy loss: + dH/dx = -d_entropy * softmax(x) * (x - E[x]) + + Where: + - d_entropy: Gradient of entropy loss (scalar per row) + - softmax(x): Softmax probabilities + - E[x]: Expected value under softmax distribution (sum_softmax_times from forward) + + Args: + dtype: Data type of input tensor (e.g., BFloat16, Float16, Float32) + N: Number of columns in input tensor (vocabulary size) + """ + + def __init__(self, dtype: Type[cutlass.Numeric], N: int): + """ + Initialize the entropy backward pass. + """ + self.dtype = dtype + self.N = N + self.vecsize = 128 // dtype.width + + def _threads_per_row(self): + """ + Determine the optimal number of threads per row for reduction. + We split by blocks of 16k for large vocabularies. + """ + N = min(self.N, 16384) # We split by blocks of 16k + for limit, threads in [(64, 8), (128, 16), (3072, 32), (6144, 64), (16384, 128)]: + if N <= limit: + return threads + return 256 + + def _get_tiled_copy(self, vecsize: int): + """ + Get tiled copy configuration for efficient memory access. + + Args: + vecsize: Vector size for memory loading + + Returns: + tuple: (tiled_copy, tiler_mn, threads_per_row) + """ + assert self.N % vecsize == 0, f"Input N {self.N} is not divisible by vector size {vecsize}" + N = min(self.N, 16384) + num_threads = 128 if N <= 16384 else 256 + threads_per_row = self._threads_per_row() + rows_per_block = num_threads // threads_per_row + num_blocks_N = cute.ceil_div(N // vecsize, threads_per_row) + tiler_mn = (rows_per_block, vecsize * num_blocks_N * threads_per_row) + tiled_copy = copy_utils.tiled_copy_2d( + self.dtype, threads_per_row, num_threads, num_copy_elems=vecsize + ) + return tiled_copy, tiler_mn, threads_per_row + + @cute.jit + def __call__( + self, + mX: cute.Tensor, #[M, N] + mTarget: cute.Tensor, #[M,] + mDLogProbs: cute.Tensor, #[M,] + mXMax: cute.Tensor, #[M] + mXSumExp: cute.Tensor, #[M] + stream: cuda.CUstream, + ): + """ + Launch the backward kernel to compute gradients. + + Args: + mX: Input logits tensor of shape (M, N) + mDLogProbs: Gradient of logprobs of shape (M,) + mXMax: Max values from forward pass of shape (M,) + mXSumExp: Sum of exponentials from forward pass of shape (M,) + mXSumSoftmaxTimes: Sum(softmax * x) from forward pass of shape (M,) + stream: CUDA stream for kernel execution + """ + assert mX.element_type == self.dtype + # Calculate vector size for memory loading (max 128 bits per load) + # e.g. if self.N isn't divisible by 8 for bf16, we might use 64 bits (4 elements) copy + vecsize = math.gcd(self.N, 128 // self.dtype.width) + tiled_copy, tiler_mn, threads_per_row = self._get_tiled_copy(vecsize=vecsize) + num_threads = tiled_copy.size + mDLogProbs, mXMax, mXSumExp = [ + layout_utils.expand(X, dim=1, size=self.N) for X in (mDLogProbs, mXMax, mXSumExp) + ] + # Launch kernel with appropriate grid and block configuration + self.kernel( + mX, + mTarget, + mDLogProbs, + mXMax, + mXSumExp, + mX.shape, + tiler_mn, + tiled_copy, + threads_per_row, + ).launch( + grid=[ + cute.ceil_div(mX.shape[0], tiler_mn[0]), + cute.ceil_div(mX.shape[1], tiler_mn[1]), + 1, + ], + block=[num_threads, 1, 1], + stream=stream, + ) + + @cute.kernel + def kernel( + self, + mX: cute.Tensor, #[M, N] + mTarget: cute.Tensor, #[M,] + mDLogProbs: cute.Tensor, #[M,] + mXMax: cute.Tensor, #[M] + mXSumExp: cute.Tensor, #[M] + shape: cute.Shape, + tiler_mn: cute.Shape, + tiled_copy: cute.TiledCopy, + threads_per_row: cutlass.Constexpr[int], + ): + """ + CUDA kernel for entropy backward pass. + + This kernel computes the gradient of entropy with respect to the input logits: + dx = -d_entropy * softmax(x) * (x - E[x]) + + Where E[x] is the expected value under the softmax distribution. + """ + tidx, _, _ = cute.arch.thread_idx() + bidx, bidy, _ = cute.arch.block_idx() + + # Allocate shared memory for input data + smem = cutlass.utils.SmemAllocator() + sX = smem.allocate_tensor( + mX.element_type, cute.make_ordered_layout(tiler_mn, order=(1, 0)), byte_alignment=16 + ) + + idX = cute.make_identity_tensor(shape) + gX, cX = [cute.local_tile(mT, tiler_mn, (bidx, bidy)) for mT in (mX, idX)] + + thr_copy = tiled_copy.get_slice(tidx) + + # Partition tensors for this thread + tXgX = thr_copy.partition_S(gX) # Global memory partition + tXsX = thr_copy.partition_D(sX) # Shared memory partition + tXcX = thr_copy.partition_S(cX)[(0, None), None, None] # Coordinate partition + tXcFull = thr_copy.partition_S(cX) + tXrX = cute.make_fragment_like(tXgX) # Register fragment + + # Handle non-divisible N by creating predicates + is_even_N = const_expr(shape[1] % tiler_mn[1] == 0) + tXpX = ( + None if is_even_N else copy_utils.predicate_k(thr_copy.partition_S(cX), limit=shape[1]) + ) + copy = partial(copy_utils.copy, pred=tXpX) + + row = tXcX[0][0] + + # Load input data asynchronously from global to shared memory + if row < shape[0]: + copy(tXgX, tXsX, is_async=True) + cute.arch.cp_async_commit_group() + cute.arch.cp_async_wait_group(0) + if const_expr(not is_even_N): + utils.fill_oob(tXsX, tXpX, -tXsX.element_type.inf) + # Copy from shared memory to registers and convert to Float32 + cute.autovec_copy(tXsX, tXrX) + x = tXrX.load().to(Float32) + + # Load intermediate results from forward pass + target = Int32.zero + d_log_prob = Float32.zero + x_max = Float32.zero + x_sum_exp = Float32.zero + if row < shape[0]: + target = Int32(mTarget[row]) + d_log_prob = mDLogProbs[row] + x_max = mXMax[row] + x_sum_exp = mXSumExp[row] + + # Compute softmax for gradient calculation + log2_e = math.log2(math.e) + exp_x = cute.math.exp2(x * log2_e - (x_max * log2_e), fastmath=False) + + # Compute gradient: dx = d_logprob * (softmax(x) - one_hot(target)) + probs = -exp_x / x_sum_exp # softmax(x) + prob_shifted = probs + 1.0 + mask = cute.make_fragment_like(tXrX, Boolean) + for i in cutlass.range(cute.size(tXcFull), unroll_full=True): + mask[i] = tXcFull[i][1] == target + grad = cute.where(mask.load(), prob_shifted, probs) + grad = grad * d_log_prob + + # Store gradient to registers and write back to global memory + tXrX.store(grad.to(tXrX.element_type)) + if row < shape[0]: + copy(tXrX, tXgX) + + +def _logprobs_backward( + x: torch.Tensor, + target: torch.Tensor, + d_logprobs: torch.Tensor, + x_max: torch.Tensor, + x_sum_exp: torch.Tensor, +) -> None: + """ + Entropy backward pass using fused kernel. + + This function computes the gradient of entropy with respect to the input logits. + The gradient formula is: + dx = -d_entropy * softmax(x) * (x - E[x]) + + Where E[x] is the expected value under the softmax distribution. + + Args: + x: Input logits tensor of shape (M, N), where M is batch size and N is vocabulary size + d_entropy: Gradient of entropy loss of shape (M,) + x_max: Max values from forward pass of shape (M,) + x_sum_exp: Sum of exponentials from forward pass of shape (M,) + x_sum_softmax_times: Sum(softmax * x) from forward pass of shape (M,) + + Returns: + None (gradient is computed in-place on the input tensor x) + + Raises: + AssertionError: If input validation fails (wrong dimensions, device, or dtype) + """ + d_logprobs = d_logprobs.contiguous() + # Input validation + assert x.dim() == 2, "Input must be 2D" + assert d_logprobs.dim() == 1, "d_logprobsmust be 1D" + assert x_max.dim() == 1, "x_max must be 1D" + assert x_sum_exp.dim() == 1, "x_sum_exp must be 1D" + assert x.shape[0] == d_logprobs.shape[0], "Batch dimensions must match" + assert x.shape[0] == x_max.shape[0], "Batch dimensions must match" + assert x.shape[0] == x_sum_exp.shape[0], "Batch dimensions must match" + assert x.is_cuda and target.is_cuda and d_logprobs.is_cuda and x_max.is_cuda and x_sum_exp.is_cuda, ( + "Tensors must be on CUDA device" + ) + assert x.dtype in [torch.float16, torch.bfloat16, torch.float32], "Unsupported input dtype" + assert d_logprobs.dtype == torch.float32, "d_entropy must be float32" + assert x_max.dtype == torch.float32, "max_x must be float32" + assert x_sum_exp.dtype == torch.float32, "x_sum_exp must be float32" + + N = x.size(1) + dtype = torch2cute_dtype_map[x.dtype] + target_dtype = torch2cute_dtype_map[target.dtype] + calc_dtype = Float32 + # Create compile key for caching compiled kernels + compile_key = (dtype, target_dtype, calc_dtype, N) + if compile_key not in _logprobs_backward.compile_cache: + batch_sym = cute.sym_int() + div = math.gcd(128 // dtype.width, N) + x_cute= fake_tensor(dtype, (batch_sym, N), div) + target_cute = fake_tensor(target_dtype, (batch_sym,)) + d_logprobs_cute, x_max_cute, x_sum_exp_cute = [fake_tensor(calc_dtype, (batch_sym,))] * 3 + fused_logprobs_backward_op = FusedLogProbsBackward(dtype, N) + # Compile backward kernel + _logprobs_backward.compile_cache[compile_key] = cute.compile( + fused_logprobs_backward_op, + x_cute, target_cute, d_logprobs_cute, + x_max_cute, x_sum_exp_cute, + cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), + options="--enable-tvm-ffi", + ) + # Execute backward kernel + _logprobs_backward.compile_cache[compile_key]( + x, target, d_logprobs, x_max, x_sum_exp + ) + + +_logprobs_backward.compile_cache = {} + + +def logprobs_bwd( + x: torch.Tensor, + target: torch.Tensor, + d_logprobs: torch.Tensor, + x_max: torch.Tensor, + x_sum_exp: torch.Tensor, +) : + """ + Compute gradient of entropy with respect to input logits. + + This function computes the gradient and returns the modified input tensor. + + Args: + x: Input logits tensor of shape (M, N) + d_logprobs: Gradient of logprobs of shape (M,) + x_max: Max values from forward pass of shape (M,) + x_sum_exp: Sum of exponentials from forward pass of shape (M,) + + Returns: + torch.Tensor: The input tensor x with gradient computed in-place + """ + _logprobs_backward(x, target, d_logprobs, x_max, x_sum_exp) + return x + + +def test_logprobs_bwd(M, N): + """ + Test function for entropy backward pass. + """ + x = torch.rand((M, N), dtype=torch.bfloat16).cuda() + target = torch.randint(low=0, high=N, size=(M,)).cuda() + d_logprobs = torch.rand((M), dtype=torch.float32).cuda() + x_max = torch.rand((M), dtype=torch.float32).cuda() + x_sum_exp = torch.rand((M), dtype=torch.float32).cuda() + logprobs_bwd(x, target, d_logprobs, x_max, x_sum_exp) + + +if __name__ == '__main__': + # test_logprobs_fwd(8192, 128000) + test_logprobs_bwd(8192, 128000) diff --git a/mcore_adapter/src/mcore_adapter/parallel_functions/vocab_parallel.py b/mcore_adapter/src/mcore_adapter/parallel_functions/vocab_parallel.py index 452a7092b..5b61b98bc 100644 --- a/mcore_adapter/src/mcore_adapter/parallel_functions/vocab_parallel.py +++ b/mcore_adapter/src/mcore_adapter/parallel_functions/vocab_parallel.py @@ -4,7 +4,15 @@ import torch.distributed as dist from megatron.core import mpu -from ..utils import divide +from ..utils import divide, get_logger + +try: + from .fused_logprobs import logprobs_fwd, logprobs_bwd + FUSED_KERNEL_AVAILABLE = True +except ImportError: + FUSED_KERNEL_AVAILABLE = False + +logger = get_logger(__name__) class VocabUtility: @@ -70,46 +78,94 @@ def calculate_predicted_logits( class _VocabParallelLogProbs(torch.autograd.Function): @staticmethod - def forward(ctx, vocab_parallel_logits: "torch.Tensor", target: "torch.Tensor"): - vocab_parallel_logits = vocab_parallel_logits.float() - logits_max = torch.max(vocab_parallel_logits, dim=-1)[0] + def forward(ctx, vocab_parallel_logits: "torch.Tensor", target: "torch.Tensor", use_fused_kernel: bool = True): + + # Only use fused kernel when TP=1 and use_fused_kernel is True + if use_fused_kernel: + # Check if we can use fused kernel (only for TP=1) + tp_world_size = mpu.get_tensor_model_parallel_world_size() + if tp_world_size != 1 or not FUSED_KERNEL_AVAILABLE: + logger.warning(f"Disable use_fused_kernel because {tp_world_size=} and {FUSED_KERNEL_AVAILABLE=}.") + use_fused_kernel = False + + if use_fused_kernel: + # Use fused kernel for TP=1 + # Convert 3D input [B, S, D] to 2D [B*S, D] + vocab_parallel_logits_2d = vocab_parallel_logits.view(-1, vocab_parallel_logits.shape[-1]) + target_1d = target.view(-1) + + # Call fused kernel + logprobs, x_max, x_sum_exp = logprobs_fwd(vocab_parallel_logits_2d, target_1d) + + # Convert output back to original shape + if vocab_parallel_logits.dim() == 3: + batch_size, seq_len, vocab_size = vocab_parallel_logits.shape + logprobs = logprobs.view(batch_size, seq_len) + + # Save for backward + ctx.save_for_backward(vocab_parallel_logits, target_1d, x_max, x_sum_exp) + ctx.use_fused = True + + return logprobs + else: + # Use original torch implementation for TP>1 + vocab_parallel_logits = vocab_parallel_logits.float() + logits_max = torch.max(vocab_parallel_logits, dim=-1)[0] - dist.all_reduce(logits_max, op=dist.ReduceOp.MAX, group=mpu.get_tensor_model_parallel_group()) + dist.all_reduce(logits_max, op=dist.ReduceOp.MAX, group=mpu.get_tensor_model_parallel_group()) - ( - target_mask, - masked_target_1d, - predicted_logits, - sum_exp_logits, - exp_logits, - ) = _VocabParallelHelper.calculate_predicted_logits(vocab_parallel_logits, target, logits_max) + ( + target_mask, + masked_target_1d, + predicted_logits, + sum_exp_logits, + exp_logits, + ) = _VocabParallelHelper.calculate_predicted_logits(vocab_parallel_logits, target, logits_max) - dist.all_reduce(sum_exp_logits, op=dist.ReduceOp.SUM, group=mpu.get_tensor_model_parallel_group()) - dist.all_reduce(predicted_logits, op=dist.ReduceOp.SUM, group=mpu.get_tensor_model_parallel_group()) + dist.all_reduce(sum_exp_logits, op=dist.ReduceOp.SUM, group=mpu.get_tensor_model_parallel_group()) + dist.all_reduce(predicted_logits, op=dist.ReduceOp.SUM, group=mpu.get_tensor_model_parallel_group()) - predicted_logprobs = predicted_logits - torch.log(sum_exp_logits) + predicted_logprobs = predicted_logits - torch.log(sum_exp_logits) - # Save for backward - ctx.save_for_backward(exp_logits, target_mask, sum_exp_logits, masked_target_1d) + # Save for backward + ctx.save_for_backward(exp_logits, target_mask, sum_exp_logits, masked_target_1d) + ctx.use_fused = False - return predicted_logprobs + return predicted_logprobs @staticmethod def backward(ctx, grad_output: "torch.Tensor"): - exp_logits, target_mask, sum_exp_logits, masked_target_1d = ctx.saved_tensors + use_fused = ctx.use_fused - exp_logits.div_(sum_exp_logits.unsqueeze(dim=-1)) - exp_logits.neg_() - grad_input = exp_logits - grad_2d = grad_input.view(-1, grad_input.size()[-1]) - arange_1d = torch.arange(start=0, end=grad_2d.size()[0], device=grad_input.device) - grad_2d[arange_1d, masked_target_1d] += 1 - target_mask.view(-1).float() - grad_input.mul_(grad_output.unsqueeze(dim=-1)) + if use_fused: + # Use fused kernel backward + vocab_parallel_logits, target_1d, x_max, x_sum_exp = ctx.saved_tensors - return grad_input, None + vocab_parallel_logits_2d = vocab_parallel_logits.view(-1, vocab_parallel_logits.shape[-1]) + # Call fused backward + grad_input_2d = logprobs_bwd(vocab_parallel_logits_2d, target_1d, grad_output.view(-1), x_max, x_sum_exp) -def vocab_parallel_logprobs(vocab_parallel_logits, target) -> "torch.Tensor": + # Convert gradient back to original shape + grad_input = grad_input_2d.view(vocab_parallel_logits.shape) + + return grad_input, None, None + else: + # Use original torch backward + exp_logits, target_mask, sum_exp_logits, masked_target_1d = ctx.saved_tensors + + exp_logits.div_(sum_exp_logits.unsqueeze(dim=-1)) + exp_logits.neg_() + grad_input = exp_logits + grad_2d = grad_input.view(-1, grad_input.size()[-1]) + arange_1d = torch.arange(start=0, end=grad_2d.size()[0], device=grad_input.device) + grad_2d[arange_1d, masked_target_1d] += 1 - target_mask.view(-1).float() + grad_input.mul_(grad_output.unsqueeze(dim=-1)) + + return grad_input, None, None + + +def vocab_parallel_logprobs(vocab_parallel_logits, target, use_fused_kernel: bool = True) -> "torch.Tensor": """ Get logprobs when logits are split across tensor parallel ranks @@ -118,12 +174,15 @@ def vocab_parallel_logprobs(vocab_parallel_logits, target) -> "torch.Tensor": dimension is [batch_size, sequence_length, vocab_size/num_parallel_ranks] target: correct vocab ids of dimension [batch_size, sequence_length] + + use_fused_kernel: whether to use fused kernel (only works for TP=1) + Returns: logprobs: logprobs of dimension [batch_size, sequence_length] (It's fine to change the order of sequence_length and batch_size in dimension) """ - return _VocabParallelLogProbs.apply(vocab_parallel_logits, target) + return _VocabParallelLogProbs.apply(vocab_parallel_logits, target, use_fused_kernel) def vocab_parallel_target_rank(vocab_parallel_logits: "torch.Tensor", target: "torch.Tensor") -> "torch.Tensor": diff --git a/mcore_adapter/src/mcore_adapter/trainer/dpo_trainer.py b/mcore_adapter/src/mcore_adapter/trainer/dpo_trainer.py index cbc26f7e6..ceaf36f99 100644 --- a/mcore_adapter/src/mcore_adapter/trainer/dpo_trainer.py +++ b/mcore_adapter/src/mcore_adapter/trainer/dpo_trainer.py @@ -85,7 +85,10 @@ def _post_compute_log_probs( self, labels: "torch.Tensor", loss_mask: "torch.Tensor", logits: "torch.Tensor", non_loss_data: bool = False ): batch_size = labels.size(0) // 2 - logprobs = vocab_parallel_logprobs(logits, labels) + logprobs = vocab_parallel_logprobs( + logits, labels, + use_fused_kernel=self.args.cross_entropy_loss_fusion + ) logprobs = (logprobs * loss_mask).sum(-1) if mpu.get_context_parallel_world_size() > 1: dist.all_reduce(logprobs, group=mpu.get_context_parallel_group()) @@ -145,7 +148,10 @@ def _post_compute_loss( logits: "torch.Tensor", ): batch_size = labels.size(0) // 2 - logprobs = vocab_parallel_logprobs(logits, labels) + logprobs = vocab_parallel_logprobs( + logits, labels, + use_fused_kernel=self.args.cross_entropy_loss_fusion + ) logprobs = (logprobs * loss_mask).sum(-1) response_lens = loss_mask.sum(-1) cp_size = mpu.get_context_parallel_world_size() diff --git a/mcore_adapter/src/mcore_adapter/trainer/trainer.py b/mcore_adapter/src/mcore_adapter/trainer/trainer.py index 06b87c170..1402ada02 100644 --- a/mcore_adapter/src/mcore_adapter/trainer/trainer.py +++ b/mcore_adapter/src/mcore_adapter/trainer/trainer.py @@ -996,7 +996,46 @@ def _maybe_log_save_evaluate( logs["grad_norm"] = grad_norm.detach().item() if isinstance(grad_norm, torch.Tensor) else grad_norm logs["learning_rate"] = self._get_learning_rate() logs.update(moe_losses) - logs.update(mtp_losses) + # MTP layers only exist on PP last stage, so mtp_losses is only populated there. + # Broadcast mtp_losses to all PP ranks so rank 0 (PP first stage) can log them. + if mpu.get_pipeline_model_parallel_world_size() > 1: + is_last_stage = mpu.is_pipeline_last_stage(ignore_virtual=True) + # First, broadcast the number of mtp_losses entries and their keys from last stage + num_keys = torch.tensor(len(mtp_losses), device=self.args.device, dtype=torch.long) + torch.distributed.broadcast( + num_keys, + src=mpu.get_pipeline_model_parallel_last_rank(), + group=mpu.get_pipeline_model_parallel_group(), + ) + if num_keys.item() > 0: + if is_last_stage: + keys_list = list(mtp_losses.keys()) + else: + keys_list = None + # Broadcast keys via object broadcast + keys_broadcast = [keys_list] + torch.distributed.broadcast_object_list( + keys_broadcast, + src=mpu.get_pipeline_model_parallel_last_rank(), + group=mpu.get_pipeline_model_parallel_group(), + ) + keys_list = keys_broadcast[0] + # Broadcast each value tensor + for key in keys_list: + if is_last_stage: + val = mtp_losses[key] + if not isinstance(val, torch.Tensor): + val = torch.tensor(val, device=self.args.device) + else: + val = torch.tensor(0.0, device=self.args.device) + torch.distributed.broadcast( + val, + src=mpu.get_pipeline_model_parallel_last_rank(), + group=mpu.get_pipeline_model_parallel_group(), + ) + mtp_losses[key] = val + for k, v in mtp_losses.items(): + logs[k] = v.detach().item() if isinstance(v, torch.Tensor) else v if metrics_tensors is not None and len(self.metrics_keys) > 1: # metrics except loss metrics = self.gather_metrics(metrics_tensors) metrics.pop("loss", None) @@ -1093,7 +1132,7 @@ def _save_optimizer_and_scheduler(self, output_dir): def save_model(self, output_dir: str = None, _internal_call: bool = False): output_dir = output_dir or self.args.output_dir if not (self.args.save_only_model and self.args.save_hf_model): - self.model.save_pretrained(output_dir, save_merged_model=self.args.save_merged_model) + self.model.save_pretrained(output_dir, save_merged_model=self.args.save_merged_model, ckpt_format=self.args.ckpt_format) if self.args.save_hf_model: self.model.save_pretrained_as_hf(output_dir) if self.args.should_save: diff --git a/mcore_adapter/src/mcore_adapter/training_args.py b/mcore_adapter/src/mcore_adapter/training_args.py index 412322d11..8fd07e9fe 100644 --- a/mcore_adapter/src/mcore_adapter/training_args.py +++ b/mcore_adapter/src/mcore_adapter/training_args.py @@ -52,6 +52,10 @@ class DistributingParallelArguments: default=None, metadata={"help": "Degree of expert model parallelism."}, ) + expert_tensor_parallel_size: Optional[int] = field( + default=None, + metadata={"help": "Degree of expert tensor parallelism, default to 1, avoid override by tp_size"} + ) account_for_embedding_in_pipeline_split: Optional[bool] = field( default=None, metadata={ @@ -133,6 +137,15 @@ class DistributingParallelArguments: default=False, metadata={"help": "Use fused RoPE kernel."}, ) + cross_entropy_loss_fusion: bool = field( + default=False, + metadata={ + "help": "Use fused CUDA kernels for softmax-based computations (logprobs and entropy). " + "These kernels leverage cluster block and online softmax techniques for better performance " + "on large vocabularies. Only effective when TP=1. Set to False to fall back to the " + "standard PyTorch implementation for debugging or compatibility." + }, + ) # moe moe_layer_recompute: bool = field( default=False, @@ -186,6 +199,13 @@ class DistributingParallelArguments: " Without this, the shared epxerts execute after the routed experts." }, ) + moe_enable_routing_replay: bool = field( + default=False, + metadata={ + "help": "Enable RouterReplay for MoE routing, which allows recording and replaying routing " + "decisions to ensure consistent expert assignment across inference and training." + }, + ) moe_router_dtype: Optional[str] = field( default=None, metadata={ @@ -194,7 +214,10 @@ class DistributingParallelArguments: }, ) # mtp - mtp_num_layers: Optional[int] = field(default=None, metadata={"help": "The number of mtp layers."}) + mtp_num_layers: Optional[int] = field( + default=0, + metadata={"help": "The number of mtp layers. Default to 0 to disable mtp."}, + ) # train options calculate_per_token_loss: bool = field( default=False, @@ -210,6 +233,20 @@ class DistributingParallelArguments: "choices": ["local", "transformer_engine"], }, ) + cuda_graph_impl: Optional[Literal["local", "transformer_engine"]] = field( + default=None, + metadata={ + "help": "Determines the CUDA graph capture implementation.", + "choices": ["local", "transformer_engine"], + }, + ) + cuda_graph_scope: str = field( + default="full", + metadata={ + "help": "Determines the CUDA graphs capturing scope." + "Default to full for backward compatibility." + }, + ) fp8_recipe: Optional[str] = field( default=None, metadata={ @@ -229,7 +266,7 @@ class DistributingParallelArguments: "help": "FP8 format to use. Supported formats: 'e4m3', 'hybrid'. Do not change if unsure", }, ) - additional_configs: Optional[Union[dict, str]] = field( + additional_configs: Optional[Union[str, dict]] = field( default_factory=dict, metadata={ "help": "Dictionary or Path to a JSON file containing additional configuration parameters for the model.", @@ -248,6 +285,10 @@ def __post_init__(self): if self.recompute_modules is not None and isinstance(self.recompute_modules, str): self.recompute_modules = self.recompute_modules.split(",") + logger.info(f"cuda_graph_scope before processing: {self.cuda_graph_scope}, type: {type(self.cuda_graph_scope)}") + if self.cuda_graph_scope is not None and isinstance(self.cuda_graph_scope, str): + self.cuda_graph_scope = self.cuda_graph_scope.split(",") + if self.variable_seq_lengths and self.moe_token_dispatcher_type in ["allgather"]: raise ValueError( f"Token dispatcher type: {self.moe_token_dispatcher_type} does not support " @@ -287,7 +328,7 @@ class MegatronArguments(DistributingParallelArguments): metadata={"help": "Use distributed optimizer."}, ) distrib_optim_fully_reshardable: bool = field( - default=True, + default=False, metadata={"help": "Whether optimizer states are fully reshardable."}, ) distrib_optim_fully_reshardable_mem_efficient: bool = field( @@ -336,12 +377,25 @@ class MegatronArguments(DistributingParallelArguments): save_hf_model: bool = field(default=False, metadata={"help": "Save model as hf format."}) save_merged_model: bool = field(default=False, metadata={"help": "Save merged model weights in LoRA training."}) + ckpt_format: Literal["legacy", "torch_dist"] = field( + default="legacy", + metadata={ + "help": "Checkpoint format to use. " + "`legacy` using `torch.save/load` directly, " + "`torch_dist` is a megatron built-in distributed checkpointing format" + }, + ) sequence_packing: bool = field( default=False, metadata={"help": "Enable sequence packing without cross-attention."}, ) + compile_warmup: bool = field( + default=False, + metadata={"help": "Enable compile warmup before training."}, + ) + def __post_init__(self): super().__post_init__() if self.overlap_param_gather: @@ -349,6 +403,10 @@ def __post_init__(self): assert self.overlap_grad_reduce, ( "--overlap_grad_reduce should be turned on when using --overlap_param_gather" ) + if self.expert_tensor_parallel_size is not None and self.expert_tensor_parallel_size != self.tensor_model_parallel_size: + logger.info(f"Set ckpt_format to `torch_dist` when etp_size[{self.expert_tensor_parallel_size}]" + f"is not equal to tp_size[{self.tensor_model_parallel_size}]") + self.ckpt_format = "torch_dist" @classmethod def from_json_file(cls, json_file_path) -> "MegatronArguments": diff --git a/requirements_common.txt b/requirements_common.txt index f57299c9f..e1fd476dc 100644 --- a/requirements_common.txt +++ b/requirements_common.txt @@ -1,6 +1,6 @@ ray[default,cgraph]==2.48.0 # vllm required ray[default,cgraph]>=2.48.0 numpy<2.0a0,>=1.25 -tensordict +tensordict>=0.10.0 sympy modelscope datasets==3.1.0 @@ -21,11 +21,17 @@ codetiming more_itertools pybase64 sglang-router +TransferQueue==0.1.6 wandb swanlab trackio +opentelemetry-api +opentelemetry-sdk +opentelemetry-exporter-otlp +opentelemetry-proto + math-verify openai langdetect diff --git a/requirements_em_local_debug.txt b/requirements_em_local_debug.txt index 396317258..a14ac83b3 100644 --- a/requirements_em_local_debug.txt +++ b/requirements_em_local_debug.txt @@ -4,7 +4,6 @@ torchaudio # mac py-cpuinfo -deepspeed==0.16.4 ray<=2.46.0,>=2.40.0 numpy<2.0a0,>=1.25 tensordict diff --git a/requirements_torch2100_vllm.txt b/requirements_torch2100_vllm.txt index 076bd41d6..5be078b3f 100644 --- a/requirements_torch2100_vllm.txt +++ b/requirements_torch2100_vllm.txt @@ -5,7 +5,6 @@ torchvision==0.25.0.* torchaudio==2.10.0.* transformers==5.2.0 -deepspeed==0.16.4 accelerate==1.7.0 flash-attn diff --git a/requirements_torch260_vllm.txt b/requirements_torch260_vllm.txt index 6e6cbb001..b27e5d589 100644 --- a/requirements_torch260_vllm.txt +++ b/requirements_torch260_vllm.txt @@ -7,6 +7,6 @@ torchaudio==2.6.0.* flash-attn transformer-engine[pytorch]==2.2.0 -deepspeed==0.16.4 +deepspeed==0.16.4 # only to avoid error caused by improper transformers->deepspeed->triton version accelerate==0.34.2 vllm==0.8.4 diff --git a/requirements_torch280_sglang.txt b/requirements_torch280_sglang.txt index b706d2fc5..9c1ce10c2 100644 --- a/requirements_torch280_sglang.txt +++ b/requirements_torch280_sglang.txt @@ -1,12 +1,13 @@ -r requirements_common.txt +quack-kernels==0.2.4 + torch==2.8.0.* torchvision==0.23.0.* torchaudio==2.8.0.* -deepspeed==0.16.4 accelerate==0.34.2 sglang[srt,torch-memory-saver]==0.5.2 # for GDN , eg: Qwen3Next -flash-linear-attention \ No newline at end of file +flash-linear-attention diff --git a/requirements_torch280_vllm.txt b/requirements_torch280_vllm.txt index 34eba804c..a6e74de05 100644 --- a/requirements_torch280_vllm.txt +++ b/requirements_torch280_vllm.txt @@ -1,11 +1,12 @@ -r requirements_common.txt +quack-kernels==0.2.4 + torch==2.8.0.* torchvision==0.23.0.* torchaudio==2.8.0.* transformers==4.57.0 -deepspeed==0.16.4 accelerate==0.34.2 flash-attn @@ -15,4 +16,4 @@ vllm==0.10.2 # for GDN , eg: Qwen3Next -flash-linear-attention \ No newline at end of file +flash-linear-attention diff --git a/requirements_vision.txt b/requirements_vision.txt index 7d57d7236..9d2f25aa2 100644 --- a/requirements_vision.txt +++ b/requirements_vision.txt @@ -1 +1,4 @@ -pycocotools # for detection verify \ No newline at end of file +pycocotools # for detection verify +qwen_vl_utils # for video +decord # for video +rouge_score==0.1.2 # for video-r1 reward \ No newline at end of file diff --git a/roll/configs/base_config.py b/roll/configs/base_config.py index 80949874c..12262b9f3 100644 --- a/roll/configs/base_config.py +++ b/roll/configs/base_config.py @@ -82,6 +82,16 @@ class ScheduleConfig: metadata={"help": "The router configuration, encapsulated in a RouterArguments object."}, ) +@dataclass +class TransferBackendArguments: + backend_name: str = field( + default=None, + metadata={"help": "The registered backend for transfer."} + ) + backend_config: Dict = field( + default_factory=dict, + metadata={"help": "Configuration dictionary for the backend."} + ) @dataclass class BaseConfig(ScheduleConfig): @@ -215,6 +225,10 @@ class BaseConfig(ScheduleConfig): profiler_output_dir: str = field( default="./output/profiler", metadata={"help": "Directory to write profiler logs."} ) + otlp_output_dir: str = field( + default="./output/otlp", + metadata={"help": "The output directory where the OTLP collector saved to."}, + ) system_envs: dict = field( default_factory=dict, metadata={"help": "system environment variables."} @@ -244,6 +258,11 @@ class BaseConfig(ScheduleConfig): metadata={"help": "Rollout mock configuration for precision alignment testing."} ) + transfer_backend: Optional[TransferBackendArguments] = field( + default=None, + metadata={"help": "Transfer backend configuration."} + ) + def to_dict(self): return dataclasses.asdict(self) @@ -286,6 +305,9 @@ def __post_init__(self): output_dir = self.checkpoint_config.get("output_dir") self.checkpoint_config["output_dir"] = os.path.join(output_dir, datetime.now().strftime("%Y%m%d-%H%M%S")) logger.info(f"add timestamp to output_dir {self.checkpoint_config['output_dir']}") + if self.resume_from_checkpoint is True: + logger.info("ensure async_upload=False when resume_from_checkpoint=True and auto resume from lastest ckpt") + self.checkpoint_config["async_upload"] = False for attribute_name in dir(self): attribute = getattr(self, attribute_name) @@ -319,6 +341,13 @@ def __post_init__(self): os.environ["REPORT_LENGTH_AND_REWARDS"] = "1" os.environ.update(self.system_envs) + # Auto-generate OTLP endpoint when OTel tracing is enabled + if self.system_envs.get("ROLL_OTEL_ENABLED") == "1": + from roll.utils.telemetry import resolve_otel_endpoint + endpoint = resolve_otel_endpoint() + self.system_envs["OTEL_EXPORTER_OTLP_ENDPOINT"] = endpoint + os.environ["OTEL_EXPORTER_OTLP_ENDPOINT"] = endpoint + from ..platforms import current_platform self.num_gpus_per_node = current_platform.device_count() @@ -330,6 +359,11 @@ def __post_init__(self): self.reference.system_envs.update({k: v for k, v in self.system_envs.items() if k not in self.reference.system_envs}) if hasattr(self, 'critic') and isinstance(self.critic, WorkerConfig): self.critic.system_envs.update({k: v for k, v in self.system_envs.items() if k not in self.critic.system_envs}) + # Also propagate system_envs to _reference_configs for multi-teacher + if hasattr(self, '_reference_configs'): + for ref_cfg in self._reference_configs.values(): + if isinstance(ref_cfg, WorkerConfig): + ref_cfg.system_envs.update({k: v for k, v in self.system_envs.items() if k not in ref_cfg.system_envs}) # Validate rollout_batch_size divisibility for Megatron data parallelism if hasattr(self, 'actor_train') and isinstance(self.actor_train, WorkerConfig) and self.actor_train.strategy_args is not None: @@ -367,6 +401,11 @@ def __post_init__(self): if isinstance(attribute, WorkerConfig): if attribute.device_mapping is not None: total_devices.extend(attribute.device_mapping) + # Also include _reference_configs device_mappings for multi-teacher + if hasattr(self, '_reference_configs'): + for ref_cfg in self._reference_configs.values(): + if isinstance(ref_cfg, WorkerConfig) and ref_cfg.device_mapping is not None: + total_devices.extend(ref_cfg.device_mapping) if len(total_devices) > 0: max_gpu_num = max(total_devices) + 1 if max_gpu_num <= self.num_gpus_per_node: @@ -557,10 +596,11 @@ class PPOConfig(BaseConfig): "to determine which config class and pipeline to use. " "'rlvr': RLVRConfig + RLVRPipeline, 'agentic': AgenticConfig + AgenticPipeline"} ) - teacher: WorkerConfig = field( + teacher: Union[Dict[str, WorkerConfig], WorkerConfig] = field( default_factory=WorkerConfig, - metadata={"help": "Configuration for the teacher role (used in OPD mode). " - "When is_pure_opd=True or use_opd=True, teacher is automatically mapped to reference."} + metadata={"help": "Teacher model config (OPD mode). Single: WorkerConfig; " + "Multi-teacher: Dict[str, WorkerConfig]. " + "Dict[str, WorkerConfig] must be placed first in Union for dacite."} ) student_train: WorkerConfig = field( default_factory=WorkerConfig, @@ -586,11 +626,6 @@ class PPOConfig(BaseConfig): "The OPD KL is computed as: reverse_kl = student_logp - teacher_logp, " "and added to token_level_rewards as: reward - opd_kl_coef * reverse_kl"} ) - opd_kl_coef: float = field( - default=1.0, - metadata={"help": "Coefficient for OPD KL penalty when use_opd=True. " - "Controls the weight of distillation signal relative to RL reward."} - ) def __post_init__(self): super().__post_init__() @@ -632,22 +667,13 @@ def _handle_opd_mapping(self): """ Handle OPD (On-Policy Distillation) mode configuration mapping. - Pure OPD mode (is_pure_opd=True): - - Requires: student_train, student_infer, teacher - - Forbidden: reference - - Mapping: student_train → actor_train, student_infer → actor_infer, teacher → reference - - Mixed OPD mode (use_opd=True): - - Requires: teacher - - Forbidden: reference - - Mapping: teacher → reference only - - actor_train/actor_infer are configured normally by user - This method is called at the beginning of __post_init__ before normal PPO initialization. + + Teacher is normalized to _reference_configs: Dict[str, WorkerConfig] immediately, + so all subsequent logic is unified regardless of single vs multi-teacher. """ - has_student_train = self.student_train.is_configured - has_student_infer = self.student_infer.is_configured - has_teacher = self.teacher.is_configured + has_teacher = isinstance(self.teacher, WorkerConfig) and self.teacher.is_configured + has_teachers = isinstance(self.teacher, dict) and bool(self.teacher) has_reference_configured = self.reference.is_configured # Mutual exclusion check @@ -658,45 +684,66 @@ def _handle_opd_mapping(self): "or use_opd=True for mixed mode (external rewards + Teacher KL)." ) - # ========== Pure OPD mode ========== + if has_teachers and has_teacher: + raise ValueError("Cannot configure both multi-teacher dict and single-teacher WorkerConfig.") + + # Step 1: Normalize teacher to _reference_configs dict (always Dict[str, WorkerConfig]) + if has_teachers: + self._reference_configs = self.teacher + elif has_teacher: + self._reference_configs = {"default": self.teacher} + else: + self._reference_configs = {} + + # Step 1.5: Build tag → teacher_names routing map for multi-teacher OPD + self._tag_to_teacher_names: Dict[str, List[str]] = {} + self._teacher_names_for_all_tags: List[str] = [] + for name, ref_cfg in self._reference_configs.items(): + if not ref_cfg.tag_included: + # Empty tag_included means this teacher handles all tags + self._teacher_names_for_all_tags.append(name) + else: + for tag in ref_cfg.tag_included: + if tag not in self._tag_to_teacher_names: + self._tag_to_teacher_names[tag] = [] + self._tag_to_teacher_names[tag].append(name) + + # Step 2: Pure OPD mode if self.is_pure_opd: - # Validation: all student fields and teacher must be configured - if not (has_student_train and has_student_infer and has_teacher): + if not self._reference_configs: + raise ValueError("Pure OPD requires teacher config.") + if not (self.student_train.is_configured and self.student_infer.is_configured): raise ValueError( "In pure OPD mode (is_pure_opd=True), 'student_train', 'student_infer' " "and teacher must be configured.\n" ) - - # Perform mapping for pure OPD - logger.info(f"Pure OPD mode: mapping student_train to actor_train") + logger.info(f"Pure OPD mode: mapping student_train to actor_train, " + f"student_infer to actor_infer, " + f"{len(self._reference_configs)} teacher(s) to reference") self.actor_train = self.student_train - logger.info(f"Pure OPD mode: mapping student_infer to actor_infer") self.actor_infer = self.student_infer - logger.info(f"Pure OPD mode: mapping teacher to reference") - self.reference = self.teacher - - # Enable reference for OPD mode (needed for both pure and mixed mode) + self.reference = next(iter(self._reference_configs.values())) self.enable_reference = True + # Propagate opd_kl_coef default (1.0) to each teacher if not explicitly set + for ref_cfg in self._reference_configs.values(): + if ref_cfg.opd_kl_coef is None: + ref_cfg.opd_kl_coef = 1.0 - # ========== Mixed OPD mode ========== + # Step 3: Mixed OPD mode elif self.use_opd: - # Validation: teacher must be configured, reference should NOT be configured - if not has_teacher: - raise ValueError( - "In mixed OPD mode (use_opd=True), 'teacher' must be configured.\n" - ) + if not self._reference_configs: + raise ValueError("Mixed OPD requires teacher config.") if has_reference_configured: raise ValueError( "In mixed OPD mode (use_opd=True), 'reference' should NOT be configured. " ) - - # Perform mapping for mixed OPD (only teacher → reference) - logger.info(f"Mixed OPD mode: mapping teacher to reference") - self.reference = self.teacher - # Note: actor_train and actor_infer are configured normally by user - - # Enable reference for OPD mode (needed for both pure and mixed mode) + logger.info(f"Mixed OPD mode: mapping {len(self._reference_configs)} teacher(s) to reference") + self.reference = next(iter(self._reference_configs.values())) self.enable_reference = True + # Propagate opd_kl_coef default (1.0) to each teacher if not explicitly set + for ref_cfg in self._reference_configs.values(): + if ref_cfg.opd_kl_coef is None: + ref_cfg.opd_kl_coef = 1.0 def set_max_steps(self, max_steps: int): actor_backward_batch_size = ( @@ -743,11 +790,7 @@ def set_old_logprobs_status(self): ) dp_size = 1 if self.actor_train.strategy_args is not None: - if self.actor_train.strategy_args.strategy_name == "deepspeed_train": - configured_ulysses_size = getattr(self.actor_train.model_args, 'ulysses_size', None) or 1 - cp_size = self._get_effective_cp_size_ulysses(configured_ulysses_size) - dp_size = len(self.actor_train.device_mapping) // cp_size - elif self.actor_train.strategy_args.strategy_name in ("fsdp2_train", "fsdp2_infer"): + if self.actor_train.strategy_args.strategy_name in ("fsdp2_train", "fsdp2_infer"): configured_ulysses_size = getattr(self.actor_train.model_args, 'ulysses_size', None) or 1 cp_size = self._get_effective_cp_size_ulysses(configured_ulysses_size) dp_size = len(self.actor_train.device_mapping) // cp_size @@ -779,6 +822,37 @@ def set_old_logprobs_status(self): def async_pipeline(self) -> bool: return self.async_generation_ratio > 0 + @property + def reference_configs(self) -> Dict[str, WorkerConfig]: + """Always returns Dict[str, WorkerConfig] for unified pipeline usage. + Single teacher is normalized to {"default": cfg}, multi-teacher to {name: cfg, ...}.""" + if not hasattr(self, '_reference_configs') or not self._reference_configs: + self._reference_configs = {"default": self.reference} + return self._reference_configs + + @property + def is_multi_teacher(self) -> bool: + return len(self.reference_configs) > 1 + + @property + def tag_to_teacher_names(self) -> Dict[str, List[str]]: + """Map tag -> list of teacher names that should handle it.""" + if not hasattr(self, '_tag_to_teacher_names'): + self._tag_to_teacher_names = {} + return self._tag_to_teacher_names + + @property + def teacher_names_for_all_tags(self) -> List[str]: + """Teachers with empty tag_included (handle all tags).""" + if not hasattr(self, '_teacher_names_for_all_tags'): + self._teacher_names_for_all_tags = [] + return self._teacher_names_for_all_tags + + @property + def needs_teacher_routing(self) -> bool: + """Whether any teacher has non-empty tag_included (requires routing logic).""" + return bool(self._tag_to_teacher_names) + @property def is_actor_infer_colocated(self) -> bool: """Whether actor_infer are colocated with any other clusters (exclude reward).""" @@ -799,16 +873,21 @@ def _apply_opd_config(self): Note: The mapping of student_*/teacher to actor_*/reference is already handled by _handle_opd_mapping(). This method only applies OPD-specific parameter overrides like gamma, adv_estimator, etc. - - This method handles both pure OPD mode (is_pure_opd=True) - and mixed OPD mode (use_opd=True). """ + if not (self.is_pure_opd or self.use_opd): + return + + # Set teacher worker names + if len(self._reference_configs) > 1: + for name, ref_cfg in self._reference_configs.items(): + ref_cfg.name = f"teacher-{name}" + else: + self.reference.name = "teacher" + # Pure OPD mode specific settings if self.is_pure_opd: - # Set worker names for OPD mode (override default names for both modes) self.actor_train.name = "student_train" self.actor_infer.name = "student_infer" - self.reference.name = "teacher" # gamma=0: OPD's token_level_rewards has KL penalty at every token # If gamma=1, compute_reinforce_return will accumulate KL values across entire sequence @@ -827,8 +906,5 @@ def _apply_opd_config(self): logger.info(f"Pure OPD mode configured: gamma={self.gamma}, adv_estimator={self.adv_estimator}") - # Mixed OPD mode doesn't need parameter overrides elif self.use_opd: - # Set worker names for OPD mode (override default names for both modes) - self.reference.name = "teacher" - logger.info(f"Mixed OPD mode configured: opd_kl_coef={self.opd_kl_coef}") + logger.info(f"Mixed OPD mode configured") diff --git a/roll/configs/data_args.py b/roll/configs/data_args.py index 921ecd089..7d8c0d98a 100644 --- a/roll/configs/data_args.py +++ b/roll/configs/data_args.py @@ -37,6 +37,45 @@ class DataArguments: prompt: Optional[str] = field(default=None, metadata={"help": "Which column in file to use as prompt"}) response: Optional[str] = field(default="solution", metadata={"help": "Which column in file to use as label"}) messages: Optional[str] = field(default=None, metadata={"help": "Which column in file to use as messages"}) + # args for multi-modal, corresponding to qwen-vl-utils + image_max_pixels: int = field( + default=4096 * 4096, + metadata={"help": "The maximum number of pixels of image inputs."}, + ) + image_min_pixels: int = field( + default=4 * 28 * 28, + metadata={"help": "The minimum number of pixels of image inputs."}, + ) + video_max_pixels: int = field( + default=2048 * 2048, + metadata={"help": "The maximum number of pixels of video frames."}, + ) + video_min_pixels: int = field( + default=4 * 28 * 28, + metadata={"help": "The minimum number of pixels of video frames."}, + ) + video_total_pixels: int = field( + default=16384 * 28 * 28, + metadata={"help": "The total number of pixels for a video which can be used to limit sequence length."}, + ) + video_fps: float = field( + default=2.0, + metadata={"help": "The frames to sample per second for video inputs."}, + ) + video_max_frames: int = field( + default=128, + metadata={"help": "The maximum number of sampled frames for video inputs."}, + ) + video_resize_chunk_frames: int = field( + default=64, + metadata={"help": "The frames of the chunk to resize the video."}, + ) + image_folder: Optional[str] = field(default=None, metadata={"help": "Path to the folder containing the images."}) + video_folder: Optional[str] = field(default=None, metadata={"help": "Path to the folder containing the videos."}) + audio_folder: Optional[str] = field(default=None, metadata={"help": "Path to the folder containing the audios."}) + custom_data_kwargs_func: Optional[str] = field( + default=None, metadata={"help": "Path to custom data kwargs function."} + ) def __post_init__(self): assert not ( diff --git a/roll/configs/model_args.py b/roll/configs/model_args.py index c9b8b8446..15e839497 100644 --- a/roll/configs/model_args.py +++ b/roll/configs/model_args.py @@ -93,7 +93,7 @@ class ModelArguments(LoraArguments): default="bf16", metadata={"help": "Set model dtype as fp32, bf16, or fp16, otherwise use config's torch_dtype"} ) model_type: Optional[ - Literal["auto_sequence_classification", "auto_token_classification", "trl", "diffusion_module"] + Literal["auto_sequence_classification", "auto_token_classification", "trl"] ] = field( default=None, metadata={"help": "reward model type."}, diff --git a/roll/configs/worker_config.py b/roll/configs/worker_config.py index 65aafa63b..922d787e4 100644 --- a/roll/configs/worker_config.py +++ b/roll/configs/worker_config.py @@ -7,23 +7,31 @@ logger = get_logger() +@dataclass +class RouterReplayConfig: + """Configuration for router replay functionality.""" + mode: Literal["disable", "R2", "R3"] = field( + default="disable", + metadata={ + "help": "Router replay mode. Options: 'disable' (no replay), 'R2' (Megatron infer + Megatron train), 'R3' (sglang infer + Megatron train)." + }, + ) + + @dataclass class StrategyArguments: strategy_name: Literal[ - "deepspeed_train", "hf_infer", - "deepspeed_infer", "vllm", "sglang", "megatron_infer", "megatron_train", - "diffusion_deepspeed_train", "fsdp2_train", "fsdp2_infer", ] = field( - default="deepspeed_train", + default="fsdp2_train", metadata={ - "help": "The name of the strategy. Options: 'deepspeed_train', 'diffusion_deepspeed_train', 'hf_infer', 'deepspeed_infer', 'vllm', 'sglang', " + "help": "The name of the strategy. Options: 'hf_infer', 'vllm', 'sglang', " "'megatron_infer', 'megatron_train', 'fsdp2_train', 'fsdp2_infer'." }, ) @@ -205,12 +213,20 @@ class WorkerConfig: logits_in_fp32: bool = field( - default=True, + default=False, metadata={ "help": "Force logits dtype to Float" } ) + # Router Replay Configuration + router_replay: RouterReplayConfig = field( + default_factory=RouterReplayConfig, + metadata={ + "help": "Configuration for router replay in training. Only supported for Megatron strategy." + } + ) + apply_loss_scale: bool = field( default=True, metadata={ @@ -225,8 +241,39 @@ class WorkerConfig: } ) - def __post_init__(self): + # MTP training configuration + mtp_training_mode: Optional[Literal["disabled", "standalone", "joint"]] = field( + default="disabled", + metadata={ + "help": "MTP training mode for this worker. " + "'disabled': MTP is loaded but not trained (default). " + "'standalone': MTP is trained independently with truncated gradients (no gradient flow to main model). " + "'joint': MTP participates in main model updates with full gradient flow." + }, + ) + + # OPD (On-Policy Distillation) KL coefficient for teacher workers + opd_kl_coef: Optional[float] = field( + default=None, + metadata={ + "help": "KL coefficient for this teacher in OPD mode. " + "Defaults to 1.0 if not set." + }, + ) + # Tags this teacher handles for multi-teacher OPD routing + tag_included: List[str] = field( + default_factory=list, + metadata={ + "help": "Tags this teacher handles in multi-teacher OPD mode. " + "Empty list means this teacher handles all tags (default). " + "Used to route data to specific teachers based on domain/tag." + }, + ) + + def __post_init__(self): + self.offload_nccl = False + logger.info(f"force set offload_nccl=False.") if self.strategy_args is not None: if self.strategy_args.strategy_name not in ["hf_infer", "vllm", "sglang"] and self.num_gpus_per_worker > 1: logger.info( @@ -244,6 +291,23 @@ def __post_init__(self): f"pipeline_parallel_size: {pipeline_parallel_size}" ) + # Validate router_replay configuration + if self.router_replay.mode == "R2": + raise NotImplementedError("Not support R2 now") + if self.strategy_args.strategy_name not in ["megatron_train", "megatron_infer"]: + logger.warning( + f"router_replay [R2] is only supported for megatron_train and megatron_infer strategy, " + f"but current strategy is {self.strategy_args.strategy_name}. " + f"router_replay will be ignored." + ) + elif self.router_replay.mode == "R3": + if self.strategy_args.strategy_name not in ["megatron_train", "sglang"]: + logger.warning( + f"router_replay [R3] is only supported for megatron_train and sgalng strategy, " + f"but current strategy is {self.strategy_args.strategy_name}. " + f"router_replay will be ignored." + ) + if self.device_mapping is not None: self.device_mapping = eval(self.device_mapping) assert ( @@ -268,7 +332,6 @@ def __post_init__(self): self.training_args.fp16 = True - def is_actor_infer_overlapping_with_any_cluster(actor_infer: WorkerConfig, actor_train: WorkerConfig = None, reference: WorkerConfig = None, critic: WorkerConfig = None) -> bool: """ Check if actor_infer overlaps with ANY of the provided clusters. diff --git a/roll/datasets/collator.py b/roll/datasets/collator.py index 47e3f6bb5..1e57d106e 100644 --- a/roll/datasets/collator.py +++ b/roll/datasets/collator.py @@ -1,14 +1,23 @@ +import os +import re import inspect from collections import defaultdict from dataclasses import dataclass, field +from functools import partial from typing import Any, Dict, List, Optional, Union import numpy as np +import PIL import torch from transformers import BatchFeature, PreTrainedTokenizerBase, ProcessorMixin from transformers.data.data_collator import pad_without_fast_tokenizer_warning from transformers.utils import PaddingStrategy +from roll.utils.logging import get_logger + + +logger = get_logger() + def collate_fn_to_dict_list(data_list: list[dict]) -> dict: """将list[dict]数据转成dict[list]""" @@ -110,32 +119,273 @@ def __call__(self, features: List[Dict[str, Any]]) -> Dict[str, Any]: return batch +# from transformers 4.57.0 +def is_valid_video_frame(frame): + # processor should load videos from image paths if frame is a path + return isinstance(frame, PIL.Image.Image) or ( + (isinstance(frame, np.ndarray) or isinstance(frame, torch.Tensor)) and frame.ndim == 3 + ) + +def is_valid_video(video): + if not isinstance(video, (list, tuple)): + return (isinstance(video, np.ndarray) or isinstance(video, torch.Tensor)) and video.ndim == 4 + return video and all(is_valid_video_frame(frame) for frame in video) + +is_valid_image = is_valid_video_frame + +def convert_pil_frames_to_video(videos) : + """ + Given a batch of videos, converts each video to a 4D array. If video is already in array type, + it is simply returned. We assume that all inputs in the list are in the same format, based on the type of the first element. + + Args: + videos (`VideoInput`): + Video inputs to turn into a list of videos. + """ + + if not (isinstance(videos[0], (list, tuple)) and is_valid_image(videos[0][0])): + return videos + + video_converted = [] + for video in videos: + video = [np.array(frame) for frame in video] + video = np.stack(video) + video_converted.append(video) + return video_converted + + +def make_batched_videos(videos): + # Early exit for deeply nested list of image frame paths. We shouldn't flatten them + try: + if isinstance(videos[0][0], list) and isinstance(videos[0][0][0], str): + return [image_paths for sublist in videos for image_paths in sublist] + except (IndexError, TypeError): + pass + + if isinstance(videos, str) or is_valid_video(videos): + return convert_pil_frames_to_video([videos]) + # only one frame passed, thus we unsqueeze time dim + elif is_valid_image(videos): + if isinstance(videos, PIL.Image.Image): + videos = np.array(videos) + return [videos[None, ...]] + elif not isinstance(videos, list): + raise ValueError( + f"Invalid video input. Expected either a list of video frames or an input of 4 or 5 dimensions, but got" + f" type {type(videos)}." + ) + + # Recursively flatten any nested structure + flat_videos_list = [] + for item in videos: + if isinstance(item, str) or is_valid_video(item): + flat_videos_list.append(item) + elif isinstance(item, list) and item: + flat_videos_list.extend(make_batched_videos(item)) + + flat_videos_list = convert_pil_frames_to_video(flat_videos_list) + return flat_videos_list + + +def make_batched_metadata(videos, video_metadata: dict): + if video_metadata is None: + video_metadata = [{} for video in videos] + + if isinstance(video_metadata, list): + # Flatten if nested list + if isinstance(video_metadata[0], list): + video_metadata = [dict(**metadata) for metadata_list in video_metadata for metadata in metadata_list] + # Simply wrap in VideoMetadata if simple dict + elif isinstance(video_metadata[0], dict): + video_metadata = [dict(**metadata) for metadata in video_metadata] + else: + # Create a batched list from single object, differ with hf's single element list + video_metadata = [dict(**video_metadata)] * len(videos) + return video_metadata + + +def load_videos(videos, video_metas, **kwargs): + # NOTE: qwen3-vl has different usage with qwen2.5-vl/qwen3-omni due to the new video processor + # refer to https://github.com/QwenLM/Qwen3-VL?tab=readme-ov-file#new-qwen-vl-utils-usage + # qwen2.5-vl/qwen3-omni get sampled fps from kwargs + # qwen3-vl get raw video fps from video_processor returned video_metadata + # from qwen_vl_utils import fetch_video + # use patched version to fetch video + from roll.utils.qwen_vl_utils import fetch_video + + sampled_videos, sampled_kwargs = [], [] + for video, video_meta in zip(videos, video_metas): + video = dict(video=video, **kwargs) + video.update(video_meta) + video_inputs, sample_fps = ( + fetch_video( + video, + image_patch_size=kwargs["image_patch_size"], + return_video_sample_fps=True, + return_video_metadata=True, + video_resize_chunk_frames=kwargs.get("video_resize_chunk_frames", 64), + ) + if "image_patch_size" in kwargs + else fetch_video(video, return_video_sample_fps=True, return_video_metadata=True, + video_resize_chunk_frames=kwargs.get("video_resize_chunk_frames", 64),) + ) + video, video_meta = video_inputs + sampled_videos.append(video) + sampled_kwargs.append({"fps": sample_fps, "video_metadata": video_meta}) + if "use_audio_in_video" in video_meta: # use sample level use_audio_in_video if provided + sampled_kwargs[-1]["use_audio_in_video"] = video_meta["use_audio_in_video"] + return sampled_videos, sampled_kwargs + + +def load_audios(audios, audio_metas, is_video: bool = False, **kwargs): + # refer to audio_process.py in qwen_omni_utils, maybe use process_audio_info directly + import av + import librosa + + def _check_if_video_has_audio(video_path): + container = av.open(video_path) + audio_streams = [stream for stream in container.streams if stream.type == "audio"] + if not audio_streams: + return False + return True + + sample_rate = kwargs.get("sample_rate", 16000) + sampled_audios, sampled_kwargs = [], [] + for audio, audio_meta in zip(audios, audio_metas): + assert not is_video or _check_if_video_has_audio(audio), ( + "Video must has audio track when use_audio_in_video=True" + ) + audio_start = audio_meta.get("audio_start", audio_meta.get("video_start", 0.0)) + audio_end = audio_meta.get("audio_end", audio_meta.get("video_end", None)) + sampled_audios.append( + librosa.load( + audio, + sr=sample_rate, + offset=audio_start, + duration=(audio_end - audio_start) if audio_end is not None else None, + )[0] + ) + return sampled_audios, sampled_kwargs + + +def load_images(images, image_metas, **kwargs): + from qwen_vl_utils import fetch_image + + sampled_images, sampled_kwargs = [], [] + for image, image_meta in zip(images, image_metas): + image = dict(image=image, **kwargs) + image.update(image_meta) + sampled_images.append(fetch_image(image, image_patch_size=kwargs["image_patch_size"])) + return sampled_images, sampled_kwargs + + @dataclass class DataCollatorWithPaddingForMM: tokenizer: Optional[PreTrainedTokenizerBase] = None processor: Optional[ProcessorMixin] = None + # if key exists in video_kwargs returned by video_load_sample_fn, get value from video_kwargs + processor_mm_kwargs: Optional[Dict] = field(default_factory=dict) extra_data_provider: Optional[callable] = None prompt_key: str = "prompt" + is_template_applied: bool = True answer_key: Optional[str] = "ground_truth" image_key: Optional[str] = "image" + # meta keys can include sample level fields such as images with different folders + image_meta_keys: List[str] = field(default_factory=lambda: []) + image_load_sample_fn: Optional[callable] = None + image_sample_kwargs: Optional[Dict] = field(default_factory=dict) image_flag_key: Optional[str] = "image_flag" video_key: Optional[str] = "video" + video_meta_keys: List[str] = field(default_factory=lambda: []) video_flag_key: Optional[str] = "video_flag" + # load and sample, return (video, video_kwargs) + video_load_sample_fn: Optional[callable] = None + video_sample_kwargs: Optional[Dict] = field(default_factory=dict) + audio_key: Optional[str] = "audio" + audio_meta_keys: List[str] = field(default_factory=lambda: []) + audio_flag_key: Optional[str] = "audio_flag" + audio_load_sample_fn: Optional[callable] = None + audio_sample_kwargs: Optional[Dict] = field(default_factory=dict) + # use_audio_in_video defult to False, and whether or not to use can be + # overrided by sample level value. When using sample level, `video_load_sample_fn` + # should return sampled_kwargs including use_audio_in_video + use_audio_in_video: bool = False image_placeholder: Optional[str] = None image_token: str = "<|vision_start|><|image_pad|><|vision_end|>" + image_folder: Optional[str] = None video_placeholder: Optional[str] = None video_token: str = "<|vision_start|><|video_pad|><|vision_end|>" + video_folder: Optional[str] = None + audio_placeholder: Optional[str] = None + audio_token: str = "<|audio_start|><|audio_pad|><|audio_end|>" + audio_folder: Optional[str] = None padding: Union[bool, str, PaddingStrategy] = True max_length: Optional[int] = None pad_to_multiple_of: Optional[int] = None + ### keys for outputs of collator + # padded fields from processor outputs, mainly for llm input fields, maybe rename later padded_keys: List[str] = field(default_factory=lambda: ["input_ids", "attention_mask", "labels"]) + # unused fields from processor outputs + processor_unused_keys: List[str] = field(default_factory=lambda: ["prompt", "position_ids", "rope_deltas"]) + # unpaded fields from feature extra_unpadded_keys: List[str] = field(default_factory=lambda: []) return_tensors: str = "pt" return_infer_inputs: bool = True # whether to include infer engine inputs which differs with train + return_train_inputs: bool = True # maybe set to False in multi-turn rollout to reduce overhead + # use lower precision for multi-modal feature values to reduce store and transfer overhead by ray + mm_feature_dtype: Optional[str] = None + mm_feature_names: List[str] = field( + default_factory=lambda: ["pixel_values", "pixel_values_videos", "input_features"] # image, video, audio + ) + + def __post_init__(self): + if self.video_load_sample_fn is None: + self.video_load_sample_fn = partial( + load_videos, image_patch_size=self.processor.image_processor.patch_size + ) + self._default_video_processor_mm_kwargs = { + # sample and resize have been done in default video_load_sample_fn + "do_sample_frames": False, # avoid duplicate operation in processor, qwen3-vl + "do_resize": False, # avoid duplicate operation in processor + # qwen2.5-vl/qwen3-omni use sampled fps as processor kwargs for both hf and vllm, + # qwen3-vl use raw video fps from video_meta + "fps": None, + "video_metadata": None, + } + if self.image_load_sample_fn is None: + self.image_load_sample_fn = partial( + load_images, image_patch_size=self.processor.image_processor.patch_size + ) + # avoid duplicate operation in processor, while it might conflict with video do_resize + # if one does and the other does not + self._default_image_processor_mm_kwargs = {"do_resize": False} + if self.audio_load_sample_fn is None: + self.audio_load_sample_fn = load_audios + self._default_audio_processor_mm_kwargs = {} + if self.processor.__class__.__name__.startswith("Qwen3Omni"): + self.processor_mm_kwargs.update({"use_audio_in_video": self.use_audio_in_video}) + # refer to https://github.com/vllm-project/vllm/issues/26630 + from packaging.version import Version + from transformers import __version__ as TRANSFORMERS_VERSION + + if Version(TRANSFORMERS_VERSION) < Version("4.58.0") and "truncation" not in self.processor_mm_kwargs: + self.processor_mm_kwargs["truncation"] = False + else: + if "use_audio_in_video" in self.processor_mm_kwargs: + logger.warning(f"{self.processor.__class__.__name__} not support use_audio_in_video, thus remove it here") + self.processor_mm_kwargs.pop("use_audio_in_video", None) + + if self.mm_feature_dtype: + self.mm_feature_dtype = ( + torch.float32 + if self.mm_feature_dtype == "fp32" + else torch.float16 + if self.mm_feature_dtype == "fp16" + else torch.bfloat16 + ) def __call__(self, features: List[Dict[str, Any]]) -> Dict[str, Any]: assert self.tokenizer and self.processor - # model_inputs for hf/deepspeed: input_id, attention_mask, pixel_values, image_grid_thw padded_features = defaultdict(list) un_padded_features = defaultdict(list) mm_feature_keys = set() @@ -147,52 +397,271 @@ def __call__(self, features: List[Dict[str, Any]]) -> Dict[str, Any]: prompt = feature[self.prompt_key] if not isinstance(prompt, str): prompt = self.processor.apply_chat_template(prompt, tokenize=False, add_generation_prompt=True) - if self.image_placeholder: - prompt = prompt.replace(self.image_placeholder, self.image_token) - if self.video_placeholder: - prompt = prompt.replace(self.video_placeholder, self.video_token) - # TODO: support video - model_inputs: BatchFeature = self.processor( - images=feature[self.image_key] - if self.image_key and (not self.image_flag_key or feature[self.image_flag_key]) - else None, - text=prompt, + elif not self.is_template_applied: + prompt = [{"role": "user", "content": prompt}] + prompt = self.processor.apply_chat_template(prompt, tokenize=False, add_generation_prompt=True) + valid_image = ( + True + if self.image_key + and feature.get(self.image_key, None) + and (not self.image_flag_key or feature.get(self.image_flag_key, True)) + else False + ) + valid_video = ( + True + if self.video_key + and feature.get(self.video_key, None) + and (not self.video_flag_key or feature.get(self.video_flag_key, True)) + else False ) - if not isinstance(model_inputs, BatchFeature): - model_inputs = BatchFeature(data=model_inputs) - for key in ["prompt", "position_ids", "rope_deltas"]: # remove unnecessary feature - if key in model_inputs: - model_inputs.pop(key) - for key in filter(lambda k: k in model_inputs, self.padded_keys): - padded_features[key].append(model_inputs.pop(key)[0]) - # mm feature fileds can be different because of mixed data - mm_feature_keys = mm_feature_keys.union(model_inputs.keys()) - # to tensors except padded_keys which would be converted after padding - model_inputs.convert_to_tensors(tensor_type=self.return_tensors) - if self.image_key: + valid_audio = ( + True + if self.audio_key + and feature.get(self.audio_key, None) + and (not self.audio_flag_key or feature.get(self.audio_flag_key, True)) + else False + ) + processor_kwargs = {} + images = [] + image_kwargs = [] + if valid_image: + images = ( + feature[self.image_key] if isinstance(feature[self.image_key], list) else [feature[self.image_key]] + ) + image_metas = make_batched_metadata( + images, dict((key, feature[key]) for key in self.image_meta_keys if key in feature) + ) + if self.image_placeholder: + prompt = prompt.replace(self.image_placeholder, self.image_token) + if self.image_folder and isinstance(images[0], str): + images = [os.path.join(self.image_folder, image) for image in images] + if self.image_load_sample_fn: + images, image_kwargs = self.image_load_sample_fn(images, image_metas, **self.image_sample_kwargs) + processor_kwargs.update(self._default_image_processor_mm_kwargs) + audios = [] + audio_kwargs = [] + if valid_audio: + audios = ( + feature[self.audio_key] if isinstance(feature[self.audio_key], list) else [feature[self.audio_key]] + ) + audio_metas = make_batched_metadata( + audios, dict((key, feature[key]) for key in self.audio_meta_keys if key in feature) + ) + if self.audio_placeholder: + # use_audio_in_video use video token in text + prompt = prompt.replace(self.audio_placeholder, self.audio_token) + if self.audio_folder and isinstance(audios[0], str): + audios = [os.path.join(self.audio_folder, audio) for audio in audios] + if self.audio_load_sample_fn: + audios, audio_kwargs = self.audio_load_sample_fn(audios, audio_metas, **self.audio_sample_kwargs) + processor_kwargs.update(self._default_audio_processor_mm_kwargs) + videos = [] + video_kwargs = [] + if valid_video: + videos = make_batched_videos([feature[self.video_key]]) + video_metas = make_batched_metadata( + videos, dict((key, feature[key]) for key in self.video_meta_keys if key in feature) + ) + if self.video_placeholder: + prompt = prompt.replace(self.video_placeholder, self.video_token) + # video path or video with frame paths + videos_for_audio = videos + if isinstance(videos[0], str) or (isinstance(videos[0], list) and isinstance(videos[0][0], str)): + if self.video_folder: + videos_for_audio = videos = [ + ( + os.path.join(self.video_folder, video) + if isinstance(video, str) + else [os.path.join(self.video_folder, frame) for frame in video] + ) + for video in videos + ] + # processor should load and sample if video_load_sample_fn is not provided + if self.video_load_sample_fn: + # video_kwargs stands for kwargs might be used in processor + videos, video_kwargs = self.video_load_sample_fn( + videos, video_metas, **self.video_sample_kwargs + ) + processor_kwargs.update(self._default_video_processor_mm_kwargs) + processor_kwargs.update(self.processor_mm_kwargs) + if image_kwargs: + processor_kwargs.update( + dict( + ( + key, + [kwargs[key] for kwargs in image_kwargs] + if key in image_kwargs[0] + else processor_kwargs[key], + ) + if isinstance(image_kwargs, list) and image_kwargs + else ( + key, + image_kwargs[key] + if isinstance(image_kwargs, dict) and key in image_kwargs + else processor_kwargs[key], + ) + for key in processor_kwargs + ) + ) + # video agnostic single value can always be provided in self.processor_mm_kwargs + # otherwise the value should have same length as videos + if video_kwargs: + processor_kwargs.update( + dict( + ( + key, + [kwargs[key] for kwargs in video_kwargs] + if key in video_kwargs[0] + else processor_kwargs[key], + ) + if isinstance(video_kwargs, list) and video_kwargs + else ( + key, + video_kwargs[key] + if isinstance(video_kwargs, dict) and key in video_kwargs + else processor_kwargs[key], + ) + for key in processor_kwargs + ) + ) + # compatibility for qwen2.5-vl/qwen3-vl/qwen3-omni, subject to change + # NOTE: qwen3-omni processor only supports single value fps, hack to work around temporarily + if ( + self.processor.__class__.__name__.startswith("Qwen3Omni") + and "fps" in processor_kwargs + and isinstance(processor_kwargs["fps"], list) + ): + assert all(fps == processor_kwargs["fps"][0] for fps in processor_kwargs["fps"]), ( + f"{self.processor.__class__} only support single value fps currently" + ) + processor_kwargs["fps"] = processor_kwargs["fps"][0] + + # use_audio_in_video should be single value + if "use_audio_in_video" in processor_kwargs and isinstance(processor_kwargs["use_audio_in_video"], list): + assert all( + use_audio_in_video == processor_kwargs["use_audio_in_video"][0] + for use_audio_in_video in processor_kwargs["use_audio_in_video"] + ), "only support same use_audio_in_video value for videos in one sample" + processor_kwargs["use_audio_in_video"] = processor_kwargs["use_audio_in_video"][0] + # sample level use_audio_in_video got from processor_kwargs + if processor_kwargs.get("use_audio_in_video", self.use_audio_in_video) and valid_video: + assert self.audio_load_sample_fn and isinstance(videos_for_audio[0], str) + video_audios, video_audio_kwargs = self.audio_load_sample_fn( + videos_for_audio, video_metas, is_video=True, **self.audio_sample_kwargs + ) + if valid_audio: + merged_audios = [] + merged_audio_kwargs = [] + # when use_audio_in_video, order of audios and audios from videos make effect and + # should be consistent with the order of multi-modal tokens in text, similar with + # the logic in qwen3-omni processor + # NOTE: thre is still a not supported case: when videos w/ and w/o audios mixed in + # the same sample, processor cannot handle it since no mapping between videos and + # audios can be passed to processor + special_tokens = [ + re.escape(tok) for tok in [self.processor.audio_token, self.processor.video_token] + ] + pattern = "|".join(special_tokens) + positions = sorted([(match.start(), match.group()) for match in re.finditer(pattern, prompt)]) + audio_index = video_index = 0 + for _, special_token in positions: + if special_token == self.processor.audio_token: + merged_audios.append(audios[audio_index]) + merged_audio_kwargs.append(audio_kwargs[audio_index]) + audio_index += 1 + else: + merged_audios.append(video_audios[audio_index]) + merged_audio_kwargs.append(video_audio_kwargs[audio_index]) + video_index += 1 + audios, audio_kwargs = merged_audios, merged_audio_kwargs + else: + audios.extend(video_audios) + audio_kwargs.extend(video_audio_kwargs) + + if audio_kwargs: + processor_kwargs.update( + dict( + ( + key, + [kwargs[key] for kwargs in audio_kwargs] + if key in audio_kwargs[0] + else processor_kwargs[key], + ) + if isinstance(audio_kwargs, list) and audio_kwargs + else ( + key, + audio_kwargs[key] + if isinstance(audio_kwargs, dict) and key in audio_kwargs + else processor_kwargs[key], + ) + for key in processor_kwargs + ) + ) + + # IndexError occurs in processor when using empty list + images = images if images else None + videos = videos if videos else None + audios = audios if audios else None + + if self.return_train_inputs: + # model_inputs are mainly for train engine + model_inputs: BatchFeature = self.processor( + images=images, videos=videos, audio=audios, text=prompt, **processor_kwargs + ) + if not isinstance(model_inputs, BatchFeature): + model_inputs = BatchFeature(data=model_inputs) + # TODO: maybe use processor produced position_ids + for key in self.processor_unused_keys: + if key in model_inputs: + model_inputs.pop(key) + for key in filter(lambda k: k in model_inputs, self.padded_keys): + padded_features[key].append(model_inputs.pop(key)[0]) + # mm feature fileds can be different because of mixed data + mm_feature_keys = mm_feature_keys.union(model_inputs.keys()) + # to tensors except padded_keys which would be converted after padding + model_inputs.convert_to_tensors(tensor_type=self.return_tensors) + if self.mm_feature_dtype: + for key in [name for name in model_inputs.keys() if name in self.mm_feature_names]: + model_inputs[key] = model_inputs[key].to(self.mm_feature_dtype) # allow mixed text and multi-modal data # assert model_inputs, "should have multi-modal features" # tensors in multi_modal_inputs dict have bsz=1 and should be # concat at dim=0 before model forward un_padded_features["multi_modal_inputs"].append(dict(model_inputs)) - # inputs for infer engine, not tensors - if self.return_infer_inputs: - un_padded_features["multi_modal_data"].append( - { - "prompt_token_ids": # different with input_ids - self.tokenizer.encode(prompt, add_special_tokens=False), - "multi_modal_data": { - "image": [feature[self.image_key]] - if not isinstance(feature[self.image_key], list) - else feature[self.image_key] - }, - } - if (not self.image_flag_key or feature[self.image_flag_key]) and feature[self.image_key] - else { - "prompt_token_ids": # different with input_ids - self.tokenizer.encode(prompt, add_special_tokens=False), - } - ) + + # inputs for infer engine, not tensors + # TODO: maybe rename multi_modal_data as multi_modal_infer_inputs + if self.return_infer_inputs: + # only support vllm as infer engine currently + un_padded_features["multi_modal_data"].append( + { + "prompt_token_ids": # different with input_ids + self.tokenizer.encode(prompt, add_special_tokens=False), + } + ) + multi_modal_data = {} + if images: + multi_modal_data["image"] = images + if audios: + multi_modal_data["audio"] = audios + if videos: + # compatibility for qwen2.5-vl/qwen3-vl/qwen3-omni, subject to change + # NOTE: video_mata is used as kwargs in hf while it is put into video for qwen3-vl in vllm==0.11.1, + # see: https://github.com/QwenLM/Qwen3-VL?tab=readme-ov-file#offline-inference + # hash error occurs for video_metadata when used as mm_kwargs in vllm==0.11.1 + # avoid vllm version incompatibility for video meta and only use it for qwen3-vl + # since video meta only after https://github.com/vllm-project/vllm/pull/19331 + video_metas = processor_kwargs.pop("video_metadata", video_metas) + if self.processor.__class__.__name__.startswith("Qwen3VL") and isinstance(video_metas, list): + # vllm gets video num using `n = len(data) if isinstance(data, list) else 1` for mm_uuids + # in `_maybe_build_mm_uuids` and gets video and video_meta from tuple in `_get_video_with_metadata` + videos = list(zip(videos, video_metas)) + multi_modal_data["video"] = videos + if multi_modal_data: + un_padded_features["multi_modal_data"][-1]["multi_modal_data"] = multi_modal_data + # vllm use mm_processor_kwargs to call processor and select processor output fileds by model defination + un_padded_features["multi_modal_data"][-1]["mm_processor_kwargs"] = processor_kwargs + if self.answer_key: un_padded_features[self.answer_key].append(feature[self.answer_key]) if self.extra_unpadded_keys: diff --git a/roll/datasets/dataset.py b/roll/datasets/dataset.py index 32097af0b..dfdf3ac04 100644 --- a/roll/datasets/dataset.py +++ b/roll/datasets/dataset.py @@ -30,7 +30,7 @@ def decorator(func: Callable[[list[str], str, dict], Union[Dataset, IterableData return decorator -def get_dataset(data_args: "DataArguments"): +def get_dataset(data_args: "DataArguments", **kwargs): # TODO: refactor get_dataset and create_local_dataset data_path = None data_name = data_args.file_name @@ -41,7 +41,7 @@ def get_dataset(data_args: "DataArguments"): local_path = "" else: local_path: str = os.path.join(dataset_dir, data_name) - if dataset_type in ("odps",): + if dataset_type in ("odps", "ailake"): data_path = dataset_type data_files.extend(data_name) elif os.path.isdir(local_path): @@ -71,7 +71,9 @@ def get_dataset(data_args: "DataArguments"): logger.info(f"load_data_files: {chr(10)} {chr(10).join(data_files)}") logger.info(f"prompt column: {data_args.prompt} label column: {data_args.response}") - return REGISTERED_DATASETS[data_path](data_files, split="train") + dataset_kwargs = {"num_proc": data_args.preprocessing_num_workers} + dataset_kwargs.update(kwargs) + return REGISTERED_DATASETS[data_path](data_files, split="train", **dataset_kwargs) def create_local_dataset( diff --git a/roll/datasets/vlm_dataset_utils.py b/roll/datasets/vlm_dataset_utils.py new file mode 100644 index 000000000..47ede10c8 --- /dev/null +++ b/roll/datasets/vlm_dataset_utils.py @@ -0,0 +1,416 @@ +import os +from io import BytesIO +from typing import List, Tuple, Union + +import datasets +import PIL.Image as Image +from datasets import load_from_disk +from transformers import ProcessorMixin +from transformers.models.qwen2_vl.image_processing_qwen2_vl import smart_resize +from transformers.image_utils import load_images +from roll.datasets.dataset import get_dataset +from roll.utils.logging import get_logger +from roll.utils.import_utils import safe_import_class + + +logger = get_logger() + + +def create_pipeline_data_kwargs(data_args, tokenizer, processor, is_val=False): + data_kwargs_getter = getattr(data_args, "custom_data_kwargs_func") + if data_kwargs_getter is None: + data_kwargs_getter = get_vlm_data_kwargs + elif isinstance(data_kwargs_getter, str): + data_kwargs_getter = safe_import_class(data_kwargs_getter) + return data_kwargs_getter(data_args, tokenizer, processor, is_val=is_val) + + +def format_prompt(prompt, processor, use_image=True, prompt_image_token=None): + question_template = "{Question} Output the thinking process in and final answer (number) in tags." + if isinstance(prompt, list): + messages = prompt + else: + messages = [ + { + "role": "user", + "content": [ + {"type": "image"}, + {"type": "text", "text": question_template.format(Question=prompt)}, + ] + if use_image and not prompt_image_token + else [ + {"type": "text", "text": question_template.format(Question=prompt)} + ], # image_token has been included in prompt + } + ] + text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + if prompt_image_token: + text = text.replace(prompt_image_token, "<|vision_start|><|image_pad|><|vision_end|>") + return text + + +def process_image(image: Image.Image, processor: ProcessorMixin): + # same as qwen2-vl image processor + image_processor = processor.image_processor + factor = ( + image_processor.patch_size * image_processor.merge_size + if "Qwen" in image_processor.image_processor_type + else 28 + ) + height, width = image.height, image.width + resized_height, resized_width = smart_resize( + height, + width, + factor=factor, + # Qwen2VLImageProcessorFast uses size["shortest_edge"]/size["longest_edge"] instead of min_pixels/max_pixels + # thus set min_pixels/max_pixels attrs before using min_pixels/max_pixels + min_pixels=image_processor.min_pixels, + max_pixels=image_processor.max_pixels, + ) + resized_image = image.resize((resized_width, resized_height), resample=image_processor.resample) + return resized_image + + +def process_images( + images: Union[List, Tuple, str, Image.Image], processor: ProcessorMixin +) -> Union[Image.Image, List[Image.Image], List[List[Image.Image]]]: + """Process images, handling different levels of nesting. + + Args: + images: A single image, a list of images, or a list of lists of images to load. + timeout: Timeout for loading images. + + Returns: + A single image, a list of images, a list of lists of images. + """ + if isinstance(images, (list, tuple)): + if len(images) and isinstance(images[0], (list, tuple)): + return [[process_image(image, processor=processor) for image in image_group] for image_group in images] + else: + return [process_image(image, processor=processor) for image in images] + else: + return process_image(images, processor=processor) + + +def encode_function( + data, processor, prompt_getter, ground_truth_getter, image_getter, tag_getter, prompt_image_token=None +): + image_flag = [True] * len(prompt_getter(data)) + image_list = [] + for idx, image in enumerate(image_getter(data)): + if not image: + image_flag[idx] = False + try: + if isinstance(image, bytes): # bytes data + # TODO: support multiple images + image_out = Image.open(BytesIO(image)) + else: + image_out = load_images(image if isinstance(image, (list, tuple)) else [image], timeout=None) + except Exception as e: + if isinstance(image, bytes): + image_out = [Image.new("RGB", (224, 224), (255, 255, 255))] + logger.error(f"Failed to get image with type: {type(image)}") + else: + image_out = [Image.new("RGB", (224, 224), (255, 255, 255))] * len(image) + logger.error(f"Failed to get image: {image}") + # since infer-image use pil image as input while train-engine use + # processed data, process image here to make them use same image + # refer to the following for Spatial Understanding with Qwen2.5-VL + # https://github.com/QwenLM/Qwen2.5-VL/blob/main/cookbooks/spatial_understanding.ipynb + # NOTE: process_image from qwen2.5-vl keeps aspect ratio almostly and + # bboxes would be normalized in detection verifier, thus nearly no need + # to change ground-truth bboxes + # process in collator, no need to process here + # image_out = process_images(image_out, processor) + image_list.append(image_out) + text_list = [] + for idx, instruct in enumerate(prompt_getter(data)): + # provide prompt_image_token if image_token in prompt + text = format_prompt(instruct, processor, use_image=image_flag[idx], prompt_image_token=prompt_image_token) + text_list.append(text) + encodings = { + "tag": tag_getter(data), + "images": image_list, + "prompt": text_list, + "ground_truth": ground_truth_getter(data), + "reward_model": data["reward_model"], + # for text and multi-modal mixed data usage, indicating valid image + "image_flag": image_flag, + } + return encodings + + +def get_vlm_dataset(data_args, encode_function, processor, get_eval=False): + cache_path = getattr(data_args, "cache_path", None) + if cache_path: + cache_path = os.path.join(cache_path, "val" if get_eval else "train") + if cache_path and os.path.exists(cache_path): + dataset = load_from_disk(cache_path) + return dataset + + dataset = get_dataset(data_args=data_args) + # regularized data filed + features = datasets.Features( + { + "tag": datasets.Value(dtype="string"), # from data_source + "images": datasets.Sequence(feature=datasets.Image(mode=None, decode=True)), + "prompt": datasets.Value(dtype="string"), + "ground_truth": datasets.Value(dtype="string"), + "reward_model": dataset.features["reward_model"], + # for text and multi-modal mixed data usage, indicating valid image + "image_flag": datasets.Value("bool"), + } + ) + remove_columns = list(dataset.features.keys() - features.keys()) + # suit to both VLM-RL/Ocean-R1 and MiniMax-AI/One-RL-to-See-Them-All data + prompt_getter = lambda data: data["prompt"] + ground_truth_getter = lambda data: [x["ground_truth"] for x in data["reward_model"]] + image_getter = lambda data: data["images"] + tag_getter = lambda data: data["data_source"] + processor.image_processor.min_pixels, processor.image_processor.max_pixels = ( + data_args.image_min_pixels, + data_args.image_max_pixels, + ) + print(f"Begin : {dataset}") + dataset = dataset.map( + lambda data: encode_function( + data, processor, prompt_getter, ground_truth_getter, image_getter, tag_getter, prompt_image_token="" + ), + batched=True, + num_proc=data_args.preprocessing_num_workers, + features=features, + remove_columns=remove_columns, + desc="Encoding dataset", + ) + print(f"Encoding: {dataset}") + if cache_path: + dataset.save_to_disk(cache_path) + return dataset + + +def get_vlm_data_kwargs(data_args, tokenizer, processor, is_val=False): + dataset = get_vlm_dataset(data_args, encode_function, processor, get_eval=is_val) + collect_fn_kwargs = dict( + extra_unpadded_keys=["reward_model", "tag"], + prompt_key="prompt", + answer_key="ground_truth", + image_key="images", + image_flag_key="image_flag", + image_sample_kwargs={"min_pixels": data_args.image_min_pixels, "max_pixels": data_args.image_max_pixels}, + ) + return dict(dataset=dataset, collect_fn_kwargs=collect_fn_kwargs) + + +### for video rlvr demo +# we use video-r1 116k video data with adjusted format, case for training data: +# { +# "modality_type": "video", +# "problem": "Why does the video conclude with a red screen displaying text?Options:\nA. To promote the BBC one program 'Super Cute Animals'\nB. To provide a warning message\nC. To show the credits of the video\nD. To indicate a change in scene\n", +# "solution": "A", +# "video": "videos/youtube_video_2024/ytb_HBxn56l9WcU.mp4", +# "problem_type": "multiple choice" +# } +# reference Video-R1 +QUESTION_TEMPLATE = ( + "{Question}\n" + "Please think about this question as if you were a human pondering deeply. " + "Engage in an internal dialogue using expressions such as 'let me think', 'wait', 'Hmm', 'oh, I see', 'let's break it down', etc, or other natural language thought expressions " + "It's encouraged to include self-reflection or verification in the reasoning process. " + "Provide your detailed reasoning between the tags, and then give your final answer between the tags." +) + +TYPE_TEMPLATE = { + "multiple choice": " Please provide only the single option letter (e.g., A, B, C, D, etc.) within the tags.", + "numerical": " Please provide the numerical value (e.g., 42 or 3.14) within the tags.", + "OCR": " Please transcribe text from the image/video clearly and provide your text answer within the tags.", + "free-form": " Please provide your text answer within the tags.", + "regression": " Please provide the numerical value (e.g., 42 or 3.14) within the tags.", +} + + +def video_r1_format_prompt(prompt, processor, problem_type, modality_type=None): + if isinstance(prompt, list): + messages = prompt + else: + messages = [ + { + "role": "user", + "content": [ + {"type": modality_type}, + {"type": "text", "text": QUESTION_TEMPLATE.format(Question=prompt) + TYPE_TEMPLATE[problem_type]}, + ] + if modality_type + else [ + {"type": "text", "text": QUESTION_TEMPLATE.format(Question=prompt) + TYPE_TEMPLATE[problem_type]} + ] + + TYPE_TEMPLATE[problem_type], + } + ] + text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + return text + + +def video_r1_encode_function(data, processor): + encodings = { + "prompt": [ + video_r1_format_prompt(problem, processor, problem_type, modality_type) + for problem, problem_type, modality_type in zip( + data["problem"], data["problem_type"], data["modality_type"] + ) + ], + "video": data["video"], + "ground_truth": data["solution"], + "reward_model": data["problem_type"], + "tag": data["modality_type"], + } + if "fps" in data: + encodings["fps"] = data["fps"] + return encodings + + +def video_r1_get_dataset(data_args, encode_function, processor, get_eval=False): + cache_path = getattr(data_args, "cache_path", None) + if cache_path: + cache_path = os.path.join(cache_path, "val" if get_eval else "train") + if cache_path and os.path.exists(cache_path): + dataset = load_from_disk(cache_path) + return dataset + dataset = get_dataset(data_args=data_args) + print(f"Begin : {dataset}") + dataset = dataset.map( + lambda data: encode_function(data, processor), + batched=True, + num_proc=data_args.preprocessing_num_workers, + desc="Encoding dataset", + ) + print(f"Encoding: {dataset}") + if cache_path: + dataset.save_to_disk(cache_path) + return dataset + + +def video_r1_get_data_kwargs(data_args, tokenizer, processor, is_val=False): + dataset = video_r1_get_dataset(data_args, video_r1_encode_function, processor, get_eval=is_val) + collect_fn_kwargs = dict( + extra_unpadded_keys=["reward_model", "tag"], + video_folder=data_args.video_folder, + video_sample_kwargs={ # settings same with RewatchR1 and VideoR1 + "total_pixels": data_args.video_total_pixels, + "min_pixels": data_args.video_min_pixels, + "max_pixels": data_args.video_max_pixels, + "fps": data_args.video_fps, + "max_frames": data_args.video_max_frames, + "video_resize_chunk_frames": data_args.video_resize_chunk_frames, + }, + video_meta_keys=[], + is_template_applied=True, + ) + get_data_item_kwargs = dict(use_dataloader=True, use_collect_fn=True, num_workers=4) + return dict(dataset=dataset, collect_fn_kwargs=collect_fn_kwargs, get_data_item_kwargs=get_data_item_kwargs) + + +def audio_test_format_prompt(prompt, processor, use_audio=True, prompt_audio_token=None): + question_template = "{Question} Output the thinking process in and final answer (number) in tags." + if isinstance(prompt, list): + messages = prompt + else: + messages = [ + { + "role": "user", + "content": [ + {"type": "audio"}, + {"type": "text", "text": question_template.format(Question=prompt)}, + ] + if use_audio and not prompt_audio_token + else [ + {"type": "text", "text": question_template.format(Question=prompt)} + ], # audio_token has been included in prompt + } + ] + text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + if prompt_audio_token: + text = text.replace(prompt_audio_token, "<|audio_start|><|audio_pad|><|audio_end|>") + return text + + +### for video rlvr demo +def audio_test_encode_function( + data, processor, prompt_getter, ground_truth_getter, audio_getter, tag_getter, prompt_audio_token=None +): + audio_list = audio_getter(data) + audio_flag = [True] * len(audio_list) + text_list = [] + for idx, instruct in enumerate(prompt_getter(data)): + # provide prompt_audio_token if audio_token in prompt + text = audio_test_format_prompt( + instruct, processor, use_audio=audio_flag[idx], prompt_audio_token=prompt_audio_token + ) + text_list.append(text) + encodings = { + "tag": tag_getter(data), + "audios": audio_list, + "prompt": text_list, + "ground_truth": ground_truth_getter(data), + "reward_model": data["reward_model"], + # for text and multi-modal mixed data usage, indicating valid audio + "audio_flag": audio_flag, + } + return encodings + + +def audio_test_get_dataset(data_args, encode_function, processor, get_eval=False): + cache_path = getattr(data_args, "cache_path", None) + if cache_path: + cache_path = os.path.join(cache_path, "val" if get_eval else "train") + if cache_path and os.path.exists(cache_path): + dataset = load_from_disk(cache_path) + return dataset + dataset = get_dataset(data_args=data_args) + # regularized data filed + features = datasets.Features( + { + "tag": datasets.Value(dtype="string"), # from data_source + "audios": datasets.Sequence(datasets.Value(dtype="string")), + "prompt": datasets.Value(dtype="string"), + "ground_truth": datasets.Value(dtype="string"), + "reward_model": dataset.features["reward_model"], + # for text and multi-modal mixed data usage, indicating valid audio + "audio_flag": datasets.Value("bool"), + } + ) + remove_columns = list(dataset.features.keys() - features.keys()) + prompt_getter = lambda data: data["prompt"] + ground_truth_getter = lambda data: [x["ground_truth"] for x in data["reward_model"]] + audio_getter = lambda data: data["audios"] + tag_getter = lambda data: data["data_source"] + print(f"Begin : {dataset}") + dataset = dataset.map( + lambda data: encode_function( + data, processor, prompt_getter, ground_truth_getter, audio_getter, tag_getter, prompt_audio_token=None + ), + batched=True, + num_proc=data_args.preprocessing_num_workers, + features=features, + remove_columns=remove_columns, + desc="Encoding dataset", + ) + print(f"Encoding: {dataset}") + if cache_path: + dataset.save_to_disk(cache_path) + return dataset + + +def audio_test_get_data_kwargs(data_args, tokenizer, processor, is_val=False): + dataset = audio_test_get_dataset(data_args, audio_test_encode_function, processor, get_eval=is_val) + collect_fn_kwargs = dict( + extra_unpadded_keys=["reward_model", "tag"], + prompt_key="prompt", + answer_key="ground_truth", + audio_key="audios", + audio_flag_key="audio_flag", + audio_folder=data_args.audio_folder, + audio_sample_kwargs={}, + audio_meta_keys=[], + is_template_applied=True, + ) + return dict(dataset=dataset, collect_fn_kwargs=collect_fn_kwargs) diff --git a/roll/distributed/executor/worker.py b/roll/distributed/executor/worker.py index d6e5fca5c..c0bf9aa04 100644 --- a/roll/distributed/executor/worker.py +++ b/roll/distributed/executor/worker.py @@ -18,6 +18,7 @@ from roll.utils.network_utils import collect_free_port, get_node_ip from roll.utils.offload_states import OffloadStateType from roll.utils.offload_nccl import monkey_patch_torch_dist +from roll.utils.telemetry import init_telemetry from roll.platforms import current_platform @@ -108,6 +109,16 @@ def get_free_port(): if success: return master_port raise RuntimeError(f"Can not allocate unique MASTER_PORT on {master_addr}.") + + @staticmethod + def is_port_available(port: int) -> bool: + try: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("", port)) + return True + except OSError: + return False + def get_master_addr_and_port(self): return self.master_addr, self.master_port @@ -128,6 +139,14 @@ def get_rank_info(self): def initialize(self, pipeline_config, *args, **kwargs): self.pipeline_config = pipeline_config + # Initialize OpenTelemetry tracing on the worker actor + if os.environ.get("ROLL_OTEL_ENABLED") == "1": + init_telemetry( + service_name=self.cluster_name, + otlp_endpoint=os.environ.get("OTEL_EXPORTER_OTLP_ENDPOINT"), + instance_id=self.worker_name, + ) + model_name = self.worker_config.model_args.model_name_or_path if model_name: self.worker_config.model_args.model_name_or_path = download_model(model_name) @@ -170,12 +189,6 @@ def setup_collective_group(self, *args, **kwargs): else: self.logger.warning("worker has not strategy") - def setup_p2p_collective_group(self, *args, **kwargs): - if getattr(self, "strategy", None) is not None: - self.strategy.setup_p2p_collective_group(*args, **kwargs) - else: - self.logger.warning("worker does not have a strategy") - def start_model_update(self, *args, **kwargs): metrics = {} if getattr(self, "strategy", None) is not None: @@ -194,12 +207,6 @@ def start_model_update(self, *args, **kwargs): output = DataProto(meta_info={"metrics": metrics}) return output - def model_update_set_read_done_handle(self, *args, **kwargs): - if getattr(self, "strategy", None) is not None: - self.strategy.model_update_set_read_done_handle(*args, **kwargs) - else: - self.logger.warning("worker has not strategy") - def update_parameter_in_bucket(self, *args, **kwargs): if getattr(self, "strategy", None) is not None: self.strategy.update_parameter_in_bucket(*args, **kwargs) diff --git a/roll/distributed/scheduler/async_generate_scheduler.py b/roll/distributed/scheduler/async_generate_scheduler.py deleted file mode 100644 index 9a854f857..000000000 --- a/roll/distributed/scheduler/async_generate_scheduler.py +++ /dev/null @@ -1,843 +0,0 @@ -import copy -import enum -import itertools -import math -import queue -import random -import threading -import time -from collections import defaultdict -from dataclasses import dataclass, field -from typing import Any, Dict, List, Optional, Set, Union - -import ray -import torch -from datasets import Dataset -from ray import ObjectRef -from torch.nn.utils.rnn import pad_sequence -from tqdm import tqdm -from transformers import set_seed - -from roll.distributed.executor.cluster import Cluster -from roll.distributed.scheduler.generate_scheduler import GlobalCounter -from roll.distributed.scheduler.protocol import DataProto -from roll.models.model_providers import default_tokenizer_provider -from roll.utils.constants import RAY_NAMESPACE -from roll.utils.functionals import ( - GenerateRequestType, - concatenate_input_and_output, - postprocess_generate, -) -from roll.utils.logging import get_logger - - -logger = get_logger() - - -class ExpStatus(enum.Enum): - RUNNING = enum.auto() - PAUSED = enum.auto() - FINISHED = enum.auto() - DELETED = enum.auto() - - -def is_report_data_finished(data: DataProto) -> bool: - finish_reasons = data.meta_info.get("finish_reasons", []) - # take no finish reason as finished - return not any([finish_reason is None for finish_reason in finish_reasons]) - - -@dataclass -class ExperienceItem: - request_id: str - prompt_id: int - domain: str = "default" - sampling_start_step: Optional[int] = None - running_dp_rank: Optional[int] = None # should be none when not running - data: Optional[DataProto] = None - status: ExpStatus = ExpStatus.RUNNING - - -@dataclass -class ItemsGroup: - # items with the same starting step - # item status index - start_step: int - by_request_id: Dict[str, ExperienceItem] = field(default_factory=dict) - request_ids_by_prompt_id: Dict[int, Set[str]] = field( - default_factory=lambda: defaultdict(set) - ) - running_request_ids: Set[str] = field(default_factory=set) - pause_request_ids: Set[str] = field(default_factory=set) - finished_ids_by_prompt_id: Dict[int, Set[str]] = field( - default_factory=lambda: defaultdict(set) - ) - # all requests for this prompt_id are finished, manually marked - finished_prompt_ids: Set[int] = field(default_factory=set) - - def info(self): - return (f"ItemsGroup {self.start_step}: {len(self.by_request_id)=}" - f"{len(self.request_ids_by_prompt_id)=} {len(self.running_request_ids)=}" - f"{len(self.pause_request_ids)=} {len(self.finished_prompt_ids)=}" - f"{len(self.finished_ids_by_prompt_id)=}") - - def add_request_item(self, request_item: ExperienceItem): - self.by_request_id[request_item.request_id] = request_item - self.request_ids_by_prompt_id[request_item.prompt_id].add(request_item.request_id) - request_item.status = ExpStatus.RUNNING - self.running_request_ids.add(request_item.request_id) - self.pause_request_ids.discard(request_item.request_id) - - def reset_status(self): - for request_id in list(self.running_request_ids): - item = self.by_request_id.get(request_id) - if item: - item.status = ExpStatus.PAUSED - self.pause_request_ids.add(request_id) - self.running_request_ids = set() - - def get_prompt_ids(self) -> Set[int]: - return set(self.request_ids_by_prompt_id.keys()) - - def pop_item(self, prompt_id: int) -> Optional[ExperienceItem]: - request_id = self.request_ids_by_prompt_id[prompt_id].pop() - return self.by_request_id.pop(request_id) - - def mark_prompt_id_finished(self, prompt_id: int, num_return_sequences: int): - finished_ids = self.finished_ids_by_prompt_id.get(prompt_id, set()) - - self.finished_prompt_ids.add(prompt_id) - - assert ( - sum(self.by_request_id.get(request_id).data.batch.batch_size[0] for request_id in finished_ids) - >= num_return_sequences - ), f"prompt_id {prompt_id} finished_ids {len(finished_ids)} num_return_sequences {num_return_sequences}" - all_request_ids = self.request_ids_by_prompt_id[prompt_id] - - # remove extra request_ids and return them - extra_request_ids = all_request_ids - finished_ids - extra_request_ids = extra_request_ids.union(finished_ids - set(list(finished_ids)[:num_return_sequences])) - extra_items = [] - for request_id in extra_request_ids: - item = self.by_request_id.get(request_id) - extra_items.append(item) - item.status = ExpStatus.DELETED - self.request_ids_by_prompt_id[prompt_id].discard(request_id) - - self.running_request_ids.discard(request_id) - self.pause_request_ids.discard(request_id) - self.finished_ids_by_prompt_id[prompt_id].discard(request_id) - return extra_items - - def pop_prompt_id(self, prompt_id: int): - running_items, finished_items = [], [] - request_ids = self.request_ids_by_prompt_id.pop(prompt_id, set()) - for request_id in request_ids: - item = self.by_request_id.get(request_id) - if item is None: - continue - if item.status == ExpStatus.FINISHED: - finished_items.append(item) - elif item.status == ExpStatus.RUNNING: - running_items.append(item) - item.status = ExpStatus.DELETED - self.pause_request_ids.discard(request_id) - self.running_request_ids.discard(request_id) - - self.finished_prompt_ids.discard(prompt_id) - self.finished_ids_by_prompt_id.pop(prompt_id, None) - return running_items, finished_items - - -@dataclass -class ReplayBuffer: - groups: Dict[int, ItemsGroup] = field(default_factory=dict) - prompt_id_to_start_step: Dict[int, int] = field(default_factory=dict) - - def info(self) -> str: - return ( - f"ReplayBuffer: " - f"by_request_id={sum(len(g.by_request_id) for g in self.groups.values())}, " - f"request_ids_by_prompt_id={sum(len(g.request_ids_by_prompt_id) for g in self.groups.values())}, " - f"running_request_ids={sum(len(g.running_request_ids) for g in self.groups.values())}, " - f"pause_request_ids={sum(len(g.pause_request_ids) for g in self.groups.values())}, " - f"finished_prompt_ids_by_start_step={sum(len(g.finished_prompt_ids) for g in self.groups.values())}, " - f"finished_ids_by_prompt_id={sum(len(g.finished_ids_by_prompt_id) for g in self.groups.values())}" - ) - - def add_request_item(self, request_item: ExperienceItem): - start_step = request_item.sampling_start_step - group = self.groups.setdefault(start_step, ItemsGroup(start_step=start_step)) - self.prompt_id_to_start_step[request_item.prompt_id] = start_step - group.add_request_item(request_item) - - def get_item(self, request_id: str) -> ExperienceItem: - for group in self.groups.values(): - if request_id in group.by_request_id: - return group.by_request_id[request_id] - raise ValueError(f"request_id {request_id} not found") - - def prompt_num(self) -> int: - return sum(len(group.request_ids_by_prompt_id) for group in self.groups.values()) - - def running_request_num(self) -> int: - return sum(len(group.running_request_ids) for group in self.groups.values()) - - def _get_group_by_prompt_id(self, prompt_id: int) -> Optional[ItemsGroup]: - start_step = self.prompt_id_to_start_step.get(prompt_id) - if start_step is None: - return None - return self.groups.get(start_step) - - def report_item(self, data: DataProto, is_finished: bool = True): - request_id = data.meta_info["request_id"] - item = self.get_item(request_id) - if item is None: - raise ValueError(f"request_id {request_id} not found") - - group = self.groups[item.sampling_start_step] - group.running_request_ids.discard(request_id) - - item.data = data - item.running_dp_rank = None - item.status = ExpStatus.FINISHED if is_finished else ExpStatus.PAUSED - if is_finished: - group.finished_ids_by_prompt_id[item.prompt_id].add(request_id) - else: - group.pause_request_ids.add(request_id) - - def get_finished_data_by_prompt_id(self, prompt_id: int) -> List[DataProto]: - """ - get all finished data for a prompt_id - """ - group = self._get_group_by_prompt_id(prompt_id) - request_ids = group.finished_ids_by_prompt_id.get(prompt_id, set()) - return [group.by_request_id[request_id].data for request_id in request_ids] - - def mark_prompt_id_finished(self, prompt_id: int, num_return_sequences: int) -> list[ExperienceItem]: - """ - mark a prompt_id as finished, pop extra items, to ensure only num_return_sequences items in finished_ids_by_prompt_id - """ - group = self._get_group_by_prompt_id(prompt_id) - return group.mark_prompt_id_finished(prompt_id, num_return_sequences) - - def pop_paused_item(self) -> Optional[ExperienceItem]: - """ - pop a paused item, used to resume paused requests - """ - for start_step in sorted(self.groups.keys()): - group = self.groups[start_step] - if len(group.pause_request_ids) > 0: - request_id = group.pause_request_ids.pop() - item = group.by_request_id.get(request_id) - item.status = ExpStatus.RUNNING - group.running_request_ids.add(request_id) - return item - return None - - def get_enough_finished_prompt_ids( - self, total_prompt_num: int, min_step: Optional[int] = None, min_step_prompt_num: Optional[int] = None - ) -> Optional[List[int]]: - # if not enough return None - if min_step is None: - assert min_step_prompt_num is None - if not ( - sum(len(group.finished_prompt_ids) for group in self.groups.values()) >= total_prompt_num - ): - return None - return list(set().union(*(group.finished_prompt_ids for group in self.groups.values())))[: total_prompt_num] - - min_step_group = self.groups.get(min_step, None) - if min_step_group is None: - return None - min_step_prompt_num = min(min_step_prompt_num, len(min_step_group.get_prompt_ids())) - required_prompt_ids = list(min_step_group.finished_prompt_ids)[:min_step_prompt_num] - if len(required_prompt_ids) < min_step_prompt_num: - return None - - for start_step in sorted(self.groups.keys()): - if start_step <= min_step: - continue - group = self.groups[start_step] - wanted_prompt_num = min(total_prompt_num - len(required_prompt_ids), len(group.get_prompt_ids())) - if wanted_prompt_num <= 0 or len(group.finished_prompt_ids) < wanted_prompt_num: - break - required_prompt_ids.extend(list(group.finished_prompt_ids)[:wanted_prompt_num]) - - if len(required_prompt_ids) < total_prompt_num: - return None - - assert len(required_prompt_ids) == total_prompt_num, f"{len(required_prompt_ids)=} {total_prompt_num=}" - return required_prompt_ids - - def pop_finished_items_by_prompt_ids(self, prompt_ids: List[int], num_return_sequences: int) -> List[ExperienceItem]: - items = [] - for prompt_id in prompt_ids: - _, finished_items = self.pop_prompt_id(prompt_id) - assert len(finished_items) == num_return_sequences - items.extend(finished_items) - return items - - def reset_status(self, min_step: Optional[int] = None, required_prompt_num: Optional[int] = None): - # clear not needed prompt_ids, clear all when min_step is None - for start_step in list(self.groups.keys()): - group = self.groups[start_step] - if min_step is None or start_step < min_step: - group = self.groups.pop(start_step) - del_prompt_ids = group.get_prompt_ids() - elif start_step == min_step: - del_prompt_ids = group.get_prompt_ids() - del_num = len(del_prompt_ids) - required_prompt_num - if del_num <= 0: - continue - logger.info(f"randomly delete {del_num} prompt_ids from step: {start_step}") - del_prompt_ids = random.sample(list(del_prompt_ids), del_num) - else: # start_step > min_step - continue - - for prompt_id in del_prompt_ids: - self.pop_prompt_id(prompt_id) - - for group in self.groups.values(): - group.reset_status() - - def is_prompt_id_finished(self, prompt_id: int) -> bool: - group = self._get_group_by_prompt_id(prompt_id) - return prompt_id in group.finished_prompt_ids - - def pop_prompt_id(self, prompt_id: int): - group = self._get_group_by_prompt_id(prompt_id) - self.prompt_id_to_start_step.pop(prompt_id, None) - if group is None: - return [], [] - return group.pop_prompt_id(prompt_id) - - -@ray.remote(concurrency_groups={"single_thread": 1, "multi_thread": 256}) -class AsyncDynamicSamplingScheduler: - def __init__(self, pipeline_config=None): - self.pipeline_config = pipeline_config - set_seed(seed=pipeline_config.seed) - self.progress_bar: Optional[tqdm] = None - self.request_counter = None - self.mp_rank_zero = {} - self.replay_buffer = ReplayBuffer() - self.prompt_id_counter = itertools.count() - self.response_batch_size: Optional[int] = None - self.lock = threading.Lock() - self.last_alive_check = time.time() - self.dataset_iter_count = 0 - self.exception_queue = queue.Queue() - self.dataset_epoch = 0 - - self.alive_check_interval = self.pipeline_config.alive_check_interval - self.max_additional_running_prompts = self.pipeline_config.max_additional_running_prompts - self.is_use_additional_prompts = self.pipeline_config.is_use_additional_prompts - - self.actor_cluster = None - self.reward_clusters = None - self.reward_worker_iters = None - self.dataset = None - self.indices = [] - self.batch_size = None - self.dataset_iter = None - self.collect_fn_cls = None - self.collect_fn_kwargs = None - self.collect_fn = None - self.tokenizer = None - self.response_callback_fn = None - self.generation_config = None - - self.async_sending_thread = None - self.stop_sending_requests_event = threading.Event() - - self.global_step = 0 - self.init_global_step = None - self.pre_send_dp_rank = 0 - self.filter_prompt_num = 0 - - def set_scheduler( - self, - actor_cluster: Union[Any, Cluster], - reward_clusters: Dict[str, Union[Any, Cluster]], - dataset: Dataset, - collect_fn_cls, - collect_fn_kwargs, - response_filter_fn=None, - query_filter_fn=None, - response_callback_fn=None, - state: Dict[str, Any] = None, - is_val: bool = False, - ): - self.is_val = is_val - self.actor_cluster = actor_cluster - self.reward_clusters = reward_clusters - self.reward_worker_iters = {} - for domain, cluster in reward_clusters.items(): - self.reward_worker_iters[domain] = itertools.cycle(cluster.workers) - - self.dataset = dataset - self.indices = list(range(len(dataset))) - # TODO: test resume - if state is not None and state.get("dataset_iter_count", 0) > 0: - for _ in range(state["dataset_iter_count"]): - self.get_next_dataset_item() - - self.collect_fn_cls = collect_fn_cls - self.collect_fn_kwargs = collect_fn_kwargs - self.tokenizer = default_tokenizer_provider(model_args=self.actor_cluster.worker_config.model_args) - self.collect_fn = self.collect_fn_cls(tokenizer=self.tokenizer, **self.collect_fn_kwargs) - self.response_callback_fn = response_callback_fn - dp_ranks: List[int] = [rank_info.dp_rank for rank_info in self.actor_cluster.worker_rank_info] - for i, dp_rank in enumerate(dp_ranks): - rank_info = self.actor_cluster.get_rank_info(rank=i) - if rank_info.tp_rank == 0 and rank_info.pp_rank == 0 and rank_info.cp_rank == 0: - self.mp_rank_zero[dp_rank] = self.actor_cluster.workers[i] - - # TODO: support response_filter_fn - if self.is_use_additional_prompts: - self.query_filter_fn = query_filter_fn - else: - self.query_filter_fn = lambda data_list, config: True - logger.info("use_additional_prompts is False, disable query and response filtering.") - - self.cluster_max_running_requests = self.pipeline_config.max_running_requests * self.actor_cluster.dp_size - self.request_counter = GlobalCounter.options( - name="DynamicSchedulerRequestCounter", - get_if_exists=True, - namespace=RAY_NAMESPACE, - ).remote() - - def reset_status(self): - self.exception_queue = queue.Queue() - min_start_step, required_prompt_num = self._get_min_step_and_prompt_num(for_keep=True) - with self.lock: - self.replay_buffer.reset_status(min_start_step, required_prompt_num) - - bar_name = "-".join(self.reward_clusters.keys()) - self.progress_bar = tqdm( - total=self.batch_size, - desc=f"{bar_name} generate progress(prompt)", - mininterval=int(self.batch_size * 0.1) + 1, - ) - self.stop_sending_requests_event.clear() - - def send_query_items(self, items: list[ExperienceItem]) -> list[ObjectRef]: - # items will be send to the same worker - refs = [] - dp_rank = next(self.get_available_dp_rank()) - - for req_item in items: - req_item.running_dp_rank = dp_rank - req_item.data.meta_info["response_callback_fn"] = self.response_callback_fn - req_item.data.meta_info["request_id"] = f"{req_item.request_id}_{self.global_step}" - refs.append( - self.actor_cluster.workers[dp_rank].add_request.remote( - command=GenerateRequestType.ADD, data=req_item.data - ) - ) - req_item.data.meta_info.pop("response_callback_fn") - req_item.data.meta_info["request_id"] = f"{req_item.request_id}" - return refs - - def send_paused_requests(self) -> Optional[list[ObjectRef]]: - with self.lock: - paused_item = self.replay_buffer.pop_paused_item() - if paused_item is None: - return None - return self.send_query_items([paused_item]) - - def send_next_data_item(self, data: DataProto) -> list[ObjectRef]: - # get a query from dataset - prompt_id = next(self.prompt_id_counter) - dataset_item = self.get_next_dataset_item() - domain = dataset_item.get("domain", "default") - collect_data = self.collect_fn([dataset_item]) - request_data: DataProto = DataProto.from_single_dict(collect_data, meta_info=data.meta_info) - - # replica, redundancy - request_data_list = self.expand_requests(request_data) - req_items: list[ExperienceItem] = [] - for req in request_data_list: - request_id = ray.get(self.request_counter.get_value.remote()) - req.meta_info["prompt_id"] = prompt_id - request_item = ExperienceItem( - request_id=f"{request_id}", - prompt_id=prompt_id, - sampling_start_step=self.global_step, - domain=domain, - data=req, - ) - req_items.append(request_item) - - with self.lock: - for req_item in req_items: - self.replay_buffer.add_request_item(req_item) - return self.send_query_items(req_items) - - def _get_min_step_and_prompt_num(self, for_keep=False): - if self.is_val: - return None, None - - min_start_step = self.global_step - math.ceil(self.pipeline_config.async_generation_ratio) - min_step_ratio = self.pipeline_config.async_generation_ratio % 1 - min_step_ratio = min_step_ratio if min_step_ratio > 0 else 1.0 - min_step_ratio += max(self.init_global_step - min_start_step, 0) - if not for_keep: - min_step_ratio = min(min_step_ratio, 1.0) - - required_prompt_num = round(self.batch_size * min_step_ratio) - min_start_step = max(min_start_step, self.init_global_step) - return min_start_step, required_prompt_num - - def get_batch(self, data: DataProto, batch_size: int) -> DataProto: - if self.is_val: - self.async_sending_thread.join() - self.async_sending_thread = None - self.batch_size = batch_size - global_step = data.meta_info.get("global_step", 0) - self.set_global_step(global_step) - min_start_step, min_step_prompt_num = self._get_min_step_and_prompt_num() - num_return_sequences = self.generation_config['num_return_sequences'] - - logger.info(f"get batch by {min_start_step=}, {min_step_prompt_num=} {num_return_sequences=} {batch_size=}") - while True: - with self.lock: - finished_prompt_ids = self.replay_buffer.get_enough_finished_prompt_ids( - batch_size, min_start_step, min_step_prompt_num - ) - if finished_prompt_ids is not None: - break - self.check_response_callback() - self.check_worker_alive(self.actor_cluster) - time.sleep(1) - - with self.lock: - finished_items = self.replay_buffer.pop_finished_items_by_prompt_ids(finished_prompt_ids, num_return_sequences) - return self.collect_items_as_batch(finished_items=finished_items, prompt_ids=finished_prompt_ids) - - def collect_items_as_batch(self, finished_items: List[ExperienceItem], prompt_ids: List[int]) -> DataProto: - collect_data_by_domain = defaultdict(list) - data_off_policy_step = 0.0 - for item in finished_items: - collect_data_by_domain[item.domain].append(item.data) - data_off_policy_step += self.global_step - item.sampling_start_step - data_off_policy_step = data_off_policy_step / len(finished_items) - - collect_data_by_domain = { - domain: DataProto.concat(data_list) for domain, data_list in collect_data_by_domain.items() - } - query_use_count = len(prompt_ids) - collect_data_num = sum(data.batch.batch_size[0] for data in collect_data_by_domain.values()) - logger.info( - f"total collect data: {collect_data_num}, collect queries: {len(prompt_ids)} " - f"used queries: {query_use_count} filter queries: {self.filter_prompt_num}" - ) - - batch = DataProto.concat(list(collect_data_by_domain.values())) - batch.meta_info["metrics"] = { - "scheduler/collect_query_count": len(prompt_ids), - "scheduler/query_use_count": query_use_count, - "scheduler/off_policy_ratio": data_off_policy_step, - "scheduler/filter_query_count": self.filter_prompt_num, - } - self.filter_prompt_num = 0 # report here, so refresh here - - metrics = {} - for domain, response_batch in collect_data_by_domain.items(): - sequence_score = response_batch.batch["scores"] - metrics[f"scheduler/{domain}/score/mean"] = torch.mean(sequence_score).detach().item() - metrics[f"scheduler/{domain}/score/max"] = torch.max(sequence_score).detach().item() - metrics[f"scheduler/{domain}/score/min"] = torch.min(sequence_score).detach().item() - - batch.meta_info["metrics"].update(metrics) - return batch - - def set_global_step(self, global_step: int): - self.init_global_step = ( - min(global_step, self.init_global_step) if self.init_global_step is not None else global_step - ) - self.global_step = global_step - - def pause_sampling(self, data: DataProto): - if self.async_sending_thread is not None: - self.stop_sending_requests_event.set() - logger.info("waiting for async sending thread to finish...") - self.async_sending_thread.join() - self.async_sending_thread = None - self.set_global_step(data.meta_info.get("global_step")) - - stop_refs = [] - for infer_worker in self.actor_cluster.workers: - stop_refs.append(infer_worker.add_request.remote(command=GenerateRequestType.STOP, data=None)) - ray.get(stop_refs) - logger.info("async sampling paused, waiting for all requests to be collected...") - start_time = time.time() - timeout = 120 - while True: - if self.replay_buffer.running_request_num() == 0: - break - if time.time() - start_time > timeout: - logger.warning(f"Timeout after {timeout}s waiting for running requests to complete. " - f"Remaining running requests: {self.replay_buffer.running_request_num()}") - break - self.check_response_callback() - time.sleep(1) - logger.info(f"async sampling paused, replay_buffer info: {self.replay_buffer.info()}") - - def sending_request(self, data: DataProto): - all_refs = [] - if self.is_val: - for i in range(self.batch_size): - all_refs.extend(self.send_next_data_item(data)) - ray.get(all_refs) - logger.info(f"async validation send {self.batch_size} prompts.") - return - - while True: - refs = self.send_paused_requests() - if refs is None: - break - all_refs.extend(refs) - - buffer_prompt_num = self.batch_size * max(1.0, self.pipeline_config.async_generation_ratio) - if self.is_use_additional_prompts: - buffer_prompt_num += self.max_additional_running_prompts - - send_prompt_num = 0 - while True: - if self.stop_sending_requests_event.is_set(): - break - if ( - self.replay_buffer.prompt_num() >= buffer_prompt_num - or self.replay_buffer.running_request_num() >= self.cluster_max_running_requests - ): - time.sleep(1) - continue - all_refs.extend(self.send_next_data_item(data)) - send_prompt_num += 1 - ray.get(all_refs) - logger.info(f"async sending thread send {send_prompt_num} prompts.") - - def start_sampling(self, data: DataProto, batch_size: int): - # in async training, called after model update - self.batch_size = batch_size - - global_step = data.meta_info.get("global_step", 0) - self.set_global_step(global_step) - logger.info(f"start async sampling, global_step: {global_step} {self.replay_buffer.info()}") - - self.generation_config = copy.deepcopy(data.meta_info["generation_config"]) - data.meta_info["collect_non_finish"] = True - self.reset_status() - logger.info( - f"start async sampling, batch_size: {self.batch_size}, " - f"num_return_sequences: {self.generation_config['num_return_sequences']}" - ) - self.async_sending_thread = threading.Thread(target=self.sending_request, args=(data,)) - self.async_sending_thread.start() - - @ray.method(concurrency_group="multi_thread") - def report_response(self, data: DataProto): - try: - data.meta_info["request_id"] = data.meta_info["request_id"].split("_")[0] - experience_item = self.replay_buffer.get_item(data.meta_info["request_id"]) - prompt_id = experience_item.prompt_id - num_return_sequences = self.generation_config["num_return_sequences"] - - is_finished = is_report_data_finished(data) - batch = self.postprocess_output_ids(data) if is_finished else self.postprocess_paused_data(data) - with self.lock: - if not is_finished: - self.replay_buffer.report_item(batch, is_finished=is_finished) - return - reward_worker = next(self.reward_worker_iters[experience_item.domain]) - - # call reward - rewards: DataProto = ray.get(reward_worker.compute_rewards.remote(batch)) - batch.union(rewards) - - with self.lock: - self.replay_buffer.report_item(batch, is_finished=is_finished) - if self.replay_buffer.is_prompt_id_finished(prompt_id): - return - data_list = self.replay_buffer.get_finished_data_by_prompt_id(prompt_id) - if not sum(data.batch.batch_size[0] for data in data_list) >= num_return_sequences: - return - - need_prompt = self.query_filter_fn(data_list, self.pipeline_config) - with self.lock: - if need_prompt: - abort_items = self.replay_buffer.mark_prompt_id_finished(prompt_id, num_return_sequences) - self.progress_bar.update() - else: - abort_items, _ = self.replay_buffer.pop_prompt_id(prompt_id) - self.filter_prompt_num += 1 - logger.debug(f"prompt_id {prompt_id} is filtered, abort {len(abort_items)} requests") - # abort uncompleted request - self.abort_requests(abort_items) - except Exception as e: - self.exception_queue.put(e) - - def get_next_dataset_item(self): - if self.dataset_iter is None: - random.seed(self.pipeline_config.seed + self.dataset_epoch) - random.shuffle(self.indices) - self.dataset_iter = iter(self.indices) - logger.info(f"{'-'.join(self.reward_clusters.keys())} dataset epoch: {self.dataset_epoch}") - - try: - dataset_item = self.dataset[next(self.dataset_iter)] - except StopIteration: - self.dataset_epoch += 1 - random.seed(self.pipeline_config.seed + self.dataset_epoch) - random.shuffle(self.indices) - self.dataset_iter = iter(self.indices) - dataset_item = self.dataset[next(self.dataset_iter)] - logger.info(f"{'-'.join(self.reward_clusters.keys())} dataset epoch: {self.dataset_epoch}") - self.dataset_iter_count += 1 - return dataset_item - - def get_scheduler_state(self): - return {"dataset_iter_count": self.dataset_iter_count} - - def abort_requests(self, request_items: list[ExperienceItem]): - abort_refs = [] - for item in request_items: - dp_rank = item.running_dp_rank - if dp_rank is None: - continue - abort_refs.append( - self.actor_cluster.workers[dp_rank].add_request.remote( - command=GenerateRequestType.ABORT, data=DataProto(meta_info={"request_id": item.request_id}) - ) - ) - ray.get(abort_refs) - - def postprocess_paused_data(self, data: DataProto) -> DataProto: - pre_data = self.replay_buffer.get_item(data.meta_info["request_id"]).data - if "output_token_ids" not in data.meta_info: # abort without inferred a token - # too many this log means need more infer workers - logger.info(f"received data without output_token_ids, request_id: {data.meta_info['request_id']}") - return pre_data - logger.debug(f"received paused data, request_id: {data.meta_info['request_id']}") - - assert len(data.meta_info["output_token_ids"]) == 1, ( - "async pipeline only support num_return_sequences=1 or is_num_return_sequences_expand=True" - ) - - # value: list[list[int|float]] - for key in ["output_token_ids", "output_logprobs"]: - cur_value = data.meta_info.pop(key) - pre_value = pre_data.meta_info.get(f"pre_{key}", [[]] * len(cur_value)) - assert len(pre_value) == len(cur_value) - pre_value = [pre_value[i] + cur_value[i] for i in range(len(pre_value))] - data.meta_info[f"pre_{key}"] = pre_value - new_batch = {**pre_data.batch} - - init_attention_mask = pre_data.batch.get("init_attention_mask", pre_data.batch["attention_mask"]) - new_batch["init_attention_mask"] = init_attention_mask - new_batch["init_input_ids"] = pre_data.batch.get("init_input_ids", pre_data.batch["input_ids"]) - - # concat pre output_ids and input_ids - new_input_ids = concatenate_input_and_output( - input_ids=new_batch["init_input_ids"], - output_ids=torch.LongTensor(data.meta_info["pre_output_token_ids"]), - num_return_sequences=len(data.meta_info["pre_output_token_ids"]), - ) - new_batch["input_ids"] = new_input_ids - - new_attention_mask = torch.ones_like(new_input_ids, dtype=init_attention_mask.dtype) - new_attention_mask[:, :init_attention_mask.shape[1]] = init_attention_mask - new_batch["attention_mask"] = new_attention_mask - - max_new_tokens = self.pipeline_config.sequence_length - new_input_ids.shape[1] - if max_new_tokens <= 0: - raise ValueError(f"max_new_tokens {max_new_tokens} <= 0, init_input_ids {new_batch['init_input_ids'].shape}, " - f"pre_output_token_ids {len(data.meta_info['pre_output_token_ids'][0])}") - data.meta_info["max_new_tokens"] = max_new_tokens - data = DataProto.from_dict( - new_batch, non_tensors=pre_data.non_tensor_batch, meta_info={**pre_data.meta_info, **data.meta_info} - ) - assert data.batch["init_attention_mask"].shape[1] == self.pipeline_config.prompt_length - assert data.batch["init_input_ids"].shape[1] == self.pipeline_config.prompt_length - return data - - def postprocess_output_ids(self, data: DataProto) -> DataProto: - # postprocess_generate, input_ids, attention_mask, left pad - request_id = data.meta_info["request_id"] - request: DataProto = self.replay_buffer.get_item(request_id).data - - eos_token_id = data.meta_info["eos_token_id"] - pad_token_id = data.meta_info["pad_token_id"] - input_ids = request.batch.pop("init_input_ids", request.batch["input_ids"]) - request.batch["input_ids"] = input_ids - request.batch["attention_mask"] = request.batch.pop("init_attention_mask", request.batch["attention_mask"]) - output_token_ids = data.meta_info["output_token_ids"] - pre_output_token_ids = request.meta_info.pop("pre_output_token_ids", [[]] * len(output_token_ids)) - output_token_ids = [pre_output_token_ids[i] + output_token_ids[i] for i in range(len(pre_output_token_ids))] - - output_logprobs = data.meta_info.get("output_logprobs", None) - if output_logprobs is not None: - pre_output_logprobs = request.meta_info.get("pre_output_logprobs", [[]] * len(output_token_ids)) - output_logprobs = [pre_output_logprobs[i] + output_logprobs[i] for i in range(len(pre_output_logprobs))] - - output_tokens = [torch.tensor(token_ids) for token_ids in output_token_ids] - output_tensor = pad_sequence(output_tokens, batch_first=True, padding_value=pad_token_id) - output_tensor = concatenate_input_and_output( - input_ids=input_ids, output_ids=output_tensor, num_return_sequences=len(output_tokens) - ) - output: DataProto = postprocess_generate( - prompts=request, - output=output_tensor, - num_return_sequences=len(output_tokens), - sequence_length=self.pipeline_config.sequence_length, - eos_token_id=eos_token_id, - pad_token_id=pad_token_id, - output_logprobs=output_logprobs, - ) - request_repeat = request.repeat(repeat_times=len(output_tokens)) - output.non_tensor_batch = request_repeat.non_tensor_batch - output.meta_info = request_repeat.meta_info - return output - - def expand_requests(self, data: DataProto): - generate_opt_level = self.pipeline_config.generate_opt_level - is_num_return_sequences_expand = self.pipeline_config.is_num_return_sequences_expand - num_return_sequences = self.generation_config["num_return_sequences"] - - assert generate_opt_level > 0, ( - f"generate_opt_level {generate_opt_level} should > 0, " f"in dynamic sampling scheduler." - ) - assert "generation_config" in data.meta_info, f"data {data.meta_info} should have key 'generation_config'" - generation_config = data.meta_info["generation_config"] - - target_requests = [] - if is_num_return_sequences_expand: - generation_config["num_return_sequences"] = 1 - for _ in range(num_return_sequences): - target_requests.append(copy.deepcopy(data)) - else: - generation_config["num_return_sequences"] = num_return_sequences - target_requests.append(copy.deepcopy(data)) - - return target_requests - - def check_worker_alive(self, cluster): - current_time = time.time() - if current_time - self.last_alive_check >= self.alive_check_interval: - cluster.add_request(command=GenerateRequestType.ALIVE_CHECK, data=DataProto()) - self.last_alive_check = current_time - if self.async_sending_thread is not None and not self.async_sending_thread.is_alive(): - raise RuntimeError("async sending thread is dead") - - def check_response_callback(self): - if self.exception_queue.qsize() > 0: - e = self.exception_queue.get() - logger.error(f"report_response get exception {e}") - raise e - - def get_available_dp_rank(self): - while True: - self.pre_send_dp_rank = (self.pre_send_dp_rank + 1) % len(self.mp_rank_zero) - yield self.pre_send_dp_rank diff --git a/roll/distributed/scheduler/decorator.py b/roll/distributed/scheduler/decorator.py index b36bdb0e8..84a18cddc 100644 --- a/roll/distributed/scheduler/decorator.py +++ b/roll/distributed/scheduler/decorator.py @@ -17,6 +17,8 @@ from roll.distributed.scheduler.protocol import DataProto, ObjectRefWrap from roll.utils.logging import get_logger from roll.platforms import current_platform +from roll.utils.offload_states import clear_memory +from roll.utils.telemetry import get_tracer, inject_trace_context, attach_trace_context logger = get_logger() @@ -150,6 +152,8 @@ def dispatch_dp_mp_compute(cluster, *args, **kwargs): def dispatch_dp_mp_dispatch_first(cluster, *args, **kwargs): return _dispatch_dp_mp_compute(cluster, True, *args, **kwargs) +def need_collect_dp_mp_compute(rank_info): + return rank_info.tp_rank == 0 and rank_info.is_pipeline_last_stage and rank_info.cp_rank == 0 def collect_dp_mp_compute(cluster, output): """ @@ -159,7 +163,7 @@ def collect_dp_mp_compute(cluster, output): output_in_dp = [] for global_rank in range(cluster.world_size): local_rank_info = cluster.get_rank_info(rank=global_rank) - if local_rank_info.tp_rank == 0 and local_rank_info.is_pipeline_last_stage and local_rank_info.cp_rank == 0: + if need_collect_dp_mp_compute(local_rank_info): output_in_dp.append(output[global_rank]) if isinstance(output[0], list): return list(chain.from_iterable(output_in_dp)) @@ -171,11 +175,7 @@ def collect_dp_mp_compute(cluster, output): for global_rank in range(cluster.world_size): local_rank_info = cluster.get_rank_info(rank=global_rank) collected = False - if ( - local_rank_info.tp_rank == 0 - and local_rank_info.is_pipeline_last_stage - and local_rank_info.cp_rank == 0 - ): + if need_collect_dp_mp_compute(local_rank_info): collected = True output_in_dp.append(ObjectRefWrap(output[global_rank], collected=collected)) return output_in_dp @@ -234,6 +234,14 @@ def func(*args, blocking=True, **kwargs): if method_name == "initialize": setattr(cls, "initialized", True) + # Inject trace context into DataProto.meta_info before dispatch + for arg in args: + if isinstance(arg, DataProto): + inject_trace_context(arg.meta_info) + for v in kwargs.values(): + if isinstance(v, DataProto): + inject_trace_context(v.meta_info) + args, kwargs = dispatch_fn(cls, *args, **kwargs) output = execute_fn(method_name, *args, **kwargs) if blocking: @@ -261,7 +269,14 @@ def _check_execute_mode(execute_mode): assert isinstance(execute_mode, Execute), f"execute_mode must be a Execute. Got {execute_mode}" -def register(dispatch_mode=Dispatch.ALL_TO_ALL, execute_mode=Execute.ALL, clear_cache=True): +def register( + dispatch_mode=Dispatch.ALL_TO_ALL, + execute_mode=Execute.ALL, + clear_cache=True, + trace=False, + prefetch: bool=False, # TODO shigao: support selective prefetching + to_remote: bool=False, +): _check_dispatch_mode(dispatch_mode) _check_execute_mode(execute_mode) @@ -272,12 +287,18 @@ def decorator(func): @wraps(func) async def inner_async(*args, **kwargs): try: + has_remote_batch = any( + arg._remote_batch is not None + for arg in (args + tuple(kwargs.values())) + if isinstance(arg, DataProto) + ) + assert not has_remote_batch, f"Remote data is not allowed for async function {func.__name__}" + result = await func(*args, **kwargs) if clear_cache: try: current_platform.clear_cublas_workspaces() - gc.collect() - current_platform.empty_cache() + clear_memory() except Exception as oe: pass @@ -292,20 +313,56 @@ async def inner_async(*args, **kwargs): else: @wraps(func) def inner(*args, **kwargs): - try: - result = func(*args, **kwargs) - if clear_cache: - try: - current_platform.clear_cublas_workspaces() - gc.collect() - current_platform.empty_cache() - except Exception as oe: - pass - - except Exception as e: - logger.error(str(e)) - logger.error(traceback.format_exc()) - raise e + from roll.distributed.executor.worker import Worker + + combined_args = args + tuple(kwargs.values()) + has_remote_batch = any(arg._remote_batch is not None for arg in combined_args if isinstance(arg, DataProto)) + + # Extract trace context from DataProto.meta_info + # When trace=False, _otel_meta stays {} → attach_trace_context + # suppresses all get_tracer() calls to no-op. + _otel_meta = {} + if trace: + for arg in combined_args: + if isinstance(arg, DataProto): + _otel_meta = arg.meta_info + break + + with ( + attach_trace_context(_otel_meta), + get_tracer("worker").start_as_current_span(f"worker.{func.__name__}"), + ): + try: + if prefetch and has_remote_batch: + with get_tracer("worker").start_as_current_span(f"prefetch"): + assert len(combined_args) == 2 + self, data = combined_args[0], combined_args[1] + assert isinstance(self, Worker) and isinstance(data, DataProto) + data.prefetch() + + with get_tracer("worker").start_as_current_span(f"exec"): + result = func(*args, **kwargs) + + if to_remote and has_remote_batch: + assert dispatch_mode in [Dispatch.DP_MP_COMPUTE, Dispatch.DP_MP_DISPATCH_FIRST] + with get_tracer("worker").start_as_current_span(f"to_remote"): + assert len(combined_args) == 2 + self, data = combined_args[0], combined_args[1] + assert isinstance(self, Worker) and isinstance(data, DataProto) + if need_collect_dp_mp_compute(self.get_rank_info()): + result = DataProto.to_remote(result, ref_data=data) + + if clear_cache: + try: + current_platform.clear_cublas_workspaces() + clear_memory() + except Exception as oe: + pass + + except Exception as e: + logger.error(str(e)) + logger.error(traceback.format_exc()) + raise e return result setattr(inner, BIND_WORKER_METHOD_FLAG, attrs) diff --git a/roll/distributed/scheduler/generate_scheduler.py b/roll/distributed/scheduler/generate_scheduler.py index f624aab29..672004f4d 100644 --- a/roll/distributed/scheduler/generate_scheduler.py +++ b/roll/distributed/scheduler/generate_scheduler.py @@ -1,6 +1,7 @@ import asyncio import copy import itertools +import os import random import math import uuid @@ -26,6 +27,7 @@ from roll.utils.metrics.metrics_manager import DurationTracker from roll.utils.import_utils import safe_import_class from roll.utils.logging import get_logger +from roll.utils.telemetry import get_tracer, inject_trace_context, attach_trace_context, init_telemetry logger = get_logger() @@ -442,6 +444,7 @@ def __init__( collect_fn_kwargs, state: Dict[str, Any] = None, is_val: bool = False, + get_data_item_kwargs: Optional[Dict[str, Any]] = None, ): self.pipeline_config = pipeline_config set_seed(seed=pipeline_config.seed) @@ -471,23 +474,24 @@ def __init__( self.request_id = uuid.uuid4() self.request_counter = 0 - self.dataset = dataset - self.indices = list(range(len(dataset))) - if state is not None and state.get("dataset_iter_count", 0) > 0: - for _ in range(state["dataset_iter_count"]): - self.get_next_dataset_item() - self.dataset_epoch = 0 - self.dataset_iter = None - self.dataset_iter_count = 0 - self.collect_fn_cls = collect_fn_cls self.collect_fn_kwargs = collect_fn_kwargs + self.get_data_item_kwargs = get_data_item_kwargs if get_data_item_kwargs else {} self.tokenizer = default_tokenizer_provider(model_args=self.actor_cluster.worker_config.model_args) self.processor = default_processor_provider(model_args=self.actor_cluster.worker_config.model_args) if "processor" in [f.name for f in fields(collect_fn_cls)]: collect_fn_kwargs["processor"] = self.processor self.collect_fn = self.collect_fn_cls(tokenizer=self.tokenizer, **self.collect_fn_kwargs) + self.dataset = dataset + self.indices = list(range(len(dataset))) + self.dataset_epoch = 0 + self.dataset_iter = None + self.dataset_iter_count = 0 + if state is not None and state.get("dataset_iter_count", 0) > 0: + for _ in range(state["dataset_iter_count"]): + self.get_next_dataset_item() + self.async_sending_task = None # Dynamic filter is supported no matter whether is_use_additional_prompts, @@ -513,6 +517,14 @@ def __init__( self.reward_scheduler = RewardScheduler() + # Initialize OpenTelemetry tracing on the scheduler actor + if os.environ.get("ROLL_OTEL_ENABLED") == "1": + init_telemetry( + service_name="scheduler", + otlp_endpoint=os.environ.get("OTEL_EXPORTER_OTLP_ENDPOINT"), + instance_id=f"scheduler-{'-'.join(self.reward_clusters.keys())}", + ) + async def initialize(self): await self.router_manager.initialize() self.router_client = await RouterManager.create_client(self.router_manager) @@ -567,16 +579,7 @@ async def get_batch_opt_level_0(self, data: DataProto, batch_size: int) -> DataP request_data: DataProto = DataProto.from_single_dict(collect_data, meta_info=data.meta_info) request_data.batch["prompt_id"] = torch.arange(request_data.batch.batch_size[0], device=request_data.batch.device) - generate_non_tensor_batch_keys = [] - if "multi_modal_data" in request_data.non_tensor_batch: - generate_non_tensor_batch_keys.append("multi_modal_data") - - gen_batch = request_data.pop( - batch_keys=["input_ids", "attention_mask", "position_ids"], - non_tensor_batch_keys=generate_non_tensor_batch_keys, - ) - for key in generate_non_tensor_batch_keys: - request_data.non_tensor_batch[key] = gen_batch.non_tensor_batch[key] + gen_batch = request_data.pop(batch_keys=["input_ids", "attention_mask", "position_ids"]) gen_batch.meta_info = request_data.meta_info num_return_sequences = generation_config["num_return_sequences"] request_data = request_data.repeat(repeat_times=num_return_sequences) @@ -619,6 +622,21 @@ async def get_batch(self, data: DataProto, global_step: int, batch_size: int) -> if self.pipeline_config.generate_opt_level == 0: return await self.get_batch_opt_level_0(data, batch_size) + with ( + attach_trace_context(data.meta_info), + get_tracer("scheduler").start_as_current_span( + "get_batch", + attributes={ + "global_step": global_step, + "batch_size": batch_size, + "domains": ",".join(self.reward_clusters.keys()), + }, + ), + ): + inject_trace_context(data.meta_info) + return await self._get_batch_impl(data, global_step, batch_size) + + async def _get_batch_impl(self, data: DataProto, global_step: int, batch_size: int) -> DataProto: num_return_sequences = data.meta_info["generation_config"]["num_return_sequences"] self.meta_info = copy.deepcopy(data.meta_info) self.meta_info["collect_non_finish"] = self.pipeline_config.async_generation_ratio > 0 @@ -677,6 +695,9 @@ async def get_batch(self, data: DataProto, global_step: int, batch_size: int) -> # DUMP MODE: Save merged batch (from mixin) await self._maybe_dump_batch(batch, global_step) + with get_tracer("scheduler").start_as_current_span("to_remote"): + loop = asyncio.get_running_loop() + batch = await loop.run_in_executor(None, DataProto.to_remote, batch) return batch def collect_items_as_batch(self, finished_items: List[ExperienceItem]) -> DataProto: @@ -711,6 +732,11 @@ def collect_items_as_batch(self, finished_items: List[ExperienceItem]) -> DataPr metrics[f"scheduler/{domain}/score/mean"] = torch.mean(sequence_score).detach().item() metrics[f"scheduler/{domain}/score/max"] = torch.max(sequence_score).detach().item() metrics[f"scheduler/{domain}/score/min"] = torch.min(sequence_score).detach().item() + domain_extra_metrics = response_batch.meta_info.get("metrics", {}) + for k, v in domain_extra_metrics.items(): + metrics[f"scheduler/{domain}/{k}/mean"] = np.mean(v).item() + metrics[f"scheduler/{domain}/{k}/max"] = np.max(v).item() + metrics[f"scheduler/{domain}/{k}/min"] = np.min(v).item() batch.meta_info["metrics"].update(metrics) # TODO shigao implement REPORT_LENGTH_AND_REWARDS (deleted at refactor) @@ -732,7 +758,48 @@ async def sending_request(self): await self.router_manager.abort_all() # Implicitly wait until all running tasks finished when TaskGroup context exit. + def get_next_data_item_by_dataloader(self, batch_size: int = 1, use_collect_fn: bool = False, num_workers=8): + # use Dataloader to parallel and prefetch + from torch.utils.data import DataLoader + from roll.datasets.collator import collate_fn_to_dict_list + + if getattr(self, "dataloader", None) is None: + self.dataloader = DataLoader( + self.dataset, + batch_size=batch_size, + collate_fn=self.collect_fn if use_collect_fn else collate_fn_to_dict_list, + shuffle=True, + drop_last=False, + num_workers=num_workers, + ) + self.dataloader_iter = iter(self.dataloader) + self.input_data_batch = None + if not (self.input_data_batch and self.idx_in_batch < len(list(self.input_data_batch.values())[0])): + self.idx_in_batch = 0 + try: + self.input_data_batch = next(self.dataloader_iter) + except StopIteration: + self.dataset_epoch += 1 + self.dataloader_iter = iter(self.dataloader) + self.input_data_batch = next(self.dataloader_iter) + data_item = dict( + ( + k, + v[self.idx_in_batch : self.idx_in_batch + 1] if use_collect_fn else v[self.idx_in_batch], + ) + for k, v in self.input_data_batch.items() + ) + self.idx_in_batch += 1 + self.dataset_iter_count += 1 + return data_item + def get_next_dataset_item(self): + if self.get_data_item_kwargs.get("use_dataloader", False): + return self.get_next_data_item_by_dataloader( + batch_size=self.get_data_item_kwargs.get("batch_size", 1), + use_collect_fn=self.get_data_item_kwargs.get("use_collect_fn", False), + num_workers=self.get_data_item_kwargs.get("num_workers", 8), + ) if self.dataset_iter is None: random.seed(self.pipeline_config.seed + self.dataset_epoch) random.shuffle(self.indices) @@ -770,40 +837,44 @@ async def process_new_prompt( num_return_sequences = scheduler.meta_info["generation_config"]["num_return_sequences"] context = RolloutContext(scheduler=scheduler, prompt_id=prompt_id, meta_info=scheduler.meta_info) success = False - try: - responses = await scheduler.udrl.process_new_prompt(context=context) - if responses is None: - logger.info(f"filter out prompt {prompt_id}") - raise asyncio.CancelledError # abort this prompt - responses = expand_responses(responses) - assert ( - len(responses) == num_return_sequences or scheduler.replay_buffer.is_use_additional_prompts - ), "is_use_additional_prompts is required when using dynamic num_return_sequences" - except Exception as e: - logger.warning(f"abort prompt {prompt_id} on exception {e}") - raise - else: - success = True - finally: - scheduler.running_tasks.pop(prompt_id, None) - - # commit/abort should be put at last in finally block, because commit may raise exception - if not success: - scheduler.replay_buffer.abort(prompt_id) + with attach_trace_context(scheduler.meta_info), get_tracer("scheduler").start_as_current_span( + "process_new_prompt", + attributes={"prompt_id": prompt_id}, + ): + try: + responses = await scheduler.udrl.process_new_prompt(context=context) + if responses is None: + logger.info(f"filter out prompt {prompt_id}") + raise asyncio.CancelledError # abort this prompt + responses = expand_responses(responses) + assert ( + len(responses) == num_return_sequences or scheduler.replay_buffer.is_use_additional_prompts + ), "is_use_additional_prompts is required when using dynamic num_return_sequences" + except Exception as e: + logger.warning(f"abort prompt {prompt_id} on exception {e}") + raise else: - assert context.sampling_start_step is not None - scheduler.replay_buffer.commit( - prompt_id, - [ - ExperienceItem( - prompt_id=prompt_id, - domain=context.domain, - sampling_start_step=context.sampling_start_step, - data=response, - ) - for response in responses - ], - ) + success = True + finally: + scheduler.running_tasks.pop(prompt_id, None) + + # commit/abort should be put at last in finally block, because commit may raise exception + if not success: + scheduler.replay_buffer.abort(prompt_id) + else: + assert context.sampling_start_step is not None + scheduler.replay_buffer.commit( + prompt_id, + [ + ExperienceItem( + prompt_id=prompt_id, + domain=context.domain, + sampling_start_step=context.sampling_start_step, + data=response, + ) + for response in responses + ], + ) def __init__( self, @@ -843,8 +914,14 @@ def get_request_data(self, meta_info): self.got_data: bool = True dataset_item = self._scheduler.get_next_dataset_item() - domain = dataset_item.get("domain", "default") - collect_data = self._scheduler.collect_fn([dataset_item]) + if self._scheduler.get_data_item_kwargs.get( + "use_dataloader", False + ) and self._scheduler.get_data_item_kwargs.get("use_collect_fn", False): + domain = dataset_item.get("domain", np.array(["default"], dtype=object)).tolist()[0] + collect_data = dataset_item + else: + domain = dataset_item.get("domain", "default") + collect_data = self._scheduler.collect_fn([dataset_item]) self.domain = domain return DataProto.from_single_dict(collect_data, meta_info=meta_info), domain @@ -869,8 +946,19 @@ async def generate( domain: str, ) -> DataProto: assert self._in_do_generate_and_reward - with self._scheduler.generate_timer[domain].track(): - request_id = self._scheduler.next_request_id() + request_id = self._scheduler.next_request_id() + with ( + self._scheduler.generate_timer[domain].track(), + get_tracer("scheduler").start_as_current_span( + "rollout.generate", + attributes={ + "domain": domain, + "prompt_id": self.prompt_id, + "request_id": request_id, + }, + ), + ): + inject_trace_context(req.meta_info) req.meta_info["request_id"] = request_id logger.debug(f"generate_and_reward: {self.prompt_id=} {request_id} generate_request") self.inflight_requests.add(request_id) @@ -887,7 +975,17 @@ async def compute_rewards( # reward worker得能支持单条数据计算, dynamic sampling对需要batch计算reward的需要注意... # 多域的时候,llm as judge, 需要单独为reward worker分配gpu assert self._in_do_generate_and_reward - with self._scheduler.reward_timer[domain].track(): + with ( + self._scheduler.reward_timer[domain].track(), + get_tracer("scheduler").start_as_current_span( + "rollout.compute_rewards", + attributes={ + "domain": domain, + "prompt_id": self.prompt_id, + }, + ), + ): + inject_trace_context(req.meta_info) reward_worker = next(self._scheduler.reward_worker_iters[domain]) logger.debug(f"generate_and_reward: {self.prompt_id=} compute_rewards") output_count = req.batch.batch_size[0] diff --git a/roll/distributed/scheduler/initialize.py b/roll/distributed/scheduler/initialize.py index 877e4ef18..6482def2c 100644 --- a/roll/distributed/scheduler/initialize.py +++ b/roll/distributed/scheduler/initialize.py @@ -1,7 +1,9 @@ import os +import socket import subprocess import sys import time +import atexit import ray @@ -22,6 +24,47 @@ from roll.platforms import current_platform logger = get_logger() +log_monitor_listener = None + +def wait_for_head_node_ready(master_addr: str, master_port: str, timeout: int = 600, check_interval: int = 2): + """Wait for Ray head node GCS to become available. + + Args: + master_addr: Head node address + master_port: Head node GCS port + timeout: Maximum time to wait in seconds (default: 10 minutes) + check_interval: Interval between connection attempts in seconds (default: 2s) + + Raises: + RuntimeError: If head node doesn't become available within timeout + """ + start_time = time.time() + elapsed = 0 + + logger.info(f"Waiting for Ray head node at {master_addr}:{master_port} to become available...") + + while elapsed < timeout: + try: + # Try to connect to the GCS port + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.settimeout(check_interval) + result = sock.connect_ex((master_addr, int(master_port))) + sock.close() + + if result == 0: + logger.info(f"Ray head node at {master_addr}:{master_port} is ready (took {elapsed:.1f}s)") + return + except (socket.timeout, socket.error, OSError) as e: + pass + + elapsed = time.time() - start_time + if elapsed < timeout: + time.sleep(check_interval) + + raise RuntimeError( + f"Ray head node at {master_addr}:{master_port} did not become available within {timeout}s. " + f"Please check if the head node is starting correctly." + ) def start_ray_cluster(): @@ -39,8 +82,8 @@ def start_ray_cluster(): if rank == 0: cmd = f"ray start --head --port={master_port} --node-name={node_name} --dashboard-port={dashboard_port}" else: - # fix: 处理大规模下可能会出现的head/worker node创建顺序不一致问题 - time.sleep(5) + # Wait for head node to be ready before starting worker + wait_for_head_node_ready(master_addr, master_port) cmd = f"ray start --address={master_addr}:{master_port} --node-name={node_name} --dashboard-port={dashboard_port}" logger.info(f"Starting ray cluster: {cmd}") @@ -53,6 +96,12 @@ def start_ray_cluster(): return True +def stop_handler(): + global log_monitor_listener + if log_monitor_listener is not None: + log_monitor_listener.stop() + + def init(): rank = get_driver_rank() world_size = get_driver_world_size() @@ -77,8 +126,10 @@ def init(): if manual_start: wait_for_nodes(expected=world_size) - listener = LogMonitorListener() - listener.start() + atexit.register(stop_handler) + global log_monitor_listener + log_monitor_listener = LogMonitorListener() + log_monitor_listener.start() logger.info(f"Current ray cluster resources: {ray.available_resources()}") diff --git a/roll/distributed/scheduler/log_monitor.py b/roll/distributed/scheduler/log_monitor.py index ff64646b5..98c7197e8 100644 --- a/roll/distributed/scheduler/log_monitor.py +++ b/roll/distributed/scheduler/log_monitor.py @@ -235,8 +235,6 @@ def stop(self): subprocess.run(cmd, shell=True, capture_output=True) def start(self): - atexit.register(self.stop) - if self.rank == 0: self.exception_monitor = ExceptionMonitor.options( name=EXCEPTION_MONITOR_ACTOR_NAME, get_if_exists=True, namespace=RAY_NAMESPACE diff --git a/roll/distributed/scheduler/protocol.py b/roll/distributed/scheduler/protocol.py index c79d33377..a3cd487e2 100644 --- a/roll/distributed/scheduler/protocol.py +++ b/roll/distributed/scheduler/protocol.py @@ -6,9 +6,9 @@ import copy import os +import uuid from collections import defaultdict -from dataclasses import dataclass, field -from typing import Any, Dict, List, Optional, Union, Set +from typing import Dict, List, Optional, Union, Set import numpy as np import ray @@ -16,7 +16,10 @@ import torch from tensordict import TensorDict from torch.utils.data import DataLoader +from codetiming import Timer +from roll.distributed.scheduler.remote_protocol import RemoteBatch, BatchProxy +from roll.distributed.scheduler import transfer_backend from roll.utils.functionals import union_two_dict, divide_by_chunk_size from roll.platforms import current_platform from roll.utils.logging import get_logger @@ -70,9 +73,13 @@ def union_tensor_dict(tensor_dict1: TensorDict, tensor_dict2: TensorDict) -> Ten if key not in tensor_dict1.keys(): tensor_dict1[key] = tensor_dict2[key] else: - assert tensor_dict1[key].equal( - tensor_dict2[key] - ), f"{key} in tensor_dict1 and tensor_dict2 are not the same object" + # Compare values - handle both tensors and NonTensorStack + val1, val2 = tensor_dict1[key], tensor_dict2[key] + assert type(val1) == type(val2), f"{key} has different types: {type(val1)} vs {type(val2)}" + if isinstance(val1, torch.Tensor): + assert val1.equal(val2), f"{key} in tensor_dict1 and tensor_dict2 are not the same" + # For NonTensorStack and other types, skip equality check + # (comparison would require iterating and checking each element) return tensor_dict1 @@ -106,6 +113,8 @@ def list_of_dict_to_dict_of_list(list_of_dict: list[dict]): output = {} for d in list_of_dict: + if d is None: + continue if not isinstance(d, dict): raise TypeError(f"Expected dict, but got {type(d)}: {d}") for k, v in d.items(): @@ -114,21 +123,11 @@ def list_of_dict_to_dict_of_list(list_of_dict: list[dict]): return output -def collate_fn(x: list["DataProtoItem"]): - batch = [] - non_tensor_batch = [] - meta_info = None - for data in x: - meta_info = data.meta_info - batch.append(data.batch) - non_tensor_batch.append(data.non_tensor_batch) - batch = torch.stack(batch).contiguous() - non_tensor_batch = list_of_dict_to_dict_of_list(non_tensor_batch) - for key, val in non_tensor_batch.items(): - non_tensor_batch[key] = np.empty(len(val), dtype=object) - non_tensor_batch[key][:] = val - return DataProto(batch=batch, non_tensor_batch=non_tensor_batch, meta_info=meta_info) - +def collate_fn(x: list["DataProto"]): + meta_info = x[-1].meta_info + data = DataProto.concat(x) + data.meta_info = meta_info + return data def move_tensors_to_device(data, device): if isinstance(data, dict): @@ -151,14 +150,6 @@ def custom_np_concatenate(val): return concatenated_array -@dataclass -class DataProtoItem: - batch: TensorDict = None - non_tensor_batch: Dict = field(default_factory=dict) - meta_info: Dict = field(default_factory=dict) - - -@dataclass class DataProto: """ A DataProto is a data structure that aims to provide a standard protocol for data exchange between functions. @@ -167,25 +158,97 @@ class DataProto: same batch size should be put inside batch. """ - batch: TensorDict = None - non_tensor_batch: Dict = field(default_factory=dict) - meta_info: Dict = field(default_factory=dict) + def __init__( + self, + batch: TensorDict = None, + non_tensor_batch: Dict = None, + remote_batch: RemoteBatch = None, + meta_info: Dict = None, + ): + if batch is None and remote_batch is not None: + batch = TensorDict({}, batch_size=[len(remote_batch)]) + self._batch = batch + self._non_tensor_batch = non_tensor_batch if non_tensor_batch is not None else {} + self._remote_batch = remote_batch + self.meta_info = meta_info if meta_info is not None else {} + self.__post_init__() + + @property + def batch(self) -> "BatchProxy": + """Hook: called before accessing batch. + Returns a BatchProxy that supports fallback lookup to remote_batch.""" + return BatchProxy(self._batch, self._remote_batch, len(self)) + + @batch.setter + def batch(self, value: TensorDict | BatchProxy): + assert isinstance(value, (TensorDict, BatchProxy)) + value = value.copy() + if self._remote_batch is not None: + for key in value.keys(): + if key in self._remote_batch: + del self._remote_batch[key] + for key in self._non_tensor_batch.keys(): + if key in value: + del value[key] + if isinstance(value, BatchProxy): + if self._remote_batch is not None: + self._remote_batch.union(value._remote_batch) + else: + self._remote_batch = value._remote_batch + self._batch = value._batch + else: + self._batch = value + self.check_consistency() + + @property + def non_tensor_batch(self) -> "BatchProxy": + """Hook: called before accessing non_tensor_batch. + Returns a BatchProxy that supports fallback lookup to remote_batch.""" + return BatchProxy(self._non_tensor_batch, self._remote_batch, len(self)) + + @non_tensor_batch.setter + def non_tensor_batch(self, value: dict | BatchProxy): + assert isinstance(value, (dict, BatchProxy)) + value = value.copy() + if self._remote_batch is not None: + for key in value.keys(): + if key in self._remote_batch: + del self._remote_batch[key] + for key in self._batch.keys(): + if key in value: + del value[key] + if isinstance(value, BatchProxy): + if self._remote_batch is not None: + self._remote_batch.union(value._remote_batch) + else: + self._remote_batch = value._remote_batch + self._non_tensor_batch = value._batch + else: + self._non_tensor_batch = value + self.check_consistency() def __post_init__(self): # perform necessary checking self.check_consistency() - - if self.batch is not None and current_platform.is_npu(): - for key, val in self.batch.items(): + + if self._batch is not None and current_platform.is_npu(): + for key, val in self._batch.items(): if isinstance(val, torch.Tensor) and val.dtype == torch.int64: logger.debug(f"[NPU] Converting Tensor {key} from int64 -> int32, shape={val.shape}") - self.batch[key] = val.to(torch.int32) + self._batch[key] = val.to(torch.int32) + + assert self._remote_batch is None or not current_platform.is_npu() + + def __repr__(self) -> str: + return f"DataProto(batch={self._batch}, non_tensor_batch={self._non_tensor_batch}, remote_batch={self._remote_batch}, meta_info={self.meta_info})" def __len__(self): - if self.batch is not None: - return self.batch.batch_size[0] - if self.non_tensor_batch is not None: - return len(next(iter(self.non_tensor_batch.values()))) + if self._batch is not None: + return len(self._batch) + if self._non_tensor_batch: + return len(next(iter(self._non_tensor_batch.values()))) + if self._remote_batch is not None: + return len(self._remote_batch) return 0 def __getitem__(self, item): @@ -199,62 +262,102 @@ def __getitem__(self, item): - list: A list of indices - numpy.ndarray: An array of indices - torch.Tensor: A tensor of indices + - str: A key to look up in batch or remote_batch Returns: - DataProto: For all indexing types except single integers - DataProtoItem: Only for single integer indices + DataProto: For slice/list/array/tensor/int indexing. + torch.Tensor | np.ndarray: For string key lookup. """ - # Case 1: Slice object - use the slice method if isinstance(item, slice): return self.slice(item.start, item.stop, item.step) - - # Case 2: List, numpy array, or torch tensor - use sel_idxs elif isinstance(item, (list, np.ndarray, torch.Tensor)): return self.select_idxs(item) - - # Case 3: Single integer - return DataProtoItem for backward compatibility elif isinstance(item, (int, np.integer)): - tensor_data = self.batch[item] - non_tensor_data = {key: val[item] for key, val in self.non_tensor_batch.items()} - return DataProtoItem(batch=tensor_data, non_tensor_batch=non_tensor_data, meta_info=self.meta_info) - - # # Case 4: Unsupported type + return self.slice(item, item + 1, 1) + elif isinstance(item, str): + # Search batch first, then remote_batch + if self._batch is not None and item in self._batch.keys(): + return self._batch[item] + elif item in self._non_tensor_batch.keys(): + return self._non_tensor_batch[item] + elif self._remote_batch is not None and item in self._remote_batch: + return self._remote_batch[item] + else: + raise KeyError(f"Key '{item}' not found in batch or remote_batch") else: raise TypeError(f"Indexing with {type(item)} is not supported") + def __setitem__(self, key: str, value) -> None: + if isinstance(key, str): + if self._remote_batch is not None and key in self._remote_batch: + del self._remote_batch[key] + if isinstance(value, torch.Tensor): + if key in self._non_tensor_batch: + del self._non_tensor_batch[key] + self._batch[key] = value + elif isinstance(value, np.ndarray): + if self._batch is not None and key in self._batch.keys(): + del self._batch[key] + self._non_tensor_batch[key] = value + else: + raise TypeError(f"Unsupported type for value: {type(value)}") + else: + raise TypeError(f"Key must be str, got {type(key)}") + + def __delitem__(self, key: str) -> None: + """Delete key from batch or remote_batch.""" + if isinstance(key, str): + if self._batch is not None and key in self._batch: + del self._batch[key] + elif key in self._non_tensor_batch: + del self._non_tensor_batch[key] + elif self._remote_batch is not None and key in self._remote_batch: + del self._remote_batch[key] + else: + raise KeyError(f"Key '{key}' not found") + else: + raise TypeError(f"Key must be str, got {type(key)}") + def __getstate__(self): import io buffer = io.BytesIO() - if tensordict.__version__ >= "0.5.0" and self.batch is not None: - self.batch = self.batch.contiguous() - self.batch = self.batch.consolidate() - torch.save(self.batch, buffer) - return buffer, self.non_tensor_batch, self.meta_info + if tensordict.__version__ >= "0.5.0" and self._batch is not None: + self._batch = self._batch.contiguous() + self._batch = self._batch.consolidate() + torch.save(self._batch, buffer) + return buffer, self._non_tensor_batch, self._remote_batch, self.meta_info def __setstate__(self, data): - batch_deserialized, non_tensor_batch, meta_info = data + batch_deserialized, non_tensor_batch, remote_batch, meta_info = data batch_deserialized.seek(0) batch = torch.load( batch_deserialized, weights_only=False, map_location="cpu" if not current_platform.is_available() else None ) - self.batch = batch - self.non_tensor_batch = non_tensor_batch + self._batch = batch + self._non_tensor_batch = non_tensor_batch + self._remote_batch = remote_batch self.meta_info = meta_info def check_consistency(self): """Check the consistency of the DataProto. Mainly for batch and non_tensor_batch We expose this function as a public one so that user can call themselves directly """ - if self.batch is not None: - assert len(self.batch.batch_size) == 1, "only support num_batch_dims=1" + existing_keys = set() + + if self._batch is not None: + assert len(self._batch.batch_size) == 1, "only support num_batch_dims=1" + existing_keys.update(self._batch.keys()) - if len(self.non_tensor_batch) != 0: + if len(self._non_tensor_batch) != 0: # TODO: we can actually lift this restriction if needed - assert len(self.batch.batch_size) == 1, "only support num_batch_dims=1 when non_tensor_batch is not empty." + assert len(self._batch.batch_size) == 1, "only support num_batch_dims=1 when non_tensor_batch is not empty." - batch_size = self.batch.batch_size[0] - for key, val in self.non_tensor_batch.items(): + assert existing_keys.isdisjoint(self._non_tensor_batch.keys()), "batch and non_tensor_batch cannot have overlapping keys" + existing_keys.update(self._non_tensor_batch.keys()) + + batch_size = self._batch.batch_size[0] + for key, val in self._non_tensor_batch.items(): assert ( isinstance(val, np.ndarray) and val.dtype == object ), "data in the non_tensor_batch must be a numpy.array with dtype=object" @@ -262,6 +365,15 @@ def check_consistency(self): val.shape[0] == batch_size ), f"key {key} length {len(val)} is not equal to batch size {batch_size}" + if self._remote_batch is not None: + assert existing_keys.isdisjoint( + self._remote_batch.keys() + ), f"batch and remote_batch cannot have overlapping keys {existing_keys} {self._remote_batch.keys()}" + if self._batch is not None: + assert ( + len(self._remote_batch) == self._batch.batch_size[0] + ), f"remote_batch length {len(self._remote_batch)} is not equal to batch size {self._batch.batch_size[0]}" + @classmethod def from_single_dict(cls, data: Dict[str, Union[torch.Tensor, np.ndarray]], meta_info=None): tensors = {} @@ -328,8 +440,10 @@ def to(self, device) -> "DataProto": DataProto: the current DataProto """ - if self.batch is not None: - self.batch = self.batch.to(device) + if self._batch is not None: + self._batch = self._batch.to(device) + if self._remote_batch is not None: + self._remote_batch = self._remote_batch.to(device) if self.meta_info is not None: self.meta_info = move_tensors_to_device(self.meta_info, device) @@ -347,10 +461,12 @@ def clone(self) -> "DataProto": independent memory. """ # Copy batch - batch_copy = self.batch.clone() if self.batch is not None else None + batch_copy = self._batch.clone() if self._batch is not None else None # Copy non-tensor objects (numpy arrays) - non_tensor_copy = {k: np.copy(v) for k, v in self.non_tensor_batch.items()} + non_tensor_copy = {k: np.copy(v) for k, v in self._non_tensor_batch.items()} + + remote_batch_copy = self._remote_batch.clone() if self._remote_batch is not None else None # Deep copy meta_info to avoid shared mutable objects meta_copy = copy.deepcopy(self.meta_info) @@ -359,9 +475,27 @@ def clone(self) -> "DataProto": return DataProto( batch=batch_copy, non_tensor_batch=non_tensor_copy, + remote_batch=remote_batch_copy, meta_info=meta_copy ) + def update(self, other: dict): + assert isinstance(other, dict) + for key, value in other.items(): + if self._remote_batch is not None and key in self._remote_batch: + del self._remote_batch[key] + if isinstance(value, torch.Tensor): + assert self._batch is not None + if key in self._non_tensor_batch: + del self._non_tensor_batch[key] + self._batch[key] = value + elif isinstance(value, np.ndarray): + if self._batch is not None and key in self._batch.keys(): + del self._batch[key] + self._non_tensor_batch[key] = value + else: + raise TypeError(f"Unsupported type {type(value)} for key '{key}': expected torch.Tensor or np.ndarray") + def select(self, batch_keys=None, non_tensor_batch_keys=None, meta_info_keys=None, deepcopy=False) -> "DataProto": """Select a subset of the DataProto via batch_keys and meta_info_keys @@ -372,19 +506,27 @@ def select(self, batch_keys=None, non_tensor_batch_keys=None, meta_info_keys=Non Returns: DataProto: the DataProto with the selected batch_keys and meta_info_keys """ + assert set(batch_keys).isdisjoint(non_tensor_batch_keys), "batch_keys and non_tensor_batch_keys cannot be overlapping" + if batch_keys is not None: batch_keys = tuple(batch_keys) - sub_batch = self.batch.select(*batch_keys) + sub_batch = self._batch.select(*batch_keys) else: - sub_batch = self.batch + batch_keys = [] + sub_batch = self._batch if non_tensor_batch_keys is not None: - non_tensor_batch = {key: val for key, val in self.non_tensor_batch.items() if key in non_tensor_batch_keys} + non_tensor_batch = {key: val for key, val in self._non_tensor_batch.items() if key in non_tensor_batch_keys} else: - non_tensor_batch = self.non_tensor_batch + non_tensor_batch_keys = [] + non_tensor_batch = self._non_tensor_batch - if deepcopy: - non_tensor_batch = copy.deepcopy(non_tensor_batch) + if self._remote_batch is not None: + assert not deepcopy, "remote_batch deepcopy is not supported yet" + # FIXME: The behavior of select is changed when batch_keys or non_tensor_batch_keys is None. + sub_remote_batch = self._remote_batch.select(batch_keys + non_tensor_batch_keys) + else: + sub_remote_batch = None if meta_info_keys is not None: sub_meta_info = {key: val for key, val in self.meta_info.items() if key in meta_info_keys} @@ -392,9 +534,13 @@ def select(self, batch_keys=None, non_tensor_batch_keys=None, meta_info_keys=Non sub_meta_info = self.meta_info if deepcopy: + non_tensor_batch = copy.deepcopy(non_tensor_batch) + sub_remote_batch = sub_remote_batch.clone() sub_meta_info = copy.deepcopy(sub_meta_info) - return DataProto(batch=sub_batch, non_tensor_batch=non_tensor_batch, meta_info=sub_meta_info) + return DataProto( + batch=sub_batch, non_tensor_batch=non_tensor_batch, remote_batch=sub_remote_batch, meta_info=sub_meta_info + ) def select_idxs(self, idxs): """ @@ -420,24 +566,33 @@ def select_idxs(self, idxs): batch_size = idxs_np.sum() if idxs_np.dtype == bool else idxs_np.shape[0] - if self.batch is not None: + if self._batch is not None: # Use TensorDict's built-in indexing capabilities selected_batch = TensorDict( - source={key: tensor[idxs_torch] for key, tensor in self.batch.items()}, batch_size=(batch_size,) + source={key: tensor[idxs_torch] for key, tensor in self._batch.items()}, batch_size=(batch_size,) ) else: selected_batch = None selected_non_tensor = {} - for key, val in self.non_tensor_batch.items(): + for key, val in self._non_tensor_batch.items(): selected_non_tensor[key] = val[idxs_np] - return type(self)(batch=selected_batch, non_tensor_batch=selected_non_tensor, meta_info=self.meta_info) + if self._remote_batch is not None: + remote_batch = self._remote_batch.select_idxs(idxs_torch) + else: + remote_batch = None + + return type(self)( + batch=selected_batch, + non_tensor_batch=selected_non_tensor, + remote_batch=remote_batch, + meta_info=self.meta_info, + ) def slice(self, start=None, end=None, step=None): """ Slice the DataProto and return a new DataProto object. - This is an improved version of direct slicing which returns a DataProtoItem. Args: start (int, optional): Start index. Defaults to None (start from beginning). @@ -459,26 +614,31 @@ def slice(self, start=None, end=None, step=None): indices = [1, 5, 10] selected_data = data_proto[indices] - # Single index still returns DataProtoItem + # Single index returns DataProto too single_item = data_proto[5] """ # Create a slice object slice_obj = slice(start, end, step) # Handle the batch data - if self.batch is not None: + if self._batch is not None: # Use TensorDict's built-in slicing capabilities - sliced_batch = self.batch[slice_obj] + sliced_batch = self._batch[slice_obj] else: sliced_batch = None # Handle the non-tensor batch data sliced_non_tensor = {} - for key, val in self.non_tensor_batch.items(): + for key, val in self._non_tensor_batch.items(): sliced_non_tensor[key] = val[slice_obj] + if self._remote_batch is not None: + remote_batch = self._remote_batch[slice_obj] + else: + remote_batch = None + # Return a new DataProto object - return type(self)(batch=sliced_batch, non_tensor_batch=sliced_non_tensor, meta_info=self.meta_info) + return type(self)(batch=sliced_batch, non_tensor_batch=sliced_non_tensor, remote_batch=remote_batch, meta_info=self.meta_info) def pop(self, batch_keys=None, non_tensor_batch_keys=None, meta_info_keys=None) -> "DataProto": """Pop a subset of the DataProto via `batch_keys` and `meta_info_keys` @@ -495,25 +655,41 @@ def pop(self, batch_keys=None, non_tensor_batch_keys=None, meta_info_keys=None) meta_info_keys = [] if non_tensor_batch_keys is None: non_tensor_batch_keys = [] + assert set(batch_keys).isdisjoint(non_tensor_batch_keys), "batch_keys and non_tensor_batch_keys cannot be overlapping" batch_keys = self.validate_input(batch_keys) non_tensor_batch_keys = self.validate_input(non_tensor_batch_keys) meta_info_keys = self.validate_input(meta_info_keys) + remote_batch_keys = set() + tensors = {} - # tensor batch for key in batch_keys: - assert key in self.batch.keys() - tensors[key] = self.batch.pop(key) + if key not in self._batch.keys(): + remote_batch_keys.add(key) + else: + tensors[key] = self._batch.pop(key) + tensors = TensorDict(tensors, batch_size=len(self)) + non_tensors = {} - # non tensor batch for key in non_tensor_batch_keys: - assert key in self.non_tensor_batch.keys() - non_tensors[key] = self.non_tensor_batch.pop(key) + if key not in self._non_tensor_batch.keys(): + remote_batch_keys.add(key) + else: + non_tensors[key] = self._non_tensor_batch.pop(key) + + remote_batch = self._remote_batch.pop(remote_batch_keys) if self._remote_batch else None + meta_info = {} for key in meta_info_keys: assert key in self.meta_info.keys() meta_info[key] = self.meta_info.pop(key) - return DataProto.from_dict(tensors=tensors, non_tensors=non_tensors, meta_info=meta_info) + + return DataProto( + batch=tensors, + non_tensor_batch=non_tensors, + remote_batch=remote_batch, + meta_info=meta_info, + ) @staticmethod def validate_input(keys): @@ -529,6 +705,8 @@ def validate_input(keys): def rename(self, old_keys=None, new_keys=None) -> "DataProto": """ Note that this function only rename the key in the batch + + WARNING: Rename will materialize the remote batch if the old key is in the remote batch. """ old_keys = self.validate_input(old_keys) @@ -539,9 +717,42 @@ def rename(self, old_keys=None, new_keys=None) -> "DataProto": f"new_keys and old_keys must have the same length, but got {len(new_keys)} and {len(old_keys)}" ) - self.batch.rename_key_(tuple(old_keys), tuple(new_keys)) + if self._remote_batch is None: + for old_key, new_key in zip(old_keys, new_keys): + self._batch.rename_key_(old_key, new_key) + self.check_consistency() + return self + else: + logger.warning(f"RemoteBatch renaming keys {old_keys} to {new_keys} is not efficient for remote data") + if self._batch is None: + self._batch = TensorDict({}, batch_size=len(self)) + + local_old_keys = [] + local_new_keys = [] + remote_old_keys = [] + remote_new_keys = [] + for old_key, new_key in zip(old_keys, new_keys): + if old_key in self._batch: + assert old_key not in self._remote_batch + local_old_keys.append(old_key) + local_new_keys.append(new_key) + elif old_key in self._remote_batch: + remote_old_keys.append(old_key) + remote_new_keys.append(new_key) + else: + raise KeyError(f"{old_key} not in batch") + + if local_old_keys: + for old_key, new_key in zip(local_old_keys, local_new_keys): + self._batch.rename_key_(old_key, new_key) - return self + self._remote_batch.materialize(remote_old_keys) + for old_key, new_key in zip(remote_old_keys, remote_new_keys): + self._batch[new_key] = self._remote_batch[old_key] + del self._remote_batch[old_key] + + self.check_consistency() + return self def union(self, other: "DataProto") -> "DataProto": """Union with another DataProto. Union batch and meta_info separately. @@ -556,9 +767,26 @@ def union(self, other: "DataProto") -> "DataProto": Returns: DataProto: the DataProto after union """ - self.batch = union_tensor_dict(self.batch, other.batch) - self.non_tensor_batch = union_numpy_dict(self.non_tensor_batch, other.non_tensor_batch) + if self._batch is not None and other._batch is not None: + self._batch = union_tensor_dict(self._batch, other._batch) + elif other._batch is not None: + self._batch = TensorDict(other._batch.to_dict(), batch_size=other._batch.batch_size) + + self._non_tensor_batch = union_numpy_dict(self._non_tensor_batch, other._non_tensor_batch) + + if self._remote_batch is not None and other._remote_batch is not None: + self._remote_batch = self._remote_batch.union(other._remote_batch) + elif self._remote_batch is None: + self._remote_batch = other._remote_batch + if self._remote_batch is not None: + existing_keys = set(self._batch.keys() if self._batch is not None else []) | set(self._non_tensor_batch.keys()) + for key in existing_keys: + # use local batch as golden source when key conflict + if key in self._remote_batch: + del self._remote_batch[key] + self.meta_info = union_two_dict(self.meta_info, other.meta_info) + self.check_consistency() return self def make_iterator(self, mini_batch_size, epochs, seed=None, dataloader_kwargs=None): @@ -574,9 +802,9 @@ def make_iterator(self, mini_batch_size, epochs, seed=None, dataloader_kwargs=No Returns: Iterator: an iterator that yields a mini-batch data at a time. The total number of iteration steps is - ``self.batch.batch_size * epochs // mini_batch_size`` + ``self._batch.batch_size * epochs // mini_batch_size`` """ - assert self.batch.batch_size[0] % mini_batch_size == 0, f"{self.batch.batch_size[0]} % {mini_batch_size} != 0" + assert self._batch.batch_size[0] % mini_batch_size == 0, f"{self._batch.batch_size[0]} % {mini_batch_size} != 0" # we can directly create a dataloader from TensorDict if dataloader_kwargs is None: dataloader_kwargs = {} @@ -587,6 +815,10 @@ def make_iterator(self, mini_batch_size, epochs, seed=None, dataloader_kwargs=No else: generator = None + # FIXME do not materialize all fields of remote batch + if self._remote_batch is not None: + self._remote_batch.materialize() + assert isinstance(dataloader_kwargs, Dict) train_dataloader = DataLoader( dataset=self, batch_size=mini_batch_size, collate_fn=collate_fn, generator=generator, **dataloader_kwargs @@ -605,7 +837,7 @@ def chunk(self, chunks: int) -> List["DataProto"]: 要求: batch_size > chunks，调用方保证，此处保证每个chunk会返回一个DataProto - np.array_split(val, chunks) 和 self.batch.chunk(chunks=chunks, dim=0) 在不能均分时行为不同 + np.array_split(val, chunks) 和 self._batch.chunk(chunks=chunks, dim=0) 在不能均分时行为不同 Args: chunks (int): the number of chunks to split on dim=0 @@ -614,29 +846,35 @@ def chunk(self, chunks: int) -> List["DataProto"]: """ chunks_sizes = None if len(self) > 0: - assert len(self) >= chunks, f"batch_size {self.batch.batch_size[0]} < chunks {chunks}" + assert len(self) >= chunks, f"batch_size {self._batch.batch_size[0]} < chunks {chunks}" index_array = np.arange(len(self)) chunks_sizes = [len(b) for b in np.array_split(index_array, chunks)] - if self.batch is not None: - batch_lst = divide_by_chunk_size(self.batch, chunk_sizes=chunks_sizes) + if self._batch is not None: + batch_lst = divide_by_chunk_size(self._batch, chunk_sizes=chunks_sizes) else: batch_lst = [None for _ in range(chunks)] non_tensor_batch_lst = [{} for _ in range(chunks)] - for key, val in self.non_tensor_batch.items(): + for key, val in self._non_tensor_batch.items(): assert isinstance(val, np.ndarray) non_tensor_lst = divide_by_chunk_size(val, chunk_sizes=chunks_sizes) assert len(non_tensor_lst) == chunks, f"len(non_tensor_lst) {len(non_tensor_lst)} != chunks {chunks}" for i in range(chunks): non_tensor_batch_lst[i][key] = non_tensor_lst[i] + if self._remote_batch: + remote_batch_lst = self._remote_batch.chunk(chunks_sizes) + else: + remote_batch_lst = [None for _ in range(chunks)] + output = [] for i in range(chunks): output.append( DataProto( batch=batch_lst[i].clone() if batch_lst[i] is not None else batch_lst[i], non_tensor_batch=non_tensor_batch_lst[i], + remote_batch=remote_batch_lst[i], meta_info=self.meta_info, ) ) @@ -668,29 +906,40 @@ def concat( A new DataProto with concatenated tensors, non-tensor data, and processed meta information. """ + if len(data) == 1: + return data[0] + global_keys = global_keys if global_keys is not None else {"metrics"} # ---------- 1. Concatenate tensor / non-tensor batches ---------- - batch_lst = [d.batch for d in data if d.batch is not None] + batch_lst = [d._batch for d in data if d._batch is not None] new_batch = torch.cat(batch_lst, dim=0) if batch_lst else None non_tensor_batch = list_of_dict_to_dict_of_list( - [d.non_tensor_batch for d in data] + [d._non_tensor_batch for d in data] ) for k, v in non_tensor_batch.items(): non_tensor_batch[k] = custom_np_concatenate(v) + remote_batch_list = [d._remote_batch for d in data if d._remote_batch is not None] + remote_batch = RemoteBatch.cat(remote_batch_list) if remote_batch_list else None + # ---------- 2. Aggregate meta information ---------- merged_meta = dict(data[0].meta_info) # start with rank-0 values for key in global_keys: - if key not in merged_meta: + # Check if any data has this key, not just the first one + has_key = any(key in d.meta_info and d.meta_info[key] is not None for d in data) + if not has_key: continue values = [d.meta_info.get(key) for d in data] + # Determine the type from first non-None value + first_non_none_value = next((v for v in values if v is not None), None) + # Case 1: dict — aggregate each sub-key across ranks - if isinstance(merged_meta[key], dict): + if isinstance(first_non_none_value, dict): sub_dict = list_of_dict_to_dict_of_list(values) for sub_key, sub_list in sub_dict.items(): try: @@ -710,6 +959,7 @@ def concat( return DataProto( batch=new_batch, non_tensor_batch=non_tensor_batch, + remote_batch=remote_batch, meta_info=merged_meta, ) @@ -720,8 +970,9 @@ def reorder(self, indices): # Ensure that indices is at least a 1-D tensor. indices = indices.view(-1) if indices.dim() == 0 else indices indices_np = indices.detach().numpy() - self.batch = self.batch[indices] - self.non_tensor_batch = {key: val[indices_np] for key, val in self.non_tensor_batch.items()} + self._batch = self._batch[indices] if self._batch is not None else None + self._non_tensor_batch = {key: val[indices_np] for key, val in self._non_tensor_batch.items()} + self._remote_batch = self._remote_batch.select_idxs(indices) if self._remote_batch else None def group_by(self, keys: Union[List[str], str]) -> Dict[str, "DataProto"]: """ @@ -740,16 +991,23 @@ def group_by(self, keys: Union[List[str], str]) -> Dict[str, "DataProto"]: keys = self.validate_input(keys) assert len(keys) > 0, "Must provide at least one grouping key" + remote_keys = [key for key in keys if self._remote_batch is not None and key in self._remote_batch] + if remote_keys and not self._remote_batch.cached(remote_keys): + self._remote_batch.materialize(remote_keys) + logger.warning(f"RemoteBatch implicit materialize key {remote_keys} for group by") + # Collect grouping values across data types group_key_values = [] for idx in range(len(self)): key_values = [] for key in keys: # Check tensor data first - if key in self.batch.keys(): - key_values.append(str(self.batch[key][idx].numpy())) - elif key in self.non_tensor_batch: - key_values.append(str(self.non_tensor_batch[key][idx])) + if self._batch is not None and key in self._batch.keys(): + key_values.append(str(self._batch[key][idx].numpy())) + elif key in self._non_tensor_batch: + key_values.append(str(self._non_tensor_batch[key][idx])) + elif self._remote_batch is not None and key in self._remote_batch: + key_values.append(str(self._remote_batch[key][idx])) else: raise KeyError(f"Grouping key '{key}' not found in tensor or non-tensor data") @@ -765,7 +1023,7 @@ def group_by(self, keys: Union[List[str], str]) -> Dict[str, "DataProto"]: # Create grouped DataProtos grouped_data = {} for group_key, indices in groups.items(): - grouped_data[group_key] = collate_fn([self[idx] for idx in indices]) + grouped_data[group_key] = self.select_idxs(indices) return grouped_data @@ -780,36 +1038,39 @@ def repeat(self, repeat_times=2, interleave=True): Returns: DataProto: A new DataProto with repeated data. """ - if self.batch is not None: + if self._batch is not None: if interleave: # Interleave the data repeated_tensors = { - key: tensor.repeat_interleave(repeat_times, dim=0) for key, tensor in self.batch.items() + key: tensor.repeat_interleave(repeat_times, dim=0) for key, tensor in self._batch.items() } else: # Stack the data repeated_tensors = { key: tensor.unsqueeze(0).expand(repeat_times, *tensor.shape).reshape(-1, *tensor.shape[1:]) - for key, tensor in self.batch.items() + for key, tensor in self._batch.items() } repeated_batch = TensorDict( source=repeated_tensors, - batch_size=(self.batch.batch_size[0] * repeat_times,), + batch_size=(self._batch.batch_size[0] * repeat_times,), ) else: repeated_batch = None repeated_non_tensor_batch = {} - for key, val in self.non_tensor_batch.items(): + for key, val in self._non_tensor_batch.items(): if interleave: repeated_non_tensor_batch[key] = np.repeat(val, repeat_times, axis=0) else: repeated_non_tensor_batch[key] = np.tile(val, (repeat_times,) + (1,) * (val.ndim - 1)) + repeated_remote_batch = self._remote_batch.repeat(repeat_times, interleave) if self._remote_batch else None + return type(self)( batch=repeated_batch, non_tensor_batch=repeated_non_tensor_batch, + remote_batch=repeated_remote_batch, meta_info=self.meta_info, ) @@ -856,6 +1117,81 @@ def materialize_concat( # Concatenate and apply global aggregation rules return DataProto.concat(data, global_keys=global_keys) + @classmethod + def to_remote(cls, data: "DataProto", partition = "train_eval", *, ref_data = None) -> "DataProto": + with Timer(name="RemoteBatch to_remote", logger=None) as timer: + batch_size = len(data) + + if ref_data is not None: + if len(ref_data) != batch_size: + logger.warning(f"RemoteBatch to_remote ref_data batch size {len(ref_data)} does not match data batch size {batch_size}, {data=}") + return data + if ref_data._remote_batch is None or ref_data._remote_batch.row_ids() is None: + logger.warning(f"RemoteBatch to_remote ref_data has no row ids") + return data + row_ids = ref_data._remote_batch.row_ids() + if len(row_ids) != batch_size: + logger.warning(f"RemoteBatch to_remote ref_data row ids batch size {len(row_ids)} does not match data batch size {batch_size}") + return data + if data._remote_batch is not None and data._remote_batch.row_ids() != row_ids: + logger.warning(f"RemoteBatch to_remote ref_data and data have different row ids") + return data + logger.info(f"RemoteBatch to_remote add to current rows, {partition=} {len(row_ids)=}") + elif data._remote_batch is not None and data._remote_batch.row_ids() is not None: + row_ids = data._remote_batch.row_ids() + else: + row_ids = [str(uuid.uuid4()) for _ in range(batch_size)] + logger.info(f"RemoteBatch to_remote add {len(row_ids)} new rows, row ids {partition=} row_ids={row_ids[:10]}...") + assert len(row_ids) == batch_size + + # Partition is used to group data of a step. + # Some backend may support drop partition or delete keys using regex. + assert isinstance(partition, str) + partition = data._remote_batch.partition if data._remote_batch is not None else partition + + data_dict: dict = data._batch.to_dict() if data._batch is not None else {} + data_dict.update(data._non_tensor_batch) + + if not data_dict: + return data + + remote_batch = transfer_backend.put(partition, row_ids, data_dict, batch_size) + if remote_batch is None: # transfer backend is not available + assert data._remote_batch is None + return data + + if data._remote_batch is not None: + remote_batch = remote_batch.union(data._remote_batch) + + logger.info(f"RemoteBatch to_remote finished in {timer.last}s") + + return DataProto( + batch=None, + non_tensor_batch={}, + remote_batch=remote_batch, + meta_info=data.meta_info, + ) + + def prefetch(self, keys: list[str] | str | None = None): + if self._remote_batch is None: + return + if keys is not None: + keys = [keys] if isinstance(keys, str) else keys + if keys is None or not self._remote_batch.cached(keys): + with Timer(name="RemoteBatch prefetch", logger=None) as timer: + self._remote_batch.materialize(keys) + logger.info(f"RemoteBatch prefetch finished in {timer.last}s") + + @classmethod + def drop(cls, data: "DataProto"): + if data._remote_batch is None: + return + with Timer(name="RemoteBatch drop", logger=None) as timer: + partition = data._remote_batch.partition + logger.info(f"RemoteBatch drop {partition=} row_ids={data._remote_batch.row_ids()}") + data._remote_batch.drop() + logger.info(f"RemoteBatch drop {partition=} finished in {timer.last}s") + class ObjectRefWrap: def __init__(self, obj_ref: ray.ObjectRef, collected=False): diff --git a/roll/distributed/scheduler/remote_protocol.py b/roll/distributed/scheduler/remote_protocol.py new file mode 100644 index 000000000..22439c8c3 --- /dev/null +++ b/roll/distributed/scheduler/remote_protocol.py @@ -0,0 +1,904 @@ +import threading +from abc import ABC, abstractmethod +from typing import Any, Optional + +import numpy as np +import torch +from tensordict import TensorDict +from tensordict.utils import LinkedList +from codetiming import Timer + +from roll.distributed.scheduler import transfer_backend +from roll.utils.logging import get_logger + +logger = get_logger() + + +class RemoteBatch: + def __init__(self, key_type: str, partition: str, device): + self.key_type = key_type + self.partition = partition + self.device = None + + def __reduce__(self): + raise NotImplementedError + + def __len__(self) -> int: + raise NotImplementedError + + def __eq__(self, other): + raise NotImplementedError + + def __hash__(self): + raise NotImplementedError + + def __getitem__(self, item): + if isinstance(item, slice): + return self.slice(item.start, item.stop, item.step) + elif isinstance(item, (list, np.ndarray, torch.Tensor)): + return self.select_idxs(item) + elif isinstance(item, str): + td = self.materialize([item]) + assert isinstance(td, TensorDict), f"Expected TensorDict, got {type(td)}" + value = td[item] + assert isinstance(value, (torch.Tensor, LinkedList)) + if isinstance(value, LinkedList): + items = list(value) + return np.array(items, dtype=object) + else: + return value + else: + raise TypeError(f"Indexing with {type(item)} is not supported") + + def __delitem__(self, key: str): + raise NotImplementedError + + def __contains__(self, key: str) -> bool: + raise NotImplementedError + + def clone(self, recurse: bool = True): + raise NotImplementedError + + def keys(self): + """ + If not specified keys and fields are the same. + (keys only reference to key to kv storage in materialize now) + """ + raise NotImplementedError + + def row_ids(self): + return None + + def to(self, device) -> "RemoteBatch": + self.device = device + return self + + def materialize(self, fields: list[str] = None) -> TensorDict: + raise NotImplementedError + + def cached(self, fields: list[str]) -> bool: + if self.cache is None: + return False + else: + return all(field in self.cache for field in fields) + + def drop(self): + raise NotImplementedError + + def select(self, fileds: list[str]) -> "RemoteBatch": + raise NotImplementedError + + def select_idxs(self, index: torch.Tensor | np.ndarray | list) -> "RemoteBatch": + raise NotImplementedError + + def slice( + self, + start: Optional[int] = None, + end: Optional[int] = None, + step: Optional[int] = None, + ) -> "RemoteBatch": + raise NotImplementedError + + def pop(self, fileds) -> "RemoteBatch": + raise NotImplementedError + + def chunk(self, chunk_sizes: list[int]) -> list["RemoteBatch"]: + raise NotImplementedError + + def repeat(self, repeat_times: int, interleave: bool) -> "RemoteBatch": + raise NotImplementedError + + def union(self, rhs: "RemoteBatch") -> "RemoteBatch": + """ + RemoteBatch.union will not check the following preconditions: + - there are conflict keys in batch and they are not equal + - the batch size of two data batch is not the same + """ + raise NotImplementedError + + @classmethod + def cat(cls, data: list["RemoteBatch"]) -> "RemoteBatch": + assert data + target_cls = type(data[0]) + assert all( + type(d) is target_cls for d in data + ), f"All batches must be of the same type, got {[type(d).__name__ for d in data]}" + return target_cls._cat(data) + + @classmethod + def _cat(cls, data: list["RemoteBatch"]) -> "RemoteBatch": + raise NotImplementedError + + +class BatchProxy: + """ + Proxy for batch that supports fallback lookup to remote_batch. + + Only support a minimal set of special methods and normal methods that works identically on + both TensorDict and dict[np.ndarray]. Raises on other special methods (__len__, __iter__, ...). + + Only support properties of TensorDict currently used in codebase for backward compatibility. + Use of properties of dict is not supported. + """ + + def __init__(self, batch: TensorDict | dict[np.ndarray] | None, remote_batch: RemoteBatch | None, batch_size: int): + assert batch is None or isinstance(batch, (TensorDict, dict)) + self._batch = batch + self._remote_batch = remote_batch + self._batch_size = batch_size + + def __getitem__(self, key: str): + if self._batch is not None and key in self._batch: + return self._batch[key] + elif self._remote_batch is not None and key in self._remote_batch: + return self._remote_batch[key] + else: + raise KeyError(f"Key '{key}' not found in batch or remote_batch") + + def __setitem__(self, key: str, value): + assert isinstance(value, (torch.Tensor, np.ndarray)) + if self._remote_batch is not None and key in self._remote_batch: + # Just delete from local, does not delete from remote server. + del self._remote_batch[key] + if self._batch is not None: + assert ( + len(value) == self._batch_size + ), f"Value length {len(value)} does not match batch length {self._batch_size}" + self._batch[key] = value + else: + raise RuntimeError("Cannot set item when batch is None") + + def __delitem__(self, key: str): + if self._batch is not None and key in self._batch: + del self._batch[key] + elif self._remote_batch is not None and key in self._remote_batch: + # Just delete from local, does not delete from remote server. + del self._remote_batch[key] + else: + raise KeyError(f"Key '{key}' not found in batch or remote_batch") + + def __contains__(self, key: str) -> bool: + in_batch = self._batch is not None and key in self._batch + in_remote = self._remote_batch is not None and key in self._remote_batch + return in_batch or in_remote + + def copy(self) -> "BatchProxy": + """Shallow copy of the BatchProxy.""" + batch_copy = self._batch.copy() if self._batch is not None else None + remote_copy = self._remote_batch.clone(recurse=False) if self._remote_batch is not None else None + return BatchProxy(batch_copy, remote_copy, self._batch_size) + + def get(self, key: str, default=None): + try: + return self[key] + except KeyError: + return default + + def keys(self): + result = set() + if self._batch is not None: + result |= set(self._batch.keys()) + if self._remote_batch is not None: + result |= set(self._remote_batch.keys()) + return result + + def items(self): + """ + WARNING: this function will materializes remote_batch if exists and does not guarantee the order of items. + """ + # Yield from _batch first + if self._batch is not None: + if isinstance(self._batch, TensorDict): + for key in self._batch.keys(): + yield (key, self._batch[key]) + else: + # dict + for key, val in self._batch.items(): + yield (key, val) + # Yield from _remote_batch for keys not in _batch + if self._remote_batch is not None: + logger.warning("RemoteBatch materializing remote batch for items()") + self._remote_batch.materialize() + for key in self._remote_batch.keys(): + yield (key, self._remote_batch[key]) + + _POP_SENTINEL = object() + + def pop(self, key: str, default=_POP_SENTINEL): + if key not in self: + if default is BatchProxy._POP_SENTINEL: + raise KeyError(f"Key '{key}' not found in batch or remote_batch") + return default + res = self[key] + if self._remote_batch is not None and key in self._remote_batch: + del self._remote_batch[key] + if self._batch is not None and key in self._batch: + del self._batch[key] + return res + + def update(self, other: dict): + assert isinstance(other, dict) + assert self._batch is not None, "Update with batch is None is not supported, use DataProto.update insted." + for key in other.keys(): + if self._remote_batch is not None and key in self._remote_batch: + del self._remote_batch[key] + self._batch.update(other) + + def __getattr__(self, name: str): + """ + Raise AttributeError for all other attributes. + Because the semantics of returning getattr(self._batch, name) is undefined. + """ + raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'") + + # ============================================================================ + # Below properties are for backward compatibility only. + # Only supported when: + # - self._batch is not None and isinstance(self._batch, TensorDict), OR + # - self._batch is None and self._remote_batch is not None + # ============================================================================ + + @property + def batch_size(self) -> torch.Size: + """Return batch size. Only supported when _batch is TensorDict or _remote_batch exists.""" + if isinstance(self._batch, TensorDict): + return self._batch.batch_size + if self._batch is None and self._remote_batch is not None: + return torch.Size([len(self._remote_batch)]) + raise AttributeError( + f"'{type(self).__name__}' has no attribute 'batch_size' " + f"(batch is {type(self._batch).__name__ if self._batch else 'None'})" + ) + + @property + def shape(self) -> torch.Size: + """Return batch shape. Only supported when _batch is TensorDict or _remote_batch exists.""" + if isinstance(self._batch, TensorDict): + return self._batch.shape + if self._batch is None and self._remote_batch is not None: + return torch.Size([len(self._remote_batch)]) + raise AttributeError( + f"'{type(self).__name__}' has no attribute 'shape' " + f"(batch is {type(self._batch).__name__ if self._batch else 'None'})" + ) + + @property + def device(self): + """Return device. Only supported when _batch is TensorDict or _remote_batch exists.""" + if isinstance(self._batch, TensorDict): + return self._batch.device + if self._batch is None and self._remote_batch is not None: + return self._remote_batch.device + raise AttributeError( + f"'{type(self).__name__}' has no attribute 'device' " + f"(batch is {type(self._batch).__name__ if self._batch else 'None'})" + ) + + +class RowRemoteBatch(RemoteBatch): + """ + A remote batch stored in a key-value store with row id as keys. + """ + + def __init__(self, partition: str, device, fields, row_ids: list[str], cache: TensorDict): + super().__init__("row", partition, device) + self.fields = set(fields) # str, stores column names + self._row_ids = row_ids.copy() + self.cache = cache.clone() if cache is not None else None + + def __reduce__(self): + return ( + RowRemoteBatch, + (self.partition, self.device, self.fields, self._row_ids, None), + ) + + def __repr__(self): + return f"RowRemoteBatch(partition={self.partition}, device={self.device}, fields={self.fields}, row_ids={self._row_ids}, cache={self.cache})" + + def __len__(self): + return len(self._row_ids) + + def __delitem__(self, key: str): + self.fields.remove(key) + if self.cache is not None and key in self.cache: + del self.cache[key] + + def __contains__(self, key: str) -> bool: + return key in self.fields + + def clone(self, recurse: bool = True): + return RowRemoteBatch( + partition=self.partition, + device=self.device, + fields=self.fields.copy(), + row_ids=self._row_ids.copy(), + cache=self.cache.clone(recurse=recurse) if self.cache is not None else None, + ) + + def keys(self): + return self.fields + + def row_ids(self): + return self._row_ids + + def to(self, device) -> "RowRemoteBatch": + super().to(device) + if self.cache is not None: + self.cache = self.cache.to(device) + return self + + def materialize(self, fields: list[str] = None) -> TensorDict: + if fields is None: + fields = self.fields + else: + assert set(fields) <= self.fields, f"Fields {set(fields)} is not subset of {self.fields}" + existing_fields = set(self.cache.keys()) if self.cache is not None else set() + fetch_fields = [field for field in fields if field not in existing_fields] + if len(fetch_fields) > 0: + with Timer(name="remote_batch_materialize", logger=None) as timer: + data: TensorDict = transfer_backend.get(partition=self.partition, keys=self._row_ids, fields=fetch_fields) + assert set(data.keys()) == set(fetch_fields) + + if self.cache is None: + self.cache = data + else: + from roll.distributed.scheduler.protocol import union_tensor_dict + + self.cache = union_tensor_dict(self.cache, data) + if self.device is not None: + self.cache.to(self.device) + logger.info(f"RemoteBatch materialize cost {timer.last}s, partition={self.partition}, new materialized {sorted(fetch_fields)}, cached fields {sorted(list(existing_fields))}") + + return self.cache.select(*fields) + + def drop(self): + transfer_backend.delete(partition=self.partition, keys=self._row_ids, fields=list(self.fields)) + + def select(self, fileds: list[str]) -> "RowRemoteBatch": + assert all(key in self.fields for key in fileds), f"Keys {fileds} not in {self.fields}" + cache = self.cache + if cache is not None: + keys_in_cache = [k for k in fileds if k in cache.keys()] + if keys_in_cache: + cache = cache.select(*keys_in_cache) + else: + cache = None + return RowRemoteBatch( + partition=self.partition, + device=self.device, + fields=fileds, + row_ids=self._row_ids, + cache=cache, + ) + + def select_idxs(self, index: torch.Tensor | np.ndarray | list) -> "RowRemoteBatch": + assert isinstance(index, (torch.Tensor, np.ndarray, list)) + if isinstance(index, np.ndarray): + index_list = index.tolist() + index = torch.from_numpy(index) + elif isinstance(index, list): + index_list = index + index = torch.tensor(index) + else: + index_list = index.tolist() + + if index.dtype == torch.bool: + selected_row_ids = [self._row_ids[i] for i, mask in enumerate(index_list) if mask] + else: + selected_row_ids = [self._row_ids[i] for i in index_list] + + cache = self.cache + if cache is not None: + cache = cache[index] + assert isinstance(cache, TensorDict) + + return RowRemoteBatch( + partition=self.partition, + device=self.device, + fields=self.fields, + row_ids=selected_row_ids, + cache=cache, + ) + + def slice( + self, + start: Optional[int] = None, + end: Optional[int] = None, + step: Optional[int] = None, + ) -> "RowRemoteBatch": + sliced_row_ids = self._row_ids[start:end:step] + + cache = self.cache + if cache is not None: + cache = cache[slice(start, end, step)] + + return RowRemoteBatch( + partition=self.partition, + device=self.device, + fields=self.fields, + row_ids=sliced_row_ids, + cache=cache, + ) + + def pop(self, filed) -> "RowRemoteBatch": + assert len(filed) == len(set(filed)), "Fields must be unique" + assert set(filed) <= self.fields, f"Fields {set(filed) - self.fields} not in batch" + ret = self.select(filed) + + self.fields -= set(filed) + if self.cache is not None: + remaining_keys = [k for k in self.fields if k in self.cache.keys()] + if remaining_keys: + self.cache = self.cache.select(*remaining_keys) + else: + self.cache = None + + return ret + + def chunk(self, chunk_sizes: list[int]) -> list["RowRemoteBatch"]: + assert sum(chunk_sizes) == len( + self + ), f"Sum of chunk_sizes {sum(chunk_sizes)} does not match batch size {len(self)}" + chunks = [] + offset = 0 + for size in chunk_sizes: + chunks.append(self.slice(offset, offset + size)) + offset += size + return chunks + + def repeat(self, repeat_times: int, interleave: bool) -> "RowRemoteBatch": + if interleave: + repeated_row_ids = [row_id for row_id in self._row_ids for _ in range(repeat_times)] + else: + repeated_row_ids = self._row_ids * repeat_times + + cache = self.cache + if cache is not None: + if interleave: + cache = cache.repeat_interleave(repeat_times) + else: + cache = cache.repeat(repeat_times) + + return RowRemoteBatch( + partition=self.partition, + device=self.device, + fields=self.fields, + row_ids=repeated_row_ids, + cache=cache, + ) + + def union(self, rhs: "RowRemoteBatch") -> "RowRemoteBatch": + assert isinstance(rhs, RemoteBatch) + assert len(self) == len(rhs), f"Two tensor dict must have identical batch size. Got {len(self)} and {len(rhs)}" + if self.cache is not None and rhs.cache is not None: + from roll.distributed.scheduler.protocol import union_tensor_dict + + union_tensor_dict(self.cache, rhs.cache) + elif self.cache is None: + self.cache = rhs.cache.clone() if rhs.cache is not None else None + + for field in rhs.fields: + if field in self.fields: + assert set(self._row_ids) == set(rhs._row_ids), f"Row ids must be the same. Got {self._row_ids} and {rhs._row_ids}" + continue + self.fields.add(field) + + return self + + @classmethod + def _cat(cls, data: list["RowRemoteBatch"]) -> "RowRemoteBatch": + assert data + if len(data) == 1: + return data[0] + + fields = data[0].fields + assert all(d.fields == fields for d in data), "All batches must have the same fields" + partition = data[0].partition + assert all(d.partition == partition for d in data), "All batches must have the same partition" + + row_ids = [row_id for d in data for row_id in d._row_ids] + + caches = [d.cache for d in data] + if all(c is not None for c in caches): + first_keys = set(caches[0].keys()) + if all(set(c.keys()) == first_keys for c in caches[1:]): + cache = TensorDict.cat(caches, dim=0) + else: + cache = None + else: + cache = None + + return RowRemoteBatch( + partition=partition, + device=data[0].device, + fields=fields, + row_ids=row_ids, + cache=cache, + ) + + +class PlanNode(ABC): + def __init__(self): + pass + + @property + @abstractmethod + def batch_size(self) -> int: + pass + + @abstractmethod + def execute(self, data): + pass + + +class SelectPlan(PlanNode): + def __init__(self, index: torch.Tensor): + super().__init__() + assert isinstance(index, torch.Tensor) and index.dim() == 1 + self.index = index + + @property + def batch_size(self): + return int(self.index.sum().item()) if self.index.dtype == torch.bool else self.index.shape[0] + + def execute(self, data): + return data[self.index.to(data.device)] + + +class SlicePlan(PlanNode): + def __init__(self, start: int, end: int, step: int, batch_size: int): + super().__init__() + self.slice_obj = slice(start, end, step) + self.source_batch_size = batch_size + + @property + def batch_size(self): + return len(range(*self.slice_obj.indices(self.source_batch_size))) + + def to_select(self): + start, stop, step = self.slice_obj.indices(self.source_batch_size) + return SelectPlan(np.arange(start, stop, step)) + + def execute(self, data): + return data[self.slice_obj] + + +class RepeatPlan(PlanNode): + def __init__(self, repeat_times: int, interleave: bool, source_batch_size: int): + super().__init__() + self.repeat_times = repeat_times + self.interleave = interleave + self.source_batch_size = source_batch_size + + @property + def batch_size(self): + return self.source_batch_size * self.repeat_times + + def execute(self, data): + assert isinstance(data, TensorDict) + if self.interleave: + return data.repeat_interleave(self.repeat_times) + else: + return data.repeat(self.repeat_times) + + +class CatPlan(PlanNode): + def __init__(self, batch_size: int): + super().__init__() + self._batch_size = batch_size + + @property + def batch_size(self): + return self._batch_size + + def execute(self, data): + assert isinstance(data, list) and all(isinstance(d, TensorDict) for d in data) + return TensorDict.cat(data, dim=0) + + +# TODO shigao: use Box to share materialized remote object +class Box: + """ + Can not used in asyncio context, threading.Lock will block event loop. + """ + + def __init__(self): + self.lock = threading.Lock() + self.value = None + + def get(self): + with self.lock: + return self.value + + def set(self, value): + with self.lock: + if self.value is None: + self.value = value + + +class ColumnRemoteBatch(RemoteBatch): + """ + A remote batch stored in a key-value store with column id as keys. + """ + + def __init__( + self, + partition: str, + device, + fields: dict[str, Any | list["ColumnRemoteBatch"]], + is_nested: bool, + cache: TensorDict, + batch_size: int, + pipeline: tuple[PlanNode] = tuple(), + ): + """ + fields contains any meta need to be hold and pass to transfer backend during get + or a list of ColumnRemoteBatch if is nested. + """ + super().__init__("column", partition, device) + self.fields = fields + self.is_nested = is_nested + self.cache = cache + self.batch_size = batch_size + self.pipeline = pipeline + + def __reduce__(self): + return ( + ColumnRemoteBatch, + (self.partition, self.device, self.fields, self.is_nested, None, self.batch_size, self.pipeline), + ) + + def __len__(self) -> int: + return self.batch_size + + def __delitem__(self, key: str): + self.fields.pop(key) + if self.cache is not None and key in self.cache: + del self.cache[key] + + def __contains__(self, key: str) -> bool: + return key in self.fields + + def clone(self, recurse: bool = True): + if self.is_nested: + cloned_fields = {k: [batch.clone() for batch in v] for k, v in self.fields.items()} + else: + cloned_fields = self.fields.copy() + return ColumnRemoteBatch( + partition=self.partition, + device=self.device, + fields=cloned_fields, + is_nested=self.is_nested, + cache=self.cache.clone(recurse=recurse) if self.cache is not None else None, + batch_size=self.batch_size, + pipeline=self.pipeline, + ) + + def keys(self): + return self.fields.keys() + + def to(self, device) -> "ColumnRemoteBatch": + super().to(device) + if self.cache is not None: + self.cache = self.cache.to(device) + return self + + def materialize(self, fields: list[str] = None) -> TensorDict: + if fields is None: + fields = self.fields.keys() + else: + assert set(fields) <= set(self.fields.keys()) + existing_fields = set(self.cache.keys()) if self.cache is not None else set() + + data = None + if self.is_nested: + assert len(self.pipeline) == 1 and isinstance(self.pipeline[0], CatPlan) + chunks: list["ColumnRemoteBatch"] = next(iter(self.fields.values())) + # TODO shigao: use batch get + # TODO: parallel gather + data: list[TensorDict] = [chunk.materialize(fields) for chunk in chunks] + else: + # pass column(key) meta back to transfer backend + fetch_fields = {field: self.fields[field] for field in fields if field not in existing_fields} + if len(fetch_fields) > 0: + data: TensorDict = transfer_backend.get( + partition=self.partition, keys=list(fetch_fields.keys()), fields=list(fetch_fields.values()) + ) + + if data is not None: + for operator in self.pipeline: + data = operator.execute(data) + assert len(data) == self.batch_size + + if self.device is not None: + data.to(self.device) + + if self.cache is None: + self.cache = data + else: + from roll.distributed.scheduler.protocol import union_tensor_dict + + self.cache = union_tensor_dict(self.cache, data) + + return self.cache.select(*fields) + + def drop(self): + transfer_backend.delete( + partition=self.partition, keys=list(self.fields.keys()), fields=list(self.fields.values()) + ) + + def select(self, fileds: list[str]) -> "ColumnRemoteBatch": + assert all(key in self.fields for key in fileds), f"Keys {fileds} not in {self.fields.keys()}" + fields = {key: self.fields[key] for key in fileds} + cache = self.cache + if cache is not None: + keys_in_cache = [k for k in fileds if k in cache.keys()] + if keys_in_cache: + cache = cache.select(*keys_in_cache) + else: + cache = None + return ColumnRemoteBatch( + partition=self.partition, + device=self.device, + fields=fields, + cache=cache, + batch_size=self.batch_size, + is_nested=self.is_nested, + pipeline=self.pipeline, + ) + + def _selection(self, plan: PlanNode) -> "ColumnRemoteBatch": + batch_size = plan.batch_size + + cache = self.cache + if cache is not None: + cache = plan.execute(cache) + assert isinstance(cache, TensorDict) + + if not self.pipeline: + pipeline = self.pipeline + (plan,) + else: + # TODO: heuristic optimization + # TODO: predicate pushdown + if isinstance(plan, SelectPlan) and isinstance(self.pipeline[-1], SelectPlan): + pipeline = self.pipeline + (plan,) + else: + pipeline = self.pipeline + (plan,) + + return ColumnRemoteBatch( + partition=self.partition, + device=self.device, + fields=self.fields, + cache=cache, + batch_size=batch_size, + is_nested=self.is_nested, + pipeline=pipeline, + ) + + def select_idxs(self, index: torch.Tensor | np.ndarray | list) -> "ColumnRemoteBatch": + assert isinstance(index, (torch.Tensor, np.ndarray, list)) + if isinstance(index, np.ndarray): + index = torch.from_numpy(index) + elif isinstance(index, list): + index = torch.tensor(index) + if index.dtype != torch.bool: + index = index.type(torch.int32) + + plan = SelectPlan(index) + return self._selection(plan) + + def slice( + self, + start: Optional[int] = None, + end: Optional[int] = None, + step: Optional[int] = None, + ) -> "ColumnRemoteBatch": + plan = SlicePlan(start, end, step, self.batch_size) + return self._selection(plan) + + def pop(self, fileds) -> "ColumnRemoteBatch": + assert len(fileds) == len(set(fileds)), "Fields must be unique" + assert set(fileds) <= self.fields.keys(), f"Fields {set(fileds) - self.fields.keys()} not in batch" + ret = self.select(fileds) + + for key in fileds: + del self.fields[key] + if self.cache is not None: + remaining_keys = [k for k in self.fields.keys() if k in self.cache.keys()] + if remaining_keys: + self.cache = self.cache.select(*remaining_keys) + else: + self.cache = None + + return ret + + def chunk(self, chunk_sizes: list[int]) -> list["ColumnRemoteBatch"]: + assert sum(chunk_sizes) == len( + self + ), f"Sum of chunk_sizes {sum(chunk_sizes)} does not match batch size {len(self)}" + chunks = [] + offset = 0 + for size in chunk_sizes: + chunks.append(self.slice(offset, offset + size)) + offset += size + return chunks + + def repeat(self, repeat_times: int, interleave: bool) -> "ColumnRemoteBatch": + plan = RepeatPlan(repeat_times, interleave, self.batch_size) + return self._selection(plan) + + def union(self, rhs: "ColumnRemoteBatch") -> "ColumnRemoteBatch": + assert isinstance(rhs, RemoteBatch) + assert len(self) == len(rhs), f"Two tensor dict must have identical batch size. Got {len(self)} and {len(rhs)}" + if self.cache is not None and rhs.cache is not None: + from roll.distributed.scheduler.protocol import union_tensor_dict + + union_tensor_dict(self.cache, rhs.cache) + elif self.cache is None: + self.cache = rhs.cache.clone() if rhs.cache is not None else None + + for field, value in rhs.fields.items(): + if field in self.fields: + # assert self.cache[field].equal(rhs.cache[field]), f"{field=}" + continue + self.fields[field] = value + + return self + + @classmethod + def _cat(cls, data: list["ColumnRemoteBatch"]) -> "ColumnRemoteBatch": + """ + ColumnRemoteBatch._cat will not check type and shape[1:] of fields. + """ + assert data + if len(data) == 1: + return data[0] + + keys = set(data[0].fields.keys()) + assert all(set(d.fields.keys()) == keys for d in data), "All batches must have the same fields" + partition = data[0].partition + assert all(d.partition == partition for d in data), "All batches must have the same partition" + device = data[0].device + + batch_size = sum(d.batch_size for d in data) + plan = CatPlan(batch_size) + + caches = [d.cache for d in data] + if all(c is not None for c in caches): + first_keys = set(caches[0].keys()) + if all(set(c.keys()) == first_keys for c in caches[1:]): + cache = plan.execute(caches) + else: + cache = None + else: + cache = None + + return ColumnRemoteBatch( + partition=partition, + device=device, + fields={field: data for field in keys}, + is_nested=True, + cache=cache, + batch_size=batch_size, + pipeline=(plan,), + ) diff --git a/roll/distributed/scheduler/resource_manager.py b/roll/distributed/scheduler/resource_manager.py index 3779827f7..ac9810f41 100644 --- a/roll/distributed/scheduler/resource_manager.py +++ b/roll/distributed/scheduler/resource_manager.py @@ -39,7 +39,6 @@ def __init__(self, num_gpus_per_node, num_nodes): if self.gpu_per_node > 0: assert self.num_gpus <= available_gpu, f"num_gpus {self.num_gpus} > available_gpu {available_gpu}" - bundles = [] for i in range(self.num_nodes): node = nodes_maybe_used[i] @@ -48,7 +47,6 @@ def __init__(self, num_gpus_per_node, num_nodes): self.placement_groups = [ray.util.placement_group([bundle]) for bundle in bundles] ray.get([pg.ready() for pg in self.placement_groups]) - gpu_ranks = ray.get([ get_visible_gpus.options( placement_group=pg, @@ -91,8 +89,7 @@ def nodes_placement_group(self, node_rank) -> PlacementGroup: return self.node2pg[node_rank] def destroy_placement_group(self): - for pg in self.placement_groups: - ray.util.remove_placement_group(pg) + [ray.util.remove_placement_group(pg) for pg in self.placement_groups] def allocate_placement_group(self, world_size, device_mapping: List[int] = None) -> List[List[Dict]]: """ diff --git a/roll/distributed/scheduler/rollout_scheduler.py b/roll/distributed/scheduler/rollout_scheduler.py index bd680e0bb..f5247b726 100644 --- a/roll/distributed/scheduler/rollout_scheduler.py +++ b/roll/distributed/scheduler/rollout_scheduler.py @@ -18,6 +18,12 @@ from roll.utils.functionals import append_to_dict from roll.utils.import_utils import safe_import_class from roll.utils.logging import get_logger +from roll.utils.telemetry import ( + attach_trace_context, + extract_trace_context, + inject_trace_context, + get_tracer, +) logger = get_logger() @@ -636,12 +642,27 @@ async def get_batch(self, data: DataProto, batch_size): if self._should_load_mock(global_step): return await self._load_mock_batch(global_step) + with ( + attach_trace_context(data.meta_info), + get_tracer("scheduler").start_as_current_span( + "get_batch", + attributes={ + "global_step": global_step, + "batch_size": batch_size, + }, + ), + ): + return await self._get_batch_impl(data, batch_size) + + async def _get_batch_impl(self, data: DataProto, batch_size): + global_step = data.meta_info["global_step"] + # start env manager if self.rollout_task is None: seed = random.randint(0, 1000000) if self.mode == "train" else self.config.seed self.rollout_task = asyncio.create_task(self._run_rollout_loop(seed)) - await asyncio.gather(*self.es_manager.update_step(global_step, blocking=False)) + await asyncio.gather(*self.es_manager.update_step(global_step, inject_trace_context({}), blocking=False)) await self.env_output_queue.advance_step.remote(global_step) await self.router_manager.resume.remote() @@ -673,6 +694,9 @@ async def get_batch(self, data: DataProto, batch_size): # DUMP MODE: Save merged batch (from mixin) await self._maybe_dump_batch(batch, global_step) + with get_tracer("scheduler").start_as_current_span("to_remote"): + loop = asyncio.get_running_loop() + batch = await loop.run_in_executor(None, DataProto.to_remote, batch) return batch async def shrink_sampler(self, target_gpus: List[int]) -> Dict[str, Any]: diff --git a/roll/distributed/scheduler/router.py b/roll/distributed/scheduler/router.py index 2f31ddd46..26c238ce0 100644 --- a/roll/distributed/scheduler/router.py +++ b/roll/distributed/scheduler/router.py @@ -70,7 +70,7 @@ def __init__(self, actor_cluster: Cluster, router_args: RouterArguments, num_gpu self.actor_cluster = actor_cluster self.workers = actor_cluster.workers - self.strategy_name = actor_cluster.worker_config.strategy_args.strategy_name + self.strategy_name = actor_cluster.worker_config.strategy_args.strategy_name self.model_path = download_model(actor_cluster.worker_config.model_args.model_name_or_path) self.tokenizer = default_tokenizer_provider(model_args=actor_cluster.worker_config.model_args) @@ -571,7 +571,17 @@ def _preprocess_generate(self, req: DataProto, request_id): if "multi_modal_data" in req.non_tensor_batch: multi_modal_data = req.non_tensor_batch["multi_modal_data"] assert len(multi_modal_data) == 1 - payload["multi_modal_data"] = multi_modal_data[0] + if 'multi_modal_data' in multi_modal_data[0] and 'video' in multi_modal_data[0]['multi_modal_data'] and self.strategy_name == 'sglang': + multi_modal_data = req.non_tensor_batch["multi_modal_inputs"] + assert len(multi_modal_data) == 1 + payload["multi_modal_data"] = {'multi_modal_data': {'video': multi_modal_data[0]}} + input_ids = req.batch["input_ids"] + attention_mask = req.batch["attention_mask"] + input_ids = gather_unpadded_input_ids(input_ids=input_ids, attention_mask=attention_mask) + payload["multi_modal_data"]["prompt_token_ids"] = input_ids[0] + else: + payload["multi_modal_data"] = multi_modal_data[0] + else: input_ids = req.batch["input_ids"] assert not collect_unfinished or input_ids.size(0) == 1 @@ -597,10 +607,19 @@ def _preprocess_generate(self, req: DataProto, request_id): def _postprocess_generate(self, req, response): output_data = DataProto(meta_info=req.meta_info) output_data.meta_info["finish_reasons"] = response["finish_reasons"] - output_data.meta_info["output_token_ids"] = response["output_token_ids"] + output_data.meta_info["output_token_ids"] = response.get("output_token_ids", None) output_data.meta_info["output_logprobs"] = response.get("output_logprobs", None) + # TODO: The size of routed_experts is [b * s * layer * topk]. + # For the 30A3 model, this data block is tens of MB in size. + # The serialization overhead of Ray transmission needs to be profiled again. + output_data.meta_info["routed_experts"] = response.get("routed_experts", None) output_data.meta_info["eos_token_id"] = [self.eos_token_id, self.pad_token_id] output_data.meta_info["pad_token_id"] = self.pad_token_id + + # Merge metrics from response (e.g., speculative decoding metrics) + if "metrics" in response: + output_data.meta_info.setdefault("metrics", {}).update(response["metrics"]) + return output_data async def generate_request(self, req: DataProto, request_id, uid): diff --git a/roll/distributed/scheduler/transfer_backend.py b/roll/distributed/scheduler/transfer_backend.py new file mode 100644 index 000000000..ca2b137ea --- /dev/null +++ b/roll/distributed/scheduler/transfer_backend.py @@ -0,0 +1,224 @@ +import os +import threading +import uuid +from typing import Any + +import ray +import torch +import numpy as np +import sys + +if sys.version_info < (3, 13): + import transfer_queue as tq +else: + tq = None +from omegaconf import OmegaConf +from tensordict import NonTensorStack, TensorDict + +from roll.configs.base_config import TransferBackendArguments +from roll.distributed.scheduler.storage import SharedStorage +from roll.utils.constants import STORAGE_NAME, RAY_NAMESPACE +from roll.utils.logging import get_logger + +logger = get_logger() + +# Global reference to keep SharedStorage actor alive +_shared_storage = None + + +def _check_transfer_queue_available(): + if tq is None: + raise ImportError( + "TransferQueue is not available on Python 3.13+. " + "Please use an alternative transfer backend or downgrade to Python <= 3.12." + ) + + +def init_transfer_backend(config: TransferBackendArguments | None): + global _shared_storage + + _shared_storage = SharedStorage.options( + name=STORAGE_NAME, get_if_exists=True, namespace=RAY_NAMESPACE + ).remote() + + if config is None: + config = TransferBackendArguments() + ray.get(_shared_storage.put.remote(key="transfer_backend_config", data=config)) + + backend_name = config.backend_name + backend_config = config.backend_config + if backend_name is None: + logger.info(f"Initialized dummy transfer backend: {config}") + elif backend_name == "TransferQueue": + _check_transfer_queue_available() + init_transfer_queue_server(backend_config) + logger.info(f"Initialized TransferQueue transfer backend: {config}") + else: + raise ValueError(f"Unsupported transfer backend: {backend_name}") + + +_client = None +_client_lock = threading.Lock() + +def reinit_after_fork(): + global _client, _client_lock + _client_lock = threading.Lock() + _client = None + +os.register_at_fork(after_in_child=reinit_after_fork) + +def init_client(): + global _client + if _client is not None: + return + with _client_lock: + if _client is not None: + return + shared_storage = ray.get_actor(name=STORAGE_NAME, namespace=RAY_NAMESPACE) + config = ray.get(shared_storage.get.remote(key="transfer_backend_config")) + assert config is not None + if config.backend_name is None: + _client = DummyClient() + elif config.backend_name == "TransferQueue": + _client = TransferQueueClient() + else: + raise ValueError(f"Unsupported transfer backend: {config.backend_name}") + logger.info(f"Initialized transfer client: {_client.__class__.__name__}") + + +def put(partition, row_ids: list[str], fields: dict[str, torch.Tensor | np.ndarray], batch_size: int): + init_client() + return _client.put(partition, row_ids, fields, batch_size) + +def get(partition, keys: list[str], fields: list[Any]): + init_client() + return _client.get(partition, keys, fields) + +def delete(partition, keys: list[str], fields: list[Any]): + init_client() + return _client.delete(partition, keys, fields) + + +def create_tensordict(fields: dict[str, torch.Tensor | np.ndarray]) -> TensorDict: + assert fields + td_dict = {} + batch_size = None + for key, val in fields.items(): + if isinstance(val, torch.Tensor): + td_dict[key] = val + elif isinstance(val, np.ndarray): + td_dict[key] = NonTensorStack(*val) + else: + raise TypeError(f"Unsupported type: {type(val)}") + if batch_size is None: + batch_size = val.shape[0] + elif batch_size != val.shape[0]: + raise ValueError("Batch size mismatch") + return TensorDict(td_dict, batch_size=[batch_size]) + + +class DummyClient: + + def put(self, partition, row_ids: list[str], fields: dict[str, torch.Tensor | np.ndarray], batch_size: int): + return None + + def get(self, partition, keys: list[str], fields: list[Any]): + raise RuntimeError("unexpected code path") + + def delete(self, partition, keys: list[str], fields: list[Any]): + raise RuntimeError("unexpected code path") + + +@ray.remote +class RayMemoryStoreServer: + def __init__(self): + super().__init__() + self.objects: dict[str, torch.Tensor | np.ndarray] = {} + + async def put(self, keys, values): + for key, data in zip(keys, values): + self.objects[key] = data + + async def get(self, keys): + return [self.objects[key] for key in keys] + + async def delete(self, keys): + for key in keys: + del self.objects[key] + + +class RayMemoryStoreClient: + def __init__(self): + self.client = RayMemoryStoreServer.options( + name="RayMemoryStore", + get_if_exists=True, + ).remote() + + def put(self, partition, row_ids: list[str], fields: dict[str, torch.Tensor | np.ndarray], batch_size: int): + # TODO move RayMemoryStoreClient to another file + from roll.distributed.scheduler.remote_protocol import ColumnRemoteBatch + + column_ids = [str(uuid.uuid4()) for _ in range(len(fields))] + ray.get(self.client.put.remote(keys=column_ids, values=list(fields.values()))) + + meta_dict = {field: column_id for field, column_id in zip(fields.keys(), column_ids)} + data = create_tensordict(fields) + assert len(data) == batch_size + return ColumnRemoteBatch( + partition=partition, + device=None, + fields=meta_dict, + is_nested=False, + cache=data, + batch_size=batch_size, + ) + + def get(self, partition, keys: list[str], fields: list[Any]): + data_list = ray.get(self.client.get.remote(fields)) + data_dict = {field: tensor for field, tensor in zip(keys, data_list)} + return create_tensordict(data_dict) + + def delete(self, partition, keys: list[str], fields: list[Any]): + pass + + +def init_transfer_queue_server(config): + # Must create enough storage units or may encounter: + # EncodeError: Can't encode Ext objects with data longer than 2**32 - 1. + # But also cannot set too many storage units that exceed the number of cores of ray cluster. + config = OmegaConf.create(config) + tq.init(config) + + +class TransferQueueClient: + def __init__(self): + _check_transfer_queue_available() + tq.init() + + def put(self, partition, row_ids: list[str], fields: dict[str, torch.Tensor | np.ndarray], batch_size: int): + # TODO move TransferQueueClient to another file + from roll.distributed.scheduler.remote_protocol import RowRemoteBatch + + data = create_tensordict(fields) + assert len(data) == batch_size + tq.kv_batch_put( + keys=row_ids, + fields=data, + partition_id=partition, + ) + return RowRemoteBatch( + partition=partition, + device=data.device, + fields=list(fields.keys()), + row_ids=row_ids, + cache=data, + ) + + def get(self, partition, keys: list[str], fields: list[Any]): + return tq.kv_batch_get(keys=keys, select_fields=fields, partition_id=partition) + + def delete(self, partition, keys: list[str], fields: list[Any]): + return tq.kv_clear(keys=keys, partition_id=partition) + + +__all__ = ["init_transfer_backend", "put", "get", "delete"] diff --git a/roll/distributed/scheduler/user_defined_rollout_loop.py b/roll/distributed/scheduler/user_defined_rollout_loop.py index aaf8756bc..8af7349bd 100644 --- a/roll/distributed/scheduler/user_defined_rollout_loop.py +++ b/roll/distributed/scheduler/user_defined_rollout_loop.py @@ -18,6 +18,7 @@ postprocess_generate, concatenate_input_and_output, ) +from roll.utils.telemetry import get_tracer from roll.utils.logging import get_logger @@ -72,6 +73,16 @@ def postprocess_paused_data(pre_data, data: DataProto, sequence_length, prompt_l data.meta_info[f"pre_{key}"] = pre_value new_batch = {**pre_data.batch} + if "routed_experts" in data.meta_info: + cur_routed_experts = data.meta_info.pop("routed_experts") + pre_routed_experts = pre_data.meta_info.get("pre_routed_experts", None) + if pre_routed_experts is not None: + # 在序列维度拼接 + data.meta_info["pre_routed_experts"] = torch.cat([pre_routed_experts, cur_routed_experts], dim=0) + else: + data.meta_info["pre_routed_experts"] = cur_routed_experts + new_batch["routed_experts"] = pre_data.batch.get("pre_routed_experts") + init_attention_mask = pre_data.batch.get("init_attention_mask", pre_data.batch["attention_mask"]) new_batch["init_attention_mask"] = init_attention_mask new_batch["init_input_ids"] = pre_data.batch.get("init_input_ids", pre_data.batch["input_ids"]) @@ -128,6 +139,15 @@ def postprocess_output_data(request, data: DataProto, sequence_length) -> DataPr pre_output_logprobs = request.meta_info.get("pre_output_logprobs", [[]] * len(output_token_ids)) output_logprobs = [pre_output_logprobs[i] + output_logprobs[i] for i in range(len(pre_output_logprobs))] + # new: process routed_experts + routed_experts = data.meta_info.get("routed_experts", None) + if routed_experts is not None: + # routed_experts : [B*S, num_expert, topk] + pre_routed_experts = request.meta_info.get("pre_routed_experts", None) + if pre_routed_experts is not None: + # B*S1, num_expert, topk] + [B*S2, num_expert, topk] -> [B*(S1+S2), num_expert, topk] + routed_experts = torch.cat([pre_routed_experts, routed_experts], dim=0) + output_tokens = [torch.tensor(token_ids) for token_ids in output_token_ids] output_tensor = pad_sequence(output_tokens, batch_first=True, padding_value=pad_token_id) output_tensor = concatenate_input_and_output( @@ -141,10 +161,14 @@ def postprocess_output_data(request, data: DataProto, sequence_length) -> DataPr eos_token_id=eos_token_id, pad_token_id=pad_token_id, output_logprobs=output_logprobs, + routed_experts=routed_experts, ) request_repeat = request.repeat(repeat_times=len(output_tokens)) output.non_tensor_batch = request_repeat.non_tensor_batch output.meta_info = request_repeat.meta_info + # Preserve metrics from data (e.g., speculative decoding metrics) + if "metrics" in data.meta_info: + output.meta_info.setdefault("metrics", {}).update(data.meta_info["metrics"]) return output # ================= example of user defined rollout loop ================= @@ -171,7 +195,7 @@ class UserDefinedRolloutLoop: The framework will only raise asyncio.CancelledError exception. (process_new_prompt will be called by scheduler as an asyncio.Task and scheduler may cancel this task if needed. User should not suppress asyncio.CancelledError exception and should handle clean up by themself.) - + User should catch all other exceptions, any other exceptions will be treat as sys.exit by framework. """ def __init__(self): @@ -218,6 +242,15 @@ async def _generate_and_reward( context: RolloutContext, req: DataProto, domain: str, + ): + with get_tracer("scheduler").start_as_current_span("generate_and_reward"): + return await self._generate_and_reward_impl(context=context, req=req, domain=domain) + + async def _generate_and_reward_impl( + self, + context: RolloutContext, + req: DataProto, + domain: str, ): responses: List[DataProto] = [] @@ -252,7 +285,10 @@ async def _generate_and_reward( req = postprocess_paused_data(req, data, context.sequence_length, context.prompt_length) rewards = await context.compute_rewards(req=req, domain=domain) + metrics = req.meta_info.pop("metrics", {}) req.union(rewards) + req_metrics = req.meta_info.pop("metrics", {}) + req.meta_info['metrics'] = {**metrics, **req_metrics} output_count = req.batch.batch_size[0] assert output_count == req.meta_info["generation_config"]["num_return_sequences"] diff --git a/roll/distributed/strategy/deepspeed_strategy.py b/roll/distributed/strategy/deepspeed_strategy.py deleted file mode 100644 index 51cf1b21d..000000000 --- a/roll/distributed/strategy/deepspeed_strategy.py +++ /dev/null @@ -1,590 +0,0 @@ -from collections import defaultdict -from contextlib import nullcontext -from datetime import timedelta -from typing import Callable, Dict, Tuple - -import deepspeed -import torch -import torch.distributed as dist -from codetiming import Timer -from deepspeed.ops.adam import DeepSpeedCPUAdam, FusedAdam -from deepspeed.runtime.zero import GatheredParameters -from deepspeed.runtime.zero.offload_config import OffloadStateTypeEnum -from peft import get_peft_model_state_dict -from transformers import get_scheduler, set_seed -from transformers.integrations import HfDeepSpeedConfig - -from roll.datasets.collator import collate_fn_to_dict_list -from roll.distributed.executor.worker import Worker -from roll.distributed.scheduler.protocol import DataProto -from roll.distributed.strategy.strategy import InferenceStrategy, TrainStrategy -from roll.models.model_providers import default_processor_provider, default_tokenizer_provider -from roll.platforms import current_platform -from roll.third_party.deepspeed.model_update import DeepSpeedWeightUpdater -from roll.third_party.deepspeed.offload_states_patch import bind_deepspeed_offload_states_func -from roll.utils.collective import collective -from roll.utils.context_parallel import get_ulysses_group, set_upg_manager -from roll.utils.deepspeed_utils import get_optimizer_grouped_parameters -from roll.utils.functionals import append_to_dict, entropy_from_logits, log_probs_from_logits -from roll.utils.constants import IGNORE_INDEX -from roll.utils.logging import get_logger -from roll.utils.offload_states import OffloadStateType - - -logger = get_logger() - - -class DeepSpeedInferStrategy(InferenceStrategy): - strategy_name = "deepspeed_infer" - - def __init__(self, worker: Worker): - super().__init__(worker) - self.worker_config.strategy_args.strategy_config["train_micro_batch_size_per_gpu"] = ( - self.worker_config.training_args.per_device_train_batch_size - ) - - # deepspeed的train_batch_size是全局batch_size - self.worker_config.strategy_args.strategy_config["train_batch_size"] = ( - self.worker_config.training_args.per_device_train_batch_size - * self.worker_config.training_args.gradient_accumulation_steps - * self.worker.world_size - ) - self.worker_config.strategy_args.strategy_config["gradient_clipping"] = ( - self.worker.pipeline_config.max_grad_norm - ) - self.ds_config = HfDeepSpeedConfig(self.worker_config.strategy_args.strategy_config) - - def initialize(self, model_provider): - set_seed(seed=self.worker.pipeline_config.seed) - - assert self.ds_config.is_zero3(), "deepspeed infer only supports zero = 3." - - deepspeed.init_distributed(timeout=timedelta(minutes=self.worker_config.backend_timeout)) - dist.all_reduce(torch.zeros(1).to(current_platform.device_type)) - - # apply Ulysses parallel - world_size = dist.get_world_size() - global_rank = dist.get_rank() - - if (cp_size := self.worker_config.model_args.ulysses_size) > 1: - if current_platform.apply_ulysses_patch() is not None: - set_upg_manager(ulysses_size=cp_size, rank=global_rank, world_size=world_size) - else: - cp_size = 1 - - self.worker.rank_info.dp_rank = global_rank // cp_size - self.worker.rank_info.dp_size = world_size // cp_size - self.worker.rank_info.cp_rank = global_rank % cp_size - self.worker.rank_info.cp_size = cp_size - - self.tokenizer = default_tokenizer_provider(model_args=self.worker_config.model_args) - self.processor = default_processor_provider(model_args=self.worker_config.model_args) - - model = model_provider(tokenizer=self.tokenizer, model_args=self.worker_config.model_args, is_trainable=False) - - try: - num_attention_heads, num_key_value_heads = model.config.num_attention_heads, model.config.num_key_value_heads - except AttributeError: - num_attention_heads, num_key_value_heads = ( - model.config.text_config.num_attention_heads, - model.config.text_config.num_key_value_heads, - ) - - assert num_attention_heads % cp_size == 0, ( - f"num_attention_heads {num_attention_heads} must be divisible by ulysses_size {cp_size}" - ) - assert num_key_value_heads % cp_size == 0 or cp_size % num_key_value_heads == 0, ( - f"num_key_value_heads {num_key_value_heads} must be divisible by ulysses_size " - f"{cp_size}or vise versa. Upon ulysses_size % num_key_value_heads == 0," - f"kv heads are repeated to ensure correctness." - ) - - logger.info(f"{self.model}") - - self.model, *_ = deepspeed.initialize( - model=model, - config=self.worker_config.strategy_args.strategy_config, - dist_init_required=True, - ) - - bind_deepspeed_offload_states_func(self.model) - - logger.info(f"{self.model}") - dist.barrier() - - def get_data_input(self, batch: DataProto): - def broadcast_obj(obj, group): - obj_list = [obj if dist.get_rank(group) == 0 else None] - src_rank = dist.get_process_group_ranks(group)[0] - dist.broadcast_object_list(obj_list, src=src_rank, group=group) - return obj_list[0] - # to avoid making side-effect on LLM, if want to broadcast non_tensor_batch, - # set _broadcast_non_tensor_batch into meta_info - broadcast_non_tensor_batch = batch.meta_info.get("_broadcast_non_tensor_batch", False) - if self.worker.rank_info.cp_size > 1: - if broadcast_non_tensor_batch: - tmp_batch = broadcast_obj(batch, get_ulysses_group()) - batch.batch = tmp_batch.batch - batch.non_tensor_batch = tmp_batch.non_tensor_batch - else: - batch.batch = broadcast_obj(batch.batch, get_ulysses_group()) - return batch - - def forward_step( - self, - batch: DataProto, - forward_func: Callable[[DataProto, torch.Tensor], Tuple[torch.Tensor, Dict[str, torch.Tensor]]], - ) -> Dict[str, torch.Tensor]: - self.model.eval() - batch_size = batch.batch.batch_size[0] - micro_batch_size = batch.meta_info["micro_batch_size"] - num_microbatches = max(batch_size // micro_batch_size, 1) - micro_batches = batch.chunk(chunks=num_microbatches) - - cp_size = self.worker.rank_info.cp_size - batch_num_tokens = self._get_batch_num_tokens(batch) - batch.meta_info['batch_num_tokens'] = {k: v // cp_size for k, v in batch_num_tokens.items()} - global_valid_tokens = self._get_global_valid_samples(batch) - batch.meta_info['global_valid_samples'] = {k: v // cp_size for k, v in global_valid_tokens.items()} - - loss_scale = num_microbatches * self.worker.rank_info.dp_size - - disable_adapter = batch.meta_info.get("disable_adapter", False) - adapter_context = self.unwrap_model().disable_adapter() if disable_adapter else nullcontext() - losses_reduced = [] - with adapter_context: - for data in micro_batches: - input_ids = data.batch["input_ids"] - attention_mask = data.batch["attention_mask"] - position_ids = data.batch["position_ids"] - forward_args = data.meta_info.get("forward_args", {}) - if position_ids.dim() == 3: - # same as megatron to be compatible with fsdp packing which change position_ids.size(1) to 4 - if position_ids.size(1) == 4: - position_ids = position_ids[:, 1:, :].contiguous() # (bsz, 4, seqlen) -> (bsz, 3, seqlen) - # qwen2vl mrope, maybe use a placeholder and let model generate position_ids - position_ids = position_ids.transpose(0, 1) # (bsz, 3, seqlen) -> (3, bsz, seqlen) - if "multi_modal_inputs" in data.non_tensor_batch: - multi_modal_inputs = data.non_tensor_batch["multi_modal_inputs"] - multi_modal_data = defaultdict(list) - # mm inputs of some samples would be empty to allow text and mm - # mixed data - for sample_mm_inputs in multi_modal_inputs: - for key in sample_mm_inputs.keys(): - multi_modal_data[key].append(sample_mm_inputs[key]) - for key in multi_modal_data.keys(): - assert key not in forward_args - # DataProto.to('cuda') in upper frame not work for non_tensor_batch - forward_args[key] = torch.concat(multi_modal_data[key], dim=0).to(input_ids.device) - forward_args.update({"force_vit_image": True}) - - if self.worker.rank_info.cp_size > 1: - splited_features = self.get_feature_on_cp_rank(input_ids, attention_mask, position_ids) - input_ids = splited_features["input_ids"] - attention_mask = splited_features["attention_mask"] - position_ids = splited_features["position_ids"] - - # set use_cache=False manually for the same reason as HfInferStrategy - output = self.model( - input_ids=input_ids, - attention_mask=attention_mask, - position_ids=position_ids, - use_cache=False, - **forward_args, - ) - loss, loss_reduced = forward_func(data, output.logits) - if self.worker_config.apply_loss_scale: - loss *= loss_scale - losses_reduced.append(loss_reduced) - results = collate_fn_to_dict_list(losses_reduced) - return results - - def get_feature_on_cp_rank( - self, input_ids: torch.Tensor, attention_mask: torch.Tensor = None, position_ids: torch.Tensor = None - ): - seqlens_in_batch = input_ids.size(1) - assert seqlens_in_batch % self.worker.rank_info.cp_size == 0, ( - f"input_length={seqlens_in_batch} not divisible by cp_size={self.worker.rank_info.cp_size}" - ) - cp_middle_rank_len = seqlens_in_batch // self.worker.rank_info.cp_size - padded_input_ids = input_ids - result = {} - start_index = cp_middle_rank_len * self.worker.rank_info.cp_rank - end_index = cp_middle_rank_len * (self.worker.rank_info.cp_rank + 1) - result["input_ids"] = padded_input_ids[:, start_index:end_index] - if attention_mask is not None: - result["attention_mask"] = attention_mask[:, start_index:end_index] - if position_ids is not None: - if position_ids.dim() == 3: - result["position_ids"] = position_ids[:, :, start_index:end_index] - else: - result["position_ids"] = position_ids[:, start_index:end_index] - return result - - def generate(self, batch: DataProto, generation_config): - input_ids = batch.batch["input_ids"] # (bs, prompt_length) - attention_mask = batch.batch["attention_mask"] # left-padded attention_mask - - output = self.model.generate( - input_ids=input_ids, - attention_mask=attention_mask, - use_cache=True, - **generation_config, - ) - - return output - - def unwrap_model(self): - return self.model.module - - # 参数同步相关接口 - def broadcast_parameter(self, model_update_name, src_pp_rank, dtype, shape, parameter_name, is_lora=False): - comm_plan = self.model_update_comm_plan[model_update_name][src_pp_rank] - weight = torch.empty(shape, dtype=dtype, device=current_platform.device_type) - collective.broadcast(tensor=weight, src_rank=0, group_name=comm_plan["group_name"]) - param = self.model.get_parameter(parameter_name) - if not self.ds_config.is_zero3(): - param.data.copy_(weight.to("cpu")) - else: - with GatheredParameters([param], modifier_rank=0): - if dist.get_rank() == 0: - param.data.copy_(weight) - del weight - - # offload/load 相关接口 - def load_states(self, include=None, non_blocking=False): - if include is not None: - ds_include = [] - if OffloadStateType.model_params in include: - ds_include.append(OffloadStateTypeEnum.lp_params) - if OffloadStateType.other_params in include: - ds_include.append(OffloadStateTypeEnum.hp_params) - ds_include.append(OffloadStateTypeEnum.lp_grads) - ds_include.append(OffloadStateTypeEnum.contiguous_grad_buffer) - if OffloadStateType.optimizer_states in include: - ds_include.append(OffloadStateTypeEnum.optim_states) - include = ds_include - self.model.reload_states(include=include, non_blocking=non_blocking) - - def offload_states(self, include=None, non_blocking=False): - if include is not None: - ds_include = [] - if OffloadStateType.model_params in include: - ds_include.append(OffloadStateTypeEnum.lp_params) - if OffloadStateType.other_params in include: - ds_include.append(OffloadStateTypeEnum.hp_params) - ds_include.append(OffloadStateTypeEnum.lp_grads) - ds_include.append(OffloadStateTypeEnum.contiguous_grad_buffer) - if OffloadStateType.optimizer_states in include: - ds_include.append(OffloadStateTypeEnum.optim_states) - include = ds_include - - self.model.offload_states(include=include, non_blocking=non_blocking) - current_platform.empty_cache() - - def op_compute_log_probs(self, logits: torch.Tensor, input_ids: torch.Tensor, attention_mask: torch.Tensor): - """ - input_ids [[p, p, r, r, r, 0, 0]] p: prompt, r: response, 0: pad - response_mask [[0, 0, 1, 1, 1, 0, 0]] - """ - labels: torch.Tensor = input_ids[:, 1:].clone() - labels[attention_mask[:, 1:] == 0] = 0 # avoid invalid token id - # TODO: don't pad here but process this shift after generation - labels = torch.cat([labels, torch.zeros_like(labels[:, :1])], dim=1) - if self.worker.rank_info.cp_size > 1: - labels = self.get_feature_on_cp_rank(labels)["input_ids"] - log_probs = log_probs_from_logits(logits, labels) - if self.worker.rank_info.cp_size > 1: - with torch.no_grad(): - all_log_probs = [torch.empty_like(log_probs) for _ in range(self.worker.rank_info.cp_size)] - dist.all_gather(all_log_probs, log_probs, group=get_ulysses_group()) - all_log_probs[self.worker.rank_info.cp_rank] = log_probs - log_probs = torch.cat(all_log_probs, dim=1) - log_probs = log_probs[:, :-1] * attention_mask[:, 1:] - return log_probs - - def op_compute_entropy(self, logits: torch.Tensor, attention_mask: torch.Tensor): - entropy = entropy_from_logits(logits) - if self.worker.rank_info.cp_size > 1: - with torch.no_grad(): - all_entropy = [torch.empty_like(entropy) for _ in range(self.worker.rank_info.cp_size)] - dist.all_gather(all_entropy, entropy, group=get_ulysses_group()) - all_entropy[self.worker.rank_info.cp_rank] = entropy - entropy = torch.cat(all_entropy, dim=1) - entropy = entropy[:, :-1] * attention_mask[:, 1:] - return entropy - - -class DeepSpeedTrainStrategy(DeepSpeedInferStrategy, TrainStrategy): - strategy_name = "deepspeed_train" - - def initialize(self, model_provider): - assert self.ds_config._stage > 0, "deepspeed train only supports zero > 0." - - set_seed(seed=self.worker.pipeline_config.seed) - deepspeed.init_distributed(timeout=timedelta(minutes=self.worker_config.backend_timeout)) - dist.all_reduce(torch.zeros(1).to(current_platform.device_type)) - - # apply Ulysses parallel - world_size = dist.get_world_size() - global_rank = dist.get_rank() - - if (cp_size := self.worker_config.model_args.ulysses_size) > 1: - current_platform.apply_ulysses_patch() - set_upg_manager(ulysses_size=cp_size, rank=global_rank, world_size=world_size) - - self.worker.rank_info.dp_rank = global_rank // cp_size - self.worker.rank_info.dp_size = world_size // cp_size - self.worker.rank_info.cp_rank = global_rank % cp_size - self.worker.rank_info.cp_size = cp_size - - - self.tokenizer = default_tokenizer_provider(model_args=self.worker_config.model_args) - self.processor = default_processor_provider(model_args=self.worker_config.model_args) - - self.weight_updaters = {} - - model = model_provider(tokenizer=self.tokenizer, model_args=self.worker_config.model_args, is_trainable=True) - - if cp_size > 1: - try: - num_attention_heads, num_key_value_heads = model.config.num_attention_heads, model.config.num_key_value_heads - except AttributeError: - num_attention_heads, num_key_value_heads = ( - model.config.text_config.num_attention_heads, - model.config.text_config.num_key_value_heads, - ) - - assert num_attention_heads % cp_size == 0, ( - f"num_attention_heads {num_attention_heads} must be divisible by ulysses_size {cp_size}" - ) - assert num_key_value_heads % cp_size == 0 or cp_size % num_key_value_heads == 0, ( - f"num_key_value_heads {num_key_value_heads} must be divisible by ulysses_size " - f"{cp_size}or vise versa. Upon ulysses_size % num_key_value_heads == 0," - f"kv heads are repeated to ensure correctness." - ) - - adam_optimizer = DeepSpeedCPUAdam if self.ds_config.is_offload() else FusedAdam - optim_params = get_optimizer_grouped_parameters( - model, weight_decay=self.worker_config.training_args.weight_decay - ) - optimizer = adam_optimizer( - optim_params, - lr=self.worker_config.training_args.learning_rate, - betas=(self.worker_config.training_args.adam_beta1, self.worker_config.training_args.adam_beta2), - ) - - logger.info(f"max steps pipeline {self.worker_config.training_args.max_steps}") - self.worker_config.training_args.max_steps = ( - self.worker_config.training_args.max_steps // self.worker.rank_info.dp_size - ) - logger.info(f"max steps worker train {self.worker_config.training_args.max_steps}") - - scheduler = get_scheduler( - self.worker_config.training_args.lr_scheduler_type, - optimizer, - num_warmup_steps=self.worker_config.training_args.get_warmup_steps( - self.worker_config.training_args.max_steps - ), - num_training_steps=self.worker_config.training_args.max_steps, - ) - - self.model, self.optimizer, _, self.scheduler = deepspeed.initialize( - model_parameters=model.parameters(), - model=model, - optimizer=optimizer, - lr_scheduler=scheduler, - config=self.worker_config.strategy_args.strategy_config, - dist_init_required=True, - ) - bind_deepspeed_offload_states_func(self.model) - - logger.info(f"{self.model}") - dist.barrier() - - def op_compute_language_loss(self, logits: torch.Tensor, labels: torch.Tensor): - """ - Override for DeepSpeed strategy: compute language loss from logits. - - In DeepSpeed strategy with HuggingFace models, the model returns logits - (not loss like in Megatron strategy where labels are passed to the model). - - Note: DataCollatorForSFT already shifts labels (shift_feature=True by default), - so logits and labels are already aligned. Do NOT shift again here. - - Args: - logits: Model output logits [batch_size, seq_len, vocab_size] - labels: Pre-shifted labels [batch_size, seq_len], already aligned with logits - - Returns: - loss: Scalar loss tensor - metrics: Dict - """ - # Labels already shifted by DataCollator, directly compute cross-entropy - loss = torch.nn.functional.cross_entropy( - logits.view(-1, logits.size(-1)), - labels.view(-1), - ignore_index=IGNORE_INDEX - ) - metrics = {f"{self.worker_config.name}/loss@sum": loss.detach().float().unsqueeze(0)} - return loss, metrics - - def train_step( - self, - batch: DataProto, - loss_func: Callable[[DataProto, torch.Tensor], Tuple[torch.Tensor, Dict[str, torch.Tensor]]], - ): - self.model.train() - mini_batch_size = self.worker_config.training_args.per_device_train_batch_size - mini_steps = batch.batch.batch_size[0] // self.worker_config.training_args.per_device_train_batch_size - - cp_size = self.worker.rank_info.cp_size - batch_num_tokens = self._get_batch_num_tokens(batch) - batch.meta_info['batch_num_tokens'] = {k: v // cp_size for k, v in batch_num_tokens.items()} - global_valid_tokens = self._get_global_valid_samples(batch) - batch.meta_info['global_valid_samples'] = {k: v // cp_size for k, v in global_valid_tokens.items()} - - loss_scale = mini_steps * self.worker.rank_info.dp_size - batch.meta_info['micro_batch_size'] = mini_batch_size - - data_iter = batch.make_iterator(mini_batch_size=mini_batch_size, epochs=1) - metrics = {} - - for step in range(mini_steps): - data: DataProto = next(data_iter) - input_ids = data.batch["input_ids"] - attention_mask = data.batch["attention_mask"] - position_ids = data.batch["position_ids"] - forward_args = data.meta_info.get("forward_args", {}) - # TODO: The offload option may be integrated into the pipeline config in the future. - is_offload_optimizer_states_in_train_step = data.meta_info.get("is_offload_optimizer_states_in_train_step", True) - if position_ids.dim() == 3: - # qwen2vl mrope, maybe use a placeholder and let model generate position_ids - position_ids = position_ids.transpose(0, 1) # (bsz, 3, seqlen) -> (3, bsz, seqlen) - if "multi_modal_inputs" in data.non_tensor_batch: - multi_modal_inputs = data.non_tensor_batch["multi_modal_inputs"] - multi_modal_data = defaultdict(list) - # mm inputs of some samples would be empty to allow text and mm - # mixed data - for sample_mm_inputs in multi_modal_inputs: - for key in sample_mm_inputs.keys(): - multi_modal_data[key].append(sample_mm_inputs[key]) - for key in multi_modal_data.keys(): - assert key not in forward_args - # DataProto.to('cuda') in upper frame not work for non_tensor_batch - forward_args[key] = torch.concat(multi_modal_data[key], dim=0).to(input_ids.device) - forward_args.update({"force_vit_image": True}) - - if self.worker.rank_info.cp_size > 1: - splited_features = self.get_feature_on_cp_rank(input_ids, attention_mask, position_ids) - input_ids = splited_features["input_ids"] - attention_mask = splited_features["attention_mask"] - position_ids = splited_features["position_ids"] - - output = self.model( - input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, **forward_args - ) - loss, loss_reduced = loss_func(data, output.logits) - append_to_dict(metrics, loss_reduced) - loss *= self.worker.rank_info.cp_size - if self.worker_config.apply_loss_scale: - loss *= loss_scale - self.model.backward(loss) - - is_gradient_accumulation_boundary = self.model.is_gradient_accumulation_boundary() - if is_gradient_accumulation_boundary: - self.load_states(include=[OffloadStateType.optimizer_states]) - self.model.step() - if is_gradient_accumulation_boundary: - # global_grad_norm is calculated in optimizer.step thus put it - # into metrics after optimizer.step - metrics.update({self.worker_config.name + "/" + "grad_norm": self.model.get_global_grad_norm().item()}) - if is_offload_optimizer_states_in_train_step: - self.offload_states(include=[OffloadStateType.optimizer_states], non_blocking=True) - return metrics - - def save_checkpoint(self, save_dir, global_step, ckpt_id, tag="checkpoint", local_state_path=None, is_last_step=None, **kwargs): - """ - save ckpt/hf model/tokenizer to local dir - save_dir/actor_train/{hf files} - save_dir/actor_train/checkpoint/{checkpoint files} - """ - logger.info(f"save_dir: {save_dir}") - if local_state_path is None: - local_state_path = save_dir - - with Timer("load") as load_timer: - self.load_states() - - if self.ds_config.is_zero3(): - if self.model.zero_gather_16bit_weights_on_model_save(): - state_dict = self.model._zero3_consolidated_16bit_state_dict() - else: - raise ValueError( - "Cannot get 16bit model weights because `stage3_gather_16bit_weights_on_model_save` in DeepSpeed config is False. " - "To save the model weights in 16bit, set `stage3_gather_16bit_weights_on_model_save` to True in DeepSpeed config file or " - "set `zero3_save_16bit_model` to True when using `accelerate config`. " - "To save the full checkpoint, run `model.save_checkpoint(save_dir)` and use `zero_to_fp32.py` to recover weights." - ) - else: - from deepspeed.checkpoint.utils import clone_tensors_for_torch_save - - state_dict = clone_tensors_for_torch_save(self.model.module.state_dict()) - - # save huggingface pretrained model - if dist.get_rank() == 0: - self.model.module.save_pretrained(save_dir, state_dict=state_dict, safe_serialization=False) - self.tokenizer.save_pretrained(save_dir) - if getattr(self, "processor", None): - self.processor.save_pretrained(save_dir) - # save tokenizer - # DeepSpeedEngine.load_checkpoint method doesn't take an is_last_step argument - kwargs.pop("is_last_step", None) - self.model.save_checkpoint(save_dir, tag=tag, **kwargs) - - if self.worker_config.checkpoint_config.get("async_upload", True) and not is_last_step: - self.thread_executor.submit(self.checkpoint_manager.upload, ckpt_id=ckpt_id, local_state_path=local_state_path) - else: - self.checkpoint_manager.upload(ckpt_id=ckpt_id, local_state_path=local_state_path) - - metrics = { - "load": load_timer.last, - } - return metrics - - def load_checkpoint(self, load_dir, tag="checkpoint", **kwargs): - logger.info(f"load checkpoint from {load_dir}") - self.model.load_checkpoint(load_dir, tag=tag, **kwargs) - - def collect_lora_params(self): - peft_model = self.unwrap_model() - if not self.ds_config.is_zero3(): - lora_state_dict = get_peft_model_state_dict(peft_model) - return lora_state_dict - - adapter_name = "default" - state_dict = peft_model.state_dict() - lora_state_dict = {k: state_dict[k] for k in state_dict if ("lora_" in k and adapter_name in k)} - - lora_params = [] - for name, param in lora_state_dict.items(): - lora_params.append((name.replace(f".{adapter_name}", ""), peft_model.get_parameter(name))) - - del lora_state_dict - return lora_params - - def setup_model_update(self, infer_cluster, model_update_name: str): - assert model_update_name not in self.weight_updaters - is_lora = self.worker_config.model_args.lora_target is not None - self.weight_updaters[model_update_name] = DeepSpeedWeightUpdater( - pipeline_config=self.worker.pipeline_config, - infer_cluster=infer_cluster, - worker_config=self.worker_config, - model_update_name=model_update_name, - model=self.unwrap_model(), - ds_config=self.ds_config, - is_lora=is_lora, - ) - - def model_update(self, model_update_name: str): - return self.weight_updaters[model_update_name].model_update() diff --git a/roll/distributed/strategy/diffusion_strategy.py b/roll/distributed/strategy/diffusion_strategy.py deleted file mode 100644 index 52a8c3e03..000000000 --- a/roll/distributed/strategy/diffusion_strategy.py +++ /dev/null @@ -1,83 +0,0 @@ -import os -import torch -import torchvision -import torch.distributed as dist - -from typing import Callable, Dict, Tuple -from codetiming import Timer - -from roll.distributed.strategy.deepspeed_strategy import DeepSpeedTrainStrategy as BaseDeepSpeedTrainStrategy -from roll.distributed.scheduler.protocol import DataProto -from roll.utils.functionals import append_to_dict -from roll.utils.logging import get_logger -from roll.utils.offload_states import OffloadStateType - - -logger = get_logger() - - -class DeepSpeedTrainStrategy(BaseDeepSpeedTrainStrategy): - - strategy_name = "diffusion_deepspeed_train" - - def train_step( - self, - batch: DataProto, - loss_func: Callable[[DataProto, torch.Tensor], Tuple[torch.Tensor, Dict[str, torch.Tensor]]], - ): - mini_batch_size = self.worker_config.training_args.per_device_train_batch_size - data_iter = batch.make_iterator(mini_batch_size=mini_batch_size, epochs=1) - mini_steps = batch.batch.batch_size[0] // self.worker_config.training_args.per_device_train_batch_size - metrics = {} - - for step in range(mini_steps): - data: DataProto = next(data_iter) - - # convert data to dict for DiffusionTrainingModule - prompt = data.non_tensor_batch["prompt"][0] - video = list(torch.unbind(data.batch["video"][0], dim=0)) - video = [torchvision.transforms.functional.to_pil_image(v) for v in video] - data = {"prompt": prompt, "video": video} - - loss, face_score, kl_loss = self.model(data) - loss, loss_reduced = loss_func(data, loss, face_score, kl_loss) - append_to_dict(metrics, loss_reduced) - - self.model.backward(loss) - - is_gradient_accumulation_boundary = self.model.is_gradient_accumulation_boundary() - if is_gradient_accumulation_boundary: - self.load_states(include=[OffloadStateType.optimizer_states]) - self.model.step() - if is_gradient_accumulation_boundary: - # global_grad_norm is calculated in optimizer.step thus put it - # into metrics after optimizer.step - metrics.update({self.worker_config.name + "/" + "grad_norm": self.model.get_global_grad_norm().item()}) - return metrics - - def offload_states(self): - pass - - def save_checkpoint(self, save_dir, global_step, ckpt_id, tag="checkpoint", local_state_path=None, **kwargs): - assert not self.ds_config.is_zero3(), "zero3 is not supported yet" - - logger.info(f"save_dir: {save_dir}") - if local_state_path is None: - local_state_path = save_dir - - with Timer("load") as load_timer: - self.load_states() - - from deepspeed.checkpoint.utils import clone_tensors_for_torch_save - state_dict = clone_tensors_for_torch_save(self.unwrap_model().state_dict()) - state_dict = self.unwrap_model().export_trainable_state_dict(state_dict, remove_prefix='pipe.dit2.') - - # save DiffusionTrainingModule - if dist.get_rank() == 0: - os.makedirs(local_state_path, exist_ok=True) - torch.save(state_dict, os.path.join(local_state_path, "diffusion_module.pth")) - - metrics = { - "load": load_timer.last, - } - return metrics diff --git a/roll/distributed/strategy/factory.py b/roll/distributed/strategy/factory.py index a83dcf0f7..5b84c0b1b 100644 --- a/roll/distributed/strategy/factory.py +++ b/roll/distributed/strategy/factory.py @@ -15,13 +15,7 @@ def create_strategy(worker: Worker, sync_wrapper: bool = False) -> Union[Inferen strategy_name = worker.worker_config.strategy_args.strategy_name # Lazy import strategy to avoid cuda initialized - if strategy_name == "deepspeed_infer": - from roll.distributed.strategy.deepspeed_strategy import DeepSpeedInferStrategy as strategy_cls - elif strategy_name == "deepspeed_train": - from roll.distributed.strategy.deepspeed_strategy import DeepSpeedTrainStrategy as strategy_cls - elif strategy_name == "diffusion_deepspeed_train": - from roll.distributed.strategy.diffusion_strategy import DeepSpeedTrainStrategy as strategy_cls - elif strategy_name == "hf_infer": + if strategy_name == "hf_infer": from roll.distributed.strategy.hf_strategy import HfInferStrategy as strategy_cls elif strategy_name == "vllm": from roll.distributed.strategy.vllm_strategy import VllmStrategy as strategy_cls diff --git a/roll/distributed/strategy/fsdp2_strategy.py b/roll/distributed/strategy/fsdp2_strategy.py index 723043313..695b0195a 100644 --- a/roll/distributed/strategy/fsdp2_strategy.py +++ b/roll/distributed/strategy/fsdp2_strategy.py @@ -4,21 +4,31 @@ from collections import defaultdict from contextlib import nullcontext from typing import Callable, Dict, Optional, Tuple +from abc import abstractmethod, ABC +import accelerate import numpy as np import ray +import transformers +from transformers import AutoConfig, get_scheduler, set_seed +from codetiming import Timer +from packaging import version import torch +from torch import optim +from torch import Tensor import torch.distributed as dist import torch.distributed.checkpoint as dcp -from codetiming import Timer -from torch import optim from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict -from torch.distributed.device_mesh import init_device_mesh +from torch.distributed.device_mesh import init_device_mesh, DeviceMesh from torch.distributed.fsdp import CPUOffloadPolicy, MixedPrecisionPolicy -from torch.distributed.tensor import DTensor, distribute_tensor +from torch.distributed.tensor import DTensor, distribute_tensor, distribute_module, Shard +from torch.distributed.tensor.parallel import parallelize_module, ParallelStyle +from torch.distributed._functional_collectives import ( + all_to_all_single, + all_to_all_single_autograd, +) from torch.nn.utils import clip_grad_norm_ from torch.nn.utils.clip_grad import _clip_grads_with_norm_, _get_total_norm -from transformers import AutoConfig, get_scheduler, set_seed from roll.datasets.collator import collate_fn_to_dict_list from roll.distributed.executor.worker import Worker @@ -36,20 +46,21 @@ from roll.utils.collective import collective from roll.utils.context_parallel import get_ulysses_group, set_upg_manager from roll.utils.context_parallel.autograd_gather import ulysses_gather -from roll.utils.context_parallel.rmpad_ulysses import ( - gather_outputs_and_unpad, - ulysses_pad_and_slice_inputs, - ulysses_pad_inputs, -) from roll.utils.fsdp_utils import ( apply_fsdp2, fsdp2_load_full_state_dict, get_init_weight_context_manager, get_shard_placement_fn, + get_shard_placement_fn_ep, + _permute, + _unpermute, + register_experts_forward_in_ExpertsInterface, + set_use_grouped_mm, ) from roll.utils.functionals import append_to_dict, log_probs_from_logits from roll.utils.logging import get_logger from roll.utils.offload_states import OffloadStateType +from roll.utils.constants import IGNORE_INDEX logger = get_logger() @@ -89,7 +100,7 @@ def _parse_dtype(dtype): return dtype -def create_device_mesh_with_ulysses(world_size: int, fsdp_size: int): +def create_device_mesh_with_ep(world_size: int, fsdp_size: int, efsdp_size: int, ep_size: int): """ Create device mesh for FSDP. """ @@ -107,12 +118,207 @@ def create_device_mesh_with_ulysses(world_size: int, fsdp_size: int): mesh_shape = (ddp_size, fsdp_size) mesh_dim_names = ["ddp", "fsdp"] - return init_device_mesh( + device_mesh = init_device_mesh( current_platform.device_type, mesh_shape=mesh_shape, mesh_dim_names=mesh_dim_names, ) + # Create device mesh for MoE + if ep_size > 1: + moe_model_size = efsdp_size * ep_size + if moe_model_size == world_size: + moe_mesh_shape = (efsdp_size, ep_size) + moe_mesh_dim_names = ("efsdp", "ep") + else: + eddp_size = world_size // moe_model_size + moe_mesh_shape = (eddp_size, efsdp_size, ep_size) + moe_mesh_dim_names = ["eddp", "efsdp", "ep"] + + moe_device_mesh = init_device_mesh( + current_platform.device_type, + mesh_shape=moe_mesh_shape, + mesh_dim_names=moe_mesh_dim_names, + ) + else: + moe_device_mesh = None + + return device_mesh, moe_device_mesh + + +class BaseExpertParallel(ParallelStyle, ABC): + """ + Mirror torchtitan's BaseExpertParallel. + Reference: https://github.com/pytorch/torchtitan/blob/main/torchtitan/distributed/expert_parallel.py #L30 #v0.2.2 + """ + @abstractmethod + def _partition_fn(self, name: str, mod: torch.nn.Module, device_mesh: DeviceMesh) -> None: + ... + + @abstractmethod + def _token_dispatch( + self, mod: torch.nn.Module, inputs: tuple, device_mesh: DeviceMesh + ) -> tuple[Tensor, Tensor]: + ... + + @abstractmethod + def _token_combine( + self, mod: torch.nn.Module, routed_output: Tensor, device_mesh: DeviceMesh + ) -> Tensor: + ... + + +class ExpertParallel(BaseExpertParallel): + """ + Mirror torchtitan's BaseExpertParallel. + Reference: https://github.com/pytorch/torchtitan/blob/main/torchtitan/distributed/expert_parallel.py #L89 #v0.2.2 + """ + def __init__(self, num_experts: int): + super().__init__() + self.input_splits = None + self.output_splits = None + self.input_shape = None + self.permuted_indices = None + self.num_experts = num_experts + + def _partition_fn(self, name: str, mod: torch.nn.Module, device_mesh: DeviceMesh) -> None: + for param_name, param in mod.named_parameters(recurse=False): + dist_param = torch.nn.Parameter(distribute_tensor(param, device_mesh, [Shard(0)])) + mod.register_parameter(param_name, dist_param) + + def _token_dispatch( + self, mod: torch.nn.Module, inputs: tuple, device_mesh: DeviceMesh + ) -> tuple[Tensor, Tensor]: + # Preprocess + x, selected_experts_indices, top_scores = inputs + self.x = x + self.top_scores = top_scores + self.bs_slen, self.dim = x.shape + self.top_k = top_scores.shape[-1] + num_tokens_per_expert = torch.histc( + selected_experts_indices.view(-1), + bins=self.num_experts, + min=0, + max=self.num_experts, + ) + self.token_indices_experts_sorted = token_indices_experts_sorted = torch.argsort( + selected_experts_indices.view(-1), stable=True + ) + top_scores_experts_sorted = top_scores.view(-1)[token_indices_experts_sorted] + # shape (bs*slen*top_k, dim) + routed_input = x[token_indices_experts_sorted // self.top_k] + + # annotate module input placements/sharding with input_layouts + ep_degree = device_mesh.shape[0] + num_local_experts = num_tokens_per_expert.shape[0] // ep_degree + + # generate the input splits and output splits for all-to-all + with torch.no_grad(): + num_tokens_per_expert_group = all_to_all_single( + num_tokens_per_expert, + None, + None, + group=device_mesh.get_group(), + ) + # Need to wait explicitly because it is used by a triton kernel later + # which doesn't realize that AsyncCollectiveTensor needs unwrapping + num_tokens_per_expert_group = torch.ops._c10d_functional.wait_tensor( + num_tokens_per_expert_group + ) + input_splits = ( + num_tokens_per_expert.view(ep_degree, -1) + .sum(dim=1) + .to(torch.device("cpu"), non_blocking=True) + ) + # NOTE: this would incur a device-to-host sync + output_splits = ( + num_tokens_per_expert_group.view(ep_degree, -1) + .sum(dim=1) + .to(torch.device("cpu"), non_blocking=False) + ) + self.input_splits = input_splits.tolist() + self.output_splits = output_splits.tolist() + + # perform all-to-all + routed_input = all_to_all_single_autograd( + routed_input, + self.output_splits, + self.input_splits, + device_mesh.get_group(), + ) + + # NOTE: After this all-to-all, the routed input is put on proper EP rank. + # However, the num_tokens_per_expert_group is not of the final target format + # [#tokens for local expert 0, #tokens for local expert 1, ...] + # Rather, it is of the format + # [#tokens for local expert 0 from EP rank 0, #tokens for local expert 1 from EP rank 0, ..., + # #tokens for local expert 0 from EP rank 1, #tokens for local expert 1 from EP rank 1, ...] + # We need to perform another shuffle to get the correct layout, via the _permute function + # below, which also does padding to make sure the number of tokens each expert gets locally + # is a multiple of TOKEN_GROUP_ALIGN_SIZE_M. + # Note that this will create side effects when wrapping the for-loop implementation + # of GroupedExperts, as it does not need padding. + + ( + self.input_shape, + routed_input, + self.permuted_indices, + num_tokens_per_expert_group, + ) = _permute( + routed_input, num_tokens_per_expert_group, ep_degree, num_local_experts + ) + + return routed_input, num_tokens_per_expert_group + + def _token_combine( + self, mod: torch.nn.Module, routed_output: Tensor, device_mesh: DeviceMesh + ) -> Tensor: + routed_output = _unpermute( + routed_output, self.input_shape, self.permuted_indices + ) + + routed_output = all_to_all_single_autograd( + routed_output, + self.input_splits, + self.output_splits, + device_mesh.get_group(), + ) + + # Postprocess + routed_output_unsorted = torch.zeros( + (self.bs_slen * self.top_k, self.dim), + dtype=routed_output.dtype, + device=routed_output.device, + ) + routed_output_unsorted[self.token_indices_experts_sorted] = routed_output + routed_output_unsorted = routed_output_unsorted.reshape( + -1, self.top_k, self.dim + ) + out_experts = ( + torch.bmm( + self.top_scores.reshape(-1, 1, self.top_k), + routed_output_unsorted.to(self.top_scores.dtype), + ) + .to(self.x.dtype) + .squeeze(1) + ) + + # Do cloning here to ensure that experts return a not-view tensor. + # View tensor + in-place modification later can break pre-backward hook of fsdp. + # For example, in qwen3.5, in-place modification is applied to the output of experts: expert_output += shared_expert_output + return out_experts.clone() + + def _apply(self, module: torch.nn.Module, device_mesh: DeviceMesh) -> torch.nn.Module: + return distribute_module( + module, + device_mesh, + partition_fn=self._partition_fn, + # pyrefly: ignore [bad-argument-type] + input_fn=self._token_dispatch, + # pyrefly: ignore [bad-argument-type] + output_fn=self._token_combine, + ) + class FSDP2StrategyBase(InferenceStrategy): def __init__(self, worker: Worker): @@ -327,9 +533,6 @@ def save_checkpoint(self, save_dir, global_step, ckpt_id, tag="checkpoint", loca with Timer("load", logger=None) as load_timer: self.load_states() - dcp_checkpoint_dir = self._get_dcp_checkpoint_dir(save_dir) - os.makedirs(dcp_checkpoint_dir, exist_ok=True) - with Timer("hf_save", logger=None) as hf_timer: full_state_options = self._get_dcp_state_dict_options(full_state_dict=True) full_model_state = get_model_state_dict( @@ -348,13 +551,23 @@ def save_checkpoint(self, save_dir, global_step, ckpt_id, tag="checkpoint", loca if getattr(self, "processor", None): self.processor.save_pretrained(save_dir) - with Timer("dcp_save", logger=None) as dcp_timer: - self._save_checkpoint_with_dcp(checkpoint_dir=dcp_checkpoint_dir, is_last_step=is_last_step) + if self.save_only_model and dist.is_initialized(): + dist.barrier() - # PumpkinComment: - # If DCP save is async, uploading (which may copy+delete the local dir) must not start - # until the async save has fully finished writing checkpoint shards. - dcp_save_future = self.checkpoint_future if (self.async_save_strategy and not is_last_step) else None + dcp_save_time = 0 + dcp_save_future = None + if not self.save_only_model: + dcp_checkpoint_dir = self._get_dcp_checkpoint_dir(save_dir) + os.makedirs(dcp_checkpoint_dir, exist_ok=True) + + with Timer("dcp_save", logger=None) as dcp_timer: + self._save_checkpoint_with_dcp(checkpoint_dir=dcp_checkpoint_dir, is_last_step=is_last_step) + dcp_save_time = dcp_timer.last + + # PumpkinComment: + # If DCP save is async, uploading (which may copy+delete the local dir) must not start + # until the async save has fully finished writing checkpoint shards. + dcp_save_future = self.checkpoint_future if (self.async_save_strategy and not is_last_step) else None checkpoint_config = getattr(self.worker_config, "checkpoint_config", None) or {} async_upload = checkpoint_config.get("async_upload", True) @@ -396,9 +609,16 @@ def _on_dcp_done(fut): keep_local_file=keep_local_file, ) + # When cpu_offload is enabled, restore optimizer states back to CPU after saving. + # Without this, optimizer.step() will fail due to device mismatch (optimizer on GPU, grads on CPU). + with Timer("offload", logger=None) as offload_timer: + if self.cpu_offload_enabled: + self.offload_states() + return { "load": load_timer.last, - "dcp_save": dcp_timer.last, + "offload": offload_timer.last, + "dcp_save": dcp_save_time, "hf_save": hf_timer.last, } @@ -610,17 +830,31 @@ def _prepare_fsdp2_model( self.worker_config.model_args.compute_dtype = torch_dtype fsdp_size = self.worker_config.strategy_args.strategy_config.get("fsdp_size", 1) - if cp_size > 1 and (fsdp_size <= 1 or fsdp_size >= world_size): - fsdp_size = world_size // cp_size - self.worker_config.strategy_args.strategy_config["fsdp_size"] = fsdp_size - if global_rank == 0: - logger.info(f"CP enabled: auto-setting fsdp_size={fsdp_size} so ddp_size==cp_size for hybrid sharding") - elif fsdp_size != world_size: - logger.warning(f"fsdp_size {fsdp_size} is not equal to world_size {world_size}, using world_size instead") - fsdp_size = world_size - - self.worker_config.strategy_args.strategy_config["fsdp_size"] = fsdp_size - self.device_mesh = create_device_mesh_with_ulysses(world_size=world_size, fsdp_size=fsdp_size) + assert fsdp_size >= cp_size, "fsdp size should be greater than cp size." + + # Get expert parallel config + efsdp_size = self.worker_config.strategy_args.strategy_config.get("efsdp_size", 1) + ep_size = self.worker_config.strategy_args.strategy_config.get("ep_size", 1) + moe_use_grouped_mm = self.worker_config.strategy_args.strategy_config.get("moe_use_grouped_mm", False) + self.ep_enabled = ep_size > 1 + + # Check efsdp, ep size + if self.ep_enabled: + assert version.parse(transformers.__version__) >= version.parse("5.2.0") + + moe_model_size = efsdp_size * ep_size + assert moe_model_size > 0 and moe_model_size <= world_size and world_size % moe_model_size == 0, \ + f"efsdp_size * ep_size {moe_model_size} must be in [1, world_size({world_size})], and divisible by world_size {world_size}" + + if moe_use_grouped_mm and self.ep_enabled: + set_use_grouped_mm(True) + + # Patch ExpertsInterface to offer more options of forward function in transformers's MoeExperts + if self.ep_enabled: + register_experts_forward_in_ExpertsInterface() + + # Create device mesh + self.device_mesh, self.moe_device_mesh = create_device_mesh_with_ep(world_size=world_size, fsdp_size=fsdp_size, efsdp_size=efsdp_size, ep_size=ep_size) model_name_or_path = download_model(self.worker_config.model_args.model_name_or_path) config = AutoConfig.from_pretrained( @@ -631,7 +865,10 @@ def _prepare_fsdp2_model( self._validate_ulysses_compat(config, cp_size) - use_meta_tensor = not getattr(config, "tie_word_embeddings", False) + use_meta_tensor = not getattr(config, "tie_word_embeddings", False) and not self.ep_enabled + # accelerate v1.7.0 don't support _is_hf_initialized which is needed by use_meta_tensor + if version.parse(accelerate.__version__) == version.parse("1.7.0"): + use_meta_tensor = False init_context = get_init_weight_context_manager( use_meta_tensor=use_meta_tensor, mesh=self.device_mesh, @@ -649,8 +886,28 @@ def _prepare_fsdp2_model( self.is_lora = self.worker_config.model_args.lora_target is not None + # Set _experts_implementation + layers = self._get_layers(model) + + if self.ep_enabled: + layers[0].mlp.experts.config._experts_implementation = "ep" + elif moe_use_grouped_mm: + assert version.parse(transformers.__version__) >= version.parse("5.2.0"), "moe_use_grouped_mm requires transformers>=5.2.0" + layers[0].mlp.experts.config._experts_implementation = "grouped_mm" + return model, torch_dtype, cp_size + def _get_layers(self, model): + if self.is_lora: # PeftModel + base_model = model.base_model.model.model + else: + base_model = model.model + + if hasattr(base_model, "layers"): + return base_model.layers + else: # multi-modal model + return base_model.language_model.layers + @staticmethod def _validate_ulysses_compat(config, cp_size: int): try: @@ -725,6 +982,10 @@ def initialize(self, model_provider): default_model_dtype=torch.bfloat16, ) + # Initialize expert parallel model (must be before FSDP2 wrapping) + if self.ep_enabled: + self.apply_moe_ep(model) + self.setup_fsdp2_configuration() self.initialize_fsdp2_model(model) @@ -737,6 +998,7 @@ def setup_fsdp2_configuration(self): self.async_save_strategy = async_save_strategy if self.async_save_strategy: self.checkpoint_future = None + self.save_only_model = self.worker_config.strategy_args.strategy_config.get("save_only_model", False) # Get mixed precision settings from config param_dtype = self.worker_config.strategy_args.strategy_config.get("param_dtype", torch.bfloat16) @@ -764,6 +1026,15 @@ def setup_fsdp2_configuration(self): offload_policy_cfg = self.worker_config.strategy_args.strategy_config.get("offload_policy", False) self.cpu_offload_enabled = bool(offload_policy_cfg) + + # Perform reduce scatter during gradient accumulation. + # Default to True, because it decrease cuda memory usage and increase training speed according to tests. + self.reduce_scatter_during_grad_accumulation = bool( + self.worker_config.strategy_args.strategy_config.get("reduce_scatter_during_grad_accumulation", True) + ) + if self.reduce_scatter_during_grad_accumulation: + logger.info("[FSDP2] reduce_scatter_during_grad_accumulation is ENABLED") + offload_policy = None if self.cpu_offload_enabled: offload_policy = CPUOffloadPolicy( @@ -781,13 +1052,40 @@ def setup_fsdp2_configuration(self): fsdp_size=self.worker_config.strategy_args.strategy_config.get("fsdp_size", 1) ), } + self.moe_fsdp_config = self.fsdp_config.copy() + if self.ep_enabled: + self.moe_fsdp_config["mesh"] = self.moe_device_mesh["eddp", "efsdp"] if "eddp" in self.moe_device_mesh.mesh_dim_names else self.moe_device_mesh["efsdp"] + self.moe_fsdp_config["shard_placement_fn"] = get_shard_placement_fn_ep( + efsdp_size=self.worker_config.strategy_args.strategy_config.get("efsdp_size", 1), + ) + + def apply_moe_ep(self, model): + layers = self._get_layers(model) + + num_experts = layers[0].mlp.experts.num_experts + experts_plan = ExpertParallel(num_experts=num_experts) + experts_mesh = self.moe_device_mesh["ep"] + for layer in layers: + if hasattr(layer.mlp, "experts"): + parallelize_module( + module=layer.mlp.experts, + device_mesh=experts_mesh, + parallelize_plan=experts_plan, + ) def initialize_fsdp2_model(self, model): offload_policy = self.fsdp_config["offload_policy"] - full_state = model.state_dict() + + # Use nn.Module.state_dict() instead of model.state_dict() to get standard FQN keys + # (e.g. 'pretrained_model.model.layers.0...') that match what PyTorch's + # _iterate_valid_model_state() returns internally. Custom state_dict() methods + # (like AutoModelForCausalLMWithValueHead's) may strip prefixes, causing key + # mismatches in _broadcast_state_dict during FSDP2 initialization. + full_state = torch.nn.Module.state_dict(model) apply_fsdp2( model, self.fsdp_config, + self.moe_fsdp_config, self.worker_config.strategy_args.strategy_config, self.is_lora, ) @@ -797,10 +1095,13 @@ def initialize_fsdp2_model(self, model): full_state, self.device_mesh, offload_policy, + self.ep_enabled, ) self.model = model + # Add torch.no_grad() to disable gradient calculation. + # torch.no_grad() in pipeline cannot work because pipeline and worker are in different processes. @torch.no_grad() def forward_step( self, @@ -851,6 +1152,15 @@ def forward_step( multi_modal_data[key].append(sample_mm_inputs[key]) for key in multi_modal_data.keys(): assert key not in forward_args + mm_data = multi_modal_data[key] + # All mm fields are not padded in collator currently and some should be padded first + # pixel_values/pixel_values_videos for images/videos are concated for packing + # input_features for audios with shape `(bs, freqs, frames)` should be padded before concat + need_padding = any(t.shape[-1] != mm_data[0].shape[-1] for t in mm_data[1:]) + if need_padding: # input_features/feature_attention_mask + max_mm_len = max(t.shape[-1] for t in mm_data) + for i, t in enumerate(mm_data): + mm_data[i] = torch.nn.functional.pad(t, (0, max_mm_len - t.shape[-1]), "constant", 0) # DataProto.to('cuda') in upper frame not work for non_tensor_batch forward_args[key] = torch.concat(multi_modal_data[key], dim=0).to(input_ids.device) forward_args.update({"force_vit_image": True}) @@ -1063,6 +1373,10 @@ def initialize(self, model_provider): ) logger.info(f"max steps worker train {self.worker_config.training_args.max_steps}") + # Initialize expert parallel model + if self.ep_enabled: + self.apply_moe_ep(model) + # Setup FSDP-2 configuration self.setup_fsdp2_configuration() @@ -1119,18 +1433,104 @@ def _requires_grad_sync_context(self, set_sync_fn): set_sync_fn(True) def _clip_grad_norm(self, max_norm: float): - if not self.cpu_offload_enabled: + if not self.cpu_offload_enabled and not self.ep_enabled: grad_norm = clip_grad_norm_( self.model.parameters(), max_norm=max_norm, ) - else: + elif not self.cpu_offload_enabled: + grad_norm = self._clip_grad_norm_with_ep(max_norm) + elif not self.ep_enabled: grad_norm = self._clip_grad_norm_cpu_offload(max_norm) + else: + grad_norm = self._clip_grad_norm_with_ep(max_norm, cpu_offload_enabled=True) if isinstance(grad_norm, DTensor): grad_norm = grad_norm.full_tensor() return grad_norm + + @torch.no_grad() + def _clip_grad_norm_with_ep(self, max_norm: float, cpu_offload_enabled: bool = False): + """ + Reference to torchtitan's _clip_grad_norm_with_ep: + Reference: https://github.com/pytorch/torchtitan/blob/main/torchtitan/distributed/utils.py #L479 #v0.2.2 + """ + parameters = list(self.model.parameters()) + ep_params = [] + non_ep_params = [] + ep_grads = [] + non_ep_grads = [] + + for p in parameters: + if p.grad is None: + continue + assert isinstance(p, DTensor) and isinstance(p.grad, DTensor) + mesh_dim_names = p.device_mesh.mesh_dim_names + assert mesh_dim_names is not None + if "efsdp" in mesh_dim_names: + # Divide the gradient of the moe parameter by ep to ensure that the gradient is calculated from data on a single rank, + # keeping it aligned with fsdp. + ep_size = self.worker_config.strategy_args.strategy_config.get("ep_size", 1) + p.grad.div_(ep_size) + ep_params.append(p) + ep_grads.append(p.grad) + else: + non_ep_params.append(p) + non_ep_grads.append(p.grad) + + # Either list can be empty depending on the parallelization strategy: + # - In torchtitan with separate dense/sparse meshes, both lists are typically non-empty + # - In autoparallel, all params may live on a single sparse mesh with "ep" dimension, + # so non_ep_grads would be empty + # - In PP + EP setups, certain PP ranks may only own EP or non-EP layers + ep_grads_total_norm = _get_total_norm( + ep_grads, + norm_type=2.0, + error_if_nonfinite=False, + foreach=None, + ) + # get_total_norm returns tensor(0.) for empty list, which is a non-DTensor + if isinstance(ep_grads_total_norm, DTensor): + ep_grads_total_norm = ep_grads_total_norm.full_tensor() + + non_ep_grads_total_norm = _get_total_norm( + non_ep_grads, + norm_type=2.0, + error_if_nonfinite=False, + foreach=None, + ) + # get_total_norm returns tensor(0.) for empty list, which is a non-DTensor + if isinstance(non_ep_grads_total_norm, DTensor): + non_ep_grads_total_norm = non_ep_grads_total_norm.full_tensor() + + # move norm scalar to GPU + if cpu_offload_enabled: + ep_grads_total_norm = ep_grads_total_norm.to(current_platform.current_device(), non_blocking=True) + non_ep_grads_total_norm = non_ep_grads_total_norm.to(current_platform.current_device(), non_blocking=True) + + # Calculate total norm + norm_type = 2.0 + total_norm = ( + ep_grads_total_norm**norm_type + non_ep_grads_total_norm**norm_type + ) + total_norm **= 1.0 / norm_type + + # Do clipping + _clip_grads_with_norm_( + ep_params, + max_norm=max_norm, + total_norm=total_norm, + foreach=None, + ) + _clip_grads_with_norm_( + non_ep_params, + max_norm=max_norm, + total_norm=total_norm, + foreach=None, + ) + + return total_norm def _clip_grad_norm_cpu_offload(self, max_norm: float): """ @@ -1217,7 +1617,15 @@ def train_step( # PumpkinComment: # model.no_sync is replaced by model.set_requires_gradient_sync(False) in FSDP2 # but also add support for model.no_sync for compatibility - sync_context = contextlib.nullcontext() + # + # When reduce_scatter_during_grad_accumulation is enabled, sync every step to decrease cuda memory usage. + # Reference: https://docs.pytorch.org/docs/stable/fsdp.html#torch.distributed.fsdp.FullyShardedDataParallel.no_sync + if self.reduce_scatter_during_grad_accumulation: + sync_context = contextlib.nullcontext() + else: + sync_context = ( + self._grad_accumulation_context() if not sync_boundary and not no_sync else contextlib.nullcontext() + ) with ( sync_context, @@ -1265,6 +1673,12 @@ def train_step( self.scheduler.step() self.optimizer.zero_grad(set_to_none=True) + + # Log cuda memory + max_memory_allocated = torch.cuda.memory.max_memory_allocated() / (1024 ** 3) + max_memory_reserved = torch.cuda.memory.max_memory_reserved() / (1024 ** 3) + metrics[f"system/max_memory_allocated@max"] = max_memory_allocated + metrics[f"system/max_memory_reserved@max"] = max_memory_reserved return metrics def setup_model_update(self, infer_cluster, model_update_name: str): @@ -1281,3 +1695,74 @@ def setup_model_update(self, infer_cluster, model_update_name: str): def model_update(self, model_update_name: str): return self.weight_updaters[model_update_name].model_update() + + def get_labels_on_cp_rank( + self, + labels: torch.Tensor, + ): + """Get labels for specific context parallel rank""" + seqlens_in_batch = labels.size(1) + assert ( + seqlens_in_batch % self.worker.rank_info.cp_size == 0 + ), f"input_length={seqlens_in_batch} not divisible by cp_size={self.worker.rank_info.cp_size}" + cp_middle_rank_len = seqlens_in_batch // self.worker.rank_info.cp_size + padded_labels = labels + start_index = cp_middle_rank_len * self.worker.rank_info.cp_rank + end_index = cp_middle_rank_len * (self.worker.rank_info.cp_rank + 1) + labels_this_cp_rank = padded_labels[:, start_index:end_index] + return labels_this_cp_rank + + def op_compute_language_loss(self, logits: torch.Tensor, labels: torch.Tensor, batch_num_tokens: int): + """ + Override for FSDP2 strategy: compute language loss from logits. + + In FSDP2 strategy with HuggingFace models, the model returns logits + (not per-token loss like in Megatron strategy where labels are passed to the model). + + Note: DataCollatorForSFT already shifts labels (shift_feature=True by default), + so logits and labels are already aligned. Do NOT shift again here. + + Args: + logits: Model output logits [batch_size, seq_len, vocab_size] + labels: Pre-shifted labels [batch_size, seq_len], already aligned with logits + batch_num_tokens: Number of valid tokens for loss normalization + + Returns: + loss: Scalar loss tensor + metrics: Dict + """ + cp_size = self.worker.rank_info.cp_size + # Slice labels to match the sharded logits + if cp_size > 1: + labels = self.get_labels_on_cp_rank(labels) + labels = labels.contiguous() + + # Compute per-token loss mask + loss_mask = (labels != IGNORE_INDEX).float() + loss_mask = loss_mask.view(-1).float() + + # Compute cross entropy loss per token (reduction='none' to get per-token losses) + per_token_losses = torch.nn.functional.cross_entropy( + logits.view(-1, logits.size(-1)), + labels.view(-1), + ignore_index=IGNORE_INDEX, + reduction='none', # Get per-token loss + ) + + # Apply loss mask and sum + losses_sum = torch.sum(per_token_losses.view(-1) * loss_mask) + + # All-reduce across CP ranks + if cp_size > 1: + loss_info = torch.cat([losses_sum.view(1)]) + dist.all_reduce( + loss_info, op=dist.ReduceOp.SUM, group=get_ulysses_group() + ) + losses_sum = loss_info[0] + + # Normalize by batch_num_tokens + loss = losses_sum.clone() / batch_num_tokens # clone to make sure loss is not a view + + metrics = {f"{self.worker_config.name}/loss@sum": loss.clone().detach().item()} + + return loss, metrics diff --git a/roll/distributed/strategy/hf_strategy.py b/roll/distributed/strategy/hf_strategy.py index a775a0cd7..26b815e2d 100644 --- a/roll/distributed/strategy/hf_strategy.py +++ b/roll/distributed/strategy/hf_strategy.py @@ -3,7 +3,6 @@ from datetime import timedelta from typing import Callable, Dict, List, Optional, Tuple -import deepspeed import torch import torch.distributed as dist from accelerate import cpu_offload_with_hook diff --git a/roll/distributed/strategy/megatron_strategy.py b/roll/distributed/strategy/megatron_strategy.py index 4eeb4cc74..a50d29b11 100644 --- a/roll/distributed/strategy/megatron_strategy.py +++ b/roll/distributed/strategy/megatron_strategy.py @@ -32,11 +32,15 @@ get_moe_layer_wise_logging_tracker, reduce_aux_losses_tracker_across_ranks, ) +from megatron.core.transformer.moe.router_replay import ( + RouterReplay, + RouterReplayAction, +) from megatron.core.transformer.multi_token_prediction import MTPLossLoggingHelper from transformers.utils import is_peft_available from mcore_adapter import TrainingArguments -from mcore_adapter.checkpointing import get_checkpoint_dir, load_state_dict_from_checkpoint +from mcore_adapter.checkpointing import generate_model_state_dict, get_checkpoint_dir, load_state_dict_from_checkpoint from mcore_adapter.parallel_functions import context_parallel_gather, vocab_parallel_logprobs from mcore_adapter.patcher import patch_torch_find_nd_overlapping_shards, patch_torch_validate_global_plan from mcore_adapter.trainer.utils import build_sharded_state_dict_metadata, get_megatron_lr_scheduler @@ -46,7 +50,9 @@ from roll.distributed.strategy.strategy import InferenceStrategy, TrainStrategy from roll.models.model_providers import default_processor_provider, default_tokenizer_provider from roll.platforms import current_platform +from roll.third_party.megatron.compile_warmup import compile_warmup_pipeline_stages from roll.third_party.megatron.model_update import MegatronWeightUpdater +from roll.third_party.megatron.mtp_patcher import patch_mtp_functions from roll.third_party.megatron.offload_states_patch import ( MegatronOffloadStateType, bind_megatron_offload_states_func, @@ -54,7 +60,13 @@ reload_megatron_no_grad_module, ) from roll.third_party.megatron.optimizer import get_megatron_optimizer +from roll.third_party.megatron.router_replay_utils import ( + RouterReplayHelper, + merge_router_topk_indices, + set_router_replay_data, +) from roll.third_party.megatron.tensor_parallel import vocab_parallel_entropy +from roll.third_party.megatron.util import unwrap_model from roll.utils.constants import ( DIST_OPTIMIZER_DIR, IGNORE_INDEX, @@ -66,7 +78,7 @@ from roll.utils.dynamic_batching import make_micro_batch_iter_for_dynamic_batching from roll.utils.functionals import adjust_sequence_length, append_to_dict, reduce_metrics from roll.utils.logging import get_logger -from roll.utils.offload_states import OffloadStateType +from roll.utils.offload_states import OffloadStateType, clear_memory from roll.utils.sequence_packing import make_micro_batch_iter_for_sequence_packing, restore_results_order @@ -85,6 +97,8 @@ class MegatronInferStrategy(InferenceStrategy): strategy_name = "megatron_infer" def __init__(self, worker: Worker): + # Apply MTP patches BEFORE model instantiation + patch_mtp_functions() #TODO remove the patches when the latest pytorch version > v2.9.1 patch_torch_find_nd_overlapping_shards() patch_torch_validate_global_plan() @@ -104,8 +118,24 @@ def __init__(self, worker: Worker): # hard to impl with offload states assert not self.megatron_train_args.overlap_param_gather, "overlap_param_gather is not supported" + # Router Replay config + self.router_replay_config = self.worker_config.router_replay + self.enable_router_replay = (self.router_replay_config.mode in ["R2", "R3"]) + self.router_replay_mode = self.router_replay_config.mode + + # Force enable moe_enable_routing_replay when router replay is enabled, + # so that RouterReplay instances are created in Megatron MoE layers. + if self.enable_router_replay: + self.megatron_train_args.moe_enable_routing_replay = True + + # store router replay data + if self.enable_router_replay and self.router_replay_mode == "R2": + self.router_topk_indices_list = [] + logger.info("Router Replay R2 mode: RECORD enabled in MegatronInferStrategy") + def initialize(self, model_provider): self.tokenizer = default_tokenizer_provider(model_args=self.worker_config.model_args) + self.processor = default_processor_provider(model_args=self.worker_config.model_args) self.model: "VirtualModels" = model_provider( tokenizer=self.tokenizer, model_args=self.worker_config.model_args, @@ -114,8 +144,14 @@ def initialize(self, model_provider): ) self.model.config.finalize_model_grads_func = finalize_model_grads + # Inject mtp_training_mode from WorkerConfig to model config + if hasattr(self.worker_config, "mtp_training_mode"): + self.model.config.mtp_training_mode = self.worker_config.mtp_training_mode + self.models_unwrapped = self.model.get_models() self.forward_backward_func = get_forward_backward_func() + self.is_multimodal = self.processor is not None + self._validate_vlm_packing_support() self.seq_length = self.worker.pipeline_config.sequence_length @@ -132,9 +168,33 @@ def initialize(self, model_provider): self.model.config.variable_seq_lengths = True logger.info("Set variable_seq_lengths to True when use dynamic batching and pipeline parallel.") + if self.enable_router_replay and self.router_replay_mode == "R2": + # R2 mode: init router_replay_action=RouterReplayAction.RECORD + RouterReplay.set_global_router_replay_action(RouterReplayAction.RECORD) + logger.info(f"{self.model.get_models()}") dist.barrier() + def _validate_vlm_packing_support(self): + """Check if the VLM model supports sequence packing via MultimodalEmbeddingMixin. + + If use_sequence_packing is enabled but the multimodal model does not + inherit MultimodalEmbeddingMixin, automatically disable packing and + log a warning. + """ + if not (self.use_sequence_packing and self.is_multimodal): + return + from mcore_adapter.models.sequence_packing_mixin import MultimodalEmbeddingMixin + model_supports_packing = isinstance(self.models_unwrapped[0], MultimodalEmbeddingMixin) + if not model_supports_packing: + logger.warning( + "use_sequence_packing is enabled but the multimodal model does not " + "inherit MultimodalEmbeddingMixin. Disabling sequence packing for " + "this model. To enable packing, the model must inherit " + "MultimodalEmbeddingMixin." + ) + self.use_sequence_packing = False + def get_data_input(self, batch: DataProto): def broadcast_obj(obj, group): obj_list = [obj if dist.get_rank(group) == 0 else None] @@ -170,6 +230,15 @@ def forward_step( forward_func: Callable[[DataProto, torch.Tensor], Tuple[torch.Tensor, Dict[str, torch.Tensor]]], ) -> Dict[str, torch.Tensor]: self.model.eval() + + # R3 mode: set router replay action for compute_log_probs forward + if self.enable_router_replay and self.router_replay_mode == "R3": + if "routed_experts" in batch.batch: + RouterReplay.set_global_router_replay_action(RouterReplayAction.REPLAY_FORWARD) + else: + # ref model or no routed_experts in batch, skip router replay + RouterReplay.clear_global_router_replay_action() + batch.meta_info['batch_num_tokens'] = self._get_batch_num_tokens(batch, dp_group=mpu.get_data_parallel_group()) batch.meta_info['global_valid_samples'] = self._get_global_valid_samples(batch, dp_group=mpu.get_data_parallel_group()) @@ -226,6 +295,11 @@ def forward_step( self.worker_config.sequence_packing_args) + # R3 mode: clear router replay action and indices after compute_log_probs forward + if self.enable_router_replay and self.router_replay_mode == "R3": + RouterReplay.clear_global_router_replay_action() + RouterReplay.clear_global_indices() + if not ( ((self.worker.rank_info.tp_rank == 0 and self.worker.rank_info.cp_rank == 0) or output_on_all_tp_cp_ranks) @@ -412,30 +486,55 @@ def inner_forward_step(self, loss_func, data_iterator: Iterator[DataProto], mode labels = data.batch["labels"] if "labels" in data.batch else None # labels is only used for sft packed_seq_params = None + # Get loss_mask early, priority: final_response_mask > response_mask > (labels != IGNORE_INDEX) > ones + # For MTP training, loss_mask must match input_ids length + if "final_response_mask" in data.batch: + loss_mask = data.batch["final_response_mask"].float() + elif "response_mask" in data.batch: + loss_mask = data.batch["response_mask"].float() + elif labels is not None: + loss_mask = (labels != IGNORE_INDEX).float() + else: + loss_mask = torch.ones_like(input_ids) + + # Ensure loss_mask length matches input_ids length + # This is important for MTP training where mtp_labels has the same length as input_ids + if loss_mask.shape[1] != input_ids.shape[1]: + if loss_mask.shape[1] < input_ids.shape[1]: + # loss_mask is shorter (e.g., sliced with [:, 1:]), pad to match + pad_length = input_ids.shape[1] - loss_mask.shape[1] + # Pad at the beginning (position 0) since slicing was [:, 1:] + loss_mask = torch.nn.functional.pad(loss_mask, (pad_length, 0), value=0.0) + else: + # loss_mask is longer, truncate to match + loss_mask = loss_mask[:, :input_ids.shape[1]] + if self.use_sequence_packing: - input_ids, packed_seq_params, cu_seqlens, cu_seqlens_padded = self._pack_sequences( + # Pack input_ids to get packed_seq_params (cu_seqlens, etc.) + packed_input_ids, packed_seq_params, cu_seqlens, cu_seqlens_padded = self._pack_sequences( input_ids, attention_mask, ) if labels is not None: labels, _, _, _ = self._pack_sequences(labels, attention_mask, pad_val=IGNORE_INDEX) - attention_mask = None + loss_mask, _, _, _ = self._pack_sequences(loss_mask, attention_mask, pad_val=0) + # Multimodal models handle CP+packing internally (similar to mbridge) via + # MultimodalEmbeddingMixin, so pass the original un-packed input_ids andattention_mask. + # Non-multimodal models use the pre-packed input_ids. + if self.is_multimodal: + # Keep original input_ids and attention_mask for model-internal packing + pass + else: + input_ids = packed_input_ids + attention_mask = None else: input_ids = self._get_feature_on_this_cp_rank(input_ids, "input_ids") attention_mask = self._get_feature_on_this_cp_rank(attention_mask, "attention_mask") if labels is not None: labels = self._get_feature_on_this_cp_rank(labels, "labels") + loss_mask = self._get_feature_on_this_cp_rank(loss_mask, "loss_mask") position_ids = None - # attention_mask: SelfAttention defalt to te DotProductAttention with - # AttnMaskType.causal in which attention_mask would not be used, pass - # it mainly for moe aux loss without pad token and it is 2D - # position_ids: not used in LLM - # While MCA Qwen2VlModel requires 4D attention_mask, and - # attention_mask and position_ids would be chunked for cp with dim 2 as - # seq dim in it if they are provided forward_args = data.meta_info.get("forward_args", {}) - if "position_ids" in data.batch.keys() and data.batch["position_ids"].dim() == 3: # qwen2vl mrope - # not support MoE VLM, not used temperarily - attention_mask = None + if "position_ids" in data.batch and data.batch["position_ids"].dim() == 3: # qwen-vl/omni mrope position_ids = data.batch["position_ids"] if position_ids.size(1) == 4: position_ids = position_ids[:, 1:, :].contiguous() # (bsz, 4, seqlen) -> (bsz, 3, seqlen) @@ -450,22 +549,69 @@ def inner_forward_step(self, loss_func, data_iterator: Iterator[DataProto], mode multi_modal_data[key].append(sample_mm_inputs[key]) for key in multi_modal_data.keys(): assert key not in forward_args + mm_data = multi_modal_data[key] + # All mm fields are not padded in collator currently and some should be padded first + # pixel_values/pixel_values_videos for images/videos are concated for packing + # input_features for audios with shape `(bs, freqs, frames)` should be padded before concat + need_padding = any(t.shape[-1] != mm_data[0].shape[-1] for t in mm_data[1:]) + if need_padding: # input_features/feature_attention_mask + max_mm_len = max(t.shape[-1] for t in mm_data) + for i, t in enumerate(mm_data): + mm_data[i] = torch.nn.functional.pad(t, (0, max_mm_len - t.shape[-1]), "constant", 0) # DataProto.to('cuda') in upper frame not work for non_tensor_batch forward_args[key] = torch.concat(multi_modal_data[key], dim=0).to(input_ids.device) forward_args.update({"force_vit_image": True}) - # megatron_llama_core need loss_mask to compute aux loss - if "loss_mask" not in forward_args: - if labels is not None: - forward_args["loss_mask"] = (labels != IGNORE_INDEX).float() + if self.enable_router_replay: + unwrapped_model = unwrap_model(model) + if hasattr(unwrapped_model, "vp_stage"): + vp_rank = unwrapped_model.vp_stage else: - forward_args["loss_mask"] = torch.ones_like(input_ids) + vp_rank = 0 + + # If the current router_replay_action is REPLAY_BACKWARD for the local router instances, + # set the action to REPLAY_FORWARD for next micro batch. + if RouterReplayHelper.is_replay_backward_action(self.model.config, vp_rank): + router_instance_list = RouterReplayHelper.get_micro_batch_router_list(self.model.config, vp_rank) + for router in router_instance_list: + router.set_router_replay_action(RouterReplayAction.REPLAY_FORWARD) + + # If the current router_replay_action is REPLAY_FORWARD, set_router_replay_data. + if RouterReplayHelper.is_replay_forward_action(self.model.config, vp_rank): + layers_topk_idx = data.batch["routed_experts"] + set_router_replay_data(layers_topk_idx, attention_mask, self.model.config, vp_rank) + + # megatron_llama_core need loss_mask to compute aux loss + forward_args["loss_mask"] = loss_mask output_tensor = model( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, labels=labels, packed_seq_params=packed_seq_params, **forward_args ) + # # 新增：R2 模式下在推理后收集 router replay 数据 + # # TODO + # if self.enable_router_replay and self.router_replay_mode == "R2": + # # 合并并收集 router topk indices + # merge_router_topk_indices( + # attention_mask=attention_mask, + # input_ids=input_ids, + # mini_layer_topk_idx_list=self.router_topk_indices_list, + # tf_config=self.megatron_train_args, + # vp_rank=None + # ) + + # if RouterReplayHelper.is_r2_record_action(self.tf_config, vp_rank): + # merge_router_topk_indices( + # attention_mask, input_ids, self.mini_layer_topk_idx_list, self.tf_config, vp_rank + # ) + + if self.enable_router_replay and RouterReplayHelper.is_replay_forward_action(self.model.config, vp_rank): + router_instance_list = RouterReplayHelper.get_micro_batch_router_list(self.model.config, vp_rank) + # Ensured the correctness of router replay during the recomputation in backward. + for router in router_instance_list: + router.set_router_replay_action(RouterReplayAction.REPLAY_BACKWARD) + if self.use_sequence_packing: cp_size = mpu.get_context_parallel_world_size() def loss_wrapper(output_tensor): @@ -539,16 +685,21 @@ def op_compute_log_probs(self, logits: torch.Tensor, input_ids: torch.Tensor, at labels = torch.cat([labels, torch.zeros_like(labels[:, :1])], dim=1) labels = self._get_feature_on_this_cp_rank(labels, "labels") # compute logprobs in remove padding token - log_probs = vocab_parallel_logprobs(logits, labels) + log_probs = vocab_parallel_logprobs( + logits, labels, + use_fused_kernel=self.megatron_train_args.cross_entropy_loss_fusion + ) if mpu.get_context_parallel_world_size() > 1: log_probs = context_parallel_gather(log_probs, parallel_dim=1) log_probs = log_probs[:, :-1] * attention_mask[:, 1:] return log_probs def op_compute_entropy(self, logits: torch.Tensor, attention_mask: torch.Tensor): - if self.worker_config.logits_in_fp32: - logits = logits.float() - entropy = vocab_parallel_entropy(logits) + entropy = vocab_parallel_entropy( + logits, + used_fp32=self.worker_config.logits_in_fp32, + use_fused_kernel=self.megatron_train_args.cross_entropy_loss_fusion + ) if mpu.get_context_parallel_world_size() > 1: entropy = context_parallel_gather(entropy, parallel_dim=1) entropy = entropy[:, :-1] * attention_mask[:, 1:] @@ -971,6 +1122,15 @@ def __init__(self, worker: Worker): self.processor = None self._validate_access_integrity = True + # 新增：Router Replay 配置（用于 R2 和 R3 的 REPLAY） + # 注意：这里会覆盖父类的配置，因为训练阶段的行为不同 + self.router_replay_config = self.worker_config.router_replay + self.enable_router_replay = (self.router_replay_config.mode in ["R2", "R3"]) + self.router_replay_mode = self.router_replay_config.mode + + if self.enable_router_replay: + logger.info(f"Router Replay {self.router_replay_mode} mode: REPLAY enabled in MegatronTrainStrategy") + def initialize(self, model_provider): self.seq_length = self.worker.pipeline_config.sequence_length self.weight_updaters: dict[str, MegatronWeightUpdater] = {} @@ -986,6 +1146,11 @@ def initialize(self, model_provider): ) self.forward_backward_func = get_forward_backward_func() self.model.config.finalize_model_grads_func = finalize_model_grads + + # Inject mtp_training_mode from WorkerConfig to model config + if hasattr(self.worker_config, "mtp_training_mode"): + self.model.config.mtp_training_mode = self.worker_config.mtp_training_mode + ddp_config = DistributedDataParallelConfig( grad_reduce_in_fp32=self.megatron_train_args.accumulate_allreduce_grads_in_fp32, overlap_grad_reduce=self.megatron_train_args.overlap_grad_reduce, @@ -1006,6 +1171,8 @@ def initialize(self, model_provider): ] self.models_unwrapped = self.model.get_models() self.model.models = self.models_wrapped + self.is_multimodal = self.processor is not None + self._validate_vlm_packing_support() params_dtype = ( torch.float16 @@ -1079,12 +1246,23 @@ def initialize(self, model_provider): self.model.config.variable_seq_lengths = True logger.info("Set variable_seq_lengths to True when use dynamic batching and pipeline parallel.") + # # In train mode, init router replay action to REPLAY_FORWARD + # if self.enable_router_replay and self.router_replay_mode == "R3": + # RouterReplay.set_global_router_replay_action(RouterReplayAction.REPLAY_FORWARD) + logger.info(f"{self.model.get_models()}") + if self.megatron_train_args.compile_warmup and self.worker.rank_info.pp_size > 1: + compile_warmup_pipeline_stages(self) + dist.barrier() def train_step(self, batch: DataProto, loss_func: Callable): self.model.train() + if self.enable_router_replay: + assert "routed_experts" in batch.batch + RouterReplay.set_global_router_replay_action(RouterReplayAction.REPLAY_FORWARD) + global_step = batch.meta_info.get("global_step", 0) is_offload_optimizer_states_in_train_step = batch.meta_info.get("is_offload_optimizer_states_in_train_step", True) batch.meta_info['batch_num_tokens'] = self._get_batch_num_tokens(batch, dp_group=mpu.get_data_parallel_group()) @@ -1129,9 +1307,13 @@ def train_step(self, batch: DataProto, loss_func: Callable): forward_only=False, ) + # clear global router replay action + if self.enable_router_replay: + RouterReplay.clear_global_router_replay_action() + RouterReplay.clear_global_indices() + # 只有step的时候需要load optimizer states self.load_states(include=[OffloadStateType.optimizer_states]) - update_successful, grad_norm, num_zeros_in_grad = self.optimizer.step() if is_offload_optimizer_states_in_train_step: self.offload_states(include=[OffloadStateType.optimizer_states], non_blocking=True) @@ -1142,6 +1324,13 @@ def train_step(self, batch: DataProto, loss_func: Callable): raise NotImplementedError("megatron optimizer step failed!") for model in self.model: + for bucket_group in model.bucket_groups + model.expert_parallel_bucket_groups: + if hasattr(bucket_group, "per_param_grad_ready_counts") and hasattr(bucket_group, "is_first_batch"): + if bucket_group.is_first_batch and len(bucket_group.per_param_grad_ready_counts) > 0: + # Fill in any missing params with count=1 so the assertion passes. + for param in bucket_group.params: + if param not in bucket_group.per_param_grad_ready_counts: + bucket_group.per_param_grad_ready_counts[param] = 1 model.zero_grad_buffer() # Offload/reload does not update cached_param_buffer_shard_list/cached_grad_buffer_shard_list, # resulting using old params in `start_param_sync`, which leads to wrong results. So we clear the cache. @@ -1232,11 +1421,17 @@ def save_checkpoint(self, save_dir, global_step, ckpt_id, tag="checkpoint", loca with Timer("load") as load_timer: self.load_states() + # Only during checkpoint: release pinned memory cache to free CPU memory + # for the upcoming serialization. Normal training keeps the cache for reuse. + clear_memory(clear_host_memory=True) + is_last_step = kwargs.get("is_last_step", False) if self.megatron_train_args.save_hf_model: self.model.save_pretrained_as_hf(save_dir) + clear_memory(clear_host_memory=True) + ckpt_format = self.megatron_train_args.ckpt_format # save model and tokenizer if len(self.models_unwrapped) == 1: if is_peft_available() and isinstance(self.models_unwrapped[0], PeftModel): @@ -1251,58 +1446,61 @@ def save_checkpoint(self, save_dir, global_step, ckpt_id, tag="checkpoint", loca ) self.models_unwrapped[0].config.save_pretrained(save_dir) else: - self.models_unwrapped[0].save_pretrained(save_dir) + self.models_unwrapped[0].save_pretrained(save_dir, ckpt_format=ckpt_format) else: - state_dict = {f"model{i}": model.state_dict_for_save_checkpoint() for i, model in - enumerate(self.models_unwrapped)} - self.models_unwrapped[0].save_pretrained(save_dir, state_dict=state_dict) + state_dict = {f"model{i}": generate_model_state_dict(model, ckpt_format=ckpt_format) for i, model in enumerate(self.models_unwrapped)} + self.models_unwrapped[0].save_pretrained(save_dir, state_dict=state_dict, ckpt_format=ckpt_format) + del state_dict if dist.get_rank() == 0: if self.tokenizer is not None: self.tokenizer.save_pretrained(save_dir) if self.processor is not None: self.processor.save_pretrained(save_dir) + clear_memory(clear_host_memory=True) # save optimizer - checkpoint_dir = get_checkpoint_dir(save_dir, - return_base_dir=self.megatron_train_args.use_distributed_optimizer) - if self.megatron_train_args.use_distributed_optimizer: - checkpoint_dir = os.path.join(checkpoint_dir, DIST_OPTIMIZER_DIR) - os.makedirs(checkpoint_dir, exist_ok=True) - if self.megatron_train_args.use_distributed_optimizer: - model_shared_state_dict = self.model.sharded_state_dict() - optimizer_state_dict = self.optimizer.sharded_state_dict( - model_shared_state_dict, metadata=self.ckpt_sharding_metadata - ) - dist_checkpointing.save( - optimizer_state_dict, - checkpoint_dir=checkpoint_dir, - sharded_strategy=self.save_strategy, - async_sharded_save=False, - validate_access_integrity=self._validate_access_integrity, - ) - self._validate_access_integrity = False - elif not dist.is_initialized() or mpu.get_expert_data_parallel_rank() == 0: - torch.save(self.optimizer.state_dict(), os.path.join(checkpoint_dir, OPTIMIZER_NAME)) - logger.info(f"Saving optimizer state to {os.path.join(checkpoint_dir, OPTIMIZER_NAME)}") - - if dist.is_initialized(): - dist.barrier() - - # save lr_scheduler - if dist.get_rank() == 0: - torch.save(self.scheduler.state_dict(), os.path.join(save_dir, SCHEDULER_NAME)) - - # save rng state - rng_states = { - "random_rng_state": random.getstate(), - "np_rng_state": np.random.get_state(), - "torch_rng_state": torch.get_rng_state(), - "cuda_rng_state": current_platform.get_rng_state(), - "rng_tracker_states": tensor_parallel.get_cuda_rng_tracker().get_states(), - } - rgn_path = os.path.join(save_dir, RNG_STATE_DIR, f"rng_state_{dist.get_rank()}.pth") - os.makedirs(os.path.dirname(rgn_path), exist_ok=True) - torch.save(rng_states, rgn_path) + if not self.megatron_train_args.save_only_model: + checkpoint_dir = get_checkpoint_dir(save_dir, + return_base_dir=self.megatron_train_args.use_distributed_optimizer) + if self.megatron_train_args.use_distributed_optimizer: + checkpoint_dir = os.path.join(checkpoint_dir, DIST_OPTIMIZER_DIR) + os.makedirs(checkpoint_dir, exist_ok=True) + if self.megatron_train_args.use_distributed_optimizer: + model_shared_state_dict = self.model.sharded_state_dict() + optimizer_state_dict = self.optimizer.sharded_state_dict( + model_shared_state_dict, metadata=self.ckpt_sharding_metadata + ) + dist_checkpointing.save( + optimizer_state_dict, + checkpoint_dir=checkpoint_dir, + sharded_strategy=self.save_strategy, + async_sharded_save=False, + validate_access_integrity=self._validate_access_integrity, + ) + del model_shared_state_dict, optimizer_state_dict + self._validate_access_integrity = False + elif not dist.is_initialized() or mpu.get_data_modulo_expert_parallel_rank() == 0: + torch.save(self.optimizer.state_dict(), os.path.join(checkpoint_dir, OPTIMIZER_NAME)) + logger.info(f"Saving optimizer state to {os.path.join(checkpoint_dir, OPTIMIZER_NAME)}") + + if dist.is_initialized(): + dist.barrier() + + # save lr_scheduler + if dist.get_rank() == 0: + torch.save(self.scheduler.state_dict(), os.path.join(save_dir, SCHEDULER_NAME)) + + # save rng state + rng_states = { + "random_rng_state": random.getstate(), + "np_rng_state": np.random.get_state(), + "torch_rng_state": torch.get_rng_state(), + "cuda_rng_state": current_platform.get_rng_state(), + "rng_tracker_states": tensor_parallel.get_cuda_rng_tracker().get_states(), + } + rng_path = os.path.join(save_dir, RNG_STATE_DIR, f"rng_state_{dist.get_rank()}.pth") + os.makedirs(os.path.dirname(rng_path), exist_ok=True) + torch.save(rng_states, rng_path) if self.worker_config.checkpoint_config.get("async_upload", True) and not is_last_step: self.thread_executor.submit(self.checkpoint_manager.upload, ckpt_id=ckpt_id, local_state_path=local_state_path) @@ -1312,6 +1510,7 @@ def save_checkpoint(self, save_dir, global_step, ckpt_id, tag="checkpoint", loca metrics = { "load": load_timer.last, } + clear_memory(clear_host_memory=True) return metrics def load_checkpoint(self, load_dir, tag="checkpoint", **kwargs): @@ -1327,8 +1526,6 @@ def load_checkpoint(self, load_dir, tag="checkpoint", **kwargs): f"Loading optimizer from {optimizer_checkpoint}, process_index: {self.megatron_train_args.process_index}" ) - self.offload_states() - if self.megatron_train_args.use_distributed_optimizer: model_shared_state_dict = self.model.sharded_state_dict() sharded_state_dict = self.optimizer.sharded_state_dict( @@ -1371,5 +1568,3 @@ def load_checkpoint(self, load_dir, tag="checkpoint", **kwargs): tensor_parallel.get_cuda_rng_tracker().set_states(checkpoint_rng_state["rng_tracker_states"]) else: logger.info(f"not load rng state, not found file: {rng_file}") - - self.load_states() diff --git a/roll/distributed/strategy/mock_strategy.py b/roll/distributed/strategy/mock_strategy.py index 81da4179b..e9a78b60c 100644 --- a/roll/distributed/strategy/mock_strategy.py +++ b/roll/distributed/strategy/mock_strategy.py @@ -3,7 +3,6 @@ from datetime import timedelta from typing import List, Optional, Callable, Dict, Tuple -import deepspeed import torch import torch.distributed as dist from accelerate import cpu_offload_with_hook diff --git a/roll/distributed/strategy/sglang_strategy.py b/roll/distributed/strategy/sglang_strategy.py index c09f65641..aab0b917b 100644 --- a/roll/distributed/strategy/sglang_strategy.py +++ b/roll/distributed/strategy/sglang_strategy.py @@ -34,7 +34,7 @@ from roll.utils.functionals import concatenate_input_and_output, gather_unpadded_input_ids from roll.utils.logging import get_logger from roll.utils.network_utils import collect_free_port -from roll.utils.offload_states import OffloadStateType +from roll.utils.offload_states import OffloadStateType, clear_memory from roll.platforms import current_platform try: @@ -94,6 +94,13 @@ async def initialize(self, model_provider): else: dtype = "auto" + self.enable_rollout_routing_replay = self.worker_config.router_replay.mode != "disable" + if self.enable_rollout_routing_replay: + from sglang import __version__ as sglang_version + from packaging import version + assert version.parse(sglang_version) >= version.parse("0.5.6.post3"), "Enable router replay requires sglang >= 0.5.6.post3" + logger.info(f"{self.enable_rollout_routing_replay=} and {self.worker_config.router_replay.mode=}") + sglang_config.setdefault("enable_memory_saver", True) sglang_config.update( { @@ -112,6 +119,8 @@ async def initialize(self, model_provider): "disable_custom_all_reduce": sglang_config.get("disable_custom_all_reduce", True), 'nnodes': nnodes, 'node_rank': 0, + # new：router replay + "enable_return_routed_experts": self.enable_rollout_routing_replay, } ) @@ -139,18 +148,18 @@ async def initialize(self, model_provider): from roll.utils.constants import RAY_NAMESPACE for i in range(1, nnodes): sglang_ray_option = { - 'scheduling_strategy': PlacementGroupSchedulingStrategy(sglang_pg_list[i]), + 'scheduling_strategy': PlacementGroupSchedulingStrategy(sglang_pg_list[i]), 'name': f'sglang-slave-{node_index_list[i]}', 'namespace': RAY_NAMESPACE, - 'runtime_env': - {'env_vars': - {'WORLD_SIZE': str(nnodes), - 'RANK': str(i), + 'runtime_env': + {'env_vars': + {'WORLD_SIZE': str(nnodes), + 'RANK': str(i), 'WORKER_NAME': f'sglang-slave-{node_index_list[i]}', - 'CUDA_VISIBLE_DEVICES': ','.join(map(str, list(range(gpu_per_worker)))), 'RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES': '1', + 'CUDA_VISIBLE_DEVICES': ','.join(map(str, list(range(gpu_per_worker)))), 'RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES': '1', 'ROLL_LOG_DIR': os.getenv("ROLL_LOG_DIR", "./output/logs/") } - }, + }, 'num_cpus': 0.01, 'num_gpus': 0.01 } @@ -188,11 +197,49 @@ async def generate_request(self, payload: dict): from sglang import __version__ as version if version < '0.5' and payload["sampling_params"]["n"] > 1: # fixed in https://github.com/sgl-project/sglang/pull/7508 payload["rid"] = None + + # new：If enable router replay，set return_routed_experts config + if self.enable_rollout_routing_replay: + payload["return_routed_experts"] = True + else: + payload["return_routed_experts"] = False + + if "multi_modal_data" in payload: + multi_modal_data = payload["multi_modal_data"] + input_ids = multi_modal_data["prompt_token_ids"] + if ( + "multi_modal_data" not in multi_modal_data + or "image" not in multi_modal_data["multi_modal_data"] + or not multi_modal_data["multi_modal_data"]["image"] + ): + image_data = None + else: + image_data = multi_modal_data["multi_modal_data"]["image"] + + if ( + "multi_modal_data" not in multi_modal_data + or "video" not in multi_modal_data["multi_modal_data"] + or not multi_modal_data["multi_modal_data"]["video"] + ): + video_data = None + else: + video_data = [dict( + multi_modal_data["multi_modal_data"]["video"], + format="processor_output" + )] + else: + input_ids = payload["input_ids"] + image_data = None + video_data = None + obj = GenerateReqInput( - input_ids=payload["input_ids"], + input_ids=input_ids, sampling_params=payload["sampling_params"], rid=payload["rid"], return_logprob=payload["return_logprob"], + return_routed_experts=payload["return_routed_experts"], + image_data=image_data, + video_data=video_data ) generator = self.model.engine.tokenizer_manager.generate_request(obj, None) chunks = None @@ -212,38 +259,47 @@ async def generate(self, batch: DataProto, generation_config): input_ids = batch.batch["input_ids"] # (bs, prompt_length) attention_mask = batch.batch["attention_mask"] # left-padded attention_mask - image_data = None + image_data = [] + video_data = [] if "multi_modal_data" in batch.non_tensor_batch: prompt_token_ids = [] - image_data = [] - # sglang enforce str(path or url)/bytes image data currently - # TODO: path image_processor.load_image with hash according to: - # https://github.com/sgl-project/sglang/pull/4915 - for data in batch.non_tensor_batch["multi_modal_data"]: - # bug exists in sglang, it only puts image str (standing for path - # or url) into list and leaves out image bytes. Thus when using - # image bytes, put it into list mannully + for i, data in enumerate(batch.non_tensor_batch["multi_modal_data"]): prompt_token_ids.append(data["prompt_token_ids"]) - # for text and multi-modal mixed data if ( "multi_modal_data" not in data or "image" not in data["multi_modal_data"] or not data["multi_modal_data"]["image"] ): image_data.append(None) - continue - image_per_sample = [] - for image in data["multi_modal_data"]["image"]: - byte_stream = io.BytesIO() - image.save(byte_stream, "png") - image_per_sample.append(byte_stream.getvalue()) - byte_stream.close() - image_data.append(image_per_sample) + else: + image_data.append(data["multi_modal_data"]["image"]) + + if ( + "multi_modal_data" not in data + or "video" not in data["video"] + or not data["multi_modal_data"]["video"] + ): + video_data.append(None) + else: + video_data.append([dict( + batch.non_tensor_batch["multi_modal_data"][i], + format="processor_output" + )]) + cur_input_ids = cur_input_ids[i].unsqueeze(0) + cur_attention_mask = attention_mask[i].unsqueeze(0) + input_ids = gather_unpadded_input_ids(input_ids=cur_input_ids, attention_mask=cur_attention_mask) + prompt_token_ids[-1] = input_ids[0] else: prompt_token_ids = gather_unpadded_input_ids(input_ids=input_ids, attention_mask=attention_mask) return_logprob = sampling_params.pop("return_logprob", False) + + if all(not x for x in image_data): + image_data = None + if all(not x for x in video_data): + video_data = None + sglang_outputs = await self.model.engine.async_generate( - input_ids=prompt_token_ids, image_data=image_data, sampling_params=sampling_params, return_logprob=return_logprob + input_ids=prompt_token_ids, image_data=image_data, video_data=video_data, sampling_params=sampling_params, return_logprob=return_logprob ) # (bs * num_return_sequences, max_response_len) @@ -321,8 +377,7 @@ async def offload_states(self, include=None, non_blocking=False): # always release all self.is_model_in_gpu, self.is_kv_cache_in_gpu = False, False - gc.collect() - current_platform.empty_cache() + clear_memory() class SglangEngine: def __init__(self): @@ -554,12 +609,35 @@ def postprocess_generate(chunks): ) output_data["output_token_ids"] = output_token_ids output_data["finish_reasons"] = [] + + output_routed_experts = [chunk["meta_info"].get("routed_experts", None) for chunk in chunks] + has_routed_experts = any(routed_experts is not None for routed_experts in output_routed_experts) + if has_routed_experts: + output_data["routed_experts"] = output_routed_experts + + # Speculative decoding metrics keys from sglang + spec_metric_keys = [ + "spec_accept_rate", + "spec_accept_length", + "spec_accept_token_num", + "spec_draft_token_num", + ] + for chunk in chunks: - finish_reason = chunk["meta_info"]["finish_reason"] + meta_info = chunk["meta_info"] + + finish_reason = meta_info["finish_reason"] if isinstance(finish_reason, dict): finish_reason = finish_reason["type"] output_data["finish_reasons"].append(finish_reason) else: output_data["finish_reasons"].append(finish_reason) + + # Extract speculative decoding metrics + for key in spec_metric_keys: + if key in meta_info: + output_data.setdefault("metrics", {}).setdefault(key, []).append(meta_info[key]) + assert len(output_data["finish_reasons"]) == len(output_data["output_token_ids"]) return output_data + diff --git a/roll/distributed/strategy/strategy.py b/roll/distributed/strategy/strategy.py index 6171b5faf..d54521b96 100644 --- a/roll/distributed/strategy/strategy.py +++ b/roll/distributed/strategy/strategy.py @@ -302,7 +302,7 @@ def op_compute_various_divergence(self, loss_callable, logits, teacher_topk_prob else: raise ValueError(f"Unsupported reduction: {reduction}. Use 'mean', 'sum', or 'none'.") - # Both megatron and deepspeed can output language loss directly. + # Both megatron and fsdp2 can output language loss directly. # This op is mainly for computing context-parallel loss. def op_compute_language_loss(self, losses: torch.Tensor, labels: torch.Tensor, batch_num_tokens: int): loss_mask = (labels != IGNORE_INDEX).float() @@ -442,71 +442,6 @@ def __init__(self, worker: "Worker"): def setup_collective_group(self, model_update_name, comm_plan, backend=None, mode="sender"): self._setup_collective_group_impl(model_update_name, comm_plan, backend, mode=mode) - - def setup_p2p_collective_group(self, model_update_name, comm_plan, backend="nccl"): - (intra_rank, info), = comm_plan.items() - collective.init_collective_group( - info["world_size"], - intra_rank, - backend=backend, - group_name=info["group_name"], - master_addr=info["master_addr"], - master_port=info["master_port"], - global_ranks=info["global_ranks"] - ) - # 可选：warm-up - collective.allreduce(torch.zeros(1).cuda(), group_name=info["group_name"]) - # 保存元数据 - if model_update_name not in self.model_update_comm_plan: - self.model_update_comm_plan[model_update_name] = {} - self.model_update_comm_plan[model_update_name][info["group_name"]] = { - "rank": intra_rank, - "world_size": info["world_size"], - "group_name": info["group_name"], - "comm_plan": comm_plan, - } - - def model_update_set_write_done_handle(self,): - """ - Set the write synchronization event required for reading and writing shared memory - """ - if not hasattr(self, "_events_inited"): - # Sender -> Receiver：Write complete - self._write_done_event = torch.cuda.Event(interprocess=True) - self._write_done_handle = self._write_done_event.ipc_handle() - # Sender <- Receiver：Read complete - self._read_done_event_remote = None - self._events_inited = True - - def model_update_set_read_done_handle(self, read_done_handles): - """ - Set the read synchronization event required for reading and writing shared memory - """ - logger.warning(f"[Rank {dist.get_rank()}] model_update_set_read_done_handle called") - read_done_handle = None - - for p2p_tgt_device in self.p2p_tgt_devices: - worker_rank = p2p_tgt_device['rank'] - local_rank = p2p_tgt_device['device']['rank'] - for read_done_handle_full_dict in read_done_handles: - if worker_rank in read_done_handle_full_dict: - read_done_handle_list = read_done_handle_full_dict[worker_rank] - for read_done_handle_dict in read_done_handle_list: - if local_rank in read_done_handle_dict: - read_done_handle = read_done_handle_dict[local_rank] - - if not hasattr(self, "_read_done_event_remote"): - if read_done_handle is not None: - logger.warning(f"[Rank {dist.get_rank()}] Creating _read_done_event_remote from handle") - self._read_done_event_remote = torch.cuda.Event.from_ipc_handle( - device=torch.cuda.current_device(), - handle=read_done_handle - ) - else: - logger.warning( - f"[Rank {dist.get_rank()}] No read_done_handle found, setting _read_done_event_remote=None") - self._read_done_event_remote = None - def train_step( self, batch: DataProto, diff --git a/roll/distributed/strategy/vllm_strategy.py b/roll/distributed/strategy/vllm_strategy.py index 87abeb945..c013d3e98 100644 --- a/roll/distributed/strategy/vllm_strategy.py +++ b/roll/distributed/strategy/vllm_strategy.py @@ -2,6 +2,7 @@ import copy import gc import os +import random from collections import deque from typing import Dict, List, Optional from packaging.version import Version @@ -27,7 +28,7 @@ gather_unpadded_input_ids, ) from roll.utils.logging import get_logger -from roll.utils.offload_states import OffloadStateType +from roll.utils.offload_states import OffloadStateType, clear_memory from roll.platforms import current_platform @@ -45,9 +46,36 @@ def __init__(self, worker: Worker): self._metrics_snapshot_interval = 1.0 # Snapshot every 1 second self._metrics_task = None + + def get_free_port_for_rank(self) -> int: + VLLM_PORT_START = 20000 + PORT_RANGE = 500 + MAX_LOCAL_WORKER_COUNT = 16 + + rank = self.worker.rank + effective_rank = rank % MAX_LOCAL_WORKER_COUNT + + range_start = VLLM_PORT_START + effective_rank * PORT_RANGE + range_end = range_start + PORT_RANGE + + for _ in range(PORT_RANGE): + port = random.randint(range_start, range_end - 1) + if self.worker.is_port_available(port): + return port + + raise RuntimeError( + f"Cannot allocate free port for rank {rank} (effective_rank={effective_rank}) in range [{range_start}, {range_end}]" + ) + async def initialize(self, model_provider): set_seed(seed=self.worker.pipeline_config.seed) vllm_config = copy.deepcopy(self.worker_config.strategy_args.strategy_config) + + # Apply GDN attention patch for mixed decode/spec-decode bug fix + # This patches vLLM versions < v0.17.2 that lack the fix + from roll.third_party.vllm.gdn_patcher import patch_gdn_attention + patch_gdn_attention() + # Must explicitly set VLLM_USE_V1 to pass this check: https://github.com/vllm-project/vllm/pull/14972 os.environ["VLLM_USE_V1"] = str(vllm_config.pop("VLLM_USE_V1", 1)) self.sleep_level = vllm_config.pop("sleep_level", 1) @@ -109,7 +137,8 @@ async def initialize(self, model_provider): # https://github.com/vllm-project/vllm/blob/releases/v0.10.0/vllm/v1/engine/coordinator.py#L72 if not data_parallel_size > 1: # set VLLM_PORT to avoid port conflict applied by vllm - vllm_port = self.worker.get_free_port() + vllm_port = self.get_free_port_for_rank() + logger.info(f"Allocated vllm_port {vllm_port} for rank {self.worker.rank}") os.environ["VLLM_PORT"] = str(vllm_port) self.model = await create_async_llm(resource_placement_groups=self.worker_config.resource_placement_groups, **vllm_config) @@ -263,9 +292,15 @@ async def generate_request(self, payload: Dict): if "multi_modal_data" in payload: multi_modal_data = payload["multi_modal_data"] prompt_token_ids = multi_modal_data["prompt_token_ids"] - multi_modal_data = (multi_modal_data["multi_modal_data"] - if "multi_modal_data" in multi_modal_data else None) - prompt = TokensPrompt(prompt_token_ids=prompt_token_ids, multi_modal_data=multi_modal_data) + prompt = TokensPrompt( + prompt_token_ids=prompt_token_ids, + multi_modal_data=multi_modal_data["multi_modal_data"] + if "multi_modal_data" in multi_modal_data + else None, + mm_processor_kwargs=multi_modal_data["mm_processor_kwargs"] + if "mm_processor_kwargs" in multi_modal_data + else None, + ) else: prompt = TokensPrompt(prompt_token_ids=payload["input_ids"]) @@ -306,12 +341,20 @@ async def generate_request(self, payload: Dict): for token_id, lps in zip(completion_output.token_ids, completion_output.logprobs) ] ) - return { + + result = { "output_token_ids": output_token_ids, "finish_reasons": finish_reasons, "output_logprobs": logprobs, } + # Add speculative metrics if available + spec_metrics = self.get_speculative_metrics() + if spec_metrics: + result["metrics"] = spec_metrics + + return result + async def abort_requests(self, request_ids): for id in request_ids: await self.model.abort(request_id=id) @@ -329,8 +372,7 @@ async def offload_states(self, include=None, non_blocking=False): if self.is_model_in_gpu and self.worker.pipeline_config.is_actor_infer_colocated: await self.model.offload_states(self.sleep_level) self.is_model_in_gpu = False - gc.collect() - current_platform.empty_cache() + clear_memory() async def process_weights_after_loading(self,*args, **kwargs): await self.model.process_weights_after_loading() @@ -351,25 +393,27 @@ async def add_lora(self, peft_config): peft_config["target_modules"] = set(self.worker_config.model_args.lora_target) await self.model.add_lora(peft_config) + # Mapping from raw vLLM metric names to internal keys + _VLLM_METRIC_MAP = { + "vllm:kv_cache_usage_perc": "vllm/kv_cache_usage_perc_max", + "vllm:num_requests_waiting": "vllm/num_requests_waiting_max", + "vllm:num_preemptions": "vllm/num_preemptions_max", + "vllm:spec_decode_num_drafts": "vllm/spec_decode_num_drafts", + "vllm:spec_decode_num_draft_tokens": "vllm/spec_decode_num_draft_tokens", + "vllm:spec_decode_num_accepted_tokens": "vllm/spec_decode_num_accepted_tokens", + } + async def _collect_metrics_snapshot(self): - """Collect metrics snapshots periodically in a background thread.""" + """Collect metrics snapshots periodically in a background task.""" from vllm.v1.metrics.reader import get_metrics_snapshot while True: raw_metrics = get_metrics_snapshot() - snapshot = { - 'vllm/kv_cache_usage_perc_max': [], - 'vllm/num_requests_waiting_max': [], - 'vllm/num_preemptions_max': [] - } + snapshot = {key: [] for key in self._VLLM_METRIC_MAP.values()} for metric in raw_metrics: - if metric.name == "vllm:kv_cache_usage_perc": - snapshot['vllm/kv_cache_usage_perc_max'].append(metric.value) - elif metric.name == "vllm:num_requests_waiting": - snapshot['vllm/num_requests_waiting_max'].append(metric.value) - elif metric.name == "vllm:num_preemptions": - snapshot['vllm/num_preemptions_max'].append(metric.value) + mapped_key = self._VLLM_METRIC_MAP.get(metric.name) + if mapped_key: + snapshot[mapped_key].append(metric.value) self._metrics_snapshots.append(snapshot) - await asyncio.sleep(self._metrics_snapshot_interval) def get_metrics(self, metric_names: Optional[List[str]] = None) -> Dict[str, float]: @@ -388,6 +432,21 @@ def get_metrics(self, metric_names: Optional[List[str]] = None) -> Dict[str, flo self._metrics_snapshots.clear() return reduce_metrics(metrics_snapshots) + def get_speculative_metrics(self) -> Dict[str, float]: + """Get speculative decoding metrics in a format aligned with SGLang.""" + metrics = self.get_metrics() + draft_tokens = metrics.get('vllm/spec_decode_num_draft_tokens', 0) + accepted_tokens = metrics.get('vllm/spec_decode_num_accepted_tokens', 0) + draft_count = metrics.get('vllm/spec_decode_num_drafts', 0) + if draft_tokens == 0: + return {} + return { + 'spec_draft_token_num': draft_tokens, + 'spec_accept_token_num': accepted_tokens, + 'spec_accept_rate': accepted_tokens / draft_tokens if draft_tokens > 0 else 0.0, + 'spec_accept_length': 1 + (accepted_tokens / draft_count) if draft_count > 0 else 1.0, + } + def gather_outputs_to_pad_tensor(request_outputs: List["RequestOutput"], pad_token_id, device=None) -> torch.Tensor: if device is None: diff --git a/roll/models/model_providers.py b/roll/models/model_providers.py index b5c432026..8ddac7ff7 100644 --- a/roll/models/model_providers.py +++ b/roll/models/model_providers.py @@ -18,7 +18,6 @@ TrainingArguments, ) from transformers.dynamic_module_utils import get_cached_module_file -from transformers.integrations import is_deepspeed_zero3_enabled from transformers.modeling_utils import is_fsdp_enabled from roll.configs import ModelArguments @@ -81,8 +80,6 @@ def prepare_automap_files(model_path: str): def default_tokenizer_provider(model_args: "ModelArguments", model_name_or_path: str = None): - if model_args.model_type == "diffusion_module": - return None if model_name_or_path is None: model_name_or_path = model_args.model_name_or_path model_name_or_path = download_model(model_name_or_path) @@ -98,8 +95,6 @@ def default_tokenizer_provider(model_args: "ModelArguments", model_name_or_path: def default_processor_provider(model_args: "ModelArguments", model_name_or_path: str = None): - if model_args.model_type == "diffusion_module": - return None if model_name_or_path is None: model_name_or_path = model_args.model_name_or_path model_name_or_path = download_model(model_args.model_name_or_path) @@ -224,13 +219,11 @@ def load_model( if model_args.moe_aux_loss_coef is not None: setattr(config, "router_aux_loss_coef", model_args.moe_aux_loss_coef) setattr(config, "output_router_logits", is_trainable) - init_kwargs["low_cpu_mem_usage"] = not is_deepspeed_zero3_enabled() and not is_fsdp2_enabled() + init_kwargs["low_cpu_mem_usage"] = not is_fsdp2_enabled() - # TODO: Shall we need the compute_dtype? Check the necessity. - if not is_deepspeed_zero3_enabled(): - init_kwargs["torch_dtype"] = model_args.compute_dtype + init_kwargs["torch_dtype"] = model_args.compute_dtype - if not is_deepspeed_zero3_enabled() and not is_fsdp_or_fsdp2_enabled(): + if not is_fsdp_or_fsdp2_enabled(): if init_kwargs["low_cpu_mem_usage"]: # device map requires low_cpu_mem_usage=True if "device_map" not in init_kwargs and model_args.device_map: init_kwargs["device_map"] = model_args.device_map @@ -285,17 +278,13 @@ def load_model( vhead_params = load_valuehead_params(model_name_or_path) if vhead_params is not None: - if is_deepspeed_zero3_enabled(): - import deepspeed # type: ignore - - params = [param for _, param in model.v_head.named_parameters(recurse=False)] - with deepspeed.zero.GatheredParameters(params, modifier_rank=0): - if torch.distributed.get_rank() == 0: - model.load_state_dict(vhead_params, strict=False) - else: - model.load_state_dict(vhead_params, strict=False) + model.load_state_dict(vhead_params, strict=False) logger.info("Loaded valuehead from checkpoint: {}".format(model_name_or_path)) + # Ensure v_head dtype matches pretrained_model to avoid FSDP2 "uniform dtype" errors + if hasattr(model, "v_head") and model_args.compute_dtype is not None: + model.v_head.to(model_args.compute_dtype) + if not is_trainable: model.requires_grad_(False) for param in model.parameters(): @@ -448,23 +437,6 @@ def forward_patch( setattr(model, "_roll_forward_patched", True) -def default_diffusion_module_provider( - tokenizer: None, - model_args: ModelArguments, - training_args: TrainingArguments = None, - is_trainable: Optional[bool] = False, -): - if model_args.model_config_kwargs["model_name"] == "wan2_2": - from roll.pipeline.diffusion.modules.wan_module import WanTrainingModule - - print(f"{model_args.model_config_kwargs=}") - training_module = WanTrainingModule(**model_args.model_config_kwargs) - else: - raise NotImplementedError(f"model_type {model_args.model_type} not implemented yet") - - return training_module - - def default_actor_model_provider( tokenizer: "PreTrainedTokenizer", model_args: "ModelArguments", @@ -486,8 +458,6 @@ def default_actor_model_provider( model = AutoModel.from_pretrained(model_args.model_name_or_path, training_args) if is_trainable: model.train() - for param in model.parameters(): - param.requires_grad = True else: model.eval() for param in model.parameters(): @@ -505,7 +475,7 @@ def default_actor_model_provider( "torch_dtype": model_args.compute_dtype, "trust_remote_code": True, } - if not is_deepspeed_zero3_enabled() and not is_fsdp2_enabled(): + if not is_fsdp2_enabled(): init_kwargs["low_cpu_mem_usage"] = True if is_trainable: init_kwargs["device_map"] = {"": current_platform.current_device()} @@ -515,7 +485,8 @@ def default_actor_model_provider( init_kwargs["device_map"] = "balanced" logger.info(f"init_kwargs: {init_kwargs}") model = load_model(model_args, is_trainable, False) - if model.config.pad_token_id is None: + # model.config may not have pad_token_id + if getattr(model.config, "pad_token_id", None) is None: model.config.pad_token_id = tokenizer.pad_token_id patch_model(model, config, use_mcore=False) @@ -553,12 +524,11 @@ class TokenClassifierOutput(ModelOutput): "torch_dtype": model_args.compute_dtype, "trust_remote_code": True, } - if not is_deepspeed_zero3_enabled(): - init_kwargs["low_cpu_mem_usage"] = True - if is_trainable: - init_kwargs["device_map"] = {"": current_platform.current_device()} - elif model_args.device_map: - init_kwargs["device_map"] = model_args.device_map + init_kwargs["low_cpu_mem_usage"] = True + if is_trainable: + init_kwargs["device_map"] = {"": current_platform.current_device()} + elif model_args.device_map: + init_kwargs["device_map"] = model_args.device_map logger.info(f"init_kwargs: {init_kwargs}") if model_args.model_type in ["auto_sequence_classification"]: logger.info(f"use AutoModelForSequenceClassification model {model_args.model_type}") @@ -606,7 +576,8 @@ class TokenClassifierOutput(ModelOutput): logger.info("patch AutoModelForCausalLMWithValueHead load_state_dict and forward") else: raise NotImplementedError - if model.config.pad_token_id is None: + # model.config may not have pad_token_id + if getattr(model.config, "pad_token_id", None) is None: model.config.pad_token_id = tokenizer.pad_token_id model_args.model_name_or_path = old_model_name_or_path @@ -638,12 +609,11 @@ def default_value_model_provider( "torch_dtype": model_args.compute_dtype, "trust_remote_code": True, } - if not is_deepspeed_zero3_enabled(): - init_kwargs["low_cpu_mem_usage"] = True - if is_trainable: - init_kwargs["device_map"] = {"": current_platform.current_device()} - elif model_args.device_map: - init_kwargs["device_map"] = model_args.device_map + init_kwargs["low_cpu_mem_usage"] = True + if is_trainable: + init_kwargs["device_map"] = {"": current_platform.current_device()} + elif model_args.device_map: + init_kwargs["device_map"] = model_args.device_map logger.info(f"init_kwargs: {init_kwargs}") if model_args.model_type in ["auto_token_classification"]: logger.info(f"use AutoModelForTokenClassification model {model_args.model_type}") @@ -677,7 +647,8 @@ def default_value_model_provider( ) else: raise NotImplementedError - if model.config.pad_token_id is None: + # model.config may not have pad_token_id + if getattr(model.config, "pad_token_id", None) is None: model.config.pad_token_id = tokenizer.pad_token_id model_args.model_name_or_path = old_model_name_or_path @@ -902,4 +873,7 @@ def _is_moe_config(cfg) -> bool: "output_router_logits", "moe_layer_freq", ) + # Use text config to check moe for multi-modal models like Qwen3.5 + if hasattr(cfg, "text_config"): + cfg = cfg.text_config return any(getattr(cfg, k, None) not in (None, 0, False) for k in moe_keys) diff --git a/roll/models/trl_patches.py b/roll/models/trl_patches.py index 749606a78..26e551269 100644 --- a/roll/models/trl_patches.py +++ b/roll/models/trl_patches.py @@ -7,7 +7,11 @@ from trl import PreTrainedModelWrapper -def value_head_load_state_dict(self: PreTrainedModelWrapper, state_dict: Dict[str, Any], strict=False) -> None: +def value_head_load_state_dict(self: PreTrainedModelWrapper, state_dict: Dict[str, Any], strict=False, **kwargs) -> None: + # When called from PyTorch's set_model_state_dict (FSDP2 path), state_dict keys use + # standard nn.Module FQN format (e.g. 'pretrained_model.model.layers.0...'). + # We remap keys to match each sub-module's expected format, then load separately. + # **kwargs must be forwarded to propagate 'assign' from PyTorch's _load_model_state_dict. for name in list(state_dict.keys()): if name.startswith("v_head."): state_dict[name] = state_dict.pop(name) @@ -15,15 +19,15 @@ def value_head_load_state_dict(self: PreTrainedModelWrapper, state_dict: Dict[st state_dict[name.replace("pretrained_model.", "")] = state_dict.pop(name) pretrained_model = getattr(self, "pretrained_model", None) if pretrained_model is not None: - pretrained_model.load_state_dict(state_dict, strict=False) + pretrained_model.load_state_dict(state_dict, strict=False, **kwargs) v_head: nn.Module = getattr(self, "v_head", None) if v_head is not None: for k in list(state_dict.keys()): if "v_head." in k: state_dict[k.replace("v_head.", "")] = state_dict.pop(k) - v_head.load_state_dict(state_dict, strict=False) + v_head.load_state_dict(state_dict, strict=False, **kwargs) else: - self.load_state_dict(state_dict, strict=False) + self.load_state_dict(state_dict, strict=False, **kwargs) def token_classifier_forward( diff --git a/roll/pipeline/agentic/agent_runner/__init__.py b/roll/pipeline/agentic/agent_runner/__init__.py new file mode 100644 index 000000000..317a617cf --- /dev/null +++ b/roll/pipeline/agentic/agent_runner/__init__.py @@ -0,0 +1,3 @@ +from roll.pipeline.agentic.agent_runner.base import AgentRunner, EpisodeResult + +__all__ = ["AgentRunner", "EpisodeResult"] diff --git a/roll/pipeline/agentic/agent_runner/base.py b/roll/pipeline/agentic/agent_runner/base.py new file mode 100644 index 000000000..ecd5a071f --- /dev/null +++ b/roll/pipeline/agentic/agent_runner/base.py @@ -0,0 +1,125 @@ +from abc import ABC, abstractmethod +from typing import Any, Dict, List, Optional, TYPE_CHECKING + +import httpx +from omegaconf import DictConfig + +if TYPE_CHECKING: + from roll.pipeline.agentic.agentic_config import EnvManagerConfig + + +class EpisodeResult: + """Structured result from a single episode execution.""" + + def __init__( + self, + status: str, + score: float, + step_scores: Optional[List[float]] = None, + agent_exit_reason: str = "", + metrics: Optional[Dict[str, Any]] = None, + ): + self.status = status + self.score = score + self.step_scores = step_scores or [] + self.agent_exit_reason = agent_exit_reason + self.metrics = metrics or {} + + def to_dict(self) -> Dict[str, Any]: + return { + "status": self.status, + "score": self.score, + "step_scores": self.step_scores, + "agent_exit_reason": self.agent_exit_reason, + **self.metrics, + } + + +class AgentRunner(ABC): + """Agent interaction loop abstraction. + + AgentRunner encapsulates how an agent interacts with an environment and + calls the LLM. It does NOT concern itself with: + - Trajectory collection (handled transparently by ProxyServer / MessageTracker) + - Training sample construction (handled by EnvManager.formulate_rollouts) + - Episode scheduling (handled by run_rollout_loop outer loop) + + Subclasses implement ``run_job(seed)``: load data, run a full episode, + return an ``EpisodeResult``. + """ + + def __init__( + self, + base_url: str, + env_id: int, + env_config: DictConfig, + worker_config: Optional["EnvManagerConfig"] = None, + **kwargs, + ): + """ + Args: + base_url: LLM inference service URL. For Proxy mode this is the + local ProxyServer address (e.g. http://127.0.0.1:8000). + env_id: Environment instance ID, fixed for the runner's lifetime. + Used by ProxyServer for routing (as Authorization Bearer token). + env_config: Full environment configuration DictConfig. + Contains top-level keys (env_type, max_steps, agent_runner_cls, ...) + and nested ``config`` / ``env_config`` dicts. Subclasses access + ``env_config.config`` for runner-level settings and + ``env_config["env_config"]`` for env constructor params. + worker_config: EnvManager-level configuration (model_args, generating_args, etc.). + """ + self.base_url = base_url + self.env_id = env_id + self.env_config: DictConfig = env_config + self.worker_config = worker_config + self.env: Optional[Any] = None + self.env_params: Dict[str, Any] = {} + if "config" in self.env_config: + self.env_params = dict(self.env_config["config"]) + self.mode: str = self.env_params.get("mode", "train") + + @abstractmethod + def run_job(self, seed: int) -> EpisodeResult: + """Run a complete episode. + + Args: + seed: Episode seed for data loading and env initialisation. + + Returns: + EpisodeResult with score, step_scores, and metrics. + """ + ... + + def setup(self) -> None: + """One-time initialisation (create env, establish connections, etc.).""" + pass + + def teardown(self) -> None: + """Clean up resources.""" + pass + + def _llm_request( + self, + client: httpx.Client, + messages: List[Dict], + tools: Optional[List[Dict]] = None, + tool_choice: str = "auto", + ) -> Dict[str, Any]: + """Send an OpenAI-compatible chat completion request to ``base_url``. + + The request is intercepted by ProxyServer, which routes it via the + ``Authorization`` header's ``self.env_id`` to the corresponding + ``EnvManager.process_request`` for inference + trajectory recording. + """ + resp = client.post( + f"{self.base_url}/v1/chat/completions", + json={ + "messages": messages, + "tools": tools or [], + "tool_choice": tool_choice if tools else "none", + }, + headers={"Authorization": f"Bearer {self.env_id}"}, + ) + resp.raise_for_status() + return resp.json() diff --git a/roll/pipeline/agentic/agent_runner/gem_runner.py b/roll/pipeline/agentic/agent_runner/gem_runner.py new file mode 100644 index 000000000..35ac759a4 --- /dev/null +++ b/roll/pipeline/agentic/agent_runner/gem_runner.py @@ -0,0 +1,196 @@ +from typing import Any, Dict, List, Union + +import httpx +from omegaconf import DictConfig + +from roll.pipeline.agentic.agent_runner.base import AgentRunner, EpisodeResult +from roll.pipeline.agentic.env import gem +from roll.utils.logging import get_logger +from roll.utils.str_utils import contains_renderable_field + + +class GEMRunner(AgentRunner): + """AgentRunner for local gem.Env environments. + + The gem.Env (e.g. SokobanEnv) returns raw text observations and env + instructions. GEMRunner is responsible for constructing the OpenAI + message history and driving the interaction loop: + + env.reset(seed) -> (obs_text, info) + loop: + build messages from history + LLM request -> response + env.step(action_text) -> (obs_text, reward, terminated, truncated, info) + + Applicable to: Sokoban, FrozenLake, and other gym-like environments. + """ + + def __init__(self, base_url: str, env_id: int, env_config: DictConfig, **kwargs): + super().__init__(base_url, env_id, env_config, **kwargs) + self.logger = get_logger() + self.max_steps: int = env_config.get("max_steps", 10) + self.http_timeout: float = float(env_config.get("http_timeout", 3600.0)) + self.system_template: str = env_config["agent_system_template"] + self.agent_template = env_config["agent_template"] + self.setup() + + def setup(self) -> None: + """Create a gem.Env instance from config. + + - ``env_type`` comes from the top-level env_config (e.g. "sokoban"). + - env constructor params come from ``env_config["config"]`` + (e.g. dim_room, num_boxes, search_depth). + - If ``tool_wrapper`` is configured, wraps the env for text-based tool support. + """ + env_type = self.env_config.get("env_type", "sokoban") + self.env = gem.make(env_id=env_type, **self.env_params) + if "tool_wrapper" in self.env_config: + from roll.pipeline.agentic.tools.tool_env_wrapper import tool_wrapper + self.env = tool_wrapper( + self.env, + wrapper_args=self.env_config.tool_wrapper.wrapper_args, + tool_configs=self.env_config.tool_wrapper.tool_configs, + ) + + def _render_observation(self, obs_text: str, turn_idx: int, actions_left: int) -> str: + """Render observation text through the agent_template with optional fields.""" + render_dict: Dict[str, Any] = {"observation": obs_text} + if contains_renderable_field(self.agent_template, "turn_idx"): + render_dict["turn_idx"] = turn_idx + if contains_renderable_field(self.agent_template, "suffix"): + render_dict["suffix"] = "" + if contains_renderable_field(self.agent_template, "actions_left"): + render_dict["actions_left"] = actions_left + if contains_renderable_field(self.agent_template, "max_response_length"): + render_dict["max_response_length"] = self.env_config.get("max_tokens_per_step", 128) + return self.agent_template.format(**render_dict) + + def run_job(self, seed: int) -> EpisodeResult: + """Run a complete text-mode episode with the local gem.Env. + + If ``tool_wrapper`` is configured, the wrapper transparently intercepts + text actions and executes tools — no change to this loop. + """ + obs_text, info = self.env.reset(seed=seed) + if obs_text is None: + return EpisodeResult(status="NoData", score=0.0) + + env_instruction = info.get("env_instruction", "") + + messages: List[Dict[str, Any]] = [ + { + "role": "system", + "content": (f"{self.system_template}\n\n{env_instruction}" + if env_instruction else self.system_template), + }, + { + "role": "user", + "content": self._render_observation(obs_text, turn_idx=1, actions_left=self.max_steps), + }, + ] + + rewards: List[float] = [] + try: + with httpx.Client(timeout=self.http_timeout) as client: + for turn_index in range(1, self.max_steps+1): + resp_json = self._llm_request(client, messages) + + if "error" in resp_json: + break + + assistant_content = resp_json["choices"][0]["message"].get("content", "") + messages.append({"role": "assistant", "content": assistant_content}) + + obs_text, reward, terminated, truncated, _ = self.env.step(assistant_content) + rewards.append(reward) + + if terminated or truncated: + break + + user_content = self._render_observation( + obs_text, + turn_idx=turn_index + 1, + actions_left=self.max_steps - turn_index, + ) + messages.append({"role": "user", "content": user_content}) + except httpx.HTTPError as e: + self.logger.error(f"[GEMRunner] HTTP error at step {len(rewards)}: {e}") + return EpisodeResult( + status="Failed", + score=0.0, + step_scores=rewards, + agent_exit_reason=f"HTTP error: {e}", + ) + + return EpisodeResult( + status="Finished", + score=float(sum(rewards)), + step_scores=rewards, + ) + + def teardown(self) -> None: + if self.env is not None: + self.env.close() + + +class ToolCallRunner(GEMRunner): + """AgentRunner for gem.Env environments using OpenAI function-calling protocol. + + Unlike GEMRunner (text-in/text-out), this runner: + - Passes ``tools`` from ``info`` to every LLM request. + - When the LLM responds with ``tool_calls``, passes the full message dict + to ``env.step(message)``; otherwise passes plain text. + - Lets the env manage the conversation history — ``env.step()`` returns + ``(messages, reward, terminated, truncated, info)``. + + Applicable to: SokobanToolCallEnv and other envs that handle tool + execution internally. + """ + + def run_job(self, seed: int) -> EpisodeResult: + """Run a complete tool-call episode with the local gem.Env.""" + messages, info = self.env.reset(seed=seed) + if messages is None: + return EpisodeResult(status="NoData", score=0.0) + assert isinstance(messages, list), ( + f"ToolCallRunner requires env.reset() to return a message list, got {type(messages)}" + ) + tools = info.get("tools") or None + + rewards: List[float] = [] + try: + with httpx.Client(timeout=self.http_timeout) as client: + for _ in range(self.max_steps): + resp_json = self._llm_request(client, messages, tools=tools) + + if "error" in resp_json: + break + + message = resp_json["choices"][0]["message"] + + # Tool calls → pass full message dict; text → pass content string + if message.get("tool_calls"): + action: Union[Dict[str, Any], str] = message + else: + action = message.get("content", "") + + # env executes tool calls internally and returns updated messages + messages, reward, terminated, truncated, _ = self.env.step(action) + rewards.append(reward) + + if terminated or truncated: + break + except httpx.HTTPError as e: + self.logger.error(f"[ToolCallRunner] HTTP error at step {len(rewards)}: {e}") + return EpisodeResult( + status="Failed", + score=0.0, + step_scores=rewards, + agent_exit_reason=f"HTTP error: {e}", + ) + + return EpisodeResult( + status="Finished", + score=float(sum(rewards)), + step_scores=rewards, + ) diff --git a/roll/pipeline/agentic/agent_runner/rock/__init__.py b/roll/pipeline/agentic/agent_runner/rock/__init__.py new file mode 100644 index 000000000..5b8948a25 --- /dev/null +++ b/roll/pipeline/agentic/agent_runner/rock/__init__.py @@ -0,0 +1,5 @@ +from roll.pipeline.agentic.agent_runner.rock.rock_agent_runner import RockAgentRunner +from roll.pipeline.agentic.agent_runner.rock.push_runner import PushModeRunner +from roll.pipeline.agentic.agent_runner.rock.pull_runner import PullModeRunner + +__all__ = ["RockAgentRunner", "PushModeRunner", "PullModeRunner"] diff --git a/roll/pipeline/agentic/agent_runner/rock/pull_runner.py b/roll/pipeline/agentic/agent_runner/rock/pull_runner.py new file mode 100644 index 000000000..34b8664b6 --- /dev/null +++ b/roll/pipeline/agentic/agent_runner/rock/pull_runner.py @@ -0,0 +1,261 @@ +"""Pull-mode Harbor runner using Rock SDK's ModelService as the LLM transport. + +Communication direction: Roll -> Rock (sandbox). +Roll actively drives each inference step via ModelService.anti_call_llm(), +instead of waiting for the agent to call back through ALB/ingress. +""" + +import asyncio +import json +import shlex +import uuid +from typing import Any, Dict, Optional + +import httpx +from omegaconf import DictConfig + +from rock.sdk.job import Job +from rock.sdk.job.operator import Operator +from rock.sdk.job.trial.harbor import HarborTrial +from rock.sdk.sandbox.model_service.base import ModelService, ModelServiceConfig + +from roll.pipeline.agentic.agent_runner.base import EpisodeResult +from roll.pipeline.agentic.agent_runner.rock.rock_agent_runner import RockAgentRunner + +# Default port avoids 8080 which is occupied by other services in the sandbox. +_DEFAULT_MODEL_SERVICE_PORT = 28080 + + +class ModelServiceHarborTrial(HarborTrial): + """HarborTrial variant that installs and starts ModelService in the sandbox. + + ``on_sandbox_ready()`` is called by JobExecutor after ``sandbox.start()`` + but before ``trial.setup()``, so the sandbox is available and api_base can + be patched before ``setup()`` writes the Harbor YAML to disk. + """ + + def __init__(self, config, model_service_port: int = _DEFAULT_MODEL_SERVICE_PORT): + super().__init__(config) + self._model_service: Optional[ModelService] = None + self._model_service_port = model_service_port + self._sandbox_host_ip: Optional[str] = None + + async def on_sandbox_ready(self, sandbox) -> None: + """Install and start ModelService; patch api_base before setup() writes Harbor YAML.""" + await super().on_sandbox_ready(sandbox) + + obs = await sandbox.arun("hostname -I 2>/dev/null | awk '{print $1}'") + self._sandbox_host_ip = obs.output.strip() or getattr(sandbox, "host_ip", None) + + if self._sandbox_host_ip and getattr(self, "_config", None) and self._config.agents: + self._config.agents[0].kwargs["api_base"] = ( + f"http://{self._sandbox_host_ip}:{self._model_service_port}/v1" + ) + + ms_config = ModelServiceConfig( + enabled=True, + type="local", + install_cmd=f"pip install {shlex.quote('rl-rock[model-service]')} --timeout 600", + start_cmd=( + f"rock model-service start --type local" + f" --host 0.0.0.0 --port {self._model_service_port}" + ), + watch_agent_cmd=( + f"rock model-service watch-agent --pid ${{pid}}" + f" --host 127.0.0.1 --port {self._model_service_port}" + ), + ) + self._model_service = ModelService(sandbox=sandbox, config=ms_config) + + await self._model_service.install() + await self._model_service.start() + + @property + def model_service(self) -> Optional[ModelService]: + return self._model_service + + @property + def sandbox_host_ip(self) -> Optional[str]: + return self._sandbox_host_ip + + +class ModelServiceOperator(Operator): + """Custom operator that creates a ModelServiceHarborTrial.""" + + def __init__(self, model_service_port: int = 8080): + self._model_service_port = model_service_port + + def apply(self, config) -> list: + return [ModelServiceHarborTrial(config, model_service_port=self._model_service_port)] + + +class PullModeRunner(RockAgentRunner): + """Pull-mode Rock agent runner. + + Roll polls the sandbox via ``ModelService.anti_call_llm()`` to serve each + LLM request the agent makes, rather than exposing an HTTP endpoint that + the agent calls back through ALB/ingress. + """ + + def __init__(self, base_url: str, env_id: int, env_config: DictConfig, **kwargs): + super().__init__(base_url, env_id, env_config, **kwargs) + self._model_service_port: int = env_config.get("config", {}).get( + "model_service_port", _DEFAULT_MODEL_SERVICE_PORT + ) + + def run_job(self, seed: int) -> EpisodeResult: + """Load data, submit a Harbor job in pull mode, and return the result.""" + data_item = self._load_data_item(seed) + if data_item is None: + return EpisodeResult(status="NoData", score=0.0) + + instance_id = data_item.get("extra_info", {}).get("instanceid", "unknown") + job_id = "ms_harbor_job_{}_{}".format(instance_id, uuid.uuid4().hex[:8]) + + task_config = self._build_task_config_from_data_item(data_item) + + self.logger.info(f"[ModelServiceRunner] Submitting Pull-mode job: {job_id}") + + metrics = self._run_async_job( + self._run_pull_job(task_config, job_id), job_id, instance_id, + ) + return self._build_result_from_metrics(metrics) + + async def _run_pull_job(self, task_config: Dict[str, Any], job_id: str) -> Dict[str, Any]: + """Core Pull-mode execution: submit -> watch_agent -> inference_loop -> wait.""" + operator = ModelServiceOperator(model_service_port=self._model_service_port) + config = self._build_job_config(task_config, job_id) + + config.agents[0].kwargs["api_base"] = ( + f"http://127.0.0.1:{self._model_service_port}/v1" + ) + + job = Job(config=config, operator=operator) + + self.logger.info(f"[ModelServiceRunner] Submitting job {job_id}...") + await job.submit() + + trial_client = job._job_client.trials[0] + trial: ModelServiceHarborTrial = trial_client.trial + model_service = trial.model_service + + if model_service is None: + raise RuntimeError(f"[ModelServiceRunner] ModelService not initialised for job {job_id}") + + harbor_pid = str(trial_client.pid) + self.logger.info(f"[ModelServiceRunner] Starting watch_agent for pid={harbor_pid}") + asyncio.ensure_future(model_service.watch_agent(pid=harbor_pid)) + + instance_id = task_config.get("data_config", {}).get("instance_id", "unknown") + self.logger.info(f"[ModelServiceRunner] Starting inference loop for {instance_id}") + async with httpx.AsyncClient(timeout=3600.0) as http_client: + await self._inference_loop(model_service, http_client) + + self.logger.info("[ModelServiceRunner] Inference loop done, waiting for Harbor result...") + result = await job.wait() + + metrics = self._extract_metrics(result, job_id, instance_id) + + try: + pass_rate = await self._read_pass_rate_from_sandbox(trial_client.sandbox, config) + if pass_rate is not None: + metrics["pass_rate"] = pass_rate + except Exception as e: + self.logger.warning(f"[ModelServiceRunner] Failed to read report.json: {e}") + + return metrics + + async def _inference_loop(self, model_service: ModelService, http_client: httpx.AsyncClient) -> None: + """Drive LLM inference for the agent via the ModelService file-based protocol.""" + self.logger.info("[InferenceLoop] Waiting for first LLM request from agent...") + current_index = 0 + current_response_payload: Optional[str] = None + + max_retries = 3 + retry_delay = 1.0 + + while True: + raw_output = None + for attempt in range(max_retries): + try: + raw_output = await model_service.anti_call_llm( + index=current_index, + response_payload=current_response_payload, + ) + break + except Exception as e: + if attempt < max_retries - 1: + self.logger.warning( + f"[InferenceLoop] anti_call_llm error at index {current_index} " + f"(attempt {attempt + 1}/{max_retries}), retrying in {retry_delay * (2 ** attempt):.1f}s: {e}" + ) + await asyncio.sleep(retry_delay * (2 ** attempt)) + else: + self.logger.error( + f"[InferenceLoop] anti_call_llm failed after {max_retries} attempts " + f"at index {current_index}: {e}" + ) + if raw_output is None: + break + + raw_output_stripped = raw_output.strip() if raw_output else "" + + if not raw_output_stripped: + await asyncio.sleep(0.5) + continue + + if "SESSION_END" in raw_output_stripped: + self.logger.info( + f"[InferenceLoop] SESSION_END received after {current_index} request(s)" + ) + break + + request_dict = self._parse_llm_request(raw_output_stripped, current_index + 1) + if request_dict is None: + self.logger.warning( + f"[InferenceLoop] Unparseable request at index {current_index + 1}, " + "sending error response" + ) + response_dict = self._make_error_response("Failed to parse request") + else: + try: + resp = await http_client.post( + f"{self.base_url}/v1/chat/completions", + json=request_dict, + headers={"Authorization": f"Bearer {self.env_id}"}, + ) + resp.raise_for_status() + response_dict = resp.json() + except Exception as e: + self.logger.error( + f"[InferenceLoop] ProxyServer request error at index {current_index + 1}: {e}" + ) + response_dict = self._make_error_response(str(e)) + + current_index += 1 + current_response_payload = json.dumps(response_dict, ensure_ascii=False) + + self.logger.info("[InferenceLoop] Exited.") + + def _parse_llm_request(self, raw_output: str, index: int) -> Optional[Dict[str, Any]]: + """Parse the request JSON returned by ModelClient.anti_call_llm().""" + try: + return json.loads(raw_output) + except Exception as e: + self.logger.warning(f"[InferenceLoop] Could not parse request at index {index}: {e}") + return None + + @staticmethod + def _make_error_response(error_msg: str) -> Dict[str, Any]: + """Build a minimal OpenAI-compatible error response.""" + return { + "id": "error", + "object": "chat.completion", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": f"[Roll inference error] {error_msg}"}, + "finish_reason": "stop", + } + ], + } diff --git a/roll/pipeline/agentic/agent_runner/rock/push_runner.py b/roll/pipeline/agentic/agent_runner/rock/push_runner.py new file mode 100644 index 000000000..2b1cb2710 --- /dev/null +++ b/roll/pipeline/agentic/agent_runner/rock/push_runner.py @@ -0,0 +1,63 @@ +import uuid +from typing import Any, Dict + +from rock.sdk.job import Job + +from roll.pipeline.agentic.agent_runner.base import EpisodeResult +from roll.pipeline.agentic.agent_runner.rock.rock_agent_runner import RockAgentRunner + + +class PushModeRunner(RockAgentRunner): + """Push-mode Rock agent runner. + + The Harbor agent calls back to Roll's ProxyServer through ALB/ingress + to perform LLM inference. ``Job.run()`` blocks until the agent finishes. + """ + + def run_job(self, seed: int) -> EpisodeResult: + """Load data, submit a Harbor job in push mode, and return the result.""" + data_item = self._load_data_item(seed) + if data_item is None: + return EpisodeResult(status="NoData", score=0.0) + + instance_id = data_item.get("extra_info", {}).get("instanceid", "unknown") + job_id = "harbor_job_{}_{}".format(instance_id, uuid.uuid4().hex[:8]) + + task_config = self._build_task_config_from_data_item(data_item) + task_config["infer_callback_url"] = self.infer_callback_url + + self.logger.info( + "Submitting Rock Harbor job: {} with agent: {}".format( + job_id, task_config["llm_config"]["model_name"] + ) + ) + + metrics = self._run_async_job( + self._run_harbor_job(task_config, job_id), job_id, instance_id, + ) + return self._build_result_from_metrics(metrics) + + async def _run_harbor_job(self, task_config: Dict[str, Any], job_id: str) -> Dict[str, Any]: + """Run a Harbor job (submit + wait) and return extracted metrics.""" + config = self._build_job_config(task_config, job_id) + + self.logger.info("[Harbor] Running job with config: {}".format(config)) + + job = Job(config=config) + self.logger.info("[Harbor] Job instance created, submitting...") + await job.submit() + result = await job.wait() + self.logger.info("[Harbor] Job completed, result: {}".format(result)) + + instance_id = task_config.get("data_config", {}).get("instance_id", "unknown") + metrics = self._extract_metrics(result, job_id, instance_id) + + try: + sandbox = job._job_client.trials[0].sandbox + pass_rate = await self._read_pass_rate_from_sandbox(sandbox, config) + if pass_rate is not None: + metrics["pass_rate"] = pass_rate + except Exception as e: + self.logger.warning(f"[Harbor] Failed to read report.json: {e}") + + return metrics diff --git a/roll/pipeline/agentic/agent_runner/rock/rock_agent_runner.py b/roll/pipeline/agentic/agent_runner/rock/rock_agent_runner.py new file mode 100644 index 000000000..b35f1d75e --- /dev/null +++ b/roll/pipeline/agentic/agent_runner/rock/rock_agent_runner.py @@ -0,0 +1,372 @@ +import json +import os +from typing import Any, Dict, Optional + +import ray +from omegaconf import DictConfig +from rock.actions import Command, ReadFileRequest +from rock.sdk.bench import AgentConfig +from rock.sdk.job import JobConfig + +from roll.datasets.global_dataset import GlobalDataset, GlobalDatasetManager +from roll.pipeline.agentic.agent_runner.base import AgentRunner, EpisodeResult +from roll.pipeline.agentic.agentic_config import EnvManagerConfig +from roll.pipeline.agentic.proxy.router import create_proxy_router +from roll.utils.constants import RAY_NAMESPACE +from roll.utils.logging import get_logger + +# Path to the default Harbor service config YAML (relative to repo root). +_DEFAULT_JOB_CONFIG_PATH = "roll/pipeline/agentic/env/rock/harbor_service_config.yaml" + + +class RockAgentRunner(AgentRunner): + """Base AgentRunner for Rock SDK sandbox environments. + + Inherits from AgentRunner and adds Rock-specific logic: + - Harbor job config building + - Metrics extraction from Harbor JobResult + - Evaluation report reading (pass rate from SWE-bench) + + Subclasses: + PushModeRunner — agent calls back to Roll via ALB/ingress + PullModeRunner — Roll polls sandbox via ModelService.anti_call_llm + """ + + def __init__( + self, + base_url: str, + env_id: int, + env_config: DictConfig, + worker_config: Optional[EnvManagerConfig] = None, + **kwargs, + ): + super().__init__(base_url, env_id, env_config, worker_config=worker_config, **kwargs) + self.logger = get_logger() + + config = dict(env_config.get("config", {})) + + # Dataset for data loading (Rock scenarios need dataset to get task instances) + assert "dataset_name" in config, "RockAgentRunner requires dataset_name in env_config.config" + dataset_name = config.get("dataset_name", "") + self.dataset = GlobalDataset.options( + name=f"{self.mode}_{dataset_name}", + get_if_exists=True, + namespace=RAY_NAMESPACE, + ).remote(dataset_name=dataset_name, mode=self.mode) + dataset_manager = GlobalDatasetManager.options( + name=f"{self.mode}_dataset_manager", + get_if_exists=True, + namespace=RAY_NAMESPACE, + ).remote() + ray.get(dataset_manager.register.remote(dataset_name=dataset_name, dataset_ref=self.dataset)) + + job_config_path = config.get("job_config_path", _DEFAULT_JOB_CONFIG_PATH) + self._base_job_config = JobConfig.from_yaml(job_config_path) + + self.listen_port = env_config.get("proxy_port", 8000) + self.infer_callback_url = self._generate_callback_url() + + # ------------------------------------------------------------------ + # Data loading + # ------------------------------------------------------------------ + + def _load_data_item(self, seed: int) -> Optional[Dict[str, Any]]: + """Load a data item from the dataset using ``seed``. + + Returns None when the dataset is exhausted. + """ + if self.dataset is None: + return None + data_item: Optional[Dict] = ray.get(self.dataset.get_data_item.remote(seed=seed)) + if data_item is not None: + data_item["seed"] = seed + return data_item + + # ------------------------------------------------------------------ + # Callback URL + # ------------------------------------------------------------------ + + def _generate_callback_url(self) -> Optional[str]: + """Build the URL that the sandbox agent uses to reach Roll's ProxyServer.""" + job_name = os.environ.get("TASK_ID") + if job_name: + try: + url = create_proxy_router().get_callback_url(job_name, self.listen_port) + if url: + self.logger.info(f"[Harbor] Generated callback URL: {url}") + return url + except Exception as e: + self.logger.error(f"Failed to get callback URL: {e}") + + config_url = self.env_config.get("config", {}).get("infer_callback_url") + if config_url: + return config_url + + self.logger.info("[Harbor] No remote callback URL available. Agent will use default LLM_BASE_URL.") + return None + + # ------------------------------------------------------------------ + # Job config + # ------------------------------------------------------------------ + + def _build_job_config(self, task_config: Dict[str, Any], job_id: str) -> JobConfig: + """Build a Rock SDK JobConfig from task_config.""" + config = self._base_job_config.model_copy(deep=True) + config.job_name = job_id + + llm_config = task_config.get("llm_config", {}) + runtime_config = task_config.get("runtime_config", {}) + agent_config = task_config.get("agent_config", {}) + data_config = task_config.get("data_config", {}) + + config.agents = [ + AgentConfig( + name="swe-agent-internal", + model_name="openai/{}".format(llm_config.get("model_name", "default_model")), + max_timeout_sec=runtime_config.get("task_timeout_sec", 300), + kwargs={ + "api_key": llm_config.get("api_key", ""), + "api_base": task_config.get("infer_callback_url", ""), + "sweagent_config": agent_config.get("scaffold_config", "anthropic"), + "max_iterations": agent_config.get("max_iterations", 15), + "temperature": llm_config.get("generation_config", {}).get("temperature", 0.99), + "max_tokens": llm_config.get("generation_config", {}).get("max_new_tokens", 2048), + "num_retries": agent_config.get("num_retries", 4), + "tools_parse_function": "function_calling", + "full_history": True, + "max_observation_length": 10000, + }, + ) + ] + + if config.datasets: + config.datasets[0].task_names = [data_config.get("instance_id")] + config.datasets[0].registry.split = data_config.get("split") + config.datasets[0].name = data_config.get("dataset") + config.datasets[0].version = data_config.get("split") + + if config.environment: + config.environment.env["INSTANCE_ID"] = data_config.get("instance_id") + config.environment.env["DATASET"] = data_config.get("dataset") + config.environment.env["SPLIT"] = data_config.get("split") + config.environment.experiment_id = os.environ["TASK_ID"] + + config.experiment_id = os.environ["TASK_ID"] + config.verifier.native_config.template.name = f"swe-agent-internal/{data_config.get('dataset')}" + + return config + + # ------------------------------------------------------------------ + # Metrics + # ------------------------------------------------------------------ + + def _extract_metrics(self, result, job_id: str, instance_id: str) -> Dict[str, Any]: + """Extract a standardised metrics dict from a Harbor JobResult.""" + metrics: Dict[str, Any] = { + "job_id": job_id, + "instance_id": instance_id, + "status": str(result.status.value if hasattr(result.status, "value") else result.status), + "exit_code": getattr(result, "exit_code", 0), + "score": 0.0, + "agent_exit_reason": "", + "time_total_sec": 0.0, + "env_setup_time_ratio": 0.0, + "agent_setup_time_ratio": 0.0, + "agent_execution_time_ratio": 0.0, + } + + if hasattr(result, "trial_results") and result.trial_results: + trial_result = result.trial_results[0] + + exception_info = getattr(trial_result, "exception_info", None) + if exception_info: + metrics["agent_exit_reason"] = str(exception_info) + + verifier_result = getattr(trial_result, "verifier_result", None) + if verifier_result and getattr(verifier_result, "rewards", None): + metrics["score"] = verifier_result.rewards.get("reward", 0.0) + + t_total = self._get_duration(trial_result) + t_env_setup = self._get_duration(getattr(trial_result, "environment_setup", None)) + t_agent_setup = self._get_duration(getattr(trial_result, "agent_setup", None)) + t_logic_exec = self._get_duration(getattr(trial_result, "agent_execution", None)) + + metrics["time_total_sec"] = round(t_total, 2) + + if t_total > 0: + metrics["env_setup_time_ratio"] = round((t_env_setup / t_total) * 100, 2) + metrics["agent_setup_time_ratio"] = round((t_agent_setup / t_total) * 100, 2) + metrics["agent_execution_time_ratio"] = round((t_logic_exec / t_total) * 100, 2) + + return metrics + + def _get_duration(self, timing_obj) -> float: + """Extract elapsed seconds from a timing object or dict.""" + if timing_obj is None: + return 0.0 + + if isinstance(timing_obj, dict): + start = timing_obj.get("started_at") + end = timing_obj.get("finished_at") + else: + start = getattr(timing_obj, "started_at", None) + end = getattr(timing_obj, "finished_at", None) + + if not start or not end: + return 0.0 + + try: + if isinstance(start, str): + from dateutil import parser + + dt_start = parser.isoparse(start) + dt_end = parser.isoparse(end) + else: + dt_start = start + dt_end = end + + if dt_start.tzinfo is not None: + dt_start = dt_start.replace(tzinfo=None) + if dt_end.tzinfo is not None: + dt_end = dt_end.replace(tzinfo=None) + + return max(0.0, (dt_end - dt_start).total_seconds()) + except Exception: + return 0.0 + + @staticmethod + def _extract_tests_status(report: Dict[str, Any]) -> Optional[Dict[str, Any]]: + """Pull ``tests_status`` out of a SWE-bench report, tolerant of nesting.""" + if not isinstance(report, dict): + return None + if isinstance(report.get("tests_status"), dict): + return report["tests_status"] + for value in report.values(): + if isinstance(value, dict) and isinstance(value.get("tests_status"), dict): + return value["tests_status"] + return None + + async def _read_pass_rate_from_sandbox(self, sandbox, config) -> Optional[float]: + """Read the FAIL_TO_PASS pass rate from Harbor's SWE-bench evaluation report. + + Returns: + FAIL_TO_PASS success ratio in [0.0, 1.0], or None when the report + is missing or carries no FAIL_TO_PASS tests. + """ + try: + job_dir = f"{config.jobs_dir}/{config.job_name}" + find_result = await sandbox.execute( + Command(command=["find", job_dir, "-path", "*/verifier/evaluation/report.json"]) + ) + files = [line.strip() for line in (find_result.stdout or "").strip().split("\n") if line.strip()] + if not files: + self.logger.info(f"[Harbor][REWARD] no evaluation report.json under {job_dir} -> fallback to binary score") + return None + + response = await sandbox.read_file(ReadFileRequest(path=files[0])) + report = json.loads(response.content) + tests_status = self._extract_tests_status(report) + if not tests_status: + self.logger.info(f"[Harbor][REWARD] report.json has no tests_status -> fallback: {files[0]}") + return None + + f2p = tests_status.get("FAIL_TO_PASS", {}) + p2p = tests_status.get("PASS_TO_PASS", {}) + f2p_success = len(f2p.get("success", [])) + f2p_total = f2p_success + len(f2p.get("failure", [])) + if f2p_total == 0: + self.logger.info(f"[Harbor][REWARD] report.json has no FAIL_TO_PASS tests -> fallback: {files[0]}") + return None + + p2p_failure = len(p2p.get("failure", [])) + if p2p_failure > 0: + self.logger.info( + "[Harbor][REWARD] pass_rate=0.0000 (PASS_TO_PASS regressions=%d) FAIL_TO_PASS=%d/%d report=%s", + p2p_failure, f2p_success, f2p_total, files[0], + ) + return 0.0 + + pass_rate = f2p_success / f2p_total + self.logger.info( + "[Harbor][REWARD] pass_rate=%.4f FAIL_TO_PASS=%d/%d PASS_TO_PASS=%d/%d report=%s", + pass_rate, f2p_success, f2p_total, + len(p2p.get("success", [])), len(p2p.get("success", [])) + p2p_failure, + files[0], + ) + return pass_rate + except Exception as e: + self.logger.warning(f"[Harbor][REWARD] failed to read report.json -> fallback to binary score: {e}") + return None + + # ------------------------------------------------------------------ + # Task config helpers + # ------------------------------------------------------------------ + + def _build_task_config_from_data_item(self, data_item: Dict[str, Any]) -> Dict[str, Any]: + """Build common task_config dict from a loaded data item. + + Note: Push-mode callers must add ``infer_callback_url`` to the returned + dict; Pull mode does not need it. + """ + instance_id = data_item.get("extra_info", {}).get("instanceid", "unknown") + return { + "metadata": {"step_id": 0}, + "data_config": { + "instance_id": instance_id, + "dataset": data_item.get("extra_info", {}).get("datasetname", "unknown"), + "split": data_item.get("extra_info", {}).get("datasettype", "test"), + }, + "llm_config": { + "model_name": "glm-5", + "api_key": str(self.env_id), + "generation_config": self.worker_config.generating_args.to_dict(), + }, + "runtime_config": { + "task_timeout_sec": self.env_config.config.get("task_timeout_sec", 1800), + }, + "agent_config": { + "max_iterations": self.env_config.config.get("max_iterations", 15), + "scaffold_config": self.env_config.config.get("scaffold_config", "anthropic"), + "num_retries": self.env_config.config.get("num_retries", 4), + }, + } + + def _build_result_from_metrics(self, metrics: Dict[str, Any]) -> EpisodeResult: + """Convert the raw metrics dict into an ``EpisodeResult``.""" + return EpisodeResult( + status=metrics.pop("status", "Unknown"), + score=float(metrics.pop("score", 0.0)), + agent_exit_reason=metrics.pop("agent_exit_reason", ""), + metrics=metrics, + ) + + def _run_async_job(self, coro, job_id: str, instance_id: str) -> Dict[str, Any]: + """Run an async coroutine on an isolated event loop and return metrics. + + Shared boilerplate for both Push and Pull runners. + """ + import asyncio + import concurrent.futures + + try: + loop = asyncio.new_event_loop() + private_executor = concurrent.futures.ThreadPoolExecutor(max_workers=4) + loop.set_default_executor(private_executor) + + try: + return loop.run_until_complete(coro) + except Exception as e: + self.logger.error(f"[RockRunner] Job {job_id} failed: {e}", exc_info=True) + return { + "instance_id": instance_id, + "status": "RunnerError", + "score": 0.0, + "agent_exit_reason": str(e), + "job_id": job_id, + } + finally: + private_executor.shutdown(wait=False) + loop.close() + except Exception as e: + self.logger.error(f"[RockRunner] Outer error for job {job_id}: {e}") + return {"status": "Failed", "score": 0.0, "agent_exit_reason": str(e)} diff --git a/roll/pipeline/agentic/agent_runner/rock/sandbox_tool_runner.py b/roll/pipeline/agentic/agent_runner/rock/sandbox_tool_runner.py new file mode 100644 index 000000000..c0a94d880 --- /dev/null +++ b/roll/pipeline/agentic/agent_runner/rock/sandbox_tool_runner.py @@ -0,0 +1,6 @@ +""" +TODO: + 手动创建sandbox，并且在sandbox中执行工具。 + 工具可以在ROLL侧创建编写，在sandbox中执行。 + AgentLoop运行在ROLL侧，可以方便地进行上下文管理。 +""" diff --git a/roll/pipeline/agentic/agentic_config.py b/roll/pipeline/agentic/agentic_config.py index 231f3859d..e792226cf 100644 --- a/roll/pipeline/agentic/agentic_config.py +++ b/roll/pipeline/agentic/agentic_config.py @@ -236,6 +236,20 @@ def __post_init__(self): # This ensures student_train/student_infer/teacher are mapped correctly self._handle_opd_mapping() + # OPD mode: build tag_2_domain/domain_2_tag from teacher tag_included for routing + if self.is_pure_opd or self.use_opd: + self.tag_2_domain = {} + self.domain_2_tag = {} + for name, ref_cfg in self._reference_configs.items(): + for tag in ref_cfg.tag_included: + if tag not in self.tag_2_domain: + self.tag_2_domain[tag] = tag + if tag not in self.domain_2_tag: + self.domain_2_tag[tag] = {tag} + else: + self.tag_2_domain = {} + self.domain_2_tag = {} + # Now safe to access actor_infer (may have been mapped from student_infer) assert self.actor_infer.generating_args or self.train_env_manager.generating_args, "must have generating_args in env_manager or actor infer." @@ -269,6 +283,11 @@ def __post_init__(self): self.reference.worker_cls = "roll.pipeline.base_worker.ActorWorker" if self.critic.worker_cls is None: self.critic.worker_cls = "roll.pipeline.base_worker.CriticWorker" + # Multi-teacher: set worker_cls for each _reference_configs entry + if hasattr(self, '_reference_configs'): + for ref_cfg in self._reference_configs.values(): + if ref_cfg.worker_cls is None: + ref_cfg.worker_cls = "roll.pipeline.base_worker.ActorWorker" if self.reward: if self.reward.worker_cls is None: self.reward.worker_cls = "roll.pipeline.base_worker.InferWorker" diff --git a/roll/pipeline/agentic/agentic_pipeline.py b/roll/pipeline/agentic/agentic_pipeline.py index e6313537f..e4652a7c4 100644 --- a/roll/pipeline/agentic/agentic_pipeline.py +++ b/roll/pipeline/agentic/agentic_pipeline.py @@ -1,7 +1,7 @@ import json import os.path import time -from concurrent.futures import ThreadPoolExecutor +from contextlib import ExitStack from typing import Any, Dict, List import numpy as np @@ -31,16 +31,20 @@ from roll.utils.dynamic_batching import dynamic_batching_shard from roll.utils.functionals import ( RunningMoments, + clusters_have_disjoint_devices, agg_loss, + build_domain_routing_context, + compute_ref_log_probs_with_routing, compute_token_reward, masked_mean, reduce_metrics, - batch_balance + batch_balance, ) from roll.utils.train_infer_corrections import apply_train_infer_correction_to_batch from roll.utils.kl_controller import get_kl_controller from roll.utils.logging import get_logger from roll.utils.offload_states import OffloadStateType +from roll.utils.telemetry import attach_trace_context, get_tracer, inject_trace_context logger = get_logger() @@ -83,13 +87,17 @@ def __init__(self, pipeline_config: AgenticConfig): download_clusters = [self.actor_train, self.actor_infer] if self.use_ref_model: - self.reference: Any = Cluster( - name=self.pipeline_config.reference.name, - worker_cls=self.pipeline_config.reference.worker_cls, - resource_manager=self.resource_manager, - worker_config=self.pipeline_config.reference, - ) - download_clusters.append(self.reference) + self.references: Dict[str, Any] = {} + for name, ref_cfg in self.pipeline_config.reference_configs.items(): + self.references[name] = Cluster( + name=ref_cfg.name, + worker_cls=ref_cfg.worker_cls, + resource_manager=self.resource_manager, + worker_config=ref_cfg, + ) + download_clusters.extend(self.references.values()) + # Backward compat: self.reference points to the first teacher + self.reference = self.references[list(self.references.keys())[0]] if self.pipeline_config.adv_estimator == "gae": @@ -180,7 +188,14 @@ def __init__(self, pipeline_config: AgenticConfig): ray.get(refs) if self.use_ref_model: - refs.extend(self.reference.initialize(pipeline_config=self.pipeline_config, blocking=True)) + if clusters_have_disjoint_devices(self.references): + ref_init_refs = [] + for ref_cluster in self.references.values(): + ref_init_refs.extend(ref_cluster.initialize(pipeline_config=self.pipeline_config, blocking=False)) + ray.get(ref_init_refs) + else: + for ref_cluster in self.references.values(): + ref_cluster.initialize(pipeline_config=self.pipeline_config, blocking=True) ray.get([self.train_rollout_scheduler.initialize.remote(), self.val_rollout_scheduler.initialize.remote()]) @@ -208,6 +223,7 @@ def __init__(self, pipeline_config: AgenticConfig): def run(self): # Calculate tokens-per-second system throughput tps_timer = _Timer(window_size=5) + tracer = get_tracer("driver") for global_step in range(self.pipeline_config.max_steps): if global_step <= self.state.step: @@ -217,7 +233,10 @@ def run(self): metrics = {} # Add overall step timing - with Timer(name="pipeline_step_total", logger=None) as step_timer: + with (Timer(name="pipeline_step_total", logger=None) as step_timer, + tracer.start_as_current_span("pipeline_step", attributes={"global_step": global_step}), + ExitStack() as defer, + ): with tps_timer: # PHASE 1: Offload States if self.pipeline_config.adv_estimator == "gae": @@ -233,7 +252,8 @@ def run(self): self.actor_infer.offload_states(include=OffloadStateType.other_params) # PHASE 3: Model Update - with Timer(name="model_update", logger=None) as model_update_timer: + with Timer(name="model_update", logger=None) as model_update_timer, \ + tracer.start_as_current_span("model_update"): model_update_metrics: Dict = self.model_update(global_step) metrics["time/step_model_update"] =model_update_timer.last metrics.update(model_update_metrics) @@ -269,39 +289,37 @@ def run(self): # PHASE 6: Validation (every eval_steps) - Async val_future = None - val_metrics = {} - with Timer(name="val", logger=None) as val_timer: - if self.pipeline_config.eval_steps > 0 and global_step % self.pipeline_config.eval_steps == 0: - # Submit val task to thread pool asynchronously - val_future = self.executor.submit(self.val, global_step) - - # PHASE 7: Rollout Get Batch - with Timer(name="rollout", logger=None) as rollout_timer: - batch = ray.get(self.train_rollout_scheduler.get_batch.remote(batch, self.pipeline_config.rollout_batch_size)) - sample_uuids = [f"{traj_id}_{i}" for i, traj_id in enumerate(batch.non_tensor_batch['traj_id'])] - batch.non_tensor_batch['sample_uuid'] = np.array(sample_uuids, dtype=object) - if "get_batch_return_start_time" in batch.meta_info: - metrics["time/get_batch_cost_train"] = time.time() - batch.meta_info.pop("get_batch_return_start_time") - actor_infer_metrics = self.actor_infer.get_metrics() - metrics.update(reduce_metrics(actor_infer_metrics.meta_info.pop("metrics", {}))) - metrics.update(compute_rollout_traj_metrics(batch)) - - dump_rollout_trajectories(self.pipeline_config.rollout_dump_dir, global_step, batch) - - metrics["time/step_rollout"] = rollout_timer.last - metrics.update(reduce_metrics(batch.meta_info.pop("metrics", {}))) - batch.meta_info["global_step"] = global_step - batch.meta_info["_broadcast_non_tensor_batch"] = True - batch.meta_info["loss_mask_keys"] = ["response_mask"] - - # PHASE 8: Stop Server Sync (sync mode only) - Wait for async val to complete - if val_future is not None: - val_metrics = val_future.result() + if self.pipeline_config.eval_steps > 0 and global_step % self.pipeline_config.eval_steps == 0: + # Submit val task to thread pool asynchronously + val_future = self.executor.submit(self.val, global_step, inject_trace_context({})) + + # PHASE 7: Rollout Get Batch + with Timer(name="rollout", logger=None) as rollout_timer, \ + tracer.start_as_current_span("rollout"): + inject_trace_context(batch.meta_info) + batch = ray.get(self.train_rollout_scheduler.get_batch.remote(batch, self.pipeline_config.rollout_batch_size)) + defer.callback(lambda b=batch: DataProto.drop(b)) + sample_uuids = [f"{traj_id}_{i}" for i, traj_id in enumerate(batch.non_tensor_batch['traj_id'])] + batch.non_tensor_batch['sample_uuid'] = np.array(sample_uuids, dtype=object) + if "get_batch_return_start_time" in batch.meta_info: + metrics["time/get_batch_cost_train"] = time.time() - batch.meta_info.pop("get_batch_return_start_time") + actor_infer_metrics = self.actor_infer.get_metrics() + metrics.update(reduce_metrics(actor_infer_metrics.meta_info.pop("metrics", {}))) + metrics.update(compute_rollout_traj_metrics(batch)) + + dump_rollout_trajectories(self.pipeline_config.rollout_dump_dir, global_step, batch) + + metrics["time/step_rollout"] = rollout_timer.last + metrics.update(reduce_metrics(batch.meta_info.pop("metrics", {}))) + batch.meta_info["global_step"] = global_step + batch.meta_info["_broadcast_non_tensor_batch"] = True + batch.meta_info["loss_mask_keys"] = ["response_mask"] - if len(val_metrics) > 0: + if val_future is not None: + val_metrics = val_future.result() metrics.update(val_metrics) - metrics["time/step_val"] = val_timer.last + # PHASE 8: Stop Server Sync (sync mode only) - Wait for async val to complete if not self.pipeline_config.async_pipeline: # Suspend scheduler before offload actor infer, because there may be # some inflight redundant trajectories. @@ -322,7 +340,8 @@ def run(self): # During training: actor_train uses freed GPUs [0,1] # Next iteration: model_update reloads actor_infer to all GPUs [0,1,2,3] elif self.partial_gpu_mode: - with Timer(name="cal_ref_log_probs", logger=None) as shrink_timer: + with Timer(name="shrink_sampler", logger=None) as shrink_timer, \ + tracer.start_as_current_span("shrink_sampler"): target_gpus = [] # Collect actor_train GPUs if hasattr(self.actor_train.worker_config, 'device_mapping') and self.actor_train.worker_config.device_mapping: @@ -344,41 +363,57 @@ def run(self): metrics.update(reduce_metrics(batch.meta_info.pop("metrics", {}))) # PHASE 11: Reference Log Probs - with Timer(name="cal_ref_log_probs", logger=None) as cal_timer: + with Timer(name="cal_ref_log_probs", logger=None) as cal_timer, \ + tracer.start_as_current_span("cal_ref_log_probs"): # TODO better the code structure, move the dynamic batching and sequence packing to worker/strategy if self.pipeline_config.enable_reference: - worker_config = self.pipeline_config.reference if self.use_ref_model else self.pipeline_config.actor_train - worker = self.reference if self.use_ref_model else self.pipeline_config.actor_train - if worker_config.use_dynamic_batching_in_infer: - batch, dynamic_batching_metrics = dynamic_batching_shard( - batch, - worker.dp_size, - worker_config.max_tokens_per_microbatch_in_infer, - worker_config.sequence_length_round_in_infer, - worker_config.strategy_args.strategy_config.get("pipeline_model_parallel_size", 1), - worker_config.strategy_args.strategy_config.get("virtual_pipeline_model_parallel_size", None), - "reference/compute_log_probs", - ) - metrics.update(dynamic_batching_metrics) if not self.use_ref_model: + # LoRA branch: use actor_train with disabled adapter + worker_config = self.pipeline_config.actor_train + batch_balance(batch, dp_size=self.actor_train.dp_size, minibatch_size=len(batch)) batch.meta_info["disable_adapter"] = True batch.meta_info["is_offload_states"] = False - batch_balance(batch, dp_size=self.actor_train.dp_size, minibatch_size=len(batch)) + if worker_config.use_dynamic_batching_in_infer: + batch, dynamic_batching_metrics = dynamic_batching_shard( + batch, + self.actor_train.dp_size, + worker_config.max_tokens_per_microbatch_in_infer, + worker_config.sequence_length_round_in_infer, + worker_config.strategy_args.strategy_config.get("pipeline_model_parallel_size", 1), + worker_config.strategy_args.strategy_config.get("virtual_pipeline_model_parallel_size", None), + "reference/compute_log_probs", + ) + metrics.update(dynamic_batching_metrics) ref_log_probs_refs: List[ray.ObjectRef] = self.actor_train.compute_log_probs(batch, blocking=False) + ref_log_probs = DataProto.materialize_concat(data_refs=ref_log_probs_refs) + ref_log_probs.rename(old_keys="log_probs", new_keys="ref_log_probs") + batch = batch.union(ref_log_probs) + avg_ref_log_prob = masked_mean(batch.batch["ref_log_probs"], batch.batch["response_mask"][:, 1:]) + metrics.update(reduce_metrics(ref_log_probs.meta_info.pop("metrics", {}))) + metrics.update({"critic/ref_log_prob/mean": avg_ref_log_prob.item()}) else: - batch_balance(batch, dp_size=self.reference.dp_size, minibatch_size=len(batch)) - ref_log_probs_refs: List[ray.ObjectRef] = self.reference.compute_log_probs(batch, blocking=False) + saved_routed_experts = batch.batch.pop("routed_experts", None) + domain_to_teacher_names, domain_values = build_domain_routing_context( + batch, self.pipeline_config, domain_key="tags" + ) + batch = compute_ref_log_probs_with_routing( + batch=batch, + references=self.references, + pipeline_config=self.pipeline_config, + domain_to_teacher_names=domain_to_teacher_names, + domain_values=domain_values, + metrics_fn=metrics.update, + ) + if saved_routed_experts is not None: + batch.batch["routed_experts"] = saved_routed_experts - ref_log_probs = DataProto.materialize_concat(data_refs=ref_log_probs_refs) - ref_log_probs.rename(old_keys="log_probs", new_keys="ref_log_probs") - batch = batch.union(ref_log_probs) - avg_ref_log_prob = masked_mean(batch.batch["ref_log_probs"], batch.batch["response_mask"][:, 1:]) - metrics.update(reduce_metrics(ref_log_probs.meta_info.pop("metrics", {}))) - metrics.update({"critic/ref_log_prob/mean": avg_ref_log_prob.item()}) + avg_ref_log_prob = masked_mean(batch.batch["ref_log_probs"], batch.batch["response_mask"][:, 1:]) + metrics.update({"critic/ref_log_prob/mean": avg_ref_log_prob.item()}) metrics["time/step_ref_log_probs_values_reward"] = cal_timer.last # PHASE 12: Old Log Probs & Values - with Timer(name="cal_old_log_probs_values", logger=None) as cal_old_logpb_timer: + with Timer(name="cal_old_log_probs_values", logger=None) as cal_old_logpb_timer, \ + tracer.start_as_current_span("cal_old_log_probs_values"): if self.pipeline_config.enable_reference and not self.use_ref_model: batch.meta_info["disable_adapter"] = False batch.meta_info["is_offload_states"] = False @@ -422,18 +457,26 @@ def run(self): batch.batch["ref_log_probs"] = batch.batch["old_log_probs"].clone() avg_ref_log_prob = masked_mean(batch.batch["ref_log_probs"], batch.batch["response_mask"][:, 1:]) metrics.update({"critic/ref_log_prob/mean": avg_ref_log_prob.item()}) + # Multi-teacher: also mock per-teacher keys + for name in self.pipeline_config.reference_configs.keys(): + batch.batch[f"ref_log_probs_{name}"] = batch.batch["old_log_probs"].clone() + batch.batch[f"ref_log_probs_{name}_mask"] = torch.ones( + batch.batch.batch_size[0], dtype=torch.bool, device=batch.batch.device + ) metrics["time/step_old_log_probs_values"] = cal_old_logpb_timer.last # TODO 当前这个还没用处 - with Timer(name="cal_response_level_mask", logger=None) as timer: + with Timer(name="cal_response_level_mask", logger=None) as timer, \ + tracer.start_as_current_span("cal_response_level_mask"): # TODO 补充完善的过滤要求，不同环境需要维持统一过滤标识 batch, mask_metrics = get_agentic_response_level_mask(batch, self.pipeline_config) metrics.update(mask_metrics) metrics["time/step_cal_response_level_mask"] = timer.last # PHASE 13: Advantage Computation - with Timer(name="cal_response_norm_rewards", logger=None) as timer: + with Timer(name="cal_response_norm_rewards", logger=None) as timer, \ + tracer.start_as_current_span("cal_response_norm_rewards"): # Rewards need to be processed after grouping # We can group by tag(env_type)/traj_group_id(group)/batch(rollout_batch)... to compute rewards / advantages # The compute_response_level_rewards function injects a response_level_rewards key into batch.batch. @@ -442,14 +485,16 @@ def run(self): metrics.update(reward_metrics) metrics["time/step_cal_norm_rewards"] = timer.last - with Timer(name="cal_token_reward", logger=None) as timer: + with Timer(name="cal_token_reward", logger=None) as timer, \ + tracer.start_as_current_span("cal_token_reward"): # Expand compute_response_level_rewards and add kl_penalty. # batch, kl_metrics = apply_kl_penalty(data=batch, kl_ctrl=self.kl_ctrl, kl_penalty=self.pipeline_config.kl_penalty) batch, token_level_metrics = compute_token_reward(batch, self.pipeline_config, self.kl_ctrl) metrics.update(token_level_metrics) metrics["time/step_cal_token_reward"] = timer.last - with Timer(name="compute_advantage", logger=None) as timer: + with Timer(name="compute_advantage", logger=None) as timer, \ + tracer.start_as_current_span("compute_advantage"): # Is the advantage calculated globally across the batch, or within each group? batch = agentic_compute_advantage( data=batch, @@ -470,7 +515,8 @@ def run(self): metrics.update(corr_metrics) # PHASE 14: Training (critic + actor) - with Timer(name="train_timer", logger=None) as train_timer: + with Timer(name="train_timer", logger=None) as train_timer, \ + tracer.start_as_current_span("train"): if self.pipeline_config.adv_estimator == "gae": critic_train_metrics_refs: List[ray.ObjectRef] = self.critic.train_step(batch, blocking=False) @@ -503,7 +549,8 @@ def run(self): tps_timer.push_units_processed(n=torch.sum(batch.batch["attention_mask"]).detach().item()) metrics["time/step_train"] = train_timer.last - with Timer(name="compute_data_metrics", logger=None) as data_metrics_timer: + with Timer(name="compute_data_metrics", logger=None) as data_metrics_timer, \ + tracer.start_as_current_span("compute_data_metrics"): data_metrics = compute_train_data_metrics(batch=batch) metrics["time/step_compute_data_metrics"] = data_metrics_timer.last @@ -517,7 +564,8 @@ def run(self): self.do_checkpoint(global_step=global_step) - with Timer(name="log", logger=None) as log_timer: + with Timer(name="log", logger=None) as log_timer, \ + tracer.start_as_current_span("log"): if self.pipeline_config.logging_steps > 0 and global_step % self.pipeline_config.logging_steps == 0: if int(os.environ.get("RAY_PROFILING", "0")): timeline_dir = os.path.join(self.pipeline_config.profiler_output_dir, "timeline") @@ -577,44 +625,52 @@ def run(self): self.val_rollout_scheduler.shutdown.remote(), ]) - logger.info("pipeline complete!") - - def val(self, global_step): + def val(self, global_step, trace_meta): + defer = ExitStack() batch = DataProto() metrics = {} batch.meta_info["is_offload_states"] = False batch.meta_info["global_step"] = global_step - ray.get(self.val_dataset_manager.reset.remote()) - eval_batch = ray.get(self.val_rollout_scheduler.get_batch.remote(batch, self.pipeline_config.val_batch_size)) - - if "get_batch_return_start_time" in eval_batch.meta_info: - metrics["time/get_batch_cost_val"] = time.time() - eval_batch.meta_info.pop("get_batch_return_start_time") - - dump_rollout_trajectories(self.pipeline_config.rollout_dump_dir, global_step, eval_batch) - eval_metrics = reduce_metrics(eval_batch.meta_info.get("metrics", {})) - eval_score = get_episode_scores(eval_batch) - eval_metrics["score/mean"] = torch.mean(eval_score).detach().item() - eval_metrics["score/max"] = torch.max(eval_score).detach().item() - eval_metrics["score/min"] = torch.min(eval_score).detach().item() - - batch_grouped = eval_batch.group_by(keys="tags") - for group_name, group_batch in batch_grouped.items(): - traj_group_scores = [] - batch_traj_grouped = group_batch.group_by(keys="traj_group_id") - for batch_traj_group_name, batch_traj_group in batch_traj_grouped.items(): - traj_group_score = get_episode_scores(batch_traj_group) - traj_group_scores.append(traj_group_score.mean().item()) - eval_score = torch.tensor(traj_group_scores, dtype=torch.float) - eval_metrics[f"{group_name}/score/mean"] = torch.mean(eval_score).detach().item() - eval_metrics[f"{group_name}/score/max"] = torch.max(eval_score).detach().item() - eval_metrics[f"{group_name}/score/min"] = torch.min(eval_score).detach().item() - - metrics.update({f"val/{k}": v for k, v in eval_metrics.items()}) - logger.info(f"val_batch_size: {len(eval_batch)}") - logger.info(f"val metrics: {metrics}") - + with ( + Timer(name="val", logger=None) as val_timer, + attach_trace_context(trace_meta), + get_tracer("driver").start_as_current_span("val"), + ): + inject_trace_context(batch.meta_info) + ray.get(self.val_dataset_manager.reset.remote()) + eval_batch = ray.get(self.val_rollout_scheduler.get_batch.remote(batch, self.pipeline_config.val_batch_size)) + defer.callback(lambda b=eval_batch: DataProto.drop(b)) + + if "get_batch_return_start_time" in eval_batch.meta_info: + metrics["time/get_batch_cost_val"] = time.time() - eval_batch.meta_info.pop("get_batch_return_start_time") + + dump_rollout_trajectories(self.pipeline_config.rollout_dump_dir, global_step, eval_batch) + eval_metrics = reduce_metrics(eval_batch.meta_info.get("metrics", {})) + eval_score = get_episode_scores(eval_batch) + eval_metrics["score/mean"] = torch.mean(eval_score).detach().item() + eval_metrics["score/max"] = torch.max(eval_score).detach().item() + eval_metrics["score/min"] = torch.min(eval_score).detach().item() + + batch_grouped = eval_batch.group_by(keys="tags") + for group_name, group_batch in batch_grouped.items(): + traj_group_scores = [] + batch_traj_grouped = group_batch.group_by(keys="traj_group_id") + for batch_traj_group_name, batch_traj_group in batch_traj_grouped.items(): + traj_group_score = get_episode_scores(batch_traj_group) + traj_group_scores.append(traj_group_score.mean().item()) + eval_score = torch.tensor(traj_group_scores, dtype=torch.float) + eval_metrics[f"{group_name}/score/mean"] = torch.mean(eval_score).detach().item() + eval_metrics[f"{group_name}/score/max"] = torch.max(eval_score).detach().item() + eval_metrics[f"{group_name}/score/min"] = torch.min(eval_score).detach().item() + + metrics.update({f"val/{k}": v for k, v in eval_metrics.items()}) + logger.info(f"val_batch_size: {len(eval_batch)}") + logger.info(f"val metrics: {metrics}") + + defer.close() + metrics["time/step_val"] = val_timer.last return metrics def adjust_batch(self, data: DataProto, mode="copy") -> DataProto: @@ -625,7 +681,14 @@ def adjust_batch(self, data: DataProto, mode="copy") -> DataProto: actor_train_infer_bsz = self.pipeline_config.actor_train.infer_batch_size * self.actor_train.dp_size ref_infer_bsz = 1 - if hasattr(self, "reference"): + if hasattr(self, "references"): + # Multi-teacher: use LCM of all teacher infer batch sizes + ref_infer_bzs = [ + ref_cfg.infer_batch_size * self.references[name].dp_size + for name, ref_cfg in self.pipeline_config.reference_configs.items() + ] + ref_infer_bsz = int(np.lcm.reduce(np.array(ref_infer_bzs))) + elif hasattr(self, "reference"): ref_infer_bsz = self.pipeline_config.reference.infer_batch_size * self.reference.dp_size critic_train_bsz = 1 critic_infer_bsz = 1 @@ -713,7 +776,13 @@ def _validate_partial_gpu_config(self) -> bool: train_devices = set(self.actor_train.worker_config.device_mapping) infer_devices = set(self.actor_infer.worker_config.device_mapping) critic_devices = set(self.critic.worker_config.device_mapping) if hasattr(self, 'critic') and self.critic else set() - ref_devices = set(self.reference.worker_config.device_mapping) if self.pipeline_config.enable_reference else set() + ref_devices = set() + if self.pipeline_config.enable_reference: + if hasattr(self, 'references'): + for ref_cluster in self.references.values(): + ref_devices.update(ref_cluster.worker_config.device_mapping) + elif hasattr(self, 'reference'): + ref_devices = set(self.reference.worker_config.device_mapping) reward_devices = set(self.reward.worker_config.device_mapping) if self.reward else set() # VAL: VAL_NON_EMPTY - ensure device_mapping not empty @@ -723,13 +792,21 @@ def _validate_partial_gpu_config(self) -> bool: f"train={list(train_devices)}, infer={list(infer_devices)}" ) - # Universal validation: Reference must always colocate with actor_train (both Model A and B) + # Universal validation: Reference must colocate with actor_train (both Model A and B) # VAL: VAL_SUBSET (exact match) - reference colocation + # For multi-teacher: relax to warning since teachers may span different GPU sets if self.pipeline_config.enable_reference: - assert ref_devices == train_devices, ( - f"Reference device_mapping must match actor_train exactly: " - f"ref={list(ref_devices)}, train={list(train_devices)}" - ) + if ref_devices != train_devices: + if self.pipeline_config.is_multi_teacher: + logger.warning( + f"Multi-teacher: Reference devices {sorted(ref_devices)} != train devices {sorted(train_devices)}, " + f"colocation may not be guaranteed for all teachers." + ) + else: + assert ref_devices == train_devices, ( + f"Reference device_mapping must match actor_train exactly: " + f"ref={list(ref_devices)}, train={list(train_devices)}" + ) # Determine configuration mode if train_devices.isdisjoint(infer_devices): diff --git a/roll/pipeline/agentic/env/__init__.py b/roll/pipeline/agentic/env/__init__.py index a983cd009..8221dcc7e 100644 --- a/roll/pipeline/agentic/env/__init__.py +++ b/roll/pipeline/agentic/env/__init__.py @@ -14,6 +14,7 @@ gem.register("roll_qa", entry_point="roll.pipeline.agentic.env.gem.qa_env:QaEnv") gem.register("sokoban_sandbox", entry_point="roll.pipeline.agentic.env.sandbox:SokobanSandboxEnv") gem.register("sokoban_native_env", entry_point="roll.pipeline.agentic.env.sokoban.native_env:SokobanNativeEnv") +gem.register("sokoban_tool_call", entry_point="roll.pipeline.agentic.env.sokoban.tool_call_env:SokobanToolCallEnv") gem.register("deepeyes", entry_point="roll.pipeline.agentic.env.deepeyes:DeepEyesEnv") gem.register("rock_tb_native_env", entry_point="roll.pipeline.agentic.env.sandbox.rock_tb_native_env:RockTBNativeEnv") diff --git a/roll/pipeline/agentic/env/cli_env/__init__.py b/roll/pipeline/agentic/env/cli_env/__init__.py new file mode 100644 index 000000000..34080dace --- /dev/null +++ b/roll/pipeline/agentic/env/cli_env/__init__.py @@ -0,0 +1,3 @@ +from .env import CLIEnv + +__all__ = ["CLIEnv"] diff --git a/roll/pipeline/agentic/env/cli_env/env.py b/roll/pipeline/agentic/env/cli_env/env.py new file mode 100644 index 000000000..7c8712da9 --- /dev/null +++ b/roll/pipeline/agentic/env/cli_env/env.py @@ -0,0 +1,340 @@ +import asyncio +import json +import re +import random +import time +from typing import Any, Dict, List, Optional, Tuple, Union + +import numpy as np +import gem +from gem import Env + +from .utils import ( + generate_random_cli_task, + validate_cli_command, + create_iflow_call, + create_iflow_search, + create_iflow_shell, + create_iflow_read_file, +) + + +class CLIEnv(Env): + + def __init__( + self, + render_mode: Optional[str] = None, + max_steps: int = 20, + workspace_dir: str = "/tmp/cli_workspace", + sandbox_image: str = "hub.docker.alibaba-inc.com/chatos/iflow-cli:4.0", # version + sandbox_base_url: str = "https://xrl-sandbox.alibaba-inc.com", + auto_clear_seconds: int = 60 * 20, # 20 minutes + format_penalty: float = -0.1, + debug_info: bool = False, + ): + """ + Initialize the CLI environment. + + Args: + render_mode: The render mode ("human" or "ansi") + max_steps: Maximum steps per episode + workspace_dir: Working directory in sandbox + sandbox_image: Docker image for sandbox + sandbox_base_url: Base URL for sandbox API + auto_clear_seconds: Auto clear timeout for sandbox + format_penalty: Format penalty + """ + self.max_steps = max_steps + self.workspace_dir = workspace_dir + self.sandbox_image = sandbox_image + self.sandbox_base_url = sandbox_base_url + self.auto_clear_seconds = auto_clear_seconds + self.debug_info = debug_info + + self.render_mode = render_mode + self.format_penalty = format_penalty + + # State tracking + self.current_step = 0 + self.xrl_client = None + self.current_observation = "" + self.task_description = "" + self.target_files = [] + + # Action mappings + self.action_names = [ + "list_directory", + "read_file", + "read_abs_path_file", + "write_file", + "run_shell_command", + "search_file_content", + "web_search", + "todo_write", + "iflow_call", + ] + + def reset(self, seed: Optional[int] = None) -> Tuple[str, Dict[str, Any]]: + Env.reset(self, seed) + + if seed is not None: + random.seed(seed) + np.random.seed(seed) + + self.current_step = 0 + + # Initialize environment (xrl client or local mode) + self._init_sandbox_sync() + + # Generate new task + ( + self.task_description, + self.target_files, + self.annotated_actions, + self.task_completion_description, + ) = self._generate_task() + + self.task_completion_description += "If completed, return task_done, otherwise return task_fail." + + # Set up initial workspace + self._setup_workspace_sync() + + # Get initial observation + self.current_observation = self._format_initial_observation(self.task_description, self.workspace_dir) + + info = { + "task": self.task_description, + "target_files": self.target_files, + "workspace": self.workspace_dir, + "step": self.current_step, + "env_instruction": "", + } + + return self.current_observation, info + + def step(self, action: str) -> Tuple[str, float, bool, bool, Dict[str, Any]]: + """ + Execute an action in the environment. + + Args: + action: action string + + Returns: + observation: Result of the action + reward: Reward for the action + terminated: Whether episode is complete + truncated: Whether episode was truncated + info: Additional information + """ + self.current_step += 1 + metrics_agg_mode = { + "action_is_effective": "mean", + "action_is_valid": "mean", + "success": "last", + "format_penalty": "mean", + } + # Parse action + action_str = str(action).strip() + match = re.search(r"(.*?)", action_str, re.DOTALL) + if match: + # 使用正则表达式抽取和中间的内容 + action_str = match.group(1).strip() + + # Map action name to index + action_name = action_str.split()[0] + if action_name not in self.action_names: + self.current_observation = f"Invalid action: {action_name}. Available: {', '.join(self.action_names)}" + metrics = { + "action_is_effective": False, + "action_is_valid": False, + "success": False, + "format_penalty": -self.format_penalty, + } + info = {"metrics": metrics, "step": self.current_step, "metrics_agg_mode": metrics_agg_mode} + return ( + self.current_observation, + -self.format_penalty, + False, + False, + info, + ) + + try: + result = self._execute_cli_action_sync(action_str) + self.current_observation = result + + terminated = self._check_task_completion(result) + + if terminated: + reward = 1.0 + else: + reward = self._calculate_reward(action_str, result) + + truncated = self.current_step >= self.max_steps + + except Exception as e: + self.current_observation = f"Error: {str(e)}" + reward = -1.0 + terminated = False + truncated = False + + metrics = { + "action_is_effective": False, + "action_is_valid": False, + "success": False, + "format_penalty": -self.format_penalty, + } + info = {"metrics": metrics, "step": self.current_step, "action": action_str, "metrics_agg_mode": metrics_agg_mode} + return self.current_observation, reward, terminated, truncated, info + + def render(self): + if self.render_mode == "human": + print(f"Step: {self.current_step}") + print(f"Task: {self.task_description}") + print(f"Current observation: {self.current_observation}") + elif self.render_mode == "ansi": + return f"Step: {self.current_step}\nTask: {self.task_description}\nObservation: {self.current_observation}" + + def close(self): + if self.xrl_client: + self.xrl_client.close() + self.xrl_client = None + + def _setup_workspace_sync(self): + if self.xrl_client: + # Use xrl client for workspace setup + self.xrl_client.run(f"mkdir -p {self.workspace_dir}") + self.xrl_client.run(f"cd {self.workspace_dir}") + else: + raise Exception("env not ready") + + def _execute_cli_action_sync(self, action_str: str) -> str: + parts = action_str.split(maxsplit=1) + action_name = parts[0] + params = parts[1] if len(parts) > 1 else "" + + if self.debug_info: + print(f"action_str: {action_str}") + print(f"action_name: {action_name}") + print(f"params: {params}") + + # TODO validate action检查，避免危险的shell的命令 + if not validate_cli_command(params): + print(f"dangerous command: {params}") + + # TODO 当前的封装感觉还是不够优雅 + if self.xrl_client: + # Use xrl client for all operations + if action_name == "list_directory": + target_dir = params if params else "." + return self.xrl_client.run(f"ls -al {target_dir}") + elif action_name == "write_file": + if not params: + return "Error: Please specify filename and content" + + parts = params.split(maxsplit=1) + if len(parts) < 2: + return "Error: Please specify both filename and content" + filename, content = parts + # Use echo for simple content writing + escaped_content = content.replace('"', '\\"') + return self.xrl_client.run(f'echo "{escaped_content}" > {self.workspace_dir}/{filename}') + elif action_name == "search_file_content": + pattern = params if params else "" + if not pattern: + return "Error: Please specify search pattern" + return self.xrl_client.run(f"grep -r '{pattern}' {self.workspace_dir} || echo 'No matches found'") + elif action_name == "read_file": + # 拼上workspace_dir + return self.xrl_client.run(create_iflow_read_file(f"{self.workspace_dir}/{params}")) + elif action_name == "read_abs_path_file": + # 直接就是绝对路径 + return self.xrl_client.run(create_iflow_read_file(params)) + elif action_name == "run_shell_command": + return self.xrl_client.run(create_iflow_shell(params)) + elif action_name == "web_search": + return self.xrl_client.run(create_iflow_search(params)) + elif action_name == "iflow_call": + return self.xrl_client.run(create_iflow_call(params)) + + else: + return f"Unknown action: {action_name}" + + raise Exception("sandbox not ready") + + def _calculate_reward(self, action: str, result: str) -> float: + """Calculate reward based on action and result.""" + if "Error" in result: + return -1.0 + elif "No matches found" in result or "No such file" in result: + return -0.5 + elif any(target in result for target in self.target_files): + return 1.0 + else: + return 0.1 # Small positive reward for successful actions + + def _check_task_completion(self, result: str) -> bool: + # # 利用iflow call的能力，所以之类就是一个agentic能力 + # task_result = self._execute_cli_action_sync( + # f"iflow_call '{self.task_completion_description}'" + # ) + """检查任务是否完成 - 使用simple_task_checker.py的能力""" + try: + from .simple_task_checker import SimpleTaskChecker + + checker = SimpleTaskChecker(self) + return checker.check_task_completion(self.task_description, self.target_files) + + except ImportError: + print("simple_task_checker.py 不可用，使用内置检查逻辑") + return self._check_task_completion_fallback(result) + except Exception as e: + print(f"任务检查错误: {str(e)}") + return False + + def _init_sandbox_sync(self, session_id: str = None): + """Initialize synchronous mode using xrl wrapper or local fallback.""" + try: + from .xrl_wrapper import create_sync_client + + self.xrl_client = create_sync_client( + sandbox_image=self.sandbox_image, + sandbox_base_url=self.sandbox_base_url, + auto_clear_seconds=self.auto_clear_seconds, + debug_info=self.debug_info, + ) + # Start the client + success = self.xrl_client.start() + if not success: + print("xrl client failed to start, using local mode") + self.xrl_client = None + except ImportError: + print("xrl wrapper not available, using local mode") + self.xrl_client = None + + def _generate_task(self) -> Tuple[str, List[str], str]: + return generate_random_cli_task() + + def _format_initial_observation(self, task_description, workspace_dir) -> str: + return ( + "Welcome to the CLI Environment!\n" + f"Task: {task_description}\n\n" + "Available actions:\n" + "0: list_directory - List the contents of the current directory\n" + "1: read_file - Read file content (filename required)\n" + "2: write_file - Write to file (filename and content required)\n" + "3: run_shell_command - Run a shell command\n" + "4: search_file_content - Search for content in files\n" + "5: web_search - Search the web for information\n" + "6: todo_write - Write a todo item\n" + "7: iflow_call - Call an iflow function\n\n" + f"Current directory: {workspace_dir}\n\n" + "Please complete the task using the available CLI commands.\n" + "Format your response as: action_name [parameters]\n" + ) + + def sample_random_action(self) -> str: + """Samples a random action given the current state.""" + import random + + return random.choice(self.action_names) diff --git a/roll/pipeline/agentic/env/cli_env/simple_task_checker.py b/roll/pipeline/agentic/env/cli_env/simple_task_checker.py new file mode 100644 index 000000000..9a9634b19 --- /dev/null +++ b/roll/pipeline/agentic/env/cli_env/simple_task_checker.py @@ -0,0 +1,206 @@ +""" +简化任务检查器 - 完全不使用iflow_call +""" + +import os +import re +from typing import List, Dict, Any + + +class SimpleTaskChecker: + """简化任务检查器，直接通过命令执行验证""" + + def __init__(self, env_instance): + self.env = env_instance + + def check_task_completion(self, task_description: str, target_files: List[str]) -> bool: + """ + 检查任务是否完成 + + Args: + task_description: 任务描述 + target_files: 目标文件列表 + + Returns: + 任务是否完成 + """ + try: + task_lower = task_description.lower() + + # 1. 检查文件存在性 + if target_files: + for target_file in target_files: + if not self._check_file_exists(target_file): + print(f"文件不存在: {target_file}") + return False + + # 2. 根据任务类型进行特定检查 + if "hello world" in task_lower: + return self._check_file_content("test.txt", "Hello World") + + elif "hello cli" in task_lower: + return self._check_python_script("hello.py", "Hello CLI") + + elif "test data" in task_lower and "test_dir" in task_lower: + return self._check_file_content("test_dir/info.txt", "Test data") + + elif "todo" in task_lower and "setup" in task_lower: + return self._check_todo_file() + + elif "search" in task_lower and "python" in task_lower: + return self._check_python_files() + + elif "directory" in task_lower: + return self._check_directory_structure(target_files) + + # 3. 默认检查：所有目标文件存在 + return True + + except Exception as e: + print(f"检查错误: {str(e)}") + return False + + def _check_file_exists(self, filename: str) -> bool: + """检查文件是否存在""" + try: + result = self.env._execute_cli_action_sync(f"read_file {filename}") + return "Error" not in result and "No such file" not in result + except: + return False + + def _check_file_content(self, filename: str, expected_content: str) -> bool: + """检查文件内容""" + try: + result = self.env._execute_cli_action_sync(f"read_file {filename}") + return expected_content in result + except: + return False + + def _check_python_script(self, filename: str, expected_output: str) -> bool: + """检查Python脚本""" + try: + # 检查文件内容 + content_result = self.env._execute_cli_action_sync(f"read_file {filename}") + if expected_output not in content_result: + return False + + # 尝试执行脚本 + exec_result = self.env._execute_cli_action_sync(f"run_shell_command python3 {filename}") + return expected_output in exec_result + except: + return False + + def _check_todo_file(self) -> bool: + """检查todo文件""" + try: + content = self.env._execute_cli_action_sync("read_file todo.txt") + keywords = ["Setup", "Develop", "Test"] + return all(keyword in content for keyword in keywords) + except: + return False + + def _check_python_files(self) -> bool: + """检查Python文件""" + try: + result = self.env._execute_cli_action_sync("run_shell_command ls *.py") + return ".py" in result and "No such file" not in result + except: + return False + + def _check_directory_structure(self, target_files: List[str]) -> bool: + """检查目录结构""" + try: + for target_file in target_files: + dir_path = os.path.dirname(target_file) + if dir_path and dir_path != ".": + dir_result = self.env._execute_cli_action_sync(f"list_directory {dir_path}") + if "No such file" in dir_result: + return False + return True + except: + return False + + def get_check_commands(self, task_description: str, target_files: List[str]) -> List[str]: + """获取检查命令列表""" + commands = ["list_directory"] + + # 添加文件检查命令 + for target_file in target_files: + commands.append(f"read_file {target_file}") + + # 根据任务类型添加特定命令 + task_lower = task_description.lower() + + if "python" in task_lower: + for target_file in target_files: + if target_file.endswith(".py"): + commands.append(f"run_shell_command python3 {target_file}") + + if "directory" in task_lower: + commands.append("run_shell_command find . -type d | head -10") + + if "search" in task_lower: + commands.append("run_shell_command ls -la") + + return commands + + +def create_simple_verification(task_description: str, target_files: List[str]) -> Dict[str, Any]: + """ + 创建简单验证配置 + + Args: + task_description: 任务描述 + target_files: 目标文件列表 + + Returns: + 验证配置字典 + """ + task_lower = task_description.lower() + + verification_map = { + "hello world": {"type": "file_content", "file": "test.txt", "expected": "Hello World"}, + "hello cli": {"type": "python_script", "file": "hello.py", "expected_output": "Hello CLI"}, + "test data": {"type": "file_content", "file": "test_dir/info.txt", "expected": "Test data"}, + "todo": {"type": "multi_content", "file": "todo.txt", "expected_keywords": ["Setup", "Develop", "Test"]}, + "search python": {"type": "file_exists", "pattern": "*.py"}, + } + + # 匹配任务类型 + for keyword, config in verification_map.items(): + if keyword in task_lower: + return config + + # 默认配置 + return {"type": "file_exists", "files": target_files} + + +# 使用示例 +if __name__ == "__main__": + """ + 使用示例： + + # 在CLIEnv中使用 + from simple_task_checker import SimpleTaskChecker + + # 创建检查器 + checker = SimpleTaskChecker(cli_env_instance) + + # 检查任务完成状态 + is_complete = checker.check_task_completion( + cli_env_instance.task_description, + cli_env_instance.target_files + ) + + # 获取验证命令 + commands = checker.get_check_commands( + cli_env_instance.task_description, + cli_env_instance.target_files + ) + + # 创建验证配置 + config = create_simple_verification( + cli_env_instance.task_description, + cli_env_instance.target_files + ) + """ diff --git a/roll/pipeline/agentic/env/cli_env/test_cli.py b/roll/pipeline/agentic/env/cli_env/test_cli.py new file mode 100644 index 000000000..d66da1800 --- /dev/null +++ b/roll/pipeline/agentic/env/cli_env/test_cli.py @@ -0,0 +1,66 @@ +#!/usr/bin/env python3 +""" +Simple test for the CLI environment with sync wrapper. +""" + +from cli_env import CLIEnv + + +def test_basic_functionality(): + """Test basic CLI environment functionality.""" + print("🧪 Testing CLI Environment...") + + # Test local mode (default) + env = CLIEnv(max_steps=5) + + print("Environment created") + + # Test reset + obs, info = env.reset() + print("Environment reset successful") + print(f"Task: {info['task']}") + print(f"Target files: {info['target_files']}") + print(f"observation: {obs}") + + # Test basic actions + # actions = [ + # # "乱七八糟的模型输出与分析balabalabala123list_directory", + # # "乱七八糟的模型输出与分析balabalabala234list_directory ~", + # # "乱七八糟的模型输出与分析balabalabala123list_directory /root/.npm-global/bin/iflow", + # # "乱七八糟的模型输出与分析balabalabala1write_file test.txt 'Hello from weixun/CLI env'", + # # "乱七八糟的模型输出与分析balabalabalaread_file test.txt", + # # "乱七八糟的模型输出与分析balabalabalaweb_search 中国的首都是哪里？", + # # # "乱七八糟的模型输出与分析balabalabalatodo_write 1 2 3 4", # not supported yet，主要是参数咋配的，还不太清楚 + # # "乱七八糟的模型输出与分析balabalabalarun_shell_command pwd", + # # "乱七八糟的模型输出与分析balabalabalalist_directory", + # "乱七八糟的模型输出与分析balabalabalaiflow_call '帮我在当前目录生成一个hello.txt文件，文件里面的内容是hello world'", + # "乱七八糟的模型输出与分析balabalabalaread_file hello.txt", + # # "乱七八糟的模型输出与分析balabalabalaiflow_call '搜一下中国的首都是哪里'", + # # "乱七八糟的模型输出与分析balabalabalarun_shell_command ls -al ~", + # ] + + actions = ( + ["乱七八糟的模型输出与分析balabalabalarun_shell_command ls -al ~"] + + env.annotated_actions + + ["乱七八糟的模型输出与分析balabalabalarun_shell_command ls -al ~"] * 3 + ) + + for i, action in enumerate(actions): + print(f"\n📋 Step {i+1}: {action}") + obs, reward, terminated, truncated, step_info = env.step(action) + print(f" Result: {obs}") + print(f" Reward: {reward}") + print(f" Terminated: {terminated}") + print(f" Truncated: {truncated}") + print(f" Step_info: {step_info}") + + # Test termination + if terminated or truncated: + break + + env.close() + print("Test completed successfully!") + + +if __name__ == "__main__": + test_basic_functionality() diff --git a/roll/pipeline/agentic/env/cli_env/test_no_iflow.py b/roll/pipeline/agentic/env/cli_env/test_no_iflow.py new file mode 100644 index 000000000..7f4087e32 --- /dev/null +++ b/roll/pipeline/agentic/env/cli_env/test_no_iflow.py @@ -0,0 +1,96 @@ +#!/usr/bin/env python3 +""" +测试无iflow_call的任务检查功能 +""" + +import os +import sys +import json +from cli_env import CLIEnv + + +def test_no_iflow_checking(): + """测试不使用iflow_call的任务检查""" + print("🧪 测试无iflow_call的任务检查...") + + # 创建环境 + env = CLIEnv(max_steps=5) + + # 重置环境获取任务 + obs, info = env.reset() + print(f"📋 任务: {info['task']}") + print(f"🎯 目标文件: {info['target_files']}") + + # 模拟完成任务 + task = info["task"] + target_files = info["target_files"] + + # 先检查 + # 测试新的检查逻辑 + print("\n🔍 测试无iflow_call检查...") + check_result = env._check_task_completion("") + print(f"检查结果: {check_result}") + + if check_result: + print("任务检查通过！") + else: + print("任务检查失败！") + + # 验证文件 + print("\n📁 验证文件状态：") + for target_file in target_files: + try: + # content = json.loads(env._execute_cli_action_sync(f"read_file {target_file}")) + content = env._execute_cli_action_sync(f"read_file {target_file}") + print(f"return content {content}") + # print(f"tool_results content: {content['tool_results']['content']}") + print( + f"{target_file}: {'存在' if 'Error' not in content and 'File not found' not in content else '不存在'}" + ) + except Exception as e: + print(f"{target_file}: 错误 - {str(e)}") + + print("\n🔍 执行任务...") + + # 根据任务类型执行相应操作 + if "Hello World" in task: + # 创建test.txt + env.step("write_file test.txt 'Hello World'") + + elif "Hello CLI" in task: + # 创建hello.py + env.step('write_file hello.py print("Hello CLI")') + + elif "test_dir" in task and "Test data" in task: + # 创建目录和文件 + env.step("run_shell_command mkdir -p test_dir") + env.step("write_file test_dir/info.txt 'Test data'") + + elif "todo" in task: + # 创建todo文件 + env.step("write_file todo.txt '1. Setup environment\n2. Develop features\n3. Test functionality'") + + print("\n🔍 测试无iflow_call检查...") + check_result = env._check_task_completion("") + print(f"检查结果: {check_result}") + + if check_result: + print("任务检查通过！") + else: + print("任务检查失败！") + + # 验证文件 + print("\n📁 验证文件状态：") + for target_file in target_files: + try: + content = env._execute_cli_action_sync(f"read_file {target_file}") + print(f"{target_file}: {'存在' if 'Error' not in content else '不存在'}") + except Exception as e: + print(f"{target_file}: 错误 - {str(e)}") + + env.close() + print("🎉 测试完成！") + + +if __name__ == "__main__": + test_no_iflow_checking() diff --git a/roll/pipeline/agentic/env/cli_env/utils.py b/roll/pipeline/agentic/env/cli_env/utils.py new file mode 100644 index 000000000..1d0e8c0d4 --- /dev/null +++ b/roll/pipeline/agentic/env/cli_env/utils.py @@ -0,0 +1,87 @@ +import random +import string +from typing import List, Tuple + + +def generate_random_cli_task() -> Tuple[str, List[str], str, str]: + """ + Generate a random CLI task with target files and verification description. + + Returns: + A tuple containing (task_description, target_files, action_commands, verification_description) + """ + tasks = [ + # ( + # "Create a test.txt file and enter Hello World", + # ["test.txt"], + # ["write_file test.txt 'Hello World'"], + # "Check that test.txt exists and contains 'Hello World'. Use read_file to verify the content.", + # ), + ( + "Create a Python script hello.py that prints 'Hello CLI'", + ["hello.py"], + ["write_file hello.py 'print(\"Hello CLI\")'"], + "Check that hello.py exists and contains Python code. Verify by running the script with python3 hello.py.", + ), + # ( + # "Create a directory named 'test_dir' and a file inside it named 'info.txt' with content 'Test data'", + # ["test_dir/info.txt"], + # [ + # "run_shell_command mkdir -p test_dir", + # "write_file test_dir/info.txt 'Test data'" + # ], + # "Check that test_dir directory exists and contains info.txt with content 'Test data'. Use list_directory and read_file to verify.", + # ), + # ( + # "Search for all Python files in the current directory", + # [], + # ["search_file_content '*.py'"], + # "Verify that the search returned Python files. Check the output contains .py file names.", + # ), + # ( + # "Create a todo list with three items: setup, develop, test", + # ["todo.txt"], + # [ + # "write_file todo.txt '1. Setup environment\n2. Develop features\n3. Test functionality'" + # ], + # "Check that todo.txt exists and contains three todo items. Use read_file to verify the content format.", + # ) + ] + + return random.choice(tasks) + + +def validate_cli_command(command: str) -> bool: + dangerous_commands = [ + "rm -rf /", + "sudo", + "shutdown", + "reboot", + ] + + command_lower = command.lower().strip() + + for dangerous in dangerous_commands: + if dangerous.lower() in command_lower: + return False + + return True + + +def create_iflow_call(prompt: str) -> str: + return f'/root/.npm-global/bin/iflow -y -p "{prompt}"' + + +def create_iflow_search(params: str) -> str: + json_parms = f'{{"tool_calls":[{{"id":"tool_search","function":{{"name":"web_search","arguments":"{{\\"query\\":\\"{params}\\"}}"}},"type":"function"}}]}}' + return f"/root/.npm-global/bin/iflow -t '{json_parms}'" + + +def create_iflow_shell(params: str) -> str: + json_parms = f'{{"tool_calls":[{{"id":"tool_shell","function":{{"name":"run_shell_command","arguments":"{{\\"command\\":\\"{params}\\"}}"}},"type":"function"}}]}}' + return f"/root/.npm-global/bin/iflow -t '{json_parms}'" + + +def create_iflow_read_file(params: str) -> str: + json_parms = f'{{"tool_calls":[{{"id":"tool_read_file","function":{{"name":"read_file","arguments":"{{\\"absolute_path\\":\\"{params}\\"}}"}},"type":"function"}}]}}' + return f"/root/.npm-global/bin/iflow -t '{json_parms}'" diff --git a/roll/pipeline/agentic/env/cli_env/xrl_wrapper.py b/roll/pipeline/agentic/env/cli_env/xrl_wrapper.py new file mode 100644 index 000000000..175df85a2 --- /dev/null +++ b/roll/pipeline/agentic/env/cli_env/xrl_wrapper.py @@ -0,0 +1,195 @@ +import time +import asyncio +import threading +from concurrent.futures import ThreadPoolExecutor + +# xrl.sdk.sandbox在使用的时候import，目的是因为后续想要兼容local模式 + + +class SyncXRLWrapper: + def __init__( + self, + sandbox_image: str, + sandbox_base_url: str, + auto_clear_seconds: int = 60 * 20, + debug_info: bool = False + ): + """ + Initialize the sync xrl wrapper. + + Args: + sandbox_image: Docker image for sandbox + sandbox_base_url: Base URL for sandbox API + auto_clear_seconds: Auto clear timeout for sandbox + """ + self.sandbox_image = sandbox_image + self.sandbox_base_url = sandbox_base_url + self.auto_clear_seconds = auto_clear_seconds + self.debug_info = debug_info + + self.sandbox = None + self.session_id = None + self.executor = ThreadPoolExecutor(max_workers=1) + + self._loop = None + self._thread = None + + def _ensure_async_loop(self): + if self._loop is None or not self._loop.is_running(): + self._loop = asyncio.new_event_loop() + + def run_loop(): + asyncio.set_event_loop(self._loop) + self._loop.run_forever() + + self._thread = threading.Thread(target=run_loop, daemon=True) + self._thread.start() + + def _run_async(self, coro): + self._ensure_async_loop() + future = asyncio.run_coroutine_threadsafe(coro, self._loop) + return future.result() + + def start_sandbox(self, session_id: str = None) -> bool: + # Start the xrl sandbox synchronously. + try: + from xrl.sdk.sandbox.client import Sandbox + from xrl.sdk.sandbox.config import SandboxConfig + from xrl.sdk.sandbox.request import CreateBashSessionRequest + + config = SandboxConfig( + image=self.sandbox_image, + base_url=self.sandbox_base_url, + auto_clear_seconds=self.auto_clear_seconds, + ) + + self.sandbox = Sandbox(config) + self._run_async(self.sandbox.start()) + + if session_id: + self.session_id = session_id + else: + self.session_id = f"sync-cli-session-{int(time.time())}" + + self._run_async(self.sandbox.create_session(CreateBashSessionRequest(session=self.session_id))) + + return True + + except ImportError: + print("xrl SDK not available") + return False + except Exception as e: + print(f"Failed to start xrl sandbox: {e}") + return False + + def execute_command(self, command: str) -> str: + if self.debug_info: + print(f"command: {command}") + if self.sandbox and self.session_id: + try: + from xrl.sdk.sandbox.request import BashAction + + response = self._run_async( + self.sandbox.run_in_session(BashAction(session=self.session_id, command=command)) + ) + if self.debug_info: + print(f"model response: {response}") + return response.output if hasattr(response, "output") else str(response) + except Exception as e: + return f"Sandbox error: {str(e)}" + else: + # Fallback to local execution + raise NotImplementedError("Local Sandbox not available") + + def stop_sandbox(self): + if self.sandbox: + try: + self._run_async(self.sandbox.stop()) + except Exception: + pass + + if self._loop and self._loop.is_running(): + self._loop.call_soon_threadsafe(self._loop.stop) + + if self.executor: + self.executor.shutdown(wait=False) + + def is_available(self) -> bool: + try: + sandbox_live = self.sandbox is not None + return sandbox_live + except: + + return False + + def __enter__(self): + self.start_sandbox() + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + """Context manager exit.""" + self.stop_sandbox() + + +class SimpleXRLClient: + """ + Simple synchronous client for xrl operations. + + This provides the simplest possible interface for xrl operations. + """ + + def __init__(self, **kwargs): + self.wrapper = SyncXRLWrapper(**kwargs) + self._started = False + self.debug_info = kwargs.get("debug_info", False) + + def start(self) -> bool: + """Start the client.""" + if not self._started: + self._started = self.wrapper.start_sandbox() + return self._started + + def run(self, command: str) -> str: + if self.debug_info: + print(f"in warp run: {command}") + if not self._started: + self.start() + return self.wrapper.execute_command(command) + + def close(self): + if self._started: + self.wrapper.stop_sandbox() + self._started = False + + def __enter__(self): + self.start() + return self + + def __exit__(self, *args): + self.close() + + +# Convenience function for simple usage +def create_sync_client( + sandbox_image: str = "hub.docker.alibaba-inc.com/chatos/iflow-cli:1.0", + sandbox_base_url: str = "https://xrl-sandbox.alibaba-inc.com", + auto_clear_seconds: int = 1200, + debug_info: bool = False +) -> SimpleXRLClient: + """ + Create a simple synchronous xrl client. + + Args: + sandbox_image: Docker image for sandbox + sandbox_base_url: Base URL for sandbox API + auto_clear_seconds: Auto clear timeout + + Returns: + SimpleXRLClient instance + """ + return SimpleXRLClient( + sandbox_image=sandbox_image, + sandbox_base_url=sandbox_base_url, + auto_clear_seconds=auto_clear_seconds, + debug_info=debug_info + ) diff --git a/roll/pipeline/agentic/env/deepeyes/env.py b/roll/pipeline/agentic/env/deepeyes/env.py index 8b3b31cfc..ea9cc270c 100644 --- a/roll/pipeline/agentic/env/deepeyes/env.py +++ b/roll/pipeline/agentic/env/deepeyes/env.py @@ -101,7 +101,6 @@ def encode_dataset(dataset, num_proc, encode_function, new_fingerprint=None): question_getter, ), batched=True, - batch_size=100, num_proc=num_proc, features=features, remove_columns=remove_columns, diff --git a/roll/pipeline/agentic/env/rock/harbor_service_config.yaml b/roll/pipeline/agentic/env/rock/harbor_service_config.yaml new file mode 100644 index 000000000..d0585babc --- /dev/null +++ b/roll/pipeline/agentic/env/rock/harbor_service_config.yaml @@ -0,0 +1,59 @@ +experiment_id: "rock_swe_test" + +datasets: + - name: "SWE-Env/SWE-Env" + registry: + split: "" + task_names: + - "" + +environment: + env: + INSTANCE_ID: "" + DATASET: "SWE-Env/SWE-Env" + SPLIT: "" + base_url: "" + xrl_authorization: "" + user_id: "" + image: "" + memory: "18g" + override_memory_mb: 16384 + cpus: 4 + startup_timeout: 1800 + cluster: "" + auto_clear_seconds: 300 + oss_deps: + "xxx": "xxx" + oss_mirror: + enabled: true + oss_bucket: "" + oss_access_key_id: "" + oss_access_key_secret: "" + oss_region: "" + oss_endpoint: "" + +agents: + - name: "swe-agent" + model_name: "openai/GLM" + kwargs: + api_key: "" + api_base: "" + sweagent_config: "anthropic" + max_iterations: 200 + temperature: "1.0" + max_tokens: 8192 + num_retries: 4 + tools_parse_function: "function_calling" + max_observation_length: 10000 + full_history: true + +orchestrator: + n_concurrent_trials: 1 + +agent_setup_timeout_multiplier: 3 + +verifier: + mode: native + native_config: + template: + name: swe-agent/SWE-Gym/SWE-Gym \ No newline at end of file diff --git a/roll/pipeline/agentic/env/rock/sanbox_manager.py b/roll/pipeline/agentic/env/rock/sanbox_manager.py new file mode 100644 index 000000000..032d574d6 --- /dev/null +++ b/roll/pipeline/agentic/env/rock/sanbox_manager.py @@ -0,0 +1,1048 @@ +import asyncio +import ast +import json +import os +import re +import shlex +import tempfile +import time +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple, Union + +import numpy as np +from jsonschema import validate +from jsonschema.exceptions import ValidationError + +from roll.pipeline.agentic.tools.iflow import iflow_config +from roll.pipeline.agentic.tools.iflow.cli_tool import IFlowCLITool +from roll.utils.logging import get_logger + +try: + from xrl.sdk.common.constants import Constants + from xrl.sdk.core.config import XRLConfig + from xrl.sdk.sandbox.client import Sandbox + from xrl.sdk.sandbox.config import SandboxConfig as XRLSandboxConfig + from xrl.sdk.sandbox.request import CreateBashSessionRequest, BashAction +except ImportError: + print("XRL SDK not available. Make sure it's installed.") + pass + +class SandboxConfig(XRLConfig): + image: str = "for-code-interpreter-registry-vpc.cn-hangzhou.cr.aliyuncs.com/chatos/python:3.11" + auto_clear_seconds: int = 60 * 60 + route_key: Optional[str] = None + startup_timeout: float = 360 + memory: str = "8g" + cpus: float = 2 + base_url: str = "https://xrl.alibaba-inc.com" + + +class RunStatus: + """Status codes for sandbox operations""" + SUCCESS = "success" + FAILED = "failed" + TIMEOUT = "timeout" + UNKNOWN_ERROR = "unknown_error" + CREATE_BOX = "create_box" + SANDBOX_START_FAILED = "sandbox_start_failed" + INFERENCE = "inference" + INFERENCE_FAILED = "inference_failed" + TEST = "test" + TEST_FAILED = "test_failed" + EXCEPTION = "exception" + + +class FailureMode: + """Failure modes for terminal-bench operations""" + NONE = "none" + UNSET = "unset" + AGENT_TIMEOUT = "agent_timeout" + UNKNOWN_AGENT_ERROR = "unknown_agent_error" + TEST_TIMEOUT = "test_timeout" + UNKNOWN_TEST_ERROR = "unknown_test_error" + PARSE_ERROR = "parse_error" + SANDBOX_START_FAILED = "sandbox_start_failed" + SANDBOX_CREATE_SESSION_FAILED = "sandbox_create_session_failed" + RUN_SANDBOX_COMMAND_FAILED = "run_sandbox_command_failed" + RUN_SANDBOX_UPLOAD_FAILED = "run_sandbox_upload_failed" + RUN_SANDBOX_EXCEPTION = "run_sandbox_exception" + RUN_CLI_TYPE_NOT_SUPPORT = "run_cli_type_not_support" + AGENT_INSTALLATION_FAILED = "agent_installation_failed" + IMAGE_NOT_FOUND_EXCEPTION = "image_not_found_exception" + + TOOL_CALL_PARSE_FAILED = "tool_call_parse_failed" + TOOL_EXECUTION_TIMEOUT = "tool_execution_timeout" + TOOL_EXECUTION_FAILED = "tool_execution_failed" + TOOL_EXECUTION_EXCEPTION = "tool_execution_exception" + TOOL_RESPONSE_PROCESSING_FAILED = "tool_response_processing_failed" + MODEL_RESPONSE_PROCESSING_EXCEPTION = "model_response_processing_exception" + + TEST_SESSION_CREATION_FAILED = "test_session_creation_failed" + TEST_DIRECTORY_CREATION_FAILED = "test_directory_creation_failed" + TEST_FILE_UPLOAD_FAILED = "test_file_upload_failed" + + IFLOW_SYSINFO_COMMAND_FAILED = "iflow_sysinfo_command_failed" + IFLOW_SYSINFO_PARSE_FAILED = "iflow_sysinfo_parse_failed" + IFLOW_SYSINFO_EXCEPTION = "iflow_sysinfo_exception" + + +class RunSessionResponse: + """Response object for sandbox session operations""" + def __init__(self): + self.exit_code = None + self.output = None + self.failure_reason = None + + + + + +class SandboxManager: + """ + Unified sandbox and session management utility. + Handles environment initialization, session management, and integrates with IFlowCLITool. + """ + def __init__(self, + sandbox_image: str, + logger, + xrl_authorization: str = "", + sandbox_base_url: str = "https://xrl.alibaba-inc.com", + run_type: str = "iflow-cli", + iflow_base_url: str = "", + iflow_api_key: str = "", + iflow_search_api_key: str = "", + iflow_selected_auth_type: str = "", + agent_version: str = "0.0.1", + run_region: str = "", + debug: bool = False, + default_timeout: float = 60.0): + + self.sandbox = None + self.sandbox_image = sandbox_image + self.logger = logger + self.xrl_authorization = xrl_authorization + self.sandbox_base_url = sandbox_base_url + + self.run_type = run_type + self.iflow_base_url = iflow_base_url + self.iflow_api_key = iflow_api_key + self.iflow_search_api_key = iflow_search_api_key + self.iflow_selected_auth_type = iflow_selected_auth_type + self.agent_version = agent_version + self.run_region = run_region + self.debug = debug + + self.active_sessions = {} + self.is_initialized = False + self.agent_session_name = "agent" + self.test_session_name = "test" + + self.max_retry = 3 + self.backoff = 2.0 + + self.image_id = sandbox_image + self.auto_clear_seconds = 60 * 60 + self.default_timeout = default_timeout + + self.failure_mode = FailureMode.NONE + self.run_status = RunStatus.SUCCESS + self.error_messages = [] + + self.is_environment_available = False + self.initialization_error = None + + self.iflow_tool = IFlowCLITool( + model_name="", + api_key=self.iflow_api_key, + base_url=self.iflow_base_url, + auth_type=self.iflow_selected_auth_type, + search_api_key=self.iflow_search_api_key, + debug=self.debug, + logger=self.logger + ) + self._initialize_sandbox_with_times() + + + + def _initialize_sandbox_with_times(self): + self.logger.info(f"[SANDBOX_INIT] START - Image ID: {self.image_id}") + self.sandbox_id = "" + + max_init_attempts = 3 + is_success = False + sandbox_ip = None + reason = "" + for attempt in range(1, max_init_attempts + 1): + self.logger.info(f"[SANDBOX_INIT] Attempt [{attempt}/{max_init_attempts}] - Initializing sandbox") + try: + is_success, sandbox_ip, reason = self._initialize_sandbox() + if is_success and sandbox_ip: + self.logger.info(f"[SANDBOX_INIT] Success on attempt {attempt}! - Sandbox started successfully with IP: {sandbox_ip}, sandbox_id: {self.sandbox_id}") + break + else: + if attempt < max_init_attempts: + wait_time = 2.0 * attempt + time.sleep(wait_time) + except Exception as e: + if attempt < max_init_attempts: + wait_time = 2.0 * attempt + time.sleep(wait_time) + + self.sandbox_ip = sandbox_ip + if is_success and sandbox_ip: + self.logger.info(f"[SANDBOX_INIT] Final Success! - Sandbox started successfully with IP: {sandbox_ip}, sandbox_id: {self.sandbox_id}") + self.is_environment_available = True + else: + self.logger.error(f"[SANDBOX_INIT] Final Failure! - Failed to start sandbox after {max_init_attempts} attempts: {reason}, sandbox_image: {self.image_id}, sandbox_ip: {self.sandbox_ip}, sandbox_id: {self.sandbox_id}") + self.is_environment_available = False + self.initialization_error = f"Failed to initialize sandbox after {max_init_attempts} attempts: {reason}, sandbox_ip: {sandbox_ip}, sandbox_image: {self.image_id}" + + + def _initialize_sandbox(self): + """Initialize sandbox and create sessions during environment construction""" + sandbox_ip = None + self.logger.info(f"[SANDBOX_START] START - Starting sandbox with image: {self.image_id}") + try: + success, sandbox_ip = self.start_sandbox( + max_retry=self.max_retry, + backoff=self.backoff + ) + if success and sandbox_ip: + self.logger.info(f"[SANDBOX_START] Success! - Sandbox environment initialized with IP: {sandbox_ip}, sandbox_id: {self.sandbox_id}") + else: + self.logger.error(f"[SANDBOX_START] Failed! - Failed to initialize sandbox environment") + self.failure_mode = FailureMode.SANDBOX_START_FAILED + return False, sandbox_ip, "Failed to start sandbox" + except Exception as e: + self.logger.error(f"[SANDBOX_START] Failed! - Error initializing sandbox environment: {e}") + self.failure_mode = FailureMode.SANDBOX_START_FAILED + return False, sandbox_ip, "Failed to start sandbox" + + + is_alive_response = asyncio.run(self.sandbox.is_alive()) + if not is_alive_response.is_alive: + self.logger.error(f"[SANDBOX_START] Failed! - Sandbox is not alive") + self.failure_mode = FailureMode.SANDBOX_START_FAILED + return False, sandbox_ip, "sandbox_util is not alive" + self.sandbox_ip = sandbox_ip + # 创建session + self.logger.info(f"[SESSION_CREATE] START - Creating session: {self.agent_session_name}") + try: + success = self.create_session(session=self.agent_session_name) + if success: + self.active_sessions[self.agent_session_name] = { + "created_at": time.time(), + "last_used": time.time() + } + self.logger.info(f"[SESSION_CREATE] Success! - Session '{self.agent_session_name}' created successfully") + else: + self.logger.error(f"[SESSION_CREATE] Failed! - Failed to create session '{self.agent_session_name}'") + self.failure_mode = FailureMode.SANDBOX_CREATE_SESSION_FAILED + return False, sandbox_ip, f"Failed to create session '{self.agent_session_name}'" + except Exception as e: + self.logger.error(f"[SESSION_CREATE] Failed! - Error creating session '{self.agent_session_name}': {e}") + self.failure_mode = FailureMode.SANDBOX_CREATE_SESSION_FAILED + return False, sandbox_ip, f"Failed to create session '{self.agent_session_name}'" + + # 初始化iflow-cli + self.logger.info(f"[AGENT_INSTALL] START - Installing IFlowCLITool") + try: + success, message = self._install_agent(self.agent_session_name) + if success: + self.is_initialized = True + self.logger.info(f"[AGENT_INSTALL] Success! - Sandbox and sessions initialized successfully") + else: + self.logger.error(f"[AGENT_INSTALL] Failed! - Agent installation failed: {message}") + self.failure_mode = FailureMode.AGENT_INSTALLATION_FAILED + return False, sandbox_ip, f"Agent installation failed: {message}" + except Exception as e: + self.logger.error(f"[AGENT_INSTALL] Failed! - Error during sandbox initialization: {e}") + self.failure_mode = FailureMode.AGENT_INSTALLATION_FAILED + return False, sandbox_ip, f"Agent installation failed: {str(e)}" + + return True, sandbox_ip, "" + + + def start_sandbox(self, max_retry: int = 3, backoff: float = 2.0): + """Start a sandbox instance""" + for attempt in range(1, max_retry + 1): + try: + start = time.time() + if self.run_region == "sg": + base_url = "http://xrl-sandbox-global.alibaba-inc.com" + else: + base_url = self.sandbox_base_url + config = SandboxConfig( + image=self.image_id, + xrl_authorization=self.xrl_authorization, + auto_clear_seconds=self.auto_clear_seconds, + startup_timeout=360, + base_url=base_url + ) + sandbox = Sandbox(config) + + asyncio.run(sandbox.start()) + cost = time.time() - start + self.logger.debug(f"image_id:{self.image_id}, sandbox_id:{sandbox.sandbox_id}, sandbox ip: {sandbox.host_ip}, start sandbox cost:{cost}") + self.sandbox = sandbox + self.sandbox_id = sandbox.sandbox_id + return True, sandbox.host_ip + except Exception as e: + self.logger.error(f"image_id:{self.image_id}, start_sandbox e:{e}") + if attempt == max_retry: + return False, None + time.sleep(backoff * attempt) + return False, None + + def create_session(self, session: str, max_retry: int = 3, backoff: float = 2.0): + """Create a session in the sandbox""" + for attempt in range(1, max_retry + 1): + try: + asyncio.run( + self.sandbox.create_session(CreateBashSessionRequest(session=session, startup_source=["/root" + "/.bashrc"], + env_enable=True, env={"HOME": "/root","IFLOW_ENV":"train","DISABLE_SEND_PV":"1", + "HF_ENDPOINT":"https://hf-mirror.com"}))) + return True + except Exception as e: + self.logger.error( + f"[{attempt}/{max_retry}] image_id:{self.image_id}, session:{session} create_session e:{e}, sandbox_id:{self.sandbox.sandbox_id}") + if attempt == max_retry: + return False + time.sleep(backoff * attempt) + return False + + def run_in_session(self, command: str, session: str, max_retry: int = 3, backoff: float = 2.0, timeout: float = None): + """ + Run a command in a session with retry logic for errors and empty outputs. + """ + self.logger.debug(f"[RUN_SESSION] START - image_id:{self.image_id}, sandbox_ip: {self.sandbox_ip}, session:{session}, command: {json.dumps(command, ensure_ascii=False)}, sandbox_id:{self.sandbox.sandbox_id}") + + response = RunSessionResponse() + last_error_msg = "" + + # Use default timeout if not specified + if timeout is None: + timeout = self.default_timeout + + for attempt in range(1, max_retry + 1): + try: + self.logger.debug(f"[RUN_SESSION] Attempt [{attempt}/{max_retry}] - Executing command in session '{session}' with timeout {timeout}s") + + async def run_with_timeout(): + return await self.sandbox.run_in_session( + BashAction(session=session, command=command, check="silent")) + + session_ret = asyncio.run( + asyncio.wait_for(run_with_timeout(), timeout=timeout)) + + response = session_ret + if session_ret.exit_code is None: + response.exit_code = -1 + if session_ret.output is None: + response.output = "" + if session_ret.failure_reason is None: + response.failure_reason = "" + + command_succeeded = response.exit_code == 0 + has_output = response.output and response.output.strip() + + self.logger.debug(f"[RUN_SESSION] Attempt [{attempt}/{max_retry}] Result - exit_code: {response.exit_code}, output_length: {len(response.output) if response.output else 0}, has_content: {bool(has_output)}") + + should_retry = False + retry_reason = "" + if "kill" in command: + self.logger.debug(f"[RUN_SESSION] Command contains 'kill', not retrying regardless of outcome") + return response + if not command_succeeded: + should_retry = True + retry_reason = f"Command failed with exit_code: {response.exit_code}" + elif not has_output and "-t" in command: + should_retry = True + retry_reason = "Command with -t succeeded but returned empty output" + + if not should_retry: + self.logger.debug(f"[RUN_SESSION] SUCCESS - Command executed successfully on attempt {attempt}/{max_retry}") + return response + + if attempt == max_retry and should_retry: + self.logger.error(f"[RUN_SESSION] FAILED - All {max_retry} attempts exhausted. Final reason: {retry_reason}, last_error_msg: {last_error_msg}, response: {json.dumps(response.output[:200], ensure_ascii=False)}") + return response + + if should_retry: + wait_time = backoff * attempt + time.sleep(wait_time) + + except asyncio.TimeoutError: + timeout_msg = f"Command execution timed out after {timeout} seconds" + if attempt == max_retry: + self.logger.error(f"[RUN_SESSION] FAILED - All {max_retry} attempts timed out. Timeout: {timeout}s, Sandbox ID: {self.sandbox.sandbox_id}, sandbox_ip: {self.sandbox_ip}, command: {command}") + response.exit_code = -1 + response.output = timeout_msg + response.failure_reason = "TIMEOUT" + return response + wait_time = backoff * attempt + time.sleep(wait_time) + except Exception as exc: + last_error_msg = str(exc) + if "/bin/bash: line 1: " in last_error_msg: + response.exit_code = -1 + response.output = last_error_msg + response.failure_reason = last_error_msg + return response + self.logger.error(f"[RUN_SESSION] EXCEPTION - Attempt [{attempt}/{max_retry}] failed with exception: {last_error_msg}") + if attempt == max_retry: + self.logger.error(f"[RUN_SESSION] FAILED - All {max_retry} attempts failed due to exceptions. Final error: {last_error_msg}") + response.exit_code = -1 + response.output = last_error_msg + response.failure_reason = last_error_msg + return response + wait_time = backoff * attempt + time.sleep(wait_time) + + self.logger.error(f"[RUN_SESSION] UNEXPECTED - Reached end of retry loop unexpectedly") + return response + + def stop_sandbox(self): + """Stop the sandbox instance""" + try: + if self.sandbox is None: + return True + asyncio.run(self.sandbox.stop()) + return True + except Exception as e: + self.logger.error(f"image_id:{self.image_id}, stop_sandbox e:{e}, sandbox_id:{self.sandbox.sandbox_id}") + return False + + def upload_file(self, file_path: Union[str, Path], target_path: str, max_retry: int = 3, backoff: float = 2.0): + """Upload a file to the sandbox""" + for attempt in range(1, max_retry + 1): + self.logger.debug( + f"[upload_file, {attempt}/{max_retry}] image_id:{self.image_id}, file_path:{file_path} target_path: {target_path}, sandbox_id:{self.sandbox.sandbox_id}") + try: + response = asyncio.run(self.sandbox.aupload(str(file_path), target_path)) + return response.success, response.message + except Exception as exc: + self.logger.error( + f"image_id:{self.image_id}, file_path:{file_path} target_path: {target_path}, upload failed: {str(exc)}, " + f"sandbox_id:{self.sandbox.sandbox_id}") + if attempt == max_retry: + return False, f"upload_file exp:{str(exc)}" + time.sleep(backoff * attempt) + return False, f"upload_file failed" + + def run_session_with_timeout(self, session_name: str, command: str, timeout, output_file, interval=10): + """Run a command with timeout and save output to a file""" + try: + start_time = time.time() + end_time = start_time + timeout + content = '' + + response = self.run_in_session(command=f"nohup {command} < /dev/null > {output_file} 2>&1 &", session=session_name) + if response.exit_code != 0: + if "511" in response.output or "/bin/bash: line 1: " in response.output: + self.logger.warning(f"HTTP 511 error or syntax error detected, attempting file-based command execution workaround") + try: + response, output_file = self._run_command_via_file(command=command, session=session_name) + + if response.exit_code == 0: + self.logger.info(f"File-based command execution succeeded") + else: + self.logger.warning(f"File-based command execution also failed") + return RunStatus.FAILED, "run command failed: " + response.output + except Exception as file_exc: + self.logger.error(f"File-based workaround failed: {str(file_exc)}") + return RunStatus.FAILED, "run command failed: " + str(file_exc) + else: + return RunStatus.FAILED, "run command failed: " + response.output + + pid = self._extract_pid(response.output) + if len(pid) == 0: + time.sleep(6) + content = self.read_content(session_name, output_file) + return RunStatus.SUCCESS, content + + + while time.time() < end_time: + content = self.read_content(session_name, output_file) + if not self.is_process_running(session_name, pid): + content = self.read_content(session_name, output_file) + return RunStatus.SUCCESS, content + + if time.time() >= end_time: + content = self.read_content(session_name, output_file) + return RunStatus.TIMEOUT, content + time.sleep(interval) + except Exception as e: + self.logger.error( + f"run_session_with_timeout exception, sandbox_id:{self.sandbox.sandbox_id}, sandbox_ip: {self.sandbox_ip}, command:{command}, exp:{str(e)}") + return RunStatus.UNKNOWN_ERROR, "run command exception: " + str(e) + return RunStatus.TIMEOUT, "run command timeout" + + def _extract_pid(self, output: str) -> str: + """Extract process ID from command output""" + lines = output.splitlines() + if not lines: + return "" + last_line = lines[-1].strip() + import re + match = re.match(r'\[\d+\]\s+(\d+)$', last_line) + return match.group(1) if match else "" + + def is_process_running(self, session_name: str, pid): + """Check if a process is still running""" + response = self.run_in_session(command=f"kill -0 {pid}", session=session_name) + return response.exit_code == 0 + + def read_content(self, session_name, output_file): + """Read content from a file in the sandbox""" + res = self.run_in_session(command=f"cat {output_file}", session=session_name) + if res.exit_code == 0: + return res.output + else: + print(f"cat {output_file} error, sandbox_id:{self.sandbox.sandbox_id}") + return '' + + def _run_command_via_file(self, command: str, session: str) -> RunSessionResponse: + """ + Workaround for 511 errors: write command to file, upload it, then execute + """ + self.logger.info(f"[FILE_WORKAROUND] START - Attempting file-based execution for command: {command[:100]}...") + response = RunSessionResponse() + temp_script_name = f"temp_command_{session}_{int(time.time())}.sh" + temp_output_file = f"temp_output_{session}_{int(time.time())}.txt" + + try: + if "tool_calls" in command: + tool_call = command + script_content = f"""#!/bin/bash +# Auto-generated script to workaround network issues +# Session: agent, Attempt: 1 + +set -e + +# Store the JSON payload in a variable using a here document +read -r -d '' JSON_PAYLOAD << 'EOF' || true +{tool_call} +EOF + +# Execute the command with the JSON payload +nohup iflow -t "$JSON_PAYLOAD" < /dev/null > {temp_output_file} 2>&1 & +""" + else: + script_content = f"""#!/bin/bash +# Auto-generated script to workaround network issues +# Session: agent, Attempt: 1 + +set -e +nohup {command} < /dev/null > {temp_output_file} 2>&1 & +""" + with tempfile.NamedTemporaryFile(mode='w', suffix='.sh', delete=False) as temp_file: + temp_file.write(script_content) + temp_file_path = temp_file.name + + self.logger.debug(f"[FILE_WORKAROUND] Created temporary script: {temp_file_path}") + + try: + self.logger.debug(f"[FILE_WORKAROUND] Uploading script to sandbox: {temp_script_name}") + upload_success, upload_message = self.upload_file( + temp_file_path, + f"{temp_script_name}" + ) + + if not upload_success: + self.logger.error(f"[FILE_WORKAROUND] Failed to upload script file: {upload_message}") + response.exit_code = -1 + response.output = f"Upload failed: {upload_message}" + response.failure_reason = "SCRIPT_UPLOAD_FAILED" + return response, temp_output_file + + self.logger.debug(f"[FILE_WORKAROUND] Successfully uploaded script to {temp_script_name}") + + chmod_command = f"chmod +x {temp_script_name}" + self.logger.debug(f"[FILE_WORKAROUND] Making script executable: {chmod_command}") + + chmod_ret = asyncio.run( + self.sandbox.run_in_session( + BashAction(session=session, command=chmod_command, check="silent"))) + + if chmod_ret.exit_code != 0: + self.logger.error(f"[FILE_WORKAROUND] chmod command failed: {chmod_ret.output}") + response.exit_code = -1 + response.output = f"chmod failed: {chmod_ret.output}" + response.failure_reason = "CHMOD_FAILED" + return response, temp_output_file + + self.logger.debug(f"[FILE_WORKAROUND] Script made executable successfully") + + exec_command = f"bash {temp_script_name}" + self.logger.info(f"[FILE_WORKAROUND] Executing script via run_session_with_timeout: {exec_command}") + session_ret = asyncio.run( + self.sandbox.run_in_session( + BashAction(session=session, command=exec_command, check="silent"))) + response = session_ret + if session_ret.exit_code is None: + response.exit_code = -1 + if session_ret.output is None: + response.output = "" + if session_ret.failure_reason is None: + response.failure_reason = "" + return response, temp_output_file + except Exception as exec_exc: + self.logger.error(f"[FILE_WORKAROUND] Failed to execute script via file method: {exec_exc}") + response.exit_code = -1 + response.output = f"Script execution failed: {str(exec_exc)}" + response.failure_reason = "SCRIPT_EXECUTION_FAILED" + return response, temp_output_file + + + except Exception as e: + self.logger.error(f"[FILE_WORKAROUND] File-based command execution failed: {str(e)}") + response.exit_code = -1 + response.output = f"File method error: {str(e)}" + response.failure_reason = "FILE_METHOD_ERROR" + return response, temp_output_file + + + def create_managed_session(self, session_name: str) -> bool: + """ + Create a new managed session in the sandbox with tracking. + """ + is_alive_response = asyncio.run(self.sandbox.is_alive()) + if not is_alive_response.is_alive: + print("sandbox_util is not alive") + return False + + try: + success = self.create_session(session=session_name) + + if success: + self.active_sessions[session_name] = { + "created_at": time.time(), + "last_used": time.time() + } + self.logger.debug(f"Session '{session_name}' created successfully") + else: + self.logger.error(f"Failed to create session '{session_name}'") + + return success + + except Exception as e: + self.logger.error(f"Error creating session '{session_name}': {e}") + return False + + def execute_command(self, session_name: str, command: str) -> RunSessionResponse: + """ + Execute a command in the specified session with session management. + """ + if not self.is_initialized: + response = RunSessionResponse() + response.exit_code = -1 + response.output = "Sandbox environment not initialized" + response.failure_reason = "ENVIRONMENT_NOT_INITIALIZED" + return response + + if session_name not in self.active_sessions: + if not self.create_managed_session(session_name): + response = RunSessionResponse() + response.exit_code = -1 + response.output = f"Failed to create session '{session_name}'" + response.failure_reason = "SESSION_CREATION_FAILED" + return response + + try: + self.active_sessions[session_name]["last_used"] = time.time() + + response = self.run_in_session( + command=command, + session=session_name, + max_retry=self.max_retry, + backoff=self.backoff + ) + + return response + + except Exception as e: + self.logger.error(f"Error executing command in session '{session_name}': {e}") + response = RunSessionResponse() + response.exit_code = -1 + response.output = f"Error executing command: {str(e)}" + response.failure_reason = "COMMAND_EXECUTION_ERROR" + return response + + def get_session_info(self, session_name: str) -> Optional[Dict]: + """ + Get information about a session. + """ + return self.active_sessions.get(session_name) + + def list_sessions(self) -> List[str]: + """ + List all active sessions. + """ + return list(self.active_sessions.keys()) + + def cleanup_session(self, session_name: str) -> bool: + """ + Clean up a specific session. + """ + try: + if session_name in self.active_sessions: + del self.active_sessions[session_name] + self.logger.debug(f"Session '{session_name}' cleaned up") + return True + + except Exception as e: + self.logger.error(f"Error cleaning up session '{session_name}': {e}") + return False + + def close(self): + """Close the sandbox environment and cleanup resources.""" + try: + self.stop_sandbox() + self.active_sessions.clear() + self.is_initialized = False + self.logger.debug("Sandbox environment closed successfully") + + except Exception as e: + self.logger.error(f"Error closing sandbox environment: {e}") + + def _install_agent(self, session_name: str) -> Tuple[bool, str]: + """Install and configure the agent in the sandbox""" + try: + response = self.run_in_session("mkdir -p /installed-agent", session_name) + if response.exit_code != 0: + print(response) + return False, "Failed to create installation directory" + + set_up_install_script = iflow_config.iflow_set_up_install_script_roll + # Qwen3-Coder是先随便设的一个值，用不到 + settings = iflow_config.get_iflow_setting_template(self.iflow_selected_auth_type, + self.iflow_api_key, + self.iflow_base_url, + "Qwen3-Coder", + self.iflow_search_api_key) + + is_success, message = self._upload_settings(set_up_install_script, + "/installed-agent", "install-agent.sh") + if not is_success: + return False, f"Failed to upload installation script: {message}" + + run_status, result = self.run_session_with_timeout( + self.agent_session_name, + "bash /installed-agent/install-agent.sh", + 300, + "install.txt" + ) + + if run_status != RunStatus.SUCCESS: + return False, f"Installation failed: {run_status}" + + settings = iflow_config.get_iflow_setting_template(self.iflow_selected_auth_type, + self.iflow_api_key, + self.iflow_base_url, + "Qwen3-Coder", + self.iflow_search_api_key) + if settings: + is_success, message = self._upload_settings(settings, "~/.iflow", "settings.json") + is_success, message = self._upload_settings(settings, "/root/.iflow", "settings.json") + if not is_success: + return False, f"Failed to upload settings: {message}" + + self.logger.debug("Agent installation completed successfully") + + version_commond = "iflow -v" + response = self.run_in_session(version_commond, self.agent_session_name) + if response.exit_code != 0: + return False, f"Failed to get agent version: {response.output}" + else: + self.logger.debug(f"!!!!!!!!!!!!!Agent version: {response.output}") + return True, "" + + except Exception as e: + error_msg = f"Error during agent installation: {e}" + self.logger.error(error_msg) + return False, error_msg + + def _upload_settings(self, content: str, directory: str, filename: str) -> tuple: + """Upload settings to the sandbox""" + + import tempfile + with tempfile.NamedTemporaryFile(mode='w', delete=False) as temp_file: + temp_file.write(content) + temp_file.write('\n') + temp_filename = temp_file.name + try: + is_success, message = self.upload_file(temp_filename, f"{directory}/{filename}") + return is_success, message + finally: + import os + os.unlink(temp_filename) + + def process_model_response(self, response: str, agent_timeout_sec: int, task_id: str, task_name: str = "") -> Tuple[str, float, bool, bool, Dict[str, Any]]: + """ + Process model response and return environment step results. + """ + self.logger.info(f"[MODEL_RESPONSE] START - Processing model response for task: {task_name}") + + self.failure_mode = FailureMode.NONE + self.error_messages.clear() + + reward = 0.0 + terminated = False + truncated = True + observation = "" + is_valid = True + + if "" in response: + response = response.split("")[-1].strip() + + if "" in response and "" in response: + try: + self.logger.info(f"[TOOL_PARSE] START - Parsing tool call from model response for task_name: {task_name}") + is_parsed, iflow_cmd = self.iflow_tool.execute_action(response, self.agent_session_name) + + if not is_parsed: + error_msg = "Could not parse tool call from response" + self.logger.error(f"[TOOL_PARSE] Failed! - {error_msg}") + self.failure_mode = FailureMode.TOOL_CALL_PARSE_FAILED + self.error_messages.append(f"Tool parse error: {error_msg}") + observation = f"Error: {error_msg}." + is_valid = False + else: + self.logger.info(f"[TOOL_PARSE] Success! - Tool call parsed successfully") + + self.logger.info(f"[TOOL_EXEC] START - Executing iflow command: {iflow_cmd[:100]}...") + run_status, result = self.run_session_with_timeout(self.agent_session_name, iflow_cmd, 360, "command.txt") + observation = "" + if run_status == RunStatus.SUCCESS: + try: + observation, is_valid = self.iflow_tool._process_tool_response(result) + if not is_valid: + error_msg = f"Tool response processing failed" + self.failure_mode = FailureMode.TOOL_RESPONSE_PROCESSING_FAILED + self.error_messages.append(f"Tool response processing error: {error_msg}") + self.logger.error(f"[TOOL_EXEC] Failed! - {error_msg}, observation: {json.dumps(observation, ensure_ascii=False)}") + else: + self.logger.info(f"[TOOL_EXEC] Success! - Tool execution completed successfully") + except Exception as process_exc: + error_msg = f"Exception during tool response processing: {str(process_exc)}" + self.failure_mode = FailureMode.TOOL_RESPONSE_PROCESSING_FAILED + self.error_messages.append(f"Tool response processing exception: {error_msg}") + self.logger.error(f"[TOOL_EXEC] Failed! - {error_msg}, observation so far: {json.dumps(observation, ensure_ascii=False)}") + observation = f"Error processing tool response: {str(process_exc)}" + is_valid = False + else: + if run_status == RunStatus.TIMEOUT: + error_msg = f"Tool execution timed out" + self.failure_mode = FailureMode.TOOL_EXECUTION_TIMEOUT + elif run_status == RunStatus.FAILED: + error_msg = f"Tool execution failed" + self.failure_mode = FailureMode.TOOL_EXECUTION_FAILED + elif run_status == RunStatus.EXCEPTION: + error_msg = f"Tool execution encountered exception" + self.failure_mode = FailureMode.TOOL_EXECUTION_EXCEPTION + else: + error_msg = f"Unknown tool execution status: {run_status}" + self.failure_mode = FailureMode.TOOL_EXECUTION_FAILED + self.error_messages.append(f"Unknown execution status: {error_msg}") + self.logger.error(f"[TOOL_EXEC] Failed! - {error_msg}") + self.error_messages.append(f"Tool execution failed: {error_msg}") + observation = f"Error: {error_msg}. Result: {result if result else 'No output'}" + is_valid = False + except Exception as e: + error_msg = f"Exception during model response processing: {str(e)}" + self.failure_mode = FailureMode.MODEL_RESPONSE_PROCESSING_EXCEPTION + self.error_messages.append(f"Model response processing exception: {error_msg}") + self.logger.error(f"[MODEL_RESPONSE] Failed! - {error_msg}") + observation = f"Error processing model response: {str(e)}" + is_valid = False + else: + terminated = True + truncated = False + self.logger.info(f"[MODEL_RESPONSE] No tool call found in response, treating as final answer") + + info = { + "action_is_valid": is_valid, + "success": is_valid + } + + self.logger.info(f"[MODEL_RESPONSE] Completed - terminated: {terminated}, valid: {is_valid}, failure_mode: {self.failure_mode}") + return observation, reward, terminated, truncated, info + + def _build_command(self, instruction: str) -> str: + """Build the command to run the agent""" + escaped_instruction = shlex.quote(instruction) + + if self.run_type == "iflow-cli": + return f"iflow -y -p {escaped_instruction}" + elif self.run_type == "qwen-code": + return f"qwen -y -p {escaped_instruction}" + elif self.run_type == "claude-code": + return f"claude --verbose --output-format stream-json -p {escaped_instruction}" + else: + return f"echo 'Unsupported run type: {self.run_type}'" + + def _extract_execution_info(self, text: str) -> tuple: + """Extract execution info from the agent output""" + if text is None: + return "", "" + + import re + pattern = r'(.*?)' + match = re.search(pattern, text, re.DOTALL) + + extracted_content = None + if match: + extracted_content = match.group(1).strip() + cleaned_text = re.sub(pattern, '', text, flags=re.DOTALL) + else: + cleaned_text = text + + return extracted_content, cleaned_text + + def run_tests(self, test_files: List[str], test_timeout_sec: int = 60, task_name: str = "") -> dict: + """Run tests for the task""" + self.logger.info(f"[TEST_SESSION] START - Creating test session: {self.test_session_name}") + test_output = "" + is_success = self.create_session(session=self.test_session_name) + if not is_success: + error_msg = "Failed to create test session" + self.logger.error(f"[TEST_SESSION] Failed! - {error_msg}") + self.failure_mode = FailureMode.TEST_SESSION_CREATION_FAILED + self.error_messages.append(f"Test session creation error: {error_msg}") + return False, FailureMode.TEST_SESSION_CREATION_FAILED, test_output + else: + self.logger.info(f"[TEST_SESSION] Success! - Test session created") + + test_dir = '/tests' + response = self.run_in_session(f"mkdir -p {test_dir}", self.test_session_name) + if response.exit_code != 0: + error_msg = f"Failed to create test directory: {response.output}" + self.logger.error(f"[TEST_SESSION] Failed! - {error_msg}") + self.failure_mode = FailureMode.TEST_DIRECTORY_CREATION_FAILED + self.error_messages.append(f"Test directory creation error: {error_msg}") + return False, FailureMode.TEST_DIRECTORY_CREATION_FAILED, test_output + else: + self.logger.info(f"[TEST_SESSION] Success! - Test directory created") + + for test_file_path in test_files: + test_path = Path(test_file_path) + if not test_path.exists(): + self.logger.warning(f"Test path not found: {test_file_path}") + continue + + if test_path.is_dir(): + tests_dir = f"{test_path}/{task_name}/tests" + if os.path.isdir(tests_dir): + for test_file_name in os.listdir(tests_dir): + test_file_path = os.path.join(tests_dir, test_file_name) + if os.path.exists(test_file_path): + is_success, message = self.upload_file(test_file_path, f"{test_dir}/{test_file_name}") + if not is_success: + error_msg = f"Failed to upload test file {test_file_name}: {message}" + self.logger.error(error_msg) + self.failure_mode = FailureMode.TEST_FILE_UPLOAD_FAILED + self.error_messages.append(f"Test file upload error: {error_msg}") + return False, FailureMode.TEST_FILE_UPLOAD_FAILED, test_output + self.logger.debug(f"Successfully uploaded test file: {test_file_name}") + else: + self.logger.warning(f"Tests path is not a directory: {tests_dir}") + run_tests_path = f"{test_path}/{task_name}/run-tests.sh" + if os.path.exists(run_tests_path): + is_success, message = self.upload_file(run_tests_path, f"{test_dir}/run-tests.sh") + if not is_success: + error_msg = f"Failed to upload test file run-tests.sh: {message}" + self.logger.error(error_msg) + self.failure_mode = FailureMode.TEST_FILE_UPLOAD_FAILED + self.error_messages.append(f"Test file upload error: {error_msg}") + return False, FailureMode.TEST_FILE_UPLOAD_FAILED, test_output + self.logger.debug(f"Successfully uploaded test file: {run_tests_path}") + else: + is_success, message = self.upload_file(str(test_path), f"{test_dir}/{test_path.name}") + if not is_success: + error_msg = f"Failed to upload test file: {message}" + self.logger.error(error_msg) + self.failure_mode = FailureMode.TEST_FILE_UPLOAD_FAILED + self.error_messages.append(f"Test file upload error: {error_msg}") + return False, FailureMode.TEST_FILE_UPLOAD_FAILED, test_output + self.logger.debug(f"Successfully uploaded test file: {test_path.name}") + self.logger.info(f"[TEST_SESSION] Success! - Test Files uploaded successfully") + + test_command = f"bash {test_dir}/run-tests.sh" + run_status, test_output = self.run_session_with_timeout( + self.test_session_name, + test_command, + 360, + "test.txt" + ) + self.logger.info(f"✅[RUN_TESTS] Completed - Test run status: {run_status}, test_output: {json.dumps(test_output, ensure_ascii=False)}...") + if run_status != RunStatus.SUCCESS: + if run_status == RunStatus.TIMEOUT: + error_msg = f"Test execution timed out" + self.failure_mode = FailureMode.TEST_TIMEOUT + self.error_messages.append(f"Test timeout error: {error_msg}") + return False, FailureMode.TEST_TIMEOUT, test_output + else: + error_msg = f"Test execution failed with status: {run_status}" + self.failure_mode = FailureMode.UNKNOWN_TEST_ERROR + self.error_messages.append(f"Test execution error: {error_msg}") + return False, FailureMode.UNKNOWN_TEST_ERROR, test_output + + with open("test_output.txt", "w") as f: + f.write(test_output) + is_resolved = self._parse_test_results(test_output) + + return is_resolved, "", test_output + + def _parse_test_results(self, test_output: str) -> bool: + """Parse test results to determine if the task is resolved""" + if "All tests passed" in test_output: + return True + if "PASS" in test_output and "FAIL" not in test_output: + return True + return False + + def get_messages(self, question: str): + """ + Get concatenated messages using iflow-cli framework. + """ + try: + escaped_question = shlex.quote(question) + if escaped_question.startswith("-"): + escaped_question = escaped_question[1:].strip() + command = f"iflow --sysinfo {escaped_question}" + + response = self.run_in_session(command, self.agent_session_name) + + if response.exit_code != 0: + error_msg = f"iflow --sysinfo command failed: {response.output}" + self.logger.error(error_msg) + self.failure_mode = FailureMode.IFLOW_SYSINFO_COMMAND_FAILED + self.error_messages.append(f"IFlow sysinfo command error: {error_msg}") + return [{"role": "user", "content": question}], response.output + + try: + response_output = response.output.split("[\r\n {\r\n \"role\": \"system")[-1] + response_output = "[\r\n {\r\n \"role\": \"system" + response_output + if isinstance(response_output, str): + try: + messages = json.loads(response_output) + except: + messages = ast.literal_eval(response_output) + else: + messages = response_output + self.logger.debug(f"iflow-cli returned {len(messages)} messages") + return messages, "" + except Exception as e: + error_msg = f"Failed to parse iflow sysinfo response: {str(e)}" + self.logger.error(f"Load messages failed: {response_output}") + self.failure_mode = FailureMode.IFLOW_SYSINFO_PARSE_FAILED + self.error_messages.append(f"IFlow sysinfo parse error: {error_msg}") + return [{"role": "user", "content": question}], f"IFLOW_SYSINFO_PARSE_FAILED, {str(e)}" + except Exception as e: + error_msg = f"Exception in get_messages: {str(e)}" + self.logger.error(f"Error in get_messages: {error_msg}") + self.failure_mode = FailureMode.IFLOW_SYSINFO_EXCEPTION + self.error_messages.append(f"IFlow sysinfo exception: {error_msg}") + return [{"role": "user", "content": question}], f"IFLOW_SYSINFO_EXCEPTION, {str(e)}" diff --git a/roll/pipeline/agentic/env/sandbox/sokoban_sandbox_env.py b/roll/pipeline/agentic/env/sandbox/sokoban_sandbox_env.py index 0b2d42680..ef8837d3c 100644 --- a/roll/pipeline/agentic/env/sandbox/sokoban_sandbox_env.py +++ b/roll/pipeline/agentic/env/sandbox/sokoban_sandbox_env.py @@ -9,9 +9,21 @@ import httpcore import httpx from gem import Env -from rock.actions import CreateBashSessionRequest -from rock.sdk.sandbox.client import Sandbox -from rock.sdk.sandbox.config import SandboxConfig + +try: + from rock.sdk.sandbox.client import Sandbox + from rock.sdk.sandbox.config import SandboxConfig +except ImportError: + print("ROCK SDK not available. Make sure it's installed.") + pass +try: + from rock.actions import CreateBashSessionRequest +except ImportError: + print("rl-rock is not available, try import from rock-rl(old-version)") + try: + from rock.sdk.sandbox.request import CreateBashSessionRequest + except ImportError: + print("rock-rl is still not available. Make sure it's installed. ") from roll.utils.logging import get_logger diff --git a/roll/pipeline/agentic/env/sokoban/tool_call_env.py b/roll/pipeline/agentic/env/sokoban/tool_call_env.py new file mode 100644 index 000000000..0055d74b7 --- /dev/null +++ b/roll/pipeline/agentic/env/sokoban/tool_call_env.py @@ -0,0 +1,494 @@ +""" +Tool call variant of SokobanNativeEnv for fork detection testing. + +Exposes Sokoban movements as OpenAI-compatible function tools instead of text-based actions. +Supports full_history mode: when enabled, invalid tool calls trigger MessageTracker fork +detection by removing the invalid assistant+tool message pair from history. +""" + +import json +from typing import Any, Dict, List, Optional, Tuple, SupportsFloat, Union + +from roll.pipeline.agentic.env.sokoban.native_env import SokobanNativeEnv +from roll.utils.logging import get_logger + + +class SokobanToolCallEnv(SokobanNativeEnv): + """ + Tool call variant of SokobanNativeEnv with fork detection support. + + This environment exposes Sokoban movements as function tools (OpenAI format) instead + of text-based tags. When full_history=True, invalid tool calls are removed + from message history to trigger MessageTracker fork detection. + + Args: + full_history: If True, remove invalid tool call rounds from history to trigger + MessageTracker fork detection. Default False. + **kwargs: Passed to SokobanNativeEnv parent class. + + Example: + >>> env = SokobanToolCallEnv(full_history=True, dim_room=(6, 6), num_boxes=1, max_steps=10) + >>> messages, info = env.reset(seed=42) + >>> tools = info['tools'] # OpenAI tool schema + >>> + >>> # Valid tool call + >>> action = { + ... "role": "assistant", + ... "tool_calls": [{ + ... "id": "call_1", + ... "type": "function", + ... "function": {"name": "move_player", "arguments": '{"direction": "right"}'} + ... }] + ... } + >>> messages, reward, done, trunc, info = env.step(action) + >>> # messages now includes assistant + tool result + >>> + >>> # Invalid tool call (direction not in VALID_DIRECTIONS) + >>> action = { + ... "role": "assistant", + ... "tool_calls": [{ + ... "id": "call_2", + ... "type": "function", + ... "function": {"name": "move_player", "arguments": '{"direction": "invalid"}'} + ... }] + ... } + >>> messages, reward, done, trunc, info = env.step(action) + >>> # With full_history=True, assistant+tool are removed, user warning added + >>> # This triggers fork detection in MessageTracker + """ + + def __init__( + self, + full_history: bool = False, + **kwargs + ): + """ + Initialize tool call environment. + + Args: + full_history: If True, remove invalid tool call rounds from history. + **kwargs: Passed to SokobanNativeEnv parent. + """ + # Use a dummy pattern that won't match (will be overridden anyway) + if 'action_pattern' not in kwargs: + kwargs['action_pattern'] = r'__TOOL_CALL_MODE__' + + env_instruction = ( + "You are solving the Sokoban puzzle. " + "You are the player and you need to push all boxes to targets. " + "When you are right next to a box, you can push it by moving in the same direction. " + "You cannot push a box through a wall, and you cannot pull a box. " + "Use the move_player tool to move the player." + ) + super().__init__(env_instruction=env_instruction, **kwargs) + + self.full_history = full_history + self.logger = get_logger() + + # Valid directions for parameter-level validation + self.VALID_DIRECTIONS = ["up", "down", "left", "right"] + + # Tool definitions (returned in reset()) + self._tools = self._create_tool_definitions() + + self.observation_suffix = "" + + def _create_tool_definitions(self) -> List[Dict]: + """ + Create OpenAI-compatible tool definitions for Sokoban movements. + + Returns: + List of tool definitions in OpenAI format. + """ + return [{ + "type": "function", + "function": { + "name": "move_player", + "description": "Move the player in the specified direction in the Sokoban game", + "parameters": { + "type": "object", + "properties": { + "direction": { + "type": "string", + "enum": ["up", "down", "left", "right"], + "description": "Direction to move: up, down, left, or right" + } + }, + "required": ["direction"] + } + } + }] + + def reset(self, seed=None) -> Tuple[List[Dict], Dict]: + """ + Reset the environment and return initial observation with tools. + + Args: + seed: Random seed for environment reset. + + Returns: + Tuple of (messages, info) where: + - messages: Initial conversation with system + user messages + - info: Dictionary containing tools and other metadata + """ + messages, info = super().reset(seed) + + # Override system message to instruct tool usage + system_content = ( + f"{self.system_template}\n\n" + f"{info.get('env_instruction', self.get_instructions())}\n\n" + "IMPORTANT: Use the move_player tool to take actions. " + "Do NOT use text-based tags. " + "Always call the tool with a valid direction parameter." + ) + messages[0]['content'] = system_content + + # Add tools to info dict for ProxyEnvManager + info['tools'] = self._tools + + return messages, info + + def step( + self, + action: Union[str, Dict] + ) -> Tuple[Union[List[Dict], str], SupportsFloat, bool, bool, Dict[str, Any]]: + """ + Execute a tool call action. + + Args: + action: Either: + - Dict with OpenAI message format: {"role": "assistant", "tool_calls": [...]} + - String (fallback for compatibility, will trigger format error) + + Returns: + Tuple of (messages, reward, terminated, truncated, info): + - messages: Updated conversation history + - reward: Step reward (may include format penalty) + - terminated: Whether episode is done (max steps) + - truncated: Whether episode was cut short + - info: Metrics and metadata + """ + self.current_step += 1 + + # Parse tool call from action + tool_call, parsed_action_text, is_valid_structure = self._parse_tool_call(action) + + if not is_valid_structure: + # Invalid structure: format penalty, no history modification + return self._handle_invalid_format(action) + + # Validate direction parameter (parameter-level validation only) + is_valid_direction = self._is_valid_direction_param(tool_call) + + # Add assistant message with tool call + assistant_msg = { + "role": "assistant", + "content": action.get("content", "") if isinstance(action, dict) else action, + "tool_calls": [tool_call] if tool_call else [] + } + self.message_history.append(assistant_msg) + + # Execute the action using parent's game logic (if direction valid) + if is_valid_direction: + message_history = self.message_history + self.message_history = [] + _, reward, terminated, truncated, info = super().step(parsed_action_text) + # Use action_is_effective (player position changed) rather than action_is_valid (format valid) + action_is_effective = info.get("metrics", {}).get("action_is_effective", False) + self.message_history = message_history + text_obs = self.render(mode="text") + else: + # Invalid direction parameter: apply format penalty, no state change + text_obs = self.render(mode="text") + reward = self.format_penalty + terminated = False + truncated = False + action_is_effective = False + info = self._create_invalid_info() + + # Add tool result message + # is_valid means: direction param valid AND action was effective (position changed) + tool_result_msg = { + "role": "tool", + "tool_call_id": tool_call["id"] if tool_call else "unknown", + "content": self._format_tool_result(text_obs, info, is_valid_direction and action_is_effective) + } + self.message_history.append(tool_result_msg) + + # FORK DETECTION TRIGGER: Remove invalid round if full_history enabled + if self.full_history and not is_valid_direction: + self.logger.info( + f"[SokobanToolCall] Removing invalid tool call round (step {self.current_step}) to trigger fork detection" + ) + # Remove last 2 messages (assistant + tool) + self.message_history = self.message_history[:-2] + + # Add warning message to inform LLM + warning_msg = { + "role": "user", + "content": ( + f"\n\n\n" + f"IMPORTANT: The last tool call was invalid (direction parameter must be one of up/down/left/right). " + f"That conversation round has been removed. Please retry with correct parameters.\n" + f"\n\n" + f"Current game state:\n{text_obs}\n\n" + f"{self.observation_suffix}" + ) + } + self.message_history.append(warning_msg) + else: + # Normal flow: add new observation + user_msg = { + "role": "user", + "content": f"Current game state:\n{text_obs}\n\n{self.observation_suffix}" + } + self.message_history.append(user_msg) + + # Update info with fork detection flag + if 'metrics' not in info: + info['metrics'] = {} + # Add metrics aggregation mode + if 'metrics_agg_mode' not in info: + info['metrics_agg_mode'] = {} + + return self.message_history, reward, terminated, truncated, info + + def _parse_tool_call( + self, + action: Union[str, Dict] + ) -> Tuple[Optional[Dict], str, bool]: + """ + Parse action into tool call dict and text for parent validation. + + Args: + action: Action from LLM (dict with tool_calls or string) + + Returns: + Tuple of (tool_call_dict, action_text_for_parent, is_valid_structure): + - tool_call_dict: Extracted tool call or None + - action_text_for_parent: Text format for parent's step() method + - is_valid_structure: Whether the tool call structure is valid + """ + # Case 1: Dict format (expected) + if isinstance(action, dict): + tool_calls = action.get("tool_calls", []) + if not tool_calls: + self.logger.warning("[SokobanToolCall] Received dict without tool_calls") + return None, "", False + + tool_call = tool_calls[0] # Take first tool call + + # Validate tool call structure + if not self._validate_tool_call_structure(tool_call): + return None, "", False + + # Extract direction parameter + func = tool_call.get("function", {}) + args_str = func.get("arguments", "{}") + + try: + args = json.loads(args_str) if isinstance(args_str, str) else args_str + direction = args.get("direction", "").lower() + except json.JSONDecodeError: + self.logger.error(f"[SokobanToolCall] JSON parse failed: {args_str}") + return tool_call, "", False + + # Map to parent's expected format + action_text = f"{direction.capitalize()}" + return tool_call, action_text, True + + # Case 2: String format (fallback for backward compatibility) + elif isinstance(action, str): + self.logger.warning( + f"[SokobanToolCall] Received string action (expected dict): {action[:100]}" + ) + return None, action, False + + return None, "", False + + def _validate_tool_call_structure(self, tool_call: Dict) -> bool: + """ + Validate tool call has required fields. + + Args: + tool_call: Tool call dict to validate + + Returns: + True if structure is valid, False otherwise + """ + if not isinstance(tool_call, dict): + return False + if "id" not in tool_call or "function" not in tool_call: + self.logger.warning("[SokobanToolCall] tool_call missing required fields (id or function)") + return False + func = tool_call["function"] + if "name" not in func: + self.logger.warning("[SokobanToolCall] function missing name field") + return False + if func["name"] != "move_player": + self.logger.warning(f"[SokobanToolCall] Unknown tool name: {func['name']}") + return False + return True + + def _is_valid_direction_param(self, tool_call: Optional[Dict]) -> bool: + """ + Check if direction parameter is valid (parameter-level validation only). + + This only checks if direction is in VALID_DIRECTIONS, not whether the move + is valid in the game state (e.g., hitting walls is OK at this level). + + Args: + tool_call: Tool call dict + + Returns: + True if direction parameter is valid, False otherwise + """ + if not tool_call: + return False + + func = tool_call.get("function", {}) + args_str = func.get("arguments", "{}") + + try: + args = json.loads(args_str) if isinstance(args_str, str) else args_str + direction = args.get("direction", "").lower() + is_valid = direction in self.VALID_DIRECTIONS + + if not is_valid: + self.logger.info( + f"[SokobanToolCall] Invalid direction parameter: {direction} " + f"(must be one of {self.VALID_DIRECTIONS})" + ) + + return is_valid + except (json.JSONDecodeError, AttributeError): + return False + + def _format_tool_result( + self, + observation: str, + info: Dict, + is_valid: bool + ) -> str: + """ + Format tool execution result message. + + Args: + observation: Current game state text + info: Step info from parent + is_valid: Whether the action was valid and effective + + Returns: + Formatted tool result message + """ + if is_valid: + action_desc = info.get("action_desc", "Action executed") + return f"Success: {action_desc}\n\nNew state:\n{observation}" + else: + # Invalid action: either bad parameters or ineffective move (e.g., hit wall) + action_desc = info.get("action_desc", "This move is invalid.") + return f"Error: {action_desc}\n\nCurrent state:\n{observation}" + + def _handle_invalid_format( + self, + action: Any + ) -> Tuple[List[Dict], float, bool, bool, Dict]: + """ + Handle structurally invalid actions (format penalty, no state change). + + Args: + action: The invalid action + + Returns: + Tuple of (messages, reward, terminated, truncated, info) + """ + self.message_history.append({ + "role": "assistant", + "content": str(action)[:200] # Truncate for safety + }) + + # Use parent's render to get current state + text_obs = self.render(mode="text") + + user_msg = { + "role": "user", + "content": ( + f"\n\n\n" + f"CRITICAL: You must use the move_player tool. " + f"Do not send text-based actions.\n" + f"\n\n" + f"Current state:\n{text_obs}\n\n" + f"{self.observation_suffix}" + ) + } + self.message_history.append(user_msg) + + info = self._create_invalid_info() + + return self.message_history, self.format_penalty, False, False, info + + def _create_invalid_info(self) -> Dict: + """ + Create info dict for invalid actions. + + Returns: + Info dict with invalid action metrics + """ + metrics = { + "action_is_valid": False, + "action_is_effective": False, + "success": False, + "format_penalty": self.format_penalty, + } + + metrics_agg_mode = { + "action_is_valid": "mean", + "action_is_effective": "mean", + "success": "last", + "format_penalty": "mean", + } + + return { + "metrics": metrics, + "metrics_agg_mode": metrics_agg_mode, + "action_desc": "Invalid action (parameter validation failed)", + } + + def parse_action(self, text: str) -> Dict: + """ + Override parent's parse_action to handle tool call format. + + This method is called by parent's step() when we pass action_text. + We use a special format that includes direction mapping. + + Args: + text: Action text in format "Direction" + + Returns: + Action info dict with action code + """ + # Extract direction from tags (our internal format) + import re + match = re.search(r'(.*?)', text, re.IGNORECASE) + if match: + direction = match.group(1).strip().lower() + # Map direction to action code + direction_map = { + "up": 1, + "down": 2, + "left": 3, + "right": 4 + } + action_code = direction_map.get(direction) + return { + "action": action_code, + "action_content": direction, + "think_content": "" + } + + # Invalid format + return { + "action": None, + "action_content": text[:50], + "think_content": "" + } diff --git a/roll/pipeline/agentic/env_manager/agent_native_env_manager.py b/roll/pipeline/agentic/env_manager/agent_native_env_manager.py index 5aff115ab..d987d1616 100644 --- a/roll/pipeline/agentic/env_manager/agent_native_env_manager.py +++ b/roll/pipeline/agentic/env_manager/agent_native_env_manager.py @@ -202,8 +202,7 @@ def format_messages(self, rollout_cache: RolloutCache) -> DataProto: prompt_ids = self.tokenizer.apply_chat_template(convert_list_content_str(messages, parse_tool_call_parameter_to_dict=self.pipeline_config.parse_tool_call_parameter_to_dict), tools=self.tools, - tokenize=True, add_generation_prompt=True, enable_thinking=False, - return_dict=False) + tokenize=True, add_generation_prompt=True, enable_thinking=False, return_dict=False) input_ids = torch.tensor(prompt_ids, dtype=torch.long).unsqueeze(0) attention_mask = torch.tensor([1] * input_ids.shape[1], dtype=torch.long).unsqueeze(0) # Huggingface Transformers prefer position_ids to be 0-based. @@ -519,4 +518,4 @@ def filter(self, group_id: int, episode_id: int, group: list[DataProto]): return False self.global_filter_stats["filtered"] += 1 - return True + return True \ No newline at end of file diff --git a/roll/pipeline/agentic/env_manager/base_env_manager.py b/roll/pipeline/agentic/env_manager/base_env_manager.py index 643a08e30..8403a827b 100644 --- a/roll/pipeline/agentic/env_manager/base_env_manager.py +++ b/roll/pipeline/agentic/env_manager/base_env_manager.py @@ -25,6 +25,7 @@ def __init__(self, *args, **kwargs): self.current_step = -1 self.running = False self.env: gem.Env + self.trace_meta = None @abstractmethod def run_rollout_loop(self, data: DataProto): @@ -57,8 +58,9 @@ def format_messages(self, history: List[Dict]) -> DataProto: def formulate_rollouts(self, rollout_cache: RolloutCache) -> DataProto: pass - def update_step(self, global_step): + def update_step(self, global_step, trace_meta=None): self.current_step = global_step + self.trace_meta = trace_meta def stop(self): self.running = False \ No newline at end of file diff --git a/roll/pipeline/agentic/env_manager/message_tracker.py b/roll/pipeline/agentic/env_manager/message_tracker.py new file mode 100644 index 000000000..86512959f --- /dev/null +++ b/roll/pipeline/agentic/env_manager/message_tracker.py @@ -0,0 +1,525 @@ +""" +Message tracking for agentic training with prefix aggregation. + +Manages incremental tokenization, hash-chain message deduplication, fork detection, +and traj-wise training data generation. +""" + +import json +import hashlib +import logging +import os +from dataclasses import dataclass +from typing import List, Dict, Optional, Any + +from pydantic import TypeAdapter +from transformers import PreTrainedTokenizer + +from sglang.srt.entrypoints.openai.protocol import ChatCompletionMessageParam + +from roll.pipeline.agentic.env_manager.token_mask_utils import convert_list_content_str + +logger = logging.getLogger(__file__) +logger.setLevel(os.getenv('ROLL_LOGGING_LEVEL', 'WARN')) + +DEFAULT_HASH = "default_hash" + +ChatCompletionMessageAdapter = TypeAdapter(ChatCompletionMessageParam) + +# Dummy prefix for single-message tokenization +BASE_CHAT_HISTORY = [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "I am a user."}, +] + + +@dataclass +class MessageItem: + """A single message with its tokenized representation.""" + index: int + message: dict + token_ids: Optional[List[int]] = None + logprobs: Optional[List[float]] = None + # Multimodal placeholders + image_data: Optional[Any] = None + video_data: Optional[Any] = None + audio_data: Optional[Any] = None + modalities: Optional[List[str]] = None + + @property + def role(self) -> str: + return self.message["role"] + + +# ========================= Message Unification ========================= + + +def unify_content(msg: dict) -> None: + """Normalize content to list-of-dicts format with merged text.""" + if msg["content"] is None or msg["content"] == "": + msg["content"] = [] + elif isinstance(msg["content"], str): + text = msg["content"] + msg["content"] = [{"type": "text", "text": text}] + + # Merge consecutive text items + content: List[dict] = [] + text = "" + has_text = False + for item in msg["content"]: + if item["type"] == "text": + text += item["text"] + has_text = True + else: + content.append(item) + if has_text: + content.append({"type": "text", "text": text}) + msg["content"] = content + + +def unify_tool_calls(msg: dict, tool_calls_dict: Dict[str, dict]) -> None: + """Replace tool_calls with canonical versions from tool_calls_dict.""" + if ( + msg["role"] == "assistant" + and "tool_calls" in msg + and isinstance(msg["tool_calls"], list) + ): + unified = [] + for item in msg["tool_calls"]: + canonical = tool_calls_dict[item["id"]] + unified.append(canonical) + msg["tool_calls"] = unified + + +def get_unified_message(msg: dict, tool_calls_dict: Optional[Dict[str, dict]]) -> dict: + """Validate and normalize a raw message dict via pydantic + unify.""" + chat_msg = ChatCompletionMessageAdapter.validate_python(msg).model_dump() + unify_content(chat_msg) + if chat_msg["role"] in ("system", "assistant", "tool"): + chat_msg["reasoning_content"] = ( + "" if chat_msg.get("reasoning_content") is None else chat_msg["reasoning_content"] + ) + chat_msg["tool_calls"] = ( + [] if chat_msg.get("tool_calls") is None else chat_msg["tool_calls"] + ) + if tool_calls_dict is not None: + unify_tool_calls(chat_msg, tool_calls_dict) + return chat_msg + + +# ========================= Hash Computation ========================= + + +def compute_message_hash(index: int, message: dict, pre_hash: str) -> str: + """SHA-256 hash of a unified message, chained to previous hash.""" + if message["role"] == "user": + data = { + "index": index, + "role": message["role"], + "content": message["content"], + "pre_hash": pre_hash, + } + elif message["role"] in ("system", "assistant", "tool"): + data = { + "index": index, + "role": message["role"], + "content": message["content"], + "reasoning_content": message.get("reasoning_content", ""), + "tool_calls": message.get("tool_calls", []), + "pre_hash": pre_hash, + } + else: + data = { + "index": index, + "role": message["role"], + "content": message["content"], + "pre_hash": pre_hash, + } + json_str = json.dumps(data, sort_keys=True, separators=(',', ':')) + return hashlib.sha256(json_str.encode('utf-8')).hexdigest() + + +# ========================= MessageTracker ========================= + + +class MessageTracker: + """Per-episode message tracker with incremental tokenization and fork detection.""" + + def __init__( + self, + tokenizer: PreTrainedTokenizer, + enable_thinking: bool = False, + enable_fork: bool = False, + ): + self.tokenizer = tokenizer + self.enable_thinking = enable_thinking + self.enable_fork = enable_fork + + # Hash-chain data structures + self.msg_item_dict: Dict[str, MessageItem] = {} + self.prev_hash_dict: Dict[str, str] = {} + self.tool_calls_dict: Dict[str, dict] = {} + self.new_line_token_id: int = tokenizer.encode("\n")[0] + + # Dummy prefix for single-message tokenization + self.base_chat_history: List[dict] = list(BASE_CHAT_HISTORY) + self.tools: Optional[List[dict]] = None + self.base_offset: Optional[int] = None + + # Fork tracking + self.step_response_hashes: List[str] = [] + self.last_recorded_msg_count: int = 0 + + # Cache current pre_hash after process_messages for record_response + self.current_pre_hash: str = DEFAULT_HASH + + # Per-step cache for step mode + self.pending_messages: Optional[List[dict]] = None + # Completed step data: list of {response_ids, logprobs, messages} + self.step_records: List[dict] = [] + + # -------------------- Tools / Base Offset -------------------- + + def compute_base_offset(self, tools: Optional[List[dict]]) -> int: + """Compute the token length of the dummy prefix.""" + ids = self.tokenizer.apply_chat_template( + self.base_chat_history, + tokenize=True, + add_generation_prompt=False, + tools=tools, + enable_thinking=self.enable_thinking, + return_dict=False + ) + return len(ids) + + def update_tools(self, tools: List[dict]) -> None: + """Update tools and recompute base_offset if tools changed.""" + if self.tools != tools: + self.tools = tools + self.base_offset = self.compute_base_offset(tools) + + def ensure_base_offset(self) -> None: + """Lazily compute base_offset on first use.""" + if self.base_offset is None: + self.base_offset = self.compute_base_offset(self.tools) + + # -------------------- Single Message Tokenization -------------------- + + def prepare_message_for_tokenize(self, msg: dict) -> dict: + """ + Prepare a unified message for tokenization. + - Restore content from list-of-dicts to str for jinja template compatibility. + - Convert tool_calls arguments from str to dict for proper jinja template rendering. + """ + return convert_list_content_str([msg], parse_tool_call_parameter_to_dict=True)[0] + + def tokenize_single_message( + self, + msg: dict, + add_generation_prompt: bool = False, + tools: Optional[List[dict]] = None, + ) -> List[int]: + """Tokenize a single message using dummy prefix + slice.""" + prepared_msg = self.prepare_message_for_tokenize(msg) + + if msg["role"] == "system": + # System messages: tokenize directly + return self.tokenizer.apply_chat_template( + [prepared_msg], + tokenize=True, + add_generation_prompt=False, + tools=tools, + enable_thinking=self.enable_thinking, + return_dict=False + ) + else: + # user/tool/assistant: dummy prefix + slice + self.ensure_base_offset() + full_ids = self.tokenizer.apply_chat_template( + [*self.base_chat_history, prepared_msg], + tokenize=True, + add_generation_prompt=add_generation_prompt, + tools=tools, + enable_thinking=self.enable_thinking, + return_dict=False + ) + return full_ids[self.base_offset:] + + # -------------------- Message Processing -------------------- + + def process_messages( + self, + messages: List[dict], + tools: Optional[List[dict]] = None, + add_generation_prompt: bool = True, + ) -> List[int]: + """Incrementally process messages, return full input_ids for inference.""" + # 1. Update tools and base_offset + if tools is not None: + self.update_tools(tools) + + # 2. Unify messages + unified_msgs = [ + get_unified_message(m, self.tool_calls_dict) for m in messages + ] + + # 3. Find already-recorded prefix + pre_hash = DEFAULT_HASH + pre_msg_length = 0 + for idx, msg in enumerate(unified_msgs): + cur_hash = compute_message_hash(idx, msg, pre_hash) + if cur_hash in self.msg_item_dict: + pre_msg_length += 1 + pre_hash = cur_hash + else: + break + + # 3.5. Fork detection + if self.last_recorded_msg_count > 0 and pre_msg_length < self.last_recorded_msg_count: + logger.warning( + f"History fork detected: matched {pre_msg_length} msgs, " + f"but last step had {self.last_recorded_msg_count} msgs" + ) + + # 4. Tokenize new messages + new_msgs = unified_msgs[pre_msg_length:] + for i, msg in enumerate(new_msgs): + msg_idx = pre_msg_length + i + is_last = (msg_idx == len(unified_msgs) - 1) + + if msg["role"] == "assistant": + logger.warning( + f"Assistant message at index {msg_idx} being tokenized " + f"(normally should come from record_response)" + ) + + token_ids = self.tokenize_single_message( + msg, + add_generation_prompt=(add_generation_prompt and is_last), + tools=self.tools, + ) + + item = MessageItem(index=msg_idx, message=msg, token_ids=token_ids) + cur_hash = compute_message_hash(msg_idx, msg, pre_hash) + self.msg_item_dict[cur_hash] = item + self.prev_hash_dict[cur_hash] = pre_hash + pre_hash = cur_hash + + # Cache pre_hash for record_response + self.current_pre_hash = pre_hash + + # 5. Build full input_ids + input_ids = self.build_input_ids(unified_msgs) + + # Cache messages for step mode pairing with record_response + self.pending_messages = messages + + return input_ids + + def build_input_ids(self, unified_msgs: List[dict]) -> List[int]: + """Concatenate all message token_ids with newline insertion after assistant.""" + input_ids: List[int] = [] + pre_hash = DEFAULT_HASH + for idx, msg in enumerate(unified_msgs): + cur_hash = compute_message_hash(idx, msg, pre_hash) + item = self.msg_item_dict[cur_hash] + if idx >= 1 and unified_msgs[idx - 1]["role"] == "assistant": + input_ids.append(self.new_line_token_id) + input_ids.extend(item.token_ids) + pre_hash = cur_hash + return input_ids + + # -------------------- Response Recording -------------------- + + def record_response( + self, + msg_pos: int, + resp_msg: dict, + response_ids: List[int], + logprobs: Optional[List[float]] = None, + ) -> None: + """Record an LLM-generated assistant response with its token_ids from inference.""" + # Record tool_calls for future unification + if "tool_calls" in resp_msg and isinstance(resp_msg["tool_calls"], list): + for tc_item in resp_msg["tool_calls"]: + if "id" in tc_item: + self.tool_calls_dict[tc_item["id"]] = tc_item + + unified_msg = get_unified_message(resp_msg, self.tool_calls_dict) + item = MessageItem( + index=msg_pos, + message=unified_msg, + token_ids=response_ids, + logprobs=logprobs, + ) + + cur_hash = compute_message_hash(msg_pos, unified_msg, self.current_pre_hash) + self.msg_item_dict[cur_hash] = item + self.prev_hash_dict[cur_hash] = self.current_pre_hash + + # Update fork tracking + self.step_response_hashes.append(cur_hash) + self.last_recorded_msg_count = msg_pos + 1 + self.current_pre_hash = cur_hash + + # Save step record for step mode extraction + if self.pending_messages is not None: + self.step_records.append({ + "response_ids": response_ids, + "logprobs": logprobs, + "messages": self.pending_messages, + }) + self.pending_messages = None + + # -------------------- Trajectory Data Extraction -------------------- + + def get_trajectory_data(self, mode: str = "traj") -> List[dict]: + """ + Extract training data as a list of sample dicts. + + Args: + mode: "traj" returns aggregated trajectory samples (one per branch); + "step" returns one sample per assistant response step. + + Each sample dict contains: + - token_ids: List[int] + - response_masks: List[int] (1 for model-generated assistant tokens) + - logprobs: List[float] + - messages: List[dict] + """ + if mode == "step": + return self._get_step_samples() + + branches = self.identify_branches() + if not branches: + return [] + + if not self.enable_fork: + # Use the longest branch only + longest = max(branches, key=len) + branches = [longest] + + results = [] + for branch in branches: + leaf_hash = branch[-1] + path = self.trace_path_to_root(leaf_hash) + sample = self.build_sample_from_path(path) + results.append(sample) + return results + + def _get_step_samples(self) -> List[dict]: + """Extract per-step training samples from cached step records. + + For each step, re-tokenizes the history messages via apply_chat_template + to get prompt_ids, and uses response_ids from the inference engine. + """ + results: List[dict] = [] + for record in self.step_records: + messages = record.get("messages", []) + response_ids = record["response_ids"] + resp_logprobs = record["logprobs"] + + # Re-tokenize history messages as a whole prompt + prompt_ids = self.tokenizer.apply_chat_template( + convert_list_content_str(messages, parse_tool_call_parameter_to_dict=True), + tools=self.tools, + tokenize=True, + add_generation_prompt=True, + enable_thinking=self.enable_thinking, + return_dict=False + ) + + token_ids = prompt_ids + response_ids + response_masks = [0] * len(prompt_ids) + [1] * len(response_ids) + logprobs = [0.0] * len(prompt_ids) + (resp_logprobs if resp_logprobs else [0.0] * len(response_ids)) + + results.append({ + "token_ids": token_ids, + "response_masks": response_masks, + "logprobs": logprobs, + "messages": messages, + }) + return results + + def identify_branches(self) -> List[List[str]]: + """ + Group step_response_hashes into branches. + Consecutive steps belong to the same branch iff the previous step's + response_hash is an ancestor of the current step's hash chain. + """ + if not self.step_response_hashes: + return [] + + branches: List[List[str]] = [[self.step_response_hashes[0]]] + + for i in range(1, len(self.step_response_hashes)): + prev_hash = self.step_response_hashes[i - 1] + curr_hash = self.step_response_hashes[i] + + if self.is_ancestor(prev_hash, curr_hash): + branches[-1].append(curr_hash) + else: + branches.append([curr_hash]) + + return branches + + def is_ancestor(self, ancestor_hash: str, descendant_hash: str) -> bool: + """Check if ancestor_hash is on the path from descendant_hash back to root.""" + h = descendant_hash + while h != DEFAULT_HASH: + if h == ancestor_hash: + return True + h = self.prev_hash_dict.get(h, DEFAULT_HASH) + return False + + def trace_path_to_root(self, leaf_hash: str) -> List[str]: + """Walk from leaf to root, return ordered hash list (root -> leaf).""" + path: List[str] = [] + h = leaf_hash + while h != DEFAULT_HASH: + path.append(h) + h = self.prev_hash_dict[h] + path.reverse() + return path + + def build_sample_from_path(self, path: List[str]) -> dict: + """Build a training sample from an ordered hash path.""" + token_ids: List[int] = [] + response_masks: List[int] = [] + logprobs: List[float] = [] + messages: List[dict] = [] + + prev_role: Optional[str] = None + for hash_key in path: + item = self.msg_item_dict[hash_key] + + # Insert newline after assistant + if prev_role == "assistant": + token_ids.append(self.new_line_token_id) + response_masks.append(0) + logprobs.append(0.0) + + token_ids.extend(item.token_ids) + + # response_mask: 1 only for model-generated assistant (has logprobs) + if item.role == "assistant" and item.logprobs is not None: + response_masks.extend([1] * len(item.token_ids)) + else: + response_masks.extend([0] * len(item.token_ids)) + + if item.logprobs is not None: + logprobs.extend(item.logprobs) + else: + logprobs.extend([0.0] * len(item.token_ids)) + + # Collect messages + messages.append(item.message) + + prev_role = item.role + + return { + "token_ids": token_ids, + "response_masks": response_masks, + "logprobs": logprobs, + "messages": messages, + } diff --git a/roll/pipeline/agentic/env_manager/proxy_env_manager.py b/roll/pipeline/agentic/env_manager/proxy_env_manager.py new file mode 100644 index 000000000..d70457434 --- /dev/null +++ b/roll/pipeline/agentic/env_manager/proxy_env_manager.py @@ -0,0 +1,563 @@ +import asyncio +from datetime import datetime +from contextlib import nullcontext +import json +import re +from threading import Lock +from typing import Any, Dict, List, Optional +import time +import uuid + +import logging +import numpy as np +import ray +import torch +from omegaconf import DictConfig +from tensordict import List, TensorDict +from transformers import PreTrainedTokenizer +from fastapi import Request +from fastapi.responses import JSONResponse + +from roll.pipeline.agentic.llm_proxy import create_llm_proxy, BaseLLMProxy +from roll.distributed.scheduler.protocol import DataProto +from roll.pipeline.agentic.env_manager.base_env_manager import BaseEnvManager +from roll.pipeline.agentic.env_manager.message_tracker import MessageTracker +from roll.distributed.scheduler.rollout_scheduler import GroupQueueManager +from roll.distributed.scheduler.router import RouterManager +from roll.pipeline.agentic.agentic_config import EnvManagerConfig, AgenticConfig +from roll.utils.functionals import pad_to_length +from roll.utils.logging import get_logger + +from roll.pipeline.agentic.agent_runner.base import AgentRunner +from roll.utils.import_utils import safe_import_class + +from sglang.srt.function_call.function_call_parser import FunctionCallParser +from sglang.srt.entrypoints.openai.protocol import Tool + + +logging.getLogger("httpx").setLevel(logging.WARNING) +logging.getLogger("httpcore").setLevel(logging.WARNING) +logging.getLogger("rock").setLevel(logging.WARNING) + + +class ProxyEnvManager(BaseEnvManager): + def __init__(self, + worker_config: EnvManagerConfig, + pipeline_config: AgenticConfig, + env_config: DictConfig, + tokenizer: PreTrainedTokenizer, + generate_scheduler, + output_queue: GroupQueueManager, + thread_lock: Lock, + mode='train', + *args, **kwargs): + super().__init__() + self.logger = get_logger() + self.worker_config: EnvManagerConfig = worker_config + self.pipeline_config = pipeline_config + self.env_config: DictConfig = env_config + self.tokenizer: PreTrainedTokenizer = tokenizer + self.output_queue = output_queue + self.mode = mode + self.generate_scheduler: RouterManager = generate_scheduler + + # EnvManager states + self.rollout_cache = None + self.group_seed = None + self.episode_id = None + self.running = False + self.use_thread_lock = self.env_config.get("use_thread_lock", False) # 避免同时执行大量cpu操作, 可以通过env_config配置 + self.thread_lock = thread_lock if self.use_thread_lock else nullcontext() + + self.llm_proxy: BaseLLMProxy = create_llm_proxy( + generate_scheduler=self.generate_scheduler, + llm_proxy_config=self.worker_config.llm_proxy, + tokenizer=self.tokenizer, + env=None + ) + self.tool_call_parser = self.env_config.config.get("tool_call_parser", "qwen25") + + runner_cls_path = self.env_config["agent_runner_cls"] + runner_cls = safe_import_class(runner_cls_path) + if runner_cls is None: + raise ImportError(f"Failed to import agent_runner_cls: {runner_cls_path}") + base_url = f"http://127.0.0.1:{self.env_config.get('proxy_port', 8000)}" + self.agent_runner: AgentRunner = runner_cls( + base_url=base_url, + env_id=self.env_config["env_id"], + env_config=self.env_config, + worker_config=self.worker_config, + ) + self.reward_granularity: str = self.env_config.config.get("reward_granularity", "pass_rate") # "pass_rate" | "binary" + self.history = [] + self.trajectory_mode: str = self.env_config.config.get("trajectory_mode", "traj") # "step" | "traj" + self.message_tracker: Optional[MessageTracker] = None + self.reset_stats() + + def run_rollout_loop(self, data: DataProto): + """ + Continuously play episodes until data collection is complete. + + Seed update logic: + group_seed = base_seed + group_id + episode_seed = group_seed + episode_id + + trajectory_id: f"{group_id}_{episode_id}_{episode_seed}" + """ + assert "seed" in data.meta_info + self.running = True + self.group_seed = data.meta_info['seed'] + self.env_config['group_seed'] + + while True: + self.episode_id = ray.get(self.output_queue.get_episode_id.remote(self.env_config['group_id'], self.env_config['env_id'])) + start_step = self.current_step + + if self.episode_id is None: + break + + self.history = [] + self.message_tracker = MessageTracker( + self.tokenizer, + enable_thinking=self.env_config.config.get("enable_thinking", False), + enable_fork=self.env_config.config.get("enable_fork", False), + ) + self.reset_stats() + + seed = self.group_seed + self.episode_id + + episode_result = self.agent_runner.run_job(seed) + if episode_result.status == "NoData": + break + result = episode_result.to_dict() + + self.running = False + rollout = self.formulate_rollouts(result) + + ray.get(self.output_queue.put.remote(self.env_config['group_id'], self.episode_id, start_step, rollout, self.env_config['env_id'])) + self.running = True + ray.get(self.output_queue.put.remote(self.env_config['group_id'], self.episode_id, start_step, None, self.env_config['env_id'])) + + def reset_stats(self): + """Reset per-episode statistics before each new episode.""" + self.traj_start_time = time.time() + self.last_response_finish_time = None + self.tools = None + self.log_stats = { + "step_rt": [], # total round-trip time per step (request in → response out) + "pure_infer_time": [], # pure LLM generation time + "env_exec_time": [], # sandbox execution time (last response end → next request start) + "proxy_overhead": [], # proxy logic overhead (tokenizer, parser, etc.) + "response_length": [], + "tokens_per_second": [], + "current_step": [], + } + self.message_tracker = MessageTracker( + self.tokenizer, + enable_thinking=self.env_config.config.get("enable_thinking", False), + enable_fork=self.env_config.config.get("enable_fork", False), + ) + + async def process_request(self, request: Request): + """ + Handle HTTP requests from the Harbor agent via ProxyServer (Push mode). + + Delegates inference to _process_request_dict and wraps the result in + a JSONResponse. Sequence-length errors surface as 400; other errors as 500. + + TODO: TITO, prefix agg + """ + if not self.running: + self.logger.warning("Dropped a late request because the episode is finalizing.") + return JSONResponse(status_code=410, content={"error": "Episode finished"}) + + try: + body = await request.json() + response_dict = await self._process_request_dict(body) + return JSONResponse(content=response_dict) + except ValueError as e: + return JSONResponse(status_code=400, content={"error": str(e), "stop_reason": "max_length"}) + except Exception as e: + self.logger.error(f"Inference process error: {e}", exc_info=True) + return JSONResponse(status_code=500, content={"error": str(e)}) + + async def _process_request_dict(self, body: Dict) -> Dict: + """ + Core inference handler shared by Push (HTTP) and Pull (ModelService) modes. + + Args: + body: Parsed JSON dict with messages, tools, tool_choice fields. + + Returns: + OpenAI-compatible chat completion dict. + + Raises: + ValueError: If the prompt exceeds sequence_length. + """ + self.logger.info(f"Received request from Harbor Agent for env_id: {self.env_config['env_id']}") + step_start_time = time.time() + + if self.last_response_finish_time: + env_work_time = step_start_time - self.last_response_finish_time + self.log_stats["env_exec_time"].append(env_work_time) + + messages = body.get('messages', []) + raw_tools = body.get('tools', []) + # Use "none" to get raw text output instead of parsed tool_calls + tool_choice = body.get('tool_choice', 'auto') + + if raw_tools: + self.tools = raw_tools + + tools_obj = [Tool.model_validate(t) for t in raw_tools] if raw_tools else [] + + prompt_ids = self.message_tracker.process_messages( + messages, tools=self.tools, add_generation_prompt=True, + ) + input_ids = torch.tensor(prompt_ids, dtype=torch.long).unsqueeze(0) + attention_mask = torch.tensor([1] * input_ids.shape[1], dtype=torch.long).unsqueeze(0) + # Huggingface Transformers prefer position_ids to be 0-based. + position_ids = attention_mask.cumsum(dim=-1) - 1 + + lm_input = DataProto() + lm_input.batch = TensorDict({ + "input_ids": input_ids, + "attention_mask": attention_mask, + "position_ids": position_ids, + }, batch_size=input_ids.shape[0]) + lm_input.meta_info["tools"] = self.tools + lm_input.meta_info["tool_choice"] = tool_choice + + if input_ids.shape[1] >= self.pipeline_config.sequence_length: + self.logger.warning( + f"sequence_length={self.pipeline_config.sequence_length} " + f"input_ids length={input_ids.shape[1]}, maybe you should increase the response_length" + ) + raise ValueError("Sequence length exceeded") + + max_new_tokens = min( + self.env_config["max_tokens_per_step"], + self.worker_config.generating_args.max_new_tokens, + self.pipeline_config.sequence_length - input_ids.shape[1], + ) + generation_config = self.worker_config.generating_args.to_dict() + generation_config["max_new_tokens"] = min(max_new_tokens, self.pipeline_config.sequence_length) + generation_config["enable_thinking"] = self.env_config.config.get("enable_thinking", False) + lm_input.meta_info["src_rank"] = self.env_config["env_id"] + + infer_start = time.time() + lm_output = await asyncio.to_thread( + self.llm_proxy.generate, + messages=messages, + lm_input=lm_input, + generation_config=generation_config, + ) + infer_end = time.time() + pure_infer_duration = infer_end - infer_start + + if lm_output is None: + self.logger.error("LLM Proxy returned None (Inference Aborted), retrying...") + lm_input.meta_info["src_rank"] = self.env_config["env_id"] + lm_output = await asyncio.to_thread( + self.llm_proxy.generate, + messages=messages, + lm_input=lm_input, + generation_config=generation_config, + ) + + if lm_output is None: + raise RuntimeError("LLM Proxy failed after retry — check model name and API key in llm_proxy config") + + response_ids = lm_output.batch['responses'][0].tolist() + resp_len = len(response_ids) + + total_step_rt = time.time() - step_start_time + overhead = total_step_rt - pure_infer_duration + + self.log_stats["step_rt"].append(total_step_rt) + self.log_stats["pure_infer_time"].append(pure_infer_duration) + self.log_stats["proxy_overhead"].append(overhead) + self.log_stats["response_length"].append(resp_len) + self.log_stats["current_step"].append(len(self.history)) + if total_step_rt > 0.01: + self.log_stats["tokens_per_second"].append(resp_len / total_step_rt) + + infer_logprobs = None + if "infer_logprobs" in lm_output.batch.keys(): + infer_logprobs = lm_output.batch['infer_logprobs'][0][-resp_len:].tolist() + + text = self.tokenizer.decode(response_ids, skip_special_tokens=True) + finish_reason = "stop" + tool_calls = None + + if tool_choice != "none" and tools_obj: + # qwen25 parser requires \n{...}\n (with newlines). + # Some external models emit the no-newline variant; normalize before parsing. + if "" in text: + text = re.sub(r"(?!\n)", "\n", text) + text = re.sub(r"(?", "\n", text) + tool_calls, text, finish_reason = self.process_tool_calls(text, tools_obj, finish_reason) + + message_out = {"role": "assistant", "content": text} + if tool_calls: + message_out["tool_calls"] = tool_calls + + # Record assistant response in MessageTracker (unified for both step and traj modes) + self.message_tracker.record_response( + msg_pos=len(messages), + resp_msg=message_out, + response_ids=response_ids, + logprobs=infer_logprobs, + ) + + with self.thread_lock: + self.history.append({ + "prompt_messages": messages, + "prompt_ids": prompt_ids, + "response_ids": response_ids, + "response_message": message_out, + "tools": self.tools, + "logprobs": infer_logprobs, + "finish_reason": finish_reason, + "timestamp": time.time(), + }) + + self.last_response_finish_time = time.time() + + return { + "id": f"chatcmpl-{uuid.uuid4().hex[:12]}", + "object": "chat.completion", + "created": int(time.time()), + "choices": [{ + "index": 0, + "message": message_out, + "finish_reason": finish_reason, + }], + } + + + def process_tool_calls( + self, + text: str, + tools: List[Any], + finish_reason: str, + ) -> tuple[Optional[List[Dict]], str, str]: + """Parse tool calls from model output using sglang FunctionCallParser.""" + parser = FunctionCallParser(tools, self.tool_call_parser) + + if parser.has_tool_call(text): + try: + text, call_info_list = parser.parse_non_stream(text) + tool_calls = [] + for call_info in call_info_list: + tool_id = f"call_{uuid.uuid4().hex[:24]}" + args = call_info.parameters + if isinstance(args, dict): + args = json.dumps(args) + tool_calls.append({ + "id": tool_id, + "type": "function", + "function": { + "name": call_info.name, + "arguments": args if args else "{}", + }, + }) + return tool_calls, text, "tool_calls" + except Exception as e: + self.logger.error(f"Tool call parsing error: {e}") + return None, text, finish_reason + + return None, text, finish_reason + + + def _select_episode_score(self, harbor_result: Dict) -> float: + """Pick the training reward from the runner result based on ``reward_granularity``. + + ``reward_granularity`` names which result field to use as the reward (the runner only + reports facts; this is the single decision point): + - "pass_rate": the continuous FAIL_TO_PASS rate, added by the runner when the evaluation + report.json is available; falls back to the 0/1 ``score`` when it is missing. + - "binary" (or any value without a matching field): falls back to the 0/1 ``score``, + which is the binary verifier reward itself. + """ + return float(harbor_result.get(self.reward_granularity, harbor_result.get("score", 0.0))) + + def formulate_rollouts(self, harbor_result: Dict) -> DataProto: + """ + Convert Harbor execution results to training data format. + Uses MessageTracker.get_trajectory_data(mode) to unify step/traj modes. + """ + tag = self.env_config.get('tag', 'proxy') + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f") + traj_group_id = f"{tag}_{self.env_config['group_id']}_{self.episode_id}_{self.group_seed}" + traj_id = f"{traj_group_id}_{self.env_config['env_id']}_{timestamp}" + + episode_score = self._select_episode_score(harbor_result) + if harbor_result.get("agent_exit_reason"): + self.logger.warning(f"[Harbor] agent_exit_reason: {harbor_result['agent_exit_reason']}") + trajectory_samples = self.message_tracker.get_trajectory_data(mode=self.trajectory_mode) + + if not trajectory_samples: + self.logger.info(f"[PLACEHOLDER_ROLLOUT] Creating for episode_id: {self.episode_id}") + trajectory_samples = [{ + "token_ids": [self.tokenizer.pad_token_id], + "response_masks": [0], + "logprobs": [0.0], + "messages": [], + }] + episode_score = 0.0 + + # Step mode: reward only on last step; traj mode: episode_score on each branch + is_step_mode = self.trajectory_mode == "step" + step_rewards: List[float] = [] + if is_step_mode: + step_rewards = [0.0] * (len(trajectory_samples) - 1) + [episode_score] + + samples: List[DataProto] = [] + seq_len = self.pipeline_config.sequence_length + + logic_pct = float(harbor_result.get('agent_execution_time_ratio', 0.0)) + env_setup_pct = float(harbor_result.get('env_setup_time_ratio', 0.0)) + agent_setup_pct = float(harbor_result.get('agent_setup_time_ratio', 0.0)) + total_sandbox_duration = float(harbor_result.get('time_total_sec', 0.0)) + + for idx, traj_data in enumerate(trajectory_samples): + token_ids_list = traj_data["token_ids"] + response_masks_list = traj_data["response_masks"] + logprobs_list = traj_data["logprobs"] + messages = traj_data.get("messages", []) + + reward = step_rewards[idx] if is_step_mode else episode_score + + # Build tensors + input_ids = torch.tensor(token_ids_list, dtype=torch.long).unsqueeze(0) + attention_mask = torch.ones(1, len(token_ids_list), dtype=torch.long) + response_mask = torch.tensor(response_masks_list, dtype=torch.bool).unsqueeze(0) + + first_resp_idx = ( + response_masks_list.index(1) if 1 in response_masks_list else len(response_masks_list) + ) + prompt_masks = [1] * first_resp_idx + [0] * (len(token_ids_list) - first_resp_idx) + prompt_mask = torch.tensor(prompt_masks, dtype=torch.bool).unsqueeze(0) + + score_tensor = torch.zeros(1, len(token_ids_list), dtype=torch.float) + score_tensor[0, -1] = reward + + position_ids = attention_mask.cumsum(dim=-1) - 1 + + # Padding + input_ids = pad_to_length(input_ids, length=seq_len, pad_value=self.tokenizer.pad_token_id) + attention_mask = pad_to_length(attention_mask, length=seq_len, pad_value=0) + position_ids = pad_to_length(position_ids, length=seq_len, pad_value=0) + response_mask = pad_to_length(response_mask, length=seq_len, pad_value=0) + prompt_mask = pad_to_length(prompt_mask, length=seq_len, pad_value=0) + score_tensor = pad_to_length(score_tensor, length=seq_len, pad_value=0) + + + sample_traj_id = f"{traj_id}_b{idx}" if len(trajectory_samples) > 1 else traj_id + state_hash = f"step_{idx}" + trajectory_data_dict = { + "trajectory_id": sample_traj_id, + "timestamp": timestamp, + "num_actions": len(self.history), + "branch_or_step_idx": idx, + "num_branches_or_steps": len(trajectory_samples), + "reward_info": {"episode_reward": episode_score, "step_reward": reward}, + "messages": messages, + "finish_reason": harbor_result.get('finish_reason', 'unknown'), + } + messages_json = json.dumps(messages) + + lm_input = DataProto( + batch=TensorDict({ + "input_ids": input_ids, + "attention_mask": attention_mask, + "position_ids": position_ids, + "response_mask": response_mask, + "prompt_mask": prompt_mask, + "scores": score_tensor, + }, batch_size=input_ids.shape[0]), + non_tensor_batch={ + "env_ids": np.array([self.env_config['env_id']], dtype=object), + "group_ids": np.array([self.env_config['group_id']], dtype=object), + "traj_group_id": np.array([traj_group_id], dtype=object), + "traj_id": np.array([sample_traj_id], dtype=object), + "tags": np.array([tag], dtype=object), + "step_scores": np.array([reward], dtype=object), + "episode_scores": np.array([episode_score], dtype=object), + "state_hash": np.array([state_hash], dtype=object), + "step": np.array([idx], dtype=object), + "trajectory_data": np.array([json.dumps(trajectory_data_dict)], dtype=object), + "messages": np.array([messages_json], dtype=object), + "tools": np.array([json.dumps(self.tools)], dtype=object), + "exp_name": np.array([self.pipeline_config.exp_name], dtype=object), + }, + ) + + # Logprobs + logprobs_tensor = torch.tensor(logprobs_list, dtype=torch.float).unsqueeze(0) + logprobs_tensor = pad_to_length(logprobs_tensor, length=seq_len, pad_value=0) + lm_input.batch["infer_logprobs"] = logprobs_tensor[:, 1:] + + samples.append(lm_input) + + batch = DataProto.concat(samples) + + # Metrics + total_resp = sum(sum(s["response_masks"]) for s in trajectory_samples) + avg_resp_len = float(total_resp / len(trajectory_samples)) + + # 性能分位数统计 (P99 Analysis) + def get_stats(data_list): + if not data_list: return 0.0, 0.0, 0.0 + return float(np.mean(data_list)), float(np.percentile(data_list, 99)), float(np.max(data_list)) + + avg_rt, p99_rt, max_rt = get_stats(self.log_stats["step_rt"]) + inf_avg, inf_p99, inf_max = get_stats(self.log_stats["pure_infer_time"]) + env_avg, env_p99, env_max = get_stats(self.log_stats["env_exec_time"]) + ovh_avg, ovh_p99, ovh_max = get_stats(self.log_stats["proxy_overhead"]) + + env_metric: Dict[str, float] = { + "num_actions": len(self.history), + "response_length": avg_resp_len, + "episode_score": episode_score, + "traj_time_total": round(float(time.time() - self.traj_start_time), 4), + # --- 宏观利用率 --- + "efficiency/logic_pct": logic_pct, # 推理+沙箱执行时间占比 + "efficiency/env_setup_pct": env_setup_pct, # Harbor基础设施准备时间占比 + "efficiency/agent_setup_pct": agent_setup_pct, # Agent内部依赖安装时间占比 + # --- 微观性能指标 --- + "latency/pure_inf_avg": round(inf_avg, 4), + "latency/pure_inf_p99": round(inf_p99, 4), + "latency/pure_inf_max": round(inf_max, 4), + "latency/env_exec_avg": round(env_avg, 4), + "latency/env_exec_p99": round(env_p99, 4), + "latency/env_exec_max": round(env_max, 4), + "latency/proxy_ovh_avg": round(ovh_avg, 4), + "latency/proxy_ovh_p99": round(ovh_p99, 4), + "latency/proxy_ovh_max": round(ovh_max, 4), + "latency/step_rt_avg": round(avg_rt, 4), + "latency/step_rt_p99": round(p99_rt, 4), + "latency/step_rt_max": round(max_rt, 4), + "throughput/step_per_sec": round(len(trajectory_samples) / (time.time() - self.traj_start_time), 4) + } + prompt_lengths = [len(s["token_ids"]) - sum(s["response_masks"]) for s in trajectory_samples] + resp_lengths = [sum(s["response_masks"]) for s in trajectory_samples] + total_tokens = sum(len(s["token_ids"]) for s in trajectory_samples) + env_metric["prompt_length"] = float(np.mean(prompt_lengths)) if prompt_lengths else 0 + env_metric["response_length"] = float(np.mean(resp_lengths)) if resp_lengths else 0 + env_metric["total_tokens"] = float(total_tokens) + env_metric["num_branches_or_steps"] = len(trajectory_samples) + batch.meta_info = { + "metrics": {f"env/{tag}/{k}": v for k, v in env_metric.items()} + } + batch.meta_info["metrics"]["env/response_length"] = avg_resp_len + batch.meta_info["COLUMMNS_CONFIG"] = [ + ["trajectory_data", "string"], + ["messages", "string"], + ["tools", "string"], + ["exp_name", "string"], + ] + + return batch + + diff --git a/roll/pipeline/agentic/env_manager/traj_env_manager.py b/roll/pipeline/agentic/env_manager/traj_env_manager.py index eef9330e6..e3d98b20a 100644 --- a/roll/pipeline/agentic/env_manager/traj_env_manager.py +++ b/roll/pipeline/agentic/env_manager/traj_env_manager.py @@ -1,5 +1,5 @@ import copy -from contextlib import nullcontext +from contextlib import nullcontext, ExitStack from threading import Lock from typing import Optional @@ -25,6 +25,7 @@ from roll.utils.functionals import pad_to_length, aggregate_metrics from roll.utils.logging import get_logger from roll.utils.str_utils import contains_renderable_field +from roll.utils.telemetry import attach_trace_context, get_tracer class TrajEnvManager(BaseEnvManager): @@ -49,6 +50,7 @@ def __init__(self, self.output_queue = output_queue self.mode = mode self.generate_scheduler: RouterManager = generate_scheduler + self.trace_stack = None # EnvManager states self.rollout_cache: Optional[RolloutCache] = None @@ -111,13 +113,15 @@ def run_rollout_loop(self, data: DataProto): while self.running and rollout_cache is not None: - with Timer(name="generate", logger=None) as generate_timer: + with Timer(name="generate", logger=None) as generate_timer, \ + get_tracer("env_manager").start_as_current_span("generate"): lm_output: DataProto = self.make_decision(rollout_cache) stop_reason = lm_output.meta_info.pop("stop_reason") log_stats["current_step"].append(self.current_step) log_stats["generate_time"].append(generate_timer.last) - with Timer(name="step", logger=None) as step_timer: + with Timer(name="step", logger=None) as step_timer, \ + get_tracer("env_manager").start_as_current_span("step"): if stop_reason == GenerateStopReason.FINISH: rollout_cache: RolloutCache = self.step(lm_output) log_stats["step_time"].append(step_timer.last) @@ -131,6 +135,7 @@ def run_rollout_loop(self, data: DataProto): rollout.non_tensor_batch["traj_group_id"] = np.array([traj_group_id] * rollout.batch.batch_size[0], dtype=object) rollout.non_tensor_batch["traj_id"] = np.array([traj_id] * rollout.batch.batch_size[0], dtype=object) ray.get(self.output_queue.put.remote(self.env_config['group_id'], self.episode_id, start_step, rollout, self.env_config['env_id'])) + self.trace_stack.close() rollout_cache = self.reset() start_step = self.current_step @@ -151,6 +156,23 @@ def reset(self) -> RolloutCache: return None seed = self.group_seed + self.episode_id + traj_group_id = f"{self.rollout_cache.tag}_{self.rollout_cache.group_id}_{self.episode_id}_{self.group_seed}" + traj_id = f"{traj_group_id}_{self.rollout_cache.env_id}" + self.trace_stack = ExitStack() + self.trace_stack.enter_context(attach_trace_context(self.trace_meta)) + self.trace_stack.enter_context( + get_tracer("env_manager").start_as_current_span( + "trajectory", + attributes={ + "traj_id": traj_id, + "traj_group_id": traj_group_id, + "env_id": self.env_config["env_id"], + "group_id": self.env_config["group_id"], + "episode_id": self.episode_id, + }, + ) + ) + with self.thread_lock, self.env_step_limiter: # `observation` describes the current game-state prompt; # `info["suffix"]` carries the current environment-specific state string. @@ -371,4 +393,4 @@ def formulate_rollouts(self, rollout_cache: RolloutCache): env_metric = {f"env/{rollout_cache.tag}/{k}": v for k, v in env_metric.items()} env_metric["env/response_length"] = response_length lm_input.meta_info = {"metrics": env_metric} - return lm_input \ No newline at end of file + return lm_input diff --git a/roll/pipeline/agentic/env_manager/vl_traj_env_manager.py b/roll/pipeline/agentic/env_manager/vl_traj_env_manager.py index 57e951601..892dfbc84 100644 --- a/roll/pipeline/agentic/env_manager/vl_traj_env_manager.py +++ b/roll/pipeline/agentic/env_manager/vl_traj_env_manager.py @@ -370,7 +370,7 @@ def replace_placeholder(text): assert messages, f"empty messages with {history=}" add_generation_prompt = False if messages[-1]["role"] == "assistant" else True lm_input_texts = self.tokenizer.apply_chat_template( - messages, add_generation_prompt=add_generation_prompt, tokenize=False + messages, add_generation_prompt=add_generation_prompt, tokenize=False, return_dict=False ) feature = { self.collator.prompt_key: lm_input_texts, diff --git a/roll/pipeline/agentic/environment_worker.py b/roll/pipeline/agentic/environment_worker.py index bd4ede7b6..2289114b9 100644 --- a/roll/pipeline/agentic/environment_worker.py +++ b/roll/pipeline/agentic/environment_worker.py @@ -5,9 +5,9 @@ from concurrent.futures import ThreadPoolExecutor, as_completed from typing import Dict, Optional -from codetiming import Timer from transformers import PreTrainedTokenizer, ProcessorMixin +from roll.pipeline.agentic.proxy.router import create_proxy_router from roll.utils.context_managers import local_profiler from roll.pipeline.agentic.env_manager.base_env_manager import BaseEnvManager @@ -18,6 +18,7 @@ from roll.pipeline.agentic.agentic_config import EnvManagerConfig from roll.utils.checkpoint_manager import download_model from roll.utils.import_utils import safe_import_class +from roll.pipeline.agentic.proxy.proxy_server import ProxyServer class EnvironmentWorker(Worker): @@ -40,6 +41,7 @@ def __init__(self, worker_config: EnvManagerConfig): self.thread_lock = threading.Lock() self.output_queue = None self.mode = "train" + self.proxy_server: Optional[ProxyServer] = None @register(dispatch_mode=Dispatch.ONE_TO_ALL, clear_cache=False) async def initialize(self, @@ -52,6 +54,38 @@ async def initialize(self, self.output_queue = output_queue self.mode = mode + + needs_proxy = any( + "proxy_env_manager" in cfg.get('env_manager_cls', '') for cfg in self.env_configs.values() + ) + has_rocknative = any(cfg.get('env_type') == "rocknative" for cfg in self.env_configs.values()) + # Pull-mode runners poll the sandbox directly and do not need ALB/ingress registration. + needs_proxy_router = has_rocknative and not any( + "PullModeRunner" in cfg.get('agent_runner_cls', '') for cfg in self.env_configs.values() + ) + + if needs_proxy: + # Create and start proxy server for any env that uses ProxyEnvManager + proxy_port = self.get_free_port() + self.proxy_server = ProxyServer(host="0.0.0.0", port=proxy_port) + self.proxy_server.start(timeout=30) + self.logger.info(f"Proxy server started at: {self.proxy_server.url}") + + for env_id, env_config in self.env_configs.items(): + env_config['proxy_port'] = proxy_port + + if needs_proxy_router: + # Register proxy router only for push-mode rocknative environments + proxy_url = self.proxy_server.url if self.proxy_server else None + if proxy_url is None: + raise RuntimeError("rocknative env_type requires a proxy server, but no ProxyEnvManager is configured") + backend = create_proxy_router() + result = backend.register_servers( + addresses=[proxy_url], + job_id=os.environ.get('TASK_ID'), + port=proxy_port, + ) + self.logger.info(f'[Network Register] Auto registration result: {result}') model_name_or_path = download_model(self.worker_config.model_args.model_name_or_path) self.tokenizer = default_tokenizer_provider(self.worker_config.model_args, model_name_or_path) self.processor = default_processor_provider(self.worker_config.model_args, model_name_or_path) @@ -91,6 +125,15 @@ def create_env_manager(env_id, env_config): self.logger.error(f"Failed to initialize env_manager: {e}", exc_info=True) raise e + # Register handlers for each env_manager + if self.proxy_server: + for env_id, env_manager in self.env_managers.items(): + if hasattr(env_manager, 'process_request'): + self.proxy_server.register_handler(env_id, env_manager.process_request) + self.logger.info(f"Registered env_id={env_id} with proxy server") + else: + self.logger.warning(f"env_manager for env_id={env_id} does not have process_request method") + @register(dispatch_mode=Dispatch.ONE_TO_ALL, clear_cache=False) async def run_rollout_loop(self, seed): # Set environment variables for profiler context @@ -119,11 +162,16 @@ def run_without_profiler(env_manager, data_proto): pool.shutdown() @register(dispatch_mode=Dispatch.ONE_TO_ALL, clear_cache=False) - async def update_step(self, global_step): + async def update_step(self, global_step, trace_meta): for env_manager in self.env_managers.values(): - env_manager.update_step(global_step) + env_manager.update_step(global_step, trace_meta) @register(dispatch_mode=Dispatch.ONE_TO_ALL, clear_cache=False) async def stop(self): for env_manager in self.env_managers.values(): env_manager.stop() + + # Stop proxy server + if self.proxy_server: + self.proxy_server.stop() + self.logger.info("Proxy server stopped") diff --git a/roll/pipeline/agentic/llm_proxy/openai_proxy.py b/roll/pipeline/agentic/llm_proxy/openai_proxy.py index 3c2147426..16a7356df 100644 --- a/roll/pipeline/agentic/llm_proxy/openai_proxy.py +++ b/roll/pipeline/agentic/llm_proxy/openai_proxy.py @@ -83,11 +83,14 @@ def generate(self, extra_body = {} if top_k is not None: extra_body["top_k"] = top_k - if enable_thinking: + if not enable_thinking: # According to the example, enable_thinking goes under extend_fields - extra_body.setdefault("extend_fields", {})["chat_template_kwargs"] = {"enable_thinking": enable_thinking} + #extra_body.setdefault("extend_fields", {})["chat_template_kwargs"] = {"enable_thinking": enable_thinking} + extra_body["chat_template_kwargs"] = {"enable_thinking": enable_thinking} attempt = 0 + tools = lm_input.meta_info.get("tools", []) + tool_choice = "none" while attempt < self.max_retries: try: logger.debug(f"Attempt {attempt + 1}/{self.max_retries}: Calling OpenAI API for model '{model_name}'...") @@ -96,20 +99,54 @@ def generate(self, completion = self.client.chat.completions.create( model=model_name, messages=messages, + tools=tools, + tool_choice=tool_choice, temperature=temperature, max_tokens=max_tokens, top_p=top_p, # presence_penalty=presence_penalty, # Pass extra_body only if it's not empty - extra_body=extra_body if extra_body else None + extra_body=extra_body if extra_body else None, + logprobs=True, + top_logprobs=1, # Return top 1 token logprobs for each position ) - if completion.choices is None: + if not completion.choices: response_text = "OpenAI API returned no choices." - else: + elif completion.choices[0].message.content is not None: response_text = completion.choices[0].message.content + else: + # Model used native function calling (content=None, tool_calls=[...]). + # Convert to XML so proxy_env_manager's sglang parser can handle it. + native_tcs = completion.choices[0].message.tool_calls or [] + if native_tcs: + import json as _json + parts = [] + for tc in native_tcs: + try: + args = _json.loads(tc.function.arguments) + except Exception: + args = tc.function.arguments + payload = _json.dumps({"name": tc.function.name, "arguments": args}, ensure_ascii=False) + parts.append(f"\n{payload}\n") + response_text = "\n".join(parts) + else: + response_text = "OpenAI API returned no choices." + + # Extract logprobs if available + infer_logprobs = None + if completion.choices and completion.choices[0].logprobs and completion.choices[0].logprobs.content: + # Extract the logprob value for each token (matching the format expected by proxy_env_manager) + infer_logprobs = [token_data.logprob for token_data in completion.choices[0].logprobs.content] + responses = self.tokenizer([response_text], return_tensors="pt") lm_input.batch["responses"] = responses["input_ids"] lm_input.non_tensor_batch["response_text"] = np.array([response_text], dtype=object) + + # Store logprobs in batch as tensor if available + if infer_logprobs is not None: + import torch + lm_input.batch["infer_logprobs"] = torch.tensor([infer_logprobs], dtype=torch.float32) + return lm_input except OpenAIError as e: diff --git a/roll/pipeline/agentic/llm_proxy/proxy_utils.py b/roll/pipeline/agentic/llm_proxy/proxy_utils.py index 3e2598f29..87a4be230 100644 --- a/roll/pipeline/agentic/llm_proxy/proxy_utils.py +++ b/roll/pipeline/agentic/llm_proxy/proxy_utils.py @@ -83,7 +83,8 @@ def generate_by_proxy( lm_input_texts = tokenizer.apply_chat_template( messages, add_generation_prompt=True, - tokenize=False + tokenize=False, + return_dict=False ) # Build feature dict diff --git a/roll/pipeline/agentic/proxy/proxy_server.py b/roll/pipeline/agentic/proxy/proxy_server.py new file mode 100644 index 000000000..0b54e06aa --- /dev/null +++ b/roll/pipeline/agentic/proxy/proxy_server.py @@ -0,0 +1,277 @@ +import asyncio +import logging +import socket +import threading +import time +from typing import Callable, Dict, Optional + +import uvicorn +from fastapi import FastAPI, Request +from fastapi.responses import JSONResponse + + +logger = logging.getLogger(__name__) + + +class ProxyServer: + """ + Proxy server that routes requests to environment managers. + + Each EnvironmentWorker creates a ProxyServer instance that: + 1. Starts a FastAPI server on a specified port + 2. Exposes the server URL (ip:port) for external access + 3. Routes requests to registered environment managers by env_id + + Attributes: + host (str): Server host address + port (int): Server port + url (str): Full server URL (http://host:port) + app (FastAPI): FastAPI application instance + """ + + def __init__(self, host: str = "0.0.0.0", port: int = 0): + """ + Initialize the proxy server. + + Args: + host (str): Host address to bind to (default: 0.0.0.0) + port (int): Port to listen on (0 for auto-assigned port) + """ + self.host = host + self.port = self._find_available_port(port) if port == 0 else port + self.url = f"http://{self._get_public_ip()}:{self.port}" + + self.app: Optional[FastAPI] = None + self.server_instance: Optional[uvicorn.Server] = None + self.server_thread: Optional[threading.Thread] = None + + # env_id -> request handler mapping + self._handlers: Dict[int, Callable] = {} + + def _find_available_port(self, preferred_port: int = 0) -> int: + """ + Find an available port. + + Args: + preferred_port (int): Preferred port (0 for any available port) + + Returns: + int: Available port number + """ + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + if preferred_port == 0: + s.bind(("", 0)) + s.listen(1) + port = s.getsockname()[1] + else: + s.bind(("", preferred_port)) + port = preferred_port + return port + + def _get_public_ip(self) -> str: + """ + Get the public IP address of this machine. + + Returns: + str: Public IP address or hostname + """ + try: + hostname = socket.gethostname() + ip_address = socket.gethostbyname(hostname) + return ip_address + except Exception as e: + logger.warning(f"Failed to get public IP: {e}, using 0.0.0.0") + return "0.0.0.0" + + def register_handler(self, env_id: int, handler: Callable): + """ + Register a request handler for a specific environment. + + Args: + env_id (int): Environment identifier + handler (Callable): Async function to handle requests for this env_id + Should accept (Request) and return Response + """ + self._handlers[env_id] = handler + logger.info(f"Registered handler for env_id={env_id}") + + def unregister_handler(self, env_id: int): + """ + Unregister a handler for a specific environment. + + Args: + env_id (int): Environment identifier + """ + if env_id in self._handlers: + del self._handlers[env_id] + logger.info(f"Unregistered handler for env_id={env_id}") + + async def _route_request(self, request: Request): + """ + Route request to the appropriate environment handler. + + Args: + env_id (int): Environment identifier from request + request (Request): FastAPI request object + + Returns: + Response: Handler response or error response + """ + env_id = await self._parser_from_request(request) + if env_id not in self._handlers: + logger.warning(f"No handler registered for env_id={env_id}") + return JSONResponse( + status_code=404, + content={"error": f"No handler found for env_id={env_id}"} + ) + + handler = self._handlers[env_id] + try: + return await handler(request) + except Exception as e: + logger.error(f"Handler error for env_id={env_id}: {e}", exc_info=True) + return JSONResponse( + status_code=500, + content={"error": f"Handler error: {str(e)}"} + ) + + async def _health_check(self): + """Health check endpoint.""" + return JSONResponse( + content={ + "status": "healthy", + "url": self.url, + "registered_env_ids": list(self._handlers.keys()) + } + ) + + async def _state_endpoint(self): + """State endpoint to query registered handlers.""" + return JSONResponse( + content={ + "url": self.url, + "host": self.host, + "port": self.port, + "registered_env_ids": list(self._handlers.keys()), + "num_handlers": len(self._handlers) + } + ) + + async def _start_async_server(self): + """Start the FastAPI server asynchronously.""" + self.app = FastAPI(title="Environment Proxy Server") + + # Register routes + self.app.add_api_route("/v1/chat/completions", self._route_request, methods=["POST"]) + self.app.add_api_route("/health", self._health_check, methods=["GET", "POST"]) + self.app.add_api_route("/state", self._state_endpoint, methods=["GET", "POST"]) + self.app.add_api_route("/v1/state", self._state_endpoint, methods=["GET", "POST"]) + + config = uvicorn.Config( + self.app, + host=self.host, + port=self.port, + log_level="info" + ) + self.server_instance = uvicorn.Server(config) + await self.server_instance.serve() + + def start(self, timeout: int = 30): + """ + Start the proxy server in a background thread. + + Args: + timeout (int): Maximum time to wait for server startup (seconds) + + Raises: + RuntimeError: If server fails to start within timeout period + """ + def run_server_in_new_loop(): + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + loop.run_until_complete(self._start_async_server()) + loop.close() + + self.server_thread = threading.Thread(target=run_server_in_new_loop) + self.server_thread.daemon = True + self.server_thread.start() + + # Wait for server to become ready + interval = 0.5 + start_time = time.time() + server_ready = False + + logger.info(f"Waiting for proxy server to start on {self.host}:{self.port}...") + + while time.time() - start_time < timeout: + try: + with socket.create_connection(('localhost', self.port), timeout=interval): + logger.info(f"Proxy server started successfully on port {self.port}") + logger.info(f"Server URL: {self.url}") + server_ready = True + break + except (ConnectionRefusedError, socket.timeout, OSError): + time.sleep(interval) + + if not server_ready: + raise RuntimeError( + f"Proxy server failed to start within {timeout} seconds on port {self.port}" + ) + + def stop(self, timeout: int = 5): + """ + Stop the proxy server and wait for thread to exit. + + Args: + timeout (int): Maximum time to wait for server shutdown (seconds) + """ + if self.server_instance: + logger.info(f"Stopping proxy server on port {self.port}") + self.server_instance.should_exit = True + + # Wait for server thread to finish + if self.server_thread and self.server_thread.is_alive(): + self.server_thread.join(timeout=timeout) + if self.server_thread.is_alive(): + logger.warning(f"Server thread did not exit within {timeout} seconds") + else: + logger.info("Proxy server stopped successfully") + + + async def _parser_from_request(self, request: Request) -> int: + auth_header = request.headers.get("Authorization", "") + if not auth_header: + return -1 + + try: + # 假设 Agent 发送的是 "Bearer 1001" + token = auth_header.split()[-1] + return int(token) # 得到 1001 + except (ValueError, IndexError): + return -1 + + +async def mock_handler(request: Request): + """模拟一个环境处理逻辑""" + return JSONResponse(content={"message": "Hello from Environment Handler!", "env_id": 1}) + +if __name__ == "__main__": + logging.basicConfig(level=logging.INFO) + + # 1. 初始化服务器（使用固定端口方便调试，如 8000） + server = ProxyServer(host="127.0.0.1", port=8000) + + # 2. 注册一个处理器 (env_id=1) + server.register_handler(env_id=1001, handler=mock_handler) + + # 3. 启动服务器 + try: + server.start() + print(f"服务器已启动: {server.url}") + + # 保持主线程运行 + while True: + time.sleep(1) + except KeyboardInterrupt: + server.stop() + print("服务器已停止") diff --git a/roll/pipeline/agentic/proxy/router/__init__.py b/roll/pipeline/agentic/proxy/router/__init__.py new file mode 100644 index 000000000..a69df7111 --- /dev/null +++ b/roll/pipeline/agentic/proxy/router/__init__.py @@ -0,0 +1,9 @@ +from .base import ProxyRouter, batch_cleanup, create_proxy_router +from .alb_proxy_router import ALBProxyRouter + +try: + from .ingress_proxy_router import IngressProxyRouter +except ImportError: + IngressProxyRouter = None + +__all__ = ["ProxyRouter", "ALBProxyRouter", "IngressProxyRouter", "batch_cleanup", "create_proxy_router"] diff --git a/roll/pipeline/agentic/proxy/router/alb_proxy_router.py b/roll/pipeline/agentic/proxy/router/alb_proxy_router.py new file mode 100644 index 000000000..2794cf533 --- /dev/null +++ b/roll/pipeline/agentic/proxy/router/alb_proxy_router.py @@ -0,0 +1,1247 @@ +# -*- coding: utf-8 -*- +""" +Alibaba Cloud ALB (Application Load Balancer) registration utilities. + +This module provides functionality to register rollout server instances +to Alibaba Cloud ALB server groups for load balancing. +""" +import hashlib +import os +import random +import re +import uuid +from typing import List, Optional, Tuple + + +def _get_logger(): + import logging + return logging.getLogger(__name__) + +logger = _get_logger() + + +def parse_address(address: str) -> Tuple[Optional[str], Optional[int]]: + """ + Parse IP and port from address string. + + Args: + address: Address string like "http://192.168.1.1:8080" + + Returns: + Tuple of (ip, port) or (None, None) if parsing fails + """ + match = re.match(r'https?://([^:]+):(\d+)', address) + if match: + return match.group(1), int(match.group(2)) + return None, None + + +def get_alb_load_balancer_id_by_hash( + job_name: str, + listen_port: int, + alb_ids_env_var: str = 'ALB_LOAD_BALANCER_IDS', +) -> Tuple[str, str]: + """ + Select an ALB load balancer ID from a comma-separated list based on hash of job_name and listen_port. + + This function distributes jobs across multiple ALBs using consistent hashing to ensure + stable routing for the same job_name/listen_port combination. + + Args: + job_name: Name of the job (used for consistent hashing) + listen_port: Port number being listened on (used for consistent hashing) + alb_ids_env_var: Environment variable name containing comma-separated ALB IDs + + Returns: + Selected ALB load balancer ID + + Environment variables: + - ALB_LOAD_BALANCER_IDS: Comma-separated list of ALB load balancer IDs + (e.g., "alb-abc,alb-def,alb-ghi") + + Raises: + ValueError: If ALB_LOAD_BALANCER_IDS environment variable is not set or empty + """ + alb_ids = os.environ.get(alb_ids_env_var, None) + alb_lsn_ids = os.environ.get('ALB_LOAD_BALANCER_LSN_IDS', None) + if not alb_ids: + raise ValueError(f"Environment variable '{alb_ids_env_var}' is not set") + + alb_list = [alb.strip() for alb in alb_ids.split(',')] + alb_lsn_list = [alb_lsn.strip() for alb_lsn in alb_lsn_ids.split(',')] + if not alb_list: + raise ValueError(f"No ALB IDs found in '{alb_ids_env_var}'") + + # Use consistent hashing to select ALB ID with fixed seed for consistency + hash_value = get_deterministic_hash(job_name, listen_port) + selected_index = hash_value % len(alb_list) + print(f"alb_list: {alb_list}, alb_lsn_list: {alb_lsn_list}") + print(f"job name {job_name}, listen_port {listen_port}, hash {hash_value}, selected_index: {selected_index}") + return alb_list[selected_index], alb_lsn_list[selected_index] + + +def get_deterministic_hash(job_name, listen_port): + # 将变量转换为字符串并拼接 + data = f"{job_name}:{listen_port}" + # 使用 SHA-256 生成哈希 + hash_obj = hashlib.sha256(data.encode('utf-8')) + # 取前几位作为整数，或者直接返回 hex 字符串 + return int(hash_obj.hexdigest()[:8], 16) + + +def discover_and_set_alb_load_balancers( + vpc_id: str = None, + alb_ids_env_var: str = 'ALB_LOAD_BALANCER_IDS', + region: str = None, + client: object = None +) -> str: + """ + Discover all ALBs in the current region and set them to environment variable ALB_LOAD_BALANCER_IDS. + + This function queries Alibaba Cloud ALB service to find all load balancers in the specified region + and sets their IDs to the ALB_LOAD_BALANCER_IDS environment variable. + + Args: + vpc_id: VPC ID to filter ALBs (optional, if None uses default VPC from ALB_VPC_ID env var) + alb_ids_env_var: Environment variable name to set ALB IDs (default: ALB_LOAD_BALANCER_IDS) + region: Region to query ALBs (optional, if None uses default from ALB_REGION env var) + + Returns: + Comma-separated string of ALB IDs + + Environment variables: + - ALIBABA_CLOUD_ACCESS_KEY_ID: Access key ID + - ALIBABA_CLOUD_ACCESS_KEY_SECRET: Access key secret + - ALB_REGION: Optional, region to query (default: cn-hangzhou) + - ALB_VPC_ID: Optional, VPC ID to filter ALBs + + Raises: + ImportError: If alibabacloud SDK is not installed + ValueError: If required environment variables are not set + Exception: If ALB discovery fails + """ + from alibabacloud_alb20200616 import models as alb_20200616_models + from alibabacloud_tea_util import models as util_models + + # Determine region + if region is None: + region = os.environ.get('ALB_REGION', None) + + assert region is not None, "region is required for querying ALBs" + + list_load_balancers_request = alb_20200616_models.ListLoadBalancersRequest( + zone_id='cn-hangzhou-a' + ) + + runtime = util_models.RuntimeOptions() + print(f"start List load balancers: {client}") + response = client.list_load_balancers_with_options(list_load_balancers_request, runtime) + + # Extract ALB IDs + alb_ids = [] + if response.body.load_balancers: + for lb in response.body.load_balancers: + alb_ids.append(lb.load_balancer_id) + + if not alb_ids: + raise ValueError(f"No ALBs found in region {region} with VPC ID {vpc_id}") + + # Join ALB IDs with comma + alb_ids_str = ','.join(alb_ids) + + # Set environment variable + os.environ[alb_ids_env_var] = alb_ids_str + + logger.info(f"Discovered {len(alb_ids)} ALBs in region {region}: {alb_ids_str}") + + return alb_ids_str + + +def create_alb_client(endpoint: str = None) -> object: + """ + Create ALB client using credentials from environment variables. + + Environment variables required: + - ALIBABA_CLOUD_ACCESS_KEY_ID + - ALIBABA_CLOUD_ACCESS_KEY_SECRET + - ALB_ENDPOINT (optional, default: alb.cn-hangzhou.aliyuncs.com) + + Returns: + Alb20200616Client instance + + Raises: + ImportError: If alibabacloud SDK is not installed + ValueError: If required environment variables are not set + """ + if endpoint is None: + alb_region = os.environ.get('ALB_REGION') + endpoint = f'alb-vpc.{alb_region}.aliyuncs.com' if alb_region else None + try: + from alibabacloud_alb20200616.client import Client as Alb20200616Client + from alibabacloud_credentials.client import Client as CredentialClient + from alibabacloud_credentials.models import Config as CredentialConfig + from alibabacloud_tea_openapi import models as open_api_models + except ImportError: + raise ImportError( + "Please install alibabacloud SDK: " + "pip install alibabacloud_alb20200616 alibabacloud_credentials alibabacloud_tea_openapi" + ) + + access_key_id = os.environ.get('ALIBABA_CLOUD_ACCESS_KEY_ID') + access_key_secret = os.environ.get('ALIBABA_CLOUD_ACCESS_KEY_SECRET') + + if not access_key_id or not access_key_secret: + raise ValueError( + "ALIBABA_CLOUD_ACCESS_KEY_ID and ALIBABA_CLOUD_ACCESS_KEY_SECRET " + "environment variables must be set" + ) + + credentials_config = CredentialConfig( + type='access_key', + access_key_id=access_key_id, + access_key_secret=access_key_secret + ) + credentials_client = CredentialClient(credentials_config) + + config = open_api_models.Config(credential=credentials_client) + config.endpoint = endpoint or os.environ.get( + 'ALB_ENDPOINT', 'alb-vpc.cn-hangzhou.aliyuncs.com' + ) + print(f"config.endpoint: {config.endpoint} config.access_key_id: {config.access_key_id} create Alb20200616Client") + return Alb20200616Client(config) + + +def delete_server_group(client, server_group_id: str) -> dict: + """ + Delete an existing ALB server group. + + Args: + client: ALB client instance + server_group_id: Server group ID to delete + + Returns: + API response dict + """ + from alibabacloud_alb20200616 import models as alb_20200616_models + from alibabacloud_tea_util import models as util_models + + print(f"start Delete server group: {server_group_id}") + delete_server_group_request = alb_20200616_models.DeleteServerGroupRequest( + server_group_id=server_group_id + ) + runtime = util_models.RuntimeOptions() + response = client.delete_server_group_with_options(delete_server_group_request, runtime) + + print(f"Deleted server group: {server_group_id}") + return response.body.to_map() if hasattr(response.body, 'to_map') else {} + + +def delete_listener(client, listener_id: str) -> dict: + """ + Delete an existing ALB listener. + + Args: + client: ALB client instance + listener_id: Listener ID to delete + + Returns: + API response dict + """ + from alibabacloud_alb20200616 import models as alb_20200616_models + from alibabacloud_tea_util import models as util_models + + print(f"start Delete listener: {listener_id}") + delete_listener_request = alb_20200616_models.DeleteListenerRequest( + listener_id=listener_id + ) + runtime = util_models.RuntimeOptions() + response = client.delete_listener_with_options(delete_listener_request, runtime) + + print(f"Deleted listener: {listener_id}") + return response.body.to_map() if hasattr(response.body, 'to_map') else {} + + +def create_server_group( + client, + server_group_name: str, + vpc_id: str, + server_type: str = 'Ip', + resource_group_id: str = 'rg-aeky6sewmnhmbyi', + health_check_enabled: bool = False, +) -> str: + """ + Create a new ALB server group. + + Args: + client: ALB client instance + server_group_name: Name for the server group + vpc_id: VPC ID where the server group will be created + health_check_enabled: Whether to enable health check + + Returns: + Server group ID + """ + from alibabacloud_alb20200616 import models as alb_20200616_models + from alibabacloud_tea_util import models as util_models + + print(f"start Created server group: {server_group_name} {vpc_id}") + + health_check_config = alb_20200616_models.CreateServerGroupRequestHealthCheckConfig( + health_check_enabled=False + ) + sticky_session_config = alb_20200616_models.CreateServerGroupRequestStickySessionConfig( + sticky_session_enabled=False + ) + + create_server_group_request = alb_20200616_models.CreateServerGroupRequest( + server_group_type=server_type, + server_group_name=server_group_name, + resource_group_id=resource_group_id, + vpc_id=vpc_id, + health_check_config=health_check_config, + sticky_session_config=sticky_session_config + ) + + runtime = util_models.RuntimeOptions() + response = client.create_server_group_with_options(create_server_group_request, runtime) + + server_group_id = response.body.server_group_id + print(f"Created server group success: {server_group_id}") + return server_group_id + + +def add_servers_to_server_group( + client, + server_group_id: str, + addresses: List[str], + server_type: str = 'Ip', + weight: int = 100, +) -> dict: + """ + Add servers to an existing ALB server group. + + Args: + client: ALB client instance + server_group_id: Target server group ID + addresses: List of server addresses (e.g., ["http://192.168.1.1:8080"]) + server_type: Server type, default is 'Ip' + weight: Server weight for load balancing (0-100) + + Returns: + API response dict + """ + from alibabacloud_alb20200616 import models as alb_20200616_models + from alibabacloud_tea_util import models as util_models + + servers = [] + for address in addresses: + ip, port = parse_address(address) + if ip is None or port is None: + logger.warning(f"Cannot parse address: {address}, skipping") + continue + + server = alb_20200616_models.AddServersToServerGroupRequestServers( + server_id=ip, + server_type=server_type, + server_ip=ip, + remote_ip_enabled=True, + port=port + ) + print(f"Added server {ip}:{port} to server group {server_group_id}") + servers.append(server) + + if not servers: + logger.warning("No valid servers to add to server group") + return {} + + request = alb_20200616_models.AddServersToServerGroupRequest( + server_group_id=server_group_id, + servers=servers + ) + + runtime = util_models.RuntimeOptions() + response = client.add_servers_to_server_group_with_options(request, runtime) + + print(f"Added {len(servers)} servers to server group {server_group_id}") + return response.body.to_map() if hasattr(response.body, 'to_map') else {} + + +def start_listener(client, listener_id: str) -> dict: + """ + Start a stopped ALB listener. + + Args: + client: ALB client instance + listener_id: Listener ID to start + + Returns: + API response dict + """ + from alibabacloud_alb20200616 import models as alb_20200616_models + from alibabacloud_tea_util import models as util_models + + start_listener_request = alb_20200616_models.StartListenerRequest( + listener_id=listener_id + ) + runtime = util_models.RuntimeOptions() + response = client.start_listener_with_options(start_listener_request, runtime) + + print(f"Started listener: {listener_id}") + return response.body.to_map() if hasattr(response.body, 'to_map') else {} + + +def poll_listener_status( + client, + listener_id: str, + target_status: str = 'Running', + max_attempts: int = 10, + interval_seconds: float = 2.0 +) -> dict: + """ + Poll listener status until it reaches target state. + + Args: + client: ALB client instance + listener_id: Listener ID to monitor + target_status: Target status to wait for (default: 'Running') + max_attempts: Maximum number of polling attempts + interval_seconds: Interval between polls in seconds + + Returns: + Final listener status dict + + Raises: + TimeoutError: If target status is not reached within max_attempts + """ + from alibabacloud_alb20200616 import models as alb_20200616_models + from alibabacloud_tea_util import models as util_models + import time + + for attempt in range(max_attempts): + list_listeners_request = alb_20200616_models.ListListenersRequest( + listener_ids=[listener_id] + ) + runtime = util_models.RuntimeOptions() + try: + resp = client.list_listeners_with_options(list_listeners_request, runtime) + if resp.body and resp.body.listeners: + status = resp.body.listeners[0].listener_status + print(f"Poll attempt {attempt + 1}/{max_attempts}: Status={status}") + + if status == target_status: + print(f"Listener {listener_id} reached target status '{target_status}'") + return resp.body.to_map() if hasattr(resp.body, 'to_map') else {} + + if status not in ['Provisioning', 'Running', 'Stopping', 'Configuring', 'Stopped', 'Starting', + 'Deleting', 'Deleted']: + raise RuntimeError(f"Listener {listener_id} entered invalid state: {status}") + else: + print(f"Poll attempt {attempt + 1}/{max_attempts}: Listener not found") + except Exception as e: + print(f"Poll attempt {attempt + 1}/{max_attempts}: Error checking status - {e}") + + if attempt < max_attempts - 1: + time.sleep(interval_seconds) + + raise TimeoutError( + f"Listener {listener_id} did not reach status '{target_status}' " + f"within {max_attempts} attempts" + ) + + +def poll_server_group_status( + client, + server_group_id: str, + target_status: str = 'Available', + max_attempts: int = 10, + interval_seconds: float = 2.0 +) -> dict: + """ + Poll server group status until it reaches target state. + + Args: + client: ALB client instance + server_group_id: Server group ID to monitor + target_status: Target status to wait for (default: 'Active') + max_attempts: Maximum number of polling attempts + interval_seconds: Interval between polls in seconds + + Returns: + Final server group status dict + + Raises: + TimeoutError: If target status is not reached within max_attempts + """ + from alibabacloud_alb20200616 import models as alb_20200616_models + from alibabacloud_tea_util import models as util_models + import time + + for attempt in range(max_attempts): + list_server_groups_request = alb_20200616_models.ListServerGroupsRequest( + server_group_ids=[server_group_id] + ) + runtime = util_models.RuntimeOptions() + try: + resp = client.list_server_groups_with_options(list_server_groups_request, runtime) + if resp.body and resp.body.server_groups: + status = resp.body.server_groups[0].server_group_status + print(f"Poll attempt {attempt + 1}/{max_attempts}: Status={status}") + + if status == target_status: + print(f"Server group {server_group_id} reached target status '{target_status}'") + return resp.body.to_map() if hasattr(resp.body, 'to_map') else {} + + if status not in ['Creating', 'Available', 'Configuring']: + raise RuntimeError(f"Server group {server_group_id} entered invalid state: {status}") + else: + print(f"Poll attempt {attempt + 1}/{max_attempts}: Server group not found") + except Exception as e: + print(f"Poll attempt {attempt + 1}/{max_attempts}: Error checking status - {e}") + + if attempt < max_attempts - 1: + time.sleep(interval_seconds) + + raise TimeoutError( + f"Server group {server_group_id} did not reach status '{target_status}' " + f"within {max_attempts} attempts" + ) + + +def add_rules( + client, + listener_id: str, + server_group_id: str, + port: int, + weight: int = 100, +) -> dict: + """ + Add a forwarding rule to a listener. + + Args: + client: ALB client instance + listener_id: Listener ID to add rule to + server_group_id: Server group ID to forward traffic to + path_pattern: Path pattern for matching (e.g., '*job_id/port/*') + priority: Rule priority (lower number = higher priority) + weight: Server group weight for load balancing + + Returns: + API response dict + """ + import time + + retry_count = 0 + while retry_count < 30: + try: + random_seed = str(uuid.uuid4()) + priority = hash(random_seed) % 10000 + from alibabacloud_alb20200616 import models as alb_20200616_models + from alibabacloud_tea_util import models as util_models + import json + + job_id = os.environ.get('TASK_ID') + print(f"add rules to : {listener_id} {job_id} {port} {priority}") + + # Create server group tuple + server_group_tuple = alb_20200616_models.CreateRulesRequestRulesRuleActionsForwardGroupConfigServerGroupTuples( + server_group_id=server_group_id, + weight=weight + ) + + # Create forward group config + forward_group_config = alb_20200616_models.CreateRulesRequestRulesRuleActionsForwardGroupConfig( + server_group_tuples=[server_group_tuple] + ) + + rule_actions_0rewrite_config = alb_20200616_models.CreateRuleRequestRuleActionsRewriteConfig( + host='${host}', + path='/${2}', + query='${query}' + ) + + rule_actions_rewrite = alb_20200616_models.CreateRuleRequestRuleActions( + order=1, + rewrite_config=rule_actions_0rewrite_config, + type='Rewrite' + ) + + # Create rule action + rule_action = alb_20200616_models.CreateRuleRequestRuleActions( + order=2, + forward_group_config=forward_group_config, + type='ForwardGroup' + ) + path_pattern = f"~/(.*)/{job_id}/{port}/(.*)" + # Create path config for condition + path_config = alb_20200616_models.CreateRulesRequestRulesRuleConditionsPathConfig( + values=[path_pattern] + ) + + # Create rule condition + rule_condition = alb_20200616_models.CreateRulesRequestRulesRuleConditions( + type='Path', + path_config=path_config + ) + + # Generate rule name based on listener_id and priority + rule_name = f"rule-{job_id}-{port}" + + # Create rule + rule = alb_20200616_models.CreateRulesRequestRules( + priority=priority, + direction='Request', + rule_name=rule_name, + rule_conditions=[rule_condition], + rule_actions=[rule_actions_rewrite, rule_action] + ) + + # Create rules request + create_rules_request = alb_20200616_models.CreateRulesRequest( + listener_id=listener_id, + rules=[rule] + ) + + runtime = util_models.RuntimeOptions() + + resp = client.create_rules_with_options(create_rules_request, runtime) + print(f"Successfully added rule: {rule_name}") + print(f"Path pattern: {path_pattern}") + print(f"Server group: {server_group_id}") + print(f"Add Rules Response: {resp}") + print(f"Add Rules Response Body: {resp.body}") + print(f"Add Rules Response Body List: {resp.body.rule_ids}") + + # 成功执行后跳出循环 + break + + except Exception as e: + # 只有在异常代码为Conflict.Priority时才重试 + retry_count += 1 + print(f"Retry attempt {retry_count}/20: {e}") + if retry_count >= 20: + raise Exception(f"Failed to add rules after {retry_count} attempts. Last error: {e}") + # 等待一段时间再重试 + sleep_time = random.randint(10, 30) + time.sleep(sleep_time) + + # Extract rule IDs from response + assert resp.body.rule_ids + rule_ids = [rule_id.rule_id for rule_id in resp.body.rule_ids] + + return { + 'rule_name': rule_name, + 'rule_ids': rule_ids, + 'response': resp.body.to_map() if hasattr(resp.body, 'to_map') else {} + } + + +def cleanup_alb_resources(): + """ + Scan ALB_RESOURCE_CLEAN_DIR and delete corresponding resources. + + Reads files in ALB_RESOURCE_CLEAN_DIR, extracts resource IDs, + and deletes the corresponding server groups/listeners. + + Files should follow the naming pattern: + - server-group_{id}.txt + - listener_{id}.txt + + Environment variables: + - ALB_RESOURCE_CLEAN_DIR: Directory containing resource ID files + - ALIBABA_CLOUD_ACCESS_KEY_ID: Access key ID + - ALIBABA_CLOUD_ACCESS_KEY_SECRET: Access key secret + - ALB_ENDPOINT: Optional, ALB API endpoint + + Returns: + dict: Summary of deleted resources + """ + clean_parent_dir = os.environ.get('ALB_RESOURCE_CLEAN_DIR') + task_dir = os.environ.get('TASK_ID') + clean_dir = os.path.join(clean_parent_dir, task_dir) + if not clean_dir: + print("ALB_RESOURCE_CLEAN_DIR not set, nothing to clean up") + return {"cleaned_up": 0, "errors": []} + + if not os.path.exists(clean_dir): + print(f"Directory {clean_dir} does not exist") + return {"cleaned_up": 0, "errors": []} + + client = create_alb_client() + cleaned_up = 0 + errors = [] + filenames = sorted(os.listdir(clean_dir), key=lambda x: ( + 0 if x.startswith('rule_') else 1, + 0 if x.startswith('server-group_') else 1, + x + )) + + for filename in filenames: + filepath = os.path.join(clean_dir, filename) + print(f"Scanning file: {filepath}") + # Skip directories and non-text files + if not os.path.isfile(filepath) or not filename.endswith('.txt'): + continue + + # Parse resource type and ID from filename + parts = filename.split('_') + if len(parts) < 2: + continue + + resource_type = parts[0] + resource_id = parts[1][:-4] # Remove .txt extension for other types + print(f"attempt to delete Resource type: {resource_type}, ID: {resource_id}") + + # Use the new helper function + success, error_msg = delete_resource_with_retry(client, resource_type, resource_id, 5) + if not success: + errors.append(error_msg) + + # Remove the file after attempting deletion + if success: + safe_remove_file(filepath) + cleaned_up += 1 + + # Check if clean_dir is empty and remove it if so + try: + if os.path.exists(clean_dir) and not os.listdir(clean_dir): + os.rmdir(clean_dir) + print(f"Removed empty directory: {clean_dir}") + except Exception as e: + print(f"Warning: Failed to remove empty directory {clean_dir}: {e}") + + return { + "cleaned_up": cleaned_up, + "errors": errors + } + + +def delete_resource_with_retry(client, resource_type: str, resource_id: str, max_retries: int = 5) -> tuple[bool, str]: + """ + 删除资源并带重试机制 + + Args: + client: ALB客户端实例 + resource_type: 资源类型 ('server-group', 'listener', 'rule') + resource_id: 资源ID + max_retries: 最大重试次数 + + Returns: + tuple: (success: bool, error_message: str) + """ + import random + import time + + for attempt in range(max_retries): + try: + # Random sleep before each attempt (1-15 seconds) + if attempt > 0: + delay = random.randint(1, 15) + print(f"Retry attempt {attempt + 1}/{max_retries}, sleeping {delay:.2f} seconds...") + time.sleep(delay) + + success = False + if resource_type == 'server-group': + delete_server_group(client, resource_id) + print(f"Deleted server group: {resource_id}") + success = True + elif resource_type == 'listener': + delete_listener(client, resource_id) + print(f"Deleted listener: {resource_id}") + success = True + elif resource_type == 'rule': + delete_rule(client, resource_id) + print(f"Initiated deletion of rule: {resource_id}") + success = True + else: + print(f"Unknown resource type: {resource_type}") + return False, f"Unknown resource type: {resource_type}" + + if success: + return True, "" # Success + + except Exception as e: + if '404' in f"{e}": + print(f"Resource {resource_id} not found, skipping") + return True, "" # Consider 404 as success + error_msg = f"Failed to delete {resource_type} {resource_id} (attempt {attempt + 1}/{max_retries}): {e}" + print(error_msg) + + if attempt == max_retries - 1: + # Final attempt failed + return False, error_msg + else: + # Will retry + print(f"Will retry after random delay...") + + return False, "" + + +def safe_remove_file(filepath: str) -> bool: + """ + Safely remove a file without throwing any exceptions. + + Args: + filepath: Path to the file to be removed + + Returns: + True if file was successfully removed, False otherwise + """ + try: + if os.path.exists(filepath): + os.remove(filepath) + return True + else: + # File doesn't exist, consider as successfully "removed" + return True + except Exception as e: + # Log the error but don't propagate it + print(f"Warning: Failed to remove file {filepath}: {e}") + return False + + +def save_resource_id_to_file(resource_id: str, resource_type: str = 'server-group') -> str: + """ + Save server group ID or listener ID to a file in ALB_RESOURCE_CLEAN_DIR. + + Args: + resource_id: Resource ID to save (server_group_id or listener_id) + resource_type: Type of resource ('server_group' or 'listener') + + Returns: + File path where the resource ID was saved + + Raises: + ValueError: If ALB_RESOURCE_CLEAN_DIR environment variable is not set + IOError: If file cannot be created or written + """ + clean_parent_dir = os.environ.get('ALB_RESOURCE_CLEAN_DIR') + task_dir = os.environ.get('TASK_ID') + clean_dir = os.path.join(clean_parent_dir, task_dir) + if not clean_dir: + raise ValueError("ALB_RESOURCE_CLEAN_DIR environment variable is not set") + + # Ensure directory exists + os.makedirs(clean_dir, exist_ok=True) + + # Generate filename based on resource type and ID + filename = f"{resource_type}_{resource_id}.txt" + filepath = os.path.join(clean_dir, filename) + + # Write resource ID to file + try: + with open(filepath, 'w') as f: + f.write(resource_id) + print(f"Saved {resource_type} ID {resource_id} to file: {filepath}") + + # Verify file was created and contains correct content + with open(filepath, 'r') as f: + content = f.read().strip() + if content != resource_id: + raise IOError(f"File verification failed: expected '{resource_id}', got '{content}'") + + return filepath + except IOError as e: + raise IOError(f"Failed to save {resource_type} ID to file: {e}") + + +def register_rollout_servers_to_alb( + addresses: List[str], + load_balancer_id: str = None, + server_group_name: str = None, + vpc_id: str = None, + listener_port: int = None, + endpoint: str = None, +) -> dict: + """ + Register rollout servers to Alibaba Cloud ALB. + + This is the main entry point for registering rollout servers. + It creates a new server group (with random name) and adds servers to it. + + Args: + addresses: List of rollout server addresses + load_balancer_id: Load balancer ID (required for creating listener) + server_group_name: Name for new server group (auto-generated if not provided) + vpc_id: VPC ID (required for creating server group) + listener_port: Port for listener + endpoint: ALB API endpoint (default from env or cn-hangzhou) + + Returns: + Dict containing created resource IDs + + Environment variables: + - ALIBABA_CLOUD_ACCESS_KEY_ID: Access key ID + - ALIBABA_CLOUD_ACCESS_KEY_SECRET: Access key secret + - ALB_ENDPOINT: Optional, ALB API endpoint + - ALB_LOAD_BALANCER_IDS: Optional, default load balancer ID + - ALB_VPC_ID: Optional, default VPC ID + """ + # Get defaults from environment variables + vpc_id = vpc_id or os.environ.get('ALB_VPC_ID') + resource_group_id = os.environ.get('ALB_RESOURCE_GROUP_ID') + alb_ids_str = os.environ.get('ALB_LOAD_BALANCER_IDS') + result = { + 'server_group_id': None, + 'listener_id': None, + 'registered_servers': [] + } + + try: + alb_region = os.environ.get('ALB_REGION') + if endpoint is None: + endpoint = os.environ.get('ALB_ENDPOINT') or f'alb.{alb_region}.aliyuncs.com' + client = create_alb_client(endpoint) + if not alb_ids_str: + discover_and_set_alb_load_balancers(vpc_id=vpc_id, region=alb_region, client=client) + job_id = os.environ.get('TASK_ID') + # Get listener_port from first address if not provided + if not listener_port and addresses and addresses[0]: + _, listener_port = parse_address(addresses[0]) + + load_balancer_id, lsn_id = get_alb_load_balancer_id_by_hash(job_id, listener_port) + + # Always create a new server group with name based on first address and port + if not server_group_name: + # Extract IP from first address (e.g., "http://10.71.240.61:20000" -> "10-71-240-61") + if addresses: + first_ip, _ = parse_address(addresses[0]) + if first_ip: + ip_suffix = first_ip.replace('.', '-') + server_group_name = f"roll-rollout-{ip_suffix}-{listener_port}" + else: + server_group_name = f"roll-rollout-{len(addresses)}nodes-{listener_port}" + else: + server_group_name = f"roll-rollout-empty-{listener_port}" + if not vpc_id: + raise ValueError("vpc_id is required for creating server group") + + server_group_id = create_server_group( + client=client, + server_group_name=server_group_name, + vpc_id=vpc_id, + resource_group_id=resource_group_id, + ) + result['server_group_id'] = server_group_id + + poll_server_group_status( + client=client, + server_group_id=server_group_id, + ) + # record server group id + save_resource_id_to_file(server_group_id, resource_type='server-group') + # Add servers to server group + add_response = add_servers_to_server_group( + client=client, + server_group_id=server_group_id, + addresses=addresses + ) + result['registered_servers'] = addresses + + poll_server_group_status( + client=client, + server_group_id=server_group_id, + ) + + assert listener_port is not None, "listener_port is required for creating listener" + # Create listener rules with retry logic + print( + f"prepare to add rules load_balancer_id: {load_balancer_id}, lsn_id: {lsn_id} listener_port: {listener_port}") + assert listener_port and load_balancer_id and lsn_id + + # Add rules to listener and save rule ids to file + try: + rule_result = add_rules( + client=client, + listener_id=lsn_id, + server_group_id=server_group_id, + port=listener_port + ) + + # Save rule IDs to file + if rule_result.get('rule_ids'): + poll_rule_status(client=client, rule_ids=rule_result['rule_ids']) + rule_ids_str = ','.join(rule_result['rule_ids']) + save_resource_id_to_file(rule_ids_str, resource_type='rule') + print(f"Saved rule IDs: {rule_ids_str}") + else: + print("Warning: No rule IDs returned from add_rules") + + except Exception as e: + print(f"Failed to add rules to listener: {e}") + raise + + logger.info(f"Successfully registered rollout servers to ALB: {result}") + return result + + except Exception as e: + logger.error(f"Failed to register rollout servers to ALB: {e}") + error_message = e.data.get("Recommend") if hasattr(e, 'data') else str(e) + logger.error(f"Failed to register rollout servers to ALB: {error_message}") + + raise + + +def poll_rule_status( + client, + rule_ids: List[str], + target_status: str = 'Available', + max_attempts: int = 10, + interval_seconds: float = 2.0 +) -> dict: + """ + Poll rule status until it reaches target state. + + Args: + client: ALB client instance + rule_id: Rule ID to monitor + target_status: Target status to wait for (default: 'Available') + max_attempts: Maximum number of polling attempts + interval_seconds: Interval between polls in seconds + + Returns: + Final rule status dict + + Raises: + TimeoutError: If target status is not reached within max_attempts + """ + from alibabacloud_alb20200616 import models as alb_20200616_models + from alibabacloud_tea_util import models as util_models + import time + + for attempt in range(max_attempts): + # Since ALB API doesn't seem to have a direct ListRulesRequest with rule_ids parameter, + # we'll need to list rules by listener. However, we don't have listener_id here. + # Let's assume there's a way to list rules directly by ID or we need to adapt the approach. + try: + # Attempt to use a generic listing approach + # This would need to be adapted based on actual ALB API capabilities + list_rules_request = alb_20200616_models.ListRulesRequest( + rule_ids=rule_ids + ) + runtime = util_models.RuntimeOptions() + + # This will likely fail without proper parameters, but shows the intended structure + resp = client.list_rules_with_options(list_rules_request, runtime) + + if resp.body and hasattr(resp.body, 'rules'): + for rule in resp.body.rules: + status = rule.rule_status if hasattr(rule, 'rule_status') else 'Unknown' + print(f"Poll attempt {attempt + 1}/{max_attempts}: Status={status}") + + if status == target_status: + print(f"Rule {rule.rule_id} reached target status '{target_status}'") + return resp.body.to_map() if hasattr(resp.body, 'to_map') else {} + + if status not in ['Provisioning', 'Configuring', 'Available', ]: + raise RuntimeError(f"Rule {rule.rule_id} entered invalid state: {status}") + + break + else: + print(f"Poll attempt {attempt + 1}/{max_attempts}: Rule {rule.rule_id} not found in response") + else: + print(f"Poll attempt {attempt + 1}/{max_attempts}: No rules found in response") + + except Exception as e: + print(f"Poll attempt {attempt + 1}/{max_attempts}: Error checking status - {e}") + + if attempt < max_attempts - 1: + time.sleep(interval_seconds) + + raise TimeoutError( + f"Rule {rule_ids} did not reach status '{target_status}' " + f"within {max_attempts} attempts" + ) + + +def delete_rule(client, rule_id: str) -> dict: + """ + Delete a forwarding rule from a listener. + + Args: + client: ALB client instance + rule_id: Rule ID to delete + + Returns: + API response dict + """ + from alibabacloud_alb20200616 import models as alb_20200616_models + from alibabacloud_tea_util import models as util_models + + print(f"Deleting rule: {rule_id}") + + delete_rule_request = alb_20200616_models.DeleteRuleRequest( + rule_id=rule_id + ) + + runtime = util_models.RuntimeOptions() + resp = client.delete_rule_with_options(delete_rule_request, runtime) + + print(f"Successfully initiated deletion of rule: {rule_id}") + print(f"Job ID: {resp.body.job_id}") + + return { + 'job_id': resp.body.job_id, + 'response': resp.body.to_map() if hasattr(resp.body, 'to_map') else {} + } + + +def list_rules() -> dict: + """ + 循环列举指定 listener_ids 下的所有规则，并返回 rule_name 到 rule_id 的映射 + + Args: + client: ALB client instance + listener_ids: Listener IDs 列表 + max_results: 每页最大返回结果数，默认100 + + Returns: + 包含 rule_name 到 rule_id 映射的字典 + """ + from alibabacloud_alb20200616 import models as alb_20200616_models + from alibabacloud_tea_util import models as util_models + client = create_alb_client() + max_results: int = 100 + all_rules = [] + listener_ids = os.environ.get('ALB_LOAD_BALANCER_LSN_IDS').split(',') + next_token = None + + while True: + list_rules_request = alb_20200616_models.ListRulesRequest( + listener_ids=listener_ids, + max_results=max_results + ) + + if next_token: + list_rules_request.next_token = next_token + + runtime = util_models.RuntimeOptions() + + try: + resp = client.list_rules_with_options(list_rules_request, runtime) + + if hasattr(resp.body, 'rules') and resp.body.rules: + all_rules.extend(resp.body.rules) + + next_token = resp.body.next_token if hasattr(resp.body, 'next_token') else None + + if not next_token: + break + + except Exception as error: + print(f"Error occurred while listing rules: {error.message if hasattr(error, 'message') else str(error)}") + if hasattr(error, 'data') and error.data: + print(error.data.get("Recommend")) + raise + + rule_name_to_id_dict = {} + for rule in all_rules: + if hasattr(rule, 'rule_name') and hasattr(rule, 'rule_id'): + rule_name_to_id_dict[rule.rule_name] = rule.rule_id + + return { + 'rules': rule_name_to_id_dict, + 'total_count': len(rule_name_to_id_dict), + 'response': {'rules': rule_name_to_id_dict} + } + + +if __name__ == "__main__": + # Example usage for testing register_rollout_servers_to_alb method + import sys + + # Set environment variables for testing (you need to provide your own values) + os.environ.setdefault('ALB_VPC_ID', '') + os.environ.setdefault('ALB_LOAD_BALANCER_IDS', '') + os.environ.setdefault('ALB_LOAD_BALANCER_LSN_IDS', '') + os.environ.setdefault('ALB_ENDPOINT', '') + os.environ.setdefault('ALIBABA_CLOUD_ACCESS_KEY_ID', '') + os.environ.setdefault('ALIBABA_CLOUD_ACCESS_KEY_SECRET', '') + os.environ.setdefault('ALB_RESOURCE_CLEAN_DIR', '') + os.environ.setdefault('ALB_RESOURCE_GROUP_ID', '') + os.environ.setdefault('ALB_REGION', '') + os.environ.setdefault('TASK_ID', 'xdl-job-123456') + + + # Test addresses (replace with actual rollout server addresses) + test_addresses = [ + "http://10.17.70.145:8000/", + ] + + print("Testing register_rollout_servers_to_alb...") + + try: + result = register_rollout_servers_to_alb( + addresses=test_addresses, + vpc_id=os.environ.get('ALB_VPC_ID'), + listener_port=20000, + ) + print(f"Success! Result: {result}") + result = cleanup_alb_resources() + print(f"Cleanup result: {result}") + except Exception as e: + print(f"Error: {e}") + sys.exit(1) + + +# --------------------------------------------------------------------------- +# ALBProxyRouter — ProxyRouter implementation wrapping the functions above +# --------------------------------------------------------------------------- + +from .base import ProxyRouter + + +class ALBProxyRouter(ProxyRouter): + """Alibaba Cloud ALB router implementing the ProxyRouter interface. + + Wraps the standalone ALB registration functions so that callers can switch + between ALB and K8s backends through a common interface. + + Environment variables consumed (same as the underlying functions): + TASK_ID, ALB_VPC_ID, ALB_REGION, ALB_LOAD_BALANCER_IDS, + ALB_LOAD_BALANCER_LSN_IDS, ALB_RESOURCE_CLEAN_DIR, + ALIBABA_CLOUD_ACCESS_KEY_ID, ALIBABA_CLOUD_ACCESS_KEY_SECRET + """ + + def register_servers(self, addresses: List[str], job_id: str, port: int) -> dict: + """Register rollout servers to ALB via server-group + listener rule.""" + os.environ['TASK_ID'] = job_id + return register_rollout_servers_to_alb(addresses=addresses, listener_port=port) + + def cleanup_resources(self, job_id: str) -> dict: + """Delete ALB server-group and rules recorded under job_id's resource dir.""" + os.environ['TASK_ID'] = job_id + result = cleanup_alb_resources() + return { + "cleaned_up": result.get("cleaned_up", 0), + "errors": result.get("errors", []), + } + + def list_active_job_ids(self) -> List[str]: + """Return job IDs parsed from active ALB listener rule names. + + Rule names follow the pattern ``rule-{job_id}-{port}``, so the job_id + is extracted as everything between the first and last '-'-delimited token. + """ + rules_dict = list_rules().get('rules', {}) + job_ids = [] + for rule_name in rules_dict: + # rule_name format: "rule-{job_id}-{port}" + parts = rule_name.split('-') + if len(parts) >= 3: + # job_id may itself contain '-', so take everything between first and last token + job_ids.append('-'.join(parts[1:-1])) + return job_ids + + def get_callback_url(self, job_id: str, port: int) -> str | None: + """Build the ALB callback URL for agents to reach the rollout server. + + Reads EP_ENDPOINT, ALB_LISTEN_PORT, ALB_REGION from the environment. + Returns None if any required variable is missing. + """ + ep_endpoint = os.environ.get('EP_ENDPOINT') + alb_listen_port = os.environ.get('ALB_LISTEN_PORT') + alb_region = os.environ.get('ALB_REGION') + if not all([ep_endpoint, alb_listen_port, alb_region]): + return None + alb_id, _ = get_alb_load_balancer_id_by_hash(job_id, port) + return f"http://{ep_endpoint}:{alb_listen_port}/{alb_region}/{alb_id}/{job_id}/{port}/v1" diff --git a/roll/pipeline/agentic/proxy/router/base.py b/roll/pipeline/agentic/proxy/router/base.py new file mode 100644 index 000000000..0f9ca9a06 --- /dev/null +++ b/roll/pipeline/agentic/proxy/router/base.py @@ -0,0 +1,106 @@ +import os +from abc import ABC, abstractmethod +from typing import List, Optional + +from roll.utils.logging import get_logger + +logger = get_logger() + + +def batch_cleanup(backend: "ProxyRouter") -> dict: + """Clean up resources for all jobs registered in the backend. + + Args: + backend: Any ProxyRouter implementation (ALB or Ingress) + + Returns: + Dict with ``cleaned_up`` (int), ``errors`` (list), and + ``checked_jobs`` (list of job IDs that were evaluated) + """ + job_ids = list(dict.fromkeys(backend.list_active_job_ids())) # deduplicate, preserve order + logger.info(f"Found {len(job_ids)} job(s) to clean up: {job_ids}") + + cleaned, errors, checked = 0, [], [] + + for job_id in job_ids: + checked.append(job_id) + logger.info(f"Cleaning up resources for job {job_id}") + result = backend.cleanup_resources(job_id) + cleaned += result.get('cleaned_up', 0) + errors.extend(result.get('errors', [])) + + return {'cleaned_up': cleaned, 'errors': errors, 'checked_jobs': checked} + + +class ProxyRouter(ABC): + """Abstract base class for proxy router registration. + + Implementations register rollout server IP:port backends to a load balancer + (ALB, Ingress, etc.) and clean them up when the job finishes. + """ + + @abstractmethod + def register_servers(self, addresses: List[str], job_id: str, port: int) -> dict: + """Register rollout servers to the load balancer backend. + + Args: + addresses: Server addresses, e.g. ["http://10.0.0.1:8080"] + job_id: Unique job identifier used for resource naming and routing + port: Backend port number + + Returns: + Dict describing created resources + """ + ... + + @abstractmethod + def cleanup_resources(self, job_id: str) -> dict: + """Clean up all backend resources created for a job. + + Args: + job_id: Job identifier whose resources should be removed + + Returns: + Dict with keys ``cleaned_up`` (int) and ``errors`` (list) + """ + ... + + @abstractmethod + def list_active_job_ids(self) -> List[str]: + """Return job IDs that currently have registered resources in this backend. + + Used by batch cleanup to discover which jobs still have live resources, + so their status can be checked and stale ones cleaned up. + + Returns: + List of job ID strings (may contain duplicates if multiple ports + are registered per job; callers should deduplicate if needed) + """ + ... + + @abstractmethod + def get_callback_url(self, job_id: str, port: int) -> Optional[str]: + """Return the URL that agents should use to reach the rollout server. + + Args: + job_id: Job identifier whose resources are already registered + port: Backend port number + + Returns: + Full callback URL string, or None if the URL cannot be determined + """ + ... + + +def create_proxy_router() -> "ProxyRouter": + """Instantiate the proxy router selected by the PROXY_ROUTER_TYPE env var. + + Supported values: ``ingress`` (Kubernetes nginx Ingress), ``alb`` (default, Alibaba Cloud ALB). + """ + if os.environ.get('PROXY_ROUTER_TYPE') == 'ingress': + from .ingress_proxy_router import IngressProxyRouter + logger.info("[ProxyRouter] Using Ingress router") + return IngressProxyRouter() + logger.info("[ProxyRouter] Using ALB router") + from .alb_proxy_router import ALBProxyRouter + return ALBProxyRouter() diff --git a/roll/pipeline/agentic/proxy/router/cleanup_alb_resources_batch.py b/roll/pipeline/agentic/proxy/router/cleanup_alb_resources_batch.py new file mode 100644 index 000000000..f328ee979 --- /dev/null +++ b/roll/pipeline/agentic/proxy/router/cleanup_alb_resources_batch.py @@ -0,0 +1,124 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +Batch cleanup script for ALB resources. + +This script lists all ALB rules, checks if the XDL job ID in each rule name +is in a cleanable state, and deletes the rule if it is. +""" + +import os + +from .alb_proxy_router import ( + ALBProxyRouter, + create_alb_client, + delete_resource_with_retry, +) +from .base import batch_cleanup + + +def batch_cleanup_alb_server_groups(): + """List all ALB server groups and delete them directly without validation.""" + print("Listing all ALB server groups...") + + try: + resource_group_id = os.environ.get('ALB_RESOURCE_GROUP_ID') + assert resource_group_id, "ALB_RESOURCE_GROUP_ID is not set" + server_groups_response = list_server_groups(resource_group_id) + server_groups_dict = server_groups_response.get('server_groups', {}) + + if not server_groups_dict: + print("No server groups found") + return + + print(f"Found {len(server_groups_dict)} server groups") + client = create_alb_client() + deleted_count = 0 + + for sg_name, sg_id in server_groups_dict.items(): + print(f"\nDeleting server group: {sg_name} ({sg_id})") + try: + success, error_msg = delete_resource_with_retry(client, 'server-group', sg_id, 5) + if success: + print(f"Successfully deleted server group: {sg_name}") + deleted_count += 1 + else: + print(f"Failed to delete server group {sg_name}: {error_msg}") + except Exception as e: + print(f"Error deleting server group {sg_name}: {e}") + continue + + print(f"\nFinished processing server groups. Deleted {deleted_count} server groups.") + + except Exception as e: + print(f"Error listing or processing server groups: {e}") + return + + +def list_server_groups(resource_group_id: str = None) -> dict: + """ + List all Server Groups and return a mapping of server_group_name -> server_group_id. + + Only includes server groups not attached to any ALB. + """ + from alibabacloud_alb20200616 import models as alb_20200616_models + from alibabacloud_tea_util import models as util_models + + try: + client = create_alb_client() + max_results = 100 + all_server_groups = [] + next_token = None + + while True: + request = alb_20200616_models.ListServerGroupsRequest( + max_results=max_results, + resource_group_id=resource_group_id + ) + if next_token: + request.next_token = next_token + + runtime = util_models.RuntimeOptions() + resp = client.list_server_groups_with_options(request, runtime) + + if hasattr(resp.body, 'server_groups') and resp.body.server_groups: + all_server_groups.extend(resp.body.server_groups) + + next_token = resp.body.next_token if hasattr(resp.body, 'next_token') else None + if not next_token: + break + + name_to_id = {} + print(f"Found {len(all_server_groups)} server groups") + for sg in all_server_groups: + print(f"Server group: {sg.server_group_id} {sg.server_group_name} {sg.related_load_balancer_ids}") + if sg.server_group_name and sg.server_group_id and sg.related_load_balancer_ids is None: + name_to_id[sg.server_group_name] = sg.server_group_id + + print(f"Found {len(name_to_id)} server groups not related to any ALB") + return {"server_groups": name_to_id} + + except Exception as error: + print(f"Error listing server groups: {error}") + if hasattr(error, 'message'): + print(f"Error message: {error.message}") + if hasattr(error, 'data') and error.data: + print(f"Recommendation: {error.data.get('Recommend')}") + return {"server_groups": {}} + + +def main(): + # Set environment variables for testing (you need to provide your own values) + os.environ.setdefault('ALB_LOAD_BALANCER_LSN_IDS', '') + os.environ.setdefault('ALB_ENDPOINT', '') + os.environ.setdefault('ALIBABA_CLOUD_ACCESS_KEY_ID', '') + os.environ.setdefault('ALIBABA_CLOUD_ACCESS_KEY_SECRET', '') + os.environ.setdefault('ALB_RESOURCE_GROUP_ID', '') + print("Starting batch ALB resource cleanup...") + batch_cleanup(ALBProxyRouter()) + batch_cleanup_alb_server_groups() + print("Batch cleanup completed") + + +if __name__ == "__main__": + main() diff --git a/roll/pipeline/agentic/proxy/router/ingress_proxy_router.py b/roll/pipeline/agentic/proxy/router/ingress_proxy_router.py new file mode 100644 index 000000000..ab9cf446f --- /dev/null +++ b/roll/pipeline/agentic/proxy/router/ingress_proxy_router.py @@ -0,0 +1,372 @@ +""" +Kubernetes Service + Ingress backend for rollout server registration. + +Each call to register_servers creates three resources in the ``xdl`` namespace: + +* Headless Service — provides a stable DNS name for the backend +* Endpoints — manually points to the rollout worker Pod IPs +* Ingress — nginx path rule ``/{job_id}/{port}(/|$)(.*)`` + rewritten to ``/$2`` and forwarded to the Service + +Cleanup is label-based: all resources share the label ``roll.ai/job-id`` +so they can be found and deleted without tracking instance state. +""" +import hashlib +import logging +import os +import re +from typing import List, Optional, Tuple + +from kubernetes import client, config +from kubernetes.client.rest import ApiException + +from .base import ProxyRouter +from roll.utils.logging import get_logger + +logger = get_logger() + +NAMESPACE = "xdl" +LABEL_MANAGED_BY = "app.kubernetes.io/managed-by" +LABEL_JOB_ID = "roll.ai/job-id" +LABEL_PORT = "roll.ai/port" + + +def _safe_k8s_name(job_id: str, port: int) -> str: + """Return a DNS-label-safe K8s resource name (max 63 chars).""" + raw = f"roll-{job_id}-{port}" + if len(raw) <= 63: + return raw + # Truncate and append a short hash to keep uniqueness + h = hashlib.sha256(f"{job_id}:{port}".encode()).hexdigest()[:8] + return f"roll-{job_id}"[:50] + f"-{h}" + + +def _parse_address(address: str) -> Tuple[Optional[str], Optional[int]]: + match = re.match(r'https?://([^:/]+):(\d+)', address) + if match: + return match.group(1), int(match.group(2)) + return None, None + + +class IngressProxyRouter(ProxyRouter): + """Kubernetes Service + nginx Ingress backend. + + Args: + namespace: Kubernetes namespace (default: ``roll``). Created automatically + if it does not exist. + ingress_class_name: Ingress controller class (default: ``nginx``). + Written to ``spec.ingressClassName`` as recommended since K8s 1.18. + tls_secret_name: Name of the TLS Secret for HTTPS. When set, ``spec.tls`` + is populated so the Ingress terminates TLS. Leave ``None`` for plain HTTP. + """ + + def __init__( + self, + namespace: str = NAMESPACE, + ingress_class_name: str = "nginx", + tls_secret_name: str = None, + ): + self.namespace = namespace + self.ingress_host = os.environ.get('EP_ENDPOINT') + self.ingress_class_name = ingress_class_name + self.tls_secret_name = tls_secret_name + + config.load_incluster_config() + # The cluster's API server certificate may not cover the internal service IP, + # so disable SSL verification for in-cluster communication. + configuration = client.Configuration.get_default_copy() + configuration.verify_ssl = False + client.Configuration.set_default(configuration) + self._core = client.CoreV1Api() + self._networking = client.NetworkingV1Api() + + self._ensure_namespace() + + # ------------------------------------------------------------------ + # ProxyRouter interface + # ------------------------------------------------------------------ + + def register_servers(self, addresses: List[str], job_id: str, port: int) -> dict: + """Create Service, Endpoints, and Ingress for the given rollout workers. + + Args: + addresses: Worker addresses, e.g. ``["http://10.0.0.1:8080"]`` + job_id: Unique job identifier used in resource names and Ingress path + port: Backend port number + + Returns: + Dict with ``service``, ``ingress``, ``namespace``, and ``registered_servers`` + """ + name = _safe_k8s_name(job_id, port) + labels = { + LABEL_MANAGED_BY: "roll", + LABEL_JOB_ID: job_id, + LABEL_PORT: str(port), + } + self._ensure_service(name, labels, port) + self._ensure_endpoints(name, labels, addresses, port) + self._ensure_ingress(name, labels, job_id, port) + logger.info("Registered %d servers: service/ingress=%s namespace=%s", len(addresses), name, self.namespace) + return { + "service": name, + "ingress": name, + "namespace": self.namespace, + "registered_servers": addresses, + } + + def list_active_job_ids(self) -> List[str]: + """Return job IDs read from the ``roll.ai/job-id`` label on active Ingress objects. + + Lists all Ingress resources managed by roll in this namespace and + extracts unique job IDs from their labels. + """ + label_selector = f"{LABEL_MANAGED_BY}=roll" + try: + items = self._networking.list_namespaced_ingress( + self.namespace, label_selector=label_selector + ).items + except ApiException as e: + logger.error("Failed to list Ingress objects for active job IDs: %s", e) + return [] + + job_ids = [] + for item in items: + job_id = (item.metadata.labels or {}).get(LABEL_JOB_ID) + if job_id: + job_ids.append(job_id) + return job_ids + + def get_callback_url(self, job_id: str, port: int) -> Optional[str]: + """Build the K8s Ingress callback URL for agents to reach the rollout server.""" + if not self.ingress_host: + return None + return f"http://{self.ingress_host}/{job_id}/{port}/v1" + + def cleanup_resources(self, job_id: str) -> dict: + """Delete all Ingress, Endpoints, and Service objects labelled with job_id. + + Deletion order: Ingress → Endpoints → Service, so routing rules are + removed before the backend disappears. + + Args: + job_id: Job whose resources should be removed + + Returns: + Dict with ``cleaned_up`` count and ``errors`` list + """ + label_selector = f"{LABEL_JOB_ID}={job_id}" + cleaned, errors = 0, [] + + # (list_fn, delete_fn, kind) — order matters + steps = [ + ( + lambda: self._networking.list_namespaced_ingress(self.namespace, label_selector=label_selector), + lambda n: self._networking.delete_namespaced_ingress(n, self.namespace), + "Ingress", + ), + ( + lambda: self._core.list_namespaced_endpoints(self.namespace, label_selector=label_selector), + lambda n: self._core.delete_namespaced_endpoints(n, self.namespace), + "Endpoints", + ), + ( + lambda: self._core.list_namespaced_service(self.namespace, label_selector=label_selector), + lambda n: self._core.delete_namespaced_service(n, self.namespace), + "Service", + ), + ] + + for list_fn, delete_fn, kind in steps: + try: + for item in list_fn().items: + name = item.metadata.name + try: + delete_fn(name) + logger.info("Deleted %s: %s", kind, name) + cleaned += 1 + except ApiException as e: + if e.status == 404: + logger.info("%s/%s already gone", kind, name) + else: + errors.append(f"{kind}/{name}: {e}") + except ApiException as e: + errors.append(f"List {kind}: {e}") + + return {"cleaned_up": cleaned, "errors": errors} + + # ------------------------------------------------------------------ + # Private helpers + # ------------------------------------------------------------------ + + def _ensure_namespace(self) -> None: + """Create the namespace if it does not already exist. + + Namespace is a cluster-scoped resource. If the ServiceAccount only has + namespace-scoped permissions (Role/RoleBinding), read_namespace and + create_namespace will raise 403. In that case we log a warning and + continue — the namespace must have been pre-created by an admin. + """ + try: + self._core.read_namespace(self.namespace) + except ApiException as e: + if e.status == 404: + try: + ns = client.V1Namespace( + metadata=client.V1ObjectMeta( + name=self.namespace, + labels={LABEL_MANAGED_BY: "roll"}, + ) + ) + self._core.create_namespace(ns) + logger.info("Created namespace: %s", self.namespace) + except ApiException as ce: + if ce.status == 403: + logger.warning( + "No permission to create namespace %s (ClusterRole required). " + "Assuming it was pre-created by an admin.", + self.namespace, + ) + else: + raise + elif e.status == 403: + logger.warning( + "No permission to read namespace %s (ClusterRole required). " + "Assuming it exists and was pre-created by an admin.", + self.namespace, + ) + else: + raise + + def _ensure_service(self, name: str, labels: dict, port: int) -> None: + svc = client.V1Service( + metadata=client.V1ObjectMeta(name=name, namespace=self.namespace, labels=labels), + spec=client.V1ServiceSpec( + cluster_ip="None", # headless — no kube-proxy load balancing + ports=[client.V1ServicePort(port=port, protocol="TCP")], + ), + ) + try: + self._core.create_namespaced_service(self.namespace, svc) + logger.info("Created Service: %s", name) + except ApiException as e: + if e.status == 409: + logger.warning("Service %s already exists, skipping", name) + else: + raise + + def _ensure_endpoints(self, name: str, labels: dict, addresses: List[str], port: int) -> None: + endpoint_addrs = [] + for addr in addresses: + ip, _ = _parse_address(addr) + if ip: + endpoint_addrs.append(client.V1EndpointAddress(ip=ip)) + else: + logger.warning("Cannot parse address %s, skipping", addr) + + if not endpoint_addrs: + raise ValueError(f"No valid addresses parsed from: {addresses}") + + ep = client.V1Endpoints( + metadata=client.V1ObjectMeta(name=name, namespace=self.namespace, labels=labels), + subsets=[ + client.V1EndpointSubset( + addresses=endpoint_addrs, + ports=[client.CoreV1EndpointPort(port=port, protocol="TCP")], + ) + ], + ) + try: + self._core.create_namespaced_endpoints(self.namespace, ep) + logger.info("Created Endpoints: %s (%d addresses)", name, len(endpoint_addrs)) + except ApiException as e: + if e.status == 409: + logger.warning("Endpoints %s already exists, patching", name) + self._core.patch_namespaced_endpoints(name, self.namespace, ep) + else: + raise + + def _ensure_ingress(self, name: str, labels: dict, job_id: str, port: int) -> None: + # Escape job_id so special regex chars (e.g. '.') don't break the pattern. + # K8s names only contain [a-z0-9-] but the raw job_id passed by the caller + # may not have been sanitised yet. + safe_job_id = re.escape(job_id) + # Capture group $2 holds the remainder after /{job_id}/{port}/ + path_pattern = f"/{safe_job_id}/{port}(/|$)(.*)" + + tls = ( + [client.V1IngressTLS(hosts=[self.ingress_host] if self.ingress_host else [], secret_name=self.tls_secret_name)] + if self.tls_secret_name + else None + ) + + ingress = client.V1Ingress( + metadata=client.V1ObjectMeta( + name=name, + namespace=self.namespace, + labels=labels, + annotations={ + "nginx.ingress.kubernetes.io/use-regex": "true", + "nginx.ingress.kubernetes.io/rewrite-target": "/$2", + }, + ), + spec=client.V1IngressSpec( + ingress_class_name=self.ingress_class_name, + tls=tls, + rules=[ + client.V1IngressRule( + host=self.ingress_host, + http=client.V1HTTPIngressRuleValue( + paths=[ + client.V1HTTPIngressPath( + path=path_pattern, + path_type="ImplementationSpecific", + backend=client.V1IngressBackend( + service=client.V1IngressServiceBackend( + name=name, + port=client.V1ServiceBackendPort(number=port), + ) + ), + ) + ] + ), + ) + ], + ), + ) + try: + self._networking.create_namespaced_ingress(self.namespace, ingress) + logger.info("Created Ingress: %s host=%s path=%s", name, self.ingress_host, path_pattern) + except ApiException as e: + if e.status == 409: + logger.warning("Ingress %s already exists, patching", name) + self._networking.patch_namespaced_ingress(name, self.namespace, ingress) + else: + raise + + +if __name__ == "__main__": + import sys + import socket + + os.environ.setdefault('EP_ENDPOINT', '') + os.environ.setdefault('TASK_ID', 'ingress-test-123') + + # 用本机 IP 作为测试 Endpoints 地址 + hostname = socket.gethostname() + local_ip = socket.gethostbyname(hostname) + test_addresses = [f"http://{local_ip}:19876"] + job_id = os.environ['TASK_ID'] + + print(f"Testing IngressProxyRouter: addresses={test_addresses} job_id={job_id}") + + try: + backend = IngressProxyRouter() + result = backend.register_servers(addresses=test_addresses, job_id=job_id, port=19876) + print(f"register_servers result: {result}") + url = backend.get_callback_url(job_id, 19876) + print(f"callback_url: {url}") + cleanup = backend.cleanup_resources(job_id) + print(f"cleanup result: {cleanup}") + except Exception as e: + print(f"Error: {e}") + sys.exit(1) diff --git a/roll/pipeline/agentic/utils.py b/roll/pipeline/agentic/utils.py index 820327acf..95e96f967 100644 --- a/roll/pipeline/agentic/utils.py +++ b/roll/pipeline/agentic/utils.py @@ -476,21 +476,49 @@ def agentic_compute_advantage( # Check OPD config is_pure_opd = getattr(pipeline_config, "is_pure_opd", False) if pipeline_config else False use_opd = getattr(pipeline_config, "use_opd", False) if pipeline_config else False - opd_kl_coef = getattr(pipeline_config, "opd_kl_coef", 1.0) if pipeline_config else 1.0 + reference_configs = getattr(pipeline_config, "reference_configs", {"default": None}) if pipeline_config else {"default": None} - # Compute KL divergence for OPD modes - kld = None if is_pure_opd or use_opd: - kld = compute_approx_kl( - log_probs=data.batch["old_log_probs"] if getattr(pipeline_config, "enable_old_logprobs_recompute", False) else data.batch["infer_logprobs"], - log_probs_base=data.batch["ref_log_probs"], - action_mask=response_mask, - kl_penalty=getattr(pipeline_config, "kl_penalty", "kl"), - ) - - # For pure OPD mode, advantage is directly -kld + log_probs_key = "old_log_probs" if getattr(pipeline_config, "enable_old_logprobs_recompute", False) else "infer_logprobs" + kl_penalty = getattr(pipeline_config, "kl_penalty", "kl") + total_weighted_kld = torch.zeros_like(response_mask, dtype=torch.float32) + for name, ref_cfg in reference_configs.items(): + ref_key = f"ref_log_probs_{name}" + mask_key = f"ref_log_probs_{name}_mask" + + # Check if this teacher has a routing mask (set by pipeline when needs_teacher_routing) + if mask_key in data.batch: + teacher_mask = data.batch[mask_key] # (batch_size,), bool + else: + # No mask = this teacher handles all tags (backward compatible) + teacher_mask = None + + kl_coef_i = ref_cfg.opd_kl_coef if (ref_cfg and ref_cfg.opd_kl_coef is not None) else 1.0 + + if teacher_mask is not None and teacher_mask.sum() < teacher_mask.numel(): + # Only compute KL on routed subset to avoid meaningless computation against zeros + subset_idx = teacher_mask.nonzero(as_tuple=True)[0] + kld_subset = compute_approx_kl( + log_probs=data.batch[log_probs_key][subset_idx], + log_probs_base=data.batch[ref_key][subset_idx], + action_mask=response_mask[subset_idx], + kl_penalty=kl_penalty, + ) + kld_i = torch.zeros_like(response_mask, dtype=torch.float32) + kld_i[subset_idx] = kld_subset + total_weighted_kld += kl_coef_i * kld_i + else: + kld_i = compute_approx_kl( + log_probs=data.batch[log_probs_key], + log_probs_base=data.batch[ref_key], + action_mask=response_mask, + kl_penalty=kl_penalty, + ) + total_weighted_kld += kl_coef_i * kld_i + + # For pure OPD mode, advantage is directly -total_weighted_kld if is_pure_opd: - advantages = -kld + advantages = -total_weighted_kld returns = advantages data.batch["raw_advantages"] = advantages else: @@ -518,9 +546,9 @@ def agentic_compute_advantage( data.batch["raw_advantages"] = advantages - # Apply mixed OPD mode + # Apply mixed OPD mode: advantages = rl_advantages - total_weighted_kld if use_opd: - advantages = advantages - opd_kl_coef * kld + advantages = advantages - total_weighted_kld if whiten_advantages: # TODO whiten过程中是否要考虑response的长度？ diff --git a/roll/pipeline/base_pipeline.py b/roll/pipeline/base_pipeline.py index 5c4d67e78..c2bbb97ea 100644 --- a/roll/pipeline/base_pipeline.py +++ b/roll/pipeline/base_pipeline.py @@ -14,9 +14,11 @@ from roll.distributed.executor.model_update_group import ModelUpdateGroup from roll.distributed.scheduler.protocol import DataProto from roll.distributed.scheduler.resource_manager import ResourceManager -from roll.utils.checkpoint_manager import CheckpointManager, download_model +from roll.distributed.scheduler.transfer_backend import init_transfer_backend +from roll.utils.checkpoint_manager import CheckpointManager, download_model, get_latest_ckpt from roll.utils.functionals import reduce_metrics from roll.utils.logging import get_logger +from roll.utils.telemetry import init_telemetry, shutdown_telemetry, start_otel_collector from roll.utils.tracking import create_tracker from roll.utils.worker_state import WorkerState @@ -40,13 +42,36 @@ def __init__(self, pipeline_config): config=self.pipeline_config.to_dict(), **self.pipeline_config.tracker_kwargs, ) + + # Initialize OpenTelemetry tracing on driver if enabled + if os.environ.get("ROLL_OTEL_ENABLED") == "1": + otlp_endpoint = os.environ["OTEL_EXPORTER_OTLP_ENDPOINT"] + start_otel_collector( + endpoint=otlp_endpoint, + output_dir=self.pipeline_config.otlp_output_dir, + ) + init_telemetry( + service_name="driver", + otlp_endpoint=otlp_endpoint, + ) + + init_transfer_backend(pipeline_config.transfer_backend) + self.resume_from_checkpoint = False self.executor: futures.ThreadPoolExecutor = futures.ThreadPoolExecutor(max_workers=5) self.resume_futures = [] if self.pipeline_config.resume_from_checkpoint: - self.resume_from_checkpoint = download_model(self.pipeline_config.resume_from_checkpoint) - + if self.pipeline_config.resume_from_checkpoint is True: + latest_ckpt = get_latest_ckpt(self.pipeline_config.checkpoint_config) + if latest_ckpt: + self.resume_from_checkpoint = download_model(latest_ckpt) + else: + self.resume_from_checkpoint = False + elif isinstance(self.pipeline_config.resume_from_checkpoint, str): + self.resume_from_checkpoint = download_model(self.pipeline_config.resume_from_checkpoint) + + if self.resume_from_checkpoint: logger.info(f"resume_from_checkpoint: {self.resume_from_checkpoint}") load_dir = os.path.join(self.resume_from_checkpoint, "pipeline") self.state = WorkerState.load_from_json(load_dir=load_dir, tag="pipeline") @@ -99,7 +124,11 @@ def do_checkpoint(self, global_step, is_last_step=None): save_dir = os.path.join(self.pipeline_config.output_dir, "pipeline", ckpt_id, "pipeline") self.state.save_to_json(save_dir=save_dir, tag="pipeline") self.state.save_rng_state(save_dir=save_dir, tag="pipeline") - self.checkpoint_manager.upload(ckpt_id=ckpt_id, local_state_path=pipeline_save_dir) + + if self.pipeline_config.checkpoint_config.get("async_upload", True) and not is_last_step: + self.executor.submit(self.checkpoint_manager.upload, ckpt_id=ckpt_id, local_state_path=pipeline_save_dir) + else: + self.checkpoint_manager.upload(ckpt_id=ckpt_id, local_state_path=pipeline_save_dir) # Clean up old checkpoints if max_ckpt_to_keep is set self._cleanup_old_checkpoints() diff --git a/roll/pipeline/base_worker.py b/roll/pipeline/base_worker.py index ccc69ec09..3e522bea7 100644 --- a/roll/pipeline/base_worker.py +++ b/roll/pipeline/base_worker.py @@ -17,12 +17,11 @@ from roll.distributed.strategy.strategy import InferenceStrategy, TrainStrategy from roll.models.model_providers import ( default_actor_model_provider, - default_diffusion_module_provider, default_reward_model_provider, default_value_model_provider, ) from roll.platforms import current_platform -from roll.utils.checkpoint_manager import download_model +from roll.utils.checkpoint_manager import download_model, get_latest_ckpt from roll.utils.context_managers import state_offload_manger, log_gpu_memory_usage from roll.utils.dynamic_batching import make_mini_batch_iter_for_dynamic_batching from roll.utils.functionals import agg_loss, append_to_dict, compute_approx_kl, flatten_sum, masked_mean, postprocess_generate, reduce_metrics @@ -43,20 +42,24 @@ def initialize(self, pipeline_config): self.strategy = create_strategy(worker=self) - if self.worker_config.model_args.model_type == "diffusion_module": - self.strategy.initialize(model_provider=default_diffusion_module_provider) - else: - self.strategy.initialize(model_provider=default_actor_model_provider) + self.strategy.initialize(model_provider=default_actor_model_provider) self.tokenizer = self.strategy.tokenizer if self.pipeline_config.resume_from_checkpoint: - load_dir = download_model(self.pipeline_config.resume_from_checkpoint) - self.strategy.load_checkpoint(load_dir=load_dir, tag="checkpoint") + load_dir = None + if self.pipeline_config.resume_from_checkpoint is True: + latest_ckpt = get_latest_ckpt(self.pipeline_config.checkpoint_config) + if latest_ckpt: + load_dir = download_model(latest_ckpt) + elif isinstance(self.pipeline_config.resume_from_checkpoint, str): + load_dir = download_model(self.pipeline_config.resume_from_checkpoint) + if load_dir: + self.strategy.load_checkpoint(load_dir=load_dir, tag="checkpoint") self.logger.info(f"{self.worker_name} initialized") self.strategy.offload_states() - @register(dispatch_mode=Dispatch.DP_MP_DISPATCH_FIRST) + @register(dispatch_mode=Dispatch.DP_MP_DISPATCH_FIRST, trace=True, prefetch=False, to_remote=False) def train_step(self, data: DataProto): """ return DataProto(meta_info={'metrics': metrics}) @@ -75,6 +78,8 @@ def train_step(self, data: DataProto): ): data = data.to(current_platform.device_type) data = self.strategy.get_data_input(data) + data.prefetch() + data = data.to(current_platform.device_type) per_device_train_batch_size = self.worker_config.training_args.per_device_train_batch_size backward_batch_size = ( per_device_train_batch_size * self.worker_config.training_args.gradient_accumulation_steps @@ -121,7 +126,7 @@ def train_step(self, data: DataProto): output = DataProto(meta_info={"metrics": metrics}) return output - @register(dispatch_mode=Dispatch.DP_MP_DISPATCH_FIRST) + @register(dispatch_mode=Dispatch.DP_MP_DISPATCH_FIRST, trace=True, prefetch=False, to_remote=True) def compute_log_probs(self, data: DataProto): """ return DataProto.from_dict(tensors={'log_probs': output}) @@ -137,6 +142,7 @@ def compute_log_probs(self, data: DataProto): load_kwargs={"include": [OffloadStateType.model_params]}, ): data = self.strategy.get_data_input(data) + data.prefetch() data = data.to(current_platform.device_type) data.meta_info["micro_batch_size"] = self.worker_config.infer_batch_size @@ -501,7 +507,7 @@ async def generate(self, data: DataProto): global_step = data.meta_info.get("global_step", 0) self.logger.info(f"{self.worker_name} generate global step {global_step}") - data = data.to(current_platform.device_type) + data = data.to("cuda") data.meta_info["micro_batch_size"] = self.worker_config.infer_batch_size output = await self.strategy.generate(batch=data, generation_config=generation_config) @@ -555,14 +561,23 @@ def initialize(self, pipeline_config): self.tokenizer = self.strategy.tokenizer if self.pipeline_config.resume_from_checkpoint: - load_dir = os.path.join(download_model(self.pipeline_config.resume_from_checkpoint), self.cluster_name) - self.strategy.load_checkpoint(load_dir=load_dir, tag="checkpoint") + load_dir = None + if self.pipeline_config.resume_from_checkpoint is True: + latest_ckpt = get_latest_ckpt(self.pipeline_config.checkpoint_config) + if latest_ckpt: + load_dir = download_model(latest_ckpt) + elif isinstance(self.pipeline_config.resume_from_checkpoint, str): + load_dir = download_model(self.pipeline_config.resume_from_checkpoint) + load_dir = os.path.join(load_dir, self.cluster_name) + + if load_dir: + self.strategy.load_checkpoint(load_dir=load_dir, tag="checkpoint") self.logger.info(f"{self.worker_name} initialized") self.strategy.offload_states() - @register(dispatch_mode=Dispatch.DP_MP_COMPUTE) + @register(dispatch_mode=Dispatch.DP_MP_COMPUTE, trace=True, prefetch=True, to_remote=True) def compute_values(self, data: DataProto): """ return DataProto.from_dict(tensors={'values': values}) @@ -591,7 +606,7 @@ def compute_values(self, data: DataProto): output.meta_info = {"metrics": metrics} return output - @register(dispatch_mode=Dispatch.DP_MP_COMPUTE) + @register(dispatch_mode=Dispatch.DP_MP_COMPUTE, trace=True, prefetch=True, to_remote=True) def train_step(self, data: DataProto): """ return DataProto(meta_info={'metrics': metrics}) diff --git a/roll/pipeline/diffusion/__init__.py b/roll/pipeline/diffusion/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/roll/pipeline/diffusion/modules/__init__.py b/roll/pipeline/diffusion/modules/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/roll/pipeline/diffusion/modules/wan_module.py b/roll/pipeline/diffusion/modules/wan_module.py deleted file mode 100644 index a80198a01..000000000 --- a/roll/pipeline/diffusion/modules/wan_module.py +++ /dev/null @@ -1,554 +0,0 @@ -import numpy as np -import torch -import types -import json -import gc -import os -import imageio -import queue - -from concurrent.futures import ThreadPoolExecutor -from torchvision.io import write_video -from typing import List, Optional -from datetime import datetime -from einops import reduce -from peft import LoraConfig, inject_adapter_in_model, get_peft_model - -from diffsynth.pipelines.wan_video_new import WanVideoPipeline, TeaCache, TemporalTiler_BCTHW -from diffsynth.models.wan_video_motion_controller import WanMotionControllerModel -from diffsynth.models.wan_video_vace import VaceWanModel - -from diffsynth.trainers.utils import DiffusionTrainingModule -from diffsynth.utils import ModelConfig, PipelineUnit -from diffsynth.models.wan_video_dit import WanModel, sinusoidal_embedding_1d - -from roll.pipeline.diffusion.reward_fl.face_tools import FaceAnalysis, Face -from roll.pipeline.diffusion.reward_fl.wan_video_vae import WanVideoVAE -from roll.pipeline.diffusion.reward_fl.euler import EulerScheduler -from roll.platforms import current_platform - - -def vae_output_to_videotensor(vae_output, pattern="B C T H W", min_value=-1, max_value=1): - # process vae_output to videotensor - if pattern != "T H W C": - vae_output = reduce(vae_output, f"{pattern} -> T H W C", reduction="mean") - video_tensor = ((vae_output - min_value) * (255 / (max_value - min_value))).clip(0, 255) - return video_tensor - - -def training_loss(self, **inputs): - self.scheduler.set_timesteps(num_inference_steps=self.num_inference_steps) - - inputs["latents"] = self.generate_noise(inputs["latents"].shape, seed=24, rand_device=self.device) - - timesteps = self.scheduler.timesteps - models = {name: getattr(self, name) for name in self.in_iteration_models} - kl_loss_total = 0 - kl_steps = 0 - for i, timestep in enumerate(timesteps[:]): - current_timestep = timestep.expand([1]) - current_sigma = current_timestep / self.scheduler.num_train_timesteps - # switch dit if necessary - if timestep.item() < 0.9 * self.scheduler.num_train_timesteps and self.dit2 is not None and not models["dit"] is self.dit2: - models["dit"] = self.dit2 - elif timestep.item() >= 0.9 * self.scheduler.num_train_timesteps and self.dit is not None and not models["dit"] is self.dit: - models["dit"] = self.dit - - # Timestep - timestep = timestep.unsqueeze(0).to(dtype=torch.float32, device=self.device) - inputs["timestep"] = timestep - inputs["less_mid_step"] = True - - # Inference - if i < self.mid_timestep: - inputs["less_mid_step"] = True - with torch.no_grad(): - model_pred = self.model_fn(**models, **inputs) - else: - inputs["less_mid_step"] = False - model_pred = self.model_fn(**models, **inputs) - with torch.no_grad(): - inputs["less_mid_step"] = True - models["dit"].disable_adapter_layers() - model_pred_no_lora = self.model_fn(**models, **inputs) - models["dit"].enable_adapter_layers() - kl_loss_step = torch.mean( - (model_pred.float() - model_pred_no_lora.float()) ** 2 / (2 * current_sigma.float() ** 2) - ) - kl_loss_total += kl_loss_step - kl_steps += 1 - - noise_pred = model_pred - - # Scheduler denoise - if i < self.final_timestep: - inputs["latents"] = self.scheduler.step(noise_pred, timestep, inputs["latents"]).to(torch.bfloat16) - else: - inputs["latents"] = self.scheduler.step(noise_pred, timestep, inputs["latents"]).to(torch.bfloat16) - break - - if "first_frame_latents" in inputs: - inputs["latents"][:, :, 0:1] = inputs["first_frame_latents"] - - kl_loss = kl_loss_total / kl_steps if kl_steps > 0 else 0.0 - video_decoded = self.vae.decode(inputs["latents"], device=self.device, tiled=True) - return video_decoded, kl_loss - - -class WanTrainingModule(DiffusionTrainingModule): - def __init__( - self, - model_paths, - reward_model_path, - tokenizer_path, - trainable_models, - model_id_with_origin_paths=None, - lora_base_model=None, - lora_target_modules="q,k,v,o,ffn.0,ffn.2", - lora_rank=32, - use_gradient_checkpointing=True, - use_gradient_checkpointing_offload=False, - extra_inputs=None, - max_timestep_boundary=1.0, - min_timestep_boundary=0.9, - num_inference_steps=8, - mid_timestep=4, - final_timestep=7, - **kwargs - ): - super().__init__() - # Load models - model_configs : List[ModelConfig] = [] - if model_paths is not None: - with open(model_paths, 'r', encoding='utf-8') as f: - model_paths = json.load(f) - model_configs += [ModelConfig(path=path) for path in model_paths] - if model_id_with_origin_paths is not None: - model_id_with_origin_paths = model_id_with_origin_paths.split(",") - model_configs += [ModelConfig(model_id=i.split(":")[0], origin_file_pattern=i.split(":")[1]) for i in model_id_with_origin_paths] - self.pipe = WanVideoPipeline.from_pretrained( - torch_dtype=torch.bfloat16, - device="cpu", - model_configs=model_configs, - tokenizer_config=ModelConfig(path=tokenizer_path), - redirect_common_files=False - ) - - self.apply_patches() - - face_model = FaceAnalysis(root=reward_model_path, device=current_platform.device_type) - - # 将冻结模型存入一个普通字典中 PyTorch 不会注册普通字典中的 nn.Module - self.frozen_dependencies = { - 'face_model': face_model, - } - - # Reset training scheduler - self.pipe.scheduler.set_timesteps(1000) - - # Freeze untrainable models - self.pipe.freeze_except([] if trainable_models is None else trainable_models.split(",")) - - # Add LoRA to the base models - if lora_base_model is not None: - model = self.add_lora_to_model( - getattr(self.pipe, lora_base_model), - target_modules=lora_target_modules.split(","), - lora_rank=lora_rank - ) - setattr(self.pipe, lora_base_model, model) - - # Store other configs - self.use_gradient_checkpointing = use_gradient_checkpointing - self.use_gradient_checkpointing_offload = use_gradient_checkpointing_offload - self.extra_inputs = extra_inputs.split(",") if extra_inputs is not None else [] - self.max_timestep_boundary = max_timestep_boundary - self.min_timestep_boundary = min_timestep_boundary - self.pipe.num_inference_steps = num_inference_steps - self.pipe.mid_timestep = mid_timestep - self.pipe.final_timestep = final_timestep - self.io_executor = ThreadPoolExecutor(max_workers=1) - self.io_queue = queue.Queue() - self.output_path = './output_video/' - - self.global_step = 0 - - def apply_patches(self): - - # apply patches - self.pipe.units.append(WanVideoUnit_Face()) - self.pipe.scheduler = EulerScheduler(num_train_timesteps=1000, shift=5, device=current_platform.device_type) - vae_state_dict = self.pipe.vae.state_dict() - self.pipe.vae = WanVideoVAE() - self.pipe.vae.load_state_dict(vae_state_dict, strict=True) - self.pipe.model_fn = model_fn_wan_video - self.pipe.training_loss = types.MethodType(training_loss, self.pipe) - - - def forward_preprocess(self, data): - # CFG-sensitive parameters - inputs_posi = {"prompt": data["prompt"]} - inputs_nega = {} - - # CFG-unsensitive parameters - inputs_shared = { - # Assume you are using this pipeline for inference, - # please fill in the input parameters. - "input_video": data["video"], - "height": data["video"][0].size[1], - "width": data["video"][0].size[0], - "num_frames": len(data["video"]), - "face_model": self.frozen_dependencies['face_model'], - # Please do not modify the following parameters - # unless you clearly know what this will cause. - "cfg_scale": 1, - "tiled": False, - "rand_device": self.pipe.device, - "use_gradient_checkpointing": self.use_gradient_checkpointing, - "use_gradient_checkpointing_offload": self.use_gradient_checkpointing_offload, - "cfg_merge": False, - "vace_scale": 1, - "max_timestep_boundary": self.max_timestep_boundary, - "min_timestep_boundary": self.min_timestep_boundary, - } - - # Extra inputs - for extra_input in self.extra_inputs: - if extra_input == "input_image": - inputs_shared["input_image"] = data["video"][0] - elif extra_input == "end_image": - inputs_shared["end_image"] = data["video"][-1] - elif extra_input == "reference_image" or extra_input == "vace_reference_image": - inputs_shared[extra_input] = data[extra_input][0] - else: - inputs_shared[extra_input] = data[extra_input] - - # Pipeline units will automatically process the input parameters. - for unit in self.pipe.units: - inputs_shared, inputs_posi, inputs_nega = self.pipe.unit_runner(unit, self.pipe, inputs_shared, inputs_posi, inputs_nega) - return {**inputs_shared, **inputs_posi} - - - def forward(self, data, inputs=None): - if inputs is None: - inputs = self.forward_preprocess(data) - - face_embeddings = inputs['face_embeddings'].to(device=self.pipe.device, dtype=self.pipe.torch_dtype) - - # step1: forward latents + vae decode - video_decoded, kl_loss = self.pipe.training_loss(**inputs) - - # step2: get video_tensor - video_tensor = vae_output_to_videotensor(video_decoded) - - # step3: true video submit - self.vae_video2_submit(video_tensor, self.output_path) - - video_decoded = video_decoded[0].permute(1, 0, 2, 3) # (C, T, H, W) -> (T, C, H, W) - print(f'video decode shape: {video_decoded.shape}') - - video_decoded = torch.clamp(video_decoded, min=-1, max=1) - self.frozen_dependencies['face_model'].detection_model.torch_model.to(self.pipe.device) - self.frozen_dependencies['face_model'].arcface_model.torch_model.to(self.pipe.device) - - id_embeds, id_masked = [], [] - face_num = 0 - for f in video_decoded: - f = f.float() - bboxes, kpss = self.frozen_dependencies['face_model'].detection_model.detect(f) - if bboxes.shape[0] > 0: - indexed_bboxes = [(i, x) for i, x in enumerate(bboxes)] - sorted_bboxes = sorted(indexed_bboxes, key=lambda item: (item[1][2] - item[1][0]) * (item[1][3] - item[1][1])) - max_index, max_bbox = sorted_bboxes[-1] - kps = kpss[max_index] - face = Face(bbox=bboxes[max_index][0:4], kps=kps, det_score=bboxes[max_index][4]) - id_embeds.append(self.frozen_dependencies['face_model'].arcface_model.get(f, face)) - id_masked.append(1) - face_num += 1 - else: - id_embeds.append(torch.zeros(512).to(self.pipe.device)) - id_masked.append(0) - assert face_num > 0, f"face_num must be greater than 0" - - id_embeds = torch.stack(id_embeds).unsqueeze(0) - id_masked = torch.tensor(id_masked).unsqueeze(0).to(self.pipe.device) - - face_score = self.frozen_dependencies['face_model'].pool_embedding_loss(id_embeds, face_embeddings, id_masked) - print(f"{face_score=}") - - face_score = face_score.to(self.pipe.device) - - del video_tensor, video_decoded - gc.collect() - current_platform.empty_cache() - - loss = -(face_score.bfloat16()-0.54)/0.16 * 0.1 - loss += kl_loss - - loss = loss.to(self.pipe.device) - - print(f'loss: {loss.float().detach().cpu().item()}') - self.global_step = self.global_step + 1 - - gc.collect() - current_platform.empty_cache() - - return loss, face_score, kl_loss - - def vae_video2_submit(self, video_tensor, output_path): - rank = int(os.environ.get("RANK", 0)) - if rank == 0: - while not self.io_queue.empty(): - try: - future = self.io_queue.get_nowait() - if not future.done(): - self.io_queue.put(future) - return - except queue.Empty: - break - - step = self.global_step - video_tmp = video_tensor.clone().detach().cpu().float().numpy() - video_tmp = video_tmp.round().astype(np.uint8) - - - future = self.io_executor.submit( - self._save_video_background, video_tmp, output_path, rank, step - ) - self.io_queue.put(future) - - - def _save_video_background(self, video_data, output_path, rank, step): - try: - timestamp = datetime.now().strftime("%Y%m%d%H%M%S") - save_oss_dir = os.path.join(output_path, 'decode_videos') - os.makedirs(save_oss_dir, exist_ok=True) - video_filename = f'video_rank{rank}_iter{step}_time{timestamp}.mp4' - save_video(video_data, video_filename, save_oss_dir, save_to_oss=True) - except Exception as e: - print(f"Error during background video saving: {e}") - raise e - - def add_lora_to_model(self, model, target_modules, lora_rank, lora_alpha=None): - if lora_alpha is None: - lora_alpha = lora_rank - lora_config = LoraConfig(r=lora_rank, lora_alpha=lora_alpha, target_modules=target_modules) - model = get_peft_model(model, lora_config) - return model - -class WanVideoUnit_Face(PipelineUnit): - def __init__(self): - super().__init__(input_params=("input_video", "face_model")) - - def process(self, pipe: WanVideoPipeline, input_video, face_model): - input_video = pipe.preprocess_video(input_video) # (1, 3, F, H, W) - input_video = input_video[0].transpose(0, 1) - - face_embeds, face_masked= [], [] - # input_video (F, 3, H, W) 数值范围(-1, 1) - for f in input_video: - f = f.float().to(pipe.device) #(3, h, w) - face_model.detection_model.torch_model.to(pipe.device) - face_model.arcface_model.torch_model.to(pipe.device) - bboxes, kpss = face_model.detection_model.detect(f) - if bboxes.shape[0] > 0: - indexed_bboxes = [(i, x) for i, x in enumerate(bboxes)] - sorted_bboxes = sorted(indexed_bboxes, key=lambda item: (item[1][2] - item[1][0]) * (item[1][3] - item[1][1])) - max_index, max_bbox = sorted_bboxes[-1] - kps = kpss[max_index] - face = Face(bbox=bboxes[max_index][0:4], kps=kps, det_score=bboxes[max_index][4]) - embedding = face_model.arcface_model.get(f, face) - face_embeds.append(embedding.cpu()) - face_masked.append(1) - else: - face_embeds.append(torch.zeros(512)) - face_masked.append(0) - face_embeds = torch.stack(face_embeds).unsqueeze(0) - face_masked = torch.tensor(face_masked).unsqueeze(0).to(pipe.device) - return {"face_embeddings": face_embeds, "face_masked":face_masked} - - -def save_video(video_frames, save_video_basename, output_oss_dir, save_to_oss=True): - if video_frames.shape[0] == 1: # T=1时保存为图像 - local_output_path = f'{save_video_basename}.png' if not save_video_basename.endswith('.png') else save_video_basename - oss_output = f'{output_oss_dir}/{local_output_path}' - imageio.imwrite(oss_output, video_frames[0]) # 取单帧保存 - else: - local_output_path = f'{save_video_basename}.mp4' if not save_video_basename.endswith('.mp4') else save_video_basename - oss_output = f'{output_oss_dir}/{local_output_path}' - write_video(local_output_path, video_frames, fps=16, options={'crf': '10'}) - if save_to_oss: - os.system(f'cp {local_output_path} {oss_output}') - os.system(f'rm -rf {local_output_path}') - - -def model_fn_wan_video( - dit: WanModel, - motion_controller: WanMotionControllerModel = None, - vace: VaceWanModel = None, - latents: torch.Tensor = None, - timestep: torch.Tensor = None, - context: torch.Tensor = None, - clip_feature: Optional[torch.Tensor] = None, - y: Optional[torch.Tensor] = None, - reference_latents = None, - vace_context = None, - vace_scale = 1.0, - tea_cache: TeaCache = None, - less_mid_step: bool = True, - use_unified_sequence_parallel: bool = False, - motion_bucket_id: Optional[torch.Tensor] = None, - sliding_window_size: Optional[int] = None, - sliding_window_stride: Optional[int] = None, - cfg_merge: bool = False, - use_gradient_checkpointing: bool = False, - use_gradient_checkpointing_offload: bool = False, - control_camera_latents_input = None, - fuse_vae_embedding_in_latents: bool = False, - recompute_num_layers: Optional[int] = 1, - **kwargs, -): - if sliding_window_size is not None and sliding_window_stride is not None: - model_kwargs = dict( - dit=dit, - motion_controller=motion_controller, - vace=vace, - latents=latents, - timestep=timestep, - context=context, - clip_feature=clip_feature, - y=y, - reference_latents=reference_latents, - vace_context=vace_context, - vace_scale=vace_scale, - tea_cache=tea_cache, - use_unified_sequence_parallel=use_unified_sequence_parallel, - motion_bucket_id=motion_bucket_id, - ) - return TemporalTiler_BCTHW().run( - model_fn_wan_video, - sliding_window_size, sliding_window_stride, - latents.device, latents.dtype, - model_kwargs=model_kwargs, - tensor_names=["latents", "y"], - batch_size=2 if cfg_merge else 1 - ) - - if use_unified_sequence_parallel: - import torch.distributed as dist - from xfuser.core.distributed import (get_sequence_parallel_rank, - get_sequence_parallel_world_size, - get_sp_group) - - # Timestep - if dit.seperated_timestep and fuse_vae_embedding_in_latents: - timestep = torch.concat([ - torch.zeros((1, latents.shape[3] * latents.shape[4] // 4), dtype=latents.dtype, device=latents.device), - torch.ones((latents.shape[2] - 1, latents.shape[3] * latents.shape[4] // 4), dtype=latents.dtype, device=latents.device) * timestep - ]).flatten() - t = dit.time_embedding(sinusoidal_embedding_1d(dit.freq_dim, timestep).unsqueeze(0)) - t_mod = dit.time_projection(t).unflatten(2, (6, dit.dim)) - else: - with torch.amp.autocast(current_platform.device_type, dtype=torch.float32): - t = dit.time_embedding(sinusoidal_embedding_1d(dit.freq_dim, timestep)) - t_mod = dit.time_projection(t).unflatten(1, (6, dit.dim)) - assert t.dtype == torch.float32 and t_mod.dtype == torch.float32 - t, t_mod = t.to(dtype=torch.bfloat16), t_mod.to(dtype=torch.bfloat16) - assert t.dtype == torch.bfloat16 and t_mod.dtype == torch.bfloat16 - - # Motion Controller - if motion_bucket_id is not None and motion_controller is not None: - t_mod = t_mod + motion_controller(motion_bucket_id).unflatten(1, (6, dit.dim)) - context = dit.text_embedding(context) - - x = latents - # Merged cfg - if x.shape[0] != context.shape[0]: - x = torch.concat([x] * context.shape[0], dim=0) - if timestep.shape[0] != context.shape[0]: - timestep = torch.concat([timestep] * context.shape[0], dim=0) - - # Image Embedding - if y is not None and dit.require_vae_embedding: - x = torch.cat([x, y], dim=1) - if clip_feature is not None and dit.require_clip_embedding: - clip_embdding = dit.img_emb(clip_feature) - context = torch.cat([clip_embdding, context], dim=1) - - # Add camera control - x, (f, h, w) = dit.patchify(x, control_camera_latents_input) - - # Reference image - if reference_latents is not None: - if len(reference_latents.shape) == 5: - reference_latents = reference_latents[:, :, 0] - reference_latents = dit.ref_conv(reference_latents).flatten(2).transpose(1, 2) - x = torch.concat([reference_latents, x], dim=1) - f += 1 - - freqs = torch.cat([ - dit.freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1), - dit.freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1), - dit.freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1) - ], dim=-1).reshape(f * h * w, 1, -1).to(x.device) - - # TeaCache - if tea_cache is not None: - tea_cache_update = tea_cache.check(dit, x, t_mod) - else: - tea_cache_update = False - - if vace_context is not None: - vace_hints = vace(x, vace_context, context, t_mod, freqs) - - # blocks - if use_unified_sequence_parallel: - if dist.is_initialized() and dist.get_world_size() > 1: - chunks = torch.chunk(x, get_sequence_parallel_world_size(), dim=1) - pad_shape = chunks[0].shape[1] - chunks[-1].shape[1] - chunks = [torch.nn.functional.pad(chunk, (0, 0, 0, chunks[0].shape[1]-chunk.shape[1]), value=0) for chunk in chunks] - x = chunks[get_sequence_parallel_rank()] - if tea_cache_update: - x = tea_cache.update(x) - else: - def create_custom_forward(module): - def custom_forward(*inputs): - return module(*inputs) - return custom_forward - - for block_id, block in enumerate(dit.blocks): - if use_gradient_checkpointing_offload: - with torch.autograd.graph.save_on_cpu(): - x = torch.utils.checkpoint.checkpoint( - create_custom_forward(block), - x, context, t_mod, freqs, - use_reentrant=False, - ) - elif use_gradient_checkpointing: - x = torch.utils.checkpoint.checkpoint( - create_custom_forward(block), - x, context, t_mod, freqs, - use_reentrant=False, - ) - else: - x = block(x, context, t_mod, freqs) - if vace_context is not None and block_id in vace.vace_layers_mapping: - current_vace_hint = vace_hints[vace.vace_layers_mapping[block_id]] - if use_unified_sequence_parallel and dist.is_initialized() and dist.get_world_size() > 1: - current_vace_hint = torch.chunk(current_vace_hint, get_sequence_parallel_world_size(), dim=1)[get_sequence_parallel_rank()] - current_vace_hint = torch.nn.functional.pad(current_vace_hint, (0, 0, 0, chunks[0].shape[1] - current_vace_hint.shape[1]), value=0) - x = x + current_vace_hint * vace_scale - if tea_cache is not None: - tea_cache.store(x) - - x = dit.head(x, t) - if use_unified_sequence_parallel: - if dist.is_initialized() and dist.get_world_size() > 1: - x = get_sp_group().all_gather(x, dim=1) - x = x[:, :-pad_shape] if pad_shape > 0 else x - # Remove reference latents - if reference_latents is not None: - x = x[:, reference_latents.shape[1]:] - f -= 1 - x = dit.unpatchify(x, (f, h, w)) - return x diff --git a/roll/pipeline/diffusion/reward_fl/__init__.py b/roll/pipeline/diffusion/reward_fl/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/roll/pipeline/diffusion/reward_fl/actor_worker.py b/roll/pipeline/diffusion/reward_fl/actor_worker.py deleted file mode 100644 index 085b0f338..000000000 --- a/roll/pipeline/diffusion/reward_fl/actor_worker.py +++ /dev/null @@ -1,60 +0,0 @@ -import numpy as np -import torch - -from tqdm import tqdm - -from roll.distributed.scheduler.decorator import Dispatch, register -from roll.distributed.scheduler.protocol import DataProto -from roll.pipeline.base_worker import ActorWorker as BaseActorWorker -from roll.utils.functionals import append_to_dict -from roll.platforms import current_platform - - -class ActorWorker(BaseActorWorker): - - @register(dispatch_mode=Dispatch.DP_MP_DISPATCH_FIRST, clear_cache=False) - def train_step(self, data: DataProto): - - global_step = data.meta_info.get("global_step", 0) - metrics = {} - - data = data.to(current_platform.device_type) - - per_device_train_batch_size = self.worker_config.training_args.per_device_train_batch_size - backward_batch_size = ( - per_device_train_batch_size * self.worker_config.training_args.gradient_accumulation_steps - ) - - dataloader = data.make_iterator( - mini_batch_size=backward_batch_size, - epochs=1, - dataloader_kwargs={"shuffle": False}, - ) - - for batch_idx, data in tqdm( - enumerate(dataloader), - desc=f"{self.worker_name} train global step {global_step}", - total=data.batch.batch_size[0] // backward_batch_size, - ): - pg_metrics = self.strategy.train_step(batch=data, loss_func=self.loss_func) - append_to_dict(metrics, pg_metrics) - - metrics["actor/loss"] = np.mean(metrics["actor/loss"]) - metrics["actor/face_score"] = np.mean(metrics["actor/face_score"]) - metrics["actor/kl_loss"] = np.mean(metrics["actor/kl_loss"]) - data.to("cpu") - - output = DataProto(meta_info={"metrics": metrics}) - return output - - def loss_func(self, data: DataProto, loss: torch.Tensor, face_score: torch.Tensor, kl_loss: torch.Tensor): - """ - data: DataProto, from train_step - """ - metrics = { - "actor/loss": loss.float().detach().cpu().item(), - "actor/face_score": face_score.float().detach().cpu().item(), - "actor/kl_loss": kl_loss.float().detach().cpu().item(), - } - - return loss, metrics diff --git a/roll/pipeline/diffusion/reward_fl/euler.py b/roll/pipeline/diffusion/reward_fl/euler.py deleted file mode 100644 index 90878a7fa..000000000 --- a/roll/pipeline/diffusion/reward_fl/euler.py +++ /dev/null @@ -1,89 +0,0 @@ -from diffusers import FlowMatchEulerDiscreteScheduler -from torch import Tensor - -import numpy as np -import torch - - -def unsqueeze_to_ndim(in_tensor: Tensor, tgt_n_dim: int): - if in_tensor.ndim > tgt_n_dim: - return in_tensor - if in_tensor.ndim < tgt_n_dim: - in_tensor = in_tensor[(...,) + (None,) * (tgt_n_dim - in_tensor.ndim)] - return in_tensor - - -def get_timesteps(num_steps, max_steps: int = 1000): - return np.linspace(max_steps, 0, num_steps + 1, dtype=np.float32) - - -def timestep_shift(timesteps, shift: float = 1.0): - return shift * timesteps / (1 + (shift - 1) * timesteps) - - -class EulerScheduler(FlowMatchEulerDiscreteScheduler): - def __init__( - self, - num_train_timesteps: int, - shift: float = 1.0, - device: torch.device | str = "cuda", - **kwargs - ) -> None: - super().__init__(num_train_timesteps=num_train_timesteps, shift=shift,**kwargs) - self.init_noise_sigma = 1.0 - self.num_train_timesteps = num_train_timesteps - self._shift = shift - self.init_noise_sigma = 1.0 - self.device = device - self.training = True - self.set_timesteps(num_inference_steps=num_train_timesteps) - - - def set_shift(self, shift: float = 1.0): - self.sigmas = self.timesteps_ori / self.num_train_timesteps - self.sigmas = timestep_shift(self.sigmas, shift=shift) - self.timesteps = self.sigmas * self.num_train_timesteps - self._shift = shift - - - def set_timesteps( - self, num_inference_steps: int, device: torch.device | str | int | None = None - ): - timesteps = get_timesteps( - num_steps=num_inference_steps, max_steps=self.num_train_timesteps - ) - self.timesteps = torch.from_numpy(timesteps).to( - dtype=torch.float32, device=device or self.device - ) - self.timesteps_ori = self.timesteps.clone() - self.set_shift(self._shift) - self._step_index = None - self._begin_index = None - - - def step( - self, - model_output: torch.FloatTensor, - timestep: float | torch.FloatTensor, - sample: torch.FloatTensor, - **kwargs, - ) -> tuple: - if ( - isinstance(timestep, int) - or isinstance(timestep, torch.IntTensor) - or isinstance(timestep, torch.LongTensor) - ): - raise ValueError( - ( - "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to" - " `EulerDiscreteScheduler.step()` is not supported. Make sure to pass" - " one of the `scheduler.timesteps` as a timestep." - ), - ) - if self.step_index is None: - self._init_step_index(timestep) - sigma = unsqueeze_to_ndim(self.sigmas[self.step_index], sample.ndim).to(sample.device) - sigma_next = unsqueeze_to_ndim(self.sigmas[self.step_index + 1], sample.ndim).to(sample.device) - x_t_next = sample + (sigma_next - sigma) * model_output - self._step_index += 1 - return x_t_next \ No newline at end of file diff --git a/roll/pipeline/diffusion/reward_fl/face_tools.py b/roll/pipeline/diffusion/reward_fl/face_tools.py deleted file mode 100644 index 96b2e2868..000000000 --- a/roll/pipeline/diffusion/reward_fl/face_tools.py +++ /dev/null @@ -1,508 +0,0 @@ -"""A simple, flexible implementation of a face analysis tool. - -Inspired by https://github.com/deepinsight/insightface -""" -import torch -import numpy as np -import onnx -import cv2 -import math -from onnx2torch import convert -from torchvision.transforms.functional import to_tensor, resize -import torch.nn.functional as F -from skimage import transform as trans -import time -import os -import torchvision.ops as ops - -from roll.platforms import current_platform - -arcface_dst = torch.tensor( - [[38.2946, 51.6963], [73.5318, 51.5014], [56.0252, 71.7366], - [41.5493, 92.3655], [70.7299, 92.2041]]).float() - -def distance2bbox(points, distance, max_shape=None): - """Decode distance prediction to bounding box. - - Args: - points (Tensor): Shape (n, 2), [x, y]. - distance (Tensor): Distance from the given point to 4 - boundaries (left, top, right, bottom). - max_shape (tuple): Shape of the image. - - Returns: - Tensor: Decoded bboxes. - """ - x1 = points[:, 0] - distance[:, 0] - y1 = points[:, 1] - distance[:, 1] - x2 = points[:, 0] + distance[:, 2] - y2 = points[:, 1] + distance[:, 3] - if max_shape is not None: - x1 = x1.clamp(min=0, max=max_shape[1]) - y1 = y1.clamp(min=0, max=max_shape[0]) - x2 = x2.clamp(min=0, max=max_shape[1]) - y2 = y2.clamp(min=0, max=max_shape[0]) - return torch.stack([x1, y1, x2, y2], axis=-1) - -def distance2kps(points, distance, max_shape=None): - """Decode distance prediction to bounding box. - - Args: - points (Tensor): Shape (n, 2), [x, y]. - distance (Tensor): Distance from the given point to 4 - boundaries (left, top, right, bottom). - max_shape (tuple): Shape of the image. - - Returns: - Tensor: Decoded bboxes. - """ - preds = [] - for i in range(0, distance.shape[1], 2): - px = points[:, i%2] + distance[:, i] - py = points[:, i%2+1] + distance[:, i+1] - if max_shape is not None: - px = px.clamp(min=0, max=max_shape[1]) - py = py.clamp(min=0, max=max_shape[0]) - preds.append(px) - preds.append(py) - return torch.stack(preds, axis=-1) - - -def face_transform(data, center, output_size, scale, rotation, device): - def to_homogeneous(mat): - """将 2x3 仿射矩阵转为 3x3 齐次矩阵""" - return torch.vstack([mat, torch.tensor([0., 0., 1.])]) - scale_ratio = scale - rot = float(rotation) * np.pi / 180.0 - cx = center[0] * scale_ratio - cy = center[1] * scale_ratio - - C, H, W = data.shape - - # 构建各个变换矩阵 - t1 = to_homogeneous(torch.tensor([ - [scale_ratio, 0, 0], - [0, scale_ratio, 0] - ])).float() - t2 = to_homogeneous(torch.tensor([ - [1, 0, -cx], - [0, 1, -cy] - ])).float() - cos_theta = math.cos(rot) - sin_theta = math.sin(rot) - t3 = to_homogeneous(torch.tensor([ - [cos_theta, -sin_theta, 0], - [sin_theta, cos_theta, 0] - ])).float() - t4 = to_homogeneous(torch.tensor([ - [1, 0, output_size / 2], - [0, 1, output_size / 2] - ])).float() - M_homogeneous = t4 @ t3 @ t2 @ t1 - M = M_homogeneous[:2, :] # 提取前两行作为 2x3 仿射矩阵 - # 应用仿射变换 - T = torch.tensor([[2 / W, 0, -1], - [0, 2 / H, -1], - [0, 0, 1]]) - theta = torch.inverse(T @ M_homogeneous @ torch.inverse(T)) - theta = theta[:2, :].unsqueeze(0).to(device) - # theta = M.unsqueeze(0) # 添加 batch 维度 (1, 2, 3) - grid = F.affine_grid(theta, data.unsqueeze(0).size(), align_corners=True) - transformed = F.grid_sample(data.unsqueeze(0), grid, align_corners=True) - cropped = transformed[0] - cropped = cropped[:,:output_size,:output_size] - # crop_map = torch.zeros(3, output_size, output_size) - # crop_map[:, :cropped.shape[1],:cropped.shape[2]] = cropped - return cropped.unsqueeze(0), M - -def trans_points2d(pts, M): - ones = torch.ones((pts.shape[0], 1), dtype=pts.dtype, device=pts.device) - points_hom = torch.cat([pts, ones], dim=1) # shape: (n, 3) - points_hom = points_hom.unsqueeze(-1) # shape: (n, 3, 1) - transformed_hom = torch.matmul(M, points_hom) # shape: (n, 3, 1) - transformed = transformed_hom[:, :2, :].squeeze(-1) # shape: (n, 2) - return transformed - -def estimate_norm(lmk, image_size=112,mode='arcface'): - assert lmk.shape == (5, 2) - assert image_size%112==0 or image_size%128==0 - if image_size%112==0: - ratio = float(image_size)/112.0 - diff_x = 0 - else: - ratio = float(image_size)/128.0 - diff_x = 8.0*ratio - dst = arcface_dst * ratio - dst[:,0] += diff_x - tform = trans.SimilarityTransform() - tform.estimate(lmk, dst) - M = torch.from_numpy(tform.params).float() - return M - -def norm_crop(img, landmark, image_size=112, mode='arcface'): - M_homogeneous = estimate_norm(landmark, image_size, mode) - C, H, W = img.shape - img = img.unsqueeze(0) - T = torch.tensor([[2 / W, 0, -1], - [0, 2 / H, -1], - [0, 0, 1]]) - T_inv = torch.inverse(T) - theta = torch.inverse(T @ M_homogeneous @ T_inv) - theta = theta[:2, :].unsqueeze(0).to(img.device) - grid = F.affine_grid(theta, img.size(), align_corners=True) - transformed = F.grid_sample(img, grid, align_corners=True) - cropped = transformed[0] - warped = cropped[:,:image_size,:image_size] - return warped - -def invert_affine_transform(matrix): - L = matrix[..., :2] # Shape: (*, 2, 2) - T = matrix[..., 2:] # Shape: (*, 2, 1) - a, b = L[..., 0, 0], L[..., 0, 1] - c, d = L[..., 1, 0], L[..., 1, 1] - det = a * d - b * c - inv_det = 1.0 / det - inv_L = torch.stack([ - torch.stack([d * inv_det, -b * inv_det], dim=-1), - torch.stack([-c * inv_det, a * inv_det], dim=-1) - ], dim=-2) - inv_T = -torch.matmul(inv_L, T) - inv_matrix = torch.cat([inv_L, inv_T], dim=-1) - return inv_matrix - -class Face(dict): - - def __init__(self, d=None, **kwargs): - if d is None: - d = {} - if kwargs: - d.update(**kwargs) - for k, v in d.items(): - setattr(self, k, v) - # Class attributes - #for k in self.__class__.__dict__.keys(): - # if not (k.startswith('__') and k.endswith('__')) and not k in ('update', 'pop'): - # setattr(self, k, getattr(self, k)) - - def __setattr__(self, name, value): - if isinstance(value, (list, tuple)): - value = [self.__class__(x) - if isinstance(x, dict) else x for x in value] - elif isinstance(value, dict) and not isinstance(value, self.__class__): - value = self.__class__(value) - super(Face, self).__setattr__(name, value) - super(Face, self).__setitem__(name, value) - - __setitem__ = __setattr__ - - def __getattr__(self, name): - return None - - @property - def embedding_norm(self): - if self.embedding is None: - return None - return torch.norm(self.embedding) - - @property - def normed_embedding(self): - if self.embedding is None: - return None - return self.embedding / self.embedding_norm - - -class SCRFD: - def __init__(self, model_file=None, device="cuda"): - self.model_file = model_file - self.device = device - self.center_cache = {} - self.nms_thresh = 0.4 - self.det_thresh = 0.5 - model = onnx.load(self.model_file) - self.torch_model = convert(model) - self.torch_model.eval() - self.torch_model.requires_grad_(False) - self.torch_model.to(self.device) - self.use_kps = True - self.fmc = 3 - self._num_anchors = 2 - self._feat_stride_fpn = [8, 16, 32] - self.input_size = (640, 640) - - def forward(self, det_img, threshold=0.5): - input_height = det_img.shape[2] - input_width = det_img.shape[3] - scores_list = [] - bboxes_list = [] - kpss_list = [] - net_outs = self.torch_model(det_img.float()) - - for idx, stride in enumerate(self._feat_stride_fpn): - scores = net_outs[idx].cpu() - bbox_preds = net_outs[idx + self.fmc].cpu() - bbox_preds = bbox_preds * stride - if self.use_kps: - kps_preds = net_outs[idx + self.fmc * 2].cpu() * stride - - height = input_height // stride - width = input_width // stride - K = height * width - key = (height, width, stride) - if key in self.center_cache: - anchor_centers = self.center_cache[key] - else: - rows = torch.arange(height) - cols = torch.arange(width) - grid_y, grid_x = torch.meshgrid(rows, cols, indexing='ij') - anchor_centers = torch.stack([grid_x, grid_y], dim=-1).float() - anchor_centers = (anchor_centers * stride).reshape((-1, 2)) - if self._num_anchors>1: - anchor_centers = torch.stack([anchor_centers]*self._num_anchors, axis=1).reshape( (-1,2) ) - if len(self.center_cache)<100: - self.center_cache[key] = anchor_centers - - pos_inds = np.where(scores>=threshold)[0] - # print(bbox_preds.shape) - bboxes = distance2bbox(anchor_centers, bbox_preds) - pos_scores = scores[pos_inds] - pos_bboxes = bboxes[pos_inds] - scores_list.append(pos_scores) - bboxes_list.append(pos_bboxes) - if self.use_kps: - kpss = distance2kps(anchor_centers, kps_preds) - #kpss = kps_preds - kpss = kpss.reshape( (kpss.shape[0], -1, 2) ) - pos_kpss = kpss[pos_inds] - kpss_list.append(pos_kpss) - - return scores_list, bboxes_list, kpss_list - - @torch.no_grad() - def detect(self, image, input_size = None, max_num = 0, metric='default'): - assert input_size is not None or self.input_size is not None - input_size = self.input_size if input_size is None else input_size - - im_ratio = float(image.shape[1]) / image.shape[2] - model_ratio = float(input_size[1]) / input_size[0] - if im_ratio>model_ratio: - new_height = input_size[1] - new_width = int(new_height / im_ratio) - else: - new_width = input_size[0] - new_height = int(new_width * im_ratio) - det_scale = float(new_height) / image.shape[1] - resized_img = resize(image, (new_height, new_width), antialias=True) - det_img = torch.zeros( (3, input_size[1], input_size[0]),device=self.device) - det_img[:, :new_height, :new_width] = resized_img - det_img = det_img.unsqueeze(0) - scores_list, bboxes_list, kpss_list = self.forward(det_img, self.det_thresh) - - scores = torch.vstack(scores_list) - scores_ravel = scores.flatten() - order = torch.argsort(scores_ravel, descending=True) - bboxes = torch.vstack(bboxes_list) / det_scale - if self.use_kps: - kpss = torch.vstack(kpss_list) / det_scale - - pre_det = torch.cat((bboxes, scores), dim=1).float() - pre_det = pre_det[order] - keep = self.nms(pre_det) - det = pre_det[keep, :] - - if self.use_kps: - kpss = kpss[order,:,:] - kpss = kpss[keep,:,:] - else: - kpss = None - return det, kpss - - def nms(self, dets): - boxes = dets[:, :4] - scores = dets[:, 4] - keep = ops.nms(boxes, scores, iou_threshold=self.nms_thresh) - return keep.tolist() - - -class ArcFace: - def __init__(self, model_file=None, device="cuda"): - self.model_file = model_file - self.device = device - model = onnx.load(self.model_file) - self.torch_model = convert(model) - self.torch_model.eval() - self.torch_model.to(self.device) - self.torch_model.requires_grad_(False) - self.taskname = 'recognition' - self.input_size = (112, 112) - - def get(self, img, face, input_size=(112, 112)): - aimg = norm_crop(img, landmark=face.kps, image_size=self.input_size[0]) - im_ratio = float(aimg.shape[1]) / aimg.shape[2] - model_ratio = float(input_size[1]) / input_size[0] - if im_ratio>model_ratio: - new_height = input_size[1] - new_width = int(new_height / im_ratio) - else: - new_width = input_size[0] - new_height = int(new_width * im_ratio) - resized_img = resize(aimg, (new_height, new_width), antialias=True) - face.embedding = self.get_feat(resized_img.unsqueeze(0)).flatten() - return face.embedding - - def compute_sim(self, feat1, feat2): - feat1 = feat1.ravel() - feat2 = feat2.ravel() - sim = torch.dot(feat1, feat2) / (torch.norm(feat1) * torch.norm(feat2)) - return sim - - def get_feat(self, imgs): - imgs = imgs[:,[2,1,0],:,:] - net_out = self.torch_model(imgs) - return net_out - -class Landmark: - def __init__(self, model_file=None, device="cuda"): - self.model_file = model_file - self.device = device - model = onnx.load(self.model_file) - self.torch_model = convert(model) - self.torch_model.eval() - self.torch_model.to(device) - self.torch_model.requires_grad_(False) - self.lmk_dim = 2 - self.lmk_num = 106 - self.taskname = 'landmark_%dd_%d'%(self.lmk_dim, self.lmk_num) - self.input_size = (192, 192) - - def get(self, img, face, input_size=(192, 192)): - bbox = face.bbox - w, h = (bbox[2] - bbox[0]), (bbox[3] - bbox[1]) - center = (bbox[2] + bbox[0]) / 2, (bbox[3] + bbox[1]) / 2 - rotate = 0 - _scale = self.input_size[0] / (max(w, h)*1.5) - aimg, M = face_transform(img, center, self.input_size[0], _scale, rotate, img.device) - aimg = (aimg + 1)/2 * 255. # [1, 3, 192, 192] - aimg = aimg[:,[2,1,0],:,:] - - input_size = self.input_size if input_size is None else input_size - im_ratio = float(aimg.shape[2]) / aimg.shape[3] - model_ratio = float(input_size[1]) / input_size[0] - if im_ratio>model_ratio: - new_height = input_size[1] - new_width = int(new_height / im_ratio) - else: - new_width = input_size[0] - new_height = int(new_width * im_ratio) - det_scale = float(new_height) / aimg.shape[2] - resized_img = resize(aimg, (new_height, new_width), antialias=True) - det_img = torch.zeros( (aimg.shape[0], 3, input_size[1], input_size[0]), device=self.device) - det_img[:, :, :new_height, :new_width] = resized_img - - pred = self.torch_model(det_img)[0] #输入图像应为RGB，不能是BGR - pred = pred.reshape((-1, 2)) - if self.lmk_num < pred.shape[0]: - pred = pred[self.lmk_num*-1:,:] - pred[:, 0:2] += 1 - pred[:, 0:2] *= (self.input_size[0] // 2) - - IM = invert_affine_transform(M).to(img.device) - pred = trans_points2d(pred, IM) - face[self.taskname] = pred - return pred - -class FaceAnalysis: - def __init__(self, root="~/.insightface", device="cuda"): - self.root = root - self.device = device - self.detection_root = os.path.join(root, "scrfd_10g_bnkps.onnx") - self.landmark_root = os.path.join(root, "2d106det.onnx") - self.arcface_root = os.path.join(root, "glintr100.onnx") - self.detection_model = SCRFD(self.detection_root, self.device) - self.landmark_model = Landmark(self.landmark_root, self.device) - self.arcface_model = ArcFace(self.arcface_root, self.device) - - def landmark_loss(self, id_landmark=None, gt_landmark=None, mask=None): - # id_landmark: [B, F, 106, 2] - # mask: [B, F] - mask = mask.unsqueeze(-1).unsqueeze(-1) # [B, F] -> [B, F, 1, 1] - error = torch.abs(id_landmark - gt_landmark) * mask - valid_frame_count = mask.sum() + 1e-8 # 避免除零 - loss = error.sum() / valid_frame_count / id_landmark.shape[-2] - return loss - - def embedding_loss(self, id_embedding=None, gt_embedding=None, mask=None): - # edbedding: [B, F, C] - # mask: [B, F] - cos_sim = F.cosine_similarity(id_embedding, gt_embedding, dim=2) #[B, F, C] - cos_loss = (1-cos_sim) * mask - valid_frame_count = mask.sum() + 1e-8 # 避免除零 - loss = cos_loss.sum() / valid_frame_count - return loss - - def pool_embedding_loss(self, id_embedding=None, gt_embedding=None, id_mask=None): - # edbedding: [B, F, C] - # mask: [B, F] - id_emb_expanded = id_embedding.unsqueeze(2) - gt_emb_expanded = gt_embedding.unsqueeze(1) - gt_mask = torch.ones(gt_embedding.shape[0], gt_embedding.shape[1]).to(id_mask.device) - gt_mask[:, 0] = 0 - is_all_zero = (gt_embedding == 0).all(dim=-1) - gt_mask[is_all_zero] = 0 - - cos_sim_all = F.cosine_similarity(id_emb_expanded, gt_emb_expanded, dim=3) - valid_mask = id_mask.unsqueeze(2) * gt_mask.unsqueeze(1) # [B, F, F] - - gt_valid_count = gt_mask.sum(dim=1) + 1e-8 - weight_matrix = valid_mask / (gt_valid_count.unsqueeze(1).unsqueeze(2) + 1e-8) - mean_similarities = (cos_sim_all * weight_matrix).sum(dim=2) - # mean_similarities = cos_sim_all.mean(dim=2) - cos_loss = mean_similarities * id_mask - valid_frame_count = id_mask.sum() + 1e-8 # 避免除零 - loss = cos_loss.sum() / valid_frame_count - return loss - -if __name__ == '__main__': - face_app = FaceAnalysis(root = "/data/models/antelopev2/",device=current_platform.device_type) - from decord import VideoReader - import time - import torch.nn.functional as F - vr = VideoReader("./video1.mp4") - print(len(vr)) - frames = [f.asnumpy() for f in vr] - print(frames[0].shape, frames[0].dtype, frames[0].max(), frames[0].min()) - h, w = frames[0].shape[:2] - id_landmark = [] - id_embedding = [] - id_mask = [] - index = 0 - index1 = 0 - all_start = time.time() - for f in frames: - f = torch.from_numpy(2*(f/255.)-1).permute(2,0,1).float().to(current_platform.device_type) #(3, h, w) - start = time.time() - bboxes, kpss = face_app.detection_model.detect(f) - end = time.time() - index += end-start - if bboxes.shape[0] > 0: - indexed_bboxes = [(i, x) for i, x in enumerate(bboxes)] - sorted_bboxes = sorted(indexed_bboxes, key=lambda item: (item[1][2] - item[1][0]) * (item[1][3] - item[1][1])) - max_index, max_bbox = sorted_bboxes[-1] - kps = kpss[max_index] - start = time.time() - face = Face(bbox=bboxes[max_index][0:4], kps=kps, det_score=bboxes[max_index][4]) - id_embedding.append(face_app.arcface_model.get(f, face)) - end = time.time() - index1 += end-start - id_mask.append(1) - else: - # id_landmark.append(torch.zeros(106, 2)) - id_embedding.append(torch.zeros(512)) - id_mask.append(0) - all_end = time.time() - print(all_end-all_start) - print(index) - id_embedding = torch.stack(id_embedding).unsqueeze(0) - face_embeddings = torch.randn_like(id_embedding) - id_mask = torch.tensor(id_mask).unsqueeze(0).to(id_embedding.device) - face_score = face_app.pool_embedding_loss(id_embedding, face_embeddings, id_mask) - \ No newline at end of file diff --git a/roll/pipeline/diffusion/reward_fl/reward_fl_config.py b/roll/pipeline/diffusion/reward_fl/reward_fl_config.py deleted file mode 100644 index ae42b9371..000000000 --- a/roll/pipeline/diffusion/reward_fl/reward_fl_config.py +++ /dev/null @@ -1,47 +0,0 @@ -import dataclasses -from dataclasses import dataclass, field - -from roll.configs.base_config import BaseConfig -from roll.configs.worker_config import WorkerConfig -from roll.utils.logging import get_logger - -logger = get_logger() - - -@dataclass -class RewardFLConfig(BaseConfig): - # global - global_template: str = field(default=None, metadata={"help": "The template of the global."}) - - train_batch_size: int = field( - default=8, - metadata={"help": "batch_size for one train step"}, - ) - - max_grad_norm: float = field(default=1.0, metadata={"help": "Maximum norm"}) - - actor_train: WorkerConfig = field( - default_factory=WorkerConfig, metadata={"help": "Configuration for the actor's training role."} - ) - - # reward_fl related - def __post_init__(self): - BaseConfig.__post_init__(self) - - # default worker_cls - if self.actor_train.worker_cls is None: - self.actor_train.worker_cls = "roll.pipeline.diffusion.reward_fl.actor_worker.ActorWorker" - - self.actor_train.training_args.output_dir = self.output_dir - - self.actor_train.name = "actor_train" - - def set_max_steps(self, max_steps: int): - self.max_steps = max_steps - self.actor_train.training_args.max_steps = max_steps - - logger.info(f"pipeline max_steps: {self.max_steps} to {max_steps}") - logger.info(f"actor train max_steps without dp_size: {self.actor_train.training_args.max_steps}") - - def to_dict(self): - return dataclasses.asdict(self) diff --git a/roll/pipeline/diffusion/reward_fl/reward_fl_pipeline.py b/roll/pipeline/diffusion/reward_fl/reward_fl_pipeline.py deleted file mode 100644 index 6f3d15fd3..000000000 --- a/roll/pipeline/diffusion/reward_fl/reward_fl_pipeline.py +++ /dev/null @@ -1,120 +0,0 @@ -from typing import Any, Dict, List -import os -import ray -import torch -import torchvision -from codetiming import Timer -from tqdm import tqdm -import numpy as np - -from roll.distributed.executor.cluster import Cluster -from roll.distributed.scheduler.protocol import DataProto -from roll.pipeline.base_pipeline import BasePipeline -from roll.pipeline.diffusion.reward_fl.reward_fl_config import RewardFLConfig -from roll.utils.logging import get_logger -from roll.utils.metrics.metrics_manager import MetricsManager - -from diffsynth.trainers.unified_dataset import UnifiedDataset - -logger = get_logger() - - -def collate_fn(examples): - video = torch.stack([ - torch.stack([ - torchvision.transforms.functional.to_tensor(frame) - for frame in example['video']], - dim=0 - ) - for example in examples - ], dim=0) - prompt = np.array([example['prompt'] for example in examples], dtype=object) - return {'video': video, 'prompt': prompt} - - -class RewardFLPipeline(BasePipeline): - def __init__(self, pipeline_config: RewardFLConfig): - super().__init__(pipeline_config) - self.pipeline_config = pipeline_config - - assert self.pipeline_config.max_steps > 0, "max_steps must be greater than 0" - self.pipeline_config.set_max_steps(max_steps=self.pipeline_config.max_steps) - - self.actor_train: Any = Cluster( - name=self.pipeline_config.actor_train.name, - worker_cls=self.pipeline_config.actor_train.worker_cls, - resource_manager=self.resource_manager, - worker_config=self.pipeline_config.actor_train, - ) - metadata_path = self.pipeline_config.actor_train.data_args.file_name - base_path = os.path.dirname(metadata_path) - dataset = UnifiedDataset( - base_path=base_path, - metadata_path=metadata_path, - data_file_keys=("video", "image"), - repeat=100, - main_data_operator=UnifiedDataset.default_video_operator( - base_path=base_path, - max_pixels=480*480, height=None, width=None, - height_division_factor=16, width_division_factor=16, - ), - ) - self.dataloader = torch.utils.data.DataLoader(dataset, batch_size=pipeline_config.train_batch_size, collate_fn=collate_fn) - refs: List[ray.ObjectRef] = [] - refs.extend(self.actor_train.initialize(pipeline_config=self.pipeline_config, blocking=False)) - ray.get(refs) - - self.set_checkpoint_clusters(self.actor_train) - - @torch.no_grad() - def run(self): - global_step = 0 - metrics_mgr = MetricsManager() - - for epoch in range(int(self.pipeline_config.actor_train.training_args.num_train_epochs)): - logger.info(f"epoch {epoch} start...") - for batch_dict in tqdm(self.dataloader): - if global_step <= self.state.step: - global_step += 1 - continue - - logger.info(f"pipeline step {global_step} start...") - metrics_mgr.clear_metrics() - - with Timer(name="step_total", logger=None) as step_total_timer: - batch_dict: Dict - batch: DataProto = DataProto.from_single_dict(batch_dict) - batch.meta_info = {"global_step": global_step, "is_offload_states": False, - "is_offload_optimizer_states_in_train_step": False, "loss_mask_keys": []} - - with Timer(name="actor_train", logger=None) as actor_train_timer: - actor_train_refs = self.actor_train.train_step(batch, blocking=False) - actor_train_refs: DataProto = DataProto.materialize_concat(data_refs=actor_train_refs) - metrics_mgr.add_reduced_metrics(actor_train_refs.meta_info.pop("metrics", {})) - - metrics_mgr.add_metric("time/actor_train", actor_train_timer.last) - - metrics_mgr.add_metric("time/step_total", step_total_timer.last) - - metrics = metrics_mgr.get_metrics() - - self.state.step = global_step - self.state.log_history.append(metrics) - self.tracker.log(values=metrics, step=global_step) - self.do_checkpoint(global_step=global_step) - - timeline_dir = os.path.join(self.pipeline_config.profiler_output_dir, "timeline") - os.makedirs(timeline_dir, exist_ok=True) - ray.timeline( - filename=os.path.join(timeline_dir, f"timeline-step-{global_step}.json"), - ) - - logger.info(f"pipeline step {global_step} finished") - global_step += 1 - if global_step >= self.pipeline_config.max_steps: - break - - if global_step >= self.pipeline_config.max_steps: - break - - logger.info("pipeline complete!") diff --git a/roll/pipeline/diffusion/reward_fl/wan_video_vae.py b/roll/pipeline/diffusion/reward_fl/wan_video_vae.py deleted file mode 100644 index bb92eb023..000000000 --- a/roll/pipeline/diffusion/reward_fl/wan_video_vae.py +++ /dev/null @@ -1,1416 +0,0 @@ -from einops import rearrange, repeat - -import torch -import torch.nn as nn -import torch.nn.functional as F -from tqdm import tqdm - -CACHE_T = 2 - - -def check_is_instance(model, module_class): - if isinstance(model, module_class): - return True - if hasattr(model, "module") and isinstance(model.module, module_class): - return True - return False - - -def block_causal_mask(x, block_size): - # params - b, n, s, _, device = *x.size(), x.device - assert s % block_size == 0 - num_blocks = s // block_size - - # build mask - mask = torch.zeros(b, n, s, s, dtype=torch.bool, device=device) - for i in range(num_blocks): - mask[:, :, - i * block_size:(i + 1) * block_size, :(i + 1) * block_size] = 1 - return mask - - -class CausalConv3d(nn.Conv3d): - """ - Causal 3d convolusion. - """ - - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - self._padding = (self.padding[2], self.padding[2], self.padding[1], - self.padding[1], 2 * self.padding[0], 0) - self.padding = (0, 0, 0) - # 获取时间轴上的 kernel size) - self.time_kernel_size = self.kernel_size[0] - - def forward(self, x, cache_x=None): - padding = list(self._padding) - - # wan cache 策略 - if cache_x is not None and self._padding[4] > 0: - cache_x = cache_x.to(x.device) - x = torch.cat([cache_x, x], dim=2) - padding[4] -= cache_x.shape[2] - x = F.pad(x, padding) - - return super().forward(x) - - -class RMS_norm(nn.Module): - - def __init__(self, dim, channel_first=True, images=True, bias=False): - super().__init__() - broadcastable_dims = (1, 1, 1) if not images else (1, 1) - shape = (dim, *broadcastable_dims) if channel_first else (dim,) - - self.channel_first = channel_first - self.scale = dim**0.5 - self.gamma = nn.Parameter(torch.ones(shape)) - self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0. - - def forward(self, x): - return F.normalize( - x, dim=(1 if self.channel_first else - -1)) * self.scale * self.gamma + self.bias - - -class Upsample(nn.Upsample): - - def forward(self, x): - """ - Fix bfloat16 support for nearest neighbor interpolation. - """ - return super().forward(x.float()).type_as(x) - - -class Resample(nn.Module): - - def __init__(self, dim, mode): - assert mode in ('none', 'upsample2d', 'upsample3d', 'downsample2d', - 'downsample3d') - super().__init__() - self.dim = dim - self.mode = mode - - # layers - if mode == 'upsample2d': - self.resample = nn.Sequential( - Upsample(scale_factor=(2., 2.), mode='nearest-exact'), - nn.Conv2d(dim, dim // 2, 3, padding=1)) - elif mode == 'upsample3d': - self.resample = nn.Sequential( - Upsample(scale_factor=(2., 2.), mode='nearest-exact'), - nn.Conv2d(dim, dim // 2, 3, padding=1)) - causal_conv3d = CausalConv3d(dim, - dim * 2, (3, 1, 1), - padding=(1, 0, 0)) - - # decoder no cache set - if hasattr(self, 'decoder') and getattr(self, 'decoder'): - setattr(causal_conv3d, 'decoder', True) - self.time_conv = causal_conv3d - elif mode == 'downsample2d': - self.resample = nn.Sequential( - nn.ZeroPad2d((0, 1, 0, 1)), - nn.Conv2d(dim, dim, 3, stride=(2, 2))) - elif mode == 'downsample3d': - self.resample = nn.Sequential( - nn.ZeroPad2d((0, 1, 0, 1)), - nn.Conv2d(dim, dim, 3, stride=(2, 2))) - # decoder no cache set - causal_conv3d = CausalConv3d(dim, - dim, (3, 1, 1), - stride=(2, 1, 1), - padding=(0, 0, 0)) - - # decoder no cache set - if hasattr(self, 'decoder') and getattr(self, 'decoder'): - setattr(causal_conv3d, 'decoder', True) - self.time_conv = causal_conv3d - else: - self.resample = nn.Identity() - - def forward(self, x, feat_cache=None, feat_idx=[0]): - b, c, t, h, w = x.size() - # upsample3d 时间维度膨胀 - if self.mode == 'upsample3d': - # decoder only no cache - if hasattr(self, 'decoder') and getattr(self, 'decoder'): - # 判断是否为首帧, 如果是 decoder 首帧, 不做任何处理即可, 当即退出, 非首帧做其他处理 - x1 = x[:,:,1:] - x1 = self.time_conv(x1) - x1 = x1.reshape(b, 2, c, t-1, h, w) - x1 = torch.stack((x1[:, 0, :, :, :, :], x1[:, 1, :, :, :, :]), 3) - x1 = x1.reshape(b, c, (t-1) * 2, h, w) - x = torch.cat([x[:,:,0:1], x1], dim=2) - elif feat_cache is not None: - idx = feat_idx[0] - if feat_cache[idx] is None: - feat_cache[idx] = 'Rep' - feat_idx[0] += 1 - else: - - cache_x = x[:, :, -CACHE_T:, :, :].clone() - if cache_x.shape[2] < 2 and feat_cache[ - idx] is not None and feat_cache[idx] != 'Rep': - # cache last frame of last two chunk - cache_x = torch.cat([ - feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( - cache_x.device), cache_x - ], - dim=2) - if cache_x.shape[2] < 2 and feat_cache[ - idx] is not None and feat_cache[idx] == 'Rep': - cache_x = torch.cat([ - torch.zeros_like(cache_x).to(cache_x.device), - cache_x - ], - dim=2) - if feat_cache[idx] == 'Rep': - x = self.time_conv(x) - else: - x = self.time_conv(x, feat_cache[idx]) - feat_cache[idx] = cache_x - feat_idx[0] += 1 - - x = x.reshape(b, 2, c, t, h, w) - x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), - 3) - x = x.reshape(b, c, t * 2, h, w) - # resample 采样, 空间维度膨胀 - t = x.shape[2] - x = rearrange(x, 'b c t h w -> (b t) c h w') - x = self.resample(x) - x = rearrange(x, '(b t) c h w -> b c t h w', t=t) - - if self.mode == 'downsample3d': - if hasattr(self, 'decoder') and getattr(self, 'decoder'): - x = self.time_conv(x) - elif feat_cache is not None: - idx = feat_idx[0] - if feat_cache[idx] is None: - feat_cache[idx] = x.clone() - feat_idx[0] += 1 - else: - cache_x = x[:, :, -1:, :, :].clone() - x = self.time_conv( - torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2)) - feat_cache[idx] = cache_x - feat_idx[0] += 1 - return x - - def init_weight(self, conv): - conv_weight = conv.weight - nn.init.zeros_(conv_weight) - c1, c2, t, h, w = conv_weight.size() - one_matrix = torch.eye(c1, c2) - init_matrix = one_matrix - nn.init.zeros_(conv_weight) - conv_weight.data[:, :, 1, 0, 0] = init_matrix - conv.weight.data.copy_(conv_weight) - nn.init.zeros_(conv.bias.data) - - def init_weight2(self, conv): - conv_weight = conv.weight.data - nn.init.zeros_(conv_weight) - c1, c2, t, h, w = conv_weight.size() - init_matrix = torch.eye(c1 // 2, c2) - conv_weight[:c1 // 2, :, -1, 0, 0] = init_matrix - conv_weight[c1 // 2:, :, -1, 0, 0] = init_matrix - conv.weight.data.copy_(conv_weight) - nn.init.zeros_(conv.bias.data) - - - -def patchify(x, patch_size): - if patch_size == 1: - return x - if x.dim() == 4: - x = rearrange(x, "b c (h q) (w r) -> b (c r q) h w", q=patch_size, r=patch_size) - elif x.dim() == 5: - x = rearrange(x, - "b c f (h q) (w r) -> b (c r q) f h w", - q=patch_size, - r=patch_size) - else: - raise ValueError(f"Invalid input shape: {x.shape}") - return x - - -def unpatchify(x, patch_size): - if patch_size == 1: - return x - if x.dim() == 4: - x = rearrange(x, "b (c r q) h w -> b c (h q) (w r)", q=patch_size, r=patch_size) - elif x.dim() == 5: - x = rearrange(x, - "b (c r q) f h w -> b c f (h q) (w r)", - q=patch_size, - r=patch_size) - return x - - -class Resample38(Resample): - - def __init__(self, dim, mode): - assert mode in ( - "none", - "upsample2d", - "upsample3d", - "downsample2d", - "downsample3d", - ) - super(Resample, self).__init__() - self.dim = dim - self.mode = mode - - # layers - if mode == "upsample2d": - self.resample = nn.Sequential( - Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), - nn.Conv2d(dim, dim, 3, padding=1), - ) - elif mode == "upsample3d": - self.resample = nn.Sequential( - Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), - nn.Conv2d(dim, dim, 3, padding=1), - ) - self.time_conv = CausalConv3d(dim, dim * 2, (3, 1, 1), padding=(1, 0, 0)) - elif mode == "downsample2d": - self.resample = nn.Sequential( - nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2)) - ) - elif mode == "downsample3d": - self.resample = nn.Sequential( - nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2)) - ) - self.time_conv = CausalConv3d( - dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0) - ) - else: - self.resample = nn.Identity() - -class ResidualBlock(nn.Module): - - def __init__(self, in_dim, out_dim, dropout=0.0): - super().__init__() - self.in_dim = in_dim - self.out_dim = out_dim - - # layers - causal_conv3d_block1 = CausalConv3d(in_dim, out_dim, 3, padding=1) - causal_conv3d_block2 = CausalConv3d(out_dim, out_dim, 3, padding=1) - causal_conv3d_shortcut = CausalConv3d(in_dim, out_dim, 1) - if hasattr(self, 'decoder') and getattr(self, 'decoder'): - setattr(causal_conv3d_block1, 'decoder', True) - setattr(causal_conv3d_block2, 'decoder', True) - setattr(causal_conv3d_shortcut, 'decoder', True) - self.residual = nn.Sequential( - RMS_norm(in_dim, images=False), nn.SiLU(), - causal_conv3d_block1, - RMS_norm(out_dim, images=False), nn.SiLU(), nn.Dropout(dropout), - causal_conv3d_block2) - self.shortcut = causal_conv3d_shortcut \ - if in_dim != out_dim else nn.Identity() - - # decoder 部分 ResidualBlock 内 CausalConv3d 做帧填充, feature_cache = None - def forward(self, x, feat_cache=None, feat_idx=[0]): - h = self.shortcut(x) - for layer in self.residual: - if hasattr(self, 'decoder') and getattr(self, 'decoder'): - x = layer(x) - elif check_is_instance(layer, CausalConv3d) and feat_cache is not None: - idx = feat_idx[0] - cache_x = x[:, :, -CACHE_T:, :, :].clone() - if cache_x.shape[2] < 2 and feat_cache[idx] is not None: - # cache last frame of last two chunk - cache_x = torch.cat([ - feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( - cache_x.device), cache_x - ], - dim=2) - x = layer(x, feat_cache[idx]) - feat_cache[idx] = cache_x - feat_idx[0] += 1 - else: - x = layer(x) - return x + h - - -class AttentionBlock(nn.Module): - """ - Causal self-attention with a single head. - """ - - def __init__(self, dim): - super().__init__() - self.dim = dim - - # layers - self.norm = RMS_norm(dim) - self.to_qkv = nn.Conv2d(dim, dim * 3, 1) - self.proj = nn.Conv2d(dim, dim, 1) - - # zero out the last layer params - nn.init.zeros_(self.proj.weight) - - def forward(self, x): - identity = x - b, c, t, h, w = x.size() - x = rearrange(x, 'b c t h w -> (b t) c h w') - x = self.norm(x) - # compute query, key, value - q, k, v = self.to_qkv(x).reshape(b * t, 1, c * 3, -1).permute( - 0, 1, 3, 2).contiguous().chunk(3, dim=-1) - - # apply attention - x = F.scaled_dot_product_attention( - q, - k, - v, - #attn_mask=block_causal_mask(q, block_size=h * w) - ) - x = x.squeeze(1).permute(0, 2, 1).reshape(b * t, c, h, w) - - # output - x = self.proj(x) - x = rearrange(x, '(b t) c h w-> b c t h w', t=t) - return x + identity - - -class AvgDown3D(nn.Module): - def __init__( - self, - in_channels, - out_channels, - factor_t, - factor_s=1, - ): - super().__init__() - self.in_channels = in_channels - self.out_channels = out_channels - self.factor_t = factor_t - self.factor_s = factor_s - self.factor = self.factor_t * self.factor_s * self.factor_s - - assert in_channels * self.factor % out_channels == 0 - self.group_size = in_channels * self.factor // out_channels - - def forward(self, x: torch.Tensor) -> torch.Tensor: - pad_t = (self.factor_t - x.shape[2] % self.factor_t) % self.factor_t - pad = (0, 0, 0, 0, pad_t, 0) - x = F.pad(x, pad) - B, C, T, H, W = x.shape - x = x.view( - B, - C, - T // self.factor_t, - self.factor_t, - H // self.factor_s, - self.factor_s, - W // self.factor_s, - self.factor_s, - ) - x = x.permute(0, 1, 3, 5, 7, 2, 4, 6).contiguous() - x = x.view( - B, - C * self.factor, - T // self.factor_t, - H // self.factor_s, - W // self.factor_s, - ) - x = x.view( - B, - self.out_channels, - self.group_size, - T // self.factor_t, - H // self.factor_s, - W // self.factor_s, - ) - x = x.mean(dim=2) - return x - - -class DupUp3D(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: int, - factor_t, - factor_s=1, - ): - super().__init__() - self.in_channels = in_channels - self.out_channels = out_channels - - self.factor_t = factor_t - self.factor_s = factor_s - self.factor = self.factor_t * self.factor_s * self.factor_s - - assert out_channels * self.factor % in_channels == 0 - self.repeats = out_channels * self.factor // in_channels - - def forward(self, x: torch.Tensor, first_chunk=False) -> torch.Tensor: - x = x.repeat_interleave(self.repeats, dim=1) - x = x.view( - x.size(0), - self.out_channels, - self.factor_t, - self.factor_s, - self.factor_s, - x.size(2), - x.size(3), - x.size(4), - ) - x = x.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous() - x = x.view( - x.size(0), - self.out_channels, - x.size(2) * self.factor_t, - x.size(4) * self.factor_s, - x.size(6) * self.factor_s, - ) - if first_chunk: - x = x[:, :, self.factor_t - 1 :, :, :] - return x - - -class Down_ResidualBlock(nn.Module): - def __init__( - self, in_dim, out_dim, dropout, mult, temperal_downsample=False, down_flag=False - ): - super().__init__() - - # Shortcut path with downsample - self.avg_shortcut = AvgDown3D( - in_dim, - out_dim, - factor_t=2 if temperal_downsample else 1, - factor_s=2 if down_flag else 1, - ) - - # Main path with residual blocks and downsample - downsamples = [] - for _ in range(mult): - downsamples.append(ResidualBlock(in_dim, out_dim, dropout)) - in_dim = out_dim - - # Add the final downsample block - if down_flag: - mode = "downsample3d" if temperal_downsample else "downsample2d" - downsamples.append(Resample38(out_dim, mode=mode)) - - self.downsamples = nn.Sequential(*downsamples) - - def forward(self, x, feat_cache=None, feat_idx=[0]): - x_copy = x.clone() - for module in self.downsamples: - x = module(x, feat_cache, feat_idx) - - return x + self.avg_shortcut(x_copy) - - -class Up_ResidualBlock(nn.Module): - def __init__( - self, in_dim, out_dim, dropout, mult, temperal_upsample=False, up_flag=False - ): - super().__init__() - # Shortcut path with upsample - if up_flag: - self.avg_shortcut = DupUp3D( - in_dim, - out_dim, - factor_t=2 if temperal_upsample else 1, - factor_s=2 if up_flag else 1, - ) - else: - self.avg_shortcut = None - - # Main path with residual blocks and upsample - upsamples = [] - for _ in range(mult): - upsamples.append(ResidualBlock(in_dim, out_dim, dropout)) - in_dim = out_dim - - # Add the final upsample block - if up_flag: - mode = "upsample3d" if temperal_upsample else "upsample2d" - upsamples.append(Resample38(out_dim, mode=mode)) - - self.upsamples = nn.Sequential(*upsamples) - - def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False): - x_main = x.clone() - for module in self.upsamples: - x_main = module(x_main, feat_cache, feat_idx) - if self.avg_shortcut is not None: - x_shortcut = self.avg_shortcut(x, first_chunk) - return x_main + x_shortcut - else: - return x_main - - -class Encoder3d(nn.Module): - - def __init__(self, - dim=128, - z_dim=4, - dim_mult=[1, 2, 4, 4], - num_res_blocks=2, - attn_scales=[], - temperal_downsample=[True, True, False], - dropout=0.0): - super().__init__() - self.dim = dim - self.z_dim = z_dim - self.dim_mult = dim_mult - self.num_res_blocks = num_res_blocks - self.attn_scales = attn_scales - self.temperal_downsample = temperal_downsample - - # dimensions - dims = [dim * u for u in [1] + dim_mult] - scale = 1.0 - - # init block - self.conv1 = CausalConv3d(3, dims[0], 3, padding=1) - - # downsample blocks - downsamples = [] - for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): - # residual (+attention) blocks - for _ in range(num_res_blocks): - downsamples.append(ResidualBlock(in_dim, out_dim, dropout)) - if scale in attn_scales: - downsamples.append(AttentionBlock(out_dim)) - in_dim = out_dim - - # downsample block - if i != len(dim_mult) - 1: - mode = 'downsample3d' if temperal_downsample[ - i] else 'downsample2d' - downsamples.append(Resample(out_dim, mode=mode)) - scale /= 2.0 - self.downsamples = nn.Sequential(*downsamples) - - # middle blocks - self.middle = nn.Sequential(ResidualBlock(out_dim, out_dim, dropout), - AttentionBlock(out_dim), - ResidualBlock(out_dim, out_dim, dropout)) - - # output blocks - self.head = nn.Sequential(RMS_norm(out_dim, images=False), nn.SiLU(), - CausalConv3d(out_dim, z_dim, 3, padding=1)) - - def forward(self, x, feat_cache=None, feat_idx=[0]): - if feat_cache is not None: - idx = feat_idx[0] - cache_x = x[:, :, -CACHE_T:, :, :].clone() - if cache_x.shape[2] < 2 and feat_cache[idx] is not None: - # cache last frame of last two chunk - cache_x = torch.cat([ - feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( - cache_x.device), cache_x - ], - dim=2) - x = self.conv1(x, feat_cache[idx]) - feat_cache[idx] = cache_x - feat_idx[0] += 1 - else: - x = self.conv1(x) - - ## downsamples - for layer in self.downsamples: - if feat_cache is not None: - x = layer(x, feat_cache, feat_idx) - else: - x = layer(x) - - ## middle - for layer in self.middle: - if check_is_instance(layer, ResidualBlock) and feat_cache is not None: - x = layer(x, feat_cache, feat_idx) - else: - x = layer(x) - - ## head - for layer in self.head: - if check_is_instance(layer, CausalConv3d) and feat_cache is not None: - idx = feat_idx[0] - cache_x = x[:, :, -CACHE_T:, :, :].clone() - if cache_x.shape[2] < 2 and feat_cache[idx] is not None: - # cache last frame of last two chunk - cache_x = torch.cat([ - feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( - cache_x.device), cache_x - ], - dim=2) - x = layer(x, feat_cache[idx]) - feat_cache[idx] = cache_x - feat_idx[0] += 1 - else: - x = layer(x) - return x - - -class Encoder3d_38(nn.Module): - - def __init__(self, - dim=128, - z_dim=4, - dim_mult=[1, 2, 4, 4], - num_res_blocks=2, - attn_scales=[], - temperal_downsample=[False, True, True], - dropout=0.0): - super().__init__() - self.dim = dim - self.z_dim = z_dim - self.dim_mult = dim_mult - self.num_res_blocks = num_res_blocks - self.attn_scales = attn_scales - self.temperal_downsample = temperal_downsample - - # dimensions - dims = [dim * u for u in [1] + dim_mult] - scale = 1.0 - - # init block - self.conv1 = CausalConv3d(12, dims[0], 3, padding=1) - - # downsample blocks - downsamples = [] - for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): - t_down_flag = ( - temperal_downsample[i] if i < len(temperal_downsample) else False - ) - downsamples.append( - Down_ResidualBlock( - in_dim=in_dim, - out_dim=out_dim, - dropout=dropout, - mult=num_res_blocks, - temperal_downsample=t_down_flag, - down_flag=i != len(dim_mult) - 1, - ) - ) - scale /= 2.0 - self.downsamples = nn.Sequential(*downsamples) - - # middle blocks - self.middle = nn.Sequential( - ResidualBlock(out_dim, out_dim, dropout), - AttentionBlock(out_dim), - ResidualBlock(out_dim, out_dim, dropout), - ) - - # # output blocks - self.head = nn.Sequential( - RMS_norm(out_dim, images=False), - nn.SiLU(), - CausalConv3d(out_dim, z_dim, 3, padding=1), - ) - - - def forward(self, x, feat_cache=None, feat_idx=[0]): - - if feat_cache is not None: - idx = feat_idx[0] - cache_x = x[:, :, -CACHE_T:, :, :].clone() - if cache_x.shape[2] < 2 and feat_cache[idx] is not None: - cache_x = torch.cat( - [ - feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), - cache_x, - ], - dim=2, - ) - x = self.conv1(x, feat_cache[idx]) - feat_cache[idx] = cache_x - feat_idx[0] += 1 - else: - x = self.conv1(x) - - ## downsamples - for layer in self.downsamples: - if feat_cache is not None: - x = layer(x, feat_cache, feat_idx) - else: - x = layer(x) - - ## middle - for layer in self.middle: - if isinstance(layer, ResidualBlock) and feat_cache is not None: - x = layer(x, feat_cache, feat_idx) - else: - x = layer(x) - - ## head - for layer in self.head: - if isinstance(layer, CausalConv3d) and feat_cache is not None: - idx = feat_idx[0] - cache_x = x[:, :, -CACHE_T:, :, :].clone() - if cache_x.shape[2] < 2 and feat_cache[idx] is not None: - cache_x = torch.cat( - [ - feat_cache[idx][:, :, -1, :, :] - .unsqueeze(2) - .to(cache_x.device), - cache_x, - ], - dim=2, - ) - x = layer(x, feat_cache[idx]) - feat_cache[idx] = cache_x - feat_idx[0] += 1 - else: - x = layer(x) - - return x - - -class Decoder3d(nn.Module): - - def __init__(self, - dim=128, - z_dim=4, - dim_mult=[1, 2, 4, 4], - num_res_blocks=2, - attn_scales=[], - temperal_upsample=[False, True, True], - dropout=0.0): - super().__init__() - self.dim = dim - self.z_dim = z_dim - self.dim_mult = dim_mult - self.num_res_blocks = num_res_blocks - self.attn_scales = attn_scales - self.temperal_upsample = temperal_upsample - - # dimensions - dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]] - scale = 1.0 / 2**(len(dim_mult) - 2) - - # init block - self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1) - setattr(self.conv1, "decoder", True) - - # middle blocks - res_block_1 = ResidualBlock(dims[0], dims[0], dropout) - attn_block_2 = AttentionBlock(dims[0]) - res_block_3 = ResidualBlock(dims[0], dims[0], dropout) - setattr(res_block_1, 'decoder', True) - setattr(attn_block_2, 'decoder', True) - setattr(res_block_3, 'decoder', True) - self.middle = nn.Sequential(res_block_1, - attn_block_2, - res_block_3) - - # upsample blocks - upsamples = [] - for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): - # residual (+attention) blocks - if i == 1 or i == 2 or i == 3: - in_dim = in_dim // 2 - for _ in range(num_res_blocks + 1): - res_block = ResidualBlock(in_dim, out_dim, dropout) - attn_block = AttentionBlock(out_dim) - setattr(res_block, 'decoder', True) - setattr(attn_block, 'decoder', True) - upsamples.append(res_block) - if scale in attn_scales: - upsamples.append(attn_block) - in_dim = out_dim - - # upsample block - if i != len(dim_mult) - 1: - mode = 'upsample3d' if temperal_upsample[i] else 'upsample2d' - resample = Resample(out_dim, mode=mode) - setattr(resample, 'decoder', True) - upsamples.append(resample) - scale *= 2.0 - self.upsamples = nn.Sequential(*upsamples) - - # output blocks - causal_conv3d = CausalConv3d(out_dim, 3, 3, padding=1) - setattr(causal_conv3d, 'decoder', True) - self.head = nn.Sequential(RMS_norm(out_dim, images=False), nn.SiLU(), - causal_conv3d) - - # docoder3d 修改为 no cache 逻辑 - def forward(self, x): - ## conv1 - x = self.conv1(x) - - ## middle - for layer in self.middle: - x = layer(x) - - ## unsample - for layer in self.upsamples: - x = layer(x) - - ## head - for layer in self.head: - x = layer(x) - return x - - - -class Decoder3d_38(nn.Module): - - def __init__(self, - dim=128, - z_dim=4, - dim_mult=[1, 2, 4, 4], - num_res_blocks=2, - attn_scales=[], - temperal_upsample=[False, True, True], - dropout=0.0): - super().__init__() - self.dim = dim - self.z_dim = z_dim - self.dim_mult = dim_mult - self.num_res_blocks = num_res_blocks - self.attn_scales = attn_scales - self.temperal_upsample = temperal_upsample - - # dimensions - dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]] - scale = 1.0 / 2 ** (len(dim_mult) - 2) - # init block - self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1) - - # middle blocks - self.middle = nn.Sequential(ResidualBlock(dims[0], dims[0], dropout), - AttentionBlock(dims[0]), - ResidualBlock(dims[0], dims[0], dropout)) - - # upsample blocks - upsamples = [] - for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): - t_up_flag = temperal_upsample[i] if i < len(temperal_upsample) else False - upsamples.append( - Up_ResidualBlock(in_dim=in_dim, - out_dim=out_dim, - dropout=dropout, - mult=num_res_blocks + 1, - temperal_upsample=t_up_flag, - up_flag=i != len(dim_mult) - 1)) - self.upsamples = nn.Sequential(*upsamples) - - # output blocks - self.head = nn.Sequential(RMS_norm(out_dim, images=False), nn.SiLU(), - CausalConv3d(out_dim, 12, 3, padding=1)) - - - def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False): - if feat_cache is not None: - idx = feat_idx[0] - cache_x = x[:, :, -CACHE_T:, :, :].clone() - if cache_x.shape[2] < 2 and feat_cache[idx] is not None: - cache_x = torch.cat( - [ - feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), - cache_x, - ], - dim=2, - ) - x = self.conv1(x, feat_cache[idx]) - feat_cache[idx] = cache_x - feat_idx[0] += 1 - else: - x = self.conv1(x) - - for layer in self.middle: - if check_is_instance(layer, ResidualBlock) and feat_cache is not None: - x = layer(x, feat_cache, feat_idx) - else: - x = layer(x) - - ## upsamples - for layer in self.upsamples: - if feat_cache is not None: - x = layer(x, feat_cache, feat_idx, first_chunk) - else: - x = layer(x) - - ## head - for layer in self.head: - if check_is_instance(layer, CausalConv3d) and feat_cache is not None: - idx = feat_idx[0] - cache_x = x[:, :, -CACHE_T:, :, :].clone() - if cache_x.shape[2] < 2 and feat_cache[idx] is not None: - cache_x = torch.cat( - [ - feat_cache[idx][:, :, -1, :, :] - .unsqueeze(2) - .to(cache_x.device), - cache_x, - ], - dim=2, - ) - x = layer(x, feat_cache[idx]) - feat_cache[idx] = cache_x - feat_idx[0] += 1 - else: - x = layer(x) - return x - - -def count_conv3d(model): - count = 0 - for m in model.modules(): - if isinstance(m, CausalConv3d): - count += 1 - return count - - -class VideoVAE_(nn.Module): - - def __init__(self, - dim=96, - z_dim=16, - dim_mult=[1, 2, 4, 4], - num_res_blocks=2, - attn_scales=[], - temperal_downsample=[False, True, True], - dropout=0.0): - super().__init__() - self.dim = dim - self.z_dim = z_dim - self.dim_mult = dim_mult - self.num_res_blocks = num_res_blocks - self.attn_scales = attn_scales - self.temperal_downsample = temperal_downsample - self.temperal_upsample = temperal_downsample[::-1] - - # modules - self.encoder = Encoder3d(dim, z_dim * 2, dim_mult, num_res_blocks, - attn_scales, self.temperal_downsample, dropout) - self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1) - self.conv2 = CausalConv3d(z_dim, z_dim, 1) - self.decoder = Decoder3d(dim, z_dim, dim_mult, num_res_blocks, - attn_scales, self.temperal_upsample, dropout) - - def forward(self, x): - mu, log_var = self.encode(x) - z = self.reparameterize(mu, log_var) - x_recon = self.decode(z) - return x_recon, mu, log_var - - def encode(self, x, scale): - self.clear_cache() - ## cache - t = x.shape[2] - iter_ = 1 + (t - 1) // 4 - - for i in range(iter_): - self._enc_conv_idx = [0] - if i == 0: - out = self.encoder(x[:, :, :1, :, :], - feat_cache=self._enc_feat_map, - feat_idx=self._enc_conv_idx) - else: - out_ = self.encoder(x[:, :, 1 + 4 * (i - 1):1 + 4 * i, :, :], - feat_cache=self._enc_feat_map, - feat_idx=self._enc_conv_idx) - out = torch.cat([out, out_], 2) - mu, log_var = self.conv1(out).chunk(2, dim=1) - if isinstance(scale[0], torch.Tensor): - scale = [s.to(dtype=mu.dtype, device=mu.device) for s in scale] - mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view( - 1, self.z_dim, 1, 1, 1) - else: - scale = scale.to(dtype=mu.dtype, device=mu.device) - mu = (mu - scale[0]) * scale[1] - return mu - - def decode(self, z, scale): - # z: [b,c,t,h,w] - if isinstance(scale[0], torch.Tensor): - scale = [s.to(dtype=z.dtype, device=z.device) for s in scale] - z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view( - 1, self.z_dim, 1, 1, 1) - else: - scale = scale.to(dtype=z.dtype, device=z.device) - z = z / scale[1] + scale[0] - - # iter_ = z.shape[2] - # x = self.conv2(z) - # for i in range(iter_): - # self._conv_idx = [0] - # if i == 0: - # out = self.decoder(x[:, :, i:i + 1, :, :], - # feat_cache=self._feat_map, - # feat_idx=self._conv_idx) - # else: - # out_ = self.decoder(x[:, :, i:i + 1, :, :], - # feat_cache=self._feat_map, - # feat_idx=self._conv_idx) - # out = torch.cat([out, out_], 2) # may add tensor offload - - # 修改此处 decoder forward 中 cache 计算逻辑 - def create_custom_forward(module): - def custom_forward(*inputs): - return module(*inputs) - return custom_forward - - x = torch.utils.checkpoint.checkpoint( - create_custom_forward(self.conv2), - z, - use_reentrant=False, - ) - - out = torch.utils.checkpoint.checkpoint( - create_custom_forward(self.decoder), - x, use_reentrant=False, - ) - return out - - def reparameterize(self, mu, log_var): - std = torch.exp(0.5 * log_var) - eps = torch.randn_like(std) - return eps * std + mu - - def sample(self, imgs, deterministic=False): - mu, log_var = self.encode(imgs) - if deterministic: - return mu - std = torch.exp(0.5 * log_var.clamp(-30.0, 20.0)) - return mu + std * torch.randn_like(std) - - def clear_cache(self): - self._conv_num = count_conv3d(self.decoder) - self._conv_idx = [0] - self._feat_map = [None] * self._conv_num - # cache encode - self._enc_conv_num = count_conv3d(self.encoder) - self._enc_conv_idx = [0] - self._enc_feat_map = [None] * self._enc_conv_num - - -class WanVideoVAE(nn.Module): - - def __init__(self, z_dim=16): - super().__init__() - - mean = [ - -0.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653, -0.1517, 1.5508, - 0.4134, -0.0715, 0.5517, -0.3632, -0.1922, -0.9497, 0.2503, -0.2921 - ] - std = [ - 2.8184, 1.4541, 2.3275, 2.6558, 1.2196, 1.7708, 2.6052, 2.0743, - 3.2687, 2.1526, 2.8652, 1.5579, 1.6382, 1.1253, 2.8251, 1.9160 - ] - self.mean = torch.tensor(mean) - self.std = torch.tensor(std) - self.scale = [self.mean, 1.0 / self.std] - - # init model - self.model = VideoVAE_(z_dim=z_dim).eval().requires_grad_(False) - self.upsampling_factor = 8 - self.z_dim = z_dim - - - def build_1d_mask(self, length, left_bound, right_bound, border_width): - x = torch.ones((length,)) - if not left_bound: - x[:border_width] = (torch.arange(border_width) + 1) / border_width - if not right_bound: - x[-border_width:] = torch.flip((torch.arange(border_width) + 1) / border_width, dims=(0,)) - return x - - - def build_mask(self, data, is_bound, border_width): - _, _, _, H, W = data.shape - h = self.build_1d_mask(H, is_bound[0], is_bound[1], border_width[0]) - w = self.build_1d_mask(W, is_bound[2], is_bound[3], border_width[1]) - - h = repeat(h, "H -> H W", H=H, W=W) - w = repeat(w, "W -> H W", H=H, W=W) - - mask = torch.stack([h, w]).min(dim=0).values - mask = rearrange(mask, "H W -> 1 1 1 H W") - return mask - - - def tiled_decode(self, hidden_states, device, tile_size, tile_stride): - _, _, T, H, W = hidden_states.shape - size_h, size_w = tile_size - stride_h, stride_w = tile_stride - - # Split tasks - tasks = [] - for h in range(0, H, stride_h): - if (h-stride_h >= 0 and h-stride_h+size_h >= H): continue - for w in range(0, W, stride_w): - if (w-stride_w >= 0 and w-stride_w+size_w >= W): continue - h_, w_ = h + size_h, w + size_w - tasks.append((h, h_, w, w_)) - - computation_device = device - - out_T = T * 4 - 3 - weight = torch.zeros((1, 1, out_T, H * self.upsampling_factor, W * self.upsampling_factor), dtype=hidden_states.dtype, device=device) - values = torch.zeros((1, 3, out_T, H * self.upsampling_factor, W * self.upsampling_factor), dtype=hidden_states.dtype, device=device) - - for h, h_, w, w_ in tqdm(tasks, desc="VAE decoding"): - hidden_states_batch = hidden_states[:, :, :, h:h_, w:w_].to(computation_device) - hidden_states_batch = self.model.decode(hidden_states_batch, self.scale).to(device) - - mask = self.build_mask( - hidden_states_batch, - is_bound=(h==0, h_>=H, w==0, w_>=W), - border_width=((size_h - stride_h) * self.upsampling_factor, (size_w - stride_w) * self.upsampling_factor) - ).to(dtype=hidden_states.dtype, device=device) - - target_h = h * self.upsampling_factor - target_w = w * self.upsampling_factor - values[ - :, - :, - :, - target_h:target_h + hidden_states_batch.shape[3], - target_w:target_w + hidden_states_batch.shape[4], - ] += hidden_states_batch * mask - weight[ - :, - :, - :, - target_h: target_h + hidden_states_batch.shape[3], - target_w: target_w + hidden_states_batch.shape[4], - ] += mask - values = values / weight - values = values.clamp_(-1, 1) - return values - - - def tiled_encode(self, video, device, tile_size, tile_stride): - _, _, T, H, W = video.shape - size_h, size_w = tile_size - stride_h, stride_w = tile_stride - - # Split tasks - tasks = [] - for h in range(0, H, stride_h): - if (h-stride_h >= 0 and h-stride_h+size_h >= H): continue - for w in range(0, W, stride_w): - if (w-stride_w >= 0 and w-stride_w+size_w >= W): continue - h_, w_ = h + size_h, w + size_w - tasks.append((h, h_, w, w_)) - - data_device = "cpu" - computation_device = device - - out_T = (T + 3) // 4 - weight = torch.zeros((1, 1, out_T, H // self.upsampling_factor, W // self.upsampling_factor), dtype=video.dtype, device=data_device) - values = torch.zeros((1, self.z_dim, out_T, H // self.upsampling_factor, W // self.upsampling_factor), dtype=video.dtype, device=data_device) - - for h, h_, w, w_ in tqdm(tasks, desc="VAE encoding"): - hidden_states_batch = video[:, :, :, h:h_, w:w_].to(computation_device) - hidden_states_batch = self.model.encode(hidden_states_batch, self.scale).to(data_device) - - mask = self.build_mask( - hidden_states_batch, - is_bound=(h==0, h_>=H, w==0, w_>=W), - border_width=((size_h - stride_h) // self.upsampling_factor, (size_w - stride_w) // self.upsampling_factor) - ).to(dtype=video.dtype, device=data_device) - - target_h = h // self.upsampling_factor - target_w = w // self.upsampling_factor - values[ - :, - :, - :, - target_h:target_h + hidden_states_batch.shape[3], - target_w:target_w + hidden_states_batch.shape[4], - ] += hidden_states_batch * mask - weight[ - :, - :, - :, - target_h: target_h + hidden_states_batch.shape[3], - target_w: target_w + hidden_states_batch.shape[4], - ] += mask - values = values / weight - return values - - - def single_encode(self, video, device): - video = video.to(device) - x = self.model.encode(video, self.scale) - return x - - - def single_decode(self, hidden_state, device): - hidden_state = hidden_state.to(device) - video = self.model.decode(hidden_state, self.scale) - return video.clamp_(-1, 1) - - - def encode(self, videos, device, tiled=False, tile_size=(34, 34), tile_stride=(18, 16)): - - videos = [video.to("cpu") for video in videos] - hidden_states = [] - for video in videos: - video = video.unsqueeze(0) - if tiled: - tile_size = (tile_size[0] * self.upsampling_factor, tile_size[1] * self.upsampling_factor) - tile_stride = (tile_stride[0] * self.upsampling_factor, tile_stride[1] * self.upsampling_factor) - hidden_state = self.tiled_encode(video, device, tile_size, tile_stride) - else: - hidden_state = self.single_encode(video, device) - hidden_state = hidden_state.squeeze(0) - hidden_states.append(hidden_state) - hidden_states = torch.stack(hidden_states) - return hidden_states - - - def decode(self, hidden_states, device, tiled=False, tile_size=(34, 34), tile_stride=(18, 16)): - if tiled: - video = self.tiled_decode(hidden_states, device, tile_size, tile_stride) - else: - video = self.single_decode(hidden_states, device) - return video - - def forward(self, hidden_states, device, tiled=False, tile_size=(34, 34), tile_stride=(18, 16)): - if tiled: - video = self.tiled_decode(hidden_states, device, tile_size, tile_stride) - else: - video = self.single_decode(hidden_states, device) - return video - - @staticmethod - def state_dict_converter(): - return WanVideoVAEStateDictConverter() - - -class WanVideoVAEStateDictConverter: - - def __init__(self): - pass - - def from_civitai(self, state_dict): - state_dict_ = {} - if 'model_state' in state_dict: - state_dict = state_dict['model_state'] - for name in state_dict: - state_dict_['model.' + name] = state_dict[name] - return state_dict_ - - -class VideoVAE38_(VideoVAE_): - - def __init__(self, - dim=160, - z_dim=48, - dec_dim=256, - dim_mult=[1, 2, 4, 4], - num_res_blocks=2, - attn_scales=[], - temperal_downsample=[False, True, True], - dropout=0.0): - super(VideoVAE_, self).__init__() - self.dim = dim - self.z_dim = z_dim - self.dim_mult = dim_mult - self.num_res_blocks = num_res_blocks - self.attn_scales = attn_scales - self.temperal_downsample = temperal_downsample - self.temperal_upsample = temperal_downsample[::-1] - - # modules - self.encoder = Encoder3d_38(dim, z_dim * 2, dim_mult, num_res_blocks, - attn_scales, self.temperal_downsample, dropout) - self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1) - self.conv2 = CausalConv3d(z_dim, z_dim, 1) - self.decoder = Decoder3d_38(dec_dim, z_dim, dim_mult, num_res_blocks, - attn_scales, self.temperal_upsample, dropout) - - - def encode(self, x, scale): - self.clear_cache() - x = patchify(x, patch_size=2) - t = x.shape[2] - iter_ = 1 + (t - 1) // 4 - for i in range(iter_): - self._enc_conv_idx = [0] - if i == 0: - out = self.encoder(x[:, :, :1, :, :], - feat_cache=self._enc_feat_map, - feat_idx=self._enc_conv_idx) - else: - out_ = self.encoder(x[:, :, 1 + 4 * (i - 1):1 + 4 * i, :, :], - feat_cache=self._enc_feat_map, - feat_idx=self._enc_conv_idx) - out = torch.cat([out, out_], 2) - mu, log_var = self.conv1(out).chunk(2, dim=1) - if isinstance(scale[0], torch.Tensor): - scale = [s.to(dtype=mu.dtype, device=mu.device) for s in scale] - mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view( - 1, self.z_dim, 1, 1, 1) - else: - scale = scale.to(dtype=mu.dtype, device=mu.device) - mu = (mu - scale[0]) * scale[1] - self.clear_cache() - return mu - - - def decode(self, z, scale): - self.clear_cache() - if isinstance(scale[0], torch.Tensor): - scale = [s.to(dtype=z.dtype, device=z.device) for s in scale] - z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view( - 1, self.z_dim, 1, 1, 1) - else: - scale = scale.to(dtype=z.dtype, device=z.device) - z = z / scale[1] + scale[0] - iter_ = z.shape[2] - x = self.conv2(z) - for i in range(iter_): - self._conv_idx = [0] - if i == 0: - out = self.decoder(x[:, :, i:i + 1, :, :], - feat_cache=self._feat_map, - feat_idx=self._conv_idx, - first_chunk=True) - else: - out_ = self.decoder(x[:, :, i:i + 1, :, :], - feat_cache=self._feat_map, - feat_idx=self._conv_idx) - out = torch.cat([out, out_], 2) - out = unpatchify(out, patch_size=2) - self.clear_cache() - return out - - -class WanVideoVAE38(WanVideoVAE): - - def __init__(self, z_dim=48, dim=160): - super(WanVideoVAE, self).__init__() - - mean = [ - -0.2289, -0.0052, -0.1323, -0.2339, -0.2799, 0.0174, 0.1838, 0.1557, - -0.1382, 0.0542, 0.2813, 0.0891, 0.1570, -0.0098, 0.0375, -0.1825, - -0.2246, -0.1207, -0.0698, 0.5109, 0.2665, -0.2108, -0.2158, 0.2502, - -0.2055, -0.0322, 0.1109, 0.1567, -0.0729, 0.0899, -0.2799, -0.1230, - -0.0313, -0.1649, 0.0117, 0.0723, -0.2839, -0.2083, -0.0520, 0.3748, - 0.0152, 0.1957, 0.1433, -0.2944, 0.3573, -0.0548, -0.1681, -0.0667 - ] - std = [ - 0.4765, 1.0364, 0.4514, 1.1677, 0.5313, 0.4990, 0.4818, 0.5013, - 0.8158, 1.0344, 0.5894, 1.0901, 0.6885, 0.6165, 0.8454, 0.4978, - 0.5759, 0.3523, 0.7135, 0.6804, 0.5833, 1.4146, 0.8986, 0.5659, - 0.7069, 0.5338, 0.4889, 0.4917, 0.4069, 0.4999, 0.6866, 0.4093, - 0.5709, 0.6065, 0.6415, 0.4944, 0.5726, 1.2042, 0.5458, 1.6887, - 0.3971, 1.0600, 0.3943, 0.5537, 0.5444, 0.4089, 0.7468, 0.7744 - ] - self.mean = torch.tensor(mean) - self.std = torch.tensor(std) - self.scale = [self.mean, 1.0 / self.std] - - # init model - self.model = VideoVAE38_(z_dim=z_dim, dim=dim).eval().requires_grad_(False) - self.upsampling_factor = 16 - self.z_dim = z_dim diff --git a/roll/pipeline/distill/distill_pipeline.py b/roll/pipeline/distill/distill_pipeline.py index 3860ea3be..3da0e0e28 100644 --- a/roll/pipeline/distill/distill_pipeline.py +++ b/roll/pipeline/distill/distill_pipeline.py @@ -180,6 +180,7 @@ def __init__(self, pipeline_config: DistillConfig): dataset = datasets.load_dataset('json', data_files=dataset_paths)['train'] val_dataset = None + self.val_dataloader = None if self.pipeline_config.validation and self.pipeline_config.validation.data_args: val_dataset_paths = self.pipeline_config.validation.data_args.file_name if not val_dataset_paths: diff --git a/roll/pipeline/distill/distill_vlm_pipeline.py b/roll/pipeline/distill/distill_vlm_pipeline.py index 330ae8658..db47ec2f2 100644 --- a/roll/pipeline/distill/distill_vlm_pipeline.py +++ b/roll/pipeline/distill/distill_vlm_pipeline.py @@ -24,7 +24,9 @@ from roll.utils.metrics.metrics_manager import MetricsManager from roll.pipeline.distill.logits_transfer_group import LogitsTransferGroup -from roll.pipeline.rlvr.rlvr_vlm_pipeline import process_images, get_extra_data_provider +from roll.models.model_providers import get_extra_data_provider +from roll.datasets.vlm_dataset_utils import process_images + logger = get_logger() diff --git a/roll/pipeline/distill/distill_worker.py b/roll/pipeline/distill/distill_worker.py index 4cba33fe3..7f42ae67e 100644 --- a/roll/pipeline/distill/distill_worker.py +++ b/roll/pipeline/distill/distill_worker.py @@ -92,7 +92,7 @@ def train_step(self, data: DataProto): data.batch['teacher_inf_mask'] = self.inf_mask_cache.pop_full_logits() if "labels" in data.batch.keys(): # rename key: labels -> labels_for_loss - data.batch.rename_key_("labels", "labels_for_loss") + data.rename("labels", "labels_for_loss") self.logger.info(f"global_step: {data.meta_info.get('global_step',0)}") student_metrics = self.strategy.train_step(batch=data, loss_func=self.loss_func) @@ -150,7 +150,7 @@ def val_step(self, data: DataProto): data = self.strategy.get_data_input(data) if "labels" in data.batch.keys(): # rename key: labels -> labels_for_loss - data.batch.rename_key_("labels", "labels_for_loss") + data.rename("labels", "labels_for_loss") metrics = self.strategy.forward_step(batch=data, forward_func=self.loss_func_for_eval) output = DataProto(meta_info={"metrics": metrics}).to("cpu") return output @@ -479,7 +479,7 @@ def forward(self, data: DataProto): data = self.strategy.get_data_input(data) if "labels" in data.batch.keys(): # rename key: labels -> labels_for_loss - data.batch.rename_key_("labels", "labels_for_loss") + data.rename("labels", "labels_for_loss") is_offload_states = data.meta_info.get("is_offload_states", False) metrics = {} with state_offload_manger( diff --git a/roll/pipeline/rlvr/rewards/__init__.py b/roll/pipeline/rlvr/rewards/__init__.py index 26b9beebd..861878404 100644 --- a/roll/pipeline/rlvr/rewards/__init__.py +++ b/roll/pipeline/rlvr/rewards/__init__.py @@ -3,4 +3,5 @@ from roll.pipeline.rlvr.rewards.general_val_rule_reward_worker import GeneralValRuleRewardWorker from roll.pipeline.rlvr.rewards.ifeval_rule_reward_worker import GeneralRuleRewardWorker from roll.pipeline.rlvr.rewards.llm_judge_reward_worker import LLMJudgeRewardWorker -from roll.pipeline.rlvr.rewards.math_rule_reward_worker import MathRuleRewardWorker \ No newline at end of file +from roll.pipeline.rlvr.rewards.math_rule_reward_worker import MathRuleRewardWorker +from roll.pipeline.rlvr.rewards.remote_reward_system_worker import RemoteRewardSystemWorker diff --git a/roll/pipeline/rlvr/rewards/llm_judge_reward_worker.py b/roll/pipeline/rlvr/rewards/llm_judge_reward_worker.py index 8dd9b5806..46446b617 100644 --- a/roll/pipeline/rlvr/rewards/llm_judge_reward_worker.py +++ b/roll/pipeline/rlvr/rewards/llm_judge_reward_worker.py @@ -1,5 +1,4 @@ from typing import Optional, Union, Dict, List, Any -import asyncio import json import re import torch @@ -8,8 +7,6 @@ import traceback import numpy as np from functools import partial -import uuid -import ray import tensordict from tensordict import TensorDict from roll.configs.worker_config import WorkerConfig @@ -20,22 +17,15 @@ from roll.distributed.strategy.strategy import InferenceStrategy, TrainStrategy from roll.models.model_providers import default_tokenizer_provider, default_reward_model_provider from roll.platforms import current_platform -from roll.utils.constants import RAY_NAMESPACE from roll.utils.logging import get_logger from roll.utils.context_managers import state_offload_manger from roll.utils.prompt import * from roll.datasets.chat_template import get_chat_template -from roll.distributed.scheduler.router import RouterManager class LLMJudgeRewardWorker(Worker): """ Reward Worker that uses LLM-as-judge to compute rewards. - - Supports three judge_model_type modes: - - "api": calls an external OpenAI-compatible API - - "inference": runs a local model via InferenceStrategy (GPU) - - "cluster": delegates to a shared reward model cluster via Ray RequestScheduler (CPU-only) """ def __init__(self, worker_config: WorkerConfig): @@ -45,7 +35,7 @@ def __init__(self, worker_config: WorkerConfig): self.tokenizer = None self.strategy: Optional[Union[InferenceStrategy, TrainStrategy]] = None - # LLM judge config + # LLM judge相关配置 self.judge_prompt = self.worker_config.judge_prompt if hasattr(self.worker_config, "judge_prompt") else None self.judge_prompt = prompt_maps.get(self.judge_prompt, None) self.judge_model_type = ( @@ -57,54 +47,28 @@ def __init__(self, worker_config: WorkerConfig): self.judge_api_url = self.worker_config.judge_api_url if hasattr(self.worker_config, "judge_api_url") else None self.judge_api_key = self.worker_config.judge_api_key if hasattr(self.worker_config, "judge_api_key") else None - # Cluster mode state (populated in _initialize_cluster_mode) - self.reward_tokenizer = None - self.chat_template_func = None - self.reward_scheduler = None - @register(dispatch_mode=Dispatch.ONE_TO_ALL) def initialize(self, pipeline_config): super().initialize(pipeline_config) self.actor_tokenizer = default_tokenizer_provider(pipeline_config.actor_train.model_args) if self.judge_model_type == "api": - self._initialize_api_mode() + self.tokenizer = default_tokenizer_provider(model_args=self.worker_config.model_args) + print(f"{self.worker_name} initialized with API model") + elif self.judge_model_type == "inference": - self._initialize_inference_mode() - elif self.judge_model_type == "cluster": - self._initialize_cluster_mode(pipeline_config) + async_strategy = self.worker_config.strategy_args.strategy_name in ["vllm", "sglang"] + if self.worker_config.strategy_args.strategy_name == "sglang": # not weight sync, need backup weights + self.worker_config.strategy_args.strategy_config["enable_weights_cpu_backup"] = True + if self.worker_config.strategy_args.strategy_name == "vllm": + self.worker_config.strategy_args.strategy_config["sleep_level"] = 1 + self.strategy = create_strategy(worker=self, sync_wrapper=async_strategy) + self.strategy.initialize(model_provider=default_reward_model_provider) + self.tokenizer = self.strategy.tokenizer + print(f"{self.worker_name} initialized with inference model") + self.strategy.offload_states() + current_platform.init() else: - raise ValueError(f"Unsupported judge_model_type: {self.judge_model_type}") - - def _initialize_api_mode(self): - self.tokenizer = default_tokenizer_provider(model_args=self.worker_config.model_args) - print(f"{self.worker_name} initialized with API model") - - def _initialize_inference_mode(self): - async_strategy = self.worker_config.strategy_args.strategy_name in ["vllm", "sglang"] - if self.worker_config.strategy_args.strategy_name == "sglang": # not weight sync, need backup weights - self.worker_config.strategy_args.strategy_config["enable_weights_cpu_backup"] = True - if self.worker_config.strategy_args.strategy_name == "vllm": - self.worker_config.strategy_args.strategy_config["sleep_level"] = 1 - self.strategy = create_strategy(worker=self, sync_wrapper=async_strategy) - self.strategy.initialize(model_provider=default_reward_model_provider) - self.tokenizer = self.strategy.tokenizer - print(f"{self.worker_name} initialized with inference model") - self.strategy.offload_states() - current_platform.init() - - def _initialize_cluster_mode(self, pipeline_config): - if pipeline_config.reward_model is None: - raise ValueError( - "judge_model_type='cluster' requires pipeline_config.reward_model to be configured" - ) - self.reward_tokenizer = default_tokenizer_provider(pipeline_config.reward_model.model_args) - template_name = pipeline_config.reward_model.data_args.template - self.chat_template_func = get_chat_template(template_name, self.reward_tokenizer) - - scheduler_name = f"RewardModelScheduler-{pipeline_config.reward_model.name}" - reward_scheduler = ray.get_actor(scheduler_name, namespace=RAY_NAMESPACE) - self.reward_scheduler = RouterManager.create_client_sync(reward_scheduler) - self.logger.info(f"{self.worker_name} initialized, connected to scheduler: {scheduler_name}") + raise ValueError(f"Unsupported model type: {self.judge_model_type}") def _call_api_model(self, messages: Dict, retry_times=3) -> str: from openai import OpenAI @@ -171,42 +135,52 @@ def _run_local_inference(self, messages: Dict) -> str: self.logger.info(f"judge model inference output: {str(output)}") return output.strip() - def _parse_score(self, response: str) -> float: - """Parse score from judge response. Supports 'Score: X' format and yes/no.""" - response_lower = response.lower().strip() - - # Try "Score: X" format first - match = re.search(r"Score:\s*([0-9.]+)", response, re.IGNORECASE) - if match: - score = float(match.group(1)) - # Normalize to [0, 1] if score > 1 - if score > 1.0: - score = score / 10.0 - return min(max(score, 0.0), 1.0) - - # Try yes/no format - if "yes" in response_lower: - return 1.0 - if "no" in response_lower: - return 0.0 - - self.logger.warning(f"Could not parse score from judge response: {response[:200]}") - return 0.0 + def _extract_score(self, response: str) -> float: + try: + match = re.search("Score: ([0-9.]+)", response) + if match: + score = float(match.group(1)) + normalized_score = score / 10 + return normalized_score + else: + self.logger.warning(f"Could not extract score from response: {response}") + return 0.5 + except Exception as e: + self.logger.error(f"Error extracting score: {e}") + return 0.5 + + def _extract_score_v2(self, response: str) -> float: + response = response.lower() + try: + if "yes" in response: + return 1 + elif "no" in response: + return 0 + else: + self.logger.warning(f"Could not extract score from response: {response}") + return 0 + except Exception as e: + self.logger.error(f"Error extracting score: {e}") + return 0 def _format_judge_prompt(self, prompt: str, response: str, reference: str = None) -> str: if "user\n" in prompt: prompt = prompt.split("user\n")[-1].strip() if not self.judge_prompt: - formatted_prompt = ( - f"You are an expert judge evaluating the quality of a response to a given prompt.\n\n" - f"Prompt: {prompt}\n\n" - f"Response: {response}\n\n" - f"Reference: {reference}\n\n" - f"Please evaluate the response on a scale from 0 to 10.\n" - f"Consider factors such as correctness, completeness, clarity, and relevance to the prompt.\n" - f"Your evaluation should be a single number between 0 and 10.\n" - f"Note output your score in the following format: Score: your score." - ) + formatted_prompt = f""" + You are an expert judge evaluating the quality of a response to a given prompt. + + Prompt: {prompt} + + Response: {response} + + Reference: {reference} + + Please evaluate the response on a scale from 0 to 10. + Consider factors such as correctness, completeness, clarity, and relevance to the prompt. + Your evaluation should be a single number between 0 and 10. + Note output your score in the following format: Score: your score. + """ else: formatted_prompt = self.judge_prompt.format(question=prompt, response=response, reference=reference) messages = [{"role": "user", "content": formatted_prompt}] @@ -222,7 +196,7 @@ def _get_llm_judgment(self, prompt_id: str, prompt: str, response: str, referenc else: raise ValueError(f"Unsupported model type: {self.judge_model_type}") - score = self._parse_score(llm_response) + score = self._extract_score_v2(llm_response) info = { "prompt_id": prompt_id, "score": score, @@ -234,100 +208,6 @@ def _get_llm_judgment(self, prompt_id: str, prompt: str, response: str, referenc } return score, info - def _tokenize_single(self, messages: List[Dict]) -> DataProto: - """Tokenize a single judge prompt into a DataProto for RequestScheduler (cluster mode).""" - text = self.chat_template_func(messages) - tokenized = self.reward_tokenizer(text, return_tensors="pt") - input_ids = tokenized["input_ids"] - attention_mask = tokenized["attention_mask"] - position_ids = torch.arange(input_ids.shape[1]).unsqueeze(0) - - data = DataProto( - batch=TensorDict( - { - "input_ids": input_ids, - "attention_mask": attention_mask, - "position_ids": position_ids, - }, - batch_size=input_ids.shape[0], - ) - ) - return data - - def _compute_rewards_cluster(self, data: DataProto, metrics: Dict) -> DataProto: - """Compute rewards via the shared reward model cluster (concurrent async requests).""" - prompts_text = self.actor_tokenizer.batch_decode(data.batch["prompts"], skip_special_tokens=True) - responses_text = self.actor_tokenizer.batch_decode(data.batch["responses"], skip_special_tokens=True) - - ground_truths = data.non_tensor_batch.get("ground_truth", [None] * len(prompts_text)) - prompt_ids = data.non_tensor_batch.get("id", [str(i) for i in range(len(prompts_text))]) - - # Prepare generation config for judge model - generation_config = self.worker_config.generating_args.to_dict() - - # Format judge prompts, tokenize, and send concurrent async requests - async def generate_all(): - tasks = [] - for i, (prompt, response, reference) in enumerate(zip(prompts_text, responses_text, ground_truths)): - messages = self._format_judge_prompt(prompt, response, reference) - single_data = self._tokenize_single(messages) - single_data.meta_info = { - "src_rank": self.rank_info.rank, - "pad_to_seq_len": False, - "generation_config": generation_config, - } - request_id = f"reward_{self.rank_info.rank}_{uuid.uuid4().hex[:8]}_{i}" - - # Use RouterClient's async method for concurrent processing - task = self.reward_scheduler.generate_request( - req=single_data, request_id=request_id, uid=self.rank_info.rank - ) - tasks.append(task) - - # Gather all results concurrently - return await asyncio.gather(*tasks) - - # Run async tasks in sync context - results = asyncio.run(generate_all()) - - # Parse scores from judge responses - scores = [] - for i, result in enumerate(results): - if result is None: - self.logger.warning(f"Sample {prompt_ids[i]}: judge request returned None, scoring 0.0") - scores.append(0.0) - continue - - judge_text = self.reward_tokenizer.batch_decode(result.meta_info["output_token_ids"], skip_special_tokens=True)[0] - score = self._parse_score(judge_text) - scores.append(score) - - self.logger.info( - json.dumps( - { - "prompt_id": prompt_ids[i], - "score": score, - "judge_response": judge_text[:500], - }, - ensure_ascii=False, - ) - ) - - # Return reward DataProto - scores_tensor = torch.tensor(scores, dtype=torch.float16) - token_level_rewards = torch.zeros_like(data.batch["responses"], dtype=torch.float16) - - output = DataProto.from_dict( - tensors={ - "token_level_rewards": token_level_rewards, - "response_level_rewards": scores_tensor, - "scores": scores_tensor, - } - ) - output.meta_info = {"metrics": metrics} - self.logger.info(f"Computed rewards for {len(scores)} samples via reward model cluster") - return output - @register(dispatch_mode=Dispatch.DP_MP_COMPUTE, clear_cache=False) def compute_rewards(self, data: DataProto): global_step = data.meta_info.get("global_step", 0) @@ -342,8 +222,6 @@ def compute_rewards(self, data: DataProto): is_offload_states=is_offload_states, ): return self._compute_rewards_impl(data, metrics) - elif self.judge_model_type == "cluster": - return self._compute_rewards_cluster(data, metrics) else: return self._compute_rewards_impl(data, metrics) diff --git a/roll/pipeline/rlvr/rewards/multiple_choice_boxed_rule_reward_worker.py b/roll/pipeline/rlvr/rewards/multiple_choice_boxed_rule_reward_worker.py index 929268191..450483755 100644 --- a/roll/pipeline/rlvr/rewards/multiple_choice_boxed_rule_reward_worker.py +++ b/roll/pipeline/rlvr/rewards/multiple_choice_boxed_rule_reward_worker.py @@ -146,7 +146,7 @@ def compute_rewards(self, data: DataProto): token_level_rewards = torch.zeros_like(data.batch["responses"], dtype=torch.float16) scores = torch.tensor(scores, dtype=torch.float16) - response_level_rewards = scores + response_level_rewards = torch.zeros_like(scores, dtype=torch.float16) # 5) 将这些张量打包进同一个字典 # TODO: 不同的reward worker的output是否需要统一output，或者有没有自适应的办法，避免在新增监控量时每个worker都需要修改 output_tensors = { diff --git a/roll/pipeline/rlvr/rewards/remote_reward_system_worker.py b/roll/pipeline/rlvr/rewards/remote_reward_system_worker.py new file mode 100644 index 000000000..1d9fcf9c3 --- /dev/null +++ b/roll/pipeline/rlvr/rewards/remote_reward_system_worker.py @@ -0,0 +1,529 @@ +# Copyright 2025 alibaba-inc. and/or its affiliates + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import time +import torch +import requests + +import concurrent.futures +from datetime import datetime +from typing import Optional, Union, Dict, List, Any, Tuple + +from codetiming import Timer + +from roll.configs.worker_config import WorkerConfig +from roll.distributed.executor.worker import Worker +from roll.distributed.scheduler.decorator import Dispatch, register +from roll.distributed.scheduler.protocol import DataProto +from roll.models.model_providers import default_tokenizer_provider +from roll.utils.logging import get_logger + +logger = get_logger() + + +class RewardSystemHandler: + def __init__(self, + reward_system_config: Dict, + reward_manager_config: Dict, + reward_server_task_timeout: int = 300): + """ + 初始化奖励系统处理器。 + + Args: + reward_system_config: 包含 xrl_authorization、URL、模式等配置。 + reward_manager_config: 奖励管理器配置，透传给后端。 + reward_server_task_timeout: 任务最大服务时间，单位为秒。超过这个时间后，任务会被服务端取消。 + """ + + xrl_authorization = reward_system_config.get("xrl_authorization", + "t-29imkiykio27pmju") + submit_url = reward_system_config.get("submit_url", + "http://xrl.alibaba-inc.com/apis/custom/reward-system/v1/submit_task") + get_result_url = reward_system_config.get("get_result_url", + "http://xrl.alibaba-inc.com/apis/custom/reward-system/v1/task_result") + + mode = reward_system_config.get("reward_system_type", "online") + assert mode in ["online", "pre"], f"Invalid xrl mode: {mode}" + + project_id = reward_system_config.get("project_id", "roll_default_task_id") + experiment_id = reward_system_config.get("experiment_id", "roll_default_task_id") + global_step = reward_system_config.get("global_step", 0) + priority = reward_system_config.get("priority", 1) + + self.xrl_authorization = xrl_authorization + + self.submit_url = submit_url + self.get_result_url = get_result_url + + self.headers = { + "Content-Type": "application/json", + "XRL-Authorization": f"Bearer {xrl_authorization}", + "BusinessType": mode + } + + # TODO: add dynamic priority + self.priority = priority + self.project_id = project_id + self.experiment_id = experiment_id + self.global_step = global_step + self.reward_server_task_timeout = reward_server_task_timeout + + self.reward_manager_config = reward_manager_config + + def submit_task( + self, + prompt: str, + response: str, + reward_info: Dict, + data_id: str, + rollout_id: str, + project_id: str = None, + experiment_id: str = None, + global_step: str = None, + priority: int = None, + max_retries: int = 5, + retry_delay: float = 1.0, + timeout: int = 20 + ) -> Optional[str]: + """ + 提交 reward 异步计算任务。 + Args: + prompt (str): _description_ + response (str): _description_ + reward_info (Dict): _description_ + data_id (str): _description_ + rollout_id (str): _description_ + project_id (str, optional): _description_. Defaults to None. + experiment_id (str, optional): _description_. Defaults to None. + global_step (str, optional): _description_. Defaults to None. + priority (int, optional): _description_. Defaults to None. + max_retries (int, optional): _description_. Defaults to 5. + retry_delay (float, optional): _description_. Defaults to 1.0. + timeout (int, optional): _description_. Defaults to 20. + + Returns: + str: reward_request_id，失败时返回 None。 + """ + if priority is None: + priority = self.priority + if project_id is None: + project_id = self.project_id + if experiment_id is None: + experiment_id = self.experiment_id + if global_step is None: + global_step = self.global_step + + + payload = { + "priority": priority, + "project_id": project_id, + "experiment_id": experiment_id, + "global_step": global_step, + "data_id": data_id, + "rollout_id": rollout_id, + "prompt": prompt, + "response": response, + "reward_info": reward_info, + "reward_manager": self.reward_manager_config, + "timeout": self.reward_server_task_timeout + } + + reward_request_id = None + task_time_start = 0 + + for attempt in range(1, max_retries + 1): + try: + task_time_start = time.time() + response = requests.post( + self.submit_url, + headers=self.headers, + json=payload, + timeout=timeout + ) + response.raise_for_status() # 如果状态码不是 2xx，会抛出 HTTPError + response_json = response.json() + reward_request_id = response_json.get("reward_request_id") + errcode = int(response_json.get("errcode", -1)) + + if errcode == 0: + break + else: + errmsg = response_json.get("errmsg", "未知错误") + logger.warning(f"Try {attempt} out of {max_retries} times, {errcode=}, {errmsg=}; {payload=}") + + except Exception as e: + logger.warning(f"Try {attempt} out of {max_retries} times, {e}") + + # 如果不是最后一次尝试，则等待后重试 + if attempt < max_retries: + time.sleep(retry_delay) + if reward_request_id is None: + debug_info = { + "project_id": project_id, + "experiment_id": experiment_id, + "global_step": global_step, + "data_id": data_id, + "rollout_id": rollout_id, + "prompt": prompt, + "response": response + } + logger.warning(f"Reward system submit task failed. {debug_info=}") + + return reward_request_id, task_time_start + + def get_task_result( + self, + reward_request_id: str, + task_time_start: float, + retry_delay: float = 2.0, + timeout: int = 20 + ) -> Tuple[Optional[float], Optional[Dict[str, float]]]: + """ + 轮训查询任务结果 + """ + def extract_time_costs(task_time_start, response_json): + time_info = { + "local_request_time_cost": 0, + "server_task_time_cost": 0, + "plugin_time_cost": [] + } + + time_info["local_request_time_cost"] = time.time() - task_time_start + if response_json: + server_log_start_time = response_json.get("start_time", None) + server_log_end_time = response_json.get("end_time", None) + if server_log_start_time and server_log_end_time: + try: + server_log_start_time = datetime.fromisoformat(server_log_start_time) + server_log_end_time = datetime.fromisoformat(server_log_end_time) + time_info["server_task_time_cost"] = (server_log_end_time - server_log_start_time).total_seconds() + except Exception as e: + logger.warning(f"Get Reward System server time cost failed. {e}") + dimention_reward_score = response_json.get("dimention_reward_score", {}) + + for key, value in dimention_reward_score.items(): + debug_info = value.get("debug_info", {}) + for key, value in debug_info.items(): + time_info["plugin_time_cost"].append({key: value.get('time_cost', 0)}) + + return time_info + + payload = {"reward_request_id": reward_request_id} + + reward_score = None + dimention_reward_score = None + local_request_time_cost = 0 + attempt = 0 + + # 预留 5s 查询余量，避免临界超时任务漏查。 + while local_request_time_cost < self.reward_server_task_timeout + 5: + time_info = None + response_json = None + attempt += 1 + + try: + response = requests.get( + self.get_result_url, + headers=self.headers, + json=payload, + timeout=timeout + ) + time_info = extract_time_costs(task_time_start, response_json) + local_request_time_cost = time_info["local_request_time_cost"] + response.raise_for_status() + response_json = response.json() + errcode = int(response_json.get("errcode", -1)) + errmsg = response_json.get("errmsg", "unknown error") + if errcode == 0: + reward_score = response_json.get("reward_score") + dimention_reward_score = response_json.get("dimention_reward_score") + break + # 1 表示任务还在执行中 + elif errcode == 1: + pass + else: + logger.warning(f"Try {attempt} times," + f"{time_info=}s; {errcode=}, {errmsg=}; {payload=}") + except Exception as e: + logger.warning(f"Try {attempt} times, {time_info=}; {e}") + + if time_info is None: + time_info = extract_time_costs(task_time_start, response_json) + if local_request_time_cost > self.reward_server_task_timeout: + logger.warning(f"Reward system get task result timeout. {time_info=}; {reward_request_id=}") + break + + time.sleep(retry_delay) + + if reward_score is None: + logger.warning(f"Reward system reward calculation failed. reward_server_task_timeout: " + f"{self.reward_server_task_timeout}; {time_info=}; {reward_request_id=}") + + return reward_score, dimention_reward_score + + def get_reward( + self, + prompt: str, + response: str, + reward_info: Dict, + data_id: str, + rollout_id: str, + project_id: str = None, + experiment_id: str = None, + global_step: str = None, + priority: int = None, + ) -> Tuple[Optional[float], Optional[Dict[str, float]]]: + """ + 提交任务并等待结果（同步阻塞） + + Args: + prompt (str): _description_ + response (str): _description_ + reward_info (Dict): _description_ + data_id (str): _description_ + rollout_id (str): _description_ + project_id (str, optional): _description_. Defaults to None. + experiment_id (str, optional): _description_. Defaults to None. + global_step (str, optional): _description_. Defaults to None. + priority (int, optional): _description_. Defaults to None. + + Returns: + _type_: _description_ + """ + reward_request_id, task_time_start = self.submit_task( + prompt=prompt, + response=response, + reward_info=reward_info, + data_id=data_id, + rollout_id=rollout_id, + project_id=project_id, + experiment_id=experiment_id, + global_step=global_step, + priority=priority, + ) + if reward_request_id is None: + return None, None + return self.get_task_result(reward_request_id, task_time_start) + + def multi_thread_request_reward_system(self, input_list) -> List[float]: + """ + 多线程批量请求奖励分数。 + + input_list 元素格式: + (project_id, experiment_id, global_step, data_id, rollout_id, reward_info, prompt, response, response_token_str) + 注意：response_token_str 当前未使用，保留以备扩展。 + """ + def thread_request_worker( + index, + project_id, + experiment_id, + global_step, + data_id, + rollout_id, + reward_info, + prompt, + response + ): + + for attempt in range(5): + score, dimention_reward = self.get_reward( + prompt=prompt, + response=response, + reward_info=reward_info, + data_id=data_id, + rollout_id=rollout_id, + project_id=project_id, + experiment_id=experiment_id, + global_step=global_step, + ) + + if score is not None: + break + + if score is None: + debug_info = { + "project_id": project_id, + "experiment_id": experiment_id, + "global_step": global_step, + "data_id": data_id, + "rollout_id": rollout_id, + "prompt": prompt, + "response": response, + "reward_info": reward_info + } + logger.warning(f"Failed to get reward from rewared system. debug_info: {debug_info}") + score = 0 + + return index, score + + output_list = [None] * len(input_list) + max_workers = max(min(32, len(input_list)), 1) + + with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor: + # 提交所有任务 + future_to_index = { + executor.submit(thread_request_worker, index, *item): index + for index, item in enumerate(input_list) + } + # 等待完成并按索引填充结果 + for future in concurrent.futures.as_completed(future_to_index): + index, result = future.result() + output_list[index] = result + + for index, result in enumerate(output_list): + if result is None: + output_list[index] = 0 + + return output_list + + +class RemoteRewardSystemWorker(Worker): + """ + Reward Worker that uses Reward System to compute rewards. + """ + + def __init__(self, worker_config: WorkerConfig): + super().__init__(worker_config=worker_config) + self.rank_info.dp_rank = self.rank_info.rank + self.rank_info.dp_size = self.rank_info.world_size + + self.tokenizer = default_tokenizer_provider(model_args=self.worker_config.model_args) + self.reward_manager_config = self.worker_config.reward_manager_config + + + @register(dispatch_mode=Dispatch.ONE_TO_ALL) + def initialize(self, pipeline_config): + pass + + def _uniform_reward_ref_info_adaper(self, data: DataProto) -> List[Dict]: + """ + Normalize various input data formats into a uniform reward reference format. + + Args: + data (DataProto): Input data containing non-tensor metadata. + + Returns: + List[Dict[str, Any]]: List of reward reference info dicts. + + Raises: + ValueError: If data format is not supported. + """ + def is_code_format_data(reward_manager_config, non_tensor_batch): + """Check if data is in code evaluation format.""" + if ( + non_tensor_batch.get("test_cases", [""])[0] + and non_tensor_batch.get("case_type", [""])[0] + ): + return True + for item in reward_manager_config: + if "code" in item["plugin_name"]: + return True + return False + + non_tensor_batch = data.non_tensor_batch + # Case 1: Already in standard format + if "reward_ref_info" in non_tensor_batch: + reward_ref_info = non_tensor_batch["reward_ref_info"] + # Case 2: Code-related data (detected by fields or config) + elif is_code_format_data(self.reward_manager_config, non_tensor_batch): + reward_ref_info = [ + {"case_type": case_type, "test_cases": test_cases} for case_type, test_cases in zip( + non_tensor_batch["case_type"], + non_tensor_batch["test_cases"], + ) + ] + # Case 3: Ground truth available + elif "ground_truth" in non_tensor_batch: + reward_ref_info = [ + {"ground_truth": ground_truth} for ground_truth in non_tensor_batch["ground_truth"] + ] + # Unsupported format + else: + known_keys = list(non_tensor_batch.keys()) + raise ValueError( + f"Unsupported data format for reward system. " + f"Available keys: {known_keys}. " + f"Expected one of: 'reward_ref_info', 'ground_truth', or code fields ('case_type', 'test_cases')." + ) + + return reward_ref_info + + @register(dispatch_mode=Dispatch.DP_MP_COMPUTE, clear_cache=False) + def compute_rewards(self, data: DataProto) -> DataProto: + """ + Compute rewards using remote reward system. + + Args: + data (DataProto): Input batch with prompts, responses, and metadata. + + Returns: + DataProto: Output with token-level and response-level rewards. + """ + global_step = data.meta_info.get("global_step", 0) + reward_system_config = data.meta_info.get("reward_system_config", {}) + project_id = reward_system_config.get("project_id", 'roll_default_experiment_id') + experiment_id = reward_system_config.get("experiment_id", 'roll_default_task_id') + + reward_system_handler = RewardSystemHandler(reward_system_config, self.reward_manager_config) + + prompts_text_list = self.tokenizer.batch_decode(data.batch["prompts"], skip_special_tokens=True) + response_text_list = self.tokenizer.batch_decode(data.batch["responses"], skip_special_tokens=True) + + data_id_list = data.non_tensor_batch["id"] + rollout_id_list = data.non_tensor_batch["rollout_id"] + + reward_ref_info_list = self._uniform_reward_ref_info_adaper(data) + + + reward_timer = Timer("reward_timer") + input_list = [] + + with reward_timer: + for data_id, prompt, resp_text, reward_ref_info, rollout_id in zip( + data_id_list, + prompts_text_list, + response_text_list, + reward_ref_info_list, + rollout_id_list, + ): + prompt = prompt.replace("<|endoftext|>", "").replace("", "") + prompt = prompt.replace("assistant\n", "") + + input_list.append([ + project_id, + experiment_id, + global_step, + data_id, + rollout_id, + reward_ref_info, + prompt, + resp_text + ]) + + scores = reward_system_handler.multi_thread_request_reward_system(input_list) + + scores_tensor = torch.tensor(scores, dtype=torch.float16) + token_level_rewards = torch.zeros_like(data.batch["responses"], dtype=torch.float16) + response_level_rewards = scores_tensor + + output = DataProto.from_dict( + tensors={ + "token_level_rewards": token_level_rewards, + "response_level_rewards": response_level_rewards, + "scores": scores_tensor, + } + ) + + return output + diff --git a/roll/pipeline/rlvr/rewards/video_r1_reward_worker.py b/roll/pipeline/rlvr/rewards/video_r1_reward_worker.py new file mode 100644 index 000000000..aa21e94ca --- /dev/null +++ b/roll/pipeline/rlvr/rewards/video_r1_reward_worker.py @@ -0,0 +1,184 @@ +""" +reference: https://github.com/tulerfeng/Video-R1/blob/main/src/r1-v/src/open_r1/grpo.py +""" + +import re +import torch +from collections import defaultdict +from typing import Dict, List + +from rouge_score import rouge_scorer + +from roll.configs.worker_config import WorkerConfig +from roll.distributed.executor.worker import Worker +from roll.distributed.scheduler.decorator import Dispatch, register +from roll.distributed.scheduler.protocol import DataProto +from roll.models.model_providers import default_tokenizer_provider +from roll.utils.logging import get_logger + + +logger = get_logger() + + +def extract_boxed_content(text): + pattern = r"\\boxed{(.*?)}" + matches = re.findall(pattern, text) + if len(matches) > 0: + return matches[-1] + else: + return text + + +def accuracy_reward(completion, solution, question_type="acc"): + def extract_answer(text): + pattern = r"\s*(.*?)\s*" + match = re.search(pattern, text, re.DOTALL) + if match: + return match.group(1).strip() + return text + + def normalize_number(num_str): + try: + num_str = num_str.replace(",", "") + return float(num_str) + except Exception as e: + print(f"Error converting '{num_str}' to float: {e}") + return None + + def wer(reference, hypothesis): + ref_words = reference.split() + hyp_words = hypothesis.split() + m = len(ref_words) + n = len(hyp_words) + d = [[0] * (n + 1) for _ in range(m + 1)] + for i in range(m + 1): + d[i][0] = i + for j in range(n + 1): + d[0][j] = j + for i in range(1, m + 1): + for j in range(1, n + 1): + if ref_words[i - 1] == hyp_words[j - 1]: + d[i][j] = d[i - 1][j - 1] + else: + d[i][j] = 1 + min(d[i - 1][j], d[i][j - 1], d[i - 1][j - 1]) + return d[m][n] / max(1, m) + + def compute_rouge_score(reference, hypothesis, use_stemmer=True): + scorer = rouge_scorer.RougeScorer(["rouge1", "rouge2", "rougeL"], use_stemmer=use_stemmer) + scores = scorer.score(reference, hypothesis) + average_fmeasure = (scores["rouge1"].fmeasure + scores["rouge2"].fmeasure + scores["rougeL"].fmeasure) / 3 + return average_fmeasure + + output_ans = extract_answer(completion) + gt_ans = extract_answer(solution) + if question_type == "multiple choice": + output_ans = extract_boxed_content(output_ans) + gt_ans = extract_boxed_content(gt_ans) + reward = 1.0 if (output_ans.strip().strip(".")[:1] == gt_ans.strip().strip(".")[:1]) else 0.0 + elif question_type == "free-form": + score = compute_rouge_score(gt_ans, output_ans) + reward = max(0.0, min(1.0, score)) + elif question_type == "numerical": + gt_has_decimal = ("." in gt_ans) or ("," in gt_ans) + out_has_decimal = ("." in output_ans) or ("," in output_ans) + if gt_has_decimal != out_has_decimal: + reward = 0.0 + else: + gt_number = normalize_number(gt_ans) + out_number = normalize_number(output_ans) + if gt_number is None or out_number is None: + reward = 0.0 + else: + reward = 1.0 if round(gt_number, 2) == round(out_number, 2) else 0.0 + elif question_type == "OCR": + error_rate = wer(gt_ans, output_ans) + reward = 1 - error_rate + reward = max(0.0, min(1.0, reward)) + elif question_type == "regression": + gt_number = normalize_number(gt_ans) + out_number = normalize_number(output_ans) + if gt_number is None or out_number is None: + reward = 0.0 + rel_diff = (abs(out_number - gt_number) + 1e-9) / (abs(gt_number) + 1e-9) + rel_diff = min(1.0, max(0.0, rel_diff)) + reward = 1 - rel_diff + else: + reward = 0.0 + + return reward + + +def format_reward(completion): + """Reward function that checks if the completion has a specific format.""" + pattern = r".*?\s*.*?" + match = re.fullmatch(pattern, completion, re.DOTALL) + return 1.0 if match else 0.0 + + +class VideoR1RewardWorker(Worker): + def __init__(self, worker_config: WorkerConfig): + super().__init__(worker_config=worker_config) + self.rank_info.dp_rank = self.rank_info.rank + self.rank_info.dp_size = self.rank_info.world_size + + @register(dispatch_mode=Dispatch.ONE_TO_ALL) + def initialize(self, pipeline_config): + super().initialize(pipeline_config) + self.tokenizer = default_tokenizer_provider(pipeline_config.actor_train.model_args) + + def judge_responses(self, prompts: List[Dict[str, str]]) -> List[Dict[str, float]]: + """Judge a batch of responses using vLLM""" + scores = [] + for prompt in prompts: + model_response = prompt["response"] + ground_truth = prompt["ground_truth"] + question_type = prompt["question_type"] + acc_reward = accuracy_reward(model_response, ground_truth, question_type=question_type) + fmt_reward = format_reward(model_response) + overall = acc_reward + fmt_reward + scores.append({"overall": overall, "acc_reward": acc_reward, "fmt_reward": fmt_reward}) + return scores + + @register(dispatch_mode=Dispatch.DP_MP_COMPUTE, clear_cache=False) + def compute_rewards(self, data: DataProto): + # some initialization need after work.initialize + response_ids = data.batch["responses"] + response_length = data.batch["response_mask"].sum(dim=-1) + judge_prompts = [] + for i in range(len(data)): + valid_response_ids = response_ids[i][: response_length[i]] + # NOTE: `re.fullmatch(pattern)` in format_reward seems not suitable for qwen3-vl + # thinking model, `` is the prompt ending suffix thus adjust format_reward + response_str = self.tokenizer.decode(valid_response_ids, skip_special_tokens=True) + ground_truth = ( + data.non_tensor_batch["ground_truth"][i] + if "ground_truth" in data.non_tensor_batch + else data.non_tensor_batch["answer"][i] + ) + question_type = data.non_tensor_batch["reward_model"][i] + judge_prompt = { + "question_type": question_type, + "ground_truth": ground_truth, + "response": response_str, + } + judge_prompts.append(judge_prompt) + scores = self.judge_responses(judge_prompts) + rewards = [] + reward_metrics = defaultdict(list) + for i, score in enumerate(scores): + rewards.append(float(score["overall"])) + for key, value in score.items(): + reward_metrics[key].append(value) + scores_tensor = torch.tensor(rewards, dtype=torch.float16) + token_level_rewards = torch.zeros_like(response_ids, dtype=torch.float16) + response_level_rewards = scores_tensor + output = DataProto.from_dict( + tensors={ + "token_level_rewards": token_level_rewards, + "response_level_rewards": response_level_rewards, + "scores": scores_tensor, + } + ) + logger.debug(f"{judge_prompt=}, {reward_metrics=}") + output.meta_info = {"metrics": reward_metrics} + return output diff --git a/roll/pipeline/rlvr/rlvr_config.py b/roll/pipeline/rlvr/rlvr_config.py index e00f6e4e4..c67100848 100644 --- a/roll/pipeline/rlvr/rlvr_config.py +++ b/roll/pipeline/rlvr/rlvr_config.py @@ -84,6 +84,10 @@ class RLVRConfig(PPOConfig): global_template: str = field( default=None, metadata={"help": "The template of the global."}) + tag_to_template: Dict[str, str] = field( + default_factory=dict, + metadata={"help": "Mapping from tag to template name. Tags not in this dict fall back to global_template."} + ) dataset_filter: DatasetFilterConfig = field( default_factory=DatasetFilterConfig, metadata={"help": "Configuration for filtering dataset by source and difficulty"}, @@ -170,6 +174,11 @@ def __post_init__(self): self.critic.worker_cls = "roll.pipeline.base_worker.CriticWorker" if self.reward_model is not None and self.reward_model.worker_cls is None: self.reward_model.worker_cls = "roll.pipeline.base_worker.InferWorker" + # Multi-teacher: set worker_cls for each _reference_configs entry + if hasattr(self, '_reference_configs'): + for ref_cfg in self._reference_configs.values(): + if ref_cfg.worker_cls is None: + ref_cfg.worker_cls = "roll.pipeline.rlvr.actor_worker.ActorWorker" if self.router_args is None: self.router_args = RouterArguments(router_name="PromptAffinityRouter", router_config=dict()) @@ -185,6 +194,30 @@ def __post_init__(self): tag: key for key, worker_config in self.rewards.items() for tag in worker_config.tag_included } + # OPD mode: include teacher tag_included in tag_2_domain mapping + # so that teacher tags are also mapped to domains for routing + if self.is_pure_opd or self.use_opd: + if self.tag_2_domain is None: + self.tag_2_domain = {} + self.domain_2_tag = {} + for name, ref_cfg in self._reference_configs.items(): + for tag in ref_cfg.tag_included: + if tag not in self.tag_2_domain: + # If no reward covers this tag, map tag to itself as domain + self.tag_2_domain[tag] = tag + if tag not in self.domain_2_tag: + self.domain_2_tag[tag] = {tag} + + # Validate tag_to_template entries + if self.tag_to_template: + from roll.datasets.chat_template import chat_templates + for tag, tmpl_name in self.tag_to_template.items(): + if tmpl_name not in chat_templates: + raise ValueError( + f"Template '{tmpl_name}' for tag '{tag}' in tag_to_template not found. " + f"Available templates: {list(chat_templates.keys())}" + ) + if self.async_pipeline: assert self.async_generation_ratio >= 1.0, "async_generation_ratio must be >= 1.0" infer_devices = self.actor_infer.device_mapping @@ -228,6 +261,11 @@ def __post_init__(self): for worker_config in self.rewards.values(): if worker_config.device_mapping is not None: total_devices.extend(worker_config.device_mapping) + # Also include _reference_configs device_mappings for multi-teacher + if hasattr(self, '_reference_configs'): + for ref_cfg in self._reference_configs.values(): + if isinstance(ref_cfg, WorkerConfig) and ref_cfg.device_mapping is not None: + total_devices.extend(ref_cfg.device_mapping) if len(total_devices) > 0: max_gpu_num = max(total_devices) + 1 if max_gpu_num <= self.num_gpus_per_node: diff --git a/roll/pipeline/rlvr/rlvr_pipeline.py b/roll/pipeline/rlvr/rlvr_pipeline.py index 5a133e9ad..1f3dca2a5 100644 --- a/roll/pipeline/rlvr/rlvr_pipeline.py +++ b/roll/pipeline/rlvr/rlvr_pipeline.py @@ -3,6 +3,7 @@ import os import time import uuid +from contextlib import ExitStack from datetime import datetime from functools import partial from typing import Any, Dict, List, Optional @@ -32,19 +33,23 @@ from roll.utils.dynamic_batching import dynamic_batching_shard from roll.utils.functionals import ( RunningMoments, + clusters_have_disjoint_devices, agg_loss, + build_domain_routing_context, compute_advantage, + compute_ref_log_probs_with_routing, compute_token_reward, get_sample_level_mask, reduce_metrics, reward_postprocess, - batch_balance + batch_balance, ) from roll.utils.train_infer_corrections import apply_train_infer_correction_to_batch from roll.utils.kl_controller import get_kl_controller from roll.utils.logging import get_logger from roll.utils.metrics.metrics_manager import MetricsManager from roll.utils.offload_states import OffloadStateType +from roll.utils.telemetry import get_tracer, inject_trace_context logger = get_logger() @@ -55,6 +60,7 @@ def is_lora_training(pipeline_config: RLVRConfig) -> bool: def preprocess_dataset(dataset, prompt_len, encode_function, data_args): + # 处理数据 print(f"Begin : {dataset}") dataset = dataset.map( @@ -75,16 +81,26 @@ def preprocess_dataset(dataset, prompt_len, encode_function, data_args): return dataset -def get_encode_function(template_name, tokenizer, data_args): - chat_template_func = get_chat_template(template_name, tokenizer) +def get_encode_function(template_name, tokenizer, data_args, tag_to_template=None): + # Pre-materialize all template functions for early validation + template_funcs = {template_name: get_chat_template(template_name, tokenizer)} + if tag_to_template: + for tmpl_name in set(tag_to_template.values()): + if tmpl_name not in template_funcs: + template_funcs[tmpl_name] = get_chat_template(tmpl_name, tokenizer) + + tag_key = getattr(data_args, "tag", "tag") def encode_function(data_i): text_list = [] if (message_key := getattr(data_args, "messages", "messages")) in data_i: - for messages in data_i[message_key]: + messages_list = data_i[message_key] + tags = data_i.get(tag_key, [None] * len(messages_list)) + for messages, tag in zip(messages_list, tags): if isinstance(messages, str): messages = json.loads(messages) - text_list.append(chat_template_func(messages)) + tmpl_name = tag_to_template.get(tag, template_name) if tag_to_template else template_name + text_list.append(template_funcs[tmpl_name](messages)) elif (prompt_key := getattr(data_args, "prompt", "prompt")) in data_i: for prompt in data_i[prompt_key]: text_list.append(prompt) @@ -127,7 +143,7 @@ def __init__(self, pipeline_config: RLVRConfig): if self.pipeline_config.global_template else self.pipeline_config.actor_train.data_args.template ) - encode_function = get_encode_function(template_name, self.tokenizer, self.pipeline_config.actor_train.data_args) + encode_function = get_encode_function(template_name, self.tokenizer, self.pipeline_config.actor_train.data_args, tag_to_template=self.pipeline_config.tag_to_template) dataset = preprocess_dataset( dataset, @@ -193,13 +209,17 @@ def __init__(self, pipeline_config: RLVRConfig): download_clusters = [self.actor_train, self.actor_infer] # use unwrapped model as reference for lora training if self.use_ref_model: - self.reference: Any = Cluster( - name=self.pipeline_config.reference.name, - worker_cls=self.pipeline_config.reference.worker_cls, - resource_manager=self.resource_manager, - worker_config=self.pipeline_config.reference, - ) - download_clusters.append(self.reference) + self.references: Dict[str, Any] = {} + for name, ref_cfg in self.pipeline_config.reference_configs.items(): + self.references[name] = Cluster( + name=ref_cfg.name, + worker_cls=ref_cfg.worker_cls, + resource_manager=self.resource_manager, + worker_config=ref_cfg, + ) + download_clusters.extend(self.references.values()) + # Backward compat: self.reference points to the first teacher + self.reference = self.references[list(self.references.keys())[0]] if self.pipeline_config.adv_estimator == "gae": self.critic: Any = Cluster( name=self.pipeline_config.critic.name, @@ -315,7 +335,14 @@ def __init__(self, pipeline_config: RLVRConfig): ray.get(refs) if self.use_ref_model: - refs.extend(self.reference.initialize(pipeline_config=self.pipeline_config, blocking=True)) + if clusters_have_disjoint_devices(self.references): + ref_init_refs = [] + for ref_cluster in self.references.values(): + ref_init_refs.extend(ref_cluster.initialize(pipeline_config=self.pipeline_config, blocking=False)) + ray.get(ref_init_refs) + else: + for ref_cluster in self.references.values(): + ref_cluster.initialize(pipeline_config=self.pipeline_config, blocking=True) refs = [] for key, cluster in self.rewards.items(): @@ -434,6 +461,7 @@ def run(self): # 创建一个专门管理监控指标的类 metrics_mgr = MetricsManager() + tracer = get_tracer("driver") tps_timer = _Timer(window_size=5) actor_infer_timer = _Timer(window_size=5) @@ -459,7 +487,10 @@ def run(self): logger.info(f"pipeline step {global_step} start...") metrics_mgr.clear_metrics() - with tps_timer, Timer(name="step_total", logger=None) as step_total_timer: + with (tps_timer, Timer(name="step_total", logger=None) as step_total_timer, + tracer.start_as_current_span("pipeline_step", attributes={"global_step": global_step}), + ExitStack() as defer, + ): # if global_step > self.state.step + 1: logger.info(f"pre_step_total_time: {pre_step_total_time}") metrics_mgr.add_metric("time/step_total", pre_step_total_time) @@ -475,13 +506,15 @@ def run(self): self.critic.offload_states(blocking=True) self.actor_train.offload_states(blocking=True) - with Timer(name="step_stop_server", logger=None) as step_stop_server_timer: + with Timer(name="step_stop_server", logger=None) as step_stop_server_timer, \ + tracer.start_as_current_span("stop_server"): if self.pipeline_config.async_pipeline: ray.get([scheduler.pause_sampling.remote() for scheduler in self.generate_schedulers.values()]) self.actor_infer.offload_states(include=OffloadStateType.other_params) metrics_mgr.add_metric("time/step_stop_server", step_stop_server_timer.last) - with Timer(name="step_model_update", logger=None) as step_model_update_timer: + with Timer(name="step_model_update", logger=None) as step_model_update_timer, \ + tracer.start_as_current_span("model_update"): model_update_metrics: Dict = self.model_update(global_step) metrics_mgr.add_metrics(model_update_metrics) batch.meta_info["generation_config"] = self.get_generation_config() @@ -491,11 +524,10 @@ def run(self): if not self.pipeline_config.async_pipeline: for reward_cluster in self.rewards.values(): reward_cluster.load_states() - if self.reward_model_cluster: - self.reward_model_cluster.load_states() if self.val_dataset and global_step % self.pipeline_config.eval_steps == 0: - with Timer(name="val_step", logger=None) as val_step_timer: + with Timer(name="val_step", logger=None) as val_step_timer, \ + tracer.start_as_current_span("validation"): val_metrics = self.val(global_step=global_step) metrics_mgr.add_metrics(val_metrics) metrics_mgr.add_metric("time/val_step", val_step_timer.last) @@ -505,10 +537,12 @@ def run(self): actor_infer_timer, actor_infer_response_timer, Timer(name="step_generate", logger=None) as step_generate_timer, + tracer.start_as_current_span("generate"), ): domain_batches = {} scheduler_refs = {} for domain, scheduler in self.generate_schedulers.items(): + inject_trace_context(batch.meta_info) scheduler_refs[domain] = scheduler.get_batch.remote( data=batch, global_step=global_step, batch_size=self.domain_batch_size[domain] ) @@ -532,41 +566,56 @@ def run(self): metrics_mgr.add_metric("time/step_generate", step_generate_timer.last) batch = generate_output + defer.callback(lambda b=batch: DataProto.drop(b)) batch.meta_info["global_step"] = global_step batch.meta_info["_broadcast_non_tensor_batch"] = True batch.meta_info["loss_mask_keys"] = ['response_mask', 'final_response_mask'] batch.non_tensor_batch['sample_uuid'] = np.array([str(uuid.uuid4()) for _ in range(batch.batch.shape[0])], dtype=object) batch.batch["prompt_id"] = torch.arange(batch.batch.batch_size[0], device=batch.batch.device) - with Timer(name="cal_ref_log_probs", logger=None) as cal_ref_log_probs_timer: + with Timer(name="cal_ref_log_probs", logger=None) as cal_ref_log_probs_timer, \ + tracer.start_as_current_span("ref_log_probs"): if self.pipeline_config.enable_reference: - worker_config = self.pipeline_config.reference if self.use_ref_model else self.pipeline_config.actor_train - worker = self.reference if self.use_ref_model else self.actor_train - if worker_config.use_dynamic_batching_in_infer: - batch, dynamic_batching_metrics = dynamic_batching_shard( - batch, - worker.dp_size, - worker_config.max_tokens_per_microbatch_in_infer, - worker_config.sequence_length_round_in_infer, - worker_config.strategy_args.strategy_config.get("pipeline_model_parallel_size", 1), - worker_config.strategy_args.strategy_config.get("virtual_pipeline_model_parallel_size", None), - "reference/compute_log_probs", - ) - metrics_mgr.add_metrics(dynamic_batching_metrics) if not self.use_ref_model: + # LoRA branch: use actor_train with disabled adapter + worker_config = self.pipeline_config.actor_train + batch_balance(batch, dp_size=self.actor_train.dp_size, minibatch_size=len(batch)) batch.meta_info["disable_adapter"] = True batch.meta_info["is_offload_states"] = False - batch_balance(batch, dp_size=self.actor_train.dp_size, minibatch_size=len(batch)) + if worker_config.use_dynamic_batching_in_infer: + batch, dynamic_batching_metrics = dynamic_batching_shard( + batch, + self.actor_train.dp_size, + worker_config.max_tokens_per_microbatch_in_infer, + worker_config.sequence_length_round_in_infer, + worker_config.strategy_args.strategy_config.get("pipeline_model_parallel_size", 1), + worker_config.strategy_args.strategy_config.get("virtual_pipeline_model_parallel_size", None), + "reference/compute_log_probs", + ) + metrics_mgr.add_metrics(dynamic_batching_metrics) ref_log_probs = self.actor_train.compute_log_probs(batch, blocking=True) + ref_log_probs.rename(old_keys="log_probs", new_keys="ref_log_probs") + batch = batch.union(ref_log_probs) + metrics_mgr.add_reduced_metrics(ref_log_probs.meta_info.pop("metrics", {})) else: - batch_balance(batch, dp_size=self.reference.dp_size, minibatch_size=len(batch)) - ref_log_probs = self.reference.compute_log_probs(batch, blocking=True) - metrics_mgr.add_reduced_metrics(ref_log_probs.meta_info.pop("metrics", {})) - ref_log_probs.rename(old_keys="log_probs", new_keys="ref_log_probs") - batch = batch.union(ref_log_probs) + saved_routed_experts = batch.batch.pop("routed_experts", None) + domain_to_teacher_names, domain_values = build_domain_routing_context( + batch, self.pipeline_config, domain_key="domain" + ) + batch = compute_ref_log_probs_with_routing( + batch=batch, + references=self.references, + pipeline_config=self.pipeline_config, + domain_to_teacher_names=domain_to_teacher_names, + domain_values=domain_values, + metrics_fn=metrics_mgr.add_metrics, + ) + if saved_routed_experts is not None: + batch.batch["routed_experts"] = saved_routed_experts metrics_mgr.add_metric("time/ref_log_probs_values", cal_ref_log_probs_timer.last) - with Timer(name="cal_old_log_probs_values", logger=None) as cal_old_logpb_timer: + with Timer(name="cal_old_log_probs_values", logger=None) as cal_old_logpb_timer, \ + tracer.start_as_current_span("old_log_probs_and_values"): if self.pipeline_config.enable_reference and not self.use_ref_model: batch.meta_info["disable_adapter"] = False batch.meta_info["is_offload_states"] = False @@ -619,6 +668,12 @@ def run(self): # Mock ref_log_probs using old_log_probs if reference is disabled if not self.pipeline_config.enable_reference: batch.batch["ref_log_probs"] = batch.batch["old_log_probs"].clone() + # Multi-teacher: also mock per-teacher keys (compute_advantage reads them) + for name in self.pipeline_config.reference_configs.keys(): + batch.batch[f"ref_log_probs_{name}"] = batch.batch["old_log_probs"].clone() + batch.batch[f"ref_log_probs_{name}_mask"] = torch.ones( + batch.batch.batch_size[0], dtype=torch.bool, device=batch.batch["attention_mask"].device + ) metrics_mgr.add_metric("time/old_log_probs", cal_old_logpb_timer.last) # 要按domain group by处理reward @@ -627,13 +682,15 @@ def run(self): batch_list = [] for domain, domain_batch in batch_grouped.items(): # 1. 处理mask相关策略， 获取sample level mask - with Timer(name="get_sample_level_mask", logger=None) as get_sample_level_mask_timer: + with Timer(name="get_sample_level_mask", logger=None) as get_sample_level_mask_timer, \ + tracer.start_as_current_span("get_sample_level_mask"): domain_batch, mask_metrics = get_sample_level_mask(domain_batch, self.pipeline_config) metrics_mgr.add_domain_metrics(domain, mask_metrics) metrics_mgr.add_domain_metrics(domain, {"time/get_sample_level_mask": get_sample_level_mask_timer.last}) # 2. 处理reward相关策略 - with Timer(name="reward_postprocess", logger=None) as reward_postprocess_timer: + with Timer(name="reward_postprocess", logger=None) as reward_postprocess_timer, \ + tracer.start_as_current_span("reward_postprocess"): domain_batch, response_level_metrics = reward_postprocess( domain_batch, self.pipeline_config, self.running ) @@ -641,7 +698,8 @@ def run(self): metrics_mgr.add_domain_metrics(domain, {"time/reward_postprocess": reward_postprocess_timer.last}) # 3. 计算token level rewards - with Timer(name="get_token_reward", logger=None) as get_token_reward_timer: + with Timer(name="get_token_reward", logger=None) as get_token_reward_timer, \ + tracer.start_as_current_span("compute_token_reward"): domain_batch, token_level_metrics = compute_token_reward( domain_batch, self.pipeline_config, self.kl_ctrl ) @@ -650,7 +708,8 @@ def run(self): # 4. 计算advantage final_response_mask = domain_batch.batch["final_response_mask"].clone() - with Timer(name="compute_advantage", logger=None) as compute_advantage_timer: + with Timer(name="compute_advantage", logger=None) as compute_advantage_timer, \ + tracer.start_as_current_span("compute_advantage"): domain_batch = compute_advantage( data=domain_batch, gamma=self.pipeline_config.gamma, @@ -705,7 +764,8 @@ def run(self): 'loss_mask_keys']) metrics_mgr.add_metrics(corr_metrics) - with Timer(name="step_train", logger=None) as step_train_timer: + with Timer(name="step_train", logger=None) as step_train_timer, \ + tracer.start_as_current_span("train"): if self.pipeline_config.adv_estimator == "gae": critic_train_metrics_refs: List[ray.ObjectRef] = self.critic.train_step(batch, blocking=False) @@ -789,10 +849,14 @@ def run(self): @torch.no_grad() def val(self, global_step): + defer = ExitStack() val_metrics_mgr = MetricsManager() + tracer = get_tracer("driver") batch = DataProto() - with Timer(name="step_generate", logger=None) as step_generate_timer: + with Timer(name="step_generate", logger=None) as step_generate_timer, \ + tracer.start_as_current_span("val_generate"): + inject_trace_context(batch.meta_info) batch.meta_info = { "is_offload_states": False, "generation_config": self.pipeline_config.validation.generating_args.to_dict(), @@ -808,6 +872,7 @@ def val(self, global_step): val_metrics_mgr.add_metric("time/step_generate", step_generate_timer.last) batch = generate_output + defer.callback(lambda b=batch: DataProto.drop(b)) val_correct_mean = (batch.batch["scores"] == 1).detach().float().mean().item() val_metrics_mgr.add_metric("val_correct/all/mean", val_correct_mean) logger.info(json.dumps({"val_correct/all/mean": val_correct_mean}, ensure_ascii=False)) @@ -822,4 +887,5 @@ def val(self, global_step): "val_correct", {f"{group_key}/mean": (group_batch.batch["scores"] == 1).detach().float().mean().item()} ) + defer.close() return val_metrics_mgr.get_metrics() diff --git a/roll/pipeline/rlvr/rlvr_rollout_pipeline.py b/roll/pipeline/rlvr/rlvr_rollout_pipeline.py index e55b14d24..6f3ddddc2 100644 --- a/roll/pipeline/rlvr/rlvr_rollout_pipeline.py +++ b/roll/pipeline/rlvr/rlvr_rollout_pipeline.py @@ -53,7 +53,7 @@ def __init__(self, pipeline_config: RLVRConfig): if self.pipeline_config.global_template else self.pipeline_config.actor_train.data_args.template ) - encode_function = get_encode_function(template_name, self.tokenizer, self.pipeline_config.actor_train.data_args) + encode_function = get_encode_function(template_name, self.tokenizer, self.pipeline_config.actor_train.data_args, tag_to_template=self.pipeline_config.tag_to_template) self.val_dataset = preprocess_dataset( self.val_dataset, self.pipeline_config.prompt_length, diff --git a/roll/pipeline/rlvr/rlvr_vlm_pipeline.py b/roll/pipeline/rlvr/rlvr_vlm_pipeline.py index dc7af456c..d894c2216 100644 --- a/roll/pipeline/rlvr/rlvr_vlm_pipeline.py +++ b/roll/pipeline/rlvr/rlvr_vlm_pipeline.py @@ -1,39 +1,41 @@ +import asyncio import copy import json import os import uuid +from contextlib import ExitStack from functools import partial -from io import BytesIO -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any, Dict, List, Optional import datasets import numpy as np -import PIL.Image as Image import ray import torch from codetiming import Timer -from datasets import load_from_disk from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy from ray.util.timer import _Timer -from transformers import AutoConfig, ProcessorMixin -from transformers.image_utils import load_images -from transformers.models.qwen2_vl.image_processing_qwen2_vl import smart_resize from roll.configs import GeneratingArguments from roll.datasets.collator import DataCollatorWithPaddingForMM -from roll.datasets.dataset import get_dataset +from roll.datasets.vlm_dataset_utils import create_pipeline_data_kwargs from roll.distributed.executor.cluster import Cluster from roll.distributed.scheduler.generate_scheduler import DynamicSamplingScheduler +from roll.distributed.scheduler.user_defined_rollout_loop import ( + UserDefinedRolloutLoop, + RolloutContext, + expand_requests, + query_filter, +) from roll.distributed.scheduler.protocol import DataProto from roll.models.model_providers import default_processor_provider, get_extra_data_provider from roll.pipeline.base_pipeline import BasePipeline from roll.pipeline.rlvr.rlvr_config import RLVRConfig from roll.pipeline.rlvr.rlvr_pipeline import update_dataset_domain from roll.pipeline.rlvr.utils import dump_rollout_to_specific_path -from roll.utils.checkpoint_manager import download_model from roll.utils.functionals import ( RunningMoments, agg_loss, + batch_balance, compute_advantage, compute_token_reward, get_sample_level_mask, @@ -44,186 +46,100 @@ from roll.utils.logging import get_logger from roll.utils.metrics.metrics_manager import MetricsManager from roll.utils.offload_states import OffloadStateType +from roll.utils.telemetry import get_tracer, inject_trace_context from roll.utils.train_infer_corrections import apply_train_infer_correction_to_batch logger = get_logger() -def format_prompt(prompt, processor, use_image=True, prompt_image_token=None): - question_template = "{Question} Output the thinking process in and final answer (number) in tags." - if isinstance(prompt, list): - messages = prompt - else: - messages = [ - { - "role": "user", - "content": [ - {"type": "image"}, - {"type": "text", "text": question_template.format(Question=prompt)}, - ] - if use_image and not prompt_image_token - else [ - {"type": "text", "text": question_template.format(Question=prompt)} - ], # image_token has been included in prompt - } - ] - text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) - if prompt_image_token: - text = text.replace(prompt_image_token, "<|vision_start|><|image_pad|><|vision_end|>") - return text - - -def process_image(image: Image.Image, processor: ProcessorMixin): - # same as qwen2-vl image processor - image_processor = processor.image_processor - factor = ( - image_processor.patch_size * image_processor.merge_size - if "Qwen" in image_processor.image_processor_type - else 28 - ) - height, width = image.height, image.width - resized_height, resized_width = smart_resize( - height, - width, - factor=factor, - min_pixels=image_processor.min_pixels, - max_pixels=image_processor.max_pixels, - ) - resized_image = image.resize((resized_width, resized_height), resample=image_processor.resample) - return resized_image - - -def process_images( - images: Union[List, Tuple, str, Image.Image], processor: ProcessorMixin -) -> Union[Image.Image, List[Image.Image], List[List[Image.Image]]]: - """Process images, handling different levels of nesting. - - Args: - images: A single image, a list of images, or a list of lists of images to load. - timeout: Timeout for loading images. - - Returns: - A single image, a list of images, a list of lists of images. +class FiltDataRolloutLoop(UserDefinedRolloutLoop): """ - if isinstance(images, (list, tuple)): - if len(images) and isinstance(images[0], (list, tuple)): - return [[process_image(image, processor=processor) for image in image_group] for image_group in images] - else: - return [process_image(image, processor=processor) for image in images] - else: - return process_image(images, processor=processor) - - -def encode_function( - data, processor, prompt_getter, ground_truth_getter, image_getter, tag_getter, prompt_image_token=None -): - image_flag = [True] * len(prompt_getter(data)) - image_list = [] - for idx, image in enumerate(image_getter(data)): - if not image: - image_flag[idx] = False - try: - if isinstance(image, bytes): # bytes data - # TODO: support multiple images - image_out = Image.open(BytesIO(image)) - else: - image_out = load_images(image if isinstance(image, (list, tuple)) else [image], timeout=None) - except Exception as e: - if isinstance(image, bytes): - image_out = [Image.new("RGB", (224, 224), (255, 255, 255))] - logger.error(f"Failed to get image with type: {type(image)}") - else: - image_out = [Image.new("RGB", (224, 224), (255, 255, 255))] * len(image) - logger.error(f"Failed to get image: {image}") - # since infer-image use pil image as input while train-engine use - # processed data, process image here to make them use same image - # refer to the following for Spatial Understanding with Qwen2.5-VL - # https://github.com/QwenLM/Qwen2.5-VL/blob/main/cookbooks/spatial_understanding.ipynb - # NOTE: process_image from qwen2.5-vl keeps aspect ratio almostly and - # bboxes would be normalized in detection verifier, thus nearly no need - # to change ground-truth bboxes - image_out = process_images(image_out, processor) - image_list.append(image_out) - text_list = [] - for idx, instruct in enumerate(prompt_getter(data)): - # provide prompt_image_token if image_token in prompt - text = format_prompt(instruct, processor, use_image=image_flag[idx], prompt_image_token=prompt_image_token) - text_list.append(text) - encodings = { - "tag": tag_getter(data), - "images": image_list, - "prompt": text_list, - "ground_truth": ground_truth_getter(data), - "reward_model": data["reward_model"], - # for text and multi-modal mixed data usage, indicating valid image - "image_flag": image_flag, - } - return encodings - - -def get_vlm_dataset(data_args, encode_function, processor, get_eval=False): - cache_path = getattr(data_args, "cache_path", None) - if cache_path: - cache_path = os.path.join(cache_path, "val" if get_eval else "train") - if cache_path and os.path.exists(cache_path): - dataset = load_from_disk(cache_path) - return dataset - - dataset = get_dataset(data_args=data_args) - # regularized data filed - features = datasets.Features( - { - "tag": datasets.Value(dtype="string"), # from data_source - "images": datasets.Sequence(feature=datasets.Image(mode=None, decode=True)), - "prompt": datasets.Value(dtype="string"), - "ground_truth": datasets.Value(dtype="string"), - "reward_model": dataset.features["reward_model"], - # for text and multi-modal mixed data usage, indicating valid image - "image_flag": datasets.Value("bool"), - } - ) - remove_columns = list(dataset.features.keys() - features.keys()) - # suit to both VLM-RL/Ocean-R1 and MiniMax-AI/One-RL-to-See-Them-All data - prompt_getter = lambda data: data["prompt"] - ground_truth_getter = lambda data: [x["ground_truth"] for x in data["reward_model"]] - image_getter = lambda data: data["images"] - tag_getter = lambda data: data["data_source"] - print(f"Begin : {dataset}") - dataset = dataset.map( - lambda data: encode_function( - data, processor, prompt_getter, ground_truth_getter, image_getter, tag_getter, prompt_image_token="" - ), - batched=True, - batch_size=100, - num_proc=data_args.preprocessing_num_workers, - features=features, - remove_columns=remove_columns, - desc="Encoding dataset", - ) - print(f"Encoding: {dataset}") - if cache_path: - dataset.save_to_disk(cache_path) - return dataset + custom to filter data whose length is larger than prompt_length + """ + async def process_new_prompt(self, context: RolloutContext) -> Optional[DataProto | List[DataProto]]: + num_return_sequences = context.meta_info["generation_config"]["num_return_sequences"] + is_num_return_sequences_expand = context.is_num_return_sequences_expand + + ################# STEP 1: get and filter dataset + request_data, domain = context.get_request_data(meta_info=context.meta_info) + if request_data.batch["input_ids"].shape[1] > context.prompt_length: + logger.error( + f"prompt_id {context.prompt_id} is filtered, " + f"since input length={request_data.batch['input_ids'].shape[1]} is larger than prompt_length={context.prompt_length}" + ) + return + request_data_list = expand_requests( + data=request_data, + num_return_sequences=num_return_sequences, + is_num_return_sequences_expand=is_num_return_sequences_expand, + ) + + ################# STEP 2: spawn tasks to process requests, including generate, reward, and filter at response level + # Must run inside RolloutContext.do_generate_and_reward context. + # RolloutContext.do_generate_and_reward will wait until can send new request (controlled by LoadBalancer). + # And at exit, RolloutContext will enforce there is no running requests. + async with context.do_generate_and_reward(max_concurrency=num_return_sequences): + responses_list: List[List[DataProto]] = await asyncio.gather( + *[self._generate_and_reward(context=context, req=req, domain=domain) for req in request_data_list] + ) + responses: List[DataProto] = [item for sublist in responses_list for item in sublist] + # some quick methods to reduce store and transfer overhead of multi-modal data by ray + # 1. remove multi_modal_inputs which is for training before generate and add it back after reward + # 2. remove multi_modal_data which is for inference after generate + # 3. change dtype of features in multi_modal_inputs to model dtype which uses lower bytes + for response in responses: + response.non_tensor_batch.pop("multi_modal_data", None) + # User can call RolloutContext.abort_running_requests to abort any running generate requests (generate will return a response + # with finish_reason=="abort", user should distinguish this from partial rollout to avoid dead loop). + # assert there is no running requests outside do_generate_and_reward context. + + ################# STEP 3: prompt level filter + if not context.is_val and not query_filter(responses, context.pipeline_config): + # TODO add metrics (query_filter_count) + logger.debug(f"prompt_id {context.prompt_id} is filtered") + return + + ################# STEP 4: return responses to commit to ReplayBuffer + return responses class RLVRVLMPipeline(BasePipeline): def __init__(self, pipeline_config: RLVRConfig): super().__init__(pipeline_config) self.pipeline_config = pipeline_config + pipeline_config.user_defined_rollout_loop_cls = f"{self.__class__.__module__}.FiltDataRolloutLoop" self.processor = default_processor_provider(self.pipeline_config.actor_train.model_args) - # set max_pixels to avoid image token num is larger than prompt length - self.processor.image_processor.max_pixels, self.processor.image_processor.min_pixels = ( - getattr(self.pipeline_config.actor_train.model_args, "max_pixels", 1024 * 1024), - getattr(self.pipeline_config.actor_train.model_args, "min_pixels", 56 * 56), - ) self.tokenizer = self.processor.tokenizer self.tokenizer.padding_side = "left" - dataset = get_vlm_dataset( - self.pipeline_config.actor_train.data_args, encode_function, self.processor, get_eval=False + # prepare dataset and collect_fn_kwargs + train_data_kwargs = create_pipeline_data_kwargs( + self.pipeline_config.actor_train.data_args, tokenizer=self.tokenizer, processor=self.processor ) + # pipeline related data args + def _data_kwargs_helper(data_kwargs): + dataset, collect_fn_kwargs = data_kwargs["dataset"], data_kwargs["collect_fn_kwargs"] + assert "tag" in dataset.features, "dataset should include tag field to get domain" + collect_fn_kwargs["extra_unpadded_keys"] = list( + set(collect_fn_kwargs.get("extra_unpadded_keys", []) + ["domain"]) + ) + collect_fn_kwargs["extra_data_provider"] = collect_fn_kwargs.get( + "extra_data_provider", + get_extra_data_provider( + self.pipeline_config.actor_train.model_args.model_name_or_path, processor=self.processor + ), + ) + collect_fn_kwargs["max_length"] = collect_fn_kwargs.get("max_length", self.pipeline_config.prompt_length) + collect_fn_kwargs["padding"] = collect_fn_kwargs.get("padding", "max_length") + collect_fn_kwargs["mm_feature_dtype"] = collect_fn_kwargs.get( + "mm_feature_dtype", self.pipeline_config.actor_train.model_args.dtype + ) + return data_kwargs + + train_data_kwargs = _data_kwargs_helper(train_data_kwargs) + dataset, collect_fn_kwargs = train_data_kwargs["dataset"], train_data_kwargs["collect_fn_kwargs"] # update domain field, DynamicSamplingScheduler requires dataset = dataset.map( partial(update_dataset_domain, self.pipeline_config.tag_2_domain), @@ -243,9 +159,15 @@ def __init__(self, pipeline_config: RLVRConfig): self.val_dataset = None if self.pipeline_config.validation and self.pipeline_config.validation.data_args: - self.val_dataset = get_vlm_dataset( - self.pipeline_config.validation.data_args, encode_function, self.processor, get_eval=True + val_data_kwargs = create_pipeline_data_kwargs( + self.pipeline_config.validation.data_args, + tokenizer=self.tokenizer, + processor=self.processor, + is_val=True, ) + val_data_kwargs = _data_kwargs_helper(val_data_kwargs) + val_dataset, val_collect_fn_kwargs = val_data_kwargs["dataset"], val_data_kwargs["collect_fn_kwargs"] + self.val_dataset = val_dataset self.val_dataset = self.val_dataset.map( partial(update_dataset_domain, self.pipeline_config.tag_2_domain), num_proc=self.pipeline_config.actor_train.data_args.preprocessing_num_workers, @@ -326,21 +248,10 @@ def __init__(self, pipeline_config: RLVRConfig): reward_clusters={domain: self.rewards[domain]}, dataset=self.domain_datasets[domain], collect_fn_cls=DataCollatorWithPaddingForMM, - collect_fn_kwargs=dict( - # tokenizer passed by DynamicSamplingScheduler.set_scheduler - # tokenizer=self.tokenizer, - extra_unpadded_keys=["domain", "reward_model"], - extra_data_provider=get_extra_data_provider( - self.pipeline_config.actor_train.model_args.model_name_or_path, processor=self.processor - ), - prompt_key="prompt", - answer_key="ground_truth", - image_key="images", - image_flag_key="image_flag", - max_length=self.pipeline_config.prompt_length, - padding="max_length", - ), + collect_fn_kwargs=collect_fn_kwargs, state=self.state.kv.get(f"scheduler_state_{domain}", None), + # enable the following line to use dataloader to speedup video loading + get_data_item_kwargs=train_data_kwargs.get("get_data_item_kwargs", None), ) self.generate_schedulers[domain] = generate_scheduler self.domain_batch_size[domain] = domain_batch_size @@ -363,22 +274,10 @@ def __init__(self, pipeline_config: RLVRConfig): reward_clusters=self.rewards, dataset=self.val_dataset, collect_fn_cls=DataCollatorWithPaddingForMM, - collect_fn_kwargs=dict( - # tokenizer passed by DynamicSamplingScheduler.set_scheduler - # tokenizer=self.tokenizer, - # val metrics are grouped by tag rather than domain - extra_unpadded_keys=["domain", "reward_model", "tag"], - extra_data_provider=get_extra_data_provider( - self.pipeline_config.actor_train.model_args.model_name_or_path, processor=self.processor - ), - prompt_key="prompt", - answer_key="ground_truth", - image_key="images", - image_flag_key="image_flag", - max_length=self.pipeline_config.prompt_length, - padding="max_length", - ), + collect_fn_kwargs=val_collect_fn_kwargs, is_val=True, + # enable the following line to use dataloader to speedup video loading + get_data_item_kwargs=val_data_kwargs.get("get_data_item_kwargs", None), ) refs = [] @@ -429,6 +328,7 @@ def get_generation_config(self, generating_args: Optional[GeneratingArguments] = @torch.no_grad() def run(self): metrics_mgr = MetricsManager() + tracer = get_tracer("driver") tps_timer = _Timer(window_size=5) actor_infer_timer = _Timer(window_size=5) @@ -453,7 +353,10 @@ def run(self): logger.info(f"pipeline step {global_step} start...") metrics_mgr.clear_metrics() - with tps_timer, Timer(name="step_total", logger=None) as step_total_timer: + with (tps_timer, Timer(name="step_total", logger=None) as step_total_timer, + tracer.start_as_current_span("pipeline_step", attributes={"global_step": global_step}), + ExitStack() as defer, + ): logger.info(f"pre_step_total_time: {pre_step_total_time}") metrics_mgr.add_metric("time/step_total", pre_step_total_time) batch: DataProto = DataProto( @@ -469,13 +372,15 @@ def run(self): self.critic.offload_states(blocking=True) self.actor_train.offload_states(blocking=True) - with Timer(name="step_stop_server", logger=None) as step_stop_server_timer: + with Timer(name="step_stop_server", logger=None) as step_stop_server_timer, \ + tracer.start_as_current_span("stop_server"): if self.pipeline_config.async_pipeline: ray.get([scheduler.pause_sampling.remote() for scheduler in self.generate_schedulers.values()]) self.actor_infer.offload_states(include=OffloadStateType.other_params) metrics_mgr.add_metric("time/step_stop_server", step_stop_server_timer.last) - with Timer(name="step_model_update", logger=None) as step_model_update_timer: + with Timer(name="step_model_update", logger=None) as step_model_update_timer, \ + tracer.start_as_current_span("model_update"): model_update_metrics: Dict = self.model_update(global_step) metrics_mgr.add_metrics(model_update_metrics) batch.meta_info["generation_config"] = self.get_generation_config() @@ -488,7 +393,8 @@ def run(self): reward_cluster.load_states() if self.val_dataset and global_step % self.pipeline_config.eval_steps == 0: - with Timer(name="val_step", logger=None) as val_step_timer: + with Timer(name="val_step", logger=None) as val_step_timer, \ + tracer.start_as_current_span("validation"): val_metrics = self.val(global_step=global_step) metrics_mgr.add_metrics(val_metrics) metrics_mgr.add_metric("time/val_step", val_step_timer.last) @@ -496,10 +402,11 @@ def run(self): # 要按domain group by生成对应的batch with actor_infer_timer, actor_infer_response_timer, Timer( name="step_generate", logger=None - ) as step_generate_timer: + ) as step_generate_timer, tracer.start_as_current_span("generate"): domain_batches = {} scheduler_refs = {} for domain, scheduler in self.generate_schedulers.items(): + inject_trace_context(batch.meta_info) scheduler_refs[domain] = scheduler.get_batch.remote( data=batch, global_step=global_step, batch_size=self.domain_batch_size[domain] ) @@ -523,25 +430,32 @@ def run(self): metrics_mgr.add_metric("time/step_generate", step_generate_timer.last) batch = generate_output + defer.callback(lambda b=batch: DataProto.drop(b)) + # mark here to make megatron get_data_input broadcast with non_batch_tensor batch.meta_info["_broadcast_non_tensor_batch"] = True batch.meta_info["loss_mask_keys"] = ["response_mask", "final_response_mask"] - batch.non_tensor_batch['sample_uuid'] = np.array([str(uuid.uuid4()) for _ in range(batch.batch.shape[0])], dtype=object) - with Timer(name="cal_ref_log_probs", logger=None) as cal_ref_log_probs_timer: + batch.batch["prompt_id"] = torch.arange(batch.batch.batch_size[0], device=batch.batch.device) + + with Timer(name="cal_ref_log_probs", logger=None) as cal_ref_log_probs_timer, \ + tracer.start_as_current_span("cal_ref_log_probs"): if self.pipeline_config.enable_reference: + batch_balance(batch, dp_size=self.reference.dp_size, minibatch_size=len(batch)) ref_log_probs = self.reference.compute_log_probs(batch, blocking=True) metrics_mgr.add_reduced_metrics(ref_log_probs.meta_info.pop("metrics", {})) ref_log_probs.rename(old_keys="log_probs", new_keys="ref_log_probs") batch = batch.union(ref_log_probs) metrics_mgr.add_metric("time/ref_log_probs_values", cal_ref_log_probs_timer.last) - with Timer(name="cal_old_log_probs_values", logger=None) as cal_old_logpb_timer: + with Timer(name="cal_old_log_probs_values", logger=None) as cal_old_logpb_timer, \ + tracer.start_as_current_span("cal_old_log_probs_values"): batch.meta_info["is_offload_states"] = False if self.pipeline_config.adv_estimator == "gae": values_refs: List[ray.ObjectRef] = self.critic.compute_values(batch, blocking=False) if self.pipeline_config.enable_old_logprobs_recompute: + batch_balance(batch, dp_size=self.actor_train.dp_size, minibatch_size=len(batch)) old_log_probs_refs: List[ray.ObjectRef] = self.actor_train.compute_log_probs(batch, blocking=False) old_log_probs = DataProto.materialize_concat(data_refs=old_log_probs_refs) agg_entropy = agg_loss( @@ -568,18 +482,22 @@ def run(self): metrics_mgr.add_metric("time/old_log_probs", cal_old_logpb_timer.last) # group by domain to process reward - batch.batch["prompt_id"] = torch.arange(batch.batch.batch_size[0], device=batch.batch.device) + # Restore original generate order before group_by, so that same-prompt responses + # remain contiguous for correct GRPO group reward normalization (reshape by n_sample). + batch.reorder(indices=torch.argsort(batch.batch["prompt_id"])) batch_grouped: Dict[str, DataProto] = batch.group_by("domain") batch_list = [] for domain, domain_batch in batch_grouped.items(): # 1. get sample level mask - with Timer(name="get_sample_level_mask", logger=None) as get_sample_level_mask_timer: + with Timer(name="get_sample_level_mask", logger=None) as get_sample_level_mask_timer, \ + tracer.start_as_current_span("get_sample_level_mask"): domain_batch, mask_metrics = get_sample_level_mask(domain_batch, self.pipeline_config) metrics_mgr.add_domain_metrics(domain, mask_metrics) metrics_mgr.add_metric("time/get_sample_level_mask", get_sample_level_mask_timer.last) # 2. process reward - with Timer(name="reward_postprocess", logger=None) as reward_postprocess_timer: + with Timer(name="reward_postprocess", logger=None) as reward_postprocess_timer, \ + tracer.start_as_current_span("reward_postprocess"): domain_batch, response_level_metrics = reward_postprocess( domain_batch, self.pipeline_config, self.running ) @@ -587,7 +505,8 @@ def run(self): metrics_mgr.add_domain_metrics(domain, {"time/reward_postprocess": reward_postprocess_timer.last}) # 3. compute token level rewards - with Timer(name="get_token_reward", logger=None) as get_token_reward_timer: + with Timer(name="get_token_reward", logger=None) as get_token_reward_timer, \ + tracer.start_as_current_span("get_token_reward"): domain_batch, token_level_metrics = compute_token_reward( domain_batch, self.pipeline_config, self.kl_ctrl ) @@ -596,7 +515,8 @@ def run(self): # 4. compute advantage final_response_mask = domain_batch.batch["final_response_mask"].clone() - with Timer(name="compute_advantage", logger=None) as compute_advantage_timer: + with Timer(name="compute_advantage", logger=None) as compute_advantage_timer, \ + tracer.start_as_current_span("compute_advantage"): domain_batch = compute_advantage( data=domain_batch, gamma=self.pipeline_config.gamma, @@ -647,13 +567,24 @@ def run(self): batch, corr_metrics = apply_train_infer_correction_to_batch(self.pipeline_config, batch) metrics_mgr.add_metrics(corr_metrics) - with Timer(name="step_train", logger=None) as step_train_timer: + with Timer(name="step_train", logger=None) as step_train_timer, \ + tracer.start_as_current_span("train"): if self.pipeline_config.adv_estimator == "gae": critic_train_metrics_refs: List[ray.ObjectRef] = self.critic.train_step(batch, blocking=False) with actor_train_timer: # implement critic warmup if self.pipeline_config.critic_warmup <= global_step: + # Reorder data for DP rank load balancing + batch_balance_metrics = batch_balance( + batch, + dp_size=self.actor_train.dp_size, + minibatch_size=self.pipeline_config.actor_train.training_args.per_device_train_batch_size + * self.pipeline_config.actor_train.training_args.gradient_accumulation_steps + * self.actor_train.dp_size, + logging_prefix="global_seqlen/actor_train", + ) + metrics_mgr.add_metrics(batch_balance_metrics) # update actor actor_train_metrics_refs = self.actor_train.train_step(batch, blocking=False) actor_train_metrics: DataProto = DataProto.materialize_concat( @@ -713,10 +644,14 @@ def run(self): @torch.no_grad() def val(self, global_step): + defer = ExitStack() val_metrics_mgr = MetricsManager() + tracer = get_tracer("driver") batch = DataProto() - with Timer(name="step_generate", logger=None) as step_generate_timer: + with Timer(name="step_generate", logger=None) as step_generate_timer, \ + tracer.start_as_current_span("val_generate"): + inject_trace_context(batch.meta_info) batch.meta_info["is_offload_states"] = False batch.meta_info["generation_config"] = self.pipeline_config.validation.generating_args.to_dict() batch.meta_info.update( @@ -730,6 +665,8 @@ def val(self, global_step): val_metrics_mgr.add_metric("time/step_generate", step_generate_timer.last) batch = generate_output + defer.callback(lambda b=batch: DataProto.drop(b)) + val_score_mean = batch.batch["scores"].detach().float().mean().item() val_metrics_mgr.add_metric("val_score/all/mean", val_score_mean) logger.info(json.dumps({"val_score/all/mean": val_score_mean}, ensure_ascii=False)) @@ -744,4 +681,5 @@ def val(self, global_step): "val_score", {f"{group_key}/mean": group_batch.batch["scores"].detach().float().mean().item()} ) + defer.close() return val_metrics_mgr.get_metrics() diff --git a/roll/pipeline/rlvr/utils.py b/roll/pipeline/rlvr/utils.py index 11e60b567..48d897a76 100644 --- a/roll/pipeline/rlvr/utils.py +++ b/roll/pipeline/rlvr/utils.py @@ -55,7 +55,7 @@ def json_checker(path:str): def dump_rollout_to_specific_path(path: str, global_step: int, data: DataProto, tokenizer): if not path: return - write_data = copy.deepcopy(data.non_tensor_batch) + write_data = {k: copy.deepcopy(v) for k, v in data.non_tensor_batch.items()} responses = tokenizer.batch_decode(data.batch['responses'], skip_special_tokens=True) data_cnt = len(responses) write_data['responses'] = responses diff --git a/roll/pipeline/sft/sft_pipeline.py b/roll/pipeline/sft/sft_pipeline.py index 97bd7b6a3..5cfbf361f 100644 --- a/roll/pipeline/sft/sft_pipeline.py +++ b/roll/pipeline/sft/sft_pipeline.py @@ -165,8 +165,8 @@ def __init__(self, pipeline_config: SFTConfig): self.pipeline_config.sequence_length, encode_function, num_proc=self.pipeline_config.sft_train.data_args.preprocessing_num_workers) - - global_val_batch_size = dp_size * self.pipeline_config.sft_train.infer_batch_size + + global_val_batch_size = dp_size * ga_steps * self.pipeline_config.sft_train.infer_batch_size self.val_dataloader = DataLoader( dataset=self.val_dataset, batch_size=global_val_batch_size, @@ -225,7 +225,7 @@ def run(self): logger.info(f"metrics: {metrics}") # Update tqdm progress bar - loss = metrics.get("sft_train/loss", 0) + loss = metrics.get("sft_train/loss@sum", 0) pbar.set_postfix({"loss": f"{loss:.4f}", "step": f"{global_step}/{total_steps}"}) self.state.step = global_step diff --git a/roll/platforms/rocm.py b/roll/platforms/rocm.py index 8b4d6f9cf..ff803ab8a 100644 --- a/roll/platforms/rocm.py +++ b/roll/platforms/rocm.py @@ -37,10 +37,10 @@ def get_custom_env_vars(cls) -> dict: "VLLM_ALLOW_INSECURE_SERIALIZATION": "1", # These VLLM related enviroment variables are related to backend. maybe used afterwards. # "VLLM_USE_TRITON_FLASH_ATTN":"0", - "VLLM_ROCM_USE_AITER":"1", - "VLLM_ROCM_USE_AITER_MOE":"1", + "VLLM_ROCM_USE_AITER":"0", + "VLLM_ROCM_USE_AITER_MOE":"0", # "VLLM_ROCM_USE_AITER_ASMMOE":"1", - "VLLM_ROCM_USE_AITER_PAGED_ATTN":"1", + "VLLM_ROCM_USE_AITER_PAGED_ATTN":"0", # "RAY_DEBUG": "legacy", "VLLM_USE_V1": "0", "TORCHINDUCTOR_COMPILE_THREADS": "2", diff --git a/roll/third_party/deepspeed/__init__.py b/roll/third_party/deepspeed/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/roll/third_party/deepspeed/model_update.py b/roll/third_party/deepspeed/model_update.py deleted file mode 100644 index b6452902c..000000000 --- a/roll/third_party/deepspeed/model_update.py +++ /dev/null @@ -1,205 +0,0 @@ -import ray -import torch.distributed as dist -from deepspeed.runtime.zero import GatheredParameters -from peft import get_peft_model_state_dict - -from roll.configs.base_config import PPOConfig -from roll.configs.worker_config import is_actor_infer_overlapping_with_any_cluster -from roll.utils.collective import collective -from roll.utils.logging import get_logger -from roll.utils.network_utils import collect_free_port, get_node_ip -from roll.utils.send_recv_utils import serialize_named_weights - - -logger = get_logger() - - -def _get_ds_param_size(param): - if hasattr(param, "ds_numel"): - ds_numel = param.ds_numel - else: - ds_numel = param.numel() - return ds_numel * param.element_size() - - -def _gather_weights(is_zero3, named_params): - if not is_zero3: - return [(n, p.data) for n, p in named_params] - with GatheredParameters([p for _, p in named_params]): - return [(n, p.data) for n, p in named_params] - - -def gather_deepspeed_weights(model, ds_config, buffer_size): - is_zero3 = ds_config.is_zero3() - named_params = [(name, param) for name, param in model.named_parameters()] - - waiting_params, waiting_params_size = [], 0 - for name, param in named_params: - if waiting_params and waiting_params_size + _get_ds_param_size(param) > buffer_size: - yield _gather_weights(is_zero3, waiting_params) - waiting_params, waiting_params_size = [], 0 - waiting_params_size += _get_ds_param_size(param) - waiting_params.append((name, param)) - - if waiting_params: - yield _gather_weights(is_zero3, waiting_params) - - -class DeepSpeedWeightUpdater: - def __init__(self, pipeline_config: PPOConfig, infer_cluster, worker_config, model_update_name: str, model, ds_config, is_lora): - self.pipeline_config = pipeline_config - self.worker_config = worker_config - self.model_update_name = model_update_name - self.model = model - self.ds_config = ds_config - self.model_update_infer_workers = infer_cluster.workers - self._model_update_buffer_size = pipeline_config.model_update_buffer_size_mb * 1024 * 1024 # Convert MB to bytes - self.is_lora = is_lora - self.infer_worker_config = infer_cluster.worker_config - self.infer_cluster = infer_cluster - self.is_colocated = is_actor_infer_overlapping_with_any_cluster(infer_cluster.worker_config, actor_train=worker_config) - - # Colocated mode attributes - self._infer_parallel_cpu_group = None - self._co_infer_worker = None - self._buffer_num = None - self._broadcast_workers = None - - # Separated mode attributes - self.model_update_group_name = None - self._model_update_locker = None - - if self.is_colocated: - self._setup_colocated_model_update() - else: - self._setup_separated_model_update() - - def model_update(self): - if self.is_colocated: - return self._colocated_model_update() - return self._separated_model_update() - - def _setup_colocated_model_update(self): - logger.info(f"RANK {dist.get_rank()} Setup colocated model update") - infer_worker_devices_num = self.infer_worker_config.num_gpus_per_worker - train_world_size = dist.get_world_size() - - device_start_diff = min(self.worker_config.device_mapping) - min(self.infer_worker_config.device_mapping) - device_end_diff = max(self.worker_config.device_mapping) - max(self.infer_worker_config.device_mapping) - - assert device_start_diff % infer_worker_devices_num == 0 - assert device_end_diff % infer_worker_devices_num == 0 - - for start_rank in range(0, train_world_size, infer_worker_devices_num): - end_rank = start_rank + infer_worker_devices_num - assert end_rank <= train_world_size - group_ranks = list(range(start_rank, end_rank)) - new_group = dist.new_group(ranks=group_ranks, backend="gloo") - if dist.get_rank() in group_ranks: - self._infer_parallel_cpu_group = new_group - infer_worker_idx = dist.get_rank() + device_start_diff // infer_worker_devices_num - self._co_infer_worker = None - if 0 <= infer_worker_idx < len(self.model_update_infer_workers): - self._co_infer_worker = self.model_update_infer_workers[infer_worker_idx] - - # rank0 broadcast to mismatch workers - if dist.get_rank() == 0 and (device_start_diff > 0 or device_end_diff < 0): - self._broadcast_workers = [] - if device_start_diff > 0: - self._broadcast_workers.extend(self.model_update_infer_workers[: device_start_diff // infer_worker_devices_num]) - if device_end_diff < 0: - self._broadcast_workers.extend(self.model_update_infer_workers[device_end_diff // infer_worker_devices_num :]) - self._setup_broadcast_group() - - def _setup_separated_model_update(self): - if dist.get_rank() != 0: - return - - self._broadcast_workers = self.model_update_infer_workers - self._setup_broadcast_group() - - def _setup_broadcast_group(self): - if not self._broadcast_workers: - return - self.model_update_group_name = f"{self.model_update_name}_deepspeed" - num_gpus_per_infer_worker = self.infer_worker_config.num_gpus_per_worker - infer_device_num = num_gpus_per_infer_worker * len(self._broadcast_workers) - master_address, master_port = get_node_ip(), collect_free_port() - - refs = [ - infer_worker.setup_collective_group.remote( - master_address=master_address, - master_port=master_port, - group_name=self.model_update_group_name, - rank_offset=i * num_gpus_per_infer_worker + 1, - world_size=infer_device_num + 1, - ) - for i, infer_worker in enumerate(self._broadcast_workers) - ] - collective.init_collective_group( - infer_device_num + 1, - 0, - group_name=self.model_update_group_name, - master_addr=master_address, - master_port=master_port, - ) - ray.get(refs) - - logger.info(f"Init weights update group {self.model_update_group_name}") - - def _colocated_model_update(self): - refs = [] - for named_weights in gather_deepspeed_weights( - self.model, self.ds_config, buffer_size=self._model_update_buffer_size - ): - serialized_tensors = serialize_named_weights( - named_weights, infer_strategy=self.infer_worker_config.strategy_args.strategy_name - ) - infer_parallel_size = dist.get_world_size(self._infer_parallel_cpu_group) - co_infer_rank = dist.get_rank(self._infer_parallel_cpu_group) - infer_parallel_tensors = [serialized_tensors] # tensors for each infer parallel rank - if infer_parallel_size > 1: - infer_parallel_tensors = [None] * infer_parallel_size if co_infer_rank == 0 else None - dist.gather_object( - serialized_tensors, infer_parallel_tensors, group_dst=0, group=self._infer_parallel_cpu_group - ) - if refs: - ray.get(refs) - refs = [] - if co_infer_rank == 0 and self._co_infer_worker is not None: - refs.append(self._co_infer_worker.update_parameter_in_bucket.remote(infer_parallel_tensors)) - if self._broadcast_workers: - refs.extend(self._broadcast_to_infer_workers(named_weights)) - if refs: - ray.get(refs) - return {} - - def _broadcast_to_infer_workers(self, named_weights) -> list[ray.ObjectRef]: - if not self._broadcast_workers: - return [] - refs = [ - worker.broadcast_parameter.remote( - group_name=self.model_update_group_name, - names=[n for n, _ in named_weights], - dtypes=[w.dtype for _, w in named_weights], - shapes=[w.shape for _, w in named_weights], - ) - for worker in self._broadcast_workers - ] - handles = [] - for _, weight in named_weights: - handles.append( - collective.broadcast(tensor=weight, src_rank=0, group_name=self.model_update_group_name, async_op=True) - ) - for handle in handles: - handle.wait() - return refs - - def _separated_model_update(self): - logger.info(f"start broadcast model update {self.model_update_group_name}") - for named_weights in gather_deepspeed_weights( - self.model, self.ds_config, buffer_size=self._model_update_buffer_size - ): - refs = self._broadcast_to_infer_workers(named_weights) - ray.get(refs) - return {} diff --git a/roll/third_party/deepspeed/offload_states.py b/roll/third_party/deepspeed/offload_states.py deleted file mode 100644 index ede868eb5..000000000 --- a/roll/third_party/deepspeed/offload_states.py +++ /dev/null @@ -1,89 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -from typing import Set -import torch - -from deepspeed.accelerator import get_accelerator -from deepspeed.runtime.zero.offload_config import OffloadStateTypeEnum - -from deepspeed.utils.tensor_fragment import safe_get_local_fp32_param, safe_get_local_optimizer_state - - -def _make_offload_state_key(key): - return f"{key}_offload_buffer" - - -def offload_adam_states(optimizer, device, pin_memory: bool = False, non_blocking: bool = False): - """Move optimizer states to device. Note that this assumes the state structure of DeepSpeed Adam.""" - - def move_key(state, key): - setattr(state[key], "origin_device", state[key].device) - if state[key].device.type == device: - return - offload_buf_key = _make_offload_state_key(key) - if offload_buf_key not in state: - state[offload_buf_key] = torch.empty_like(state[key], device=device) - if pin_memory: - state[offload_buf_key] = get_accelerator().pin_memory(state[offload_buf_key]) - state[offload_buf_key].copy_(state[key], non_blocking=non_blocking) - state[key].data = state[offload_buf_key] - - for _, state in optimizer.state.items(): - if "exp_avg" in state: - move_key(state, "exp_avg") - if "exp_avg_sq" in state: - move_key(state, "exp_avg_sq") - - -def reload_adam_states(optimizer, device, non_blocking: bool = False): - """Move optimizer states to device. Note that this assumes the state structure of DeepSpeed Adam.""" - - def move_back_key(state, key): - if _make_offload_state_key(key) in state: - state[key].data = state[_make_offload_state_key(key)].to(getattr(state[key], "origin_device", device), - non_blocking=non_blocking) - - for _, state in optimizer.state.items(): - if "exp_avg" in state: - move_back_key(state, "exp_avg") - if "exp_avg_sq" in state: - move_back_key(state, "exp_avg_sq") - - -def get_state_devices(model, state: OffloadStateTypeEnum) -> Set[torch.device]: - """Retrieve the devices of the specified state of the model. - - Args: - model (DeepSpeedEngine): The model whose device allocations are to be checked. - state (OffloadStateTypeEnum): The specific state for which the devices should be retrieved. - - Returns: - Set[torch.device]: A set of devices of the specified state. - - """ - trainable_params = [param for param in model.parameters() if param.requires_grad] - if state == OffloadStateTypeEnum.hp_params: - return set(safe_get_local_fp32_param(p).device for p in trainable_params) - elif state == OffloadStateTypeEnum.lp_params: - return set(p.ds_tensor.device for p in model.parameters()) - elif state == OffloadStateTypeEnum.lp_grads: - return {model.optimizer.grad_partitions_flat_buffer.device} - elif state == OffloadStateTypeEnum.optim_states: - return set(safe_get_local_optimizer_state(p, "exp_avg").device for p in trainable_params) | \ - set(safe_get_local_optimizer_state(p, "exp_avg_sq").device for p in trainable_params) - elif state == OffloadStateTypeEnum.contiguous_grad_buffer: - if model.optimizer._DeepSpeedZeroOptimizer_Stage3__ipg_bucket_flat_buffer == None: - return {} - return {model.optimizer._DeepSpeedZeroOptimizer_Stage3__ipg_bucket_flat_buffer.device} - - -def needs_offload(target, include, offloaded_states): - # return True - return target not in offloaded_states and (include == None or target in include) - - -def needs_reload(target, include, offloaded_states): - return (include == None or target in include) and (target in offloaded_states) diff --git a/roll/third_party/deepspeed/offload_states_patch.py b/roll/third_party/deepspeed/offload_states_patch.py deleted file mode 100644 index 62f320921..000000000 --- a/roll/third_party/deepspeed/offload_states_patch.py +++ /dev/null @@ -1,491 +0,0 @@ -import gc -import itertools -import types -from typing import Callable, Dict, Union, Iterable, Container, List - -from deepspeed import DeepSpeedEngine -from deepspeed.accelerator import get_accelerator -from deepspeed.runtime.zero.offload_config import OffloadStateTypeEnum, OffloadDeviceEnum -from deepspeed.runtime.zero.parameter_offload import DeepSpeedZeRoOffload -from deepspeed.runtime.zero.stage3 import DeepSpeedZeroOptimizer_Stage3 -from deepspeed.runtime.zero.stage_1_and_2 import DeepSpeedZeroOptimizer -from deepspeed.runtime.zero.utils import get_mapping_to_flat_buffer -from deepspeed.utils import logger - -import torch -from torch import Tensor - -from roll.third_party.deepspeed.offload_states import needs_offload, needs_reload, offload_adam_states, \ - reload_adam_states - - -def bind_deepspeed_offload_states_func(engine: DeepSpeedEngine): - if hasattr(engine, 'offload_states') and hasattr(engine, 'reload_states'): - setattr(engine, 'old_offload_states', engine.offload_states) - setattr(engine, 'old_reload_states', engine.reload_states) - engine.offload_states = types.MethodType(engine_offload_states, engine) - engine.reload_states = types.MethodType(engine_reload_states, engine) - engine.offload_states = types.MethodType(engine_offload_states, engine) - engine.offload_non_trainable_parameters = types.MethodType(engine_offload_non_trainable_parameters, engine) - engine.reload_non_trainable_parameters = types.MethodType(engine_reload_non_trainable_parameters, engine) - - assert isinstance(engine.optimizer, (DeepSpeedZeroOptimizer, DeepSpeedZeRoOffload, DeepSpeedZeroOptimizer_Stage3)), \ - f"offload_states only supports optimizer [DeepSpeedZeroOptimizer, DeepSpeedZeRoOffload, DeepSpeedZeroOptimizer_Stage3], but get {type(engine.optimizer)}" - - optimizer = engine.optimizer - if hasattr(optimizer, 'offload_states') and hasattr(optimizer, 'reload_states'): - setattr(optimizer, 'old_offload_states', optimizer.offload_states) - setattr(optimizer, 'old_reload_states', optimizer.reload_states) - - if isinstance(optimizer, DeepSpeedZeroOptimizer): - optimizer.offload_states = types.MethodType(stage_1_and_2_offload_states, optimizer) - optimizer.reload_states = types.MethodType(stage_1_and_2_reload_states, optimizer) - elif isinstance(optimizer, DeepSpeedZeRoOffload): - optimizer.offload_states = types.MethodType(parameter_offload_offload_states, optimizer) - optimizer.reload_states = types.MethodType(parameter_offload_reload_states, optimizer) - elif isinstance(optimizer, DeepSpeedZeroOptimizer_Stage3): - optimizer.offload_states = types.MethodType(stage_3_offload_states, optimizer) - optimizer.reload_states = types.MethodType(stage_3_reload_states, optimizer) - else: - raise RuntimeError(f'engine {engine} does not support offload_states func') - - -def engine_offload_non_trainable_parameters(self: DeepSpeedEngine, - include: Container[OffloadStateTypeEnum] = None, - device: OffloadDeviceEnum = OffloadDeviceEnum.cpu, - pin_memory: bool = True, - non_blocking: bool = False): - self.offloaded_states = getattr(self, "offloaded_states", set()) - if needs_offload(OffloadStateTypeEnum.lp_params, include, self.offloaded_states): - if not hasattr(self, "non_trainable_params"): - self.non_trainable_params: List[Tensor] = [] - self.non_trainable_params_meta: List[Tensor] = [] - for param in self.module.parameters(): - if not param.requires_grad: - if hasattr(param, "ds_id"): - self.non_trainable_params.append(param.ds_tensor) - else: - self.non_trainable_params.append(param) - self.non_trainable_params_meta.append( - torch.zeros_like(self.non_trainable_params[-1], device="meta")) - if len(self.non_trainable_params) > 0: - self.non_trainable_lp_param_buffer = self.flatten(self.non_trainable_params) - - if len(self.non_trainable_params) == 0: - return - - if pin_memory: - if not hasattr(self, "non_trainable_lp_param_contiguous_pin_buffer"): - self.non_trainable_lp_param_contiguous_pin_buffer = get_accelerator().pin_memory( - torch.empty_like(self.non_trainable_lp_param_buffer, device=device)) - self.non_trainable_lp_param_contiguous_pin_buffer.copy_(self.non_trainable_lp_param_buffer, - non_blocking=non_blocking) - self.non_trainable_lp_param_buffer.data = self.non_trainable_lp_param_contiguous_pin_buffer - else: - self.non_trainable_lp_param_buffer.data = self.non_trainable_lp_param_buffer.to(device, - non_blocking=non_blocking) - unflatten_tensor = self.unflatten(self.non_trainable_lp_param_buffer, self.non_trainable_params_meta) - for src_tensor, dist_tensor in zip(unflatten_tensor, self.non_trainable_params): - dist_tensor.data = src_tensor.data - - self.offloaded_states.add(OffloadStateTypeEnum.lp_params) - - -def engine_reload_non_trainable_parameters(self: DeepSpeedEngine, include=None, non_blocking: bool = False): - device = get_accelerator().current_device_name() - self.offloaded_states = getattr(self, "offloaded_states", set()) - if needs_reload(OffloadStateTypeEnum.lp_params, include, self.offloaded_states) and len( - self.non_trainable_params) > 0: - lp_param_cpu_buffer = self.non_trainable_lp_param_contiguous_pin_buffer if hasattr( - self, "non_trainable_lp_param_contiguous_pin_buffer") else self.non_trainable_lp_param_buffer - self.non_trainable_lp_param_buffer.data = lp_param_cpu_buffer.data.to(device, non_blocking=non_blocking) - - unflatten_tensor = self.unflatten(self.non_trainable_lp_param_buffer, self.non_trainable_params_meta) - for src_tensor, dist_tensor in zip(unflatten_tensor, self.non_trainable_params): - dist_tensor.data = src_tensor.data - self.offloaded_states.remove(OffloadStateTypeEnum.lp_params) - - -def engine_offload_states(self: DeepSpeedEngine, - include: Container[OffloadStateTypeEnum] = None, - device: OffloadDeviceEnum = OffloadDeviceEnum.cpu, - pin_memory: bool = True, - non_blocking: bool = False) -> None: - """Offload the engine's states to the specified device. - - Arguments: - include: Optional. The set of states to offload. If not provided, all states are offloaded. - device: Optional. The device to move the ZeRO optimizer buffers to. Currently only `OffloadDeviceEnum.cpu` is supported. - pin_memory: Optional. Whether to pin the memory of the offloaded states. - non_blocking: Optional. Whether to offload the states asynchronously. - """ - param_offload_config = self.zero_offload_param() - assert param_offload_config is None or param_offload_config.device == OffloadDeviceEnum.none, "Moving states across devices is not supported for offloaded parameters." - - assert isinstance(self.optimizer, (DeepSpeedZeroOptimizer, DeepSpeedZeRoOffload, DeepSpeedZeroOptimizer_Stage3)), \ - f"offload_states only supports optimizer [DeepSpeedZeroOptimizer, DeepSpeedZeRoOffload, DeepSpeedZeroOptimizer_Stage3], but get {type(self.optimizer)}" - - assert self.zero_optimization_stage() > 0, "offload_states only supports optimizer stage > 0." - - if device == OffloadDeviceEnum.none: - logger.warning("No device specified for offloading states.") - return - - if device == OffloadDeviceEnum.nvme: - raise ValueError("NVMe offload is not supported for offloading states.") - - self.offload_non_trainable_parameters(include=include, device=device, pin_memory=pin_memory, - non_blocking=non_blocking) - self.optimizer.offload_states(include=include, device=device, pin_memory=pin_memory, non_blocking=non_blocking) - - gc.collect() - get_accelerator().empty_cache() - - -def engine_reload_states(self, include: Container[OffloadStateTypeEnum] = None, non_blocking: bool = False) -> None: - """Reload the engine states to the original device. - - Arguments: - include: Optional. The set of states to offload. If not provided, all states are reloaded. - non_blocking: Optional. Whether to offload the states asynchronously. - """ - self.reload_non_trainable_parameters(include=include, non_blocking=non_blocking) - self.optimizer.reload_states(include=include, non_blocking=non_blocking) - - -def stage_1_and_2_offload_states(self: DeepSpeedZeroOptimizer, - include: Container[OffloadStateTypeEnum] = None, - device: OffloadDeviceEnum = OffloadDeviceEnum.cpu, - pin_memory: bool = True, - non_blocking: bool = False): - device = device.value - self.offloaded_states = getattr(self, "offloaded_states", set()) - # HP param - if needs_offload(OffloadStateTypeEnum.hp_params, include, self.offloaded_states) and not \ - self.single_partition_of_fp32_groups[0].is_cpu: - if pin_memory: - if not hasattr(self, "hp_params_pin_buffers"): - self.hp_params_pin_buffers = [ - get_accelerator().pin_memory(torch.empty_like(t, device=device)) - for t in self.single_partition_of_fp32_groups - ] - for src_tensor, dest_buf in zip(self.single_partition_of_fp32_groups, self.hp_params_pin_buffers): - dest_buf.copy_(src_tensor, non_blocking=non_blocking) - src_tensor.data = dest_buf - else: - for buf in self.single_partition_of_fp32_groups: - buf.data = buf.data.to(device, non_blocking=non_blocking) - self.offloaded_states.add(OffloadStateTypeEnum.hp_params) - - # LP param - if needs_offload(OffloadStateTypeEnum.lp_params, include, self.offloaded_states) and not self.bit16_groups_flat[ - 0].is_cpu: - # NOTE: 这里只支持offload optimizer 里的参数部分 - if pin_memory: - if not hasattr(self, "lp_params_pin_buffers"): - self.lp_params_pin_buffers = [ - get_accelerator().pin_memory(torch.empty_like(t, device=device)) - for t in self.bit16_groups_flat - ] - for src_tensor, dest_buf in zip(self.bit16_groups_flat, self.lp_params_pin_buffers): - dest_buf.copy_(src_tensor, non_blocking=non_blocking) - src_tensor.data = dest_buf.data - else: - for buf in self.bit16_groups_flat: - buf.data = buf.data.to(device, non_blocking=non_blocking) - for i in range(len(self.bit16_groups)): - self._update_model_bit16_weights(i) - - self.parallel_partitioned_bit16_groups = [] - self.offloaded_states.add(OffloadStateTypeEnum.lp_params) - - # LP grad - # NOTE: 这里好像没有 grad 缓存 - if needs_offload(OffloadStateTypeEnum.lp_grads, include, self.offloaded_states): - pass - - # contiguous bucket - if needs_offload(OffloadStateTypeEnum.contiguous_grad_buffer, include, self.offloaded_states): - self.ipg_buffer_meta = [] - if hasattr(self, "ipg_buffer") and self.ipg_buffer is not None: - # Record properties like shape, strides, etc. as a meta tensor - for buffer in self.ipg_buffer: - self.ipg_buffer_meta.append(buffer.to("meta")) - self.ipg_buffer = None - if hasattr(self, "temp_grad_buffer_for_gpu_offload"): - if pin_memory: - if not hasattr(self, "temp_grad_buffer_for_gpu_offload_pin_buffer"): - self.temp_grad_buffer_for_gpu_offload_pin_buffer = get_accelerator().pin_memory( - torch.empty_like(self.temp_grad_buffer_for_gpu_offload, device=device)) - self.temp_grad_buffer_for_gpu_offload_pin_buffer.copy_(self.temp_grad_buffer_for_gpu_offload, - non_blocking=non_blocking) - self.temp_grad_buffer_for_gpu_offload.data = self.temp_grad_buffer_for_gpu_offload_pin_buffer - else: - self.temp_grad_buffer_for_gpu_offload.data = self.temp_grad_buffer_for_gpu_offload.data.to(device, - non_blocking=non_blocking) - self.averaged_gradients = {} - self.offloaded_states.add(OffloadStateTypeEnum.contiguous_grad_buffer) - - # Adam - if needs_offload(OffloadStateTypeEnum.optim_states, include, self.offloaded_states): - offload_adam_states(self.optimizer, device, pin_memory=pin_memory, non_blocking=non_blocking) - self.offloaded_states.add(OffloadStateTypeEnum.optim_states) - - # NOTE: 清理额外引用，hp_mapping里包含了一份对全部flat tensor的引用 - for group in self.bit16_groups: - for param in group: - param._hp_mapping = None - self._link_all_hp_params() - - -def stage_1_and_2_reload_states(self: DeepSpeedZeroOptimizer, include=None, non_blocking: bool = False): - device = get_accelerator().current_device_name() - self.offloaded_states = getattr(self, "offloaded_states", set()) - # HP param - if needs_reload(OffloadStateTypeEnum.hp_params, include, self.offloaded_states): - if hasattr(self, "hp_params_pin_buffers"): - for src, dest in zip(self.hp_params_pin_buffers, self.single_partition_of_fp32_groups): - dest.data = src.to(device, non_blocking=non_blocking) - else: - for buf in self.single_partition_of_fp32_groups: - buf.data = buf.data.to(device, non_blocking=non_blocking) - self.offloaded_states.remove(OffloadStateTypeEnum.hp_params) - - # LP Param - if needs_reload(OffloadStateTypeEnum.lp_params, include, self.offloaded_states): - if hasattr(self, "lp_params_pin_buffers"): - for src, dest in zip(self.lp_params_pin_buffers, self.bit16_groups_flat): - dest.data = src.to(device, non_blocking=non_blocking) - else: - for buf in self.bit16_groups_flat: - buf.data = buf.data.to(device, non_blocking=non_blocking) - for i in range(len(self.bit16_groups)): - self._update_model_bit16_weights(i) - data_parallel_partitions = self.get_data_parallel_partitions(self.bit16_groups_flat[i], i) - self.parallel_partitioned_bit16_groups.append(data_parallel_partitions) - self.offloaded_states.remove(OffloadStateTypeEnum.lp_params) - - # LP grad - if needs_reload(OffloadStateTypeEnum.lp_grads, include, self.offloaded_states): - pass - - # contiguous bucket - if needs_reload(OffloadStateTypeEnum.contiguous_grad_buffer, include, self.offloaded_states): - self.ipg_buffer = [] - for buffer in self.ipg_buffer_meta: - self.ipg_buffer.append(torch.empty_like(buffer, device=device)) - if hasattr(self, "temp_grad_buffer_for_gpu_offload"): - cpu_buffer = self.temp_grad_buffer_for_gpu_offload_pin_buffer if hasattr( - self, "temp_grad_buffer_for_gpu_offload_pin_buffer") else self.temp_grad_buffer_for_gpu_offload - self.temp_grad_buffer_for_gpu_offload.data = cpu_buffer.data.to(device, non_blocking=non_blocking) - self.averaged_gradients = {} - self.offloaded_states.remove(OffloadStateTypeEnum.contiguous_grad_buffer) - - # Adam - if needs_reload(OffloadStateTypeEnum.optim_states, include, self.offloaded_states): - reload_adam_states(self.optimizer, device, non_blocking=non_blocking) - self.offloaded_states.remove(OffloadStateTypeEnum.optim_states) - - # NOTE: 恢复link - for group in self.bit16_groups: - for param in group: - param._hp_mapping = None - self._link_all_hp_params() - - if non_blocking: - get_accelerator().synchronize() - - -def stage_3_offload_states(self: DeepSpeedZeroOptimizer_Stage3, - include: Container[OffloadStateTypeEnum] = None, - device: OffloadDeviceEnum = OffloadDeviceEnum.cpu, - pin_memory: bool = True, - non_blocking: bool = False): - device = device.value - self.offloaded_states = getattr(self, "offloaded_states", set()) - self.empty_partition_cache() - - # HP param - if needs_offload(OffloadStateTypeEnum.hp_params, include, self.offloaded_states) and not \ - self.fp32_partitioned_groups_flat[0].is_cpu: - if pin_memory: - if not hasattr(self, "hp_params_pin_buffers"): - self.hp_params_pin_buffers = [ - get_accelerator().pin_memory(torch.empty_like(t, device=device)) - for t in self.fp32_partitioned_groups_flat - ] - - for src_tensor, dest_buf in zip(self.fp32_partitioned_groups_flat, self.hp_params_pin_buffers): - dest_buf.copy_(src_tensor, non_blocking=non_blocking) - src_tensor.data = dest_buf - else: - for buf in self.fp32_partitioned_groups_flat: - buf.data = buf.data.to(device, non_blocking=non_blocking) - self.offloaded_states.add(OffloadStateTypeEnum.hp_params) - - # LP param - if needs_offload(OffloadStateTypeEnum.lp_params, include, self.offloaded_states): - if pin_memory: - if not hasattr(self, "lp_param_contiguous_pin_buffer"): - self.lp_param_contiguous_pin_buffer = get_accelerator().pin_memory( - torch.empty_like(self.lp_param_buffer, device=device)) - # NOTE: lp_param_buffer保存了由optimizer里取到的参数顺序 - # offload的时候先将 lp_param_buffer.cpu() - # 然后将tensor.data cp给model 的tensor.data，这一步也会有顺序不一致问题 - self.lp_param_contiguous_pin_buffer.copy_(self.lp_param_buffer, non_blocking=non_blocking) - cpu_buffer = self.lp_param_contiguous_pin_buffer - else: - cpu_buffer = self.lp_param_buffer.to(device, non_blocking=non_blocking) - - self.lp_param_buffer.data = cpu_buffer - parameter_partitions: List[Tensor] = [param.ds_tensor for sub_group in self.fp16_groups for param in - sub_group] - for tensor, offset, tensor_numel in get_mapping_to_flat_buffer(parameter_partitions): - tensor.data = cpu_buffer.narrow(0, offset, tensor_numel) - - self.fp16_partitioned_groups_flat.clear() - self.offloaded_states.add(OffloadStateTypeEnum.lp_params) - - # LP grad - if needs_offload(OffloadStateTypeEnum.lp_grads, include, - self.offloaded_states) and not self.grad_partitions_flat_buffer.is_cpu: - if pin_memory: - if not hasattr(self, "lp_grad_partitions_flat_pin_buffers"): - self.lp_grad_partitions_flat_pin_buffers = get_accelerator().pin_memory( - torch.empty_like(self.grad_partitions_flat_buffer, device=device)) - self.lp_grad_partitions_flat_pin_buffers.copy_(self.grad_partitions_flat_buffer, - non_blocking=non_blocking) - self.grad_partitions_flat_buffer.data = self.lp_grad_partitions_flat_pin_buffers - else: - self.grad_partitions_flat_buffer.data = self.grad_partitions_flat_buffer.data.to(device) - self.averaged_gradients = {} - # NOTE: self.__param_id_to_grad_partition里存了一份对grad_partitions_flat_buffer的引用，patch修改需要使用名称修饰 - setattr(self, "_DeepSpeedZeroOptimizer_Stage3__param_id_to_grad_partition", {}) - - self.offloaded_states.add(OffloadStateTypeEnum.lp_grads) - - # contiguous bucket - if needs_offload(OffloadStateTypeEnum.contiguous_grad_buffer, include, self.offloaded_states): - if hasattr(self, "_DeepSpeedZeroOptimizer_Stage3__ipg_bucket_flat_buffer"): - # Record properties like shape, strides, etc. as a meta tensor - self.grad_buffer_meta = getattr(self, "_DeepSpeedZeroOptimizer_Stage3__ipg_bucket_flat_buffer").to("meta") - setattr(self, "_DeepSpeedZeroOptimizer_Stage3__ipg_bucket_flat_buffer", None) - self.offloaded_states.add(OffloadStateTypeEnum.contiguous_grad_buffer) - - # Adam - if needs_offload(OffloadStateTypeEnum.optim_states, include, self.offloaded_states): - offload_adam_states(self.optimizer, device, pin_memory=pin_memory, non_blocking=non_blocking) - self.offloaded_states.add(OffloadStateTypeEnum.optim_states) - - -def stage_3_reload_states(self: DeepSpeedZeroOptimizer_Stage3, include=None, non_blocking: bool = False): - device = get_accelerator().current_device_name() - self.offloaded_states = getattr(self, "offloaded_states", set()) - # HP param - if needs_reload(OffloadStateTypeEnum.hp_params, include, self.offloaded_states): - if hasattr(self, "hp_params_pin_buffers"): - for src, dest in zip(self.hp_params_pin_buffers, self.fp32_partitioned_groups_flat): - dest.data = src.to(device, non_blocking=non_blocking) - else: - for buf in self.fp32_partitioned_groups_flat: - buf.data = buf.data.to(device, non_blocking=non_blocking) - self.offloaded_states.remove(OffloadStateTypeEnum.hp_params) - - # LP Param - if needs_reload(OffloadStateTypeEnum.lp_params, include, self.offloaded_states): - cpu_buffer = self.lp_param_contiguous_pin_buffer if hasattr( - self, "lp_param_contiguous_pin_buffer") else self.lp_param_buffer - self.lp_param_buffer.data = cpu_buffer.data.to(device, non_blocking=non_blocking) - self._set_fp16_partitioned_groups_flat() - - # NOTE: 这里遍历的是self.module.parameters()， 而lp_param_buffer里的是fp16 group里取到的，这里参数的顺序不一致 - # 这里[p.ds_tensor for p in self.module.parameters()]需要按self.fp16_groups的顺序reorder一下 - parameter_partitions: List[Tensor] = [param.ds_tensor for sub_group in self.fp16_groups for param in sub_group] - for tensor, offset, tensor_numel in get_mapping_to_flat_buffer(parameter_partitions): - tensor.data = self.lp_param_buffer.data.narrow(0, offset, tensor_numel) - self.offloaded_states.remove(OffloadStateTypeEnum.lp_params) - - # LP grad - if needs_reload(OffloadStateTypeEnum.lp_grads, include, self.offloaded_states): - if hasattr(self, "lp_grad_partitions_flat_pin_buffers"): - self.grad_partitions_flat_buffer.data = self.lp_grad_partitions_flat_pin_buffers.to( - device, non_blocking=non_blocking) - else: - self.grad_partitions_flat_buffer.data = self.grad_partitions_flat_buffer.data.to( - device, non_blocking=non_blocking) - self.averaged_gradients = {} - - offset = 0 - all_params = list(itertools.chain.from_iterable(self.fp16_groups)) - - param_id_to_grad_partition = getattr(self, "_DeepSpeedZeroOptimizer_Stage3__param_id_to_grad_partition") - for param in all_params: - param_id_to_grad_partition[param.ds_id] = self.grad_partitions_flat_buffer.narrow( - 0, offset, param.partition_numel()) - offset += param.partition_numel() - - self.offloaded_states.remove(OffloadStateTypeEnum.lp_grads) - - # contiguous bucket - if needs_reload(OffloadStateTypeEnum.contiguous_grad_buffer, include, self.offloaded_states): - setattr(self, "_DeepSpeedZeroOptimizer_Stage3__ipg_bucket_flat_buffer", torch.empty_like(self.grad_buffer_meta, device=device)) - self.offloaded_states.remove(OffloadStateTypeEnum.contiguous_grad_buffer) - - # Adam - if needs_reload(OffloadStateTypeEnum.optim_states, include, self.offloaded_states): - reload_adam_states(self.optimizer, device, non_blocking=non_blocking) - self.offloaded_states.remove(OffloadStateTypeEnum.optim_states) - - if non_blocking: - get_accelerator().synchronize() - - -def parameter_offload_offload_states(self: DeepSpeedZeRoOffload, - include: Container[OffloadStateTypeEnum] = None, - device: OffloadDeviceEnum = OffloadDeviceEnum.cpu, - pin_memory: bool = True, - non_blocking: bool = False): - device = device.value - self.empty_partition_cache() - self.offloaded_states = getattr(self, "offloaded_states", set()) - - if needs_offload(OffloadStateTypeEnum.lp_params, include, self.offloaded_states): - # NOTE: 这里不会执行了non_trainable_params都在engine里处理了 - if not hasattr(self, "trainable_params"): - self.trainable_params = [param.ds_tensor for param in self.module.parameters() if param.requires_grad] - if len(self.trainable_params) == 0: - return - - if not hasattr(self, "lp_param_buffer"): - from deepspeed.runtime.zero.stage3 import DeepSpeedZeroOptimizer_Stage3 - self.lp_param_buffer = DeepSpeedZeroOptimizer_Stage3.defragment(self.trainable_params) - - if pin_memory: - if not hasattr(self, "lp_param_contiguous_pin_buffer"): - self.lp_param_contiguous_pin_buffer = get_accelerator().pin_memory( - torch.empty_like(self.lp_param_buffer, device=device)) - self.lp_param_contiguous_pin_buffer.copy_(self.lp_param_buffer, non_blocking=non_blocking) - lp_param_cpu_buffer = self.lp_param_contiguous_pin_buffer - else: - lp_param_cpu_buffer = self.lp_param_buffer.to(device, non_blocking=non_blocking) - self.lp_param_buffer.data = lp_param_cpu_buffer - for tensor, offset, tensor_numel in get_mapping_to_flat_buffer( - [p.ds_tensor for p in self.module.parameters() if p.requires_grad]): - tensor.data = self.lp_param_buffer.data.narrow(0, offset, tensor_numel) - - self.offloaded_states.add(OffloadStateTypeEnum.lp_params) - - -def parameter_offload_reload_states(self: DeepSpeedZeRoOffload, include=None, non_blocking: bool = False): - device = get_accelerator().current_device_name() - self.offloaded_states = getattr(self, "offloaded_states", set()) - if needs_reload(OffloadStateTypeEnum.lp_params, include, self.offloaded_states): - lp_param_cpu_buffer = self.lp_param_contiguous_pin_buffer if hasattr( - self, "lp_param_contiguous_pin_buffer") else self.lp_param_buffer - self.lp_param_buffer.data = lp_param_cpu_buffer.data.to(device, non_blocking=non_blocking) - - for tensor, offset, tensor_numel in get_mapping_to_flat_buffer( - [p.ds_tensor for p in self.module.parameters() if p.requires_grad]): - tensor.data = self.lp_param_buffer.narrow(0, offset, tensor_numel) - - self.offloaded_states.remove(OffloadStateTypeEnum.lp_params) diff --git a/roll/third_party/fsdp2/model_update.py b/roll/third_party/fsdp2/model_update.py index b2dffb4c7..59676d383 100644 --- a/roll/third_party/fsdp2/model_update.py +++ b/roll/third_party/fsdp2/model_update.py @@ -3,6 +3,8 @@ import ray import torch +import transformers +from packaging import version import torch.distributed as dist from torch.distributed.tensor import DTensor @@ -13,6 +15,7 @@ from roll.utils.logging import get_logger from roll.utils.network_utils import collect_free_port, get_node_ip from roll.utils.send_recv_utils import serialize_named_weights +from roll.third_party.fsdp2.qwen3_moe_patch import _iter_convert_fsdps_moe_weights logger = get_logger() @@ -21,6 +24,11 @@ def gather_fsdp2_weights(model, buffer_size, is_lora=False): """ Gather FSDP2 weights for model update. For FSDP2, we need to get the full tensor from the sharded parameters. + Also converts FSDP MoE weights from FSDP format to HF/vLLM expected format. + + Yields batches of (name, tensor) pairs whose total size does not exceed + buffer_size. Expert DTensor parameters are gathered lazily (one layer at a + time) to avoid materializing all expert weights on GPU simultaneously. """ if is_lora: from peft.utils import get_peft_model_state_dict @@ -31,7 +39,12 @@ def gather_fsdp2_weights(model, buffer_size, is_lora=False): named_params = [(name, param) for name, param in model.named_parameters()] waiting_params, waiting_params_size = [], 0 - for name, param in named_params: + need_convert_moe = ( + version.parse(transformers.__version__) >= version.parse("5.2.0") + and getattr(model.config, "model_type", "") in ("qwen3_moe", "qwen3_next") + ) + params_iter = _iter_convert_fsdps_moe_weights(named_params) if need_convert_moe else named_params + for name, param in params_iter: full_tensor_size = param.numel() * param.element_size() if waiting_params and waiting_params_size + full_tensor_size > buffer_size: yield [(n, p.data if not isinstance(p.data, DTensor) else p.data.full_tensor()) for n, p in waiting_params] @@ -223,7 +236,7 @@ def _colocated_model_update(self): infer_parallel_tensors = [None] * infer_parallel_size if co_infer_rank == 0 else None global_dst_rank = dist.get_global_rank(self._infer_parallel_cpu_group, 0) dist.gather_object( - serialized_tensors, infer_parallel_tensors, dst=global_dst_rank, group=self._infer_parallel_cpu_group + serialized_tensors, infer_parallel_tensors, global_dst_rank, group=self._infer_parallel_cpu_group ) if refs: ray.get(refs) diff --git a/roll/third_party/fsdp2/qwen3_moe_patch.py b/roll/third_party/fsdp2/qwen3_moe_patch.py index 5686a8495..e9dfea948 100644 --- a/roll/third_party/fsdp2/qwen3_moe_patch.py +++ b/roll/third_party/fsdp2/qwen3_moe_patch.py @@ -1,5 +1,6 @@ import torch import torch.nn.functional as F +from torch.distributed.tensor import DTensor # force each expert to participate in computation graph so FSDP could gather all expert outputs @@ -34,3 +35,55 @@ def qwen3_moe_forward(self, hidden_states: torch.Tensor) -> torch.Tensor: final_hidden_states = final_hidden_states.reshape(batch_size, sequence_length, hidden_dim) return final_hidden_states, router_logits + + +def _iter_convert_fsdps_moe_weights(named_params): + """ + Lazily convert FSDP MoE weights from FSDP format to HF/vLLM expected format. + + Yields (name, tensor, full_tensor_size) one at a time so that the caller can + control memory by batching into buffers. For expert parameters stored as + DTensors (EP mode), full_tensor() is called here to gather the shards, but + because this is a generator the full tensor is only materialized when the + caller consumes the item. + + FSDP format: + - layers.0.mlp.experts.gate_up_proj: [num_experts, 2 * intermediate_dim, hidden_dim] + - layers.0.mlp.experts.down_proj: [num_experts, hidden_dim, intermediate_dim] + + Expected format: + - layers.0.mlp.experts.0.gate_proj.weight: [intermediate_dim, hidden_dim] + - layers.0.mlp.experts.0.up_proj.weight: [intermediate_dim, hidden_dim] + - layers.0.mlp.experts.0.down_proj.weight: [hidden_dim, intermediate_dim] + """ + for name, param in named_params: + if "experts.gate_up_proj" in name: + if isinstance(param, DTensor): + param = param.full_tensor() + + gate_proj, up_proj = torch.chunk(param, 2, dim=1) + num_experts = gate_proj.shape[0] + + for expert_idx in range(num_experts): + gate_name = name.replace("gate_up_proj", f"{expert_idx}.gate_proj.weight") + up_name = name.replace("gate_up_proj", f"{expert_idx}.up_proj.weight") + yield gate_name, gate_proj[expert_idx] + yield up_name, up_proj[expert_idx] + + # Allow GC to reclaim the full tensor once all slices are yielded + del param, gate_proj, up_proj + + elif "experts.down_proj" in name: + if isinstance(param, DTensor): + param = param.full_tensor() + + num_experts = param.shape[0] + + for expert_idx in range(num_experts): + down_name = name.replace("down_proj", f"{expert_idx}.down_proj.weight") + yield down_name, param[expert_idx] + + del param + + else: + yield name, param diff --git a/roll/third_party/megatron/compile_warmup.py b/roll/third_party/megatron/compile_warmup.py new file mode 100644 index 000000000..fada90ae9 --- /dev/null +++ b/roll/third_party/megatron/compile_warmup.py @@ -0,0 +1,146 @@ +import random +from contextlib import contextmanager +from typing import TYPE_CHECKING, Any + +import numpy as np +import torch +import torch.distributed as dist +from megatron.core import mpu, tensor_parallel + +from mcore_adapter.trainer.utils import get_ltor_masks_and_position_ids +from roll.platforms import current_platform +from roll.utils.constants import IGNORE_INDEX +from roll.utils.logging import get_logger + + +if TYPE_CHECKING: + from roll.distributed.strategy.megatron_strategy import MegatronTrainStrategy + + +logger = get_logger() + + +@contextmanager +def _preserve_rng_states(): + """Save and restore all RNG states to prevent warmup from perturbing training randomness.""" + saved = { + "random": random.getstate(), + "numpy": np.random.get_state(), + "torch": torch.get_rng_state(), + "cuda": current_platform.get_rng_state(), + "tracker": tensor_parallel.get_cuda_rng_tracker().get_states(), + } + try: + yield + finally: + random.setstate(saved["random"]) + np.random.set_state(saved["numpy"]) + torch.set_rng_state(saved["torch"]) + current_platform.set_rng_state(saved["cuda"]) + tensor_parallel.get_cuda_rng_tracker().set_states(saved["tracker"]) + + +def _build_warmup_inputs(strategy: "MegatronTrainStrategy", seq_length: int) -> dict[str, Any]: + """Build synthetic inputs that mirror the shape/structure produced by _prepare_train_inputs.""" + batch_size = strategy.worker_config.training_args.per_device_train_batch_size + input_ids = torch.arange(seq_length, device=strategy.megatron_train_args.device, dtype=torch.long).repeat(batch_size, 1) + attention_mask = torch.ones_like(input_ids) + labels = torch.full_like(input_ids, IGNORE_INDEX) + inputs = {"input_ids": input_ids, "attention_mask": attention_mask, "labels": labels} + + if strategy.worker_config.use_sequence_packing: + input_ids, packed_seq_params, cu_seqlens, cu_seqlens_padded = strategy._pack_sequences( + input_ids, attention_mask, + ) + labels, _, _, _ = strategy._pack_sequences(labels, attention_mask, pad_val=IGNORE_INDEX) + attention_mask = None + position_ids = None + inputs = {"input_ids": input_ids, "attention_mask": attention_mask, "labels": labels, "packed_seq_params": packed_seq_params} + else: + attention_mask, position_ids = get_ltor_masks_and_position_ids( + input_ids, + build_attention_mask=strategy.model.config.transformer_impl != "transformer_engine", + attn_mask_1D=attention_mask, + ) + if not strategy.model.config.num_moe_experts and strategy.model.config.transformer_impl == "transformer_engine": + attention_mask = None + inputs["attention_mask"] = attention_mask + + inputs["position_ids"] = position_ids if strategy.model.config.mtp_num_layers else None + return strategy.models_unwrapped[0].get_batch_on_this_cp_rank(inputs, dim3_keys=[]) + + +def _build_decoder_input( + strategy: "MegatronTrainStrategy", inputs: dict[str, Any], module: torch.nn.Module +) -> torch.Tensor: + """Build a zero-filled hidden-state tensor for non-first pipeline stages.""" + config = strategy.model.config + cp_size = strategy.worker.rank_info.cp_size + seq_length = inputs["input_ids"].shape[1] // cp_size + if config.sequence_parallel: + seq_length = max(1, seq_length // mpu.get_tensor_model_parallel_world_size()) + batch_size = inputs["input_ids"].shape[0] + dtype = ( + getattr(config, "pipeline_dtype", None) + or getattr(config, "params_dtype", None) + or next(module.parameters()).dtype + ) + return torch.zeros((seq_length, batch_size, config.hidden_size), device=strategy.megatron_train_args.device, dtype=dtype) + + +@torch.no_grad() +def compile_warmup_pipeline_stages(strategy: "MegatronTrainStrategy"): + """Run a local forward pass on each pipeline chunk to trigger torch.compile ahead of real training. + + Without this, the first training step serializes compilation across pipeline stages + because P2P communication blocks until upstream stages finish compiling. + """ + if mpu.get_pipeline_model_parallel_world_size() <= 1: + return + + seq_length = strategy.seq_length or 8192 + cp_size = strategy.worker.rank_info.cp_size + tp_size = strategy.worker.rank_info.tp_size + chunk = 128 + if cp_size > 1: + # _get_batch_on_this_cp_rank will split by cp_size*2; ensure seq_length is divisible. + chunk *= cp_size + if tp_size > 1 and strategy.megatron_train_args.sequence_parallel: + chunk *= tp_size + seq_length = ((seq_length + chunk - 1) // chunk) * chunk + + logger.info( + "Running local pipeline compile warmup with batch_size=%s, seq_length=%s.", + strategy.worker_config.training_args.per_device_train_batch_size, + seq_length, + ) + + original_vp_rank = mpu.get_virtual_pipeline_model_parallel_rank() + with _preserve_rng_states(): + try: + for model in strategy.models_wrapped: + model.train() + model.zero_grad_buffer() + strategy.optimizer.zero_grad() + + for model in strategy.models_wrapped: + module = getattr(model, "module", model) + vp_stage = getattr(module, "vp_stage", None) + if vp_stage is not None: + mpu.set_virtual_pipeline_model_parallel_rank(vp_stage) + + inputs = _build_warmup_inputs(strategy, seq_length) + if not getattr(module, "pre_process", True): + decoder_input = _build_decoder_input(strategy, inputs, module) + module.set_input_tensor(decoder_input) + inputs["decoder_input"] = decoder_input + + model(**inputs) + finally: + if original_vp_rank is not None: + mpu.set_virtual_pipeline_model_parallel_rank(original_vp_rank) + for model in strategy.models_wrapped: + model.zero_grad_buffer() + strategy.optimizer.zero_grad() + if dist.is_initialized(): + dist.barrier() diff --git a/roll/third_party/megatron/fused_entropy.py b/roll/third_party/megatron/fused_entropy.py new file mode 100644 index 000000000..d0c4529fa --- /dev/null +++ b/roll/third_party/megatron/fused_entropy.py @@ -0,0 +1,941 @@ +""" +Fused Entropy Kernel Implementation + +This module provides CUDA fused kernel implementations for entropy computation, +optimized for tensor parallel size 1 (TP=1). It accelerates the entropy calculation +used in tensor_parallel.py by fusing multiple operations into a single kernel pass. + +The entropy forward computation is split into two stages: +- Stage 1: Computes max_x and sum_exp(x - max_x) for each row +- Stage 2: Computes the final entropy value: H = max_x + log(sum_exp) - sum(softmax(x) * x) + +Note: Currently only supports TP=1 (single GPU tensor parallelism). +""" + +import math +import operator +from functools import partial +from typing import Optional, Tuple, Type, Literal, Callable + +import torch +from torch import Tensor + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +from cutlass import Int32, Int64, Float32, Boolean, const_expr, BFloat16 + +import quack.utils as utils +import quack.copy_utils as copy_utils +import quack.layout_utils as layout_utils +from quack.compile_utils import make_fake_tensor as fake_tensor +from quack.reduce import row_reduce, online_softmax_reduce +from quack.reduction_base import ReductionBase +from quack.cute_dsl_utils import torch2cute_dtype_map + + +class FusedEntropyForwardStage1(ReductionBase): + """ + First stage of the fused entropy forward pass. + + This kernel computes the intermediate values needed for entropy computation: + - max_x: Maximum value in each row (for numerical stability) + - sum_exp: Sum of exp(x - max_x) for each row (softmax denominator) + + Args: + dtype: Data type of input tensor (e.g., BFloat16, Float16, Float32) + N: Number of columns in input tensor (vocabulary size) + online_softmax: If True, uses online softmax algorithm for better performance + """ + + def __init__(self, dtype: Type[cutlass.Numeric], N: int, online_softmax: bool = True): + """ + Initialize the first stage of entropy forward pass. + """ + # 2 stages: 1 for max, 1 for sum when not using online softmax + # 1 stage when using online softmax (computes both max and sum together) + super().__init__( + dtype, + N, + stage=2 if not online_softmax else 1, + reduction_dtype=Float32 if not online_softmax else Int64, + ) + self.online_softmax = online_softmax + # For large N (>16384), reload data from shared memory to avoid register spilling + self.reload_from = None if N <= 16384 or online_softmax else "smem" + + def _threads_per_row(self): + """ + Determine the optimal number of threads per row for reduction. + """ + N = self.N + for limit, threads in [(64, 8), (128, 16), (3072, 32), (6144, 64), (16384, 128)]: + if N <= limit: + return threads + return 256 + + def _set_cluster_n(self): + """ + Determine the optimal cluster size in the N dimension. + """ + N = self.N + if const_expr(self.dtype.width == 16): + thresholds = [(16 * 1024, 1), (32 * 1024, 2), (64 * 1024, 4), (128 * 1024, 8)] + else: + thresholds = [(16 * 1024, 1), (64 * 1024, 2), (128 * 1024, 4), (256 * 1024, 8)] + for limit, cluster in thresholds: + if N <= limit: + self.cluster_n = cluster + return + self.cluster_n = 16 + + @cute.jit + def __call__( + self, + mX: cute.Tensor, # (M, N) in + mXMax: cute.Tensor, # (M,) out + mXSumExp: cute.Tensor, # (M,) out + stream: cuda.CUstream, + ): + """ + Launch the Stage 1 kernel to compute max and sum_exp for each row. + + Args: + mX: Input logits tensor of shape (M, N) + mXMax: Output tensor for max values of shape (M,) + mXSumExp: Output tensor for sum of exponentials of shape (M,) + stream: CUDA stream for kernel execution + """ + assert mX.element_type == self.dtype + self._set_cluster_n() + # Calculate vector size for memory loading (max 128 bits per load) + largest_dtype_width = const_expr(mX.element_type.width) + vecsize = math.gcd(self.N, 128 // largest_dtype_width) # 8 for bf16/fp16 + tiled_copy, tiler_mn, threads_per_row = self._get_tiled_copy(vecsize=vecsize) + num_threads = tiled_copy.size + # Launch kernel with appropriate grid, block, and cluster configuration + self.kernel( + mX, + mXMax, + mXSumExp, + tiler_mn, + tiled_copy, + threads_per_row, + ).launch( + grid=[cute.ceil_div(mX.shape[0], tiler_mn[0]), self.cluster_n, 1], + block=[num_threads, 1, 1], + cluster=[1, self.cluster_n, 1] if const_expr(self.cluster_n > 1) else None, + stream=stream, + ) + + @cute.kernel + def kernel( + self, + mX: cute.Tensor, # (M, N) in + mXMax: cute.Tensor, # (M,) out + mXSumExp: cute.Tensor, # (M,) out + tiler_mn: cute.Shape, + tiled_copy: cute.TiledCopy, + threads_per_row: cutlass.Constexpr[int], + ): + """ + CUDA kernel for Stage 1: Compute max and sum_exp for each row. + + This kernel loads input data, performs reduction operations to compute + the maximum value and sum of exponentials for each row, and writes + the results to global memory. + + Args: + mX: Input logits tensor of shape (M, N) + mXMax: Output tensor for max values of shape (M,) + mXSumExp: Output tensor for sum of exponentials of shape (M,) + tiler_mn: Tiling configuration + tiled_copy: Tiled copy configuration for efficient memory access + threads_per_row: Number of threads working on each row + """ + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + cluster_y = const_expr(0) if const_expr(self.cluster_n == 1) else cute.arch.block_idx()[1] + tv_layout = tiled_copy.layout_tv_tiled + + shape = mX.shape + idX = cute.make_identity_tensor(shape) + gX, cX = [cute.local_tile(mT, tiler_mn, (bidx, cluster_y)) for mT in (mX, idX)] + + smem = cutlass.utils.SmemAllocator() + sX = smem.allocate_tensor( + mX.element_type, cute.make_ordered_layout(tiler_mn, order=(1, 0)), byte_alignment=16 + ) + reduction_buffer, mbar_ptr = self._allocate_reduction_buffer_and_mbar(smem, tv_layout) + + thr_copy = tiled_copy.get_slice(tidx) + + # Partition tensors for this thread + tXgX = thr_copy.partition_S(gX) # Global memory partition + tXsX = thr_copy.partition_D(sX) # Shared memory partition + tXcX = thr_copy.partition_S(cX)[(0, None), None, None] # Coordinate partition + tXrX = cute.make_fragment_like(tXgX) # Register fragment + + # Handle non-divisible N by creating predicates + is_even_N = const_expr(shape[1] == tiler_mn[1] * self.cluster_n) + tXpX = None if is_even_N else copy_utils.predicate_k(thr_copy.partition_S(cX), limit=shape[1]) + copy = partial(copy_utils.copy, pred=tXpX) + + # Initialize cluster synchronization + num_warps = cute.size(tiled_copy) // cute.arch.WARP_SIZE + self._initialize_cluster(tidx, mbar_ptr, num_warps) + + row = tXcX[0][0] + + # Load input data asynchronously from global to shared memory + if row < shape[0]: + copy(tXgX, tXsX, is_async=True) + cute.arch.cp_async_commit_group() + cute.arch.cp_async_wait_group(0) + # Fill out-of-bounds values with -inf (for numerical stability in max reduction) + if const_expr(not is_even_N): + # utils.fill_oob(tXsX, tXpX, -tXsX.element_type.inf) + utils.fill_oob(tXsX, tXpX, BFloat16(-(2**15))) + # Copy from shared memory to registers and convert to Float32 + cute.autovec_copy(tXsX, tXrX) + x = tXrX.load().to(Float32) + + # Compute max and sum_exp using either two-pass or online softmax + if const_expr(not self.online_softmax): + max_x = row_reduce( + x, + cute.ReductionOp.MAX, + threads_per_row, + reduction_buffer[None, None, 0], + mbar_ptr + 0 if const_expr(self.cluster_n > 1) else None, + init_val=-Float32.inf, + hook_fn=cute.arch.cluster_wait if const_expr(self.cluster_n > 1) else None, + ) + if const_expr(self.reload_from == "smem"): + cute.autovec_copy(tXsX, tXrX) + x = tXrX.load().to(Float32) + log2_e = math.log2(math.e) + exp_x = cute.math.exp2(x * log2_e - (max_x * log2_e), fastmath=False) + denom = row_reduce( + exp_x, + cute.ReductionOp.ADD, + threads_per_row, + reduction_buffer[None, None, 1], + mbar_ptr + 1 if const_expr(self.cluster_n > 1) else None, + init_val=0.0, + ) + else: + # Online softmax: compute both max and sum in a single pass + max_x, denom, _ = online_softmax_reduce( + x, + threads_per_row, + reduction_buffer[None, None, 0], + mbar_ptr, + hook_fn=cute.arch.cluster_wait if const_expr(self.cluster_n > 1) else None, + return_exp_x=False + ) + + # Write results to global memory (only first thread in row, and first block in cluster) + if ( + tXcX[0][1] == 0 + and row < shape[0] + and (self.cluster_n == 1 or cute.arch.block_idx_in_cluster() == 0) + ): + mXMax[row] = max_x + mXSumExp[row] = denom + + +class FusedEntropyForwardStage2(ReductionBase): + """ + Second stage of the fused entropy forward pass. + + This kernel computes the final entropy value using the intermediate results + from Stage 1. The entropy formula is: + H = max_x + log(sum_exp) - sum(softmax(x) * x) + + Where: + - max_x: Maximum value from Stage 1 + - sum_exp: Sum of exponentials from Stage 1 + - sum(softmax(x) * x): Expected value under the softmax distribution + + Args: + dtype: Data type of input tensor (e.g., BFloat16, Float16, Float32) + N: Number of columns in input tensor (vocabulary size) + """ + + def __init__(self, dtype: Type[cutlass.Numeric], N: int): + """ + Initialize the second stage of entropy forward pass. + """ + # 1 stage: for computing sum(softmax * x) + super().__init__( + dtype, + N, + stage=1, + reduction_dtype=Float32, + ) + # For large N (>16384), reload data from shared memory to avoid register spilling + self.reload_from = None if N <= 16384 else "smem" + + def _threads_per_row(self): + """ + Determine the optimal number of threads per row for reduction. + """ + N = self.N + for limit, threads in [(64, 8), (128, 16), (3072, 32), (6144, 64), (16384, 128)]: + if N <= limit: + return threads + return 256 + + def _set_cluster_n(self): + """ + Determine the optimal cluster size in the N dimension. + """ + N = self.N + if const_expr(self.dtype.width == 16): + thresholds = [(16 * 1024, 1), (32 * 1024, 2), (64 * 1024, 4), (128 * 1024, 8)] + else: + thresholds = [(16 * 1024, 1), (64 * 1024, 2), (128 * 1024, 4), (256 * 1024, 8)] + for limit, cluster in thresholds: + if N <= limit: + self.cluster_n = cluster + return + self.cluster_n = 16 + + @cute.jit + def __call__( + self, + mX: cute.Tensor, # (M, N) in + mXMax: cute.Tensor, # (M,) in + mXSumExp: cute.Tensor, # (M,) in + mXSumSoftmaxTimes: cute.Tensor, # (M,) out + mEntropy: cute.Tensor, # (M,) out + stream: cuda.CUstream, + ): + """ + Launch the Stage 2 kernel to compute final entropy values. + + Args: + mX: Input logits tensor of shape (M, N) + mXMax: Input tensor with max values from Stage 1 of shape (M,) + mXSumExp: Input tensor with sum of exponentials from Stage 1 of shape (M,) + mXSumSoftmaxTimes: Output tensor for sum(softmax * x) of shape (M,) + mEntropy: Output tensor for entropy values of shape (M,) + stream: CUDA stream for kernel execution + """ + assert mX.element_type == self.dtype + self._set_cluster_n() + largest_dtype_width = const_expr(mX.element_type.width) + vecsize = math.gcd(self.N, 128 // largest_dtype_width) # 8 for bf16/fp16 + tiled_copy, tiler_mn, threads_per_row = self._get_tiled_copy(vecsize=vecsize) + num_threads = tiled_copy.size + self.kernel( + mX, mXMax, mXSumExp, + mXSumSoftmaxTimes, mEntropy, + tiler_mn, tiled_copy, threads_per_row, + ).launch( + grid=[cute.ceil_div(mX.shape[0], tiler_mn[0]), self.cluster_n, 1], + block=[num_threads, 1, 1], + cluster=[1, self.cluster_n, 1] if const_expr(self.cluster_n > 1) else None, + stream=stream, + ) + + @cute.kernel + def kernel( + self, + mX: cute.Tensor, # (M, N) in + mXMax: cute.Tensor, # (M,) in + mXSumExp: cute.Tensor, # (M,) in + mXSumSoftmaxTimes: cute.Tensor, # (M,) out + mEntropy: cute.Tensor, # (M,) out + tiler_mn: cute.Shape, + tiled_copy: cute.TiledCopy, + threads_per_row: cutlass.Constexpr[int], + ): + """ + CUDA kernel for Stage 2: Compute final entropy values. + + This kernel computes the entropy using the formula: + H = max_x + log(sum_exp) - sum(softmax(x) * x) + """ + # Get thread and block indices + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + cluster_y = const_expr(0) if const_expr(self.cluster_n == 1) else cute.arch.block_idx()[1] + tv_layout = tiled_copy.layout_tv_tiled + + shape = mX.shape + idX = cute.make_identity_tensor(shape) + gX, cX = [cute.local_tile(mT, tiler_mn, (bidx, cluster_y)) for mT in (mX, idX)] + + # Allocate shared memory for input data and reduction buffers + smem = cutlass.utils.SmemAllocator() + sX = smem.allocate_tensor( + mX.element_type, cute.make_ordered_layout(tiler_mn, order=(1, 0)), byte_alignment=16 + ) + # Allocate reduction buffer and memory barrier for cluster synchronization + reduction_buffer, mbar_ptr = self._allocate_reduction_buffer_and_mbar(smem, tv_layout) + + thr_copy = tiled_copy.get_slice(tidx) + + # Partition tensors for this thread + tXgX = thr_copy.partition_S(gX) # Global memory partition + tXsX = thr_copy.partition_D(sX) # Shared memory partition + tXcX = thr_copy.partition_S(cX)[(0, None), None, None] # Coordinate partition + tXrX = cute.make_fragment_like(tXgX) # Register fragment + + # Handle non-divisible N by creating predicates + is_even_N = const_expr(shape[1] == tiler_mn[1] * self.cluster_n) + tXpX = None if is_even_N else copy_utils.predicate_k(thr_copy.partition_S(cX), limit=shape[1]) + copy = partial(copy_utils.copy, pred=tXpX) + + num_warps = cute.size(tiled_copy) // cute.arch.WARP_SIZE + self._initialize_cluster(tidx, mbar_ptr, num_warps) + + row = tXcX[0][0] + + # Load input data asynchronously from global to shared memory + if row < shape[0]: + copy(tXgX, tXsX, is_async=True) + cute.arch.cp_async_commit_group() + cute.arch.cp_async_wait_group(0) + if const_expr(not is_even_N): + # utils.fill_oob(tXsX, tXpX, -tXsX.element_type.inf) + utils.fill_oob(tXsX, tXpX, -BFloat16(2**15)) + # Copy from shared memory to registers and convert to Float32 + cute.autovec_copy(tXsX, tXrX) + x = tXrX.load().to(Float32) + + # Load intermediate results from Stage 1 + max_x = mXMax[row] + denom = mXSumExp[row] + + # Compute softmax: softmax(x) = exp(x - max_x) / sum_exp + # rcp_approx: 1.0 / denom (reciprocal approximation for faster division) + log2_e = math.log2(math.e) + exp_x = cute.math.exp2(x * log2_e - (max_x * log2_e), fastmath=False) + softmax_x = exp_x * cute.arch.rcp_approx(denom) + + # Compute sum(softmax * x) via reduction + sum_softmax_times_x = row_reduce( + softmax_x * x, + cute.ReductionOp.ADD, + threads_per_row, + reduction_buffer[None, None, 0], + mbar_ptr if const_expr(self.cluster_n > 1) else None, + hook_fn=cute.arch.cluster_wait if const_expr(self.cluster_n > 1) else None, + init_val=0.0, + ) + + # Write results to global memory (only first thread in row, and first block in cluster) + if ( + tXcX[0][1] == 0 + and row < shape[0] + and (self.cluster_n == 1 or cute.arch.block_idx_in_cluster() == 0) + ): + mXSumSoftmaxTimes[row] = sum_softmax_times_x + # Entropy formula: H = max_x + log(sum_exp) - sum(softmax * x) + mEntropy[row] = max_x + cute.math.log(denom, fastmath=True) - sum_softmax_times_x + + +@torch.library.custom_op("roll_kernel::entropy_fwd_out", mutates_args={"entropy", "x_max", "x_sum_exp", "x_sum_softmax_times"}) +def entropy_fwd_out( + x: Tensor, + entropy: Tensor, + x_max: Tensor, + x_sum_exp: Tensor, + x_sum_softmax_times: Tensor, +) -> None: + """ + Fused entropy forward pass with in-place output tensors. + + This function computes the entropy of each row in the input tensor using + a two-stage fused kernel approach: + - Stage 1: Computes max_x and sum_exp for each row + - Stage 2: Computes the final entropy: H = max_x + log(sum_exp) - sum(softmax * x) + + The function uses a compilation cache to avoid recompiling kernels for + the same configuration (dtype, target_dtype, N). + + Args: + x: Input logits tensor of shape (M, N), where M is batch size and N is vocabulary size + entropy: Output tensor for entropy values of shape (M,) [mutated in-place] + x_max: Output tensor for max values of shape (M,) [mutated in-place] + x_sum_exp: Output tensor for sum of exponentials of shape (M,) [mutated in-place] + x_sum_softmax_times: Output tensor for sum(softmax * x) of shape (M,) [mutated in-place] + + Returns: + None (output tensors are mutated in-place) + + Raises: + AssertionError: If input validation fails (wrong dimensions, device, or dtype) + """ + assert x.dim() == 2, "Input must be 2D" + assert x.is_cuda, "Tensors must be on CUDA device" + assert x.dtype in [torch.float16, torch.bfloat16, torch.float32], "Unsupported input dtype" + assert entropy.dtype in [torch.float16, torch.bfloat16, torch.float32], "Unsupported input dtype" + + N = x.size(1) + + dtype = torch2cute_dtype_map[x.dtype] + target_dtype = torch2cute_dtype_map[entropy.dtype] + # Create compile key for caching compiled kernels + compile_key = (dtype, target_dtype, N) + # Compile kernels if not in cache + if compile_key not in entropy_fwd_out.stage1_compile_cache: + # Create symbolic batch size for compilation + batch_sym = cute.sym_int() + div = math.gcd(128 // dtype.width, N) + # Create fake tensors for compilation + x_cute = fake_tensor(dtype, (batch_sym, N), div) + entropy_cute = fake_tensor(target_dtype, (batch_sym,)) + x_max_cute = fake_tensor(target_dtype, (batch_sym,)) + x_sum_exp_cute = fake_tensor(target_dtype, (batch_sym,)) + x_sum_softmax_times_cute = fake_tensor(target_dtype, (batch_sym,)) + + # Initialize Stage 1 and Stage 2 operators with online softmax enabled + # Online softmax computes both max and sum in a single pass for better performance + entropy_stage1_op = FusedEntropyForwardStage1(dtype, N, online_softmax=True) + entropy_stage2_op = FusedEntropyForwardStage2(dtype, N) + + # Compile Stage 1 kernel (computes max_x and sum_exp) + entropy_fwd_out.stage1_compile_cache[compile_key] = cute.compile( + entropy_stage1_op, + x_cute, + x_max_cute, + x_sum_exp_cute, + cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), + options="--enable-tvm-ffi", + ) + + # Compile Stage 2 kernel (computes entropy) + entropy_fwd_out.stage2_compile_cache[compile_key] = cute.compile( + entropy_stage2_op, + x_cute, + x_max_cute, + x_sum_exp_cute, + x_sum_softmax_times_cute, + entropy_cute, + cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), + options="--enable-tvm-ffi", + ) + + # Execute Stage 1: compute max_x and sum_exp + entropy_fwd_out.stage1_compile_cache[compile_key]( + x, x_max, x_sum_exp + ) + # Execute Stage 2: compute entropy using results from Stage 1 + entropy_fwd_out.stage2_compile_cache[compile_key]( + x, x_max, x_sum_exp, x_sum_softmax_times, entropy + ) + + +# Compilation cache for Stage 1 and Stage 2 kernels +entropy_fwd_out.stage1_compile_cache = {} +entropy_fwd_out.stage2_compile_cache = {} + + +def entropy_fwd( + x: torch.Tensor, +) -> torch.Tensor | tuple[torch.Tensor]: + """ + Compute entropy of each row in the input tensor. + + This function allocates output tensors and calls the fused entropy forward pass. + It returns the entropy value along with intermediate results that are needed + for the backward pass. + + Args: + x: Input logits tensor of shape (M, N), where M is batch size and N is vocabulary size + + Returns: + tuple: A tuple containing: + - entropy: Tensor of shape (M,) with entropy values + - x_max: Tensor of shape (M,) with max values (for backward pass) + - x_sum_exp: Tensor of shape (M,) with sum of exponentials (for backward pass) + - x_sum_softmax_times: Tensor of shape (M,) with sum(softmax * x) (for backward pass) + """ + M = x.size(0) + device = x.device + calc_type = torch.float32 + + # Allocate output tensors + entropy = torch.zeros(M, dtype=calc_type, device=device) + x_max = torch.zeros(M, dtype=calc_type, device=device) + x_sum_exp = torch.zeros(M, dtype=calc_type, device=device) + x_sum_softmax_times = torch.zeros(M, dtype=calc_type, device=device) + + # Compute entropy using fused kernel + entropy_fwd_out(x, entropy, x_max, x_sum_exp, x_sum_softmax_times) + return entropy, x_max, x_sum_exp, x_sum_softmax_times + + +def test_entropy_fwd(M, N): + entropy_fwd(x) + + +class FusedEntropyBackward: + """ + Backward pass for entropy computation. + + This kernel computes the gradient of entropy with respect to the input logits. + The gradient formula is derived from the entropy loss: + dH/dx = -d_entropy * softmax(x) * (x - E[x]) + + Where: + - d_entropy: Gradient of entropy loss (scalar per row) + - softmax(x): Softmax probabilities + - E[x]: Expected value under softmax distribution (sum_softmax_times from forward) + + Args: + dtype: Data type of input tensor (e.g., BFloat16, Float16, Float32) + N: Number of columns in input tensor (vocabulary size) + """ + + def __init__(self, dtype: Type[cutlass.Numeric], N: int): + """ + Initialize the entropy backward pass. + """ + self.dtype = dtype + self.N = N + self.vecsize = 128 // dtype.width + + def _threads_per_row(self): + """ + Determine the optimal number of threads per row for reduction. + We split by blocks of 16k for large vocabularies. + """ + N = min(self.N, 16384) # We split by blocks of 16k + for limit, threads in [(64, 8), (128, 16), (3072, 32), (6144, 64), (16384, 128)]: + if N <= limit: + return threads + return 256 + + def _get_tiled_copy(self, vecsize: int): + """ + Get tiled copy configuration for efficient memory access. + + Args: + vecsize: Vector size for memory loading + + Returns: + tuple: (tiled_copy, tiler_mn, threads_per_row) + """ + assert self.N % vecsize == 0, f"Input N {self.N} is not divisible by vector size {vecsize}" + N = min(self.N, 16384) + num_threads = 128 if N <= 16384 else 256 + threads_per_row = self._threads_per_row() + rows_per_block = num_threads // threads_per_row + num_blocks_N = cute.ceil_div(N // vecsize, threads_per_row) + tiler_mn = (rows_per_block, vecsize * num_blocks_N * threads_per_row) + tiled_copy = copy_utils.tiled_copy_2d( + self.dtype, threads_per_row, num_threads, num_copy_elems=vecsize + ) + return tiled_copy, tiler_mn, threads_per_row + + @cute.jit + def __call__( + self, + mX: cute.Tensor, #[M, N] + mDEntropy: cute.Tensor, #[M, N] + mXMax: cute.Tensor, #[M] + mXSumExp: cute.Tensor, #[M] + mXSumSoftmaxTimes: cute.Tensor, #[M] + stream: cuda.CUstream, + ): + """ + Launch the backward kernel to compute gradients. + + Args: + mX: Input logits tensor of shape (M, N) + mDEntropy: Gradient of entropy loss of shape (M,) + mXMax: Max values from forward pass of shape (M,) + mXSumExp: Sum of exponentials from forward pass of shape (M,) + mXSumSoftmaxTimes: Sum(softmax * x) from forward pass of shape (M,) + stream: CUDA stream for kernel execution + """ + assert mX.element_type == self.dtype + # Calculate vector size for memory loading (max 128 bits per load) + # e.g. if self.N isn't divisible by 8 for bf16, we might use 64 bits (4 elements) copy + vecsize = math.gcd(self.N, 128 // self.dtype.width) + tiled_copy, tiler_mn, threads_per_row = self._get_tiled_copy(vecsize=vecsize) + num_threads = tiled_copy.size + mDEntropy, mXMax, mXSumExp, mXSumSoftmaxTimes= [ + layout_utils.expand(X, dim=1, size=self.N) for X in (mDEntropy, mXMax, mXSumExp, mXSumSoftmaxTimes) + ] + # Launch kernel with appropriate grid and block configuration + self.kernel( + mX, + mDEntropy, + mXMax, + mXSumExp, + mXSumSoftmaxTimes, + mX.shape, + tiler_mn, + tiled_copy, + threads_per_row, + ).launch( + grid=[ + cute.ceil_div(mX.shape[0], tiler_mn[0]), + cute.ceil_div(mX.shape[1], tiler_mn[1]), + 1, + ], + block=[num_threads, 1, 1], + stream=stream, + ) + + @cute.kernel + def kernel( + self, + mX: cute.Tensor, #[M, N] + mDEntropy: cute.Tensor, #[M, N] + mXMax: cute.Tensor, #[M] + mXSumExp: cute.Tensor, #[M] + mXSumSoftmaxTimes: cute.Tensor, #[M] + shape: cute.Shape, + tiler_mn: cute.Shape, + tiled_copy: cute.TiledCopy, + threads_per_row: cutlass.Constexpr[int], + ): + """ + CUDA kernel for entropy backward pass. + + This kernel computes the gradient of entropy with respect to the input logits: + dx = -d_entropy * softmax(x) * (x - E[x]) + + Where E[x] is the expected value under the softmax distribution. + """ + tidx, _, _ = cute.arch.thread_idx() + bidx, bidy, _ = cute.arch.block_idx() + + # Allocate shared memory for input data + smem = cutlass.utils.SmemAllocator() + sX = smem.allocate_tensor( + mX.element_type, cute.make_ordered_layout(tiler_mn, order=(1, 0)), byte_alignment=16 + ) + + idX = cute.make_identity_tensor(shape) + gX, cX = [cute.local_tile(mT, tiler_mn, (bidx, bidy)) for mT in (mX, idX)] + + thr_copy = tiled_copy.get_slice(tidx) + + # Partition tensors for this thread + tXgX = thr_copy.partition_S(gX) # Global memory partition + tXsX = thr_copy.partition_D(sX) # Shared memory partition + tXcX = thr_copy.partition_S(cX)[(0, None), None, None] # Coordinate partition + tXrX = cute.make_fragment_like(tXgX) # Register fragment + + # Handle non-divisible N by creating predicates + is_even_N = const_expr(shape[1] % tiler_mn[1] == 0) + tXpX = ( + None if is_even_N else copy_utils.predicate_k(thr_copy.partition_S(cX), limit=shape[1]) + ) + copy = partial(copy_utils.copy, pred=tXpX) + + row = tXcX[0][0] + + # Load input data asynchronously from global to shared memory + if row < shape[0]: + copy(tXgX, tXsX, is_async=True) + cute.arch.cp_async_commit_group() + cute.arch.cp_async_wait_group(0) + if const_expr(not is_even_N): + utils.fill_oob(tXsX, tXpX, -tXsX.element_type.inf) + # Copy from shared memory to registers and convert to Float32 + cute.autovec_copy(tXsX, tXrX) + x = tXrX.load().to(Float32) + + # Load intermediate results from forward pass + d_entropy = Float32.zero + x_max = Float32.zero + x_sum_exp = Float32.zero + x_sum_softmax_times = Float32.zero + if row < shape[0]: + d_entropy = mDEntropy[row] + x_max = mXMax[row] + x_sum_exp = mXSumExp[row] + x_sum_softmax_times = mXSumSoftmaxTimes[row] + + # Compute softmax for gradient calculation + log2_e = math.log2(math.e) + exp_x = cute.math.exp2(x * log2_e - (x_max * log2_e), fastmath=False) + softmax_x = exp_x * cute.arch.rcp_approx(x_sum_exp) + + # Compute gradient: dx = -d_entropy * softmax(x) * (x - E[x]) + x = -(x - x_sum_softmax_times) * softmax_x * d_entropy + + # Store gradient to registers and write back to global memory + tXrX.store(x.to(tXrX.element_type)) + if row < shape[0]: + copy(tXrX, tXgX) + + +def _entropy_backward( + x: torch.Tensor, + d_entropy: torch.Tensor, + x_max: torch.Tensor, + x_sum_exp: torch.Tensor, + x_sum_softmax_times: torch.Tensor, +) -> None: + """ + Entropy backward pass using fused kernel. + + This function computes the gradient of entropy with respect to the input logits. + The gradient formula is: + dx = -d_entropy * softmax(x) * (x - E[x]) + + Where E[x] is the expected value under the softmax distribution. + + Args: + x: Input logits tensor of shape (M, N), where M is batch size and N is vocabulary size + d_entropy: Gradient of entropy loss of shape (M,) + x_max: Max values from forward pass of shape (M,) + x_sum_exp: Sum of exponentials from forward pass of shape (M,) + x_sum_softmax_times: Sum(softmax * x) from forward pass of shape (M,) + + Returns: + None (gradient is computed in-place on the input tensor x) + + Raises: + AssertionError: If input validation fails (wrong dimensions, device, or dtype) + """ + d_entropy = d_entropy.contiguous() + # Input validation + assert x.dim() == 2, "Input must be 2D" + assert d_entropy.dim() == 1, "d_entropy must be 1D" + assert x_max.dim() == 1, "x_max must be 1D" + assert x_sum_exp.dim() == 1, "x_sum_exp must be 1D" + assert x_sum_softmax_times.dim() == 1, "x_sum_softmax_timesmust be 1D" + assert x.shape[0] == d_entropy .shape[0], "Batch dimensions must match" + assert x.shape[0] == x_max.shape[0], "Batch dimensions must match" + assert x.shape[0] == x_sum_exp.shape[0], "Batch dimensions must match" + assert x.shape[0] == x_sum_softmax_times.shape[0], "Batch dimensions must match" + assert x.is_cuda and d_entropy.is_cuda and x_max.is_cuda and x_sum_exp.is_cuda and x_sum_softmax_times.is_cuda, ( + "Tensors must be on CUDA device" + ) + assert x.dtype in [torch.float16, torch.bfloat16, torch.float32], "Unsupported input dtype" + assert d_entropy.dtype == torch.float32, "d_entropy must be float32" + assert x_max.dtype == torch.float32, "max_x must be float32" + assert x_sum_exp.dtype == torch.float32, "x_sum_exp must be float32" + assert x_sum_softmax_times.dtype == torch.float32, "x_sum_softmax_times must be float32" + + N = x.size(1) + dtype = torch2cute_dtype_map[x.dtype] + calc_dtype = Float32 + # Create compile key for caching compiled kernels + compile_key = (dtype, calc_dtype, N) + if compile_key not in _entropy_backward.compile_cache: + batch_sym = cute.sym_int() + div = math.gcd(128 // dtype.width, N) + x_cute= fake_tensor(dtype, (batch_sym, N), div) + d_entropy_cute, x_max_cute, x_sum_exp_cute, x_sum_softmax_times_cute = [fake_tensor(calc_dtype, (batch_sym,))] * 4 + fused_entropy_backward_op = FusedEntropyBackward(dtype, N) + # Compile backward kernel + _entropy_backward.compile_cache[compile_key] = cute.compile( + fused_entropy_backward_op, + x_cute, d_entropy_cute, + x_max_cute, x_sum_exp_cute, x_sum_softmax_times_cute, + cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), + options="--enable-tvm-ffi", + ) + # Execute backward kernel + _entropy_backward.compile_cache[compile_key]( + x, d_entropy, x_max, x_sum_exp, x_sum_softmax_times, + ) + + +_entropy_backward.compile_cache = {} + + +# @torch.library.custom_op("quack::entropy_bwd_out", mutates_args={"x", "d_entropy", "x_max", "x_sum_exp", "x_sum_softmax_times"}) +# def entropy_bwd_out( + # x: torch.Tensor, + # d_entropy: torch.Tensor, + # x_max: torch.Tensor, + # x_sum_exp: torch.Tensor, + # x_sum_softmax_times: torch.Tensor, +# ) -> None: + # """ + # Fused entropy backward pass with in-place gradient computation. + + # This function computes the gradient of entropy with respect to the input logits + # and stores it in the input tensor x. + + # Args: + # x: Input logits tensor of shape (M, N) [mutated in-place with gradient] + # d_entropy: Gradient of entropy loss of shape (M,) + # x_max: Max values from forward pass of shape (M,) + # x_sum_exp: Sum of exponentials from forward pass of shape (M,) + # x_sum_softmax_times: Sum(softmax * x) from forward pass of shape (M,) + + # Returns: + # None (gradient is computed in-place on x) + # """ + # _entropy_backward(x, d_entropy, x_max, x_sum_exp, x_sum_softmax_times) + + +def entropy_bwd( + x: torch.Tensor, + d_entropy: torch.Tensor, + x_max: torch.Tensor, + x_sum_exp: torch.Tensor, + x_sum_softmax_times: torch.Tensor, +) : + """ + Compute gradient of entropy with respect to input logits. + + This function computes the gradient and returns the modified input tensor. + + Args: + x: Input logits tensor of shape (M, N) + d_entropy: Gradient of entropy loss of shape (M,) + x_max: Max values from forward pass of shape (M,) + x_sum_exp: Sum of exponentials from forward pass of shape (M,) + x_sum_softmax_times: Sum(softmax * x) from forward pass of shape (M,) + + Returns: + torch.Tensor: The input tensor x with gradient computed in-place + """ + _entropy_backward(x, d_entropy, x_max, x_sum_exp, x_sum_softmax_times) + return x + + +def test_entropy_bwd(M, N): + """ + Test function for entropy backward pass. + """ + x = torch.rand((M, N), dtype=torch.bfloat16).cuda() + d_entropy = torch.rand((M), dtype=torch.float32).cuda() + x_max = torch.rand((M), dtype=torch.float32).cuda() + x_sum_exp = torch.rand((M), dtype=torch.float32).cuda() + x_sum_softmax_times = torch.rand((M), dtype=torch.float32).cuda() + entropy_bwd(x, d_entropy, x_max, x_sum_exp, x_sum_softmax_times) + + +class EntropyFunction(torch.autograd.Function): + """ + PyTorch autograd function for entropy computation with fused kernels. + """ + + @staticmethod + def forward(ctx, x): + entropy, x_max, x_sum_exp, x_sum_softmax_times = entropy_fwd( + x, target, ignore_index=ignore_index, return_lse=True + ) + + ctx.save_for_backward(x, x_max, x_sum_exp, x_sum_softmax_times) + return entropy + + @staticmethod + def backward(ctx, d_entropy): + x, x_max, x_sum_exp, x_sum_softmax_times = ctx.saved_tensors + dx = entropy_bwd( + x, d_entropy, x_max, x_sum_exp, x_sum_softmax_times + ) + return dx + + +if __name__ == '__main__': + test_entropy_fwd(8192, 128000) + test_entropy_bwd(8192, 128000) diff --git a/roll/third_party/megatron/model_update.py b/roll/third_party/megatron/model_update.py index 8903b8c63..ab2bb7e43 100644 --- a/roll/third_party/megatron/model_update.py +++ b/roll/third_party/megatron/model_update.py @@ -51,28 +51,21 @@ def gather_and_convert_weights( gathered_named_weights.append((mcore_name, gathered_weights)) handles.append(dist.all_gather(gathered_weights, weight, group=ep_group, async_op=True)) - def extract_suffix_number(s): - import re - - match = re.search(r"\d+$", s) - return match.group() if match else None - hf_named_weights = [] for handle, (mcore_name, weights) in zip(handles, gathered_named_weights): if handle is not None: handle.wait() - local_moe_index = extract_suffix_number(mcore_name) for ep_rank, weight in enumerate(weights): - global_moe_index = model_converter.dist_converter.num_layers_for_expert * ep_rank + int( - local_moe_index - ) - name = mcore_name[: -len(local_moe_index)] + str(global_moe_index) + # Temporarily set EP rank so the converter computes the correct global expert index + saved_ep_rank = model_converter.dist_converter.expert_model_parallel_rank + model_converter.dist_converter.expert_model_parallel_rank = ep_rank converted_weights = ( model_converter.convert_to_hf( - {name: [weight]}, layer_index_preprocessed=True, moe_index_preprocessed=True, **kwargs + {mcore_name: [weight]}, layer_index_preprocessed=True, **kwargs ) or {} ) + model_converter.dist_converter.expert_model_parallel_rank = saved_ep_rank hf_named_weights.extend([(name, weight) for name, weight in converted_weights.items()]) return hf_named_weights @@ -438,13 +431,16 @@ def _colocated_model_update(self): for hf_named_weights in gather_all_hf_weights( self.models_unwrapped, buffer_size=self._model_update_buffer_size, weights_meta=self._weights_meta ): + if not hf_named_weights: + continue if self._co_infer_worker is not None: serialized_tensors = serialize_named_weights( hf_named_weights, infer_strategy=self.infer_worker_config.strategy_args.strategy_name ) infer_parallel_tensors = [None] * infer_parallel_size if co_infer_rank == 0 else None + global_dst_rank = dist.get_global_rank(self._infer_parallel_cpu_group, 0) dist.gather_object( - serialized_tensors, infer_parallel_tensors, group_dst=0, group=self._infer_parallel_cpu_group + serialized_tensors, infer_parallel_tensors, global_dst_rank, group=self._infer_parallel_cpu_group ) if refs: @@ -467,8 +463,6 @@ def _colocated_model_update(self): return {} def _separated_model_update(self): - if not mpu.get_expert_data_parallel_rank() == 0: - return {} logger.info(f"start broadcast model update {self.model_update_name}") for hf_named_weights in gather_pp_stage_hf_weights( @@ -476,6 +470,8 @@ def _separated_model_update(self): ): if not self._broadcast_workers: continue + if mpu.get_expert_data_parallel_rank() != 0: + continue while not ray.get(self._model_update_locker.acquire.remote()): time.sleep(0.1) refs = self._broadcast_to_infer_workers(hf_named_weights) diff --git a/roll/third_party/megatron/mtp_patcher.py b/roll/third_party/megatron/mtp_patcher.py new file mode 100644 index 000000000..0944cfb42 --- /dev/null +++ b/roll/third_party/megatron/mtp_patcher.py @@ -0,0 +1,293 @@ +""" +Monkey patches for Megatron Multi-Token Prediction (MTP) training. + +This module provides patches for MTP-related functions in Megatron to add +extra operations during MTP training. + +MTP training mode is read from `self.config.mtp_training_mode`: +- 'disabled': MTP is loaded but not trained (default) +- 'standalone': MTP is trained independently with truncated gradients +- 'joint': MTP participates in main model updates with full gradient flow +""" + +from typing import TYPE_CHECKING, Callable, Optional + +import torch + +if TYPE_CHECKING: + from megatron.core.models.gpt.gpt_model import GPTModel + from megatron.core.transformer.multi_token_prediction import MultiTokenPredictionLayer + +def patch_mtp_functions(): + """ + Apply monkey patches for MTP training. + + This function patches the following methods: + 1. GPTModel.forward - for forward pass modifications + 2. GPTModel._postprocess - for postprocess modifications + 3. MultiTokenPredictionLayer._get_embeddings - for embedding modifications + """ + from collections import OrderedDict + + from megatron.core import parallel_state + from megatron.core.config_logger import has_config_logger_enabled, log_config_to_disk + from megatron.core.models.gpt.gpt_model import GPTModel + from megatron.core.tensor_parallel import gather_from_sequence_parallel_region + from megatron.core.transformer.multi_token_prediction import ( + MTPLossAutoScaler, + MTPLossLoggingHelper, + MultiTokenPredictionLayer, + roll_tensor, + ) + + # Save original methods for potential use in patched versions + original_gpt_postprocess = GPTModel._postprocess + original_mtp_get_embeddings = MultiTokenPredictionLayer._get_embeddings + + # ============================================================================ + # GPTModel._postprocess + # ============================================================================ + def patched_gpt_postprocess( + self: "GPTModel", + hidden_states, + input_ids, + position_ids, + labels, + rotary_pos_emb, + rotary_pos_cos, + rotary_pos_sin, + mtp_in_postprocess=None, + loss_mask=None, + decoder_input=None, + attention_mask=None, + inference_params=None, + packed_seq_params=None, + sequence_len_offset=None, + runtime_gather_output=None, + extra_block_kwargs=None, + inference_context=None, + **kwargs, + ): + """ + Patched _postprocess method for GPTModel with MTP support. + Reads mtp_training_mode from self.config. + """ + # signature of mtp related methods differs between megatron_core 0.16/0.17 and dev. + # currently, dev is needed for GDN cp and packing which includes padding_mask, + # and use kwargs to be compatible + if "padding_mask" in kwargs: + if extra_block_kwargs: + extra_block_kwargs["padding_mask"] = kwargs["padding_mask"] + else: + extra_block_kwargs = {"padding_mask": kwargs["padding_mask"]} + mtp_training_mode = getattr(self.config, "mtp_training_mode", "disabled") + + in_inference_mode = inference_context is not None and not self.training + if in_inference_mode: + assert runtime_gather_output, "Inference must always gather TP logits" + + # logits and loss + output_weight = None + if self.share_embeddings_and_output_weights: + output_weight = self.shared_embedding_or_output_weight() + if mtp_in_postprocess and mtp_training_mode != "disabled": + hidden_states = self.mtp( + input_ids=input_ids, + position_ids=position_ids, + hidden_states=hidden_states, + attention_mask=attention_mask, + inference_params=inference_params, + rotary_pos_emb=rotary_pos_emb, + rotary_pos_cos=rotary_pos_cos, + rotary_pos_sin=rotary_pos_sin, + packed_seq_params=packed_seq_params, + sequence_len_offset=sequence_len_offset, + embedding=self.embedding, + **(extra_block_kwargs or {}), + ) + + if not self.post_process: + return hidden_states + + if self.config.mtp_num_layers is not None and mtp_training_mode != "disabled": + mtp_labels, _ = roll_tensor( + input_ids, + shifts=-1, + dims=-1, + cp_group=self.cp_group, + packed_seq_params=packed_seq_params, + ) + + hidden_states_list = torch.chunk(hidden_states, 1 + self.config.mtp_num_layers, dim=0) + hidden_states = hidden_states_list[0] + if loss_mask is None: + # if loss_mask is not provided, use all ones as loss_mask + loss_mask = torch.ones_like(mtp_labels) + + for mtp_layer_number in range(self.config.mtp_num_layers): + # output + mtp_logits, _ = self.output_layer( + hidden_states_list[mtp_layer_number + 1], + weight=output_weight.detach() if mtp_training_mode == "standalone" and output_weight is not None else output_weight, + runtime_gather_output=runtime_gather_output, + ) + # Calc loss for the current Multi-Token Prediction (MTP) layers. + mtp_labels, _ = roll_tensor( + mtp_labels, + shifts=-1, + dims=-1, + cp_group=self.cp_group, + packed_seq_params=packed_seq_params, + ) + loss_mask, num_tokens = roll_tensor( + loss_mask, + shifts=-1, + dims=-1, + cp_group=self.cp_group, + packed_seq_params=packed_seq_params, + ) + + mtp_loss = self.compute_language_model_loss(mtp_labels, mtp_logits) + mtp_loss = loss_mask * mtp_loss + + num_tokens = max(num_tokens, 1) + + if self.training: + MTPLossLoggingHelper.save_loss_to_tracker( + torch.sum(mtp_loss) / num_tokens, + mtp_layer_number, + self.config.mtp_num_layers, + avg_group=parallel_state.get_data_parallel_group( + with_context_parallel=True + ), + ) + mtp_loss_scale = self.config.mtp_loss_scaling_factor / self.config.mtp_num_layers + if self.config.calculate_per_token_loss: + hidden_states = MTPLossAutoScaler.apply( + hidden_states, mtp_loss_scale * mtp_loss + ) + else: + hidden_states = MTPLossAutoScaler.apply( + hidden_states, mtp_loss_scale * mtp_loss / num_tokens + ) + sequence_parallel_override = False + + if in_inference_mode and inference_context.materialize_only_last_token_logits: + if inference_context.is_static_batching(): + hidden_states = hidden_states[-1:, :, :] + else: + if self.output_layer.sequence_parallel: + hidden_states = gather_from_sequence_parallel_region( + hidden_states, group=self.pg_collection.tp + ) + self.output_layer.sequence_parallel = False + sequence_parallel_override = True + + hidden_states = inference_context.last_token_logits( + hidden_states.squeeze(1).unsqueeze(0) + ).unsqueeze(1) + + logits, _ = self.output_layer( + hidden_states, weight=output_weight, runtime_gather_output=runtime_gather_output + ) + + # Restore sequence parallel execution to the output layer if necessary. + if sequence_parallel_override: + assert ( + in_inference_mode + and inference_context.is_dynamic_batching() + and inference_context.materialize_only_last_token_logits + ) + self.output_layer.sequence_parallel = True + + if has_config_logger_enabled(self.config): + payload = OrderedDict( + { + 'input_ids': input_ids, + 'position_ids': position_ids, + 'attention_mask': attention_mask, + 'decoder_input': decoder_input, + 'logits': logits, + } + ) + log_config_to_disk(self.config, payload, prefix='input_and_logits') + + if labels is None: + # [s b h] => [b s h] + return logits.transpose(0, 1).contiguous() + + loss = self.compute_language_model_loss(labels, logits) + return loss + + # ============================================================================ + # MultiTokenPredictionLayer._get_embeddings + # ============================================================================ + def patched_mtp_get_embeddings( + self: "MultiTokenPredictionLayer", + input_ids: "Tensor", + position_ids: "Tensor", + embedding: Callable, + hidden_states: "Tensor", + packed_seq_params=None, + **kwargs, + ): + """ + Patched _get_embeddings method for MultiTokenPredictionLayer. + Reads mtp_training_mode from self.config. + """ + from megatron.core.utils import make_viewless_tensor + + # signature of mtp related methods differs between megatron_core 0.16/0.17 and dev. + # currently, dev is needed for GDN cp and packing which includes padding_mask, + # and use kwargs to be compatible + padding_mask = kwargs.get("padding_mask", None) + mtp_training_mode = getattr(self.config, "mtp_training_mode", "disabled") + + # Calc logits for the current Multi-Token Prediction (MTP) layers. + input_ids, _ = roll_tensor( + input_ids, + shifts=-1, + dims=-1, + cp_group=self.cp_group, + packed_seq_params=packed_seq_params, + ) + position_ids, _ = roll_tensor( + position_ids, + shifts=-1, + dims=-1, + cp_group=self.cp_group, + packed_seq_params=packed_seq_params, + ) + if padding_mask is not None: + padding_mask, _ = roll_tensor( + padding_mask, + shifts=-1, + dims=-1, + cp_group=self.cp_group, + packed_seq_params=packed_seq_params, + ) + # embedding + decoder_input = embedding(input_ids=input_ids, position_ids=position_ids) + + # + if mtp_training_mode == "standalone": + decoder_input = decoder_input.detach() + hidden_states = make_viewless_tensor(inp=hidden_states, requires_grad=True, keep_graph=False) + else: + hidden_states = make_viewless_tensor(inp=hidden_states, requires_grad=True, keep_graph=True) + + if "padding_mask" in kwargs: + result = input_ids, position_ids, padding_mask, decoder_input, hidden_states + else: + result = input_ids, position_ids, decoder_input, hidden_states + return result + + # ============================================================================ + # Apply patches + # ============================================================================ + GPTModel._postprocess = patched_gpt_postprocess + MultiTokenPredictionLayer._get_embeddings = patched_mtp_get_embeddings + + # Store original methods for potential restoration + GPTModel._original_postprocess = original_gpt_postprocess + MultiTokenPredictionLayer._original_get_embeddings = original_mtp_get_embeddings diff --git a/roll/third_party/megatron/offload_states_patch.py b/roll/third_party/megatron/offload_states_patch.py index ffef202c5..9009fae09 100644 --- a/roll/third_party/megatron/offload_states_patch.py +++ b/roll/third_party/megatron/offload_states_patch.py @@ -28,7 +28,7 @@ from torch import Tensor from roll.platforms import current_platform -from roll.utils.offload_states import move_tensors_to_device_buffer, move_device_buffer_to_tensors +from roll.utils.offload_states import move_tensors_to_device_buffer, move_device_buffer_to_tensors, clear_memory def bind_megatron_offload_states_func(optimizer: MegatronOptimizer): @@ -124,9 +124,7 @@ def float16_optimizer_with_float16_params_offload_states(self: Float16OptimizerW offload_adam_states(self.optimizer, device, pin_memory=pin_memory, non_blocking=non_blocking) self.offloaded_states.add(MegatronOffloadStateType.optimizer_states) - current_platform.synchronize() - gc.collect() - current_platform.empty_cache() + clear_memory() def float16_optimizer_with_float16_params_reload_states(self: Float16OptimizerWithFloat16Params, @@ -142,12 +140,16 @@ def float16_optimizer_with_float16_params_reload_states(self: Float16OptimizerWi move_device_buffer_to_tensors(tensors=float16_weights, device_buffer=getattr(self, "float16_groups_cpu_buffer").to(device, non_blocking=non_blocking)) + self.float16_groups_cpu_buffer = None + fp32_weights: List[Tensor] = [param for sub_group in self.fp32_from_fp32_groups for param in sub_group] if getattr(self, "float32_groups_cpu_buffer") is not None: move_device_buffer_to_tensors(tensors=fp32_weights, device_buffer=getattr(self, "float32_groups_cpu_buffer").to(device, non_blocking=non_blocking)) + self.float32_groups_cpu_buffer = None + self.offloaded_states.remove(MegatronOffloadStateType.model_params) if needs_reload(MegatronOffloadStateType.other_params, include, self.offloaded_states): @@ -161,6 +163,7 @@ def float16_optimizer_with_float16_params_reload_states(self: Float16OptimizerWi move_device_buffer_to_tensors(tensors=fp32_from_float16_weights, device_buffer=getattr(self, "fp32_from_float16_groups_cpu_buffer").to(device, non_blocking=non_blocking)) + self.fp32_from_float16_groups_cpu_buffer = None self.offloaded_states.remove(MegatronOffloadStateType.other_params) @@ -325,9 +328,7 @@ def distributed_optimizer_offload_states(self: DistributedOptimizer, offload_adam_states(self.optimizer, device, pin_memory=pin_memory, non_blocking=non_blocking) self.offloaded_states.add(MegatronOffloadStateType.optimizer_states) - current_platform.synchronize() - gc.collect() - current_platform.empty_cache() + clear_memory() def distributed_optimizer_reload_states(self: DistributedOptimizer, @@ -354,6 +355,7 @@ def distributed_optimizer_reload_states(self: DistributedOptimizer, move_device_buffer_to_tensors(tensors=shard_fp32_from_float16_weights, device_buffer=getattr(self, "shard_fp32_from_float16_groups_cpu_buffer").to(device, non_blocking=non_blocking), ) + self.shard_fp32_from_float16_groups_cpu_buffer = None self.offloaded_states.remove(MegatronOffloadStateType.other_params) @@ -396,9 +398,7 @@ def fp32_optimizer_offload_states(self: FP32Optimizer, # offload optimizer states offload_adam_states(self.optimizer, device, pin_memory=pin_memory, non_blocking=non_blocking) self.offloaded_states.add(MegatronOffloadStateType.optimizer_states) - current_platform.synchronize() - gc.collect() - current_platform.empty_cache() + clear_memory() def fp32_optimizer_reload_states(self: FP32Optimizer, @@ -415,6 +415,7 @@ def fp32_optimizer_reload_states(self: FP32Optimizer, move_device_buffer_to_tensors(tensors=float32_weights, device_buffer=getattr(self, "optimizer_param_groups_cpu_buffer").to(device, non_blocking=non_blocking), ) + self.optimizer_param_groups_cpu_buffer = None self.offloaded_states.remove(MegatronOffloadStateType.model_params) @@ -505,10 +506,11 @@ def reload_megatron_no_grad_module(model_chunks: List[Union[DistributedDataParal offloaded_states=model_chunk.offloaded_states): param_dtype_to_params = getattr(model_chunk, "param_dtype_to_params", {}) for param_dtype, params in param_dtype_to_params.items(): - if getattr(model_chunk, f"{param_dtype}_ddp_no_grad_groups_cpu_buffer") is not None: + buffer_attr = f"{param_dtype}_ddp_no_grad_groups_cpu_buffer" + if getattr(model_chunk, buffer_attr, None) is not None: move_device_buffer_to_tensors(tensors=params, - device_buffer=getattr(model_chunk, - f"{param_dtype}_ddp_no_grad_groups_cpu_buffer").to(device, non_blocking=non_blocking)) + device_buffer=getattr(model_chunk, buffer_attr).to(device, non_blocking=non_blocking)) + setattr(model_chunk, buffer_attr, None) model_chunk.offloaded_states.remove(MegatronOffloadStateType.model_params) @@ -523,7 +525,7 @@ def needs_reload(target, include, offloaded_states): def offload_adam_states(optimizer, device, pin_memory: bool = False, non_blocking: bool = False): - """Move optimizer states to device. Note that this assumes the state structure of DeepSpeed Adam.""" + """Move optimizer states to device.""" state_tensors = [] for _, state in optimizer.state.items(): if "exp_avg" in state: @@ -536,7 +538,7 @@ def offload_adam_states(optimizer, device, pin_memory: bool = False, non_blockin def reload_adam_states(optimizer, device, non_blocking: bool = False): - """Move optimizer states to device. Note that this assumes the state structure of DeepSpeed Adam.""" + """Move optimizer states to device.""" state_tensors = [] for _, state in optimizer.state.items(): if "exp_avg" in state: @@ -546,3 +548,6 @@ def reload_adam_states(optimizer, device, non_blocking: bool = False): if getattr(optimizer, "optimizer_states_cpu_buffers", None) is not None: move_device_buffer_to_tensors(tensors=state_tensors, device_buffer=getattr(optimizer, "optimizer_states_cpu_buffers").to(device, non_blocking=non_blocking),) + optimizer.optimizer_states_cpu_buffers = None + + diff --git a/roll/third_party/megatron/router_replay_utils.py b/roll/third_party/megatron/router_replay_utils.py new file mode 100644 index 000000000..7ddbd3cbc --- /dev/null +++ b/roll/third_party/megatron/router_replay_utils.py @@ -0,0 +1,358 @@ +""" +Router Replay Utilities +Utilities for handling router replay functionality in Megatron models. +ref from https://github.com/verl-project/verl/blob/cb236075dbf1f9b89660d5e2f28e30f3268ec7ee/verl/utils/megatron/router_replay_utils.py +""" + +from typing import Optional + +import torch +from megatron.core import parallel_state as mpu +from megatron.core.pipeline_parallel.schedules import get_schedule_table +from megatron.core.pipeline_parallel.utils import is_vp_first_stage, is_vp_last_stage +from megatron.core.tensor_parallel import gather_from_sequence_parallel_region, scatter_to_sequence_parallel_region +from megatron.core.transformer.transformer_config import TransformerConfig +from megatron.core.transformer.transformer_layer import get_transformer_layer_offset +from megatron.core.transformer.transformer_block import get_num_layers_to_build + +from roll.third_party.megatron.util import postprocess_packed_seqs, preprocess_packed_seqs +# from roll.third_party.megatron.router_replay_patch import RouterReplay, RouterReplayAction +from megatron.core.transformer.moe.router_replay import ( + RouterReplay, + RouterReplayAction, +) + + +def get_routed_experts_dtype(max_expert_idx: int) -> torch.dtype: + """Select the minimal dtype for storing routed expert indices. + + Uses uint8 when all indices fit in 0-255 (1 byte per element), + otherwise falls back to int16 (2 bytes per element, range 0-32767). + + Args: + max_expert_idx: Maximum expert index value. + + Returns: + torch.dtype: uint8 if max_expert_idx <= 255, int16 otherwise. + """ + if max_expert_idx <= 255: + return torch.uint8 + assert max_expert_idx <= 32767, ( + f"Expert index {max_expert_idx} exceeds int16 range (0-32767). " + f"Consider using a larger dtype." + ) + return torch.int16 + + +def get_device_name() -> str: + """Get the device type string based on available accelerators. + + Detects the available accelerator and returns the corresponding PyTorch + device type string. Currently supports CUDA, Ascend NPU, and CPU. + + Returns: + str: Device type string ('cuda', or 'cpu'). + """ + if torch.cuda.is_available(): + device = "cuda" + else: + device = "cpu" + return device + + +def merge_router_topk_indices(attention_mask, input_ids, mini_layer_topk_idx_list, tf_config, vp_rank=None): + """ + Merge recorded router top-k indices across sequence-parallel ranks for all router instances, + then pack/unpack them to align with the original (batch, seq_len) layout and append the result. + + Args: + attention_mask (torch.Tensor): Attention mask of shape [batch_size, seq_len]. Used to determine + the valid token positions during pack/unpack. + input_ids (torch.Tensor): Input token IDs of shape [batch_size, seq_len]. Used together with + attention_mask for sequence packing/unpacking. + mini_layer_topk_idx_list (list): A Python list to which the merged top-k indices tensor will be appended. + tf_config: Megatron/Transformer engine configuration object. Used to locate router instances for + the current micro-batch. + vp_rank (Optional[int]): Virtual pipeline stage rank override. If None, the current VP rank from + Megatron parallel state will be used. + + Returns: + None: The function has side effects only; it appends a tensor of shape + [1, dynamic_bs_all, layer_num, topk] to mini_layer_topk_idx_list. + """ + with torch.no_grad(): + router_instances_list = RouterReplayHelper.get_micro_batch_router_list(tf_config, vp_rank) + max_expert_idx = max(int(r.recorded_topk_idx.max().item()) for r in router_instances_list) + expert_dtype = get_routed_experts_dtype(max_expert_idx) + layers_topk_idx = [] + for router in router_instances_list: + layers_topk_idx.append(router.recorded_topk_idx.to(expert_dtype)) + + # layer_num, dynamic_bs, topk -> dynamic_bs, layer_num, topk + layers_topk_idx = torch.stack(layers_topk_idx).permute(1, 0, 2).to(device_name) + # dynamic_bs, layer_num, topk -> 1, dynamic_bs_all, layer_num, topk + layers_topk_idx = ( + gather_from_sequence_parallel_region(layers_topk_idx, tensor_parallel_output_grad=False) + .unsqueeze(0) + .contiguous() + ) + + batch_size, seq_len = attention_mask.shape[:2] + _, packed_seq_params = preprocess_packed_seqs(input_ids, attention_mask, pre_process=True) + layers_topk_idx = postprocess_packed_seqs( + layers_topk_idx, packed_seq_params, attention_mask, batch_size, seq_len, post_process=True + ) + mini_layer_topk_idx_list.append(layers_topk_idx.cpu()) + + +def set_router_replay_data(layers_topk_idx, attention_mask, tf_config, vp_rank=None): + """ + Scatter the packed router top-k indices back to sequence-parallel ranks and update each local + RouterReplay instance with target indices for replay mode. + + This function prepares the per-layer, per-sample top-k routing decisions (recorded during an earlier + forward) so that subsequent replay passes can follow exactly the same routing. + + Args: + layers_topk_idx (torch.Tensor): Router top-k indices with shape [bs, max_seq_len, layer_num, topk]. + This should be the merged output produced by merge_router_topk_indices. + attention_mask (torch.Tensor): Attention mask [batch_size, seq_len] used for pack/unpack alignment. + tf_config: Megatron/Transformer engine configuration object. + vp_rank (Optional[int]): Virtual pipeline stage rank override. If None, the current VP rank from + Megatron parallel state will be used. + + Returns: + None: The function updates internal RouterReplay instances in-place. + """ + with torch.no_grad(): + # layers_topk_idx_rmpad, _ = preprocess_packed_seqs(layers_topk_idx, attention_mask, pre_process=True) + # layers_topk_idx_rmpad = layers_topk_idx_rmpad.contiguous() # 1, dynamic_bs_all, layer_num, topk + + # # 1, dynamic_bs_split, layer_num, topk + # layers_topk_idx_rmpad_split = scatter_to_sequence_parallel_region( + # layers_topk_idx_rmpad.to(device_name).squeeze(dim=0) + # ).unsqueeze(dim=0) + + # # dynamic_bs_split, layer_num, topk -> layer_num, dynamic_bs_split, topk + # layers_topk_idx_reshape = layers_topk_idx_rmpad_split.permute(0, 2, 1, 3).squeeze( + # dim=0 + # ) # layer_num, dynamic_bs_all, topk + local_rank_info = get_current_rank_layer_info(tf_config, vp_rank) + offset = local_rank_info["start"] + + router_instances_list = RouterReplayHelper.get_micro_batch_router_list(tf_config, vp_rank) + + # 1. [bsz, seq_len, layer_num, topk] -> [seq_len, bsz, layer_num, topk] + layers_topk_idx = layers_topk_idx.permute(1, 0, 2, 3).contiguous() + + # 2. SP split along seq_len -> [seq_len/sp, bsz, layer_num, topk] + layers_topk_idx = scatter_to_sequence_parallel_region(layers_topk_idx.to(get_device_name())) + + # 3. reshape -> [seq_len/sp * bsz, layer_num, topk] + # flatten order: seq_len outer, bsz inner — matches router's logits.view(-1, num_experts) + layers_topk_idx = layers_topk_idx.reshape(-1, layers_topk_idx.shape[2], layers_topk_idx.shape[3]) + + # 4. permute -> [layer_num, seq_len/sp * bsz, topk] + layers_topk_idx = layers_topk_idx.permute(1, 0, 2).contiguous() + + # 5. Per layer: [seq_len/sp * bsz, topk] — aligned with scores [seq_len/sp * bsz, num_experts] + for i, router in enumerate(router_instances_list): + router.set_target_indices(layers_topk_idx[i + offset].to(torch.int64)) + + +def reorder_and_merge_vpp_layers( + micro_batch_tensor_list, + num_microbatches: int, + vpp_size: int, + microbatch_group_size_per_vp_stage: int, +) -> torch.Tensor: + """ + Reorder and merge per-VPP layer blocks into a contiguous layer dimension. + + Given a tensor shaped as [bs*vpp_size, max_token_len, layer_num_per_vpp, topk], this function: + 1) Builds the schedule table for virtual microbatches and reorders the first dimension so that entries + belonging to the same model chunk (VPP stage) become contiguous. + 2) Reshapes and merges the (vpp_size, layer_num_per_vpp) into a single layer dimension, producing + [bs, max_token_len, layer_num, topk]. + + Args: + micro_batch_tensor_list : the list of Input tensor. + num_microbatches (int): Number of microbatches per pipeline stage (bs). + vpp_size (int): Virtual pipeline parallel size (number of model chunks). + microbatch_group_size_per_vp_stage (int): Number of consecutive microbatches processed per VPP stage. + + Returns: + torch.Tensor: Output tensor of shape [bs, max_token_len, layer_num, topk]. + + Raises: + ValueError: If input tensor dimensionality or expected sizes do not match. + RuntimeError: If the computed output shape is unexpected or the schedule length mismatches. + """ + # 1) Build schedule table: map each virtual_microbatch_id -> (microbatch_id, model_chunk_id) + schedule_table = get_schedule_table(num_microbatches, vpp_size, microbatch_group_size_per_vp_stage) + + # 2) Group by model_chunk_id to build reorder indices so entries of the same chunk become contiguous along dim 0 + tensor_by_chunk = [[] for _ in range(vpp_size)] + mini_tensor_list = [] + + for vidx, (_mb, chunk_id) in enumerate(schedule_table): + tensor_by_chunk[chunk_id].append(micro_batch_tensor_list[vidx]) + + for chunk_id in range(vpp_size): + mini_tensor_list.append(torch.cat(tensor_by_chunk[chunk_id], dim=0)) + + out = torch.cat(mini_tensor_list, dim=2) + return out + + +def get_current_rank_layer_info(tf_config, vp_rank=None): + # When vp_rank is None, default to the current VP rank (or 0 if VP is disabled). + """Return the local layer range/count for the current process and the full assignment table. + + Args: + tf_config: Configuration object used by compute_pipeline_layer_assignment. + vp_rank (Optional[int]): Explicit virtual pipeline stage rank to query. If None, uses + mpu.get_virtual_pipeline_model_parallel_rank() when VP is enabled; otherwise 0. + + Returns: + Tuple[dict, dict]: A tuple of (local_assignment, all_assignments) where local_assignment contains + keys {"start", "end", "count"} for the current (pp_rank, vp_stage). + """ + if vp_rank is None: + vp_rank = 0 + num_layers_to_build = get_num_layers_to_build(tf_config, vp_stage=vp_rank) + offset = get_transformer_layer_offset(tf_config, vp_stage=vp_rank) + local = {} + local["start"] = offset + local["end"] = offset + num_layers_to_build + local["count"] = num_layers_to_build + return local + + +def pp_gather(local_layers_router_map, tf_config): + # TODO: Consider non-uniform layer allocation cases. + """ + Gather local router maps from all PP ranks into a global router map. + pp_gather 是为 R2 模式（Megatron 推理 + Megatron 训练）设计的辅助函数，用于： + 在 Pipeline Parallel（PP）场景下，将各 PP rank 上记录的局部 router map 汇聚成全局 router map。 + + Args: + local_layers_router_map (torch.Tensor): Local router map of shape + [bs, max_seq_len, local_num_layers, topk]. + tf_config: Configuration providing pipeline_model_parallel_size. + + Returns: + torch.Tensor: Global router map of shape [bs, max_seq_len, num_layers, topk] placed on CPU. + """ + pp_size = tf_config.pipeline_model_parallel_size + if pp_size <= 1: + return local_layers_router_map + + pp_group = mpu.get_pipeline_model_parallel_group() + world_size = torch.distributed.get_world_size(pp_group) + local_layers_router_map = local_layers_router_map.to(device_name) + layers_topk_idx_global_list = [ + torch.empty( + size=local_layers_router_map.shape, + dtype=local_layers_router_map.dtype, + device=local_layers_router_map.device, + ) + for _ in range(world_size) + ] + torch.distributed.all_gather( + tensor=local_layers_router_map, + tensor_list=layers_topk_idx_global_list, + group=pp_group, + async_op=False, + ) + vp_size = tf_config.virtual_pipeline_model_parallel_size + if vp_size is not None: + vpp_router_map_offset = [[] for _ in range(pp_size)] + for pp_stage in range(pp_size): + vpp_router_map_offset[pp_stage].append(0) + for vp_stage in range(vp_size): + num_layers_to_build = get_num_layers_to_build(tf_config, vp_stage, pp_stage) + vpp_router_map_offset[pp_stage].append(num_layers_to_build + vpp_router_map_offset[pp_stage][-1]) + layers_topk_idx_global = [] + for vp_stage in range(vp_size): + for pp_stage in range(pp_size): + piece = slice(vpp_router_map_offset[pp_stage][vp_stage], vpp_router_map_offset[pp_stage][vp_stage + 1]) + layers_topk_idx_global.append(layers_topk_idx_global_list[pp_stage][:, :, piece, :]) + global_router_map = torch.cat(layers_topk_idx_global, dim=2).to("cpu") + else: + global_router_map = torch.cat(layers_topk_idx_global_list, dim=2).to("cpu") + + return global_router_map + + +class RouterReplayHelper: + """Helper class to query router replay state and locate local RouterReplay instances.""" + + @staticmethod + def get_micro_batch_router_list(tf_config, vp_rank=None): + """ + Return the list of RouterReplay instances corresponding to the current micro-batch and local + (pp_rank, vp_stage) layer range. + + When virtual pipeline (VPP) is enabled, the local range for the PP rank is expanded to include + all VP stages by multiplying the per-VP count by vp_size. The returned slice is taken from the + global RouterReplay.router_instances list. + + Args: + tf_config: Configuration object used to compute layer assignments. + vp_rank (Optional[int]): Explicit virtual pipeline stage to query. If None, the current VP + rank from Megatron parallel state is used when available. + Returns: + list: A contiguous sublist of RouterReplay.router_instances for the local layer range. + """ + vp_size = tf_config.virtual_pipeline_model_parallel_size + if vp_size is not None: + vp_rank = 0 if vp_rank is None else vp_rank + offset = 0 + for pre_vp_stage in range(vp_size): + if pre_vp_stage == vp_rank: + break + num_layers_to_build = get_num_layers_to_build(tf_config, pre_vp_stage) + offset += num_layers_to_build + else: + offset = 0 + + num_layers_to_build = get_num_layers_to_build(tf_config, vp_rank) + router_instances_list = RouterReplay.global_router_replay_instances[offset : offset + num_layers_to_build] + return router_instances_list + + @staticmethod + def is_r2_record_action(tf_config, vp_rank=None) -> bool: + """Return True if the current router_replay_action is RECORD (R2) for the local router instances. + + This inspects the first local RouterReplay instance's router_replay_action and compares it to + RouterReplayAction.RECORD. + """ + router_instances_list = RouterReplayHelper.get_micro_batch_router_list(tf_config, vp_rank) + return router_instances_list and router_instances_list[0].router_replay_action == RouterReplayAction.RECORD + + @staticmethod + def is_replay_forward_action(tf_config, vp_rank=None) -> bool: + """Return True if the current router_replay_action is REPLAY_FORWARD for the local router instances. + + This inspects the first local RouterReplay instance's router_replay_action and compares it to + RouterReplayAction.REPLAY_FORWARD. + """ + router_instances_list = RouterReplayHelper.get_micro_batch_router_list(tf_config, vp_rank) + return ( + router_instances_list and router_instances_list[0].router_replay_action == RouterReplayAction.REPLAY_FORWARD + ) + + @staticmethod + def is_replay_backward_action(tf_config, vp_rank=None) -> bool: + """Return True if the current router_replay_action is REPLAY_BACKWARD for the local router instances. + + This inspects the first local RouterReplay instance's router_replay_action and compares it to + RouterReplayAction.REPLAY_BACKWARD. + """ + router_instances_list = RouterReplayHelper.get_micro_batch_router_list(tf_config, vp_rank) + return ( + router_instances_list + and router_instances_list[0].router_replay_action == RouterReplayAction.REPLAY_BACKWARD + ) + diff --git a/roll/third_party/megatron/tensor_parallel.py b/roll/third_party/megatron/tensor_parallel.py index 69648296e..602f30f6f 100644 --- a/roll/third_party/megatron/tensor_parallel.py +++ b/roll/third_party/megatron/tensor_parallel.py @@ -2,49 +2,120 @@ import torch.distributed as dist from megatron.core import parallel_state as mpu +from roll.utils.logging import get_logger + +try: + from .fused_entropy import entropy_fwd, entropy_bwd + FUSED_KERNEL_AVAILABLE = True +except ImportError: + FUSED_KERNEL_AVAILABLE = False + +logger = get_logger() + class _VocabParallelEntropy(torch.autograd.Function): @staticmethod - def forward(ctx, vocab_parallel_logits: torch.Tensor) -> torch.Tensor: - @torch.compile(dynamic=True) - def mul_reduce(a, b): - return (a * b).sum(dim=-1, keepdim=True) - - logits_max = vocab_parallel_logits.max(dim=-1, keepdim=True).values - dist.all_reduce(logits_max, op=dist.ReduceOp.MAX, group=mpu.get_tensor_model_parallel_group()) - normalized_vocab_parallel_logits = vocab_parallel_logits - logits_max - normalized_exp_logits = normalized_vocab_parallel_logits.exp_() - normalized_sum_exp_logits = normalized_exp_logits.sum(dim=-1, keepdim=True) - dist.all_reduce(normalized_sum_exp_logits, group=mpu.get_tensor_model_parallel_group()) - softmax_logits = normalized_exp_logits.div_(normalized_sum_exp_logits) - sum_softmax_times_logits = mul_reduce(softmax_logits, vocab_parallel_logits) - dist.all_reduce(sum_softmax_times_logits, group=mpu.get_tensor_model_parallel_group()) - entropy = logits_max + normalized_sum_exp_logits.log() - sum_softmax_times_logits - ctx.save_for_backward(vocab_parallel_logits, softmax_logits, sum_softmax_times_logits) - return entropy.squeeze(dim=-1) + def forward( + ctx, + vocab_parallel_logits: torch.Tensor, + used_fp32: bool = True, + use_fused_kernel: bool = True + ) -> torch.Tensor: + + # Only use fused kernel when TP=1 and use_fused_kernel is True + if use_fused_kernel: + # Get tensor parallel world size + tp_world_size = mpu.get_tensor_model_parallel_world_size() + if tp_world_size != 1 or not FUSED_KERNEL_AVAILABLE: + logger.warning(f"Disable use_fused_kernel because {tp_world_size=} and {FUSED_KERNEL_AVAILABLE=}.") + use_fused_kernel = False + + if use_fused_kernel: + vocab_parallel_logits_2d = vocab_parallel_logits.view(-1, vocab_parallel_logits.shape[-1]) + + # Use fused kernel implementation (only for TP=1) + entropy, x_max, x_sum_exp, x_sum_softmax_times = entropy_fwd(vocab_parallel_logits_2d) + + # Convert output back to original shape + if vocab_parallel_logits.dim() == 3: + batch_size, seq_len, vocab_size = vocab_parallel_logits.shape + entropy = entropy.view(batch_size, seq_len) + + # Save for backward: vocab_parallel_logits and intermediate results + ctx.save_for_backward(vocab_parallel_logits, x_max, x_sum_exp, x_sum_softmax_times) + ctx.use_fused_kernel = True + # ctx.original_shape = original_shape + return entropy + else: + # Original implementation for TP>1 or when fused kernel is not available + @torch.compile(dynamic=True) + def mul_reduce(a, b): + return (a * b).sum(dim=-1, keepdim=True) + + ctx.input_dtype = vocab_parallel_logits.dtype + + if used_fp32: + vocab_parallel_logits = vocab_parallel_logits.float() + + logits_max = vocab_parallel_logits.max(dim=-1, keepdim=True).values + dist.all_reduce(logits_max, op=dist.ReduceOp.MAX, group=mpu.get_tensor_model_parallel_group()) + normalized_vocab_parallel_logits = vocab_parallel_logits - logits_max + normalized_exp_logits = normalized_vocab_parallel_logits.exp_() + normalized_sum_exp_logits = normalized_exp_logits.sum(dim=-1, keepdim=True) + dist.all_reduce(normalized_sum_exp_logits, group=mpu.get_tensor_model_parallel_group()) + softmax_logits = normalized_exp_logits.div_(normalized_sum_exp_logits) + sum_softmax_times_logits = mul_reduce(softmax_logits, vocab_parallel_logits) + dist.all_reduce(sum_softmax_times_logits, group=mpu.get_tensor_model_parallel_group()) + entropy = logits_max + normalized_sum_exp_logits.log() - sum_softmax_times_logits + ctx.save_for_backward(vocab_parallel_logits, softmax_logits, sum_softmax_times_logits) + ctx.use_fused_kernel = False + + return entropy.squeeze(dim=-1) @staticmethod def backward(ctx, grad_output: torch.Tensor) -> torch.Tensor: - vocab_parallel_logits, softmax_logits, sum_softmax_times_logits = ctx.saved_tensors - # reuse softmax_logits as grad - vocab_parallel_logits.sub_(sum_softmax_times_logits) - softmax_logits.mul_(vocab_parallel_logits) - softmax_logits.mul_(grad_output.unsqueeze(dim=-1)) - # recover vocab_parallel_logits - vocab_parallel_logits.add_(sum_softmax_times_logits) - softmax_logits.mul_(-1) - return softmax_logits - - -def vocab_parallel_entropy(vocab_parallel_logits: torch.Tensor) -> torch.Tensor: + if ctx.use_fused_kernel: + # Use fused kernel backward with recomputation (only for TP=1) + vocab_parallel_logits, x_max, x_sum_exp, x_sum_softmax_times = ctx.saved_tensors + + vocab_parallel_logits_2d = vocab_parallel_logits.view(-1, vocab_parallel_logits.shape[-1]) + + # Call fused backward kernel (performs recomputation internally) + grad_input_2d = entropy_bwd( + vocab_parallel_logits_2d, + grad_output.view(-1), + x_max, + x_sum_exp, + x_sum_softmax_times + ) + + grad_input = grad_input_2d.view_as(vocab_parallel_logits) + + return grad_input, None, None + else: + # Original implementation for TP>1 + vocab_parallel_logits, softmax_logits, sum_softmax_times_logits = ctx.saved_tensors + # reuse softmax_logits as grad + vocab_parallel_logits.sub_(sum_softmax_times_logits) + softmax_logits.mul_(vocab_parallel_logits) + softmax_logits.mul_(grad_output.unsqueeze(dim=-1)) + # recover vocab_parallel_logits + vocab_parallel_logits.add_(sum_softmax_times_logits) + softmax_logits.mul_(-1) + softmax_logits = softmax_logits.to(ctx.input_dtype) + + return softmax_logits, None, None + +def vocab_parallel_entropy(vocab_parallel_logits: torch.Tensor, used_fp32=True, use_fused_kernel: bool = True) -> torch.Tensor: """ ref: https://github.com/volcengine/verl/blob/78532923368aeb058f62201489546d013df47710/verl/utils/megatron/tensor_parallel.py#L109 Compute entropy when the logits are sharded in tp ranks Args: vocab_parallel_logits: (total_nnz, vocab_size // tp_size) + use_fused_kernel: whether to use fused kernel implementation (default: True) Returns: (total_nnz,) """ - return _VocabParallelEntropy.apply(vocab_parallel_logits) + return _VocabParallelEntropy.apply(vocab_parallel_logits, used_fp32, use_fused_kernel) diff --git a/roll/third_party/megatron/util.py b/roll/third_party/megatron/util.py new file mode 100644 index 000000000..26142aee0 --- /dev/null +++ b/roll/third_party/megatron/util.py @@ -0,0 +1,181 @@ +# https://github.com/verl-project/verl/blob/cb236075dbf1f9b89660d5e2f28e30f3268ec7ee/verl/models/mcore/util.py +import math +import torch + +from megatron.core import parallel_state as mpu +from megatron.core.packed_seq_params import PackedSeqParams +from megatron.core.distributed import DistributedDataParallel as DDP +from megatron.core.transformer.module import Float16Module + + +ALL_MODULE_WRAPPER_CLASSNAMES = (DDP, Float16Module) + + +def unwrap_model(model, module_instances=ALL_MODULE_WRAPPER_CLASSNAMES): + return_list = True + if not isinstance(model, list): + model = [model] + return_list = False + unwrapped_model = [] + for model_module in model: + while isinstance(model_module, module_instances): + model_module = model_module.module + unwrapped_model.append(model_module) + if not return_list: + return unwrapped_model[0] + return unwrapped_model + + +def preprocess_packed_seqs( + input_ids: torch.Tensor, attention_mask: torch.Tensor, pre_process: bool = True, use_fp8_padding=False +) -> tuple[torch.Tensor, PackedSeqParams]: + """ + Preprocess packed sequences + CP splits sequence into CP*2 chunks, and each GPU gets 2 chunks (GPU0 gets first and last chunks, GPU1 + gets second and second last chunks, and so on), this is for load balancing with causal masking. + See https://github.com/NVIDIA/TransformerEngine/issues/1368 + """ + batch_size = input_ids.shape[0] + + seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) + tp_size = mpu.get_tensor_model_parallel_world_size() + cp_size = mpu.get_context_parallel_world_size() + cp_rank = mpu.get_context_parallel_rank() + align_size = tp_size * cp_size * 2 if cp_size > 1 else tp_size + if use_fp8_padding: + # if fp8 is enabled, ensure the sequence is padded to multiples of 16 for better performance + original_align_size = align_size + align_size = math.lcm(16, align_size) + + pad_size = (align_size - seqlens_in_batch % align_size) % align_size + seqlens_in_batch_padded = seqlens_in_batch + pad_size + + cu_seqlens = torch.zeros(batch_size + 1, dtype=torch.int32, device=input_ids.device) + cu_seqlens[1:] = torch.cumsum(seqlens_in_batch, dim=0) + cu_seqlens_padded = torch.zeros(batch_size + 1, dtype=torch.int32, device=input_ids.device) + cu_seqlens_padded[1:] = torch.cumsum(seqlens_in_batch_padded, dim=0) + + if use_fp8_padding: + # make sure all the sequences are padded to multiples of 128 for TE compatibility + align_size_last = original_align_size * 128 + pad_size_last = (align_size_last - cu_seqlens_padded[-1] % align_size_last) % align_size_last + cu_seqlens_padded[-1] += pad_size_last + seqlens_in_batch_padded[-1] += pad_size_last + + # ---------------------------------------------------------------------------- + # Move the index information needed in the subsequent loop to the CPU at once, + # to avoid frequent .item() calls in the loop that cause D2H synchronization + # ---------------------------------------------------------------------------- + seqlens_in_batch_cpu: list[int] = seqlens_in_batch.tolist() # original valid lengths + seqlens_in_batch_padded_cpu: list[int] = seqlens_in_batch_padded.tolist() # lengths after padding + cu_seqlens_padded_cpu: list[int] = cu_seqlens_padded.tolist() # start positions (after padding) + + # Pure Python int calculation to avoid further synchronization + max_seqlen_in_batch = max(seqlens_in_batch_padded_cpu) + + shape = list(input_ids.shape[1:]) + shape[0] = sum(seqlens_in_batch_padded_cpu) // cp_size + if pre_process: + input_ids_rmpad = torch.zeros(shape, dtype=input_ids.dtype, device=input_ids.device) + for i in range(batch_size): + # Use Python int, so no GPU→CPU sync in the loop + if cp_size <= 1: + seqlen = seqlens_in_batch_cpu[i] + start_idx = cu_seqlens_padded_cpu[i] + input_ids_rmpad[start_idx : start_idx + seqlen] = input_ids[i, attention_mask[i]] + continue + + seqlen_padded_i = seqlens_in_batch_padded_cpu[i] + seqlen = seqlen_padded_i // cp_size + half_seqlen = seqlen // 2 + start_idx = cu_seqlens_padded_cpu[i] // cp_size + # split to 2 chunks + d = input_ids[i, attention_mask[i]] + input_ids_rmpad[start_idx : start_idx + half_seqlen] = d[ + half_seqlen * cp_rank : half_seqlen * (cp_rank + 1) + ] + + remain_start = seqlen_padded_i - half_seqlen * (cp_rank + 1) + remain_end = seqlen_padded_i - half_seqlen * cp_rank + remain_end = min(remain_end, d.shape[0]) + remain_len = remain_end - remain_start + if remain_len > 0: + input_ids_rmpad[start_idx + half_seqlen : start_idx + half_seqlen + remain_len] = d[ + remain_start:remain_end + ] + + packed_seq_params = PackedSeqParams( + qkv_format="thd", + cu_seqlens_q=cu_seqlens_padded, + max_seqlen_q=max_seqlen_in_batch, + cu_seqlens_kv=cu_seqlens_padded, + max_seqlen_kv=max_seqlen_in_batch, + cu_seqlens_q_padded=cu_seqlens_padded, + cu_seqlens_kv_padded=cu_seqlens_padded, + ) + if pre_process: + return input_ids_rmpad.unsqueeze(0), packed_seq_params + else: + return input_ids, packed_seq_params + + +def postprocess_packed_seqs( + output: torch.Tensor, + packed_seq_params: PackedSeqParams, + attention_mask: torch.Tensor, + batch_size: int, + seq_len: int, + post_process: bool = True, +) -> torch.Tensor: + """ + Postprocess packed sequences + """ + if not post_process: + return output + + # ------------------------------------------------------------------------- + # Move the lengths and offsets needed for subsequent Python-level indexing to the CPU in advance, + # to avoid a large number of .item() calls in the loop + # ------------------------------------------------------------------------- + cu_padded_cpu: list[int] = packed_seq_params.cu_seqlens_q_padded.tolist() + seq_lens_cpu: list[int] = attention_mask.sum(dim=1, dtype=torch.int32).cpu().tolist() + + shape = [batch_size, seq_len] + list(output.shape[2:]) # 1,packed, dim -> batch_size, seq_len, dim + output_new = torch.zeros(shape, dtype=output.dtype, device=output.device) + + cp_size = mpu.get_context_parallel_world_size() + # all gather output across context parallel group + if cp_size > 1: + # output shape: [1, packed_len, hidden_dim] + # need to gather across cp group and concatenate in sequence dimension + output_list = [torch.empty_like(output, dtype=output.dtype) for _ in range(cp_size)] + torch.distributed.all_gather(output_list, output.detach(), group=mpu.get_context_parallel_group()) + output_list[mpu.get_context_parallel_rank()] = output + else: + output_list = [output] + for i in range(batch_size): + if cp_size <= 1: + s = seq_lens_cpu[i] + start_idx = cu_padded_cpu[i] + output_new[i, attention_mask[i]] = output[0][start_idx : start_idx + s] + continue + s_len_padded_chunk = (cu_padded_cpu[i + 1] - cu_padded_cpu[i]) // cp_size + half_seqlen = s_len_padded_chunk // 2 + s_len = seq_lens_cpu[i] + s_len_padded = s_len_padded_chunk * cp_size + tmp = torch.empty(s_len_padded, *output.shape[2:], device=output.device, dtype=output.dtype) + for j in range(cp_size): + o = output_list[j][0] + # split to 2 chunks + packed_start_idx = cu_padded_cpu[i] // cp_size + o0, o1 = ( + o[packed_start_idx : packed_start_idx + half_seqlen], + o[packed_start_idx + half_seqlen : packed_start_idx + s_len_padded_chunk], + ) + tmp[j * half_seqlen : (j + 1) * half_seqlen] = o0 + tmp[s_len_padded - (j + 1) * half_seqlen : s_len_padded - j * half_seqlen] = o1 + output_new[i, attention_mask[i]] = tmp[:s_len] + + return output_new + + diff --git a/roll/third_party/sglang/__init__.py b/roll/third_party/sglang/__init__.py index 4ddf780fa..1df493906 100644 --- a/roll/third_party/sglang/__init__.py +++ b/roll/third_party/sglang/__init__.py @@ -13,10 +13,7 @@ elif sgl.__version__ == '0.5.2': from roll.third_party.sglang import v052_patch patch = v052_patch -elif sgl.__version__ == '0.5.4.post2': - from roll.third_party.sglang import v054_patch - patch = v054_patch -elif sgl.__version__ == '0.5.5.post3' or sgl.__version__ == '0.5.6.post2': +elif sgl.__version__ in ['0.5.4.post2', '0.5.5.post3', '0.5.6.post2', '0.5.7']: from roll.third_party.sglang import v054_patch patch = v054_patch else: diff --git a/roll/third_party/sglang/v046post4_patch/engine.py b/roll/third_party/sglang/v046post4_patch/engine.py index cf0a47f21..6e68525a0 100644 --- a/roll/third_party/sglang/v046post4_patch/engine.py +++ b/roll/third_party/sglang/v046post4_patch/engine.py @@ -1,6 +1,9 @@ import os import multiprocessing as mp +import warnings +from packaging.version import Version +import sglang import sglang.srt.entrypoints.engine as engine_module from sglang.srt.server_args import ServerArgs from sglang.srt.utils import ( @@ -36,8 +39,17 @@ def _set_envs_and_config(server_args: ServerArgs): mp.set_start_method("spawn", force=True) def run_scheduler_process(*args, **kwargs): - from roll.third_party.sglang import fp8 - fp8.monkey_patch_fp8() + sglang_version = Version(sglang.__version__) + if sglang_version >= Version("0.4.6.post4"): + from roll.third_party.sglang import fp8 + fp8.monkey_patch_fp8() + else: + warnings.warn( + f"sglang version {sglang.__version__} < 0.4.6.post4, " + "fp8 monkey patch is not supported. " + "Please upgrade sglang to 0.4.6.post4 or later to use fp8.", + stacklevel=2, + ) from sglang.srt.managers.scheduler import run_scheduler_process return run_scheduler_process(*args, **kwargs) diff --git a/roll/third_party/vllm/__init__.py b/roll/third_party/vllm/__init__.py index 77f67cbbb..bbc58a042 100644 --- a/roll/third_party/vllm/__init__.py +++ b/roll/third_party/vllm/__init__.py @@ -2,6 +2,7 @@ import pathlib from typing import Dict, List +import dataclasses import torch import vllm from packaging.version import Version @@ -46,6 +47,10 @@ async def create_async_llm(resource_placement_groups: List[Dict], **kwargs): kwargs["enable_sleep_mode"] = True + if "attention_config" not in kwargs and "attention_config" in { + f.name: f for f in dataclasses.fields(AsyncEngineArgs) + }: # vllm<=0.12.0 not has attention_config in AsyncEngineArgs + kwargs["attention_config"] = {"backend": "FLASH_ATTN"} if "worker_extension_cls" not in kwargs: # VLLM_USE_V1 is deprecated in vllm>=0.11.1 @@ -78,7 +83,7 @@ async def create_async_llm(resource_placement_groups: List[Dict], **kwargs): # VLLM_USE_V1 may be modified inside create_engine_config vllm_config = engine_args.create_engine_config(UsageContext.ENGINE_CONTEXT) - fp8.update_quant_config(vllm_config) + fp8.update_quant_config(config=kwargs, vllm_config=vllm_config) # change parallel_config.placement_group for CustomRayDistributedExecutor parallel_config = vllm_config.parallel_config diff --git a/roll/third_party/vllm/fp8.py b/roll/third_party/vllm/fp8.py index 762c743a5..782e8b0d3 100644 --- a/roll/third_party/vllm/fp8.py +++ b/roll/third_party/vllm/fp8.py @@ -21,25 +21,27 @@ logger = get_logger() -def update_quant_config(vllm_config): - # Use hf_overrides arguments of LLM with weight_block_size +def update_quant_config(config, vllm_config): + # Use hf_overrides arguments with weight_block_size # to enable block quantization. # e.g. # strategy_args: # strategy_name: vllm # strategy_config: + # quantization: fp8 # hf_overrides: # quantization_config: # activation_scheme: dynamic + # fmt: e4m3 # quant_method: fp8 # weight_block_size: [128, 128] - if not vllm_config.quant_config: + if "hf_overrides" not in config or "quantization_config" not in config["hf_overrides"]: return - if not isinstance(vllm_config.quant_config, Fp8Config): - return - + assert config["hf_overrides"]["quantization_config"]["quant_method"] == "fp8" + assert isinstance(vllm_config.quant_config, Fp8Config) assert vllm_config.quant_config.activation_scheme == "dynamic" - vllm_config.quant_config.is_checkpoint_fp8_serialized = True + assert vllm_config.quant_config.is_checkpoint_fp8_serialized + vllm_config.quant_config.skip_process_weights_after_loading = True logger.info(f"Using custom vLLM quantization, block size {vllm_config.quant_config.weight_block_size}") def _fp8_linear_weight_loader(layer: weakref.ReferenceType, original_weight_loader, param: torch.Tensor, loaded_weight: torch.Tensor, *args, **kwargs) -> None: @@ -108,6 +110,8 @@ def _fp8_linear_create_weights( ): _original_fp8_linear_create_weights(self, layer, input_size_per_partition, output_partition_sizes, input_size, output_size, params_dtype, **extra_weight_attrs) + if not getattr(self.quant_config, "skip_process_weights_after_loading", False): + return assert self.quant_config.is_checkpoint_fp8_serialized assert self.quant_config.activation_scheme == "dynamic" @@ -137,12 +141,16 @@ def _fp8_linear_create_weights( layer.shard_num = len(output_partition_sizes) layer.shard_loaded = 0 + layer.register_parameter("input_scale", None) + _original_fp8_linear_create_weights = Fp8LinearMethod.create_weights Fp8LinearMethod.create_weights = _fp8_linear_create_weights def _fp8_linear_process_weights_after_loading(self, layer: Module) -> None: - pass + if not getattr(self.quant_config, "skip_process_weights_after_loading", False): + _original_fp8_linear_process_weights_after_loading(self, layer) +_original_fp8_linear_process_weights_after_loading = Fp8LinearMethod.process_weights_after_loading Fp8LinearMethod.process_weights_after_loading = _fp8_linear_process_weights_after_loading def _fp8_moe_w13_weight_loader(layer: weakref.ReferenceType, original_weight_loader, param: torch.Tensor, loaded_weight: torch.Tensor, *args, **kwargs) -> None: @@ -176,6 +184,8 @@ def _fp8_moe_create_weights(self, layer: Module, num_experts: int, hidden_size: params_dtype: torch.dtype, **extra_weight_attrs): _original_fp8_moe_create_weights(self, layer, num_experts, hidden_size, intermediate_size_per_partition, params_dtype, **extra_weight_attrs) + if not getattr(self.quant_config, "skip_process_weights_after_loading", False): + return assert self.quant_config.is_checkpoint_fp8_serialized assert self.quant_config.activation_scheme == "dynamic" @@ -192,6 +202,25 @@ def _fp8_moe_create_weights(self, layer: Module, num_experts: int, hidden_size: # TODO: support ep assert layer.local_num_experts == num_experts + if getattr(self, "_setup_kernel", None): + from vllm.model_executor.layers.fused_moe.oracle.fp8 import Fp8MoeBackend + assert self.fp8_backend not in [ + Fp8MoeBackend.AITER, + Fp8MoeBackend.MARLIN, + Fp8MoeBackend.FLASHINFER_CUTLASS, + Fp8MoeBackend.FLASHINFER_TRTLLM, + # TODO: support inflight fp8 quantization for DEEPGEMM and BATCHED_DEEPGEMM + Fp8MoeBackend.DEEPGEMM, + Fp8MoeBackend.BATCHED_DEEPGEMM, + ] + assert self.fp8_backend in [ + Fp8MoeBackend.TRITON, + Fp8MoeBackend.BATCHED_TRITON, + Fp8MoeBackend.VLLM_CUTLASS, + Fp8MoeBackend.BATCHED_VLLM_CUTLASS, + Fp8MoeBackend.XPU, + ] + # store essential config in layer for custom weight loader layer.weight_block_size = self.quant_config.weight_block_size @@ -213,6 +242,26 @@ def _fp8_moe_create_weights(self, layer: Module, num_experts: int, hidden_size: Fp8MoEMethod.create_weights = _fp8_moe_create_weights def _fp8_moe_process_weights_after_loading(self, layer: Module) -> None: - pass - + if not getattr(self.quant_config, "skip_process_weights_after_loading", False): + _original_fp8_moe_process_weights_after_loading(self, layer) + else: + if getattr(self, "_setup_kernel", None): + w13 = layer.w13_weight + w2 = layer.w2_weight + w13_scale = getattr(layer, f"w13_{self.weight_scale_name}") + w2_scale = getattr(layer, f"w2_{self.weight_scale_name}") + w13_input_scale = layer.w13_input_scale + w2_input_scale = layer.w2_input_scale + assert w13_input_scale is None + assert w2_input_scale is None + # _setup_kernel will change layer.w13_weight, layer.w2_weight + self._setup_kernel( + layer, w13, w2, w13_scale, w2_scale, w13_input_scale, w2_input_scale + ) + assert w13.data_ptr() == layer.w13_weight.data_ptr() + assert w2.data_ptr() == layer.w2_weight.data_ptr() + assert w13_scale.data_ptr() == getattr(layer, f"w13_{self.weight_scale_name}").data_ptr() + assert w2_scale.data_ptr() == getattr(layer, f"w2_{self.weight_scale_name}").data_ptr() + +_original_fp8_moe_process_weights_after_loading = Fp8MoEMethod.process_weights_after_loading Fp8MoEMethod.process_weights_after_loading = _fp8_moe_process_weights_after_loading diff --git a/roll/third_party/vllm/gdn_patcher.py b/roll/third_party/vllm/gdn_patcher.py new file mode 100644 index 000000000..2c3a39c16 --- /dev/null +++ b/roll/third_party/vllm/gdn_patcher.py @@ -0,0 +1,332 @@ +""" +Monkey patches for vLLM GDN attention to fix mixed decode/spec-decode crash. + +This module patches the GDNAttentionMetadataBuilder.build method to add +reclassification logic for non-spec decodes when spec decodes exist. + +Bug: https://github.com/vllm-project/vllm/pull/34871 +Fix commit: 116ed130f + +The bug occurs in vLLM versions < v0.17.2 when processing batches containing +both regular decode requests and speculative decode requests with GDN attention. + +Error: AssertionError: num_decodes: X, num_spec_decodes: Y +""" + +# TODO: This file should be removed when upgrading vLLM to version >= v0.17.2 + +import logging + +logger = logging.getLogger(__name__) + +_patch_applied = False + + +def patch_gdn_attention(): + """ + Apply monkey patch for GDN attention mixed decode/spec-decode bug. + + This function patches GDNAttentionMetadataBuilder.build to add + reclassification logic that converts non-spec decodes to prefills + when spec decodes exist in the same batch. + + The patch should be applied before any vLLM inference is performed. + """ + global _patch_applied + + if _patch_applied: + logger.debug("[GDN PATCH] GDN attention patch already applied, skipping") + return + + try: + from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadataBuilder + + # Save original build method BEFORE checking + original_build = GDNAttentionMetadataBuilder.build + + # Check if already patched (has _original_build attribute) + if hasattr(GDNAttentionMetadataBuilder, '_original_build'): + # Already patched, use the real original + original_build = GDNAttentionMetadataBuilder._original_build + + # Check if the fix is already in the source + import inspect + original_source = inspect.getsource(original_build) + + # Check for various signs that the fix is already present + has_fix = ( + "num_decodes and num_spec_decodes are mutually exclusive" in original_source or + ("num_prefills += num_decodes" in original_source and "num_decodes = 0" in original_source) or + ("num_decodes = 0" in original_source and "num_decode_tokens = 0" in original_source) + ) + + if has_fix: + logger.info("[GDN PATCH] GDN attention already has the fix, skipping patch") + _patch_applied = True + return + + def patched_build( + self, + common_prefix_len: int, + common_attn_metadata, + num_accepted_tokens=None, + num_decode_draft_tokens_cpu=None, + fast_build: bool = False, + ): + """Patched build method with catch + fallback pattern.""" + try: + return original_build( + self, + common_prefix_len, + common_attn_metadata, + num_accepted_tokens, + num_decode_draft_tokens_cpu, + fast_build, + ) + except AssertionError as e: + error_msg = str(e) + if "num_decodes" in error_msg and "num_spec_decodes" in error_msg: + logger.warning( + f"[GDN PATCH] GDN attention assertion caught, applying workaround: {error_msg}" + ) + return _build_with_fix( + self, + common_prefix_len, + common_attn_metadata, + num_accepted_tokens, + num_decode_draft_tokens_cpu, + fast_build, + ) + else: + raise + + # Apply the patch + GDNAttentionMetadataBuilder.build = patched_build + GDNAttentionMetadataBuilder._original_build = original_build + + _patch_applied = True + logger.info("[GDN PATCH] GDN attention patch applied successfully") + + except ImportError as e: + logger.debug(f"[GDN PATCH] vLLM GDN attention not available, skipping patch: {e}") + except Exception as e: + logger.warning(f"[GDN PATCH] Failed to apply GDN attention patch: {e}") + + +def _build_with_fix( + self, + common_prefix_len: int, + common_attn_metadata, + num_accepted_tokens=None, + num_decode_draft_tokens_cpu=None, + fast_build: bool = False, +): + """ + Build method with the reclassification fix inline. + + This is a copy of the original vLLM build method with ONLY the + reclassification logic added. All other logic is identical to + the original implementation. + + The fix is at lines 273-278: reclassify num_decodes as num_prefills + when both num_decodes > 0 and num_spec_decodes > 0. + """ + import torch + from vllm.v1.attention.backends.gdn_attn import ( + GDNAttentionMetadata, + mamba_get_block_table_tensor, + ) + from vllm.v1.attention.backends.utils import ( + split_decodes_and_prefills, + compute_causal_conv1d_metadata, + ) + + m = common_attn_metadata + + query_start_loc = m.query_start_loc + query_start_loc_cpu = m.query_start_loc_cpu + context_lens_tensor = m.compute_num_computed_tokens() + # Initialize these to None at the start (same as original vLLM) + nums_dict, batch_ptr, token_chunk_offset_ptr = None, None, None + block_table_tensor = mamba_get_block_table_tensor( + m.block_table_tensor, + m.seq_lens, + self.kv_cache_spec, + self.vllm_config.cache_config.mamba_cache_mode, + ) + + spec_sequence_masks_cpu = None + if ( + not self.use_spec_decode + or num_decode_draft_tokens_cpu is None + or num_decode_draft_tokens_cpu[num_decode_draft_tokens_cpu >= 0] + .sum() + .item() + == 0 + ): + spec_sequence_masks = None + num_spec_decodes = 0 + else: + spec_sequence_masks_cpu = num_decode_draft_tokens_cpu >= 0 + num_spec_decodes = spec_sequence_masks_cpu.sum().item() + if num_spec_decodes == 0: + spec_sequence_masks = None + spec_sequence_masks_cpu = None + else: + spec_sequence_masks = spec_sequence_masks_cpu.to( + query_start_loc.device, non_blocking=True + ) + + if spec_sequence_masks is None: + num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = ( + split_decodes_and_prefills(m, decode_threshold=1) + ) + num_spec_decode_tokens = 0 + spec_token_indx = None + non_spec_token_indx = None + spec_state_indices_tensor = None + non_spec_state_indices_tensor = block_table_tensor[:, 0] + spec_query_start_loc = None + non_spec_query_start_loc = query_start_loc + non_spec_query_start_loc_cpu = query_start_loc_cpu + num_accepted_tokens = None + else: + query_lens = query_start_loc[1:] - query_start_loc[:-1] + assert spec_sequence_masks_cpu is not None + query_lens_cpu = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1] + + non_spec_query_lens_cpu = query_lens_cpu[~spec_sequence_masks_cpu] + num_decodes = (non_spec_query_lens_cpu == 1).sum().item() + num_zero_len = (non_spec_query_lens_cpu == 0).sum().item() + num_prefills = non_spec_query_lens_cpu.size(0) - num_decodes - num_zero_len + num_decode_tokens = num_decodes + num_prefill_tokens = ( + non_spec_query_lens_cpu.sum().item() - num_decode_tokens + ) + num_spec_decode_tokens = ( + query_lens_cpu.sum().item() - num_prefill_tokens - num_decode_tokens + ) + + # ================================================================ + # THE FIX: Reclassify non-spec decodes as prefills when spec decodes exist + # This is the ONLY change from the original vLLM code. + # ================================================================ + if num_decodes > 0 and num_spec_decodes > 0: + num_prefills += num_decodes + num_prefill_tokens += num_decode_tokens + num_decodes = 0 + num_decode_tokens = 0 + # ================================================================ + + if num_prefills == 0 and num_decodes == 0: + spec_token_size = min( + num_spec_decodes * (self.num_spec + 1), + query_start_loc_cpu[-1].item(), + ) + spec_token_indx = torch.arange( + spec_token_size, + dtype=torch.int32, + device=query_start_loc.device, + ) + non_spec_token_indx = torch.empty( + 0, dtype=torch.int32, device=query_start_loc.device + ) + spec_state_indices_tensor = block_table_tensor[ + spec_sequence_masks_cpu, : self.num_spec + 1 + ] + non_spec_state_indices_tensor = None + spec_query_start_loc = query_start_loc[: num_spec_decodes + 1] + non_spec_query_start_loc = None + non_spec_query_start_loc_cpu = None + else: + spec_token_masks = torch.repeat_interleave( + spec_sequence_masks, query_lens + ) + index = torch.argsort(spec_token_masks, stable=True) + num_non_spec_tokens = num_prefill_tokens + num_decode_tokens + non_spec_token_indx = index[:num_non_spec_tokens] + spec_token_indx = index[num_non_spec_tokens:] + + spec_state_indices_tensor = block_table_tensor[ + spec_sequence_masks_cpu, : self.num_spec + 1 + ] + non_spec_state_indices_tensor = block_table_tensor[ + ~spec_sequence_masks_cpu, 0 + ] + + spec_query_start_loc = torch.zeros( + num_spec_decodes + 1, + dtype=torch.int32, + device=query_start_loc.device, + ) + torch.cumsum( + query_lens[spec_sequence_masks_cpu], dim=0, out=spec_query_start_loc[1:] + ) + non_spec_query_start_loc = torch.zeros( + query_lens.size(0) - num_spec_decodes + 1, + dtype=torch.int32, + device=query_start_loc.device, + ) + torch.cumsum( + query_lens[~spec_sequence_masks_cpu], + dim=0, + out=non_spec_query_start_loc[1:], + ) + non_spec_query_start_loc_cpu = torch.zeros( + query_lens_cpu.size(0) - num_spec_decodes + 1, + dtype=torch.int32, + ) + torch.cumsum( + query_lens_cpu[~spec_sequence_masks_cpu], + dim=0, + out=non_spec_query_start_loc_cpu[1:], + ) + + assert num_accepted_tokens is not None + num_accepted_tokens = num_accepted_tokens[spec_sequence_masks_cpu] + + if num_prefills > 0: + has_initial_state = context_lens_tensor > 0 + if spec_sequence_masks is not None: + has_initial_state = has_initial_state[~spec_sequence_masks] + assert non_spec_query_start_loc_cpu is not None + nums_dict, batch_ptr, token_chunk_offset_ptr = ( + compute_causal_conv1d_metadata( + non_spec_query_start_loc_cpu, + device=query_start_loc.device, + ) + ) + else: + has_initial_state = None + + # NOTE: We skip the cudagraph optimization logic here (lines 314-382 in original) + # because our fallback path is only triggered for mixed batches, which don't + # benefit from cudagraph anyway. This simplifies the code and avoids potential + # issues with cudagraph tensor management. + + return GDNAttentionMetadata( + num_prefills=num_prefills, + num_prefill_tokens=num_prefill_tokens, + num_decodes=num_decodes, + num_decode_tokens=num_decode_tokens, + num_spec_decodes=num_spec_decodes, + num_spec_decode_tokens=num_spec_decode_tokens, + num_actual_tokens=m.num_actual_tokens, + has_initial_state=has_initial_state, + spec_query_start_loc=spec_query_start_loc, + non_spec_query_start_loc=non_spec_query_start_loc, + spec_state_indices_tensor=spec_state_indices_tensor, + non_spec_state_indices_tensor=non_spec_state_indices_tensor, + spec_sequence_masks=spec_sequence_masks, + spec_token_indx=spec_token_indx, + non_spec_token_indx=non_spec_token_indx, + num_accepted_tokens=num_accepted_tokens, + nums_dict=nums_dict, + batch_ptr=batch_ptr, + token_chunk_offset_ptr=token_chunk_offset_ptr, + ) + + +if __name__ == "__main__": + # Apply patch when run directly + patch_gdn_attention() diff --git a/roll/third_party/vllm/worker.py b/roll/third_party/vllm/worker.py index 0348a2172..2aaf46320 100644 --- a/roll/third_party/vllm/worker.py +++ b/roll/third_party/vllm/worker.py @@ -14,8 +14,15 @@ from roll.utils.collective import collective from roll.utils.cuda_ipc_utils import MultiprocessingSerializer from roll.utils.logging import get_logger +from roll.utils.offload_states import clear_memory from roll.utils.send_recv_utils import monkey_patch_torch_reductions, named_tensors_from_bucket +# Apply GDN attention patch at module import time. +# This ensures the patch is applied in EngineCore subprocesses as well, +# since this module is imported when worker_extension_cls is loaded. +from roll.third_party.vllm.gdn_patcher import patch_gdn_attention +patch_gdn_attention() + logger = get_logger() @@ -71,7 +78,19 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): from roll.third_party.vllm.vllm_utils import patch_vllm_moe_model_weight_loader patch_vllm_moe_model_weight_loader(self.model_runner.model) - self.model_runner.model.load_weights(weights=weights) + + # Convert to list for multiple iterations (draft model also needs the same weights) + weights_list = list(weights) + + # Update target model (skips mtp.* prefixed weights via skip_prefixes) + self.model_runner.model.load_weights(weights=weights_list) + + # Update drafter model (MTP/EAGLE) if exists + # Drafter models like EagleProposer and DraftModelProposer have a model attribute + # that needs to be updated separately with the same weights + if hasattr(self.model_runner, "drafter") and hasattr(self.model_runner.drafter, "model"): + logger.info("Updating drafter (MTP/EAGLE) model weights...") + self.model_runner.drafter.model.load_weights(weights=weights_list) def load_states(self): self.reload_model() @@ -99,8 +118,7 @@ def offload_states(self, level): self.kv_cache_loaded = False if hasattr(self, "recv_manager"): self.recv_manager.clear() - gc.collect() - current_platform.empty_cache() + clear_memory() def setup_collective_group(self, master_address, master_port, rank_offset, world_size, group_name, backend): group_rank = self.rank + rank_offset diff --git a/roll/utils/checkpoint_manager.py b/roll/utils/checkpoint_manager.py index 19b32225a..87dafb6f3 100644 --- a/roll/utils/checkpoint_manager.py +++ b/roll/utils/checkpoint_manager.py @@ -52,6 +52,19 @@ def wrapper(model_name_or_path: str, local_dir: Optional[str] = None): return cached_path return wrapper +def get_latest_ckpt(checkpoint_config): + logger.info(f"checkpoint_config: {checkpoint_config}") + if not checkpoint_config: + return None + upload_type = checkpoint_config.get("type", None) + if not upload_type: + return None + if upload_type not in uploader_registry: + raise ValueError(f"Unknown tracker name: {upload_type}, total registered trackers: {uploader_registry.keys()}") + uploader_cls = uploader_registry[upload_type] + uploader = uploader_cls(**checkpoint_config) + return uploader.get_latest_ckpt() + @model_path_cache def download_model(model_name_or_path: str, local_dir: Optional[str] = None): diff --git a/roll/utils/deepspeed_utils.py b/roll/utils/deepspeed_utils.py deleted file mode 100644 index fe7887748..000000000 --- a/roll/utils/deepspeed_utils.py +++ /dev/null @@ -1,43 +0,0 @@ -from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus -""" -ref: https://github.com/OpenRLHF/OpenRLHF/blob/494850f50342ed38d5ae76ef45a3207f3523b582/openrlhf/utils/deepspeed/deepspeed_utils.py#L104 -""" - -def get_optimizer_grouped_parameters( - model, - weight_decay, - no_decay_name_list=["bias", "layer_norm.weight", "layernorm.weight", "norm.weight", "ln_f.weight"], -): - """ - 用于生成优化器的参数组列表，参数分组，为了在训练时可以对不同的参数应用不同的权重衰减策略 - Args: - model: - weight_decay: - no_decay_name_list: - - Returns: - - """ - optimizer_grouped_parameters = [ - { - "params": [ - p - for n, p in model.named_parameters() - if (not any(nd in n for nd in no_decay_name_list) and p.requires_grad) - ], - "weight_decay": weight_decay, - }, - { - "params": [ - p - for n, p in model.named_parameters() - if (any(nd in n for nd in no_decay_name_list) and p.requires_grad) - ], - "weight_decay": 0.0, - }, - ] - return optimizer_grouped_parameters - - -def _z3_params_to_fetch(param_list): - return [p for p in param_list if hasattr(p, "ds_id") and p.ds_status == ZeroParamStatus.NOT_AVAILABLE] diff --git a/roll/utils/fsdp_utils.py b/roll/utils/fsdp_utils.py index 16273da7f..7acf15002 100644 --- a/roll/utils/fsdp_utils.py +++ b/roll/utils/fsdp_utils.py @@ -2,23 +2,56 @@ import dataclasses from abc import ABC from contextlib import contextmanager +from typing import Callable import torch import torch.distributed as dist import torch.nn as nn from torch.distributed.fsdp import CPUOffloadPolicy, fully_shard -from torch.distributed.tensor import Shard +from torch.distributed.tensor import Shard, DTensor +import torch.nn.functional as F +from packaging import version +import transformers from roll.models.model_providers import _is_moe_config from roll.platforms import current_platform +from roll.utils.logging import get_logger + +logger = get_logger() try: from torch.distributed.device_mesh import DeviceMesh except ImportError: DeviceMesh = None +# Optional Triton dependency: enables GPU-side fill_indices kernel for EP permutation. +# Falls back to CPU implementation transparently when triton is unavailable. +try: + import triton + import triton.language as tl + _TRITON_AVAILABLE = True +except ImportError: + triton = None + tl = None + _TRITON_AVAILABLE = False + fully_shard_module = torch.distributed.fsdp._fully_shard._fully_shard +# TOKEN_GROUP_ALIGN_SIZE_M is set as 8 in torchtitan while quantization is not used +TOKEN_GROUP_ALIGN_SIZE_M = 8 + + +use_grouped_mm = False + +def set_use_grouped_mm(moe_use_grouped_mm): + global use_grouped_mm + use_grouped_mm = moe_use_grouped_mm + + +def get_use_grouped_mm(): + global use_grouped_mm + return use_grouped_mm + @contextmanager def maybe_patch_fsdp_module(model): @@ -70,6 +103,26 @@ def shard_placement_fn(param): return shard_placement_fn +def get_shard_placement_fn_ep(efsdp_size): + """ + Choose the dimension that can divide fsdp_size to avoid padding + Reference: https://github.com/volcengine/verl/blob/main/verl/utils/fsdp_utils.py + + """ + + def shard_placement_fn(param): + if isinstance(param, DTensor): + shape = list(param.data._local_tensor.shape) + else: + shape = list(param.shape) + for i in range(len(shape)): + if shape[i] % efsdp_size == 0: + return Shard(i) + return Shard(0) + + return shard_placement_fn + + def _clone_mp_policy(mp_policy, **overrides): if mp_policy is None: return None @@ -109,7 +162,7 @@ def _fsdp_kwargs_for_module(fsdp_kwargs: dict, module: nn.Module) -> dict: return new_kwargs -def apply_fsdp2(model, fsdp_kwargs, config, is_lora=False): +def apply_fsdp2(model, fsdp_kwargs, moe_fsdp_kwargs, config, is_lora=False): """ model: AutoModelForCausalLM @@ -122,6 +175,9 @@ def apply_fsdp2(model, fsdp_kwargs, config, is_lora=False): model_cfg = getattr(model, "config", None) is_moe = _is_moe_config(model_cfg) apply_expert_patch = bool(config.get("apply_expert_patch", False)) + if version.parse(transformers.__version__) >= version.parse("5.2.0"): + apply_expert_patch = False + logger.warning("[apply_fsdp2] apply_expert_patch was set as False automatically, because transformers>=5.2.0 uses fused experts") if is_moe and apply_expert_patch: from transformers.models.qwen3_moe.modeling_qwen3_moe import Qwen3MoeSparseMoeBlock @@ -181,8 +237,8 @@ def _already_fully_sharded(mod: nn.Module) -> bool: # Experts are invoked conditionally per-rank (based on routing), # so wrapping `experts.*` as separate FSDP modules can deadlock collectives when # different ranks activate different experts. Therefor we only wrap experts - # if we apply the expert patch. - if is_moe and config.get("apply_expert_patch", False): + # if we apply the expert patch or expert parallelism is enabled. + if is_moe and (config.get("apply_expert_patch", False) or config.get("ep_size", 1) > 1): moe_block = config.get("wrap_policy", {}).get("moe_experts", None) if isinstance(moe_block, str): moe_block = [moe_block] @@ -235,7 +291,7 @@ def _wrap_once(mod: nn.Module, kwargs: dict): # 3. MoE for idx, module in enumerate(moe_modules): - _wrap_once(module, fsdp_kwargs) + _wrap_once(module, moe_fsdp_kwargs) # 4. Transformers Layers for idx, module in enumerate(modules): @@ -256,6 +312,7 @@ def fsdp2_load_full_state_dict( full_state: dict, device_mesh=None, cpu_offload=None, + ep_enabled=False, ): """ Reference: https://github1s.com/volcengine/verl/blob/main/verl/utils/fsdp_utils.py @@ -272,18 +329,21 @@ def fsdp2_load_full_state_dict( device_id = current_platform.current_device() - if dist.get_rank() == 0: - model = model.to(device=device_id, non_blocking=True) + if not ep_enabled: + if dist.get_rank() == 0: + model = model.to(device=device_id, non_blocking=True) + else: + model = model.to_empty(device=device_id) + + cpu_offload = cpu_offload is not None + options = StateDictOptions( + full_state_dict=True, + cpu_offload=cpu_offload, + broadcast_from_rank0=True, + ) + set_model_state_dict(model, full_state, options=options) else: - model = model.to_empty(device=device_id) - - cpu_offload = cpu_offload is not None - options = StateDictOptions( - full_state_dict=True, - cpu_offload=cpu_offload, - broadcast_from_rank0=True, - ) - set_model_state_dict(model, full_state, options=options) + model = model.to(device=device_id, non_blocking=True) # rotary_emb is not in state_dict, so we need to broadcast it manually for name, buf in model.named_buffers(): @@ -294,3 +354,333 @@ def fsdp2_load_full_state_dict( model.to("cpu", non_blocking=True) for buf in model.buffers(): buf.data = buf.data.to(device_id) + +def _permute(x, num_tokens_per_expert, ep_degree, num_local_experts): + """ + Reference: https://github.com/pytorch/torchtitan/blob/73a0e6979dd10b6b1904098eb3c8f62c18ab87ce/torchtitan/models/moe/utils.py #42 #v0.2.2 + """ + global TOKEN_GROUP_ALIGN_SIZE_M + x_padded_per_expert = x.shape[0] + num_local_experts * TOKEN_GROUP_ALIGN_SIZE_M + padded_max_len = _round_up(x_padded_per_expert, TOKEN_GROUP_ALIGN_SIZE_M) + with torch.no_grad(): + (permuted_indices, num_tokens_per_expert, _offsets,) = generate_permute_indices( + num_tokens_per_expert, + num_local_experts, + ep_degree, + padded_max_len, + TOKEN_GROUP_ALIGN_SIZE_M, + ) + + x = torch.vstack((x, x.new_zeros((x.shape[-1])))) + input_shape = x.shape + x = x[permuted_indices, :] + + return input_shape, x, permuted_indices, num_tokens_per_expert + + +def _unpermute(out, input_shape, permuted_indices): + """ + Reference: https://github.com/pytorch/torchtitan/blob/73a0e6979dd10b6b1904098eb3c8f62c18ab87ce/torchtitan/models/moe/utils.py #62 #v0.2.2 + """ + out_unpermuted = out.new_empty(input_shape) + out_unpermuted[permuted_indices, :] = out + out = out_unpermuted[:-1] + return out + +def _round_up(x: int, y: int) -> int: + """ + Round up x to the nearest multiple of y. + Reference: https://github.com/pytorch/torchtitan/blob/73a0e6979dd10b6b1904098eb3c8f62c18ab87ce/torchtitan/tools/utils.py #228 #v0.2.2 + """ + x_ceil_div_y = (x + y - 1) // y + return x_ceil_div_y * y + +def generate_permute_indices( + tokens_per_expert_group: torch.Tensor, + experts_per_rank: int, + num_ranks: int, + max_len: int, + alignment: int, + use_cpu: bool = False, +): + """ + Prepare permutation indices and the number of tokens for each expert. + + Args: + tokens_per_expert_group: number of tokens for each expert from all ranks. + experts_per_rank: number of experts per rank. + num_ranks: number of ranks. + max_len: maximum length of the output index vector. + alignment: alignment for each returned element in `m_sizes` and padding min for zero token experts. + use_cpu: whether to use CPU implementation. + + + Returns: + permuted_indices: Tensor of indices that map original token order to the expert-grouped order. + m_sizes: aligned number of tokens for each expert (padded to alignment boundary). + m_offsets: Cumulative sum of m_sizes. The exclusive ending position for each expert's tokens. + + Explanatory details: + `tokens_per_expert_group` is of shape (num_ranks * experts_per_rank,), for example: + From: | rank 0 | rank 1 | + To: | E0 | E1 | E2 | E3 | E0 | E1 | E2 | E3 | + | 4 | 2 | 1 | 3 | 1 | 2 | 3 | 4 | + + Reference: https://github.com/pytorch/torchtitan/blob/73a0e6979dd10b6b1904098eb3c8f62c18ab87ce/torchtitan/models/moe/kernels.py #143 #v0.2.2 + """ + + # prefix sum to get start index of each expert (parallel scan kernel in future?) + start_index_values = ( + torch.cumsum(tokens_per_expert_group, 0) - tokens_per_expert_group + ) + + # total tokens for each expert (sum over ranks) + total_tokens_per_expert = tokens_per_expert_group.view(num_ranks, -1).sum(0) + + # pad out empty experts to alignment requirement + total_tokens_per_expert = torch.clamp_min(total_tokens_per_expert, alignment) + + # align the chunk sizes (cdiv) + m_sizes = ((total_tokens_per_expert + alignment - 1) // alignment * alignment).to( + torch.int32 + ) + + # additional prefix sum to get write offset of each expert in permuted_indices + # write offsets is per local expert, not global + m_offsets = torch.cumsum(m_sizes, 0) + write_offsets = m_offsets - m_sizes + + if use_cpu or not _TRITON_AVAILABLE: + permuted_indices = fill_indices_cpu( + tokens_per_expert_group, + start_index_values, + write_offsets, + experts_per_rank, + num_ranks, + max_len, + ) + else: + permuted_indices = fill_indices_wrapper( + tokens_per_expert_group, + start_index_values, + write_offsets, + experts_per_rank, + num_ranks, + max_len, + ) + + return permuted_indices, m_sizes, m_offsets.to(torch.int32) + +# ============================================================================ +# Triton GPU implementation of fill_indices. +# Avoids hundreds of D2H .item() syncs per MoE layer that the CPU fallback +# incurs. The kernel and wrapper below are guarded by `_TRITON_AVAILABLE` and +# only defined when triton can be imported, so environments without triton +# (e.g. CPU-only debug runs) keep working through the CPU fallback path. +# Reference: https://github.com/pytorch/torchtitan/blob/main/torchtitan/models/moe/kernels.py +# ============================================================================ +if _TRITON_AVAILABLE: + + @triton.jit + def _fill_indices_kernel( + tokens_per_expert_group_ptr, + start_index_values_ptr, + write_offsets_ptr, + output_ptr, + experts_per_rank: tl.constexpr, + num_ranks: tl.constexpr, + BLOCK_SIZE: tl.constexpr, # Number of threads per block + ): + pid = tl.program_id(axis=0) + num_programs = tl.num_programs(axis=0) + + # map programs (blocks) to the experts and loop (grid stride) if needed + for expert_id in range(pid, experts_per_rank, num_programs): + # read this expert's write offset + write_offset = tl.load(write_offsets_ptr + expert_id) + + for r in range(num_ranks): + # index into tokens_per_expert_group array + i = r * experts_per_rank + expert_id + + # load start index and number of tokens for this expert-rank pair + start_index = tl.load(start_index_values_ptr + i) + length = tl.load(tokens_per_expert_group_ptr + i) + + # each thread in block processes tokens in parallel + offsets = tl.arange(0, BLOCK_SIZE) + + # tokens are processed in chunks of BLOCK_SIZE + for chunk_start in range(0, length, BLOCK_SIZE): + chunk_offsets = chunk_start + offsets + mask = chunk_offsets < length + values = start_index + chunk_offsets + dest_indices = write_offset + chunk_offsets + tl.store(output_ptr + dest_indices, values, mask=mask) + + # update write offset for next rank + write_offset += length + +def fill_indices_wrapper( + tokens_per_expert_group: torch.Tensor, + start_index_values: torch.Tensor, + write_offsets: torch.Tensor, + experts_per_rank: int, + num_ranks: int, + max_len: int, + block_size: int = 128, + max_blocks: int = 1024, # cap on total number of blocks to launch +): + """ + GPU implementation of fill_indices via a Triton kernel. All inputs must be + on the same CUDA device; the returned tensor lives on that device too. + + Reference: https://github.com/pytorch/torchtitan/blob/main/torchtitan/models/moe/kernels.py + """ + assert _TRITON_AVAILABLE, "fill_indices_wrapper requires triton to be installed." + # preallocate output on the same device as the inputs (GPU) + permuted_indices = torch.full( + (max_len,), -1, dtype=torch.int32, device=tokens_per_expert_group.device, + ) + + # one block per local expert (capped to avoid launching huge grids on + # configurations with very many experts per rank) + num_blocks = min(experts_per_rank, max_blocks) + grid = (num_blocks,) + + _fill_indices_kernel[grid]( + tokens_per_expert_group, + start_index_values, + write_offsets, + permuted_indices, + experts_per_rank, + num_ranks, + BLOCK_SIZE=block_size, + ) + return permuted_indices + +def fill_indices_cpu( + tokens_per_expert_group: torch.Tensor, + start_index_values: torch.Tensor, + write_offsets: torch.Tensor, + experts_per_rank: int, + num_ranks: int, + max_len: int, +): + """ + Reference: https://github.com/pytorch/torchtitan/blob/73a0e6979dd10b6b1904098eb3c8f62c18ab87ce/torchtitan/models/moe/kernels.py #106 #v0.2.2 + """ + # We need to preallocate the output - we ignore device and force it on cpu + # device = tokens_per_expert_group.device + permuted_indices = torch.full( + (max_len,), + -1, + dtype=torch.int32, + ) # device=device) + # Fill the permuted indices + # For each local expert + for e in range(experts_per_rank): + write_start = write_offsets[e].item() + # For each remote rank + for r in range(num_ranks): + i = r * experts_per_rank + e + start_index = start_index_values[i].item() + length = tokens_per_expert_group[i].item() + # Fill in the indices + if length > 0: + end_idx = min(write_start + length, max_len) + permuted_indices[write_start:end_idx] = torch.arange( + start_index, + start_index + (end_idx - write_start), + dtype=torch.int32, + # device=device, + ) + write_start += length + return permuted_indices + +def register_experts_forward_in_ExpertsInterface(): + from transformers.integrations.moe import ExpertsInterface + ExpertsInterface.register("ep", ep_experts_forward) + +def ep_experts_forward( + self, + x: torch.Tensor, + num_tokens_per_expert: torch.Tensor, +) -> torch.Tensor: + """ + Reference: https://github.com/pytorch/torchtitan/blob/73a0e6979dd10b6b1904098eb3c8f62c18ab87ce/torchtitan/models/moe/moe.py #L148 #v0.2.2 + """ + if isinstance(self.gate_up_proj, DTensor): + # Convert parameters from DTensors to plain Tensors, to work with + # dynamic-shape inputs in EP which cannot be easily expressed as DTensors. + gate_up_proj = self.gate_up_proj.to_local() + down_proj = self.down_proj.to_local() + else: + gate_up_proj = self.gate_up_proj + down_proj = self.down_proj + + use_grouped_mm = get_use_grouped_mm() + if use_grouped_mm: + return _run_experts_grouped_mm(gate_up_proj, down_proj, x, num_tokens_per_expert, self.act_fn) + else: + return _run_experts_for_loop(gate_up_proj, down_proj, x, num_tokens_per_expert, self.act_fn) + +def _run_experts_grouped_mm( + gate_up_proj: torch.Tensor, + down_proj: torch.Tensor, + x: torch.Tensor, + num_tokens_per_expert: torch.Tensor, + act_fn: Callable +) -> torch.Tensor: + """ + Reference: https://github.com/pytorch/torchtitan/blob/73a0e6979dd10b6b1904098eb3c8f62c18ab87ce/torchtitan/models/moe/moe.py #L113 #v0.2.2 + """ + offsets = torch.cumsum(num_tokens_per_expert, dim=0, dtype=torch.int32) + + gate, up = torch._grouped_mm( + x.bfloat16(), gate_up_proj.bfloat16().transpose(-2, -1), offs=offsets + ).chunk(2, dim=-1) + h = act_fn(gate) * up + out = torch._grouped_mm( + h, down_proj.bfloat16().transpose(-2, -1), offs=offsets + ).type_as(x) + + return out + +def _run_experts_for_loop( + gate_up_proj: torch.Tensor, + down_proj: torch.Tensor, + x: torch.Tensor, + num_tokens_per_expert: torch.Tensor, + act_fn: Callable, +) -> torch.Tensor: + """ + Reference: https://github.com/pytorch/torchtitan/blob/73a0e6979dd10b6b1904098eb3c8f62c18ab87ce/torchtitan/models/moe/moe.py #L78 #v0.2.2 + """ + # NOTE: this would incur a synchronization between device and host + num_tokens_per_expert_list = num_tokens_per_expert.tolist() + + # side-effect code due to the usage of generate_permute_indices + num_padding = x.shape[0] - sum(num_tokens_per_expert_list) + + # a tuple of tensors indexed by experts + # each with shape (tokens_per_expert(varying), dim) + x_splits = torch.split( + x[: sum(num_tokens_per_expert_list)], + split_size_or_sections=num_tokens_per_expert_list, + dim=0, + ) + out_experts_splits = [] + for expert_idx, x_expert in enumerate(x_splits): + gate, up = nn.functional.linear(x_expert, gate_up_proj[expert_idx]).chunk(2, dim=-1) + h = act_fn(gate) * up + h = nn.functional.linear(h, down_proj[expert_idx]) + + # h shape (tokens_per_expert(varying), dim) + out_experts_splits.append(h) + out = torch.cat(out_experts_splits, dim=0) + + # side-effect code due to the usage of generate_permute_indices + out = torch.vstack((out, out.new_zeros((num_padding, out.shape[-1])))) + + return out \ No newline at end of file diff --git a/roll/utils/functionals.py b/roll/utils/functionals.py index 6e251a092..66ebe9b34 100644 --- a/roll/utils/functionals.py +++ b/roll/utils/functionals.py @@ -8,7 +8,8 @@ import enum import traceback import heapq -from typing import Dict, List, Optional, Tuple, Union +import math +from typing import Callable, Dict, List, Optional, Tuple, Union import numpy as np import torch @@ -460,9 +461,11 @@ def pad_to_length(tensor: torch.Tensor, length, pad_value, dim=-1): else: pad_size = list(tensor.shape) pad_size[dim] = length - tensor.size(dim) - return torch.cat( - [tensor, pad_value * torch.ones(*pad_size, dtype=tensor.dtype, device=tensor.device)], dim=dim - ) + # Use torch.full to strictly honor tensor.dtype. `pad_value * torch.ones(..., dtype=bool)` + # would type-promote a bool mask to int64, causing TransferQueue dtype mismatch when + # some samples need padding and others don't. + pad_block = torch.full(pad_size, pad_value, dtype=tensor.dtype, device=tensor.device) + return torch.cat([tensor, pad_block], dim=dim) def concatenate_input_and_output(input_ids, output_ids, num_return_sequences): @@ -770,6 +773,315 @@ def apply_kl_penalty(data: "DataProto", kl_ctrl: AdaptiveKLController, kl_penalt return data, metrics +def build_domain_to_teacher_names( + domain_values: np.ndarray, + domain_2_tag: Dict[str, set], + tag_to_teacher: Dict[str, List[str]], + all_tag_teachers: List[str], +) -> Dict[str, List[str]]: + """Build mapping from domain -> list of teacher names that should handle it. + + Args: + domain_values: Array of domain strings for each sample in the batch. + domain_2_tag: Mapping from domain -> set of tags. + tag_to_teacher: Mapping from tag -> list of teacher names. + all_tag_teachers: Teacher names that handle all tags (empty tag_included). + + Returns: + Dict mapping domain -> list of teacher names responsible for that domain. + """ + domain_to_teacher_names: Dict[str, List[str]] = {} + for domain in set(domain_values): + teachers_for_domain = list(all_tag_teachers) + tags = domain_2_tag.get(domain, {domain}) + for tag in tags: + teachers_for_domain.extend(tag_to_teacher.get(tag, [])) + domain_to_teacher_names[domain] = list(set(teachers_for_domain)) + return domain_to_teacher_names + + +def build_domain_routing_context( + batch: "DataProto", + pipeline_config, + domain_key: str = "domain", +) -> Tuple[Optional[Dict[str, List[str]]], Optional[np.ndarray]]: + """Build domain routing context from batch and pipeline config. + + Returns: + (domain_to_teacher_names, domain_values) — both None if no routing needed. + """ + if not pipeline_config.needs_teacher_routing: + return None, None + domain_2_tag = getattr(pipeline_config, "domain_2_tag", None) or {} + domain_values = batch.non_tensor_batch.get(domain_key, None) + if domain_values is None: + domain_values = np.array(["default"] * batch.batch.batch_size[0]) + domain_to_teacher_names = build_domain_to_teacher_names( + domain_values=domain_values, + domain_2_tag=domain_2_tag, + tag_to_teacher=pipeline_config.tag_to_teacher_names, + all_tag_teachers=pipeline_config.teacher_names_for_all_tags, + ) + return domain_to_teacher_names, domain_values + + +def extract_real_results_from_routed( + ref_result: "DataProto", + sub_batch: "DataProto", + name: str, +) -> "DataProto": + """Extract real (non-padded) samples from ref_result in original order. + + Must be called INSIDE the compute_log_probs closure, where the final + sub_batch (after batch_balance and dynamic_batching_shard reordering) + is available. The tracking fields _position_idx and _is_real_sample + in sub_batch have been reordered along with the data, so they correctly + identify real samples and their original positions. + + After extraction, cleans up the tracking fields from sub_batch. + + Args: + ref_result: DataProto from compute_log_probs, containing ref_log_probs_{name}. + sub_batch: The final sub_batch (after all reorders inside the closure). + name: Teacher name (used as key in ref_result). + + Returns: + ref_result with ref_log_probs_{name} containing only real samples + in the original (pre-reorder) order. + """ + if "_is_real_sample" not in sub_batch.batch: + return ref_result + + is_real = sub_batch.batch["_is_real_sample"] + position_idx = sub_batch.batch["_position_idx"] + + # Select real results (in whatever order sub_batch is currently in) + real_results_reordered = ref_result.batch[f"ref_log_probs_{name}"][is_real] + real_positions = position_idx[is_real] + + # Sort by original position to restore the original order + sort_order = torch.argsort(real_positions) + real_ref_log_probs = real_results_reordered[sort_order] + + # Clean up tracking fields + sub_batch.batch.pop("_position_idx", None) + sub_batch.batch.pop("_is_real_sample", None) + + # Update ref_result to contain only real results in original order + ref_result.batch[f"ref_log_probs_{name}"] = real_ref_log_probs + return ref_result + + +def merge_ref_log_probs_compat( + batch: "DataProto", + references: Dict[str, object], + is_routing: bool, +) -> None: + """Set backward-compat ``ref_log_probs`` from per-teacher results. + + In routing mode, merge per-teacher ref_log_probs so each sample gets + the value from the teacher that was responsible for it. + In non-routing mode, just use the first teacher's result. + """ + if is_routing: + ref_log_probs_compat = torch.zeros_like( + batch.batch["response_mask"][:, 1:], dtype=torch.float32 + ) + filled = torch.zeros( + batch.batch.batch_size[0], dtype=torch.bool, device=batch.batch.device + ) + for tname in references: + mask_key = f"ref_log_probs_{tname}_mask" + if mask_key in batch.batch: + teacher_mask = batch.batch[mask_key] + fill_mask = teacher_mask & ~filled + ref_log_probs_compat[fill_mask] = batch.batch[f"ref_log_probs_{tname}"][fill_mask] + filled |= fill_mask + else: + # Teacher without mask computed for all samples + ref_log_probs_compat = batch.batch[f"ref_log_probs_{tname}"].clone() + break + batch.batch["ref_log_probs"] = ref_log_probs_compat + else: + first_name = list(references.keys())[0] + batch.batch["ref_log_probs"] = batch.batch[f"ref_log_probs_{first_name}"] + + +def clusters_have_disjoint_devices(clusters: Union[Dict[str, object], List[object]]) -> bool: + """Check if all clusters have non-overlapping device_mapping. + + Args: + clusters: A dict of name->Cluster or a list of Cluster objects. + Each cluster must have ``worker_config.device_mapping``. + + Returns: + True if no two clusters share any device index. + """ + items = clusters.values() if isinstance(clusters, dict) else clusters + seen_devices: set = set() + for cluster in items: + devices = set(cluster.worker_config.device_mapping or []) + if devices & seen_devices: + return False + seen_devices |= devices + return True + + +def compute_ref_log_probs_with_routing( + batch: "DataProto", + references: Dict[str, object], + pipeline_config, + domain_to_teacher_names: Optional[Dict[str, List[str]]] = None, + domain_values: Optional[np.ndarray] = None, + metrics_fn: Optional[Callable[[dict], None]] = None, +) -> "DataProto": + """Compute ref_log_probs for each teacher, with optional tag-based routing. + + Handles the full pipeline internally: batch_balance, dynamic_batching_shard, + compute_log_probs, rename, reduce_metrics, and extract_real_results_from_routed. + + When ``domain_to_teacher_names`` is None (non-routing mode), every teacher + processes all samples. When provided, only samples whose domain maps to a + teacher are sent to that teacher (routing mode). + + When multiple teachers have non-overlapping device_mapping, their forward + passes are launched in parallel using threads to reduce wall time. + + Args: + batch: The full DataProto batch. + references: Dict mapping teacher name -> Cluster object. + pipeline_config: Pipeline config. + domain_to_teacher_names: Mapping from domain -> list of teacher names. + None means non-routing (all samples to all teachers). + domain_values: Array of domain strings for each sample. + None means non-routing. + metrics_fn: Callable that receives flat metric dicts (e.g. ``metrics_mgr.add_metrics`` + or ``dict.update``). + + Returns: + The batch with ref_log_probs_{name}, ref_log_probs_{name}_mask, and + backward-compat ref_log_probs set. + """ + from roll.distributed.scheduler.protocol import DataProto, pad_dataproto_to_divisor # noqa: F811 + from roll.utils.dynamic_batching import dynamic_batching_shard + + is_routing = domain_to_teacher_names is not None + parallel = ( + len(references) > 1 + and clusters_have_disjoint_devices(references) + ) + + teacher_tasks: Dict[str, Tuple] = {} + + for name, ref_cluster in references.items(): + if is_routing: + responsible_domains = [ + domain for domain, teachers in domain_to_teacher_names.items() + if name in teachers + ] + + if not responsible_domains or len(domain_values) == 0: + batch.batch[f"ref_log_probs_{name}"] = torch.zeros_like( + batch.batch["response_mask"][:, 1:], dtype=torch.float32 + ) + batch.batch[f"ref_log_probs_{name}_mask"] = torch.zeros( + batch.batch.batch_size[0], dtype=torch.bool + ) + continue + + domain_mask = torch.tensor( + [d in responsible_domains for d in domain_values], + device=batch.batch.device, + ) + + if domain_mask.sum() == 0: + batch.batch[f"ref_log_probs_{name}"] = torch.zeros_like( + batch.batch["response_mask"][:, 1:], dtype=torch.float32 + ) + batch.batch[f"ref_log_probs_{name}_mask"] = torch.zeros( + batch.batch.batch_size[0], dtype=torch.bool + ) + continue + else: + domain_mask = torch.ones(batch.batch.batch_size[0], dtype=torch.bool, device=batch.batch.device) + + sub_batch = batch.select_idxs(domain_mask) + sub_batch.meta_info = dict(batch.meta_info) + num_real = sub_batch.batch.batch_size[0] + + sub_batch.batch["_position_idx"] = torch.arange(num_real, device=sub_batch.batch.device) + sub_batch.batch["_is_real_sample"] = torch.ones(num_real, dtype=torch.bool, device=sub_batch.batch.device) + + dp_size = ref_cluster.dp_size + if dp_size > 0: + sub_batch, pad_size = pad_dataproto_to_divisor(sub_batch, dp_size) + if pad_size > 0: + sub_batch.batch["_is_real_sample"][-pad_size:] = False + + teacher_tasks[name] = (ref_cluster, sub_batch, domain_mask) + + def _run_teacher(name: str, ref_cluster, sub_batch) -> Tuple["DataProto", List[dict]]: + """Run a single teacher's forward pass. Returns (result, collected_metrics). + + Metrics are collected locally to avoid thread-safety issues when running + multiple teachers in parallel threads. + """ + collected_metrics: List[dict] = [] + worker_config = ref_cluster.worker_config + assert len(sub_batch) % ref_cluster.dp_size == 0, \ + f"sub_batch size {len(sub_batch)} not divisible by dp_size {ref_cluster.dp_size}" + batch_balance(sub_batch, dp_size=ref_cluster.dp_size, minibatch_size=len(sub_batch)) + if worker_config.use_dynamic_batching_in_infer: + sub_batch, db_metrics = dynamic_batching_shard( + sub_batch, + ref_cluster.dp_size, + worker_config.max_tokens_per_microbatch_in_infer, + worker_config.sequence_length_round_in_infer, + worker_config.strategy_args.strategy_config.get("pipeline_model_parallel_size", 1), + worker_config.strategy_args.strategy_config.get("virtual_pipeline_model_parallel_size", None), + f"reference-{name}/compute_log_probs", + ) + collected_metrics.append(db_metrics) + refs = ref_cluster.compute_log_probs(sub_batch, blocking=False) + ref_result = DataProto.materialize_concat(data_refs=refs) + ref_result.rename(old_keys="log_probs", new_keys=f"ref_log_probs_{name}") + collected_metrics.append(reduce_metrics(ref_result.meta_info.pop("metrics", {}))) + ref_result = extract_real_results_from_routed(ref_result, sub_batch, name) + return ref_result, collected_metrics + + def _scatter_result(name: str, ref_result: "DataProto", domain_mask: torch.Tensor) -> None: + real_ref_log_probs = ref_result.batch[f"ref_log_probs_{name}"] + full_ref_log_probs = torch.zeros_like( + batch.batch["response_mask"][:, 1:], dtype=torch.float32 + ) + full_ref_log_probs[domain_mask] = real_ref_log_probs + batch.batch[f"ref_log_probs_{name}"] = full_ref_log_probs + batch.batch[f"ref_log_probs_{name}_mask"] = domain_mask + + if parallel and teacher_tasks: + from concurrent.futures import ThreadPoolExecutor + futures = {} + with ThreadPoolExecutor(max_workers=len(teacher_tasks)) as executor: + for name, (ref_cluster, sub_batch, _) in teacher_tasks.items(): + futures[name] = executor.submit(_run_teacher, name, ref_cluster, sub_batch) + for name, future in futures.items(): + ref_result, collected_metrics = future.result() + _scatter_result(name, ref_result, teacher_tasks[name][2]) + if metrics_fn: + for m in collected_metrics: + metrics_fn(m) + else: + for name, (ref_cluster, sub_batch, domain_mask) in teacher_tasks.items(): + ref_result, collected_metrics = _run_teacher(name, ref_cluster, sub_batch) + _scatter_result(name, ref_result, domain_mask) + if metrics_fn: + for m in collected_metrics: + metrics_fn(m) + + merge_ref_log_probs_compat(batch, references, is_routing) + return batch + + @torch.no_grad() def compute_advantage( data: "DataProto", @@ -792,21 +1104,49 @@ def compute_advantage( # Check OPD config is_pure_opd = getattr(pipeline_config, "is_pure_opd", False) if pipeline_config else False use_opd = getattr(pipeline_config, "use_opd", False) if pipeline_config else False - opd_kl_coef = getattr(pipeline_config, "opd_kl_coef", 1.0) if pipeline_config else 1.0 + reference_configs = getattr(pipeline_config, "reference_configs", {"default": None}) if pipeline_config else {"default": None} - # Compute KL divergence for OPD modes - kld = None if is_pure_opd or use_opd: - kld = compute_approx_kl( - log_probs=data.batch["old_log_probs"] if getattr(pipeline_config, "enable_old_logprobs_recompute", False) else data.batch["infer_logprobs"], - log_probs_base=data.batch["ref_log_probs"], - action_mask=response_mask, - kl_penalty=getattr(pipeline_config, "kl_penalty", "kl"), - ) - - # For pure OPD mode, advantage is directly -kld + log_probs_key = "old_log_probs" if getattr(pipeline_config, "enable_old_logprobs_recompute", False) else "infer_logprobs" + kl_penalty = getattr(pipeline_config, "kl_penalty", "kl") + total_weighted_kld = torch.zeros_like(response_mask, dtype=torch.float32) + for name, ref_cfg in reference_configs.items(): + ref_key = f"ref_log_probs_{name}" + mask_key = f"ref_log_probs_{name}_mask" + + # Check if this teacher has a routing mask (set by pipeline when needs_teacher_routing) + if mask_key in data.batch: + teacher_mask = data.batch[mask_key] # (batch_size,), bool + else: + # No mask = this teacher handles all tags (backward compatible) + teacher_mask = None + + kl_coef_i = ref_cfg.opd_kl_coef if (ref_cfg and ref_cfg.opd_kl_coef is not None) else 1.0 + + if teacher_mask is not None and teacher_mask.sum() < teacher_mask.numel(): + # Only compute KL on routed subset to avoid meaningless computation against zeros + subset_idx = teacher_mask.nonzero(as_tuple=True)[0] + kld_subset = compute_approx_kl( + log_probs=data.batch[log_probs_key][subset_idx], + log_probs_base=data.batch[ref_key][subset_idx], + action_mask=response_mask[subset_idx], + kl_penalty=kl_penalty, + ) + kld_i = torch.zeros_like(response_mask, dtype=torch.float32) + kld_i[subset_idx] = kld_subset + total_weighted_kld += kl_coef_i * kld_i + else: + kld_i = compute_approx_kl( + log_probs=data.batch[log_probs_key], + log_probs_base=data.batch[ref_key], + action_mask=response_mask, + kl_penalty=kl_penalty, + ) + total_weighted_kld += kl_coef_i * kld_i + + # For pure OPD mode, advantage is directly -total_weighted_kld if is_pure_opd: - advantages = -kld + advantages = -total_weighted_kld returns = advantages data.batch["raw_advantages"] = advantages else: @@ -830,9 +1170,9 @@ def compute_advantage( data.batch["raw_advantages"] = advantages - # Apply mixed OPD mode + # Apply mixed OPD mode: advantages = rl_advantages - total_weighted_kld if use_opd: - advantages = advantages - opd_kl_coef * kld + advantages = advantages - total_weighted_kld if whiten_advantages: # TODO whiten过程中是否要考虑response的长度？ @@ -857,6 +1197,7 @@ def postprocess_generate( pad_token_id, fill_eos_token=False, output_logprobs: Optional[list[list[float]]] = None, + routed_experts: Optional[list[torch.Tensor]]=None, pad_to_seq_len=True, ) -> "DataProto": from roll.distributed.scheduler.protocol import DataProto @@ -911,11 +1252,28 @@ def postprocess_generate( assert attention_mask.any(dim=1).all(), f"has all 0 attention_mask, {attention_mask} {input_ids}" first_one = attention_mask.float().argmax(dim=1) new_response_mask = torch.zeros_like(attention_mask) # response mask for cat input_ids + + new_routed_experts = None + expert_dtype = None + if routed_experts is not None: + _, layer_num, topk = routed_experts[0].size() + # Select minimal dtype based on actual max expert index to reduce data size. + max_expert_idx = max(int(re.max().item()) for re in routed_experts) + from roll.third_party.megatron.router_replay_utils import get_routed_experts_dtype + expert_dtype = get_routed_experts_dtype(max_expert_idx) + new_routed_experts = ( + torch.zeros( + [output_batch_size, sequence_length, layer_num, topk], + dtype=expert_dtype, + ) + ) + logprobs = ( torch.zeros([output_batch_size, sequence_length - 1], dtype=torch.float32) if output_logprobs is not None else None ) + for i in range(output_batch_size): shift = first_one[i].item() if shift > 0: @@ -932,6 +1290,12 @@ def postprocess_generate( logprobs[i][prompt_len - 1 : valid_length - 1] = torch.tensor( output_logprobs[i][:response_length], dtype=logprobs.dtype ) + if new_routed_experts is not None: + # new_routed_experts[i, :valid_length - 1] = routed_experts[i] + routed_experts_len = routed_experts[i].size(0) + if routed_experts_len != (valid_length - 1): + logger.info(f"Warning: {routed_experts_len=} != {(valid_length - 1)=}.") + new_routed_experts[i, :routed_experts_len] = routed_experts[i].to(expert_dtype) if position_ids.dim() == 3 and shift > 0: # shift as output to convert to right padding # NOTE: left shift without clear right might lead to unclean values @@ -941,7 +1305,6 @@ def postprocess_generate( # like Qwen2-vl, since extra image_token would be left which might # cause error: Image features and image tokens do not match output_position_ids[i, ..., :-shift] = output_position_ids[i, ..., shift:].clone() - # only clean in VLM(qwen2-vl) to make no effect on LLM if shift > 0 and prompt_length > valid_length: output[i, -shift:] = pad_token_id @@ -969,6 +1332,8 @@ def postprocess_generate( batch["prompt_id"] = prompt_id if logprobs is not None: batch["infer_logprobs"] = logprobs + if new_routed_experts is not None: + batch["routed_experts"] = new_routed_experts return DataProto(batch=batch) @@ -1330,3 +1695,4 @@ def calculate_workload(seq_len_list): metrics.update(global_balance_stats) return metrics + diff --git a/roll/utils/multi_thread_utils.py b/roll/utils/multi_thread_utils.py deleted file mode 100644 index be71e8ca1..000000000 --- a/roll/utils/multi_thread_utils.py +++ /dev/null @@ -1,72 +0,0 @@ -import threading -import time -from collections import defaultdict - - -class ThreadSafeDict: - def __init__(self): - self._data = {} - self._conditions = defaultdict(threading.Condition) - self._lock = threading.Lock() - - def set(self, key, value): - with self._lock: - self._data[key] = value - condition = self._conditions[key] - with condition: - condition.notify_all() - - def get(self, key, timeout=None): - with self._lock: - if key in self._data: - return self._data[key] - condition = self._conditions[key] - - with condition: - waited = condition.wait(timeout=timeout) - with self._lock: - if key in self._data: - return self._data[key] - else: - raise KeyError(f"Key '{key}' was not set within the timeout period.") - - def pop(self, key, timeout=None): - ret = self.get(key, timeout=timeout) - self.remove(key) - return ret - - def contains(self, key): - with self._lock: - return key in self._data - - def remove(self, key): - with self._lock: - if key in self._data: - del self._data[key] - if key in self._conditions: - del self._conditions[key] - - def clear(self): - with self._lock: - self._data.clear() - self._conditions.clear() - - def keys(self): - with self._lock: - return self._data.keys() - - def __len__(self): - with self._lock: - return len(self._data) - - def __contains__(self, key): - return self.contains(key) - - def __getitem__(self, key): - return self.get(key) - - def __setitem__(self, key, value): - self.set(key, value) - - def __delitem__(self, key): - self.remove(key) \ No newline at end of file diff --git a/roll/utils/offload_states.py b/roll/utils/offload_states.py index 5bbf20aab..032680f91 100644 --- a/roll/utils/offload_states.py +++ b/roll/utils/offload_states.py @@ -1,3 +1,4 @@ +import gc from enum import Enum from typing import List, Tuple, Union @@ -7,6 +8,20 @@ from roll.platforms import current_platform +def clear_memory(clear_host_memory: bool = False): + """Clear GPU and CPU memory caches. + + Synchronizes CUDA, runs Python GC, and clears the GPU memory cache. + Optionally releases pinned host memory cached by PyTorch's CachingHostAllocator + back to the OS. + """ + current_platform.synchronize() + gc.collect() + current_platform.empty_cache() + if clear_host_memory: + torch._C._host_emptyCache() + + class OffloadStateType(str, Enum): """ offload/reload需要区分训练和推理阶段，在actor_train/critic用于计算log_probs和values时，只需要reload model_params diff --git a/roll/utils/otel_receiver.py b/roll/utils/otel_receiver.py new file mode 100644 index 000000000..011af12f1 --- /dev/null +++ b/roll/utils/otel_receiver.py @@ -0,0 +1,202 @@ +"""Lightweight Python-based OTLP gRPC receiver for traces. + +A minimal replacement for ``otelcol-contrib`` that receives OTLP spans +via gRPC and writes them to a JSONL file (same format as otelcol). +Requires only ``grpcio`` and ``protobuf`` (already pulled in by the OTel SDK). + +Usage:: + + from roll.utils.otel_receiver import OTLPReceiver + + receiver = OTLPReceiver(port=4317, output_path="traces/traces.jsonl") + receiver.start() + # ... application runs ... + receiver.stop() +""" + +import base64 +import json +import os +import threading +from concurrent import futures + +import grpc +from google.protobuf.json_format import MessageToJson +from opentelemetry.proto.collector.trace.v1 import ( + trace_service_pb2, + trace_service_pb2_grpc, +) + +from roll.utils.logging import get_logger + +logger = get_logger() + +# Jaeger requires numeric span kind values (protobuf enum numbers), +# not the string enum names that MessageToJson produces. +# See: opentelemetry/proto/trace/v1/trace.proto SpanKind enum +_SPAN_KIND_MAP = { + "SPAN_KIND_UNSPECIFIED": 0, + "SPAN_KIND_INTERNAL": 1, + "SPAN_KIND_SERVER": 2, + "SPAN_KIND_CLIENT": 3, + "SPAN_KIND_PRODUCER": 4, + "SPAN_KIND_CONSUMER": 5, +} + + +def _normalize_spans(data): + """Recursively normalize span fields for Jaeger compatibility. + + Performs two transformations required for Jaeger to parse OTLP JSONL: + + 1. **Span kind**: Convert string enum (e.g. ``"SPAN_KIND_INTERNAL"``) to the + corresponding protobuf integer (e.g. ``1``). Jaeger's OTLP-JSON parser + requires the numeric form; string enums cause a conversion error for + large span batches. + + 2. **Byte fields**: ``MessageToJson`` encodes protobuf ``bytes`` fields as + base64 strings (e.g. ``"a+eBrRzJ2u8rRVD7TpsVAg=="``), but the OTLP JSON + spec (and Jaeger) require ``traceId``, ``spanId``, and ``parentSpanId`` + as lowercase hex strings (e.g. ``"a7e781ad1cc9daefc...""``). + """ + if isinstance(data, dict): + result = {} + for key, value in data.items(): + if key == "kind" and isinstance(value, str) and value in _SPAN_KIND_MAP: + result[key] = _SPAN_KIND_MAP[value] + elif key in ("traceId", "spanId", "parentSpanId") and isinstance(value, str): + try: + result[key] = base64.b64decode(value).hex() + except Exception: + result[key] = value # already hex or unknown format — pass through + else: + result[key] = _normalize_spans(value) + return result + elif isinstance(data, list): + return [_normalize_spans(item) for item in data] + else: + return data + + +class _TraceServicer: + """gRPC servicer that receives spans and writes them to JSONL (matching otelcol format).""" + + def __init__(self, output_path: str, max_lines: int = 500): + self.output_path = output_path + self.max_lines = max_lines + self._lock = threading.Lock() + self._line_count = 0 + self._file_index = 0 + self._current_path = self._make_path() + os.makedirs(os.path.dirname(self._current_path) if os.path.dirname(self._current_path) else ".", exist_ok=True) + self._file = open(self._current_path, "a") + logger.info(f"OTLPReceiver: writing traces to {self._current_path}") + + def _make_path(self) -> str: + """Generate file path with index suffix.""" + if self._file_index == 0: + return self.output_path + base, ext = os.path.splitext(self.output_path) + return f"{base}.{self._file_index}{ext}" + + def _rotate_file(self): + """Close current file and open a new one if line limit reached.""" + self._file.close() + self._file_index += 1 + self._line_count = 0 + self._current_path = self._make_path() + self._file = open(self._current_path, "a") + logger.info(f"OTLPReceiver: rotated to {self._current_path}") + + def Export(self, request, context): + """Handle Export RPC call from OTLP client.""" + try: + # Convert to canonical protobuf JSON (camelCase field names, matches otelcol). + # indent=None produces compact single-line output required for JSONL. + json_str = MessageToJson( + request, + preserving_proto_field_name=False, + indent=None, + sort_keys=False, + ) + + # Normalize: hex-encode byte fields and convert span kind to numeric + data_dict = _normalize_spans(json.loads(json_str)) + json_str = json.dumps(data_dict, separators=(",", ":")) + + with self._lock: + self._file.write(json_str + "\n") + self._line_count += 1 + if self._line_count >= self.max_lines: + self._rotate_file() + self._file.flush() + except Exception as exc: + logger.warning(f"OTLPReceiver: failed to write spans: {exc}") + + return trace_service_pb2.ExportTraceServiceResponse( + partial_success=trace_service_pb2.ExportTracePartialSuccess( + rejected_spans=0, + error_message="", + ) + ) + + def close(self): + """Close the output file.""" + with self._lock: + if self._file: + self._file.close() + self._file = None + + +class OTLPReceiver: + """Lightweight OTLP gRPC receiver for traces. + + Receives spans via OTLP gRPC protocol and writes them to a JSONL file + (same format as otelcol's file exporter). Can be used as a drop-in + replacement for ``otelcol-contrib`` in development/testing scenarios. + """ + + def __init__(self, port: int, output_path: str, max_lines: int = 500): + """ + Args: + port: gRPC port to listen on. + output_path: Path to write JSONL trace data. + max_lines: Maximum number of lines per file before rotating (default: 500). + """ + self.port = port + self.output_path = output_path + self.max_lines = max_lines + self._server = None + self._servicer = None + + def start(self) -> None: + """Start the OTLP gRPC receiver.""" + self._servicer = _TraceServicer(self.output_path, self.max_lines) + self._server = grpc.server(futures.ThreadPoolExecutor(max_workers=10)) + trace_service_pb2_grpc.add_TraceServiceServicer_to_server( + self._servicer, self._server + ) + self._server.add_insecure_port(f"[::]:{self.port}") + self._server.start() + logger.info(f"Python OTLPReceiver started on port {self.port}") + + def stop(self, grace: float = 2.0) -> None: + """Stop the OTLP gRPC receiver. + + Args: + grace: Grace period in seconds for pending RPCs to complete. + """ + if self._server: + self._server.stop(grace) + self._server = None + if self._servicer: + self._servicer.close() + self._servicer = None + logger.info("Python OTLPReceiver stopped.") + + def __enter__(self): + self.start() + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + self.stop() diff --git a/roll/utils/qwen_vl_utils.py b/roll/utils/qwen_vl_utils.py new file mode 100644 index 000000000..5aed13d44 --- /dev/null +++ b/roll/utils/qwen_vl_utils.py @@ -0,0 +1,552 @@ +""" +ref: https://github.com/QwenLM/Qwen3-VL/blob/main/qwen-vl-utils/src/qwen_vl_utils/vision_process.py + +add `video_resize_chunk_frames` in `fetch_video` reducing memory usage when loading long videos +""" +import base64 +import copy +import logging +import math +import os +import sys +import time +import warnings +from functools import lru_cache +from io import BytesIO +from typing import Optional, Union, Tuple, List, Any, Dict +from concurrent.futures import ThreadPoolExecutor + +import requests +import torch +import torchvision +from packaging import version +from PIL import Image +import numpy as np +from torchvision import io, transforms +from torchvision.transforms import InterpolationMode + + +MAX_RATIO = 200 +SPATIAL_MERGE_SIZE = 2 +IMAGE_MIN_TOKEN_NUM = 4 +IMAGE_MAX_TOKEN_NUM = 16384 +VIDEO_MIN_TOKEN_NUM = 128 +VIDEO_MAX_TOKEN_NUM = 768 + +FPS = 2.0 +FRAME_FACTOR = 2 +FPS_MIN_FRAMES = 4 +FPS_MAX_FRAMES = 768 +MAX_NUM_WORKERS_FETCH_VIDEO = 8 + +MODEL_SEQ_LEN = int(float(os.environ.get('MODEL_SEQ_LEN', 128000))) +logger = logging.getLogger(__name__) + + +def round_by_factor(number: int, factor: int) -> int: + """Returns the closest integer to 'number' that is divisible by 'factor'.""" + return round(number / factor) * factor + + +def ceil_by_factor(number: int, factor: int) -> int: + """Returns the smallest integer greater than or equal to 'number' that is divisible by 'factor'.""" + return math.ceil(number / factor) * factor + + +def floor_by_factor(number: int, factor: int) -> int: + """Returns the largest integer less than or equal to 'number' that is divisible by 'factor'.""" + return math.floor(number / factor) * factor + + +def smart_resize(height: int, width: int, factor: int, min_pixels: Optional[int] = None, max_pixels: Optional[int] = None) -> Tuple[int, int]: + """ + Rescales the image so that the following conditions are met: + + 1. Both dimensions (height and width) are divisible by 'factor'. + 2. The total number of pixels is within the range ['min_pixels', 'max_pixels']. + 3. The aspect ratio of the image is maintained as closely as possible. + """ + max_pixels = max_pixels if max_pixels is not None else (IMAGE_MAX_TOKEN_NUM * factor ** 2) + min_pixels = min_pixels if min_pixels is not None else (IMAGE_MIN_TOKEN_NUM * factor ** 2) + assert max_pixels >= min_pixels, "The max_pixels of image must be greater than or equal to min_pixels." + if max(height, width) / min(height, width) > MAX_RATIO: + raise ValueError( + f"absolute aspect ratio must be smaller than {MAX_RATIO}, got {max(height, width) / min(height, width)}" + ) + h_bar = max(factor, round_by_factor(height, factor)) + w_bar = max(factor, round_by_factor(width, factor)) + if h_bar * w_bar > max_pixels: + beta = math.sqrt((height * width) / max_pixels) + h_bar = floor_by_factor(height / beta, factor) + w_bar = floor_by_factor(width / beta, factor) + elif h_bar * w_bar < min_pixels: + beta = math.sqrt(min_pixels / (height * width)) + h_bar = ceil_by_factor(height * beta, factor) + w_bar = ceil_by_factor(width * beta, factor) + return h_bar, w_bar + + +def to_rgb(pil_image: Image.Image) -> Image.Image: + if pil_image.mode == 'RGBA': + white_background = Image.new("RGB", pil_image.size, (255, 255, 255)) + white_background.paste(pil_image, mask=pil_image.split()[3]) # Use alpha channel as mask + return white_background + else: + return pil_image.convert("RGB") + + +def fetch_image(ele: Dict[str, Union[str, Image.Image]], image_patch_size: int = 14) -> Image.Image: + if "image" in ele: + image = ele["image"] + else: + image = ele["image_url"] + + image_obj = None + patch_factor = int(image_patch_size * SPATIAL_MERGE_SIZE) + if isinstance(image, Image.Image): + image_obj = image + elif image.startswith("http://") or image.startswith("https://"): + with requests.get(image, stream=True) as response: + response.raise_for_status() + with BytesIO(response.content) as bio: + image_obj = copy.deepcopy(Image.open(bio)) + elif image.startswith("file://"): + image_obj = Image.open(image[7:]) + elif image.startswith("data:image"): + if "base64," in image: + _, base64_data = image.split("base64,", 1) + data = base64.b64decode(base64_data) + with BytesIO(data) as bio: + image_obj = copy.deepcopy(Image.open(bio)) + else: + image_obj = Image.open(image) + if image_obj is None: + raise ValueError(f"Unrecognized image input, support local path, http url, base64 and PIL.Image, got {image}") + image = to_rgb(image_obj) + + ## resize + if "resized_height" in ele and "resized_width" in ele: + resized_height, resized_width = smart_resize( + ele["resized_height"], + ele["resized_width"], + factor=patch_factor, + ) + else: + width, height = image.size + min_pixels = ele.get("min_pixels", IMAGE_MIN_TOKEN_NUM * patch_factor ** 2) + max_pixels = ele.get("max_pixels", IMAGE_MAX_TOKEN_NUM * patch_factor ** 2) + resized_height, resized_width = smart_resize( + height, + width, + factor=patch_factor, + min_pixels=min_pixels, + max_pixels=max_pixels, + ) + image = image.resize((resized_width, resized_height)) + return image + + +def smart_nframes( + ele: Dict[str, Any], + total_frames: int, + video_fps: Union[int, float], +) -> int: + """calculate the number of frames for video used for model inputs. + + Args: + ele (dict): a dict contains the configuration of video. + support either `fps` or `nframes`: + - nframes: the number of frames to extract for model inputs. + - fps: the fps to extract frames for model inputs. + - min_frames: the minimum number of frames of the video, only used when fps is provided. + - max_frames: the maximum number of frames of the video, only used when fps is provided. + total_frames (int): the original total number of frames of the video. + video_fps (int | float): the original fps of the video. + + Raises: + ValueError: nframes should in interval [FRAME_FACTOR, total_frames]. + + Returns: + int: the number of frames for video used for model inputs. + """ + assert not ("fps" in ele and "nframes" in ele), "Only accept either `fps` or `nframes`" + if "nframes" in ele: + nframes = round_by_factor(ele["nframes"], FRAME_FACTOR) + else: + fps = ele.get("fps", FPS) + min_frames = ceil_by_factor(ele.get("min_frames", FPS_MIN_FRAMES), FRAME_FACTOR) + max_frames = floor_by_factor(ele.get("max_frames", min(FPS_MAX_FRAMES, total_frames)), FRAME_FACTOR) + nframes = total_frames / video_fps * fps + if nframes > total_frames: + logger.warning(f"smart_nframes: nframes[{nframes}] > total_frames[{total_frames}]") + nframes = min(min(max(nframes, min_frames), max_frames), total_frames) + nframes = floor_by_factor(nframes, FRAME_FACTOR) + if not (FRAME_FACTOR <= nframes and nframes <= total_frames): + raise ValueError(f"nframes should in interval [{FRAME_FACTOR}, {total_frames}], but got {nframes}.") + return nframes + + +def _read_video_torchvision( + ele: Dict[str, Any], +) -> Tuple[torch.Tensor, float]: + """read video using torchvision.io.read_video + + Args: + ele (dict): a dict contains the configuration of video. + support keys: + - video: the path of video. support "file://", "http://", "https://" and local path. + - video_start: the start time of video. + - video_end: the end time of video. + Returns: + torch.Tensor: the video tensor with shape (T, C, H, W). + """ + video_path = ele["video"] + if version.parse(torchvision.__version__) < version.parse("0.19.0"): + if "http://" in video_path or "https://" in video_path: + warnings.warn("torchvision < 0.19.0 does not support http/https video path, please upgrade to 0.19.0.") + if "file://" in video_path: + video_path = video_path[7:] + st = time.time() + video, audio, info = io.read_video( + video_path, + start_pts=ele.get("video_start", 0.0), + end_pts=ele.get("video_end", None), + pts_unit="sec", + output_format="TCHW", + ) + total_frames, video_fps = video.size(0), info["video_fps"] + logger.info(f"torchvision: {video_path=}, {total_frames=}, {video_fps=}, time={time.time() - st:.3f}s") + nframes = smart_nframes(ele, total_frames=total_frames, video_fps=video_fps) + idx = torch.linspace(0, total_frames - 1, nframes).round().long() + sample_fps = nframes / max(total_frames, 1e-6) * video_fps + video = video[idx] + + video_metadata = dict( + fps=video_fps, + frames_indices=idx, + total_num_frames=total_frames, + video_backend="torchvision", + ) + return video, video_metadata, sample_fps + + +def is_decord_available() -> bool: + import importlib.util + + return importlib.util.find_spec("decord") is not None + + +def calculate_video_frame_range( + ele: Dict[str, Any], + total_frames: int, + video_fps: float, +) -> Tuple[int, int, int]: + """ + Calculate the start and end frame indices based on the given time range. + + Args: + ele (dict): A dictionary containing optional 'video_start' and 'video_end' keys (in seconds). + total_frames (int): Total number of frames in the video. + video_fps (float): Frames per second of the video. + + Returns: + tuple: A tuple containing (start_frame, end_frame, frame_count). + + Raises: + ValueError: If input parameters are invalid or the time range is inconsistent. + """ + # Validate essential parameters + if video_fps <= 0: + raise ValueError("video_fps must be a positive number") + if total_frames <= 0: + raise ValueError("total_frames must be a positive integer") + + # Get start and end time in seconds + video_start = ele.get("video_start", None) + video_end = ele.get("video_end", None) + if video_start is None and video_end is None: + return 0, total_frames - 1, total_frames + + max_duration = total_frames / video_fps + # Process start frame + if video_start is not None: + video_start_clamped = max(0.0, min(video_start, max_duration)) + start_frame = math.ceil(video_start_clamped * video_fps) + else: + start_frame = 0 + # Process end frame + if video_end is not None: + video_end_clamped = max(0.0, min(video_end, max_duration)) + end_frame = math.floor(video_end_clamped * video_fps) + end_frame = min(end_frame, total_frames - 1) + else: + end_frame = total_frames - 1 + + # Validate frame order + if start_frame >= end_frame: + raise ValueError( + f"Invalid time range: Start frame {start_frame} (at {video_start_clamped if video_start is not None else 0}s) " + f"exceeds end frame {end_frame} (at {video_end_clamped if video_end is not None else max_duration}s). " + f"Video duration: {max_duration:.2f}s ({total_frames} frames @ {video_fps}fps)" + ) + + logger.info(f"calculate video frame range: {start_frame=}, {end_frame=}, {total_frames=} from {video_start=}, {video_end=}, {video_fps=:.3f}") + return start_frame, end_frame, end_frame - start_frame + 1 + + +def _read_video_decord( + ele: Dict[str, Any], +) -> Tuple[torch.Tensor, float]: + """read video using decord.VideoReader + + Args: + ele (dict): a dict contains the configuration of video. + support keys: + - video: the path of video. support "file://", "http://", "https://" and local path. + - video_start: the start time of video. + - video_end: the end time of video. + Returns: + torch.Tensor: the video tensor with shape (T, C, H, W). + """ + import decord + video_path = ele["video"] + st = time.time() + vr = decord.VideoReader(video_path) + total_frames, video_fps = len(vr), vr.get_avg_fps() + start_frame, end_frame, total_frames = calculate_video_frame_range( + ele, + total_frames, + video_fps, + ) + nframes = smart_nframes(ele, total_frames=total_frames, video_fps=video_fps) + idx = torch.linspace(start_frame, end_frame, nframes).round().long().tolist() + video = vr.get_batch(idx).asnumpy() + video = torch.tensor(video).permute(0, 3, 1, 2) # Convert to TCHW format + logger.info(f"decord: {video_path=}, {total_frames=}, {video_fps=}, time={time.time() - st:.3f}s") + sample_fps = nframes / max(total_frames, 1e-6) * video_fps + + video_metadata = dict( + fps=video_fps, + frames_indices=idx, + total_num_frames=total_frames, + video_backend="decord", + ) + return video, video_metadata, sample_fps + + +def is_torchcodec_available() -> bool: + import importlib.util + + return importlib.util.find_spec("torchcodec") is not None + + +def _read_video_torchcodec( + ele: Dict[str, Any], +) -> Tuple[torch.Tensor, float]: + """read video using torchcodec.decoders.VideoDecoder + + Args: + ele (dict): a dict contains the configuration of video. + support keys: + - video: the path of video. support "file://", "http://", "https://" and local path. + - video_start: the start time of video. + - video_end: the end time of video. + Returns: + torch.Tensor: the video tensor with shape (T, C, H, W). + """ + from torchcodec.decoders import VideoDecoder + TORCHCODEC_NUM_THREADS = int(os.environ.get('TORCHCODEC_NUM_THREADS', 8)) + logger.info(f"set TORCHCODEC_NUM_THREADS: {TORCHCODEC_NUM_THREADS}") + video_path = ele["video"] + st = time.time() + decoder = VideoDecoder(video_path, num_ffmpeg_threads=TORCHCODEC_NUM_THREADS) + video_fps = decoder.metadata.average_fps + total_frames = decoder.metadata.num_frames + start_frame, end_frame, total_frames = calculate_video_frame_range( + ele, + total_frames, + video_fps, + ) + nframes = smart_nframes(ele, total_frames=total_frames, video_fps=video_fps) + idx = torch.linspace(start_frame, end_frame, nframes).round().long().tolist() + sample_fps = nframes / max(total_frames, 1e-6) * video_fps + video = decoder.get_frames_at(indices=idx).data + logger.info(f"torchcodec: {video_path=}, {total_frames=}, {video_fps=}, time={time.time() - st:.3f}s") + + video_metadata = dict( + fps=video_fps, + frames_indices=idx, + total_num_frames=total_frames, + video_backend="torchcodec", + ) + return video, video_metadata, sample_fps + + +VIDEO_READER_BACKENDS = { + "decord": _read_video_decord, + "torchvision": _read_video_torchvision, + "torchcodec": _read_video_torchcodec, +} + +FORCE_QWENVL_VIDEO_READER = os.getenv("FORCE_QWENVL_VIDEO_READER", None) + + +@lru_cache(maxsize=1) +def get_video_reader_backend() -> str: + if FORCE_QWENVL_VIDEO_READER is not None: + video_reader_backend = FORCE_QWENVL_VIDEO_READER + elif is_torchcodec_available(): + video_reader_backend = "torchcodec" + elif is_decord_available(): + video_reader_backend = "decord" + else: + video_reader_backend = "torchvision" + print(f"qwen-vl-utils using {video_reader_backend} to read video.", file=sys.stderr) + return video_reader_backend + + +def fetch_video(ele: Dict[str, Any], image_patch_size: int = 14, return_video_sample_fps: bool = False, + return_video_metadata: bool = False, video_resize_chunk_frames: int = None) -> Union[torch.Tensor, List[Image.Image]]: + image_factor = image_patch_size * SPATIAL_MERGE_SIZE + VIDEO_FRAME_MIN_PIXELS = VIDEO_MIN_TOKEN_NUM * image_factor * image_factor + VIDEO_FRAME_MAX_PIXELS = VIDEO_MAX_TOKEN_NUM * image_factor * image_factor + if isinstance(ele["video"], str): + video_reader_backend = get_video_reader_backend() + video, video_metadata, sample_fps = VIDEO_READER_BACKENDS[video_reader_backend](ele) + else: + # The input is a list of frames + assert isinstance(ele["video"], (list, tuple)) + process_info = ele.copy() + process_info.pop("type", None) + process_info.pop("video", None) + # use ThreadPoolExecutor to parallel process frames + max_workers = min(MAX_NUM_WORKERS_FETCH_VIDEO, len(ele["video"])) + with ThreadPoolExecutor(max_workers=max_workers) as executor: + futures = [ + executor.submit(fetch_image, {"image": video_element, **process_info}, image_factor) + for video_element in ele["video"] + ] + image_list = [future.result() for future in futures] + + nframes = ceil_by_factor(len(image_list), FRAME_FACTOR) + if len(image_list) < nframes: + image_list.extend([image_list[-1]] * (nframes - len(image_list))) + + sample_fps = ele.get("sample_fps", 2.0) + video = torch.stack([ + torch.from_numpy(np.array(image).transpose(2, 0, 1)) + for image in image_list + ]) + + # fake video metadata + raw_fps = process_info.pop("raw_fps", sample_fps) + video_metadata = dict( + fps=raw_fps, + frames_indices=[i for i in range(len(video))], + total_num_frames=(nframes / sample_fps) * raw_fps, + ) + + nframes, _, height, width = video.shape + min_pixels = ele.get("min_pixels", VIDEO_FRAME_MIN_PIXELS) + total_pixels = ele.get("total_pixels", MODEL_SEQ_LEN * image_factor * image_factor * 0.9) + max_pixels = max(min(VIDEO_FRAME_MAX_PIXELS, total_pixels / nframes * FRAME_FACTOR), int(min_pixels * 1.05)) + max_pixels_supposed = ele.get("max_pixels", max_pixels) + if max_pixels_supposed > max_pixels: + logger.warning(f"The given max_pixels[{max_pixels_supposed}] exceeds limit[{max_pixels}].") + max_pixels = min(max_pixels_supposed, max_pixels) + if "resized_height" in ele and "resized_width" in ele: + resized_height, resized_width = smart_resize( + ele["resized_height"], + ele["resized_width"], + factor=image_factor, + ) + else: + resized_height, resized_width = smart_resize( + height, + width, + factor=image_factor, + min_pixels=min_pixels, + max_pixels=max_pixels, + ) + + if video_resize_chunk_frames and video_resize_chunk_frames < nframes: + resized_chunks = [] + for i in range(0, nframes, video_resize_chunk_frames): + chunk = video[i:i + video_resize_chunk_frames] # 切片，共享内存 + chunk_resized = transforms.functional.resize( + chunk, + [resized_height, resized_width], + interpolation=InterpolationMode.BICUBIC, + antialias=True, + ).float() + resized_chunks.append(chunk_resized) + video = torch.cat(resized_chunks, dim=0) + del resized_chunks + else: + video = transforms.functional.resize( + video, + [resized_height, resized_width], + interpolation=InterpolationMode.BICUBIC, + antialias=True, + ).float() + + final_video = (video, video_metadata) if return_video_metadata else video + if return_video_sample_fps: + return final_video, sample_fps + return final_video + + +def extract_vision_info(conversations: Union[List[Dict[str, Any]], List[List[Dict[str, Any]]]]) -> List[Dict[str, Any]]: + vision_infos = [] + if isinstance(conversations[0], dict): + conversations = [conversations] + for conversation in conversations: + for message in conversation: + if isinstance(message["content"], list): + for ele in message["content"]: + if ( + "image" in ele + or "image_url" in ele + or "video" in ele + or ele.get("type", "text") in ("image", "image_url", "video") + ): + vision_infos.append(ele) + return vision_infos + + +def process_vision_info( + conversations: Union[List[Dict[str, Any]], List[List[Dict[str, Any]]]], + return_video_kwargs: bool = False, + return_video_metadata: bool = False, + image_patch_size: int = 14, + video_resize_chunk_frames: int = None, +) -> Tuple[Optional[List[Image.Image]], Optional[List[Union[torch.Tensor, List[Image.Image]]]], Optional[Dict[str, Any]]]: + + vision_infos = extract_vision_info(conversations) + ## Read images or videos + image_inputs = [] + video_inputs = [] + video_sample_fps_list = [] + for vision_info in vision_infos: + if "image" in vision_info or "image_url" in vision_info: + image_inputs.append(fetch_image(vision_info, image_patch_size=image_patch_size)) + elif "video" in vision_info: + video_input, video_sample_fps = fetch_video(vision_info, return_video_sample_fps=True, + image_patch_size=image_patch_size, return_video_metadata=return_video_metadata, + video_resize_chunk_frames=video_resize_chunk_frames) + video_sample_fps_list.append(video_sample_fps) + video_inputs.append(video_input) + else: + raise ValueError("image, image_url or video should in content.") + if len(image_inputs) == 0: + image_inputs = None + if len(video_inputs) == 0: + video_inputs = None + + video_kwargs = {'do_sample_frames': False} + if not return_video_metadata: # BC for qwen2.5vl + video_kwargs.update({'fps': video_sample_fps_list}) + + if return_video_kwargs: + return image_inputs, video_inputs, video_kwargs + return image_inputs, video_inputs \ No newline at end of file diff --git a/roll/utils/taskgroups.py b/roll/utils/taskgroups.py index d7fd737c0..59ac5560a 100644 --- a/roll/utils/taskgroups.py +++ b/roll/utils/taskgroups.py @@ -214,7 +214,7 @@ def create_task(self, coro, *, name=None, context=None): task = self._loop.create_task(coro) else: task = self._loop.create_task(coro, context=context) - tasks._set_task_name(task, name) + task.set_name(name) # Always schedule the done callback even if the task is # already done (e.g. if the coro was able to complete eagerly), diff --git a/roll/utils/telemetry.py b/roll/utils/telemetry.py new file mode 100644 index 000000000..c77f3c512 --- /dev/null +++ b/roll/utils/telemetry.py @@ -0,0 +1,304 @@ +"""OpenTelemetry tracing utilities for ROLL pipelines. + +Provides cross-actor trace context propagation via W3C TraceContext. +The driver auto-starts an OpenTelemetry Collector and all processes +(driver and actors) export spans via OTLP gRPC to it. + +**Collector Selection** (automatic): +1. **otelcol-contrib/otelcol** (preferred) - Production-ready system binary. +2. **Python OTLP Receiver** (fallback) - Pure Python, no system dependencies. + +Configuration is done via ``system_envs`` in the pipeline config:: + + system_envs: + ROLL_OTEL_ENABLED: "1" + +The ``OTEL_EXPORTER_OTLP_ENDPOINT`` is auto-generated in +``BaseConfig.__post_init__`` (picks a free port on the driver node). + +Usage:: + + # Propagate context to remote actors: + inject_trace_context(data_proto.meta_info) + + # Attach parent context in a receiver: + with attach_trace_context(data_proto.meta_info): + with get_tracer("worker").start_as_current_span("my_span"): + ... + + # At shutdown: + shutdown_telemetry() +""" + +import atexit +import os +import shutil +import subprocess +import textwrap +import time +from contextlib import contextmanager +from contextvars import ContextVar +from typing import Optional +from urllib.parse import urlparse + +from opentelemetry import trace +from opentelemetry.context import Context, attach, detach +from opentelemetry.propagate import inject, extract +from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter +from opentelemetry.sdk.resources import Resource +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import BatchSpanProcessor +from opentelemetry.sdk.trace.id_generator import IdGenerator + +from roll.utils.otel_receiver import OTLPReceiver +from roll.utils.network_utils import get_node_ip, collect_free_port +from roll.utils.logging import get_logger + +logger = get_logger() + +_telemetry_initialized = False +_otel_collector_process = None +_otel_receiver = None + +_OTEL_CTX_KEY = "_otel_ctx" +_trace_suppressed: ContextVar[bool] = ContextVar("_trace_suppressed", default=False) +_noop_tracer_provider = trace.NoOpTracerProvider() + + +class _OsUrandomIdGenerator(IdGenerator): + """Span/trace ID generator immune to ``random.seed()``. + + The default ``RandomIdGenerator`` uses ``random.getrandbits()`` which + is affected by ML training code calling ``random.seed(42)`` for + reproducibility. When every worker shares the same seed, they all + produce **identical** span IDs, corrupting the trace tree. + + This generator uses ``os.urandom()`` (kernel CSPRNG) which is never + affected by ``random.seed()``. + """ + + def generate_span_id(self) -> int: + return int.from_bytes(os.urandom(8), byteorder="big") & ((1 << 64) - 1) + + def generate_trace_id(self) -> int: + return int.from_bytes(os.urandom(16), byteorder="big") & ((1 << 128) - 1) + + +def inject_trace_context(meta_info: dict, context: Optional[Context] = None) -> dict: + """Inject current span context into a *meta_info* dict. + + The context is stored under the ``_otel_ctx`` key as a plain ``dict`` + containing W3C ``traceparent`` / ``tracestate`` headers. This dict is + safe to pass through Ray serialisation and ``DataProto.meta_info``. + """ + carrier: dict = {} + inject(carrier, context) # writes traceparent/tracestate into carrier + meta_info[_OTEL_CTX_KEY] = carrier + return meta_info + + +def extract_trace_context(meta_info: Optional[dict]) -> Context: + """Extract parent span context from a *meta_info* dict. + + Returns a :class:`~opentelemetry.context.Context` that can be passed as + the ``context`` argument of ``tracer.start_as_current_span()``. + If *meta_info* is ``None`` or contains no ``_otel_ctx`` key, an empty + context is returned. + """ + if meta_info is None: + return Context() + carrier = meta_info.get(_OTEL_CTX_KEY, {}) + return extract(carrier) + + +@contextmanager +def attach_trace_context(meta_info: Optional[dict]): + """Context manager: extract parent context from *meta_info* and attach it. + + If *meta_info* is ``None`` or contains no trace context (no ``_otel_ctx`` + key), tracing is suppressed for the duration of the block — all + :func:`get_tracer` calls will return a no-op tracer. + + Automatically detaches when the block exits. Usage:: + + with attach_trace_context(data.meta_info): + with get_tracer("worker").start_as_current_span("my_span"): + ... + """ + ctx = extract_trace_context(meta_info) + token = attach(ctx) + suppressed = meta_info is None or _OTEL_CTX_KEY not in meta_info + suppress_token = _trace_suppressed.set(suppressed) + try: + yield ctx + finally: + _trace_suppressed.reset(suppress_token) + detach(token) + + +def resolve_otel_endpoint() -> str: + """Auto-detect driver IP and pick a free port for the OTLP gRPC endpoint.""" + ip = get_node_ip() + port = collect_free_port() + endpoint = f"http://{ip}:{port}" + logger.info(f"Resolved OTel endpoint: {endpoint}") + return endpoint + + +def start_otel_collector(endpoint: str, output_dir: str) -> None: + """Start an OpenTelemetry Collector subprocess on the driver. + + Tries ``otelcol-contrib`` / ``otelcol`` first (production-ready). + Falls back to a pure-Python OTLP receiver if no system binary is available. + + The collector listens for OTLP gRPC on the port embedded in *endpoint* + and writes traces to ``/traces/traces.jsonl``. + """ + global _otel_collector_process, _otel_receiver + if _otel_collector_process is not None or _otel_receiver is not None: + logger.warning("OTel Collector already running, skipping.") + return + + parsed = urlparse(endpoint) + port = parsed.port or 4317 + + traces_dir = os.path.join(output_dir, "traces") + os.makedirs(traces_dir, exist_ok=True) + + # Try system otelcol binary first (production-ready) + binary = shutil.which("otelcol-contrib") or shutil.which("otelcol") + if binary is not None: + config_content = textwrap.dedent(f"""\ + receivers: + otlp: + protocols: + grpc: + endpoint: "0.0.0.0:{port}" + exporters: + file: + path: "{traces_dir}/traces.jsonl" + service: + telemetry: + metrics: + level: none + pipelines: + traces: + receivers: [otlp] + exporters: [file] + """) + + config_path = os.path.join(traces_dir, "otelcol-config.yaml") + with open(config_path, "w") as f: + f.write(config_content) + + stdout_f = open(os.path.join(traces_dir, "otelcol.stdout.log"), "w") + stderr_f = open(os.path.join(traces_dir, "otelcol.stderr.log"), "w") + p = subprocess.Popen( + [binary, "--config", config_path], + stdout=stdout_f, + stderr=stderr_f, + ) + time.sleep(1) + if p.poll() is not None: + raise RuntimeError( + f"OTel Collector exited immediately with code {p.returncode}. " + f"Check {traces_dir}/otelcol.stderr.log for details." + ) + _otel_collector_process = p + logger.info(f"OTel Collector started (PID {_otel_collector_process.pid}) on port {port}") + else: + # Fallback to Python-based receiver (lightweight, no system dependency) + output_path = os.path.join(traces_dir, "traces.jsonl") + receiver = OTLPReceiver(port=port, output_path=output_path) + receiver.start() + _otel_receiver = receiver + logger.info(f"Python OTLPReceiver started on port {port}") + + atexit.register(shutdown_telemetry) + + +def stop_otel_collector() -> None: + """Terminate the collector subprocess started by :func:`start_otel_collector`.""" + global _otel_collector_process, _otel_receiver + if _otel_collector_process is not None: + _otel_collector_process.terminate() + try: + _otel_collector_process.wait(timeout=5) + except subprocess.TimeoutExpired: + _otel_collector_process.kill() + logger.info("OTel Collector stopped.") + _otel_collector_process = None + if _otel_receiver is not None: + _otel_receiver.stop() + logger.info("Python OTLPReceiver stopped.") + _otel_receiver = None + +def init_telemetry( + service_name: str, + otlp_endpoint: Optional[str] = None, + instance_id: Optional[str] = None, +) -> None: + """Initialize OpenTelemetry tracing. + + Requires an OTLP endpoint (provided directly or via the + ``OTEL_EXPORTER_OTLP_ENDPOINT`` environment variable, typically set + through ``system_envs`` in the pipeline config). + + Args: + service_name: The service name for trace identification. + otlp_endpoint: Optional OTLP gRPC endpoint (e.g. ``http://localhost:4317``). + If not set, checks ``OTEL_EXPORTER_OTLP_ENDPOINT``. + instance_id: Optional identifier (e.g. worker name) added as + the ``service.instance.id`` resource attribute for Jaeger. + """ + global _telemetry_initialized + if _telemetry_initialized: + logger.warning("Telemetry already initialized, skipping.") + return + + endpoint = otlp_endpoint or os.environ.get("OTEL_EXPORTER_OTLP_ENDPOINT") + + if not endpoint: + logger.info("Telemetry: no OTLP endpoint available, skipping init.") + return + + resource_attrs = {"service.name": service_name} + if instance_id: + resource_attrs["service.instance.id"] = instance_id + resource = Resource.create(resource_attrs) + provider = TracerProvider(resource=resource, id_generator=_OsUrandomIdGenerator()) + + otlp_exporter = OTLPSpanExporter(endpoint=endpoint, insecure=True) + provider.add_span_processor(BatchSpanProcessor(otlp_exporter)) + + trace.set_tracer_provider(provider) + _telemetry_initialized = True + log_msg = f"Telemetry initialized: service={service_name}, endpoint={endpoint}" + if instance_id: + log_msg += f", instance={instance_id}" + logger.info(log_msg) + + +def get_tracer(name: str = "roll") -> trace.Tracer: + """Return a tracer from the global TracerProvider. + + Returns a no-op tracer if telemetry has not been initialized or if + the current context has tracing suppressed (i.e. inside an + ``attach_trace_context({})`` block with no trace context). + """ + if _trace_suppressed.get(): + return _noop_tracer_provider.get_tracer(name) + return trace.get_tracer(name) + + +def shutdown_telemetry() -> None: + """Flush pending spans, shut down the tracer provider, and stop the collector.""" + global _telemetry_initialized + if not _telemetry_initialized: + return + provider = trace.get_tracer_provider() + provider.force_flush() + provider.shutdown() + _telemetry_initialized = False + logger.info("Telemetry shut down.") + stop_otel_collector() diff --git a/roll/utils/train_infer_corrections.py b/roll/utils/train_infer_corrections.py index 99107df98..c1737ee37 100644 --- a/roll/utils/train_infer_corrections.py +++ b/roll/utils/train_infer_corrections.py @@ -69,9 +69,14 @@ def compute_train_infer_correction( infer_log_probs: torch.Tensor, # [B, T] global_valid_samples: Optional[int] = None, # Number of valid sequences global_valid_tokens: Optional[int] = None, # Total number of valid tokens - apply_filters: bool = True, ) -> Tuple[torch.Tensor, torch.Tensor, Dict[str, float]]: - """Compute importance sampling weights and apply optional filters based on train-infer divergence.""" + """Compute importance sampling weights and filters based on train-infer divergence. + + Metrics are always recorded for monitoring. The `enabled` flags in config control + whether the corrections are actually applied: + - `is_weight.enabled`: controls whether IS weight is computed (vs. all ones) + - `filters[i].enabled`: controls whether each filter is applied + """ metrics: Dict[str, float] = {} base_mask = response_mask.long() @@ -88,6 +93,21 @@ def compute_train_infer_correction( ratio = stats["ratio"] diff = stats["diff"] + # Always record ratio and diff metrics at all granularities for monitoring + for agg in ["token", "geometric", "sequence", "segment"]: + metrics[f"actor/train_infer_ratio_{agg}_mean@sum"] = agg_loss( + loss_mat=ratio[agg], + loss_mask=base_mask, + loss_agg_mode="seq-mean-token-mean", + global_valid_samples=global_valid_samples, + ).detach().item() + metrics[f"actor/train_infer_diff_{agg}_mean@sum"] = agg_loss( + loss_mat=diff[agg], + loss_mask=base_mask, + loss_agg_mode="seq-mean-token-mean", + global_valid_samples=global_valid_samples, + ).detach().item() + # 1) Importance Sampling (IS) Weight Handling if cfg.is_weight.enabled: is_weight = ratio[cfg.is_weight.weight_type] @@ -109,67 +129,43 @@ def compute_train_infer_correction( # 2) Apply Filters (if enabled) filter_mask = torch.ones_like(base_mask) - recorded_val_metrics = set() # Avoid duplicate metric logging for the same granularity - - if apply_filters: - for i, f in enumerate(cfg.filters): - if not f.enabled: - continue - - agg = f.agg_type - - # --- Ratio-based Filter --- - if f.ratio_enabled: - m_ratio = (ratio[agg] >= f.ratio_low).float() * (ratio[agg] <= f.ratio_high).float() - - # Log pass rate of this filter over currently active tokens - metrics[f"actor/train_infer_{agg}_ratio_mask_mean@sum"] = agg_loss( - loss_mat=m_ratio, - loss_mask=base_mask, - loss_agg_mode='token-mean', - batch_num_tokens=global_valid_tokens, - ).detach().item() - - # Log mean value of the ratio at this granularity (for monitoring) - val_key = f"actor/train_infer_ratio_{agg}_mean@sum" - if val_key not in recorded_val_metrics: - metrics[val_key] = agg_loss( - loss_mat=ratio[agg], - loss_mask=base_mask, - loss_agg_mode="seq-mean-token-mean", - global_valid_samples=global_valid_samples, - ).detach().item() - recorded_val_metrics.add(val_key) - - filter_mask = filter_mask * m_ratio - - # --- Difference-based Filter --- - if f.diff_enabled: - m_diff = (diff[agg] >= f.diff_low).float() * (diff[agg] <= f.diff_high).float() - - # Log pass rate of this filter - metrics[f"actor/train_infer_{agg}_diff_mask_mean"] = agg_loss( - loss_mat=m_diff, - loss_mask=base_mask, - loss_agg_mode='token-mean', - batch_num_tokens=global_valid_tokens, - ).detach().item() - - # Log mean value of the difference at this granularity - val_key = f"actor/train_infer_diff_{agg}_mean@sum" - if val_key not in recorded_val_metrics: - metrics[val_key] = agg_loss( - loss_mat=diff[agg], - loss_mask=base_mask, - loss_agg_mode="seq-mean-token-mean", - global_valid_samples=global_valid_samples, - ).detach().item() - recorded_val_metrics.add(val_key) - - filter_mask = filter_mask * m_diff + + for i, f in enumerate(cfg.filters): + if not f.enabled: + continue + + agg = f.agg_type + + # --- Ratio-based Filter --- + if f.ratio_enabled: + m_ratio = (ratio[agg] >= f.ratio_low).float() * (ratio[agg] <= f.ratio_high).float() + + # Log pass rate of this filter over currently active tokens + metrics[f"actor/train_infer_{agg}_ratio_mask_mean@sum"] = agg_loss( + loss_mat=m_ratio, + loss_mask=base_mask, + loss_agg_mode='token-mean', + batch_num_tokens=global_valid_tokens, + ).detach().item() + + filter_mask = filter_mask * m_ratio + + # --- Difference-based Filter --- + if f.diff_enabled: + m_diff = (diff[agg] >= f.diff_low).float() * (diff[agg] <= f.diff_high).float() + + # Log pass rate of this filter + metrics[f"actor/train_infer_{agg}_diff_mask_mean@sum"] = agg_loss( + loss_mat=m_diff, + loss_mask=base_mask, + loss_agg_mode='token-mean', + batch_num_tokens=global_valid_tokens, + ).detach().item() + + filter_mask = filter_mask * m_diff # 3) Final overall pass rate after all filters - if apply_filters and cfg.filters: + if cfg.filters: metrics["actor/train_infer_final_mask_mean"] = masked_mean( base_mask*filter_mask.float(), base_mask ).detach().item() @@ -189,6 +185,11 @@ def apply_train_infer_correction_to_batch( original shape [B, T]. It handles slicing internally and updates the original masks with the computed filter mask. + Metrics are always recorded for monitoring. The `enabled` flags in config control + whether corrections are applied: + - `is_weight.enabled`: controls whether IS weight is computed (vs. all ones) + - `filters[i].enabled`: controls whether each filter is applied + Args: pipeline_config: Pipeline configuration containing train_infer_correction config batch: DataProto batch to modify @@ -231,7 +232,6 @@ def apply_train_infer_correction_to_batch( infer_log_probs=infer_lp, global_valid_samples=None, # Will be inferred from stat_mask global_valid_tokens=None, # Will be inferred from stat_mask - apply_filters=True, ) # Set train_infer_is_weight diff --git a/roll/utils/upload_utils.py b/roll/utils/upload_utils.py index e8819c377..e2695fc6d 100644 --- a/roll/utils/upload_utils.py +++ b/roll/utils/upload_utils.py @@ -1,4 +1,5 @@ import os +import re import shutil from roll.utils.logging import get_logger @@ -22,6 +23,7 @@ class FileSystemUploader: def __init__(self, output_dir, *args, **kwargs): self.output_dir = output_dir logger.info(f"use FileSystemUploader to upload {output_dir}") + self._re_checkpoint = re.compile(r"^" + "checkpoint" + r"\-(\d+)$") def upload(self, ckpt_id: str, local_state_path: str, **kwargs): ckpt_id_output_dir = os.path.join(self.output_dir, ckpt_id) @@ -30,5 +32,16 @@ def upload(self, ckpt_id: str, local_state_path: str, **kwargs): shutil.copytree(local_state_path, ckpt_id_output_dir, dirs_exist_ok=True) logger.info(f"{local_state_path} save to {ckpt_id_output_dir}, done...") + def get_latest_ckpt(self): + content = os.listdir(self.output_dir) + checkpoints = [ + path + for path in content + if self._re_checkpoint.search(path) is not None and os.path.isdir(os.path.join(self.output_dir, path)) + ] + if len(checkpoints) == 0: + return None + return os.path.join(self.output_dir, max(checkpoints, key=lambda x: int(self._re_checkpoint.search(x).groups()[0]))) + uploader_registry['file_system'] = FileSystemUploader diff --git a/setup.py b/setup.py index d41c5fa23..afce68abc 100644 --- a/setup.py +++ b/setup.py @@ -2,7 +2,7 @@ setup( name="roll", - version="0.1.0", + version="0.3.0", description="ROLL - Reinforcement Learning Optimization for Large-Scale Learning", packages=find_packages(include=["roll", "roll.*"]), python_requires=">=3.10",